diff --git a/services/cloud-agent-next/wrapper/src/control/control-event-ordering.test.ts b/services/cloud-agent-next/wrapper/src/control/control-event-ordering.test.ts index 80eb327558..948060dd53 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-event-ordering.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-event-ordering.test.ts @@ -162,7 +162,7 @@ describe('control event publication ordering', () => { } ); - it('re-arms a woken reservation if another Session consumes the available bytes', async () => { + it('keeps a woken reservation isolated from another root lane', async () => { const outbox = createControlEventOutbox({ publish: async () => {}, onFailure: mock() }); try { for (let index = 0; index < 8; index += 1) @@ -182,18 +182,7 @@ describe('control event publication ordering', () => { outbox.prepare({ event: 'session.event', session: other, payload: medium }) ) ).toBe(true); - expect(outbox.enqueue(older)).toBe(false); - const waitingAgain = outbox.waitForSpace(older); - expect(waitingAgain).not.toBe(waiting); - expect(outbox.waitForSpace(older)).toBe(waitingAgain); - let settled = false; - void waitingAgain.then(() => { - settled = true; - }); - await Promise.resolve(); - expect(settled).toBe(false); - outbox.close(); - expect(await waitingAgain).toBe(false); + expect(outbox.enqueue(older)).toBe(true); } finally { outbox.close(); } diff --git a/services/cloud-agent-next/wrapper/src/control/control-event-outbox.test.ts b/services/cloud-agent-next/wrapper/src/control/control-event-outbox.test.ts index 0d1375e798..55b01f1a34 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-event-outbox.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-event-outbox.test.ts @@ -12,7 +12,356 @@ const session = { rootKiloSessionId: 'ses_root', }; +async function waitFor(condition: () => boolean, attempts = 100): Promise { + for (let attempt = 0; attempt < attempts; attempt += 1) { + if (condition()) return; + await Bun.sleep(1); + } + throw new Error('Timed out waiting for outbox publication'); +} + describe('control event outbox', () => { + it('does not let a retryable root A head delay root B', async () => { + const published: ControlEventPublication[] = []; + let attemptsA = 0; + const outbox = createControlEventOutbox({ + publish: async publication => { + published.push(publication); + const root = publication.session.rootKiloSessionId ?? publication.session.kiloSessionId; + if (root === 'root_a' && attemptsA++ === 0) + throw new ControlDeliveryError('root A is not attached', true); + }, + onFailure: mock(), + }); + try { + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session: { ...session, kiloSessionId: 'root_a', rootKiloSessionId: 'root_a' }, + payload: { type: 'session.idle', properties: {} }, + }) + ); + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session: { ...session, kiloSessionId: 'root_b', rootKiloSessionId: 'root_b' }, + payload: { type: 'session.idle', properties: {} }, + }) + ); + + expect(await outbox.resume()).toBe(false); + expect( + published.map(item => item.session.rootKiloSessionId ?? item.session.kiloSessionId) + ).toEqual(['root_a', 'root_b']); + } finally { + outbox.close(); + } + }); + + it('wakes an active cycle when a new root is admitted during a pending receipt', async () => { + const startedA = Promise.withResolvers(); + const releaseA = Promise.withResolvers(); + const rootA = { ...session, kiloSessionId: 'root_a', rootKiloSessionId: 'root_a' }; + const rootB = { ...session, kiloSessionId: 'root_b', rootKiloSessionId: 'root_b' }; + let releasedA = false; + let startedB = false; + const outbox = createControlEventOutbox({ + publish: async publication => { + const root = publication.session.rootKiloSessionId ?? publication.session.kiloSessionId; + if (root === 'root_a') { + startedA.resolve(); + await releaseA.promise; + } else { + startedB = true; + } + }, + onFailure: mock(), + }); + try { + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session: rootA, + payload: { type: 'session.idle' }, + }) + ); + const draining = outbox.resume(); + await startedA.promise; + expect( + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session: rootB, + payload: { type: 'session.idle' }, + }) + ) + ).toBe(true); + await waitFor(() => startedB); + expect(releasedA).toBe(false); + releasedA = true; + releaseA.resolve(); + expect(await draining).toBe(true); + } finally { + releaseA.resolve(); + outbox.close(); + } + }); + + it('wakes a retry-ready root while another root receipt remains pending', async () => { + const releaseB = Promise.withResolvers(); + const rootA = { ...session, kiloSessionId: 'root_a', rootKiloSessionId: 'root_a' }; + const rootB = { ...session, kiloSessionId: 'root_b', rootKiloSessionId: 'root_b' }; + let attemptsA = 0; + let startedB = false; + let retriedA = false; + const outbox = createControlEventOutbox({ + publish: async publication => { + const root = publication.session.rootKiloSessionId ?? publication.session.kiloSessionId; + if (root === 'root_a') { + attemptsA += 1; + if (attemptsA === 1) throw new ControlDeliveryError('root A is not attached', true); + retriedA = true; + return; + } + startedB = true; + await releaseB.promise; + }, + onFailure: mock(), + }); + try { + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session: rootA, + payload: { type: 'session.idle' }, + }) + ); + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session: rootB, + payload: { type: 'session.idle' }, + }) + ); + const draining = outbox.resume(); + await waitFor(() => startedB); + await waitFor(() => retriedA, 500); + releaseB.resolve(); + expect(await draining).toBe(true); + } finally { + releaseB.resolve(); + outbox.close(); + } + }); + + it('keeps entry, byte, and waiter limits independent per root', async () => { + const outbox = createControlEventOutbox({ publish: async () => {}, onFailure: mock() }); + const rootA = { ...session, kiloSessionId: 'root_a', rootKiloSessionId: 'root_a' }; + const rootB = { ...session, kiloSessionId: 'root_b', rootKiloSessionId: 'root_b' }; + const small = { type: 'session.idle', properties: {} }; + try { + for (let index = 0; index < 256; index += 1) + expect( + outbox.enqueue(outbox.prepare({ event: 'session.event', session: rootA, payload: small })) + ).toBe(true); + expect( + outbox.enqueue(outbox.prepare({ event: 'session.event', session: rootA, payload: small })) + ).toBe(false); + expect( + outbox.enqueue(outbox.prepare({ event: 'session.event', session: rootB, payload: small })) + ).toBe(true); + + outbox.pause(); + outbox.close(); + const bytesOutbox = createControlEventOutbox({ publish: async () => {}, onFailure: mock() }); + try { + const medium = { + type: 'message.updated', + properties: { text: 'm'.repeat(Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES * 0.45)) }, + }; + for (let index = 0; index < 256; index += 1) { + if ( + !bytesOutbox.enqueue( + bytesOutbox.prepare({ event: 'session.event', session: rootA, payload: medium }) + ) + ) + break; + } + expect( + bytesOutbox.enqueue( + bytesOutbox.prepare({ event: 'session.event', session: rootA, payload: medium }) + ) + ).toBe(false); + expect( + bytesOutbox.enqueue( + bytesOutbox.prepare({ event: 'session.event', session: rootB, payload: medium }) + ) + ).toBe(true); + } finally { + bytesOutbox.close(); + } + + const waiterOutbox = createControlEventOutbox({ publish: async () => {}, onFailure: mock() }); + try { + for (let index = 0; index < 256; index += 1) + expect( + waiterOutbox.enqueue( + waiterOutbox.prepare({ event: 'session.event', session: rootA, payload: small }) + ) + ).toBe(true); + const rootWaiters: Array> = []; + for (let index = 0; index < 256; index += 1) { + const publication = waiterOutbox.prepare({ + event: 'session.event', + session: rootA, + payload: small, + }); + expect(waiterOutbox.enqueue(publication)).toBe(false); + rootWaiters.push(waiterOutbox.waitForSpace(publication)); + } + for (let index = 0; index < 256; index += 1) + expect( + waiterOutbox.enqueue( + waiterOutbox.prepare({ event: 'session.event', session: rootB, payload: small }) + ) + ).toBe(true); + const rootBPublication = waiterOutbox.prepare({ + event: 'session.event', + session: rootB, + payload: small, + }); + expect(waiterOutbox.enqueue(rootBPublication)).toBe(false); + const rootBWaiter = waiterOutbox.waitForSpace(rootBPublication); + let settled = false; + void rootBWaiter.then(() => { + settled = true; + }); + await Promise.resolve(); + expect(settled).toBe(false); + waiterOutbox.close(); + expect(await rootBWaiter).toBe(false); + expect(await Promise.all(rootWaiters)).toEqual(Array.from({ length: 256 }, () => false)); + } finally { + waiterOutbox.close(); + } + } finally { + outbox.close(); + } + }); + + it('sends roots fairly while keeping one publication in flight per root', async () => { + const started: ControlEventPublication[] = []; + const active = new Map(); + const maximum = new Map(); + const releases = new Map>(); + const outbox = createControlEventOutbox({ + publish: async publication => { + const root = publication.session.rootKiloSessionId ?? publication.session.kiloSessionId; + if (!root) throw new Error('Missing root'); + const count = (active.get(root) ?? 0) + 1; + active.set(root, count); + maximum.set(root, Math.max(maximum.get(root) ?? 0, count)); + started.push(publication); + const release = Promise.withResolvers(); + releases.set(`${root}:${publication.sequence}`, release); + await release.promise; + active.set(root, count - 1); + }, + onFailure: mock(), + }); + const rootA = { ...session, kiloSessionId: 'root_a', rootKiloSessionId: 'root_a' }; + const rootB = { ...session, kiloSessionId: 'root_b', rootKiloSessionId: 'root_b' }; + try { + for (const [root, count] of [ + [rootA, 2], + [rootB, 2], + ] as const) + for (let index = 0; index < count; index += 1) + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session: root, + payload: { type: 'session.idle', properties: {} }, + }) + ); + + const draining = outbox.resume(); + await waitFor(() => started.length === 2); + expect(started.map(item => item.session.rootKiloSessionId)).toEqual(['root_a', 'root_b']); + expect(maximum).toEqual( + new Map([ + ['root_a', 1], + ['root_b', 1], + ]) + ); + + releases.get('root_a:1')?.resolve(); + await waitFor(() => started.length === 3); + expect(started[2]?.session.rootKiloSessionId).toBe('root_a'); + expect(maximum.get('root_a')).toBe(1); + expect(maximum.get('root_b')).toBe(1); + + releases.get('root_b:3')?.resolve(); + await waitFor(() => started.length === 4); + expect(started[3]?.session.rootKiloSessionId).toBe('root_b'); + releases.get('root_a:2')?.resolve(); + releases.get('root_b:4')?.resolve(); + expect(await draining).toBe(true); + } finally { + outbox.close(); + } + }); + + it('keeps root and child publications in one FIFO lane', async () => { + const published: ControlEventPublication[] = []; + const outbox = createControlEventOutbox({ + publish: async publication => { + published.push(publication); + }, + onFailure: mock(), + }); + try { + const child = { ...session, kiloSessionId: 'ses_child' }; + const sibling = { + ...session, + kiloSessionId: 'ses_sibling', + rootKiloSessionId: 'ses_sibling', + }; + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session: session, + payload: { type: 'session.idle', properties: {} }, + }) + ); + outbox.enqueue( + outbox.prepare({ + event: 'session.preparing', + session: child, + payload: { action: 'step_started' }, + }) + ); + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session: sibling, + payload: { type: 'session.idle', properties: {} }, + }) + ); + expect(await outbox.resume()).toBe(true); + const rootPublished = published.filter( + item => (item.session.rootKiloSessionId ?? item.session.kiloSessionId) === 'ses_root' + ); + expect(rootPublished.map(item => item.session.kiloSessionId)).toEqual([ + 'ses_root', + 'ses_child', + ]); + expect(rootPublished.map(item => item.sequence)).toEqual([1, 2]); + } finally { + outbox.close(); + } + }); + it('snapshots native lifetime before replacement', async () => { const published: ControlEventPublication[] = []; const outbox = createControlEventOutbox({ @@ -175,6 +524,36 @@ describe('control event outbox', () => { } }); + it('does not publish a deferred attempt after immediate close', async () => { + const published = mock(async () => {}); + const outbox = createControlEventOutbox({ publish: published, onFailure: mock() }); + outbox.enqueue( + outbox.prepare({ event: 'session.event', session, payload: { type: 'session.idle' } }) + ); + const draining = outbox.resume(); + outbox.close(); + expect(await draining).toBe(false); + expect(published).not.toHaveBeenCalled(); + }); + + it('does not publish a deferred attempt after immediate pause', async () => { + const published = mock(async () => {}); + const outbox = createControlEventOutbox({ publish: published, onFailure: mock() }); + try { + outbox.enqueue( + outbox.prepare({ event: 'session.event', session, payload: { type: 'session.idle' } }) + ); + const pumping = outbox.resume(); + outbox.pause(); + expect(await pumping).toBe(false); + expect(published).not.toHaveBeenCalled(); + expect(await outbox.resume()).toBe(true); + expect(published).toHaveBeenCalledTimes(1); + } finally { + outbox.close(); + } + }); + it('settles an in-flight pump on close without reporting expiry or publishing queued events', async () => { const started = Promise.withResolvers(); const held = Promise.withResolvers(); diff --git a/services/cloud-agent-next/wrapper/src/control/control-event-outbox.ts b/services/cloud-agent-next/wrapper/src/control/control-event-outbox.ts index cb8b6a35ce..905a175de3 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-event-outbox.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-event-outbox.ts @@ -38,67 +38,166 @@ export type ControlEventOutbox = { close(): void; }; +type RootKey = string | undefined; + +type SpaceWaiter = { + promise: Promise; + ready: boolean; + resolve: (available: boolean) => void; + timeout: ReturnType; +}; + +type Lane = { + root: RootKey; + entries: PreparedControlEventPublication[]; + spaceWaiters: Map; + waitingBytes: number; + bytes: number; + pending?: Promise; + expirePending?: () => void; + retryAt?: number; + wakeup?: ReturnType; +}; + +type PumpCycle = { + promise: Promise; + resolve: (drained: boolean) => void; + wake: Promise; + resolveWake: () => void; + wakeSignaled: boolean; +}; + +function rootFor(publication: PreparedControlEventPublication): RootKey { + return publication.session.rootKiloSessionId ?? publication.session.kiloSessionId; +} + +function isRetryable(error: unknown): boolean { + return ( + typeof error === 'object' && error !== null && 'retryable' in error && error.retryable === true + ); +} + export function createControlEventOutbox(options: { publish: (publication: ControlEventPublication, deadlineAt: number) => Promise; onFailure: (failure: ControlEventOutboxFailure) => void; }): ControlEventOutbox { - const entries: PreparedControlEventPublication[] = []; - const spaceWaiters = new Map< - PreparedControlEventPublication, - { - promise: Promise; - ready: boolean; - resolve: (available: boolean) => void; - timeout: ReturnType; - } - >(); - let waitingBytes = 0; - let bytes = 0; + const lanes = new Map(); let paused = true; let closed = false; - let pending: Promise | undefined; + let cycle: PumpCycle | undefined; let nextSequence = 0; - let retryAt: number | undefined; - let wakeup: ReturnType | undefined; - let expirePending: (() => void) | undefined; + let lastScheduledLane: Lane | undefined; + + const getLane = (root: RootKey): Lane => { + const existing = lanes.get(root); + if (existing) return existing; + const lane: Lane = { + root, + entries: [], + spaceWaiters: new Map(), + waitingBytes: 0, + bytes: 0, + }; + lanes.set(root, lane); + return lane; + }; + + const clearWakeup = (lane: Lane): void => { + clearTimeout(lane.wakeup); + lane.wakeup = undefined; + }; + + const cleanupLane = (lane: Lane): void => { + if (lane.pending || lane.entries.length > 0 || lane.spaceWaiters.size > 0) return; + clearWakeup(lane); + if (lanes.get(lane.root) === lane) lanes.delete(lane.root); + }; - const clearWakeup = () => { - clearTimeout(wakeup); - wakeup = undefined; + const reportFailure = (failure: ControlEventOutboxFailure): void => { + try { + options.onFailure(failure); + } catch { + // Failure reporting must not strand the other root lanes. + } }; - const hasSpaceFor = (publication: PreparedControlEventPublication): boolean => { - if (entries.length >= MAX_EVENTS || bytes + publication.bytes > MAX_BYTES) return false; - const root = publication.session.rootKiloSessionId ?? publication.session.kiloSessionId; - for (const reserved of spaceWaiters.keys()) { - if ( - reserved.sequence < publication.sequence && - Date.now() < reserved.deadlineAt && - reserved.session.directory === publication.session.directory && - (reserved.session.rootKiloSessionId ?? reserved.session.kiloSessionId) === root - ) - return false; + const hasSpaceFor = (lane: Lane, publication: PreparedControlEventPublication): boolean => { + if (lane.entries.length >= MAX_EVENTS || lane.bytes + publication.bytes > MAX_BYTES) + return false; + const now = Date.now(); + for (const reserved of lane.spaceWaiters.keys()) { + if (reserved.sequence < publication.sequence && now < reserved.deadlineAt) return false; } return true; }; - const releaseSpaceWaiter = (publication: PreparedControlEventPublication) => { - const waiter = spaceWaiters.get(publication); - if (!waiter) return; + const releaseSpaceWaiter = ( + lane: Lane, + publication: PreparedControlEventPublication + ): SpaceWaiter | undefined => { + const waiter = lane.spaceWaiters.get(publication); + if (!waiter) return undefined; clearTimeout(waiter.timeout); - spaceWaiters.delete(publication); - waitingBytes -= publication.bytes; + lane.spaceWaiters.delete(publication); + lane.waitingBytes -= publication.bytes; + cleanupLane(lane); return waiter; }; - const notifySpace = (available: boolean) => { + const notifySpace = (lane: Lane, available: boolean): void => { const now = Date.now(); - for (const [publication, waiter] of spaceWaiters) { - if (!available || now >= publication.deadlineAt) releaseSpaceWaiter(publication); - else if (!hasSpaceFor(publication)) continue; - waiter.ready = true; - waiter.resolve(available); + for (const [publication, waiter] of lane.spaceWaiters) { + if (!available || now >= publication.deadlineAt) { + const released = releaseSpaceWaiter(lane, publication); + released?.resolve(available ? true : false); + } else if (hasSpaceFor(lane, publication)) { + waiter.ready = true; + waiter.resolve(true); + } + } + cleanupLane(lane); + }; + + const removeHead = (lane: Lane, entry: PreparedControlEventPublication): boolean => { + if (lane.entries[0] !== entry) return false; + lane.entries.shift(); + lane.bytes -= entry.bytes; + lane.retryAt = undefined; + notifySpace(lane, true); + scheduleWakeup(lane); + cleanupLane(lane); + return true; + }; + + const expireHead = (lane: Lane): void => { + if (lane.pending) return; + while (lane.entries[0] && Date.now() >= lane.entries[0].deadlineAt) { + const entry = lane.entries[0]; + if (!entry || !removeHead(lane, entry)) return; + reportFailure({ reason: 'expired', publication: entry }); } + scheduleWakeup(lane); + }; + + const scheduleWakeup = (lane: Lane): void => { + clearWakeup(lane); + const entry = lane.entries[0]; + if (closed || lane.pending || !entry) { + cleanupLane(lane); + return; + } + const now = Date.now(); + const nextAt = paused ? entry.deadlineAt : Math.min(lane.retryAt ?? now, entry.deadlineAt); + lane.wakeup = setTimeout( + () => { + lane.wakeup = undefined; + expireHead(lane); + if (!paused) void pump(); + else scheduleWakeup(lane); + }, + Math.max(1, nextAt - now) + ); + lane.wakeup.unref(); }; const prepare = ( @@ -127,129 +226,201 @@ export function createControlEventOutbox(options: { return { ...publication, bytes, deadlineAt: Date.now() + 30_000 }; }; - const removeHead = (entry: PreparedControlEventPublication) => { - entries.shift(); - bytes -= entry.bytes; - retryAt = undefined; - notifySpace(true); + const nextRunnableLane = (): Lane | undefined => { + if (lanes.size === 0) return undefined; + const available = [...lanes.values()]; + const previousIndex = available.findIndex(lane => lane === lastScheduledLane); + const start = previousIndex === -1 ? 0 : (previousIndex + 1) % available.length; + const now = Date.now(); + for (let offset = 0; offset < available.length; offset += 1) { + const lane = available[(start + offset) % available.length]; + if (!lane || lane.pending || !lane.entries[0]) continue; + if (lane.retryAt !== undefined) { + if (now < lane.retryAt) continue; + lane.retryAt = undefined; + } + if (now >= lane.entries[0].deadlineAt) expireHead(lane); + const entry = lane.entries[0]; + if (!entry || lane.pending || Date.now() >= entry.deadlineAt) continue; + lastScheduledLane = lane; + return lane; + } + return undefined; }; - const expireHead = () => { - while (entries[0] && Date.now() >= entries[0].deadlineAt) { - const entry = entries[0]; - removeHead(entry); - options.onFailure({ reason: 'expired', publication: entry }); - } + const queuedEntries = (): boolean => { + for (const lane of lanes.values()) if (lane.entries.length > 0) return true; + return false; }; - const scheduleWakeup = () => { - clearWakeup(); - const entry = entries[0]; - if (closed || pending || !entry) return; - const nextAt = paused ? entry.deadlineAt : Math.min(retryAt ?? Date.now(), entry.deadlineAt); - wakeup = setTimeout( - () => { - wakeup = undefined; - expireHead(); - if (!paused) void pump(); - else scheduleWakeup(); - }, - Math.max(1, nextAt - Date.now()) - ); - wakeup.unref(); + const signalCycle = (active: PumpCycle | undefined = cycle): void => { + if (!active || active.wakeSignaled) return; + active.wakeSignaled = true; + active.resolveWake(); }; - const pump = (): Promise => { - if (pending) return pending; - if (closed) return Promise.resolve(false); - expireHead(); - if (paused || (retryAt !== undefined && Date.now() < retryAt)) { - scheduleWakeup(); - return Promise.resolve(entries.length === 0); + const resetCycleWake = (active: PumpCycle): void => { + if (!active.wakeSignaled) return; + const wake = Promise.withResolvers(); + active.wake = wake.promise; + active.resolveWake = wake.resolve; + active.wakeSignaled = false; + }; + + const runAttempt = async (lane: Lane, entry: PreparedControlEventPublication): Promise => { + if (closed || paused || lane.entries[0] !== entry) return; + const expired = Promise.withResolvers(); + lane.expirePending = expired.resolve; + const timeout = setTimeout(expired.resolve, Math.max(1, entry.deadlineAt - Date.now())); + timeout.unref(); + let published: Promise; + try { + published = Promise.resolve( + options.publish( + { + event: entry.event, + receiptId: entry.receiptId, + sequence: entry.sequence, + session: entry.session, + payload: entry.payload, + }, + entry.deadlineAt + ) + ); + } catch (error) { + published = Promise.reject(error); } - clearWakeup(); - pending = Promise.resolve() - .then(async () => { - while (!paused && !closed) { - expireHead(); - const entry = entries[0]; - if (!entry) return true; - let timeout: ReturnType | undefined; - const published = options.publish( - { - event: entry.event, - receiptId: entry.receiptId, - sequence: entry.sequence, - session: entry.session, - payload: entry.payload, - }, - entry.deadlineAt - ); - try { - const expired = new Promise(resolve => { - expirePending = resolve; - timeout = setTimeout(resolve, Math.max(1, entry.deadlineAt - Date.now())); - timeout.unref(); - }); - await Promise.race([published, expired]); - } catch (error) { - if (closed || Date.now() >= entry.deadlineAt) { - void published.catch(() => undefined); - if (closed) return false; - continue; - } - if (!(error instanceof Error) || !('retryable' in error) || error.retryable !== true) { - removeHead(entry); - options.onFailure({ reason: 'rejected', publication: entry }); - continue; - } - retryAt = Date.now() + RETRY_DELAY_MS; - return false; - } finally { - clearTimeout(timeout); - expirePending = undefined; - } - if (closed || Date.now() >= entry.deadlineAt) { - void published.catch(() => undefined); - if (closed) return false; - continue; - } - removeHead(entry); + try { + try { + await Promise.race([published, expired.promise]); + } catch (error) { + if (closed || Date.now() >= entry.deadlineAt) { + void published.catch(() => undefined); + return; + } + if (isRetryable(error)) { + lane.retryAt = Date.now() + RETRY_DELAY_MS; + return; } - return entries.length === 0; + if (removeHead(lane, entry)) reportFailure({ reason: 'rejected', publication: entry }); + return; + } + + if (closed || Date.now() >= entry.deadlineAt) { + void published.catch(() => undefined); + if (!closed && removeHead(lane, entry)) + reportFailure({ reason: 'expired', publication: entry }); + return; + } + removeHead(lane, entry); + } finally { + clearTimeout(timeout); + if (lane.expirePending === expired.resolve) lane.expirePending = undefined; + } + }; + + const startAttempt = (lane: Lane): void => { + const entry = lane.entries[0]; + if (!entry || lane.pending) return; + const pending = Promise.resolve() + .then(() => runAttempt(lane, entry)) + .catch(() => { + if (closed || !removeHead(lane, entry)) return; + reportFailure({ reason: 'rejected', publication: entry }); }) - .finally(() => { - pending = undefined; - scheduleWakeup(); + .then(() => { + if (lane.pending === pending) lane.pending = undefined; + scheduleWakeup(lane); + cleanupLane(lane); }); - return pending; + lane.pending = pending; + }; + + const runCycle = async (active: PumpCycle): Promise => { + try { + while (true) { + for (const lane of lanes.values()) expireHead(lane); + if (closed) { + active.resolve(false); + return; + } + if (paused) { + active.resolve(!queuedEntries()); + return; + } + + const lane = nextRunnableLane(); + if (lane) { + startAttempt(lane); + continue; + } + + const pending = [...lanes.values()] + .map(item => item.pending) + .filter((item): item is Promise => item !== undefined); + if (pending.length > 0) { + const wake = active.wake; + await Promise.race([...pending, wake]); + if (active.wake === wake) resetCycleWake(active); + continue; + } + active.resolve(!queuedEntries()); + return; + } + } finally { + if (cycle === active) cycle = undefined; + for (const lane of lanes.values()) scheduleWakeup(lane); + } + }; + + const pump = (): Promise => { + if (cycle) { + signalCycle(cycle); + return cycle.promise; + } + if (closed) return Promise.resolve(false); + const next = Promise.withResolvers(); + const wake = Promise.withResolvers(); + const active: PumpCycle = { + promise: next.promise, + resolve: next.resolve, + wake: wake.promise, + resolveWake: wake.resolve, + wakeSignaled: false, + }; + cycle = active; + void runCycle(active); + return active.promise; }; return { prepare, enqueue(publication) { if (closed) return false; + const lane = getLane(rootFor(publication)); if (Date.now() >= publication.deadlineAt) { - releaseSpaceWaiter(publication)?.resolve(true); - notifySpace(true); - options.onFailure({ reason: 'expired', publication }); + releaseSpaceWaiter(lane, publication)?.resolve(true); + notifySpace(lane, true); + reportFailure({ reason: 'expired', publication }); + cleanupLane(lane); return true; } - if (!hasSpaceFor(publication)) return false; - entries.push({ ...publication }); - bytes += publication.bytes; - releaseSpaceWaiter(publication)?.resolve(true); - notifySpace(true); + if (!hasSpaceFor(lane, publication)) return false; + lane.entries.push({ ...publication }); + lane.bytes += publication.bytes; + releaseSpaceWaiter(lane, publication)?.resolve(true); + notifySpace(lane, true); if (!paused) void pump(); - else scheduleWakeup(); + else scheduleWakeup(lane); return true; }, waitForSpace(publication) { if (closed) return Promise.resolve(false); + const lane = getLane(rootFor(publication)); if (Date.now() >= publication.deadlineAt) return Promise.resolve(true); - const existing = spaceWaiters.get(publication); + const existing = lane.spaceWaiters.get(publication); if (existing) { - if (existing.ready && !hasSpaceFor(publication)) { + if (existing.ready && !hasSpaceFor(lane, publication)) { const { promise, resolve } = Promise.withResolvers(); existing.promise = promise; existing.resolve = resolve; @@ -257,38 +428,44 @@ export function createControlEventOutbox(options: { } return existing.promise; } - if (spaceWaiters.size >= MAX_EVENTS || waitingBytes + publication.bytes > MAX_BYTES) + if (lane.spaceWaiters.size >= MAX_EVENTS || lane.waitingBytes + publication.bytes > MAX_BYTES) return Promise.resolve(false); const { promise, resolve } = Promise.withResolvers(); const timeout = setTimeout( () => { - releaseSpaceWaiter(publication)?.resolve(true); - notifySpace(true); + const released = releaseSpaceWaiter(lane, publication); + released?.resolve(true); + notifySpace(lane, true); }, Math.max(1, publication.deadlineAt - Date.now()) ); timeout.unref(); - const ready = hasSpaceFor(publication); - spaceWaiters.set(publication, { promise, resolve, timeout, ready }); - waitingBytes += publication.bytes; + const ready = hasSpaceFor(lane, publication); + lane.spaceWaiters.set(publication, { promise, resolve, timeout, ready }); + lane.waitingBytes += publication.bytes; if (ready) resolve(true); return promise; }, pause() { paused = true; - scheduleWakeup(); + for (const lane of lanes.values()) scheduleWakeup(lane); }, resume() { paused = false; return pump(); }, close() { + if (closed) return; closed = true; - clearWakeup(); - expirePending?.(); - entries.length = 0; - bytes = 0; - notifySpace(false); + signalCycle(); + for (const lane of lanes.values()) { + clearWakeup(lane); + lane.expirePending?.(); + lane.entries.length = 0; + lane.bytes = 0; + notifySpace(lane, false); + cleanupLane(lane); + } }, }; } diff --git a/services/cloud-agent-next/wrapper/src/control/main.ts b/services/cloud-agent-next/wrapper/src/control/main.ts index fe4c9992a1..f3742fd7ba 100644 --- a/services/cloud-agent-next/wrapper/src/control/main.ts +++ b/services/cloud-agent-next/wrapper/src/control/main.ts @@ -74,7 +74,7 @@ function main( onDiagnostic: diagnostics.onDiagnostic, onRootDisappeared: notifyRootDisappeared, onRootRetirement: settleRootRetirement, - onEvent: async (runtime, event) => { + onEvent: (runtime, event) => { mutationNotifications.observe(runtime, event); const identity = sessionEventIdentity({ ...event, @@ -89,21 +89,24 @@ function main( identity.rootKiloSessionId, event.properties ); - const published = - control?.publishSessionEvent === undefined - ? false - : await control.publishSessionEvent( - { type: event.type, properties: event.properties }, - identity - ); - if (!published) { - try { - await retirePublicationFailure(runtime, identity, 'Session event delivery failed'); - } catch { - diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'failed' }); - } + let publication: Promise; + try { + publication = Promise.resolve( + control?.publishSessionEvent?.( + { type: event.type, properties: event.properties }, + identity + ) ?? false + ); + } catch { + reportPublicationAdmissionFailure(runtime, identity); return; } + void publication.then( + published => { + if (!published) reportPublicationAdmissionFailure(runtime, identity); + }, + () => reportPublicationAdmissionFailure(runtime, identity) + ); }, onUnexpectedClose: failure => { logToFile(`Kilo worktree retired reason=${failure.reason} directory=${failure.directory}`); @@ -210,6 +213,21 @@ function main( return attempt.physical; } + function reportPublicationAdmissionFailure( + runtime: WorktreeKiloRuntime, + identity: SessionEventIdentity + ): void { + try { + void retirePublicationFailure(runtime, identity, 'Session event delivery failed').catch( + () => { + diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'failed' }); + } + ); + } catch { + diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'failed' }); + } + } + function startPublicationFailure(identity: SessionEventIdentity, reason: string): void { const runtime = kiloRuntimes.get(identity.directory); if (!runtime) return; diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-feed.test.ts b/services/cloud-agent-next/wrapper/src/control/worktree-feed.test.ts index 4eb9df8655..053f659a71 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-feed.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-feed.test.ts @@ -55,7 +55,7 @@ function nativeFeed(options: FeedOptions) { }; } -function fixture(rejectReconnections = false) { +function fixture(rejectReconnections = false, onEvent?: (event: KiloFeedEvent) => unknown) { const source = { scopeId: 'worktree_a', runtimeId: crypto.randomUUID(), @@ -79,7 +79,7 @@ function fixture(rejectReconnections = false) { source, isCurrent: (runtimeId, kiloClient) => runtimeId === source.runtimeId && kiloClient === source.kiloClient, - onEvent: event => events.push(event), + onEvent: onEvent ?? (event => events.push(event)), onFailure: reason => failures.push(reason), onDiagnostic: (_event, fields) => diagnostics.push({ @@ -97,6 +97,37 @@ afterEach(() => { }); describe('createWorktreeFeed', () => { + it('does not let a slow event callback delay the next feed callback', async () => { + const firstEntered = Promise.withResolvers(); + const releaseFirst = Promise.withResolvers(); + const received: string[] = []; + const h = fixture(false, async event => { + received.push(String(event.properties.id)); + if (received.length === 1) { + firstEntered.resolve(); + await releaseFirst.promise; + } + }); + cleanups.push(() => { + h.feed.close(); + h.start.mockRestore(); + }); + await h.feed.open(); + const first = h.attempts[0]; + if (!first) throw new Error('Missing native feed'); + await first.emit({ + directory: h.source.directory, + payload: { type: 'message.updated', properties: { id: 'first' } }, + }); + await firstEntered.promise; + await first.emit({ + directory: h.source.directory, + payload: { type: 'message.updated', properties: { id: 'second' } }, + }); + expect(received).toEqual(['first', 'second']); + releaseFirst.resolve(); + }); + it.each([true, false])( 'preserves producer lifetime with event receipts=%s', async eventReceipts => { diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-feed.ts b/services/cloud-agent-next/wrapper/src/control/worktree-feed.ts index 2e9529fe76..e94070203b 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-feed.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-feed.ts @@ -135,7 +135,7 @@ export function createWorktreeFeed(options: { consume: async stream => { for await (const event of unfilteredKiloEvents(stream)) { if (!isCurrentAttempt(attempt)) return; - await options.onEvent?.({ ...event, nativeRuntimeId: runtimeId }); + void options.onEvent?.({ ...event, nativeRuntimeId: runtimeId }); } }, onUnexpectedClose: error => {