diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 5774f549ca208..2c795a97a35a9 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -1,4 +1,4 @@ -import { QueryClient } from "react-query"; +import { QueryClient, QueryObserver } from "react-query"; import { describe, expect, it, vi } from "vitest"; import { API } from "#/api/api"; import type * as TypesGen from "#/api/typesGenerated"; @@ -8,9 +8,13 @@ import { SUCCESS_STATUSES, } from "#/pages/AgentsPage/components/RightPanel/DebugPanel/debugPanelUtils"; import { MockChatMessage } from "#/testHelpers/chatEntities"; +import { createDeferred } from "#/testHelpers/deferred"; import { buildOptimisticEditedMessage } from "./chatMessageEdits"; import { addChildToParentInCache, + applyChatArchiveStateToCaches, + applyWatchedChatArchived, + applyWatchedChatCreatedOrUnarchived, archiveChat, type ChatListInput, cancelChatEntity, @@ -27,6 +31,7 @@ import { chatDebugRunKey, chatDebugRunsKey, chatDiffContentsKey, + chatEntitiesFamilyKey, chatEntityKey, chatListFamilyKey, chatListKey, @@ -63,6 +68,7 @@ import { removeChildFromParentInCache, reorderPinnedChat, replaceChatMessagesHistory, + resetUnloadedChatEntity, setChatGroupRole, setChatUserRole, shouldInvalidateChatSearches, @@ -172,6 +178,34 @@ const createTestQueryClient = (): QueryClient => }, }); +const observeChatWithDeferredFirstFetch = ( + queryClient: QueryClient, + staleChat: TypesGen.Chat, + durableChat: TypesGen.Chat, +) => { + const firstFetch = createDeferred(); + const durableResult = createDeferred(); + let fetchCount = 0; + const observer = new QueryObserver(queryClient, { + queryKey: chatEntityKey(staleChat.id), + queryFn: () => { + fetchCount++; + return fetchCount === 1 ? firstFetch.promise : durableChat; + }, + }); + const unsubscribe = observer.subscribe((result) => { + if (result.data === durableChat) { + durableResult.resolve(result.data); + } + }); + return { + durableResult, + firstFetch, + fetchCount: () => fetchCount, + unsubscribe, + }; +}; + describe("advisor config query factories", () => { it("builds the advisor config query and delegates to the API", async () => { const advisorConfig: TypesGen.AdvisorConfig = { @@ -532,6 +566,25 @@ describe("archiveChat optimistic update", () => { expect(readInfiniteChats(queryClient, { archived: false })).toEqual([]); }); + it("removes loaded search rows after success", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + const unrelatedRow = makeChat("chat-2"); + queryClient.setQueryData(chatSearch({ q: "alpha" }).queryKey, [ + makeChat(chatId, { pin_order: 2 }), + unrelatedRow, + ]); + + const mutation = archiveChat(queryClient); + mutation.onSuccess(undefined, chatId); + + const rows = queryClient.getQueryData( + chatSearch({ q: "alpha" }).queryKey, + ); + expect(rows?.find((row) => row.id === chatId)).toBeUndefined(); + expect(rows?.find((row) => row.id === "chat-2")).toEqual(unrelatedRow); + }); + it("rolls back the chats list on error by invalidating", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; @@ -694,6 +747,25 @@ describe("unarchiveChat optimistic update", () => { }); }); + it("removes loaded search rows after success", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + const unrelatedRow = makeChat("chat-2", { archived: true }); + queryClient.setQueryData(chatSearch({ q: "alpha" }).queryKey, [ + makeChat(chatId, { archived: true }), + unrelatedRow, + ]); + + const mutation = unarchiveChat(queryClient); + mutation.onSuccess(undefined, chatId); + + const rows = queryClient.getQueryData( + chatSearch({ q: "alpha" }).queryKey, + ); + expect(rows?.find((row) => row.id === chatId)).toBeUndefined(); + expect(rows?.find((row) => row.id === "chat-2")).toEqual(unrelatedRow); + }); + it("rolls back both caches on error", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; @@ -3336,6 +3408,28 @@ describe("semantic cache operations: cancellation", () => { }); }); + it("resetUnloadedChatEntity resets exactly when detail data is absent", async () => { + const queryClient = createTestQueryClient(); + const resetSpy = vi.spyOn(queryClient, "resetQueries"); + + await resetUnloadedChatEntity(queryClient, "chat-1"); + + expect(resetSpy).toHaveBeenCalledWith({ + queryKey: chatEntityKey("chat-1"), + exact: true, + }); + }); + + it("resetUnloadedChatEntity is a no-op when detail data exists", async () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + const resetSpy = vi.spyOn(queryClient, "resetQueries"); + + await resetUnloadedChatEntity(queryClient, "chat-1"); + + expect(resetSpy).not.toHaveBeenCalled(); + }); + it("cancelChatMessages cancels the exact messages entry", async () => { const queryClient = createTestQueryClient(); const cancelSpy = vi.spyOn(queryClient, "cancelQueries"); @@ -3739,3 +3833,344 @@ describe("message upsert fan-out and history replacement", () => { ).toBeUndefined(); }); }); + +describe("chatEntitiesFamilyKey shape", () => { + // chatEntityKey builds on this prefix, so every entity detail entry + // shares the family root and can be addressed as a group. + it("prefixes every chat entity key", () => { + expect(chatEntityKey("chat-1")).toEqual([ + ...chatEntitiesFamilyKey, + "chat-1", + ]); + }); +}); + +describe("applyChatArchiveStateToCaches search rows", () => { + it("removes matching rows from every cached search and preserves unrelated rows by reference", () => { + const queryClient = createTestQueryClient(); + const unrelatedRow = makeChat("chat-2"); + queryClient.setQueryData(chatSearch({ q: "alpha" }).queryKey, [ + makeChat("chat-1", { pin_order: 2 }), + unrelatedRow, + ]); + queryClient.setQueryData(chatSearch({ q: "archived:true" }).queryKey, [ + makeChat("chat-1", { archived: true }), + ]); + + applyChatArchiveStateToCaches(queryClient, "chat-1", true); + + expect( + queryClient + .getQueryData(chatSearch({ q: "alpha" }).queryKey) + ?.find((row) => row.id === "chat-1"), + ).toBeUndefined(); + expect( + queryClient.getQueryData( + chatSearch({ q: "archived:true" }).queryKey, + ), + ).toEqual([]); + expect( + queryClient + .getQueryData(chatSearch({ q: "alpha" }).queryKey) + ?.find((row) => row.id === "chat-2"), + ).toEqual(unrelatedRow); + }); + + it("preserves the previous array reference when the row is not cached", () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatSearch({ q: "alpha" }).queryKey, [ + makeChat("chat-2"), + ]); + const before = queryClient.getQueryData( + chatSearch({ q: "alpha" }).queryKey, + ); + + applyChatArchiveStateToCaches(queryClient, "chat-1", true); + + expect(queryClient.getQueryData(chatSearch({ q: "alpha" }).queryKey)).toBe( + before, + ); + }); + + it("leaves per-chat sub-resource entries untouched", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + queryClient.setQueryData(chatMessagesKey(chatId), []); + queryClient.setQueryData(chatPromptsKey(chatId), { prompts: [] }); + queryClient.setQueryData(chatACLKey(chatId), {}); + queryClient.setQueryData(chatDiffContentsKey(chatId), { files: [] }); + + applyChatArchiveStateToCaches(queryClient, chatId, true); + + for (const [label, key] of [ + ["messages", chatMessagesKey(chatId)], + ["prompts", chatPromptsKey(chatId)], + ["acl", chatACLKey(chatId)], + ["diff-contents", chatDiffContentsKey(chatId)], + ] as const) { + expect( + queryClient.getQueryData(key), + `${label} entry should survive`, + ).toBeDefined(); + expect( + queryClient.getQueryState(key)?.isInvalidated, + `${label} entry should NOT be invalidated`, + ).not.toBe(true); + } + }); +}); + +describe("applyWatchedChatArchived", () => { + it("restarts an active initial entity fetch so stale data cannot overwrite archive", async () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + const staleChat = makeChat(chatId, { archived: false }); + const durableChat = makeChat(chatId, { archived: true }); + const fetch = observeChatWithDeferredFirstFetch( + queryClient, + staleChat, + durableChat, + ); + + applyWatchedChatArchived(queryClient, durableChat); + fetch.firstFetch.resolve(staleChat); + await fetch.durableResult.promise; + + expect(fetch.fetchCount()).toBe(2); + expect( + queryClient.getQueryData(chatEntityKey(chatId))?.archived, + ).toBe(true); + fetch.unsubscribe(); + }); + + it("does not create an entity query when no observer is mounted", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + + applyWatchedChatArchived(queryClient, makeChat(chatId, { archived: true })); + + expect(queryClient.getQueryState(chatEntityKey(chatId))).toBeUndefined(); + }); + + it("patches the entity in place instead of removing it", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + queryClient.setQueryData( + chatEntityKey(chatId), + makeChat(chatId, { pin_order: 2 }), + ); + + applyWatchedChatArchived(queryClient, makeChat(chatId, { archived: true })); + + expect( + queryClient.getQueryData(chatEntityKey(chatId)), + ).toMatchObject({ + archived: true, + pin_order: 0, + }); + }); + + it("drops the chat from active lists and patches it in archived lists", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + seedInfiniteChats(queryClient, [makeChat(chatId), makeChat("chat-2")], { + archived: false, + }); + seedInfiniteChats(queryClient, [makeChat(chatId, { pin_order: 3 })], { + archived: true, + }); + + applyWatchedChatArchived(queryClient, makeChat(chatId, { archived: true })); + + expect( + readInfiniteChats(queryClient, { archived: false })?.map( + (chat) => chat.id, + ), + ).toEqual(["chat-2"]); + expect( + readInfiniteChats(queryClient, { archived: true })?.[0], + ).toMatchObject({ + id: chatId, + archived: true, + pin_order: 0, + }); + }); + + it("removes search rows and removes the by-workspace mapping", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + queryClient.setQueryData(chatSearch({ q: "alpha" }).queryKey, [ + makeChat(chatId), + ]); + queryClient.setQueryData(chatsByWorkspace(["ws-1"]).queryKey, { + "ws-1": chatId, + }); + + applyWatchedChatArchived(queryClient, makeChat(chatId, { archived: true })); + + expect( + queryClient.getQueryData( + chatSearch({ q: "alpha" }).queryKey, + ), + ).toEqual([]); + expect( + queryClient.getQueryData(chatsByWorkspace(["ws-1"]).queryKey), + ).toEqual({}); + }); + + it("invalidates the list, by-workspace, and search families but not per-chat sub-resources", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + seedInfiniteChats(queryClient, [makeChat(chatId)]); + queryClient.setQueryData(chatEntityKey(chatId), makeChat(chatId)); + queryClient.setQueryData(chatMessagesKey(chatId), []); + queryClient.setQueryData(chatSearch({ q: "alpha" }).queryKey, []); + queryClient.setQueryData(chatsByWorkspace(["ws-1"]).queryKey, {}); + + applyWatchedChatArchived(queryClient, makeChat(chatId, { archived: true })); + + for (const [label, key] of [ + ["chat list", infiniteChatsTestKey], + ["chat search", chatSearch({ q: "alpha" }).queryKey], + ["by-workspace", chatsByWorkspace(["ws-1"]).queryKey], + ] as const) { + expect( + queryClient.getQueryState(key)?.isInvalidated, + `${label} entry should be invalidated`, + ).toBe(true); + } + expect( + queryClient.getQueryData(chatMessagesKey(chatId)), + "messages entry should survive", + ).toBeDefined(); + expect( + queryClient.getQueryState(chatMessagesKey(chatId))?.isInvalidated, + "messages entry should NOT be invalidated", + ).not.toBe(true); + expect( + queryClient.getQueryData(chatEntityKey(chatId)), + "entity entry should survive", + ).toBeDefined(); + }); +}); + +describe("applyWatchedChatCreatedOrUnarchived", () => { + it("restarts an active initial entity fetch so stale data cannot overwrite unarchive", async () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + const staleChat = makeChat(chatId, { archived: true }); + const durableChat = makeChat(chatId, { archived: false }); + const fetch = observeChatWithDeferredFirstFetch( + queryClient, + staleChat, + durableChat, + ); + + applyWatchedChatCreatedOrUnarchived(queryClient, durableChat); + fetch.firstFetch.resolve(staleChat); + await fetch.durableResult.promise; + + expect(fetch.fetchCount()).toBe(2); + expect( + queryClient.getQueryData(chatEntityKey(chatId))?.archived, + ).toBe(false); + fetch.unsubscribe(); + }); + + it("restarts an active initial fetch for a genuinely new chat", async () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-new"; + const staleChat = makeChat(chatId, { title: "stale" }); + const durableChat = makeChat(chatId, { title: "durable" }); + const fetch = observeChatWithDeferredFirstFetch( + queryClient, + staleChat, + durableChat, + ); + + applyWatchedChatCreatedOrUnarchived(queryClient, durableChat); + fetch.firstFetch.resolve(staleChat); + await fetch.durableResult.promise; + + expect(fetch.fetchCount()).toBe(2); + expect( + queryClient.getQueryData(chatEntityKey(chatId))?.title, + ).toBe("durable"); + fetch.unsubscribe(); + }); + + it("flips a cached archived entity back to active and repairs list rows", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + queryClient.setQueryData( + chatEntityKey(chatId), + makeChat(chatId, { archived: true }), + ); + seedInfiniteChats(queryClient, [makeChat(chatId, { archived: true })], { + archived: true, + }); + seedInfiniteChats(queryClient, [makeChat(chatId, { archived: true })], { + archived: false, + }); + + applyWatchedChatCreatedOrUnarchived(queryClient, makeChat(chatId)); + + expect( + queryClient.getQueryData(chatEntityKey(chatId))?.archived, + ).toBe(false); + expect(readInfiniteChats(queryClient, { archived: true })).toEqual([]); + expect( + readInfiniteChats(queryClient, { archived: false })?.[0].archived, + ).toBe(false); + }); + + it("only invalidates the collection families for a truly new chat", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-new"; + seedInfiniteChats(queryClient, [makeChat("chat-2")]); + queryClient.setQueryData(chatSearch({ q: "alpha" }).queryKey, []); + queryClient.setQueryData(chatsByWorkspace(["ws-1"]).queryKey, {}); + + applyWatchedChatCreatedOrUnarchived(queryClient, makeChat(chatId)); + + expect(queryClient.getQueryData(chatEntityKey(chatId))).toBeUndefined(); + expect(queryClient.getQueryState(chatEntityKey(chatId))).toBeUndefined(); + for (const [label, key] of [ + ["chat list", infiniteChatsTestKey], + ["chat search", chatSearch({ q: "alpha" }).queryKey], + ["by-workspace", chatsByWorkspace(["ws-1"]).queryKey], + ] as const) { + expect( + queryClient.getQueryState(key)?.isInvalidated, + `${label} entry should be invalidated`, + ).toBe(true); + } + }); +}); + +describe("archive mutation entity retention", () => { + it.each([ + { name: "archiveChat", factory: archiveChat, archived: true }, + { name: "unarchiveChat", factory: unarchiveChat, archived: false }, + ])("$name onSuccess never removes the entity family", ({ + factory, + archived, + }) => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + queryClient.setQueryData( + chatEntityKey(chatId), + makeChat(chatId, { archived: !archived }), + ); + queryClient.setQueryData(chatMessagesKey(chatId), []); + + const mutation = factory(queryClient); + mutation.onSuccess(undefined, chatId); + mutation.onSettled(undefined, undefined, chatId); + + expect( + queryClient.getQueryData(chatEntityKey(chatId)), + ).toMatchObject({ archived }); + expect(queryClient.getQueryData(chatMessagesKey(chatId))).toBeDefined(); + }); +}); diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 83e031392bd93..94c7752b7614e 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -28,8 +28,10 @@ const chatsByWorkspaceFamilyKey = [ "by-workspace", ] as const; +export const chatEntitiesFamilyKey = ["chats", "entities"] as const; + export const chatEntityKey = (chatId: string) => - ["chats", "entities", chatId] as const; + [...chatEntitiesFamilyKey, chatId] as const; export const chatFilesKey = ["chats", "files"] as const; @@ -131,15 +133,23 @@ export const updateInfiniteChatsCache = ( * in the cache, but only if the chat doesn't already exist in any * page. This avoids the per-page duplication that would occur if * a prepend updater were passed to updateInfiniteChatsCache, which - * runs independently on each page. + * runs independently on each page. Lists whose archived filter + * conflicts with the chat's archive state are skipped, so an active + * chat is never inserted into an archived-only list. */ export const prependToInfiniteChatsCache = ( queryClient: QueryClient, chat: TypesGen.Chat, ) => { - queryClient.setQueriesData( - { queryKey: chatListFamilyKey }, - (prev) => { + const queries = queryClient.getQueriesData({ + queryKey: chatListFamilyKey, + }); + for (const [queryKey] of queries) { + const archivedFilter = archivedFilterForChatListKey(queryKey); + if (archivedFilter !== undefined && archivedFilter !== chat.archived) { + continue; + } + queryClient.setQueryData(queryKey, (prev) => { if (!prev?.pages) return prev; // Check across ALL pages to avoid duplicates. const exists = prev.pages.some((page) => @@ -151,8 +161,8 @@ export const prependToInfiniteChatsCache = ( i === 0 ? [chat, ...page] : page, ); return { ...prev, pages: nextPages }; - }, - ); + }); + } }; /** @@ -298,9 +308,15 @@ const patchChatArchiveState = ( }; /** - * Applies an accepted archive state to loaded sidebar and detail caches. - * Removes the chat from any filtered list whose archived filter conflicts - * with the new state, and resets pin_order to 0 when archiving. + * Applies an accepted archive state to loaded sidebar, search, and + * detail caches. Removes the chat from any filtered list whose archived + * filter conflicts with the new state, and resets pin_order to 0 when + * archiving. + * + * Search rows are removed rather than patched: a cached row matched its + * query's archived filter before the change, so after the change it + * belongs to a different result set. Search invalidations issued by the + * callers repopulate any result set that still matches. */ export const applyChatArchiveStateToCaches = ( queryClient: QueryClient, @@ -367,6 +383,68 @@ export const applyChatArchiveStateToCaches = ( return changed ? { ...prev, pages } : prev; }); } + + const searchQueries = queryClient.getQueriesData({ + queryKey: chatSearchFamilyKey, + }); + for (const [queryKey] of searchQueries) { + queryClient.setQueryData(queryKey, (prev) => { + if (!prev) { + return prev; + } + const next = prev.filter((row) => row.id !== chatId); + return next.length === prev.length ? prev : next; + }); + } +}; + +/** + * Watch-event effect for the `deleted` kind, which the server publishes + * once per family member when a chat family is archived. Archive is a + * patch, never an eviction: the entity and its sub-resources stay + * cached so an open route flips to the archived read-only state without + * a loading flash or a zombie render. + */ +export const applyWatchedChatArchived = ( + queryClient: QueryClient, + chat: TypesGen.Chat, +) => { + void cancelChatListRefetches(queryClient); + if (queryClient.getQueryData(chatEntityKey(chat.id)) === undefined) { + void resetUnloadedChatEntity(queryClient, chat.id); + } else { + void cancelLoadedChatEntityRefetch(queryClient, chat.id); + } + applyChatArchiveStateToCaches(queryClient, chat.id, true); + removeChatFromChatsByWorkspace(queryClient, chat.id); + void invalidateChatListQueries(queryClient); + void invalidateChatsByWorkspace(queryClient); + void invalidateChatSearches(queryClient); +}; + +/** + * Watch-event effect for a root `created` event, which the server + * publishes both for new chats and for unarchive transitions (one event + * per family member). A cached entity marked archived identifies the + * unarchive case; a truly new chat only needs the family invalidations + * and never gets a speculative entity entry. The caller remains + * responsible for list prepend and child insertion. + */ +export const applyWatchedChatCreatedOrUnarchived = ( + queryClient: QueryClient, + chat: TypesGen.Chat, +) => { + const cachedChat = queryClient.getQueryData( + chatEntityKey(chat.id), + ); + if (cachedChat === undefined) { + void resetUnloadedChatEntity(queryClient, chat.id); + } else if (cachedChat.archived) { + applyChatArchiveStateToCaches(queryClient, chat.id, false); + } + void invalidateChatListQueries(queryClient); + void invalidateChatsByWorkspace(queryClient); + void invalidateChatSearches(queryClient); }; const parseUpdatedAtInstant = (updatedAt: string) => { @@ -765,6 +843,23 @@ export const cancelLoadedChatEntityRefetch = ( }); }; +/** + * Restarts an active first-time fetch after a durable watch transition. + * Invalidation reuses the stale initial promise when no data is loaded. + */ +export const resetUnloadedChatEntity = ( + queryClient: QueryClient, + chatId: string, +) => { + if (queryClient.getQueryData(chatEntityKey(chatId)) !== undefined) { + return; + } + return queryClient.resetQueries({ + queryKey: chatEntityKey(chatId), + exact: true, + }); +}; + export const cancelChatMessages = (queryClient: QueryClient, chatId: string) => queryClient.cancelQueries({ queryKey: chatMessagesKey(chatId), diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.stories.tsx b/site/src/pages/AgentsPage/AgentsPageLayout.stories.tsx index 8e77cfe51356a..c452dd94aea35 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.stories.tsx +++ b/site/src/pages/AgentsPage/AgentsPageLayout.stories.tsx @@ -13,6 +13,12 @@ import { } from "storybook/test"; import { reactRouterParameters } from "storybook-addon-remix-react-router"; import { API } from "#/api/api"; +import { getAuthorizationKey } from "#/api/queries/authCheck"; +import { + chatEntityKey, + chatMessagesKey, + chatPromptsKey, +} from "#/api/queries/chats"; import type * as TypesGen from "#/api/typesGenerated"; import type { Chat } from "#/api/typesGenerated"; import { DeleteDialog } from "#/components/Dialog/DeleteDialog/DeleteDialog"; @@ -25,9 +31,11 @@ import { import { withAuthProvider, withDashboardProvider, + withProxyProvider, withWebSocket, } from "#/testHelpers/storybook"; import { CoderAgentsPageView } from "../AISettingsPage/CoderAgentsPage/CoderAgentsPageView"; +import AgentChatPage, { RIGHT_PANEL_OPEN_KEY } from "./AgentChatPage"; import AgentCreatePage from "./AgentCreatePage"; import AgentSettingsCompactionPage from "./AgentSettingsCompactionPage"; import AgentSettingsGeneralPage from "./AgentSettingsGeneralPage"; @@ -932,6 +940,139 @@ export const WithAgentSelected: Story = { }, }; +// --------------------------------------------------------------------------- +// Watch-event archive semantics: these stories mount the real AgentChatPage +// under the layout's :agentId route so the layout's chat-watch socket drives +// the page through the production watch path. +// --------------------------------------------------------------------------- + +const agentsWithAgentChatPageRouting = { + ...agentsRouting, + children: agentsRouting.children.map((route) => + "path" in route && route.path === ":agentId" + ? { ...route, element: } + : route, + ), +}; + +const WATCHED_CHAT_ID = "chat-watched"; + +// MockChat is owned by MockUserOwner, so the page renders the owner view +// (composer enabled unless archived) instead of the other-user banner. +const watchedChat = (overrides: Partial = {}): Chat => ({ + ...MockChat, + id: WATCHED_CHAT_ID, + title: "Watched agent", + last_model_config_id: defaultModelConfigID, + created_at: oneWeekAgo, + updated_at: oneWeekAgo, + ...overrides, +}); + +const watchedChatQueries = (chat: Chat) => [ + { key: chatEntityKey(chat.id), data: chat }, + { + key: chatMessagesKey(chat.id), + data: { + pages: [{ messages: [], queued_messages: [], has_more: false }], + pageParams: [undefined], + }, + }, + { key: chatPromptsKey(chat.id), data: { prompts: [] } }, + { + key: getAuthorizationKey({ + checks: { + canShareChat: { + object: { + resource_type: "chat", + owner_id: chat.owner_id, + organization_id: chat.organization_id, + }, + action: "share", + }, + }, + }), + data: { canShareChat: true }, + }, +]; + +const chatWatchEvent = (kind: TypesGen.ChatWatchEventKind, chat: Chat) => ({ + event: "message" as const, + data: JSON.stringify({ kind, chat } satisfies TypesGen.ChatWatchEvent), +}); + +const watchedChatPageParameters = ( + chat: Chat, + watchEvents: readonly ReturnType[], +) => ({ + queries: watchedChatQueries(chat), + webSocket: { + "/chats/watch": [...watchEvents], + }, + reactRouter: reactRouterParameters({ + location: { + path: `/agents/${WATCHED_CHAT_ID}`, + pathParams: { agentId: WATCHED_CHAT_ID }, + }, + routing: [agentsWithAgentChatPageRouting, aiSettingsRouting], + }), +}); + +const mockAgentChatPageAPIs = () => { + localStorage.removeItem(RIGHT_PANEL_OPEN_KEY); + spyOn(API, "getApiKey").mockRejectedValue(new Error("missing API key")); + spyOn(API.experimental, "updateChat").mockResolvedValue(); + return () => localStorage.removeItem(RIGHT_PANEL_OPEN_KEY); +}; + +export const ArchiveWatchEventKeepsOpenChatMounted: Story = { + decorators: [withProxyProvider()], + beforeEach: () => { + mockChats([watchedChat()]); + return mockAgentChatPageAPIs(); + }, + parameters: watchedChatPageParameters(watchedChat(), [ + chatWatchEvent("deleted", watchedChat({ archived: true })), + ]), + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect( + await canvas.findByText("This agent has been archived and is read-only."), + ).toBeVisible(); + await waitFor(() => { + expect(canvas.getByRole("textbox")).toHaveAttribute( + "aria-disabled", + "true", + ); + }); + expect(canvas.queryByText("Chat not found")).not.toBeInTheDocument(); + }, +}; + +export const UnarchiveWatchEventRecoversArchivedChat: Story = { + decorators: [withProxyProvider()], + beforeEach: () => { + mockChats([watchedChat({ archived: true })]); + return mockAgentChatPageAPIs(); + }, + parameters: watchedChatPageParameters(watchedChat({ archived: true }), [ + chatWatchEvent("created", watchedChat({ archived: false })), + ]), + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await waitFor(() => { + expect(canvas.getByRole("textbox")).not.toHaveAttribute( + "aria-disabled", + "true", + ); + }); + expect( + canvas.queryByText("This agent has been archived and is read-only."), + ).not.toBeInTheDocument(); + expect(canvas.queryByText("Chat not found")).not.toBeInTheDocument(); + }, +}; + // Error reasons surface via each chat's last_error, which the // layout turns into sidebar error badges. export const WithErrorReasons: Story = { diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.tsx b/site/src/pages/AgentsPage/AgentsPageLayout.tsx index 16f66ca937a73..41f79d6b280cd 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.tsx +++ b/site/src/pages/AgentsPage/AgentsPageLayout.tsx @@ -18,6 +18,8 @@ import { getErrorMessage } from "#/api/errors"; import { addChildToParentInCache, applyChatArchiveStateToCaches, + applyWatchedChatArchived, + applyWatchedChatCreatedOrUnarchived, archiveChat, cancelChatListRefetches, cancelLoadedChatEntityRefetch, @@ -36,9 +38,7 @@ import { prependToInfiniteChatsCache, proposeChatTitle, readInfiniteChatsCache, - removeChatEntity, removeChatFromChatsByWorkspace, - removeChildFromParentInCache, reorderPinnedChat, shouldInvalidateChatSearches, shouldInvalidateChatsByWorkspace, @@ -611,20 +611,12 @@ const AgentsPageLayout: FC = () => { } if (chatEvent.kind === "deleted") { - // Drop the chat from the flat root list (root or - // cascade via root_chat_id) and from any parent's - // embedded children (individual child archive). - updateInfiniteChatsCache(queryClient, (chats) => - chats.filter( - (c) => - c.id !== updatedChat.id && c.root_chat_id !== updatedChat.id, - ), - ); - removeChildFromParentInCache(queryClient, updatedChat.id); - removeChatEntity(queryClient, updatedChat.id); - removeChatFromChatsByWorkspace(queryClient, updatedChat.id); - void invalidateChatsByWorkspace(queryClient); - void invalidateChatSearches(queryClient); + // The server publishes `deleted` when a chat is + // archived (one event per family member); there is + // no hard-delete wire event. Patch archive state in + // place so an open route stays mounted and flips to + // its read-only state. + applyWatchedChatArchived(queryClient, updatedChat); return; } if (chatEvent.kind === "diff_status_change") { @@ -657,11 +649,23 @@ const AgentsPageLayout: FC = () => { updatedChat, updatedChat.parent_chat_id, ); + // A family unarchive and a new sub-agent with a + // mounted initial fetch both need entity recovery. + const cachedChat = queryClient.getQueryData( + chatEntityKey(updatedChat.id), + ); + if ( + cachedChat?.archived || + (cachedChat === undefined && + queryClient.getQueryState(chatEntityKey(updatedChat.id)) !== + undefined) + ) { + applyWatchedChatCreatedOrUnarchived(queryClient, updatedChat); + } } else { + // `created` also fires for unarchive transitions. + applyWatchedChatCreatedOrUnarchived(queryClient, updatedChat); prependToInfiniteChatsCache(queryClient, updatedChat); - void invalidateChatListQueries(queryClient); - void invalidateChatsByWorkspace(queryClient); - void invalidateChatSearches(queryClient); } } else { mergeWatchedChatIntoCaches(queryClient, updatedChat, {