From 521c383f6bd845664bbff0cf47eedb3ddbf9e03f Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Sun, 16 Aug 2026 20:30:32 +0200 Subject: [PATCH] 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 | 3 + coderd/database/queries/workspaceagents.sql | 1 + coderd/exp_chats.go | 60 ++++++++++++---- coderd/exp_chats_internal_test.go | 64 +++++++++++++++-- site/src/api/queries/chats.test.ts | 42 +++++++++++ site/src/api/queries/chats.ts | 10 ++- .../AgentsPage/AgentChatPage.stories.tsx | 71 +++++++++++++++++++ .../pages/AgentsPage/AgentChatPage.test.ts | 66 +++++++++++++++++ site/src/pages/AgentsPage/AgentChatPage.tsx | 51 +++++++++++++ 9 files changed, 345 insertions(+), 23 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 31c1331426..f1029f44a1 100644 --- a/coderd/database/queries.sql.go +++ b/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, diff --git a/coderd/database/queries/workspaceagents.sql b/coderd/database/queries/workspaceagents.sql index e5280252da..1d3ff2c4ac 100644 --- a/coderd/database/queries/workspaceagents.sql +++ b/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 diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index eed66f53e3..f1b07c3f2b 100644 --- a/coderd/exp_chats.go +++ b/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) diff --git a/coderd/exp_chats_internal_test.go b/coderd/exp_chats_internal_test.go index 85ad42dbdd..da8ba38286 100644 --- a/coderd/exp_chats_internal_test.go +++ b/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) }) } diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index fde7d34325..05d0397d5d 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -2486,6 +2486,48 @@ describe("mergeWatchedChatSummary", () => { ).toBe(context); }); + it("keeps the repaired build_id when the event snapshot carries a stale binding", () => { + const cachedChat = makeChat("chat-1", { + workspace_id: "workspace-1", + agent_id: "agent-new", + build_id: "build-new", + updated_at: "2025-01-01T00:00:00.000Z", + }); + const watchedChat = makeChat("chat-1", { + workspace_id: "workspace-1", + agent_id: "agent-old", + build_id: "build-old", + updated_at: "2025-01-01T00:05:00.000Z", + }); + + const merged = mergeWatchedChatSummary(cachedChat, watchedChat, { + eventKind: "diff_status_change", + }); + expect(merged.build_id).toBe("build-new"); + expect(merged.agent_id).toBe("agent-new"); + }); + + it("adopts a fresh build_id when the event snapshot agrees on the agent", () => { + const cachedChat = makeChat("chat-1", { + workspace_id: "workspace-1", + agent_id: "agent-1", + build_id: "build-old", + updated_at: "2025-01-01T00:00:00.000Z", + }); + const watchedChat = makeChat("chat-1", { + workspace_id: "workspace-1", + agent_id: "agent-1", + build_id: "build-new", + updated_at: "2025-01-01T00:05:00.000Z", + }); + + expect( + mergeWatchedChatSummary(cachedChat, watchedChat, { + eventKind: "status_change", + }).build_id, + ).toBe("build-new"); + }); + it("merges fresh status updates without clobbering a newer title snapshot", () => { const cachedChat = makeChat("chat-1", { status: "waiting", diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index a092b55249..e323151392 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -551,9 +551,13 @@ export const mergeWatchedChatSummary = ( const nextWorkspaceId = isFreshEnough ? (watchedChat.workspace_id ?? cachedChat.workspace_id) : cachedChat.workspace_id; - const nextBuildId = isFreshEnough - ? (watchedChat.build_id ?? cachedChat.build_id) - : cachedChat.build_id; + // Single-chat reads repair agent/build bindings response-only, so watch + // events can replay stale DB pairs. Adopting build_id with a mismatched + // agent would split the repaired pair because merge never adopts agent_id. + const nextBuildId = + isFreshEnough && watchedChat.agent_id === cachedChat.agent_id + ? (watchedChat.build_id ?? cachedChat.build_id) + : cachedChat.build_id; // All event types carry the current model config from the DB. const nextLastModelConfigId = isFreshEnough ? watchedChat.last_model_config_id diff --git a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx index b220e2ca06..cc6b3784b6 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx @@ -34,6 +34,7 @@ import { MockOrganizationMember2, MockUserOwner, MockWorkspace, + MockWorkspaceAgent, mockApiError, } from "#/testHelpers/entities"; import { @@ -2205,6 +2206,76 @@ export const SidebarWithSingleRepo: Story = { }, }, }; + +const rebuiltWorkspaceAgent: TypesGen.WorkspaceAgent = { + ...MockWorkspaceAgent, + id: "rebuilt-agent-1", +}; +const rebuiltWorkspace: TypesGen.Workspace = { + ...mockWorkspace, + latest_build: { + ...mockWorkspace.latest_build, + id: "rebuilt-build-1", + resources: [ + { + ...mockWorkspace.latest_build.resources[0], + agents: [rebuiltWorkspaceAgent], + }, + ], + }, +}; +const rebuildRecoveryChat: TypesGen.Chat = { + id: CHAT_ID, + ...baseChatFields, + agent_id: "stale-agent-1", + title: "Rebuild recovery", + status: "waiting", +}; + +export const RecoversSidebarAfterWorkspaceRebuild: Story = { + beforeEach: () => { + localStorage.setItem(RIGHT_PANEL_OPEN_KEY, "true"); + spyOn(API.experimental, "getChat").mockResolvedValue({ + ...rebuildRecoveryChat, + agent_id: rebuiltWorkspaceAgent.id, + }); + return () => localStorage.removeItem(RIGHT_PANEL_OPEN_KEY); + }, + parameters: { + queries: [ + ...withoutQuery( + buildQueries( + rebuildRecoveryChat, + { messages: [], queued_messages: [], has_more: false }, + { diffUrl: undefined }, + ), + workspaceByIdKey(mockWorkspace.id), + ), + { key: workspaceByIdKey(mockWorkspace.id), data: rebuiltWorkspace }, + ], + webSocket: { + "watch-ws": [ + { + event: "message", + data: JSON.stringify({ + type: "data", + data: rebuiltWorkspace, + } satisfies TypesGen.ServerSentEvent), + }, + ], + }, + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const terminalTab = await canvas.findByRole( + "tab", + { name: "Terminal" }, + { timeout: 5000 }, + ); + expect(terminalTab).toBeVisible(); + }, +}; + /** * Streaming reasoning part via WebSocket, renders inline text. */ diff --git a/site/src/pages/AgentsPage/AgentChatPage.test.ts b/site/src/pages/AgentsPage/AgentChatPage.test.ts index 46331761ad..eea803a843 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.test.ts +++ b/site/src/pages/AgentsPage/AgentChatPage.test.ts @@ -23,6 +23,7 @@ import { draftInputStorageKeyPrefix, getPersistedDraftInputValue, getWorkspaceOptionsWithLinkedWorkspace, + isChatAgentBindingUnresolved, isWatchedWorkspaceViewUnchanged, reconcilePromotedQueueHead, restoreOptimisticRequestSnapshot, @@ -1437,4 +1438,69 @@ describe("isWatchedWorkspaceViewUnchanged", () => { ), ).toBe(false); }); + + it("is false when the latest build changes", () => { + const next: Workspace = { + ...MockWorkspace, + latest_build: { ...MockWorkspace.latest_build, id: "new-build-id" }, + }; + + expect( + isWatchedWorkspaceViewUnchanged( + MockWorkspace, + next, + MockWorkspaceAgent.id, + ), + ).toBe(false); + }); +}); + +describe("isChatAgentBindingUnresolved", () => { + it("is true when the bound agent is missing from the running build", () => { + expect(isChatAgentBindingUnresolved(MockWorkspace, "stale-agent-id")).toBe( + true, + ); + }); + + it("is true when the chat has no binding yet", () => { + expect(isChatAgentBindingUnresolved(MockWorkspace, undefined)).toBe(true); + }); + + it("is false when the bound agent resolves", () => { + expect( + isChatAgentBindingUnresolved(MockWorkspace, MockWorkspaceAgent.id), + ).toBe(false); + }); + + it("is false when the workspace is not running", () => { + const stopped: Workspace = { + ...MockWorkspace, + latest_build: { ...MockWorkspace.latest_build, status: "stopped" }, + }; + + expect(isChatAgentBindingUnresolved(stopped, "stale-agent-id")).toBe(false); + }); + + it("is false when the running build has no agents", () => { + const noAgents: Workspace = { + ...MockWorkspace, + latest_build: { + ...MockWorkspace.latest_build, + resources: MockWorkspace.latest_build.resources.map((resource) => ({ + ...resource, + agents: [], + })), + }, + }; + + expect(isChatAgentBindingUnresolved(noAgents, "stale-agent-id")).toBe( + false, + ); + }); + + it("is false while the workspace is loading", () => { + expect(isChatAgentBindingUnresolved(undefined, "stale-agent-id")).toBe( + false, + ); + }); }); diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 705b47c8ae..da472f7d64 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -73,6 +73,7 @@ import { isMobileViewport } from "#/utils/mobile"; import { pageTitle } from "#/utils/page"; import { rewriteLocalhostURL } from "#/utils/portForward"; import { createReconnectingWebSocket } from "#/utils/reconnectingWebSocket"; +import { getWorkspaceAgents } from "#/utils/workspace"; import { AgentChatPageErrorView } from "./AgentChatPageErrorView"; import { AgentChatPageLoadingView, @@ -130,6 +131,9 @@ import { export const RIGHT_PANEL_OPEN_KEY = "agents.right-panel-open"; const lastModelConfigIDStorageKey = "agents.last-model-config-id"; + +const AGENT_BINDING_REPAIR_POLL_MS = 30_000; + class CompactCommandPendingError extends Error {} /** @internal Exported for testing. */ @@ -445,6 +449,7 @@ export const isWatchedWorkspaceViewUnchanged = ( const prevApps = prevAgent?.apps ?? []; const nextApps = nextAgent?.apps ?? []; return ( + prev.latest_build.id === next.latest_build.id && prev.latest_build.status === next.latest_build.status && prev.health.healthy === next.health.healthy && prev.name === next.name && @@ -464,6 +469,24 @@ export const isWatchedWorkspaceViewUnchanged = ( ); }; +/** + * True when a running workspace has agents but the chat's agent ID is absent + * from the latest build (stale after a rebuild, or not yet persisted). Chat + * reads can return a repaired ID, so callers should refetch the chat. + * + * @internal Exported for testing. + */ +export const isChatAgentBindingUnresolved = ( + workspace: TypesGen.Workspace | undefined, + chatAgentId: string | undefined, +): boolean => { + if (!workspace || workspace.latest_build.status !== "running") { + return false; + } + const agents = getWorkspaceAgents(workspace); + return agents.length > 0 && !agents.some((agent) => agent.id === chatAgentId); +}; + const buildAttachmentMediaTypes = ( attachments?: readonly PendingAttachment[], ): ReadonlyMap | undefined => { @@ -913,6 +936,21 @@ const AgentChatPage: FC = () => { const chatQuery = useQuery({ ...chat(agentId ?? ""), enabled: Boolean(agentId), + // Poll while the binding is unresolved: repair happens on chat reads + // and watch events cannot be relied on for retries because an idle + // workspace publishes none. + refetchInterval: ({ state }) => { + const workspaceId = state.data?.workspace_id; + const workspace = workspaceId + ? queryClient.getQueryData( + workspaceByIdKey(workspaceId), + ) + : undefined; + return isChatAgentBindingUnresolved(workspace, state.data?.agent_id) + ? AGENT_BINDING_REPAIR_POLL_MS + : false; + }, + refetchIntervalInBackground: false, }); const chatMessagesQuery = useInfiniteQuery({ ...chatMessagesForInfiniteScroll(agentId ?? ""), @@ -995,6 +1033,7 @@ const AgentChatPage: FC = () => { chatModelsQuery.data, ); + const agentBindingRefetchKeyRef = useRef(undefined); // Subscribe to live workspace updates so that agent status changes // (e.g. connected/disconnected) are reflected without a page refresh. const applyWatchedWorkspaceUpdate = useEffectEvent( @@ -1015,6 +1054,18 @@ const AgentChatPage: FC = () => { return next; }, ); + // Refetch once per chat/build/binding key for immediate repair + // after a rebuild; the chat query's refetchInterval owns retries + // when repair fails, so the latch never blocks recovery. + if (!agentId || !isChatAgentBindingUnresolved(next, chatAgentId)) { + return; + } + const refetchKey = `${agentId}:${next.latest_build.id}:${chatAgentId ?? ""}`; + if (agentBindingRefetchKeyRef.current === refetchKey) { + return; + } + agentBindingRefetchKeyRef.current = refetchKey; + void invalidateChatEntity(queryClient, agentId); }, ); useEffect(() => {