diff --git a/services/cloud-agent-next/src/execution/errors.ts b/services/cloud-agent-next/src/execution/errors.ts index 68ef6cee9c..f930db5b81 100644 --- a/services/cloud-agent-next/src/execution/errors.ts +++ b/services/cloud-agent-next/src/execution/errors.ts @@ -14,7 +14,8 @@ export type RetryableErrorCode = | 'SANDBOX_CONNECT_FAILED' // Sandbox may be waking up or network issue | 'WORKSPACE_SETUP_FAILED' // Git clone/network transient failure | 'KILO_SERVER_FAILED' // Kilo server starting up - | 'WRAPPER_START_FAILED'; // Wrapper process starting + | 'WRAPPER_START_FAILED' // Wrapper process starting + | 'WRAPPER_FINALIZING'; // Wrapper sealed the current run before this delivery /** * Error codes for non-retryable failures (4xx/5xx). diff --git a/services/cloud-agent-next/src/execution/orchestrator.test.ts b/services/cloud-agent-next/src/execution/orchestrator.test.ts index 85c0f6894a..30b7105f2c 100644 --- a/services/cloud-agent-next/src/execution/orchestrator.test.ts +++ b/services/cloud-agent-next/src/execution/orchestrator.test.ts @@ -1,6 +1,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { AgentSandbox, WrapperInstanceLease } from '../agent-sandbox/protocol.js'; import type { Env } from '../types.js'; +import { WrapperError } from '../kilo/wrapper-client.js'; import type { ExecutionError } from './errors.js'; import type { FencedWrapperDispatchRequest } from './types.js'; import { @@ -236,6 +237,30 @@ describe('ExecutionOrchestrator AgentSandbox delivery', () => { } satisfies Partial); }); + it('preserves a finalizing error from wrapper startup', async () => { + const { orchestrator, ensureWrapper } = createOrchestrator(); + const finalizingError = new WrapperError( + 'Wrapper batch is finalizing', + 'WRAPPER_FINALIZING', + 409 + ); + ensureWrapper.mockRejectedValueOnce(finalizingError); + + await expect(orchestrator.execute(basePlan)).rejects.toBe(finalizingError); + }); + + it('preserves a finalizing error from wrapper dispatch', async () => { + const { orchestrator, prompt } = createOrchestrator(); + const finalizingError = new WrapperError( + 'Wrapper batch is finalizing', + 'WRAPPER_FINALIZING', + 409 + ); + prompt.mockRejectedValueOnce(finalizingError); + + await expect(orchestrator.execute(basePlan)).rejects.toBe(finalizingError); + }); + it('does not recover the shared sandbox for plain capacity admission rejection', async () => { const { orchestrator, ensureWrapper, deleteSandbox } = createOrchestrator(); ensureWrapper.mockRejectedValueOnce( diff --git a/services/cloud-agent-next/src/execution/orchestrator.ts b/services/cloud-agent-next/src/execution/orchestrator.ts index 0531fdd278..51eda4ffc7 100644 --- a/services/cloud-agent-next/src/execution/orchestrator.ts +++ b/services/cloud-agent-next/src/execution/orchestrator.ts @@ -141,6 +141,7 @@ export class ExecutionOrchestrator { }); } catch (error) { if (error instanceof ExecutionError) throw error; + if (error instanceof WrapperError && error.code === 'WRAPPER_FINALIZING') throw error; throw ExecutionError.wrapperStartFailed( `Failed to start wrapper: ${error instanceof Error ? error.message : String(error)}`, error @@ -210,6 +211,9 @@ export class ExecutionOrchestrator { if (error.code === 'KILO_SERVER_FAILED') { throw ExecutionError.kiloServerFailed(error.message, error); } + if (error.code === 'WRAPPER_FINALIZING') { + throw error; + } } if (error instanceof ExecutionError) throw error; throw ExecutionError.wrapperStartFailed( diff --git a/services/cloud-agent-next/src/execution/types.ts b/services/cloud-agent-next/src/execution/types.ts index e5ed9b895d..de490346b5 100644 --- a/services/cloud-agent-next/src/execution/types.ts +++ b/services/cloud-agent-next/src/execution/types.ts @@ -225,7 +225,8 @@ export type RetryableResultCode = | 'SANDBOX_CONNECT_FAILED' | 'WORKSPACE_SETUP_FAILED' | 'KILO_SERVER_FAILED' - | 'WRAPPER_START_FAILED'; + | 'WRAPPER_START_FAILED' + | 'WRAPPER_FINALIZING'; export type AdmissionFailure = { success: false; diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts index 9f89af9ba1..3ff7d2ba41 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts @@ -637,6 +637,7 @@ function promptAdmissionError( case 'WORKSPACE_SETUP_FAILED': case 'KILO_SERVER_FAILED': case 'WRAPPER_START_FAILED': + case 'WRAPPER_FINALIZING': return facadeError(503, result.code, result.error); case 'INTERNAL': return facadeError(500, 'KILO_PROMPT_ADMISSION_FAILED', result.error); diff --git a/services/cloud-agent-next/src/kilo/wrapper-client.test.ts b/services/cloud-agent-next/src/kilo/wrapper-client.test.ts index bdfa4a5fb0..d809dd46a9 100644 --- a/services/cloud-agent-next/src/kilo/wrapper-client.test.ts +++ b/services/cloud-agent-next/src/kilo/wrapper-client.test.ts @@ -12,6 +12,7 @@ import { WrapperClient, WrapperContainerClient, WrapperError, + WrapperFinalizingError, WrapperNotReadyError, WrapperNoJobError, WrapperJobConflictError, @@ -361,6 +362,14 @@ describe('WrapperClient', () => { expect(result.lastError).toBeDefined(); expect(result.lastError?.code).toBe('INFLIGHT_TIMEOUT'); }); + + it('returns finalizing status', async () => { + const statusResponse: JobStatus = { state: 'finalizing', sessionId: 'kilo_456' }; + const session = createMockSession(createSuccessResponse(statusResponse)); + const client = new WrapperClient({ session, port: defaultPort }); + + await expect(client.status()).resolves.toEqual(statusResponse); + }); }); // ------------------------------------------------------------------------- @@ -396,6 +405,41 @@ describe('WrapperClient', () => { expect(result.messageId).toBeUndefined(); }); + it('parses finalizing responses with the wrapper run identity', async () => { + const session = createMockSession({ + exitCode: 0, + stdout: JSON.stringify({ + error: 'WRAPPER_FINALIZING', + message: 'Wrapper batch is finalizing', + wrapperRunId: 'wr_old', + }), + }); + const client = new WrapperClient({ session, port: defaultPort }); + + await expect(client.prompt(createPromptOptions())).rejects.toEqual( + expect.objectContaining({ + name: 'WrapperFinalizingError', + code: 'WRAPPER_FINALIZING', + wrapperRunId: 'wr_old', + }) + ); + }); + + it('parses legacy finalizing responses without a wrapper run identity', async () => { + const session = createMockSession( + createErrorResponse('WRAPPER_FINALIZING', 'Wrapper batch is finalizing') + ); + const client = new WrapperClient({ session, port: defaultPort }); + + await expect(client.prompt(createPromptOptions())).rejects.toEqual( + expect.objectContaining({ + name: 'WrapperFinalizingError', + code: 'WRAPPER_FINALIZING', + wrapperRunId: undefined, + }) + ); + }); + it('sends prompt text', async () => { const session = createMockSession( createSuccessResponse({ status: 'sent', messageId: 'msg_1' }) @@ -2074,6 +2118,11 @@ describe('WrapperClient', () => { // ------------------------------------------------------------------------- describe('error classes', () => { + it('WrapperFinalizingError carries optional wrapper run identity', () => { + expect(new WrapperFinalizingError('finalizing', 'wr_old').wrapperRunId).toBe('wr_old'); + expect(new WrapperFinalizingError('legacy').wrapperRunId).toBeUndefined(); + }); + it('WrapperError has correct properties', () => { const error = new WrapperError('Test message', 'TEST_CODE', 500); diff --git a/services/cloud-agent-next/src/kilo/wrapper-client.ts b/services/cloud-agent-next/src/kilo/wrapper-client.ts index 582f4fa66a..2b9208528b 100644 --- a/services/cloud-agent-next/src/kilo/wrapper-client.ts +++ b/services/cloud-agent-next/src/kilo/wrapper-client.ts @@ -131,7 +131,7 @@ export type WrapperPty = { }; export type JobStatus = { - state: 'idle' | 'active'; + state: 'idle' | 'active' | 'finalizing'; sessionId?: string; lastError?: { code: string; @@ -168,6 +168,16 @@ export class WrapperError extends Error { } } +export class WrapperFinalizingError extends WrapperError { + constructor( + message: string, + public readonly wrapperRunId?: string + ) { + super(message, 'WRAPPER_FINALIZING', 409); + this.name = 'WrapperFinalizingError'; + } +} + export class WrapperNotReadyError extends WrapperError { constructor(message: string, options?: ErrorOptions) { super(message, 'NOT_READY', 503, options); @@ -198,6 +208,7 @@ const ERROR_STATUS_CODES: Record = { WORKSPACE_SETUP_FAILED: 503, KILO_SERVER_FAILED: 503, SEND_ERROR: 500, + WRAPPER_FINALIZING: 409, }; // --------------------------------------------------------------------------- @@ -469,6 +480,7 @@ export class WrapperClient { error?: string; message?: string; retryable?: boolean; + wrapperRunId?: string; }; // Check for error response @@ -485,6 +497,12 @@ export class WrapperClient { if (errorCode === 'JOB_CONFLICT') { throw new WrapperJobConflictError(parsed.message ?? 'Job conflict'); } + if (errorCode === 'WRAPPER_FINALIZING') { + throw new WrapperFinalizingError( + parsed.message ?? 'Wrapper batch is finalizing', + parsed.wrapperRunId + ); + } throw new WrapperError(parsed.message ?? errorCode, errorCode, statusCode); } diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index 9767a7fe4f..aed7e68e8a 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -98,6 +98,8 @@ import { clearWrapperRuntimeIdentity, getWrapperLease, getWrapperRuntimeState, + isWrapperDeliveryHeld, + isWrapperRunFinalizing, nextWrapperCleanupDeadline, nextWrapperLeaseDeadline, } from '../session/wrapper-runtime-state.js'; @@ -494,10 +496,7 @@ export class CloudAgentSession extends DurableObject { ...event, }); }, - hasObservedWrapperIdle: async () => { - const state = await getWrapperRuntimeState(this.ctx.storage); - return state.lastWrapperIdleAt !== undefined; - }, + hasObservedWrapperIdle: async () => false, requestAlarmAtOrBefore: deadline => this.scheduleAlarmAtOrBefore(deadline), getSessionIdForLogs: () => this.sessionId, }); @@ -549,6 +548,8 @@ export class CloudAgentSession extends DurableObject { hasActiveIngestConnection: async params => (await this.getIngestHandler()).hasActiveConnection(params), clearInterruptRequest: () => this.executionQueries.clearInterrupt(), + ensureAcceptedMessageBeforeTerminal: (messageId, wrapperRunId) => + this.ensureAcceptedMessageBeforeTerminal(messageId, wrapperRunId), stopWrappers: async request => { if (this.physicalWrapperStopper) return this.physicalWrapperStopper(request); if (this.orchestrator || (!this.env.Sandbox && !this.env.SandboxSmall)) { @@ -604,6 +605,8 @@ export class CloudAgentSession extends DurableObject { return retryAt === undefined ? null : { retryAt }; }, deliver: plan => this.executeDirectly(plan), + isDeliveryHeld: async () => + isWrapperRunFinalizing(await getWrapperRuntimeState(this.ctx.storage)), ensureQueuedMessageEvent: event => { this.ensureQueuedMessageEvent({ executionId: '' as EventSourceId, @@ -2013,9 +2016,14 @@ export class CloudAgentSession extends DurableObject { nextAlarmAt = existingAlarm; } + const pendingDeliveryHeld = isWrapperDeliveryHeld( + await getWrapperRuntimeState(this.ctx.storage), + await getWrapperLease(this.ctx.storage) + ); if ( pendingFlushRetryAt === undefined && pendingCount > 0 && + !pendingDeliveryHeld && currentTime + PENDING_FLUSH_DEBOUNCE_MS < nextAlarmAt ) { nextAlarmAt = currentTime + PENDING_FLUSH_DEBOUNCE_MS; @@ -2918,6 +2926,7 @@ export class CloudAgentSession extends DurableObject { status: 'completed' | 'failed' | 'interrupted'; error?: string; gateResult?: 'pass' | 'fail'; + messageIds?: string[]; }): Promise { await this.resolveSessionId(); await this.getWrapperSupervisor().onTerminalEvent(params); diff --git a/services/cloud-agent-next/src/session/agent-runtime.test.ts b/services/cloud-agent-next/src/session/agent-runtime.test.ts index 7035e8ccd5..1f3cb4bf4b 100644 --- a/services/cloud-agent-next/src/session/agent-runtime.test.ts +++ b/services/cloud-agent-next/src/session/agent-runtime.test.ts @@ -6,6 +6,7 @@ import type { MessageDeliveryRequest, WorkspaceReady, } from '../execution/types.js'; +import { WrapperFinalizingError } from '../kilo/wrapper-client.js'; import { createAgentRuntime } from './agent-runtime.js'; import { getWrapperLease, getWrapperRuntimeState } from './wrapper-runtime-state.js'; import type { SessionMetadata } from '../persistence/session-metadata.js'; @@ -48,6 +49,7 @@ function createMetadata(): SessionMetadata { }, auth: { kiloSessionId: 'kilo_runtime', + kilocodeToken: 'kilo_token', }, lifecycle: { version: 1, @@ -188,6 +190,46 @@ describe('AgentRuntime', () => { expect(wrapperState.nextPingAt).toEqual(expect.any(Number)); }); + it('fences the dispatching message until acceptance bookkeeping completes', async () => { + const storage = createMemoryStorage(); + const messageId = createPlan().turn.messageId; + const runtime = createAgentRuntime({ + storage, + env: {} as Env, + getMetadata: async () => createMetadata(), + getOrchestratorOverride: () => ({ + execute: async (_plan, options) => { + await expect(getWrapperRuntimeState(storage)).resolves.not.toHaveProperty( + 'dispatchingMessageId' + ); + await options?.onWorkspaceReady?.(createWorkspaceReady()); + await expect(getWrapperRuntimeState(storage)).resolves.toMatchObject({ + dispatchingMessageId: messageId, + }); + return { kiloSessionId: 'kilo_runtime' }; + }, + }), + getSessionIdForLogs: () => 'agent_runtime', + sendToWrapper: () => false, + createAgentSandbox: () => + ({ + discoverSessionWrappers: vi.fn().mockResolvedValue({ status: 'absent' }), + }) as unknown as AgentSandbox, + }); + + await runtime.send(createPlan(), { + onAccepted: async () => { + await expect(getWrapperRuntimeState(storage)).resolves.toMatchObject({ + dispatchingMessageId: messageId, + }); + }, + }); + + await expect(getWrapperRuntimeState(storage)).resolves.not.toHaveProperty( + 'dispatchingMessageId' + ); + }); + it('keeps an accepted new delivery supervised when physical lease acceptance persistence fails', async () => { let rejectedAcceptedLeaseWrite = false; const storage = createMemoryStorage(undefined, (key, value) => { @@ -239,6 +281,180 @@ describe('AgentRuntime', () => { }); }); + it('holds a typed finalizing race from the real orchestrator without clearing the current run', async () => { + const storage = createMemoryStorage([ + [ + 'wrapper_runtime_state', + { wrapperGeneration: 3, wrapperConnectionId: 'conn_hot', wrapperRunId: 'wr_hot' }, + ], + [ + 'wrapper_lease', + { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_hot', instanceGeneration: 1 }, + }, + ], + ]); + const prompt = vi + .fn() + .mockRejectedValue(new WrapperFinalizingError('Wrapper batch is finalizing', 'wr_hot')); + const wrapper = { + ensureSessionReady: vi.fn().mockResolvedValue({ kiloSessionId: 'kilo_runtime' }), + prompt, + command: vi.fn(), + }; + const sandbox = { + discoverSessionWrappers: vi.fn().mockResolvedValue({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'wrapper-hot', + port: 5_000, + instanceId: 'instance_hot', + instanceGeneration: 1, + }, + ], + }), + ensureWrapper: vi.fn().mockResolvedValue({ status: 'wrapper-running', client: wrapper }), + delete: vi.fn(), + } as unknown as AgentSandbox; + const runtime = createAgentRuntime({ + storage, + env: { WORKER_URL: 'https://worker.example.com' } as Env, + getMetadata: async () => createMetadata(), + getSessionIdForLogs: () => 'agent_runtime', + sendToWrapper: () => false, + createAgentSandbox: () => sandbox, + }); + + await expect(runtime.send(createPlan())).resolves.toEqual({ + success: false, + code: 'WRAPPER_FINALIZING', + error: 'Wrapper batch is finalizing', + }); + expect(prompt).toHaveBeenCalledOnce(); + await expect(getWrapperRuntimeState(storage)).resolves.toMatchObject({ + wrapperRunId: 'wr_hot', + finalizingWrapperRunId: 'wr_hot', + }); + }); + + it('does not mark a fresh run finalizing when the old wrapper rejects it during close delay', async () => { + const storage = createMemoryStorage([ + [ + 'wrapper_runtime_state', + { wrapperGeneration: 4, wrapperConnectionId: 'conn_fresh', wrapperRunId: 'wr_fresh' }, + ], + [ + 'wrapper_lease', + { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_hot', instanceGeneration: 1 }, + }, + ], + ]); + const wrapper = { + ensureSessionReady: vi.fn().mockResolvedValue({ kiloSessionId: 'kilo_runtime' }), + prompt: vi + .fn() + .mockRejectedValue(new WrapperFinalizingError('Wrapper batch is finalizing', 'wr_old')), + command: vi.fn(), + }; + const sandbox = { + discoverSessionWrappers: vi.fn().mockResolvedValue({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'wrapper-closing', + port: 5_000, + instanceId: 'instance_hot', + instanceGeneration: 1, + }, + ], + }), + ensureWrapper: vi.fn().mockResolvedValue({ status: 'wrapper-running', client: wrapper }), + delete: vi.fn(), + } as unknown as AgentSandbox; + const runtime = createAgentRuntime({ + storage, + env: { WORKER_URL: 'https://worker.example.com' } as Env, + getMetadata: async () => createMetadata(), + getSessionIdForLogs: () => 'agent_runtime', + sendToWrapper: () => false, + createAgentSandbox: () => sandbox, + }); + + await expect(runtime.send(createPlan())).resolves.toEqual({ + success: false, + code: 'WRAPPER_FINALIZING', + error: 'Wrapper batch is finalizing', + }); + const runtimeState = await getWrapperRuntimeState(storage); + expect(runtimeState.wrapperRunId).toBe('wr_fresh'); + expect(runtimeState.finalizingWrapperRunId).toBeUndefined(); + }); + + it('preserves older finalizing errors without a wrapper run identity conservatively', async () => { + const storage = createMemoryStorage([ + [ + 'wrapper_runtime_state', + { wrapperGeneration: 3, wrapperConnectionId: 'conn_hot', wrapperRunId: 'wr_hot' }, + ], + [ + 'wrapper_lease', + { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_hot', instanceGeneration: 1 }, + }, + ], + ]); + const sandbox = { + discoverSessionWrappers: vi.fn().mockResolvedValue({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'wrapper-hot', + port: 5_000, + instanceId: 'instance_hot', + instanceGeneration: 1, + }, + ], + }), + ensureWrapper: vi.fn().mockResolvedValue({ + status: 'wrapper-running', + client: { + ensureSessionReady: vi.fn().mockResolvedValue({ kiloSessionId: 'kilo_runtime' }), + prompt: vi.fn().mockRejectedValue(new WrapperFinalizingError('legacy finalizing')), + command: vi.fn(), + }, + }), + delete: vi.fn(), + } as unknown as AgentSandbox; + const runtime = createAgentRuntime({ + storage, + env: { WORKER_URL: 'https://worker.example.com' } as Env, + getMetadata: async () => createMetadata(), + getSessionIdForLogs: () => 'agent_runtime', + sendToWrapper: () => false, + createAgentSandbox: () => sandbox, + }); + + await expect(runtime.send(createPlan())).resolves.toMatchObject({ + success: false, + code: 'WRAPPER_FINALIZING', + }); + await expect(getWrapperRuntimeState(storage)).resolves.toMatchObject({ + wrapperRunId: 'wr_hot', + finalizingWrapperRunId: 'wr_hot', + }); + }); + it('preserves accepted-message liveness when a hot follow-up fails before acceptance', async () => { const storage = createMemoryStorage([ [ diff --git a/services/cloud-agent-next/src/session/agent-runtime.ts b/services/cloud-agent-next/src/session/agent-runtime.ts index 82737569c1..06e45bfa5d 100644 --- a/services/cloud-agent-next/src/session/agent-runtime.ts +++ b/services/cloud-agent-next/src/session/agent-runtime.ts @@ -13,6 +13,7 @@ import type { WorkspaceReady, } from '../execution/types.js'; import { logger } from '../logger.js'; +import { WrapperFinalizingError } from '../kilo/wrapper-client.js'; import type { SessionMetadata } from '../persistence/session-metadata.js'; import type { WrapperCommand } from '../shared/protocol.js'; import type { Env as WorkerEnv } from '../types.js'; @@ -20,13 +21,17 @@ import { WrapperCleanupBlockedError } from './wrapper-cleanup-blocked-error.js'; import { allocateWrapperRuntimeState, clearAllocatedWrapperRuntimeState, + clearWrapperDispatchingMessage, clearWrapperRuntimeIdentity, getWrapperLease, getWrapperRuntimeState, + isWrapperRunFinalizing, + markWrapperFinalizing, nextWrapperCleanupDeadline, putWrapperLease, READY_ONLY_IDLE_MS, recordWrapperAcceptedMessage, + recordWrapperDispatchingMessage, recordWrapperReadyLease, reduceWrapperLease, type WrapperLease, @@ -268,6 +273,14 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen ): Promise { const { sessionId } = plan.scope; const { turn, agent } = plan; + const currentRuntimeState = await getWrapperRuntimeState(storage); + if (isWrapperRunFinalizing(currentRuntimeState)) { + return { + success: false, + code: 'WRAPPER_FINALIZING', + error: 'Wrapper batch is finalizing', + }; + } const { leasedInstance, allocatedPhysicalInstance, requiresFreshRunFence } = await authorizePhysicalWrapper(plan); const previousRuntimeState = await getWrapperRuntimeState(storage); @@ -336,6 +349,7 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen }) .info('AgentRuntime wrapper workspace reported ready'); await hooks.onWorkspaceReady?.(ready); + await recordWrapperDispatchingMessage(storage, wrapperRuntimeState, turn.messageId); }, }); @@ -350,6 +364,7 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen acceptedAt, wrapperRunId: wrapperRuntimeState.wrapperRunId, }); + await clearWrapperDispatchingMessage(storage, wrapperRuntimeState, turn.messageId); try { const acceptedLease = await getWrapperLease(storage); await putWrapperLease( @@ -381,6 +396,20 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen .info('AgentRuntime wrapper accepted pending session message'); return buildRuntimeAcceptanceResult(turn.messageId, wrapperRuntimeState.wrapperRunId); } catch (error) { + await clearWrapperDispatchingMessage(storage, wrapperRuntimeState, turn.messageId); + if (error instanceof WrapperFinalizingError) { + if ( + error.wrapperRunId === undefined || + error.wrapperRunId === wrapperRuntimeState.wrapperRunId + ) { + await markWrapperFinalizing(storage, wrapperRuntimeState.wrapperRunId); + } + return { + success: false, + code: 'WRAPPER_FINALIZING', + error: error.message, + }; + } logger .withFields({ sessionId, diff --git a/services/cloud-agent-next/src/session/message-settlement-outbox.test.ts b/services/cloud-agent-next/src/session/message-settlement-outbox.test.ts index 633d4a141e..001e414ac5 100644 --- a/services/cloud-agent-next/src/session/message-settlement-outbox.test.ts +++ b/services/cloud-agent-next/src/session/message-settlement-outbox.test.ts @@ -9,6 +9,7 @@ import { createMessageSettlementOutbox, type MessageSettlementOutboxStorage, } from './message-settlement-outbox.js'; +import { createPendingSessionMessage, storePendingSessionMessage } from './pending-messages.js'; import { getSessionMessageState, putSessionMessageState, @@ -500,6 +501,83 @@ describe('MessageSettlementOutbox', () => { }); }); + it('finalizes a terminal wrapper-run callback while the next run remains pending', async () => { + const harness = createHarness(); + await putSessionMessageState( + harness.storage, + acceptedMessageState(firstMessageId, { url: 'https://example.com/sealed-batch' }) + ); + await storePendingSessionMessage( + harness.storage, + createPendingSessionMessage({ + messageId: secondMessageId, + role: 'user', + content: 'next prompt', + createdAt: 3_000, + }) + ); + + await harness.outbox.terminalizeSessionMessageOnce(firstMessageId, { + kind: 'completed', + completionSource: 'idle_reconciliation', + }); + expect(harness.callbackJobs).toHaveLength(0); + + await harness.outbox.finalizeTerminalWrapperRunCallbackIfReady('wr_outbox'); + + expect(harness.callbackJobs).toHaveLength(1); + expect(harness.callbackJobs[0].payload.messageId).toBe(firstMessageId); + }); + + it('keeps a wrapper-run callback blocked while that run has a nonterminal accepted message', async () => { + const harness = createHarness(); + await putSessionMessageState( + harness.storage, + acceptedMessageState(firstMessageId, { url: 'https://example.com/sealed-batch' }) + ); + await putSessionMessageState(harness.storage, acceptedMessageState(secondMessageId)); + + await harness.outbox.terminalizeSessionMessageOnce(firstMessageId, { + kind: 'completed', + completionSource: 'idle_reconciliation', + }); + await harness.outbox.finalizeTerminalWrapperRunCallbackIfReady('wr_outbox'); + + expect(harness.callbackJobs).toHaveLength(0); + }); + + it('preserves gate-result waits when finalizing a terminal wrapper-run callback', async () => { + const harness = createHarness({ + metadata: { ...metadata, finalization: { gateThreshold: 'warning' } }, + }); + await putSessionMessageState( + harness.storage, + acceptedMessageState(firstMessageId, { url: 'https://example.com/sealed-batch-gate' }) + ); + await storePendingSessionMessage( + harness.storage, + createPendingSessionMessage({ + messageId: secondMessageId, + role: 'user', + content: 'next prompt', + createdAt: 3_000, + }) + ); + + await harness.outbox.terminalizeSessionMessageOnce(firstMessageId, { + kind: 'completed', + completionSource: 'idle_reconciliation', + }); + await harness.outbox.finalizeTerminalWrapperRunCallbackIfReady('wr_outbox'); + expect(harness.callbackJobs).toHaveLength(0); + + await harness.outbox.observeWrapperTerminalForIdleBatch('pass'); + await harness.outbox.finalizeTerminalWrapperRunCallbackIfReady('wr_outbox'); + + expect(harness.callbackJobs).toHaveLength(1); + expect(harness.callbackJobs[0].payload.gateResult).toBe('pass'); + }); + it('includes a persisted completed message gate result in callback jobs', async () => { const harness = createHarness(); await putSessionMessageState( diff --git a/services/cloud-agent-next/src/session/message-settlement-outbox.ts b/services/cloud-agent-next/src/session/message-settlement-outbox.ts index ef80843ff4..5e94277f1c 100644 --- a/services/cloud-agent-next/src/session/message-settlement-outbox.ts +++ b/services/cloud-agent-next/src/session/message-settlement-outbox.ts @@ -77,6 +77,7 @@ export type MessageSettlementOutbox = { releaseWrapperTerminalWaitForIdleBatchForWrapperRun(wrapperRunId: string): Promise; isWaitingForWrapperTerminalGateResult(): Promise; finalizeIdleBatchCallbackIfReady(options?: FinalizeIdleBatchCallbackOptions): Promise; + finalizeTerminalWrapperRunCallbackIfReady(wrapperRunId: string): Promise; repairTerminalEffects(): Promise; retryPendingCallbacks(now: number): Promise; nextCallbackDeadline(): Promise; @@ -547,6 +548,33 @@ export function createMessageSettlementOutbox( return gateThreshold !== undefined && gateThreshold !== 'off'; } + async function finalizeIdleBatchCallbackState( + batch: IdleBatchCallbackState, + allowWithoutObservedIdle: boolean + ): Promise { + const now = Date.now(); + const finalized: IdleBatchCallbackState = { + ...batch, + finalizedAt: now, + updatedAt: now, + }; + await storage.put(idleBatchCallbackKey(finalized.batchId), finalized); + await storage.delete(CURRENT_IDLE_BATCH_CALLBACK_KEY); + logger + .withFields({ + sessionId: getSessionIdForLogs(), + batchId: finalized.batchId, + representativeMessageId: finalized.representativeMessageId, + allowWithoutObservedIdle, + }) + .info('Finalized idle-batch callback state'); + + if (!finalized.representativeMessageId) return; + const representative = await getSessionMessageState(storage, finalized.representativeMessageId); + if (!representative?.callbackRequired || representative.callbackEnqueuedAt) return; + await enqueueMessageCallbackNotification(representative); + } + async function finalizeIdleBatchCallbackIfReady( options?: FinalizeIdleBatchCallbackOptions ): Promise { @@ -559,7 +587,8 @@ export function createMessageSettlementOutbox( const acceptedMessages = await listNonTerminalAcceptedMessages(storage); if (acceptedMessages.length > 0) return; - if (!options?.allowWithoutObservedIdle && !(await hasObservedWrapperIdle())) { + const allowWithoutObservedIdle = options?.allowWithoutObservedIdle ?? false; + if (!allowWithoutObservedIdle && !(await hasObservedWrapperIdle())) { return; } @@ -567,26 +596,24 @@ export function createMessageSettlementOutbox( return; } - const finalized: IdleBatchCallbackState = { - ...batch, - finalizedAt: Date.now(), - updatedAt: Date.now(), - }; - await storage.put(idleBatchCallbackKey(finalized.batchId), finalized); - await storage.delete(CURRENT_IDLE_BATCH_CALLBACK_KEY); - logger - .withFields({ - sessionId: getSessionIdForLogs(), - batchId: finalized.batchId, - representativeMessageId: finalized.representativeMessageId, - allowWithoutObservedIdle: options?.allowWithoutObservedIdle ?? false, - }) - .info('Finalized idle-batch callback state'); + await finalizeIdleBatchCallbackState(batch, allowWithoutObservedIdle); + } - if (!finalized.representativeMessageId) return; - const representative = await getSessionMessageState(storage, finalized.representativeMessageId); - if (!representative?.callbackRequired || representative.callbackEnqueuedAt) return; - await enqueueMessageCallbackNotification(representative); + async function finalizeTerminalWrapperRunCallbackIfReady(wrapperRunId: string): Promise { + const batch = await getCurrentIdleBatchCallbackState(); + if (!batch?.representativeMessageId) return; + + const representative = await getSessionMessageState(storage, batch.representativeMessageId); + if (representative?.wrapperRunId !== wrapperRunId) return; + + const acceptedMessages = await listNonTerminalAcceptedMessages(storage, wrapperRunId); + if (acceptedMessages.length > 0) return; + + if (await shouldWaitForWrapperGateResult(batch)) { + return; + } + + await finalizeIdleBatchCallbackState(batch, true); } async function listPendingIdleBatchCallbacks(): Promise { @@ -770,6 +797,7 @@ export function createMessageSettlementOutbox( releaseWrapperTerminalWaitForIdleBatchForWrapperRun, isWaitingForWrapperTerminalGateResult, finalizeIdleBatchCallbackIfReady, + finalizeTerminalWrapperRunCallbackIfReady, repairTerminalEffects, retryPendingCallbacks, nextCallbackDeadline, diff --git a/services/cloud-agent-next/src/session/pending-messages.ts b/services/cloud-agent-next/src/session/pending-messages.ts index 1da5eb69ce..f166908972 100644 --- a/services/cloud-agent-next/src/session/pending-messages.ts +++ b/services/cloud-agent-next/src/session/pending-messages.ts @@ -68,6 +68,7 @@ const PendingFlushFailureCodeSchema = z.enum([ 'WORKSPACE_SETUP_FAILED', 'KILO_SERVER_FAILED', 'WRAPPER_START_FAILED', + 'WRAPPER_FINALIZING', 'NOT_FOUND', 'BAD_REQUEST', 'INTERNAL', diff --git a/services/cloud-agent-next/src/session/queries/events.ts b/services/cloud-agent-next/src/session/queries/events.ts index 622cdfccbc..4ec65dc357 100644 --- a/services/cloud-agent-next/src/session/queries/events.ts +++ b/services/cloud-agent-next/src/session/queries/events.ts @@ -340,10 +340,6 @@ export function createEventQueries(db: DrizzleSqliteDODatabase, rawSql: SqlStora AND json_extract(payload, '$.properties.info.role') = 'assistant' AND json_extract(payload, '$.properties.info.sessionID') = ? AND json_extract(payload, '$.properties.info.parentID') = ? - AND ( - json_extract(payload, '$.properties.info.time.completed') IS NOT NULL - OR json_extract(payload, '$.properties.info.error') IS NOT NULL - ) LIMIT 1 `, sessionId, @@ -378,10 +374,6 @@ export function createEventQueries(db: DrizzleSqliteDODatabase, rawSql: SqlStora AND json_extract(payload, '$.properties.info.role') = 'assistant' AND json_extract(payload, '$.properties.info.sessionID') = ? AND json_extract(payload, '$.properties.info.parentID') = ? - AND ( - json_extract(payload, '$.properties.info.time.completed') IS NOT NULL - OR json_extract(payload, '$.properties.info.error') IS NOT NULL - ) ORDER BY id DESC LIMIT 1 `, diff --git a/services/cloud-agent-next/src/session/queue-message.ts b/services/cloud-agent-next/src/session/queue-message.ts index 189c8b8b8e..2a3acbe955 100644 --- a/services/cloud-agent-next/src/session/queue-message.ts +++ b/services/cloud-agent-next/src/session/queue-message.ts @@ -28,6 +28,7 @@ const RETRYABLE_CODES: readonly RetryableResultCode[] = [ 'WORKSPACE_SETUP_FAILED', 'KILO_SERVER_FAILED', 'WRAPPER_START_FAILED', + 'WRAPPER_FINALIZING', ] as const; function isRetryableCode(code: string): code is RetryableResultCode { diff --git a/services/cloud-agent-next/src/session/session-message-queue.test.ts b/services/cloud-agent-next/src/session/session-message-queue.test.ts index b9997a92be..5ea77dc32f 100644 --- a/services/cloud-agent-next/src/session/session-message-queue.test.ts +++ b/services/cloud-agent-next/src/session/session-message-queue.test.ts @@ -417,6 +417,37 @@ describe('flushNextPendingSessionMessage', () => { }); }); + it('holds typed wrapper finalizing without consuming a flush attempt', async () => { + const storage = createMemoryStorage(); + await storePendingSessionMessage( + storage, + createPendingSessionMessage({ + messageId: FIRST_MESSAGE_ID, + role: 'user', + content: 'hold during finalization', + createdAt: 1, + }) + ); + const deliver = vi.fn().mockResolvedValue({ + success: false, + code: 'WRAPPER_FINALIZING', + error: 'Wrapper batch is finalizing', + }); + + const result = await flushNextPendingSessionMessage({ + storage, + now: 10, + getDeliveryContext: async () => createContext(), + validateModeAgainstRuntimeAgents: () => null, + deliver, + }); + + expect(result).toEqual({ type: 'held', remainingCount: 1 }); + const [pending] = await listPendingSessionMessages(storage); + expect(pending?.flushAttempts).toBeUndefined(); + expect(pending?.lastFlushError).toBeUndefined(); + }); + it('delivers the next current message without execution-runtime blocking', async () => { const storage = createMemoryStorage(); await storePendingSessionMessage( diff --git a/services/cloud-agent-next/src/session/session-message-queue.ts b/services/cloud-agent-next/src/session/session-message-queue.ts index c4bfb83efb..48d652fce1 100644 --- a/services/cloud-agent-next/src/session/session-message-queue.ts +++ b/services/cloud-agent-next/src/session/session-message-queue.ts @@ -81,7 +81,16 @@ export type PendingFlushDelivered = { remainingCount: number; }; -export type PendingFlushResult = PendingFlushFailure | PendingFlushSkipped | PendingFlushDelivered; +export type PendingFlushHeld = { + type: 'held'; + remainingCount: number; +}; + +export type PendingFlushResult = + | PendingFlushFailure + | PendingFlushSkipped + | PendingFlushDelivered + | PendingFlushHeld; type PersistedQueuedMessageEvent = { sessionId: string; @@ -127,6 +136,7 @@ export type SessionMessageQueueDependencies = { getDeliveryContext: () => Promise; getDeliveryBlock: () => Promise<{ retryAt: number } | null>; deliver: (plan: MessageDeliveryRequest) => Promise; + isDeliveryHeld?: () => Promise; ensureQueuedMessageEvent: (event: PersistedQueuedMessageEvent & { entityId: string }) => void; reportQueuedState?: (state: SessionMessageState) => void; ensureAcceptedMessageEffects: (messageId: string) => Promise; @@ -182,6 +192,7 @@ function classifyDeliveryFailure(code: PendingFlushFailureCode | undefined): { return { failureStage: 'pre_dispatch', failureCode: 'invalid_delivery_request' }; case 'MODEL_MISSING': return { failureStage: 'pre_dispatch', failureCode: 'model_missing' }; + case 'WRAPPER_FINALIZING': case 'INTERNAL': case 'UNKNOWN': case undefined: @@ -274,6 +285,7 @@ export async function flushNextPendingSessionMessage(params: { getDeliveryBlock?: SessionMessageQueueDependencies['getDeliveryBlock']; validateModeAgainstRuntimeAgents: SessionMessageQueueDependencies['validateModeAgainstRuntimeAgents']; deliver: (plan: MessageDeliveryRequest) => Promise; + isDeliveryHeld?: () => Promise; repairQueuedMessageEffects?: (intent: SessionMessageIntent) => Promise; prepareQueuedMessageDelivery?: (intent: SessionMessageIntent) => Promise; ensureAcceptedMessageEffects?: (messageId: string) => Promise; @@ -297,6 +309,9 @@ export async function flushNextPendingSessionMessage(params: { if (!message) { return { type: 'skipped', remainingCount: 0 }; } + if (await params.isDeliveryHeld?.()) { + return { type: 'held', remainingCount: totalCount }; + } const existingState = await getSessionMessageState(params.storage, message.messageId); if (existingState?.status === 'accepted') { @@ -394,6 +409,9 @@ export async function flushNextPendingSessionMessage(params: { params.validateModeAgainstRuntimeAgents ); const startResult = await params.deliver(plan); + if (!startResult.success && startResult.code === 'WRAPPER_FINALIZING') { + return { type: 'held', remainingCount: totalCount }; + } if (!startResult.success) { const failure = await recordPendingFlushFailure( params.storage, @@ -477,6 +495,7 @@ export function createSessionMessageQueue( getDeliveryContext, getDeliveryBlock, deliver, + isDeliveryHeld, ensureQueuedMessageEvent, reportQueuedState, ensureAcceptedMessageEffects, @@ -854,6 +873,7 @@ export function createSessionMessageQueue( getDeliveryBlock, validateModeAgainstRuntimeAgents, deliver, + isDeliveryHeld, repairQueuedMessageEffects: repairQueuedAdmissionEffects, prepareQueuedMessageDelivery, ensureAcceptedMessageEffects, @@ -873,6 +893,10 @@ export function createSessionMessageQueue( }; } + if (flushResult.type === 'held') { + return { remainingPendingCount: flushResult.remainingCount }; + } + if (flushResult.type === 'delivered') { logger .withFields({ diff --git a/services/cloud-agent-next/src/session/session-message-state.test.ts b/services/cloud-agent-next/src/session/session-message-state.test.ts index 4c83bb8400..e696d8c732 100644 --- a/services/cloud-agent-next/src/session/session-message-state.test.ts +++ b/services/cloud-agent-next/src/session/session-message-state.test.ts @@ -10,6 +10,7 @@ import { markMessageInterrupted, terminalizeMessageOnce, listNonTerminalAcceptedMessages, + listMessagesForWrapperRun, listMessagesWithPendingCallbacks, isTerminalMessageState, type SessionMessageState, @@ -19,10 +20,13 @@ import type { SessionMessageIntent } from '../execution/types.js'; function createFakeStorage(): SessionMessageStorage & { store: Map; + listPrefixes: string[]; } { const store = new Map(); + const listPrefixes: string[] = []; return { store, + listPrefixes, async get(key: string): Promise { return store.get(key) as T | undefined; }, @@ -30,6 +34,7 @@ function createFakeStorage(): SessionMessageStorage & { store.set(key, value); }, async list(options: { prefix: string }): Promise> { + listPrefixes.push(options.prefix); const result = new Map(); for (const [key, value] of store.entries()) { if (key.startsWith(options.prefix)) { @@ -613,6 +618,134 @@ describe('listNonTerminalAcceptedMessages', () => { }); }); +describe('listMessagesForWrapperRun', () => { + it('lists accepted and terminal messages fenced to the wrapper run', async () => { + const storage = createFakeStorage(); + const acceptedId = 'msg_0123456789abAAAAAAAAAAAAAA'; + const completedId = 'msg_0123456789abBBBBBBBBBBBBBB'; + const failedId = 'msg_0123456789abCCCCCCCCCCCCCC'; + const queuedId = 'msg_0123456789abDDDDDDDDDDDDDD'; + const otherRunId = 'msg_0123456789abEEEEEEEEEEEEEE'; + + for (const [messageId, status, wrapperRunId] of [ + [acceptedId, 'accepted', 'wr_run1'], + [completedId, 'completed', 'wr_run1'], + [failedId, 'failed', 'wr_run1'], + [queuedId, 'queued', undefined], + [otherRunId, 'completed', 'wr_run2'], + ] as const) { + await putSessionMessageState(storage, { + ...createQueuedSessionMessageState(createIntent(messageId, messageId), undefined, 1000), + status, + acceptedAt: status === 'queued' ? undefined : 2000, + terminalAt: status === 'accepted' || status === 'queued' ? undefined : 3000, + wrapperRunId, + }); + } + + const messages = await listMessagesForWrapperRun(storage, 'wr_run1'); + + expect(messages.map(message => message.messageId)).toEqual([acceptedId, completedId, failedId]); + expect(storage.listPrefixes).toContain('session_message:'); + }); + + it('falls back to history after an older writer drops the live run index version', async () => { + const storage = createFakeStorage(); + const legacyId = 'msg_0123456789abAAAAAAAAAAAAAA'; + const indexedId = 'msg_0123456789abBBBBBBBBBBBBBB'; + + await storage.put('wrapper_runtime_state', { + wrapperGeneration: 1, + wrapperConnectionId: 'conn_live', + wrapperRunId: 'wr_live', + messageIndexVersion: 1, + }); + await putSessionMessageState(storage, { + ...createQueuedSessionMessageState(createIntent(indexedId, 'indexed'), undefined, 1000), + status: 'accepted', + acceptedAt: 2000, + wrapperRunId: 'wr_live', + }); + await storage.put('wrapper_runtime_state', { + wrapperGeneration: 1, + wrapperConnectionId: 'conn_live', + wrapperRunId: 'wr_live', + }); + await storage.put(`session_message:${legacyId}`, { + ...createQueuedSessionMessageState(createIntent(legacyId, 'legacy'), undefined, 1000), + status: 'accepted', + acceptedAt: 2000, + wrapperRunId: 'wr_live', + }); + + const messages = await listMessagesForWrapperRun(storage, 'wr_live'); + + expect(messages.map(message => message.messageId).sort()).toEqual([legacyId, indexedId]); + expect(storage.listPrefixes).toContain('session_message:'); + }); + + it('reads only an initialized wrapper run index and preserves accepted filtering', async () => { + const storage = createFakeStorage(); + const acceptedId = 'msg_0123456789abAAAAAAAAAAAAAA'; + const completedId = 'msg_0123456789abBBBBBBBBBBBBBB'; + const historicalId = 'msg_0123456789abCCCCCCCCCCCCCC'; + await storage.put('wrapper_runtime_state', { + wrapperGeneration: 1, + wrapperConnectionId: 'conn_current', + wrapperRunId: 'wr_current', + messageIndexVersion: 1, + }); + + for (const [messageId, status, wrapperRunId] of [ + [acceptedId, 'accepted', 'wr_current'], + [completedId, 'completed', 'wr_current'], + [historicalId, 'completed', 'wr_historical'], + ] as const) { + await putSessionMessageState(storage, { + ...createQueuedSessionMessageState(createIntent(messageId, messageId), undefined, 1000), + status, + acceptedAt: 2000, + terminalAt: status === 'completed' ? 3000 : undefined, + wrapperRunId, + }); + } + storage.listPrefixes.length = 0; + + await expect(listNonTerminalAcceptedMessages(storage, 'wr_current')).resolves.toMatchObject([ + { messageId: acceptedId, status: 'accepted' }, + ]); + expect(storage.listPrefixes).toEqual(['session_message_wrapper_run:wr_current:']); + }); + + it('does not expose an indexed member whose primary state write failed', async () => { + const storage = createFakeStorage(); + const messageId = 'msg_0123456789abAAAAAAAAAAAAAA'; + await storage.put('wrapper_runtime_state', { + wrapperGeneration: 1, + wrapperConnectionId: 'conn_current', + wrapperRunId: 'wr_current', + messageIndexVersion: 1, + }); + const put = storage.put.bind(storage); + storage.put = async (key, value) => { + if (key === `session_message:${messageId}`) throw new Error('primary write failed'); + await put(key, value); + }; + + await expect( + putSessionMessageState(storage, { + ...createQueuedSessionMessageState(createIntent(messageId, 'message'), undefined, 1000), + status: 'accepted', + acceptedAt: 2000, + wrapperRunId: 'wr_current', + }) + ).rejects.toThrow('primary write failed'); + storage.put = put; + + await expect(listMessagesForWrapperRun(storage, 'wr_current')).resolves.toEqual([]); + }); +}); + describe('listMessagesWithPendingCallbacks', () => { it('lists terminal messages with callbackRequired but no callbackEnqueuedAt', async () => { const storage = createFakeStorage(); diff --git a/services/cloud-agent-next/src/session/session-message-state.ts b/services/cloud-agent-next/src/session/session-message-state.ts index 821d11e9bf..4e5ba6ad21 100644 --- a/services/cloud-agent-next/src/session/session-message-state.ts +++ b/services/cloud-agent-next/src/session/session-message-state.ts @@ -5,8 +5,13 @@ import type { ExecutionMode, SessionMessageIntent } from '../execution/types.js' import { renderExecutionTurnContent } from '../execution/types.js'; import { AttachmentsSchema } from '../persistence/schemas.js'; import { MESSAGE_ID_FORMAT_DESCRIPTION, MESSAGE_ID_PATTERN } from './message-id.js'; +import { + getWrapperRuntimeState, + hasCompleteWrapperRunMessageIndex, +} from './wrapper-runtime-state.js'; const SESSION_MESSAGE_STATE_PREFIX = 'session_message:'; +const WRAPPER_RUN_MESSAGE_INDEX_PREFIX = 'session_message_wrapper_run:'; export type SessionMessageStatus = 'queued' | 'accepted' | 'completed' | 'failed' | 'interrupted'; @@ -269,6 +274,14 @@ function sessionMessageStateKey(messageId: string): string { return `${SESSION_MESSAGE_STATE_PREFIX}${messageId}`; } +function wrapperRunMessageIndexPrefix(wrapperRunId: string): string { + return `${WRAPPER_RUN_MESSAGE_INDEX_PREFIX}${encodeURIComponent(wrapperRunId)}:`; +} + +function wrapperRunMessageIndexKey(wrapperRunId: string, messageId: string): string { + return `${wrapperRunMessageIndexPrefix(wrapperRunId)}${messageId}`; +} + function normalizeLegacyAdmissionConstraints( constraints: z.infer['legacyAdmissionConstraints'] ): LegacyAdmissionConstraints | undefined { @@ -404,10 +417,14 @@ export async function putSessionMessageState( storage: SessionMessageStorage, state: SessionMessageState ): Promise { - await storage.put( - sessionMessageStateKey(state.messageId), - SessionMessageStateSchema.parse(state) - ); + const parsedState = SessionMessageStateSchema.parse(state); + if (parsedState.wrapperRunId) { + await storage.put( + wrapperRunMessageIndexKey(parsedState.wrapperRunId, parsedState.messageId), + true + ); + } + await storage.put(sessionMessageStateKey(parsedState.messageId), parsedState); } export function createQueuedSessionMessageState( @@ -564,12 +581,40 @@ export async function listNonTerminalAcceptedMessages( storage: SessionMessageStorage, wrapperRunId?: string ): Promise { - const entries = await listSessionMessageStates(storage); - return entries.filter( - state => - state.status === 'accepted' && - (wrapperRunId === undefined || state.wrapperRunId === wrapperRunId) + const entries = + wrapperRunId === undefined + ? await listSessionMessageStates(storage) + : await listMessagesForWrapperRun(storage, wrapperRunId); + return entries.filter(state => state.status === 'accepted'); +} + +async function listIndexedMessagesForWrapperRun( + storage: SessionMessageStorage, + wrapperRunId: string +): Promise { + const indexPrefix = wrapperRunMessageIndexPrefix(wrapperRunId); + const index = await storage.list({ prefix: indexPrefix }); + const messages = await Promise.all( + Array.from(index.keys()).map(async key => { + const messageId = key.slice(indexPrefix.length); + if (!messageId) return undefined; + const state = await getSessionMessageState(storage, messageId); + return state?.wrapperRunId === wrapperRunId ? state : undefined; + }) ); + return messages.filter(state => state !== undefined); +} + +export async function listMessagesForWrapperRun( + storage: SessionMessageStorage, + wrapperRunId: string +): Promise { + if (hasCompleteWrapperRunMessageIndex(await getWrapperRuntimeState(storage), wrapperRunId)) { + return listIndexedMessagesForWrapperRun(storage, wrapperRunId); + } + + const entries = await listSessionMessageStates(storage); + return entries.filter(state => state.wrapperRunId === wrapperRunId); } type NeverAcceptedTerminalQueuedMessageState = SessionMessageState & { diff --git a/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts b/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts index 0cd3299d49..5863127189 100644 --- a/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts +++ b/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts @@ -1,10 +1,18 @@ import { describe, expect, it } from 'vitest'; import { + allocateWrapperRuntimeState, + clearCurrentWrapperRuntimeLivenessState, emptyWrapperLease, getWrapperLease, + getWrapperRuntimeState, + hasCompleteWrapperRunMessageIndex, + isWrapperDeliveryHeld, + markWrapperFinalizing, nextWrapperCleanupDeadline, nextWrapperLeaseDeadline, putWrapperLease, + recordMeaningfulWrapperOutput, + recordWrapperAcceptedMessage, reduceWrapperLease, } from './wrapper-runtime-state.js'; @@ -197,6 +205,146 @@ describe('WrapperLease', () => { expect(nextWrapperLeaseDeadline(retrying)).toBe(5_200); }); + it('marks a newly allocated wrapper run as maintaining its message index', async () => { + const storage = createMemoryStorage(); + + const { state } = await allocateWrapperRuntimeState(storage, 1_000); + + expect(state.messageIndexVersion).toBe(1); + await recordWrapperAcceptedMessage(storage, state, 5_000, 4_000); + await expect(getWrapperRuntimeState(storage)).resolves.toMatchObject({ + messageIndexVersion: 1, + }); + }); + + it('preserves the message index version while clearing current liveness deadlines', async () => { + const storage = createMemoryStorage(); + const { state } = await allocateWrapperRuntimeState(storage, 1_000); + await storage.put('wrapper_runtime_state', { + ...state, + nextPingAt: 2_000, + noOutputDeadlineAt: 3_000, + }); + + await clearCurrentWrapperRuntimeLivenessState( + storage, + state.wrapperGeneration, + state.wrapperConnectionId + ); + + await expect(getWrapperRuntimeState(storage)).resolves.toMatchObject({ + wrapperRunId: state.wrapperRunId, + messageIndexVersion: 1, + }); + }); + + it('preserves runtime identity while treating a future message index version as untrusted', async () => { + const storage = createMemoryStorage(); + await storage.put('wrapper_runtime_state', { + wrapperGeneration: 2, + wrapperConnectionId: 'conn_future_index', + wrapperRunId: 'wr_future_index', + messageIndexVersion: 2, + }); + + const state = await getWrapperRuntimeState(storage); + + expect(state).toMatchObject({ + wrapperGeneration: 2, + wrapperConnectionId: 'conn_future_index', + wrapperRunId: 'wr_future_index', + messageIndexVersion: 2, + }); + expect(hasCompleteWrapperRunMessageIndex(state, 'wr_future_index')).toBe(false); + }); + + it('persists finalizing only for the matching current wrapper run', async () => { + const storage = createMemoryStorage(); + await storage.put('wrapper_runtime_state', { + wrapperGeneration: 2, + wrapperConnectionId: 'conn_finalizing', + wrapperRunId: 'wr_finalizing', + }); + + await expect(markWrapperFinalizing(storage, 'wr_stale')).resolves.toBeNull(); + await expect(markWrapperFinalizing(storage, 'wr_finalizing')).resolves.toMatchObject({ + finalizingWrapperRunId: 'wr_finalizing', + }); + }); + + it('preserves run-level finalizing while refreshing liveness output', async () => { + const storage = createMemoryStorage(); + await storage.put('wrapper_runtime_state', { + wrapperGeneration: 2, + wrapperConnectionId: 'conn_housekeeping', + wrapperRunId: 'wr_housekeeping', + finalizingWrapperRunId: 'wr_housekeeping', + wrapperIdleDeadlineAt: 3_000, + }); + + await recordMeaningfulWrapperOutput(storage, 2, 'conn_housekeeping', 1_500, 4_000, 5_000); + + await expect(getWrapperRuntimeState(storage)).resolves.toMatchObject({ + finalizingWrapperRunId: 'wr_housekeeping', + lastWrapperMessageAt: 1_500, + noOutputDeadlineAt: 5_000, + nextPingAt: 4_000, + wrapperIdleDeadlineAt: 3_000, + }); + }); + + it('does not clear finalizing when an acceptance write races after complete', async () => { + const storage = createMemoryStorage(); + const allocated = { + wrapperGeneration: 2, + wrapperConnectionId: 'conn_new_work', + wrapperRunId: 'wr_new_work', + finalizingWrapperRunId: 'wr_new_work', + wrapperIdleDeadlineAt: 3_000, + }; + await storage.put('wrapper_runtime_state', allocated); + + await recordWrapperAcceptedMessage(storage, allocated, 5_000, 4_000); + + await expect(getWrapperRuntimeState(storage)).resolves.toMatchObject({ + finalizingWrapperRunId: 'wr_new_work', + noOutputDeadlineAt: 5_000, + nextPingAt: 4_000, + }); + }); + + it.each(['stop_needed', 'stopping'] as const)( + 'holds delivery while physical cleanup is %s', + state => { + const runtime = { wrapperGeneration: 2 }; + const lease = + state === 'stop_needed' + ? reduceWrapperLease(emptyWrapperLease(), { + type: 'request_stop', + target: { kind: 'session' }, + reason: 'terminal-failed', + now: 1_000, + }) + : reduceWrapperLease( + reduceWrapperLease(emptyWrapperLease(), { + type: 'request_stop', + target: { kind: 'session' }, + reason: 'terminal-failed', + now: 1_000, + }), + { + type: 'begin_stop_attempt', + attemptId: 'attempt', + now: 1_000, + attemptDeadlineAt: 2_000, + } + ); + + expect(isWrapperDeliveryHeld(runtime, lease)).toBe(true); + expect(isWrapperDeliveryHeld(runtime, emptyWrapperLease())).toBe(false); + } + ); + it('validates the separately persisted physical ownership record', async () => { const storage = createMemoryStorage(); await expect(getWrapperLease(storage)).resolves.toEqual(emptyWrapperLease()); 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 e5191604ee..f0f8e94cf1 100644 --- a/services/cloud-agent-next/src/session/wrapper-runtime-state.ts +++ b/services/cloud-agent-next/src/session/wrapper-runtime-state.ts @@ -6,6 +6,7 @@ import type { } from '../agent-sandbox/protocol.js'; const WRAPPER_RUNTIME_STATE_KEY = 'wrapper_runtime_state'; +const WRAPPER_RUN_MESSAGE_INDEX_VERSION = 1; const WRAPPER_LEASE_KEY = 'wrapper_lease'; const wrapperInstanceLeaseSchema = z.object({ @@ -196,7 +197,6 @@ export function nextWrapperLeaseDeadline(lease: WrapperLease): number | undefine return lease.startupDeadlineAt ?? lease.keepWarmUntil; } -export const IDLE_RECONCILIATION_GRACE_MS = 15_000; export const IDLE_KEEP_WARM_MS = 5 * 60 * 1000; export const READY_ONLY_IDLE_MS = 60_000; @@ -204,11 +204,12 @@ const wrapperRuntimeStateSchema = z.object({ wrapperGeneration: z.number().int().nonnegative(), wrapperConnectionId: z.string().optional(), wrapperRunId: z.string().optional(), + messageIndexVersion: z.number().int().nonnegative().optional(), + dispatchingMessageId: z.string().optional(), lastWrapperConnectedAt: z.number().int().nonnegative().optional(), lastWrapperMessageAt: z.number().int().nonnegative().optional(), lastWrapperPongAt: z.number().int().nonnegative().optional(), - lastWrapperIdleAt: z.number().int().nonnegative().optional(), - idleReconcileAfter: z.number().int().nonnegative().optional(), + finalizingWrapperRunId: z.string().optional(), wrapperIdleDeadlineAt: z.number().int().nonnegative().optional(), pingDeadlineAt: z.number().int().nonnegative().optional(), nextPingAt: z.number().int().nonnegative().optional(), @@ -228,6 +229,16 @@ export function isActiveWrapperRuntimeState( return Boolean(state.wrapperConnectionId && state.wrapperRunId); } +export function hasCompleteWrapperRunMessageIndex( + state: WrapperRuntimeState, + wrapperRunId: string +): boolean { + return ( + state.wrapperRunId === wrapperRunId && + state.messageIndexVersion === WRAPPER_RUN_MESSAGE_INDEX_VERSION + ); +} + export function hasCompleteWrapperIdentity(state: WrapperRuntimeState): boolean { const hasIdentityField = Boolean(state.wrapperConnectionId || state.wrapperRunId); return !hasIdentityField || Boolean(state.wrapperConnectionId && state.wrapperRunId); @@ -237,8 +248,12 @@ export const emptyWrapperRuntimeState = (): WrapperRuntimeState => ({ wrapperGeneration: 0, }); +type WrapperRuntimeStateReader = { + get(key: string): Promise; +}; + export async function getWrapperRuntimeState( - storage: DurableObjectStorage + storage: WrapperRuntimeStateReader ): Promise { const stored = await storage.get(WRAPPER_RUNTIME_STATE_KEY); const parsed = wrapperRuntimeStateSchema.safeParse(stored); @@ -277,6 +292,7 @@ export async function allocateWrapperRuntimeState( wrapperGeneration: current.wrapperGeneration + 1, wrapperConnectionId: crypto.randomUUID(), wrapperRunId: `wr_${crypto.randomUUID().replace(/-/g, '')}`, + messageIndexVersion: WRAPPER_RUN_MESSAGE_INDEX_VERSION, lastWrapperConnectedAt: now, } satisfies ActiveWrapperRuntimeState; await storage.put(WRAPPER_RUNTIME_STATE_KEY, next); @@ -371,6 +387,37 @@ async function updateIfCurrent( return next; } +export async function recordWrapperDispatchingMessage( + storage: DurableObjectStorage, + allocated: ActiveWrapperRuntimeState, + messageId: string +): Promise { + await updateIfCurrent( + storage, + allocated.wrapperGeneration, + allocated.wrapperConnectionId, + current => ({ ...current, dispatchingMessageId: messageId }) + ); +} + +export async function clearWrapperDispatchingMessage( + storage: DurableObjectStorage, + allocated: ActiveWrapperRuntimeState, + messageId: string +): Promise { + await updateIfCurrent( + storage, + allocated.wrapperGeneration, + allocated.wrapperConnectionId, + current => { + if (current.dispatchingMessageId !== messageId) return current; + const next = { ...current }; + delete next.dispatchingMessageId; + return next; + } + ); +} + export async function recordWrapperAcceptedMessage( storage: DurableObjectStorage, allocated: ActiveWrapperRuntimeState, @@ -388,8 +435,6 @@ export async function recordWrapperAcceptedMessage( noOutputDeadlineAt, nextPingAt: current.pingDeadlineAt === undefined ? (current.nextPingAt ?? nextPingAt) : undefined, - lastWrapperIdleAt: undefined, - idleReconcileAfter: undefined, wrapperIdleDeadlineAt: undefined, }) ); @@ -429,44 +474,30 @@ export async function recordWrapperPong( })); } -/** - * Record that the wrapper received a root session.idle event. - * Stores the idle timestamp, the reconciliation deadline, and the - * keep-warm cleanup deadline. - */ -export async function recordRootSessionIdle( +export async function markWrapperFinalizing( storage: DurableObjectStorage, - wrapperGeneration: number, - wrapperConnectionId: string, - now = Date.now(), - idleReconcileAfter = now + IDLE_RECONCILIATION_GRACE_MS, - wrapperIdleDeadlineAt = now + IDLE_KEEP_WARM_MS + wrapperRunId: string ): Promise { - return updateIfCurrent(storage, wrapperGeneration, wrapperConnectionId, current => ({ - ...current, - lastWrapperIdleAt: now, - idleReconcileAfter, - wrapperIdleDeadlineAt, - })); + const current = await getWrapperRuntimeState(storage); + if (current.wrapperRunId !== wrapperRunId) return null; + if (current.finalizingWrapperRunId === wrapperRunId) return current; + + const next = { ...current, finalizingWrapperRunId: wrapperRunId } satisfies WrapperRuntimeState; + await storage.put(WRAPPER_RUNTIME_STATE_KEY, next); + return next; } -/** - * Record a meaningful output event from the wrapper. - * - * Refreshes `noOutputDeadlineAt` so mid-execution stalls are caught: without - * this, the deadline would be cleared forever after the first event, and a - * wrapper whose kilo-server SSE subscription silently stalls would remain - * live without failing its accepted messages. - * - * This intentionally does NOT clear the idle-reconciliation fields - * (`lastWrapperIdleAt`/`idleReconcileAfter`/`wrapperIdleDeadlineAt`). Those are - * only ever armed by `recordRootSessionIdle` once the root session goes idle, - * so any wrapper output observed afterwards is post-completion infrastructure - * work (autocommit, condense, log upload) — not a new agent turn. Clearing the - * idle deadline here would disarm the reconciler that finalizes the in-flight - * message, stranding it non-terminal and hanging the callback. A genuinely new - * turn clears idle via `recordWrapperAcceptedMessage` instead. - */ +export function isWrapperRunFinalizing(state: WrapperRuntimeState): boolean { + return Boolean(state.wrapperRunId && state.finalizingWrapperRunId === state.wrapperRunId); +} + +export function isWrapperDeliveryHeld(state: WrapperRuntimeState, lease: WrapperLease): boolean { + return ( + isWrapperRunFinalizing(state) || lease.state === 'stop_needed' || lease.state === 'stopping' + ); +} + +/** Record meaningful output while accepted work remains supervised. */ export async function recordMeaningfulWrapperOutput( storage: DurableObjectStorage, wrapperGeneration: number, @@ -505,28 +536,16 @@ export async function clearCurrentWrapperRuntimeLivenessState( wrapperGeneration: current.wrapperGeneration, wrapperConnectionId: current.wrapperConnectionId, wrapperRunId: current.wrapperRunId, + messageIndexVersion: current.messageIndexVersion, + dispatchingMessageId: current.dispatchingMessageId, lastWrapperConnectedAt: current.lastWrapperConnectedAt, lastWrapperMessageAt: current.lastWrapperMessageAt, lastWrapperPongAt: current.lastWrapperPongAt, - lastWrapperIdleAt: current.lastWrapperIdleAt, - idleReconcileAfter: current.idleReconcileAfter, + finalizingWrapperRunId: current.finalizingWrapperRunId, wrapperIdleDeadlineAt: current.wrapperIdleDeadlineAt, })); } -export async function clearWrapperIdleState( - storage: DurableObjectStorage, - wrapperGeneration: number, - wrapperConnectionId: string -): Promise { - return updateIfCurrent(storage, wrapperGeneration, wrapperConnectionId, current => ({ - ...current, - lastWrapperIdleAt: undefined, - idleReconcileAfter: undefined, - wrapperIdleDeadlineAt: undefined, - })); -} - export async function clearCurrentWrapperRuntimeFailureState( storage: DurableObjectStorage, wrapperGeneration: number, 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 2ac07ab53d..a981375b32 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts @@ -5,6 +5,7 @@ import { createMessageSettlementOutbox, type MessageSettlementOutboxStorage, } from './message-settlement-outbox.js'; +import { storePendingSessionMessage } from './pending-messages.js'; import { getSessionMessageState, putSessionMessageState, @@ -103,17 +104,22 @@ function createHarness( initialEntries?: Array<[string, unknown]>, options?: { metadata?: SessionMetadata; + storage?: MemoryStorage; storageHooks?: { beforeList?: (prefix: string) => Promise }; getAssistantMessageForUserMessage?: ( sessionId: string, kiloSessionId: string, parentMessageId: string ) => LatestAssistantMessage | null; + ensureAcceptedMessageBeforeTerminal?: ( + messageId: string, + wrapperRunId: string + ) => Promise; } ) { const getAssistantMessageForUserMessage = options?.getAssistantMessageForUserMessage ?? (() => null); - const storage = createMemoryStorage(initialEntries, options?.storageHooks); + const storage = options?.storage ?? createMemoryStorage(initialEntries, options?.storageHooks); const events: MessageEvent[] = []; const callbackJobs: CallbackJob[] = []; const sentPings: string[] = []; @@ -155,6 +161,8 @@ function createHarness( getAssistantMessageForUserMessage, hasActiveIngestConnection: async () => false, clearInterruptRequest: async () => {}, + ensureAcceptedMessageBeforeTerminal: + options?.ensureAcceptedMessageBeforeTerminal ?? (async () => {}), stopWrappers, requestAlarmAtOrBefore: async deadline => { requestedAlarms.push(deadline); @@ -226,6 +234,37 @@ describe('WrapperSupervisor', () => { await expect(harness.storage.get('disconnect_grace')).resolves.toBeUndefined(); }); + it('starts disconnect grace for a finalizing current run without accepted messages', async () => { + const harness = createHarness([ + liveRuntimeState({ finalizingWrapperRunId: WRAPPER_RUN_ID }), + OWNED_WRAPPER_LEASE, + ]); + + await harness.supervisor.onDisconnected({ + disconnected: { + wrapperRunId: WRAPPER_RUN_ID, + wrapperGeneration: 4, + wrapperConnectionId: WRAPPER_CONNECTION_ID, + }, + wsCloseCode: 1006, + wsCloseReason: 'socket closed while finalizing', + }); + + const grace = await harness.storage.get<{ disconnectedAt: number }>('disconnect_grace'); + expect(grace).toBeDefined(); + if (!grace) throw new Error('Expected finalizing disconnect grace'); + + await harness.supervisor.runMaintenance(grace.disconnectedAt + 10_001); + + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'unhealthy-wrapper', + }); + await expect(getWrapperRuntimeState(harness.storage)).resolves.toEqual({ + wrapperGeneration: 5, + }); + }); + it('starts disconnect grace while a completed gate callback still waits for wrapper terminal state', async () => { const harness = createHarness([liveRuntimeState()], { metadata: { @@ -257,8 +296,8 @@ describe('WrapperSupervisor', () => { expect(harness.callbackJobs).toHaveLength(0); }); - it('releases a gate-waiting callback after disconnect grace expires without a terminal event', async () => { - const harness = createHarness([liveRuntimeState()], { + it('releases a finalizing run callback after disconnect grace despite a queued follow-up', async () => { + const harness = createHarness([liveRuntimeState({ finalizingWrapperRunId: WRAPPER_RUN_ID })], { metadata: { ...createMetadata(), finalization: { gateThreshold: 'warning' }, @@ -273,6 +312,15 @@ describe('WrapperSupervisor', () => { kind: 'completed', completionSource: 'assistant_message_event', }); + await storePendingSessionMessage(harness.storage, { + messageId: NEWER_MESSAGE_ID, + content: 'queued follow-up', + createdAt: 4_000, + intent: { + turn: { type: 'prompt', messageId: NEWER_MESSAGE_ID, prompt: 'queued follow-up' }, + agent: { mode: 'code', model: 'test-model' }, + }, + }); await harness.supervisor.onDisconnected({ disconnected: { wrapperRunId: WRAPPER_RUN_ID, @@ -520,7 +568,42 @@ describe('WrapperSupervisor', () => { }, }); expect(harness.stops).toEqual([]); - expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); + expect(harness.requestPendingDrainIfNeeded).not.toHaveBeenCalled(); + } + ); + + it.each([ + { status: 'failed' as const, expected: 'failed' as const }, + { status: 'interrupted' as const, expected: 'interrupted' as const }, + ])( + 'repairs dispatching work before retiring a current $status terminal event', + async ({ status, expected }) => { + const storage = createMemoryStorage([ + liveRuntimeState({ dispatchingMessageId: MESSAGE_ID }), + OWNED_WRAPPER_LEASE, + ]); + const ensureAccepted = vi.fn(async (messageId: string, wrapperRunId: string) => { + await putSessionMessageState(storage, { + ...acceptedMessage(messageId), + wrapperRunId, + dispatchAcceptanceKind: 'inferred_from_terminal', + }); + }); + const harness = createHarness(undefined, { + storage, + ensureAcceptedMessageBeforeTerminal: ensureAccepted, + }); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status, + error: 'terminal event', + }); + + expect(ensureAccepted).toHaveBeenCalledWith(MESSAGE_ID, WRAPPER_RUN_ID); + await expect(getSessionMessageState(storage, MESSAGE_ID)).resolves.toMatchObject({ + status: expected, + }); } ); @@ -586,12 +669,76 @@ describe('WrapperSupervisor', () => { expect(harness.events.map(event => event.streamEventType)).toEqual(['cloud.message.failed']); }); - it('aggregates concurrent physical, liveness, disconnect, and idle deadlines', async () => { + it('releases a finalizing callback on liveness expiry but holds a queued follow-up until physical absence', async () => { + const harness = createHarness( + [ + liveRuntimeState({ + finalizingWrapperRunId: WRAPPER_RUN_ID, + noOutputDeadlineAt: 9_000, + nextPingAt: 30_000, + }), + OWNED_WRAPPER_LEASE, + ], + { + metadata: { + ...createMetadata(), + finalization: { gateThreshold: 'warning' }, + }, + } + ); + await putSessionMessageState(harness.storage, { + ...acceptedMessage(), + status: 'completed', + terminalAt: 3_000, + completionSource: 'assistant_message_event', + callbackRequired: true, + callbackTarget: { url: 'https://example.com/finalizing-liveness-release' }, + }); + await harness.settlementOutbox.terminalizeSessionMessageOnce(MESSAGE_ID, { + kind: 'completed', + completionSource: 'assistant_message_event', + }); + await storePendingSessionMessage(harness.storage, { + messageId: NEWER_MESSAGE_ID, + content: 'queued follow-up', + createdAt: 4_000, + intent: { + turn: { type: 'prompt', messageId: NEWER_MESSAGE_ID, prompt: 'queued follow-up' }, + agent: { mode: 'code', model: 'test-model' }, + }, + }); + + await harness.supervisor.runMaintenance(10_000); + + expect(harness.callbackJobs).toHaveLength(1); + expect(harness.callbackJobs[0].payload).toMatchObject({ + messageId: MESSAGE_ID, + status: 'completed', + }); + expect(harness.callbackJobs[0].payload.gateResult).toBeUndefined(); + const cleanupHold = await getWrapperLease(harness.storage); + expect(cleanupHold).toMatchObject({ + state: 'stop_needed', + reason: 'unhealthy-wrapper', + }); + if (cleanupHold.state !== 'stop_needed') throw new Error('Expected physical cleanup hold'); + expect(harness.requestPendingDrainIfNeeded).not.toHaveBeenCalled(); + + await harness.supervisor.runMaintenance(cleanupHold.nextAttemptAt); + + expect(harness.stopWrappers).toHaveBeenCalledOnce(); + await expect(getWrapperLease(harness.storage)).resolves.toEqual({ + state: 'none', + nextInstanceGeneration: 2, + }); + expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); + }); + + it('aggregates concurrent physical, liveness, disconnect, and cleanup deadlines', async () => { const harness = createHarness([ liveRuntimeState({ nextPingAt: 20_000, noOutputDeadlineAt: 50_000, - idleReconcileAfter: 30_000, wrapperIdleDeadlineAt: 40_000, }), [ @@ -619,120 +766,585 @@ describe('WrapperSupervisor', () => { const deadlines = await harness.supervisor.nextMaintenanceDeadlines(); - expect(deadlines).toHaveLength(5); - expect(deadlines).toEqual(expect.arrayContaining([60_000, 20_000, 15_000, 30_000, 40_000])); + expect(deadlines).toHaveLength(4); + expect(deadlines).toEqual(expect.arrayContaining([60_000, 20_000, 15_000, 40_000])); expect(Math.min(...deadlines)).toBe(15_000); }); - it('reconciles accepted idle work after its root-idle deadline', async () => { + it('persists finalizing only for the current wrapper run', async () => { + const harness = createHarness([liveRuntimeState()]); + + await harness.supervisor.observeFinalizing('wr_stale'); + await expect(getWrapperRuntimeState(harness.storage)).resolves.not.toHaveProperty( + 'finalizingWrapperRunId' + ); + + await harness.supervisor.observeFinalizing(WRAPPER_RUN_ID); + await expect(getWrapperRuntimeState(harness.storage)).resolves.toMatchObject({ + finalizingWrapperRunId: WRAPPER_RUN_ID, + }); + }); + + it('raw idle maintenance never settles accepted work', async () => { + const harness = createHarness([liveRuntimeState({ wrapperIdleDeadlineAt: 50_000 })]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + await harness.supervisor.runMaintenance(10_000); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'accepted', + }); + }); + + it('settles a still-accepted successful reply when the wrapper completes', async () => { + // Wrapper complete is the normal turn boundary, independent of assistant + // message completion markers emitted during the turn. + const assistantMessageId = 'ase_complete_reconcile'; + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE], { + getAssistantMessageForUserMessage: () => + ({ + info: { id: assistantMessageId, role: 'assistant' }, + parts: [], + }) as unknown as LatestAssistantMessage, + }); + await putSessionMessageState(harness.storage, { + ...acceptedMessage(), + callbackRequired: true, + callbackTarget: { url: 'https://example.com/complete-reconcile' }, + }); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + messageIds: [MESSAGE_ID], + }); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'completed', + completionSource: 'idle_reconciliation', + assistantMessageId, + }); + expect(harness.callbackJobs).toHaveLength(1); + expect(harness.callbackJobs[0].payload).toMatchObject({ + messageId: MESSAGE_ID, + status: 'completed', + }); + expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); + }); + + it('settles current accepted work from a legacy complete without sealed membership', async () => { + const assistantMessageId = 'ase_legacy_complete'; + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE], { + getAssistantMessageForUserMessage: () => + ({ + info: { id: assistantMessageId, role: 'assistant' }, + parts: [], + }) as unknown as LatestAssistantMessage, + }); + await putSessionMessageState(harness.storage, acceptedMessage()); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + }); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'completed', + completionSource: 'idle_reconciliation', + assistantMessageId, + }); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'owns_wrapper', + keepWarmUntil: expect.any(Number), + }); + }); + + it('rejects missing sealed membership from a current indexed wrapper run', async () => { const harness = createHarness([ - liveRuntimeState({ - lastWrapperIdleAt: 1_000, - idleReconcileAfter: 9_000, - wrapperIdleDeadlineAt: 50_000, - }), + liveRuntimeState({ messageIndexVersion: 1 }), + OWNED_WRAPPER_LEASE, ]); await putSessionMessageState(harness.storage, acceptedMessage()); - await harness.supervisor.runMaintenance(10_000); + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + }); await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ status: 'failed', - failureReason: 'missing_assistant_reply', - error: 'No assistant reply found during reconciliation', - completionSource: 'idle_reconciliation', - failureStage: 'post_dispatch_no_activity', - failureCode: 'missing_assistant_reply', + failureReason: 'wrapper_protocol_error', + }); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'terminal-failed', }); }); - it('does not clear idle reconciliation fields after wrapper identity changes', async () => { - let replacedRuntimeState = false; - const storageForHook: { current?: MemoryStorage } = {}; + it('deduplicates accepted dispatching membership from a legacy complete', async () => { const harness = createHarness( - [ - liveRuntimeState({ - lastWrapperIdleAt: 1_000, - idleReconcileAfter: 9_000, - wrapperIdleDeadlineAt: 50_000, - }), - ], + [liveRuntimeState({ dispatchingMessageId: MESSAGE_ID }), OWNED_WRAPPER_LEASE], { - storageHooks: { - beforeList: async prefix => { - if (replacedRuntimeState || !prefix.startsWith('session_message:')) return; - replacedRuntimeState = true; - await storageForHook.current?.put('wrapper_runtime_state', { - wrapperGeneration: 5, - wrapperConnectionId: 'conn_replacement', - wrapperRunId: WRAPPER_RUN_ID, - lastWrapperIdleAt: 2_000, - idleReconcileAfter: 12_000, - wrapperIdleDeadlineAt: 60_000, - }); - }, - }, + getAssistantMessageForUserMessage: () => + ({ + info: { id: 'ase_legacy_accepted_race', role: 'assistant' }, + parts: [], + }) as unknown as LatestAssistantMessage, } ); - storageForHook.current = harness.storage; + await putSessionMessageState(harness.storage, acceptedMessage()); - await harness.supervisor.runMaintenance(10_000); + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + }); - await expect(getWrapperRuntimeState(harness.storage)).resolves.toMatchObject({ - wrapperGeneration: 5, - wrapperConnectionId: 'conn_replacement', - lastWrapperIdleAt: 2_000, - idleReconcileAfter: 12_000, - wrapperIdleDeadlineAt: 60_000, + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'completed', + assistantMessageId: 'ase_legacy_accepted_race', + }); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'owns_wrapper', + keepWarmUntil: expect.any(Number), }); }); - it('settles a still-accepted message and enqueues its callback when the wrapper completes without a terminal assistant event', async () => { - // Reproduces the webhook hang: the final assistant message.updated never - // carried time.completed (so assistant_message_event never settled it), and - // post-idle autocommit refreshed liveness while clearing the idle-reconcile - // deadline (idle fields absent below). The race-free `complete` event must - // still terminalize the message and release the gated callback. - const assistantMessageId = 'ase_complete_reconcile'; + it('repairs legacy complete-before-acceptance from the dispatching message fence', async () => { + const storage = createMemoryStorage([ + liveRuntimeState({ dispatchingMessageId: MESSAGE_ID }), + OWNED_WRAPPER_LEASE, + ]); + const ensureAccepted = vi.fn(async (messageId: string, wrapperRunId: string) => { + await putSessionMessageState(storage, { + ...acceptedMessage(messageId), + wrapperRunId, + dispatchAcceptanceKind: 'inferred_from_terminal', + }); + }); + const harness = createHarness(undefined, { + storage, + getAssistantMessageForUserMessage: () => + ({ + info: { id: 'ase_legacy_race_reply', role: 'assistant' }, + parts: [], + }) as unknown as LatestAssistantMessage, + ensureAcceptedMessageBeforeTerminal: ensureAccepted, + }); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + }); + + expect(ensureAccepted).toHaveBeenCalledWith(MESSAGE_ID, WRAPPER_RUN_ID); + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'completed', + assistantMessageId: 'ase_legacy_race_reply', + }); + }); + + it('repairs complete-before-acceptance from durable pending intent', async () => { const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE], { getAssistantMessageForUserMessage: () => ({ - info: { id: assistantMessageId, role: 'assistant' }, + info: { id: 'ase_race_reply', role: 'assistant' }, parts: [], }) as unknown as LatestAssistantMessage, }); await putSessionMessageState(harness.storage, { ...acceptedMessage(), - callbackRequired: true, - callbackTarget: { url: 'https://example.com/complete-reconcile' }, + status: 'queued', + acceptedAt: undefined, + wrapperRunId: undefined, + }); + const ensureAccepted = vi.fn(async (messageId: string, wrapperRunId: string) => { + await putSessionMessageState(harness.storage, { + ...acceptedMessage(messageId), + wrapperRunId, + dispatchAcceptanceKind: 'inferred_from_terminal', + }); + }); + const supervisor = createWrapperSupervisor({ + storage: harness.storage, + agentRuntime: { sendPing: () => {} }, + messageSettlementOutbox: harness.settlementOutbox, + sessionMessageQueue: { requestPendingDrainIfNeeded: harness.requestPendingDrainIfNeeded }, + getMetadata: async () => createMetadata(), + getAssistantMessageForUserMessage: () => + ({ + info: { id: 'ase_race_reply', role: 'assistant' }, + parts: [], + }) as unknown as LatestAssistantMessage, + hasActiveIngestConnection: async () => false, + clearInterruptRequest: async () => {}, + ensureAcceptedMessageBeforeTerminal: ensureAccepted, + requestAlarmAtOrBefore: async () => {}, + getSessionIdForLogs: () => 'agent_supervisor', }); + await supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + messageIds: [MESSAGE_ID], + }); + + expect(ensureAccepted).toHaveBeenCalledWith(MESSAGE_ID, WRAPPER_RUN_ID); + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'completed', + assistantMessageId: 'ase_race_reply', + }); + }); + + it('fails a command-free prompt with no reply when the wrapper completes', async () => { + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE]); + await putSessionMessageState(harness.storage, acceptedMessage()); + await harness.supervisor.onTerminalEvent({ wrapperRunId: WRAPPER_RUN_ID, status: 'completed', + messageIds: [MESSAGE_ID], }); await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + failureReason: 'missing_assistant_reply', + completionSource: 'idle_reconciliation', + }); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'terminal-failed', + }); + await expect(getWrapperRuntimeState(harness.storage)).resolves.toEqual({ + wrapperGeneration: 4, + }); + }); + + it.each([ + { + label: 'current', + messageShape: { + admissionSnapshot: { + turn: { + type: 'command' as const, + messageId: MESSAGE_ID, + command: 'compact', + arguments: '', + }, + agent: { mode: 'code', model: 'test-model' }, + }, + }, + }, + { + label: 'legacy', + messageShape: { + legacyAdmissionConstraints: { + turn: { + type: 'command' as const, + messageId: MESSAGE_ID, + command: 'compact', + arguments: '', + }, + }, + }, + }, + ])( + 'completes a successful $label command without an assistant reply', + async ({ messageShape }) => { + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE]); + await putSessionMessageState(harness.storage, { + ...acceptedMessage(), + ...messageShape, + }); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + messageIds: [MESSAGE_ID], + }); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'completed', + completionSource: 'idle_reconciliation', + assistantMessageId: undefined, + }); + } + ); + + it('fails a still-accepted errored reply when the wrapper completes', async () => { + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE], { + getAssistantMessageForUserMessage: () => + ({ + info: { + id: 'ase_complete_error', + role: 'assistant', + error: { data: { message: 'provider failed during completion' } }, + }, + parts: [], + }) as unknown as LatestAssistantMessage, + }); + await putSessionMessageState(harness.storage, acceptedMessage()); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, status: 'completed', + messageIds: [MESSAGE_ID], + }); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + failureReason: 'assistant_error', + error: 'provider failed during completion', completionSource: 'idle_reconciliation', - assistantMessageId, }); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'terminal-failed', + }); + await expect(getWrapperRuntimeState(harness.storage)).resolves.toEqual({ + wrapperGeneration: 4, + }); + }); + + it('includes the gate result when wrapper completion releases a reconciled callback', async () => { + const assistantMessageId = 'ase_complete_gate'; + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE], { + metadata: { + ...createMetadata(), + finalization: { gateThreshold: 'warning' }, + }, + getAssistantMessageForUserMessage: () => + ({ + info: { id: assistantMessageId, role: 'assistant' }, + parts: [], + }) as unknown as LatestAssistantMessage, + }); + await putSessionMessageState(harness.storage, { + ...acceptedMessage(), + callbackRequired: true, + callbackTarget: { url: 'https://example.com/complete-gate' }, + }); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + gateResult: 'pass', + messageIds: [MESSAGE_ID], + }); + expect(harness.callbackJobs).toHaveLength(1); expect(harness.callbackJobs[0].payload).toMatchObject({ messageId: MESSAGE_ID, status: 'completed', + gateResult: 'pass', + }); + }); + + it('fails accepted work omitted from sealed complete membership', async () => { + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + messageIds: [], + }); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + failureReason: 'wrapper_protocol_error', + }); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'terminal-failed', }); - expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); }); - it('keeps the idle-reconcile deadline armed when post-completion output is observed', async () => { - // Layer 2: autocommit/condense events after session.idle must refresh - // liveness without disarming the reconciler that finalizes the in-flight - // message. Before the fix, observeMeaningfulOutput cleared these fields. + it('rejects complete membership that omits an already terminal message fenced to the run', async () => { + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE]); + await putSessionMessageState(harness.storage, acceptedMessage()); + await putSessionMessageState(harness.storage, { + ...acceptedMessage(NEWER_MESSAGE_ID), + status: 'completed', + terminalAt: 3_000, + completionSource: 'idle_reconciliation', + }); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + messageIds: [MESSAGE_ID], + }); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + failureReason: 'wrapper_protocol_error', + }); + await expect(getSessionMessageState(harness.storage, NEWER_MESSAGE_ID)).resolves.toMatchObject({ + status: 'completed', + }); + }); + + it.each(['failed', 'interrupted'] as const)( + 'treats an already %s sealed member fenced to the run as batch failure', + async status => { + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE]); + await putSessionMessageState(harness.storage, { + ...acceptedMessage(), + status, + terminalAt: 3_000, + completionSource: status === 'failed' ? 'wrapper_failure' : 'interrupt', + }); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + messageIds: [MESSAGE_ID], + }); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status, + }); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'terminal-failed', + }); + await expect(getWrapperRuntimeState(harness.storage)).resolves.toEqual({ + wrapperGeneration: 4, + }); + } + ); + + it('accepts an already completed sealed member fenced to the run as a no-op', async () => { + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE]); + await putSessionMessageState(harness.storage, { + ...acceptedMessage(), + status: 'completed', + terminalAt: 3_000, + completionSource: 'idle_reconciliation', + }); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + messageIds: [MESSAGE_ID], + }); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'completed', + terminalAt: 3_000, + }); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'owns_wrapper', + keepWarmUntil: expect.any(Number), + }); + await expect(getWrapperRuntimeState(harness.storage)).resolves.toEqual({ + wrapperGeneration: 5, + }); + }); + + it('fails every accepted message before retiring a duplicate complete membership', async () => { + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE]); + await putSessionMessageState(harness.storage, acceptedMessage()); + await putSessionMessageState(harness.storage, acceptedMessage(NEWER_MESSAGE_ID)); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + messageIds: [MESSAGE_ID, MESSAGE_ID], + }); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + failureReason: 'wrapper_protocol_error', + }); + await expect(getSessionMessageState(harness.storage, NEWER_MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + failureReason: 'wrapper_protocol_error', + }); + await expect(getWrapperRuntimeState(harness.storage)).resolves.toEqual({ + wrapperGeneration: 4, + }); + }); + + it('fails complete-before-acceptance work before retiring duplicate membership', async () => { + const storage = createMemoryStorage([ + liveRuntimeState({ dispatchingMessageId: MESSAGE_ID }), + OWNED_WRAPPER_LEASE, + ]); + const ensureAccepted = vi.fn(async (messageId: string, wrapperRunId: string) => { + await putSessionMessageState(storage, { + ...acceptedMessage(messageId), + wrapperRunId, + dispatchAcceptanceKind: 'inferred_from_terminal', + }); + }); + const harness = createHarness(undefined, { + storage, + ensureAcceptedMessageBeforeTerminal: ensureAccepted, + }); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + messageIds: [MESSAGE_ID, MESSAGE_ID], + }); + + expect(ensureAccepted).toHaveBeenCalledWith(MESSAGE_ID, WRAPPER_RUN_ID); + await expect(getSessionMessageState(storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + failureReason: 'wrapper_protocol_error', + }); + }); + + it('rejects a terminal sealed message fenced to another wrapper run', async () => { + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE]); + await putSessionMessageState(harness.storage, acceptedMessage()); + await putSessionMessageState(harness.storage, { + ...acceptedMessage(NEWER_MESSAGE_ID), + status: 'completed', + terminalAt: 3_000, + wrapperRunId: 'wr_other_run', + completionSource: 'idle_reconciliation', + }); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + messageIds: [NEWER_MESSAGE_ID], + }); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + failureReason: 'wrapper_protocol_error', + }); + await expect(getSessionMessageState(harness.storage, NEWER_MESSAGE_ID)).resolves.toMatchObject({ + status: 'completed', + wrapperRunId: 'wr_other_run', + }); + }); + + it('keeps supervising finalizing work after all accepted messages settle', async () => { + const harness = createHarness([ + liveRuntimeState({ + finalizingWrapperRunId: WRAPPER_RUN_ID, + nextPingAt: 2_000, + noOutputDeadlineAt: 50_000, + }), + ]); + + await harness.supervisor.runMaintenance(2_001); + + expect(harness.sentPings).toEqual([WRAPPER_RUN_ID]); + await expect(getWrapperRuntimeState(harness.storage)).resolves.toMatchObject({ + finalizingWrapperRunId: WRAPPER_RUN_ID, + pingDeadlineAt: 32_001, + }); + }); + + it('keeps finalizing fenced while post-processing output refreshes liveness', async () => { const harness = createHarness([ liveRuntimeState({ - lastWrapperIdleAt: 1_000, - idleReconcileAfter: 9_000, + finalizingWrapperRunId: WRAPPER_RUN_ID, wrapperIdleDeadlineAt: 50_000, }), ]); @@ -740,10 +1352,10 @@ describe('WrapperSupervisor', () => { await harness.supervisor.observeMeaningfulOutput(4, WRAPPER_CONNECTION_ID, 2_000); await expect(getWrapperRuntimeState(harness.storage)).resolves.toMatchObject({ - lastWrapperIdleAt: 1_000, - idleReconcileAfter: 9_000, + finalizingWrapperRunId: WRAPPER_RUN_ID, wrapperIdleDeadlineAt: 50_000, lastWrapperMessageAt: 2_000, + noOutputDeadlineAt: 302_000, }); }); @@ -753,6 +1365,7 @@ describe('WrapperSupervisor', () => { await harness.supervisor.onTerminalEvent({ wrapperRunId: WRAPPER_RUN_ID, status: 'completed', + messageIds: [], }); await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ @@ -788,6 +1401,34 @@ describe('WrapperSupervisor', () => { expect(harness.stops).toEqual([]); }); + it('skips expired keep-warm cleanup while the current wrapper run is finalizing', async () => { + const harness = createHarness([ + liveRuntimeState({ + finalizingWrapperRunId: WRAPPER_RUN_ID, + wrapperIdleDeadlineAt: 9_000, + }), + [ + 'wrapper_lease', + { + ...(OWNED_WRAPPER_LEASE[1] as object), + keepWarmUntil: 9_000, + }, + ], + ]); + + await harness.supervisor.runMaintenance(10_000); + + await expect(getWrapperRuntimeState(harness.storage)).resolves.toMatchObject({ + wrapperConnectionId: WRAPPER_CONNECTION_ID, + wrapperGeneration: 4, + finalizingWrapperRunId: WRAPPER_RUN_ID, + }); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'owns_wrapper', + keepWarmUntil: 9_000, + }); + }); + it('turns an expired startup allowance into verified cleanup work', async () => { const harness = createHarness([ [ @@ -876,6 +1517,7 @@ describe('WrapperSupervisor', () => { state: 'none', nextInstanceGeneration: 2, }); + expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); }); it.each([ diff --git a/services/cloud-agent-next/src/session/wrapper-supervisor.ts b/services/cloud-agent-next/src/session/wrapper-supervisor.ts index ac080e74e0..57f36b12ef 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.ts @@ -12,26 +12,29 @@ import type { MessageSettlementOutbox } from './message-settlement-outbox.js'; import { countPendingSessionMessages, type SessionQueueStorage } from './pending-messages.js'; import type { SessionMessageQueue } from './session-message-queue.js'; import { + listMessagesForWrapperRun, listNonTerminalAcceptedMessages, + type SessionMessageState, type SessionMessageStorage, } from './session-message-state.js'; import type { LatestAssistantMessage } from './types.js'; import { clearCurrentWrapperRuntimeFailureState, clearCurrentWrapperRuntimeLivenessState, - clearWrapperIdleState, clearWrapperRuntimeIdentity, getWrapperLease, getWrapperRuntimeState, hasCompleteWrapperIdentity, + hasCompleteWrapperRunMessageIndex, IDLE_KEEP_WARM_MS, - IDLE_RECONCILIATION_GRACE_MS, isCurrentWrapperConnection, + isWrapperDeliveryHeld, + isWrapperRunFinalizing, + markWrapperFinalizing, markWrapperPingSent, nextWrapperLeaseDeadline, putWrapperLease, recordMeaningfulWrapperOutput, - recordRootSessionIdle, recordWrapperPong, reduceWrapperLease, type WrapperConnectionFence, @@ -85,6 +88,11 @@ export type WrapperTerminalEvent = { status: 'completed' | 'failed' | 'interrupted'; error?: string; gateResult?: 'pass' | 'fail'; + messageIds?: string[]; +}; + +type SealedBatchSettlementResult = { + failedTerminalObserved: boolean; }; export type WrapperSupervisorStorage = DurableObjectStorage & @@ -101,11 +109,7 @@ export type WrapperSupervisor = { wrapperConnectionId: string, now: number ): Promise; - observeRootIdle( - wrapperGeneration: number, - wrapperConnectionId: string, - now: number - ): Promise; + observeFinalizing(wrapperRunId: string): Promise; onDisconnected(input: WrapperDisconnectedInput): Promise; onTerminalEvent(params: WrapperTerminalEvent): Promise; requestPhysicalWrapperStop(reason: WrapperStopReason, target?: WrapperStopTarget): Promise; @@ -125,6 +129,7 @@ export type WrapperSupervisorDependencies = { | 'releaseWrapperTerminalWaitForIdleBatchForWrapperRun' | 'isWaitingForWrapperTerminalGateResult' | 'finalizeIdleBatchCallbackIfReady' + | 'finalizeTerminalWrapperRunCallbackIfReady' >; sessionMessageQueue: Pick; getMetadata: () => Promise; @@ -140,6 +145,7 @@ export type WrapperSupervisorDependencies = { wrapperConnectionId: string; }) => Promise; clearInterruptRequest: () => Promise; + ensureAcceptedMessageBeforeTerminal: (messageId: string, wrapperRunId: string) => Promise; stopWrappers?: (request: { target: WrapperStopTarget; attemptId: string; @@ -208,6 +214,7 @@ export function createWrapperSupervisor( observeCorrelatedAgentActivity, hasActiveIngestConnection, clearInterruptRequest, + ensureAcceptedMessageBeforeTerminal, stopWrappers, requestAlarmAtOrBefore, getSessionIdForLogs, @@ -256,9 +263,7 @@ export function createWrapperSupervisor( ); if (!released) return; - await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady({ - allowWithoutObservedIdle: true, - }); + await messageSettlementOutbox.finalizeTerminalWrapperRunCallbackIfReady(wrapperRunId); } async function checkReconnect(input: WrapperReconnectInput): Promise { @@ -344,19 +349,8 @@ export function createWrapperSupervisor( await requestAlarmAtOrBefore?.(now + IDLE_KEEP_WARM_MS); } - async function observeRootIdle( - wrapperGeneration: number, - wrapperConnectionId: string, - now: number - ): Promise { - await recordRootSessionIdle( - storage, - wrapperGeneration, - wrapperConnectionId, - now, - now + IDLE_RECONCILIATION_GRACE_MS - ); - await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady(); + async function observeFinalizing(wrapperRunId: string): Promise { + await markWrapperFinalizing(storage, wrapperRunId); } async function startDisconnectGrace(input: WrapperDisconnectedInput): Promise { @@ -401,7 +395,13 @@ export function createWrapperSupervisor( ); const isWaitingForWrapperTerminalGateResult = await messageSettlementOutbox.isWaitingForWrapperTerminalGateResult(); - if (acceptedMessages.length === 0 && !isWaitingForWrapperTerminalGateResult) return; + if ( + acceptedMessages.length === 0 && + !isWaitingForWrapperTerminalGateResult && + !isWrapperRunFinalizing(state) + ) { + return; + } await startDisconnectGrace(input); } @@ -458,9 +458,13 @@ export function createWrapperSupervisor( }); } await messageSettlementOutbox.releaseWrapperTerminalWaitForIdleBatch(); - await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady({ - allowWithoutObservedIdle: true, - }); + if (isWrapperRunFinalizing(state) && state.wrapperRunId) { + await messageSettlementOutbox.finalizeTerminalWrapperRunCallbackIfReady(state.wrapperRunId); + } else { + await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady({ + allowWithoutObservedIdle: true, + }); + } if (state.wrapperConnectionId) { await clearCurrentWrapperRuntimeFailureState( @@ -507,7 +511,7 @@ export function createWrapperSupervisor( } const acceptedMessages = await listNonTerminalAcceptedMessages(storage, wrapperRunId); - if (acceptedMessages.length === 0) { + if (acceptedMessages.length === 0 && !isWrapperRunFinalizing(state)) { logger .withFields({ wrapperRunId }) .info('No accepted messages during grace period - skipping failure'); @@ -518,7 +522,7 @@ export function createWrapperSupervisor( logger .withFields({ wrapperRunId, messageCount: acceptedMessages.length }) - .warn('Grace period expired - failing accepted messages'); + .warn('Grace period expired - failing supervised wrapper work'); await requestPhysicalWrapperStop('unhealthy-wrapper'); await storage.delete(DISCONNECT_GRACE_KEY); for (const message of acceptedMessages) { @@ -540,10 +544,11 @@ export function createWrapperSupervisor( }, { incrementGeneration: true } ); - await releaseWrapperTerminalWaitForIdleBatch(); + await releaseWrapperTerminalWaitForIdleBatchForWrapperRun(wrapperRunId); } async function hasActiveWrapperWork(state: WrapperRuntimeState): Promise { + if (isWrapperRunFinalizing(state)) return true; return (await listNonTerminalAcceptedMessages(storage, state.wrapperRunId)).length > 0; } @@ -645,92 +650,102 @@ export function createWrapperSupervisor( return false; } - /** - * Terminalize the wrapper run's still-accepted messages against the latest - * assistant reply, mirroring how a finished turn would have settled them. - * - * Shared by two triggers: - * - `idle`: the idle-reconciliation grace deadline elapsed. - * - `wrapper_complete`: the wrapper emitted its terminal `complete` event, - * which is the race-free "fully done" signal (it fires only after all - * post-completion work). This is the authoritative backstop for the case - * where the assistant `message.updated` arrived without `time.completed`, - * so neither the assistant-event nor the idle path settled the message. - * - * `terminalizeSessionMessageOnce` is idempotent, so re-running here after a - * partial settlement is safe. - */ - async function reconcileAcceptedMessages( - now: number, - state: WrapperRuntimeState, - metadata: SessionMetadata, - acceptedMessages: Awaited>, - trigger: 'idle' | 'wrapper_complete' + function isPromptMessage(message: SessionMessageState): boolean { + const turn = message.admissionSnapshot?.turn ?? message.legacyAdmissionConstraints?.turn; + return turn?.type !== 'command'; + } + + async function failAcceptedMessagesForProtocolError( + acceptedMessages: SessionMessageState[], + error: string ): Promise { - logger - .withFields({ - sessionId: metadata.identity.sessionId, - wrapperRunId: state.wrapperRunId, - acceptedMessageCount: acceptedMessages.length, - hasKiloSessionId: metadata.auth.kiloSessionId !== undefined, - trigger, - }) - .info('Reconciling accepted messages'); - - let failedTerminalObserved = !metadata.auth.kiloSessionId; - if (!failedTerminalObserved && metadata.auth.kiloSessionId) { - failedTerminalObserved = acceptedMessages.some(message => { - const assistantMessage = getAssistantMessageForUserMessage( - metadata.identity.sessionId, - metadata.auth.kiloSessionId ?? '', - message.messageId - ); - return ( - !assistantMessage || getAssistantErrorMessage(assistantMessage.info.error) !== undefined - ); + for (const message of acceptedMessages) { + await messageSettlementOutbox.terminalizeSessionMessageOnce(message.messageId, { + kind: 'failed', + reason: 'wrapper_protocol_error', + error, + completionSource: 'wrapper_failure', + failureStage: 'agent_activity', + failureCode: 'wrapper_error_after_activity', }); } - if (failedTerminalObserved) { + } + + async function settleSealedBatch( + wrapperRunId: string, + messageIds: string[], + dispatchingMessageId?: string, + membershipProtocolError?: string + ): Promise { + const sealedMessageIds = [...new Set(messageIds)]; + const repairMessageIds = [ + ...new Set([...sealedMessageIds, ...(dispatchingMessageId ? [dispatchingMessageId] : [])]), + ]; + for (const messageId of repairMessageIds) { + await ensureAcceptedMessageBeforeTerminal(messageId, wrapperRunId); + } + + const wrapperRunMessages = await listMessagesForWrapperRun(storage, wrapperRunId); + const wrapperRunMessagesById = new Map( + wrapperRunMessages.map(message => [message.messageId, message]) + ); + const acceptedMessages = wrapperRunMessages.filter(message => message.status === 'accepted'); + const earlyProtocolError = + membershipProtocolError ?? + (sealedMessageIds.length !== messageIds.length + ? 'Wrapper complete contained duplicate sealed batch membership' + : undefined); + if (earlyProtocolError) { await requestPhysicalWrapperStop('terminal-failed'); + await failAcceptedMessagesForProtocolError(acceptedMessages, earlyProtocolError); + return { failedTerminalObserved: true }; } - for (const message of acceptedMessages) { - if (!metadata.auth.kiloSessionId) { - failedTerminalObserved = true; - await messageSettlementOutbox.terminalizeSessionMessageOnce(message.messageId, { - kind: 'failed', - reason: 'missing_assistant_reply', - error: 'No assistant reply found during reconciliation', - completionSource: 'idle_reconciliation', - failureStage: 'post_dispatch_no_activity', - failureCode: 'missing_assistant_reply', - }); - continue; - } + const invalidMessageIds: string[] = []; + for (const messageId of sealedMessageIds) { + const state = wrapperRunMessagesById.get(messageId); + if (!state || state.status === 'queued') invalidMessageIds.push(messageId); + } + const sealedSet = new Set(sealedMessageIds); + const omittedMessages = wrapperRunMessages.filter(message => !sealedSet.has(message.messageId)); + const protocolFailure = invalidMessageIds.length > 0 || omittedMessages.length > 0; - const assistantMessage = getAssistantMessageForUserMessage( - metadata.identity.sessionId, - metadata.auth.kiloSessionId, - message.messageId + if (protocolFailure) { + await requestPhysicalWrapperStop('terminal-failed'); + await failAcceptedMessagesForProtocolError( + acceptedMessages, + 'Wrapper complete contained invalid sealed batch membership' ); - if (!assistantMessage) { - failedTerminalObserved = true; - await messageSettlementOutbox.terminalizeSessionMessageOnce(message.messageId, { - kind: 'failed', - reason: 'missing_assistant_reply', - error: 'No assistant reply found during reconciliation', - completionSource: 'idle_reconciliation', - failureStage: 'post_dispatch_no_activity', - failureCode: 'missing_assistant_reply', - }); - continue; - } + logger + .withFields({ + wrapperRunId, + invalidMessageIds, + omittedMessageIds: omittedMessages.map(message => message.messageId), + }) + .warn('Wrapper complete contained invalid sealed batch membership'); + return { failedTerminalObserved: true }; + } - await observeCorrelatedAgentActivity?.(message.messageId); - const assistantError = getAssistantErrorMessage(assistantMessage.info.error); + const metadata = await getMetadata(); + if (!metadata) return null; + const kiloSessionId = metadata.auth.kiloSessionId; + let failedTerminalObserved = wrapperRunMessages.some( + message => + sealedSet.has(message.messageId) && + (message.status === 'failed' || message.status === 'interrupted') + ); + + for (const messageId of sealedMessageIds) { + const message = wrapperRunMessagesById.get(messageId); + if (!message || message.status !== 'accepted') continue; + const assistantMessage = kiloSessionId + ? getAssistantMessageForUserMessage(metadata.identity.sessionId, kiloSessionId, messageId) + : null; + const assistantError = getAssistantErrorMessage(assistantMessage?.info.error); if (assistantError !== undefined) { failedTerminalObserved = true; - await messageSettlementOutbox.terminalizeSessionMessageOnce(message.messageId, { + await observeCorrelatedAgentActivity?.(messageId); + await messageSettlementOutbox.terminalizeSessionMessageOnce(messageId, { kind: 'failed', reason: 'assistant_error', error: assistantError, @@ -738,71 +753,40 @@ export function createWrapperSupervisor( failureStage: 'agent_activity', failureCode: 'assistant_error', }); - continue; - } - - await messageSettlementOutbox.terminalizeSessionMessageOnce(message.messageId, { - kind: 'completed', - assistantMessageId: assistantMessage.info.id, - completionSource: 'idle_reconciliation', - }); - } - - if (failedTerminalObserved) { - if (state.wrapperConnectionId) { - await clearWrapperRuntimeIdentity(storage, { - wrapperGeneration: state.wrapperGeneration, - wrapperConnectionId: state.wrapperConnectionId, + } else if (assistantMessage) { + await observeCorrelatedAgentActivity?.(messageId); + await messageSettlementOutbox.terminalizeSessionMessageOnce(messageId, { + kind: 'completed', + assistantMessageId: assistantMessage.info.id, + completionSource: 'idle_reconciliation', + }); + } else if (!isPromptMessage(message)) { + await messageSettlementOutbox.terminalizeSessionMessageOnce(messageId, { + kind: 'completed', + completionSource: 'idle_reconciliation', + }); + } else { + failedTerminalObserved = true; + await messageSettlementOutbox.terminalizeSessionMessageOnce(messageId, { + kind: 'failed', + reason: 'missing_assistant_reply', + error: 'No assistant reply found during wrapper completion', + completionSource: 'idle_reconciliation', + failureStage: 'post_dispatch_no_activity', + failureCode: 'missing_assistant_reply', }); } - } else { - await retainPhysicalWrapperWarm(now); } - await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady(); - logger - .withFields({ - sessionId: metadata.identity.sessionId, - wrapperRunId: state.wrapperRunId, - acceptedMessageCount: acceptedMessages.length, - trigger, - }) - .info('Reconciliation pass completed'); - } - - async function checkIdleReconciliation(now: number): Promise { - const metadata = await getMetadata(); - if (!metadata) return; - - const state = await getWrapperRuntimeState(storage); - if (!state.wrapperRunId) return; - const acceptedMessages = await listNonTerminalAcceptedMessages(storage, state.wrapperRunId); - if (acceptedMessages.length === 0) { - if ( - state.wrapperConnectionId && - (state.lastWrapperIdleAt !== undefined || state.idleReconcileAfter !== undefined) - ) { - await clearWrapperIdleState(storage, state.wrapperGeneration, state.wrapperConnectionId); - } - return; - } - - if (state.idleReconcileAfter !== undefined) { - if (now < state.idleReconcileAfter) return; - } else { - const hasRecentOutput = - state.lastWrapperMessageAt !== undefined && - now - state.lastWrapperMessageAt < WRAPPER_NO_OUTPUT_TIMEOUT_MS; - if (hasRecentOutput) return; - } - - await reconcileAcceptedMessages(now, state, metadata, acceptedMessages, 'idle'); + if (failedTerminalObserved) await requestPhysicalWrapperStop('terminal-failed'); + return { failedTerminalObserved }; } async function checkKeepWarmCleanup(now: number): Promise { const lease = await getWrapperLease(storage); if (lease.state === 'owns_wrapper' && lease.startupDeadlineAt !== undefined) return; const wrapperState = await getWrapperRuntimeState(storage); + if (isWrapperRunFinalizing(wrapperState)) return; const keepWarmUntil = lease.state === 'owns_wrapper' ? lease.keepWarmUntil : wrapperState.wrapperIdleDeadlineAt; if (keepWarmUntil === undefined || keepWarmUntil > now) return; @@ -812,16 +796,7 @@ export function createWrapperSupervisor( storage, wrapperState.wrapperRunId ); - if (pendingCount > 0 || acceptedMessages.length > 0) { - if (wrapperState.wrapperConnectionId) { - await clearWrapperIdleState( - storage, - wrapperState.wrapperGeneration, - wrapperState.wrapperConnectionId - ); - } - return; - } + if (pendingCount > 0 || acceptedMessages.length > 0) return; logger .withFields({ @@ -912,10 +887,11 @@ export function createWrapperSupervisor( const latest = await getWrapperLease(storage); if (result.status === 'absent') { - await putWrapperLease( - storage, - reduceWrapperLease(latest, { type: 'stop_absent', attemptId }) - ); + const cleaned = reduceWrapperLease(latest, { type: 'stop_absent', attemptId }); + await putWrapperLease(storage, cleaned); + if (!isWrapperDeliveryHeld(await getWrapperRuntimeState(storage), cleaned)) { + await sessionMessageQueue.requestPendingDrainIfNeeded(); + } return; } const error = @@ -934,7 +910,7 @@ export function createWrapperSupervisor( } async function onTerminalEvent(params: WrapperTerminalEvent): Promise { - const { wrapperRunId, status, error, gateResult } = params; + const { wrapperRunId, status, error, gateResult, messageIds } = params; const sessionId = getSessionIdForLogs(); const state = await getWrapperRuntimeState(storage); if ( @@ -956,6 +932,7 @@ export function createWrapperSupervisor( status, error, gateResult, + messageCount: messageIds?.length, }) .info('Wrapper terminal event received by supervisor'); @@ -963,6 +940,9 @@ export function createWrapperSupervisor( await requestPhysicalWrapperStop( status === 'failed' ? 'terminal-failed' : 'terminal-interrupted' ); + if (state.dispatchingMessageId) { + await ensureAcceptedMessageBeforeTerminal(state.dispatchingMessageId, wrapperRunId); + } const acceptedMessages = await listNonTerminalAcceptedMessages(storage, wrapperRunId); for (const message of acceptedMessages) { if (status === 'failed') { @@ -991,33 +971,48 @@ export function createWrapperSupervisor( } if (status === 'completed') { - const acceptedMessages = await listNonTerminalAcceptedMessages(storage, wrapperRunId); - if (acceptedMessages.length === 0) { - await retainPhysicalWrapperWarm(Date.now()); - await clearInterruptRequest(); + const currentRunRequiresMembership = hasCompleteWrapperRunMessageIndex(state, wrapperRunId); + const missingRequiredMembership = messageIds === undefined && currentRunRequiresMembership; + const sealedMessageIds = + messageIds ?? + (missingRequiredMembership + ? [] + : [ + ...new Set([ + ...(await listMessagesForWrapperRun(storage, wrapperRunId)).map( + message => message.messageId + ), + ...(state.dispatchingMessageId ? [state.dispatchingMessageId] : []), + ]), + ]); + const settlement = await settleSealedBatch( + wrapperRunId, + sealedMessageIds, + state.dispatchingMessageId, + missingRequiredMembership + ? 'Current wrapper complete omitted sealed batch membership' + : undefined + ); + if (!settlement) { + await requestPhysicalWrapperStop('terminal-failed'); + const acceptedMessages = await listNonTerminalAcceptedMessages(storage, wrapperRunId); + await failAcceptedMessagesForProtocolError( + acceptedMessages, + 'Wrapper complete omitted sealed batch membership' + ); + await clearWrapperRuntimeIdentity(storage, { + wrapperGeneration: state.wrapperGeneration, + wrapperConnectionId: state.wrapperConnectionId, + }); + } else if (settlement.failedTerminalObserved) { + await clearWrapperRuntimeIdentity(storage, { + wrapperGeneration: state.wrapperGeneration, + wrapperConnectionId: state.wrapperConnectionId, + }); } else { - // `complete` is the race-free "fully done" signal: it is emitted only - // after all post-completion work. If messages are still accepted here, - // the assistant-event and idle-reconcile paths both missed them (e.g. - // the final assistant `message.updated` lacked `time.completed`, and - // post-idle autocommit refreshed liveness). Settle them now rather than - // leaving the callback gated indefinitely. - const metadata = await getMetadata(); - if (metadata) { - await reconcileAcceptedMessages( - Date.now(), - state, - metadata, - acceptedMessages, - 'wrapper_complete' - ); - } else { - logger - .withFields({ sessionId, wrapperRunId, acceptedMessageCount: acceptedMessages.length }) - .warn('Wrapper complete with accepted messages but no session metadata to reconcile'); - } - await clearInterruptRequest(); + await retainPhysicalWrapperWarm(Date.now()); } + await clearInterruptRequest(); } else { await clearWrapperRuntimeIdentity(storage, { wrapperGeneration: state.wrapperGeneration, @@ -1028,17 +1023,18 @@ export function createWrapperSupervisor( await clearDisconnectGrace(); await messageSettlementOutbox.observeWrapperTerminalForIdleBatch(gateResult); - await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady({ - allowWithoutObservedIdle: true, - }); - await sessionMessageQueue.requestPendingDrainIfNeeded(); + await messageSettlementOutbox.finalizeTerminalWrapperRunCallbackIfReady(wrapperRunId); + if ( + !isWrapperDeliveryHeld(await getWrapperRuntimeState(storage), await getWrapperLease(storage)) + ) { + await sessionMessageQueue.requestPendingDrainIfNeeded(); + } } async function runMaintenance(now: number): Promise { await reconcilePhysicalCleanup(now); await checkDisconnectGrace(now); await checkWrapperLiveness(now); - await checkIdleReconciliation(now); await checkKeepWarmCleanup(now); } @@ -1059,9 +1055,6 @@ export function createWrapperSupervisor( } const wrapperState = await getWrapperRuntimeState(storage); - if (wrapperState.idleReconcileAfter !== undefined) { - deadlines.push(wrapperState.idleReconcileAfter); - } if (wrapperState.wrapperIdleDeadlineAt !== undefined) { deadlines.push(wrapperState.wrapperIdleDeadlineAt); } @@ -1075,7 +1068,7 @@ export function createWrapperSupervisor( isCurrentConnection, observePong, observeMeaningfulOutput, - observeRootIdle, + observeFinalizing, onDisconnected, onTerminalEvent, requestPhysicalWrapperStop, diff --git a/services/cloud-agent-next/src/shared/protocol.ts b/services/cloud-agent-next/src/shared/protocol.ts index b6a6dd0f32..3d417ac574 100644 --- a/services/cloud-agent-next/src/shared/protocol.ts +++ b/services/cloud-agent-next/src/shared/protocol.ts @@ -5,7 +5,7 @@ import type { SlashCommandInfo } from './slash-commands.js'; * * From wrapper -> DO: * started, kilocode, output, status, heartbeat, pong, error, interrupted, complete, wrapper_resumed, - * autocommit_started, autocommit_completed, cloud.message.completed + * wrapper_finalizing, autocommit_started, autocommit_completed, cloud.message.completed * * From DO -> /stream clients: * All of the above, plus wrapper_disconnected, wrapper_reconnected, preparing, @@ -21,7 +21,8 @@ export type StreamEventType = | 'pong' // Response to ping command from DO | 'error' // Error occurred { error: string, fatal: boolean } | 'interrupted' // User/signal interrupt - | 'complete' // Execution finished { exitCode, currentBranch? } + | 'complete' // Execution finished { exitCode, currentBranch?, messageIds } + | 'wrapper_finalizing' // Wrapper sealed the current run batch before post-processing | '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) @@ -65,6 +66,7 @@ export type CompleteEventData = { exitCode: number; currentBranch?: string; // Omitted if detached HEAD gateResult?: 'pass' | 'fail'; + messageIds?: string[]; }; /** diff --git a/services/cloud-agent-next/src/shared/wrapper-bootstrap.ts b/services/cloud-agent-next/src/shared/wrapper-bootstrap.ts index 689293c257..3894556fc8 100644 --- a/services/cloud-agent-next/src/shared/wrapper-bootstrap.ts +++ b/services/cloud-agent-next/src/shared/wrapper-bootstrap.ts @@ -148,9 +148,14 @@ export type WrapperSessionReadySuccessResponse = { export type WrapperSessionReadyErrorResponse = { status: 'error'; error: { - code: 'INVALID_REQUEST' | 'WORKSPACE_SETUP_FAILED' | 'KILO_SERVER_FAILED'; + code: + | 'INVALID_REQUEST' + | 'WRAPPER_FINALIZING' + | 'WORKSPACE_SETUP_FAILED' + | 'KILO_SERVER_FAILED'; message: string; retryable?: boolean; + wrapperRunId?: string; }; }; diff --git a/services/cloud-agent-next/src/shared/wrapper-version.ts b/services/cloud-agent-next/src/shared/wrapper-version.ts index c0d3173841..5c407c0576 100644 --- a/services/cloud-agent-next/src/shared/wrapper-version.ts +++ b/services/cloud-agent-next/src/shared/wrapper-version.ts @@ -1 +1 @@ -export const WRAPPER_VERSION = '2.2.0'; +export const WRAPPER_VERSION = '2.3.0'; diff --git a/services/cloud-agent-next/src/websocket/ingest.test.ts b/services/cloud-agent-next/src/websocket/ingest.test.ts index 303fb96224..0fcaf2d547 100644 --- a/services/cloud-agent-next/src/websocket/ingest.test.ts +++ b/services/cloud-agent-next/src/websocket/ingest.test.ts @@ -48,7 +48,7 @@ function createFakeDOContext(): IngestDOContext { isCurrentConnection: vi.fn().mockResolvedValue(true), observePong: vi.fn().mockResolvedValue(undefined), observeMeaningfulOutput: vi.fn().mockResolvedValue(undefined), - observeRootIdle: vi.fn().mockResolvedValue(undefined), + observeFinalizing: vi.fn().mockResolvedValue(undefined), onTerminalEvent: vi.fn().mockResolvedValue(undefined), }, }; @@ -891,7 +891,7 @@ describe('createIngestHandler', () => { expect(doContext.terminalizeSessionMessageOnce).not.toHaveBeenCalled(); }); - it('terminalizes on terminal assistant message.updated (with time.completed)', async () => { + it('observes activity without terminalizing on completed assistant message.updated', async () => { const state = createFakeState(); const doContext = createNewPathDOContext(); const handler = createIngestHandler( @@ -922,15 +922,8 @@ describe('createIngestHandler', () => { await handler.handleIngestMessage(ws, message); - expect(doContext.terminalizeSessionMessageOnce).toHaveBeenCalledWith( - 'msg_user_222', - expect.objectContaining({ - kind: 'completed', - assistantMessageId: 'asst_222', - completionSource: 'assistant_message_event', - }), - WRAPPER_RUN_ID - ); + expect(doContext.observeCorrelatedAgentActivity).toHaveBeenCalledWith('msg_user_222'); + expect(doContext.terminalizeSessionMessageOnce).not.toHaveBeenCalled(); }); it('terminalizes on wrapper cloud.message.completed control event', async () => { @@ -1124,7 +1117,7 @@ describe('createIngestHandler', () => { ); }); - it('ingests duplicate terminal updates without error (idempotency at DO level)', async () => { + it('ingests duplicate completed assistant updates as non-terminal activity', async () => { const state = createFakeState(); const doContext = createNewPathDOContext(); doContext.terminalizeSessionMessageOnce = vi.fn().mockResolvedValue(undefined); @@ -1158,13 +1151,25 @@ describe('createIngestHandler', () => { await handler.handleIngestMessage(ws, makeMessage()); await handler.handleIngestMessage(ws, makeMessage()); - // Both events trigger terminalizeSessionMessageOnce; the DO handles idempotency - expect(doContext.terminalizeSessionMessageOnce).toHaveBeenCalledTimes(2); - expect(doContext.terminalizeSessionMessageOnce).toHaveBeenCalledWith( - 'msg_user_444', - expect.objectContaining({ kind: 'completed' }), - WRAPPER_RUN_ID + expect(doContext.observeCorrelatedAgentActivity).toHaveBeenCalledTimes(2); + expect(doContext.observeCorrelatedAgentActivity).toHaveBeenCalledWith('msg_user_444'); + expect(doContext.terminalizeSessionMessageOnce).not.toHaveBeenCalled(); + }); + + it('marks the current run finalizing from wrapper control event', async () => { + const doContext = createNewPathDOContext(); + const handler = createIngestHandler( + createFakeState(), + createFakeEventQueries(), + SESSION_ID, + vi.fn(), + doContext ); + const ws = createFakeWebSocket(makeNewPathAttachment()); + + await handler.handleIngestMessage(ws, makeStreamMessage('wrapper_finalizing')); + + expect(doContext.wrapperSupervisor.observeFinalizing).toHaveBeenCalledWith(WRAPPER_RUN_ID); }); it('does NOT terminalize on wrapper complete event (new path)', async () => { @@ -1182,7 +1187,7 @@ describe('createIngestHandler', () => { const message = JSON.stringify({ streamEventType: 'complete', - data: { exitCode: 0 }, + data: { exitCode: 0, messageIds: ['msg_user_complete'] }, timestamp: new Date().toISOString(), }); @@ -1193,6 +1198,27 @@ describe('createIngestHandler', () => { expect.objectContaining({ status: 'completed', wrapperRunId: WRAPPER_RUN_ID }) ); }); + + it('forwards legacy wrapper complete events without sealed membership', async () => { + const doContext = createNewPathDOContext(); + const handler = createIngestHandler( + createFakeState(), + createFakeEventQueries(), + SESSION_ID, + vi.fn(), + doContext + ); + const ws = createFakeWebSocket(makeNewPathAttachment()); + + await handler.handleIngestMessage(ws, makeStreamMessage('complete', { exitCode: 0 })); + + expect(doContext.wrapperSupervisor.onTerminalEvent).toHaveBeenCalledWith({ + status: 'completed', + wrapperRunId: WRAPPER_RUN_ID, + gateResult: undefined, + messageIds: undefined, + }); + }); }); describe('handleIngestRequest — new-path sessionId validation', () => { diff --git a/services/cloud-agent-next/src/websocket/ingest.ts b/services/cloud-agent-next/src/websocket/ingest.ts index 4da06d4834..7e88e28934 100644 --- a/services/cloud-agent-next/src/websocket/ingest.ts +++ b/services/cloud-agent-next/src/websocket/ingest.ts @@ -41,6 +41,7 @@ const completeEventSchema = z.object({ exitCode: z.number(), currentBranch: z.string().optional(), gateResult: z.enum(['pass', 'fail']).optional(), + messageIds: z.array(z.string()).optional(), }); const kilocodeEventSchema = z @@ -147,7 +148,7 @@ export type IngestDOContext = { | 'isCurrentConnection' | 'observePong' | 'observeMeaningfulOutput' - | 'observeRootIdle' + | 'observeFinalizing' | 'onTerminalEvent' >; keepContainerAlive?: () => void; @@ -192,6 +193,7 @@ export function createIngestHandler( status: 'completed' | 'failed' | 'interrupted'; error?: string; gateResult?: 'pass' | 'fail'; + messageIds?: string[]; }): Promise { await doContext.wrapperSupervisor.onTerminalEvent(params); } @@ -614,49 +616,31 @@ export function createIngestHandler( } } - if (isSessionIdle && wrapperGeneration !== undefined && wrapperConnectionId) { - await doContext.wrapperSupervisor.observeRootIdle( - wrapperGeneration, - wrapperConnectionId, - now - ); - } - - // Terminalize user messages from terminal assistant message.updated events only. - // Partial updates (no time.completed or error) must not terminalize. + // Assistant completion can be an intermediate tool-loop step. Record + // correlated activity here and let a turn-level lifecycle boundary settle success. if (eventType === 'kilocode') { const data = ingestEvent.data as Record; const eventName = data.event; if (eventName === 'message.updated') { const properties = data.properties as Record | undefined; const info = properties?.info as Record | undefined; - const time = info?.time as Record | undefined; - const isCompleted = Boolean(time?.completed); const assistantError = getAssistantErrorMessage(info?.error); - const hasError = assistantError !== undefined; - const isTerminal = isCompleted || hasError; const parentMessageId = info?.role === 'assistant' && typeof info.parentID === 'string' ? info.parentID : undefined; if (parentMessageId !== undefined) { await doContext.observeCorrelatedAgentActivity?.(parentMessageId); - if (isTerminal) { + if (assistantError !== undefined) { await doContext.terminalizeSessionMessageOnce( parentMessageId, - hasError - ? { - kind: 'failed', - assistantMessageId: typeof info?.id === 'string' ? info.id : undefined, - reason: 'assistant_error', - error: assistantError, - completionSource: 'assistant_message_event', - } - : { - kind: 'completed', - assistantMessageId: typeof info?.id === 'string' ? info.id : undefined, - completionSource: 'assistant_message_event', - }, + { + kind: 'failed', + assistantMessageId: typeof info?.id === 'string' ? info.id : undefined, + reason: 'assistant_error', + error: assistantError, + completionSource: 'assistant_message_event', + }, wrapperRunId ); } @@ -693,6 +677,10 @@ export function createIngestHandler( }); } + if (eventType === 'wrapper_finalizing') { + await doContext.wrapperSupervisor.observeFinalizing(wrapperRunId); + } + if (eventType === 'complete') { broadcastFn({ id: 0 as EventId, @@ -718,6 +706,7 @@ export function createIngestHandler( wrapperRunId, status: 'completed', gateResult: parsedComplete.data.gateResult, + messageIds: parsedComplete.data.messageIds, }); logger .withFields({ @@ -726,6 +715,7 @@ export function createIngestHandler( wrapperGeneration, wrapperConnectionId, gateResult: parsedComplete.data.gateResult, + messageCount: parsedComplete.data.messageIds?.length, }) .info('Wrapper complete event forwarded to session coordinator'); } diff --git a/services/cloud-agent-next/src/websocket/types.ts b/services/cloud-agent-next/src/websocket/types.ts index de6586f012..86f87dadbc 100644 --- a/services/cloud-agent-next/src/websocket/types.ts +++ b/services/cloud-agent-next/src/websocket/types.ts @@ -29,6 +29,7 @@ export type StreamEventType = | 'metadata' // execution metadata updates | 'error' // error occurred during execution | 'complete' // execution completed successfully + | 'wrapper_finalizing' // wrapper sealed its current admitted batch | 'interrupted' // execution was interrupted | 'started' // execution started | 'progress' // progress update (e.g., tokens consumed) diff --git a/services/cloud-agent-next/test/e2e/smoke.ts b/services/cloud-agent-next/test/e2e/smoke.ts index ec321463fe..80fd525615 100644 --- a/services/cloud-agent-next/test/e2e/smoke.ts +++ b/services/cloud-agent-next/test/e2e/smoke.ts @@ -56,10 +56,7 @@ const DEFAULT_MATRIX: Case[] = [ { lifecycle: 'unknown-model', conversation: '_' }, { lifecycle: 'waiters-clean', conversation: '_' }, - // Callback delivery via the outbound HTTP fetch from workerd. - { lifecycle: 'callback-completion', conversation: 'echo:done' }, - { lifecycle: 'callback-batch-followup', conversation: '_' }, - { lifecycle: 'callback-interrupt', conversation: '_' }, + // Callback scenarios remain manual because callbackTarget uses the internal legacy API. // Legacy-API sanity: one cold boot plus the same reused hot turn sequence. { lifecycle: 'cold-hot', conversation: 'echo:legacy', api: 'legacy' }, diff --git a/services/cloud-agent-next/test/fixtures/tool-loop-turn-events.ts b/services/cloud-agent-next/test/fixtures/tool-loop-turn-events.ts new file mode 100644 index 0000000000..4f14399d3d --- /dev/null +++ b/services/cloud-agent-next/test/fixtures/tool-loop-turn-events.ts @@ -0,0 +1,178 @@ +import type { IngestEvent } from '../../src/websocket/types.js'; + +export type ToolLoopTurnFixture = { + label: string; + rootKiloSessionId: string; + childKiloSessionId: string; + userMessageId: string; + intermediateAssistantMessageId: string; + finalAssistantMessageId: string; + finalText: string; + eventsBeforeIdle: IngestEvent[]; + childIdle: IngestEvent; + rootIdle: IngestEvent; + wrapperComplete: IngestEvent; +}; + +type ToolLoopFixtureInput = { + label: string; + intermediateText?: string; +}; + +const ROOT_KILO_SESSION_ID = 'ses_synthetic_tool_loop_root'; +const CHILD_KILO_SESSION_ID = 'ses_synthetic_tool_loop_child'; +const BASE_TIME_MS = Date.now() - 60_000; + +function timestamp(offsetMs: number): string { + return new Date(BASE_TIME_MS + offsetMs).toISOString(); +} + +function kilocode(offsetMs: number, data: Record): IngestEvent { + return { + streamEventType: 'kilocode', + timestamp: timestamp(offsetMs), + data, + }; +} + +function createToolLoopTurnFixture(input: ToolLoopFixtureInput): ToolLoopTurnFixture { + const suffix = input.label.replace(/[^a-z]/g, '_'); + const userMessageId = + input.label === 'sonnet_preamble' + ? 'msg_0123456789abSonnetUserAbCd' + : 'msg_abcdef012345FreeToolUsrAbC'; + const intermediateAssistantMessageId = `msg_synthetic_${suffix}_intermediate`; + const finalAssistantMessageId = `msg_synthetic_${suffix}_final`; + const finalText = `Synthetic final answer for ${input.label}`; + const intermediateParts: IngestEvent[] = []; + + if (input.intermediateText) { + intermediateParts.push( + kilocode(1_100, { + event: 'message.part.updated', + properties: { + part: { + id: `part_${suffix}_preamble`, + messageID: intermediateAssistantMessageId, + sessionID: ROOT_KILO_SESSION_ID, + type: 'text', + text: input.intermediateText, + }, + }, + }) + ); + } + + intermediateParts.push( + kilocode(1_200, { + event: 'message.part.updated', + properties: { + part: { + id: `part_${suffix}_tool`, + messageID: intermediateAssistantMessageId, + sessionID: ROOT_KILO_SESSION_ID, + type: 'tool', + tool: 'read', + state: { + status: 'completed', + input: { filePath: '/synthetic/read-only-input.txt' }, + output: 'sanitized placeholder output', + }, + }, + }, + }) + ); + + return { + label: input.label, + rootKiloSessionId: ROOT_KILO_SESSION_ID, + childKiloSessionId: CHILD_KILO_SESSION_ID, + userMessageId, + intermediateAssistantMessageId, + finalAssistantMessageId, + finalText, + eventsBeforeIdle: [ + kilocode(1_000, { + event: 'message.updated', + properties: { + info: { + id: intermediateAssistantMessageId, + role: 'assistant', + sessionID: ROOT_KILO_SESSION_ID, + parentID: userMessageId, + time: { completed: BASE_TIME_MS + 1_000 }, + }, + }, + }), + ...intermediateParts, + kilocode(2_000, { + event: 'session.status', + properties: { + sessionID: ROOT_KILO_SESSION_ID, + status: { type: 'busy' }, + }, + }), + kilocode(3_000, { + event: 'message.updated', + properties: { + info: { + id: finalAssistantMessageId, + role: 'assistant', + sessionID: ROOT_KILO_SESSION_ID, + parentID: userMessageId, + time: { completed: BASE_TIME_MS + 3_000 }, + }, + }, + }), + kilocode(3_100, { + event: 'message.part.updated', + properties: { + part: { + id: `part_${suffix}_final_text`, + messageID: finalAssistantMessageId, + sessionID: ROOT_KILO_SESSION_ID, + type: 'text', + text: finalText, + }, + }, + }), + ], + childIdle: kilocode(4_000, { + event: 'session.idle', + properties: { sessionID: CHILD_KILO_SESSION_ID }, + }), + rootIdle: kilocode(5_000, { + event: 'session.idle', + properties: { sessionID: ROOT_KILO_SESSION_ID }, + }), + wrapperComplete: { + streamEventType: 'complete', + timestamp: timestamp(6_000), + data: { exitCode: 0, messageIds: [userMessageId] }, + }, + }; +} + +// Reconstructed and sanitized from production-shaped tool-loop timelines. These +// fixtures contain no raw SDK archive events or production-specific values. +export const reconstructedToolLoopTurnFixtures = [ + createToolLoopTurnFixture({ + label: 'sonnet_preamble', + intermediateText: 'I will inspect the synthetic workspace before answering.', + }), + createToolLoopTurnFixture({ label: 'free_tool_only' }), +] satisfies ToolLoopTurnFixture[]; + +export const productionFixtureDenylist = [ + /agent_[0-9a-f]{8}-[0-9a-f-]{27,}/, + /ses_[a-z0-9]{20,}/i, + /msg_e8d[a-z0-9]+/i, + /wr_[0-9a-f]{32}/, + /[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}/, + /[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}/i, + /\b(?:gh[opurs]_[A-Za-z0-9_]{20,}|github_pat_[A-Za-z0-9_]{20,})\b/, + /\bBearer\s+\S+/i, + /\/Users\/[^/]+\//, + /na2-org\//, + /\/tmp\/cloud-agent-cli-logs\//, +] as const; diff --git a/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts b/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts index 67591f3999..c58fdabeaf 100644 --- a/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts +++ b/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts @@ -117,6 +117,7 @@ describe('Disconnect handling and compatibility execution RPCs', () => { await instance.handleWrapperTerminalEvent({ wrapperRunId: 'wr_idle_warm_reuse', status: 'completed', + messageIds: [], }); await instance.alarm(); diff --git a/services/cloud-agent-next/test/integration/session/events.test.ts b/services/cloud-agent-next/test/integration/session/events.test.ts index fe2a33682a..0633dcbb72 100644 --- a/services/cloud-agent-next/test/integration/session/events.test.ts +++ b/services/cloud-agent-next/test/integration/session/events.test.ts @@ -492,4 +492,68 @@ describe('Event Storage', () => { expect(result.root?.info.id).toBe('msg_00000000000000000000000002'); expect(result.missingRoot).toBeNull(); }); + + it('should select and hydrate an assistant without time.completed', async () => { + const id = env.CLOUD_AGENT_SESSION.idFromName('user_1:sess_9'); + const stub = env.CLOUD_AGENT_SESSION.get(id); + + const result = await runInDurableObject(stub, async (_instance, state) => { + const db = drizzle(state.storage, { logger: false }); + const events = createEventQueries(db, state.storage.sql); + const parentMessageId = 'msg_user_0000000000000000000001'; + const olderAssistantId = 'msg_assistant_000000000000000001'; + const latestAssistantId = 'msg_assistant_000000000000000002'; + + events.upsert({ + executionId: 'exc_1', + sessionId: 'sess_1', + streamEventType: 'kilocode', + payload: JSON.stringify({ + event: 'message.updated', + properties: { + info: { + id: olderAssistantId, + role: 'assistant', + sessionID: 'ses_root', + parentID: parentMessageId, + time: { completed: 1 }, + }, + }, + }), + timestamp: 1, + entityId: `message/${olderAssistantId}`, + }); + events.upsert({ + executionId: 'exc_1', + sessionId: 'sess_1', + streamEventType: 'kilocode', + payload: JSON.stringify({ + event: 'message.updated', + properties: { + info: { + id: latestAssistantId, + role: 'assistant', + sessionID: 'ses_root', + parentID: parentMessageId, + }, + }, + }), + timestamp: 2, + entityId: `message/${latestAssistantId}`, + }); + + return { + selected: events.getAssistantMessageForUserMessage('sess_1', 'ses_root', parentMessageId), + hydrated: events.getAssistantMessageById( + 'sess_1', + 'ses_root', + latestAssistantId, + parentMessageId + ), + }; + }); + + expect(result.selected?.info.id).toBe('msg_assistant_000000000000000002'); + expect(result.hydrated?.info.id).toBe('msg_assistant_000000000000000002'); + }); }); diff --git a/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts b/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts index a2b1219249..51f9538712 100644 --- a/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts +++ b/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts @@ -336,6 +336,7 @@ describe('handleWrapperTerminalEvent — new-path identity and message preservat await instance.handleWrapperTerminalEvent({ wrapperRunId: wrapperRunId!, status: 'completed', + messageIds: [messageId], }); const wrapperRuntimeState = await getWrapperRuntimeState(instance.ctx.storage); diff --git a/services/cloud-agent-next/test/integration/session/hot-delivery.test.ts b/services/cloud-agent-next/test/integration/session/hot-delivery.test.ts index 0d6d1ec8e1..c16a38d79a 100644 --- a/services/cloud-agent-next/test/integration/session/hot-delivery.test.ts +++ b/services/cloud-agent-next/test/integration/session/hot-delivery.test.ts @@ -8,33 +8,31 @@ * All tests follow red-green discipline. */ -import { env, runInDurableObject, listDurableObjectIds } from 'cloudflare:test'; +import { env, runInDurableObject } from 'cloudflare:test'; import { drizzle } from 'drizzle-orm/durable-sqlite'; -import { describe, it, expect, beforeEach } from 'vitest'; +import { describe, expect, it } from 'vitest'; import { listPendingSessionMessages } from '../../../src/session/pending-messages.js'; import { createPendingSessionMessage, storePendingSessionMessage, } from '../../../src/session/pending-messages.js'; import { + getSessionMessageState, listNonTerminalAcceptedMessages, putSessionMessageState, type SessionMessageState, } from '../../../src/session/session-message-state.js'; import { allocateWrapperRuntimeState, + getWrapperRuntimeState, + markWrapperFinalizing, recordMeaningfulWrapperOutput, } from '../../../src/session/wrapper-runtime-state.js'; import type { FencedWrapperDispatchRequest } from '../../../src/execution/types.js'; import { registerReadySession } from '../../helpers/session-setup.js'; describe('hot delivery — DO integration', () => { - beforeEach(async () => { - const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); - expect(ids).toHaveLength(0); - }); - - it('queued follow-up hot-delivers to a warm wrapper', async () => { + it('holds a queued follow-up while the current wrapper run finalizes', async () => { const userId = 'user_hot_deliv'; const sessionId = 'agent_hot_deliv'; const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); @@ -100,6 +98,7 @@ describe('hot delivery — DO integration', () => { wrapperRunId: wrapperState.wrapperRunId!, }; await putSessionMessageState(instance.ctx.storage, acceptedMsg); + await markWrapperFinalizing(instance.ctx.storage, wrapperState.wrapperRunId); // Queue a follow-up message const pendingMsg = createPendingSessionMessage({ @@ -110,8 +109,9 @@ describe('hot delivery — DO integration', () => { }); await storePendingSessionMessage(instance.ctx.storage, pendingMsg); - // Flush via alarm — should hot-deliver because wrapper is warm + // Flush via alarm — finalizing must hold the pending message without retry churn. await instance.alarm(); + const runtimeAfterDelivery = await getWrapperRuntimeState(instance.ctx.storage); const pending = await listPendingSessionMessages(instance.ctx.storage); const acceptedMessages = await listNonTerminalAcceptedMessages( @@ -120,28 +120,197 @@ describe('hot delivery — DO integration', () => { ); const executions = await instance.getExecutions(); - return { capturedPlans, pending, acceptedMessages, executions }; + return { capturedPlans, pending, acceptedMessages, executions, runtimeAfterDelivery }; }); - // Follow-up message was delivered (orchestrator received the plan) - expect(result.capturedPlans.length).toBeGreaterThanOrEqual(1); - const deliveredPlan = result.capturedPlans.find( + expect(result.capturedPlans).toHaveLength(0); + expect(result.executions).toHaveLength(0); + expect(result.pending).toHaveLength(1); + expect(result.pending[0]?.messageId).toBe(followUpMessageId); + expect(result.pending[0]?.flushAttempts).toBeUndefined(); + expect(result.runtimeAfterDelivery.finalizingWrapperRunId).toBe( + result.runtimeAfterDelivery.wrapperRunId + ); + expect(result.acceptedMessages.map(message => message.messageId)).not.toContain( + followUpMessageId + ); + }); + + it('normal wrapper completion releases the finalizing hold and drains one follow-up under a fresh run', async () => { + const userId = 'user_complete_drain'; + const sessionId = 'agent_complete_drain'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + const currentMessageId = 'msg_018f1e2d3c4bCmplteDrain001'; + const followUpMessageId = 'msg_018f1e2d3c4bCmplteDrain002'; + + const result = await runInDurableObject(stub, async instance => { + const capturedPlans: FencedWrapperDispatchRequest[] = []; + instance['orchestrator'] = { + execute: async (plan: FencedWrapperDispatchRequest) => { + capturedPlans.push(plan); + return { messageId: plan.turn.messageId, kiloSessionId: 'kilo_complete_drain' }; + }, + }; + instance['physicalWrapperObserver'] = async () => ({ + status: 'present' as const, + observed: [ + { + representation: 'process' as const, + id: 'wrapper-complete-drain', + port: 4_173, + instanceId: 'instance_complete_drain', + instanceGeneration: 1, + }, + ], + }); + + await registerReadySession(instance, { + sessionId, + userId, + orgId: 'org_complete_drain', + kiloSessionId: '33333333-3333-4333-8333-333333333333', + prompt: 'initial prompt', + mode: 'code', + model: 'test-model', + kilocodeToken: 'token-complete-drain', + }); + + const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); + await instance.ctx.storage.put('wrapper_lease', { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_complete_drain', instanceGeneration: 1 }, + }); + await putSessionMessageState(instance.ctx.storage, { + messageId: currentMessageId, + status: 'accepted', + prompt: '/status', + admissionSnapshot: { + turn: { type: 'command', messageId: currentMessageId, command: 'status', arguments: '' }, + agent: { mode: 'code', model: 'test-model' }, + }, + createdAt: 1, + acceptedAt: 1, + wrapperRunId: wrapperState.wrapperRunId, + }); + await markWrapperFinalizing(instance.ctx.storage, wrapperState.wrapperRunId); + await storePendingSessionMessage( + instance.ctx.storage, + createPendingSessionMessage({ + messageId: followUpMessageId, + role: 'user', + content: 'follow up after normal completion', + createdAt: 2, + }) + ); + + await instance.alarm(); + const plansBeforeComplete = capturedPlans.length; + await instance.handleWrapperTerminalEvent({ + wrapperRunId: wrapperState.wrapperRunId, + status: 'completed', + messageIds: [currentMessageId], + }); + const runtimeAfterComplete = await getWrapperRuntimeState(instance.ctx.storage); + await instance.alarm(); + await instance.alarm(); + + return { + currentWrapperRunId: wrapperState.wrapperRunId, + plansBeforeComplete, + capturedPlans, + runtimeAfterComplete, + remaining: await listPendingSessionMessages(instance.ctx.storage), + currentMessage: await getSessionMessageState(instance.ctx.storage, currentMessageId), + }; + }); + + const followUpPlans = result.capturedPlans.filter( plan => plan.turn.messageId === followUpMessageId ); - expect(deliveredPlan).not.toBeUndefined(); - expect(deliveredPlan?.turn.messageId).toBe(followUpMessageId); + expect(result.plansBeforeComplete).toBe(0); + expect(result.currentMessage).toMatchObject({ + status: 'completed', + completionSource: 'idle_reconciliation', + }); + expect(result.runtimeAfterComplete.finalizingWrapperRunId).toBeUndefined(); + expect(followUpPlans).toHaveLength(1); + expect(followUpPlans[0]?.wrapper.fence.wrapperRunId).not.toBe(result.currentWrapperRunId); + expect(result.remaining).toHaveLength(0); + }); - // Phase 5 Slice 2: warm-followup flush plan has no executionId - expect(deliveredPlan).not.toHaveProperty('executionId'); + it('holds pending delivery through physical cleanup and drains after confirmed absence', async () => { + const userId = 'user_cleanup_hold'; + const sessionId = 'agent_cleanup_hold'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + const followUpMessageId = 'msg_018f1e2d3c4bCleanupHold001'; - // Phase 5 Slice 2: warm-followup flush does not create an execution row - expect(result.executions).toHaveLength(0); + const result = await runInDurableObject(stub, async instance => { + const capturedPlans: FencedWrapperDispatchRequest[] = []; + instance['orchestrator'] = { + execute: async (plan: FencedWrapperDispatchRequest) => { + capturedPlans.push(plan); + return { messageId: plan.turn.messageId, kiloSessionId: 'kilo_cleanup_hold' }; + }, + }; + instance['physicalWrapperStopper'] = async () => ({ status: 'absent' as const }); + await registerReadySession(instance, { + sessionId, + userId, + orgId: 'org_cleanup_hold', + kiloSessionId: '22222222-2222-4222-8222-222222222222', + prompt: 'initial prompt', + mode: 'code', + model: 'test-model', + kilocodeToken: 'token-cleanup-hold', + }); + await instance.ctx.storage.put('wrapper_lease', { + state: 'stop_needed', + nextInstanceGeneration: 2, + target: { kind: 'session' }, + reason: 'terminal-failed', + requestedAt: 1, + nextAttemptAt: Date.now() + 60_000, + attempts: 0, + }); + await storePendingSessionMessage( + instance.ctx.storage, + createPendingSessionMessage({ + messageId: followUpMessageId, + role: 'user', + content: 'follow after cleanup', + createdAt: 1, + }) + ); + + await instance.alarm(); + const held = await listPendingSessionMessages(instance.ctx.storage); + await instance.ctx.storage.put('wrapper_lease', { + state: 'stop_needed', + nextInstanceGeneration: 2, + target: { kind: 'session' }, + reason: 'terminal-failed', + requestedAt: 1, + nextAttemptAt: 1, + attempts: 0, + }); + await instance.alarm(); + await instance.alarm(); - // Pending queue is drained - expect(result.pending).toHaveLength(0); + return { + held, + remaining: await listPendingSessionMessages(instance.ctx.storage), + capturedPlans, + }; + }); - // Both messages are now accepted - expect(result.acceptedMessages.length).toBeGreaterThanOrEqual(1); - expect(result.acceptedMessages.map(m => m.messageId)).toContain(followUpMessageId); + expect(result.held).toHaveLength(1); + expect(result.held[0]?.flushAttempts).toBeUndefined(); + expect(result.remaining).toHaveLength(0); + expect(result.capturedPlans.map(plan => plan.turn.messageId)).toContain(followUpMessageId); }); }); diff --git a/services/cloud-agent-next/test/integration/session/idle-reconciliation.test.ts b/services/cloud-agent-next/test/integration/session/idle-reconciliation.test.ts index 75d5f24609..671820f87c 100644 --- a/services/cloud-agent-next/test/integration/session/idle-reconciliation.test.ts +++ b/services/cloud-agent-next/test/integration/session/idle-reconciliation.test.ts @@ -1,30 +1,16 @@ -/** - * Integration tests for idle reconciliation scheduling. - * - * Phase 6 remediation: root session.idle must record lastWrapperIdleAt and - * idleReconcileAfter, drive checkIdleReconciliation by that deadline, and - * reschedule the alarm. Meaningful output after idle must clear the idle - * state so reconciliation is cancelled. - */ - -import { env, runInDurableObject, listDurableObjectIds } from 'cloudflare:test'; -import { describe, it, expect, beforeEach } from 'vitest'; -import { drizzle } from 'drizzle-orm/durable-sqlite'; -import { createEventQueries } from '../../../src/session/queries/events.js'; +import { env, listDurableObjectIds, runInDurableObject } from 'cloudflare:test'; +import { beforeEach, describe, expect, it } from 'vitest'; import { - getWrapperLease, - getWrapperRuntimeState, allocateWrapperRuntimeState, + getWrapperRuntimeState, } from '../../../src/session/wrapper-runtime-state.js'; import { getSessionMessageState, - listNonTerminalAcceptedMessages, putSessionMessageState, - type SessionMessageState, } from '../../../src/session/session-message-state.js'; import { registerReadySession } from '../../helpers/session-setup.js'; -describe('idle reconciliation scheduling', () => { +describe('idle lifecycle integration', () => { beforeEach(async () => { const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); await Promise.all( @@ -36,573 +22,71 @@ describe('idle reconciliation scheduling', () => { ); }); - it('root session.idle records lastWrapperIdleAt and idleReconcileAfter', async () => { - const userId = 'user_idle_schedule'; - const sessionId = 'agent_idle_schedule'; - const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); - const stub = env.CLOUD_AGENT_SESSION.get(doId); - - const result = await runInDurableObject(stub, async instance => { - await registerReadySession(instance, { - sessionId, - userId, - orgId: 'org_idle_schedule', - kiloSessionId: 'aaaaaaaa-aaaa-4aaa-aaaa-aaaaaaaaaaaa', - prompt: 'initial prompt', - mode: 'code', - model: 'test-model', - kilocodeToken: 'token-idle-schedule', - }); - - const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); - const { wrapperRunId, wrapperConnectionId, wrapperGeneration } = wrapperState; - - const handler = await (instance as any).getIngestHandler(); - const ws = { - deserializeAttachment: () => ({ - wrapperRunId, - sessionId, - connectedAt: Date.now(), - kiloSessionState: { captured: false }, - lastHeartbeatUpdate: Date.now(), - lastEventAtUpdate: Date.now(), - wrapperGeneration, - wrapperConnectionId, - }), - serializeAttachment: () => {}, - send: () => {}, - } as unknown as WebSocket; - - const beforeIngest = Date.now(); - await handler.handleIngestMessage( - ws, - JSON.stringify({ - streamEventType: 'kilocode', - data: { - event: 'session.idle', - properties: { - sessionID: 'aaaaaaaa-aaaa-4aaa-aaaa-aaaaaaaaaaaa', - }, - }, - timestamp: new Date().toISOString(), - }) - ); - - const runtimeState = await getWrapperRuntimeState(instance.ctx.storage); - - return { runtimeState, beforeIngest, afterIngest: Date.now() }; - }); - - expect(result.runtimeState.lastWrapperIdleAt).toBeDefined(); - expect(result.runtimeState.lastWrapperIdleAt).toBeGreaterThanOrEqual(result.beforeIngest); - expect(result.runtimeState.lastWrapperIdleAt).toBeLessThanOrEqual(result.afterIngest); - - expect(result.runtimeState.idleReconcileAfter).toBeDefined(); - expect(result.runtimeState.idleReconcileAfter).toBeGreaterThanOrEqual( - result.beforeIngest + 10_000 - ); - expect(result.runtimeState.idleReconcileAfter).toBeLessThanOrEqual(result.afterIngest + 20_000); - }); - - it('idle reconciliation does not run before idleReconcileAfter', async () => { - const userId = 'user_idle_before'; - const sessionId = 'agent_idle_before'; - const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); - const stub = env.CLOUD_AGENT_SESSION.get(doId); - - const result = await runInDurableObject(stub, async instance => { - await registerReadySession(instance, { - sessionId, - userId, - orgId: 'org_idle_before', - kiloSessionId: 'bbbbbbbb-bbbb-4bbb-bbbb-bbbbbbbbbbbb', - prompt: 'initial prompt', - mode: 'code', - model: 'test-model', - kilocodeToken: 'token-idle-before', - }); - - const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); - const { wrapperRunId, wrapperConnectionId } = wrapperState; - - const acceptedMessage: SessionMessageState = { - messageId: 'msg_018f1e2d3c4b00000000000001', - status: 'accepted', - prompt: 'hello', - createdAt: Date.now(), - acceptedAt: Date.now(), - wrapperRunId: wrapperRunId!, - }; - await putSessionMessageState(instance.ctx.storage, acceptedMessage); - - const future = Date.now() + 60_000; - await instance.ctx.storage.put('wrapper_runtime_state', { - wrapperGeneration: wrapperState.wrapperGeneration, - wrapperConnectionId, - wrapperRunId, - lastWrapperIdleAt: Date.now(), - idleReconcileAfter: future, - }); - - await instance.alarm(); - - const nonTerminalMessages = await listNonTerminalAcceptedMessages( - instance.ctx.storage, - wrapperRunId! - ); - - return { nonTerminalMessages }; - }); - - expect(result.nonTerminalMessages).toHaveLength(1); - expect(result.nonTerminalMessages[0]?.status).toBe('accepted'); - }); - - it('idle reconciliation fails accepted messages with missing_assistant_reply after idleReconcileAfter', async () => { - const userId = 'user_idle_after'; - const sessionId = 'agent_idle_after'; - const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); - const stub = env.CLOUD_AGENT_SESSION.get(doId); - - const result = await runInDurableObject(stub, async (instance, state) => { - await registerReadySession(instance, { - sessionId, - userId, - orgId: 'org_idle_after', - kiloSessionId: 'cccccccc-cccc-4ccc-cccc-cccccccccccc', - prompt: 'initial prompt', - mode: 'code', - model: 'test-model', - kilocodeToken: 'token-idle-after', - }); - const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); - const messageId = 'msg_018f1e2d3c4b00000000000002'; - await putSessionMessageState(instance.ctx.storage, { - messageId, - status: 'accepted', - prompt: 'hello', - createdAt: Date.now(), - acceptedAt: Date.now(), - wrapperRunId: wrapperState.wrapperRunId!, - }); - const past = Date.now() - 1; - await instance.ctx.storage.put('wrapper_runtime_state', { - wrapperGeneration: wrapperState.wrapperGeneration, - wrapperConnectionId: wrapperState.wrapperConnectionId, - wrapperRunId: wrapperState.wrapperRunId, - lastWrapperIdleAt: past - 15_000, - idleReconcileAfter: past, - }); - await instance.alarm(); - - const nonTerminalMessages = await listNonTerminalAcceptedMessages( - instance.ctx.storage, - wrapperState.wrapperRunId! - ); - const failedMessage = await getSessionMessageState(instance.ctx.storage, messageId); - const events = createEventQueries( - drizzle(state.storage, { logger: false }), - state.storage.sql - ); - - return { - nonTerminalMessages, - failedMessage, - failedEvents: events.findByFilters({ eventTypes: ['cloud.message.failed'] }), - lease: await getWrapperLease(instance.ctx.storage), - }; - }); - - expect(result.nonTerminalMessages).toHaveLength(0); - expect(result.failedMessage).toMatchObject({ - failureStage: 'post_dispatch_no_activity', - failureCode: 'missing_assistant_reply', - }); - - expect(result.failedEvents).toHaveLength(1); - expect(JSON.parse(result.failedEvents[0].payload)).toMatchObject({ - messageId: 'msg_018f1e2d3c4b00000000000002', - status: 'failed', - error: 'No assistant reply found during reconciliation', - delivery: 'sent', - accepted: true, - completionSource: 'idle_reconciliation', - }); - expect(result.lease).toMatchObject({ state: 'stop_needed', reason: 'terminal-failed' }); - }); - - it('idle reconciliation treats object-shaped assistant errors as failed replies', async () => { - const userId = 'user_idle_object_error'; - const sessionId = 'agent_idle_object_error'; + it('persists raw root idle without using it as a success boundary', async () => { + const userId = 'user_idle_no_fallback'; + const sessionId = 'agent_idle_no_fallback'; const stub = env.CLOUD_AGENT_SESSION.get( env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) ); - const result = await runInDurableObject(stub, async (instance, state) => { + const result = await runInDurableObject(stub, async instance => { await registerReadySession(instance, { sessionId, userId, - orgId: 'org_idle_object_error', - kiloSessionId: 'edededed-eded-4ede-8ede-edededededed', + kiloSessionId: 'aaaaaaaa-aaaa-4aaa-aaaa-aaaaaaaaaaaa', prompt: 'initial prompt', mode: 'code', model: 'test-model', - kilocodeToken: 'token-idle-object-error', + kilocodeToken: 'token-idle-no-fallback', }); const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); - const messageId = 'msg_018f1e2d3c4bObjErrIdleAbCd'; + const messageId = 'msg_018f1e2d3c4bNoIdleFallbkAB'; await putSessionMessageState(instance.ctx.storage, { messageId, status: 'accepted', - prompt: 'hello', + prompt: 'remain accepted', createdAt: Date.now(), acceptedAt: Date.now(), - wrapperRunId: wrapperState.wrapperRunId!, - }); - const events = createEventQueries( - drizzle(state.storage, { logger: false }), - state.storage.sql - ); - events.upsert({ - executionId: '', - sessionId, - streamEventType: 'kilocode', - payload: JSON.stringify({ - event: 'message.updated', - properties: { - info: { - id: 'assistant_obj_error_idle', - role: 'assistant', - sessionID: 'edededed-eded-4ede-8ede-edededededed', - parentID: messageId, - error: { name: 'UnknownError', data: { message: 'provider failed during idle' } }, - }, - }, - }), - timestamp: Date.now(), - entityId: 'message/assistant_obj_error_idle', - }); - const past = Date.now() - 1; - await instance.ctx.storage.put('wrapper_runtime_state', { - wrapperGeneration: wrapperState.wrapperGeneration, - wrapperConnectionId: wrapperState.wrapperConnectionId, wrapperRunId: wrapperState.wrapperRunId, - lastWrapperIdleAt: past - 15_000, - idleReconcileAfter: past, - }); - await instance.alarm(); - return { - failedEvents: events.findByFilters({ eventTypes: ['cloud.message.failed'] }), - completedEvents: events.findByFilters({ eventTypes: ['cloud.message.completed'] }), - lease: await getWrapperLease(instance.ctx.storage), - }; - }); - - expect(result.failedEvents).toHaveLength(1); - expect(JSON.parse(result.failedEvents[0].payload)).toMatchObject({ - messageId: 'msg_018f1e2d3c4bObjErrIdleAbCd', - status: 'failed', - error: 'provider failed during idle', - completionSource: 'idle_reconciliation', - }); - expect(result.completedEvents).toHaveLength(0); - expect(result.lease).toMatchObject({ state: 'stop_needed', reason: 'terminal-failed' }); - }); - - it('meaningful wrapper output refreshes liveness without clearing idle state', async () => { - const userId = 'user_idle_output'; - const sessionId = 'agent_idle_output'; - const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); - const stub = env.CLOUD_AGENT_SESSION.get(doId); - - const result = await runInDurableObject(stub, async instance => { - await registerReadySession(instance, { - sessionId, - userId, - orgId: 'org_idle_output', - kiloSessionId: 'dddddddd-dddd-4ddd-dddd-dddddddddddd', - prompt: 'initial prompt', - mode: 'code', - model: 'test-model', - kilocodeToken: 'token-idle-output', - }); - - const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); - const { wrapperRunId, wrapperConnectionId, wrapperGeneration } = wrapperState; - - // Set idle state - const idleAt = Date.now(); - const reconcileAt = idleAt + 15_000; - const keepWarmAt = idleAt + 5 * 60 * 1000; - await instance.ctx.storage.put('wrapper_runtime_state', { - wrapperGeneration, - wrapperConnectionId, - wrapperRunId, - lastWrapperIdleAt: idleAt, - idleReconcileAfter: reconcileAt, - wrapperIdleDeadlineAt: keepWarmAt, }); - const handler = await (instance as any).getIngestHandler(); + const handler = await instance['getIngestHandler'](); const ws = { deserializeAttachment: () => ({ - wrapperRunId, + wrapperRunId: wrapperState.wrapperRunId, sessionId, connectedAt: Date.now(), kiloSessionState: { captured: false }, lastHeartbeatUpdate: Date.now(), lastEventAtUpdate: Date.now(), - wrapperGeneration, - wrapperConnectionId, + wrapperGeneration: wrapperState.wrapperGeneration, + wrapperConnectionId: wrapperState.wrapperConnectionId, }), serializeAttachment: () => {}, send: () => {}, } as unknown as WebSocket; - // Simulate a non-fatal error event. This is post-completion infrastructure - // output: it must refresh wrapper liveness but must NOT disarm the idle - // reconciler that is about to finalize the in-flight message. - await handler.handleIngestMessage( - ws, - JSON.stringify({ - streamEventType: 'error', - data: { fatal: false, error: 'something happened' }, - timestamp: new Date().toISOString(), - }) - ); - - const runtimeState = await getWrapperRuntimeState(instance.ctx.storage); - return { runtimeState, idleAt, reconcileAt, keepWarmAt }; - }); - - // Idle reconciliation fields stay armed so the reconciler still fires. - expect(result.runtimeState.lastWrapperIdleAt).toBe(result.idleAt); - expect(result.runtimeState.idleReconcileAfter).toBe(result.reconcileAt); - expect(result.runtimeState.wrapperIdleDeadlineAt).toBe(result.keepWarmAt); - // Liveness is refreshed by the output. - expect(result.runtimeState.lastWrapperMessageAt).toBeDefined(); - }); - - it('root session.idle records wrapperIdleDeadlineAt for keep-warm', async () => { - const userId = 'user_idle_warm'; - const sessionId = 'agent_idle_warm'; - const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); - const stub = env.CLOUD_AGENT_SESSION.get(doId); - - const result = await runInDurableObject(stub, async instance => { - await registerReadySession(instance, { - sessionId, - userId, - orgId: 'org_idle_warm', - kiloSessionId: 'eeeeeeee-eeee-4eee-eeee-eeeeeeeeeeee', - prompt: 'initial prompt', - mode: 'code', - model: 'test-model', - kilocodeToken: 'token-idle-warm', - }); - - const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); - const { wrapperRunId, wrapperConnectionId, wrapperGeneration } = wrapperState; - - const handler = await (instance as any).getIngestHandler(); - const ws = { - deserializeAttachment: () => ({ - wrapperRunId, - sessionId, - connectedAt: Date.now(), - kiloSessionState: { captured: false }, - lastHeartbeatUpdate: Date.now(), - lastEventAtUpdate: Date.now(), - wrapperGeneration, - wrapperConnectionId, - }), - serializeAttachment: () => {}, - send: () => {}, - } as unknown as WebSocket; - - const beforeIngest = Date.now(); await handler.handleIngestMessage( ws, JSON.stringify({ streamEventType: 'kilocode', data: { event: 'session.idle', - properties: { - sessionID: 'eeeeeeee-eeee-4eee-eeee-eeeeeeeeeeee', - }, + properties: { sessionID: 'aaaaaaaa-aaaa-4aaa-aaaa-aaaaaaaaaaaa' }, }, timestamp: new Date().toISOString(), }) ); - - const runtimeState = await getWrapperRuntimeState(instance.ctx.storage); - - return { runtimeState, beforeIngest, afterIngest: Date.now() }; - }); - - expect(result.runtimeState.wrapperIdleDeadlineAt).toBeDefined(); - expect(result.runtimeState.wrapperIdleDeadlineAt).toBeGreaterThanOrEqual( - result.beforeIngest + 4 * 60 * 1000 - ); - expect(result.runtimeState.wrapperIdleDeadlineAt).toBeLessThanOrEqual( - result.afterIngest + 6 * 60 * 1000 - ); - }); - - it('keep-warm cleanup does not run before wrapperIdleDeadlineAt', async () => { - const userId = 'user_warm_before'; - const sessionId = 'agent_warm_before'; - const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); - const stub = env.CLOUD_AGENT_SESSION.get(doId); - - let stopWrapperCalled = false; - - const result = await runInDurableObject(stub, async instance => { - instance['physicalWrapperStopper'] = async () => { - stopWrapperCalled = true; - return { status: 'absent' }; - }; - - await registerReadySession(instance, { - sessionId, - userId, - orgId: 'org_warm_before', - kiloSessionId: 'ffffffff-ffff-4fff-ffff-ffffffffffff', - prompt: 'initial prompt', - mode: 'code', - model: 'test-model', - kilocodeToken: 'token-warm-before', - }); - - const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); - const { wrapperRunId, wrapperConnectionId } = wrapperState; - - const future = Date.now() + 60_000; - await instance.ctx.storage.put('wrapper_runtime_state', { - wrapperGeneration: wrapperState.wrapperGeneration, - wrapperConnectionId, - wrapperRunId, - wrapperIdleDeadlineAt: future, - idleReconcileAfter: future, - }); - await instance.alarm(); - const runtimeState = await getWrapperRuntimeState(instance.ctx.storage); - return { runtimeState, stopWrapperCalled }; - }); - - expect(result.stopWrapperCalled).toBe(false); - expect(result.runtimeState.wrapperConnectionId).toBe(result.runtimeState.wrapperConnectionId); - }); - - it('keep-warm cleanup tears down idle wrapper after deadline with no work', async () => { - const userId = 'user_warm_teardown'; - const sessionId = 'agent_warm_teardown'; - const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); - const stub = env.CLOUD_AGENT_SESSION.get(doId); - - let stopWrapperCalled = false; - - const result = await runInDurableObject(stub, async instance => { - instance['physicalWrapperStopper'] = async () => { - stopWrapperCalled = true; - return { status: 'absent' }; - }; - - await registerReadySession(instance, { - sessionId, - userId, - orgId: 'org_warm_teardown', - kiloSessionId: '11111111-1111-4111-1111-111111111111', - prompt: 'initial prompt', - mode: 'code', - model: 'test-model', - kilocodeToken: 'token-warm-teardown', - }); - - const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); - const { wrapperRunId, wrapperConnectionId } = wrapperState; - - const past = Date.now() - 1; - await instance.ctx.storage.put('wrapper_runtime_state', { - wrapperGeneration: wrapperState.wrapperGeneration, - wrapperConnectionId, - wrapperRunId, - wrapperIdleDeadlineAt: past, - }); - - await instance.alarm(); - - const runtimeState = await getWrapperRuntimeState(instance.ctx.storage); return { - runtimeState, - lease: await getWrapperLease(instance.ctx.storage), - stopWrapperCalled, - }; - }); - - expect(result.stopWrapperCalled).toBe(false); - expect(result.runtimeState.wrapperConnectionId).toBeUndefined(); - expect(result.lease).toMatchObject({ state: 'stop_needed', reason: 'keep-warm-expired' }); - }); - - it('keep-warm cleanup clears idle state when work exists after deadline', async () => { - const userId = 'user_warm_work'; - const sessionId = 'agent_warm_work'; - const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); - const stub = env.CLOUD_AGENT_SESSION.get(doId); - - let stopWrapperCalled = false; - - const result = await runInDurableObject(stub, async instance => { - instance['physicalWrapperStopper'] = async () => { - stopWrapperCalled = true; - return { status: 'absent' }; - }; - - await registerReadySession(instance, { - sessionId, - userId, - orgId: 'org_warm_work', - kiloSessionId: '22222222-2222-4222-2222-222222222222', - prompt: 'initial prompt', - mode: 'code', - model: 'test-model', - kilocodeToken: 'token-warm-work', - }); - - const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); - const { wrapperRunId, wrapperConnectionId } = wrapperState; - - // Place an accepted message so work exists - const acceptedMessage: SessionMessageState = { - messageId: 'msg_018f1e2d3c4b00000000000099', - status: 'accepted', - prompt: 'hello', - createdAt: Date.now(), - acceptedAt: Date.now(), - wrapperRunId: wrapperRunId!, + message: await getSessionMessageState(instance.ctx.storage, messageId), + runtime: await getWrapperRuntimeState(instance.ctx.storage), }; - await putSessionMessageState(instance.ctx.storage, acceptedMessage); - - const past = Date.now() - 1; - const future = Date.now() + 60_000; - await instance.ctx.storage.put('wrapper_runtime_state', { - wrapperGeneration: wrapperState.wrapperGeneration, - wrapperConnectionId, - wrapperRunId, - wrapperIdleDeadlineAt: past, - idleReconcileAfter: future, - }); - - await instance.alarm(); - - const runtimeState = await getWrapperRuntimeState(instance.ctx.storage); - return { runtimeState, stopWrapperCalled }; }); - expect(result.stopWrapperCalled).toBe(false); - expect(result.runtimeState.wrapperIdleDeadlineAt).toBeUndefined(); + expect(result.message?.status).toBe('accepted'); + expect(result.runtime).not.toHaveProperty('lastWrapperIdleAt'); + expect(result.runtime).not.toHaveProperty('idleReconcileAfter'); }); }); diff --git a/services/cloud-agent-next/test/integration/session/tool-loop-terminalization.test.ts b/services/cloud-agent-next/test/integration/session/tool-loop-terminalization.test.ts new file mode 100644 index 0000000000..7290064db6 --- /dev/null +++ b/services/cloud-agent-next/test/integration/session/tool-loop-terminalization.test.ts @@ -0,0 +1,312 @@ +import { env, listDurableObjectIds, runInDurableObject } from 'cloudflare:test'; +import { drizzle } from 'drizzle-orm/durable-sqlite'; +import { beforeEach, describe, expect, it } from 'vitest'; +import type { CloudAgentSession } from '../../../src/persistence/CloudAgentSession.js'; +import { createEventQueries } from '../../../src/session/queries/events.js'; +import { + allocateWrapperRuntimeState, + recordWrapperDispatchingMessage, + type ActiveWrapperRuntimeState, +} from '../../../src/session/wrapper-runtime-state.js'; +import { storePendingSessionMessage } from '../../../src/session/pending-messages.js'; +import { putSessionMessageState } from '../../../src/session/session-message-state.js'; +import type { IngestAttachment, IngestHandler } from '../../../src/websocket/ingest.js'; +import type { IngestEvent } from '../../../src/websocket/types.js'; +import { + productionFixtureDenylist, + reconstructedToolLoopTurnFixtures, + type ToolLoopTurnFixture, +} from '../../fixtures/tool-loop-turn-events.js'; +import { registerReadySession } from '../../helpers/session-setup.js'; + +async function resetSessions(): Promise { + const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); + await Promise.all( + ids.map(id => + runInDurableObject(env.CLOUD_AGENT_SESSION.get(id), instance => + instance.ctx.storage.deleteAll() + ) + ) + ); +} + +function createIngestSocket(sessionId: string, wrapperState: ActiveWrapperRuntimeState): WebSocket { + const attachment: IngestAttachment = { + wrapperRunId: wrapperState.wrapperRunId, + sessionId, + connectedAt: Date.now(), + kiloSessionState: { captured: false }, + lastHeartbeatUpdate: Date.now(), + lastEventAtUpdate: Date.now(), + wrapperGeneration: wrapperState.wrapperGeneration, + wrapperConnectionId: wrapperState.wrapperConnectionId, + }; + return { + deserializeAttachment: () => attachment, + serializeAttachment: () => {}, + send: () => {}, + close: () => {}, + } as unknown as WebSocket; +} + +async function getIngestHandler(instance: CloudAgentSession): Promise { + const privateAccess = instance as unknown as { + getIngestHandler(): Promise; + }; + return privateAccess.getIngestHandler(); +} + +async function replayEvents( + handler: IngestHandler, + ws: WebSocket, + events: IngestEvent[] +): Promise { + for (const event of events) { + await handler.handleIngestMessage(ws, JSON.stringify(event)); + } +} + +async function seedAcceptedTurn( + instance: CloudAgentSession, + fixture: ToolLoopTurnFixture, + sessionId: string +) { + await registerReadySession(instance, { + sessionId, + userId: `user_${fixture.label}`, + orgId: `org_${fixture.label}`, + kiloSessionId: fixture.rootKiloSessionId, + prompt: 'Inspect the synthetic workspace without modifying files.', + mode: 'code', + model: 'test-model', + kilocodeToken: `token_${fixture.label}`, + }); + const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); + await putSessionMessageState(instance.ctx.storage, { + messageId: fixture.userMessageId, + status: 'accepted', + prompt: 'Inspect the synthetic workspace without modifying files.', + createdAt: Date.now(), + acceptedAt: Date.now(), + wrapperRunId: wrapperState.wrapperRunId, + }); + return wrapperState; +} + +function completedEvents(state: DurableObjectState) { + return createEventQueries( + drizzle(state.storage, { logger: false }), + state.storage.sql + ).findByFilters({ + eventTypes: ['cloud.message.completed'], + }); +} + +describe('tool-loop terminalization replay', () => { + beforeEach(resetSessions); + + it('contains only reconstructed sanitized fixture values', () => { + const serializedFixtures = JSON.stringify(reconstructedToolLoopTurnFixtures); + for (const forbiddenPattern of productionFixtureDenylist) { + expect(serializedFixtures).not.toMatch(forbiddenPattern); + } + }); + + for (const fixture of reconstructedToolLoopTurnFixtures) { + it(`settles ${fixture.label} to the final assistant only after wrapper completion`, async () => { + const sessionId = `agent_synthetic_${fixture.label}`; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`user_${fixture.label}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async (instance, state) => { + const wrapperState = await seedAcceptedTurn(instance, fixture, sessionId); + const handler = await getIngestHandler(instance); + const ws = createIngestSocket(sessionId, wrapperState); + + await replayEvents(handler, ws, fixture.eventsBeforeIdle.slice(0, 1)); + const afterIntermediate = await instance.getMessageResult(fixture.userMessageId); + const completedAfterIntermediate = completedEvents(state).length; + + await replayEvents(handler, ws, fixture.eventsBeforeIdle.slice(1)); + await replayEvents(handler, ws, [fixture.childIdle]); + const afterChildIdle = await instance.getMessageResult(fixture.userMessageId); + await replayEvents(handler, ws, [fixture.rootIdle]); + const afterRootIdle = await instance.getMessageResult(fixture.userMessageId); + const completedAfterRootIdle = completedEvents(state).length; + + await replayEvents(handler, ws, [fixture.wrapperComplete]); + const afterWrapperComplete = await instance.getMessageResult(fixture.userMessageId); + const completedAfterWrapperComplete = completedEvents(state).length; + + return { + afterIntermediate, + completedAfterIntermediate, + afterChildIdle, + afterRootIdle, + completedAfterRootIdle, + afterWrapperComplete, + completedAfterWrapperComplete, + }; + }); + + expect(result.afterIntermediate).toMatchObject({ + type: 'found', + result: { status: 'running' }, + }); + expect(result.completedAfterIntermediate).toBe(0); + expect(result.afterChildIdle).toMatchObject({ + type: 'found', + result: { status: 'running' }, + }); + expect(result.afterRootIdle).toMatchObject({ + type: 'found', + result: { status: 'running' }, + }); + expect(result.completedAfterRootIdle).toBe(0); + expect(result.afterWrapperComplete).toMatchObject({ + type: 'found', + result: { + status: 'completed', + completionSource: 'idle_reconciliation', + assistant: { + messageId: fixture.finalAssistantMessageId, + text: fixture.finalText, + }, + }, + }); + expect(result.completedAfterWrapperComplete).toBe(1); + }); + } + + it('settles the final assistant from bare wrapper completion when root idle is lost', async () => { + const fixture = reconstructedToolLoopTurnFixtures[0]; + if (!fixture) throw new Error('Expected a reconstructed tool-loop fixture'); + const sessionId = 'agent_synthetic_bare_complete'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`user_synthetic_bare_complete:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async (instance, state) => { + const wrapperState = await seedAcceptedTurn(instance, fixture, sessionId); + const handler = await getIngestHandler(instance); + const ws = createIngestSocket(sessionId, wrapperState); + + await replayEvents(handler, ws, fixture.eventsBeforeIdle); + await replayEvents(handler, ws, [fixture.wrapperComplete]); + + return { + messageResult: await instance.getMessageResult(fixture.userMessageId), + completedEventCount: completedEvents(state).length, + }; + }); + + expect(result.messageResult).toMatchObject({ + type: 'found', + result: { + status: 'completed', + completionSource: 'idle_reconciliation', + assistant: { + messageId: fixture.finalAssistantMessageId, + text: fixture.finalText, + }, + }, + }); + expect(result.completedEventCount).toBe(1); + }); + + it('repairs a legacy complete that wins the acceptance persistence race', async () => { + const fixture = reconstructedToolLoopTurnFixtures[0]; + if (!fixture) throw new Error('Expected a reconstructed tool-loop fixture'); + const sessionId = 'agent_synthetic_legacy_complete_race'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`user_synthetic_legacy_complete_race:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId: 'user_synthetic_legacy_complete_race', + kiloSessionId: fixture.rootKiloSessionId, + prompt: 'Inspect the synthetic workspace without modifying files.', + mode: 'code', + model: 'test-model', + kilocodeToken: 'token_synthetic_legacy_complete_race', + }); + const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); + await instance.ctx.storage.put('wrapper_runtime_state', { + wrapperGeneration: wrapperState.wrapperGeneration, + wrapperConnectionId: wrapperState.wrapperConnectionId, + wrapperRunId: wrapperState.wrapperRunId, + lastWrapperConnectedAt: wrapperState.lastWrapperConnectedAt, + }); + await recordWrapperDispatchingMessage( + instance.ctx.storage, + wrapperState, + fixture.userMessageId + ); + await storePendingSessionMessage(instance.ctx.storage, { + messageId: fixture.userMessageId, + content: 'Inspect the synthetic workspace without modifying files.', + createdAt: Date.now(), + intent: { + turn: { + type: 'prompt', + messageId: fixture.userMessageId, + prompt: 'Inspect the synthetic workspace without modifying files.', + }, + agent: { mode: 'code', model: 'test-model' }, + }, + }); + const handler = await getIngestHandler(instance); + const ws = createIngestSocket(sessionId, wrapperState); + + await replayEvents(handler, ws, fixture.eventsBeforeIdle); + await replayEvents(handler, ws, [ + { + streamEventType: 'complete', + data: { exitCode: 0 }, + timestamp: new Date().toISOString(), + }, + ]); + + return instance.getMessageResult(fixture.userMessageId); + }); + + expect(result).toMatchObject({ + type: 'found', + result: { + status: 'completed', + completionSource: 'idle_reconciliation', + assistant: { messageId: fixture.finalAssistantMessageId }, + }, + }); + }); + + it('does not settle accepted work from raw idle maintenance', async () => { + const fixture = reconstructedToolLoopTurnFixtures[1]; + if (!fixture) throw new Error('Expected a reconstructed tool-loop fixture'); + const sessionId = 'agent_synthetic_no_idle_fallback'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`user_synthetic_no_idle_fallback:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async (instance, state) => { + const wrapperState = await seedAcceptedTurn(instance, fixture, sessionId); + const handler = await getIngestHandler(instance); + const ws = createIngestSocket(sessionId, wrapperState); + + await replayEvents(handler, ws, fixture.eventsBeforeIdle); + await replayEvents(handler, ws, [fixture.rootIdle]); + await instance.alarm(); + + return { + messageResult: await instance.getMessageResult(fixture.userMessageId), + completedEventCount: completedEvents(state).length, + }; + }); + + expect(result.messageResult).toMatchObject({ type: 'found', result: { status: 'running' } }); + expect(result.completedEventCount).toBe(0); + }); +}); diff --git a/services/cloud-agent-next/test/unit/wrapper/batch-admission.test.ts b/services/cloud-agent-next/test/unit/wrapper/batch-admission.test.ts new file mode 100644 index 0000000000..d7103da5c1 --- /dev/null +++ b/services/cloud-agent-next/test/unit/wrapper/batch-admission.test.ts @@ -0,0 +1,407 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { createLifecycleManager } from '../../../wrapper/src/lifecycle.js'; +import type { WrapperKiloClient } from '../../../wrapper/src/kilo-api.js'; +import { + bindSessionContext, + createCommandHandler, + createPromptHandler, + type ServerConfig, + type SessionBinding, +} from '../../../wrapper/src/server.js'; +import { WrapperState } from '../../../wrapper/src/state.js'; + +const config: ServerConfig = { + port: 5000, + workspacePath: '/workspace', + version: 'test', + sessionId: 'kilo_session', + agentSessionId: 'agent_session', + userId: 'user', +}; + +const binding: SessionBinding = { + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', + wrapperRunId: 'run_1', + wrapperGeneration: 1, + wrapperConnectionId: 'connection_1', +}; + +function request(body: unknown): Request { + return new Request('http://wrapper.test/job', { + method: 'POST', + body: JSON.stringify(body), + headers: { 'content-type': 'application/json' }, + }); +} + +function createDeps(state: WrapperState) { + return { + state, + kiloClient: { + sendPromptAsync: vi.fn().mockResolvedValue(undefined), + sendCommand: vi.fn().mockResolvedValue(undefined), + summarizeSession: vi.fn().mockResolvedValue(true), + } as WrapperKiloClient, + openConnection: vi.fn().mockResolvedValue(undefined), + closeConnection: vi.fn().mockResolvedValue(undefined), + setAborted: vi.fn(), + resetLifecycle: vi.fn(), + onDeliveryAcknowledged: vi.fn(), + configureCommitCoAuthor: vi.fn().mockResolvedValue(undefined), + }; +} + +describe('wrapper batch admission', () => { + afterEach(() => { + vi.useRealTimers(); + }); + + it('records only successfully admitted prompts and latest admitted finalization config', async () => { + const state = new WrapperState(); + const deps = createDeps(state); + const handler = createPromptHandler(config, deps); + + await handler( + request({ + message: { id: 'message-1', prompt: 'first' }, + finalization: { autoCommit: true }, + session: binding, + }) + ); + await handler( + request({ + message: { id: 'message-2', prompt: 'second' }, + finalization: { condenseOnComplete: true }, + session: binding, + }) + ); + + expect(state.pendingMessageIds).toEqual(['message-1', 'message-2']); + expect(state.batchFinalizationConfig).toMatchObject({ + autoCommit: false, + condenseOnComplete: true, + }); + expect(state.deliveryAcknowledgementsInFlight).toBe(0); + expect(deps.onDeliveryAcknowledged).toHaveBeenLastCalledWith('async-prompt'); + }); + + it('preserves the previously admitted batch when prompt admission fails', async () => { + const state = new WrapperState(); + const deps = createDeps(state); + const handler = createPromptHandler(config, deps); + + await handler(request({ message: { id: 'message-1', prompt: 'first' }, session: binding })); + vi.mocked(deps.kiloClient.sendPromptAsync).mockRejectedValueOnce(new Error('rejected')); + + const response = await handler( + request({ message: { id: 'message-2', prompt: 'second' }, session: binding }) + ); + + expect(response.status).toBe(500); + expect(state.pendingMessageIds).toEqual(['message-1']); + expect(state.deliveryAcknowledgementsInFlight).toBe(0); + expect(deps.onDeliveryAcknowledged).toHaveBeenLastCalledWith('failed'); + }); + + it('preserves prior prompt membership and config when its retry fails', async () => { + const state = new WrapperState(); + const deps = createDeps(state); + const handler = createPromptHandler(config, deps); + + await handler( + request({ + message: { id: 'message-1', prompt: 'first' }, + finalization: { autoCommit: true }, + session: binding, + }) + ); + vi.mocked(deps.kiloClient.sendPromptAsync).mockRejectedValueOnce(new Error('rejected')); + + const response = await handler( + request({ + message: { id: 'message-1', prompt: 'retry' }, + finalization: { condenseOnComplete: true }, + session: binding, + }) + ); + + expect(response.status).toBe(500); + expect(state.pendingMessageIds).toEqual(['message-1']); + expect(state.getMessageConfig('message-1')).toEqual({ + autoCommit: true, + condenseOnComplete: false, + upstreamBranch: undefined, + }); + expect(state.batchFinalizationConfig).toEqual({ + autoCommit: true, + condenseOnComplete: false, + upstreamBranch: undefined, + }); + }); + + it('preserves prior command membership and config when its retry fails', async () => { + const state = new WrapperState(); + const deps = createDeps(state); + const handler = createCommandHandler(config, deps); + + await handler( + request({ + command: 'first', + messageId: 'message-1', + autoCommit: true, + session: binding, + }) + ); + vi.mocked(deps.kiloClient.sendCommand).mockRejectedValueOnce(new Error('rejected')); + + const response = await handler( + request({ + command: 'retry', + messageId: 'message-1', + condenseOnComplete: true, + session: binding, + }) + ); + + expect(response.status).toBe(500); + expect(state.pendingMessageIds).toEqual(['message-1']); + expect(state.getMessageConfig('message-1')).toEqual({ + autoCommit: true, + condenseOnComplete: false, + upstreamBranch: undefined, + }); + expect(state.batchFinalizationConfig).toEqual({ + autoCommit: true, + condenseOnComplete: false, + upstreamBranch: undefined, + }); + }); + + it('successful async prompt admission cancels an armed idle candidate until later root idle', async () => { + vi.useFakeTimers(); + const state = new WrapperState(); + const deps = createDeps(state); + const sendToIngest = vi.fn(); + state.setSendToIngestFn(sendToIngest); + await bindSessionContext(binding, config, deps); + state.acceptMessage('message-1', { autoCommit: false, condenseOnComplete: false }); + const lifecycle = createLifecycleManager( + { workspacePath: '/workspace' }, + { + state, + kiloClient: deps.kiloClient, + closeConnections: vi.fn().mockResolvedValue(undefined), + isConnected: () => true, + reconnectEventSubscription: vi.fn(), + } + ); + deps.onDeliveryAcknowledged.mockImplementation(kind => lifecycle.onDeliveryAcknowledged(kind)); + + lifecycle.onSessionIdle(); + await vi.advanceTimersByTimeAsync(2_000); + + const response = await createPromptHandler( + config, + deps + )(request({ message: { id: 'message-2', prompt: 'later' }, session: binding })); + expect(response.status).toBe(200); + + await vi.advanceTimersByTimeAsync(1_001); + expect(sendToIngest).not.toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'wrapper_finalizing' }) + ); + + lifecycle.onSessionIdle(); + await vi.advanceTimersByTimeAsync(3_000); + await vi.waitFor(() => { + expect(sendToIngest).toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'complete' }) + ); + }); + + expect(sendToIngest).toHaveBeenCalledWith({ + streamEventType: 'complete', + data: expect.objectContaining({ messageIds: ['message-1', 'message-2'] }), + timestamp: expect.any(String), + }); + lifecycle.stop(); + }); + + it('preserves root idle observed while async prompt acknowledgement is in flight', async () => { + vi.useFakeTimers(); + const state = new WrapperState(); + const deps = createDeps(state); + const sendToIngest = vi.fn(); + state.setSendToIngestFn(sendToIngest); + await bindSessionContext(binding, config, deps); + let resolvePrompt: (() => void) | undefined; + vi.mocked(deps.kiloClient.sendPromptAsync).mockImplementationOnce( + () => + new Promise(resolve => { + resolvePrompt = resolve; + }) + ); + const lifecycle = createLifecycleManager( + { workspacePath: '/workspace' }, + { + state, + kiloClient: deps.kiloClient, + closeConnections: vi.fn().mockResolvedValue(undefined), + isConnected: () => true, + reconnectEventSubscription: vi.fn(), + } + ); + deps.onDeliveryAcknowledged.mockImplementation(kind => lifecycle.onDeliveryAcknowledged(kind)); + + const responsePromise = createPromptHandler( + config, + deps + )(request({ message: { id: 'message-1', prompt: 'first' }, session: binding })); + await vi.advanceTimersByTimeAsync(0); + expect(deps.kiloClient.sendPromptAsync).toHaveBeenCalledOnce(); + + lifecycle.onSessionIdle(); + await vi.advanceTimersByTimeAsync(3_000); + expect(sendToIngest).not.toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'wrapper_finalizing' }) + ); + + resolvePrompt?.(); + await expect(responsePromise).resolves.toMatchObject({ status: 200 }); + await vi.advanceTimersByTimeAsync(3_000); + + await vi.waitFor(() => { + expect(sendToIngest).toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'complete' }) + ); + }); + lifecycle.stop(); + }); + + it('rejects prompt immediately after an unsealed aborted drain starts', async () => { + vi.useFakeTimers(); + const state = new WrapperState(); + const deps = createDeps(state); + await bindSessionContext(binding, config, deps); + const lifecycle = createLifecycleManager( + { workspacePath: '/workspace' }, + { + state, + kiloClient: deps.kiloClient, + closeConnections: vi.fn().mockResolvedValue(undefined), + isConnected: () => true, + reconnectEventSubscription: vi.fn(), + } + ); + + state.clearAllMessages(); + lifecycle.setAborted(); + lifecycle.triggerDrainAndClose(); + + const response = await createPromptHandler( + config, + deps + )(request({ message: { id: 'new-prompt', prompt: 'later' }, session: binding })); + + expect(await response.json()).toMatchObject({ + error: 'WRAPPER_FINALIZING', + wrapperRunId: 'run_1', + }); + expect(deps.kiloClient.sendPromptAsync).not.toHaveBeenCalled(); + lifecycle.stop(); + }); + + it('continues accepting work after lifecycle monitoring stops for a runtime restart', async () => { + const state = new WrapperState(); + const deps = createDeps(state); + await bindSessionContext(binding, config, deps); + const lifecycle = createLifecycleManager( + { workspacePath: '/workspace' }, + { + state, + kiloClient: deps.kiloClient, + closeConnections: vi.fn().mockResolvedValue(undefined), + isConnected: () => true, + reconnectEventSubscription: vi.fn(), + } + ); + + lifecycle.stop(); + + const response = await createPromptHandler( + config, + deps + )(request({ message: { id: 'post-restart-prompt', prompt: 'continue' }, session: binding })); + + expect(response.status).toBe(200); + expect(deps.kiloClient.sendPromptAsync).toHaveBeenCalledOnce(); + }); + + it('accepts a fresh wrapper run after finalizing clears its session', async () => { + const state = new WrapperState(); + const deps = createDeps(state); + await bindSessionContext(binding, config, deps); + state.blockAdmissions(); + state.clearSession(); + + const staleResponse = await createPromptHandler( + config, + deps + )(request({ message: { id: 'stale-prompt', prompt: 'late' }, session: binding })); + + expect(await staleResponse.json()).toMatchObject({ + error: 'WRAPPER_FINALIZING', + wrapperRunId: 'run_1', + }); + expect(deps.kiloClient.sendPromptAsync).not.toHaveBeenCalled(); + + const freshResponse = await createPromptHandler( + config, + deps + )( + request({ + message: { id: 'fresh-prompt', prompt: 'later' }, + session: { ...binding, wrapperRunId: 'run_2' }, + }) + ); + + expect(freshResponse.status).toBe(200); + expect(state.currentSession?.wrapperRunId).toBe('run_2'); + expect(state.admissionsBlocked).toBe(false); + expect(deps.kiloClient.sendPromptAsync).toHaveBeenCalledOnce(); + }); + + it('rejects prompt, command, and ordinary rebind while finalizing', async () => { + const state = new WrapperState(); + const deps = createDeps(state); + await bindSessionContext(binding, config, deps); + state.acceptMessage('sealed-message', { autoCommit: false, condenseOnComplete: false }); + expect(state.beginFinalizing()).toBe(true); + + const promptResponse = await createPromptHandler( + config, + deps + )(request({ message: { id: 'new-prompt', prompt: 'later' }, session: binding })); + const commandResponse = await createCommandHandler( + config, + deps + )(request({ command: 'test', messageId: 'new-command', session: binding })); + const rebindResponse = await bindSessionContext(binding, config, deps); + + expect(await promptResponse.json()).toMatchObject({ + error: 'WRAPPER_FINALIZING', + wrapperRunId: 'run_1', + }); + expect(await commandResponse.json()).toMatchObject({ + error: 'WRAPPER_FINALIZING', + wrapperRunId: 'run_1', + }); + expect(await rebindResponse?.json()).toMatchObject({ + error: 'WRAPPER_FINALIZING', + wrapperRunId: 'run_1', + }); + expect(deps.kiloClient.sendPromptAsync).not.toHaveBeenCalled(); + expect(deps.kiloClient.sendCommand).not.toHaveBeenCalled(); + }); +}); diff --git a/services/cloud-agent-next/test/unit/wrapper/batch-lifecycle.test.ts b/services/cloud-agent-next/test/unit/wrapper/batch-lifecycle.test.ts new file mode 100644 index 0000000000..046cacb0a8 --- /dev/null +++ b/services/cloud-agent-next/test/unit/wrapper/batch-lifecycle.test.ts @@ -0,0 +1,177 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { createLifecycleManager } from '../../../wrapper/src/lifecycle.js'; +import type { WrapperKiloClient } from '../../../wrapper/src/kilo-api.js'; +import { WrapperState } from '../../../wrapper/src/state.js'; + +vi.mock('../../../wrapper/src/auto-commit.js', () => ({ + runAutoCommit: vi.fn().mockResolvedValue({ success: true }), +})); + +vi.mock('../../../wrapper/src/condense-on-complete.js', () => ({ + runCondenseOnComplete: vi.fn().mockResolvedValue({ wasAborted: false, success: true }), +})); + +vi.mock('../../../wrapper/src/utils.js', () => ({ + getCurrentBranch: vi.fn().mockResolvedValue('main'), + logToFile: vi.fn(), +})); + +const messageConfig = { + autoCommit: false, + condenseOnComplete: false, +}; + +function createKiloClient(): WrapperKiloClient { + return { + createSession: vi.fn(), + getSession: vi.fn(), + sendPromptAsync: vi.fn(), + abortSession: vi.fn(), + summarizeSession: vi.fn(), + sendCommand: vi.fn(), + answerPermission: vi.fn(), + answerQuestion: vi.fn(), + rejectQuestion: vi.fn(), + generateCommitMessage: vi.fn(), + getSessionStatuses: vi.fn(), + getQuestions: vi.fn(), + getPermissions: vi.fn(), + getNetworkWaits: vi.fn(), + resumeNetworkWait: vi.fn(), + subscribeEvents: vi.fn(), + serverUrl: 'http://127.0.0.1:0', + } as WrapperKiloClient; +} + +function bindRun(state: WrapperState): void { + state.bindSession({ + kiloSessionId: 'kilo_session', + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', + wrapperRunId: 'run_1', + wrapperGeneration: 1, + wrapperConnectionId: 'connection_1', + }); +} + +describe('sealed wrapper batch lifecycle', () => { + let state: WrapperState; + let sendToIngest: ReturnType; + let closeConnections: ReturnType; + let manager: ReturnType; + + beforeEach(() => { + vi.useFakeTimers(); + state = new WrapperState(); + bindRun(state); + sendToIngest = vi.fn(); + state.setSendToIngestFn(sendToIngest); + closeConnections = vi.fn().mockResolvedValue(undefined); + manager = createLifecycleManager( + { workspacePath: '/workspace' }, + { + state, + kiloClient: createKiloClient(), + closeConnections, + isConnected: () => true, + reconnectEventSubscription: vi.fn(), + } + ); + }); + + afterEach(() => { + manager.stop(); + vi.useRealTimers(); + }); + + it('seals exact admitted membership after three seconds of stable root idle', async () => { + state.acceptMessage('message-1', messageConfig); + state.acceptMessage('message-2', messageConfig); + + manager.onSessionIdle(); + await vi.advanceTimersByTimeAsync(2_999); + + expect(sendToIngest).not.toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'wrapper_finalizing' }) + ); + + await vi.advanceTimersByTimeAsync(1); + + expect(sendToIngest).toHaveBeenCalledWith({ + streamEventType: 'wrapper_finalizing', + data: { wrapperRunId: 'run_1' }, + timestamp: expect.any(String), + }); + expect(sendToIngest).toHaveBeenCalledWith({ + streamEventType: 'complete', + data: expect.objectContaining({ + exitCode: 0, + kiloSessionId: 'kilo_session', + messageIds: ['message-1', 'message-2'], + }), + timestamp: expect.any(String), + }); + }); + + it('requires a later root idle after root activity', async () => { + state.acceptMessage('message-1', messageConfig); + + manager.onSessionIdle(); + await vi.advanceTimersByTimeAsync(2_000); + manager.onRootSessionActivity(); + await vi.advanceTimersByTimeAsync(3_000); + + expect(sendToIngest).not.toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'wrapper_finalizing' }) + ); + + manager.onSessionIdle(); + await vi.advanceTimersByTimeAsync(3_000); + + expect(sendToIngest).toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'wrapper_finalizing' }) + ); + }); + + it('keeps repeated root idle and trailing turn close on the existing candidate', async () => { + state.acceptMessage('message-1', messageConfig); + + manager.onSessionIdle(); + await vi.advanceTimersByTimeAsync(1_000); + manager.onSessionIdle(); + await vi.advanceTimersByTimeAsync(2_000); + + expect(sendToIngest).toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'wrapper_finalizing' }) + ); + }); + + it('does not seal while a delivery acknowledgement is in flight', async () => { + state.acceptMessage('message-1', messageConfig); + state.beginDeliveryAcknowledgement(); + + manager.onSessionIdle(); + await vi.advanceTimersByTimeAsync(3_000); + + expect(sendToIngest).not.toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'wrapper_finalizing' }) + ); + + state.endDeliveryAcknowledgement(); + manager.onDeliveryAcknowledged('sync-command'); + await vi.advanceTimersByTimeAsync(3_000); + + expect(sendToIngest).toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'wrapper_finalizing' }) + ); + }); + + it('blocks admissions immediately when drain starts without a sealed batch', () => { + state.clearAllMessages(); + + manager.triggerDrainAndClose(); + + expect(state.beginDeliveryAcknowledgement()).toBe(false); + expect(state.isFinalizing).toBe(false); + }); +}); 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 5799257dc6..cc927e757c 100644 --- a/services/cloud-agent-next/test/unit/wrapper/connection.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/connection.test.ts @@ -8,8 +8,6 @@ import { describe, expect, it } from 'vitest'; import { buildIngestConnectionFailureMessage, isSessionIdleEvent, - isAssistantMessageCompleted, - getCompletedAssistantParentID, trimIngestEvent, } from '../../../wrapper/src/connection.js'; @@ -147,125 +145,3 @@ describe('isSessionIdleEvent', () => { expect(isSessionIdleEvent(data)).toBe(false); }); }); - -// --------------------------------------------------------------------------- -// isAssistantMessageCompleted -// --------------------------------------------------------------------------- - -describe('isAssistantMessageCompleted', () => { - const completedAssistant = { - event: 'message.updated' as const, - properties: { - info: { - role: 'assistant', - parentID: 'msg_018f1e2d3c4bParentMsg123456', - id: 'assistant_msg_001', - time: { completed: 1700000000000 }, - }, - }, - }; - - it('returns true for a completed assistant message.updated event', () => { - expect(isAssistantMessageCompleted(completedAssistant)).toBe(true); - }); - - it('returns parentID via getCompletedAssistantParentID', () => { - expect(getCompletedAssistantParentID(completedAssistant)).toBe( - 'msg_018f1e2d3c4bParentMsg123456' - ); - }); - - it('returns true for object-shaped terminal assistant errors', () => { - const failedAssistant = { - event: 'message.updated' as const, - properties: { - info: { - role: 'assistant', - parentID: 'msg_018f1e2d3c4bParentMsg123456', - id: 'assistant_msg_error', - error: { name: 'UnknownError', data: { message: 'provider failed' } }, - }, - }, - }; - - expect(isAssistantMessageCompleted(failedAssistant)).toBe(true); - expect(getCompletedAssistantParentID(failedAssistant)).toBe('msg_018f1e2d3c4bParentMsg123456'); - }); - - it('returns false when event is not message.updated', () => { - expect( - isAssistantMessageCompleted({ - event: 'session.idle', - properties: completedAssistant.properties, - }) - ).toBe(false); - }); - - it('returns false when role is not assistant', () => { - expect( - isAssistantMessageCompleted({ - event: 'message.updated', - properties: { - info: { ...completedAssistant.properties.info, role: 'user' }, - }, - }) - ).toBe(false); - }); - - it('returns false when parentID is missing', () => { - const data = { - event: 'message.updated' as const, - properties: { - info: { - role: 'assistant', - id: 'assistant_msg_001', - time: { completed: 1700000000000 }, - }, - }, - }; - expect(isAssistantMessageCompleted(data)).toBe(false); - }); - - it('returns false when time.completed is missing', () => { - const data = { - event: 'message.updated' as const, - properties: { - info: { - role: 'assistant', - parentID: 'msg_018f1e2d3c4bParentMsg123456', - id: 'assistant_msg_001', - time: {}, - }, - }, - }; - expect(isAssistantMessageCompleted(data)).toBe(false); - }); - - it('returns false when time is missing entirely', () => { - const data = { - event: 'message.updated' as const, - properties: { - info: { - role: 'assistant', - parentID: 'msg_018f1e2d3c4bParentMsg123456', - id: 'assistant_msg_001', - }, - }, - }; - expect(isAssistantMessageCompleted(data)).toBe(false); - }); - - it('returns false when data is null', () => { - expect(isAssistantMessageCompleted(null)).toBe(false); - }); - - it('returns false when data is not an object', () => { - expect(isAssistantMessageCompleted('message.updated')).toBe(false); - expect(isAssistantMessageCompleted(undefined)).toBe(false); - }); - - it('getCompletedAssistantParentID returns undefined for non-matching data', () => { - expect(getCompletedAssistantParentID(null)).toBeUndefined(); - expect(getCompletedAssistantParentID({ event: 'other' })).toBeUndefined(); - }); -}); diff --git a/services/cloud-agent-next/test/unit/wrapper/lifecycle.test.ts b/services/cloud-agent-next/test/unit/wrapper/lifecycle.test.ts index 50a1e22ec4..216260df33 100644 --- a/services/cloud-agent-next/test/unit/wrapper/lifecycle.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/lifecycle.test.ts @@ -1,28 +1,15 @@ -/** - * Unit tests for lifecycle management. - * - * Tests timer logic with mocked state for: - * - SSE transport timer (15s reconnect) - * - Drain period - * - Post-completion task triggering - */ - -import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest'; -import { - createLifecycleManager, - type LifecycleConfig, - type LifecycleDependencies, - type LifecycleManager, -} from '../../../wrapper/src/lifecycle.js'; -import { WrapperState, type SessionContext } from '../../../wrapper/src/state.js'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { runAutoCommit } from '../../../wrapper/src/auto-commit.js'; +import { createLifecycleManager } from '../../../wrapper/src/lifecycle.js'; import type { WrapperKiloClient } from '../../../wrapper/src/kilo-api.js'; +import { WrapperState } from '../../../wrapper/src/state.js'; vi.mock('../../../wrapper/src/auto-commit.js', () => ({ - runAutoCommit: vi.fn(), + runAutoCommit: vi.fn().mockResolvedValue({ success: true }), })); vi.mock('../../../wrapper/src/condense-on-complete.js', () => ({ - runCondenseOnComplete: vi.fn(), + runCondenseOnComplete: vi.fn().mockResolvedValue({ wasAborted: false, success: true }), })); vi.mock('../../../wrapper/src/utils.js', () => ({ @@ -30,1064 +17,164 @@ vi.mock('../../../wrapper/src/utils.js', () => ({ logToFile: vi.fn(), })); -import { runAutoCommit } from '../../../wrapper/src/auto-commit.js'; -import { runCondenseOnComplete } from '../../../wrapper/src/condense-on-complete.js'; - -const mockRunAutoCommit = vi.mocked(runAutoCommit); -const mockRunCondenseOnComplete = vi.mocked(runCondenseOnComplete); -// --------------------------------------------------------------------------- -// Test Helpers -// --------------------------------------------------------------------------- - -const createMockKiloClient = (): WrapperKiloClient => ({ - createSession: vi.fn().mockResolvedValue({ id: 'kilo_sess' }), - getSession: vi.fn().mockResolvedValue({ id: 'kilo_sess' }), - sendPromptAsync: vi.fn().mockResolvedValue(undefined), - abortSession: vi.fn().mockResolvedValue(true), - summarizeSession: vi.fn().mockResolvedValue(true), - sendCommand: vi.fn().mockResolvedValue(undefined), - answerPermission: vi.fn().mockResolvedValue(true), - answerQuestion: vi.fn().mockResolvedValue(true), - rejectQuestion: vi.fn().mockResolvedValue(true), - getSessionStatuses: vi.fn().mockResolvedValue({}), - getQuestions: vi.fn().mockResolvedValue([]), - getPermissions: vi.fn().mockResolvedValue([]), - getNetworkWaits: vi.fn().mockResolvedValue([]), - resumeNetworkWait: vi.fn().mockResolvedValue(true), - generateCommitMessage: vi.fn().mockResolvedValue({ message: 'test commit' }), - getSessionStatuses: vi.fn().mockResolvedValue({}), - getQuestions: vi.fn().mockResolvedValue([]), - getPermissions: vi.fn().mockResolvedValue([]), - subscribeEvents: vi.fn().mockResolvedValue({ stream: undefined }), - serverUrl: 'http://127.0.0.1:0', -}); - -type MockConnectionFns = { - closeConnections: ReturnType; - isConnected: ReturnType; - reconnectEventSubscription: ReturnType; -}; - -const createMockConnectionFns = (): MockConnectionFns => ({ - closeConnections: vi.fn().mockResolvedValue(undefined), - isConnected: vi.fn().mockReturnValue(false), - reconnectEventSubscription: vi.fn(), -}); - -const createDefaultConfig = (overrides: Partial = {}): LifecycleConfig => ({ - workspacePath: '/workspace', - ...overrides, -}); +function createKiloClient(): WrapperKiloClient { + return { serverUrl: 'http://127.0.0.1:0' } as WrapperKiloClient; +} -const createSessionContext = (overrides: Partial = {}): SessionContext => ({ - kiloSessionId: 'kilo_sess_456', - ingestUrl: 'wss://ingest.example.com', - ingestToken: 'token_secret', - workerAuthToken: 'kilo_token_789', - ...overrides, -}); - -const createDeferred = () => { - let resolve: (value: T) => void = () => {}; - const promise = new Promise(res => { - resolve = res; +function bindRun(state: WrapperState): void { + state.bindSession({ + kiloSessionId: 'kilo_session', + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', + wrapperRunId: 'run_1', + wrapperGeneration: 1, + wrapperConnectionId: 'connection_1', }); - return { promise, resolve }; -}; +} -// --------------------------------------------------------------------------- -// Tests -// --------------------------------------------------------------------------- +const config = { autoCommit: false, condenseOnComplete: false }; -describe('createLifecycleManager', () => { +describe('wrapper sealed batch lifecycle', () => { let state: WrapperState; - let kiloClient: WrapperKiloClient; - let connectionFns: MockConnectionFns; - let config: LifecycleConfig; - let manager: LifecycleManager; + let sendToIngest: ReturnType; + let closeConnections: ReturnType; + let manager: ReturnType; beforeEach(() => { vi.useFakeTimers(); + vi.mocked(runAutoCommit).mockClear(); state = new WrapperState(); - kiloClient = createMockKiloClient(); - connectionFns = createMockConnectionFns(); - config = createDefaultConfig(); - mockRunAutoCommit.mockResolvedValue({ success: true }); - mockRunCondenseOnComplete.mockResolvedValue({ wasAborted: false, success: true }); - }); - - afterEach(() => { - if (manager) { - manager.stop(); - } - vi.useRealTimers(); - vi.clearAllMocks(); - }); - - const createManager = (overrides: Partial = {}): LifecycleManager => { + bindRun(state); + sendToIngest = vi.fn(); + state.setSendToIngestFn(sendToIngest); + closeConnections = vi.fn().mockResolvedValue(undefined); manager = createLifecycleManager( - { ...config, ...overrides }, + { workspacePath: '/workspace' }, { state, - kiloClient, - closeConnections: connectionFns.closeConnections, - isConnected: connectionFns.isConnected, - reconnectEventSubscription: connectionFns.reconnectEventSubscription, + kiloClient: createKiloClient(), + closeConnections, + isConnected: () => true, + reconnectEventSubscription: vi.fn(), } ); - return manager; - }; - - // ------------------------------------------------------------------------- - // Basic Lifecycle - // ------------------------------------------------------------------------- - - describe('basic lifecycle', () => { - it('returns a manager with expected methods', () => { - const mgr = createManager(); - - expect(mgr).toHaveProperty('start'); - expect(mgr).toHaveProperty('stop'); - expect(mgr).toHaveProperty('onMessageComplete'); - expect(mgr).toHaveProperty('onRootSessionActivity'); - expect(mgr).toHaveProperty('triggerDrainAndClose'); - expect(mgr).toHaveProperty('onSseEvent'); - expect(mgr).toHaveProperty('signalCompletion'); - expect(mgr).toHaveProperty('setAborted'); - }); - - it('has onSseEvent method', () => { - const mgr = createManager(); - expect(typeof mgr.onSseEvent).toBe('function'); - }); }); - // ------------------------------------------------------------------------- - // Message Completion - // ------------------------------------------------------------------------- - - describe('onMessageComplete', () => { - it('sets state to idle after completing last message', () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - mgr.onMessageComplete('msg_1'); - expect(state.isIdle).toBe(true); - }); - - it('triggers drain after the final message and root session idle', async () => { - const mgr = createManager(); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - mgr.onMessageComplete('msg_1'); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(500); - expect(connectionFns.closeConnections).toHaveBeenCalled(); - }); - - it('defers drain while ingest reconnects and resumes it when connectivity is restored', async () => { - const mgr = createManager(); - (connectionFns.isConnected as ReturnType).mockReturnValue(false); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - mgr.onMessageComplete('msg_1'); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(500); - expect(connectionFns.closeConnections).not.toHaveBeenCalled(); - - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - mgr.onConnectionRestored(); - await vi.advanceTimersByTimeAsync(500); - expect(connectionFns.closeConnections).toHaveBeenCalled(); - }); - - it('does not trigger drain when pending messages remain', async () => { - const mgr = createManager(); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_2', { autoCommit: false, condenseOnComplete: false }); - mgr.onMessageComplete('msg_1'); - await vi.advanceTimersByTimeAsync(500); - expect(connectionFns.closeConnections).not.toHaveBeenCalled(); - expect(state.isActive).toBe(true); - expect(state.activeMessageId).toBe('msg_2'); - }); - - it('handles unknown messageId gracefully', () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - expect(() => mgr.onMessageComplete('unknown_msg')).not.toThrow(); - }); + afterEach(() => { + manager.stop(); + vi.useRealTimers(); }); - // ------------------------------------------------------------------------- - // Root idle freshness barrier - // ------------------------------------------------------------------------- - - describe('root idle freshness barrier', () => { - it('waits for a fresh root idle after post-idle activity before finalizing', async () => { - const mgr = createManager(); - const sendToIngestSpy = vi.fn(); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - state.setSendToIngestFn(sendToIngestSpy); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - mgr.onSessionIdle(); - mgr.onRootSessionActivity(); - mgr.onMessageComplete('msg_1'); - - await vi.advanceTimersByTimeAsync(1000); - - expect(sendToIngestSpy).not.toHaveBeenCalledWith( - expect.objectContaining({ streamEventType: 'complete' }) - ); - expect(connectionFns.closeConnections).not.toHaveBeenCalled(); - - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(1000); - - expect(sendToIngestSpy).toHaveBeenCalledTimes(1); - expect(sendToIngestSpy).toHaveBeenCalledWith( - expect.objectContaining({ streamEventType: 'complete' }) - ); - expect(connectionFns.closeConnections).toHaveBeenCalledTimes(1); - }); - - it('finalizes once after repeated stale idle and root activity cycles', async () => { - const mgr = createManager(); - const sendToIngestSpy = vi.fn(); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - state.setSendToIngestFn(sendToIngestSpy); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_cycle', { autoCommit: false, condenseOnComplete: false }); - - mgr.onSessionIdle(); - mgr.onRootSessionActivity(); - mgr.onSessionIdle(); - mgr.onRootSessionActivity(); - mgr.onMessageComplete('msg_cycle'); - - await vi.advanceTimersByTimeAsync(1000); + it('emits finalizing then complete once with exact admitted membership', async () => { + state.acceptMessage('message-1', config); + state.acceptMessage('message-2', config); - expect(sendToIngestSpy).not.toHaveBeenCalledWith( - expect.objectContaining({ streamEventType: 'complete' }) - ); - expect(connectionFns.closeConnections).not.toHaveBeenCalled(); + manager.onSessionIdle(); + await vi.advanceTimersByTimeAsync(3_300); - mgr.onSessionIdle(); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(1000); - - expect(sendToIngestSpy).toHaveBeenCalledTimes(1); - expect(sendToIngestSpy).toHaveBeenCalledWith( - expect.objectContaining({ streamEventType: 'complete' }) - ); - expect(connectionFns.closeConnections).toHaveBeenCalledTimes(1); + expect(sendToIngest).toHaveBeenCalledWith({ + streamEventType: 'wrapper_finalizing', + data: { wrapperRunId: 'run_1' }, + timestamp: expect.any(String), }); - }); - - it('includes the observed gate result in the final complete event', async () => { - const mgr = createManager(); - const sendToIngestSpy = vi.fn(); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - state.setSendToIngestFn(sendToIngestSpy); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_gate_result', { autoCommit: false, condenseOnComplete: false }); - state.observeGateResult('pass'); - - mgr.onMessageComplete('msg_gate_result'); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(1000); - - expect(sendToIngestSpy).toHaveBeenCalledWith( - expect.objectContaining({ - streamEventType: 'complete', - data: expect.objectContaining({ gateResult: 'pass' }), - }) - ); - }); - - it('does not carry an observed gate result into a follow-up message accepted during drain', async () => { - const mgr = createManager(); - const sendToIngestSpy = vi.fn(); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - state.setSendToIngestFn(sendToIngestSpy); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_first', { autoCommit: false, condenseOnComplete: false }); - state.observeGateResult('fail'); - - mgr.onMessageComplete('msg_first'); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(0); - - expect(sendToIngestSpy).toHaveBeenCalledWith( - expect.objectContaining({ - streamEventType: 'complete', - data: expect.objectContaining({ gateResult: 'fail' }), - }) - ); - - state.acceptMessage('msg_followup', { autoCommit: false, condenseOnComplete: false }); - await vi.advanceTimersByTimeAsync(300); - - expect(connectionFns.closeConnections).not.toHaveBeenCalled(); - - mgr.onMessageComplete('msg_followup'); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(1000); - - const completeEvents = sendToIngestSpy.mock.calls + const completeEvents = sendToIngest.mock.calls .map(([event]) => event) .filter(event => event.streamEventType === 'complete'); - - expect(completeEvents).toHaveLength(2); - expect(completeEvents[1].data).not.toHaveProperty('gateResult'); - }); - - // ------------------------------------------------------------------------- - // Drain and Close - // ------------------------------------------------------------------------- - - describe('triggerDrainAndClose', () => { - it('closes connection after drain delay', async () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - - mgr.triggerDrainAndClose(); - - // Before delay - not closed - expect(connectionFns.closeConnections).not.toHaveBeenCalled(); - - // After 250ms drain delay - await vi.advanceTimersByTimeAsync(300); - - expect(connectionFns.closeConnections).toHaveBeenCalled(); - }); - - it('is idempotent - multiple calls do not queue multiple drains', async () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - - mgr.triggerDrainAndClose(); - mgr.triggerDrainAndClose(); - mgr.triggerDrainAndClose(); - - await vi.advanceTimersByTimeAsync(1000); - - // Close should only be called once - expect(connectionFns.closeConnections).toHaveBeenCalledTimes(1); - }); - }); - - // ------------------------------------------------------------------------- - // Stop - // ------------------------------------------------------------------------- - - describe('stop', () => { - it('clears all timers', async () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - - mgr.start(); - - await vi.advanceTimersByTimeAsync(3000); - mgr.stop(); - - await vi.advanceTimersByTimeAsync(20000); - - expect(state.hasSession).toBe(true); - }); - - it('cancels pending drain', async () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - - mgr.triggerDrainAndClose(); - - // Stop before drain completes - vi.advanceTimersByTime(100); - mgr.stop(); - - // Advance past drain delay - vi.advanceTimersByTime(500); - - // Close should not have been called - expect(connectionFns.closeConnections).not.toHaveBeenCalled(); - }); - - it('sets aborted flag', () => { - const mgr = createManager(); - - mgr.stop(); - - // This is internal state, verified by behavior in post-completion tests - // The stop() method sets isAborted = true + expect(completeEvents).toHaveLength(1); + expect(completeEvents[0].data).toMatchObject({ + kiloSessionId: 'kilo_session', + messageIds: ['message-1', 'message-2'], }); + expect(closeConnections).toHaveBeenCalledOnce(); }); - // ------------------------------------------------------------------------- - // setAborted - // ------------------------------------------------------------------------- - - describe('setAborted', () => { - it('prevents post-completion tasks from running', async () => { - const mgr = createManager({}, { autoCommit: true }); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: true, condenseOnComplete: false }); + it('does not seal until root idle is stable for three seconds', async () => { + state.acceptMessage('message-1', config); - mgr.setAborted(); + manager.onSessionIdle(); + await vi.advanceTimersByTimeAsync(2_999); + expect(sendToIngest).not.toHaveBeenCalled(); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - mgr.onMessageComplete('msg_1'); + manager.onRootSessionActivity(); + await vi.advanceTimersByTimeAsync(3_000); + expect(sendToIngest).not.toHaveBeenCalled(); - await vi.advanceTimersByTimeAsync(1000); - }); + manager.onSessionIdle(); + await vi.advanceTimersByTimeAsync(3_000); + expect(sendToIngest).toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'wrapper_finalizing' }) + ); }); - // ------------------------------------------------------------------------- - // reset - // ------------------------------------------------------------------------- - - describe('reset', () => { - it('reset clears aborted flag - allows complete event after reset', async () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - mgr.setAborted(); - mgr.reset(); - - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - const sendToIngestSpy = vi.fn(); - state.setSendToIngestFn(sendToIngestSpy); - - mgr.onMessageComplete('msg_1'); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(1000); - - expect(sendToIngestSpy).toHaveBeenCalledWith( - expect.objectContaining({ - streamEventType: 'complete', - }) - ); - }); - - it('reset clears draining flag - allows new drain after reset', async () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - // First drain — complete the message to simulate realistic flow - mgr.onMessageComplete('msg_1'); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(500); + it('preserves the idle candidate through transient ingest reconnect', async () => { + state.acceptMessage('message-1', config); + manager.onSessionIdle(); - expect(connectionFns.closeConnections).toHaveBeenCalledTimes(1); + await vi.advanceTimersByTimeAsync(2_000); + manager.onConnectionRestored(); + await vi.advanceTimersByTimeAsync(1_000); - // Reset clears isDraining so a second drain can happen - mgr.reset(); - - // Start a fresh session - state.clearSession(); - state.bindSession(createSessionContext({ kiloSessionId: 'kilo_sess_second' })); - state.acceptMessage('msg_2', { autoCommit: false, condenseOnComplete: false }); - - // Completing last active message triggers a new drain - mgr.onMessageComplete('msg_2'); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(1000); - - expect(connectionFns.closeConnections).toHaveBeenCalledTimes(2); - }); - - it('reset enables post-completion flow after previous abort', async () => { - const mgr = createManager({}, { autoCommit: false }); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - mgr.setAborted(); - mgr.reset(); - - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - const sendToIngestSpy = vi.fn(); - state.setSendToIngestFn(sendToIngestSpy); - - mgr.onMessageComplete('msg_1'); - mgr.signalCompletion(); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(1000); - - expect(sendToIngestSpy).toHaveBeenCalledWith( - expect.objectContaining({ - streamEventType: 'complete', - }) - ); - }); + expect(sendToIngest).toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'wrapper_finalizing' }) + ); }); - // ------------------------------------------------------------------------- - // signalCompletion - // ------------------------------------------------------------------------- + it('does not seal while delivery acknowledgement remains in flight', async () => { + state.acceptMessage('message-1', config); + state.beginDeliveryAcknowledgement(); + manager.onSessionIdle(); + await vi.advanceTimersByTimeAsync(3_000); - describe('signalCompletion', () => { - it('can be called without error', () => { - const mgr = createManager(); - - // Should not throw - expect(() => mgr.signalCompletion()).not.toThrow(); - }); + expect(sendToIngest).not.toHaveBeenCalled(); - // Integration test: signalCompletion resolves waitForCompletion in runPostCompletionTasks - // This is tested more thoroughly in integration tests + state.endDeliveryAcknowledgement(); + manager.onDeliveryAcknowledged('failed'); + await vi.advanceTimersByTimeAsync(3_000); + expect(sendToIngest).toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'wrapper_finalizing' }) + ); }); - // ------------------------------------------------------------------------- - // Post-Completion Tasks - // ------------------------------------------------------------------------- - - describe('post-completion tasks', () => { - it('runs auto-commit when enabled', async () => { - const mgr = createManager({}, { autoCommit: true }); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { - autoCommit: true, - condenseOnComplete: false, - commitCoAuthor: { name: 'kiloconnect[bot]', email: 'bot@example.com' }, - }); - - mgr.start(); - mgr.onMessageComplete('msg_1'); - mgr.onSessionIdle(); - - mgr.signalCompletion(); - - await vi.advanceTimersByTimeAsync(1000); - expect(mockRunAutoCommit).toHaveBeenCalledWith( - expect.objectContaining({ - commitCoAuthor: { name: 'kiloconnect[bot]', email: 'bot@example.com' }, + it('waits for enabled post-processing before complete and close', async () => { + let resolveAutoCommit: ((result: { success: boolean }) => void) | undefined; + vi.mocked(runAutoCommit).mockImplementationOnce( + () => + new Promise(resolve => { + resolveAutoCommit = resolve; }) - ); - }); - - it('runs condense when enabled', async () => { - const mgr = createManager({}, { condenseOnComplete: true }); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: true }); - - mgr.start(); - mgr.onMessageComplete('msg_1'); - mgr.onSessionIdle(); - - mgr.signalCompletion(); - - await vi.advanceTimersByTimeAsync(1000); - }); - - it('aborts auto-commit when the lifecycle timeout fires', async () => { - const sendToIngestSpy = vi.fn(); - state.setSendToIngestFn(sendToIngestSpy); - mockRunAutoCommit.mockImplementation( - ({ signal }) => - new Promise(resolve => { - signal?.addEventListener( - 'abort', - () => resolve({ success: false, error: 'exec aborted' }), - { - once: true, - } - ); - }) - ); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: true, condenseOnComplete: false }); - - mgr.onMessageComplete('msg_1'); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(120_000); - await vi.advanceTimersByTimeAsync(300); - - const autoCommitCall = mockRunAutoCommit.mock.calls[0]?.[0]; - expect(autoCommitCall?.signal?.aborted).toBe(true); - expect(sendToIngestSpy).toHaveBeenCalledWith( - expect.objectContaining({ - streamEventType: 'error', - data: { error: 'Auto-commit timed out', fatal: false }, - }) - ); - expect(sendToIngestSpy).toHaveBeenCalledWith( - expect.objectContaining({ streamEventType: 'complete' }) - ); - }); - - it('does not report lifecycle timeout when auto-commit wins the timeout race', async () => { - const sendToIngestSpy = vi.fn(); - state.setSendToIngestFn(sendToIngestSpy); - mockRunAutoCommit.mockImplementation( - ({ signal }) => - new Promise(resolve => { - signal?.addEventListener('abort', () => resolve({ success: true }), { - once: true, - }); - }) - ); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: true, condenseOnComplete: false }); - - mgr.onMessageComplete('msg_1'); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(120_000); - await vi.advanceTimersByTimeAsync(300); - - expect(sendToIngestSpy).not.toHaveBeenCalledWith( - expect.objectContaining({ - streamEventType: 'error', - data: { error: 'Auto-commit timed out', fatal: false }, - }) - ); - expect(sendToIngestSpy).toHaveBeenCalledWith( - expect.objectContaining({ streamEventType: 'complete' }) - ); - }); - - it('awaits final message auto-commit before log upload, complete event, and close', async () => { - const order: string[] = []; - const autoCommit = createDeferred<{ success: true }>(); - mockRunAutoCommit.mockImplementation(async () => { - order.push('auto-commit:start'); - const result = await autoCommit.promise; - order.push('auto-commit:finish'); - return result; - }); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - connectionFns.closeConnections.mockImplementation(async () => { - order.push('close'); - }); - state.setLogUploader({ - start: vi.fn(), - uploadNow: vi.fn(async () => { - order.push('upload'); - }), - stop: vi.fn(), - }); - state.setSendToIngestFn(event => { - if (event.streamEventType === 'complete') { - order.push('complete'); - } - }); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: true, condenseOnComplete: false }); - - mgr.onMessageComplete('msg_1'); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(0); - - expect(order).toEqual(['auto-commit:start']); - - autoCommit.resolve({ success: true }); - await vi.advanceTimersByTimeAsync(300); - - expect(order).toEqual([ - 'auto-commit:start', - 'auto-commit:finish', - 'upload', - 'complete', - 'close', - ]); - }); - }); - - // ------------------------------------------------------------------------- - // Edge Cases - // ------------------------------------------------------------------------- - - describe('edge cases', () => { - it('handles connection not connected during drain', async () => { - const mgr = createManager(); - (connectionFns.isConnected as ReturnType).mockReturnValue(false); - - state.bindSession(createSessionContext()); - - mgr.triggerDrainAndClose(); - - await vi.advanceTimersByTimeAsync(500); - - expect(connectionFns.closeConnections).toHaveBeenCalled(); - }); - }); - - // ------------------------------------------------------------------------- - // onSseEvent / transport timer - // ------------------------------------------------------------------------- - - describe('onSseEvent / transport timer', () => { - it('fires reconnectEventSubscription after 15s of no events', async () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - - mgr.onSseEvent(); - - await vi.advanceTimersByTimeAsync(15_000); - - expect(connectionFns.reconnectEventSubscription).toHaveBeenCalledTimes(1); - }); - - it('resets timer on each onSseEvent call', async () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - - mgr.onSseEvent(); - - await vi.advanceTimersByTimeAsync(10_000); - expect(connectionFns.reconnectEventSubscription).not.toHaveBeenCalled(); - - mgr.onSseEvent(); - - await vi.advanceTimersByTimeAsync(10_000); - expect(connectionFns.reconnectEventSubscription).not.toHaveBeenCalled(); - - await vi.advanceTimersByTimeAsync(5_000); - expect(connectionFns.reconnectEventSubscription).toHaveBeenCalledTimes(1); - }); - - it('does not fire timer when no session context', async () => { - const mgr = createManager(); - - mgr.onSseEvent(); - - await vi.advanceTimersByTimeAsync(15_000); - expect(connectionFns.reconnectEventSubscription).not.toHaveBeenCalled(); - }); - - it('transport timer is cleared on stop', async () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - - mgr.onSseEvent(); - mgr.stop(); - - await vi.advanceTimersByTimeAsync(15_000); - expect(connectionFns.reconnectEventSubscription).not.toHaveBeenCalled(); - }); - - it('transport timer is cleared on reset', async () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - - mgr.onSseEvent(); - mgr.reset(); - - await vi.advanceTimersByTimeAsync(15_000); - expect(connectionFns.reconnectEventSubscription).not.toHaveBeenCalled(); - }); - - it('initial arming triggers reconnect when stream never yields', async () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - - mgr.onSseEvent(); - - await vi.advanceTimersByTimeAsync(15_000); - expect(connectionFns.reconnectEventSubscription).toHaveBeenCalledTimes(1); - }); - }); - - // ------------------------------------------------------------------------- - // Post-completion exactly-once - // ------------------------------------------------------------------------- - - describe('post-completion exactly-once', () => { - it('final message auto-commit runs exactly once, not twice', async () => { - mockRunAutoCommit.mockResolvedValue({ success: true }); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_final', { autoCommit: true, condenseOnComplete: false }); - - mgr.onMessageComplete('msg_final'); - mgr.signalCompletion(); - mgr.onSessionIdle(); - - await vi.advanceTimersByTimeAsync(1000); - - expect(mockRunAutoCommit).toHaveBeenCalledTimes(1); - }); - - it('defers non-final message auto-commit while wrapper remains active', async () => { - mockRunAutoCommit.mockResolvedValue({ success: true }); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: true, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); - - mgr.onMessageComplete('msg_a'); - mgr.signalCompletion(); - - await vi.advanceTimersByTimeAsync(1000); - - expect(mockRunAutoCommit).not.toHaveBeenCalled(); - expect(state.activeMessageId).toBe('msg_b'); - expect(connectionFns.closeConnections).not.toHaveBeenCalled(); - }); - - it('finalizes a multi-message sequence once after the batch reaches idle', async () => { - mockRunAutoCommit.mockResolvedValue({ success: true }); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: true, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: true, condenseOnComplete: false }); - - mgr.onMessageComplete('msg_a'); - mgr.signalCompletion(); - await vi.advanceTimersByTimeAsync(100); - - expect(mockRunAutoCommit).not.toHaveBeenCalled(); - - mgr.onMessageComplete('msg_b'); - mgr.signalCompletion(); - mgr.onSessionIdle(); - await vi.advanceTimersByTimeAsync(1000); - - expect(mockRunAutoCommit).toHaveBeenCalledTimes(1); - }); - - it('abort path does not run auto-commit', async () => { - mockRunAutoCommit.mockResolvedValue({ success: true }); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: true, condenseOnComplete: false }); - - mgr.setAborted(); - mgr.triggerDrainAndClose(); - - await vi.advanceTimersByTimeAsync(1000); - - expect(mockRunAutoCommit).not.toHaveBeenCalled(); - }); - - it('final message condense runs exactly once', async () => { - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_final', { autoCommit: false, condenseOnComplete: true }); - - mgr.onMessageComplete('msg_final'); - mgr.signalCompletion(); - mgr.onSessionIdle(); - - await vi.advanceTimersByTimeAsync(1000); - - expect(mockRunCondenseOnComplete).toHaveBeenCalledTimes(1); - }); - }); - - // ------------------------------------------------------------------------- - // Idle-batch post-completion tasks - // ------------------------------------------------------------------------- - - describe('onMessageComplete — idle-batch post-completion', () => { - it('runs auto-commit after the final message reaches drain', async () => { - mockRunAutoCommit.mockResolvedValue({ success: true }); - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_comp_1', { autoCommit: true, condenseOnComplete: false }); - - mgr.onMessageComplete('msg_comp_1'); - mgr.onSessionIdle(); - - // Auto-commit runs asynchronously after completion - await vi.advanceTimersByTimeAsync(100); - - expect(mockRunAutoCommit).toHaveBeenCalledTimes(1); - }); - - it('does not trigger drain when onMessageComplete and more messages are pending', async () => { - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_first', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_second', { autoCommit: false, condenseOnComplete: false }); - - // Complete first message — should NOT drain because second is still pending - mgr.onMessageComplete('msg_first'); - - await vi.advanceTimersByTimeAsync(100); - - expect(connectionFns.closeConnections).not.toHaveBeenCalled(); - expect(state.activeMessageId).toBe('msg_second'); - }); - - it('promotes next accepted message after completing active', async () => { - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_c', { autoCommit: false, condenseOnComplete: false }); - - expect(state.activeMessageId).toBe('msg_a'); - - mgr.onMessageComplete('msg_a'); - expect(state.activeMessageId).toBe('msg_b'); - - mgr.onMessageComplete('msg_b'); - expect(state.activeMessageId).toBe('msg_c'); - - mgr.onMessageComplete('msg_c'); - expect(state.activeMessageId).toBeNull(); - expect(state.isIdle).toBe(true); - }); - }); - - // ------------------------------------------------------------------------- - // Message-ID-gated completion (stale/duplicate protection) - // ------------------------------------------------------------------------- - - describe('onMessageComplete — stale and duplicate protection', () => { - it('duplicate completion for already-completed message does not affect the new active message', async () => { - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - mockRunAutoCommit.mockResolvedValue({ success: true }); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); - - mgr.onMessageComplete('msg_a'); - expect(state.activeMessageId).toBe('msg_b'); - - mgr.onMessageComplete('msg_a'); - expect(state.activeMessageId).toBe('msg_b'); - expect(state.isActive).toBe(true); - - await vi.advanceTimersByTimeAsync(500); - expect(connectionFns.closeConnections).not.toHaveBeenCalled(); - }); - - it('completion for accepted-but-not-active message is ignored', async () => { - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); - - mgr.onMessageComplete('msg_b'); - - expect(state.activeMessageId).toBe('msg_a'); - expect(state.isActive).toBe(true); - - await vi.advanceTimersByTimeAsync(500); - expect(connectionFns.closeConnections).not.toHaveBeenCalled(); - }); - - it('completion for unknown message ID is ignored', async () => { - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - - mgr.onMessageComplete('msg_nonexistent'); - - expect(state.activeMessageId).toBe('msg_a'); - expect(state.isActive).toBe(true); - - await vi.advanceTimersByTimeAsync(500); - expect(connectionFns.closeConnections).not.toHaveBeenCalled(); - }); - - it('stale completion does not run post-completion tasks for the wrong message', async () => { - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - mockRunAutoCommit.mockResolvedValue({ success: true }); + ); + state.acceptMessage('message-1', { ...config, autoCommit: true }); - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: true, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); + manager.onSessionIdle(); + await vi.advanceTimersByTimeAsync(3_000); - mgr.onMessageComplete('msg_a'); - await vi.advanceTimersByTimeAsync(100); + expect(sendToIngest).toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'wrapper_finalizing' }) + ); + expect(sendToIngest).not.toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'complete' }) + ); + expect(closeConnections).not.toHaveBeenCalled(); - const autoCommitCallsAfterFirst = mockRunAutoCommit.mock.calls.length; + resolveAutoCommit?.({ success: true }); + await vi.advanceTimersByTimeAsync(0); - mgr.onMessageComplete('msg_a'); - await vi.advanceTimersByTimeAsync(100); + expect(sendToIngest).toHaveBeenCalledWith( + expect.objectContaining({ streamEventType: 'complete' }) + ); + expect(closeConnections).not.toHaveBeenCalled(); - expect(mockRunAutoCommit.mock.calls.length).toBe(autoCommitCallsAfterFirst); - }); + await vi.advanceTimersByTimeAsync(250); + expect(closeConnections).toHaveBeenCalledOnce(); }); - // ------------------------------------------------------------------------- - // Drain guard — drain does not close over newly accepted prompts - // ------------------------------------------------------------------------- - - describe('drain guard', () => { - it('drain does not close connections when a new message is accepted during drain', async () => { - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_last', { autoCommit: false, condenseOnComplete: false }); - - // Complete last message and observe root idle — triggers drain - mgr.onMessageComplete('msg_last'); - mgr.onSessionIdle(); - - // Before drain delay expires, accept a new message - state.acceptMessage('msg_new', { autoCommit: false, condenseOnComplete: false }); + it('uses latest admitted finalization config', async () => { + state.acceptMessage('message-1', { ...config, autoCommit: true }); + state.acceptMessage('message-2', config); - // Advance past drain delay - await vi.advanceTimersByTimeAsync(1000); + manager.onSessionIdle(); + await vi.advanceTimersByTimeAsync(3_300); - // Drain should NOT have closed connections because a new message is active - expect(connectionFns.closeConnections).not.toHaveBeenCalled(); - - // The session should still be active - expect(state.activeMessageId).toBe('msg_new'); - expect(state.isIdle).toBe(false); - }); - - it('drain still closes after delay when no new messages arrive', async () => { - (connectionFns.isConnected as ReturnType).mockReturnValue(true); - - const mgr = createManager(); - state.bindSession(createSessionContext()); - state.acceptMessage('msg_only', { autoCommit: false, condenseOnComplete: false }); - - // Complete the only message and observe root idle — triggers drain - mgr.onMessageComplete('msg_only'); - mgr.onSessionIdle(); - - // Advance past drain delay - await vi.advanceTimersByTimeAsync(1000); - - // Drain should close connections - expect(connectionFns.closeConnections).toHaveBeenCalledTimes(1); - }); + expect(runAutoCommit).not.toHaveBeenCalled(); }); }); diff --git a/services/cloud-agent-next/test/unit/wrapper/network-resume.test.ts b/services/cloud-agent-next/test/unit/wrapper/network-resume.test.ts index aa1dbea65a..53df0ae17b 100644 --- a/services/cloud-agent-next/test/unit/wrapper/network-resume.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/network-resume.test.ts @@ -102,7 +102,6 @@ const createCallbacks = (): ConnectionCallbacks & { onDisconnect: ReturnType; onTerminalError: ReturnType; } => ({ - onMessageComplete: vi.fn(), onTerminalError: vi.fn(), onCommand: vi.fn(), onDisconnect: vi.fn(), diff --git a/services/cloud-agent-next/test/unit/wrapper/reconnection.test.ts b/services/cloud-agent-next/test/unit/wrapper/reconnection.test.ts index d988be03e4..ea704b13d3 100644 --- a/services/cloud-agent-next/test/unit/wrapper/reconnection.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/reconnection.test.ts @@ -141,7 +141,6 @@ const createCallbacks = (): ConnectionCallbacks & { onRootSessionActivity: ReturnType; onSseEvent: ReturnType; } => ({ - onMessageComplete: vi.fn(), onTerminalError: vi.fn(), onCommand: vi.fn(), onDisconnect: vi.fn(), @@ -1143,7 +1142,26 @@ describe('ingest WS reconnection', () => { expect(state.observedGateResult).toBe('fail'); }); - it('reports root activity before a terminal root message completion after idle', async () => { + it('does not treat trailing root session.turn.close as resumed execution', async () => { + callbacks.onSessionIdle = vi.fn(); + const kiloClient = createMockKiloClient({ + subscribeEvents: vi.fn().mockResolvedValue({ + stream: createEventStream([ + { type: 'session.idle', properties: { sessionID: 'kilo_sess_456' } }, + { type: 'session.turn.close', properties: { sessionID: 'kilo_sess_456' } }, + ]), + }), + }); + + const manager = createManagerWithClient(kiloClient); + await openConnection(manager); + await vi.advanceTimersByTimeAsync(0); + + expect(callbacks.onSessionIdle).toHaveBeenCalledOnce(); + expect(callbacks.onRootSessionActivity).not.toHaveBeenCalled(); + }); + + it('does not treat the post-idle root session epilogue as resumed execution', async () => { const callbackOrder: string[] = []; callbacks.onSessionIdle = vi.fn(() => { callbackOrder.push('idle'); @@ -1151,17 +1169,14 @@ describe('ingest WS reconnection', () => { callbacks.onRootSessionActivity.mockImplementation(() => { callbackOrder.push('activity'); }); - callbacks.onMessageComplete = vi.fn((messageId: string) => { - callbackOrder.push(`complete:${messageId}`); - }); const kiloClient = createMockKiloClient({ subscribeEvents: vi.fn().mockResolvedValue({ stream: createEventStream([ - { - type: 'session.idle', - properties: { sessionID: 'kilo_sess_456' }, - }, + { type: 'session.idle', properties: { sessionID: 'kilo_sess_456' } }, + { type: 'session.turn.close', properties: { sessionID: 'kilo_sess_456' } }, + { type: 'session.updated', properties: { sessionID: 'kilo_sess_456' } }, + { type: 'session.diff', properties: { sessionID: 'kilo_sess_456' } }, { type: 'message.updated', properties: { @@ -1183,10 +1198,30 @@ describe('ingest WS reconnection', () => { await vi.advanceTimersByTimeAsync(0); await vi.advanceTimersByTimeAsync(0); - expect(callbacks.onSessionIdle).toHaveBeenCalledTimes(1); - expect(callbacks.onRootSessionActivity).toHaveBeenCalledTimes(1); - expect(callbacks.onMessageComplete).toHaveBeenCalledWith('msg_root_user_123'); - expect(callbackOrder).toEqual(['idle', 'activity', 'complete:msg_root_user_123']); + expect(callbacks.onSessionIdle).toHaveBeenCalledOnce(); + expect(callbacks.onRootSessionActivity).not.toHaveBeenCalled(); + expect(callbacks.onCompletionSignal).toHaveBeenCalledTimes(2); + expect(callbackOrder).toEqual(['idle']); + }); + + it('treats root turn-open and busy-status events as resumed execution', async () => { + const kiloClient = createMockKiloClient({ + subscribeEvents: vi.fn().mockResolvedValue({ + stream: createEventStream([ + { type: 'session.turn.open', properties: { sessionID: 'kilo_sess_456' } }, + { + type: 'session.status', + properties: { sessionID: 'kilo_sess_456', status: { type: 'busy' } }, + }, + ]), + }), + }); + + const manager = createManagerWithClient(kiloClient); + await openConnection(manager); + await vi.advanceTimersByTimeAsync(0); + + expect(callbacks.onRootSessionActivity).toHaveBeenCalledTimes(2); }); it('rejects real-time code-review questions without disconnecting', async () => { @@ -1212,7 +1247,6 @@ describe('ingest WS reconnection', () => { expect(questionEvents).toHaveLength(0); expect(rejectQuestion).toHaveBeenCalledWith('q_123'); expect(callbacks.onDisconnect).not.toHaveBeenCalled(); - expect(callbacks.onMessageComplete).not.toHaveBeenCalled(); }); it('rejects real-time code-review permissions without disconnecting', async () => { @@ -1245,7 +1279,6 @@ describe('ingest WS reconnection', () => { CODE_REVIEW_PERMISSION_REJECTION_MESSAGE ); expect(callbacks.onDisconnect).not.toHaveBeenCalled(); - expect(callbacks.onMessageComplete).not.toHaveBeenCalled(); }); it.each(['question', 'permission'])( @@ -1281,7 +1314,6 @@ describe('ingest WS reconnection', () => { }); expect(statusEvents).toHaveLength(0); expect(callbacks.onDisconnect).not.toHaveBeenCalled(); - expect(callbacks.onMessageComplete).not.toHaveBeenCalled(); } ); @@ -1331,7 +1363,6 @@ describe('ingest WS reconnection', () => { }); expect(callbacks.onTerminalError).toHaveBeenCalledWith('Insufficient credits'); expect(callbacks.onDisconnect).not.toHaveBeenCalled(); - expect(callbacks.onMessageComplete).not.toHaveBeenCalled(); }); // ------------------------------------------------------------------------- diff --git a/services/cloud-agent-next/test/unit/wrapper/server.test.ts b/services/cloud-agent-next/test/unit/wrapper/server.test.ts index 6d559b0935..a192380c65 100644 --- a/services/cloud-agent-next/test/unit/wrapper/server.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/server.test.ts @@ -60,7 +60,7 @@ function createMockDeps(state: WrapperState) { closeConnection: vi.fn().mockResolvedValue(undefined), setAborted: vi.fn(), resetLifecycle: vi.fn(), - onMessageComplete: vi.fn(), + onDeliveryAcknowledged: vi.fn(), readySession: vi.fn(), materializePromptAttachments: vi.fn(async prompt => prompt), configureCommitCoAuthor: vi.fn().mockResolvedValue(undefined), @@ -968,7 +968,7 @@ describe('createCommandHandler', () => { }, timestamp: expect.any(String), }); - expect(deps.onMessageComplete).toHaveBeenCalledWith('msg_compact'); + expect(deps.onDeliveryAcknowledged).toHaveBeenCalledWith('sync-command'); expect(state.getMessageConfig('msg_compact')).toEqual({ autoCommit: false, condenseOnComplete: false, @@ -1024,7 +1024,7 @@ describe('createCommandHandler', () => { message: 'Failed to send command: summarize failed', }); expect(sendToIngest).not.toHaveBeenCalled(); - expect(deps.onMessageComplete).not.toHaveBeenCalled(); + expect(deps.onDeliveryAcknowledged).toHaveBeenCalledWith('failed'); expect(state.getMessageConfig('msg_compact')).toBeNull(); }); }); @@ -1075,6 +1075,45 @@ describe('createSessionReadyHandler', () => { expect(deps.readySession).toHaveBeenCalledOnce(); }); + it('preserves finalizing retry metadata from readiness', async () => { + const state = new WrapperState(); + const deps = createMockDeps(state); + deps.readySession.mockResolvedValue({ + status: 'error', + error: { + code: 'WRAPPER_FINALIZING', + message: 'Wrapper batch is finalizing', + retryable: true, + wrapperRunId: 'run_finalizing', + }, + }); + const handler = createSessionReadyHandler(deps); + + const response = await handler( + jsonRequest({ + agentSessionId: 'agent_test', + userId: 'user_test', + sandboxId: 'usr-test', + kiloSessionId: 'kilo_sess_1', + workspace: { + workspacePath: '/workspace', + sessionHome: '/home/agent_test', + branchName: 'main', + }, + materialized: { env: { KILOCODE_TOKEN: 'kilo-token' } }, + session: completeBinding, + }) + ); + + expect(response.status).toBe(503); + await expect(readJson(response)).resolves.toEqual({ + error: 'WRAPPER_FINALIZING', + message: 'Wrapper batch is finalizing', + retryable: true, + wrapperRunId: 'run_finalizing', + }); + }); + it('maps typed bootstrap errors to JSON error responses', async () => { const state = new WrapperState(); const deps = createMockDeps(state); diff --git a/services/cloud-agent-next/test/unit/wrapper/snapshot.test.ts b/services/cloud-agent-next/test/unit/wrapper/snapshot.test.ts index a0dd531d61..0b5edb162f 100644 --- a/services/cloud-agent-next/test/unit/wrapper/snapshot.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/snapshot.test.ts @@ -115,7 +115,6 @@ const createCodeReviewSessionContext = (): SessionContext => createSessionContext({ platform: 'code-review' }); const createCallbacks = (): ConnectionCallbacks => ({ - onMessageComplete: vi.fn(), onTerminalError: vi.fn(), onCommand: vi.fn(), onDisconnect: vi.fn(), diff --git a/services/cloud-agent-next/test/unit/wrapper/state.test.ts b/services/cloud-agent-next/test/unit/wrapper/state.test.ts index f3315fe331..3e5a1e6dfc 100644 --- a/services/cloud-agent-next/test/unit/wrapper/state.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/state.test.ts @@ -1,1025 +1,96 @@ -/** - * Unit tests for WrapperState class. - * - * Tests state transitions, invariants, and edge cases for the wrapper's - * centralized state management. - */ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { WrapperState } from '../../../wrapper/src/state.js'; -import { describe, expect, it, vi, beforeEach } from 'vitest'; -import { WrapperState, type SessionContext } from '../../../wrapper/src/state.js'; -import type { IngestEvent } from '../../../src/shared/protocol.js'; +const config = { autoCommit: false, condenseOnComplete: false }; -// --------------------------------------------------------------------------- -// Test Helpers -// --------------------------------------------------------------------------- - -const createSessionContext = (overrides: Partial = {}): SessionContext => ({ - kiloSessionId: 'kilo_sess_456', - ingestUrl: 'wss://ingest.example.com', - ingestToken: 'token_secret', - workerAuthToken: 'kilo_token_789', - ...overrides, -}); - -// --------------------------------------------------------------------------- -// Tests -// --------------------------------------------------------------------------- - -describe('WrapperState', () => { +describe('WrapperState sealed batch', () => { let state: WrapperState; beforeEach(() => { state = new WrapperState(); }); - // ------------------------------------------------------------------------- - // Initial State - // ------------------------------------------------------------------------- - - describe('initial state', () => { - it('starts in idle state', () => { - expect(state.isIdle).toBe(true); - expect(state.isActive).toBe(false); - }); - - it('has no session context', () => { - expect(state.hasSession).toBe(false); - expect(state.currentSession).toBeNull(); - }); - - it('is not active', () => { - expect(state.isActive).toBe(false); - }); - - it('is not connected', () => { - expect(state.isConnected).toBe(false); - }); - - it('has no last error', () => { - expect(state.getLastError()).toBeNull(); - }); - }); - - // ------------------------------------------------------------------------- - // setActive - // ------------------------------------------------------------------------- - - describe('setActive', () => { - it('transitions state to active', () => { - expect(state.isIdle).toBe(true); - state.setActive(true); - expect(state.isActive).toBe(true); - expect(state.isIdle).toBe(false); - }); - - it('transitions state back to idle', () => { - state.setActive(true); - state.setActive(false); - expect(state.isIdle).toBe(true); - expect(state.isActive).toBe(false); - }); - - it('updates activity timestamp when activating', () => { - const before = Date.now(); - state.setActive(true); - const after = Date.now(); - const idleMs = state.getIdleMs(after); - expect(idleMs).toBeLessThanOrEqual(after - before + 1); - }); - - it('is idempotent for same value', () => { - state.setActive(true); - state.setActive(true); - expect(state.isActive).toBe(true); - - state.setActive(false); - state.setActive(false); - expect(state.isIdle).toBe(true); - }); - }); - - // ------------------------------------------------------------------------- - // Activity Tracking - // ------------------------------------------------------------------------- - - describe('activity tracking', () => { - it('updateActivity updates timestamp', () => { - const before = Date.now(); - state.updateActivity(); - const after = Date.now(); - - const idleMs = state.getIdleMs(after); - expect(idleMs).toBeLessThanOrEqual(after - before + 1); - }); - - it('getIdleMs returns time since last activity', async () => { - state.updateActivity(); - const activityTime = Date.now(); - - // Wait a bit - await new Promise(resolve => setTimeout(resolve, 50)); - - const now = Date.now(); - const idleMs = state.getIdleMs(now); - - // Should be at least 50ms but not much more - expect(idleMs).toBeGreaterThanOrEqual(now - activityTime - 5); - expect(idleMs).toBeLessThan(200); - }); - }); - - // ------------------------------------------------------------------------- - // Error Tracking - // ------------------------------------------------------------------------- - - describe('error tracking', () => { - it('setLastError stores error', () => { - const error = { - code: 'TEST_ERROR', - message: 'Something went wrong', - timestamp: Date.now(), - }; - - state.setLastError(error); - - expect(state.getLastError()).toEqual(error); - }); - - it('setLastError with messageId', () => { - const error = { - code: 'INFLIGHT_TIMEOUT', - messageId: 'msg_123', - message: 'Timeout', - timestamp: Date.now(), - }; - - state.setLastError(error); - - expect(state.getLastError()).toEqual(error); - }); - - it('clearLastError removes error', () => { - state.setLastError({ - code: 'TEST_ERROR', - message: 'Error', - timestamp: Date.now(), - }); - - state.clearLastError(); - - expect(state.getLastError()).toBeNull(); - }); - }); - - // ------------------------------------------------------------------------- - // Connection Management - // ------------------------------------------------------------------------- - - describe('connection management', () => { - it('isConnected returns false with no WebSocket', () => { - expect(state.isConnected).toBe(false); - }); - - it('setConnections stores WebSocket and AbortController', () => { - const mockWs = { readyState: WebSocket.OPEN, close: vi.fn() } as unknown as WebSocket; - const mockAbort = new AbortController(); - - state.setConnections(mockWs, mockAbort); - - expect(state.ingestWs).toBe(mockWs); - expect(state.sseAbortController).toBe(mockAbort); - }); - - it('isConnected returns true when WebSocket is OPEN', () => { - const mockWs = { readyState: WebSocket.OPEN, close: vi.fn() } as unknown as WebSocket; - state.setConnections(mockWs, new AbortController()); - - expect(state.isConnected).toBe(true); - }); - - it('isConnected returns false when WebSocket is not OPEN', () => { - const mockWs = { readyState: WebSocket.CLOSED, close: vi.fn() } as unknown as WebSocket; - state.setConnections(mockWs, new AbortController()); - - expect(state.isConnected).toBe(false); - }); - - it('clearConnectionRefs nulls references without closing or aborting', () => { - const mockClose = vi.fn(); - const mockWs = { readyState: WebSocket.OPEN, close: mockClose } as unknown as WebSocket; - const mockAbort = new AbortController(); - const abortSpy = vi.spyOn(mockAbort, 'abort'); - - state.setConnections(mockWs, mockAbort); - state.clearConnectionRefs(); - - // Refs are nulled - expect(state.ingestWs).toBeNull(); - expect(state.sseAbortController).toBeNull(); - - // clearConnectionRefs is purely passive — close/abort owned by connection.ts - expect(mockClose).not.toHaveBeenCalled(); - expect(abortSpy).not.toHaveBeenCalled(); - }); - }); - - // ------------------------------------------------------------------------- - // Send to Ingest - // ------------------------------------------------------------------------- - - describe('sendToIngest', () => { - it('does nothing when no send function set', () => { - const event: IngestEvent = { - streamEventType: 'status', - data: { message: 'test' }, - timestamp: new Date().toISOString(), - }; - - // Should not throw - expect(() => state.sendToIngest(event)).not.toThrow(); + it('tracks exact admitted membership and latest admitted finalization config', () => { + state.acceptMessage('message-1', { ...config, autoCommit: true, model: 'first-model' }); + state.acceptMessage('message-2', { + ...config, + condenseOnComplete: true, + model: 'latest-model', }); - it('calls send function when set', () => { - const mockSend = vi.fn(); - state.setSendToIngestFn(mockSend); - - const event: IngestEvent = { - streamEventType: 'status', - data: { message: 'test' }, - timestamp: new Date().toISOString(), - }; - - state.sendToIngest(event); - - expect(mockSend).toHaveBeenCalledWith(event); - }); - - it('setSendToIngestFn can clear function', () => { - const mockSend = vi.fn(); - state.setSendToIngestFn(mockSend); - state.setSendToIngestFn(null); - - const event: IngestEvent = { - streamEventType: 'status', - data: {}, - timestamp: new Date().toISOString(), - }; - - state.sendToIngest(event); - - expect(mockSend).not.toHaveBeenCalled(); + expect(state.pendingMessageIds).toEqual(['message-1', 'message-2']); + expect(state.batchFinalizationConfig).toEqual({ + autoCommit: false, + condenseOnComplete: true, + model: 'latest-model', }); }); - // ------------------------------------------------------------------------- - // Status API - // ------------------------------------------------------------------------- - - describe('getStatus', () => { - it('returns idle state with no session', () => { - const status = state.getStatus(); - - expect(status).toEqual({ - state: 'idle', - sessionId: undefined, - pendingMessages: [], - lastError: undefined, - }); - }); - - it('returns idle state with session but no inflight', () => { - state.bindSession( - createSessionContext({ - kiloSessionId: 'kilo_456', - }) - ); - - const status = state.getStatus(); - - expect(status).toEqual({ - state: 'idle', - sessionId: 'kilo_456', - pendingMessages: [], - lastError: undefined, - }); - }); - - it('returns active state when active', () => { - state.setActive(true); - state.bindSession( - createSessionContext({ - kiloSessionId: 'kilo_456', - }) - ); - const status = state.getStatus(); - expect(status).toEqual({ - state: 'active', - sessionId: 'kilo_456', - pendingMessages: [], - lastError: undefined, - }); - }); - - it('includes lastError when present', () => { - state.bindSession(createSessionContext()); - const error = { - code: 'INFLIGHT_TIMEOUT', - messageId: 'msg_123', - message: 'Timeout', - timestamp: Date.now(), - }; - state.setLastError(error); - - const status = state.getStatus(); - - expect(status.lastError).toEqual(error); - }); - - it('includes pendingMessages from message tracking', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - const status = state.getStatus(); - - expect(status.pendingMessages).toEqual(['msg_1']); - expect(status.state).toBe('active'); - }); - - it('uses sessionId from session', () => { - state.bindSession(createSessionContext({ kiloSessionId: 'session_id' })); + it('guards finalization while a delivery acknowledgement is in flight', () => { + state.acceptMessage('message-1', config); + expect(state.beginDeliveryAcknowledgement()).toBe(true); - const status = state.getStatus(); + expect(state.beginFinalizing()).toBe(false); - expect(status.sessionId).toBe('session_id'); - }); - }); - - // ------------------------------------------------------------------------- - // Session Context - // ------------------------------------------------------------------------- - - describe('session context', () => { - describe('bindSession', () => { - it('stores context on first bind', () => { - const context = createSessionContext(); - state.bindSession(context); - - expect(state.hasSession).toBe(true); - expect(state.currentSession).toEqual(context); - }); - - it('returns { changed: true } on first bind', () => { - const result = state.bindSession(createSessionContext()); - expect(result).toEqual({ changed: true }); - }); - - it('clears previous error on first bind', () => { - state.setLastError({ - code: 'TEST_ERROR', - message: 'previous error', - timestamp: Date.now(), - }); - - state.bindSession(createSessionContext()); - - expect(state.getLastError()).toBeNull(); - }); - - it('returns { changed: false } when same context', () => { - state.bindSession(createSessionContext()); - const result = state.bindSession(createSessionContext()); - expect(result).toEqual({ changed: false }); - }); - - it('returns { changed: true } when connection fields change', () => { - state.bindSession(createSessionContext()); - const result = state.bindSession( - createSessionContext({ ingestUrl: 'wss://new-ingest.example.com' }) - ); - expect(result).toEqual({ changed: true }); - }); - - it('returns { changed: true } when wrapperGeneration changes', () => { - state.bindSession(createSessionContext()); - const result = state.bindSession(createSessionContext({ wrapperGeneration: 2 })); - expect(result).toEqual({ changed: true }); - }); - - it('returns { changed: true } when wrapperConnectionId changes', () => { - state.bindSession(createSessionContext()); - const result = state.bindSession(createSessionContext({ wrapperConnectionId: 'conn_new' })); - expect(result).toEqual({ changed: true }); - }); - - it('returns { changed: true } when ingestToken changes', () => { - state.bindSession(createSessionContext()); - const result = state.bindSession(createSessionContext({ ingestToken: 'new_token' })); - expect(result).toEqual({ changed: true }); - }); - - it('returns { changed: true } when workerAuthToken changes', () => { - state.bindSession(createSessionContext()); - const result = state.bindSession( - createSessionContext({ workerAuthToken: 'new_auth_token' }) - ); - expect(result).toEqual({ changed: true }); - }); - }); - - describe('clearSession', () => { - it('clears session context', () => { - state.bindSession(createSessionContext()); - state.clearSession(); - - expect(state.hasSession).toBe(false); - expect(state.currentSession).toBeNull(); - }); - - it('clears all messages', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - state.clearSession(); - - expect(state.hasPendingMessages).toBe(false); - expect(state.pendingMessageIds).toEqual([]); - }); - - it('clears active message', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - state.clearSession(); - - expect(state.activeMessageId).toBeNull(); - }); - - it('sets state to idle', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - expect(state.isActive).toBe(true); - - state.clearSession(); - - expect(state.isActive).toBe(false); - expect(state.isIdle).toBe(true); - }); - - it('clears last assistant message id', () => { - state.bindSession(createSessionContext()); - state.setLastAssistantMessageId('assistant_msg_1'); - - state.clearSession(); - - expect(state.lastAssistantMessageId).toBeNull(); - }); - }); - - describe('hasSession / currentSession', () => { - it('hasSession is false initially', () => { - expect(state.hasSession).toBe(false); - }); - - it('currentSession is null initially', () => { - expect(state.currentSession).toBeNull(); - }); - }); + state.endDeliveryAcknowledgement(); + expect(state.beginFinalizing()).toBe(true); + expect(state.beginDeliveryAcknowledgement()).toBe(false); }); - // ------------------------------------------------------------------------- - // Message Tracking - // ------------------------------------------------------------------------- + it('keeps admissions blocked when tracked messages are cleared during drain', () => { + state.blockAdmissions(); + state.clearAllMessages(); - describe('message tracking', () => { - describe('acceptMessage', () => { - it('adds message in active state when no active message', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); + expect(state.beginDeliveryAcknowledgement()).toBe(false); - expect(state.activeMessageId).toBe('msg_1'); - expect(state.isActive).toBe(true); - }); - - it('adds message in accepted state when active message exists', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_2', { autoCommit: false, condenseOnComplete: false }); - - expect(state.activeMessageId).toBe('msg_1'); - expect(state.pendingMessageIds).toEqual(['msg_1', 'msg_2']); - }); - - it('stores config per message', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { - autoCommit: true, - condenseOnComplete: true, - model: 'claude-3', - upstreamBranch: 'main', - commitCoAuthor: { name: 'kiloconnect[bot]', email: 'bot@example.com' }, - }); - - const config = state.activeMessageConfig; - expect(config).toEqual({ - autoCommit: true, - condenseOnComplete: true, - model: 'claude-3', - upstreamBranch: 'main', - commitCoAuthor: { name: 'kiloconnect[bot]', email: 'bot@example.com' }, - }); - }); - - it('makes wrapper active when first message is accepted', () => { - state.bindSession(createSessionContext()); - expect(state.isActive).toBe(false); - - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - expect(state.isActive).toBe(true); - expect(state.isIdle).toBe(false); - }); - }); - - describe('completeActiveMessage', () => { - it('returns null when no active message', () => { - expect(state.completeActiveMessage()).toBeNull(); - }); - - it('returns completed messageId', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - const completed = state.completeActiveMessage(); - - expect(completed?.messageId).toBe('msg_1'); - }); - - it('transitions next accepted message to active', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_2', { autoCommit: false, condenseOnComplete: false }); - - state.completeActiveMessage(); - - expect(state.activeMessageId).toBe('msg_2'); - expect(state.isActive).toBe(true); - }); - - it('sets idle when no more accepted messages', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - state.completeActiveMessage(); - - expect(state.activeMessageId).toBeNull(); - expect(state.isActive).toBe(false); - expect(state.isIdle).toBe(true); - }); - - it('completes messages in FIFO order', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_2', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_3', { autoCommit: false, condenseOnComplete: false }); - - expect(state.completeActiveMessage()?.messageId).toBe('msg_1'); - expect(state.activeMessageId).toBe('msg_2'); - - expect(state.completeActiveMessage()?.messageId).toBe('msg_2'); - expect(state.activeMessageId).toBe('msg_3'); - - expect(state.completeActiveMessage()?.messageId).toBe('msg_3'); - expect(state.activeMessageId).toBeNull(); - expect(state.isActive).toBe(false); - }); - }); - - describe('activeMessageId', () => { - it('returns null initially', () => { - expect(state.activeMessageId).toBeNull(); - }); - - it('returns the currently active messageId', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - expect(state.activeMessageId).toBe('msg_1'); - }); - }); - - describe('hasPendingMessages', () => { - it('returns false initially', () => { - expect(state.hasPendingMessages).toBe(false); - }); - - it('returns true when messages are pending', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - expect(state.hasPendingMessages).toBe(true); - }); - - it('returns false when all messages are completed', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - state.completeActiveMessage(); - - expect(state.hasPendingMessages).toBe(false); - }); - }); - - describe('pendingMessageIds', () => { - it('returns empty array initially', () => { - expect(state.pendingMessageIds).toEqual([]); - }); - - it('returns all non-completed message IDs', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_2', { autoCommit: false, condenseOnComplete: false }); - - expect(state.pendingMessageIds).toEqual(['msg_1', 'msg_2']); - }); - - it('excludes completed messages', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_2', { autoCommit: false, condenseOnComplete: false }); - - state.completeActiveMessage(); - - expect(state.pendingMessageIds).toEqual(['msg_2']); - }); - }); - - describe('activeMessageConfig', () => { - it('returns null when no active message', () => { - expect(state.activeMessageConfig).toBeNull(); - }); - - it('returns config for the active message', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { - autoCommit: true, - condenseOnComplete: false, - model: 'gpt-4', - }); - - expect(state.activeMessageConfig).toEqual({ - autoCommit: true, - condenseOnComplete: false, - model: 'gpt-4', - upstreamBranch: undefined, - }); - }); - - it('returns null when active message is completed', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - state.completeActiveMessage(); - - expect(state.activeMessageConfig).toBeNull(); - }); - }); - - describe('updateMessageConfig', () => { - it('updates config for an existing message', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { - autoCommit: false, - condenseOnComplete: false, - }); - - state.updateMessageConfig('msg_1', { autoCommit: true }); - - expect(state.activeMessageConfig?.autoCommit).toBe(true); - }); - - it('does nothing for unknown message', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - state.updateMessageConfig('msg_unknown', { autoCommit: true }); - - expect(state.activeMessageConfig?.autoCommit).toBe(false); - }); - - it('updates model and upstreamBranch', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { - autoCommit: false, - condenseOnComplete: false, - }); - - state.updateMessageConfig('msg_1', { - model: 'claude-3', - upstreamBranch: 'develop', - }); - - expect(state.activeMessageConfig?.model).toBe('claude-3'); - expect(state.activeMessageConfig?.upstreamBranch).toBe('develop'); - }); - }); - - describe('removeMessage', () => { - it('removes a message from tracking', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_2', { autoCommit: false, condenseOnComplete: false }); - - state.removeMessage('msg_2'); - - expect(state.pendingMessageIds).toEqual(['msg_1']); - }); - - it('clears active state when removing active message', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - state.removeMessage('msg_1'); - - expect(state.activeMessageId).toBeNull(); - expect(state.isActive).toBe(false); - }); - - it('does not affect other messages when removing non-active', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_2', { autoCommit: false, condenseOnComplete: false }); - - state.removeMessage('msg_2'); - - expect(state.activeMessageId).toBe('msg_1'); - expect(state.isActive).toBe(true); - }); - - it('promotes next accepted message when removing the active message', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_c', { autoCommit: false, condenseOnComplete: false }); - - state.removeMessage('msg_a'); - - expect(state.activeMessageId).toBe('msg_b'); - expect(state.isActive).toBe(true); - expect(state.pendingMessageIds).toEqual(['msg_b', 'msg_c']); - }); - - it('sets idle when removing the active message with no more accepted messages', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - - state.removeMessage('msg_a'); - - expect(state.activeMessageId).toBeNull(); - expect(state.isIdle).toBe(true); - expect(state.hasPendingMessages).toBe(false); - }); - }); - - describe('clearAllMessages', () => { - it('clears all messages', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_2', { autoCommit: false, condenseOnComplete: false }); - - state.clearAllMessages(); - - expect(state.hasPendingMessages).toBe(false); - expect(state.pendingMessageIds).toEqual([]); - }); - - it('clears active message', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - state.clearAllMessages(); - - expect(state.activeMessageId).toBeNull(); - }); - - it('sets state to idle', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - state.clearAllMessages(); - - expect(state.isActive).toBe(false); - expect(state.isIdle).toBe(true); - }); - - it('does not clear session context', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_1', { autoCommit: false, condenseOnComplete: false }); - - state.clearAllMessages(); - - expect(state.hasSession).toBe(true); - }); + state.bindSession({ + kiloSessionId: 'kilo-session', + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', }); + expect(state.beginDeliveryAcknowledgement()).toBe(true); }); - // ------------------------------------------------------------------------- - // Multi-message Flow Tests - // ------------------------------------------------------------------------- - - describe('multi-message flow', () => { - it('accept A → A is active', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - - expect(state.activeMessageId).toBe('msg_a'); - expect(state.isActive).toBe(true); - expect(state.pendingMessageIds).toEqual(['msg_a']); - }); - - it('accept B → B is accepted, A still active', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); - - expect(state.activeMessageId).toBe('msg_a'); - expect(state.pendingMessageIds).toEqual(['msg_a', 'msg_b']); - }); - - it('complete A → B becomes active, not idle', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); - - const completed = state.completeActiveMessage(); - - expect(completed?.messageId).toBe('msg_a'); - expect(state.activeMessageId).toBe('msg_b'); - expect(state.isActive).toBe(true); - expect(state.isIdle).toBe(false); - }); - - it('complete B → idle', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); - - state.completeActiveMessage(); - const completed = state.completeActiveMessage(); - - expect(completed?.messageId).toBe('msg_b'); - expect(state.activeMessageId).toBeNull(); - expect(state.isActive).toBe(false); - expect(state.isIdle).toBe(true); - expect(state.hasPendingMessages).toBe(false); - }); + it('clears failed admission without changing the prior finalization config', () => { + state.acceptMessage('message-1', { ...config, autoCommit: true }); + state.acceptMessage('message-2', { ...config, condenseOnComplete: true }); - it('finalization config is per-message', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { - autoCommit: true, - condenseOnComplete: false, - model: 'claude-3', - }); - state.acceptMessage('msg_b', { - autoCommit: false, - condenseOnComplete: true, - model: 'gpt-4', - }); + state.removeMessage('message-2'); - expect(state.activeMessageConfig).toEqual({ - autoCommit: true, - condenseOnComplete: false, - model: 'claude-3', - upstreamBranch: undefined, - }); - - state.completeActiveMessage(); - - expect(state.activeMessageConfig).toEqual({ - autoCommit: false, - condenseOnComplete: true, - model: 'gpt-4', - upstreamBranch: undefined, - }); - }); - - it('abort (clearAllMessages) clears all messages mid-flow', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); - expect(state.isActive).toBe(true); - - state.clearAllMessages(); - - expect(state.activeMessageId).toBeNull(); - expect(state.hasPendingMessages).toBe(false); - expect(state.pendingMessageIds).toEqual([]); - expect(state.isActive).toBe(false); - expect(state.isIdle).toBe(true); - expect(state.hasSession).toBe(true); - }); + expect(state.pendingMessageIds).toEqual(['message-1']); + expect(state.batchFinalizationConfig).toEqual({ ...config, autoCommit: true }); }); - // ------------------------------------------------------------------------- - // completeMessage (message-ID-gated completion) - // ------------------------------------------------------------------------- - - describe('completeMessage', () => { - it('returns null for unknown message ID without changing state', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - - const result = state.completeMessage('msg_unknown'); - - expect(result).toBeNull(); - expect(state.activeMessageId).toBe('msg_a'); - expect(state.isActive).toBe(true); - }); - - it('returns null for already-completed message and does not affect the active message', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); - - const first = state.completeMessage('msg_a'); - expect(first?.messageId).toBe('msg_a'); - - const duplicate = state.completeMessage('msg_a'); - expect(duplicate).toBeNull(); - expect(state.activeMessageId).toBe('msg_b'); - expect(state.isActive).toBe(true); - }); - - it('returns null for accepted-but-not-active message', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); - - const result = state.completeMessage('msg_b'); - - expect(result).toBeNull(); - expect(state.activeMessageId).toBe('msg_a'); - expect(state.pendingMessageIds).toEqual(['msg_a', 'msg_b']); - }); - - it('completes the active message when messageId matches and promotes next accepted', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); - - const completed = state.completeMessage('msg_a'); - - expect(completed?.messageId).toBe('msg_a'); - expect(completed?.state).toBe('completed'); - expect(state.activeMessageId).toBe('msg_b'); - expect(state.isActive).toBe(true); - }); - - it('sets idle when completing the last active message', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - - const completed = state.completeMessage('msg_a'); - - expect(completed?.messageId).toBe('msg_a'); - expect(state.activeMessageId).toBeNull(); - expect(state.isIdle).toBe(true); + it('reports finalizing status and clears the batch with the session', () => { + state.bindSession({ + kiloSessionId: 'kilo-session', + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', }); + state.acceptMessage('message-1', config); + state.beginFinalizing(); - it('completes messages in FIFO order via ID gating', () => { - state.bindSession(createSessionContext()); - state.acceptMessage('msg_a', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_b', { autoCommit: false, condenseOnComplete: false }); - state.acceptMessage('msg_c', { autoCommit: false, condenseOnComplete: false }); - - expect(state.completeMessage('msg_a')?.messageId).toBe('msg_a'); - expect(state.activeMessageId).toBe('msg_b'); - - expect(state.completeMessage('msg_b')?.messageId).toBe('msg_b'); - expect(state.activeMessageId).toBe('msg_c'); - - expect(state.completeMessage('msg_c')?.messageId).toBe('msg_c'); - expect(state.activeMessageId).toBeNull(); - expect(state.isActive).toBe(false); + expect(state.getStatus()).toMatchObject({ + state: 'finalizing', + pendingMessages: ['message-1'], }); - it('returns null when no messages are tracked', () => { - expect(state.completeMessage('msg_anything')).toBeNull(); - }); + state.clearSession(); + expect(state.pendingMessageIds).toEqual([]); + expect(state.isFinalizing).toBe(false); }); - // ------------------------------------------------------------------------- - // Edge Cases and Invariants - // ------------------------------------------------------------------------- + it('sends ingest events through the current sender', () => { + const send = vi.fn(); + state.setSendToIngestFn(send); + const event = { + streamEventType: 'status' as const, + data: { message: 'test' }, + timestamp: new Date().toISOString(), + }; - describe('edge cases and invariants', () => { - it('state is IDLE when not active and ACTIVE when active', () => { - expect(state.isIdle).toBe(true); - expect(state.isActive).toBe(false); + state.sendToIngest(event); - state.setActive(true); - expect(state.isIdle).toBe(false); - expect(state.isActive).toBe(true); - - state.setActive(false); - expect(state.isIdle).toBe(true); - expect(state.isActive).toBe(false); - }); + expect(send).toHaveBeenCalledWith(event); }); }); diff --git a/services/cloud-agent-next/wrapper/src/connection.ts b/services/cloud-agent-next/wrapper/src/connection.ts index 5766ba3b6a..3686c2e9e8 100644 --- a/services/cloud-agent-next/wrapper/src/connection.ts +++ b/services/cloud-agent-next/wrapper/src/connection.ts @@ -81,12 +81,9 @@ function isRootSessionActivity( properties: Record, rootSessionID: string | undefined ): boolean { - if (!rootSessionID || eventType === 'session.idle') return false; - if (eventType === 'session.status' && statusTypeFromProperties(properties) === 'idle') { - return false; - } - - return getActivitySessionID(eventType, properties) === rootSessionID; + if (!rootSessionID || getActivitySessionID(eventType, properties) !== rootSessionID) return false; + if (eventType === 'session.turn.open') return true; + return eventType === 'session.status' && statusTypeFromProperties(properties) === 'busy'; } export const CODE_REVIEW_PERMISSION_REJECTION_MESSAGE = @@ -151,36 +148,10 @@ export function isSessionIdleEvent( return isRecord(props) && typeof props.sessionID === 'string'; } -/** - * Type guard for message.updated events where the assistant message is - * terminal (has time.completed or an error) and has a resolvable parentID. - * Used by the wrapper to detect per-message completion in the new-path - * keep-warm model. - */ -export function isAssistantMessageCompleted(data: unknown): data is { - event: 'message.updated'; - properties: { info: { role: string; parentID: string; time?: { completed?: number } } }; -} { - if (!isRecord(data)) return false; - if (data.event !== 'message.updated') return false; - const info = isRecord(data.properties) ? data.properties.info : undefined; - if (!isRecord(info)) return false; - if (info.role !== 'assistant') return false; - if (typeof info.parentID !== 'string') return false; +function isAssistantCompletionSignal(info: unknown): boolean { + if (!isRecord(info) || info.role !== 'assistant') return false; const time = isRecord(info.time) ? info.time : undefined; - const isCompleted = time !== undefined && typeof time.completed === 'number'; - const hasError = info.error !== undefined && info.error !== null; - return isCompleted || hasError; -} - -/** - * Extracts the parentID of a completed assistant message from a - * message.updated event. Returns undefined if the event is not a - * terminal assistant message update. - */ -export function getCompletedAssistantParentID(data: unknown): string | undefined { - if (!isAssistantMessageCompleted(data)) return undefined; - return (data.properties.info as { parentID: string }).parentID; + return typeof time?.completed === 'number' || (info.error !== undefined && info.error !== null); } // --------------------------------------------------------------------------- @@ -192,8 +163,6 @@ export type ConnectionConfig = { }; export type ConnectionCallbacks = { - /** Called when a completion event is detected for a message */ - onMessageComplete: (messageId: string) => void; /** Called when a terminal error is detected */ onTerminalError: (reason: string) => void; /** Called when a command is received from DO */ @@ -1025,13 +994,7 @@ export function createConnectionManager( state.setLastAssistantMessageId(messageInfo.id); } } - - // Detect terminal assistant messages for per-message completion in - // the new-path keep-warm model. - const data = { event: eventType as 'message.updated', properties }; - const parentID = getCompletedAssistantParentID(data); - if (parentID) { - callbacks.onMessageComplete(parentID); + if (isAssistantCompletionSignal(messageInfo)) { callbacks.onCompletionSignal(); } } diff --git a/services/cloud-agent-next/wrapper/src/lifecycle.test.ts b/services/cloud-agent-next/wrapper/src/lifecycle.test.ts index ad764ed9ee..234821f6be 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.test.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.test.ts @@ -40,19 +40,19 @@ describe('wrapper lifecycle drain races', () => { autoCommit: false, condenseOnComplete: false, }); - lifecycle.setAborted(); state.clearAllMessages(); + lifecycle.setAborted(); lifecycle.triggerDrainAndClose(); + lifecycle.reset(); state.acceptMessage('message-2', { autoCommit: false, condenseOnComplete: false, }); await wait(300); - lifecycle.onMessageComplete('message-2'); lifecycle.onSessionIdle(); - await wait(50); + await wait(3_050); expect(events.map(event => event.streamEventType)).toContain('complete'); }); diff --git a/services/cloud-agent-next/wrapper/src/lifecycle.ts b/services/cloud-agent-next/wrapper/src/lifecycle.ts index 31ac73965b..1f1ce3eb73 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.ts @@ -1,115 +1,67 @@ -/** - * Lifecycle management for the long-running wrapper. - * - * Handles: - * - SSE transport timer (15s reconnect on inactivity) - * - Drain period (grace period before closing connections) - * - Auto-commit and condense on completion - */ - import type { WrapperState } from './state.js'; -import type { WrapperCommitCoAuthor } from '../../src/shared/wrapper-bootstrap.js'; import type { WrapperKiloClient } from './kilo-api.js'; import { runAutoCommit } from './auto-commit.js'; import { runCondenseOnComplete } from './condense-on-complete.js'; import { getCurrentBranch, logToFile } from './utils.js'; -// --------------------------------------------------------------------------- -// Constants -// --------------------------------------------------------------------------- - -/** Grace period before closing connections after inflight hits 0 (250ms) */ const DRAIN_DELAY_MS = 250; - -/** If no SSE event arrives within this window, reconnect the event subscription (15s) */ +const STABLE_ROOT_IDLE_MS = 3_000; const SSE_TRANSPORT_TIMEOUT_MS = 15_000; - -/** Overall timeout for auto-commit operation (2 minutes) */ const AUTO_COMMIT_TIMEOUT_MS = 120_000; -// --------------------------------------------------------------------------- -// Types -// --------------------------------------------------------------------------- - export type LifecycleConfig = { - /** Workspace path for auto-commit/condense (session-stable) */ workspacePath: string; }; -/** - * Per-turn options set by the prompt handler. - * These override defaults for the current turn and are consumed - * when session.idle fires (auto-commit, condense, timeout). - */ -export type PerTurnConfig = { - autoCommit: boolean; - condenseOnComplete: boolean; - model?: string; - upstreamBranch?: string; - commitCoAuthor?: WrapperCommitCoAuthor; -}; - export type LifecycleDependencies = { state: WrapperState; kiloClient: WrapperKiloClient; - /** Close all connections (ingest WS + event subscription) */ closeConnections: () => Promise; - /** Check if ingest WS is currently connected */ isConnected: () => boolean; - /** Abort and restart the SDK event subscription */ reconnectEventSubscription: () => void; }; export type LifecycleManager = { - /** Start lifecycle monitoring */ start: () => void; - /** Stop lifecycle monitoring */ stop: () => void; - /** Called when a message completes - checks if idle */ - onMessageComplete: (messageId: string) => void; - /** Called when the root Kilo session reaches idle. */ onSessionIdle: () => void; - /** Called when the root Kilo session emits activity after idle. */ onRootSessionActivity: () => void; - /** Called when ingest connectivity is restored after a reconnect. */ + onDeliveryAcknowledged: (kind: 'async-prompt' | 'sync-command' | 'failed') => void; onConnectionRestored: () => void; - /** Called to trigger drain and close sequence */ triggerDrainAndClose: () => void; - /** Signal completion for post-processing waiters (called by connection on completion events) */ signalCompletion: () => void; - /** Set the aborted flag to prevent post-completion tasks from running */ setAborted: () => void; - /** Reset lifecycle state for a new execution (clears isAborted, isDraining, etc.) */ reset: () => void; - /** Called by connection manager on every SSE event to reset the transport timer */ onSseEvent: () => void; }; -// --------------------------------------------------------------------------- -// Lifecycle Manager -// --------------------------------------------------------------------------- - export function createLifecycleManager( config: LifecycleConfig, deps: LifecycleDependencies ): LifecycleManager { const { state, kiloClient } = deps; - let sseTransportTimer: ReturnType | null = null; + let stableIdleTimer: ReturnType | null = null; let drainTimeout: ReturnType | null = null; let isDraining = false; let isAborted = false; - let rootSessionIdleBarrierPresent = false; - - // Completion waiter for post-processing tasks (auto-commit, condense) + let rootIdleCandidatePresent = false; + let idleObservedDuringDelivery = false; let postProcessingResolve: (() => void) | null = null; let postProcessingCompleted = false; function clearSseTransportTimer(): void { - if (sseTransportTimer) { - clearTimeout(sseTransportTimer); - sseTransportTimer = null; - } + if (!sseTransportTimer) return; + clearTimeout(sseTransportTimer); + sseTransportTimer = null; + } + + function clearStableIdleCandidate(): void { + rootIdleCandidatePresent = false; + idleObservedDuringDelivery = false; + if (!stableIdleTimer) return; + clearTimeout(stableIdleTimer); + stableIdleTimer = null; } function resetSseTransportTimer(): void { @@ -121,44 +73,23 @@ export function createLifecycleManager( }, SSE_TRANSPORT_TIMEOUT_MS); } - /** - * Signal that a completion event was received (called by connection manager). - * This resolves any pending waitForCompletion() promises used by post-processing tasks. - */ function signalCompletion(): void { postProcessingCompleted = true; - if (postProcessingResolve) { - postProcessingResolve(); - postProcessingResolve = null; - } + postProcessingResolve?.(); + postProcessingResolve = null; } - /** - * Run post-completion tasks (auto-commit, condense). - */ async function runPostCompletionTasks(): Promise { const session = state.currentSession; - if (!session) return; - - const msgConfig = state.completedMessageConfig ?? state.activeMessageConfig; - if (!msgConfig) { - logToFile('no message config for post-completion tasks — skipping'); - return; - } - - if (isAborted) { - logToFile('skipping post-completion tasks — session was aborted'); - return; - } + const msgConfig = state.batchFinalizationConfig; + if (!session || !msgConfig || isAborted) return; if (msgConfig.autoCommit) { - logToFile('running auto-commit'); try { const autoCommitController = new AbortController(); let autoCommitTimedOut = false; const timeout = setTimeout(() => { autoCommitTimedOut = true; - logToFile('auto-commit lifecycle timeout reached; aborting in-flight work'); autoCommitController.abort(); }, AUTO_COMMIT_TIMEOUT_MS); const result = await runAutoCommit({ @@ -171,44 +102,33 @@ export function createLifecycleManager( signal: autoCommitController.signal, }).finally(() => clearTimeout(timeout)); if (autoCommitTimedOut && !result.success) { - logToFile('auto-commit aborted by lifecycle timeout'); state.sendToIngest({ streamEventType: 'error', data: { error: 'Auto-commit timed out', fatal: false }, timestamp: new Date().toISOString(), }); - } else { - logToFile( - `auto-commit complete: success=${result.success} skipped=${result.skipped ?? false} error=${result.error ?? '(none)'}` - ); } } catch (error) { - const msg = error instanceof Error ? error.message : String(error); - logToFile(`auto-commit error: ${msg}`); + const message = error instanceof Error ? error.message : String(error); state.sendToIngest({ streamEventType: 'error', - data: { error: `Auto-commit failed: ${msg}`, fatal: false }, + data: { error: `Auto-commit failed: ${message}`, fatal: false }, timestamp: new Date().toISOString(), }); } } - const expectCompletion = () => { - postProcessingCompleted = false; - postProcessingResolve = null; - }; - - const waitForCompletion = (): Promise => { - if (postProcessingCompleted) return Promise.resolve(); - return new Promise(resolve => { - postProcessingResolve = resolve; - }); - }; - - const wasAborted = () => isAborted; - if (msgConfig.condenseOnComplete) { - logToFile('running condense'); + const expectCompletion = () => { + postProcessingCompleted = false; + postProcessingResolve = null; + }; + const waitForCompletion = (): Promise => { + if (postProcessingCompleted) return Promise.resolve(); + return new Promise(resolve => { + postProcessingResolve = resolve; + }); + }; try { await runCondenseOnComplete({ workspacePath: config.workspacePath, @@ -218,89 +138,74 @@ export function createLifecycleManager( kiloClient, expectCompletion, waitForCompletion, - wasAborted, + wasAborted: () => isAborted, }); - logToFile('condense complete'); } catch (error) { - const msg = error instanceof Error ? error.message : String(error); - logToFile(`condense error: ${msg}`); + const message = error instanceof Error ? error.message : String(error); state.sendToIngest({ streamEventType: 'error', - data: { error: `Condense failed: ${msg}`, fatal: false }, + data: { error: `Condense failed: ${message}`, fatal: false }, timestamp: new Date().toISOString(), }); } } } - /** - * Trigger drain period and close connections. - * Runs post-completion tasks (auto-commit, condense), sends complete event, then closes after drain delay. - */ function triggerDrainAndClose(): void { + state.blockAdmissions(); if (isDraining) return; isDraining = true; + clearStableIdleCandidate(); + const sealedMessageIds = state.pendingMessageIds; + const session = state.currentSession; - logToFile(`starting drain period (isAborted=${isAborted})`); + if (session && !isAborted) { + state.sendToIngest({ + streamEventType: 'wrapper_finalizing', + data: { wrapperRunId: session.wrapperRunId }, + timestamp: new Date().toISOString(), + }); + } - // Run the full drain sequence as a single async flow. - // Order matters: post-completion tasks (autocommit/condense) → log upload → - // complete event → drain delay → close connections. void (async () => { try { await runPostCompletionTasks(); - - // Final log upload const uploader = state.logUploader; if (uploader) { try { await uploader.uploadNow(); - } catch (err) { + } catch (error) { logToFile( - `final log upload failed: ${err instanceof Error ? err.message : String(err)}` + `final log upload failed: ${error instanceof Error ? error.message : String(error)}` ); } uploader.stop(); } } finally { - // 3. Send complete event (always runs, even if upload/post-processing failed) - const session = state.currentSession; - if (session && !isAborted) { + const currentSession = state.currentSession; + if (currentSession && !isAborted) { const currentBranch = await getCurrentBranch(config.workspacePath, 10_000).catch( () => '' ); const gateResult = state.consumeObservedGateResult(); - logToFile( - `sending complete event for kiloSessionId=${session.kiloSessionId} branch=${currentBranch || '(none)'}` - ); state.sendToIngest({ streamEventType: 'complete', data: { exitCode: 0, - kiloSessionId: session.kiloSessionId, + kiloSessionId: currentSession.kiloSessionId, + messageIds: sealedMessageIds, ...(currentBranch ? { currentBranch } : {}), ...(gateResult ? { gateResult } : {}), }, timestamp: new Date().toISOString(), }); - } else if (session && isAborted) { - logToFile('skipping complete event — session was aborted'); } - // 4. Drain delay, then close connections (if no new messages arrived) drainTimeout = setTimeout(() => { - if (state.isActive) { - logToFile(`drain aborted — wrapper became active again during drain`); - isDraining = false; - isAborted = false; - drainTimeout = null; - return; - } - logToFile('drain complete, closing connections'); - deps + void deps .closeConnections() - .catch(err => - logToFile(`close failed: ${err instanceof Error ? err.message : String(err)}`) + .catch(error => + logToFile(`close failed: ${error instanceof Error ? error.message : String(error)}`) ) .finally(() => { isDraining = false; @@ -312,88 +217,81 @@ export function createLifecycleManager( })(); } - function maybeFinalizeIdleBatch(): void { - if (!state.isIdle || !deps.isConnected() || !rootSessionIdleBarrierPresent) { + function trySealIdleBatch(): void { + stableIdleTimer = null; + if (!rootIdleCandidatePresent || state.deliveryAcknowledgementsInFlight > 0) { return; } - triggerDrainAndClose(); - } - - /** - * Handle message completion. - */ - function onMessageComplete(messageId: string): void { - const completedConfig = state.getMessageConfig(messageId); - const completedInfo = state.completeMessage(messageId); - if (completedInfo) { - logToFile( - `message complete: messageId=${completedInfo.messageId}, pending: [${state.pendingMessageIds.join(',')}]` - ); + if (!deps.isConnected()) { + armStableIdleCandidate(); + return; } - - if (completedInfo && completedConfig && state.hasSession) { - state.setCompletedMessageConfig(completedConfig); + if (state.beginFinalizing()) { + triggerDrainAndClose(); } - - maybeFinalizeIdleBatch(); - } - - function onSessionIdle(): void { - rootSessionIdleBarrierPresent = true; - maybeFinalizeIdleBatch(); } - function onRootSessionActivity(): void { - rootSessionIdleBarrierPresent = false; + function armStableIdleCandidate(): void { + if (!rootIdleCandidatePresent || stableIdleTimer || !state.hasPendingMessages) return; + stableIdleTimer = setTimeout(trySealIdleBatch, STABLE_ROOT_IDLE_MS); } - function onConnectionRestored(): void { - maybeFinalizeIdleBatch(); + function restartStableIdleCandidate(): void { + if (stableIdleTimer) clearTimeout(stableIdleTimer); + stableIdleTimer = null; + armStableIdleCandidate(); } return { - start: () => { - logToFile('lifecycle started (transport timer is event-driven)'); - }, - + start: () => logToFile('lifecycle started (transport timer is event-driven)'), stop: () => { - logToFile('stopping lifecycle'); isAborted = true; clearSseTransportTimer(); - - if (drainTimeout) { - clearTimeout(drainTimeout); - drainTimeout = null; + clearStableIdleCandidate(); + if (drainTimeout) clearTimeout(drainTimeout); + drainTimeout = null; + }, + onSessionIdle: () => { + rootIdleCandidatePresent = true; + if (state.deliveryAcknowledgementsInFlight > 0) idleObservedDuringDelivery = true; + armStableIdleCandidate(); + }, + onRootSessionActivity: clearStableIdleCandidate, + onDeliveryAcknowledged: kind => { + if (kind === 'async-prompt') { + if (!idleObservedDuringDelivery) { + clearStableIdleCandidate(); + return; + } + if (state.deliveryAcknowledgementsInFlight > 0) return; + idleObservedDuringDelivery = false; + restartStableIdleCandidate(); + return; + } + if (state.deliveryAcknowledgementsInFlight === 0) idleObservedDuringDelivery = false; + if (kind === 'sync-command') { + rootIdleCandidatePresent = true; } + armStableIdleCandidate(); }, - - onMessageComplete, - onSessionIdle, - onRootSessionActivity, - onConnectionRestored, + onConnectionRestored: armStableIdleCandidate, triggerDrainAndClose, signalCompletion, - setAborted: () => { isAborted = true; + state.blockAdmissions(); + clearStableIdleCandidate(); }, - reset: () => { isAborted = false; isDraining = false; - rootSessionIdleBarrierPresent = false; + clearStableIdleCandidate(); postProcessingCompleted = false; postProcessingResolve = null; - state.clearCompletedMessageConfig(); clearSseTransportTimer(); - if (drainTimeout) { - clearTimeout(drainTimeout); - drainTimeout = null; - } - }, - - onSseEvent: () => { - resetSseTransportTimer(); + if (drainTimeout) clearTimeout(drainTimeout); + drainTimeout = null; }, + onSseEvent: resetSseTransportTimer, }; } diff --git a/services/cloud-agent-next/wrapper/src/main.ts b/services/cloud-agent-next/wrapper/src/main.ts index 4b1fe63100..1040da4545 100644 --- a/services/cloud-agent-next/wrapper/src/main.ts +++ b/services/cloud-agent-next/wrapper/src/main.ts @@ -260,7 +260,8 @@ async function main() { closeConnection: () => connectionManager?.close() ?? Promise.resolve(), setAborted: () => lifecycleManager?.setAborted(), resetLifecycle: () => lifecycleManager?.reset(), - onMessageComplete: (messageId: string) => lifecycleManager?.onMessageComplete(messageId), + onDeliveryAcknowledged: (kind: 'async-prompt' | 'sync-command' | 'failed') => + lifecycleManager?.onDeliveryAcknowledged(kind), readySession: readySession, updateRuntimeEnvironment: updateRuntimeEnvironment, materializePromptAttachments, @@ -392,9 +393,6 @@ async function main() { state, { kiloClient: nextKiloClient }, { - onMessageComplete: (messageId: string) => { - lifecycleManager?.onMessageComplete(messageId); - }, onTerminalError: (reason: string) => { logToFile(`terminal error: ${reason}`); state.sendToIngest({ @@ -455,7 +453,6 @@ async function main() { nextKiloClient.abortSession({ sessionId: targetSessionId }).catch(() => {}); } lifecycleManager?.setAborted(); - state.setActive(false); lifecycleManager?.triggerDrainAndClose(); }, onCompletionSignal: () => { @@ -517,10 +514,6 @@ async function main() { `session/ready received agentSessionId=${request.agentSessionId} kiloSessionId=${request.kiloSessionId} preferSnapshot=${request.workspace.preferSnapshot} workspacePath=${request.workspace.workspacePath} sessionHome=${request.workspace.sessionHome} branchName=${request.workspace.branchName} strictBranch=${request.workspace.strictBranch ?? false} repoKind=${request.repo?.kind ?? '(none)'} setupCommandCount=${request.materialized.setupCommands?.length ?? 0} runtimeSkillCount=${request.materialized.runtimeSkills?.length ?? 0} platform=${request.materialized.env.KILO_PLATFORM ?? process.env.KILO_PLATFORM ?? '(unset)'} stateConnected=${state.isConnected}` ); try { - serverConfig.workspacePath = request.workspace.workspacePath; - serverConfig.sessionId = request.kiloSessionId; - serverConfig.platform = request.materialized.env.KILO_PLATFORM ?? process.env.KILO_PLATFORM; - const bindError = await bindSessionContext( request.session, serverConfig, @@ -528,20 +521,31 @@ async function main() { 'close-until-runtime-ready' ); if (bindError) { - const error = (await bindError.json()) as { error?: string; message?: string }; + const error = (await bindError.json()) as { + error?: string; + message?: string; + wrapperRunId?: string; + }; + const code = + error.error === 'WRAPPER_FINALIZING' ? 'WRAPPER_FINALIZING' : 'INVALID_REQUEST'; logToFile( `session/ready binding rejected kiloSessionId=${request.kiloSessionId} status=${bindError.status} message=${error.message ?? error.error ?? 'Invalid session binding'} elapsedMs=${Date.now() - readyStartedAt}` ); return { status: 'error', error: { - code: 'INVALID_REQUEST', + code, message: error.message ?? error.error ?? 'Invalid session binding', - retryable: false, + retryable: code === 'WRAPPER_FINALIZING', + ...(error.wrapperRunId ? { wrapperRunId: error.wrapperRunId } : {}), }, }; } + serverConfig.workspacePath = request.workspace.workspacePath; + serverConfig.sessionId = request.kiloSessionId; + serverConfig.platform = request.materialized.env.KILO_PLATFORM ?? process.env.KILO_PLATFORM; + if (!state.isConnected) { progressChannel = await openIngestProgressChannel(state); } diff --git a/services/cloud-agent-next/wrapper/src/server.test.ts b/services/cloud-agent-next/wrapper/src/server.test.ts index 94400c5748..62a824c2b8 100644 --- a/services/cloud-agent-next/wrapper/src/server.test.ts +++ b/services/cloud-agent-next/wrapper/src/server.test.ts @@ -91,7 +91,6 @@ function createTestFetch(overrides?: { closeConnection: async () => {}, setAborted: () => {}, resetLifecycle: () => {}, - setPerTurnConfig: () => {}, updateRuntimeEnvironment: async env => { runtimeEnvironmentUpdates.push(env); }, @@ -215,7 +214,6 @@ describe('wrapper PTY routes', () => { closeConnection: async () => {}, setAborted: () => {}, resetLifecycle: () => {}, - setPerTurnConfig: () => {}, }, () => {} ); @@ -329,6 +327,49 @@ describe('wrapper Kilo proxy route', () => { }); describe('wrapper session binding', () => { + it('rejects even the current binding while the wrapper is finalizing', async () => { + const state = new WrapperState(); + const sessionBinding = { + kiloSessionId: 'kilo_sess_test', + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', + wrapperRunId: 'run_1', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_1', + agentSessionId: 'agent_00000000-0000-0000-0000-000000000000', + }; + state.bindSession(sessionBinding); + state.blockAdmissions(); + + const response = await bindSessionContext( + sessionBinding, + { + port: 5000, + workspacePath: '/workspace/repo', + version: 'test', + sessionId: 'kilo_sess_test', + agentSessionId: 'agent_00000000-0000-0000-0000-000000000000', + userId: 'user_test', + }, + { + state, + kiloClient: {} as WrapperKiloClient, + openConnection: async () => {}, + closeConnection: async () => {}, + setAborted: () => {}, + resetLifecycle: () => {}, + }, + 'close-until-runtime-ready' + ); + + if (!response) throw new Error('Expected finalizing binding rejection'); + expect(response.status).toBe(409); + expect(await response.json()).toMatchObject({ + error: 'WRAPPER_FINALIZING', + wrapperRunId: 'run_1', + }); + }); + it('keeps bootstrap rebindings close-only until runtime readiness is verified', async () => { const state = new WrapperState(); state.bindSession({ diff --git a/services/cloud-agent-next/wrapper/src/server.ts b/services/cloud-agent-next/wrapper/src/server.ts index 526ed78942..b65172575a 100644 --- a/services/cloud-agent-next/wrapper/src/server.ts +++ b/services/cloud-agent-next/wrapper/src/server.ts @@ -15,7 +15,6 @@ import type { WrapperState, SessionContext } from './state.js'; import type { WrapperKiloClient, WrapperPtySize } from './kilo-api.js'; -import type { PerTurnConfig } from './lifecycle.js'; import { createLogUploader } from './log-uploader.js'; import { configureCommitCoAuthorHook } from './commit-co-author-hook.js'; import { logToFile } from './utils.js'; @@ -62,10 +61,8 @@ export type ServerDependencies = { setAborted: () => void; /** Reset lifecycle state for a new execution */ resetLifecycle: () => void; - /** Mark a submitted message complete when the wrapper handles a synchronous session action. */ - onMessageComplete?: (messageId: string) => void; - /** Compatibility hook for callers that still construct wrapper server test deps. */ - setPerTurnConfig?: (config: PerTurnConfig) => void; + /** Notify lifecycle after an acknowledgement guard clears. */ + onDeliveryAcknowledged?: (kind: 'async-prompt' | 'sync-command' | 'failed') => void; /** Workspace/Kilo readiness path */ readySession?: (request: WrapperSessionReadyRequest) => Promise; /** Apply refreshed runtime variables to the active Kilo runtime. */ @@ -160,6 +157,17 @@ function errorResponse(error: string, message: string, status: number): Response return jsonResponse({ error, message }, status); } +function wrapperFinalizingResponse(state: WrapperState): Response { + return jsonResponse( + { + error: 'WRAPPER_FINALIZING', + message: 'Wrapper batch is finalizing', + wrapperRunId: state.finalizingWrapperRunId, + }, + 409 + ); +} + async function applyCommitAttribution( workspacePath: string, commitCoAuthor: WrapperCommitCoAuthor | undefined, @@ -245,6 +253,17 @@ export async function bindSessionContext( feedPolicy: SessionBoundFeedPolicy = 'restart' ): Promise { const { state } = deps; + const blockedWrapperRunId = state.finalizingWrapperRunId; + const isFreshRunAfterFinalization = + state.admissionsBlocked && + !state.hasSession && + blockedWrapperRunId !== undefined && + binding?.wrapperRunId !== undefined && + binding.wrapperRunId !== blockedWrapperRunId; + + if (state.admissionsBlocked && !isFreshRunAfterFinalization) { + return wrapperFinalizingResponse(state); + } if (!binding) { if (!state.hasSession) { @@ -397,9 +416,17 @@ export function createPromptHandler(config: ServerConfig, deps: ServerDependenci const binding = body.session; const bindError = await bindSessionContext(binding, config, deps); if (bindError) return bindError; + if (!state.beginDeliveryAcknowledgement()) { + return wrapperFinalizingResponse(state); + } + const acknowledgeDelivery = (kind: 'async-prompt' | 'failed') => { + state.endDeliveryAcknowledgement(); + deps.onDeliveryAcknowledged?.(kind); + }; const session = state.currentSession; if (!session) { + acknowledgeDelivery('failed'); return errorResponse('NO_SESSION', 'No session context available', 400); } const messageId = body.message.id; @@ -413,6 +440,7 @@ export function createPromptHandler(config: ServerConfig, deps: ServerDependenci } catch (error) { const msg = error instanceof Error ? error.message : String(error); logToFile(`job/prompt: failed to materialize attachments: ${msg}`); + acknowledgeDelivery('failed'); return errorResponse('SEND_ERROR', `Failed to materialize attachments: ${msg}`, 500); } } @@ -422,7 +450,10 @@ export function createPromptHandler(config: ServerConfig, deps: ServerDependenci prompt.finalization?.commitCoAuthor, deps.configureCommitCoAuthor ?? configureCommitCoAuthorHook ); - if (attributionError) return attributionError; + if (attributionError) { + acknowledgeDelivery('failed'); + return attributionError; + } if (!state.isConnected) { try { @@ -431,11 +462,12 @@ export function createPromptHandler(config: ServerConfig, deps: ServerDependenci } catch (error) { const msg = error instanceof Error ? error.message : String(error); logToFile(`job/prompt: failed to open connection: ${msg}`); + acknowledgeDelivery('failed'); return errorResponse('CONNECTION_ERROR', `Failed to open connection: ${msg}`, 500); } } - state.acceptMessage(messageId, { + const addedMessage = state.acceptMessage(messageId, { autoCommit: prompt.finalization?.autoCommit ?? false, condenseOnComplete: prompt.finalization?.condenseOnComplete ?? false, model: prompt.agent?.model?.modelID, @@ -458,8 +490,10 @@ export function createPromptHandler(config: ServerConfig, deps: ServerDependenci tools: prompt.agent?.tools, }); logToFile(`job/prompt: sent messageId=${messageId}`); + acknowledgeDelivery('async-prompt'); } catch (error) { - state.removeMessage(messageId); + if (addedMessage) state.removeMessage(messageId); + acknowledgeDelivery('failed'); const msg = error instanceof Error ? error.message : String(error); logToFile(`job/prompt: failed to send: ${msg}`); return errorResponse('SEND_ERROR', `Failed to send prompt: ${msg}`, 500); @@ -482,17 +516,27 @@ export function createCommandHandler(config: ServerConfig, deps: ServerDependenc const bindError = await bindSessionContext(body.session ?? body.execution, config, deps); if (bindError) return bindError; + if (!state.beginDeliveryAcknowledgement()) { + return wrapperFinalizingResponse(state); + } + const acknowledgeDelivery = (kind: 'sync-command' | 'failed') => { + state.endDeliveryAcknowledgement(); + deps.onDeliveryAcknowledged?.(kind); + }; const session = state.currentSession; if (!session) { + acknowledgeDelivery('failed'); return errorResponse('NO_SESSION', 'No session context available', 400); } if (!body.command) { + acknowledgeDelivery('failed'); return errorResponse('INVALID_REQUEST', 'command is required', 400); } const compactModel = body.command === 'compact' ? body.agent?.model : undefined; if (body.command === 'compact' && !compactModel?.modelID) { + acknowledgeDelivery('failed'); return errorResponse('INVALID_REQUEST', 'model is required for compact', 400); } @@ -501,19 +545,22 @@ export function createCommandHandler(config: ServerConfig, deps: ServerDependenc body.commitCoAuthor, deps.configureCommitCoAuthor ?? configureCommitCoAuthorHook ); - if (attributionError) return attributionError; + if (attributionError) { + acknowledgeDelivery('failed'); + return attributionError; + } const binding = body.session ?? body.execution; const messageId = body.messageId; - if (messageId) { - state.acceptMessage(messageId, { - autoCommit: body.autoCommit ?? false, - condenseOnComplete: body.condenseOnComplete ?? false, - model: body.agent?.model?.modelID, - upstreamBranch: binding?.upstreamBranch, - ...(body.commitCoAuthor ? { commitCoAuthor: body.commitCoAuthor } : {}), - }); - } + const addedMessage = messageId + ? state.acceptMessage(messageId, { + autoCommit: body.autoCommit ?? false, + condenseOnComplete: body.condenseOnComplete ?? false, + model: body.agent?.model?.modelID, + upstreamBranch: binding?.upstreamBranch, + ...(body.commitCoAuthor ? { commitCoAuthor: body.commitCoAuthor } : {}), + }) + : false; if (!state.isConnected) { try { @@ -522,7 +569,8 @@ export function createCommandHandler(config: ServerConfig, deps: ServerDependenc } catch (error) { const msg = error instanceof Error ? error.message : String(error); logToFile(`job/command: failed to open connection: ${msg}`); - if (messageId) state.removeMessage(messageId); + if (messageId && addedMessage) state.removeMessage(messageId); + acknowledgeDelivery('failed'); return errorResponse('CONNECTION_ERROR', `Failed to open connection: ${msg}`, 500); } } @@ -543,7 +591,6 @@ export function createCommandHandler(config: ServerConfig, deps: ServerDependenc }, timestamp: new Date().toISOString(), }); - deps.onMessageComplete?.(messageId); } } else { result = await kiloClient.sendCommand({ @@ -555,9 +602,11 @@ export function createCommandHandler(config: ServerConfig, deps: ServerDependenc } state.updateActivity(); logToFile(`job/command: sent command=${body.command}`); + acknowledgeDelivery('sync-command'); return jsonResponse({ status: 'sent', result }); } catch (error) { - if (messageId) state.removeMessage(messageId); + if (messageId && addedMessage) state.removeMessage(messageId); + acknowledgeDelivery('failed'); const msg = error instanceof Error ? error.message : String(error); logToFile(`job/command: failed: ${msg}`); return errorResponse('COMMAND_ERROR', `Failed to send command: ${msg}`, 500); @@ -932,6 +981,7 @@ export function createSessionReadyHandler(deps: ServerDependencies) { error: result.error.code, message: result.error.message, ...(result.error.retryable !== undefined ? { retryable: result.error.retryable } : {}), + ...(result.error.wrapperRunId ? { wrapperRunId: result.error.wrapperRunId } : {}), }, status ); diff --git a/services/cloud-agent-next/wrapper/src/state.ts b/services/cloud-agent-next/wrapper/src/state.ts index e5e865559a..9787008b29 100644 --- a/services/cloud-agent-next/wrapper/src/state.ts +++ b/services/cloud-agent-next/wrapper/src/state.ts @@ -1,26 +1,8 @@ -/** - * WrapperState - Single source of truth for wrapper state. - * - * All wrapper state is centralized here. Other modules receive a WrapperState - * instance and interact with it through methods. This makes state transitions - * explicit, simplifies testing, and prevents scattered state bugs. - * - * Session-level multi-message model: - * - Session context: shared connection parameters across messages - * - Message tracking: each message has state 'accepted' | 'active' | 'completed' - * - At most one 'active' message at a time; others are 'accepted' (queued) - * - Wrapper is idle only when no messages are pending - */ - import type { IngestEvent } from '../../src/shared/protocol.js'; import type { WrapperCommitCoAuthor } from '../../src/shared/wrapper-bootstrap.js'; import type { LogUploader } from './log-uploader.js'; export type { LogUploader } from './log-uploader.js'; -// --------------------------------------------------------------------------- -// Types -// --------------------------------------------------------------------------- - export type SessionContext = { kiloSessionId: string; ingestUrl: string; @@ -33,11 +15,7 @@ export type SessionContext = { agentSessionId?: string; }; -export type MessageState = 'accepted' | 'active' | 'completed'; - -export type MessageInfo = { - messageId: string; - state: MessageState; +export type FinalizationConfig = { autoCommit: boolean; condenseOnComplete: boolean; model?: string; @@ -53,86 +31,68 @@ export type LastError = { }; export type WrapperStatus = { - state: 'idle' | 'active'; + state: 'idle' | 'active' | 'finalizing'; sessionId?: string; pendingMessages: string[]; lastError?: LastError; }; -// --------------------------------------------------------------------------- -// WrapperState Class -// --------------------------------------------------------------------------- - export class WrapperState { - private _isActive = false; - - // Session-level context (set on first bind, updated on subsequent binds) private session: SessionContext | null = null; - - // Per-message tracking - private messages: Map = new Map(); - - // The currently active messageId (Kilo is processing it) - private _activeMessageId: string | null = null; - - // Connection state - managed externally, stored here for reference + private admittedMessages = new Map(); + private latestAdmittedFinalizationConfig: FinalizationConfig | null = null; + private _deliveryAcknowledgementsInFlight = 0; + private _admissionsBlocked = false; + private _blockedWrapperRunId: string | undefined; + private _isFinalizing = false; private _ingestWs: WebSocket | null = null; private _sseAbortController: AbortController | null = null; - - // Activity tracking private lastActivityAt = Date.now(); private _lastError: LastError | null = null; - - // Last root-session assistant message ID (tracked from message.updated kilocode events) private _lastAssistantMessageId: string | null = null; - private _observedGateResult: 'pass' | 'fail' | null = null; - - // Config of the most recently completed message, captured before active message advances - private _completedMessageConfig: { - autoCommit: boolean; - condenseOnComplete: boolean; - model?: string; - upstreamBranch?: string; - commitCoAuthor?: WrapperCommitCoAuthor; - } | null = null; - - // Callbacks for sending events to ingest private _sendToIngestFn: ((event: IngestEvent) => void) | null = null; - - // Log uploader (set per-job, cleared on job end) private _logUploader: LogUploader | null = null; - // --------------------------------------------------------------------------- - // State Queries - // --------------------------------------------------------------------------- - get isIdle(): boolean { - return !this._isActive; + return !this.isActive; } get isActive(): boolean { - return this._isActive; + return ( + this.hasPendingMessages || this._deliveryAcknowledgementsInFlight > 0 || this._isFinalizing + ); } - // --------------------------------------------------------------------------- - // Active State Management - // --------------------------------------------------------------------------- + get isFinalizing(): boolean { + return this._isFinalizing; + } - /** - * Set whether the wrapper is actively processing a prompt. - * Replaces addInflight/removeInflight — only one prompt is active at a time. - */ - setActive(active: boolean): void { - this._isActive = active; - if (active) { - this.updateActivity(); + get admissionsBlocked(): boolean { + return this._admissionsBlocked; + } + + beginFinalizing(): boolean { + if ( + this._isFinalizing || + !this.hasPendingMessages || + this._deliveryAcknowledgementsInFlight > 0 + ) { + return false; } + this._isFinalizing = true; + this.blockAdmissions(); + return true; } - // --------------------------------------------------------------------------- - // Connection Management - // --------------------------------------------------------------------------- + blockAdmissions(): void { + this._admissionsBlocked = true; + this._blockedWrapperRunId = this.session?.wrapperRunId; + } + + get finalizingWrapperRunId(): string | undefined { + return this._blockedWrapperRunId ?? this.session?.wrapperRunId; + } get isConnected(): boolean { return this._ingestWs !== null && this._ingestWs.readyState === WebSocket.OPEN; @@ -146,46 +106,24 @@ export class WrapperState { return this._sseAbortController; } - /** - * Store connection references. Actual connection management is in connection.ts. - */ setConnections(ws: WebSocket, sseAbortController: AbortController): void { this._ingestWs = ws; this._sseAbortController = sseAbortController; } - /** - * Clear connection references. Does NOT close or abort — connection.ts - * exclusively owns close semantics and calls this after its own cleanup. - */ clearConnectionRefs(): void { this._sseAbortController = null; this._ingestWs = null; } - /** - * Set the function used to send events to ingest. - * This is set by connection.ts when connection is established. - */ setSendToIngestFn(fn: ((event: IngestEvent) => void) | null): void { this._sendToIngestFn = fn; } - /** - * Send an event to ingest WebSocket. - * Silently drops the event if not connected (events are buffered in ConnectionManager). - */ sendToIngest(event: IngestEvent): void { - if (!this._sendToIngestFn) { - return; - } - this._sendToIngestFn(event); + this._sendToIngestFn?.(event); } - // --------------------------------------------------------------------------- - // Log Uploader - // --------------------------------------------------------------------------- - get logUploader(): LogUploader | null { return this._logUploader; } @@ -195,65 +133,30 @@ export class WrapperState { this._logUploader = uploader; } - // --------------------------------------------------------------------------- - // Activity Tracking - // --------------------------------------------------------------------------- - - /** - * Update last activity timestamp. Called on any meaningful action. - */ updateActivity(): void { this.lastActivityAt = Date.now(); } - /** - * Get milliseconds since last activity. - */ getIdleMs(now: number): number { return now - this.lastActivityAt; } - // --------------------------------------------------------------------------- - // Error Tracking - // --------------------------------------------------------------------------- - - /** - * Set the last error. This is cached for Worker to poll via /job/status. - */ setLastError(error: LastError): void { this._lastError = error; } - /** - * Get the last error. - */ getLastError(): LastError | null { return this._lastError; } - /** - * Clear the last error. - */ clearLastError(): void { this._lastError = null; } - // --------------------------------------------------------------------------- - // Assistant Message ID Tracking - // --------------------------------------------------------------------------- - - /** - * Get the last root-session assistant message ID. - * Tracked from message.updated kilocode events for autocommit association. - */ get lastAssistantMessageId(): string | null { return this._lastAssistantMessageId; } - /** - * Update the last assistant message ID. - * Called by connection.ts when a message.updated event with role=assistant is seen. - */ setLastAssistantMessageId(messageId: string): void { this._lastAssistantMessageId = messageId; } @@ -272,23 +175,15 @@ export class WrapperState { return gateResult; } - // --------------------------------------------------------------------------- - // Status for API Responses - // --------------------------------------------------------------------------- - getStatus(): WrapperStatus { return { - state: this.isActive ? 'active' : 'idle', + state: this._isFinalizing ? 'finalizing' : this.isActive ? 'active' : 'idle', sessionId: this.session?.kiloSessionId, pendingMessages: this.pendingMessageIds, lastError: this._lastError ?? undefined, }; } - // --------------------------------------------------------------------------- - // Session Context - // --------------------------------------------------------------------------- - get hasSession(): boolean { return this.session !== null; } @@ -300,6 +195,8 @@ export class WrapperState { bindSession(context: SessionContext): { changed: boolean } { if (!this.session) { this.session = context; + this._admissionsBlocked = false; + this._blockedWrapperRunId = undefined; this._lastError = null; this.updateActivity(); return { changed: true }; @@ -323,218 +220,65 @@ export class WrapperState { this._logUploader?.stop(); this._logUploader = null; this.session = null; - this.messages.clear(); - this._activeMessageId = null; - this._isActive = false; + this.clearAllMessages(); this._lastAssistantMessageId = null; - this._observedGateResult = null; - this._completedMessageConfig = null; } - // --------------------------------------------------------------------------- - // Message Tracking - // --------------------------------------------------------------------------- - - acceptMessage( - messageId: string, - config: { - autoCommit: boolean; - condenseOnComplete: boolean; - model?: string; - upstreamBranch?: string; - commitCoAuthor?: WrapperCommitCoAuthor; - } - ): void { - const info: MessageInfo = { - messageId, - state: !this._activeMessageId ? 'active' : 'accepted', - ...config, - }; - this.messages.set(messageId, info); - if (!this._activeMessageId) { - this._activeMessageId = messageId; - this._isActive = true; - } + beginDeliveryAcknowledgement(): boolean { + if (this._admissionsBlocked || this._isFinalizing) return false; + this._deliveryAcknowledgementsInFlight++; this.updateActivity(); + return true; } - completeActiveMessage(): MessageInfo | null { - if (!this._activeMessageId) return null; - const info = this.messages.get(this._activeMessageId); - if (info) info.state = 'completed'; - const completedInfo = info ?? null; - - let nextActive: string | null = null; - for (const [id, msg] of this.messages) { - if (msg.state === 'accepted') { - nextActive = id; - break; - } - } - - this._activeMessageId = nextActive; - if (nextActive) { - const nextInfo = this.messages.get(nextActive); - if (nextInfo) nextInfo.state = 'active'; - } else { - this._isActive = false; + endDeliveryAcknowledgement(): void { + if (this._deliveryAcknowledgementsInFlight > 0) { + this._deliveryAcknowledgementsInFlight--; } - - return completedInfo; } - completeMessage(messageId: string): MessageInfo | null { - if (messageId !== this._activeMessageId) return null; - const info = this.messages.get(messageId); - if (!info) return null; - - info.state = 'completed'; - - let nextActive: string | null = null; - for (const [id, msg] of this.messages) { - if (msg.state === 'accepted') { - nextActive = id; - break; - } - } - - this._activeMessageId = nextActive; - if (nextActive) { - const nextInfo = this.messages.get(nextActive); - if (nextInfo) nextInfo.state = 'active'; - } else { - this._isActive = false; - } - - return info; + get deliveryAcknowledgementsInFlight(): number { + return this._deliveryAcknowledgementsInFlight; } - get activeMessageId(): string | null { - return this._activeMessageId; + acceptMessage(messageId: string, config: FinalizationConfig): boolean { + if (this.admittedMessages.has(messageId)) return false; + this.admittedMessages.set(messageId, config); + this.latestAdmittedFinalizationConfig = config; + this.updateActivity(); + return true; } get hasPendingMessages(): boolean { - for (const msg of this.messages.values()) { - if (msg.state !== 'completed') return true; - } - return false; + return this.admittedMessages.size > 0; } get pendingMessageIds(): string[] { - const ids: string[] = []; - for (const msg of this.messages.values()) { - if (msg.state !== 'completed') ids.push(msg.messageId); - } - return ids; - } - - get activeMessageConfig(): { - autoCommit: boolean; - condenseOnComplete: boolean; - model?: string; - upstreamBranch?: string; - commitCoAuthor?: WrapperCommitCoAuthor; - } | null { - if (!this._activeMessageId) return null; - const info = this.messages.get(this._activeMessageId); - if (!info) return null; - return { - autoCommit: info.autoCommit, - condenseOnComplete: info.condenseOnComplete, - model: info.model, - upstreamBranch: info.upstreamBranch, - ...(info.commitCoAuthor ? { commitCoAuthor: info.commitCoAuthor } : {}), - }; - } - - get completedMessageConfig(): { - autoCommit: boolean; - condenseOnComplete: boolean; - model?: string; - upstreamBranch?: string; - commitCoAuthor?: WrapperCommitCoAuthor; - } | null { - return this._completedMessageConfig; + return [...this.admittedMessages.keys()]; } - setCompletedMessageConfig(config: { - autoCommit: boolean; - condenseOnComplete: boolean; - model?: string; - upstreamBranch?: string; - commitCoAuthor?: WrapperCommitCoAuthor; - }): void { - this._completedMessageConfig = config; + get batchFinalizationConfig(): FinalizationConfig | null { + return this.latestAdmittedFinalizationConfig; } - clearCompletedMessageConfig(): void { - this._completedMessageConfig = null; - } - - getMessageConfig(messageId: string): { - autoCommit: boolean; - condenseOnComplete: boolean; - model?: string; - upstreamBranch?: string; - commitCoAuthor?: WrapperCommitCoAuthor; - } | null { - const info = this.messages.get(messageId); - if (!info) return null; - return { - autoCommit: info.autoCommit, - condenseOnComplete: info.condenseOnComplete, - model: info.model, - upstreamBranch: info.upstreamBranch, - ...(info.commitCoAuthor ? { commitCoAuthor: info.commitCoAuthor } : {}), - }; - } - - updateMessageConfig( - messageId: string, - config: { - autoCommit?: boolean; - condenseOnComplete?: boolean; - model?: string; - upstreamBranch?: string; - commitCoAuthor?: WrapperCommitCoAuthor; - } - ): void { - const info = this.messages.get(messageId); - if (!info) return; - if (config.autoCommit !== undefined) info.autoCommit = config.autoCommit; - if (config.condenseOnComplete !== undefined) - info.condenseOnComplete = config.condenseOnComplete; - if (config.model !== undefined) info.model = config.model; - if (config.upstreamBranch !== undefined) info.upstreamBranch = config.upstreamBranch; - if (config.commitCoAuthor !== undefined) info.commitCoAuthor = config.commitCoAuthor; + getMessageConfig(messageId: string): FinalizationConfig | null { + return this.admittedMessages.get(messageId) ?? null; } removeMessage(messageId: string): void { - this.messages.delete(messageId); - if (this._activeMessageId === messageId) { - let nextActive: string | null = null; - for (const [id, msg] of this.messages) { - if (msg.state === 'accepted') { - nextActive = id; - break; - } - } - if (nextActive) { - this._activeMessageId = nextActive; - const nextInfo = this.messages.get(nextActive); - if (nextInfo) nextInfo.state = 'active'; - } else { - this._activeMessageId = null; - this._isActive = false; - } + const removedLatest = + this.latestAdmittedFinalizationConfig === this.admittedMessages.get(messageId); + this.admittedMessages.delete(messageId); + if (removedLatest) { + this.latestAdmittedFinalizationConfig = [...this.admittedMessages.values()].at(-1) ?? null; } } clearAllMessages(): void { - this.messages.clear(); - this._activeMessageId = null; - this._isActive = false; + this.admittedMessages.clear(); + this.latestAdmittedFinalizationConfig = null; + this._deliveryAcknowledgementsInFlight = 0; + this._isFinalizing = false; this._observedGateResult = null; - this._completedMessageConfig = null; } }