Keep live delivery reads tracked through socket closure

server/src/__tests__/live-events-toast.test.ts

@@ -8,6 +8,7 @@ import { describe, expect, it, vi } from "vitest";
 import { setupLiveEventsWebSocketServer } from "../realtime/live-events-ws.js";
 import { publishLiveEvent } from "../services/live-events.js";
 import { taskToastEligible } from "../realtime/task-toast-eligibility.js";
+import { idleWorkSnapshot } from "../services/task-admission.js";
 
 vi.mock("../realtime/task-toast-eligibility.js", () => ({ taskToastEligible: vi.fn(async () => false) }));
 
@@ -39,6 +40,48 @@ function database(onPredicate: (kind: "where" | "join", condition: SQL) => void
   } };
 }
 describe("authenticated activity notifications", () => {
+  it("retains live delivery idle work until the identifier read settles after the last socket closes", async () => {
+    const baseline = idleWorkSnapshot().active;
+    let release!: () => void, entered!: () => void;
+    const pending = new Promise<void>(resolve => { release = resolve; });
+    const started = new Promise<void>(resolve => { entered = resolve; });
+    const db = database(), select = db.select;
+    db.select = fields => {
+      const query = select(fields);
+      if (fields.identifier) {
+        const limit = query.limit;
+        query.limit = async () => { entered(); await pending; return limit(); };
+      }
+      return query;
+    };
+    const server = createServer();
+    const wss = setupLiveEventsWebSocketServer(server, db as never, { deploymentMode: "local_trusted" });
+    const closed = new Promise<void>(resolve => wss.on("connection", socket => socket.on("close", resolve)));
+    server.listen(0, "127.0.0.1");
+    await once(server, "listening");
+    const socket = new WebSocket(`ws://127.0.0.1:${(server.address() as { port: number }).port}/api/companies/company/events/ws`);
+    try {
+      await once(socket, "open");
+      await vi.waitFor(() => expect(idleWorkSnapshot().active).toBe(baseline));
+      // System activity has no actor lookup: this gate is the new identifier SQL.
+      publishLiveEvent({ companyId: "company", type: "activity.logged", payload: {
+        action: "issue.updated", entityType: "issue", entityId: "visible", actorType: "system",
+      } });
+      await started;
+      socket.terminate();
+      await closed;
+      expect(wss.clients.size).toBe(0);
+      expect(idleWorkSnapshot().active).toBe(baseline + 1);
+      release();
+      await vi.waitFor(() => expect(idleWorkSnapshot().active).toBe(baseline));
+    } finally {
+      release();
+      socket.terminate();
+      await new Promise<void>(resolve => server.close(() => resolve()));
+      await vi.waitFor(() => expect(idleWorkSnapshot().active).toBe(baseline));
+    }
+  });
+
   it.each(["local_trusted", "authenticated"] as const)("keeps authorized notifications when optional identifier reads fail (%s)", async deploymentMode => {
     const db = database();
     const select = db.select;

server/src/realtime/live-event-delivery.test.ts

@@ -1,6 +1,7 @@
 import { afterEach, expect, it, vi } from "vitest";
 import type { LiveEvent } from "@paperclipai/shared";
 import { createLiveEventDelivery, LIVE_DELIVERY_MAX_BYTES, LIVE_DELIVERY_MAX_ITEMS } from "./live-event-delivery.js";
+import { idleWorkSnapshot } from "../services/task-admission.js";
 
 const event = (seq = 1, runId = "run-1"): LiveEvent => ({ id: seq, companyId: "company", type: "heartbeat.run.log",
   createdAt: "2026-10-24T00:00:00Z", payload: { runId, agentId: "agent", seq, chunk: "output" } });
@@ -116,7 +117,7 @@ it.each([false, true])("preserves the overflow deadline across resubscription un
     await vi.advanceTimersByTimeAsync(1_000);
     if (settled) expect(f.socket.close).not.toHaveBeenCalled();
     else expect(f.socket.close).toHaveBeenCalledWith(1013, "persistent event backlog");
-  } finally { f.delivery.dispose(); }
+  } finally { f.delivery.dispose(); release(event()); await settle(); }
   expect(vi.getTimerCount()).toBe(0);
 });
 it("recovers failed delivery reads through resync rather than silently claiming healthy streaming", async () => {
@@ -146,6 +147,7 @@ it("rejects unbounded interest sets and stale revision commands", () => {
 it.each([false, true].flatMap(negotiated => ["authorize", "decorate"].map(stage => ({ negotiated, stage }))))(
   "closes a stalled $stage read below the queue bound (negotiated=$negotiated)", async ({ negotiated, stage }) => {
     vi.useFakeTimers();
+    const baseline = idleWorkSnapshot().active;
     let release!: (value: LiveEvent) => void;
     const pending = new Promise<LiveEvent>(resolve => { release = resolve; });
     const f = fixture(stage === "authorize" ? vi.fn(() => pending) : undefined);
@@ -168,11 +170,13 @@ it.each([false, true].flatMap(negotiated => ["authorize", "decorate"].map(stage
       expect(f.socket.close).toHaveBeenCalledWith(1013, "persistent event backlog");
       expect(f.delivery.metrics().retainedItems).toBe(1);
       expect(vi.getTimerCount()).toBe(0);
+      expect(idleWorkSnapshot().active).toBe(baseline + 1);
       release(source); await settle();
+      expect(idleWorkSnapshot().active).toBe(baseline);
       expect(f.delivery.metrics()).toMatchObject({ retainedItems: 0, retainedBytes: 0 });
       expect(f.frames.filter(frame => frame.type.startsWith("heartbeat."))).toEqual([]);
       expect(f.authorize).toHaveBeenCalledOnce();
-    } finally { f.delivery.dispose(); other.delivery.dispose(); }
+    } finally { f.delivery.dispose(); other.delivery.dispose(); release(event()); await settle(); }
   });
 
 it.each(["resolve", "dispose"])("clears a pending-read deadline on %s without idle work", async ending => {

server/src/realtime/live-event-delivery.ts

@@ -1,4 +1,5 @@
 import type { LiveEvent } from "@paperclipai/shared";
+import { beginIdleTrackedWork } from "../services/task-admission.js";
 
 export const LIVE_DELIVERY_MAX_ITEMS = 1_000;
 export const LIVE_DELIVERY_MAX_BYTES = 1024 * 1024;
@@ -90,6 +91,7 @@ export function createLiveEventDelivery(input: {
   async function drain() {
     if (draining) return;
     draining = true;
+    const finishIdleWork = beginIdleTrackedWork();
     try {
       while (queued.length && isOpen()) {
         const item = queued.shift()!;
@@ -121,7 +123,10 @@ export function createLiveEventDelivery(input: {
           retainedBytes -= item.bytes; retainedItems -= 1;
         }
       }
-    } finally { draining = false; resync(); }
+    } finally {
+      try { draining = false; resync(); }
+      finally { finishIdleWork(); }
+    }
   }
   return {
     receive(raw: string) {