perf(sessions): coalesce each flush batch into one store append per notification run

Every 16ms flush handed its buffered events to handleSessionEvent one at a
time, so a 15-event burst did 15 setState calls through the session store —
each one an O(transcript) walk (session-map copy + subscriber notification
fan-out). On a 30k-event transcript that is ~248ms of main-thread time per
second of heavy streaming, burned before React paints anything.

handleSessionEvents now coalesces runs of plain notifications (agent message
chunks, tool-call updates — the bulk of a turn) into a single appendEvents
call per flush. Events carrying a JSON-RPC id (user-prompt echoes,
stop-reason responses) still go through handleSessionEvent individually, so
turn-lifecycle semantics are untouched; global order is preserved by flushing
the pending run before each id-carrying event. The notification side effects
(config/usage/adapter/compaction updates) split into applyNotificationEffects
so both paths share them.

Measured with scripts/measure-append-hotpath.ts (real store module, V8 via
tsx): 60 flushes x 15 events over a 30k-event transcript drops from ~248ms to
~26ms per streamed second (~10x); store subscriber notifications drop from
900/s to 60/s.

Generated-By: PostHog Code
Task-Id: e3bcc511-15af-4f46-a3c4-7b29e65f026b

Author: KarloAldrete

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

@@ -32,6 +32,35 @@ function chunkText(event: AcpMessage): string {
   return params.update.content.text;
 }
 
+/** The renderer-side echo of a user prompt — a JSON-RPC request (carries an
+ * id), so it must not be coalesced with plain notifications. */
+function promptEcho(id: number, text: string): AcpMessage {
+  return {
+    ts: 1,
+    message: {
+      jsonrpc: "2.0",
+      id,
+      method: "session/prompt",
+      params: {
+        sessionId: RUN_ID,
+        prompt: [{ type: "text", text }],
+      },
+    },
+  } as unknown as AcpMessage;
+}
+
+/** The agent's terminal response for a prompt — completes the turn. */
+function stopResponse(id: number): AcpMessage {
+  return {
+    ts: 2,
+    message: {
+      jsonrpc: "2.0",
+      id,
+      result: { stopReason: "end_turn" },
+    },
+  } as unknown as AcpMessage;
+}
+
 function createHarness() {
   const sessions: Record<string, AgentSession> = {
     [RUN_ID]: {
@@ -61,8 +90,11 @@ function createHarness() {
       sessions[session.taskRunId] = session;
     },
     updateSession: (taskRunId: string, updates: Partial<AgentSession>) => {
+      // Like the real store: produce a NEW session object per update, so a
+      // reference captured before the update keeps its pre-update values
+      // (handleSessionEvent's stop-reason check depends on that).
       const session = sessions[taskRunId];
-      if (session) Object.assign(session, updates);
+      if (session) sessions[taskRunId] = { ...session, ...updates };
     },
     appendEvents,
     replaceOptimisticWithEvent: vi.fn(),
@@ -80,10 +112,11 @@ function createHarness() {
     debug: vi.fn(),
   };
 
+  const notifyPromptComplete = vi.fn();
   const deps = {
     store,
     log: noopLog,
-    notifyPromptComplete: vi.fn(),
+    notifyPromptComplete,
     notifyPermissionRequest: vi.fn(),
     taskViewedApi: { markActivity: vi.fn() },
     getPersistedConfigOptions: () => undefined,
@@ -120,6 +153,7 @@ function createHarness() {
   return {
     service,
     appendEvents,
+    notifyPromptComplete,
     emit: (event: AcpMessage) => onEvent?.(event),
     events: () => sessions[RUN_ID].events,
   };
@@ -149,6 +183,60 @@ describe("streamed event batching", () => {
     expect(h.events().map(chunkText)).toEqual(["a", "b", "c"]);
   });
 
+  it("coalesces a run of plain notifications into one appendEvents call", () => {
+    const h = createHarness();
+
+    h.emit(chunk("a"));
+    h.emit(chunk("b"));
+    h.emit(chunk("c"));
+    vi.advanceTimersByTime(FLUSH_MS);
+
+    // One store write for the whole run — not one per event. Each write costs
+    // O(transcript) (map copy + subscriber fan-out), so this is the batching
+    // that keeps big transcripts smooth during streaming.
+    expect(h.appendEvents).toHaveBeenCalledTimes(1);
+    expect(h.appendEvents.mock.calls[0][1].map(chunkText)).toEqual([
+      "a",
+      "b",
+      "c",
+    ]);
+  });
+
+  it("an id-carrying event splits the run but preserves global order", () => {
+    const h = createHarness();
+
+    h.emit(chunk("a"));
+    h.emit(chunk("b"));
+    h.emit(promptEcho(7, "user steer-less prompt"));
+    h.emit(chunk("c"));
+    vi.advanceTimersByTime(FLUSH_MS);
+
+    // Two coalesced appends around the individually-handled echo.
+    expect(h.appendEvents).toHaveBeenCalledTimes(2);
+    expect(h.appendEvents.mock.calls[0][1].map(chunkText)).toEqual(["a", "b"]);
+    expect(h.appendEvents.mock.calls[1][1].map(chunkText)).toEqual(["c"]);
+    expect(h.events().map(chunkText)).toEqual(["a", "b", "c"]);
+  });
+
+  it("a stop-reason response batched with chunks still completes the turn", () => {
+    const h = createHarness();
+
+    // Turn starts: the echo claims currentPromptId…
+    h.emit(promptEcho(9, "do the thing"));
+    // …streams some content…
+    h.emit(chunk("a"));
+    h.emit(chunk("b"));
+    // …and finishes, all within one flush window.
+    h.emit(stopResponse(9));
+    vi.advanceTimersByTime(FLUSH_MS);
+
+    // Transcript keeps everything appendable: the two chunks plus the
+    // response itself (the echo goes through replaceOptimisticWithEvent).
+    expect(h.events()).toHaveLength(3);
+    expect(h.events().slice(0, 2).map(chunkText)).toEqual(["a", "b"]);
+    expect(h.notifyPromptComplete).toHaveBeenCalledTimes(1);
+  });
+
   it("flushes buffered events synchronously on teardown", () => {
     const h = createHarness();
 

packages/core/src/sessions/sessionService.ts

@@ -1462,9 +1462,7 @@ export class SessionService {
     const batches = this.pendingSessionEvents;
     this.pendingSessionEvents = new Map();
     for (const [taskRunId, events] of batches) {
-      for (const acpMsg of events) {
-        this.handleSessionEvent(taskRunId, acpMsg);
-      }
+      this.handleSessionEvents(taskRunId, events);
     }
   }
 
@@ -1474,9 +1472,7 @@ export class SessionService {
     const events = this.pendingSessionEvents.get(taskRunId);
     if (!events) return;
     this.pendingSessionEvents.delete(taskRunId);
-    for (const acpMsg of events) {
-      this.handleSessionEvent(taskRunId, acpMsg);
-    }
+    this.handleSessionEvents(taskRunId, events);
   }
 
   // --- Transcript residency (memory eviction) ---
@@ -1886,6 +1882,50 @@ export class SessionService {
     }
   }
 
+  /** Apply one flush batch for a task. Runs of plain notifications (message
+   * chunks, tool-call updates — the bulk of a streamed turn) are appended to
+   * the store in a single `appendEvents` call: every store write costs
+   * O(transcript) (session-map copy plus subscriber notification fan-out), so
+   * writing per event made a 15-event flush walk the whole transcript 15
+   * times. Events carrying a JSON-RPC `id` (user-prompt echoes, stop-reason
+   * responses) keep going through `handleSessionEvent` one at a time, so
+   * turn-lifecycle semantics are untouched. Global event order is preserved:
+   * a run is flushed before any id-carrying event is handled. */
+  private handleSessionEvents(taskRunId: string, acpMsgs: AcpMessage[]): void {
+    let run: AcpMessage[] = [];
+
+    const flushRun = () => {
+      if (run.length === 0) return;
+      const batch = run;
+      run = [];
+      const session = this.d.store.getSessions()[taskRunId];
+      if (!session) return;
+      // Once the agent starts responding, clear initialPrompt so that retry
+      // reconnects to this session instead of creating a new one.
+      if (session.initialPrompt?.length) {
+        this.d.store.updateSession(taskRunId, { initialPrompt: undefined });
+      }
+      this.d.store.appendEvents(taskRunId, batch);
+      this.updatePromptStateFromEvents(taskRunId, batch, { isLive: true });
+      for (const acpMsg of batch) {
+        this.applyNotificationEffects(taskRunId, acpMsg, session);
+      }
+    };
+
+    for (const acpMsg of acpMsgs) {
+      // Notifications carry no JSON-RPC id; requests and responses do. Only
+      // id-less events are safe to coalesce: they are pure appends plus
+      // field-level side effects, with no turn-lifecycle branching.
+      if ("id" in acpMsg.message) {
+        flushRun();
+        this.handleSessionEvent(taskRunId, acpMsg);
+      } else {
+        run.push(acpMsg);
+      }
+    }
+    flushRun();
+  }
+
   private handleSessionEvent(taskRunId: string, acpMsg: AcpMessage): void {
     const session = this.d.store.getSessions()[taskRunId];
     if (!session) return;
@@ -1947,6 +1987,20 @@ export class SessionService {
       this.d.taskViewedApi.markActivity(session.taskId);
     }
 
+    this.applyNotificationEffects(taskRunId, acpMsg, session);
+  }
+
+  /** Side effects carried by session notifications (config, usage, adapter,
+   * compaction status). Split out of `handleSessionEvent` so a coalesced run
+   * of notifications (see `handleSessionEvents`) shares it; no-op for
+   * requests/responses, which carry no `method`. */
+  private applyNotificationEffects(
+    taskRunId: string,
+    acpMsg: AcpMessage,
+    session: AgentSession,
+  ): void {
+    const msg = acpMsg.message;
+
     if ("method" in msg && msg.method === "session/update" && "params" in msg) {
       const params = msg.params as {
         update?: {

scripts/measure-append-hotpath.ts

@@ -0,0 +1,88 @@
+// Mide el hot path real de appendEvents (módulo real, sin réplicas):
+// 1 segundo de streaming pesado = 60 flushes x 15 eventos sobre un transcript
+// de N eventos. Correr en main y en la rama para comparar.
+//   bun scripts/measure-append-hotpath.ts
+import type { AcpMessage, AgentSession } from "@posthog/shared";
+import { sessionStoreSetters } from "../packages/core/src/sessions/sessionStore";
+
+const RUN_ID = "run-bench";
+let seq = 0;
+function makeEvent(): AcpMessage {
+  seq++;
+  const message =
+    seq % 12 === 0
+      ? {
+          jsonrpc: "2.0",
+          method: "session/update",
+          params: {
+            sessionId: RUN_ID,
+            update: {
+              sessionUpdate: "tool_call_update",
+              toolCallId: `tool-${seq % 40}`,
+              status: "in_progress",
+              rawInput: { file_path: "/src/app.ts", content: "x".repeat(2048) },
+            },
+          },
+        }
+      : {
+          jsonrpc: "2.0",
+          method: "session/update",
+          params: {
+            sessionId: RUN_ID,
+            update: {
+              sessionUpdate: "agent_message_chunk",
+              content: { type: "text", text: "palabra ".repeat(10) },
+            },
+          },
+        };
+  return { ts: 1700000000000 + seq, message } as unknown as AcpMessage;
+}
+
+function seedSession(n: number): void {
+  const events: AcpMessage[] = [];
+  for (let i = 0; i < n; i++) events.push(Object.freeze(makeEvent()));
+  sessionStoreSetters.setSession({
+    taskRunId: RUN_ID,
+    taskId: "task-bench",
+    events,
+    messageQueue: [],
+    pendingPermissions: new Map(),
+    status: "connected",
+  } as unknown as AgentSession);
+}
+
+const FRAMES = 60;
+const PER_FRAME = 15;
+
+function run(perEvent: boolean): number {
+  const batches: AcpMessage[][] = [];
+  for (let f = 0; f < FRAMES; f++) {
+    const b: AcpMessage[] = [];
+    for (let e = 0; e < PER_FRAME; e++) b.push(makeEvent());
+    batches.push(b);
+  }
+  const t0 = performance.now();
+  for (const batch of batches) {
+    if (perEvent) {
+      for (const ev of batch) sessionStoreSetters.appendEvents(RUN_ID, [ev]);
+    } else {
+      sessionStoreSetters.appendEvents(RUN_ID, batch);
+    }
+  }
+  return performance.now() - t0;
+}
+
+for (const n of [10_000, 30_000]) {
+  for (const perEvent of [true, false]) {
+    const label = perEvent ? "per-evento x15" : "batcheado x1  ";
+    const times: number[] = [];
+    for (let rep = 0; rep < 7; rep++) {
+      seedSession(n);
+      times.push(run(perEvent));
+    }
+    times.sort((a, b) => a - b);
+    console.log(
+      `transcript ${n.toLocaleString()} | ${label} | mediana ${times[3].toFixed(1)} ms/seg-streaming (min ${times[0].toFixed(1)}, max ${times[6].toFixed(1)})`,
+    );
+  }
+}