| 1 | package plugin |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "crypto/sha256" |
| 6 | "encoding/json" |
| 7 | "fmt" |
| 8 | "log/slog" |
| 9 | "sort" |
| 10 | "strings" |
| 11 | "sync" |
| 12 | "sync/atomic" |
| 13 | "time" |
| 14 | |
| 15 | "reasonix/internal/provider" |
| 16 | "reasonix/internal/tool" |
| 17 | ) |
| 18 | |
| 19 | const ( |
| 20 | toolListRefreshDebounce = 300 * time.Millisecond |
| 21 | toolListRefreshMaxBackoff = 5 * time.Second |
| 22 | toolListRefreshMaxAttempts = 4 |
| 23 | ) |
| 24 | |
| 25 | type toolListRefreshWait func(context.Context, time.Duration) error |
| 26 | |
| 27 | type toolListRefreshState struct { |
| 28 | mu sync.Mutex |
| 29 | ctx context.Context |
| 30 | cancel context.CancelFunc |
| 31 | closed bool |
| 32 | running bool |
| 33 | cycleDone chan struct{} |
| 34 | onChanged func([]tool.Tool) |
| 35 | stopNotifications func() |
| 36 | wait toolListRefreshWait |
| 37 | noticeRevision atomic.Uint64 |
| 38 | publishedRevision atomic.Uint64 |
| 39 | } |
| 40 | |
| 41 | type auxiliaryRefreshState struct { |
| 42 | revision atomic.Uint64 |
| 43 | applied atomic.Uint64 |
| 44 | running atomic.Bool |
| 45 | } |
| 46 | |
| 47 | type auxiliaryListRefreshState struct { |
| 48 | prompt auxiliaryRefreshState |
| 49 | resource auxiliaryRefreshState |
| 50 | } |
| 51 | |
| 52 | type clientCapabilities struct { |
| 53 | tools bool |
| 54 | prompts bool |
| 55 | resources bool |
| 56 | toolsListChanged bool |
| 57 | promptsListChanged bool |
| 58 | resourcesListChanged bool |
| 59 | // serverExtensions records extension IDs the server declared in its |
| 60 | // initialize capabilities. MCP Apps needs two-way agreement: the client |
| 61 | // declares io.modelcontextprotocol/ui, the server answers with it. |
| 62 | serverExtensions map[string]bool |
| 63 | } |
| 64 | |
| 65 | // appsUI reports two-way MCP Apps agreement for this server. |
| 66 | func (cc clientCapabilities) appsUI() bool { return cc.serverExtensions[AppsUIExtensionID] } |
| 67 | |
| 68 | func (c *Client) appsNegotiated() bool { |
| 69 | return c != nil && c.profile.Capabilities().AppsUI && c.capabilities.appsUI() |
| 70 | } |
| 71 | |
| 72 | // toolCatalogSnapshot is immutable after publication. All slices are built off |
| 73 | // lock and replaced together under toolsMu, so readers see either the complete |
| 74 | // old catalog or the complete new catalog. |
| 75 | type toolCatalogSnapshot struct { |
| 76 | listed bool |
| 77 | generation uint64 |
| 78 | fingerprint [sha256.Size]byte |
| 79 | infos []ToolInfo |
| 80 | adapters []tool.Tool |
| 81 | // appAdapters holds App-callable tools (visibility contains "app"), |
| 82 | // including app-only ones that never enter the model catalog. |
| 83 | appAdapters []tool.Tool |
| 84 | } |
| 85 | |
| 86 | type toolCatalogFingerprintEntry struct { |
| 87 | RawName string `json:"raw_name"` |
| 88 | VisibleName string `json:"visible_name"` |
| 89 | Description string `json:"description"` |
| 90 | Visibility []string `json:"visibility,omitempty"` |
| 91 | Schema json.RawMessage `json:"schema,omitempty"` |
| 92 | OutputSchema json.RawMessage `json:"output_schema,omitempty"` |
| 93 | SchemaError string `json:"schema_error,omitempty"` |
| 94 | ReadOnlyHint bool `json:"read_only_hint,omitempty"` |
| 95 | DestructiveHint bool `json:"destructive_hint,omitempty"` |
| 96 | } |
| 97 | |
| 98 | type toolListSubscriptions struct { |
| 99 | nextID atomic.Uint64 |
| 100 | subscribers map[uint64]*toolListSubscriber |
| 101 | } |
| 102 | |
| 103 | type toolListSubscriber struct { |
| 104 | ctx context.Context |
| 105 | callback func(Spec, []tool.Tool) |
| 106 | deliveryMu sync.Mutex |
| 107 | } |
| 108 | |
| 109 | type toolListReplay struct { |
| 110 | spec Spec |
| 111 | tools []tool.Tool |
| 112 | } |
| 113 | |
| 114 | // SubscribeToolListChanges receives refreshed live tool sets after a connected |
| 115 | // server sends notifications/tools/list_changed. The returned function removes |
| 116 | // the subscription; ctx cancellation does the same automatically. |
| 117 | func (h *Host) SubscribeToolListChanges(ctx context.Context, callback func(Spec, []tool.Tool)) func() { |
| 118 | return h.subscribeToolListChanges(ctx, callback, false) |
| 119 | } |
| 120 | |
| 121 | // SubscribeToolListChangesWithReplay first delivers the complete catalogs of |
| 122 | // already-connected servers, then every later list_changed publication. Replay |
| 123 | // and live deliveries are serialized per subscriber so an update cannot |
| 124 | // overtake the snapshot used to initialize a late session runtime. |
| 125 | func (h *Host) SubscribeToolListChangesWithReplay(ctx context.Context, callback func(Spec, []tool.Tool)) func() { |
| 126 | return h.subscribeToolListChanges(ctx, callback, true) |
| 127 | } |
| 128 | |
| 129 | func (h *Host) subscribeToolListChanges(ctx context.Context, callback func(Spec, []tool.Tool), replay bool) func() { |
| 130 | if h == nil || callback == nil { |
| 131 | return func() {} |
| 132 | } |
| 133 | if ctx == nil { |
| 134 | ctx = context.Background() |
| 135 | } |
| 136 | subscriber := &toolListSubscriber{ctx: ctx, callback: callback} |
| 137 | if replay { |
| 138 | // Publications can discover the subscriber as soon as h.mu is released. |
| 139 | // Hold its delivery gate until every captured catalog has been replayed so |
| 140 | // those later publications cannot apply first and then be overwritten. |
| 141 | subscriber.deliveryMu.Lock() |
| 142 | } |
| 143 | id := h.toolListChanges.nextID.Add(1) |
| 144 | h.mu.Lock() |
| 145 | if h.closed { |
| 146 | h.mu.Unlock() |
| 147 | if replay { |
| 148 | subscriber.deliveryMu.Unlock() |
| 149 | } |
| 150 | return func() {} |
| 151 | } |
| 152 | if h.toolListChanges.subscribers == nil { |
| 153 | h.toolListChanges.subscribers = map[uint64]*toolListSubscriber{} |
| 154 | } |
| 155 | h.toolListChanges.subscribers[id] = subscriber |
| 156 | var snapshots []toolListReplay |
| 157 | if replay { |
| 158 | snapshots = make([]toolListReplay, 0, len(h.clients)) |
| 159 | for _, client := range h.clients { |
| 160 | if client == nil { |
| 161 | continue |
| 162 | } |
| 163 | if tools, ok := client.cachedTools(); ok { |
| 164 | snapshots = append(snapshots, toolListReplay{spec: client.spec, tools: tools}) |
| 165 | } |
| 166 | } |
| 167 | } |
| 168 | h.mu.Unlock() |
| 169 | if replay { |
| 170 | for _, snapshot := range snapshots { |
| 171 | if ctx.Err() != nil { |
| 172 | break |
| 173 | } |
| 174 | callback(snapshot.spec, append([]tool.Tool(nil), snapshot.tools...)) |
| 175 | } |
| 176 | subscriber.deliveryMu.Unlock() |
| 177 | } |
| 178 | |
| 179 | var once sync.Once |
| 180 | unsubscribe := func() { |
| 181 | once.Do(func() { |
| 182 | h.mu.Lock() |
| 183 | delete(h.toolListChanges.subscribers, id) |
| 184 | h.mu.Unlock() |
| 185 | }) |
| 186 | } |
| 187 | stop := context.AfterFunc(ctx, unsubscribe) |
| 188 | return func() { |
| 189 | stop() |
| 190 | unsubscribe() |
| 191 | } |
| 192 | } |
| 193 | |
| 194 | func (h *Host) bindToolListChanges(c *Client) { |
| 195 | if h == nil || c == nil { |
| 196 | return |
| 197 | } |
| 198 | c.setToolsChangedCallback(func(tools []tool.Tool) { |
| 199 | h.publishToolListChange(c, tools) |
| 200 | }) |
| 201 | c.watchToolListChanges() |
| 202 | h.watchAuxiliaryListChanges(c) |
| 203 | } |
| 204 | |
| 205 | func (h *Host) watchAuxiliaryListChanges(c *Client) { |
| 206 | t, ok := c.t.(notificationTransport) |
| 207 | if !ok { |
| 208 | return |
| 209 | } |
| 210 | c.refresh.mu.Lock() |
| 211 | ctx := c.refresh.ctx |
| 212 | closed := c.refresh.closed |
| 213 | c.refresh.mu.Unlock() |
| 214 | if closed { |
| 215 | return |
| 216 | } |
| 217 | var stops []func() |
| 218 | if c.capabilities.promptsListChanged { |
| 219 | stops = append(stops, t.registerNotification("notifications/prompts/list_changed", func(json.RawMessage) { |
| 220 | h.requestPromptRefresh(ctx, c) |
| 221 | })) |
| 222 | } |
| 223 | if c.capabilities.resourcesListChanged { |
| 224 | stops = append(stops, t.registerNotification("notifications/resources/list_changed", func(json.RawMessage) { |
| 225 | h.requestResourceRefresh(ctx, c) |
| 226 | })) |
| 227 | } |
| 228 | c.surfaceStopsMu.Lock() |
| 229 | if c.closed.Load() { |
| 230 | c.surfaceStopsMu.Unlock() |
| 231 | for _, stop := range stops { |
| 232 | stop() |
| 233 | } |
| 234 | return |
| 235 | } |
| 236 | c.surfaceStops = append(c.surfaceStops, stops...) |
| 237 | c.surfaceStopsMu.Unlock() |
| 238 | } |
| 239 | |
| 240 | func (h *Host) requestPromptRefresh(ctx context.Context, c *Client) { |
| 241 | c.auxiliaryRefresh.prompt.revision.Add(1) |
| 242 | h.startPromptRefresh(ctx, c) |
| 243 | } |
| 244 | |
| 245 | func (h *Host) startPromptRefresh(ctx context.Context, c *Client) { |
| 246 | if !c.auxiliaryRefresh.prompt.running.CompareAndSwap(false, true) { |
| 247 | return |
| 248 | } |
| 249 | if !h.goSurface(func() { |
| 250 | defer func() { |
| 251 | c.auxiliaryRefresh.prompt.running.Store(false) |
| 252 | if ctx.Err() == nil && !c.closed.Load() && c.auxiliaryRefresh.prompt.applied.Load() < c.auxiliaryRefresh.prompt.revision.Load() { |
| 253 | h.startPromptRefresh(ctx, c) |
| 254 | } |
| 255 | }() |
| 256 | for ctx.Err() == nil && !c.closed.Load() { |
| 257 | target := c.auxiliaryRefresh.prompt.revision.Load() |
| 258 | h.fetchPrompts(ctx, c, nil) |
| 259 | c.auxiliaryRefresh.prompt.applied.Store(target) |
| 260 | if c.auxiliaryRefresh.prompt.revision.Load() == target { |
| 261 | return |
| 262 | } |
| 263 | } |
| 264 | }) { |
| 265 | c.auxiliaryRefresh.prompt.running.Store(false) |
| 266 | } |
| 267 | } |
| 268 | |
| 269 | func (h *Host) requestResourceRefresh(ctx context.Context, c *Client) { |
| 270 | c.auxiliaryRefresh.resource.revision.Add(1) |
| 271 | h.startResourceRefresh(ctx, c) |
| 272 | } |
| 273 | |
| 274 | func (h *Host) startResourceRefresh(ctx context.Context, c *Client) { |
| 275 | if !c.auxiliaryRefresh.resource.running.CompareAndSwap(false, true) { |
| 276 | return |
| 277 | } |
| 278 | if !h.goSurface(func() { |
| 279 | defer func() { |
| 280 | c.auxiliaryRefresh.resource.running.Store(false) |
| 281 | if ctx.Err() == nil && !c.closed.Load() && c.auxiliaryRefresh.resource.applied.Load() < c.auxiliaryRefresh.resource.revision.Load() { |
| 282 | h.startResourceRefresh(ctx, c) |
| 283 | } |
| 284 | }() |
| 285 | for ctx.Err() == nil && !c.closed.Load() { |
| 286 | target := c.auxiliaryRefresh.resource.revision.Load() |
| 287 | h.fetchResources(ctx, c, nil) |
| 288 | c.auxiliaryRefresh.resource.applied.Store(target) |
| 289 | if c.auxiliaryRefresh.resource.revision.Load() == target { |
| 290 | return |
| 291 | } |
| 292 | } |
| 293 | }) { |
| 294 | c.auxiliaryRefresh.resource.running.Store(false) |
| 295 | } |
| 296 | } |
| 297 | |
| 298 | func (h *Host) registerStartedClient(c *Client, tools []tool.Tool) ([]tool.Tool, error) { |
| 299 | h.mu.Lock() |
| 300 | err := h.noteClientLocked(c, nil) |
| 301 | h.mu.Unlock() |
| 302 | if err != nil { |
| 303 | return nil, err |
| 304 | } |
| 305 | if cached, ok := c.cachedTools(); ok { |
| 306 | return cached, nil |
| 307 | } |
| 308 | return tools, nil |
| 309 | } |
| 310 | |
| 311 | func (h *Host) publishToolListChange(c *Client, tools []tool.Tool) { |
| 312 | if h == nil || c == nil { |
| 313 | return |
| 314 | } |
| 315 | h.mu.RLock() |
| 316 | if h.closed || h.lookupClientLocked(c.name) != c { |
| 317 | h.mu.RUnlock() |
| 318 | return |
| 319 | } |
| 320 | spec := c.spec |
| 321 | subscribers := make([]*toolListSubscriber, 0, len(h.toolListChanges.subscribers)) |
| 322 | for _, subscriber := range h.toolListChanges.subscribers { |
| 323 | subscribers = append(subscribers, subscriber) |
| 324 | } |
| 325 | h.mu.RUnlock() |
| 326 | for _, subscriber := range subscribers { |
| 327 | deliverToolListChange(subscriber, spec, tools) |
| 328 | } |
| 329 | } |
| 330 | |
| 331 | func deliverToolListChange(subscriber *toolListSubscriber, spec Spec, tools []tool.Tool) { |
| 332 | if subscriber == nil { |
| 333 | return |
| 334 | } |
| 335 | subscriber.deliveryMu.Lock() |
| 336 | defer subscriber.deliveryMu.Unlock() |
| 337 | if subscriber.ctx.Err() != nil { |
| 338 | return |
| 339 | } |
| 340 | subscriber.callback(spec, append([]tool.Tool(nil), tools...)) |
| 341 | } |
| 342 | |
| 343 | func (c *Client) watchToolListChanges() { |
| 344 | if !c.capabilities.toolsListChanged { |
| 345 | return |
| 346 | } |
| 347 | t, ok := c.t.(notificationTransport) |
| 348 | if !ok { |
| 349 | return |
| 350 | } |
| 351 | c.refresh.mu.Lock() |
| 352 | defer c.refresh.mu.Unlock() |
| 353 | if c.refresh.closed || c.refresh.stopNotifications != nil { |
| 354 | return |
| 355 | } |
| 356 | c.refresh.stopNotifications = t.registerNotification("notifications/tools/list_changed", func(json.RawMessage) { |
| 357 | c.requestToolsRefresh() |
| 358 | }) |
| 359 | } |
| 360 | |
| 361 | func (c *Client) requestToolsRefresh() { |
| 362 | c.refresh.mu.Lock() |
| 363 | if c.refresh.closed { |
| 364 | c.refresh.mu.Unlock() |
| 365 | return |
| 366 | } |
| 367 | c.refresh.noticeRevision.Add(1) |
| 368 | if c.refresh.running { |
| 369 | c.refresh.mu.Unlock() |
| 370 | return |
| 371 | } |
| 372 | c.refresh.running = true |
| 373 | c.refresh.cycleDone = make(chan struct{}) |
| 374 | c.refresh.mu.Unlock() |
| 375 | go c.runToolsRefreshes() |
| 376 | } |
| 377 | |
| 378 | // ensureToolsRefresh lets a user retry recover a dirty catalog after a bounded |
| 379 | // refresh cycle failed or exhausted its catch-up attempts. It never |
| 380 | // invents a new notice revision and therefore cannot create a self-sustaining |
| 381 | // loop without a real notification or call attempt. |
| 382 | func (c *Client) ensureToolsRefresh() { |
| 383 | if c == nil || !c.toolCatalogStale() { |
| 384 | return |
| 385 | } |
| 386 | c.refresh.mu.Lock() |
| 387 | if c.refresh.closed || c.refresh.running { |
| 388 | c.refresh.mu.Unlock() |
| 389 | return |
| 390 | } |
| 391 | c.refresh.running = true |
| 392 | c.refresh.cycleDone = make(chan struct{}) |
| 393 | c.refresh.mu.Unlock() |
| 394 | go c.runToolsRefreshes() |
| 395 | } |
| 396 | |
| 397 | func (c *Client) runToolsRefreshes() { |
| 398 | c.refresh.mu.Lock() |
| 399 | ctx := c.refresh.ctx |
| 400 | wait := c.refresh.wait |
| 401 | c.refresh.mu.Unlock() |
| 402 | if ctx == nil { |
| 403 | ctx = context.Background() |
| 404 | } |
| 405 | if wait == nil { |
| 406 | wait = sleepContext |
| 407 | } |
| 408 | finished := false |
| 409 | defer func() { |
| 410 | if !finished { |
| 411 | c.finishToolsRefreshCycle() |
| 412 | } |
| 413 | }() |
| 414 | |
| 415 | // Bound catch-up because some servers self-notify on every tools/list. |
| 416 | // Exhaustion stays stale and fail-closed; a later notice or user retry starts |
| 417 | // a fresh bounded cycle. |
| 418 | refreshDelay := toolListRefreshDebounce |
| 419 | for range toolListRefreshMaxAttempts { |
| 420 | if err := wait(ctx, refreshDelay); err != nil { |
| 421 | return |
| 422 | } |
| 423 | c.refresh.mu.Lock() |
| 424 | closed := c.refresh.closed |
| 425 | c.refresh.mu.Unlock() |
| 426 | if closed { |
| 427 | return |
| 428 | } |
| 429 | targetRevision := c.refresh.noticeRevision.Load() |
| 430 | refreshCtx, cancel := context.WithTimeout(ctx, c.toolListRefreshTimeout()) |
| 431 | tools, changed, err := c.refreshTools(refreshCtx, targetRevision) |
| 432 | cancel() |
| 433 | if err != nil { |
| 434 | if ctx.Err() == nil && !c.closed.Load() { |
| 435 | slog.Warn("plugin: refresh tools after list_changed failed", "server", c.name, "err", err) |
| 436 | } |
| 437 | return |
| 438 | } |
| 439 | if changed { |
| 440 | c.refresh.mu.Lock() |
| 441 | callback := c.refresh.onChanged |
| 442 | closed = c.refresh.closed |
| 443 | c.refresh.mu.Unlock() |
| 444 | if !closed && callback != nil { |
| 445 | callback(append([]tool.Tool(nil), tools...)) |
| 446 | } |
| 447 | } |
| 448 | c.refresh.mu.Lock() |
| 449 | if c.refresh.noticeRevision.Load() == targetRevision { |
| 450 | c.finishToolsRefreshCycleLocked() |
| 451 | finished = true |
| 452 | c.refresh.mu.Unlock() |
| 453 | return |
| 454 | } |
| 455 | c.refresh.mu.Unlock() |
| 456 | refreshDelay = nextToolListRefreshDelay(refreshDelay) |
| 457 | } |
| 458 | } |
| 459 | |
| 460 | func nextToolListRefreshDelay(current time.Duration) time.Duration { |
| 461 | if current <= 0 { |
| 462 | return toolListRefreshDebounce |
| 463 | } |
| 464 | if current >= toolListRefreshMaxBackoff/2 { |
| 465 | return toolListRefreshMaxBackoff |
| 466 | } |
| 467 | return current * 2 |
| 468 | } |
| 469 | |
| 470 | func (c *Client) toolListRefreshTimeout() time.Duration { |
| 471 | startupTimeout := c.spec.ResolvedStartupTimeout() |
| 472 | callTimeout := c.callTimeout("tools/list", map[string]any{}) |
| 473 | if startupTimeout <= 0 { |
| 474 | return callTimeout |
| 475 | } |
| 476 | if callTimeout <= 0 || startupTimeout < callTimeout { |
| 477 | return startupTimeout |
| 478 | } |
| 479 | return callTimeout |
| 480 | } |
| 481 | |
| 482 | func (c *Client) finishToolsRefreshCycle() { |
| 483 | c.refresh.mu.Lock() |
| 484 | c.finishToolsRefreshCycleLocked() |
| 485 | c.refresh.mu.Unlock() |
| 486 | } |
| 487 | |
| 488 | func (c *Client) finishToolsRefreshCycleLocked() { |
| 489 | c.refresh.running = false |
| 490 | if c.refresh.cycleDone != nil { |
| 491 | close(c.refresh.cycleDone) |
| 492 | c.refresh.cycleDone = nil |
| 493 | } |
| 494 | } |
| 495 | |
| 496 | func (c *Client) setToolsChangedCallback(callback func([]tool.Tool)) { |
| 497 | c.refresh.mu.Lock() |
| 498 | if c.refresh.closed { |
| 499 | c.refresh.mu.Unlock() |
| 500 | return |
| 501 | } |
| 502 | c.refresh.onChanged = callback |
| 503 | c.refresh.mu.Unlock() |
| 504 | } |
| 505 | |
| 506 | func (c *Client) listTools(ctx context.Context) ([]tool.Tool, error) { |
| 507 | if tools, ok := c.cachedTools(); ok { |
| 508 | return tools, nil |
| 509 | } |
| 510 | c.toolListFetchMu.Lock() |
| 511 | defer c.toolListFetchMu.Unlock() |
| 512 | if tools, ok := c.cachedTools(); ok { |
| 513 | return tools, nil |
| 514 | } |
| 515 | targetRevision := c.refresh.noticeRevision.Load() |
| 516 | candidate, err := c.fetchToolCatalog(ctx, true) |
| 517 | if err != nil { |
| 518 | return nil, err |
| 519 | } |
| 520 | tools, _, err := c.publishToolCatalog(candidate, targetRevision) |
| 521 | return tools, err |
| 522 | } |
| 523 | |
| 524 | func (c *Client) refreshTools(ctx context.Context, targetRevision uint64) ([]tool.Tool, bool, error) { |
| 525 | c.toolListFetchMu.Lock() |
| 526 | defer c.toolListFetchMu.Unlock() |
| 527 | candidate, err := c.fetchToolCatalog(ctx, false) |
| 528 | if err != nil { |
| 529 | return nil, false, err |
| 530 | } |
| 531 | return c.publishToolCatalog(candidate, targetRevision) |
| 532 | } |
| 533 | |
| 534 | // fetchToolCatalog performs MCP I/O and constructs a complete candidate off |
| 535 | // lock. Startup keeps the existing empty-list settling window; notifications |
| 536 | // use one tools/list per attempt so retries remain bounded by the scheduler. |
| 537 | func (c *Client) fetchToolCatalog(ctx context.Context, settleEmpty bool) (toolCatalogSnapshot, error) { |
| 538 | var ( |
| 539 | out []mcpTool |
| 540 | err error |
| 541 | ) |
| 542 | if settleEmpty { |
| 543 | out, err = c.listToolsRawSettled(ctx) |
| 544 | } else { |
| 545 | out, err = c.listToolsRaw(ctx) |
| 546 | } |
| 547 | if err != nil { |
| 548 | return toolCatalogSnapshot{}, err |
| 549 | } |
| 550 | if err := validateMCPToolNames(out); err != nil { |
| 551 | return toolCatalogSnapshot{}, fmt.Errorf("plugin %q: %w", c.name, err) |
| 552 | } |
| 553 | |
| 554 | toolInfos := make([]ToolInfo, 0, len(out)) |
| 555 | tools := make([]tool.Tool, 0, len(out)) |
| 556 | appTools := make([]tool.Tool, 0, len(out)) |
| 557 | fingerprintEntries := make([]toolCatalogFingerprintEntry, 0, len(out)) |
| 558 | normalizedSchemas := make(map[string]json.RawMessage, len(out)) |
| 559 | for _, candidate := range out { |
| 560 | schema, err := normalizeAndValidateToolSchema(candidate.InputSchema) |
| 561 | if err == nil { |
| 562 | normalizedSchemas[candidate.Name] = schema |
| 563 | } |
| 564 | } |
| 565 | for _, candidate := range out { |
| 566 | readOnlyHint := candidate.Annotations != nil && candidate.Annotations.ReadOnlyHint |
| 567 | destructiveHint := candidate.Annotations != nil && candidate.Annotations.DestructiveHint |
| 568 | modelVisible, appVisible := candidate.Meta.appVisibility() |
| 569 | appCallable := appVisible && c.appsNegotiated() |
| 570 | info := ToolInfo{Name: candidate.Name, Description: candidate.Description, ReadOnlyHint: readOnlyHint, DestructiveHint: destructiveHint} |
| 571 | visibleName := candidate.Name |
| 572 | if c.spec.StripRawPrefix != "" { |
| 573 | visibleName = strings.TrimPrefix(visibleName, c.spec.StripRawPrefix) |
| 574 | } |
| 575 | schema, ok := normalizedSchemas[candidate.Name] |
| 576 | if !ok { |
| 577 | if _, err := normalizeAndValidateToolSchema(candidate.InputSchema); err != nil { |
| 578 | info.SchemaError = schemaValidationError(err) |
| 579 | } |
| 580 | if modelVisible { |
| 581 | toolInfos = append(toolInfos, info) |
| 582 | } |
| 583 | fingerprintEntries = append(fingerprintEntries, toolCatalogFingerprintEntry{ |
| 584 | RawName: candidate.Name, VisibleName: visibleName, Description: candidate.Description, |
| 585 | SchemaError: info.SchemaError, ReadOnlyHint: readOnlyHint, DestructiveHint: destructiveHint, |
| 586 | Visibility: candidate.Meta.visibilityCopy(), |
| 587 | }) |
| 588 | continue |
| 589 | } |
| 590 | if modelVisible { |
| 591 | toolInfos = append(toolInfos, info) |
| 592 | } |
| 593 | outputSchema := append(json.RawMessage(nil), candidate.OutputSchema...) |
| 594 | fingerprintOutputSchema := canonicalizeCatalogJSON(outputSchema) |
| 595 | uiResourceURI, uiCSP := candidate.Meta.uiResource() |
| 596 | adapter := &remoteTool{ |
| 597 | client: c, |
| 598 | name: toolName(c.name, visibleName), |
| 599 | rawName: candidate.Name, |
| 600 | visibleName: visibleName, |
| 601 | desc: candidate.Description, |
| 602 | schema: schema, |
| 603 | outputSchema: outputSchema, |
| 604 | declaredReadOnly: readOnlyHint, |
| 605 | readOnly: readOnlyHint, |
| 606 | destructive: destructiveHint, |
| 607 | visibility: candidate.Meta.visibilityCopy(), |
| 608 | appCallable: appCallable, |
| 609 | uiResourceURI: uiResourceURI, |
| 610 | uiCSP: uiCSP, |
| 611 | } |
| 612 | if modelVisible { |
| 613 | tools = append(tools, adapter) |
| 614 | } |
| 615 | if appCallable { |
| 616 | appTools = append(appTools, adapter) |
| 617 | } |
| 618 | fingerprintEntries = append(fingerprintEntries, toolCatalogFingerprintEntry{ |
| 619 | RawName: candidate.Name, VisibleName: visibleName, Description: candidate.Description, |
| 620 | Schema: schema, OutputSchema: fingerprintOutputSchema, ReadOnlyHint: readOnlyHint, DestructiveHint: destructiveHint, |
| 621 | Visibility: candidate.Meta.visibilityCopy(), |
| 622 | }) |
| 623 | } |
| 624 | sort.SliceStable(toolInfos, func(i, j int) bool { return toolInfos[i].Name < toolInfos[j].Name }) |
| 625 | sort.SliceStable(fingerprintEntries, func(i, j int) bool { return fingerprintEntries[i].RawName < fingerprintEntries[j].RawName }) |
| 626 | sortedTools := sortToolsByName(tools) |
| 627 | fingerprintJSON, err := json.Marshal(fingerprintEntries) |
| 628 | if err != nil { |
| 629 | return toolCatalogSnapshot{}, fmt.Errorf("plugin %q: fingerprint tools/list: %w", c.name, err) |
| 630 | } |
| 631 | return toolCatalogSnapshot{ |
| 632 | listed: true, fingerprint: sha256.Sum256(fingerprintJSON), |
| 633 | infos: toolInfos, adapters: sortedTools, appAdapters: sortToolsByName(appTools), |
| 634 | }, nil |
| 635 | } |
| 636 | |
| 637 | func canonicalizeCatalogJSON(raw json.RawMessage) json.RawMessage { |
| 638 | if len(raw) == 0 { |
| 639 | return nil |
| 640 | } |
| 641 | var value any |
| 642 | if err := json.Unmarshal(raw, &value); err != nil { |
| 643 | return append(json.RawMessage(nil), raw...) |
| 644 | } |
| 645 | canonical, err := json.Marshal(value) |
| 646 | if err != nil { |
| 647 | return append(json.RawMessage(nil), raw...) |
| 648 | } |
| 649 | return canonical |
| 650 | } |
| 651 | |
| 652 | // publishToolCatalog waits for already-authorized calls, then atomically makes |
| 653 | // a complete candidate current. publishedRevision is advanced under the same |
| 654 | // dispatch gate, making stale checks linearizable with call admission. |
| 655 | func (c *Client) publishToolCatalog(candidate toolCatalogSnapshot, targetRevision uint64) ([]tool.Tool, bool, error) { |
| 656 | c.toolDispatchMu.Lock() |
| 657 | defer c.toolDispatchMu.Unlock() |
| 658 | if c.closed.Load() { |
| 659 | return nil, false, fmt.Errorf("MCP server %q is closed", c.name) |
| 660 | } |
| 661 | |
| 662 | c.toolsMu.Lock() |
| 663 | changed := !c.toolCatalog.listed || c.toolCatalog.fingerprint != candidate.fingerprint |
| 664 | if changed { |
| 665 | oldFingerprints := catalogSchemaFingerprints(c.toolCatalog.adapters) |
| 666 | c.catalogGeneration++ |
| 667 | candidate.generation = c.catalogGeneration |
| 668 | for _, adapters := range [][]tool.Tool{candidate.adapters, candidate.appAdapters} { |
| 669 | for _, adapter := range adapters { |
| 670 | if remote, ok := adapter.(*remoteTool); ok { |
| 671 | remote.generation = candidate.generation |
| 672 | } |
| 673 | } |
| 674 | } |
| 675 | c.toolCatalog = candidate |
| 676 | tool.InvalidateArgumentSchemas(oldFingerprints) |
| 677 | } |
| 678 | tools := append([]tool.Tool(nil), c.toolCatalog.adapters...) |
| 679 | c.toolsMu.Unlock() |
| 680 | c.advancePublishedToolListRevision(targetRevision) |
| 681 | return tools, changed, nil |
| 682 | } |
| 683 | |
| 684 | func catalogSchemaFingerprints(adapters []tool.Tool) []string { |
| 685 | out := make([]string, 0, len(adapters)) |
| 686 | for _, adapter := range adapters { |
| 687 | if adapter == nil { |
| 688 | continue |
| 689 | } |
| 690 | out = append(out, tool.SchemaFingerprint(adapter.Schema())) |
| 691 | } |
| 692 | return out |
| 693 | } |
| 694 | |
| 695 | func (c *Client) advancePublishedToolListRevision(revision uint64) { |
| 696 | for { |
| 697 | current := c.refresh.publishedRevision.Load() |
| 698 | if revision <= current || c.refresh.publishedRevision.CompareAndSwap(current, revision) { |
| 699 | return |
| 700 | } |
| 701 | } |
| 702 | } |
| 703 | |
| 704 | func (c *Client) toolCatalogStale() bool { |
| 705 | return c != nil && c.refresh.publishedRevision.Load() < c.refresh.noticeRevision.Load() |
| 706 | } |
| 707 | |
| 708 | func normalizeAndValidateToolSchema(raw json.RawMessage) (json.RawMessage, error) { |
| 709 | schema := canonicalizeSchema(raw) |
| 710 | if err := provider.ValidateToolSchema(schema); err != nil { |
| 711 | return nil, err |
| 712 | } |
| 713 | return schema, nil |
| 714 | } |
| 715 | |
| 716 | func schemaValidationError(err error) string { |
| 717 | const maxRunes = 512 |
| 718 | msg := strings.TrimSpace(err.Error()) |
| 719 | runes := []rune(msg) |
| 720 | if len(runes) > maxRunes { |
| 721 | msg = string(runes[:maxRunes]) + "..." |
| 722 | } |
| 723 | return "invalid input schema: " + msg |
| 724 | } |
| 725 | |
| 726 | func (c *Client) listToolsRaw(ctx context.Context) ([]mcpTool, error) { |
| 727 | res, err := c.call(ctx, "tools/list", map[string]any{}) |
| 728 | if err != nil { |
| 729 | return nil, err |
| 730 | } |
| 731 | var out struct { |
| 732 | Tools []mcpTool `json:"tools"` |
| 733 | } |
| 734 | if err := json.Unmarshal(res, &out); err != nil { |
| 735 | return nil, fmt.Errorf("plugin %q: decode tools/list: %w", c.name, err) |
| 736 | } |
| 737 | return out.Tools, nil |
| 738 | } |
| 739 | |
| 740 | // listToolsRawSettled gives dynamically registering servers a bounded startup |
| 741 | // window before their initial tool catalog is considered complete. |
| 742 | func (c *Client) listToolsRawSettled(ctx context.Context) ([]mcpTool, error) { |
| 743 | out, err := c.listToolsRaw(ctx) |
| 744 | if err != nil || !c.capabilities.tools || len(out) > 0 { |
| 745 | return out, err |
| 746 | } |
| 747 | for _, delay := range advertisedToolsEmptyListRetryDelays { |
| 748 | if err := sleepContext(ctx, delay); err != nil { |
| 749 | return nil, err |
| 750 | } |
| 751 | out, err = c.listToolsRaw(ctx) |
| 752 | if err != nil || len(out) > 0 { |
| 753 | return out, err |
| 754 | } |
| 755 | } |
| 756 | return out, nil |
| 757 | } |
| 758 | |
| 759 | func validateMCPToolNames(tools []mcpTool) error { |
| 760 | seen := make(map[string]bool, len(tools)) |
| 761 | for _, candidate := range tools { |
| 762 | name := strings.TrimSpace(candidate.Name) |
| 763 | if name == "" { |
| 764 | return fmt.Errorf("tools/list returned an empty tool name") |
| 765 | } |
| 766 | if seen[candidate.Name] { |
| 767 | return fmt.Errorf("tools/list returned duplicate tool name %q", candidate.Name) |
| 768 | } |
| 769 | seen[candidate.Name] = true |
| 770 | } |
| 771 | return nil |
| 772 | } |
| 773 | |
| 774 | func sleepContext(ctx context.Context, delay time.Duration) error { |
| 775 | if delay <= 0 { |
| 776 | return nil |
| 777 | } |
| 778 | timer := time.NewTimer(delay) |
| 779 | defer timer.Stop() |
| 780 | select { |
| 781 | case <-ctx.Done(): |
| 782 | return ctx.Err() |
| 783 | case <-timer.C: |
| 784 | return nil |
| 785 | } |
| 786 | } |
| 787 | |
| 788 | func (c *Client) cachedTools() ([]tool.Tool, bool) { |
| 789 | c.toolsMu.RLock() |
| 790 | defer c.toolsMu.RUnlock() |
| 791 | if !c.toolCatalog.listed { |
| 792 | return nil, false |
| 793 | } |
| 794 | return append([]tool.Tool(nil), c.toolCatalog.adapters...), true |
| 795 | } |
| 796 |