fix(coderd/x/chatd): discover workspace MCP tools mid-turn after create_workspace (#25169)

## 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 ``` <details> <summary>Decision log</summary> - 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. </details> --- _This pull request was generated by Coder Agents._
2026-09-24 15:04:27 +08:00 · 2026-05-12 00:30:56 -04:00
parent 5a5cd79c4c
commit 376fc80451
4 changed files with 541 additions and 61 deletions
@@ -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
@@ -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()