feat: allow customising context via callback on transport servers #32
Conversation
WalkthroughThis pull request introduces new types and methods for customizing context in both the SSE and Stdio servers. It adds a function type for server options to the Changes
Possibly related PRs
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
server/stdio.go (1)
194-209: Consider allowing contextFunc in the ServeStdio helper function.The
ServeStdiohelper function doesn't provide a way to set a context function on the created server. Consider adding a parameter or making the created server accessible before starting it.-func ServeStdio(server *MCPServer) error { +func ServeStdio(server *MCPServer, contextFunc StdioContextFunc) error { s := NewStdioServer(server) s.SetErrorLogger(log.New(os.Stderr, "", log.LstdFlags)) + if contextFunc != nil { + s.SetContextFunc(contextFunc) + }Alternatively, you could return the server and have a separate start function:
-func ServeStdio(server *MCPServer) error { +func PrepareStdioServer(server *MCPServer) *StdioServer { s := NewStdioServer(server) s.SetErrorLogger(log.New(os.Stderr, "", log.LstdFlags)) + return s +} + +func ServeStdio(server *MCPServer) error { + return PrepareStdioServer(server).Listen(context.Background(), os.Stdin, os.Stdout) }server/sse.go (1)
51-72: Update NewTestServer to support context function.The
NewTestServerfunction doesn't provide a way to set a context function. For consistency with the regular server, consider adding this capability.func NewTestServer(server *MCPServer) *httptest.Server { sseServer := &SSEServer{ server: server, } testServer := httptest.NewServer( http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/sse": sseServer.handleSSE(w, r) case "/message": sseServer.handleMessage(w, r) default: http.NotFound(w, r) } }), ) sseServer.baseURL = testServer.URL return testServer } +// NewTestServerWithContext creates a test server with a custom context function +func NewTestServerWithContext(server *MCPServer, contextFunc SSEContextFunc) *httptest.Server { + testServer := NewTestServer(server) + // Extract the SSE server to set the context function + // This is a bit hacky but works for testing purposes + sseServer := &SSEServer{ + server: server, + } + sseServer.SetContextFunc(contextFunc) + + return testServer +}Alternatively, you could modify
NewTestServerto accept an optional context function:-func NewTestServer(server *MCPServer) *httptest.Server { +func NewTestServer(server *MCPServer, contextFunc SSEContextFunc) *httptest.Server { sseServer := &SSEServer{ server: server, } + + if contextFunc != nil { + sseServer.SetContextFunc(contextFunc) + }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
server/sse.go(3 hunks)server/stdio.go(3 hunks)
🔇 Additional comments (8)
server/stdio.go (4)
17-21: Well-designed type definition with clear documentation.The
StdioContextFunctype is appropriately defined with a clear signature and good documentation that explains its purpose and potential use cases.
27-29: Clean addition of the contextFunc field to the struct.The new field is properly added to the
StdioServerstruct, following the existing style and organization of the struct.
51-56: Good method implementation with helpful documentation.The
SetContextFuncmethod is well-documented, clearly explaining that this context function will only be called once per server instance since the stdio server uses the same context for all requests. This is important information for users of this API.
72-75: Correctly implemented context customization logic.The implementation properly checks if the context function is set before applying it and correctly updates the context. The comment is clear and descriptive.
server/sse.go (4)
15-18: Well-designed type definition with appropriate parameters.The
SSEContextFunctype is well-defined with appropriate parameters for the SSE server context. The function receives both the context and HTTP request, allowing for modification based on request attributes like headers. The documentation is clear and helpful.
23-27: Clean addition of the contextFunc field to the struct.The new field is properly added to the
SSEServerstruct, maintaining the structure's organization.
45-49: Good method implementation with appropriate documentation.The
SetContextFuncmethod is implemented correctly with helpful documentation.
187-189: Correctly implemented context customization logic.The implementation properly checks if the context function is set before applying it and correctly updates the context with both the existing context and the HTTP request.
|
@sd2k Thanks for this. Would it be possible to add a simple test or two that demonstrate how these new functions would work in practice? |
|
I think what's confusing to me is that we are describing a function which has different scopes depending on transport, so I agree motivating concrete use cases or tests would help. For example, I assume that the request-scoped SSE allows you to overcome JSON-RPC which has no headers. So this way you could make tracing middleware. That said due to the same problem, I would expect folks to make an extra json-rpc field like "traceparent" and then instrumentation would be portable. So, request-scoped is neat, but certainly some concrete use cases would help. I'm more curious about stdio, such as what value we are providing with the context func since it is both listener scope and has no input except for the input context? Should we assume this is to provide a deadline scope? If that's the case, how would that be different than the initial context (knowing there's no other parameter to know more?). Or are you trying to generalize a plugin? or is this just to have some sort of feature parity with sse? |
|
Sorry I was focused elsewhere last week so didn't reply to this! Happy to put some tests together. The basic thing I want to support is allowing tools to access things like configuration without having to know where the config is coming from. For example, in the Grafana MCP server, tools need to know what the current Grafana URL is. In url := GrafanaURLFromContext(ctx)The stdio case is certainly less interesting since we can't provide anything, but continuing the above example I would expect the context func to do something like: if u := os.Getenv("GRAFANA_URL"); u != "" {
ctx = WithGrafanaURL(ctx, u)
}
return ctxI think this broadly comes under the category of 'generalizing a plugin', yep. |
|
Cool answers my question. I think as long as we aren't promising strong api compatibility, this doesn't hurt. Maybe your tests, add a comment saying this shows how you can use the feature to add context scoped configuration, such as the endpoint used in a http based tool. Later, if we read those comments ;) and have a better solution for dynamic config we can revisit the pattern |
This commit adds an callback function to the transport servers that allows developers to inject context values into the server context. This can be used to inject context values extracted from environment variables (in stdio mode) or from headers (in sse mode), and access them in tools using the provided context.
sd2k
force-pushed
the
allow-lifecycle-via-context
branch
from
77436b9 to
37d5406
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
server/sse.go (3)
24-28: Consider clarifying concurrency in the documentation.
The doc comment is concise, but it could explicitly mention thread safety or concurrency implications for context usage, especially if multiple goroutines might invoke this function simultaneously.
39-39: Highlight potential concurrency concerns for the contextFunc field.
If the field is modified while the server is running, there's a risk of data races. Consider clarifying that it should be set before the server starts, or add synchronization to avoid race conditions.
102-106: Document timing requirements for SetContextFunc usage.
It would be good to note that existing connections won't use the updated contextFunc if set after they have been established. This minor documentation enhancement can help avoid confusion.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
server/sse.go(4 hunks)server/stdio.go(3 hunks)server/sse.go(1 hunks)server/sse_test.go(2 hunks)server/stdio_test.go(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- server/stdio.go
- server/sse.go
🔇 Additional comments (5)
server/sse.go (1)
245-248: Ensure proper concurrency handling when accessing contextFunc.
While calling the function here is straightforward, verify that no concurrent reassignment to s.contextFunc occurs during request handling. If such updates are possible, store it in a local variable first to avoid race conditions.server/stdio_test.go (2)
9-12: New imports for environment use and MCP integration look fine.
Importing "os" and the MCP package is appropriate for the new context functionality and test coverage.
117-272: Comprehensive test for custom context function.
This sub-test robustly validates the environment-based context injection. The usage of t.Cleanup to restore the environment variable is a nice touch, ensuring minimal side effects when tests run in parallel.server/sse_test.go (2)
14-15: Import statement for MCP is appropriate.
Bringing in the MCP package is necessary for the new test tools. No issues here.
474-617: Good demonstration of SSE context injection.
The new sub-test thoroughly shows how a custom context can inject header-derived data into the server. This helps confirm that SSE and JSON-RPC interactions benefit from the same context injection mechanism as stdio mode.
|
I've added a test for the SSE and stdin cases. I'd like to add an example of using a tool where config is required. I was thinking something like a minimal GitHub server with a single |
|
@k33g I think you use this library (I don't yet, but will soon when I try yours ;) ). wdyt? |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
server/stdio_test.go (2)
117-135: Well-structured context customization approach.The test effectively demonstrates the context customization capability with a clean implementation of a custom context key type and helper functions. The approach of using environment variables as a context source is a good real-world example.
Consider enhancing type safety in the
testValFromContextfunction:func testValFromContext(ctx context.Context) string { val := ctx.Value(testContextKey{}) if val == nil { return "" } - return val.(string) + str, ok := val.(string) + if !ok { + return "" + } + return str }This prevents panics if the context value is not of the expected type.
218-258: Thorough testing of tool call functionality with context values.The test effectively validates that the context function works by:
- Setting an environment variable
- Configuring a context function to inject this value
- Calling a tool that retrieves the value
- Verifying the tool received the correct value
For improved readability in verification, consider restructuring the deeply nested map access:
- if response["result"].(map[string]interface{})["content"].([]interface{})[0].(map[string]interface{})["text"] != "test_value" { - t.Errorf("Expected result 'test_value', got %v", response["result"]) - } + result, ok := response["result"].(map[string]interface{}) + if !ok { + t.Fatalf("Expected result to be a map, got %T", response["result"]) + } + content, ok := result["content"].([]interface{}) + if !ok || len(content) == 0 { + t.Fatalf("Expected content to be a non-empty array, got %v", result["content"]) + } + textObj, ok := content[0].(map[string]interface{}) + if !ok { + t.Fatalf("Expected content[0] to be a map, got %T", content[0]) + } + text, ok := textObj["text"].(string) + if !ok { + t.Fatalf("Expected text to be a string, got %T", textObj["text"]) + } + if text != "test_value" { + t.Errorf("Expected result 'test_value', got %q", text) + }This approach provides clearer error messages and better type safety.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
server/stdio_test.go(2 hunks)
🔇 Additional comments (4)
server/stdio_test.go (4)
9-12: New imports look good.The addition of the
ospackage andgithub.com/mark3labs/mcp-go/mcppackage is appropriate for the new test functionality that uses environment variables and MCP tooling.
143-147: Tool implementation effectively demonstrates context usage.The tool implementation clearly shows how to access context values in a transport-agnostic way, which aligns perfectly with the PR objective of allowing tools to access configuration values without needing to know their source.
149-150: Good demonstration of context function configuration.Line 150 demonstrates how to set the custom context function on the stdio server, which is the core new functionality being tested.
117-269: Overall, excellent test coverage for the new context function feature.The test thoroughly validates the new context customization capability, providing a clear example of how it can be used to inject environment-based configuration into the context. The comments also clearly explain the real-world use case for this feature.
This test effectively addresses the reviewer's request for tests demonstrating the functionality described in the PR objectives.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (4)
server/sse_test.go (1)
473-490: Context value extraction strategy is well-structured.
- Defining a custom
testContextKeytype (line 476) and using it intestValFromContexthelps avoid key collisions in the context.- Storing the request header in the context (line 489) is a straightforward way to pass domain-specific data to the tools.
Consider adding some type assertions in
testValFromContextto gracefully handle unexpected context values if the key might be overwritten somewhere else.examples/custom_context/main.go (3)
30-33: Environment variable fallback is appropriate.
authFromEnvfits use cases where the token may come from external sources. Consider validating the presence ofAPI_KEYif needed.Add a small log message if
API_KEYis missing, to improve debug visibility.
53-74: Robust external request logic.
makeRequestsets the"Authorization"header, making it easy to see how context-based tokens become real HTTP credentials.Consider capturing HTTP status codes or implementing retries for more robust error handling.
134-164: Main function covers multiple modes.
Switch logic is straightforward. Logging on invalid transport types is user friendly.A usage example or help message (e.g.,
-h) might further assist users, though flag defaults serve as a minimal guide.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
examples/custom_context/main.go(1 hunks)server/sse.go(5 hunks)server/sse_test.go(2 hunks)server/stdio.go(4 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- server/sse.go
🔇 Additional comments (20)
server/sse_test.go (5)
14-14: No meaningful change at line 14.
15-15: New import statement is appropriate.
The import for"github.com/mark3labs/mcp-go/mcp"aligns with the newly introduced features referencingmcpserver functionality.
496-500: Transport-agnostic tool injection is commendable.
UsingtestValFromContext(ctx)within the handler ensures the logic does not depend on any particular transport, which is a clean separation of concerns.
590-591: Header-based context injection tested thoroughly.
Setting the custom test header to"test_value"aligns with the logic insetTestValFromRequest, following a clear end-to-end test path.
609-611: Validation of expected result text.
The test confirms that the returned content matches"test_value", ensuring correctness of context propagation and tool usage.server/stdio.go (8)
17-21: Clear function signature for customizing context.
StdioContextFuncneatly explains how context can be manipulated before the server starts processing messages.
27-30: Introducing contextFunc in StdioServer.
Storing a function pointer in the server struct is a flexible design for dynamic context injection.
32-34: Add-on struct configuration approach is idiomatic.
StdioOptionfosters composability by letting callers chain multiple configuration functions.
35-40: Logger configuration.
ProvidingWithErrorLoggeras an option cleanly separates concerns of logging from the server’s default behavior.
42-49: Context initialization is consistent with the SSE approach.
Naming is descriptive, and the docstring clarifies that this function runs once on start, which is an important detail for the Stdio transport.
70-76: SetContextFunc is a helpful runtime configuration method.
Exposing a setter method (beyond constructor options) further extends flexibility for advanced usage patterns.
92-95: Safely applying the context function.
Applyings.contextFunc(ctx)if non-nil integrates well with the existing lifecycle, letting the server unify environment-based or other custom data.
213-219: ServeStdio now elegantly handles chained options.
This design helps unify server configuration during startup, and ensures consistent usage with SSE’s pattern.examples/custom_context/main.go (7)
17-23: Context key usage.
authKeyis effectively isolated from potential collisions, andwithAuthKeyis straightforward to set the value.
25-29: Request-based token injection.
authFromRequestreading"Authorization"is a standard pattern. This keeps logic minimal and easy to maintain.
35-44: Safe retrieval of context token.
The graceful error path intokenFromContextensures tools can handle missing auth keys.
76-96: Tool ensures separation of concerns for request data.
handleMakeAuthenticatedRequestTooldecodes a short parameter, obtains token, and delegates tomakeRequest()—this layering is good.
98-120: MCPServer and tool registration logic is well done.
The server’s configuration includes resource, prompt, and tool capabilities for flexible usage.
123-128: SSE server usage.
ApplyingWithSSEContextFunc(authFromRequest)is consistent with the approach in the test modules, providing a single location for SSE-based context injection.
130-132: Stdio-based environment token injection.
Mirrors the SSE approach, but uses environment-based context function. This consistency is beneficial for cross-transport usage.
|
This is awesome! Thanks. Merging. |
mark3labs/mcp-go#32 was merged so we can use the latest release now.
|
Nice, thanks! FYI at Grafana we're migrating our Python MCP server to Go using this library so appreciate the fast response to this PR 🫶 |
mark3labs/mcp-go#32 was merged so we can use the latest release now.
|
@sd2k awesome. just wrote you a love letter https://www.linkedin.com/feed/update/urn:li:activity:7305754295721689088/ |
This commit adds an callback function to the transport servers that allows developers to inject context values into the server context.
This can be used to inject context values extracted from environment variables (in stdio mode) or from headers (in SSE mode), and access them in tools using the provided context.
Closes #27.
Summary by CodeRabbit
New Features
Tests