diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index 5363b33321..c71071089a 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -951,6 +951,17 @@ export class CloudAgentSession extends DurableObject { const disconnected = await ingestHandler.handleIngestClose(ws); if (disconnected) { + if (code === 1009) { + logger + .withFields({ + sessionId: this.sessionId, + wrapperRunId: disconnected.wrapperRunId, + wrapperGeneration: disconnected.wrapperGeneration, + wrapperConnectionId: disconnected.wrapperConnectionId, + logTag: 'wrapper_ingest_closed_message_too_large', + }) + .warn('Wrapper ingest closed because a message exceeded the platform size limit'); + } const wrapperSupervisor = this.getWrapperSupervisor(); await wrapperSupervisor.onDisconnected({ disconnected, diff --git a/services/cloud-agent-next/src/session/wrapper-runtime-state.ts b/services/cloud-agent-next/src/session/wrapper-runtime-state.ts index 363eaf07bf..1e7eda3254 100644 --- a/services/cloud-agent-next/src/session/wrapper-runtime-state.ts +++ b/services/cloud-agent-next/src/session/wrapper-runtime-state.ts @@ -742,6 +742,34 @@ export async function recordWrapperPong( })); } +/** + * Reset liveness deadlines that went stale with the previous socket after an + * accepted reconnect. A ping (or its pong) lost with the old socket can never + * be satisfied, so a fresh ping is scheduled instead of letting the stale + * deadline expire into a wrapper_ping_timeout. The no-output deadline kept + * ticking while delivery was impossible, so it is extended to a fresh window + * rather than firing wrapper_no_output before buffered output can drain. + */ +export async function resetWrapperLivenessAfterReconnect( + storage: DurableObjectStorage, + wrapperGeneration: number, + wrapperConnectionId: string, + nextPingAt: number, + noOutputDeadlineAt: number +): Promise { + return updateIfCurrent(storage, wrapperGeneration, wrapperConnectionId, current => { + const next = { ...current }; + if (next.pingDeadlineAt !== undefined) { + next.pingDeadlineAt = undefined; + next.nextPingAt = nextPingAt; + } + if (next.noOutputDeadlineAt !== undefined) { + next.noOutputDeadlineAt = Math.max(next.noOutputDeadlineAt, noOutputDeadlineAt); + } + return next; + }); +} + export async function markWrapperFinalizing( storage: DurableObjectStorage, wrapperRunId: string diff --git a/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts b/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts index 0d4e9417cb..ec6fc666d9 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts @@ -208,6 +208,24 @@ function liveRuntimeState(overrides?: Record): [string, unknown ]; } +function disconnectGraceForCurrentConnection( + disconnectedAt: number, + overrides?: Record +): [string, unknown] { + return [ + 'disconnect_grace', + { + wrapperRunId: WRAPPER_RUN_ID, + disconnectedAt, + wsCloseCode: 1009, + wsCloseReason: 'message too large', + wrapperGeneration: 4, + wrapperConnectionId: WRAPPER_CONNECTION_ID, + ...overrides, + }, + ]; +} + describe('WrapperSupervisor', () => { it('starts disconnect grace for current accepted work and cancels it after an approved fenced reconnect', async () => { const harness = createHarness([liveRuntimeState()]); @@ -928,6 +946,175 @@ describe('WrapperSupervisor', () => { }); }); + it('defers liveness failure while disconnect grace is active for the current connection', async () => { + const pingDeadlineAt = 92_000; + const noOutputDeadlineAt = 332_000; + // Grace active from 90_000 through 100_000; ping deadline (92_000) already expired. + const harness = createHarness([ + liveRuntimeState({ pingDeadlineAt, noOutputDeadlineAt }), + OWNED_WRAPPER_LEASE, + disconnectGraceForCurrentConnection(90_000), + ]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + await harness.supervisor.runMaintenance(95_000); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'accepted', + }); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'owns_wrapper', + }); + expect(harness.events).toHaveLength(0); + }); + + it('does not defer liveness for a stale disconnect grace left by a previous wrapper run', async () => { + const pingDeadlineAt = 92_000; + const harness = createHarness([ + liveRuntimeState({ + wrapperRunId: 'wr_newer_current', + wrapperGeneration: 5, + wrapperConnectionId: 'conn_newer', + pingDeadlineAt, + noOutputDeadlineAt: 332_000, + }), + OWNED_WRAPPER_LEASE, + disconnectGraceForCurrentConnection(90_000, { + wrapperRunId: 'wr_stale_run', + wrapperGeneration: 4, + wrapperConnectionId: WRAPPER_CONNECTION_ID, + }), + ]); + await putSessionMessageState(harness.storage, { + ...acceptedMessage(), + wrapperRunId: 'wr_newer_current', + }); + + await harness.supervisor.runMaintenance(pingDeadlineAt); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + failureCode: 'wrapper_ping_timeout', + }); + }); + + it('clears grace without terminalizing accepted work when the same wrapper reconnects during grace', async () => { + const harness = createHarness([ + liveRuntimeState({ noOutputDeadlineAt: 332_000, nextPingAt: 200_000 }), + OWNED_WRAPPER_LEASE, + disconnectGraceForCurrentConnection(90_000), + ]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + const decision = await harness.supervisor.checkReconnect({ + wrapperRunId: WRAPPER_RUN_ID, + wrapperGeneration: 4, + wrapperConnectionId: WRAPPER_CONNECTION_ID, + }); + expect(decision).toEqual({ accepted: true } satisfies WrapperReconnectDecision); + await harness.supervisor.recordReconnectAccepted({ + wrapperGeneration: 4, + wrapperConnectionId: WRAPPER_CONNECTION_ID, + }); + + await expect(harness.storage.get('disconnect_grace')).resolves.toBeUndefined(); + // Grace is cleared, so liveness is no longer suppressed. No deadline has + // expired yet (nextPingAt 200_000, noOutput 332_000), so work stays accepted. + await harness.supervisor.runMaintenance(95_000); + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'accepted', + }); + }); + + it('clears a stale in-flight ping when a reconnect is accepted during grace', async () => { + const harness = createHarness([ + liveRuntimeState({ pingDeadlineAt: 92_000, noOutputDeadlineAt: 332_000 }), + OWNED_WRAPPER_LEASE, + disconnectGraceForCurrentConnection(90_000), + ]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + await harness.supervisor.recordReconnectAccepted( + { wrapperGeneration: 4, wrapperConnectionId: WRAPPER_CONNECTION_ID }, + 95_000 + ); + + // The ping (or its pong) died with the old socket and can never resolve on + // the new one; a fresh ping is scheduled instead of letting the stale + // expired deadline fire wrapper_ping_timeout right after reconnecting. + await expect(getWrapperRuntimeState(harness.storage)).resolves.toMatchObject({ + pingDeadlineAt: undefined, + nextPingAt: 155_000, + }); + await harness.supervisor.runMaintenance(95_000); + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'accepted', + }); + }); + + it('extends an expired no-output deadline when a reconnect is accepted during grace', async () => { + const harness = createHarness([ + liveRuntimeState({ noOutputDeadlineAt: 94_000, nextPingAt: 200_000 }), + OWNED_WRAPPER_LEASE, + disconnectGraceForCurrentConnection(90_000), + ]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + await harness.supervisor.recordReconnectAccepted( + { wrapperGeneration: 4, wrapperConnectionId: WRAPPER_CONNECTION_ID }, + 95_000 + ); + + // Output could not be delivered while the socket was down, so the stale + // deadline gets a fresh window instead of firing wrapper_no_output on the + // first maintenance tick after the reconnect. + await expect(getWrapperRuntimeState(harness.storage)).resolves.toMatchObject({ + noOutputDeadlineAt: 95_000 + WRAPPER_NO_OUTPUT_TIMEOUT_MS, + }); + await harness.supervisor.runMaintenance(95_000); + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'accepted', + }); + }); + + it('terminalizes accepted work as wrapper_disconnected once grace expires without reconnect', async () => { + const harness = createHarness([ + liveRuntimeState({ pingDeadlineAt: 92_000, noOutputDeadlineAt: 332_000 }), + OWNED_WRAPPER_LEASE, + disconnectGraceForCurrentConnection(90_000), + ]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + // Grace expired at 100_000; liveness was deferred while it was active. + await harness.supervisor.runMaintenance(100_001); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + failureReason: 'wrapper_disconnected', + failureCode: 'wrapper_disconnected', + }); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'unhealthy-wrapper', + }); + }); + + it('includes the disconnect grace expiry deadline even when the ping deadline is earlier', async () => { + const harness = createHarness([ + liveRuntimeState({ pingDeadlineAt: 92_000, noOutputDeadlineAt: 332_000 }), + OWNED_WRAPPER_LEASE, + disconnectGraceForCurrentConnection(90_000), + ]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + const deadlines = await harness.supervisor.nextMaintenanceDeadlines(); + + // Grace expiry (100_000) must remain scheduled alongside the earlier + // expired ping deadline (92_000) so reconnect gets its full window. + expect(deadlines).toContain(100_000); + expect(deadlines).toContain(92_000); + }); + it('releases a finalizing callback on liveness expiry but holds a queued follow-up until physical absence', async () => { const harness = createHarness( [ diff --git a/services/cloud-agent-next/src/session/wrapper-supervisor.ts b/services/cloud-agent-next/src/session/wrapper-supervisor.ts index a5c8648a1e..da9a7895ea 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.ts @@ -53,6 +53,7 @@ import { recordWrapperPong, reduceSandboxRecoveryState, reduceWrapperLease, + resetWrapperLivenessAfterReconnect, type WrapperConnectionFence, type WrapperRuntimeState, WRAPPER_STOP_MAX_ATTEMPTS, @@ -124,7 +125,7 @@ export type WrapperSupervisorStorage = DurableObjectStorage & export type WrapperSupervisor = { checkReconnect(input: WrapperReconnectInput): Promise; - recordReconnectAccepted(fence: WrapperConnectionFence): Promise; + recordReconnectAccepted(fence: WrapperConnectionFence, now?: number): Promise; isCurrentConnection(wrapperGeneration: number, wrapperConnectionId: string): Promise; observePong(wrapperGeneration: number, wrapperConnectionId: string, now: number): Promise; observeMeaningfulOutput( @@ -388,6 +389,28 @@ export function createWrapperSupervisor( await storage.delete(DISCONNECT_GRACE_KEY); } + /** + * A disconnect grace is "active for the current connection" when it has not + * yet expired and still describes the wrapper runtime we are supervising. + * While active, liveness checks are deferred so a reconnectable transport + * failure (e.g. a `1009` oversize close) gets its full grace window instead + * of immediately being converted into a `wrapper_ping_timeout` physical stop. + * + * A grace left behind by a stale wrapper run/generation does not suppress + * liveness for a newer current connection. + */ + async function isActiveDisconnectGraceForCurrentConnection(now: number): Promise { + const graceState = await readDisconnectGrace(); + if (!graceState) return false; + if (now - graceState.disconnectedAt >= DISCONNECT_GRACE_MS) return false; + const state = await getWrapperRuntimeState(storage); + return ( + state.wrapperRunId === graceState.wrapperRunId && + state.wrapperGeneration === graceState.wrapperGeneration && + state.wrapperConnectionId === graceState.wrapperConnectionId + ); + } + async function releaseWrapperTerminalWaitForIdleBatch(): Promise { await messageSettlementOutbox.releaseWrapperTerminalWaitForIdleBatch(); await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady({ @@ -428,8 +451,27 @@ export function createWrapperSupervisor( return { accepted: true }; } - async function recordReconnectAccepted(fence: WrapperConnectionFence): Promise { + async function recordReconnectAccepted( + fence: WrapperConnectionFence, + now = Date.now() + ): Promise { await cancelDisconnectGrace(fence); + // A ping in flight when the previous socket closed can never be answered on + // the new socket (the wrapper only pongs in response to a ping command), so + // its stale deadline would expire into a wrapper_ping_timeout right after a + // successful reconnect. Clear it and schedule a fresh ping instead. The + // no-output deadline kept ticking while output could not be delivered, so + // it gets a fresh window too instead of firing wrapper_no_output before + // the wrapper's buffered events can drain. + if (fence.wrapperGeneration !== undefined && fence.wrapperConnectionId !== undefined) { + await resetWrapperLivenessAfterReconnect( + storage, + fence.wrapperGeneration, + fence.wrapperConnectionId, + now + WRAPPER_PING_INTERVAL_MS, + now + WRAPPER_NO_OUTPUT_TIMEOUT_MS + ); + } } async function isCurrentConnection( @@ -1425,7 +1467,21 @@ export function createWrapperSupervisor( await reconcilePhysicalCleanup(now); await reconcileSharedSandboxFailover(); await checkDisconnectGrace(now); - await checkWrapperLiveness(now); + // While the current wrapper connection is inside its disconnect grace + // window, defer liveness checks. Otherwise an already-expired ping/no-output + // deadline would immediately terminalize accepted work as + // `wrapper_ping_timeout`/`wrapper_no_output` before the wrapper gets a + // chance to reconnect, defeating the purpose of the grace period. + if (await isActiveDisconnectGraceForCurrentConnection(now)) { + logger + .withFields({ + sessionId: getSessionIdForLogs(), + logTag: 'wrapper_disconnect_grace_liveness_deferred', + }) + .debug('Deferring wrapper liveness check while disconnect grace is active'); + } else { + await checkWrapperLiveness(now); + } await checkKeepWarmCleanup(now); } diff --git a/services/cloud-agent-next/src/shared/ingest-frame.test.ts b/services/cloud-agent-next/src/shared/ingest-frame.test.ts new file mode 100644 index 0000000000..19ebbffd46 --- /dev/null +++ b/services/cloud-agent-next/src/shared/ingest-frame.test.ts @@ -0,0 +1,305 @@ +import { describe, expect, it } from 'vitest'; +import { + MAX_INGEST_BUFFERED_BYTES, + MAX_INGEST_EVENT_BYTES, + IngestEventBuffer, + isLifecycleIngestEvent, + prepareIngestFrame, +} from './ingest-frame.js'; +import type { IngestEvent } from './protocol.js'; + +describe('prepareIngestFrame', () => { + it('passes small events through unchanged', () => { + const event: IngestEvent = { + streamEventType: 'kilocode', + timestamp: '2026-04-14T08:00:00.000Z', + data: { event: 'session.status', type: 'session.status', properties: { type: 'idle' } }, + }; + const frame = prepareIngestFrame(event); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.compacted).toBe(false); + expect(frame.bytes).toBeLessThanOrEqual(MAX_INGEST_EVENT_BYTES); + }); + + it('applies payload trimming to file parts before serialization', () => { + const rawDataUrl = 'data:image/png;base64,wrapper-private-image'; + const rawSourceText = 'wrapper private source text'; + + const frame = prepareIngestFrame({ + streamEventType: 'kilocode', + data: { + event: 'message.part.updated', + type: 'message.part.updated', + part: { + type: 'file', + url: rawDataUrl, + source: { text: { value: rawSourceText } }, + }, + }, + timestamp: '2026-04-14T08:00:00.000Z', + }); + + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.serialized).not.toContain(rawDataUrl); + expect(frame.serialized).not.toContain(rawSourceText); + }); + + it('truncates oversized output before send via the single serialization path', () => { + const event: IngestEvent = { + streamEventType: 'output', + timestamp: '2026-04-14T08:00:00.000Z', + data: { content: 'x'.repeat(2_000_000), source: 'stdout' }, + }; + const frame = prepareIngestFrame(event); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.bytes).toBeLessThanOrEqual(MAX_INGEST_EVENT_BYTES); + const sent = JSON.parse(frame.serialized); + expect((sent.data.content as string).length).toBeLessThan(20_000); + }); + + it('compacts oversized message.part.updated tool output under the byte budget', () => { + const event: IngestEvent = { + streamEventType: 'kilocode', + timestamp: '2026-04-14T08:00:00.000Z', + data: { + event: 'message.part.updated', + type: 'message.part.updated', + properties: { + sessionID: 'sess_1', + part: { + type: 'tool', + id: 'part_1', + messageID: 'msg_1', + state: { + status: 'completed', + output: 'x'.repeat(50_000), + customDiagnostics: 'y'.repeat(2_000_000), + }, + }, + }, + part: { + type: 'tool', + id: 'part_1', + messageID: 'msg_1', + state: { status: 'completed', output: 'x'.repeat(50_000) }, + }, + }, + }; + const frame = prepareIngestFrame(event); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.compacted).toBe(true); + expect(frame.bytes).toBeLessThanOrEqual(MAX_INGEST_EVENT_BYTES); + expect(frame.serialized).not.toContain('yyyyy'); + const sent = JSON.parse(frame.serialized); + expect(sent.streamEventType).toBe('kilocode'); + const part = sent.data.properties.part; + expect(part.id).toBe('part_1'); + expect(part.messageID).toBe('msg_1'); + expect(part.state.status).toBe('completed'); + }); + + it('reduces oversized message.updated to a metadata-only event', () => { + const event: IngestEvent = { + streamEventType: 'kilocode', + timestamp: '2026-04-14T08:00:00.000Z', + data: { + event: 'message.updated', + type: 'message.updated', + properties: { + info: { + id: 'msg_1', + parentID: 'msg_0', + sessionID: 'sess_1', + role: 'assistant', + time: { started: 1, completed: 2 }, + parts: [{ type: 'text', text: 'z'.repeat(2_000_000) }], + }, + }, + }, + }; + const frame = prepareIngestFrame(event); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.compacted).toBe(true); + expect(frame.bytes).toBeLessThanOrEqual(MAX_INGEST_EVENT_BYTES); + expect(frame.serialized).not.toContain('zzzzz'); + const sent = JSON.parse(frame.serialized); + const info = sent.data.properties.info; + expect(info.id).toBe('msg_1'); + expect(info.parentID).toBe('msg_0'); + expect(info.role).toBe('assistant'); + expect(info.parts).toBeUndefined(); + }); + + it('still sends a compact fatal event for a terminal error with huge diagnostics', () => { + const event: IngestEvent = { + streamEventType: 'error', + timestamp: '2026-04-14T08:00:00.000Z', + data: { + fatal: true, + errorSource: 'assistant', + error: 'Provider failed: ' + 'z'.repeat(2_000_000), + failureCode: 'payment_required', + modelNotFoundRuntimeDiagnostics: { availableModels: 'q'.repeat(2_000_000).split('') }, + }, + }; + const frame = prepareIngestFrame(event); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.compacted).toBe(true); + expect(frame.bytes).toBeLessThanOrEqual(MAX_INGEST_EVENT_BYTES); + // Heavy diagnostics are dropped; a safe truncated error text is preserved. + expect(frame.serialized).not.toContain('qqqqq'); + const sent = JSON.parse(frame.serialized); + expect(sent.streamEventType).toBe('error'); + expect(sent.data.fatal).toBe(true); + expect(sent.data.errorSource).toBe('assistant'); + expect(sent.data.failureCode).toBe('payment_required'); + expect((sent.data.error as string).length).toBeLessThan(20_000); + expect(sent.data.modelNotFoundRuntimeDiagnostics).toBeUndefined(); + }); + + it('replaces a generic oversized non-terminal event with a wrapper_event_truncated surrogate', () => { + const event: IngestEvent = { + streamEventType: 'status', + timestamp: '2026-04-14T08:00:00.000Z', + data: { message: 'z'.repeat(2_000_000) }, + }; + const frame = prepareIngestFrame(event); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.compacted).toBe(true); + expect(frame.bytes).toBeLessThanOrEqual(MAX_INGEST_EVENT_BYTES); + const sent = JSON.parse(frame.serialized); + expect(sent.streamEventType).toBe('wrapper_event_truncated'); + expect(sent.data.originalStreamEventType).toBe('status'); + expect(sent.data.reason).toBe('oversized_ingest_event'); + expect(sent.data.originalBytes).toBeGreaterThan(MAX_INGEST_EVENT_BYTES); + }); + + it('degrades an unserializable non-terminal event to a surrogate instead of throwing', () => { + const circular: Record = { message: 'boom' }; + circular.self = circular; + const frame = prepareIngestFrame({ + streamEventType: 'status', + timestamp: '2026-04-14T08:00:00.000Z', + data: circular, + }); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.compacted).toBe(true); + const sent = JSON.parse(frame.serialized); + expect(sent.streamEventType).toBe('wrapper_event_truncated'); + expect(sent.data.originalStreamEventType).toBe('status'); + }); + + it('still sends a compact terminal error when its data is unserializable', () => { + const circular: Record = { fatal: true, error: 'boom' }; + circular.self = circular; + const frame = prepareIngestFrame({ + streamEventType: 'error', + timestamp: '2026-04-14T08:00:00.000Z', + data: circular, + }); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.compacted).toBe(true); + const sent = JSON.parse(frame.serialized); + expect(sent.streamEventType).toBe('error'); + expect(sent.data.fatal).toBe(true); + expect(sent.data.error).toBe('boom'); + }); + + it('falls back to the minimal form when even the compacted event is unserializable', () => { + const circular: Record = {}; + circular.self = circular; + // compactMessagePartUpdated copies `type` verbatim, so the compacted event + // is still circular and only the minimal surrogate can be serialized. + const frame = prepareIngestFrame({ + streamEventType: 'kilocode', + timestamp: '2026-04-14T08:00:00.000Z', + data: { event: 'message.part.updated', type: circular }, + }); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.compacted).toBe(true); + const sent = JSON.parse(frame.serialized); + expect(sent.streamEventType).toBe('wrapper_event_truncated'); + expect(sent.data.kiloEventName).toBe('message.part.updated'); + }); +}); + +describe('IngestEventBuffer', () => { + it('drops non-lifecycle events under byte pressure and marks eventsLost', () => { + const buffer = new IngestEventBuffer(); + const big = 2_000_000; + expect(buffer.push({ serialized: 'a', bytes: big, lifecycle: false })).toBe(true); + expect(buffer.push({ serialized: 'b', bytes: big, lifecycle: false })).toBe(true); + // Third non-lifecycle frame does not fit; oldest non-lifecycle is evicted. + expect(buffer.push({ serialized: 'c', bytes: big, lifecycle: false })).toBe(true); + + expect(buffer.isOverflowed).toBe(true); + expect(buffer.bytes).toBeLessThanOrEqual(MAX_INGEST_BUFFERED_BYTES); + expect(buffer.length).toBe(2); + const drained = buffer.drain(); + expect(drained.map(f => f.serialized)).toEqual(['b', 'c']); + expect(buffer.isOverflowed).toBe(false); + }); + + it('prioritizes lifecycle events by evicting non-lifecycle frames to make room', () => { + const buffer = new IngestEventBuffer(); + expect(buffer.push({ serialized: 'a', bytes: 2_000_000, lifecycle: false })).toBe(true); + expect(buffer.push({ serialized: 'b', bytes: 2_000_000, lifecycle: false })).toBe(true); + // Buffer is at the byte budget; lifecycle frame must still be accepted. + const accepted = buffer.push({ serialized: 'L', bytes: 1_000_000, lifecycle: true }); + expect(accepted).toBe(true); + expect(buffer.isOverflowed).toBe(true); + const drained = buffer.drain(); + expect(drained.some(f => f.serialized === 'L')).toBe(true); + expect(drained.some(f => f.serialized === 'a')).toBe(false); + }); + + it('drops a non-lifecycle frame instead of evicting lifecycle frames', () => { + const buffer = new IngestEventBuffer(); + expect(buffer.push({ serialized: 'L1', bytes: 2_000_000, lifecycle: true })).toBe(true); + expect(buffer.push({ serialized: 'L2', bytes: 2_000_000, lifecycle: true })).toBe(true); + // Buffer full of lifecycle frames: a non-lifecycle frame is dropped. + const accepted = buffer.push({ serialized: 'N', bytes: 1_000_000, lifecycle: false }); + expect(accepted).toBe(false); + expect(buffer.isOverflowed).toBe(true); + expect(buffer.length).toBe(2); + const drained = buffer.drain(); + expect(drained.map(f => f.serialized)).toEqual(['L1', 'L2']); + }); + + it('classifies terminal stream event types as lifecycle', () => { + for (const streamEventType of [ + 'complete', + 'interrupted', + 'error', + 'cloud.message.completed', + 'wrapper_finalizing', + 'autocommit_started', + 'autocommit_completed', + ] as const) { + expect( + isLifecycleIngestEvent({ + streamEventType, + timestamp: '2026-04-14T08:00:00.000Z', + data: {}, + }) + ).toBe(true); + } + expect( + isLifecycleIngestEvent({ + streamEventType: 'kilocode', + timestamp: '2026-04-14T08:00:00.000Z', + data: {}, + }) + ).toBe(false); + }); +}); diff --git a/services/cloud-agent-next/src/shared/ingest-frame.ts b/services/cloud-agent-next/src/shared/ingest-frame.ts new file mode 100644 index 0000000000..ad0a65c3b2 --- /dev/null +++ b/services/cloud-agent-next/src/shared/ingest-frame.ts @@ -0,0 +1,472 @@ +/** + * Ingest frame budgeting for the wrapper -> Durable Object `/ingest` WebSocket. + * + * Cloudflare closes a WebSocket with `1009` when a single received message + * exceeds its per-message limit, and Durable Object SQLite has a much smaller + * practical limit for persisted row/string/blob values. Both failures are + * avoided by never sending an unbounded serialized event. + * + * Every ingest send goes through `prepareIngestFrame`, which: + * 1. applies the existing payload trimming, + * 2. serializes once and measures UTF-8 byte length, + * 3. sends when under the safe budget, + * 4. compacts (or replaces with a small internal surrogate) when oversized, + * 5. drops only as a last resort. + * + * Terminal/lifecycle signals are never silently dropped: oversized terminal + * events are reduced to a compact form that preserves message IDs, terminal + * status, failure code, and safe error text. + */ + +import { trimPayload } from './trim-payload.js'; +import type { IngestEvent, StreamEventType, WrapperEventTruncatedData } from './protocol.js'; + +/** Cloudflare's per-WebSocket-message receive limit. Documentation/logging only. */ +export const CLOUDFLARE_WEBSOCKET_RECEIVE_LIMIT_BYTES = 32 * 1024 * 1024; + +/** + * Enforced per-frame send budget. Kept well below the platform 32 MiB limit + * because the DO may persist and replay the payload, and SQLite row/string/blob + * limits are much closer to ~2 MB. + */ +export const MAX_INGEST_EVENT_BYTES = 1 * 1024 * 1024; + +/** Disconnected-buffer memory budget (count cap is secondary to this). */ +export const MAX_INGEST_BUFFERED_BYTES = 4 * 1024 * 1024; + +/** Secondary disconnected-buffer count cap. */ +export const MAX_INGEST_BUFFER_COUNT = 1000; + +const encoder = new TextEncoder(); + +function byteLength(serialized: string): number { + return encoder.encode(serialized).byteLength; +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function safeString(value: unknown, max: number): string | undefined { + if (typeof value !== 'string') return undefined; + if (value.length <= max) return value; + return value.slice(0, max) + '\n\n[…truncated]'; +} + +/** + * Stream event types that carry terminal/lifecycle meaning and must survive + * compaction rather than being replaced by a truncation surrogate. + */ +export const TERMINAL_INGEST_STREAM_EVENT_TYPES: ReadonlySet = new Set([ + 'complete', + 'interrupted', + 'error', + 'cloud.message.completed', + 'wrapper_finalizing', + 'autocommit_started', + 'autocommit_completed', +]); + +export function isLifecycleIngestEvent(event: IngestEvent): boolean { + return TERMINAL_INGEST_STREAM_EVENT_TYPES.has(event.streamEventType); +} + +export function kiloEventNameOf(data: unknown): string | undefined { + if (!isRecord(data)) return undefined; + const name = data.event ?? data.type; + return typeof name === 'string' ? name : undefined; +} + +function compactTerminalData(streamEventType: StreamEventType, data: unknown): unknown { + if (!isRecord(data)) return {}; + switch (streamEventType) { + case 'complete': { + const out: Record = { + exitCode: typeof data.exitCode === 'number' ? data.exitCode : 0, + }; + if (typeof data.currentBranch === 'string') out.currentBranch = data.currentBranch; + if (data.gateResult === 'pass' || data.gateResult === 'fail') + out.gateResult = data.gateResult; + if (Array.isArray(data.messageIds)) { + out.messageIds = data.messageIds.filter((id): id is string => typeof id === 'string'); + } + return out; + } + case 'interrupted': { + const out: Record = {}; + const reason = safeString(data.reason, 1000); + if (reason !== undefined) out.reason = reason; + if (typeof data.exitCode === 'number') out.exitCode = data.exitCode; + if (data.interruptionSource === 'container_shutdown') + out.interruptionSource = data.interruptionSource; + return out; + } + case 'error': { + const safeError = + safeString(data.error, 1000) ?? safeString(data.message, 1000) ?? 'Agent wrapper failed'; + const out: Record = { + error: safeError, + message: safeError, + }; + if (typeof data.fatal === 'boolean') out.fatal = data.fatal; + if (data.errorSource === 'assistant') out.errorSource = data.errorSource; + if (typeof data.failureCode === 'string') out.failureCode = data.failureCode; + // Drop heavy modelNotFoundRuntimeDiagnostics on compaction. + return out; + } + case 'cloud.message.completed': { + const out: Record = { + completionSource: 'manual_compact_summarize', + }; + if (typeof data.messageId === 'string') out.messageId = data.messageId; + if (typeof data.assistantMessageId === 'string') + out.assistantMessageId = data.assistantMessageId; + return out; + } + case 'autocommit_started': { + const out: Record = { message: safeString(data.message, 1000) ?? '' }; + if (typeof data.messageId === 'string') out.messageId = data.messageId; + return out; + } + case 'autocommit_completed': { + const out: Record = { + success: typeof data.success === 'boolean' ? data.success : false, + message: safeString(data.message, 1000) ?? '', + }; + if (typeof data.messageId === 'string') out.messageId = data.messageId; + if (typeof data.skipped === 'boolean') out.skipped = data.skipped; + if (typeof data.commitHash === 'string') out.commitHash = data.commitHash; + // Drop commitMessage (may be large) on compaction. + return out; + } + case 'wrapper_finalizing': { + const out: Record = {}; + if (typeof data.wrapperRunId === 'string') out.wrapperRunId = data.wrapperRunId; + return out; + } + default: + return {}; + } +} + +function compactMessagePartUpdated(data: Record): Record { + const out: Record = { event: data.event, type: data.type }; + const properties = data.properties; + if (isRecord(properties)) { + const compactProps: Record = {}; + if (typeof properties.sessionID === 'string') compactProps.sessionID = properties.sessionID; + const part = properties.part; + if (isRecord(part)) { + const compactPart: Record = {}; + if (typeof part.type === 'string') compactPart.type = part.type; + if (typeof part.id === 'string') compactPart.id = part.id; + if (typeof part.messageID === 'string') compactPart.messageID = part.messageID; + if (typeof part.status === 'string') compactPart.status = part.status; + const state = part.state; + if (isRecord(state) && typeof state.status === 'string') { + compactPart.state = { status: state.status }; + } + compactProps.part = compactPart; + } + out.properties = compactProps; + } + return out; +} + +function compactMessageUpdated(data: Record): Record { + const out: Record = { event: data.event, type: data.type }; + const properties = data.properties; + if (isRecord(properties)) { + const info = properties.info; + if (isRecord(info)) { + const compactInfo: Record = {}; + if (typeof info.id === 'string') compactInfo.id = info.id; + if (typeof info.parentID === 'string') compactInfo.parentID = info.parentID; + if (typeof info.sessionID === 'string') compactInfo.sessionID = info.sessionID; + if (typeof info.role === 'string') compactInfo.role = info.role; + if (isRecord(info.time)) compactInfo.time = info.time; + const safeError = safeString(info.error, 2000); + if (safeError !== undefined) compactInfo.error = safeError; + out.properties = { info: compactInfo }; + } + } + return out; +} + +function compactCommandsAvailable(data: Record): Record { + const commands = data.commands; + if (!Array.isArray(commands)) return { commands: [] }; + return { + commands: commands + .filter((c): c is Record => isRecord(c) && typeof c.name === 'string') + .map(c => ({ name: c.name as string })), + }; +} + +function buildTruncationSurrogate( + event: IngestEvent, + originalBytes: number, + compactedBytes?: number +): IngestEvent { + const data: WrapperEventTruncatedData = { + originalStreamEventType: event.streamEventType, + originalBytes, + reason: 'oversized_ingest_event', + }; + const kiloEventName = + event.streamEventType === 'kilocode' ? kiloEventNameOf(event.data) : undefined; + if (kiloEventName !== undefined) data.kiloEventName = kiloEventName; + if (compactedBytes !== undefined) data.compactedBytes = compactedBytes; + return { + streamEventType: 'wrapper_event_truncated', + timestamp: event.timestamp, + data, + }; +} + +function compactIngestEvent(event: IngestEvent, originalBytes: number): IngestEvent { + const { streamEventType, timestamp, data } = event; + if (TERMINAL_INGEST_STREAM_EVENT_TYPES.has(streamEventType)) { + return { streamEventType, timestamp, data: compactTerminalData(streamEventType, data) }; + } + if (streamEventType === 'kilocode') { + const kiloEventName = kiloEventNameOf(data); + if (kiloEventName === 'message.part.updated' && isRecord(data)) { + return { streamEventType, timestamp, data: compactMessagePartUpdated(data) }; + } + if (kiloEventName === 'message.updated' && isRecord(data)) { + return { streamEventType, timestamp, data: compactMessageUpdated(data) }; + } + } + if (streamEventType === 'commands.available' && isRecord(data)) { + return { streamEventType, timestamp, data: compactCommandsAvailable(data) }; + } + return buildTruncationSurrogate(event, originalBytes); +} + +function minimalIngestEvent(event: IngestEvent, originalBytes: number): IngestEvent { + const { streamEventType, timestamp, data } = event; + if (TERMINAL_INGEST_STREAM_EVENT_TYPES.has(streamEventType)) { + let minimalData: unknown = {}; + switch (streamEventType) { + case 'complete': + minimalData = { + exitCode: + typeof data === 'object' && + data && + 'exitCode' in data && + typeof (data as Record).exitCode === 'number' + ? (data as Record).exitCode + : 0, + }; + break; + case 'interrupted': + minimalData = { reason: 'Wrapper interrupted' }; + break; + case 'error': { + const d = data as Record; + minimalData = { + fatal: typeof d?.fatal === 'boolean' ? d.fatal : true, + error: safeString(d?.error, 1000) ?? 'Agent wrapper failed', + }; + break; + } + case 'cloud.message.completed': { + const d = data as Record; + minimalData = + typeof d?.messageId === 'string' + ? { messageId: d.messageId, completionSource: 'manual_compact_summarize' } + : { completionSource: 'manual_compact_summarize' }; + break; + } + case 'autocommit_started': + minimalData = { message: '' }; + break; + case 'autocommit_completed': + minimalData = { success: false, message: '' }; + break; + case 'wrapper_finalizing': + minimalData = {}; + break; + } + return { streamEventType, timestamp, data: minimalData }; + } + return buildTruncationSurrogate(event, originalBytes); +} + +export type PreparedIngestFrame = + | { + kind: 'send'; + serialized: string; + bytes: number; + compacted: boolean; + originalBytes: number; + } + | { kind: 'dropped'; originalBytes: number; reason: string }; + +/** + * Serialization must never crash the send path: unserializable data (e.g. a + * circular reference) degrades to the same compaction/drop ladder as an + * oversized event. + */ +function tryStringify(event: IngestEvent): string | undefined { + try { + return JSON.stringify(event); + } catch { + return undefined; + } +} + +/** + * Trim, serialize, measure, and (if needed) compact an ingest event so its + * serialized UTF-8 byte length stays under `MAX_INGEST_EVENT_BYTES`. + */ +export function prepareIngestFrame(event: IngestEvent): PreparedIngestFrame { + const trimmed: IngestEvent = { + ...event, + data: trimPayload(event.streamEventType, event.data), + }; + const originalSerialized = tryStringify(trimmed); + const originalBytes = originalSerialized === undefined ? 0 : byteLength(originalSerialized); + if (originalSerialized !== undefined && originalBytes <= MAX_INGEST_EVENT_BYTES) { + return { + kind: 'send', + serialized: originalSerialized, + bytes: originalBytes, + compacted: false, + originalBytes, + }; + } + + const compactedEvent = compactIngestEvent(trimmed, originalBytes); + const compactedSerialized = tryStringify(compactedEvent); + if (compactedSerialized !== undefined) { + const compactedBytes = byteLength(compactedSerialized); + if (compactedBytes <= MAX_INGEST_EVENT_BYTES) { + return { + kind: 'send', + serialized: compactedSerialized, + bytes: compactedBytes, + compacted: true, + originalBytes, + }; + } + } + + const minimalEvent = minimalIngestEvent(trimmed, originalBytes); + const minimalSerialized = tryStringify(minimalEvent); + if (minimalSerialized !== undefined) { + const minimalBytes = byteLength(minimalSerialized); + if (minimalBytes <= MAX_INGEST_EVENT_BYTES) { + return { + kind: 'send', + serialized: minimalSerialized, + bytes: minimalBytes, + compacted: true, + originalBytes, + }; + } + } + + return { + kind: 'dropped', + originalBytes, + reason: + originalSerialized === undefined ? 'unserializable_ingest_event' : 'oversized_ingest_event', + }; +} + +export type BufferedIngestFrame = { + serialized: string; + bytes: number; + lifecycle: boolean; +}; + +/** + * Disconnected-ingest buffer with a byte budget and a secondary count cap. + * + * Lifecycle/terminal frames are protected: when buffer pressure would force an + * eviction, older non-lifecycle frames are dropped first. `isOverflowed` tracks + * whether any frame was lost so the `wrapper_resumed` marker can report + * `eventsLost`. + */ +export class IngestEventBuffer { + private frames: BufferedIngestFrame[] = []; + private totalBytes = 0; + private overflowed = false; + + get length(): number { + return this.frames.length; + } + + get bytes(): number { + return this.totalBytes; + } + + get isOverflowed(): boolean { + return this.overflowed; + } + + /** + * Push a prepared frame. Returns `true` when the frame was buffered. + * Non-lifecycle frames that cannot fit (even after evicting older + * non-lifecycle frames) are dropped and `false` is returned. + */ + push(frame: BufferedIngestFrame): boolean { + // Evict older non-lifecycle frames first to protect lifecycle signals. + while ( + (this.totalBytes + frame.bytes > MAX_INGEST_BUFFERED_BYTES || + this.frames.length >= MAX_INGEST_BUFFER_COUNT) && + this.frames.some(f => !f.lifecycle) + ) { + const idx = this.frames.findIndex(f => !f.lifecycle); + this.totalBytes -= this.frames[idx].bytes; + this.frames.splice(idx, 1); + this.overflowed = true; + } + + if ( + this.totalBytes + frame.bytes <= MAX_INGEST_BUFFERED_BYTES && + this.frames.length < MAX_INGEST_BUFFER_COUNT + ) { + this.frames.push(frame); + this.totalBytes += frame.bytes; + return true; + } + + if (!frame.lifecycle) { + this.overflowed = true; + return false; + } + + // Buffer is saturated with lifecycle frames. Evict the oldest frames until + // it fits so growth stays bounded; this loses an older terminal signal, + // which is unavoidable when the buffer is saturated with lifecycle events. + while ( + this.frames.length > 0 && + (this.totalBytes + frame.bytes > MAX_INGEST_BUFFERED_BYTES || + this.frames.length >= MAX_INGEST_BUFFER_COUNT) + ) { + const oldest = this.frames.shift(); + if (oldest) this.totalBytes -= oldest.bytes; + this.overflowed = true; + } + this.frames.push(frame); + this.totalBytes += frame.bytes; + return true; + } + + drain(): BufferedIngestFrame[] { + const out = this.frames; + this.frames = []; + this.totalBytes = 0; + this.overflowed = false; + return out; + } + + clear(): void { + this.frames = []; + this.totalBytes = 0; + this.overflowed = false; + } +} diff --git a/services/cloud-agent-next/src/shared/protocol.ts b/services/cloud-agent-next/src/shared/protocol.ts index f8c08ec935..0d28d29aae 100644 --- a/services/cloud-agent-next/src/shared/protocol.ts +++ b/services/cloud-agent-next/src/shared/protocol.ts @@ -26,6 +26,8 @@ export type StreamEventType = | 'wrapper_resumed' // Wrapper reconnected after disconnect (may have lost events) | 'autocommit_started' // Auto-commit process began | 'autocommit_completed' // Auto-commit finished (success, skip, or failure) + // Wrapper -> DO (internal size-safety signal, not persisted or broadcast) + | 'wrapper_event_truncated' // A wrapper ingest event was compacted/replaced to fit the frame budget // DO -> /stream clients (connection status) | 'wrapper_disconnected' // Wrapper WebSocket closed unexpectedly | 'wrapper_reconnected' // Wrapper reconnected successfully @@ -101,6 +103,22 @@ export type AutocommitCompletedData = { commitMessage?: string; }; +/** + * Data included in internal 'wrapper_event_truncated' events. + * + * Emitted by the wrapper when an ingest event had to be compacted or replaced + * to stay under the per-frame byte budget. The DO treats this as a + * liveness-relevant signal but does not persist or broadcast it to /stream + * clients; it is an internal size-safety diagnostic only. + */ +export type WrapperEventTruncatedData = { + originalStreamEventType: StreamEventType; + kiloEventName?: string; + originalBytes: number; + compactedBytes?: number; + reason: string; +}; + /** * Preparation step identifiers for async preparation progress events. */ diff --git a/services/cloud-agent-next/src/websocket/ingest.ts b/services/cloud-agent-next/src/websocket/ingest.ts index c26a40b8c2..ab4b564b40 100644 --- a/services/cloud-agent-next/src/websocket/ingest.ts +++ b/services/cloud-agent-next/src/websocket/ingest.ts @@ -81,6 +81,14 @@ const cloudMessageCompletedEventSchema = z.object({ completionSource: z.literal('manual_compact_summarize'), }); +const wrapperEventTruncatedSchema = z.object({ + originalStreamEventType: z.string(), + kiloEventName: z.string().optional(), + originalBytes: z.number(), + compactedBytes: z.number().optional(), + reason: z.string(), +}); + const wrapperGenerationParamSchema = z.coerce.number().int().nonnegative(); function getAssistantErrorMessage(error: unknown): string | undefined { @@ -454,10 +462,10 @@ export function createIngestHandler( state.acceptWebSocket(server, [ingestTag]); server.serializeAttachment(attachment); - await doContext.wrapperSupervisor.recordReconnectAccepted({ - wrapperGeneration, - wrapperConnectionId, - }); + await doContext.wrapperSupervisor.recordReconnectAccepted( + { wrapperGeneration, wrapperConnectionId }, + now + ); doContext.keepContainerAlive?.(); logger @@ -530,11 +538,46 @@ export function createIngestHandler( : Date.now(); const eventType = ingestEvent.streamEventType; + const now = Date.now(); + + // Internal size-safety signal: the wrapper compacted/replaced an event + // to fit the frame budget. Count it as liveness-relevant (the wrapper + // is actively sending) but do not persist or broadcast it to /stream + // clients; it is an internal diagnostic only. + if (eventType === 'wrapper_event_truncated') { + const parsedTruncation = wrapperEventTruncatedSchema.safeParse(ingestEvent.data); + logger + .withFields({ + sessionId, + wrapperRunId, + wrapperGeneration, + wrapperConnectionId, + ...(parsedTruncation.success + ? { + originalStreamEventType: parsedTruncation.data.originalStreamEventType, + kiloEventName: parsedTruncation.data.kiloEventName, + originalBytes: parsedTruncation.data.originalBytes, + compactedBytes: parsedTruncation.data.compactedBytes, + reason: parsedTruncation.data.reason, + } + : { parseError: parsedTruncation.error.message }), + logTag: 'wrapper_event_truncated', + }) + .warn('Wrapper ingest event was truncated to fit the frame budget'); + if (wrapperGeneration !== undefined && wrapperConnectionId) { + await doContext.wrapperSupervisor.observeMeaningfulOutput( + wrapperGeneration, + wrapperConnectionId, + now + ); + } + return; + } + const publicEventData = sanitizePublicEventData(eventType, ingestEvent.data ?? {}); const payload = JSON.stringify(publicEventData); const eventTypeStr: string = eventType; - const now = Date.now(); const kiloEventName = eventType === 'kilocode' ? ((ingestEvent.data as Record | undefined)?.event as diff --git a/services/cloud-agent-next/src/websocket/types.ts b/services/cloud-agent-next/src/websocket/types.ts index 86f87dadbc..0201e0054d 100644 --- a/services/cloud-agent-next/src/websocket/types.ts +++ b/services/cloud-agent-next/src/websocket/types.ts @@ -43,7 +43,8 @@ export type StreamEventType = | 'cloud.message.queued' // user message accepted into pending queue | 'cloud.message.sent' // queued user message delivered to Kilo | 'cloud.message.completed' // accepted user message completed execution - | 'cloud.message.failed'; // user message delivery failed or was canceled + | 'cloud.message.failed' // user message delivery failed or was canceled + | 'wrapper_event_truncated'; // internal size-safety signal (not persisted or broadcast) // --------------------------------------------------------------------------- // Server -> Client Events (/stream endpoint) diff --git a/services/cloud-agent-next/test/unit/wrapper/connection.test.ts b/services/cloud-agent-next/test/unit/wrapper/connection.test.ts index cc927e757c..4904f90ca9 100644 --- a/services/cloud-agent-next/test/unit/wrapper/connection.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/connection.test.ts @@ -1,14 +1,13 @@ /** * Unit tests for connection module. * - * Tests connection diagnostics, event trimming, and session.idle filtering logic. + * Tests connection diagnostics and session.idle filtering logic. */ import { describe, expect, it } from 'vitest'; import { buildIngestConnectionFailureMessage, isSessionIdleEvent, - trimIngestEvent, } from '../../../wrapper/src/connection.js'; // --------------------------------------------------------------------------- @@ -55,36 +54,6 @@ describe('buildIngestConnectionFailureMessage', () => { // isSessionIdleEvent // --------------------------------------------------------------------------- -describe('trimIngestEvent', () => { - it('trims top-level file parts before ingest serialization', () => { - const rawDataUrl = 'data:image/png;base64,wrapper-private-image'; - const rawSourceText = 'wrapper private source text'; - - const event = trimIngestEvent({ - streamEventType: 'kilocode', - data: { - event: 'message.part.updated', - type: 'message.part.updated', - part: { - type: 'file', - url: rawDataUrl, - source: { text: { value: rawSourceText } }, - }, - }, - timestamp: '2026-04-14T08:00:00.000Z', - }); - - const payload = event.data as { - part: { url: string; source: { text: { value: string } } }; - }; - - expect(payload.part.url).toBe(''); - expect(payload.part.source.text.value).toBe(''); - expect(JSON.stringify(event)).not.toContain(rawDataUrl); - expect(JSON.stringify(event)).not.toContain(rawSourceText); - }); -}); - describe('isSessionIdleEvent', () => { it('returns true for a valid session.idle event with sessionID', () => { const data = { diff --git a/services/cloud-agent-next/wrapper/src/connection.ts b/services/cloud-agent-next/wrapper/src/connection.ts index 4f110d1768..814b13985f 100644 --- a/services/cloud-agent-next/wrapper/src/connection.ts +++ b/services/cloud-agent-next/wrapper/src/connection.ts @@ -15,7 +15,15 @@ import type { WrapperCommand, WrapperTerminalFailureCode, } from '../../src/shared/protocol.js'; -import { trimPayload } from '../../src/shared/trim-payload.js'; +import { + MAX_INGEST_EVENT_BYTES, + MAX_INGEST_BUFFERED_BYTES, + IngestEventBuffer, + isLifecycleIngestEvent, + kiloEventNameOf, + prepareIngestFrame, + type PreparedIngestFrame, +} from '../../src/shared/ingest-frame.js'; import { logToFile } from './utils.js'; import type { KiloEvent, WrapperKiloClient } from './kilo-api.js'; import type { ModelNotFoundRuntimeDiagnostics } from '../../src/shared/runtime-model-diagnostics.js'; @@ -144,13 +152,67 @@ function rejectCodeReviewPermission( }); } -export function trimIngestEvent(event: IngestEvent): IngestEvent { +function kiloEventNameForLog(event: IngestEvent): string | undefined { + return event.streamEventType === 'kilocode' ? kiloEventNameOf(event.data) : undefined; +} + +function sessionLogFields(state: WrapperState): Record { + const session = state.currentSession; + if (!session) return {}; return { - ...event, - data: trimPayload(event.streamEventType, event.data), + agentSessionId: session.agentSessionId, + kiloSessionId: session.kiloSessionId, + wrapperRunId: session.wrapperRunId, + ...(session.wrapperGeneration !== undefined + ? { wrapperGeneration: session.wrapperGeneration } + : {}), + ...(session.wrapperConnectionId ? { wrapperConnectionId: session.wrapperConnectionId } : {}), }; } +function logIngestFrameMessage( + state: WrapperState, + message: string, + fields: Record +): void { + logToFile(JSON.stringify({ message, ...sessionLogFields(state), ...fields })); +} + +/** + * Prepare an ingest event for sending: trim, serialize, measure, and compact + * when oversized. Emits the structured size-safety logs. Returns the prepared + * frame (or a `dropped` frame that callers must not send). + */ +function prepareFrameForSend(event: IngestEvent, state: WrapperState): PreparedIngestFrame { + const frame = prepareIngestFrame(event); + if (frame.originalBytes > MAX_INGEST_EVENT_BYTES) { + logIngestFrameMessage(state, 'ingest_event_oversized_before_send', { + streamEventType: event.streamEventType, + kiloEventName: kiloEventNameForLog(event), + originalBytes: frame.originalBytes, + limitBytes: MAX_INGEST_EVENT_BYTES, + }); + } + if (frame.kind === 'dropped') { + logIngestFrameMessage(state, 'ingest_event_dropped', { + streamEventType: event.streamEventType, + kiloEventName: kiloEventNameForLog(event), + originalBytes: frame.originalBytes, + reason: frame.reason, + }); + return frame; + } + if (frame.compacted) { + logIngestFrameMessage(state, 'ingest_event_compacted', { + streamEventType: event.streamEventType, + kiloEventName: kiloEventNameForLog(event), + originalBytes: frame.originalBytes, + compactedBytes: frame.bytes, + }); + } + return frame; +} + /** * Type guard for session.idle events. * Kilo server sends: {type: "session.idle", properties: {sessionID: "..."}} @@ -311,8 +373,10 @@ export async function openIngestProgressChannel( }); const sendProgressEvent = (event: IngestEvent): void => { - if (active && ws.readyState === WebSocket.OPEN) { - ws.send(JSON.stringify(event)); + if (!active || ws.readyState !== WebSocket.OPEN) return; + const frame = prepareFrameForSend(event, state); + if (frame.kind === 'send') { + ws.send(frame.serialized); } }; @@ -425,55 +489,68 @@ export function createConnectionManager( let reconnectTimer: ReturnType | null = null; let generation = 0; - // Event buffer for disconnection periods - const MAX_BUFFER_SIZE = 1000; - const eventBuffer: IngestEvent[] = []; - let bufferOverflowed = false; + // Event buffer for disconnection periods. Byte budget is primary; the count + // cap is secondary. Lifecycle/terminal frames are protected from eviction. + const eventBuffer = new IngestEventBuffer(); /** * Send an event to the ingest WebSocket. - * Buffers events if disconnected. + * Buffers the prepared frame if disconnected. */ function sendToIngest(event: IngestEvent): void { + const frame = prepareFrameForSend(event, state); + if (frame.kind === 'dropped') return; + if (ingestWs && ingestWs.readyState === WebSocket.OPEN) { - ingestWs.send(JSON.stringify(event)); - } else { - // Buffer events while disconnected - if (eventBuffer.length < MAX_BUFFER_SIZE) { - eventBuffer.push(event); - } else { - bufferOverflowed = true; - } + ingestWs.send(frame.serialized); + return; + } + + const accepted = eventBuffer.push({ + serialized: frame.serialized, + bytes: frame.bytes, + lifecycle: isLifecycleIngestEvent(event), + }); + if (!accepted) { + logIngestFrameMessage(state, 'ingest_buffer_event_dropped', { + streamEventType: event.streamEventType, + kiloEventName: kiloEventNameForLog(event), + bytes: frame.bytes, + bufferBytes: eventBuffer.bytes, + bufferCount: eventBuffer.length, + limitBytes: MAX_INGEST_BUFFERED_BYTES, + }); } } /** - * Flush buffered events after reconnection. + * Clear buffered events (e.g. on close). */ function clearBuffer(): void { - eventBuffer.length = 0; - bufferOverflowed = false; + eventBuffer.clear(); } function flushBuffer(): void { if (!ingestWs || ingestWs.readyState !== WebSocket.OPEN) return; // Send resume marker so DO knows we may have lost events - if (eventBuffer.length > 0 || bufferOverflowed) { + if (eventBuffer.length > 0 || eventBuffer.isOverflowed) { ingestWs.send( JSON.stringify({ streamEventType: 'wrapper_resumed', timestamp: new Date().toISOString(), - data: { bufferedEvents: eventBuffer.length, eventsLost: bufferOverflowed }, + data: { + bufferedEvents: eventBuffer.length, + eventsLost: eventBuffer.isOverflowed, + }, }) ); } - // Flush buffer - for (const event of eventBuffer) { - ingestWs.send(JSON.stringify(event)); + // Flush buffered pre-serialized frames + for (const buffered of eventBuffer.drain()) { + ingestWs.send(buffered.serialized); } - clearBuffer(); } async function resumeNetworkWait(requestID: string): Promise { @@ -1111,14 +1188,15 @@ export function createConnectionManager( maybeResumeNetworkWait(eventType, properties); - // Build and forward ingest event - const untrimmedIngestEvent: IngestEvent = { + // Build and forward ingest event. Payload trimming and per-frame + // byte budgeting are applied inside sendToIngest/prepareIngestFrame, + // so this is the single serialization path for kilocode events. + const ingestEvent: IngestEvent = { streamEventType: 'kilocode', data: { ...properties, event: eventType, type: eventType, properties }, timestamp: new Date().toISOString(), }; - const ingestEvent = trimIngestEvent(untrimmedIngestEvent); sendToIngest(ingestEvent); callbacks.onSseEvent?.(); diff --git a/services/cloud-agent-next/wrapper/src/global-feed.test.ts b/services/cloud-agent-next/wrapper/src/global-feed.test.ts index 9d84adc752..41a00963c4 100644 --- a/services/cloud-agent-next/wrapper/src/global-feed.test.ts +++ b/services/cloud-agent-next/wrapper/src/global-feed.test.ts @@ -32,11 +32,13 @@ function asFetch( class FakeWebSocket { static instances: FakeWebSocket[] = []; static initialBufferedAmount = 0; + /** Per-instance sequence returned by `bufferedAmount`, consumed left-to-right. */ + static bufferedAmountSequence: number[] | undefined; readonly url: string; readonly options?: { headers?: Record } | string | string[]; readyState: number = WebSocket.CONNECTING; - bufferedAmount = FakeWebSocket.initialBufferedAmount; + private readonly bufferedAmountQueue: number[]; sent: string[] = []; onopen: ((event: Event) => void) | null = null; onerror: ((event: Event) => void) | null = null; @@ -45,6 +47,9 @@ class FakeWebSocket { constructor(url: string, options?: { headers?: Record } | string | string[]) { this.url = url; this.options = options; + this.bufferedAmountQueue = FakeWebSocket.bufferedAmountSequence + ? [...FakeWebSocket.bufferedAmountSequence] + : []; FakeWebSocket.instances.push(this); queueMicrotask(() => { this.readyState = WebSocket.OPEN; @@ -52,6 +57,13 @@ class FakeWebSocket { }); } + get bufferedAmount(): number { + if (this.bufferedAmountQueue.length > 0) { + return this.bufferedAmountQueue.shift() as number; + } + return FakeWebSocket.initialBufferedAmount; + } + send(data: string): void { this.sent.push(data); } @@ -86,6 +98,7 @@ function makeKiloClient(serverUrl = 'http://127.0.0.1:4321'): WrapperKiloClient afterEach(() => { FakeWebSocket.instances = []; FakeWebSocket.initialBufferedAmount = 0; + FakeWebSocket.bufferedAmountSequence = undefined; }); describe('buildKiloGlobalFeedWebSocketUrl', () => { @@ -260,6 +273,120 @@ describe('openKiloGlobalFeed', () => { expect(FakeWebSocket.instances[0].sent).toEqual([]); }); + it('drops a single oversized event without restarting the feed loop', async () => { + const state = new WrapperState(); + bindGlobalFeedSession(state); + const fetchImpl = asFetch(async () => { + return new Response( + streamFromChunks([ + 'data: {"big":"' + 'x'.repeat(1_500_000) + '"}\n\n', + 'data: {"directory":"/workspace/root","payload":{"type":"message.updated","properties":{"id":"msg_after_oversized"}}}\n\n', + ]), + { status: 200 } + ); + }); + + const connection = openKiloGlobalFeed({ + state, + kiloClient: makeKiloClient(), + fetchImpl, + WebSocketImpl: FakeWebSocket as unknown as GlobalFeedWebSocketImpl, + retryDelayMs: 60_000, + }); + + await new Promise(resolve => setTimeout(resolve, 0)); + connection.close(); + await connection.done; + + // No reconnect (single WebSocket instance); the oversized event was + // dropped and the following normal event was still forwarded. + expect(FakeWebSocket.instances.length).toBe(1); + expect(FakeWebSocket.instances[0].sent).toEqual([ + JSON.stringify({ + directory: '/workspace/root', + payload: { type: 'message.updated', properties: { id: 'msg_after_oversized' } }, + }), + ]); + }); + + it('drops an event under backpressure then forwards later events once pressure clears', async () => { + const state = new WrapperState(); + bindGlobalFeedSession(state); + // First read is backed up; the second read has drained. + FakeWebSocket.bufferedAmountSequence = [Number.MAX_SAFE_INTEGER, 0]; + const fetchImpl = asFetch(async () => { + return new Response( + streamFromChunks([ + 'data: {"directory":"/workspace/root","payload":{"type":"message.updated","properties":{"id":"msg_dropped"}}}\n\n', + 'data: {"directory":"/workspace/root","payload":{"type":"message.updated","properties":{"id":"msg_sent"}}}\n\n', + ]), + { status: 200 } + ); + }); + + const connection = openKiloGlobalFeed({ + state, + kiloClient: makeKiloClient(), + fetchImpl, + WebSocketImpl: FakeWebSocket as unknown as GlobalFeedWebSocketImpl, + retryDelayMs: 60_000, + }); + + await new Promise(resolve => setTimeout(resolve, 0)); + connection.close(); + await connection.done; + + expect(FakeWebSocket.instances.length).toBe(1); + expect(FakeWebSocket.instances[0].sent).toEqual([ + JSON.stringify({ + directory: '/workspace/root', + payload: { type: 'message.updated', properties: { id: 'msg_sent' } }, + }), + ]); + }); + + it('fails the attempt and reconnects when backpressure never clears', async () => { + const state = new WrapperState(); + bindGlobalFeedSession(state); + FakeWebSocket.initialBufferedAmount = Number.MAX_SAFE_INTEGER; + let fetchCount = 0; + let secondFetchStarted: (() => void) | undefined; + const secondFetch = new Promise(resolve => { + secondFetchStarted = resolve; + }); + const fetchImpl = asFetch(async () => { + fetchCount += 1; + if (fetchCount === 2) { + secondFetchStarted?.(); + } + return new Response( + streamFromChunks([ + 'data: {"directory":"/workspace/root","payload":{"type":"message.updated","properties":{"id":"msg_1"}}}\n\n', + 'data: {"directory":"/workspace/root","payload":{"type":"message.updated","properties":{"id":"msg_2"}}}\n\n', + ]), + { status: 200 } + ); + }); + + const connection = openKiloGlobalFeed({ + state, + kiloClient: makeKiloClient(), + fetchImpl, + WebSocketImpl: FakeWebSocket as unknown as GlobalFeedWebSocketImpl, + retryDelayMs: 1, + backpressureStallMs: 0, + }); + + // The first drop starts the stall window; the second drop exceeds it and + // fails the attempt, forcing a reconnect on a fresh socket. + await secondFetch; + connection.close(); + await connection.done; + + expect(FakeWebSocket.instances.length).toBeGreaterThanOrEqual(2); + expect(FakeWebSocket.instances[0].sent).toEqual([]); + }); + it('restarts the global feed after the private Kilo SSE stream ends', async () => { const state = new WrapperState(); bindGlobalFeedSession(state); diff --git a/services/cloud-agent-next/wrapper/src/global-feed.ts b/services/cloud-agent-next/wrapper/src/global-feed.ts index d7b6e98cc9..3c1ecdff77 100644 --- a/services/cloud-agent-next/wrapper/src/global-feed.ts +++ b/services/cloud-agent-next/wrapper/src/global-feed.ts @@ -18,11 +18,18 @@ type OpenKiloGlobalFeedOptions = { fetchImpl?: typeof fetch; WebSocketImpl?: WebSocketCtor; retryDelayMs?: number; + backpressureStallMs?: number; }; const OPEN_READY_STATE = 1; const GLOBAL_FEED_RETRY_DELAY_MS = 1_000; const MAX_GLOBAL_FEED_WEBSOCKET_BUFFERED_BYTES = 1024 * 1024; +/** + * How long outbound backpressure may persist before the attempt is failed so + * the outer loop reconnects on a fresh socket. Catches half-open sockets that + * stay OPEN but never drain, which would otherwise drop events forever. + */ +const GLOBAL_FEED_BACKPRESSURE_STALL_MS = 60_000; const encoder = new TextEncoder(); type RequiredGlobalFeedSession = SessionContext & { @@ -156,6 +163,7 @@ export function openKiloGlobalFeed(options: OpenKiloGlobalFeedOptions): KiloGlob const fetchImpl = options.fetchImpl ?? fetch; const WebSocketWithHeaders = (options.WebSocketImpl ?? WebSocket) as WebSocketCtor; const retryDelayMs = options.retryDelayMs ?? GLOBAL_FEED_RETRY_DELAY_MS; + const backpressureStallMs = options.backpressureStallMs ?? GLOBAL_FEED_BACKPRESSURE_STALL_MS; const closeController = new AbortController(); let closed = false; let attemptAbortController: AbortController | undefined; @@ -235,6 +243,7 @@ export function openKiloGlobalFeed(options: OpenKiloGlobalFeedOptions): KiloGlob throw new Error(`Kilo global event stream failed: ${response.status}`); } + let backpressureSince: number | undefined; for await (const data of parseSseDataStream(response.body)) { if (closed || abortController.signal.aborted) break; let parsed: unknown; @@ -249,14 +258,53 @@ export function openKiloGlobalFeed(options: OpenKiloGlobalFeedOptions): KiloGlob continue; } - if (ws.readyState === OPEN_READY_STATE) { - const serialized = JSON.stringify(parsed); - const pendingBytes = ws.bufferedAmount + encoder.encode(serialized).byteLength; - if (pendingBytes > MAX_GLOBAL_FEED_WEBSOCKET_BUFFERED_BYTES) { - throw new Error('Kilo global feed WebSocket buffer limit exceeded'); + if (ws.readyState !== OPEN_READY_STATE) { + continue; + } + + const serialized = JSON.stringify(parsed); + const eventBytes = encoder.encode(serialized).byteLength; + + // The global feed is best-effort: it does not own session completion, and + // the primary /ingest path remains authoritative. A single oversized + // event or transient backpressure never fatals the feed loop — drop the + // event, log a structured signal, and keep consuming so later events + // still flow once pressure clears. Backpressure that never clears is + // treated as a dead socket and fails the attempt so the outer loop + // reconnects. + + if (eventBytes > MAX_GLOBAL_FEED_WEBSOCKET_BUFFERED_BYTES) { + logToFile( + JSON.stringify({ + message: 'kilo_global_feed_event_dropped_oversized', + eventBytes, + limitBytes: MAX_GLOBAL_FEED_WEBSOCKET_BUFFERED_BYTES, + }) + ); + continue; + } + + const bufferedAmount = ws.bufferedAmount; + const pendingBytes = bufferedAmount + eventBytes; + if (pendingBytes > MAX_GLOBAL_FEED_WEBSOCKET_BUFFERED_BYTES) { + const now = Date.now(); + if (backpressureSince !== undefined && now - backpressureSince >= backpressureStallMs) { + throw new Error('Kilo global feed WebSocket backpressure stalled'); } - ws.send(serialized); + backpressureSince ??= now; + logToFile( + JSON.stringify({ + message: 'kilo_global_feed_event_dropped_backpressure', + bufferedAmount, + eventBytes, + limitBytes: MAX_GLOBAL_FEED_WEBSOCKET_BUFFERED_BYTES, + }) + ); + continue; } + + backpressureSince = undefined; + ws.send(serialized); } }