diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 1f88828c6a..a8f5035aa8 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -72,7 +72,22 @@ const ( // cold-start agent's first MCP reload can settle before // chatd gives up. workspaceMCPDiscoveryTimeout = 35 * time.Second - turnStatusLabelWriteTimeout = 5 * time.Second + // workspaceMCPPrimeMaxWait bounds the deadline used by the + // create_workspace / start_workspace post-ready cache primer + // loop. The primer checks the deadline only after each + // discoverWorkspaceMCPTools call returns, so total wall-clock + // time can exceed this by one such call (dialTimeout + + // workspaceMCPDiscoveryTimeout in the worst case). The constant + // caps when new retries can start, not when an in-flight call + // must finish. Empty results usually mean the agent's MCP + // Connect is still racing with agent startup. The agent-side + // budget is agent/x/agentmcp.connectTimeout (30s). + workspaceMCPPrimeMaxWait = 30 * time.Second + // workspaceMCPPrimeRetryInterval is the short backoff between + // re-attempts inside the primer when ListMCPTools returns an + // empty list without error. + workspaceMCPPrimeRetryInterval = 2 * time.Second + turnStatusLabelWriteTimeout = 5 * time.Second // defaultDialTimeout matches the timeout used by ~8 other // server-side AgentConn callers. defaultDialTimeout = 30 * time.Second @@ -568,6 +583,61 @@ func (p *Server) discoverWorkspaceMCPTools( return tools } +// primeWorkspaceMCPCache populates workspaceMCPToolsCache after the +// create_workspace or start_workspace tool finishes waiting for the +// workspace agent to become reachable. By the time it runs the agent +// is already Ready, so a single ListMCPTools call usually succeeds. +// When the agent's MCP server is still racing with agent startup, +// ListMCPTools may return an empty list (no error) on the first call; +// the primer retries with a short backoff up to +// workspaceMCPPrimeMaxWait so the LLM step that follows the tool call +// sees the workspace MCP tools in the cache and PrepareTools does not +// need to dial again. +// +// Returns silently on every failure mode. The chat continues without +// workspace MCP tools when the agent does not advertise any within +// the budget. The next user turn re-runs top-of-turn discovery from +// scratch. +func (p *Server) primeWorkspaceMCPCache( + ctx context.Context, + logger slog.Logger, + chatID uuid.UUID, + workspaceCtx *turnWorkspaceContext, +) { + deadline := p.clock.Now().Add(workspaceMCPPrimeMaxWait) + attempt := 0 + for { + attempt++ + tools := p.discoverWorkspaceMCPTools(ctx, logger, chatID, workspaceCtx) + if len(tools) > 0 { + logger.Debug(ctx, "primed workspace MCP cache", + slog.F("chat_id", chatID), + slog.F("tool_count", len(tools)), + slog.F("attempts", attempt), + ) + return + } + if ctx.Err() != nil { + return + } + if !p.clock.Now().Before(deadline) { + logger.Debug(ctx, + "workspace MCP cache primer gave up waiting for tools", + slog.F("chat_id", chatID), + slog.F("attempts", attempt), + ) + return + } + timer := p.clock.NewTimer(workspaceMCPPrimeRetryInterval, "chatd", "workspace-mcp-prime") + select { + case <-timer.C: + case <-ctx.Done(): + timer.Stop() + return + } + } +} + type turnWorkspaceContext struct { server *Server chatStateMu *sync.Mutex @@ -6457,6 +6527,11 @@ type rootChatToolsOptions struct { resolvePlanPath func(context.Context) (string, string, error) storeFile chattool.StoreFileFunc isPlanModeTurn bool + // primerCtx scopes the workspace MCP cache primer goroutines + // that onChatUpdated launches. runChat cancels it before + // workspaceCtx.close() so an in-flight primer cannot dial a + // fresh conn after the cached one was released. + primerCtx context.Context } func (p *Server) loadPlanModeInstructions( @@ -6520,6 +6595,50 @@ func (p *Server) appendRootChatTools( } } } + + // Prime the workspace MCP tools cache while the create_workspace + // or start_workspace tool is still running. The AgentID guard + // below restricts the primer to the post-ready callback, when + // the agent is reachable. ListMCPTools may still return an + // empty list on the first try when the agent's MCP Connect is + // racing with agent startup; primeWorkspaceMCPCache retries + // with a short backoff up to workspaceMCPPrimeMaxWait. Priming + // here lets the next LLM step's PrepareTools hit the cache + // instead of dialing again on a separate timeout budget. + // + // Run asynchronously: the tool itself must not block on the + // primer because the agent may not advertise any MCP tools at + // all (e.g. minimal templates), in which case the primer waits + // the full budget before giving up. PrepareTools on the next + // step covers the cache miss path; the primer is purely an + // optimization that warms the cache while the LLM is thinking. + // inflight tracking ensures server shutdown still waits for any + // in-progress primer. + // + // Guard on both WorkspaceID and AgentID being valid: + // create_workspace and start_workspace each fire onChatUpdated + // twice for a new build (binding before waitForAgentReady; + // post-ready after it), and stop_workspace fires it with a nil + // agent. Only the post-ready callback has a live AgentID, so + // the pre-build and stop-side firings would otherwise spawn a + // primer goroutine that dials a missing or dying agent and + // burns the full budget for nothing. + // + // Read the snapshot from workspaceCtx rather than the + // updatedChat parameter: persistInstructionFiles above runs + // ensureWorkspaceAgent which calls persistBuildAgentBinding and + // setCurrentChat, so by the time we get here the in-memory + // snapshot has the freshly bound AgentID even when the + // updatedChat parameter (read from the DB before the binding + // was persisted) does not. + snapshot := opts.workspaceCtx.currentChatSnapshot() + if snapshot.WorkspaceID.Valid && snapshot.AgentID.Valid { + p.inflight.Add(1) + go func() { + defer p.inflight.Done() + p.primeWorkspaceMCPCache(opts.primerCtx, p.logger, snapshot.ID, opts.workspaceCtx) + }() + } } tools = append(tools, @@ -6852,7 +6971,16 @@ func (p *Server) runChat( currentChat: ¤tChat, loadChatSnapshot: loadChatSnapshot, } - defer workspaceCtx.close() + // primerCtx scopes the workspace MCP cache primer goroutines that + // onChatUpdated launches. We cancel it before workspaceCtx.close() + // so an in-flight primer cannot wake from its retry backoff, + // observe a cleared cached conn, dial a fresh one, and leak it + // when no subsequent close() runs. + primerCtx, primerCancel := context.WithCancel(ctx) + defer func() { + primerCancel() + workspaceCtx.close() + }() planPathFn := func(ctx context.Context) (string, string, error) { conn, err := workspaceCtx.getWorkspaceConn(ctx) @@ -7435,6 +7563,7 @@ func (p *Server) runChat( resolvePlanPath: resolvePlanPathForTools, storeFile: storeChatAttachment, isPlanModeTurn: isPlanModeTurn, + primerCtx: primerCtx, }) } @@ -7766,11 +7895,17 @@ func (p *Server) runChat( }, 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. + // 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. + // + // create_workspace and start_workspace prime + // workspaceMCPToolsCache via onChatUpdated after + // waitForAgentReady returns, so the call below is almost + // always a cache hit. The primer's bounded wait means the + // dial fallback here only runs when priming itself failed. if workspaceMCPDiscovered || isExploreSubagent { return nil } @@ -7778,13 +7913,21 @@ func (p *Server) runChat( if !snapshot.WorkspaceID.Valid { return nil } - workspaceMCPDiscovered = true discovered := p.discoverWorkspaceMCPTools( ctx, loopLogger, chat.ID, &workspaceCtx, ) if len(discovered) == 0 { + // Leave workspaceMCPDiscovered false so a subsequent + // step retries discovery. PrepareTools fires once per + // LLM step, so retries are unbounded for the rest of + // the turn. Per-step cost is one + // GetWorkspaceAgentsInLatestBuildByWorkspaceID query + // plus one ListMCPTools RPC, both fast against a live + // conn. The primer's 30s budget applies to its own + // loop only. return nil } + workspaceMCPDiscovered = true return append(slices.Clone(currentTools), discovered...) }, PrepareMessages: func(msgs []fantasy.Message) []fantasy.Message { diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 45a2940882..e18e0574cd 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -5,6 +5,7 @@ import ( "database/sql" "encoding/json" "sync" + "sync/atomic" "testing" "time" @@ -5732,3 +5733,441 @@ func TestSubscribeToStream_FiltersBufferedParts_Integration(t *testing.T) { require.Equal(t, []string{"C-1"}, texts, "only in-progress (un-claimed) buffered parts must survive the filter") } + +// TestPrimeWorkspaceMCPCache_SuccessOnFirstAttempt verifies the +// onChatUpdated cache primer path: when create_workspace / +// start_workspace finish waitForAgentReady and the agent's MCP +// server is already advertising tools, a single ListMCPTools call +// populates the cache so the next PrepareTools step is a cache hit +// and does not need to dial. +func TestPrimeWorkspaceMCPCache_SuccessOnFirstAttempt(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + workspaceID := uuid.New() + agentID := uuid.New() + chat := database.Chat{ + ID: uuid.New(), + WorkspaceID: uuid.NullUUID{ + UUID: workspaceID, + Valid: true, + }, + AgentID: uuid.NullUUID{ + UUID: agentID, + Valid: true, + }, + } + now := time.Now() + workspaceAgent := database.WorkspaceAgent{ + ID: agentID, + FirstConnectedAt: sql.NullTime{ + Time: now.Add(-time.Minute), + Valid: true, + }, + LastConnectedAt: sql.NullTime{ + Time: now, + Valid: true, + }, + } + + db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). + Return(workspaceAgent, nil).AnyTimes() + db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). + Return([]database.WorkspaceAgent{workspaceAgent}, nil).AnyTimes() + + toolName := "workspace-mcp__echo" + conn := agentconnmock.NewMockAgentConn(ctrl) + conn.EXPECT().SetExtraHeaders(gomock.Any()).AnyTimes() + conn.EXPECT().ListMCPTools(gomock.Any()).Return(workspacesdk.ListMCPToolsResponse{ + Tools: []workspacesdk.MCPToolInfo{{ + ServerName: "workspace-mcp", + Name: toolName, + Schema: map[string]any{}, + }}, + }, nil).Times(1) + + server := &Server{ + db: db, + logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + clock: quartz.NewMock(t), + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: time.Second, + agentConnFn: func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { + return conn, func() {}, nil + }, + } + + chatStateMu := &sync.Mutex{} + currentChat := chat + workspaceCtx := turnWorkspaceContext{ + server: server, + chatStateMu: chatStateMu, + currentChat: ¤tChat, + loadChatSnapshot: func(context.Context, uuid.UUID) (database.Chat, error) { return chat, nil }, + } + t.Cleanup(workspaceCtx.close) + + server.primeWorkspaceMCPCache(ctx, server.logger, chat.ID, &workspaceCtx) + + cached, ok := server.workspaceMCPToolsCache.Load(chat.ID) + require.True(t, ok, "primer must populate the cache on success") + entry, ok := cached.(*cachedWorkspaceMCPTools) + require.True(t, ok) + require.Equal(t, agentID, entry.agentID) + require.Len(t, entry.tools, 1) + require.Equal(t, toolName, entry.tools[0].Name) +} + +// TestPrimeWorkspaceMCPCache_RetriesUntilToolsAppear simulates the +// race between agent reachability and the agent's MCP Connect: the +// first ListMCPTools call returns an empty list (no error), the +// second returns the workspace tools. The primer must retry after +// workspaceMCPPrimeRetryInterval and write the cache on the second +// attempt. +func TestPrimeWorkspaceMCPCache_RetriesUntilToolsAppear(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + workspaceID := uuid.New() + agentID := uuid.New() + chat := database.Chat{ + ID: uuid.New(), + WorkspaceID: uuid.NullUUID{ + UUID: workspaceID, + Valid: true, + }, + AgentID: uuid.NullUUID{ + UUID: agentID, + Valid: true, + }, + } + now := time.Now() + workspaceAgent := database.WorkspaceAgent{ + ID: agentID, + FirstConnectedAt: sql.NullTime{ + Time: now.Add(-time.Minute), + Valid: true, + }, + LastConnectedAt: sql.NullTime{ + Time: now, + Valid: true, + }, + } + + db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). + Return(workspaceAgent, nil).AnyTimes() + db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). + Return([]database.WorkspaceAgent{workspaceAgent}, nil).AnyTimes() + + toolName := "workspace-mcp__echo" + var listCalls atomic.Int32 + emptyOnce := make(chan struct{}, 1) + emptyOnce <- struct{}{} + conn := agentconnmock.NewMockAgentConn(ctrl) + conn.EXPECT().SetExtraHeaders(gomock.Any()).AnyTimes() + conn.EXPECT().ListMCPTools(gomock.Any()).DoAndReturn( + func(context.Context) (workspacesdk.ListMCPToolsResponse, error) { + listCalls.Add(1) + select { + case <-emptyOnce: + return workspacesdk.ListMCPToolsResponse{}, nil + default: + return workspacesdk.ListMCPToolsResponse{ + Tools: []workspacesdk.MCPToolInfo{{ + ServerName: "workspace-mcp", + Name: toolName, + Schema: map[string]any{}, + }}, + }, nil + } + }, + ).AnyTimes() + + mockClock := quartz.NewMock(t) + timerTrap := mockClock.Trap().NewTimer("chatd", "workspace-mcp-prime") + t.Cleanup(timerTrap.Close) + + server := &Server{ + db: db, + logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + clock: mockClock, + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: time.Second, + agentConnFn: func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { + return conn, func() {}, nil + }, + } + + chatStateMu := &sync.Mutex{} + currentChat := chat + workspaceCtx := turnWorkspaceContext{ + server: server, + chatStateMu: chatStateMu, + currentChat: ¤tChat, + loadChatSnapshot: func(context.Context, uuid.UUID) (database.Chat, error) { return chat, nil }, + } + t.Cleanup(workspaceCtx.close) + + done := make(chan struct{}) + go func() { + defer close(done) + server.primeWorkspaceMCPCache(ctx, server.logger, chat.ID, &workspaceCtx) + }() + + // First attempt returns empty. The primer arms a timer; release + // it and advance the clock so the second attempt fires. + call := timerTrap.MustWait(ctx) + call.MustRelease(ctx) + mockClock.Advance(workspaceMCPPrimeRetryInterval).MustWait(ctx) + + select { + case <-done: + case <-ctx.Done(): + t.Fatal("primer did not finish after second attempt") + } + + require.GreaterOrEqual(t, listCalls.Load(), int32(2), + "primer must retry after empty result") + cached, ok := server.workspaceMCPToolsCache.Load(chat.ID) + require.True(t, ok, "primer must populate the cache on retry success") + entry, ok := cached.(*cachedWorkspaceMCPTools) + require.True(t, ok) + require.Equal(t, agentID, entry.agentID) + require.Len(t, entry.tools, 1) + require.Equal(t, toolName, entry.tools[0].Name) +} + +// TestPrimeWorkspaceMCPCache_GivesUpAfterDeadline verifies the +// bounded-wait guarantee: when ListMCPTools always returns an empty +// list (e.g. the agent's MCP server never advertises tools), the +// primer stops trying at workspaceMCPPrimeMaxWait and does not cache +// the empty result. PrepareTools is then free to retry on the next +// chat step. +func TestPrimeWorkspaceMCPCache_GivesUpAfterDeadline(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + workspaceID := uuid.New() + agentID := uuid.New() + chat := database.Chat{ + ID: uuid.New(), + WorkspaceID: uuid.NullUUID{ + UUID: workspaceID, + Valid: true, + }, + AgentID: uuid.NullUUID{ + UUID: agentID, + Valid: true, + }, + } + now := time.Now() + workspaceAgent := database.WorkspaceAgent{ + ID: agentID, + FirstConnectedAt: sql.NullTime{ + Time: now.Add(-time.Minute), + Valid: true, + }, + LastConnectedAt: sql.NullTime{ + Time: now, + Valid: true, + }, + } + + db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). + Return(workspaceAgent, nil).AnyTimes() + db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). + Return([]database.WorkspaceAgent{workspaceAgent}, nil).AnyTimes() + + var listCalls atomic.Int32 + conn := agentconnmock.NewMockAgentConn(ctrl) + conn.EXPECT().SetExtraHeaders(gomock.Any()).AnyTimes() + conn.EXPECT().ListMCPTools(gomock.Any()).DoAndReturn( + func(context.Context) (workspacesdk.ListMCPToolsResponse, error) { + listCalls.Add(1) + return workspacesdk.ListMCPToolsResponse{}, nil + }, + ).AnyTimes() + + mockClock := quartz.NewMock(t) + timerTrap := mockClock.Trap().NewTimer("chatd", "workspace-mcp-prime") + t.Cleanup(timerTrap.Close) + + server := &Server{ + db: db, + logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + clock: mockClock, + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: time.Second, + agentConnFn: func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { + return conn, func() {}, nil + }, + } + + chatStateMu := &sync.Mutex{} + currentChat := chat + workspaceCtx := turnWorkspaceContext{ + server: server, + chatStateMu: chatStateMu, + currentChat: ¤tChat, + loadChatSnapshot: func(context.Context, uuid.UUID) (database.Chat, error) { return chat, nil }, + } + t.Cleanup(workspaceCtx.close) + + done := make(chan struct{}) + go func() { + defer close(done) + server.primeWorkspaceMCPCache(ctx, server.logger, chat.ID, &workspaceCtx) + }() + + // Drive the retry loop forward until the primer gives up. Each + // iteration: release the trapped NewTimer call, then advance the + // clock past the retry interval. The primer exits when + // p.clock.Now() is no longer before deadline. The loop bounds + // itself on maxIterations and uses a done-aware wait context so + // the test fails cleanly instead of hanging when the primer + // shuts down between iterations. + maxIterations := int(workspaceMCPPrimeMaxWait/workspaceMCPPrimeRetryInterval) + 2 +Loop: + for i := 0; i < maxIterations; i++ { + waitCtx, cancel := context.WithCancel(ctx) + go func() { + select { + case <-done: + cancel() + case <-waitCtx.Done(): + } + }() + call, err := timerTrap.Wait(waitCtx) + cancel() + if err != nil { + break Loop + } + call.MustRelease(ctx) + mockClock.Advance(workspaceMCPPrimeRetryInterval).MustWait(ctx) + } + + // expectedAttempts is the floor on how many times the primer + // should call discoverWorkspaceMCPTools before the deadline + // expires. The primer makes one attempt before sleeping, then + // one per workspaceMCPPrimeRetryInterval until the deadline. + // We assert a high-water mark (rather than exact equality) so + // the test is robust to off-by-one boundaries while still + // catching deadline miscomputations: a primer that exits after a + // handful of attempts would suggest the deadline was set with a + // shorter window than workspaceMCPPrimeMaxWait. + expectedAttempts := int32(workspaceMCPPrimeMaxWait/workspaceMCPPrimeRetryInterval) / 2 + require.GreaterOrEqual(t, listCalls.Load(), expectedAttempts, + "primer must retry enough times to consume the full budget") + _, ok := server.workspaceMCPToolsCache.Load(chat.ID) + require.False(t, ok, + "primer must not cache an empty result; PrepareTools needs to retry on the next step") +} + +// TestPrimeWorkspaceMCPCache_ExitsOnContextCancel verifies the +// primer's context.Done() branch: the retry loop must exit promptly +// when the chat ctx is canceled (runChat cancels its primerCtx +// before workspaceCtx.close runs to prevent a primer from re-dialing +// the freed conn). +func TestPrimeWorkspaceMCPCache_ExitsOnContextCancel(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + workspaceID := uuid.New() + agentID := uuid.New() + chat := database.Chat{ + ID: uuid.New(), + WorkspaceID: uuid.NullUUID{ + UUID: workspaceID, + Valid: true, + }, + AgentID: uuid.NullUUID{ + UUID: agentID, + Valid: true, + }, + } + now := time.Now() + workspaceAgent := database.WorkspaceAgent{ + ID: agentID, + FirstConnectedAt: sql.NullTime{ + Time: now.Add(-time.Minute), + Valid: true, + }, + LastConnectedAt: sql.NullTime{ + Time: now, + Valid: true, + }, + } + + db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). + Return(workspaceAgent, nil).AnyTimes() + db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). + Return([]database.WorkspaceAgent{workspaceAgent}, nil).AnyTimes() + + conn := agentconnmock.NewMockAgentConn(ctrl) + conn.EXPECT().SetExtraHeaders(gomock.Any()).AnyTimes() + conn.EXPECT().ListMCPTools(gomock.Any()). + Return(workspacesdk.ListMCPToolsResponse{}, nil).AnyTimes() + + mockClock := quartz.NewMock(t) + timerTrap := mockClock.Trap().NewTimer("chatd", "workspace-mcp-prime") + t.Cleanup(timerTrap.Close) + + server := &Server{ + db: db, + logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + clock: mockClock, + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: time.Second, + agentConnFn: func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { + return conn, func() {}, nil + }, + } + + chatStateMu := &sync.Mutex{} + currentChat := chat + workspaceCtx := turnWorkspaceContext{ + server: server, + chatStateMu: chatStateMu, + currentChat: ¤tChat, + loadChatSnapshot: func(context.Context, uuid.UUID) (database.Chat, error) { return chat, nil }, + } + t.Cleanup(workspaceCtx.close) + + primerCtx, primerCancel := context.WithCancel(ctx) + t.Cleanup(primerCancel) + + done := make(chan struct{}) + go func() { + defer close(done) + server.primeWorkspaceMCPCache(primerCtx, server.logger, chat.ID, &workspaceCtx) + }() + + // Let the primer arm at least one retry timer so we know it is + // blocked in the select. Canceling before this would race with + // the loop entering the retry path. + call := timerTrap.MustWait(ctx) + call.MustRelease(ctx) + + primerCancel() + + select { + case <-done: + case <-ctx.Done(): + t.Fatal("primer did not exit after context cancellation") + } + + _, ok := server.workspaceMCPToolsCache.Load(chat.ID) + require.False(t, ok, "primer must not cache anything when canceled") +} diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index ed90c601d5..88504f4aed 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -13,6 +13,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "slices" "strings" "sync" "sync/atomic" @@ -11664,3 +11665,214 @@ func TestRunChat_WorkspaceMCPDiscoveryAfterMidTurnCreateWorkspace(t *testing.T) "second call (after create_workspace) must advertise the workspace MCP tool: "+ "this is the fix for mid-turn workspace MCP discovery") } + +// TestRunChat_PrepareToolsRetriesAfterEmptyDiscovery guards the +// regression on the workspaceMCPDiscovered flag flip: the prior +// implementation set the flag to true before calling +// discoverWorkspaceMCPTools, so a single empty result permanently +// blocked retries within the turn. The fix sets the flag to true +// only after a non-empty discovery, so subsequent PrepareTools +// invocations keep retrying until tools appear. +// +// Scenario: create_workspace binds a workspace mid-turn. The first +// few ListMCPTools calls return empty (simulating the agent's MCP +// Connect still racing with agent startup); a later call returns +// the workspace MCP tool. The chat takes multiple steps before +// finishing, and we assert that one of the post-create_workspace +// streamed model calls advertises the workspace tool. +func TestRunChat_PrepareToolsRetriesAfterEmptyDiscovery(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + var ( + requestsMu sync.Mutex + requests []recordedOpenAIRequest + ) + + workspaceToolName := "workspace-empty-retry-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() + + // Step 1: trigger create_workspace. + if callIdx == 1 { + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk("create_workspace", workspaceCreateToolArgsJSON), + ) + } + // Step 2..N-1: emit empty text to keep the chatloop running so + // PrepareTools fires on each step. The chatloop ends a turn + // when the model returns a non-empty assistant message with no + // tool calls; an empty text chunk would terminate the turn, so + // we attach a dummy tool call to force another step. Use the + // LS tool because it exists for all workspaces and is cheap. + if callIdx < 6 { + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk("ls", `{"path":"/tmp"}`), + ) + } + // Final step: finish the chat. + 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-empty-retry-mcp", + Name: workspaceToolName, + Description: "workspace echo tool", + Schema: map[string]any{ + "input": map[string]any{"type": "string"}, + }, + Required: []string{"input"}, + }}, + } + + // First two ListMCPTools calls return empty (no error). One is the + // primer goroutine's only attempt before its retry timer fires; + // the other is PrepareTools on the first post-create_workspace + // step. The third and later calls return the workspace tool. The + // assertion below requires that a post-create_workspace step + // eventually advertises the tool, which can only happen if the + // PrepareTools callback retries discovery on subsequent steps. + var listCalls atomic.Int32 + 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()).DoAndReturn( + func(context.Context) (workspacesdk.ListMCPToolsResponse, error) { + n := listCalls.Add(1) + if n <= 2 { + return workspacesdk.ListMCPToolsResponse{}, nil + } + 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-empty-retry", + 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), 3, + "expected at least three streamed model calls; chat must run past the empty discovery") + + // The first call has no workspace yet; the second call is the + // first post-create_workspace step which sees an empty + // ListMCPTools result. By the third (or later) call PrepareTools + // must have retried discovery, so at least one post-step request + // must advertise the workspace tool. Without the + // workspaceMCPDiscovered flag-flip fix the flag would have been + // set true on the failed first attempt and no subsequent step + // would have re-attempted discovery. + sawWorkspaceTool := false + for i := 2; i < len(recorded); i++ { + if slices.Contains(recorded[i].Tools, workspaceToolName) { + sawWorkspaceTool = true + break + } + } + require.True(t, sawWorkspaceTool, + "PrepareTools must retry workspace MCP discovery on subsequent "+ + "steps; without the fix the first empty result would "+ + "permanently block retries within the turn") +}