From 376fc80451f94bbf1d73a1f26b471129f23b7bd1 Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Tue, 12 May 2026 00:30:56 -0400 Subject: [PATCH] fix(coderd/x/chatd): discover workspace MCP tools mid-turn after create_workspace (#25169) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem In `coderd/x/chatd/chatd.go` `runChat`, workspace MCP discovery is gated on `chat.WorkspaceID.Valid` at the start of each turn. New chats that bind their workspace mid-turn (via `create_workspace` or `start_workspace`) get an empty workspace tool list on the first step, and the model falls back to `execute` (bash) because no workspace MCP tools are advertised. **Repro:** new chat → "create a workspace and use MCP tools". No `/api/v0/mcp/tools` request hits the agent on turn 1; turn 2 in the same chat works fine. ## Fix - Add a `PrepareTools` callback to `chatloop.RunOptions`, analogous to `PrepareMessages`. It is invoked once before each LLM step with the current tool list. When it returns non-nil, the chatloop replaces `opts.Tools`, rebuilds the per-step tool definitions, and appends new tool names to `opts.ActiveTools` so newly injected tools are callable immediately. - Wire `PrepareTools` in `runChat` to trigger workspace MCP discovery the first time the chat snapshot reports a valid `WorkspaceID`. The previous top-of-turn discovery path is unchanged for chats that start with a workspace. - Extract the discovery logic into `Server.discoverWorkspaceMCPTools` so the top-of-turn and mid-turn paths share identical behavior (cache, agent resolution, `ListMCPTools` timeout, invalidation). Mid-turn discovery stays disabled in plan-mode turns and Explore subagents, matching the existing top-of-turn gate. The `workspaceMCPDiscovered` flag prevents redundant dials after the first successful discovery. ## Tests - `coderd/x/chatd/chatloop/chatloop_test.go`: two new `TestRun_PrepareTools*` cases covering injection on the next step and active-set merging when `ActiveTools` is non-empty. - `coderd/x/chatd/chatd_test.go`: `TestRunChat_WorkspaceMCPDiscoveryAfterMidTurnCreateWorkspace` drives `runChat` through a `create_workspace` tool call against a real Postgres + mocked agent conn and asserts the second streamed LLM request advertises the workspace MCP tool. Verified that the test fails (and pinpoints the missing tool) when the `PrepareTools` wiring is disabled. ## Validation ``` go test ./coderd/x/chatd/chatloop/... -count=1 go test ./coderd/x/chatd/... -count=1 make lint/emdash ```
Decision log - Chose a per-step `PrepareTools` callback over mutating `opts.Tools` in place because `chatloop.Run` builds the `fantasy.Tool` definitions once at start; a hook is required to let the LLM see new tools on the next step. - Returned `[]fantasy.AgentTool` (not also active-tool-names) and let the chatloop derive name merges via `mergeNewToolNames`. This avoids leaking plan-mode gating decisions into the callback contract. - Kept the existing top-of-turn discovery path so chats that already have a workspace at turn start pay no extra latency. - Skipped reusing `ReloadMessages` (history reload) since this is purely a tool-availability concern; coupling it to a history reload would defeat the chatloop cache prefix optimizations.
--- _This pull request was generated by Coder Agents._ --- coderd/x/chatd/chatd.go | 174 ++++++++++++------- coderd/x/chatd/chatd_test.go | 160 +++++++++++++++++ coderd/x/chatd/chatloop/chatloop.go | 58 +++++++ coderd/x/chatd/chatloop/chatloop_test.go | 210 +++++++++++++++++++++++ 4 files changed, 541 insertions(+), 61 deletions(-) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index e5703be92e..00da401d33 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -493,6 +493,81 @@ func (p *Server) loadCachedWorkspaceContext( return tools } +// discoverWorkspaceMCPTools resolves the chat's workspace agent and +// lists the workspace MCP tools advertised by that agent. Results are +// cached per chat keyed on the agent ID so subsequent calls hit the +// cache. Returns nil (and never an error) on every failure mode so the +// caller can continue without MCP tools. +// +// This helper is shared between the top-of-turn discovery path and the +// mid-turn PrepareTools path triggered after create_workspace / +// start_workspace bind a workspace to a chat that started without one. +func (p *Server) discoverWorkspaceMCPTools( + ctx context.Context, + logger slog.Logger, + chatID uuid.UUID, + workspaceCtx *turnWorkspaceContext, +) []fantasy.AgentTool { + // Fast path: check cache using the in-memory cached agent + // (ensureWorkspaceAgent is free when already loaded). This + // avoids a per-turn latest-build DB query on the common + // subsequent-turn path. + if agent, agentErr := workspaceCtx.getWorkspaceAgent(ctx); agentErr == nil { + if tools := p.loadCachedWorkspaceContext( + chatID, agent, workspaceCtx.getWorkspaceConn, + ); tools != nil { + return tools + } + } // Cache miss, agent changed, or no cache: validate + // that the workspace still has a live agent before + // attempting a dial. + _, _, agentErr := workspaceCtx.workspaceAgentIDForConn(ctx) + if agentErr != nil { + if xerrors.Is(agentErr, errChatHasNoWorkspaceAgent) { + p.workspaceMCPToolsCache.Delete(chatID) + return nil + } + logger.Warn(ctx, "failed to resolve workspace agent for MCP tools", + slog.Error(agentErr)) + return nil + } + + // List workspace MCP tools via the agent conn. + conn, connErr := workspaceCtx.getWorkspaceConn(ctx) + if connErr != nil { + logger.Warn(ctx, "failed to get workspace conn for MCP tools", + slog.Error(connErr)) + return nil + } + listCtx, cancel := context.WithTimeout(ctx, workspaceMCPDiscoveryTimeout) + defer cancel() + toolsResp, listErr := conn.ListMCPTools(listCtx) + if listErr != nil { + logger.Warn(ctx, "failed to list workspace MCP tools", + slog.Error(listErr)) + return nil + } + // Cache the result for subsequent turns. Skip caching when + // the list is empty because the agent's MCP Connect may not + // have finished yet; caching an empty list would hide tools + // permanently. + if len(toolsResp.Tools) > 0 { + if agent, agentErr := workspaceCtx.getWorkspaceAgent(ctx); agentErr == nil { + p.workspaceMCPToolsCache.Store(chatID, &cachedWorkspaceMCPTools{ + agentID: agent.ID, + tools: toolsResp.Tools, + }) + } + } + + invalidate := func() { p.workspaceMCPToolsCache.Delete(chatID) } + tools := make([]fantasy.AgentTool, 0, len(toolsResp.Tools)) + for _, t := range toolsResp.Tools { + tools = append(tools, chattool.NewWorkspaceMCPTool(t, workspaceCtx.getWorkspaceConn, invalidate)) + } + return tools +} + type turnWorkspaceContext struct { server *Server chatStateMu *sync.Mutex @@ -6875,69 +6950,14 @@ func (p *Server) runChat( } // Workspace MCP discovery stays disabled for all plan-mode turns. // Root plan mode only gets approved external MCP servers, and - // plan-mode subagents get no MCP tools. + // plan-mode subagents get no MCP tools. When the chat has no + // workspace yet, discovery happens mid-turn via the chatloop + // PrepareTools callback installed below in chatloop.Run options. if chat.WorkspaceID.Valid && !isPlanModeTurn { g2.Go(func() error { - // Fast path: check cache using the in-memory cached - // agent (ensureWorkspaceAgent is free when already - // loaded). This avoids a per-turn latest-build DB - // query on the common subsequent-turn path. - agent, agentErr := workspaceCtx.getWorkspaceAgent(ctx) - if agentErr == nil { - if workspaceMCPTools = p.loadCachedWorkspaceContext( - chat.ID, agent, workspaceCtx.getWorkspaceConn, - ); workspaceMCPTools != nil { - return nil - } - } // Cache miss, agent changed, or no cache: validate - // that the workspace still has a live agent before - // attempting a dial. - _, _, agentErr = workspaceCtx.workspaceAgentIDForConn(ctx) - if agentErr != nil { - if xerrors.Is(agentErr, errChatHasNoWorkspaceAgent) { - p.workspaceMCPToolsCache.Delete(chat.ID) - return nil - } - logger.Warn(ctx, "failed to resolve workspace agent for MCP tools", - slog.Error(agentErr)) - return nil - } - - // List workspace MCP tools via the agent conn. - conn, connErr := workspaceCtx.getWorkspaceConn(ctx) - if connErr != nil { - logger.Warn(ctx, "failed to get workspace conn for MCP tools", - slog.Error(connErr)) - return nil - } - listCtx, cancel := context.WithTimeout(ctx, workspaceMCPDiscoveryTimeout) - defer cancel() - toolsResp, listErr := conn.ListMCPTools(listCtx) - if listErr != nil { - logger.Warn(ctx, "failed to list workspace MCP tools", - slog.Error(listErr)) - return nil - } - // Cache the result for subsequent turns. Skip - // caching when the list is empty because the - // agent's MCP Connect may not have finished yet; - // caching an empty list would hide tools - // permanently. - if len(toolsResp.Tools) > 0 { - if agent, agentErr := workspaceCtx.getWorkspaceAgent(ctx); agentErr == nil { - p.workspaceMCPToolsCache.Store(chat.ID, &cachedWorkspaceMCPTools{ - agentID: agent.ID, - tools: toolsResp.Tools, - }) - } - } - - invalidate := func() { p.workspaceMCPToolsCache.Delete(chat.ID) } - for _, t := range toolsResp.Tools { - workspaceMCPTools = append(workspaceMCPTools, - chattool.NewWorkspaceMCPTool(t, workspaceCtx.getWorkspaceConn, invalidate), - ) - } + workspaceMCPTools = p.discoverWorkspaceMCPTools( + ctx, logger, chat.ID, &workspaceCtx, + ) return nil }) } @@ -6984,6 +7004,15 @@ func (p *Server) runChat( } instructionInjected := instruction != "" + // workspaceMCPDiscovered tracks whether workspace MCP discovery + // has already been attempted for this turn. The top-of-turn + // discovery path above only fires when chat.WorkspaceID is + // valid at the start of the turn. For chats that bind a + // workspace mid-turn (e.g. via create_workspace) the chatloop + // PrepareTools callback below triggers discovery on the next + // step. After discovery has run once (here or in PrepareTools), + // this flag prevents redundant dials. + workspaceMCPDiscovered := chat.WorkspaceID.Valid || isPlanModeTurn prompt = renderPlanPathPrompt(prompt, resolvePlanPathBlock(ctx)) setAdvisorPromptSnapshot(prompt) // Use the model config's context_limit as a fallback when the LLM @@ -7682,6 +7711,29 @@ func (p *Server) runChat( DisableChainMode: func() { chainModeActive = false }, + PrepareTools: func(currentTools []fantasy.AgentTool) []fantasy.AgentTool { + // Mid-turn workspace MCP discovery for chats that bind a + // workspace via create_workspace or start_workspace + // after the turn has already started. The top-of-turn + // discovery path is gated on chat.WorkspaceID.Valid; this + // callback bridges the gap so the LLM sees workspace MCP + // tools on the very next step instead of the turn after. + if workspaceMCPDiscovered || isExploreSubagent { + return nil + } + snapshot := workspaceCtx.currentChatSnapshot() + if !snapshot.WorkspaceID.Valid { + return nil + } + workspaceMCPDiscovered = true + discovered := p.discoverWorkspaceMCPTools( + ctx, loopLogger, chat.ID, &workspaceCtx, + ) + if len(discovered) == 0 { + return nil + } + return append(slices.Clone(currentTools), discovered...) + }, PrepareMessages: func(msgs []fantasy.Message) []fantasy.Message { // Skip the snapshot update when chain mode is active; // the chatloop passes in the chain-filtered prompt diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 6ec2e33920..02d8cf50d9 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -11339,3 +11339,163 @@ func TestRunChat_WorkspaceMCPDiscoveryWaitsForSlowAgent(t *testing.T) { "workspace MCP tool should reach the LLM once chatd's discovery "+ "timeout exceeds the agent's MCP reload time") } + +// TestRunChat_WorkspaceMCPDiscoveryAfterMidTurnCreateWorkspace guards the +// regression where chats that bound their workspace mid-turn (via +// create_workspace) never saw workspace MCP tools on the same turn. The +// chatloop tool list was frozen at the top of the turn, so the first +// post-create_workspace step had no workspace MCP tools and the model +// fell back to bash. See PrepareTools wiring in runChat. +func TestRunChat_WorkspaceMCPDiscoveryAfterMidTurnCreateWorkspace(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + var ( + requestsMu sync.Mutex + requests []recordedOpenAIRequest + ) + + workspaceToolName := "workspace-midturn-mcp__echo" + workspaceCreateToolArgsJSON := "" + + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + + requestsMu.Lock() + requests = append(requests, recordOpenAIRequest(req)) + callIdx := len(requests) + requestsMu.Unlock() + + if callIdx == 1 { + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk("create_workspace", workspaceCreateToolArgsJSON), + ) + } + return chattest.OpenAIStreamingResponse( + chattest.OpenAITextChunks("done")..., + ) + }) + + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + + // Seed a workspace+agent for create_workspace to bind to. + tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + OrganizationID: org.ID, + CreatedBy: user.ID, + }) + tpl := dbgen.Template(t, db, database.Template{ + CreatedBy: user.ID, + OrganizationID: org.ID, + ActiveVersionID: tv.ID, + }) + workspaceCreateToolArgsJSON = fmt.Sprintf(`{"template_id":%q}`, tpl.ID.String()) + + ws := dbgen.Workspace(t, db, database.WorkspaceTable{ + TemplateID: tpl.ID, + OwnerID: user.ID, + OrganizationID: org.ID, + }) + pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + InitiatorID: user.ID, + OrganizationID: org.ID, + CompletedAt: sql.NullTime{Valid: true, Time: dbtime.Now()}, + }) + build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + TemplateVersionID: tv.ID, + WorkspaceID: ws.ID, + JobID: pj.ID, + }) + res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + Transition: database.WorkspaceTransitionStart, + JobID: pj.ID, + }) + now := dbtime.Now() + dbAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: res.ID, + LifecycleState: database.WorkspaceAgentLifecycleStateReady, + StartedAt: sql.NullTime{Time: now, Valid: true}, + ReadyAt: sql.NullTime{Time: now, Valid: true}, + FirstConnectedAt: sql.NullTime{Time: now, Valid: true}, + LastConnectedAt: sql.NullTime{Time: now, Valid: true}, + }) + + workspaceToolsResp := workspacesdk.ListMCPToolsResponse{ + Tools: []workspacesdk.MCPToolInfo{{ + ServerName: "workspace-midturn-mcp", + Name: workspaceToolName, + Description: "workspace echo tool", + Schema: map[string]any{ + "input": map[string]any{"type": "string"}, + }, + Required: []string{"input"}, + }}, + } + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + mockConn.EXPECT().SetExtraHeaders(gomock.Any()).AnyTimes() + mockConn.EXPECT().ContextConfig(gomock.Any()). + Return(workspacesdk.ContextConfigResponse{}, xerrors.New("not supported")).AnyTimes() + mockConn.EXPECT().ListMCPTools(gomock.Any()). + Return(workspaceToolsResp, nil).AnyTimes() + mockConn.EXPECT().LS(gomock.Any(), gomock.Any(), gomock.Any()). + Return(workspacesdk.LSResponse{}, nil).AnyTimes() + mockConn.EXPECT().ReadFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(io.NopCloser(strings.NewReader("")), "", nil).AnyTimes() + mockConn.EXPECT().AwaitReachable(gomock.Any()).Return(true).AnyTimes() + + createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { + return codersdk.Workspace{ + ID: ws.ID, + Name: req.Name, + OwnerName: user.Username, + OrganizationID: org.ID, + TemplateID: tpl.ID, + LatestBuild: codersdk.WorkspaceBuild{ + ID: build.ID, + Status: codersdk.WorkspaceStatusRunning, + }, + }, nil + } + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + cfg.CreateWorkspace = createFn + }) + + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "workspace-mcp-midturn", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("Create a workspace and call the workspace MCP tool."), + }, + }) + require.NoError(t, err) + + chatResult := waitForTerminalChat(ctx, t, db, chat.ID) + if chatResult.Status == database.ChatStatusError { + require.FailNowf(t, "chat failed", "last_error=%q", + chatLastErrorMessage(chatResult.LastError)) + } + require.Equal(t, database.ChatStatusWaiting, chatResult.Status) + + requestsMu.Lock() + recorded := append([]recordedOpenAIRequest(nil), requests...) + requestsMu.Unlock() + require.GreaterOrEqual(t, len(recorded), 2, + "expected at least two streamed model calls (create_workspace + follow-up)") + require.NotContains(t, recorded[0].Tools, workspaceToolName, + "first call should not advertise workspace MCP tools because the chat has no workspace yet") + require.Contains(t, recorded[1].Tools, workspaceToolName, + "second call (after create_workspace) must advertise the workspace MCP tool: "+ + "this is the fix for mid-turn workspace MCP discovery") +} diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 21822c2311..739bfd0cbd 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -171,6 +171,20 @@ type RunOptions struct { // retry, so callbacks should avoid duplicating messages. PrepareMessages func([]fantasy.Message) []fantasy.Message + // PrepareTools is called once before each LLM step with the + // current tool list. If it returns non-nil, the returned slice + // replaces opts.Tools for this and all subsequent steps, and any + // new tool names are appended to opts.ActiveTools so they become + // callable immediately. Used to inject tools that become available + // mid-turn (e.g. workspace MCP tools discovered after + // create_workspace). + // + // The chatloop tracks whether tools have already been replaced so + // PrepareTools is not retried on subsequent steps once it has + // returned a non-nil slice. Callbacks may still be invoked on later + // steps when they previously returned nil. + PrepareTools func([]fantasy.AgentTool) []fantasy.AgentTool + // OnRetry is called before each retry attempt when the LLM // stream fails with a retryable error. It provides the attempt // number, raw error, normalized classification, and backoff @@ -392,6 +406,17 @@ func Run(ctx context.Context, opts RunOptions) error { modelName := opts.Model.Model() opts.Metrics.StepsTotal.WithLabelValues(provider, modelName).Inc() stepStart := time.Now() + if opts.PrepareTools != nil { + if updated := opts.PrepareTools(opts.Tools); updated != nil { + opts.ActiveTools = mergeNewToolNames( + opts.ActiveTools, opts.Tools, updated, + ) + opts.Tools = updated + tools = buildToolDefinitions( + opts.Tools, opts.ActiveTools, opts.ProviderTools, + ) + } + } var prepared []fantasy.Message messages, prepared = prepareMessagesForRequest( ctx, opts, messages, provider, modelName, step, totalSteps, @@ -1704,6 +1729,39 @@ func isToolActive(name string, activeTools []string) bool { return len(activeTools) == 0 || slices.Contains(activeTools, name) } +// mergeNewToolNames returns activeTools augmented with any tool names +// from newTools that are not present in oldTools and not already in +// activeTools. This keeps newly injected tools (e.g. via PrepareTools) +// callable even when activeTools is non-empty. +// +// When activeTools is empty, all tools are already active and the slice +// is returned unchanged. +func mergeNewToolNames(activeTools []string, oldTools, newTools []fantasy.AgentTool) []string { + if len(activeTools) == 0 { + return activeTools + } + old := make(map[string]struct{}, len(oldTools)) + for _, t := range oldTools { + old[t.Info().Name] = struct{}{} + } + active := make(map[string]struct{}, len(activeTools)) + for _, name := range activeTools { + active[name] = struct{}{} + } + for _, t := range newTools { + name := t.Info().Name + if _, alreadyActive := active[name]; alreadyActive { + continue + } + if _, existedBefore := old[name]; existedBefore { + continue + } + activeTools = append(activeTools, name) + active[name] = struct{}{} + } + return activeTools +} + // buildToolDefinitions converts AgentTool definitions into the // fantasy.Tool slice expected by fantasy.Call. When activeTools // is non-empty, only function tools whose name appears in the diff --git a/coderd/x/chatd/chatloop/chatloop_test.go b/coderd/x/chatd/chatloop/chatloop_test.go index 8af379e645..daa4588dcc 100644 --- a/coderd/x/chatd/chatloop/chatloop_test.go +++ b/coderd/x/chatd/chatloop/chatloop_test.go @@ -4007,6 +4007,216 @@ func TestRun_PrepareMessagesOnlyFiresOnce(t *testing.T) { require.Equal(t, 3, int(prepareCalls.Load())) } +// TestRun_PrepareToolsInjectsToolMidLoop guards the regression where a +// chat creating its workspace mid-turn (via create_workspace) saw the +// workspace MCP tools only on the next turn. Before the fix, the tool +// list was frozen at the top of the turn and the model could not call +// any workspace MCP tools until turn 2. With the fix, PrepareTools is +// invoked before every step and can inject tools that become available +// mid-loop. +func TestRun_PrepareToolsInjectsToolMidLoop(t *testing.T) { + t.Parallel() + + const injectedToolName = "workspace_mcp__echo" + + var mu sync.Mutex + var streamCalls int + var secondCallTools []fantasy.Tool + + // Step 0 calls create_workspace. Step 1 should see the + // injected workspace MCP tool. + model := &chattest.FakeModel{ + ProviderName: "fake", + StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { + mu.Lock() + step := streamCalls + streamCalls++ + mu.Unlock() + + switch step { + case 0: + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-1", ToolCallName: "create_workspace"}, + {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-1", Delta: `{}`}, + {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-1"}, + { + Type: fantasy.StreamPartTypeToolCall, + ID: "tc-1", + ToolCallName: "create_workspace", + ToolCallInput: `{}`, + }, + {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, + }), nil + default: + mu.Lock() + secondCallTools = append([]fantasy.Tool(nil), call.Tools...) + mu.Unlock() + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, + {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "done"}, + {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, + {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, + }), nil + } + }, + } + + var workspaceReady atomic.Bool + createWorkspaceTool := fantasy.NewAgentTool( + "create_workspace", + "create a workspace", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + workspaceReady.Store(true) + return fantasy.ToolResponse{}, nil + }, + ) + + var prepareCalls atomic.Int32 + err := Run(context.Background(), RunOptions{ + Model: model, + Messages: []fantasy.Message{ + textMessage(fantasy.MessageRoleUser, "create a workspace and use MCP"), + }, + Tools: []fantasy.AgentTool{createWorkspaceTool}, + ActiveTools: []string{"create_workspace"}, + MaxSteps: 5, + PersistStep: func(_ context.Context, _ PersistedStep) error { + return nil + }, + PrepareTools: func(currentTools []fantasy.AgentTool) []fantasy.AgentTool { + prepareCalls.Add(1) + if !workspaceReady.Load() { + return nil + } + return append(currentTools, newNoopTool(injectedToolName)) + }, + }) + require.NoError(t, err) + require.Equal(t, 2, streamCalls) + // PrepareTools is called before each of the 2 steps. + require.Equal(t, int32(2), prepareCalls.Load()) + + require.NotEmpty(t, secondCallTools) + var foundInjectedTool bool + for _, tool := range secondCallTools { + if tool.GetName() == injectedToolName { + foundInjectedTool = true + break + } + } + require.True(t, foundInjectedTool, + "step 1 prompt should advertise the workspace MCP tool injected by PrepareTools") +} + +// TestRun_PrepareToolsAddsNewToolToActiveSet guards the contract that +// when PrepareTools injects a tool, that tool is callable on the +// next step even when opts.ActiveTools was non-empty (and would +// otherwise filter the new tool out). +func TestRun_PrepareToolsAddsNewToolToActiveSet(t *testing.T) { + t.Parallel() + + const injectedToolName = "workspace_mcp__echo" + + var mu sync.Mutex + var streamCalls int + var injectedToolRan atomic.Bool + + model := &chattest.FakeModel{ + ProviderName: "fake", + StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + mu.Lock() + step := streamCalls + streamCalls++ + mu.Unlock() + + switch step { + case 0: + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-1", ToolCallName: "create_workspace"}, + {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-1", Delta: `{}`}, + {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-1"}, + { + Type: fantasy.StreamPartTypeToolCall, + ID: "tc-1", + ToolCallName: "create_workspace", + ToolCallInput: `{}`, + }, + {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, + }), nil + case 1: + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-2", ToolCallName: injectedToolName}, + {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-2", Delta: `{}`}, + {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-2"}, + { + Type: fantasy.StreamPartTypeToolCall, + ID: "tc-2", + ToolCallName: injectedToolName, + ToolCallInput: `{}`, + }, + {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, + }), nil + default: + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, + {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "done"}, + {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, + {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, + }), nil + } + }, + } + + var workspaceReady atomic.Bool + createWorkspaceTool := fantasy.NewAgentTool( + "create_workspace", + "create a workspace", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + workspaceReady.Store(true) + return fantasy.ToolResponse{}, nil + }, + ) + + injectedTool := fantasy.NewAgentTool( + injectedToolName, + "injected workspace MCP tool", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + injectedToolRan.Store(true) + return fantasy.ToolResponse{}, nil + }, + ) + + err := Run(context.Background(), RunOptions{ + Model: model, + Messages: []fantasy.Message{ + textMessage(fantasy.MessageRoleUser, "create a workspace and use MCP"), + }, + Tools: []fantasy.AgentTool{createWorkspaceTool}, + // Active list deliberately excludes the injected tool name; + // PrepareTools must add it so the tool is callable. + ActiveTools: []string{"create_workspace"}, + MaxSteps: 5, + PersistStep: func(_ context.Context, _ PersistedStep) error { + return nil + }, + PrepareTools: func(currentTools []fantasy.AgentTool) []fantasy.AgentTool { + if !workspaceReady.Load() { + return nil + } + for _, t := range currentTools { + if t.Info().Name == injectedToolName { + return nil + } + } + return append(currentTools, injectedTool) + }, + }) + require.NoError(t, err) + require.GreaterOrEqual(t, streamCalls, 2) + require.True(t, injectedToolRan.Load(), + "injected tool must be callable on the step after PrepareTools adds it") +} + func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { t.Parallel()