From b28da9be32791de4758452b5bf2e53400bcc35cd Mon Sep 17 00:00:00 2001 From: Dotta Date: Tue, 3 Nov 2026 14:00:38 -0600 Subject: [PATCH] Keep live delivery reads tracked through socket closure --- .../src/__tests__/live-events-toast.test.ts | 43 +++++++++++++++++++ .../src/realtime/live-event-delivery.test.ts | 8 +++- server/src/realtime/live-event-delivery.ts | 7 ++- 3 files changed, 55 insertions(+), 3 deletions(-) diff --git a/server/src/__tests__/live-events-toast.test.ts b/server/src/__tests__/live-events-toast.test.ts index c877b70f090..ffa96702e50 100644 --- a/server/src/__tests__/live-events-toast.test.ts +++ b/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(resolve => { release = resolve; }); + const started = new Promise(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(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(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; diff --git a/server/src/realtime/live-event-delivery.test.ts b/server/src/realtime/live-event-delivery.test.ts index 02e8c536c83..2f8b4c4eac4 100644 --- a/server/src/realtime/live-event-delivery.test.ts +++ b/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(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 => { diff --git a/server/src/realtime/live-event-delivery.ts b/server/src/realtime/live-event-delivery.ts index 3d7f4371779..588eb654a35 100644 --- a/server/src/realtime/live-event-delivery.ts +++ b/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) {