Skip to content

Commit d9e601a

Browse files
authored
Fix stdio session tool handler conflicts (#17)
Resolves intermittent connection failures where stdio sessions overwrote each other's tool handlers. Each session now registers tools in isolation using mcp-go's SessionWithTools interface.
1 parent 40d637d commit d9e601a

File tree

8 files changed

+678
-102
lines changed

8 files changed

+678
-102
lines changed

CLAUDE.md

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,20 @@ The animated logo creates a face-like character:
300300
- Assume developer audience
301301
- Keep it concise but complete
302302

303+
### Pull Request Guidelines
304+
305+
**PR Titles**: Use clear, descriptive titles focused on the change impact, not just restating commit messages. Examples:
306+
- ❌ "feat: add message endpoint support for SSE MCP servers"
307+
- ✅ "Add SSE message endpoint support"
308+
- ❌ "fix: implement session-specific tool registration for stdio clients"
309+
- ✅ "Fix stdio session tool handler conflicts"
310+
311+
**PR Descriptions**: Write terse prose for humans, not documentation. Avoid bullet lists unless they add genuine value. Focus on the problem solved and solution approach:
312+
- Explain what was broken and how it's fixed
313+
- Use conversational, developer-to-developer tone
314+
- Skip implementation details unless critical for review
315+
- Keep it concise but informative
316+
303317
### Common Issues
304318

305319
1. **Theme switching breaks**: Check CSS variable inheritance

internal/client/client.go

Lines changed: 55 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,22 @@ func (c *Client) AddToMCPServerWithTokenCheck(
150150
serverName string,
151151
setupBaseURL string,
152152
tokenSetup *config.TokenSetupConfig,
153+
) error {
154+
return c.AddToMCPServerWithSession(ctx, clientInfo, mcpServer, userEmail, requiresToken, tokenStore, serverName, setupBaseURL, tokenSetup, nil)
155+
}
156+
157+
// AddToMCPServerWithSession connects the client to an MCP server with optional session-specific tools
158+
func (c *Client) AddToMCPServerWithSession(
159+
ctx context.Context,
160+
clientInfo mcp.Implementation,
161+
mcpServer *server.MCPServer,
162+
userEmail string,
163+
requiresToken bool,
164+
tokenStore storage.UserTokenStore,
165+
serverName string,
166+
setupBaseURL string,
167+
tokenSetup *config.TokenSetupConfig,
168+
session server.ClientSession,
153169
) error {
154170
if c.needManualStart {
155171
err := c.client.Start(ctx)
@@ -176,7 +192,7 @@ func (c *Client) AddToMCPServerWithTokenCheck(
176192
"server": c.name,
177193
})
178194

179-
err = c.addToolsToServer(ctx, mcpServer, userEmail, requiresToken, tokenStore, serverName, setupBaseURL, tokenSetup)
195+
err = c.addToolsToServer(ctx, mcpServer, userEmail, requiresToken, tokenStore, serverName, setupBaseURL, tokenSetup, session)
180196
if err != nil {
181197
return err
182198
}
@@ -224,6 +240,7 @@ func (c *Client) addToolsToServer(
224240
serverName string,
225241
setupBaseURL string,
226242
tokenSetup *config.TokenSetupConfig,
243+
session server.ClientSession,
227244
) error {
228245
toolsRequest := mcp.ListToolsRequest{}
229246
filterFunc := func(toolName string) bool {
@@ -263,6 +280,22 @@ func (c *Client) addToolsToServer(
263280
})
264281

265282
totalTools := 0
283+
284+
var sessionWithTools server.SessionWithTools
285+
var sessionTools map[string]server.ServerTool
286+
if session != nil {
287+
var ok bool
288+
sessionWithTools, ok = session.(server.SessionWithTools)
289+
if !ok {
290+
return fmt.Errorf("session does not support session-specific tools")
291+
}
292+
sessionTools = make(map[string]server.ServerTool)
293+
internal.LogInfoWithFields("client", "Using session-specific tool registration", map[string]interface{}{
294+
"server": c.name,
295+
"sessionID": session.SessionID(),
296+
})
297+
}
298+
266299
for {
267300
tools, err := c.client.ListTools(ctx, toolsRequest)
268301
if err != nil {
@@ -286,8 +319,9 @@ func (c *Client) addToolsToServer(
286319
"description": tool.Description,
287320
})
288321
// Wrap the tool handler to check for user tokens if required
322+
var handler server.ToolHandlerFunc
289323
if requiresToken && tokenStore != nil {
290-
wrappedHandler := c.wrapToolHandler(
324+
handler = c.wrapToolHandler(
291325
c.client.CallTool,
292326
requiresToken,
293327
tokenStore,
@@ -296,9 +330,17 @@ func (c *Client) addToolsToServer(
296330
setupBaseURL,
297331
tokenSetup,
298332
)
299-
mcpServer.AddTool(tool, wrappedHandler)
300333
} else {
301-
mcpServer.AddTool(tool, c.client.CallTool)
334+
handler = c.client.CallTool
335+
}
336+
337+
if sessionTools != nil {
338+
sessionTools[tool.Name] = server.ServerTool{
339+
Tool: tool,
340+
Handler: handler,
341+
}
342+
} else {
343+
mcpServer.AddTool(tool, handler)
302344
}
303345
}
304346
}
@@ -308,6 +350,15 @@ func (c *Client) addToolsToServer(
308350
toolsRequest.Params.Cursor = tools.NextCursor
309351
}
310352

353+
if len(sessionTools) > 0 {
354+
sessionWithTools.SetSessionTools(sessionTools)
355+
internal.LogInfoWithFields("client", "Registered session-specific tools", map[string]interface{}{
356+
"server": c.name,
357+
"sessionID": session.SessionID(),
358+
"toolCount": len(sessionTools),
359+
})
360+
}
361+
311362
internal.LogInfoWithFields("client", "Tool discovery completed", map[string]interface{}{
312363
"server": c.name,
313364
"totalTools": totalTools,

0 commit comments

Comments
 (0)