diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index e15e05da2cb..4149388d6f2 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -34028,6 +34028,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 @@ -34054,6 +34055,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"` } @@ -34068,6 +34070,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 e5280252da0..1d3ff2c4ac8 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 72196d418b7..3f83fa00708 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 85ad42dbdda..da8ba38286c 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 fde7d34325c..05d0397d5d7 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 a092b552494..e3231513920 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 581d9b5ba4b..dc4fadf5b15 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx @@ -35,6 +35,7 @@ import { MockOrganizationMember2, MockUserOwner, MockWorkspace, + MockWorkspaceAgent, mockApiError, } from "#/testHelpers/entities"; import { @@ -2208,6 +2209,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 676f315e937..68add84a6f1 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, @@ -1460,4 +1461,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 4fd0b0d09a3..1df0687c03d 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -69,6 +69,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, @@ -126,6 +127,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. */ @@ -450,6 +454,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 && @@ -469,6 +474,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 => { @@ -920,6 +943,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 ?? ""), @@ -1002,6 +1040,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( @@ -1022,6 +1061,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(() => {