fix: mcptoolset: cached MCP session not validated, causes 'connection closed' errors #447
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
Summary of ChangesHello @majiayu000, I'm Gemini Code Assist1! 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 resolves an issue where the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request correctly addresses an issue where a cached MCP session is not validated, leading to 'connection closed' errors. The fix involves adding a Ping health check before reusing a session and re-establishing the connection if the check fails. The logic is sound. However, the accompanying test only covers the happy path of reusing a valid session. I've suggested adding a new test case to cover the scenario where a session is stale and needs to be recreated, which is the core of this fix.
| // TestSessionValidationOnReuse verifies that the MCP toolset validates | ||
| // cached sessions before reuse using Ping health check. | ||
| func TestSessionValidationOnReuse(t *testing.T) { | ||
| const toolDescription = "returns weather in the given city" | ||
|
|
||
| clientTransport, serverTransport := mcp.NewInMemoryTransports() | ||
|
|
||
| // Create server instance. | ||
| server := mcp.NewServer(&mcp.Implementation{Name: "weather_server", Version: "v1.0.0"}, nil) | ||
| mcp.AddTool(server, &mcp.Tool{Name: "get_weather", Description: toolDescription}, weatherFunc) | ||
| _, err := server.Connect(t.Context(), serverTransport, nil) | ||
| if err != nil { | ||
| t.Fatal(err) | ||
| } | ||
|
|
||
| ts, err := mcptoolset.New(mcptoolset.Config{ | ||
| Transport: clientTransport, | ||
| }) | ||
| if err != nil { | ||
| t.Fatalf("Failed to create MCP tool set: %v", err) | ||
| } | ||
|
|
||
| readonlyCtx := icontext.NewReadonlyContext( | ||
| icontext.NewInvocationContext( | ||
| t.Context(), | ||
| icontext.InvocationContextParams{}, | ||
| ), | ||
| ) | ||
|
|
||
| // First call should establish a session and return tools. | ||
| tools, err := ts.Tools(readonlyCtx) | ||
| if err != nil { | ||
| t.Fatalf("First Tools() call failed: %v", err) | ||
| } | ||
| if len(tools) != 1 || tools[0].Name() != "get_weather" { | ||
| t.Fatalf("Expected 1 tool named 'get_weather', got %d tools", len(tools)) | ||
| } | ||
|
|
||
| // Second call should reuse the cached session (validated via Ping). | ||
| tools2, err := ts.Tools(readonlyCtx) | ||
| if err != nil { | ||
| t.Fatalf("Second Tools() call failed: %v", err) | ||
| } | ||
| if len(tools2) != 1 || tools2[0].Name() != "get_weather" { | ||
| t.Fatalf("Expected 1 tool named 'get_weather', got %d tools", len(tools2)) | ||
| } | ||
| } |
There was a problem hiding this comment.
This test is a good start as it covers the happy path where a cached session is successfully validated and reused. However, it doesn't test the core of the fix, which is handling a stale or closed session.
To ensure the fix is robust, please consider adding another test case that simulates a connection failure. This would involve:
- Establishing a session by calling
ts.Tools()once. - Closing the connection from the server side. You can capture the server-side session from
server.Connectand callClose()on it. - Calling
ts.Tools()a second time. This call should detect the stale session viaPing, create a new session, and still succeed.
Here's an example of how you could structure such a test:
func TestStaleSessionIsRecreated(t *testing.T) {
const toolDescription = "returns weather in the given city"
clientTransport, serverTransport := mcp.NewInMemoryTransports()
// Create server instance.
server := mcp.NewServer(&mcp.Implementation{Name: "weather_server", Version: "v1.0.0"}, nil)
mcp.AddTool(server, &mcp.Tool{Name: "get_weather", Description: toolDescription}, weatherFunc)
// Capture the server-side session to be able to close it later.
serverSession, err := server.Connect(t.Context(), serverTransport, nil)
if err != nil {
t.Fatal(err)
}
ts, err := mcptoolset.New(mcptoolset.Config{
Transport: clientTransport,
})
if err != nil {
t.Fatalf("Failed to create MCP tool set: %v", err)
}
readonlyCtx := icontext.NewReadonlyContext(
icontext.NewInvocationContext(
t.Context(),
icontext.InvocationContextParams{},
),
)
// First call should establish a session and return tools.
_, err = ts.Tools(readonlyCtx)
if err != nil {
t.Fatalf("First Tools() call failed: %v", err)
}
// Simulate connection drop by closing the server-side session.
if err := serverSession.Close(); err != nil {
t.Fatalf("Failed to close server session: %v", err)
}
// Second call should detect the stale session, reconnect, and succeed.
tools, err := ts.Tools(readonlyCtx)
if err != nil {
t.Fatalf("Second Tools() call failed after session drop: %v", err)
}
if len(tools) != 1 || tools[0].Name() != "get_weather" {
t.Fatalf("Expected 1 tool named 'get_weather' after reconnect, got %d tools", len(tools))
}
}
majiayu000
force-pushed
the
fix-399-mcptoolset-cached-mcp-session--1230-1654
branch
from
4d7e35c to
431551d
Compare
Add session validation using Ping health check before reusing cached MCP sessions. If the session is stale or closed, it will be automatically recreated to prevent 'connection closed' errors. Also adds TestStaleSessionIsRecreated test to verify the fix behavior.
majiayu000
force-pushed
the
fix-399-mcptoolset-cached-mcp-session--1230-1654
branch
from
431551d to
f082621
Compare
1 participant
Summary
This PR fixes #399
Changes
tool/mcptoolset/set.gotool/mcptoolset/set_test.go