Eliminate full timeline memory load: route committed queries to DB

2026-07-29 12:01:31 +00:00 · 2026-03-25 15:10:17 +07:00
parent f27c803bf8
commit 296f283d44
7 changed files with 279 additions and 63 deletions
--- a/packages/server/src/server/agent/agent-manager.test.ts
+++ b/packages/server/src/server/agent/agent-manager.test.ts
@@ -712,7 +712,7 @@ describe("AgentManager", () => {
    expect(afterReload?.config?.title).toBeUndefined();
  });

-  test("resumeAgentFromPersistence seeds live helpers from durable rows and skips provider replay", async () => {
+  test("resumeAgentFromPersistence reads durable helpers without loading committed rows into live memory", async () => {
    const workdir = mkdtempSync(join(tmpdir(), "agent-manager-durable-seed-"));
    const storagePath = join(workdir, "agents");
    const dataDir = join(workdir, "db");
@@ -803,13 +803,8 @@ describe("AgentManager", () => {
      const resumed = await manager.resumeAgentFromPersistence(handle, undefined, snapshot.id);

      expect(resumed.id).toBe(snapshot.id);
-      expect(manager.getTimeline(snapshot.id)).toEqual([
-        {
-          type: "assistant_message",
-          text: "durable only",
-        },
-      ]);
-      expect(manager.getLastAssistantMessage(snapshot.id)).toBe("durable only");
+      expect(manager.getTimeline(snapshot.id)).toEqual([]);
+      await expect(manager.getLastAssistantMessage(snapshot.id)).resolves.toBe("durable only");
      await expect(manager.getTimelineRows(snapshot.id)).resolves.toEqual([
        {
          seq: 1,
@@ -824,12 +819,7 @@ describe("AgentManager", () => {
      await manager.hydrateTimelineFromProvider(snapshot.id);

      expect(historyReplayCount).toBe(0);
-      expect(manager.getTimeline(snapshot.id)).toEqual([
-        {
-          type: "assistant_message",
-          text: "durable only",
-        },
-      ]);
+      expect(manager.getTimeline(snapshot.id)).toEqual([]);

      await manager.closeAgent(snapshot.id);
      await manager.deleteCommittedTimeline(snapshot.id);
@@ -1261,7 +1251,7 @@ describe("AgentManager", () => {
      await durableTimelineStore.bulkInsert(snapshot.id, [durableOnlyRow]);

      expect(manager.getTimeline(snapshot.id)).toEqual([]);
-      expect(manager.getLastAssistantMessage(snapshot.id)).toBeNull();
+      await expect(manager.getLastAssistantMessage(snapshot.id)).resolves.toBe("durable only");
      await expect(manager.getTimelineRows(snapshot.id)).resolves.toEqual([durableOnlyRow]);
      await expect(
        manager.fetchTimeline(snapshot.id, {
--- a/packages/server/src/server/agent/agent-manager.ts
+++ b/packages/server/src/server/agent/agent-manager.ts
@@ -279,6 +279,7 @@ export class AgentManager {
  private readonly clients = new Map<AgentProvider, AgentClient>();
  private readonly agents = new Map<string, ActiveManagedAgent>();
  private readonly timelineStore = new InMemoryAgentTimelineStore();
+  private readonly sessionEventTails = new Map<string, Promise<void>>();
  private readonly pendingForegroundRuns = new Map<string, PendingForegroundRun>();
  private readonly subscribers = new Set<SubscriptionRecord>();
  private readonly idFactory: () => string;
@@ -526,7 +527,8 @@ export class AgentManager {
  }

  // Reconstruct an agent from provider persistence. When a durable timeline
-  // store is configured, committed history is seeded from the durable store.
+  // store is configured, the live timeline buffer only seeds seq metadata from
+  // the durable store instead of loading committed history back into memory.
  // Tests without a durable timeline store can still call
  // hydrateTimelineFromProvider() for backward compatibility.
  async resumeAgentFromPersistence(
@@ -1417,20 +1419,27 @@ export class AgentManager {
    await this.durableTimelineStore.deleteAgent(agentId);
  }

-  getLastAssistantMessage(agentId: string): string | null {
+  async getLastAssistantMessage(agentId: string): Promise<string | null> {
    const agent = this.agents.get(agentId);
    if (!agent) {
      return null;
    }

-    return this.timelineStore.getLastAssistantMessage(agentId);
+    return await this.getLastAssistantMessageFromStores(agentId);
  }

  private getLastAssistantMessageFromTimeline(
    timeline: readonly AgentTimelineItem[],
  ): string | null {
+    return this.getLastAssistantMessageSegmentFromTimeline(timeline)?.text ?? null;
+  }
+
+  private getLastAssistantMessageSegmentFromTimeline(
+    timeline: readonly AgentTimelineItem[],
+  ): { text: string; startsAtBeginning: boolean } | null {
    // Collect the last contiguous assistant messages (Claude streams chunks)
    const chunks: string[] = [];
+    let startsAtBeginning = false;
    for (let i = timeline.length - 1; i >= 0; i--) {
      const item = timeline[i];
      if (item.type !== "assistant_message") {
@@ -1440,13 +1449,65 @@ export class AgentManager {
        continue;
      }
      chunks.push(item.text);
+      startsAtBeginning = i === 0;
    }

    if (!chunks.length) {
      return null;
    }

-    return chunks.reverse().join("");
+    return {
+      text: chunks.reverse().join(""),
+      startsAtBeginning,
+    };
+  }
+
+  private async getLastAssistantMessageFromStores(agentId: string): Promise<string | null> {
+    const liveTimeline = this.timelineStore.getItems(agentId);
+    const liveSegment = this.getLastAssistantMessageSegmentFromTimeline(liveTimeline);
+    if (!this.durableTimelineStore) {
+      return liveSegment?.text ?? null;
+    }
+
+    if (!liveSegment) {
+      return await this.durableTimelineStore.getLastAssistantMessage(agentId);
+    }
+
+    if (!liveSegment.startsAtBeginning) {
+      return liveSegment.text;
+    }
+
+    const lastDurableItem = await this.durableTimelineStore.getLastItem(agentId);
+    if (lastDurableItem?.type !== "assistant_message") {
+      return liveSegment.text;
+    }
+
+    const durableMessage = await this.durableTimelineStore.getLastAssistantMessage(agentId);
+    return durableMessage ? `${durableMessage}${liveSegment.text}` : liveSegment.text;
+  }
+
+  private async getLastItemFromStores(agentId: string): Promise<AgentTimelineItem | null> {
+    const lastLiveItem = this.timelineStore.getLastItem(agentId);
+    if (lastLiveItem) {
+      return lastLiveItem;
+    }
+    if (!this.durableTimelineStore) {
+      return null;
+    }
+    return await this.durableTimelineStore.getLastItem(agentId);
+  }
+
+  private async hasCommittedUserMessageFromStores(
+    agentId: string,
+    options: { messageId: string; text: string },
+  ): Promise<boolean> {
+    if (this.timelineStore.hasCommittedUserMessage(agentId, options)) {
+      return true;
+    }
+    if (!this.durableTimelineStore) {
+      return false;
+    }
+    return await this.durableTimelineStore.hasCommittedUserMessage(agentId, options);
  }

  async waitForAgentEvent(
@@ -1466,7 +1527,7 @@ export class AgentManager {
      return {
        status: snapshot.lifecycle,
        permission: immediatePermission,
-        lastMessage: this.getLastAssistantMessage(agentId),
+        lastMessage: await this.getLastAssistantMessage(agentId),
      };
    }

@@ -1477,14 +1538,14 @@ export class AgentManager {
      return {
        status: initialStatus,
        permission: null,
-        lastMessage: this.getLastAssistantMessage(agentId),
+        lastMessage: await this.getLastAssistantMessage(agentId),
      };
    }
    if (waitForActive && !initialBusy && !hasForegroundTurn) {
      return {
        status: initialStatus,
        permission: null,
-        lastMessage: this.getLastAssistantMessage(agentId),
+        lastMessage: await this.getLastAssistantMessage(agentId),
      };
    }

@@ -1503,6 +1564,7 @@ export class AgentManager {
      let currentStatus: AgentLifecycleStatus = initialStatus;
      let hasStarted = initialBusy || hasForegroundTurn;
      let terminalStatusOverride: AgentLifecycleStatus | null = null;
+      let finished = false;

      // Bug #3 Fix: Declare unsubscribe and abortHandler upfront so cleanup can reference them
      let unsubscribe: (() => void) | null = null;
@@ -1531,12 +1593,20 @@ export class AgentManager {
      };

      const finish = (permission: AgentPermissionRequest | null) => {
+        if (finished) {
+          return;
+        }
+        finished = true;
        cleanup();
-        resolve({
-          status: currentStatus,
-          permission,
-          lastMessage: this.getLastAssistantMessage(agentId),
-        });
+        void this.getLastAssistantMessage(agentId)
+          .then((lastMessage) => {
+            resolve({
+              status: currentStatus,
+              permission,
+              lastMessage,
+            });
+          })
+          .catch(reject);
      };

      // Bug #3 Fix: Set up abort handler BEFORE subscription
@@ -1711,18 +1781,9 @@ export class AgentManager {
      return { timestamp: now.toISOString() };
    }

-    const rows = await this.durableTimelineStore.getCommittedRows(agentId);
-    if (rows.length === 0) {
-      return {
-        nextSeq: 1,
-        timestamp: now.toISOString(),
-      };
-    }
-
    return {
-      rows,
-      nextSeq: rows[rows.length - 1]!.seq + 1,
-      timestamp: rows[rows.length - 1]!.timestamp,
+      nextSeq: (await this.durableTimelineStore.getLatestCommittedSeq(agentId)) + 1,
+      timestamp: now.toISOString(),
    };
  }

@@ -1732,16 +1793,42 @@ export class AgentManager {
    }
    const agentId = agent.id;
    const unsubscribe = agent.session.subscribe((event: AgentStreamEvent) => {
-      const current = this.agents.get(agentId);
-      if (!current) {
-        return;
-      }
-      this.dispatchSessionEvent(current, event);
+      this.enqueueSessionEvent(agentId, event);
    });
    agent.unsubscribeSession = unsubscribe;
  }

-  private dispatchSessionEvent(agent: ActiveManagedAgent, event: AgentStreamEvent): void {
+  private enqueueSessionEvent(agentId: string, event: AgentStreamEvent): void {
+    const previous = this.sessionEventTails.get(agentId) ?? Promise.resolve();
+    const next = previous
+      .catch(() => undefined)
+      .then(async () => {
+        const current = this.agents.get(agentId);
+        if (!current) {
+          return;
+        }
+        await this.dispatchSessionEvent(current, event);
+      })
+      .catch((err) => {
+        this.logger.error(
+          { err, agentId, eventType: event.type },
+          "Failed to process session event",
+        );
+      });
+
+    this.sessionEventTails.set(agentId, next);
+    this.trackBackgroundTask(next);
+    void next.finally(() => {
+      if (this.sessionEventTails.get(agentId) === next) {
+        this.sessionEventTails.delete(agentId);
+      }
+    });
+  }
+
+  private async dispatchSessionEvent(
+    agent: ActiveManagedAgent,
+    event: AgentStreamEvent,
+  ): Promise<void> {
    const turnId = (event as { turnId?: string }).turnId;
    const matchingWaiters =
      turnId == null
@@ -1750,7 +1837,7 @@ export class AgentManager {
            (waiter) => waiter.turnId === turnId && !waiter.settled,
          );

-    this.handleStreamEvent(agent, event);
+    await this.handleStreamEvent(agent, event);

    for (const waiter of matchingWaiters) {
      waiter.callback(event);
@@ -1960,14 +2047,14 @@ export class AgentManager {
    }
  }

-  private handleStreamEvent(
+  private async handleStreamEvent(
    agent: ActiveManagedAgent,
    event: AgentStreamEvent,
    options?: {
      fromHistory?: boolean;
      canonicalUserMessagesById?: ReadonlyMap<string, string>;
    },
-  ): void {
+  ): Promise<void> {
    const eventTurnId = (event as { turnId?: string }).turnId;
    const isForegroundEvent = Boolean(
      eventTurnId && agent.activeForegroundTurnId === eventTurnId,
@@ -2017,7 +2104,7 @@ export class AgentManager {
          const eventText = event.item.text;
          if (eventMessageId) {
            if (
-              this.timelineStore.hasCommittedUserMessage(agent.id, {
+              await this.hasCommittedUserMessageFromStores(agent.id, {
                messageId: eventMessageId,
                text: eventText,
              })
@@ -2103,7 +2190,7 @@ export class AgentManager {
          agent.lifecycle = "error";
        }
        agent.lastError = event.error;
-        this.appendSystemErrorTimelineMessage(
+        await this.appendSystemErrorTimelineMessage(
          agent,
          event.provider,
          this.formatTurnFailedMessage(event),
@@ -2209,7 +2296,7 @@ export class AgentManager {
    }
  }

-  private appendSystemErrorTimelineMessage(
+  private async appendSystemErrorTimelineMessage(
    agent: ActiveManagedAgent,
    provider: AgentProvider,
    message: string,
@@ -2217,7 +2304,7 @@ export class AgentManager {
      fromHistory?: boolean;
      canonicalUserMessagesById?: ReadonlyMap<string, string>;
    },
-  ): void {
+  ): Promise<void> {
    if (options?.fromHistory) {
      return;
    }
@@ -2228,7 +2315,7 @@ export class AgentManager {
    }

    const text = `${SYSTEM_ERROR_PREFIX} ${normalized}`;
-    const lastItem = this.timelineStore.getLastItem(agent.id);
+    const lastItem = await this.getLastItemFromStores(agent.id);
    if (lastItem?.type === "assistant_message" && lastItem.text === text) {
      return;
    }
--- a/packages/server/src/server/agent/agent-timeline-store-types.ts
+++ b/packages/server/src/server/agent/agent-timeline-store-types.ts
@@ -49,6 +49,12 @@ export interface AgentTimelineStore {
  ): Promise<AgentTimelineFetchResult>;
  getLatestCommittedSeq(agentId: string): Promise<number>;
  getCommittedRows(agentId: string): Promise<AgentTimelineRow[]>;
+  getLastItem(agentId: string): Promise<AgentTimelineItem | null>;
+  getLastAssistantMessage(agentId: string): Promise<string | null>;
+  hasCommittedUserMessage(
+    agentId: string,
+    options: { messageId: string; text: string },
+  ): Promise<boolean>;
  deleteAgent(agentId: string): Promise<void>;
  bulkInsert(agentId: string, rows: readonly AgentTimelineRow[]): Promise<void>;
 }
--- a/packages/server/src/server/db/db-agent-timeline-store.test.ts
+++ b/packages/server/src/server/db/db-agent-timeline-store.test.ts
@@ -211,4 +211,56 @@ describe("DbAgentTimelineStore", () => {
      }),
    ]);
  });
+
+  test("getLastItem returns the latest committed item", async () => {
+    await store.bulkInsert("agent-1", [
+      createRow(1, createTimelineItem("user_message", "1")),
+      createRow(2, createTimelineItem("assistant_message", "2")),
+    ]);
+
+    await expect(store.getLastItem("agent-1")).resolves.toEqual(
+      createTimelineItem("assistant_message", "2"),
+    );
+    await expect(store.getLastItem("missing-agent")).resolves.toBeNull();
+  });
+
+  test("getLastAssistantMessage assembles the latest contiguous assistant chunks", async () => {
+    await store.bulkInsert("agent-1", [
+      createRow(1, createTimelineItem("assistant_message", "1")),
+      createRow(2, createTimelineItem("assistant_message", "2")),
+      createRow(3, { type: "reasoning", text: "separator-1" }),
+      createRow(4, createTimelineItem("assistant_message", "4")),
+      createRow(5, createTimelineItem("assistant_message", "5")),
+      createRow(6, { type: "reasoning", text: "separator-2" }),
+    ]);
+
+    await expect(store.getLastAssistantMessage("agent-1")).resolves.toBe("assistant-4assistant-5");
+    await expect(store.getLastAssistantMessage("missing-agent")).resolves.toBeNull();
+  });
+
+  test("hasCommittedUserMessage matches by normalized messageId and text", async () => {
+    await store.bulkInsert("agent-1", [
+      createRow(1, createTimelineItem("user_message", "1")),
+      createRow(2, createTimelineItem("assistant_message", "2")),
+    ]);
+
+    await expect(
+      store.hasCommittedUserMessage("agent-1", {
+        messageId: " message-1 ",
+        text: "user-1",
+      }),
+    ).resolves.toBe(true);
+    await expect(
+      store.hasCommittedUserMessage("agent-1", {
+        messageId: "message-1",
+        text: "different",
+      }),
+    ).resolves.toBe(false);
+    await expect(
+      store.hasCommittedUserMessage("agent-1", {
+        messageId: "   ",
+        text: "user-1",
+      }),
+    ).resolves.toBe(false);
+  });
 });
--- a/packages/server/src/server/db/db-agent-timeline-store.ts
+++ b/packages/server/src/server/db/db-agent-timeline-store.ts
@@ -16,6 +16,14 @@ type AgentTimelineRowInsert = typeof agentTimelineRows.$inferInsert;

 const DEFAULT_TIMELINE_FETCH_LIMIT = 200;

+function normalizeTimelineMessageId(messageId: string | undefined): string | undefined {
+  if (typeof messageId !== "string") {
+    return undefined;
+  }
+  const normalized = messageId.trim();
+  return normalized.length > 0 ? normalized : undefined;
+}
+
 function toTimelineRow(row: AgentTimelineRowRecord): AgentTimelineRow {
  return {
    seq: row.seq,
@@ -185,6 +193,76 @@ export class DbAgentTimelineStore implements AgentTimelineStore {
    return rows.map(toTimelineRow);
  }

+  async getLastItem(agentId: string): Promise<AgentTimelineItem | null> {
+    const [row] = await this.db
+      .select({ item: agentTimelineRows.item })
+      .from(agentTimelineRows)
+      .where(eq(agentTimelineRows.agentId, agentId))
+      .orderBy(desc(agentTimelineRows.seq))
+      .limit(1);
+    return row?.item ?? null;
+  }
+
+  async getLastAssistantMessage(agentId: string): Promise<string | null> {
+    const rows = await this.db
+      .select({
+        seq: agentTimelineRows.seq,
+        item: agentTimelineRows.item,
+      })
+      .from(agentTimelineRows)
+      .where(
+        and(
+          eq(agentTimelineRows.agentId, agentId),
+          eq(agentTimelineRows.itemKind, "assistant_message"),
+        ),
+      )
+      .orderBy(desc(agentTimelineRows.seq));
+
+    if (rows.length === 0) {
+      return null;
+    }
+
+    const chunks: string[] = [];
+    let previousSeq: number | null = null;
+    for (const row of rows) {
+      if (previousSeq !== null && row.seq !== previousSeq - 1) {
+        break;
+      }
+      if (row.item.type !== "assistant_message") {
+        break;
+      }
+      chunks.push(row.item.text);
+      previousSeq = row.seq;
+    }
+
+    return chunks.length > 0 ? chunks.reverse().join("") : null;
+  }
+
+  async hasCommittedUserMessage(
+    agentId: string,
+    options: { messageId: string; text: string },
+  ): Promise<boolean> {
+    const messageId = normalizeTimelineMessageId(options.messageId);
+    if (!messageId) {
+      return false;
+    }
+
+    const [row] = await this.db
+      .select({ seq: agentTimelineRows.seq })
+      .from(agentTimelineRows)
+      .where(
+        and(
+          eq(agentTimelineRows.agentId, agentId),
+          eq(agentTimelineRows.itemKind, "user_message"),
+          sql`${agentTimelineRows.item} ->> 'messageId' = ${messageId}`,
+          sql`${agentTimelineRows.item} ->> 'text' = ${options.text}`,
+        ),
+      )
+      .limit(1);
+
+    return row !== undefined;
+  }
+
  async deleteAgent(agentId: string): Promise<void> {
    await this.db.delete(agentTimelineRows).where(eq(agentTimelineRows.agentId, agentId));
  }
--- a/packages/server/src/server/websocket-server.notifications.test.ts
+++ b/packages/server/src/server/websocket-server.notifications.test.ts
@@ -63,7 +63,7 @@ function createServer(agentManagerOverrides?: Record<string, unknown>) {
  const agentManager = {
    setAgentAttentionCallback: vi.fn(),
    getAgent: vi.fn(() => null),
-    getLastAssistantMessage: vi.fn(() => null),
+    getLastAssistantMessage: vi.fn(async () => null),
    ...agentManagerOverrides,
  };

@@ -93,9 +93,9 @@ describe("VoiceAssistantWebSocketServer notification payloads", () => {
    vi.clearAllMocks();
  });

-  it("uses assistant preview text for push notifications with markdown removed", () => {
+  it("uses assistant preview text for push notifications with markdown removed", async () => {
    const getLastAssistantMessage = vi.fn(
-      () => "**Done**. Updated `README.md` and [link](https://example.com).",
+      async () => "**Done**. Updated `README.md` and [link](https://example.com).",
    );
    const { server } = createServer({
      getAgent: vi.fn(() => ({
@@ -106,7 +106,7 @@ describe("VoiceAssistantWebSocketServer notification payloads", () => {
      getLastAssistantMessage,
    });

-    (server as any).broadcastAgentAttention({
+    await (server as any).broadcastAgentAttention({
      agentId: "agent-1",
      provider: "claude",
      reason: "finished",
@@ -124,8 +124,8 @@ describe("VoiceAssistantWebSocketServer notification payloads", () => {
    expect(getLastAssistantMessage).toHaveBeenCalledWith("agent-1");
  });

-  it("sends push notifications regardless of UI label presence", () => {
-    const getLastAssistantMessage = vi.fn(() => "Done.");
+  it("sends push notifications regardless of UI label presence", async () => {
+    const getLastAssistantMessage = vi.fn(async () => "Done.");
    const { server } = createServer({
      getAgent: vi.fn(() => ({
        config: { title: null },
@@ -136,7 +136,7 @@ describe("VoiceAssistantWebSocketServer notification payloads", () => {
      getLastAssistantMessage,
    });

-    (server as any).broadcastAgentAttention({
+    await (server as any).broadcastAgentAttention({
      agentId: "agent-2",
      provider: "claude",
      reason: "finished",
--- a/packages/server/src/server/websocket-server.ts
+++ b/packages/server/src/server/websocket-server.ts
@@ -342,7 +342,9 @@ export class VoiceAssistantWebSocketServer {
    this.pushService = new PushService(pushLogger, this.pushTokenStore);

    this.agentManager.setAgentAttentionCallback((params) => {
-      this.broadcastAgentAttention(params);
+      void this.broadcastAgentAttention(params).catch((err) => {
+        this.logger.warn({ err, agentId: params.agentId }, "Failed to broadcast agent attention");
+      });
    });

    const { allowedOrigins, allowedHosts } = wsConfig;
@@ -1239,11 +1241,11 @@ export class VoiceAssistantWebSocketServer {
    };
  }

-  private broadcastAgentAttention(params: {
+  private async broadcastAgentAttention(params: {
    agentId: string;
    provider: AgentProvider;
    reason: "finished" | "error" | "permission";
-  }): void {
+  }): Promise<void> {
    const clientEntries: Array<{
      ws: WebSocketLike;
      state: ClientAttentionState;
@@ -1258,11 +1260,12 @@ export class VoiceAssistantWebSocketServer {

    const allStates = clientEntries.map((e) => e.state);
    const agent = this.agentManager.getAgent(params.agentId);
+    const assistantMessage = await this.agentManager.getLastAssistantMessage(params.agentId);
    const notification = buildAgentAttentionNotificationPayload({
      reason: params.reason,
      serverId: this.serverId,
      agentId: params.agentId,
-      assistantMessage: this.agentManager.getLastAssistantMessage(params.agentId),
+      assistantMessage,
      permissionRequest: agent ? findLatestPermissionRequest(agent.pendingPermissions) : null,
    });