| 1 | package plugin |
| 2 | |
| 3 | import ( |
| 4 | "bufio" |
| 5 | "bytes" |
| 6 | "context" |
| 7 | "encoding/json" |
| 8 | "errors" |
| 9 | "os" |
| 10 | "path/filepath" |
| 11 | "slices" |
| 12 | "strings" |
| 13 | "sync" |
| 14 | "testing" |
| 15 | "time" |
| 16 | |
| 17 | "reasonix/internal/tool" |
| 18 | ) |
| 19 | |
| 20 | type notificationToolsTransport struct { |
| 21 | mu sync.Mutex |
| 22 | refreshActive bool |
| 23 | refreshStartedCh chan struct{} |
| 24 | refreshOnce sync.Once |
| 25 | notifications notificationRouter |
| 26 | closeOnce sync.Once |
| 27 | closed chan struct{} |
| 28 | } |
| 29 | |
| 30 | func newNotificationToolsTransport() *notificationToolsTransport { |
| 31 | return ¬ificationToolsTransport{closed: make(chan struct{}), refreshStartedCh: make(chan struct{})} |
| 32 | } |
| 33 | |
| 34 | func (t *notificationToolsTransport) call(ctx context.Context, method string, _ any) (json.RawMessage, error) { |
| 35 | if method != "tools/list" { |
| 36 | return json.RawMessage(`{}`), nil |
| 37 | } |
| 38 | t.mu.Lock() |
| 39 | t.refreshActive = true |
| 40 | t.mu.Unlock() |
| 41 | t.refreshOnce.Do(func() { close(t.refreshStartedCh) }) |
| 42 | <-ctx.Done() |
| 43 | return nil, ctx.Err() |
| 44 | } |
| 45 | |
| 46 | type controlledToolsTransport struct { |
| 47 | mu sync.Mutex |
| 48 | notifications notificationRouter |
| 49 | listCalls int |
| 50 | toolCalls int |
| 51 | emitOnList bool |
| 52 | failList bool |
| 53 | blockList bool |
| 54 | blockTool bool |
| 55 | toolName string |
| 56 | listStarted chan int |
| 57 | listRelease chan struct{} |
| 58 | toolStarted chan struct{} |
| 59 | toolRelease chan struct{} |
| 60 | } |
| 61 | |
| 62 | func newControlledToolsTransport() *controlledToolsTransport { |
| 63 | return &controlledToolsTransport{ |
| 64 | toolName: "echo", |
| 65 | listStarted: make(chan int, 16), listRelease: make(chan struct{}, 16), |
| 66 | toolStarted: make(chan struct{}, 1), toolRelease: make(chan struct{}, 1), |
| 67 | } |
| 68 | } |
| 69 | |
| 70 | func (t *controlledToolsTransport) call(ctx context.Context, method string, _ any) (json.RawMessage, error) { |
| 71 | switch method { |
| 72 | case "tools/list": |
| 73 | t.mu.Lock() |
| 74 | t.listCalls++ |
| 75 | call := t.listCalls |
| 76 | block, fail, emit, toolName := t.blockList, t.failList, t.emitOnList, t.toolName |
| 77 | t.mu.Unlock() |
| 78 | t.listStarted <- call |
| 79 | if block { |
| 80 | select { |
| 81 | case <-t.listRelease: |
| 82 | case <-ctx.Done(): |
| 83 | return nil, ctx.Err() |
| 84 | } |
| 85 | } |
| 86 | if emit { |
| 87 | t.notifications.dispatchNotification("notifications/tools/list_changed", nil) |
| 88 | } |
| 89 | if fail { |
| 90 | return nil, errors.New("tools/list failed") |
| 91 | } |
| 92 | response, _ := json.Marshal(map[string]any{"tools": []map[string]any{{ |
| 93 | "name": toolName, "description": "Echo.", "inputSchema": map[string]any{"type": "object"}, |
| 94 | }}}) |
| 95 | return response, nil |
| 96 | case "tools/call": |
| 97 | t.mu.Lock() |
| 98 | t.toolCalls++ |
| 99 | block := t.blockTool |
| 100 | t.mu.Unlock() |
| 101 | if block { |
| 102 | t.toolStarted <- struct{}{} |
| 103 | select { |
| 104 | case <-t.toolRelease: |
| 105 | case <-ctx.Done(): |
| 106 | return nil, ctx.Err() |
| 107 | } |
| 108 | } |
| 109 | return json.RawMessage(`{"content":[{"type":"text","text":"ok"}]}`), nil |
| 110 | default: |
| 111 | return json.RawMessage(`{}`), nil |
| 112 | } |
| 113 | } |
| 114 | |
| 115 | func (*controlledToolsTransport) close() {} |
| 116 | func (t *controlledToolsTransport) registerNotification(method string, callback func(json.RawMessage)) func() { |
| 117 | return t.notifications.registerNotification(method, callback) |
| 118 | } |
| 119 | func (t *controlledToolsTransport) emit() { |
| 120 | t.notifications.dispatchNotification("notifications/tools/list_changed", nil) |
| 121 | } |
| 122 | func (t *controlledToolsTransport) counts() (list, calls int) { |
| 123 | t.mu.Lock() |
| 124 | defer t.mu.Unlock() |
| 125 | return t.listCalls, t.toolCalls |
| 126 | } |
| 127 | |
| 128 | func newControlledRefreshClient(t *testing.T, tr *controlledToolsTransport) (*Client, tool.Tool) { |
| 129 | t.Helper() |
| 130 | ctx, cancel := context.WithCancel(context.Background()) |
| 131 | client := &Client{ |
| 132 | name: "controlled", t: tr, spec: Spec{Name: "controlled"}, capabilities: clientCapabilities{toolsListChanged: true}, |
| 133 | refresh: toolListRefreshState{ctx: ctx, cancel: cancel, wait: func(context.Context, time.Duration) error { return nil }}, |
| 134 | } |
| 135 | tools, err := client.listTools(ctx) |
| 136 | if err != nil { |
| 137 | client.close() |
| 138 | t.Fatalf("initial listTools: %v", err) |
| 139 | } |
| 140 | <-tr.listStarted // drain the synchronous initial tools/list observation |
| 141 | client.watchToolListChanges() |
| 142 | return client, tools[0] |
| 143 | } |
| 144 | |
| 145 | func refreshDone(t *testing.T, client *Client) <-chan struct{} { |
| 146 | t.Helper() |
| 147 | client.refresh.mu.Lock() |
| 148 | defer client.refresh.mu.Unlock() |
| 149 | if client.refresh.cycleDone == nil { |
| 150 | t.Fatal("refresh cycle did not start") |
| 151 | } |
| 152 | return client.refresh.cycleDone |
| 153 | } |
| 154 | |
| 155 | func waitClosed(t *testing.T, ch <-chan struct{}, label string) { |
| 156 | t.Helper() |
| 157 | select { |
| 158 | case <-ch: |
| 159 | case <-time.After(time.Second): |
| 160 | t.Fatalf("timed out waiting for %s", label) |
| 161 | } |
| 162 | } |
| 163 | |
| 164 | func TestToolListRefreshCoalescesNotificationBurst(t *testing.T) { |
| 165 | tr := newControlledToolsTransport() |
| 166 | client, _ := newControlledRefreshClient(t, tr) |
| 167 | defer client.close() |
| 168 | |
| 169 | waitStarted := make(chan struct{}, 1) |
| 170 | releaseWait := make(chan struct{}) |
| 171 | client.refresh.mu.Lock() |
| 172 | client.refresh.wait = func(ctx context.Context, _ time.Duration) error { |
| 173 | waitStarted <- struct{}{} |
| 174 | select { |
| 175 | case <-releaseWait: |
| 176 | return nil |
| 177 | case <-ctx.Done(): |
| 178 | return ctx.Err() |
| 179 | } |
| 180 | } |
| 181 | client.refresh.mu.Unlock() |
| 182 | |
| 183 | tr.emit() |
| 184 | done := refreshDone(t, client) |
| 185 | <-waitStarted |
| 186 | for range 99 { |
| 187 | tr.emit() |
| 188 | } |
| 189 | close(releaseWait) |
| 190 | waitClosed(t, done, "coalesced refresh") |
| 191 | if lists, _ := tr.counts(); lists != 2 { |
| 192 | t.Fatalf("tools/list calls = %d, want initial + one coalesced refresh", lists) |
| 193 | } |
| 194 | if client.toolCatalogStale() { |
| 195 | t.Fatal("catalog remained stale after coalesced refresh") |
| 196 | } |
| 197 | } |
| 198 | |
| 199 | func TestToolListRefreshBacksOffAndConvergesAfterRepeatedNotices(t *testing.T) { |
| 200 | tr := newControlledToolsTransport() |
| 201 | client, _ := newControlledRefreshClient(t, tr) |
| 202 | defer client.close() |
| 203 | delays := make(chan time.Duration, 4) |
| 204 | releaseDelay := make(chan struct{}, 4) |
| 205 | client.refresh.mu.Lock() |
| 206 | client.refresh.wait = func(ctx context.Context, delay time.Duration) error { |
| 207 | select { |
| 208 | case delays <- delay: |
| 209 | case <-ctx.Done(): |
| 210 | return ctx.Err() |
| 211 | } |
| 212 | select { |
| 213 | case <-releaseDelay: |
| 214 | return nil |
| 215 | case <-ctx.Done(): |
| 216 | return ctx.Err() |
| 217 | } |
| 218 | } |
| 219 | client.refresh.mu.Unlock() |
| 220 | tr.mu.Lock() |
| 221 | tr.blockList = true |
| 222 | tr.emitOnList = true |
| 223 | tr.mu.Unlock() |
| 224 | |
| 225 | tr.emit() |
| 226 | done := refreshDone(t, client) |
| 227 | if delay := <-delays; delay != toolListRefreshDebounce { |
| 228 | t.Fatalf("first refresh delay = %s, want %s", delay, toolListRefreshDebounce) |
| 229 | } |
| 230 | releaseDelay <- struct{}{} |
| 231 | if call := <-tr.listStarted; call != 2 { |
| 232 | t.Fatalf("first refresh call = %d, want 2", call) |
| 233 | } |
| 234 | tr.listRelease <- struct{}{} |
| 235 | if delay := <-delays; delay != 2*toolListRefreshDebounce { |
| 236 | t.Fatalf("first catch-up delay = %s, want %s", delay, 2*toolListRefreshDebounce) |
| 237 | } |
| 238 | releaseDelay <- struct{}{} |
| 239 | if call := <-tr.listStarted; call != 3 { |
| 240 | t.Fatalf("catch-up refresh call = %d, want 3", call) |
| 241 | } |
| 242 | tr.listRelease <- struct{}{} |
| 243 | if delay := <-delays; delay != 4*toolListRefreshDebounce { |
| 244 | t.Fatalf("second catch-up delay = %s, want %s", delay, 4*toolListRefreshDebounce) |
| 245 | } |
| 246 | tr.mu.Lock() |
| 247 | tr.emitOnList = false |
| 248 | tr.mu.Unlock() |
| 249 | releaseDelay <- struct{}{} |
| 250 | if call := <-tr.listStarted; call != 4 { |
| 251 | t.Fatalf("second catch-up refresh call = %d, want 4", call) |
| 252 | } |
| 253 | tr.listRelease <- struct{}{} |
| 254 | waitClosed(t, done, "backed-off refresh convergence") |
| 255 | if lists, _ := tr.counts(); lists != 4 { |
| 256 | t.Fatalf("tools/list calls = %d, want initial + three converging attempts", lists) |
| 257 | } |
| 258 | if client.toolCatalogStale() { |
| 259 | t.Fatal("catalog remained stale after the self-notification stopped") |
| 260 | } |
| 261 | } |
| 262 | |
| 263 | func TestToolListRefreshBoundsPermanentSelfNotificationsAndRecoversOnRetry(t *testing.T) { |
| 264 | tr := newControlledToolsTransport() |
| 265 | client, adapter := newControlledRefreshClient(t, tr) |
| 266 | defer client.close() |
| 267 | delays := make(chan time.Duration, toolListRefreshMaxAttempts+2) |
| 268 | releaseDelay := make(chan struct{}, toolListRefreshMaxAttempts+2) |
| 269 | client.refresh.mu.Lock() |
| 270 | client.refresh.wait = func(ctx context.Context, delay time.Duration) error { |
| 271 | select { |
| 272 | case delays <- delay: |
| 273 | case <-ctx.Done(): |
| 274 | return ctx.Err() |
| 275 | } |
| 276 | select { |
| 277 | case <-releaseDelay: |
| 278 | return nil |
| 279 | case <-ctx.Done(): |
| 280 | return ctx.Err() |
| 281 | } |
| 282 | } |
| 283 | client.refresh.mu.Unlock() |
| 284 | tr.mu.Lock() |
| 285 | tr.emitOnList = true |
| 286 | tr.mu.Unlock() |
| 287 | |
| 288 | tr.emit() |
| 289 | done := refreshDone(t, client) |
| 290 | wantDelay := toolListRefreshDebounce |
| 291 | for attempt := range toolListRefreshMaxAttempts { |
| 292 | if delay := <-delays; delay != wantDelay { |
| 293 | t.Fatalf("refresh attempt %d delay = %s, want %s", attempt+1, delay, wantDelay) |
| 294 | } |
| 295 | releaseDelay <- struct{}{} |
| 296 | if call := <-tr.listStarted; call != attempt+2 { |
| 297 | t.Fatalf("refresh attempt %d tools/list call = %d, want %d", attempt+1, call, attempt+2) |
| 298 | } |
| 299 | wantDelay = nextToolListRefreshDelay(wantDelay) |
| 300 | } |
| 301 | waitClosed(t, done, "bounded self-notification refresh") |
| 302 | if lists, _ := tr.counts(); lists != 1+toolListRefreshMaxAttempts { |
| 303 | t.Fatalf("tools/list calls = %d, want initial + %d bounded attempts", lists, toolListRefreshMaxAttempts) |
| 304 | } |
| 305 | if !client.toolCatalogStale() { |
| 306 | t.Fatal("permanently self-notifying server incorrectly marked its catalog current") |
| 307 | } |
| 308 | select { |
| 309 | case delay := <-delays: |
| 310 | t.Fatalf("refresh cycle scheduled an unbounded extra delay %s", delay) |
| 311 | default: |
| 312 | } |
| 313 | |
| 314 | // A user attempt on the stale adapter fails closed and starts a fresh cycle. |
| 315 | // Once the server stops self-notifying, that bounded retry converges and the |
| 316 | // unchanged adapter becomes callable again. |
| 317 | tr.mu.Lock() |
| 318 | tr.emitOnList = false |
| 319 | tr.mu.Unlock() |
| 320 | if _, err := adapter.Execute(context.Background(), json.RawMessage(`{}`)); err == nil || !strings.Contains(err.Error(), "refresh is still pending or failed") { |
| 321 | t.Fatalf("stale adapter error = %v, want fail-closed refresh error", err) |
| 322 | } |
| 323 | retryDone := refreshDone(t, client) |
| 324 | if delay := <-delays; delay != toolListRefreshDebounce { |
| 325 | t.Fatalf("retry refresh delay = %s, want %s", delay, toolListRefreshDebounce) |
| 326 | } |
| 327 | releaseDelay <- struct{}{} |
| 328 | if call := <-tr.listStarted; call != 2+toolListRefreshMaxAttempts { |
| 329 | t.Fatalf("retry tools/list call = %d, want %d", call, 2+toolListRefreshMaxAttempts) |
| 330 | } |
| 331 | waitClosed(t, retryDone, "retry convergence") |
| 332 | if client.toolCatalogStale() { |
| 333 | t.Fatal("catalog remained stale after the server stopped self-notifying") |
| 334 | } |
| 335 | if _, err := adapter.Execute(context.Background(), json.RawMessage(`{}`)); err != nil { |
| 336 | t.Fatalf("adapter after bounded recovery: %v", err) |
| 337 | } |
| 338 | if _, calls := tr.counts(); calls != 1 { |
| 339 | t.Fatalf("tools/call count = %d, want one post-recovery dispatch", calls) |
| 340 | } |
| 341 | } |
| 342 | |
| 343 | func TestToolListRefreshTimeoutUsesResolvedBudgets(t *testing.T) { |
| 344 | tests := []struct { |
| 345 | name string |
| 346 | spec Spec |
| 347 | want time.Duration |
| 348 | }{ |
| 349 | {name: "built-in defaults", spec: Spec{}, want: defaultStartupTimeout}, |
| 350 | {name: "call timeout is stricter", spec: Spec{StartupTimeout: 45 * time.Second, CallTimeout: 30 * time.Second}, want: 30 * time.Second}, |
| 351 | {name: "startup timeout is stricter", spec: Spec{StartupTimeout: 10 * time.Second, CallTimeout: 60 * time.Second}, want: 10 * time.Second}, |
| 352 | {name: "global defaults", spec: Spec{DefaultStartupTimeout: 40 * time.Second, DefaultCallTimeout: 20 * time.Second}, want: 20 * time.Second}, |
| 353 | } |
| 354 | for _, tc := range tests { |
| 355 | t.Run(tc.name, func(t *testing.T) { |
| 356 | client := &Client{spec: tc.spec} |
| 357 | if got := client.toolListRefreshTimeout(); got != tc.want { |
| 358 | t.Fatalf("refresh timeout = %s, want %s", got, tc.want) |
| 359 | } |
| 360 | }) |
| 361 | } |
| 362 | } |
| 363 | |
| 364 | func TestServersDoesNotWaitForBlockedToolCall(t *testing.T) { |
| 365 | tr := newControlledToolsTransport() |
| 366 | client, adapter := newControlledRefreshClient(t, tr) |
| 367 | defer client.close() |
| 368 | tr.mu.Lock() |
| 369 | tr.blockTool = true |
| 370 | tr.mu.Unlock() |
| 371 | |
| 372 | callDone := make(chan error, 1) |
| 373 | go func() { |
| 374 | _, err := adapter.Execute(context.Background(), json.RawMessage(`{}`)) |
| 375 | callDone <- err |
| 376 | }() |
| 377 | <-tr.toolStarted |
| 378 | |
| 379 | host := &Host{clients: []*Client{client}} |
| 380 | statusDone := make(chan []ServerStatus, 1) |
| 381 | go func() { statusDone <- host.Servers() }() |
| 382 | select { |
| 383 | case statuses := <-statusDone: |
| 384 | if len(statuses) != 1 || statuses[0].Tools != 1 { |
| 385 | t.Fatalf("statuses = %+v, want one server with one tool", statuses) |
| 386 | } |
| 387 | case <-time.After(time.Second): |
| 388 | tr.toolRelease <- struct{}{} |
| 389 | t.Fatal("Host.Servers blocked behind tools/call") |
| 390 | } |
| 391 | tr.toolRelease <- struct{}{} |
| 392 | if err := <-callDone; err != nil { |
| 393 | t.Fatalf("blocked tool call: %v", err) |
| 394 | } |
| 395 | } |
| 396 | |
| 397 | func TestChangedCatalogPublishesAfterInFlightToolCall(t *testing.T) { |
| 398 | tr := newControlledToolsTransport() |
| 399 | client, oldAdapter := newControlledRefreshClient(t, tr) |
| 400 | defer client.close() |
| 401 | tr.mu.Lock() |
| 402 | tr.blockTool = true |
| 403 | tr.toolName = "echo_v2" |
| 404 | tr.mu.Unlock() |
| 405 | |
| 406 | callDone := make(chan error, 1) |
| 407 | go func() { |
| 408 | _, err := oldAdapter.Execute(context.Background(), json.RawMessage(`{}`)) |
| 409 | callDone <- err |
| 410 | }() |
| 411 | <-tr.toolStarted |
| 412 | tr.emit() |
| 413 | done := refreshDone(t, client) |
| 414 | <-tr.listStarted |
| 415 | select { |
| 416 | case <-done: |
| 417 | t.Fatal("changed catalog published before the admitted tool call completed") |
| 418 | default: |
| 419 | } |
| 420 | |
| 421 | tr.toolRelease <- struct{}{} |
| 422 | if err := <-callDone; err != nil { |
| 423 | t.Fatalf("admitted tool call: %v", err) |
| 424 | } |
| 425 | waitClosed(t, done, "catalog publication after tool call") |
| 426 | if _, err := oldAdapter.Execute(context.Background(), json.RawMessage(`{}`)); err == nil || !strings.Contains(err.Error(), "changed tool") { |
| 427 | t.Fatalf("old adapter after publication = %v, want changed-tool refusal", err) |
| 428 | } |
| 429 | current, ok := client.cachedTools() |
| 430 | if !ok || findToolByName(current, "mcp__controlled__echo_v2") == nil { |
| 431 | t.Fatalf("current tools = %v, want echo_v2", toolNames(current)) |
| 432 | } |
| 433 | } |
| 434 | |
| 435 | func TestServersDoesNotWaitForBlockedToolListRefresh(t *testing.T) { |
| 436 | tr := newControlledToolsTransport() |
| 437 | client, _ := newControlledRefreshClient(t, tr) |
| 438 | defer client.close() |
| 439 | tr.mu.Lock() |
| 440 | tr.blockList = true |
| 441 | tr.mu.Unlock() |
| 442 | |
| 443 | tr.emit() |
| 444 | done := refreshDone(t, client) |
| 445 | <-tr.listStarted |
| 446 | host := &Host{clients: []*Client{client}} |
| 447 | statusDone := make(chan []ServerStatus, 1) |
| 448 | go func() { statusDone <- host.Servers() }() |
| 449 | select { |
| 450 | case statuses := <-statusDone: |
| 451 | if len(statuses) != 1 || statuses[0].Tools != 1 { |
| 452 | t.Fatalf("statuses = %+v, want previous complete snapshot", statuses) |
| 453 | } |
| 454 | case <-time.After(time.Second): |
| 455 | tr.listRelease <- struct{}{} |
| 456 | t.Fatal("Host.Servers blocked behind tools/list") |
| 457 | } |
| 458 | tr.listRelease <- struct{}{} |
| 459 | waitClosed(t, done, "blocked refresh release") |
| 460 | } |
| 461 | |
| 462 | func TestToolListRefreshFailureKeepsOldAdapterFailClosed(t *testing.T) { |
| 463 | tr := newControlledToolsTransport() |
| 464 | client, adapter := newControlledRefreshClient(t, tr) |
| 465 | defer client.close() |
| 466 | tr.mu.Lock() |
| 467 | tr.blockList = true |
| 468 | tr.failList = true |
| 469 | tr.mu.Unlock() |
| 470 | |
| 471 | tr.emit() |
| 472 | done := refreshDone(t, client) |
| 473 | <-tr.listStarted |
| 474 | tr.listRelease <- struct{}{} |
| 475 | waitClosed(t, done, "failed refresh") |
| 476 | if !client.toolCatalogStale() { |
| 477 | t.Fatal("failed refresh incorrectly marked the old catalog current") |
| 478 | } |
| 479 | |
| 480 | tr.mu.Lock() |
| 481 | tr.failList = false |
| 482 | tr.mu.Unlock() |
| 483 | if _, err := adapter.Execute(context.Background(), json.RawMessage(`{}`)); err == nil || !strings.Contains(err.Error(), "refresh is still pending or failed") { |
| 484 | t.Fatalf("stale adapter error = %v, want fail-closed refresh error", err) |
| 485 | } |
| 486 | if _, calls := tr.counts(); calls != 0 { |
| 487 | t.Fatalf("stale adapter reached tools/call %d times", calls) |
| 488 | } |
| 489 | retryDone := refreshDone(t, client) |
| 490 | <-tr.listStarted |
| 491 | client.close() |
| 492 | waitClosed(t, retryDone, "cancelled retry refresh") |
| 493 | } |
| 494 | |
| 495 | func TestToolListRefreshNoOpPreservesAdapterGeneration(t *testing.T) { |
| 496 | tr := newControlledToolsTransport() |
| 497 | client, adapter := newControlledRefreshClient(t, tr) |
| 498 | defer client.close() |
| 499 | tr.mu.Lock() |
| 500 | tr.blockList = true |
| 501 | tr.mu.Unlock() |
| 502 | remote := adapter.(*remoteTool) |
| 503 | generation := remote.generation |
| 504 | changes := make(chan struct{}, 1) |
| 505 | client.setToolsChangedCallback(func([]tool.Tool) { changes <- struct{}{} }) |
| 506 | |
| 507 | tr.emit() |
| 508 | done := refreshDone(t, client) |
| 509 | <-tr.listStarted |
| 510 | tr.listRelease <- struct{}{} |
| 511 | waitClosed(t, done, "no-op refresh") |
| 512 | if client.toolCatalogStale() { |
| 513 | t.Fatal("no-op refresh did not clear the notification revision") |
| 514 | } |
| 515 | if remote.generation != generation || client.catalogGeneration != generation { |
| 516 | t.Fatalf("generation changed on no-op: adapter=%d catalog=%d want=%d", remote.generation, client.catalogGeneration, generation) |
| 517 | } |
| 518 | select { |
| 519 | case <-changes: |
| 520 | t.Fatal("no-op refresh published a change callback") |
| 521 | default: |
| 522 | } |
| 523 | if _, err := adapter.Execute(context.Background(), json.RawMessage(`{}`)); err != nil { |
| 524 | t.Fatalf("adapter after no-op refresh: %v", err) |
| 525 | } |
| 526 | } |
| 527 | |
| 528 | func TestClientIgnoresToolListChangedWithoutAdvertisedCapability(t *testing.T) { |
| 529 | tr := newControlledToolsTransport() |
| 530 | client := &Client{name: "unsupported", t: tr, spec: Spec{Name: "unsupported"}} |
| 531 | client.watchToolListChanges() |
| 532 | tr.notifications.mu.Lock() |
| 533 | listeners := len(tr.notifications.listeners["notifications/tools/list_changed"]) |
| 534 | tr.notifications.mu.Unlock() |
| 535 | if listeners != 0 { |
| 536 | t.Fatalf("notification listeners = %d, want none without tools.listChanged", listeners) |
| 537 | } |
| 538 | } |
| 539 | |
| 540 | func (t *notificationToolsTransport) registerNotification(method string, callback func(json.RawMessage)) func() { |
| 541 | return t.notifications.registerNotification(method, callback) |
| 542 | } |
| 543 | |
| 544 | func (t *notificationToolsTransport) emit(method string) { |
| 545 | t.notifications.dispatchNotification(method, nil) |
| 546 | } |
| 547 | |
| 548 | func (t *notificationToolsTransport) close() { |
| 549 | t.closeOnce.Do(func() { close(t.closed) }) |
| 550 | } |
| 551 | |
| 552 | func TestHostRefreshesToolsAfterListChangedNotification(t *testing.T) { |
| 553 | ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) |
| 554 | defer cancel() |
| 555 | |
| 556 | startCount := filepath.Join(t.TempDir(), "starts") |
| 557 | spec := Spec{ |
| 558 | Name: "dynamic", |
| 559 | Command: os.Args[0], |
| 560 | Args: []string{"-test.run=TestDynamicToolsHelperProcess", "--"}, |
| 561 | Env: map[string]string{ |
| 562 | "GO_WANT_DYNAMIC_TOOLS_HELPER": "1", |
| 563 | "GO_WANT_HELPER_START_COUNT": startCount, |
| 564 | }, |
| 565 | } |
| 566 | |
| 567 | host := NewHost() |
| 568 | initial, err := host.Add(ctx, spec) |
| 569 | if err != nil { |
| 570 | t.Fatalf("Host.Add: %v", err) |
| 571 | } |
| 572 | defer host.Close() |
| 573 | if got := toolNames(initial); !slices.Equal(got, []string{"mcp__dynamic__load_toolset"}) { |
| 574 | t.Fatalf("initial tools = %v, want load_toolset only", got) |
| 575 | } |
| 576 | |
| 577 | changes := make(chan []tool.Tool, 1) |
| 578 | unsubscribe := host.SubscribeToolListChanges(ctx, func(changed Spec, tools []tool.Tool) { |
| 579 | if MCPRuntimeSpecMatches(changed, spec) { |
| 580 | changes <- tools |
| 581 | } |
| 582 | }) |
| 583 | defer unsubscribe() |
| 584 | loader := findToolByName(initial, "mcp__dynamic__load_toolset") |
| 585 | if _, err := loader.Execute(ctx, json.RawMessage(`{}`)); err != nil { |
| 586 | t.Fatalf("load_toolset: %v", err) |
| 587 | } |
| 588 | |
| 589 | select { |
| 590 | case refreshed := <-changes: |
| 591 | if findToolByName(refreshed, "mcp__dynamic__list_schematic_components") == nil { |
| 592 | t.Fatalf("refreshed tools = %v, want list_schematic_components", toolNames(refreshed)) |
| 593 | } |
| 594 | if _, err := loader.Execute(ctx, json.RawMessage(`{}`)); err == nil || !strings.Contains(err.Error(), "changed tool") { |
| 595 | t.Fatalf("stale pre-refresh adapter error = %v, want retryable changed-tool refusal", err) |
| 596 | } |
| 597 | cached, listErr := host.ToolsFor(ctx, spec.Name) |
| 598 | if listErr != nil { |
| 599 | t.Fatalf("ToolsFor after list_changed: %v", listErr) |
| 600 | } |
| 601 | if findToolByName(cached, "mcp__dynamic__list_schematic_components") == nil { |
| 602 | t.Fatalf("cached tools = %v, want list_schematic_components", toolNames(cached)) |
| 603 | } |
| 604 | case <-time.After(2 * time.Second): |
| 605 | t.Fatal("host did not publish refreshed tools after notifications/tools/list_changed") |
| 606 | } |
| 607 | if got := readHelperCounter(t, startCount); got != 1 { |
| 608 | t.Fatalf("process starts = %d, want one persistent MCP process", got) |
| 609 | } |
| 610 | } |
| 611 | |
| 612 | func TestClientCloseCancelsBlockedToolListRefresh(t *testing.T) { |
| 613 | ctx, cancel := context.WithCancel(context.Background()) |
| 614 | tr := newNotificationToolsTransport() |
| 615 | client := &Client{ |
| 616 | name: "blocked", |
| 617 | t: tr, |
| 618 | spec: Spec{Name: "blocked"}, |
| 619 | capabilities: clientCapabilities{toolsListChanged: true}, |
| 620 | refresh: toolListRefreshState{ |
| 621 | ctx: ctx, |
| 622 | cancel: cancel, |
| 623 | wait: func(context.Context, time.Duration) error { return nil }, |
| 624 | }, |
| 625 | } |
| 626 | client.watchToolListChanges() |
| 627 | tr.emit("notifications/tools/list_changed") |
| 628 | done := refreshDone(t, client) |
| 629 | waitClosed(t, tr.refreshStartedCh, "tools/list start") |
| 630 | client.close() |
| 631 | select { |
| 632 | case <-tr.closed: |
| 633 | case <-time.After(time.Second): |
| 634 | t.Fatal("client close did not close transport") |
| 635 | } |
| 636 | waitClosed(t, done, "refresh cancellation") |
| 637 | |
| 638 | client.refresh.mu.Lock() |
| 639 | defer client.refresh.mu.Unlock() |
| 640 | if !client.refresh.closed { |
| 641 | t.Fatal("refresh state was not closed") |
| 642 | } |
| 643 | } |
| 644 | |
| 645 | // TestDynamicToolsHelperProcess serves a minimal stdio MCP whose tool catalog |
| 646 | // expands after load_toolset and advertises that change through the protocol. |
| 647 | func TestDynamicToolsHelperProcess(t *testing.T) { |
| 648 | if os.Getenv("GO_WANT_DYNAMIC_TOOLS_HELPER") != "1" { |
| 649 | return |
| 650 | } |
| 651 | defer os.Exit(0) |
| 652 | incrementHelperCounter(os.Getenv("GO_WANT_HELPER_START_COUNT")) |
| 653 | |
| 654 | loaded := false |
| 655 | in := bufio.NewReader(os.Stdin) |
| 656 | for { |
| 657 | line, err := in.ReadBytes('\n') |
| 658 | if err != nil { |
| 659 | return |
| 660 | } |
| 661 | line = bytes.TrimSpace(line) |
| 662 | var request struct { |
| 663 | ID *int `json:"id"` |
| 664 | Method string `json:"method"` |
| 665 | Params json.RawMessage `json:"params"` |
| 666 | } |
| 667 | if len(line) == 0 || json.Unmarshal(line, &request) != nil || request.ID == nil { |
| 668 | continue |
| 669 | } |
| 670 | |
| 671 | var result any |
| 672 | notifyChanged := false |
| 673 | switch request.Method { |
| 674 | case "initialize": |
| 675 | result = map[string]any{ |
| 676 | "protocolVersion": testLegacyProtocolVersion, |
| 677 | "serverInfo": map[string]any{"name": "dynamic", "version": "1"}, |
| 678 | "capabilities": map[string]any{"tools": map[string]any{"listChanged": true}}, |
| 679 | } |
| 680 | case "tools/list": |
| 681 | tools := []map[string]any{{ |
| 682 | "name": "load_toolset", "description": "Load a toolset.", |
| 683 | "inputSchema": map[string]any{"type": "object"}, |
| 684 | }} |
| 685 | if loaded { |
| 686 | tools = append(tools, map[string]any{ |
| 687 | "name": "list_schematic_components", "description": "List schematic components.", |
| 688 | "inputSchema": map[string]any{"type": "object"}, |
| 689 | }) |
| 690 | } |
| 691 | result = map[string]any{"tools": tools} |
| 692 | case "tools/call": |
| 693 | loaded = true |
| 694 | notifyChanged = true |
| 695 | result = map[string]any{"content": []map[string]any{{"type": "text", "text": "loaded"}}} |
| 696 | } |
| 697 | |
| 698 | response, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": *request.ID, "result": result}) |
| 699 | _, _ = os.Stdout.Write(append(response, '\n')) |
| 700 | if notifyChanged { |
| 701 | notification, _ := json.Marshal(map[string]any{ |
| 702 | "jsonrpc": "2.0", "method": "notifications/tools/list_changed", |
| 703 | }) |
| 704 | _, _ = os.Stdout.Write(append(notification, '\n')) |
| 705 | } |
| 706 | } |
| 707 | } |
| 708 |