feat: Add HeaderProvider to mcptoolset

[kaugesaar]

Overview

This PR aligns the Go MCP toolset with the adk-python implementation by introducing:

1. HeaderProvider support

Allows tool calls to attach context-aware HTTP headers (auth tokens, user/session IDs, tenant headers, etc). Headers are generated per call with agent.ReadonlyContext context.

2. Session pooling keyed by headers

Adds a new sessionManager that reuses MCP sessions when headers are identical and creates new ones only when necessary.

3. Updated example:

examples/mcp/main.go now supports a third mode (AGENT_MODE=custom-headers) that uses HeaderProvider to attach a GitHub PAT + optional context headers.

Summary of Changes

mcptoolset/set.go

• Added HeaderProvider to Config
• Added a sessionManager instance to set
• getSession now:
  • accepts agent.ReadonlyContext
  • gathers provider headers
  • requests a pooled session for those headers

mcptoolset/session_manager.go (new)

• Pools sessions keyed by deterministic header hash
• Wraps HTTP transports with a custom RoundTripper for static header injection
• Validates sessions using Ping
• Implements Close() to cleanly shut down all sessions

mcptoolset/session_manager_test.go (new)

• Tests for:
  • session key stability
  • pooling behavior
  • header injection
  • cleanup
  • asserts that headers do not change session key for stdio/in-memory

examples/mcp/main.go

• Added “custom-headers” mode demonstrating HeaderProvider usage.

Alignment with Python

This matches the adk-python behavior in:

• per-call dynamic headers
• session pooling keyed by headers
• stdio/in-memory = "single session regardless of headers"

Example

headerProvider := func(ctx agent.ReadonlyContext) map[string]string {
	return map[string]string{
		"Authorization": "Bearer " + os.Getenv("GITHUB_PAT"),
		"X-User-ID":     ctx.UserID(),
	}
}

ts, _ := mcptoolset.New(mcptoolset.Config{
	Transport: &mcp.StreamableClientTransport{Endpoint: "https://api.githubcopilot.com/mcp/"},
	HeaderProvider: headerProvider,
})

Known Limitation

Sessions are pooled based on headers from HeaderProvider only. Headers added by custom HTTPClient transports (e.g., oauth2.NewClient) are not visible to the session manager and do not affect pooling.

[gemini-code-assist]

Summary of Changes

Hello @kaugesaar, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the Go MCP toolset by aligning its behavior with the existing adk-python implementation, focusing on more flexible and efficient handling of HTTP headers and client sessions. It introduces a HeaderProvider for dynamically attaching context-specific headers to tool calls and a sessionManager for pooling and reusing MCP sessions based on these headers. This improves the toolset's ability to manage authenticated and context-aware interactions with MCP servers, leading to better performance and more robust integrations.

Highlights

• HeaderProvider Integration: Introduced a HeaderProvider function to allow dynamic, context-aware HTTP headers to be attached to MCP tool calls. This enables flexible handling of authentication tokens, user/session IDs, and tenant headers per request.
• Session Pooling by Headers: Implemented a new sessionManager that intelligently reuses MCP client sessions. Sessions are pooled and retrieved based on a deterministic hash of the HTTP headers, reducing the overhead of establishing new connections when headers are identical.
• Updated Example for HeaderProvider: The examples/mcp/main.go file has been updated to include a new 'custom-headers' mode. This demonstrates how to configure and use the HeaderProvider to attach dynamic headers, such as a GitHub Personal Access Token and a user ID, to MCP requests.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces HeaderProvider support to the mcptoolset, enabling dynamic, context-aware HTTP headers for tool calls, and adds session pooling based on these headers. The changes align the Go MCP toolset with the Python implementation. The implementation looks solid, with a new sessionManager to handle session pooling and a comprehensive set of tests. I have a few suggestions to improve performance, security posture, and code clarity, mainly in the new session_manager.go file. My main concerns are around the use of a weak hashing algorithm and a locking pattern that could cause performance issues under load.

[gemini-code-assist]

The use of crypto/md5 is discouraged as it's a weak cryptographic hash function. While not used for a security-critical purpose here, it's best practice to use a stronger hash function like SHA-256 to prevent potential hash collisions and to satisfy security scanners. You'll also need to update the import from crypto/md5 to crypto/sha256.

Suggested change

	h := md5.Sum([]byte(jsonStr))
	h := sha256.Sum256([]byte(jsonStr))