Keep active thread state consistent after browser resume

apps/app/src/hooks/cache-owners/active-thread-lifecycle-cache-owner.ts

@@ -0,0 +1,77 @@
+import type { QueryClient, QueryKey } from "@tanstack/react-query";
+import type { QueryClientArg } from "../cache-effect-types";
+import {
+  THREAD_CONVERSATION_OUTLINE_QUERY_KEY,
+  THREAD_DEFAULT_EXECUTION_OPTIONS_QUERY_KEY,
+  THREAD_DETAIL_BOOTSTRAP_QUERY_KEY,
+  THREAD_PENDING_INTERACTIONS_QUERY_KEY,
+  THREAD_PROMPT_HISTORY_QUERY_KEY,
+  THREAD_QUERY_KEY,
+  THREAD_QUEUED_MESSAGES_QUERY_KEY,
+  THREAD_TIMELINE_QUERY_KEY,
+  threadConversationOutlineQueryKey,
+  threadDefaultExecutionOptionsQueryKey,
+  threadDetailBootstrapQueryKey,
+  threadPendingInteractionsQueryKey,
+  threadPromptHistoryQueryKey,
+  threadQueryKey,
+  threadQueuedMessagesQueryKey,
+  threadTimelineQueryKey,
+} from "../queries/query-keys";
+
+const ACTIVE_THREAD_QUERY_ROOTS = new Set<unknown>([
+  THREAD_QUERY_KEY,
+  THREAD_DETAIL_BOOTSTRAP_QUERY_KEY,
+  THREAD_DEFAULT_EXECUTION_OPTIONS_QUERY_KEY,
+  THREAD_QUEUED_MESSAGES_QUERY_KEY,
+  THREAD_PROMPT_HISTORY_QUERY_KEY,
+  THREAD_PENDING_INTERACTIONS_QUERY_KEY,
+  THREAD_TIMELINE_QUERY_KEY,
+  THREAD_CONVERSATION_OUTLINE_QUERY_KEY,
+]);
+
+function getThreadIdFromActiveThreadQueryKey(
+  queryKey: QueryKey,
+): string | null {
+  const [queryRoot, threadId] = queryKey;
+  if (!ACTIVE_THREAD_QUERY_ROOTS.has(queryRoot)) {
+    return null;
+  }
+  return typeof threadId === "string" && threadId.length > 0
+    ? threadId
+    : null;
+}
+
+function getActiveThreadBundleQueryKeys(threadId: string): QueryKey[] {
+  return [
+    threadQueryKey(threadId),
+    threadDetailBootstrapQueryKey(threadId),
+    threadDefaultExecutionOptionsQueryKey(threadId),
+    threadQueuedMessagesQueryKey(threadId),
+    threadPromptHistoryQueryKey(threadId),
+    threadPendingInteractionsQueryKey(threadId),
+    threadTimelineQueryKey(threadId),
+    threadConversationOutlineQueryKey(threadId),
+  ];
+}
+
+function collectActiveThreadBundleThreadIds(queryClient: QueryClient): string[] {
+  const threadIds = new Set<string>();
+  for (const query of queryClient.getQueryCache().findAll({ type: "active" })) {
+    const threadId = getThreadIdFromActiveThreadQueryKey(query.queryKey);
+    if (threadId !== null) {
+      threadIds.add(threadId);
+    }
+  }
+  return Array.from(threadIds);
+}
+
+export function invalidateActiveThreadBundleQueriesAfterBrowserResume({
+  queryClient,
+}: QueryClientArg): void {
+  for (const threadId of collectActiveThreadBundleThreadIds(queryClient)) {
+    for (const queryKey of getActiveThreadBundleQueryKeys(threadId)) {
+      queryClient.invalidateQueries({ queryKey });
+    }
+  }
+}

apps/app/src/hooks/cache-owners/cache-owner-registry.test.ts

@@ -46,6 +46,24 @@ const DEPRECATED_CACHE_SHIM_MODULES = new Set([
 const QUERY_KEYS_MODULE_PATH = "hooks/queries/query-keys";
 
 const CACHE_OWNER_QUERY_KEY_IMPORTS: CacheOwnerQueryKeyImportRegistry = {
+  "hooks/cache-owners/active-thread-lifecycle-cache-owner.ts": [
+    "THREAD_CONVERSATION_OUTLINE_QUERY_KEY",
+    "THREAD_DEFAULT_EXECUTION_OPTIONS_QUERY_KEY",
+    "THREAD_DETAIL_BOOTSTRAP_QUERY_KEY",
+    "THREAD_PENDING_INTERACTIONS_QUERY_KEY",
+    "THREAD_PROMPT_HISTORY_QUERY_KEY",
+    "THREAD_QUERY_KEY",
+    "THREAD_QUEUED_MESSAGES_QUERY_KEY",
+    "THREAD_TIMELINE_QUERY_KEY",
+    "threadConversationOutlineQueryKey",
+    "threadDefaultExecutionOptionsQueryKey",
+    "threadDetailBootstrapQueryKey",
+    "threadPendingInteractionsQueryKey",
+    "threadPromptHistoryQueryKey",
+    "threadQueryKey",
+    "threadQueuedMessagesQueryKey",
+    "threadTimelineQueryKey",
+  ],
   "hooks/cache-owners/automation-cache-effects.ts": [
     "automationDetailQueryKey",
     "automationRunsQueryKey",
@@ -170,6 +188,8 @@ const CACHE_OWNER_QUERY_KEY_IMPORTS: CacheOwnerQueryKeyImportRegistry = {
     "allHostQueryKeyPrefix",
     "allProjectPathsQueryKeyPrefix",
     "allSystemExecutionOptionsQueryKeyPrefix",
+    "allThreadConversationOutlineQueryKeyPrefix",
+    "allThreadDetailBootstrapQueryKeyPrefix",
     "allThreadPendingInteractionsQueryKeyPrefix",
     "allThreadQueryKeyPrefix",
     "allThreadQueuedMessagesQueryKeyPrefix",

apps/app/src/hooks/cache-owners/system-cache-effects.ts

@@ -10,6 +10,8 @@ import {
   allHostQueryKeyPrefix,
   allProjectPathsQueryKeyPrefix,
   allSystemExecutionOptionsQueryKeyPrefix,
+  allThreadConversationOutlineQueryKeyPrefix,
+  allThreadDetailBootstrapQueryKeyPrefix,
   allThreadQueuedMessagesQueryKeyPrefix,
   allThreadPendingInteractionsQueryKeyPrefix,
   allThreadQueryKeyPrefix,
@@ -103,7 +105,9 @@ function getServerReconnectInvalidationQueryKeys(): QueryKey[] {
     threadsQueryKey(),
     threadSearchQueryKeyPrefix(),
     allThreadQueryKeyPrefix(),
+    allThreadDetailBootstrapQueryKeyPrefix(),
     allThreadTimelineQueryKeyPrefix(),
+    allThreadConversationOutlineQueryKeyPrefix(),
     allThreadTimelineTurnSummaryDetailsQueryKeyPrefix(),
     allThreadQueuedMessagesQueryKeyPrefix(),
     threadPromptHistoryQueryKeyPrefix(),

apps/app/src/hooks/queries/query-keys.ts

@@ -173,6 +173,9 @@ export type DisabledThreadListQueryKey = readonly [
 ];
 export type ThreadQueryKeyPrefix = readonly [typeof THREAD_QUERY_KEY];
 export type ThreadQueryKey = readonly [typeof THREAD_QUERY_KEY, string];
+export type ThreadDetailBootstrapQueryKeyPrefix = readonly [
+  typeof THREAD_DETAIL_BOOTSTRAP_QUERY_KEY,
+];
 export type ThreadDetailBootstrapQueryKey = readonly [
   typeof THREAD_DETAIL_BOOTSTRAP_QUERY_KEY,
   string,
@@ -635,6 +638,10 @@ export function threadDetailBootstrapQueryKey(
   return [THREAD_DETAIL_BOOTSTRAP_QUERY_KEY, threadId];
 }
 
+export function allThreadDetailBootstrapQueryKeyPrefix(): ThreadDetailBootstrapQueryKeyPrefix {
+  return [THREAD_DETAIL_BOOTSTRAP_QUERY_KEY];
+}
+
 export function allThreadQueryKeyPrefix(): ThreadQueryKeyPrefix {
   return [THREAD_QUERY_KEY];
 }

apps/app/src/hooks/system-cache-effects.test.ts

@@ -1,16 +1,21 @@
 import { describe, expect, it, vi } from "vitest";
 import { QueryObserver } from "@tanstack/react-query";
+import type { QueryClient, QueryKey } from "@tanstack/react-query";
 import { createAppQueryClient } from "@/lib/query-client";
 import {
   environmentDiffFilesQueryKey,
   environmentDiffPatchQueryKey,
   sidebarNavigationQueryKey,
   systemExecutionOptionsQueryKey,
+  threadConversationOutlineQueryKey,
   threadDefaultExecutionOptionsQueryKey,
+  threadDetailBootstrapQueryKey,
   threadPendingInteractionsQueryKey,
   threadPromptHistoryQueryKey,
+  threadQueryKey,
   threadQueuedMessagesQueryKey,
   threadSearchQueryKey,
+  threadTimelineQueryKey,
 } from "./queries/query-keys";
 import { invalidateRealtimeQueriesAfterServerReconnect } from "./cache-owners/system-cache-effects";
 
@@ -46,9 +51,47 @@ const EMPTY_EXECUTION_OPTIONS = {
   modelLoadError: null,
 };
 
+interface ObservedQuery {
+  queryFn: ReturnType<typeof vi.fn<() => Promise<string>>>;
+  unsubscribe: () => void;
+}
+
+function observeIdleQuery(
+  queryClient: QueryClient,
+  queryKey: QueryKey,
+): ObservedQuery {
+  const queryFn = vi.fn<() => Promise<string>>().mockResolvedValue("loaded");
+  const observer = new QueryObserver(queryClient, {
+    queryKey,
+    queryFn,
+    refetchOnWindowFocus: false,
+    staleTime: Infinity,
+  });
+  return {
+    queryFn,
+    unsubscribe: observer.subscribe(() => {}),
+  };
+}
+
+async function waitForQueryCalls(
+  queries: readonly ObservedQuery[],
+  expectedCallCount: number,
+): Promise<void> {
+  await vi.waitFor(() => {
+    for (const query of queries) {
+      expect(query.queryFn).toHaveBeenCalledTimes(expectedCallCount);
+    }
+  });
+}
+
 describe("system cache effects", () => {
-  it("invalidates canonical composer caches after reconnect", () => {
+  it("invalidates canonical active thread caches after reconnect", () => {
     const queryClient = createCacheEffectQueryClient();
+    const threadKey = threadQueryKey("thread-1");
+    const threadBootstrapKey = threadDetailBootstrapQueryKey("thread-1");
+    const timelineKey = threadTimelineQueryKey("thread-1");
+    const conversationOutlineKey =
+      threadConversationOutlineQueryKey("thread-1");
     const queuedMessagesKey = threadQueuedMessagesQueryKey("thread-1");
     const promptHistoryKey = threadPromptHistoryQueryKey("thread-1");
     const pendingInteractionsKey =
@@ -63,6 +106,10 @@ describe("system cache effects", () => {
       environmentId: "env-1",
     });
     const sidebarNavigationKey = sidebarNavigationQueryKey();
+    queryClient.setQueryData(threadKey, { id: "thread-1" });
+    queryClient.setQueryData(threadBootstrapKey, { id: "thread-1" });
+    queryClient.setQueryData(timelineKey, { rows: [] });
+    queryClient.setQueryData(conversationOutlineKey, { items: [] });
     queryClient.setQueryData(queuedMessagesKey, []);
     queryClient.setQueryData(promptHistoryKey, []);
     queryClient.setQueryData(pendingInteractionsKey, []);
@@ -82,6 +129,14 @@ describe("system cache effects", () => {
 
     invalidateRealtimeQueriesAfterServerReconnect({ queryClient });
 
+    expect(queryClient.getQueryState(threadKey)?.isInvalidated).toBe(true);
+    expect(queryClient.getQueryState(threadBootstrapKey)?.isInvalidated).toBe(
+      true,
+    );
+    expect(queryClient.getQueryState(timelineKey)?.isInvalidated).toBe(true);
+    expect(
+      queryClient.getQueryState(conversationOutlineKey)?.isInvalidated,
+    ).toBe(true);
     expect(queryClient.getQueryState(queuedMessagesKey)?.isInvalidated).toBe(
       true,
     );
@@ -105,6 +160,42 @@ describe("system cache effects", () => {
     );
   });
 
+  it("refetches active thread bundle queries together after reconnect", async () => {
+    const queryClient = createCacheEffectQueryClient();
+    queryClient.mount();
+    const activeThreadQueries = [
+      observeIdleQuery(queryClient, threadQueryKey("thread-1")),
+      observeIdleQuery(queryClient, threadDetailBootstrapQueryKey("thread-1")),
+      observeIdleQuery(
+        queryClient,
+        threadDefaultExecutionOptionsQueryKey("thread-1"),
+      ),
+      observeIdleQuery(queryClient, threadQueuedMessagesQueryKey("thread-1")),
+      observeIdleQuery(queryClient, threadPromptHistoryQueryKey("thread-1")),
+      observeIdleQuery(
+        queryClient,
+        threadPendingInteractionsQueryKey("thread-1"),
+      ),
+      observeIdleQuery(queryClient, threadTimelineQueryKey("thread-1")),
+      observeIdleQuery(
+        queryClient,
+        threadConversationOutlineQueryKey("thread-1"),
+      ),
+    ];
+
+    await waitForQueryCalls(activeThreadQueries, 1);
+
+    invalidateRealtimeQueriesAfterServerReconnect({ queryClient });
+
+    await waitForQueryCalls(activeThreadQueries, 2);
+
+    for (const query of activeThreadQueries) {
+      query.unsubscribe();
+    }
+    queryClient.unmount();
+    queryClient.clear();
+  });
+
   it("refetches an active diff TOC query but evicts the observer-less patch cache after reconnect", async () => {
     const queryClient = createCacheEffectQueryClient();
     const diffFilesKey = environmentDiffFilesQueryKey("env-1", "all", "main");