CloudAgentNext - budget ingest frame size and harden wrapper resume

services/cloud-agent-next/src/persistence/CloudAgentSession.ts

@@ -951,6 +951,17 @@ export class CloudAgentSession extends DurableObject<WorkerEnv> {
       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,

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<WrapperRuntimeState | null> {
+  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

services/cloud-agent-next/src/session/wrapper-supervisor.test.ts

@@ -208,6 +208,24 @@ function liveRuntimeState(overrides?: Record<string, unknown>): [string, unknown
   ];
 }
 
+function disconnectGraceForCurrentConnection(
+  disconnectedAt: number,
+  overrides?: Record<string, unknown>
+): [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(
       [

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<WrapperReconnectDecision>;
-  recordReconnectAccepted(fence: WrapperConnectionFence): Promise<void>;
+  recordReconnectAccepted(fence: WrapperConnectionFence, now?: number): Promise<void>;
   isCurrentConnection(wrapperGeneration: number, wrapperConnectionId: string): Promise<boolean>;
   observePong(wrapperGeneration: number, wrapperConnectionId: string, now: number): Promise<void>;
   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<boolean> {
+    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<void> {
     await messageSettlementOutbox.releaseWrapperTerminalWaitForIdleBatch();
     await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady({
@@ -428,8 +451,27 @@ export function createWrapperSupervisor(
     return { accepted: true };
   }
 
-  async function recordReconnectAccepted(fence: WrapperConnectionFence): Promise<void> {
+  async function recordReconnectAccepted(
+    fence: WrapperConnectionFence,
+    now = Date.now()
+  ): Promise<void> {
     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);
   }