diff --git a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts index 86d2c1afa7..9857dc844c 100644 --- a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts +++ b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts @@ -175,6 +175,7 @@ import { matchesSessionMessageReplay, nextQueuedMessageId, recordAcceptedMessageActivity, + releaseUnadmittedWaitingMessages, resolveSessionMessageIntent, streamCloudStatus, streamQueuedSnapshots, @@ -3024,7 +3025,7 @@ export class SandboxSession extends DurableObject { return; } if (wrapperInstanceId) this.retainRuntimeCleanup(metadata, wrapperInstanceId, reason); - await this.failWaitingMessages(reason, wrapperInstanceId); + await this.failDeliveryWaitingMessages(reason, wrapperInstanceId); if (this.pendingRuntimeCleanup()) await this.transferRuntimeCleanup(); } @@ -3044,6 +3045,34 @@ export class SandboxSession extends DurableObject { ) return; } + const epoch = this.terminalLifecycle.captureEpoch(); + if (epoch === null) return; + let before = this.loadMessages(); + if (!this.terminalLifecycle.isCurrent(epoch)) return; + let released = false; + if (wrapperInstanceId) { + const result = releaseUnadmittedWaitingMessages(before, wrapperInstanceId); + if (result.releasedIds.length > 0) { + before = result.messages; + released = true; + } + } + const { messages, failedIds } = applyFailWaitingMessages( + before, + reason, + wrapperInstanceId, + false + ); + if (failedIds.length === 0 && !released) return; + if (!this.saveMessages(messages, epoch)) return; + if (this.pendingRuntimeCleanup() || nextQueuedMessageId(this.loadMessages())) + await this.armQueueRetry(); + } + + private async failDeliveryWaitingMessages( + reason: string, + wrapperInstanceId?: string + ): Promise { const epoch = this.terminalLifecycle.captureEpoch(); if (epoch === null) return; const before = this.loadMessages(); diff --git a/services/cloud-agent-next/src/sandbox-session/recovery/quarantine-releases-unadmitted-queue.test.ts b/services/cloud-agent-next/src/sandbox-session/recovery/quarantine-releases-unadmitted-queue.test.ts new file mode 100644 index 0000000000..b876110b65 --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-session/recovery/quarantine-releases-unadmitted-queue.test.ts @@ -0,0 +1,191 @@ +import { describe, expect, it } from 'vitest'; +import { + ATTACH_FAILURE_LIMIT, + failWaitingMessages, + nextQueuedMessageId, + releaseUnadmittedWaitingMessages, + type SessionMessageRecord, + type SessionOperationProof, +} from '../session-message-queue.js'; + +describe('quarantine releases unadmitted queued messages', () => { + const wrapperA = 'wrapper-a'; + const wrapperB = 'wrapper-b'; + + function queued( + messageId: string, + overrides: Partial = {} + ): SessionMessageRecord { + return { + messageId, + state: 'queued' as const, + wrapperInstanceId: wrapperA, + ...overrides, + } as SessionMessageRecord; + } + + it('releases a queued message with retryable not_ready attach failure and no prompt', () => { + const messages: SessionMessageRecord[] = [queued('unadmitted', { attachFailures: 1 })]; + + const { messages: released, releasedIds } = releaseUnadmittedWaitingMessages( + messages, + wrapperA + ); + + expect(releasedIds).toEqual(['unadmitted']); + expect(released[0]).toEqual( + expect.objectContaining({ + messageId: 'unadmitted', + state: 'queued', + wrapperInstanceId: undefined, + attachFailures: 1, + }) + ); + expect(released[0].wrapperInstanceId).toBeUndefined(); + }); + + it('still fails an accepted message on the same wrapper', () => { + const messages: SessionMessageRecord[] = [ + { messageId: 'accepted', state: 'accepted', acceptedAt: 5, wrapperInstanceId: wrapperA }, + queued('unadmitted', { attachFailures: 1 }), + ]; + + const { messages: released, releasedIds } = releaseUnadmittedWaitingMessages( + messages, + wrapperA + ); + + // Accepted message is untouched by release — only queued unadmitted are released + expect(releasedIds).toEqual(['unadmitted']); + const accepted = released.find(m => m.messageId === 'accepted'); + expect(accepted?.state).toBe('accepted'); + expect(accepted?.wrapperInstanceId).toBe(wrapperA); + + // failWaitingMessages then fails the accepted message + const { failedIds } = failWaitingMessages(released, 'kilo_unhealthy', wrapperA, false); + expect(failedIds).toEqual(['accepted']); + }); + + it('does not release a queued message with a committed attach', () => { + const attachProof = { + authorization: {}, + dispatched: true, + completedAt: 100, + attachmentEpoch: 1, + } as SessionOperationProof; + const messages: SessionMessageRecord[] = [ + queued('attached', { operations: { attach: attachProof } }), + ]; + + const { releasedIds } = releaseUnadmittedWaitingMessages(messages, wrapperA); + expect(releasedIds).toEqual([]); + }); + + it('does not release a queued message that has a prompt operation', () => { + const promptProof = { + authorization: {}, + dispatched: true, + } as SessionOperationProof; + const messages: SessionMessageRecord[] = [ + queued('prompted', { operations: { prompt: promptProof } }), + ]; + + const { releasedIds } = releaseUnadmittedWaitingMessages(messages, wrapperA); + expect(releasedIds).toEqual([]); + }); + + it('does not release a queued message with exhausted attach failures', () => { + const messages: SessionMessageRecord[] = [ + queued('exhausted', { attachFailures: ATTACH_FAILURE_LIMIT }), + ]; + + const { releasedIds } = releaseUnadmittedWaitingMessages(messages, wrapperA); + expect(releasedIds).toEqual([]); + }); + + it('does not release messages bound to a different wrapper', () => { + const messages: SessionMessageRecord[] = [ + queued('other-wrapper', { wrapperInstanceId: wrapperB }), + ]; + + const { releasedIds } = releaseUnadmittedWaitingMessages(messages, wrapperA); + expect(releasedIds).toEqual([]); + }); + + it('does not release unassigned queued messages', () => { + const messages: SessionMessageRecord[] = [ + queued('unassigned', { wrapperInstanceId: undefined }), + ]; + + const { releasedIds } = releaseUnadmittedWaitingMessages(messages, wrapperA); + expect(releasedIds).toEqual([]); + }); + + it('released message is picked up by drain against wrapper B', () => { + const messages: SessionMessageRecord[] = [queued('unadmitted', { attachFailures: 1 })]; + + const { messages: released } = releaseUnadmittedWaitingMessages(messages, wrapperA); + + // After release, message is queued and unassigned — nextQueuedMessageId finds it + expect(nextQueuedMessageId(released)).toBe('unadmitted'); + + // The message can be bound to wrapper B by the normal delivery path + const rebound = released.map(m => + m.messageId === 'unadmitted' ? { ...m, wrapperInstanceId: wrapperB } : m + ); + expect(rebound[0].wrapperInstanceId).toBe(wrapperB); + }); + + it('preserves completed and failed history', () => { + const messages: SessionMessageRecord[] = [ + { messageId: 'done', state: 'completed' }, + { messageId: 'old-fail', state: 'failed', failedReason: 'prompt_exhausted' }, + queued('unadmitted'), + ]; + + const { messages: released, releasedIds } = releaseUnadmittedWaitingMessages( + messages, + wrapperA + ); + + expect(releasedIds).toEqual(['unadmitted']); + expect(released[0]).toEqual({ messageId: 'done', state: 'completed' }); + expect(released[1]).toEqual({ + messageId: 'old-fail', + state: 'failed', + failedReason: 'prompt_exhausted', + }); + }); + + it('clears preparationAttemptId and deliveryDeadlineAt on release', () => { + const messages: SessionMessageRecord[] = [ + queued('unadmitted', { + preparationAttemptId: 'attempt-1', + deliveryDeadlineAt: 999_999, + attachFailures: 1, + }), + ]; + + const { messages: released } = releaseUnadmittedWaitingMessages(messages, wrapperA); + expect(released[0].preparationAttemptId).toBeUndefined(); + expect(released[0].deliveryDeadlineAt).toBeUndefined(); + }); + + it('drops incomplete attach proofs on release', () => { + const attachProof = { + authorization: {}, + dispatched: false, + } as SessionOperationProof; + const messages: SessionMessageRecord[] = [ + queued('unadmitted', { attachFailures: 1, operations: { attach: attachProof } }), + ]; + + const { messages: released, releasedIds } = releaseUnadmittedWaitingMessages( + messages, + wrapperA + ); + + expect(releasedIds).toEqual(['unadmitted']); + expect(released[0].operations).toBeUndefined(); + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-session/session-message-queue.ts b/services/cloud-agent-next/src/sandbox-session/session-message-queue.ts index 417c6edcac..5b7ebea952 100644 --- a/services/cloud-agent-next/src/sandbox-session/session-message-queue.ts +++ b/services/cloud-agent-next/src/sandbox-session/session-message-queue.ts @@ -279,6 +279,41 @@ export function failWaitingMessages( }; } +/** + * Release queued messages bound to a dying wrapper that never reached a + * committed attach or prompt. These are safe to retry on a replacement + * runtime. Messages with a completed attach proof, a prompt operation, + * or exhausted attach failures remain bound so `failWaitingMessages` + * can fail them as today. + */ +export function releaseUnadmittedWaitingMessages( + messages: readonly SessionMessageRecord[], + wrapperInstanceId: string +): { messages: SessionMessageRecord[]; releasedIds: string[] } { + const releasedIds: string[] = []; + return { + messages: messages.map(message => { + if (message.state !== 'queued' || message.wrapperInstanceId !== wrapperInstanceId) { + return message; + } + if (message.unresolvedDispatch) return message; + if (message.operations?.prompt) return message; + if (message.operations?.attach?.completedAt !== undefined) return message; + if ((message.attachFailures ?? 0) >= ATTACH_FAILURE_LIMIT) return message; + + releasedIds.push(message.messageId); + return { + ...message, + wrapperInstanceId: undefined, + preparationAttemptId: undefined, + deliveryDeadlineAt: undefined, + operations: undefined, + }; + }), + releasedIds, + }; +} + export function incrementDeliveryFailure( messages: readonly SessionMessageRecord[], messageId: string,