From 3925d3941b242e36f89fe72ec4f0061e17664efd Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Fri, 8 May 2026 17:49:10 +0300 Subject: [PATCH] fix(coderd/x/chatd): wait long enough for cold-start workspace MCP discovery (#25035) The 5s timeout cancelled cold-start ListMCPTools calls before the agent's 30s connectTimeout could settle, so workspace MCP tools never reached the LLM. Bump to 35s and scope to ListMCPTools only. --- coderd/x/chatd/chatd.go | 21 ++++---- coderd/x/chatd/chatd_test.go | 99 ++++++++++++++++++++++++++++++++++++ 2 files changed, 109 insertions(+), 11 deletions(-) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index ff5ba1f3f4..54dfd43963 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -67,7 +67,10 @@ const ( planPathLookupTimeout = 5 * time.Second instructionCacheTTL = 5 * time.Minute workspaceDialValidationDelay = 5 * time.Second - workspaceMCPDiscoveryTimeout = 5 * time.Second + // Must exceed agent/x/agentmcp.connectTimeout (30s) so a + // cold-start agent's first MCP reload can settle before + // chatd gives up. + workspaceMCPDiscoveryTimeout = 35 * time.Second turnSummaryWriteTimeout = 5 * time.Second // defaultDialTimeout matches the timeout used by ~8 other // server-side AgentConn callers. @@ -6658,13 +6661,7 @@ func (p *Server) runChat( } // Cache miss, agent changed, or no cache: validate // that the workspace still has a live agent before // attempting a dial. - workspaceMCPCtx, cancel := context.WithTimeout( - ctx, - workspaceMCPDiscoveryTimeout, - ) - defer cancel() - - _, _, agentErr = workspaceCtx.workspaceAgentIDForConn(workspaceMCPCtx) + _, _, agentErr = workspaceCtx.workspaceAgentIDForConn(ctx) if agentErr != nil { if xerrors.Is(agentErr, errChatHasNoWorkspaceAgent) { p.workspaceMCPToolsCache.Delete(chat.ID) @@ -6676,13 +6673,15 @@ func (p *Server) runChat( } // List workspace MCP tools via the agent conn. - conn, connErr := workspaceCtx.getWorkspaceConn(workspaceMCPCtx) + conn, connErr := workspaceCtx.getWorkspaceConn(ctx) if connErr != nil { logger.Warn(ctx, "failed to get workspace conn for MCP tools", slog.Error(connErr)) return nil } - toolsResp, listErr := conn.ListMCPTools(workspaceMCPCtx) + 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)) @@ -6694,7 +6693,7 @@ func (p *Server) runChat( // caching an empty list would hide tools // permanently. if len(toolsResp.Tools) > 0 { - if agent, agentErr := workspaceCtx.getWorkspaceAgent(workspaceMCPCtx); agentErr == nil { + if agent, agentErr := workspaceCtx.getWorkspaceAgent(ctx); agentErr == nil { p.workspaceMCPToolsCache.Store(chat.ID, &cachedWorkspaceMCPTools{ agentID: agent.ID, tools: toolsResp.Tools, diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 03bf22c092..34ba65218e 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -11157,3 +11157,102 @@ func TestRecoverStaleChatsWaitingPropagatesSynthError(t *testing.T) { } } } + +// Regression for the cold-start race: chatd must wait long enough +// for ListMCPTools to return after the agent's MCP reload settles. +func TestRunChat_WorkspaceMCPDiscoveryWaitsForSlowAgent(t *testing.T) { + t.Parallel() + + const slowAgentMCPListDelay = 7 * time.Second + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + var ( + requestsMu sync.Mutex + requests []recordedOpenAIRequest + ) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + + requestsMu.Lock() + requests = append(requests, recordOpenAIRequest(req)) + requestsMu.Unlock() + + return chattest.OpenAIStreamingResponse( + chattest.OpenAITextChunks("done")..., + ) + }) + + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + workspaceToolName := "workspace-slow-mcp__echo" + workspaceToolsResp := workspacesdk.ListMCPToolsResponse{ + Tools: []workspacesdk.MCPToolInfo{{ + ServerName: "workspace-slow-mcp", + Name: workspaceToolName, + Description: "Slow 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() + // Honor ctx so the goroutine exits if chatd cancels. + mockConn.EXPECT().ListMCPTools(gomock.Any()). + DoAndReturn(func(ctx context.Context) (workspacesdk.ListMCPToolsResponse, error) { + select { + case <-time.After(slowAgentMCPListDelay): + return workspaceToolsResp, nil + case <-ctx.Done(): + return workspacesdk.ListMCPToolsResponse{}, ctx.Err() + } + }).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() + + 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 + } + }) + + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "workspace-mcp-slow-agent", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("List the workspace MCP tools."), + }, + }) + 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.Len(t, recorded, 1, "expected exactly one streamed model call") + require.Contains(t, recorded[0].Tools, workspaceToolName, + "workspace MCP tool should reach the LLM once chatd's discovery "+ + "timeout exceeds the agent's MCP reload time") +}