mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user