fix: repair stale chat agent bindings after workspace rebuild (#28152)

## Problem

When a chat is bound to a workspace, chatd persists `chats.agent_id`
pointing at a specific workspace agent, and it only rebinds on the next
chat turn. A workspace stop/start creates a new agent with a new ID in
the latest build, so the chat page resolves the stale agent ID to
`undefined` and the right sidebar silently drops Terminal, Desktop,
Browser, apps, and ports even though the workspace is running. The
existing read-time enrichment only filled nil agent IDs and skipped
stale non-nil ones, so refreshing did not help until the user sent
another message.

## Fix

- `coderd/exp_chats.go`: single-chat reads now repair agent IDs that no
longer resolve in the workspace's latest build, using the same
`agentselect.FindChatAgent` selection chatd uses. A repaired binding
also carries the latest build's ID so the response never pairs the new
agent with the previous build. Bindings that still resolve are
preserved, and repair stays best-effort and response-only (no
write-on-read). List reads keep the previous nil-fill-only behavior
because validating existing bindings would cost a per-workspace
authorization lookup per listed chat.
- `site/src/pages/AgentsPage/AgentChatPage.tsx`: the workspace watch
update handler detects when a running workspace's latest build no longer
contains the chat's bound agent and invalidates the chat query once per
chat/build/binding key for immediate recovery, and the chat query polls
every 30 seconds while the binding remains unresolved so a transiently
failed repair retries even when an idle workspace publishes no further
watch events. The watch stream replays the current workspace on every
(re)connect, so this covers rebuilds that happen while the page is open
or disconnected; page loads are covered by the server-side repair. The
workspace-watcher bailout now also keys on `latest_build.id` so a
rebuild propagates while the page is open.
- `site/src/api/queries/chats.ts`: chat watch events replay the
persisted (pre-repair) binding, so the summary merge adopts a snapshot's
`build_id` only when the snapshot agrees on `agent_id`, keeping the
repaired agent/build pair atomic in the caches.

## Testing

- `go test ./coderd -run TestEnrichChatAgentIDs` covering repair,
keep-valid, selection-error, list-mode-skips-bound, and no-workspaces
cases.
- Storybook interaction story `RecoversSidebarAfterWorkspaceRebuild`
exercising the watch-event to chat-refetch to sidebar-recovery flow
(verified red without the invalidation, green with it).
- `pnpm test AgentChatPage.test.ts` covering the binding-resolution
predicate.

> Mux created this PR on Mike's behalf.

coderd/database/queries.sql.go

@@ -34060,6 +34060,7 @@ func (q *sqlQuerier) GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx context.Co
 const getWorkspaceAgentsInLatestBuildByWorkspaceIDs = `-- name: GetWorkspaceAgentsInLatestBuildByWorkspaceIDs :many
 SELECT
 	workspace_builds.workspace_id,
+	workspace_builds.id AS build_id,
 	workspace_agents.id, workspace_agents.created_at, workspace_agents.updated_at, workspace_agents.name, workspace_agents.first_connected_at, workspace_agents.last_connected_at, workspace_agents.disconnected_at, workspace_agents.resource_id, workspace_agents.auth_token, workspace_agents.auth_instance_id, workspace_agents.architecture, workspace_agents.environment_variables, workspace_agents.operating_system, workspace_agents.instance_metadata, workspace_agents.resource_metadata, workspace_agents.directory, workspace_agents.version, workspace_agents.last_connected_replica_id, workspace_agents.connection_timeout_seconds, workspace_agents.troubleshooting_url, workspace_agents.motd_file, workspace_agents.lifecycle_state, workspace_agents.expanded_directory, workspace_agents.logs_length, workspace_agents.logs_overflowed, workspace_agents.started_at, workspace_agents.ready_at, workspace_agents.subsystems, workspace_agents.display_apps, workspace_agents.api_version, workspace_agents.display_order, workspace_agents.parent_id, workspace_agents.api_key_scope, workspace_agents.deleted
 FROM
 	workspace_agents
@@ -34086,6 +34087,7 @@ WHERE
 
 type GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow struct {
 	WorkspaceID    uuid.UUID      `db:"workspace_id" json:"workspace_id"`
+	BuildID        uuid.UUID      `db:"build_id" json:"build_id"`
 	WorkspaceAgent WorkspaceAgent `db:"workspace_agent" json:"workspace_agent"`
 }
 
@@ -34100,6 +34102,7 @@ func (q *sqlQuerier) GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(ctx context.C
 		var i GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow
 		if err := rows.Scan(
 			&i.WorkspaceID,
+			&i.BuildID,
 			&i.WorkspaceAgent.ID,
 			&i.WorkspaceAgent.CreatedAt,
 			&i.WorkspaceAgent.UpdatedAt,

coderd/database/queries/workspaceagents.sql

@@ -340,6 +340,7 @@ WHERE
 -- name: GetWorkspaceAgentsInLatestBuildByWorkspaceIDs :many
 SELECT
 	workspace_builds.workspace_id,
+	workspace_builds.id AS build_id,
 	sqlc.embed(workspace_agents)
 FROM
 	workspace_agents

coderd/exp_chats.go

@@ -498,28 +498,46 @@ func (api *API) listChats(rw http.ResponseWriter, r *http.Request) {
 	}
 
 	sdkChats := db2sdk.ChatRowsWithChildren(chatRows, childRows, diffStatusesByChatID)
-	api.enrichChatWithWorkspaceAgentIDs(ctx, sdkChats)
+	api.enrichChatsWithMissingAgentIDs(ctx, sdkChats)
 	httpapi.Write(ctx, rw, http.StatusOK, sdkChats)
 }
 
-// enrichChatWithWorkspaceAgentIDs fills missing AgentIDs for chats with a bound
-// workspace, since chatd persists the binding lazily. Best-effort and
-// response-only; on error the field stays null.
-func (api *API) enrichChatWithWorkspaceAgentIDs(ctx context.Context, chats []codersdk.Chat) {
-	missingChats := make([]*codersdk.Chat, 0, len(chats))
+// enrichChatsWithMissingAgentIDs skips existing bindings on list reads to avoid
+// one authorization check per bound workspace.
+func (api *API) enrichChatsWithMissingAgentIDs(ctx context.Context, chats []codersdk.Chat) {
+	api.enrichChatAgentIDs(ctx, chats, func(chat *codersdk.Chat) bool {
+		return chat.AgentID == nil
+	})
+}
+
+// repairChatAgentIDs handles stale bindings left by workspace rebuilds. List
+// reads skip this work to avoid authorization checks for bound workspaces.
+func (api *API) repairChatAgentIDs(ctx context.Context, chats []codersdk.Chat) {
+	api.enrichChatAgentIDs(ctx, chats, func(*codersdk.Chat) bool {
+		return true
+	})
+}
+
+// enrichChatAgentIDs performs best-effort response-only updates.
+func (api *API) enrichChatAgentIDs(ctx context.Context, chats []codersdk.Chat, shouldEnrich func(*codersdk.Chat) bool) {
+	candidateChats := make([]*codersdk.Chat, 0, len(chats))
 	var workspaceIDs []uuid.UUID
-	addMissing := func(chat *codersdk.Chat) {
-		if chat.AgentID == nil && chat.WorkspaceID != nil {
-			missingChats = append(missingChats, chat)
-			workspaceIDs = append(workspaceIDs, *chat.WorkspaceID)
+	addCandidate := func(chat *codersdk.Chat) {
+		if chat.WorkspaceID == nil || !shouldEnrich(chat) {
+			return
 		}
+		candidateChats = append(candidateChats, chat)
+		workspaceIDs = append(workspaceIDs, *chat.WorkspaceID)
 	}
 	for i := range chats {
-		addMissing(&chats[i])
+		addCandidate(&chats[i])
 		for j := range chats[i].Children {
-			addMissing(&chats[i].Children[j])
+			addCandidate(&chats[i].Children[j])
 		}
 	}
+	if len(candidateChats) == 0 {
+		return
+	}
 
 	slices.SortFunc(workspaceIDs, func(a, b uuid.UUID) int {
 		return cmp.Compare(a.String(), b.String())
@@ -531,8 +549,10 @@ func (api *API) enrichChatWithWorkspaceAgentIDs(ctx context.Context, chats []cod
 	}
 
 	agentsByWorkspace := make(map[uuid.UUID][]database.WorkspaceAgent)
+	latestBuildIDs := make(map[uuid.UUID]uuid.UUID)
 	for _, row := range rows {
 		agentsByWorkspace[row.WorkspaceID] = append(agentsByWorkspace[row.WorkspaceID], row.WorkspaceAgent)
+		latestBuildIDs[row.WorkspaceID] = row.BuildID
 	}
 	agentIDs := make(map[uuid.UUID]uuid.UUID, len(agentsByWorkspace))
 	for workspaceID, agents := range agentsByWorkspace {
@@ -544,10 +564,22 @@ func (api *API) enrichChatWithWorkspaceAgentIDs(ctx context.Context, chats []cod
 		agentIDs[workspaceID] = agent.ID
 	}
 
-	for _, chat := range missingChats {
+	for _, chat := range candidateChats {
+		// Preserve bindings that still resolve in the latest build instead
+		// of replacing them with the selected agent.
+		if chat.AgentID != nil && slices.ContainsFunc(
+			agentsByWorkspace[*chat.WorkspaceID],
+			func(agent database.WorkspaceAgent) bool { return agent.ID == *chat.AgentID },
+		) {
+			continue
+		}
 		if agentID, ok := agentIDs[*chat.WorkspaceID]; ok {
 			id := agentID
 			chat.AgentID = &id
+			// Pair the agent with its build so the response never mixes
+			// the latest build's agent with a previous build's ID.
+			buildID := latestBuildIDs[*chat.WorkspaceID]
+			chat.BuildID = &buildID
 		}
 	}
 }
@@ -1639,7 +1671,7 @@ func (api *API) getChat(rw http.ResponseWriter, r *http.Request) {
 	}
 
 	enriched := []codersdk.Chat{sdkChat}
-	api.enrichChatWithWorkspaceAgentIDs(ctx, enriched)
+	api.repairChatAgentIDs(ctx, enriched)
 	sdkChat = enriched[0]
 
 	httpapi.Write(ctx, rw, http.StatusOK, sdkChat)

coderd/exp_chats_internal_test.go

@@ -139,7 +139,7 @@ func TestGetChatCostFallsBackToParentChat(t *testing.T) {
 	require.Equal(t, int64(125), cost.TotalCostMicros)
 }
 
-func TestEnrichMissingChatAgentIDs(t *testing.T) {
+func TestEnrichChatAgentIDs(t *testing.T) {
 	t.Parallel()
 	newAPI := func(t *testing.T) (*API, *dbmock.MockStore) {
 		t.Helper()
@@ -149,9 +149,12 @@ func TestEnrichMissingChatAgentIDs(t *testing.T) {
 	}
 	workspaceID, otherWorkspaceID := uuid.New(), uuid.New()
 	rootAgentID, otherAgentID := uuid.New(), uuid.New()
+	latestBuildID, otherLatestBuildID := uuid.New(), uuid.New()
+	latestBuildIDs := map[uuid.UUID]uuid.UUID{workspaceID: latestBuildID, otherWorkspaceID: otherLatestBuildID}
 	row := func(workspaceID, id uuid.UUID, parentID uuid.NullUUID, name string) database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow {
 		return database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow{
 			WorkspaceID: workspaceID,
+			BuildID:     latestBuildIDs[workspaceID],
 			WorkspaceAgent: database.WorkspaceAgent{
 				ID:       id,
 				ParentID: parentID,
@@ -169,29 +172,78 @@ func TestEnrichMissingChatAgentIDs(t *testing.T) {
 			}, nil
 		}).Times(1)
 		chats := []codersdk.Chat{{WorkspaceID: &workspaceID, Children: []codersdk.Chat{{WorkspaceID: &workspaceID}}}, {WorkspaceID: &otherWorkspaceID}}
-		api.enrichChatWithWorkspaceAgentIDs(testutil.Context(t, testutil.WaitShort), chats)
+		api.enrichChatsWithMissingAgentIDs(testutil.Context(t, testutil.WaitShort), chats)
 		require.Equal(t, rootAgentID, *chats[0].AgentID)
 		require.Equal(t, rootAgentID, *chats[0].Children[0].AgentID)
 		require.Equal(t, otherAgentID, *chats[1].AgentID)
+		require.Equal(t, latestBuildID, *chats[0].BuildID)
+		require.Equal(t, latestBuildID, *chats[0].Children[0].BuildID)
+		require.Equal(t, otherLatestBuildID, *chats[1].BuildID)
 	})
 	t.Run("query error", func(t *testing.T) {
 		t.Parallel()
 		api, mDB := newAPI(t)
 		mDB.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(gomock.Any(), gomock.Any()).Return(nil, xerrors.New("boom"))
 		chats := []codersdk.Chat{{WorkspaceID: &workspaceID}, {WorkspaceID: &otherWorkspaceID}}
-		api.enrichChatWithWorkspaceAgentIDs(testutil.Context(t, testutil.WaitShort), chats)
+		api.enrichChatsWithMissingAgentIDs(testutil.Context(t, testutil.WaitShort), chats)
 		require.Nil(t, chats[0].AgentID)
 		require.Nil(t, chats[1].AgentID)
 	})
-	t.Run("selection error and skips bound or unbound", func(t *testing.T) {
+	t.Run("selection error keeps persisted values", func(t *testing.T) {
 		t.Parallel()
 		api, mDB := newAPI(t)
 		mDB.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(gomock.Any(), []uuid.UUID{workspaceID}).Return([]database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow{row(workspaceID, uuid.New(), uuid.NullUUID{UUID: rootAgentID, Valid: true}, "sub")}, nil)
 		bound := otherAgentID
-		chats := []codersdk.Chat{{}, {WorkspaceID: &workspaceID}, {WorkspaceID: &workspaceID, AgentID: &bound}}
-		api.enrichChatWithWorkspaceAgentIDs(testutil.Context(t, testutil.WaitShort), chats)
+		boundBuildID := uuid.New()
+		chats := []codersdk.Chat{{}, {WorkspaceID: &workspaceID}, {WorkspaceID: &workspaceID, AgentID: &bound, BuildID: &boundBuildID}}
+		api.repairChatAgentIDs(testutil.Context(t, testutil.WaitShort), chats)
 		require.Nil(t, chats[1].AgentID)
+		require.Nil(t, chats[1].BuildID)
 		require.Equal(t, bound, *chats[2].AgentID)
+		require.Equal(t, boundBuildID, *chats[2].BuildID)
+	})
+	t.Run("repairs stale and keeps valid bindings", func(t *testing.T) {
+		t.Parallel()
+		api, mDB := newAPI(t)
+		secondRootAgentID := uuid.New()
+		mDB.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(gomock.Any(), []uuid.UUID{workspaceID}).Return([]database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow{
+			row(workspaceID, rootAgentID, uuid.NullUUID{}, "a"),
+			row(workspaceID, secondRootAgentID, uuid.NullUUID{}, "b"),
+		}, nil)
+		stale, valid := uuid.New(), secondRootAgentID
+		staleBuildID, validBuildID := uuid.New(), uuid.New()
+		chats := []codersdk.Chat{
+			{WorkspaceID: &workspaceID, AgentID: &stale, BuildID: &staleBuildID},
+			{WorkspaceID: &workspaceID, AgentID: &valid, BuildID: &validBuildID},
+		}
+		api.repairChatAgentIDs(testutil.Context(t, testutil.WaitShort), chats)
+		require.Equal(t, rootAgentID, *chats[0].AgentID)
+		require.Equal(t, secondRootAgentID, *chats[1].AgentID)
+		require.Equal(t, latestBuildID, *chats[0].BuildID)
+		require.Equal(t, validBuildID, *chats[1].BuildID)
+	})
+	t.Run("list mode skips bound chats entirely", func(t *testing.T) {
+		t.Parallel()
+		api, mDB := newAPI(t)
+		mDB.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(gomock.Any(), []uuid.UUID{workspaceID}).Return([]database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow{
+			row(workspaceID, rootAgentID, uuid.NullUUID{}, "root"),
+		}, nil).Times(1)
+		stale := uuid.New()
+		chats := []codersdk.Chat{
+			{WorkspaceID: &workspaceID},
+			{WorkspaceID: &otherWorkspaceID, AgentID: &stale},
+		}
+		api.enrichChatsWithMissingAgentIDs(testutil.Context(t, testutil.WaitShort), chats)
+		require.Equal(t, rootAgentID, *chats[0].AgentID)
+		require.Equal(t, stale, *chats[1].AgentID)
+	})
+	t.Run("no bound workspaces skips the query", func(t *testing.T) {
+		t.Parallel()
+		api, _ := newAPI(t)
+		chats := []codersdk.Chat{{AgentID: &rootAgentID}, {}}
+		api.repairChatAgentIDs(testutil.Context(t, testutil.WaitShort), chats)
+		require.Equal(t, rootAgentID, *chats[0].AgentID)
+		require.Nil(t, chats[1].AgentID)
 	})
 }