fix(cloud-agent-next): squash adjacent control events

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<string, unknown>();
   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<void>();
+    const releaseFirst = Promise.withResolvers<void>();
+    const startedSecond = Promise.withResolvers<void>();
+    const releaseSecond = Promise.withResolvers<void>();
+    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<void>();
+    const releaseTarget = Promise.withResolvers<void>();
+    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 => {