fix(app): keep timeline requests runtime-scoped

packages/app/src/hooks/use-agent-initialization.test.ts

@@ -1,6 +1,7 @@
 import { afterEach, describe, expect, it, vi } from "vitest";
 import { useSessionStore } from "@/stores/session-store";
 import { TIMELINE_FETCH_PAGE_SIZE } from "@/timeline/timeline-fetch-policy";
+import type { HostRuntimeStore } from "@/runtime/host-runtime";
 import { getInitDeferred, getInitKey, resolveInitDeferred } from "@/utils/agent-initialization";
 import {
   createSetAgentInitializing,
@@ -12,15 +13,28 @@ import {
 const serverId = "server-1";
 const agentId = "agent-1";
 
-interface FakeDaemonClient {
-  fetchAgentTimeline: ReturnType<typeof vi.fn>;
-  refreshAgent: ReturnType<typeof vi.fn>;
+class FakeDaemonClient {
+  readonly refreshedAgentIds: string[] = [];
+
+  async refreshAgent(requestedAgentId: string): Promise<void> {
+    this.refreshedAgentIds.push(requestedAgentId);
+  }
 }
 
-function makeClient(): FakeDaemonClient {
-  return {
-    fetchAgentTimeline: vi.fn().mockResolvedValue(undefined),
-    refreshAgent: vi.fn().mockResolvedValue(undefined),
+class FakeTimelineRuntime {
+  readonly requests: Array<{
+    serverId: string;
+    agentId: string;
+    request: Parameters<HostRuntimeStore["fetchAgentTimeline"]>[2];
+  }> = [];
+
+  fetchAgentTimeline: HostRuntimeStore["fetchAgentTimeline"] = async (
+    requestedServerId,
+    requestedAgentId,
+    request,
+  ) => {
+    this.requests.push({ serverId: requestedServerId, agentId: requestedAgentId, request });
+    return undefined as never;
   };
 }
 
@@ -36,7 +50,8 @@ afterEach(() => {
 
 describe("ensureAgentIsInitialized", () => {
   it("requests bounded projected catch-up after the current cursor when authoritative history is loaded", () => {
-    const client = makeClient();
+    const client = new FakeDaemonClient();
+    const runtime = new FakeTimelineRuntime();
     useSessionStore.getState().initializeSession(serverId, client as never);
     useSessionStore
       .getState()
@@ -50,46 +65,63 @@ describe("ensureAgentIsInitialized", () => {
       serverId,
       agentId,
       client: client as never,
+      runtime,
       setAgentInitializing: bindSetAgentInitializing(),
     });
 
-    expect(client.fetchAgentTimeline).toHaveBeenCalledWith(agentId, {
-      direction: "after",
-      cursor: { epoch: "epoch-1", seq: 42 },
-      limit: TIMELINE_FETCH_PAGE_SIZE,
-      projection: "projected",
-    });
+    expect(runtime.requests).toEqual([
+      {
+        serverId,
+        agentId,
+        request: {
+          direction: "after",
+          cursor: { epoch: "epoch-1", seq: 42 },
+          limit: TIMELINE_FETCH_PAGE_SIZE,
+          projection: "projected",
+        },
+      },
+    ]);
     expect(getInitDeferred(getInitKey(serverId, agentId))?.requestDirection).toBe("after");
   });
 
   it("requests a bounded projected tail when no authoritative cursor is available", () => {
-    const client = makeClient();
+    const client = new FakeDaemonClient();
+    const runtime = new FakeTimelineRuntime();
     useSessionStore.getState().initializeSession(serverId, client as never);
 
     void ensureAgentIsInitialized({
       serverId,
       agentId,
       client: client as never,
+      runtime,
       setAgentInitializing: bindSetAgentInitializing(),
     });
 
-    expect(client.fetchAgentTimeline).toHaveBeenCalledWith(agentId, {
-      direction: "tail",
-      limit: TIMELINE_FETCH_PAGE_SIZE,
-      projection: "projected",
-    });
+    expect(runtime.requests).toEqual([
+      {
+        serverId,
+        agentId,
+        request: {
+          direction: "tail",
+          limit: TIMELINE_FETCH_PAGE_SIZE,
+          projection: "projected",
+        },
+      },
+    ]);
     expect(getInitDeferred(getInitKey(serverId, agentId))?.requestDirection).toBe("tail");
   });
 
   it("times out initialization after 65 seconds", async () => {
     vi.useFakeTimers();
-    const client = makeClient();
+    const client = new FakeDaemonClient();
+    const runtime = new FakeTimelineRuntime();
     useSessionStore.getState().initializeSession(serverId, client as never);
 
     const promise = ensureAgentIsInitialized({
       serverId,
       agentId,
       client: client as never,
+      runtime,
       setAgentInitializing: bindSetAgentInitializing(),
     });
 
@@ -108,7 +140,8 @@ describe("ensureAgentIsInitialized", () => {
 
   it("refreshes the initialization timeout after paged catch-up progress", async () => {
     vi.useFakeTimers();
-    const client = makeClient();
+    const client = new FakeDaemonClient();
+    const runtime = new FakeTimelineRuntime();
     useSessionStore.getState().initializeSession(serverId, client as never);
     const setAgentInitializing = bindSetAgentInitializing();
     const key = getInitKey(serverId, agentId);
@@ -117,6 +150,7 @@ describe("ensureAgentIsInitialized", () => {
       serverId,
       agentId,
       client: client as never,
+      runtime,
       setAgentInitializing,
     });
 
@@ -141,20 +175,29 @@ describe("ensureAgentIsInitialized", () => {
 
 describe("refreshAgent", () => {
   it("fetches a bounded projected tail after refreshing the agent", async () => {
-    const client = makeClient();
+    const client = new FakeDaemonClient();
+    const runtime = new FakeTimelineRuntime();
     useSessionStore.getState().initializeSession(serverId, client as never);
 
     await refreshAgent({
+      serverId,
       agentId,
       client: client as never,
+      runtime,
       setAgentInitializing: bindSetAgentInitializing(),
     });
 
-    expect(client.refreshAgent).toHaveBeenCalledWith(agentId);
-    expect(client.fetchAgentTimeline).toHaveBeenCalledWith(agentId, {
-      direction: "tail",
-      limit: TIMELINE_FETCH_PAGE_SIZE,
-      projection: "projected",
-    });
+    expect(client.refreshedAgentIds).toEqual([agentId]);
+    expect(runtime.requests).toEqual([
+      {
+        serverId,
+        agentId,
+        request: {
+          direction: "tail",
+          limit: TIMELINE_FETCH_PAGE_SIZE,
+          projection: "projected",
+        },
+      },
+    ]);
   });
 });

packages/app/src/hooks/use-agent-initialization.ts

@@ -10,7 +10,7 @@ import {
   rejectInitDeferred,
   refreshInitTimeout,
 } from "@/utils/agent-initialization";
-import { getHostRuntimeStore } from "@/runtime/host-runtime";
+import { getHostRuntimeStore, type HostRuntimeStore } from "@/runtime/host-runtime";
 import { planInitialAgentTimelineSync, planTimelineTailFetch } from "@/timeline/timeline-sync-plan";
 import { i18n } from "@/i18n/i18next";
 
@@ -38,6 +38,7 @@ export interface EnsureAgentIsInitializedInput {
   serverId: string;
   agentId: string;
   client: Pick<DaemonClient, "fetchAgentTimeline"> | null;
+  runtime: Pick<HostRuntimeStore, "fetchAgentTimeline">;
   setAgentInitializing: SetAgentInitializing;
   hostDisconnectedMessage?: string;
 }
@@ -69,25 +70,25 @@ export function ensureAgentIsInitialized(input: EnsureAgentIsInitializedInput):
     return deferred.promise;
   }
 
-  getHostRuntimeStore()
-    .fetchAgentTimeline(serverId, agentId, timelineRequest)
-    .catch((error) => {
-      setAgentInitializing(agentId, false);
-      rejectInitDeferred(key, error instanceof Error ? error : new Error(String(error)));
-    });
+  input.runtime.fetchAgentTimeline(serverId, agentId, timelineRequest).catch((error) => {
+    setAgentInitializing(agentId, false);
+    rejectInitDeferred(key, error instanceof Error ? error : new Error(String(error)));
+  });
 
   return deferred.promise;
 }
 
 export interface RefreshAgentInput {
+  serverId: string;
   agentId: string;
-  client: Pick<DaemonClient, "refreshAgent" | "fetchAgentTimeline"> | null;
+  client: Pick<DaemonClient, "refreshAgent"> | null;
+  runtime: Pick<HostRuntimeStore, "fetchAgentTimeline">;
   setAgentInitializing: SetAgentInitializing;
   hostDisconnectedMessage?: string;
 }
 
 export async function refreshAgent(input: RefreshAgentInput): Promise<void> {
-  const { agentId, client, setAgentInitializing } = input;
+  const { serverId, agentId, client, runtime, setAgentInitializing } = input;
   if (!client) {
     throw new Error(input.hostDisconnectedMessage ?? i18n.t("workspace.terminal.hostDisconnected"));
   }
@@ -95,11 +96,7 @@ export async function refreshAgent(input: RefreshAgentInput): Promise<void> {
 
   try {
     await client.refreshAgent(agentId);
-    const serverId = Object.entries(useSessionStore.getState().sessions).find(
-      ([, session]) => session.client === client,
-    )?.[0];
-    if (!serverId) throw new Error("Agent session is no longer connected");
-    await getHostRuntimeStore().fetchAgentTimeline(serverId, agentId, planTimelineTailFetch());
+    await runtime.fetchAgentTimeline(serverId, agentId, planTimelineTailFetch());
   } catch (error) {
     setAgentInitializing(agentId, false);
     throw error;
@@ -142,6 +139,7 @@ export function useAgentInitialization({
         serverId,
         agentId,
         client,
+        runtime: getHostRuntimeStore(),
         setAgentInitializing,
         hostDisconnectedMessage: t("workspace.terminal.hostDisconnected"),
       }),
@@ -151,12 +149,14 @@ export function useAgentInitialization({
   const refreshAgentCallback = useCallback(
     (agentId: string): Promise<void> =>
       refreshAgent({
+        serverId,
         agentId,
         client,
+        runtime: getHostRuntimeStore(),
         setAgentInitializing,
         hostDisconnectedMessage: t("workspace.terminal.hostDisconnected"),
       }),
-    [client, setAgentInitializing, t],
+    [client, serverId, setAgentInitializing, t],
   );
 
   return {