Merge pull request #15 from GitHubSecurityLab/restrict_gh_mcp_toolset

Add optional headers to restrict GitHub mcp server toolsets

Author: m-y-mo

mcp_utils.py

@@ -314,6 +314,19 @@ def mcp_client_params(available_toolboxes: dict, requested_toolboxes: list):
                     if headers and isinstance(headers, dict):
                         for k, v in headers.items():
                             headers[k] = swap_env(v)
+                    optional_headers = available_toolboxes[tb]['server_params'].get('optional_headers')
+                    # support {{ env SOMETHING }} for header values as well for e.g. tokens
+                    if optional_headers and isinstance(optional_headers, dict):
+                        for k, v in dict(optional_headers).items():
+                            try:
+                                optional_headers[k] = swap_env(v)
+                            except LookupError as e:
+                                del optional_headers[k]
+                    if isinstance(headers, dict):
+                        if isinstance(optional_headers, dict):
+                            headers.update(optional_headers)
+                    elif isinstance(optional_headers, dict):
+                        headers = optional_headers
                     # if None will default to float(5) in client code
                     timeout = available_toolboxes[tb]['server_params'].get('timeout')
                     server_params['url'] = available_toolboxes[tb]['server_params'].get('url')
@@ -329,6 +342,19 @@ def mcp_client_params(available_toolboxes: dict, requested_toolboxes: list):
                     if headers and isinstance(headers, dict):
                         for k, v in headers.items():
                             headers[k] = swap_env(v)
+                    optional_headers = available_toolboxes[tb]['server_params'].get('optional_headers')
+                    # support {{ env SOMETHING }} for header values as well for e.g. tokens
+                    if optional_headers and isinstance(optional_headers, dict):
+                        for k, v in dict(optional_headers).items():
+                            try:
+                                optional_headers[k] = swap_env(v)
+                            except LookupError as e:
+                                del optional_headers[k]
+                    if isinstance(headers, dict):
+                        if isinstance(optional_headers, dict):
+                            headers.update(optional_headers)
+                    elif isinstance(optional_headers, dict):
+                        headers = optional_headers
                     # if None will default to float(5) in client code
                     timeout = available_toolboxes[tb]['server_params'].get('timeout')
                     server_params['url'] = available_toolboxes[tb]['server_params'].get('url')

taskflows/examples/example_large_list_result_iter.yaml

@@ -13,6 +13,8 @@ taskflow:
         You do not need to provide a summary of the results.
       toolboxes:
         - github_official
+      env:
+        GITHUB_MCP_TOOLSETS: pull_requests
   - task:
       must_complete: true
       repeat_prompt: true

toolboxes/github_official.yaml

@@ -5,5 +5,9 @@ seclab-taskflow-agent:
 server_params:
   kind: streamable
   url: https://api.githubcopilot.com/mcp/
+  #See https://github.com/github/github-mcp-server/blob/main/docs/remote-server.md
   headers:
-    Authorization: "{{ env GITHUB_AUTH_HEADER }}"
+    Authorization: "{{ env GITHUB_AUTH_HEADER }}"
+  optional_headers:
+    X-MCP-Toolsets: "{{ env GITHUB_MCP_TOOLSETS }}"
+    X-MCP-Readonly: "{{ env GITHUB_MCP_READONLY }}"