diff --git a/mcp_utils.py b/mcp_utils.py index 3a7507c..dd64909 100644 --- a/mcp_utils.py +++ b/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') diff --git a/taskflows/examples/example_large_list_result_iter.yaml b/taskflows/examples/example_large_list_result_iter.yaml index bb0ae87..2733351 100644 --- a/taskflows/examples/example_large_list_result_iter.yaml +++ b/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 diff --git a/toolboxes/github_official.yaml b/toolboxes/github_official.yaml index e1cde32..0bfa367 100644 --- a/toolboxes/github_official.yaml +++ b/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 }}" \ No newline at end of file + Authorization: "{{ env GITHUB_AUTH_HEADER }}" + optional_headers: + X-MCP-Toolsets: "{{ env GITHUB_MCP_TOOLSETS }}" + X-MCP-Readonly: "{{ env GITHUB_MCP_READONLY }}" \ No newline at end of file