| 1 | // MCP placeholder tools. Background startup registers cheap placeholder entries |
| 2 | // in the tool registry at boot — using the on-disk schema cache when it exists — |
| 3 | // and kicks the real subprocess spawn / handshake immediately. By the time the |
| 4 | // model calls a tool, the connection is usually already up. |
| 5 | // |
| 6 | // Cache-hit placeholders are PINNED for the whole session: they present the |
| 7 | // cached names/descriptions/schemas from boot onward and forward Execute to the |
| 8 | // real tools once the handshake completes, but the registry entries themselves |
| 9 | // are never replaced. The provider request's tools array is part of the cached |
| 10 | // prompt prefix, so swapping in live tools mid-session — whenever the live |
| 11 | // handshake differed from the cache — invalidated the whole conversation's |
| 12 | // provider cache at 10x miss pricing. Live drift lands in the schema cache and |
| 13 | // surfaces next session. Only the cache-miss connect stub still swaps (there |
| 14 | // was nothing real to present), a one-time cost per server. |
| 15 | package plugin |
| 16 | |
| 17 | import ( |
| 18 | "context" |
| 19 | "encoding/json" |
| 20 | "errors" |
| 21 | "fmt" |
| 22 | "strings" |
| 23 | "sync" |
| 24 | "time" |
| 25 | |
| 26 | "reasonix/internal/tool" |
| 27 | ) |
| 28 | |
| 29 | // DefaultStartupBudget is the per-plugin latency budget used by boot when |
| 30 | // deciding whether to auto-demote (see Recommend). Kept here rather than in |
| 31 | // stats.go because it's the value boot.go pairs with each Recommend call. |
| 32 | func DefaultStartupBudget() time.Duration { return defaultStartTimeout } |
| 33 | |
| 34 | // spawnState is the lazy-spawn state machine. Transitions are: |
| 35 | // |
| 36 | // idle → inFlight → ready |
| 37 | // idle → inFlight → failed |
| 38 | // |
| 39 | // All transitions are gated by lazySpawn.mu so only one goroutine runs the |
| 40 | // handshake even when multiple Execute calls race on first use. |
| 41 | type spawnState int |
| 42 | |
| 43 | const ( |
| 44 | spawnIdle spawnState = iota |
| 45 | spawnInFlight |
| 46 | spawnReady |
| 47 | spawnFailed |
| 48 | ) |
| 49 | |
| 50 | // lazySpawn is shared by every placeholder lazyTool registered for one |
| 51 | // server: they all observe the same state machine and trigger at most one |
| 52 | // handshake. |
| 53 | type lazySpawn struct { |
| 54 | spec Spec |
| 55 | host *Host |
| 56 | reg *tool.Registry |
| 57 | ctx context.Context // session-scoped — outlives any single turn |
| 58 | generation uint64 |
| 59 | |
| 60 | mu sync.Mutex |
| 61 | state spawnState |
| 62 | real map[string]tool.Tool // namespaced name → real tool, populated on success |
| 63 | spawnErr error |
| 64 | swapped bool |
| 65 | // waitBudget bounds how long one tool call waits for a shared startup. The |
| 66 | // handshake itself uses spec.startupTimeout and continues in the background. |
| 67 | waitBudget time.Duration |
| 68 | // ready is closed when state leaves spawnInFlight so concurrent waiters can |
| 69 | // observe the result without killing a shared host process. |
| 70 | ready chan struct{} |
| 71 | // removePrefix is set for cache-miss placeholders so trySwap drops the |
| 72 | // single "<server>__connect" stub before re-registering the real tools |
| 73 | // under their actual namespaced names. Cache-hit placeholders use the |
| 74 | // same names as the real tools, so reg.Add overwrites in place and no |
| 75 | // prefix removal is needed. |
| 76 | removePrefix string |
| 77 | } |
| 78 | |
| 79 | // beginInFlight transitions idle → inFlight and creates the waiter channel. |
| 80 | // Caller must hold s.mu. Returns false when the host is closed. |
| 81 | func (s *lazySpawn) beginInFlight() bool { |
| 82 | if !s.host.beginDeferredSpawn() { |
| 83 | s.state = spawnFailed |
| 84 | s.spawnErr = fmt.Errorf("plugin host is closed") |
| 85 | s.broadcastReady() |
| 86 | return false |
| 87 | } |
| 88 | s.state = spawnInFlight |
| 89 | s.ready = make(chan struct{}) |
| 90 | return true |
| 91 | } |
| 92 | |
| 93 | // broadcastReady closes the current ready channel if any. Caller holds s.mu. |
| 94 | func (s *lazySpawn) broadcastReady() { |
| 95 | if s.ready != nil { |
| 96 | close(s.ready) |
| 97 | s.ready = nil |
| 98 | } |
| 99 | } |
| 100 | |
| 101 | // kick starts the spawn if it has not yet started. Cache-miss catalog discovery |
| 102 | // and tests may call this; cache-hit boot registration uses kick=false so the |
| 103 | // process starts on first real tool call via EnsureConnected. |
| 104 | func (s *lazySpawn) kick() { |
| 105 | s.mu.Lock() |
| 106 | defer s.mu.Unlock() |
| 107 | if s.state != spawnIdle { |
| 108 | return |
| 109 | } |
| 110 | if !s.beginInFlight() { |
| 111 | return |
| 112 | } |
| 113 | go func() { |
| 114 | defer s.host.endDeferredSpawn() |
| 115 | s.run() |
| 116 | }() |
| 117 | } |
| 118 | |
| 119 | // run does the handshake without holding mu (host.EnsureConnected can take |
| 120 | // seconds), then reacquires mu to publish the result. |
| 121 | func (s *lazySpawn) run() { |
| 122 | started := time.Now() |
| 123 | startupCtx, cancel := context.WithTimeout(s.ctx, s.spec.startupTimeout()) |
| 124 | real, err := s.host.EnsureConnectedWithLifecycle(s.ctx, startupCtx, s.spec, s.generation) |
| 125 | cancel() |
| 126 | if err != nil { |
| 127 | err = newStartupFailure("connect", started, "", err) |
| 128 | } |
| 129 | var cacheTools []tool.Tool |
| 130 | s.mu.Lock() |
| 131 | defer s.mu.Unlock() |
| 132 | if err != nil { |
| 133 | if errors.Is(err, ErrDeferredSpawnCancelled) || errors.Is(err, context.Canceled) { |
| 134 | s.state = spawnFailed |
| 135 | s.spawnErr = err |
| 136 | s.broadcastReady() |
| 137 | return |
| 138 | } |
| 139 | s.state = spawnFailed |
| 140 | s.spawnErr = err |
| 141 | s.host.RecordFailure(s.spec, err) |
| 142 | s.broadcastReady() |
| 143 | return |
| 144 | } |
| 145 | s.real = make(map[string]tool.Tool, len(real)) |
| 146 | for _, t := range real { |
| 147 | s.real[t.Name()] = t |
| 148 | } |
| 149 | s.state = spawnReady |
| 150 | s.trySwap() |
| 151 | cacheTools = real |
| 152 | s.broadcastReady() |
| 153 | // Save cache outside the critical path of tool dispatch. Register the write |
| 154 | // before this deferred spawn ends so Host.Close observes and drains it. |
| 155 | s.host.queueBackgroundWrite(func() { |
| 156 | saveLazyCachedSchema(s.spec, cacheTools) |
| 157 | }) |
| 158 | } |
| 159 | |
| 160 | func saveLazyCachedSchema(spec Spec, real []tool.Tool) { |
| 161 | _ = SaveCachedSchema(spec.Name, CachedSchema{ |
| 162 | CacheKey: SchemaCacheKey(spec), |
| 163 | Capabilities: map[string]bool{"tools": len(real) > 0}, |
| 164 | Tools: cacheableToolsOf(real), |
| 165 | }) |
| 166 | } |
| 167 | |
| 168 | // trySwap publishes the real tools after a successful spawn. Caller must hold |
| 169 | // s.mu. |
| 170 | // |
| 171 | // Cache-miss placeholders (removePrefix set) genuinely swap: the single |
| 172 | // "<server>__connect" stub is dropped and the real tools register under their |
| 173 | // own names — a one-time tool-set change per server, unavoidable because no |
| 174 | // schema existed to present earlier. |
| 175 | // |
| 176 | // Cache-hit placeholders do NOT touch the registry. The lazyTools already |
| 177 | // carry the cached names/descriptions/schemas the model has seen since boot, |
| 178 | // and Execute forwards to the real tool once ready — swapping in the live |
| 179 | // tools would rewrite the request's tools array mid-session whenever the live |
| 180 | // handshake differs from the cache (description tweaks, schema upgrades, new |
| 181 | // tools), invalidating the provider prefix cache at 10x miss pricing. The |
| 182 | // live result still lands in the schema cache (saveLazyCachedSchema), so the |
| 183 | // NEXT session presents the updated surface — freshness deferred one session |
| 184 | // in exchange for byte-stable tool bytes within this one, same trade the |
| 185 | // environment-probe snapshot makes for the system prompt. |
| 186 | func (s *lazySpawn) trySwap() { |
| 187 | if s.swapped || s.state != spawnReady { |
| 188 | return |
| 189 | } |
| 190 | if s.removePrefix != "" { |
| 191 | s.reg.RemovePrefix(s.removePrefix) |
| 192 | for _, t := range s.real { |
| 193 | s.reg.Add(t) |
| 194 | } |
| 195 | } |
| 196 | s.swapped = true |
| 197 | } |
| 198 | |
| 199 | // lazyTool is a tool.Tool placeholder backed by a shared lazySpawn. The model |
| 200 | // sees cached metadata (or a stub when no cache exists); Execute consults the |
| 201 | // state machine, kicking off the handshake on first call. |
| 202 | type lazyTool struct { |
| 203 | shared *lazySpawn |
| 204 | name string // namespaced "mcp__<server>__<tool>" |
| 205 | rawName string // original server-local tool name, when cached |
| 206 | visibleName string // raw name after configured prefix stripping |
| 207 | desc string |
| 208 | schema json.RawMessage |
| 209 | // readOnly is guarded by shared.mu because a live handshake can demote a |
| 210 | // stale cached reader before asking the model to retry. |
| 211 | readOnly bool |
| 212 | // destructive is guarded by shared.mu because a live handshake may promote |
| 213 | // a stale cached false value before asking the model to retry. |
| 214 | destructive bool |
| 215 | // hasCache true → schema is trusted, so Execute starts the shared handshake, |
| 216 | // waits briefly, and forwards in the same turn when ready. false → schema is empty, so we |
| 217 | // can't honour the model's call; we kick the spawn async and ask for a |
| 218 | // retry on the next turn, when the swap will have installed the real |
| 219 | // tools with real schemas. |
| 220 | hasCache bool |
| 221 | } |
| 222 | |
| 223 | func (lt *lazyTool) Name() string { return lt.name } |
| 224 | func (lt *lazyTool) Description() string { return lt.desc } |
| 225 | func (lt *lazyTool) ReadOnly() bool { |
| 226 | if lt.shared == nil { |
| 227 | return lt.readOnly |
| 228 | } |
| 229 | lt.shared.mu.Lock() |
| 230 | defer lt.shared.mu.Unlock() |
| 231 | return lt.readOnly |
| 232 | } |
| 233 | func (lt *lazyTool) MCPServerName() string { |
| 234 | if lt.shared == nil { |
| 235 | return "" |
| 236 | } |
| 237 | return lt.shared.spec.Name |
| 238 | } |
| 239 | func (lt *lazyTool) MCPRawToolName() string { return lt.rawName } |
| 240 | func (lt *lazyTool) MCPVisibleToolName() string { return lt.visibleName } |
| 241 | func (lt *lazyTool) MCPPackageName() string { |
| 242 | if lt.shared == nil { |
| 243 | return "" |
| 244 | } |
| 245 | return lt.shared.spec.Package |
| 246 | } |
| 247 | |
| 248 | func (lt *lazyTool) MCPServerAuthorized() bool { |
| 249 | return lt.shared != nil && lt.shared.spec.ServerAuthorized() |
| 250 | } |
| 251 | |
| 252 | func (lt *lazyTool) MCPDestructiveHint() bool { |
| 253 | if lt.shared == nil { |
| 254 | return lt.destructive |
| 255 | } |
| 256 | lt.shared.mu.Lock() |
| 257 | defer lt.shared.mu.Unlock() |
| 258 | return lt.destructive |
| 259 | } |
| 260 | func (lt *lazyTool) Schema() json.RawMessage { |
| 261 | if len(lt.schema) == 0 { |
| 262 | return json.RawMessage(`{"type":"object"}`) |
| 263 | } |
| 264 | return canonicalizeSchema(lt.schema) |
| 265 | } |
| 266 | |
| 267 | func (lt *lazyTool) Execute(ctx context.Context, args json.RawMessage) (string, error) { |
| 268 | sp := lt.shared |
| 269 | for { |
| 270 | sp.mu.Lock() |
| 271 | |
| 272 | // Catch up on a background spawn that finished while we were idle. |
| 273 | if sp.state == spawnReady && !sp.swapped { |
| 274 | sp.trySwap() |
| 275 | } |
| 276 | |
| 277 | switch sp.state { |
| 278 | case spawnReady: |
| 279 | if !lt.hasCache { |
| 280 | sp.mu.Unlock() |
| 281 | return fmt.Sprintf("MCP server %q is connected; its real tools are now available on the next turn", sp.spec.Name), nil |
| 282 | } |
| 283 | real := sp.real[lt.name] |
| 284 | safetyErr := lt.reconcileLiveSafety(real) |
| 285 | sp.mu.Unlock() |
| 286 | if real == nil { |
| 287 | return "", fmt.Errorf("MCP server %q did not expose tool %q (the cached schema may be stale)", sp.spec.Name, lt.name) |
| 288 | } |
| 289 | if safetyErr != nil { |
| 290 | return "", safetyErr |
| 291 | } |
| 292 | return real.Execute(ctx, args) |
| 293 | |
| 294 | case spawnFailed: |
| 295 | err := sp.spawnErr |
| 296 | sp.mu.Unlock() |
| 297 | return "", fmt.Errorf("MCP server %q failed to start: %w", sp.spec.Name, err) |
| 298 | |
| 299 | case spawnInFlight: |
| 300 | // Wait for the in-flight handshake (same process for every waiter). |
| 301 | // Cancelling ctx only abandons this wait; the shared host spawn continues. |
| 302 | wait := sp.ready |
| 303 | sp.mu.Unlock() |
| 304 | if wait == nil { |
| 305 | continue |
| 306 | } |
| 307 | if err := waitForLazyStartup(ctx, sp.ctx, wait, sp.waitBudget, sp.spec.Name, sp.spec.startupTimeout()); err != nil { |
| 308 | return "", err |
| 309 | } |
| 310 | continue |
| 311 | |
| 312 | case spawnIdle: |
| 313 | if !lt.hasCache { |
| 314 | // Cache-miss: we don't trust args to match a real schema, so |
| 315 | // drive the handshake async and ask the model to retry. By the |
| 316 | // next turn the swap will have installed the real tools with |
| 317 | // real schemas under different names. |
| 318 | if !sp.beginInFlight() { |
| 319 | err := sp.spawnErr |
| 320 | sp.mu.Unlock() |
| 321 | return "", fmt.Errorf("MCP server %q failed to start: %w", sp.spec.Name, err) |
| 322 | } |
| 323 | go func() { |
| 324 | defer sp.host.endDeferredSpawn() |
| 325 | sp.run() |
| 326 | }() |
| 327 | sp.mu.Unlock() |
| 328 | return "", fmt.Errorf("MCP server %q is initializing on first use — call again on the next turn for its real tools", sp.spec.Name) |
| 329 | } |
| 330 | // Cache-hit: start the shared handshake in the background. The current |
| 331 | // call waits briefly so healthy fast servers still complete in one turn; |
| 332 | // slow servers keep initializing under the session-owned lifecycle and |
| 333 | // become ready for a later retry instead of being killed and restarted. |
| 334 | if !sp.beginInFlight() { |
| 335 | err := sp.spawnErr |
| 336 | sp.mu.Unlock() |
| 337 | return "", fmt.Errorf("MCP server %q failed to start: %w", sp.spec.Name, err) |
| 338 | } |
| 339 | wait := sp.ready |
| 340 | go func() { |
| 341 | defer sp.host.endDeferredSpawn() |
| 342 | sp.run() |
| 343 | }() |
| 344 | sp.mu.Unlock() |
| 345 | if err := waitForLazyStartup(ctx, sp.ctx, wait, sp.waitBudget, sp.spec.Name, sp.spec.startupTimeout()); err != nil { |
| 346 | return "", err |
| 347 | } |
| 348 | continue |
| 349 | } |
| 350 | |
| 351 | sp.mu.Unlock() |
| 352 | return "", fmt.Errorf("deferred plugin %q in unexpected state", sp.spec.Name) |
| 353 | } |
| 354 | } |
| 355 | |
| 356 | func waitForLazyStartup(ctx, sessionCtx context.Context, ready <-chan struct{}, waitBudget time.Duration, server string, startupLimit time.Duration) error { |
| 357 | if waitBudget <= 0 { |
| 358 | waitBudget = defaultStartTimeout |
| 359 | } |
| 360 | timer := time.NewTimer(waitBudget) |
| 361 | defer timer.Stop() |
| 362 | select { |
| 363 | case <-ready: |
| 364 | return nil |
| 365 | case <-ctx.Done(): |
| 366 | return ctx.Err() |
| 367 | case <-sessionCtx.Done(): |
| 368 | return sessionCtx.Err() |
| 369 | case <-timer.C: |
| 370 | return fmt.Errorf("MCP server %q is still initializing after %s; startup continues in background (limit %s) — retry this tool on a later turn", |
| 371 | server, formatTimeout(waitBudget), formatTimeout(startupLimit)) |
| 372 | } |
| 373 | } |
| 374 | |
| 375 | // reconcileLiveSafety updates a pinned cache-hit placeholder when the live |
| 376 | // server becomes stricter. Caller must hold shared.mu. The current call always |
| 377 | // stops on a reader-to-writer demotion or destructive promotion so the next |
| 378 | // attempt re-enters the agent's Plan/read-only safety checks with current metadata. |
| 379 | func (lt *lazyTool) reconcileLiveSafety(real tool.Tool) error { |
| 380 | if real == nil { |
| 381 | return nil |
| 382 | } |
| 383 | live, err := ReconcileCachedToolSafety(lt.shared.spec.Name, lt.rawName, CachedToolSafety{ |
| 384 | ReadOnly: lt.readOnly, |
| 385 | Destructive: lt.destructive, |
| 386 | }, real) |
| 387 | lt.readOnly = live.ReadOnly |
| 388 | lt.destructive = live.Destructive |
| 389 | return err |
| 390 | } |
| 391 | |
| 392 | // LazyToolset returns the placeholder tools to register for one enabled MCP. |
| 393 | // When cs is non-nil (cache hit) the returned slice has one lazyTool per cached |
| 394 | // tool, carrying the cached schema so the model can pass real args. Execute |
| 395 | // waits briefly for EnsureConnected and completes the call in the same turn |
| 396 | // when startup is fast; slow startup continues in the background. When cs is |
| 397 | // nil (cache miss) the returned slice has a single stub named |
| 398 | // "mcp__<server>__connect": the model can call it to drive the handshake, and |
| 399 | // the real tools surface on the next turn. |
| 400 | // |
| 401 | // kick=true starts a one-shot catalog discovery immediately (used for cache-miss |
| 402 | // servers at boot). kick=false leaves the process idle until the first real |
| 403 | // tool call — the product default for cache-hit sessions. |
| 404 | // |
| 405 | // host is the Host that receives the real Client. reg is the registry where |
| 406 | // real tools land after a successful spawn. sessionCtx must outlive any |
| 407 | // single Execute (use the controller's PluginCtx) — a turn-scoped ctx would |
| 408 | // kill the stdio child between turns. |
| 409 | func LazyToolset(spec Spec, cs *CachedSchema, host *Host, reg *tool.Registry, sessionCtx context.Context, kick bool) []tool.Tool { |
| 410 | // Resolve an existing exact project grant before constructing cached |
| 411 | // placeholders. This is read-only host preparation; no MCP process or network |
| 412 | // connection starts here. |
| 413 | spec = ResolveStoredAuthorization(sessionCtx, spec) |
| 414 | spawnCtx, cancel := context.WithCancel(sessionCtx) |
| 415 | shared := &lazySpawn{ |
| 416 | spec: spec, |
| 417 | host: host, |
| 418 | reg: reg, |
| 419 | ctx: spawnCtx, |
| 420 | waitBudget: defaultStartTimeout, |
| 421 | } |
| 422 | shared.generation = host.registerDeferredCancel(spec.Name, cancel) |
| 423 | |
| 424 | var out []tool.Tool |
| 425 | // A snapshot with zero tools presents nothing the model could call, so it |
| 426 | // gets the same connect stub as a cache miss — otherwise the live tools |
| 427 | // would silently join the registry mid-session with no placeholder names |
| 428 | // reserved for them. |
| 429 | if cs == nil || len(cs.Tools) == 0 { |
| 430 | shared.removePrefix = ToolPrefix(spec.Name) |
| 431 | out = []tool.Tool{&lazyTool{ |
| 432 | shared: shared, |
| 433 | name: shared.removePrefix + "connect", |
| 434 | desc: fmt.Sprintf("Connect MCP server %q. Call this once to drive the handshake; the server's real tools become available on the next turn.", spec.Name), |
| 435 | hasCache: false, |
| 436 | }} |
| 437 | } else { |
| 438 | out = make([]tool.Tool, 0, len(cs.Tools)) |
| 439 | for _, ct := range cs.Tools { |
| 440 | visibleName := ct.Name |
| 441 | if spec.StripRawPrefix != "" { |
| 442 | visibleName = strings.TrimPrefix(visibleName, spec.StripRawPrefix) |
| 443 | } |
| 444 | out = append(out, &lazyTool{ |
| 445 | shared: shared, |
| 446 | name: toolName(spec.Name, visibleName), |
| 447 | rawName: ct.Name, |
| 448 | visibleName: visibleName, |
| 449 | desc: ct.Description, |
| 450 | schema: ct.Schema, |
| 451 | readOnly: ct.ReadOnly, |
| 452 | destructive: ct.Destructive, |
| 453 | hasCache: true, |
| 454 | }) |
| 455 | } |
| 456 | } |
| 457 | |
| 458 | if kick { |
| 459 | shared.kick() |
| 460 | } |
| 461 | return out |
| 462 | } |
| 463 |