Skip to content

Commit c1e70f3

Browse files
david-hamilton-gleandavid-hamilton-glean
authored
fix(session): Don't send tool changed notifications if session not initialized yet (#289)
AddSessionTools and DeleteSessionTools can be called before the session is registered, in particular in the RegisterSession hook. They should not attempt to send `notifications/tools/list_changed` notifications in this case as it will only generate errors.
1 parent e767652 commit c1e70f3

File tree

2 files changed

+218
-26
lines changed

2 files changed

+218
-26
lines changed

‎server/session.go

Lines changed: 36 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -243,19 +243,24 @@ func (s *MCPServer) AddSessionTools(sessionID string, tools ...ServerTool) error
243243
// Set the tools (this should be thread-safe)
244244
session.SetSessionTools(newSessionTools)
245245

246-
// Send notification only to this session
247-
if err := s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil); err != nil {
248-
// Log the error but don't fail the operation
249-
// The tools were successfully added, but notification failed
250-
if s.hooks != nil && len(s.hooks.OnError) > 0 {
251-
hooks := s.hooks
252-
go func(sID string, hooks *Hooks) {
253-
ctx := context.Background()
254-
hooks.onError(ctx, nil, "notification", map[string]any{
255-
"method": "notifications/tools/list_changed",
256-
"sessionID": sID,
257-
}, fmt.Errorf("failed to send notification after adding tools: %w", err))
258-
}(sessionID, hooks)
246+
// It only makes sense to send tool notifications to initialized sessions --
247+
// if we're not initialized yet the client can't possibly have sent their
248+
// initial tools/list message
249+
if session.Initialized() {
250+
// Send notification only to this session
251+
if err := s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil); err != nil {
252+
// Log the error but don't fail the operation
253+
// The tools were successfully added, but notification failed
254+
if s.hooks != nil && len(s.hooks.OnError) > 0 {
255+
hooks := s.hooks
256+
go func(sID string, hooks *Hooks) {
257+
ctx := context.Background()
258+
hooks.onError(ctx, nil, "notification", map[string]any{
259+
"method": "notifications/tools/list_changed",
260+
"sessionID": sID,
261+
}, fmt.Errorf("failed to send notification after adding tools: %w", err))
262+
}(sessionID, hooks)
263+
}
259264
}
260265
}
261266

@@ -296,19 +301,24 @@ func (s *MCPServer) DeleteSessionTools(sessionID string, names ...string) error
296301
// Set the tools (this should be thread-safe)
297302
session.SetSessionTools(newSessionTools)
298303

299-
// Send notification only to this session
300-
if err := s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil); err != nil {
301-
// Log the error but don't fail the operation
302-
// The tools were successfully deleted, but notification failed
303-
if s.hooks != nil && len(s.hooks.OnError) > 0 {
304-
hooks := s.hooks
305-
go func(sID string, hooks *Hooks) {
306-
ctx := context.Background()
307-
hooks.onError(ctx, nil, "notification", map[string]any{
308-
"method": "notifications/tools/list_changed",
309-
"sessionID": sID,
310-
}, fmt.Errorf("failed to send notification after deleting tools: %w", err))
311-
}(sessionID, hooks)
304+
// It only makes sense to send tool notifications to initialized sessions --
305+
// if we're not initialized yet the client can't possibly have sent their
306+
// initial tools/list message
307+
if session.Initialized() {
308+
// Send notification only to this session
309+
if err := s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil); err != nil {
310+
// Log the error but don't fail the operation
311+
// The tools were successfully deleted, but notification failed
312+
if s.hooks != nil && len(s.hooks.OnError) > 0 {
313+
hooks := s.hooks
314+
go func(sID string, hooks *Hooks) {
315+
ctx := context.Background()
316+
hooks.onError(ctx, nil, "notification", map[string]any{
317+
"method": "notifications/tools/list_changed",
318+
"sessionID": sID,
319+
}, fmt.Errorf("failed to send notification after deleting tools: %w", err))
320+
}(sessionID, hooks)
321+
}
312322
}
313323
}
314324

‎server/session_test.go

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,188 @@ func TestMCPServer_AddSessionTool(t *testing.T) {
296296
assert.Contains(t, session.GetSessionTools(), "session-tool-helper")
297297
}
298298

299+
func TestMCPServer_AddSessionToolsUninitialized(t *testing.T) {
300+
// This test verifies that adding tools to an uninitialized session works correctly.
301+
//
302+
// This scenario can occur when tools are added during the session registration hook,
303+
// before the session is fully initialized. In this case, we should:
304+
// 1. Successfully add the tools to the session
305+
// 2. Not attempt to send a notification (since the session isn't ready)
306+
// 3. Have the tools available once the session is initialized
307+
// 4. Not trigger any error hooks when adding tools to uninitialized sessions
308+
309+
// Set up error hook to track if it's called
310+
errorChan := make(chan error)
311+
hooks := &Hooks{}
312+
hooks.AddOnError(
313+
func(ctx context.Context, id any, method mcp.MCPMethod, message any, err error) {
314+
errorChan <- err
315+
},
316+
)
317+
318+
server := NewMCPServer("test-server", "1.0.0",
319+
WithToolCapabilities(true),
320+
WithHooks(hooks),
321+
)
322+
ctx := context.Background()
323+
324+
// Create an uninitialized session
325+
sessionChan := make(chan mcp.JSONRPCNotification, 1)
326+
session := &sessionTestClientWithTools{
327+
sessionID: "uninitialized-session",
328+
notificationChannel: sessionChan,
329+
initialized: false,
330+
}
331+
332+
// Register the session
333+
err := server.RegisterSession(ctx, session)
334+
require.NoError(t, err)
335+
336+
// Add session-specific tools to the uninitialized session
337+
err = server.AddSessionTools(session.SessionID(),
338+
ServerTool{Tool: mcp.NewTool("uninitialized-tool")},
339+
)
340+
require.NoError(t, err)
341+
342+
// Verify no errors
343+
select {
344+
case err := <-errorChan:
345+
t.Error("Expected no errors, but OnError called with: ", err)
346+
case <-time.After(25 * time.Millisecond): // no errors
347+
}
348+
349+
// Verify no notification was sent (channel should be empty)
350+
select {
351+
case <-sessionChan:
352+
t.Error("Expected no notification to be sent for uninitialized session")
353+
default: // no notifications
354+
}
355+
356+
// Verify tool was added to session
357+
assert.Len(t, session.GetSessionTools(), 1)
358+
assert.Contains(t, session.GetSessionTools(), "uninitialized-tool")
359+
360+
// Initialize the session
361+
session.Initialize()
362+
363+
// Now verify that subsequent tool additions will send notifications
364+
err = server.AddSessionTools(session.SessionID(),
365+
ServerTool{Tool: mcp.NewTool("initialized-tool")},
366+
)
367+
require.NoError(t, err)
368+
369+
// Verify no errors
370+
select {
371+
case err := <-errorChan:
372+
t.Error("Expected no errors, but OnError called with:", err)
373+
case <-time.After(200 * time.Millisecond): // No errors
374+
}
375+
376+
// Verify notification was sent for the initialized session
377+
select {
378+
case notification := <-sessionChan:
379+
assert.Equal(t, "notifications/tools/list_changed", notification.Method)
380+
case <-time.After(100 * time.Millisecond):
381+
t.Error("Timeout waiting for expected notifications/tools/list_changed notification")
382+
}
383+
384+
// Verify both tools are available
385+
assert.Len(t, session.GetSessionTools(), 2)
386+
assert.Contains(t, session.GetSessionTools(), "uninitialized-tool")
387+
assert.Contains(t, session.GetSessionTools(), "initialized-tool")
388+
}
389+
390+
func TestMCPServer_DeleteSessionToolsUninitialized(t *testing.T) {
391+
// This test verifies that deleting tools from an uninitialized session works correctly.
392+
//
393+
// This is a bit of a weird edge case but can happen if tools are added and
394+
// deleted during the RegisterSession hook.
395+
//
396+
// In this case, we should:
397+
// 1. Successfully delete the tools from the session
398+
// 2. Not attempt to send a notification (since the session isn't ready)
399+
// 3. Have the tools properly deleted once the session is initialized
400+
// 4. Not trigger any error hooks when deleting tools from uninitialized sessions
401+
402+
// Set up error hook to track if it's called
403+
errorChan := make(chan error)
404+
hooks := &Hooks{}
405+
hooks.AddOnError(
406+
func(ctx context.Context, id any, method mcp.MCPMethod, message any, err error) {
407+
errorChan <- err
408+
},
409+
)
410+
411+
server := NewMCPServer("test-server", "1.0.0",
412+
WithToolCapabilities(true),
413+
WithHooks(hooks),
414+
)
415+
ctx := context.Background()
416+
417+
// Create an uninitialized session with some tools
418+
sessionChan := make(chan mcp.JSONRPCNotification, 1)
419+
session := &sessionTestClientWithTools{
420+
sessionID: "uninitialized-session",
421+
notificationChannel: sessionChan,
422+
initialized: false,
423+
sessionTools: map[string]ServerTool{
424+
"tool-to-delete": {Tool: mcp.NewTool("tool-to-delete")},
425+
"tool-to-keep": {Tool: mcp.NewTool("tool-to-keep")},
426+
},
427+
}
428+
429+
// Register the session
430+
err := server.RegisterSession(ctx, session)
431+
require.NoError(t, err)
432+
433+
// Delete a tool from the uninitialized session
434+
err = server.DeleteSessionTools(session.SessionID(), "tool-to-delete")
435+
require.NoError(t, err)
436+
437+
select {
438+
case err := <-errorChan:
439+
t.Errorf("Expected error hooks not to be called, got error: %v", err)
440+
case <-time.After(25 * time.Millisecond): // No errors
441+
}
442+
443+
// Verify no notification was sent (channel should be empty)
444+
select {
445+
case <-sessionChan:
446+
t.Error("Expected no notification to be sent for uninitialized session")
447+
default:
448+
// This is the expected case - no notification should be sent
449+
}
450+
451+
// Verify tool was deleted from session
452+
assert.Len(t, session.GetSessionTools(), 1)
453+
assert.NotContains(t, session.GetSessionTools(), "tool-to-delete")
454+
assert.Contains(t, session.GetSessionTools(), "tool-to-keep")
455+
456+
// Initialize the session
457+
session.Initialize()
458+
459+
// Now verify that subsequent tool deletions will send notifications
460+
err = server.DeleteSessionTools(session.SessionID(), "tool-to-keep")
461+
require.NoError(t, err)
462+
463+
select {
464+
case err := <-errorChan:
465+
t.Errorf("Expected error hooks not to be called, got error: %v", err)
466+
case <-time.After(200 * time.Millisecond): // No errors
467+
}
468+
469+
// Verify notification was sent for the initialized session
470+
select {
471+
case notification := <-sessionChan:
472+
assert.Equal(t, "notifications/tools/list_changed", notification.Method)
473+
case <-time.After(100 * time.Millisecond):
474+
t.Error("Expected notification not received for initialized session")
475+
}
476+
477+
// Verify all tools are deleted
478+
assert.Len(t, session.GetSessionTools(), 0)
479+
}
480+
299481
func TestMCPServer_CallSessionTool(t *testing.T) {
300482
server := NewMCPServer("test-server", "1.0.0", WithToolCapabilities(true))
301483

0 commit comments

Comments
 (0)