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._
This commit is contained in:
Kyle Carberry
2026-05-12 00:30:56 -04:00
committed by GitHub
parent 5a5cd79c4c
commit 376fc80451
4 changed files with 541 additions and 61 deletions
+58
View File
@@ -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