fix(coderd/x/chatd): prime workspace MCP cache after create/start (#25298)

## Problem Mid-turn workspace MCP discovery was broken when an agent was still cold-starting. `PrepareTools` in `chatd.go` flipped `workspaceMCPDiscovered = true` *before* calling `discoverWorkspaceMCPTools`, so a failed discovery attempt permanently blocked retries within the turn. Customer-reported repro: - New chat with no pre-selected workspace. - LLM calls `create_workspace` mid-turn at `23:35:05`. - `PrepareTools` fires, dials the agent with a 30s timeout, dial times out at `23:38:15`, `discoverWorkspaceMCPTools` returns empty. - Agent connects at `23:38:29`, 14 seconds later. - `workspaceMCPDiscovered` was already true, so `PrepareTools` never retried for the rest of the turn. MCP tools only appeared on the next user message. A naive retry loop in `PrepareTools` would also miss the bigger picture: a workspace boot can take several minutes (EC2 cold start, 10 min startup scripts), and the chatloop only gets a chance to call `PrepareTools` between LLM steps. ## Fix Do the workspace MCP discovery from inside the tool that already waits for the agent. `chattool.CreateWorkspace` and `chattool.StartWorkspace` call `waitForAgentReady`, which has a 2 min agent-online budget plus a 10 min startup-script budget. By the time they fire `OnChatUpdated`, the agent is `Ready`. The chatd `onChatUpdated` callback now launches an async `primeWorkspaceMCPCache` goroutine on every bind that has a valid workspace ID: - The primer calls `discoverWorkspaceMCPTools` until it returns a non-empty list or `workspaceMCPPrimeMaxWait` (30s) elapses, with a 2s backoff between attempts. The bounded wait handles the short race between agent-online and the agent's MCP `Connect` settling. - The primer runs asynchronously so the tool itself never blocks. Some templates simply do not advertise MCP tools, in which case the primer would otherwise spend its full budget for nothing. - The primer shares the chat `ctx` (not a detached one) so it is canceled together with the chat. A dangling primer would re-dial the workspace conn after `runChat`'s deferred `workspaceCtx.close()` and leak that conn. - `inflight.Add(1)` ensures server shutdown still waits for any in-progress primer. - `PrepareTools` is simplified back to a single discovery call. It now only sets `workspaceMCPDiscovered = true` on success, so an empty result no longer permanently blocks discovery within the turn. The cache hit warmed by the primer makes that call cheap in the common case; the dial fallback handles the rare cache miss. ## Tests All in `coderd/x/chatd/chatd_internal_test.go`: - `TestPrimeWorkspaceMCPCache_SuccessOnFirstAttempt` — single `ListMCPTools` call returning tools populates the cache. - `TestPrimeWorkspaceMCPCache_RetriesUntilToolsAppear` — first call empty, second returns tools; primer retries past the backoff and writes the cache. Uses `quartz.Mock.Trap` on `NewTimer`. - `TestPrimeWorkspaceMCPCache_GivesUpAfterDeadline` — `ListMCPTools` always empty; primer stops at `workspaceMCPPrimeMaxWait` and refuses to cache the empty result so PrepareTools can retry on the next step. The existing integration test `TestRunChat_WorkspaceMCPDiscoveryAfterMidTurnCreateWorkspace` continues to pass and now also exercises the async-primer path end-to-end via the create_workspace tool. ``` go test ./coderd/x/chatd/... -count=1 go test ./coderd/x/chatd/ -race -count=1 make pre-commit ``` <details> <summary>Design notes</summary> - The first iteration of this PR added retry+cooldown+failure-cap logic inside `PrepareTools`. It worked for the customer's ~30s race window but did not help workspaces that take several minutes to boot, because `PrepareTools` only fires between LLM steps. Reviewer pointed out the right place to handle this is the tool itself; the current implementation does that. - Why async: a primer that ran synchronously inside the `OnChatUpdated` callback blocked the create_workspace tool from returning for up to `workspaceMCPPrimeMaxWait`, which broke `TestCreateWorkspaceTool_EndToEnd` and would hurt any template that does not expose MCP tools. Decoupling lets the tool return immediately and lets the primer warm the cache concurrently with the next LLM step. - Why share the chat `ctx` rather than `context.WithoutCancel(ctx)` (the title-generation pattern): the primer touches `workspaceCtx.getWorkspaceConn`, which `runChat`'s deferred `workspaceCtx.close()` invalidates. A detached primer outliving the chat would dial a fresh conn and leak it. - The constant naming distinguishes `workspaceMCPDiscoveryTimeout` (35s per-call dial budget, unchanged from #25169) from `workspaceMCPPrimeMaxWait` (30s total budget for the post-ready primer loop) and `workspaceMCPPrimeRetryInterval` (2s between empty-result retries). </details> Follow-up to #25169. --- _This pull request was generated by Coder Agents._
2026-09-24 15:04:27 +08:00 · 2026-05-18 07:55:56 -04:00
parent 7a985f8830
commit 159089686a
3 changed files with 802 additions and 8 deletions
@@ -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:      &currentChat,
 		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 {
@@ -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:      &currentChat,
+		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:      &currentChat,
+		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:      &currentChat,
+		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:      &currentChat,
+		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")
+}
@@ -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")
+}