Skip to content

Commit 739219c

Browse files
committed
Impl AddSessionTool
1 parent 36b261d commit 739219c

File tree

3 files changed

+60
-0
lines changed

3 files changed

+60
-0
lines changed

README.md

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -626,6 +626,21 @@ userSpecificTool := mcp.NewTool(
626626
"user_data",
627627
mcp.WithDescription("Access user-specific data"),
628628
)
629+
// You can use AddSessionTool (similar to AddTool)
630+
err := s.AddSessionTool(
631+
advSession.SessionID(),
632+
userSpecificTool,
633+
func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
634+
// This handler is only available to this specific session
635+
return mcp.NewToolResultText("User-specific data for " + advSession.SessionID()), nil
636+
},
637+
)
638+
if err != nil {
639+
log.Printf("Failed to add session tool: %v", err)
640+
}
641+
642+
// Or use AddSessionTools directly with ServerTool
643+
/*
629644
err := s.AddSessionTools(
630645
advSession.SessionID(),
631646
server.ServerTool{
@@ -639,6 +654,7 @@ err := s.AddSessionTools(
639654
if err != nil {
640655
log.Printf("Failed to add session tool: %v", err)
641656
}
657+
*/
642658

643659
// Delete session-specific tools when no longer needed
644660
err = s.DeleteSessionTools(advSession.SessionID(), "user_data")

server/session.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,11 @@ func (s *MCPServer) SendNotificationToSpecificClient(
199199
}
200200
}
201201

202+
// AddSessionTool adds a tool for a specific session
203+
func (s *MCPServer) AddSessionTool(sessionID string, tool mcp.Tool, handler ToolHandlerFunc) error {
204+
return s.AddSessionTools(sessionID, ServerTool{Tool: tool, Handler: handler})
205+
}
206+
202207
// AddSessionTools adds tools for a specific session
203208
func (s *MCPServer) AddSessionTools(sessionID string, tools ...ServerTool) error {
204209
sessionValue, ok := s.sessions.Load(sessionID)

server/session_test.go

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,45 @@ func TestMCPServer_AddSessionTools(t *testing.T) {
230230
assert.Contains(t, session.GetSessionTools(), "session-tool")
231231
}
232232

233+
func TestMCPServer_AddSessionTool(t *testing.T) {
234+
server := NewMCPServer("test-server", "1.0.0", WithToolCapabilities(true))
235+
ctx := context.Background()
236+
237+
// Create a session
238+
sessionChan := make(chan mcp.JSONRPCNotification, 10)
239+
session := &sessionTestClientWithTools{
240+
sessionID: "session-1",
241+
notificationChannel: sessionChan,
242+
initialized: true,
243+
}
244+
245+
// Register the session
246+
err := server.RegisterSession(ctx, session)
247+
require.NoError(t, err)
248+
249+
// Add session-specific tool using the new helper method
250+
err = server.AddSessionTool(
251+
session.SessionID(),
252+
mcp.NewTool("session-tool-helper"),
253+
func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
254+
return mcp.NewToolResultText("helper result"), nil
255+
},
256+
)
257+
require.NoError(t, err)
258+
259+
// Check that notification was sent
260+
select {
261+
case notification := <-sessionChan:
262+
assert.Equal(t, "notifications/tools/list_changed", notification.Method)
263+
case <-time.After(100 * time.Millisecond):
264+
t.Error("Expected notification not received")
265+
}
266+
267+
// Verify tool was added to session
268+
assert.Len(t, session.GetSessionTools(), 1)
269+
assert.Contains(t, session.GetSessionTools(), "session-tool-helper")
270+
}
271+
233272
func TestMCPServer_DeleteSessionTools(t *testing.T) {
234273
server := NewMCPServer("test-server", "1.0.0", WithToolCapabilities(true))
235274
ctx := context.Background()

0 commit comments

Comments
 (0)