perf(memory): bound renderer sessions, terminals and polling (#3073)

Author: charlesvien

packages/core/src/archive/archiveOrchestration.test.ts

@@ -18,9 +18,7 @@ class Harness {
     unpin: vi.fn().mockResolvedValue(undefined),
     togglePin: vi.fn().mockResolvedValue(undefined),
     navigateAwayFromTaskIfActive: vi.fn(),
-    snapshotTerminalStates: vi.fn().mockReturnValue({}),
     clearTerminalStates: vi.fn(),
-    restoreTerminalStates: vi.fn(),
     snapshotCommandCenter: vi
       .fn()
       .mockReturnValue({ index: -1, wasActive: false }),
@@ -103,6 +101,27 @@ describe("archiveTask", () => {
     expect(harness.list).toEqual([]);
     expect(harness.deps.togglePin).toHaveBeenCalledWith(TASK_ID);
   });
+
+  it("destroys terminals only after the archive succeeds", async () => {
+    let clearedWhenArchiveCalled = true;
+    harness.deps.archive = vi.fn().mockImplementation(async () => {
+      clearedWhenArchiveCalled =
+        vi.mocked(harness.deps.clearTerminalStates).mock.calls.length > 0;
+    });
+
+    await archiveTask(TASK_ID, harness.deps);
+
+    expect(clearedWhenArchiveCalled).toBe(false);
+    expect(harness.deps.clearTerminalStates).toHaveBeenCalledWith(TASK_ID);
+  });
+
+  it("keeps terminals when archive fails", async () => {
+    harness.deps.archive = vi.fn().mockRejectedValue(new Error("boom"));
+
+    await expect(archiveTask(TASK_ID, harness.deps)).rejects.toThrow("boom");
+
+    expect(harness.deps.clearTerminalStates).not.toHaveBeenCalled();
+  });
 });
 
 describe("archiveTasks", () => {

packages/core/src/archive/archiveOrchestration.ts

@@ -27,9 +27,7 @@ export interface ArchiveOrchestrationDeps {
   unpin(taskId: string): Promise<void>;
   togglePin(taskId: string): Promise<void>;
   navigateAwayFromTaskIfActive(taskId: string): void;
-  snapshotTerminalStates(taskId: string): Record<string, unknown>;
   clearTerminalStates(taskId: string): void;
-  restoreTerminalStates(states: Record<string, unknown>): void;
   snapshotCommandCenter(taskId: string): { index: number; wasActive: boolean };
   removeFromCommandCenter(taskId: string): void;
   restoreCommandCenter(
@@ -69,11 +67,9 @@ export async function archiveTask(
     deps.navigateAwayFromTaskIfActive(taskId);
   }
 
-  const terminalStatesSnapshot = deps.snapshotTerminalStates(taskId);
   const commandCenterSnapshot = deps.snapshotCommandCenter(taskId);
 
   await deps.unpin(taskId);
-  deps.clearTerminalStates(taskId);
   deps.removeFromCommandCenter(taskId);
 
   await deps.cache.cancelPathFilter();
@@ -101,6 +97,9 @@ export async function archiveTask(
   try {
     await deps.disconnectFromTask(taskId);
     await deps.archive(taskId);
+    // Destroying terminals is irreversible, so it waits for the archive to
+    // commit; a failed archive keeps its live terminals.
+    deps.clearTerminalStates(taskId);
     // Non-optimistic flows keep the row visible during the request, then remove
     // it the moment the archive succeeds.
     if (!optimistic) {
@@ -115,9 +114,6 @@ export async function archiveTask(
     if (wasPinned) {
       await deps.togglePin(taskId);
     }
-    if (Object.keys(terminalStatesSnapshot).length > 0) {
-      deps.restoreTerminalStates(terminalStatesSnapshot);
-    }
     if (commandCenterSnapshot.index !== -1) {
       deps.restoreCommandCenter(taskId, commandCenterSnapshot);
     }

packages/core/src/sessions/sessionEviction.test.ts

@@ -0,0 +1,142 @@
+import type { AgentSession } from "@posthog/shared";
+import { describe, expect, it } from "vitest";
+import { getCellCount, type LayoutPreset } from "../command-center/grid";
+import {
+  isSessionIdle,
+  MAX_CONNECTED_SESSIONS,
+  selectSessionsToEvict,
+} from "./sessionEviction";
+
+function makeSession(overrides: Partial<AgentSession>): AgentSession {
+  return {
+    taskRunId: `run-${overrides.taskId}`,
+    status: "connected",
+    isPromptPending: false,
+    pendingPermissions: new Map(),
+    messageQueue: [],
+    startedAt: 0,
+    ...overrides,
+  } as AgentSession;
+}
+
+describe("isSessionIdle", () => {
+  it.each([
+    ["connected idle local session", {}, true],
+    ["connecting session", { status: "connecting" as const }, false],
+    ["pending prompt", { isPromptPending: true }, false],
+    ["compacting session", { isCompacting: true }, false],
+    ["handoff in progress", { handoffInProgress: true }, false],
+    [
+      "pending permission",
+      { pendingPermissions: new Map([["p1", {} as never]]) },
+      false,
+    ],
+    [
+      "queued messages",
+      { messageQueue: [{ id: "m1", content: "x", queuedAt: 0 }] },
+      false,
+    ],
+    [
+      "running cloud session",
+      { isCloud: true, cloudStatus: "in_progress" as const },
+      false,
+    ],
+    [
+      "queued cloud session",
+      { isCloud: true, cloudStatus: "queued" as const },
+      false,
+    ],
+    [
+      "completed cloud session",
+      { isCloud: true, cloudStatus: "completed" as const },
+      true,
+    ],
+    ["cloud session without status", { isCloud: true }, false],
+    ["disconnected local session", { status: "disconnected" as const }, true],
+    ["errored local session", { status: "error" as const }, true],
+  ])("%s -> %s", (_name, overrides, expected) => {
+    expect(isSessionIdle(makeSession({ taskId: "t", ...overrides }))).toBe(
+      expected,
+    );
+  });
+});
+
+describe("selectSessionsToEvict", () => {
+  const lastUsedAt = (session: AgentSession) => session.startedAt;
+
+  it.each([
+    [
+      "returns nothing under the budget",
+      {
+        sessions: [makeSession({ taskId: "a" }), makeSession({ taskId: "b" })],
+        activeTaskId: "a",
+        maxSessions: 3,
+      },
+      [],
+    ],
+    [
+      "evicts the least recently used idle sessions over the budget",
+      {
+        sessions: [
+          makeSession({ taskId: "a", startedAt: 30 }),
+          makeSession({ taskId: "b", startedAt: 10 }),
+          makeSession({ taskId: "c", startedAt: 20 }),
+          makeSession({ taskId: "d", startedAt: 40 }),
+        ],
+        activeTaskId: "d",
+        maxSessions: 3,
+      },
+      ["b", "c"],
+    ],
+    [
+      "never evicts the active task or busy sessions",
+      {
+        sessions: [
+          makeSession({ taskId: "active", startedAt: 1 }),
+          makeSession({ taskId: "busy", startedAt: 2, isPromptPending: true }),
+          makeSession({ taskId: "idle", startedAt: 3 }),
+        ],
+        activeTaskId: "active",
+        maxSessions: 2,
+      },
+      ["idle"],
+    ],
+    [
+      "never evicts mounted tasks",
+      {
+        sessions: [
+          makeSession({ taskId: "a", startedAt: 1 }),
+          makeSession({ taskId: "b", startedAt: 2 }),
+          makeSession({ taskId: "c", startedAt: 3 }),
+        ],
+        activeTaskId: "c",
+        protectedTaskIds: new Set(["a"]),
+        maxSessions: 2,
+      },
+      ["b"],
+    ],
+  ])("%s", (_name, params, expected) => {
+    const evicted = selectSessionsToEvict({ ...params, lastUsedAt });
+    expect(evicted.map((s) => s.taskId)).toEqual(expected);
+  });
+});
+
+describe("MAX_CONNECTED_SESSIONS", () => {
+  it("stays above the largest Command Center grid so full layouts never evict", () => {
+    // Record<LayoutPreset, ...> forces this list to grow with the union, so a
+    // new larger preset breaks this test instead of silently churning cells.
+    const allPresets: Record<LayoutPreset, true> = {
+      "1x1": true,
+      "2x1": true,
+      "1x2": true,
+      "2x2": true,
+      "3x2": true,
+      "3x3": true,
+    };
+    const largestGrid = Math.max(
+      ...(Object.keys(allPresets) as LayoutPreset[]).map(getCellCount),
+    );
+
+    expect(MAX_CONNECTED_SESSIONS).toBeGreaterThan(largestGrid);
+  });
+});

packages/core/src/sessions/sessionEviction.ts

@@ -0,0 +1,42 @@
+import type { AgentSession } from "@posthog/shared";
+import { isTerminalStatus } from "@posthog/shared/domain-types";
+
+// Above the Command Center's 3x3 grid so fully-visible layouts never evict.
+export const MAX_CONNECTED_SESSIONS = 12;
+
+export function isSessionIdle(session: AgentSession): boolean {
+  if (session.status === "connecting") return false;
+  if (session.isPromptPending) return false;
+  if (session.isCompacting) return false;
+  if (session.handoffInProgress) return false;
+  if (session.pendingPermissions.size > 0) return false;
+  if (session.messageQueue.length > 0) return false;
+  if (session.isCloud) return isTerminalStatus(session.cloudStatus);
+  return true;
+}
+
+export function selectSessionsToEvict(params: {
+  sessions: AgentSession[];
+  activeTaskId: string;
+  protectedTaskIds?: ReadonlySet<string>;
+  lastUsedAt: (session: AgentSession) => number;
+  maxSessions?: number;
+}): AgentSession[] {
+  const { sessions, activeTaskId, protectedTaskIds, lastUsedAt } = params;
+  const maxSessions = params.maxSessions ?? MAX_CONNECTED_SESSIONS;
+
+  // Reserves a slot for the incoming session even when a resume replaces an
+  // existing one; deliberately over-evicts by one in that case.
+  const excess = sessions.length - (maxSessions - 1);
+  if (excess <= 0) return [];
+
+  return sessions
+    .filter(
+      (session) =>
+        session.taskId !== activeTaskId &&
+        !protectedTaskIds?.has(session.taskId) &&
+        isSessionIdle(session),
+    )
+    .sort((a, b) => lastUsedAt(a) - lastUsedAt(b))
+    .slice(0, excess);
+}

packages/core/src/sessions/sessionService.ts

@@ -77,6 +77,7 @@ import {
   promptReferencesAbsoluteFolder,
   shellExecutesToContextBlocks,
 } from "./sessionEvents";
+import { selectSessionsToEvict } from "./sessionEviction";
 import { createBaseSession } from "./sessionFactory";
 import { type ParsedSessionLogs, parseSessionLogContent } from "./sessionLogs";
 
@@ -513,6 +514,8 @@ export class SessionService {
   >();
   private localRepoPaths = new Map<string, string>();
   private localRecoveryAttempts = new Map<string, Promise<boolean>>();
+  private sessionLastUsedAt = new Map<string, number>();
+  private mountedTaskCounts = new Map<string, number>();
   /** Re-entrance guard for cloud queue dispatch (per taskId). */
   private dispatchingCloudQueues = new Set<string>();
   /** Coalesces deferred cloud queue flush timers (per taskId). */
@@ -610,6 +613,8 @@ export class SessionService {
     const { task } = params;
     const taskId = task.id;
     this.localRepoPaths.set(taskId, params.repoPath);
+    this.sessionLastUsedAt.set(taskId, Date.now());
+    void this.evictIdleSessions(taskId);
 
     // Return existing connection promise if already connecting
     const existingPromise = this.connectingTasks.get(taskId);
@@ -1029,7 +1034,10 @@ export class SessionService {
     }
   }
 
-  private async teardownSession(taskRunId: string): Promise<void> {
+  private async teardownSession(
+    taskRunId: string,
+    opts?: { preserveResumeState?: boolean },
+  ): Promise<void> {
     const session = this.getSessionByRunId(taskRunId);
 
     try {
@@ -1053,9 +1061,14 @@ export class SessionService {
     if (session) {
       this.localRepoPaths.delete(session.taskId);
       this.localRecoveryAttempts.delete(session.taskId);
+      this.sessionLastUsedAt.delete(session.taskId);
+    }
+    if (!opts?.preserveResumeState) {
+      // Reconnect restores the model and permission mode from these; only a
+      // permanent disconnect (archive, delete, fresh session) may drop them.
+      this.d.adapterStore.removeAdapter(taskRunId);
+      this.d.removePersistedConfigOptions(taskRunId);
     }
-    this.d.adapterStore.removeAdapter(taskRunId);
-    this.d.removePersistedConfigOptions(taskRunId);
   }
 
   /**
@@ -1358,6 +1371,51 @@ export class SessionService {
     await this.teardownSession(session.taskRunId);
   }
 
+  registerMountedTask(taskId: string): () => void {
+    this.mountedTaskCounts.set(
+      taskId,
+      (this.mountedTaskCounts.get(taskId) ?? 0) + 1,
+    );
+    this.sessionLastUsedAt.set(taskId, Date.now());
+    return () => {
+      const count = this.mountedTaskCounts.get(taskId) ?? 0;
+      if (count <= 1) {
+        this.mountedTaskCounts.delete(taskId);
+      } else {
+        this.mountedTaskCounts.set(taskId, count - 1);
+      }
+      this.sessionLastUsedAt.set(taskId, Date.now());
+    };
+  }
+
+  private async evictIdleSessions(activeTaskId: string): Promise<void> {
+    const toEvict = selectSessionsToEvict({
+      sessions: Object.values(this.d.store.getSessions()),
+      activeTaskId,
+      protectedTaskIds: new Set(this.mountedTaskCounts.keys()),
+      lastUsedAt: (session) =>
+        this.sessionLastUsedAt.get(session.taskId) ?? session.startedAt,
+    });
+
+    for (const session of toEvict) {
+      this.d.log.info("Evicting idle session to bound memory", {
+        taskId: session.taskId,
+        taskRunId: session.taskRunId,
+      });
+      this.sessionLastUsedAt.delete(session.taskId);
+      try {
+        await this.teardownSession(session.taskRunId, {
+          preserveResumeState: true,
+        });
+      } catch (error) {
+        this.d.log.error("Failed to evict idle session", {
+          taskId: session.taskId,
+          error,
+        });
+      }
+    }
+  }
+
   // --- Subscription Management ---
 
   /** Streamed events awaiting their frame flush, keyed by taskRunId. Order
@@ -1598,6 +1656,7 @@ export class SessionService {
     this.connectingTasks.clear();
     this.localRepoPaths.clear();
     this.localRecoveryAttempts.clear();
+    this.sessionLastUsedAt.clear();
     this.cloudPermissionRequestIds.clear();
     this.liveTurnContent.clear();
     this.cloudLogGapReconciler.clear();
@@ -4285,6 +4344,10 @@ export class SessionService {
     } = params;
 
     if (isCloud) {
+      // Local connects bound the session budget inside connectToTask; cloud
+      // watches would otherwise never trigger eviction.
+      this.sessionLastUsedAt.set(task.id, Date.now());
+      void this.evictIdleSessions(task.id);
       return this.reconcileCloudConnection(
         task,
         cloudAuth,