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 948060dd53..78a9ca0595 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 @@ -3,6 +3,7 @@ import { MAX_SANDBOX_CONTROL_FRAME_BYTES } from '../../../src/shared/sandbox-con import { controlEventReceiptDisposition, recordControlEventReceipt, + readControlEventReceipts, } from '../../../src/sandbox-session/control-event-receipts'; import { createControlEventTransport } from './control-event-transport'; import { createControlEventOutbox, type ControlEventPublication } from './control-event-outbox'; @@ -22,6 +23,13 @@ const large = { }; const small = { type: 'session.idle', properties: {} }; +function messageUpdatedPayload(id: string, text: string, sessionID = session.kiloSessionId) { + return { + type: 'message.updated', + properties: { info: { id, sessionID, role: 'assistant', text } }, + }; +} + function receiptStorage() { const values = new Map(); return { @@ -124,6 +132,245 @@ describe('control event publication ordering', () => { } }); + it('compares a waiting publication with the current tail when it is admitted', async () => { + const startedFirst = Promise.withResolvers(); + const releaseFirst = Promise.withResolvers(); + const startedSecond = Promise.withResolvers(); + const releaseSecond = Promise.withResolvers(); + const delivered: Array<{ publication: ControlEventPublication; deadlineAt: number }> = []; + const outbox = createControlEventOutbox({ + publish: async (publication, deadlineAt) => { + delivered.push({ publication, deadlineAt }); + if (delivered.length === 1) { + startedFirst.resolve(); + await releaseFirst.promise; + } else if (delivered.length === 2) { + startedSecond.resolve(); + await releaseSecond.promise; + } + }, + onFailure: mock(), + }); + const filler = 'm'.repeat(Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES * 0.45)); + try { + for (let index = 0; index < 8; index += 1) + expect( + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload(`filler_${index}`, filler), + }) + ) + ).toBe(true); + const old = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('target', 'old'), + }); + const waiting = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload( + 'target', + 'l'.repeat(Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES * 0.7)) + ), + }); + expect(outbox.enqueue(old)).toBe(true); + expect(outbox.enqueue(waiting)).toBe(false); + const admitted = outbox.waitForSpace(waiting); + const draining = outbox.resume(); + await startedFirst.promise; + releaseFirst.resolve(); + await startedSecond.promise; + expect(await admitted).toBe(true); + expect(outbox.enqueue(waiting)).toBe(true); + releaseSecond.resolve(); + expect(await draining).toBe(true); + + const targetPublications = delivered.filter( + item => item.publication.sequence === old.sequence + ); + expect(targetPublications).toHaveLength(1); + expect(targetPublications[0]).toMatchObject({ + publication: { + receiptId: old.receiptId, + sequence: old.sequence, + payload: waiting.payload, + }, + deadlineAt: old.deadlineAt, + }); + } finally { + releaseFirst.resolve(); + releaseSecond.resolve(); + outbox.close(); + } + }); + + it('preserves receipt high-water across a squash and same-root child successor', async () => { + const storage = receiptStorage(); + const wrapperInstanceId = crypto.randomUUID(); + const applied: ControlEventPublication[] = []; + const root = session; + const child = { ...session, kiloSessionId: 'ses_child' }; + const outbox = createControlEventOutbox({ + publish: async publication => { + const receipt = { ...publication, wrapperInstanceId }; + expect(controlEventReceiptDisposition(storage, receipt)).toBe('apply'); + recordControlEventReceipt(storage, receipt); + applied.push(publication); + }, + onFailure: mock(), + }); + try { + const first = outbox.prepare({ + event: 'session.event', + session: root, + payload: messageUpdatedPayload('msg_1', 'first'), + }); + const second = outbox.prepare({ + event: 'session.event', + session: root, + payload: messageUpdatedPayload('msg_1', 'latest'), + }); + const successor = outbox.prepare({ + event: 'session.event', + session: child, + payload: { type: 'session.idle', properties: {} }, + }); + expect(outbox.enqueue(first)).toBe(true); + expect(outbox.enqueue(second)).toBe(true); + expect(outbox.enqueue(successor)).toBe(true); + expect(second.sequence).toBe(first.sequence + 1); + expect(successor.sequence).toBe(second.sequence + 1); + + expect(await outbox.resume()).toBe(true); + expect(applied.map(publication => publication.sequence)).toEqual([ + first.sequence, + successor.sequence, + ]); + expect(applied[0]).toMatchObject({ + receiptId: first.receiptId, + sequence: first.sequence, + payload: second.payload, + }); + expect(applied[1]).toMatchObject({ + receiptId: successor.receiptId, + sequence: successor.sequence, + session: child, + payload: successor.payload, + }); + expect( + controlEventReceiptDisposition(storage, { + receiptId: second.receiptId, + sequence: second.sequence, + wrapperInstanceId, + }) + ).toBe('stale'); + expect( + controlEventReceiptDisposition(storage, { + receiptId: successor.receiptId, + sequence: successor.sequence, + wrapperInstanceId, + }) + ).toBe('duplicate'); + expect(readControlEventReceipts(storage).highWater[wrapperInstanceId]).toBe( + successor.sequence + ); + } finally { + outbox.close(); + } + }); + + it('re-evaluates a waiting publication after its matching tail is sent and removed', async () => { + const startedTarget = Promise.withResolvers(); + const releaseTarget = Promise.withResolvers(); + const delivered: Array<{ publication: ControlEventPublication; deadlineAt: number }> = []; + const outbox = createControlEventOutbox({ + publish: async (publication, deadlineAt) => { + delivered.push({ publication, deadlineAt }); + if (delivered.length === 9) { + startedTarget.resolve(); + await releaseTarget.promise; + } + }, + onFailure: mock(), + }); + const filler = 'm'.repeat(Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES * 0.45)); + try { + for (let index = 0; index < 8; index += 1) + expect( + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload(`filler_${index}`, filler), + }) + ) + ).toBe(true); + const prior = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('target', 'prior'), + }); + expect(outbox.enqueue(prior)).toBe(true); + const older = outbox.prepare({ + event: 'session.event', + session, + payload: { type: 'session.idle', properties: {} }, + }); + const olderWaiting = outbox.waitForSpace(older); + const waiting = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload( + 'target', + 'l'.repeat(Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES * 0.7)) + ), + }); + expect(waiting.bytes).toBeGreaterThan(prior.bytes); + expect(outbox.enqueue(waiting)).toBe(false); + const waitingReady = outbox.waitForSpace(waiting); + let waitingSettled = false; + void waitingReady.then(() => { + waitingSettled = true; + }); + await Promise.resolve(); + expect(waitingSettled).toBe(false); + const draining = outbox.resume(); + await startedTarget.promise; + expect(waitingSettled).toBe(false); + releaseTarget.resolve(); + expect(await draining).toBe(true); + expect(waitingSettled).toBe(false); + expect(await olderWaiting).toBe(true); + expect(outbox.enqueue(older)).toBe(true); + expect(await waitingReady).toBe(true); + expect(outbox.enqueue(waiting)).toBe(true); + expect(await outbox.resume()).toBe(true); + expect(delivered).toHaveLength(11); + expect(delivered[8]).toMatchObject({ + publication: { + receiptId: prior.receiptId, + sequence: prior.sequence, + payload: prior.payload, + }, + deadlineAt: prior.deadlineAt, + }); + expect(delivered.at(-1)).toMatchObject({ + publication: { + receiptId: waiting.receiptId, + sequence: waiting.sequence, + payload: waiting.payload, + }, + deadlineAt: waiting.deadlineAt, + }); + } finally { + releaseTarget.resolve(); + outbox.close(); + } + }); + it.each(['count', 'bytes'] as const)( 'bounds waiting reservations by %s and coalesces repeated waits', async budget => { 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 55b01f1a34..4899764fe2 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 @@ -1,6 +1,10 @@ import { describe, expect, it, mock, spyOn } from 'bun:test'; import { ControlDeliveryError } from './sandbox-control-client'; -import { createControlEventOutbox, type ControlEventPublication } from './control-event-outbox'; +import { + createControlEventOutbox, + type ControlEventPublication, + type PreparedControlEventPublication, +} from './control-event-outbox'; import { MAX_SANDBOX_CONTROL_FRAME_BYTES, sandboxEventPublicationPayloadSchema, @@ -12,6 +16,27 @@ const session = { rootKiloSessionId: 'ses_root', }; +function messageUpdatedPayload(id: string, marker: string, sessionID = session.kiloSessionId) { + return { + type: 'message.updated', + properties: { info: { id, sessionID, role: 'assistant', marker } }, + }; +} + +function partUpdatedPayload( + messageID: string, + id: string, + marker: string, + sessionID = session.kiloSessionId +) { + return { + type: 'message.part.updated', + properties: { + part: { id, messageID, sessionID, type: 'text', text: marker }, + }, + }; +} + async function waitFor(condition: () => boolean, attempts = 100): Promise { for (let attempt = 0; attempt < attempts; attempt += 1) { if (condition()) return; @@ -21,6 +46,590 @@ async function waitFor(condition: () => boolean, attempts = 100): Promise } describe('control event outbox', () => { + it('squashes adjacent same-part updates while retaining the original receipt metadata', async () => { + const delivered: Array<{ publication: ControlEventPublication; deadlineAt: number }> = []; + const outbox = createControlEventOutbox({ + publish: async (publication, deadlineAt) => { + delivered.push({ publication, deadlineAt }); + }, + onFailure: mock(), + }); + try { + const barrier = outbox.prepare({ + event: 'session.event', + session, + payload: { type: 'session.idle', properties: {} }, + }); + const first = outbox.prepare({ + event: 'session.event', + session, + payload: partUpdatedPayload('msg_1', 'part_1', 'first'), + }); + const second = outbox.prepare({ + event: 'session.event', + session, + payload: partUpdatedPayload('msg_1', 'part_1', 'latest'), + }); + expect(outbox.enqueue(barrier)).toBe(true); + expect(outbox.enqueue(first)).toBe(true); + expect(outbox.enqueue(second)).toBe(true); + + expect(await outbox.resume()).toBe(true); + expect(delivered).toHaveLength(2); + expect(delivered.map(item => item.publication.sequence)).toEqual([1, 2]); + expect(delivered[0]?.publication.payload).toEqual(barrier.payload); + expect(delivered[1]?.publication).toMatchObject({ + receiptId: first.receiptId, + sequence: first.sequence, + payload: second.payload, + }); + expect(delivered[1]?.deadlineAt).toBe(first.deadlineAt); + } finally { + outbox.close(); + } + }); + + it('squashes adjacent same-message updates to the latest payload', async () => { + const clock = spyOn(Date, 'now').mockReturnValue(1_000); + const delivered: Array<{ publication: ControlEventPublication; deadlineAt: number }> = []; + const outbox = createControlEventOutbox({ + publish: async (publication, deadlineAt) => { + delivered.push({ publication, deadlineAt }); + }, + onFailure: mock(), + }); + try { + const first = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_1', 'first'), + }); + clock.mockReturnValue(2_000); + const second = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_1', 'latest'), + }); + expect(first.deadlineAt).toBe(31_000); + expect(second.deadlineAt).toBe(32_000); + expect(first.deadlineAt).toBeLessThan(second.deadlineAt); + expect(outbox.enqueue(first)).toBe(true); + expect(outbox.enqueue(second)).toBe(true); + + expect(await outbox.resume()).toBe(true); + expect(delivered).toHaveLength(1); + expect(delivered[0]?.publication).toMatchObject({ + receiptId: first.receiptId, + sequence: first.sequence, + payload: second.payload, + }); + expect(delivered[0]?.deadlineAt).toBe(first.deadlineAt); + } finally { + outbox.close(); + clock.mockRestore(); + } + }); + + it('keeps an in-flight head unchanged while squashing an unsent tail', async () => { + const started = Promise.withResolvers(); + const release = Promise.withResolvers(); + const delivered: Array<{ publication: ControlEventPublication; deadlineAt: number }> = []; + const outbox = createControlEventOutbox({ + publish: async (publication, deadlineAt) => { + delivered.push({ publication, deadlineAt }); + if (delivered.length === 1) { + started.resolve(); + await release.promise; + } + }, + onFailure: mock(), + }); + try { + const first = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_1', 'in flight'), + }); + const second = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_2', 'queued'), + }); + const latest = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_2', 'latest'), + }); + expect(outbox.enqueue(first)).toBe(true); + expect(outbox.enqueue(second)).toBe(true); + const draining = outbox.resume(); + await started.promise; + expect(outbox.enqueue(latest)).toBe(true); + release.resolve(); + + expect(await draining).toBe(true); + expect(delivered).toHaveLength(2); + expect(delivered[0]).toMatchObject({ + publication: { + receiptId: first.receiptId, + sequence: first.sequence, + payload: first.payload, + }, + deadlineAt: first.deadlineAt, + }); + expect(delivered[1]).toMatchObject({ + publication: { + receiptId: second.receiptId, + sequence: second.sequence, + payload: latest.payload, + }, + deadlineAt: second.deadlineAt, + }); + } finally { + release.resolve(); + outbox.close(); + } + }); + + it('does not squash a newer update into a single in-flight publication', async () => { + const started = Promise.withResolvers(); + const release = Promise.withResolvers(); + const delivered: Array<{ publication: ControlEventPublication; deadlineAt: number }> = []; + const outbox = createControlEventOutbox({ + publish: async (publication, deadlineAt) => { + delivered.push({ publication, deadlineAt }); + if (delivered.length === 1) { + started.resolve(); + await release.promise; + } + }, + onFailure: mock(), + }); + try { + const first = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_1', 'first'), + }); + const second = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_1', 'latest'), + }); + expect(outbox.enqueue(first)).toBe(true); + const draining = outbox.resume(); + await started.promise; + expect(outbox.enqueue(second)).toBe(true); + release.resolve(); + + expect(await draining).toBe(true); + expect(delivered).toHaveLength(2); + expect(delivered[0]).toMatchObject({ + publication: { + receiptId: first.receiptId, + sequence: first.sequence, + payload: first.payload, + }, + deadlineAt: first.deadlineAt, + }); + expect(delivered[1]).toMatchObject({ + publication: { + receiptId: second.receiptId, + sequence: second.sequence, + payload: second.payload, + }, + deadlineAt: second.deadlineAt, + }); + expect(first.receiptId).not.toBe(second.receiptId); + expect(first.sequence).not.toBe(second.sequence); + } finally { + release.resolve(); + outbox.close(); + } + }); + + it('squashes the new head after the prior head is removed before its attempt settles', async () => { + const delivered: Array<{ publication: ControlEventPublication; deadlineAt: number }> = []; + let replacement: PreparedControlEventPublication | undefined = undefined; + let replacementAdmitted = false; + const outbox = createControlEventOutbox({ + publish: async (publication, deadlineAt) => { + delivered.push({ publication, deadlineAt }); + if (delivered.length === 1) { + throw new ControlDeliveryError('permanent failure', false); + } + }, + onFailure: failure => { + if (failure.reason === 'rejected' && replacement !== undefined) + replacementAdmitted = outbox.enqueue(replacement); + }, + }); + const preparedFirst = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_1', 'failed'), + }); + const preparedSecond = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_2', 'queued'), + }); + const preparedLatest = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_2', 'latest'), + }); + replacement = preparedLatest; + try { + expect(outbox.enqueue(preparedFirst)).toBe(true); + expect(outbox.enqueue(preparedSecond)).toBe(true); + expect(await outbox.resume()).toBe(true); + expect(replacementAdmitted).toBe(true); + expect(delivered).toHaveLength(2); + expect(delivered[1]).toMatchObject({ + publication: { + receiptId: preparedSecond.receiptId, + sequence: preparedSecond.sequence, + payload: preparedLatest.payload, + }, + deadlineAt: preparedSecond.deadlineAt, + }); + } finally { + outbox.close(); + } + }); + + it('does not squash across a message and part entity barrier', async () => { + const published: ControlEventPublication[] = []; + const outbox = createControlEventOutbox({ + publish: async publication => { + published.push(publication); + }, + onFailure: mock(), + }); + try { + const first = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_1', 'message before'), + }); + const barrier = outbox.prepare({ + event: 'session.event', + session, + payload: partUpdatedPayload('msg_1', 'part_1', 'part barrier'), + }); + const last = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_1', 'message after'), + }); + expect(outbox.enqueue(first)).toBe(true); + expect(outbox.enqueue(barrier)).toBe(true); + expect(outbox.enqueue(last)).toBe(true); + + expect(await outbox.resume()).toBe(true); + expect(published).toHaveLength(3); + expect(published.map(item => item.sequence)).toEqual([1, 2, 3]); + expect(published.map(item => item.payload)).toEqual([ + first.payload, + barrier.payload, + last.payload, + ]); + } finally { + outbox.close(); + } + }); + + it.each(['entity', 'root', 'native'] as const)( + 'does not squash when the %s is different', + async fence => { + const published: ControlEventPublication[] = []; + const outbox = createControlEventOutbox({ + publish: async publication => { + published.push(publication); + }, + onFailure: mock(), + }); + const firstNativeRuntimeId = crypto.randomUUID(); + const secondNativeRuntimeId = crypto.randomUUID(); + const firstSession = + fence === 'root' + ? { ...session, kiloSessionId: 'root_a', rootKiloSessionId: 'root_a' } + : { ...session, nativeRuntimeId: firstNativeRuntimeId }; + const secondSession = + fence === 'root' + ? { ...session, kiloSessionId: 'root_b', rootKiloSessionId: 'root_b' } + : fence === 'native' + ? { ...session, nativeRuntimeId: secondNativeRuntimeId } + : firstSession; + const first = outbox.prepare({ + event: 'session.event', + session: firstSession, + payload: messageUpdatedPayload('msg_1', 'first', firstSession.kiloSessionId), + }); + const second = outbox.prepare({ + event: 'session.event', + session: secondSession, + payload: messageUpdatedPayload( + fence === 'entity' ? 'msg_2' : 'msg_1', + 'second', + secondSession.kiloSessionId + ), + }); + try { + expect(outbox.enqueue(first)).toBe(true); + expect(outbox.enqueue(second)).toBe(true); + expect(await outbox.resume()).toBe(true); + expect(published).toHaveLength(2); + expect(published.map(item => item.payload)).toEqual([first.payload, second.payload]); + } finally { + outbox.close(); + } + } + ); + + it.each([ + [ + 'outcome', + 'session.event', + { type: 'session.message.outcome', properties: { messageId: 'msg_1', status: 'completed' } }, + ], + [ + 'removed', + 'session.event', + { type: 'message.removed', properties: { sessionID: 'ses_root', messageID: 'msg_1' } }, + ], + [ + 'lifecycle', + 'session.event', + { type: 'session.updated', properties: { info: { id: 'ses_root' } } }, + ], + [ + 'preparing', + 'session.preparing', + { + version: 2, + attemptId: 'attempt_1', + triggerMessageId: 'msg_1', + revision: 0, + timestamp: 1, + step: 'started', + message: 'preparing', + action: 'start', + }, + ], + [ + 'non-entity', + 'session.event', + { type: 'session.status', properties: { sessionID: 'ses_root' } }, + ], + ] as const)('does not squash across a %s barrier', async (_name, event, payload) => { + const published: ControlEventPublication[] = []; + const outbox = createControlEventOutbox({ + publish: async publication => { + published.push(publication); + }, + onFailure: mock(), + }); + try { + const first = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_1', 'before'), + }); + const barrier = outbox.prepare({ event, session, payload }); + const last = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_1', 'after'), + }); + expect(outbox.enqueue(first)).toBe(true); + expect(outbox.enqueue(barrier)).toBe(true); + expect(outbox.enqueue(last)).toBe(true); + expect(await outbox.resume()).toBe(true); + expect(published).toHaveLength(3); + expect(published.map(item => item.sequence)).toEqual([1, 2, 3]); + } finally { + outbox.close(); + } + }); + + it('keeps a retry-held head unchanged and queues a newer same-message update behind it', async () => { + const delivered: Array<{ publication: ControlEventPublication; deadlineAt: number }> = []; + const retried = Promise.withResolvers(); + const outbox = createControlEventOutbox({ + publish: async (publication, deadlineAt) => { + delivered.push({ publication, deadlineAt }); + if (delivered.length === 1) throw new ControlDeliveryError('not attached', true); + if (delivered.length === 2) retried.resolve(); + }, + onFailure: mock(), + }); + try { + const first = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_1', 'first'), + }); + const second = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('msg_1', 'latest'), + }); + expect(outbox.enqueue(first)).toBe(true); + expect(await outbox.resume()).toBe(false); + expect(outbox.enqueue(second)).toBe(true); + await retried.promise; + await waitFor(() => delivered.length === 3, 500); + expect(delivered[0]?.publication).toEqual(delivered[1]?.publication); + expect(delivered[0]?.deadlineAt).toBe(first.deadlineAt); + expect(delivered[2]?.publication).toMatchObject({ + receiptId: second.receiptId, + sequence: second.sequence, + payload: second.payload, + }); + expect(await outbox.resume()).toBe(true); + } finally { + outbox.close(); + } + }); + + it('uses post-replacement bytes to admit a smaller adjacent update', async () => { + const delivered: ControlEventPublication[] = []; + const outbox = createControlEventOutbox({ + publish: async publication => { + delivered.push(publication); + }, + onFailure: mock(), + }); + const filler = 'm'.repeat(Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES * 0.45)); + const oldText = 'o'.repeat(Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES * 0.35)); + const smallerText = 's'.repeat(Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES * 0.3)); + try { + for (let index = 0; index < 8; index += 1) + expect( + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload(`filler_${index}`, filler), + }) + ) + ).toBe(true); + const old = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('target', oldText), + }); + const smaller = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('target', smallerText), + }); + expect(smaller.bytes).toBeLessThan(old.bytes); + expect(outbox.enqueue(old)).toBe(true); + expect(outbox.enqueue(smaller)).toBe(true); + expect(await outbox.resume()).toBe(true); + expect(delivered).toHaveLength(9); + expect(delivered.at(-1)).toMatchObject({ + receiptId: old.receiptId, + sequence: old.sequence, + payload: smaller.payload, + }); + } finally { + outbox.close(); + } + }); + + it('rejects a larger adjacent replacement without mutating the old payload', async () => { + const delivered: ControlEventPublication[] = []; + const outbox = createControlEventOutbox({ + publish: async publication => { + delivered.push(publication); + }, + onFailure: mock(), + }); + const filler = 'm'.repeat(Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES * 0.45)); + const largerText = 'l'.repeat(Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES * 0.7)); + try { + for (let index = 0; index < 8; index += 1) + expect( + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload(`filler_${index}`, filler), + }) + ) + ).toBe(true); + const old = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('target', 'old'), + }); + const larger = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('target', largerText), + }); + expect(larger.bytes).toBeGreaterThan(old.bytes); + expect(outbox.enqueue(old)).toBe(true); + expect(outbox.enqueue(larger)).toBe(false); + expect(await outbox.resume()).toBe(true); + expect(delivered).toHaveLength(9); + expect(delivered.at(-1)).toMatchObject({ + receiptId: old.receiptId, + sequence: old.sequence, + payload: old.payload, + }); + } finally { + outbox.close(); + } + }); + + it('admits a tail replacement at the 256-entry count boundary', async () => { + const published: ControlEventPublication[] = []; + const outbox = createControlEventOutbox({ + publish: async publication => { + published.push(publication); + }, + onFailure: mock(), + }); + try { + for (let index = 0; index < 255; index += 1) + expect( + outbox.enqueue( + outbox.prepare({ + event: 'session.event', + session, + payload: { type: 'session.idle', properties: {} }, + }) + ) + ).toBe(true); + const old = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('target', 'old'), + }); + const latest = outbox.prepare({ + event: 'session.event', + session, + payload: messageUpdatedPayload('target', 'latest'), + }); + expect(outbox.enqueue(old)).toBe(true); + expect(outbox.enqueue(latest)).toBe(true); + expect(await outbox.resume()).toBe(true); + expect(published).toHaveLength(256); + expect(published.at(-1)).toMatchObject({ + receiptId: old.receiptId, + sequence: old.sequence, + payload: latest.payload, + }); + } finally { + outbox.close(); + } + }); + it('does not let a retryable root A head delay root B', async () => { const published: ControlEventPublication[] = []; let attemptsA = 0; 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 905a175de3..2601346c13 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 @@ -40,6 +40,12 @@ export type ControlEventOutbox = { type RootKey = string | undefined; +type SquashKey = { + entityId: string; + root: RootKey; + nativeRuntimeId: string | undefined; +}; + type SpaceWaiter = { promise: Promise; ready: boolean; @@ -54,6 +60,7 @@ type Lane = { waitingBytes: number; bytes: number; pending?: Promise; + pendingEntry?: PreparedControlEventPublication; expirePending?: () => void; retryAt?: number; wakeup?: ReturnType; @@ -71,6 +78,47 @@ function rootFor(publication: PreparedControlEventPublication): RootKey { return publication.session.rootKiloSessionId ?? publication.session.kiloSessionId; } +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function entityIdFor(payload: unknown): string | undefined { + if (!isRecord(payload) || !isRecord(payload.properties)) return undefined; + if (payload.type === 'message.updated') { + const info = payload.properties.info; + return isRecord(info) && typeof info.id === 'string' ? `message/${info.id}` : undefined; + } + if (payload.type === 'message.part.updated') { + const part = payload.properties.part; + if (!isRecord(part)) return undefined; + return typeof part.messageID === 'string' && typeof part.id === 'string' + ? `part/${part.messageID}/${part.id}` + : undefined; + } + return undefined; +} + +function squashKeyFor(publication: PreparedControlEventPublication): SquashKey | undefined { + if (publication.event !== 'session.event') return undefined; + const entityId = entityIdFor(publication.payload); + if (!entityId) return undefined; + return { + entityId, + root: rootFor(publication), + nativeRuntimeId: publication.session.nativeRuntimeId, + }; +} + +function sameSquashKey(left: SquashKey | undefined, right: SquashKey | undefined): boolean { + return ( + left !== undefined && + right !== undefined && + left.entityId === right.entityId && + left.root === right.root && + left.nativeRuntimeId === right.nativeRuntimeId + ); +} + function isRetryable(error: unknown): boolean { return ( typeof error === 'object' && error !== null && 'retryable' in error && error.retryable === true @@ -113,6 +161,17 @@ export function createControlEventOutbox(options: { if (lanes.get(lane.root) === lane) lanes.delete(lane.root); }; + const squashTarget = ( + lane: Lane, + publication: PreparedControlEventPublication + ): PreparedControlEventPublication | undefined => { + const previous = lane.entries.at(-1); + if (!previous) return undefined; + if (previous === lane.pendingEntry) return undefined; + if (previous === lane.entries[0] && lane.retryAt !== undefined) return undefined; + return sameSquashKey(squashKeyFor(previous), squashKeyFor(publication)) ? previous : undefined; + }; + const reportFailure = (failure: ControlEventOutboxFailure): void => { try { options.onFailure(failure); @@ -121,9 +180,14 @@ export function createControlEventOutbox(options: { } }; - const hasSpaceFor = (lane: Lane, publication: PreparedControlEventPublication): boolean => { - if (lane.entries.length >= MAX_EVENTS || lane.bytes + publication.bytes > MAX_BYTES) - return false; + const hasSpaceFor = ( + lane: Lane, + publication: PreparedControlEventPublication, + replacement = squashTarget(lane, publication) + ): boolean => { + const entryCount = lane.entries.length - (replacement ? 1 : 0); + const bytes = lane.bytes - (replacement?.bytes ?? 0); + if (entryCount >= MAX_EVENTS || 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; @@ -329,10 +393,14 @@ export function createControlEventOutbox(options: { reportFailure({ reason: 'rejected', publication: entry }); }) .then(() => { - if (lane.pending === pending) lane.pending = undefined; + if (lane.pending === pending) { + lane.pending = undefined; + lane.pendingEntry = undefined; + } scheduleWakeup(lane); cleanupLane(lane); }); + lane.pendingEntry = entry; lane.pending = pending; }; @@ -405,9 +473,20 @@ export function createControlEventOutbox(options: { cleanupLane(lane); return true; } - if (!hasSpaceFor(lane, publication)) return false; - lane.entries.push({ ...publication }); - lane.bytes += publication.bytes; + const replacement = squashTarget(lane, publication); + if (!hasSpaceFor(lane, publication, replacement)) return false; + if (replacement) { + const index = lane.entries.length - 1; + lane.entries[index] = { + ...replacement, + payload: publication.payload, + bytes: publication.bytes, + }; + lane.bytes += publication.bytes - replacement.bytes; + } else { + lane.entries.push({ ...publication }); + lane.bytes += publication.bytes; + } releaseSpaceWaiter(lane, publication)?.resolve(true); notifySpace(lane, true); if (!paused) void pump(); @@ -461,6 +540,7 @@ export function createControlEventOutbox(options: { for (const lane of lanes.values()) { clearWakeup(lane); lane.expirePending?.(); + lane.pendingEntry = undefined; lane.entries.length = 0; lane.bytes = 0; notifySpace(lane, false);