diff --git a/services/cloud-agent-next/wrapper/src/control/control-event-transport.test.ts b/services/cloud-agent-next/wrapper/src/control/control-event-transport.test.ts index a93d75e23d..ddaa4c6b43 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-event-transport.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-event-transport.test.ts @@ -121,9 +121,13 @@ describe('native-scoped control event failures', () => { const original = { runtimeId: crypto.randomUUID() }; let current = original; const retired = mock(); + const cleanup = Promise.withResolvers(); const handleFailure = createControlEventFailureHandler({ getRuntime: () => current, - onFailure: retired, + onFailure: (...args) => { + retired(...args); + return cleanup.promise; + }, }); const failures: ControlEventOutboxFailure[] = []; const published: ControlEventPublication[] = []; @@ -153,6 +157,10 @@ describe('native-scoped control event failures', () => { expect(failures).toHaveLength(2); expect(retired).toHaveBeenCalledTimes(1); expect(retired).toHaveBeenCalledWith(failures[0], original); + cleanup.resolve(); + await Promise.resolve(); + handleFailure(failures[0]); + expect(retired).toHaveBeenCalledTimes(2); current = { runtimeId: crypto.randomUUID() }; expect( await transport.publishSessionEvent(payload, { @@ -162,7 +170,7 @@ describe('native-scoped control event failures', () => { ).toBe(true); expect(await transport.resume()).toBe(true); expect(published.at(-1)?.session.nativeRuntimeId).toBe(current.runtimeId); - expect(retired).toHaveBeenCalledTimes(1); + expect(retired).toHaveBeenCalledTimes(2); const failure = failures[0]; if (!failure) throw new Error('Missing native failure'); handleFailure(failure); @@ -173,13 +181,50 @@ describe('native-scoped control event failures', () => { session: { ...session, nativeRuntimeId: current.runtimeId }, }, }); - expect(retired).toHaveBeenCalledTimes(2); - expect(retired.mock.calls[1]?.[1]).toBe(current); + expect(retired).toHaveBeenCalledTimes(3); + expect(retired.mock.calls[2]?.[1]).toBe(current); } finally { transport.close(); } }); + it('coalesces distinct roots only during registry cleanup and separates a reused native object incarnation', async () => { + const runtime = { runtimeId: 'N1' }; + const cleanup = Promise.withResolvers(); + const failures: ControlEventOutboxFailure[] = []; + const handleFailure = createControlEventFailureHandler({ + getRuntime: () => runtime, + onFailure: failure => { + failures.push(failure); + return cleanup.promise; + }, + }); + const failure = (root: string, nativeRuntimeId: string): ControlEventOutboxFailure => ({ + reason: 'rejected', + publication: { + event: 'session.event', + receiptId: `${root}-${nativeRuntimeId}`, + sequence: failures.length + 1, + session: { ...session, kiloSessionId: root, rootKiloSessionId: root, nativeRuntimeId }, + payload, + }, + }); + + handleFailure(failure('root_a', 'N1')); + handleFailure(failure('root_a', 'N1')); + handleFailure(failure('root_b', 'N1')); + expect(failures).toHaveLength(2); + cleanup.resolve(); + await Promise.resolve(); + handleFailure(failure('root_a', 'N1')); + expect(failures).toHaveLength(3); + + runtime.runtimeId = 'N2'; + handleFailure(failure('root_a', 'N1')); + handleFailure(failure('root_a', 'N2')); + expect(failures).toHaveLength(4); + }); + it('preserves a sealed result and its acknowledgement when failure retires the matching native runtime', async () => { const nativeLifetime = new AbortController(); const runtime = { diff --git a/services/cloud-agent-next/wrapper/src/control/control-event-transport.ts b/services/cloud-agent-next/wrapper/src/control/control-event-transport.ts index 7ebbe3f199..f60a63997a 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-event-transport.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-event-transport.ts @@ -9,18 +9,41 @@ type EventKind = 'session.event' | 'session.preparing'; export function createControlEventFailureHandler(options: { getRuntime: (directory: string) => Runtime | undefined; - onFailure: (failure: ControlEventOutboxFailure, runtime: Runtime) => void; + onFailure: (failure: ControlEventOutboxFailure, runtime: Runtime) => unknown; }) { - const failedRuntimes = new WeakSet(); + const inFlight = new WeakMap>(); return (failure?: ControlEventOutboxFailure): void => { if (!failure) return; if (failure.publication.event === 'session.preparing') return; const { directory, nativeRuntimeId } = failure.publication.session; - if (!nativeRuntimeId) return; + const root = + failure.publication.session.rootKiloSessionId ?? failure.publication.session.kiloSessionId; + if (!nativeRuntimeId || !root) return; const runtime = options.getRuntime(directory); - if (runtime?.runtimeId !== nativeRuntimeId || failedRuntimes.has(runtime)) return; - failedRuntimes.add(runtime); - options.onFailure(failure, runtime); + if (runtime?.runtimeId !== nativeRuntimeId) return; + const key = JSON.stringify([nativeRuntimeId, root]); + const keys = inFlight.get(runtime) ?? new Set(); + if (keys.has(key)) return; + keys.add(key); + inFlight.set(runtime, keys); + let result: unknown; + try { + result = options.onFailure(failure, runtime); + } catch { + keys.delete(key); + if (keys.size === 0) inFlight.delete(runtime); + return; + } + void Promise.resolve(result).then( + () => { + keys.delete(key); + if (keys.size === 0) inFlight.delete(runtime); + }, + () => { + keys.delete(key); + if (keys.size === 0) inFlight.delete(runtime); + } + ); }; } diff --git a/services/cloud-agent-next/wrapper/src/control/main.ts b/services/cloud-agent-next/wrapper/src/control/main.ts index 7be75bdd85..fe4c9992a1 100644 --- a/services/cloud-agent-next/wrapper/src/control/main.ts +++ b/services/cloud-agent-next/wrapper/src/control/main.ts @@ -2,9 +2,11 @@ import { heartbeatReasonFrom, sessionAttachResultSchema, type SandboxHeartbeatPayload, + type SessionEventIdentity, } from '../../../src/shared/sandbox-control-protocol.js'; import { WRAPPER_VERSION } from '../../../src/shared/wrapper-version.js'; import { logToFile } from '../utils.js'; +import { rootForSession } from './session-directories'; import { KILO_CONTROL_REQUEST_TIMEOUT_MS, maybeStartSandboxControlClient, @@ -19,7 +21,13 @@ import { } from './sandbox-control-handlers'; import { eventKiloSessionId, sessionEventIdentity, updateSessionSnapshots } from './feed'; import { createControlTerminalRuntime } from './terminal-runtime'; -import { createWorktreeKiloRuntimes } from './worktree-runtime'; +import { + createWorktreeKiloRuntimes, + type RootRuntimeDisappearance, + type RootRuntimeRetirement, + type WorktreeKiloRuntime, +} from './worktree-runtime'; +import type { NativeRetirement, RootScopedCleanupResult } from './session-operation-cleanup'; import { createControlDiagnostics, type ControlDiagnostics } from './diagnostics'; import { createControlFileLogUploader, type ControlFileLogUploader } from './file-log-uploader'; import { @@ -29,6 +37,13 @@ import { } from '../../../src/shared/control-diagnostics.js'; import { createWorktreeMutationNotifications } from './worktree-mutation-notifications'; import { createControlEventFailureHandler } from './control-event-transport'; + +type PublicationRetirementResult = RootScopedCleanupResult | NativeRetirement | 'shared'; + +type PublicationFailureAttempt = { + cleanup: Promise; + physical: Promise; +}; import type { ControlEventOutboxFailure } from './control-event-outbox'; function main( @@ -49,8 +64,16 @@ function main( let control: ReturnType = null; let shuttingDown = false; let heartbeatReason: SandboxHeartbeatPayload['kilo']['reason']; + const settleRootRetirement = (retirement: RootRuntimeRetirement): void => { + deps.operations.settleRootPublication(retirement); + }; + const notifyRootDisappeared = (disappearance: RootRuntimeDisappearance): void => { + deps.operations.notifyRootDisappeared(disappearance); + }; const kiloRuntimes = createWorktreeKiloRuntimes({ onDiagnostic: diagnostics.onDiagnostic, + onRootDisappeared: notifyRootDisappeared, + onRootRetirement: settleRootRetirement, onEvent: async (runtime, event) => { mutationNotifications.observe(runtime, event); const identity = sessionEventIdentity({ @@ -66,20 +89,16 @@ function main( identity.rootKiloSessionId, event.properties ); - if ( - !control?.publishSessionEvent || - !(await control.publishSessionEvent( - { type: event.type, properties: event.properties }, - identity - )) - ) { + const published = + control?.publishSessionEvent === undefined + ? false + : await control.publishSessionEvent( + { type: event.type, properties: event.properties }, + identity + ); + if (!published) { try { - await deps.operations.retireDirectory( - runtime.directory, - 'Session event delivery failed', - Date.now() + KILO_CONTROL_REQUEST_TIMEOUT_MS, - { runtimeId: runtime.runtimeId, client: runtime.kiloClient } - ); + await retirePublicationFailure(runtime, identity, 'Session event delivery failed'); } catch { diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'failed' }); } @@ -138,22 +157,69 @@ function main( throw new Error('Sandbox control operation result delivery unavailable'); return control.sendOperationResult(session, delivery, signal, deadlineAt); }, - emitSessionEvent: (session, payload, options) => - control?.sendEvent?.( - 'session.event', - payload, - { - directory: session.directory, - kiloSessionId: session.kiloSessionId, - rootKiloSessionId: session.kiloSessionId, - ...(options?.nativeRuntimeId ? { nativeRuntimeId: options.nativeRuntimeId } : {}), - }, - options?.retained ? { preserveConnectionOnFailure: true } : undefined - ) === true, + emitSessionEvent: (session, payload, options) => { + const identity = { + directory: session.directory, + kiloSessionId: session.kiloSessionId, + rootKiloSessionId: + rootForSession(session.kiloSessionId, session.directory) ?? session.kiloSessionId, + ...(options?.nativeRuntimeId ? { nativeRuntimeId: options.nativeRuntimeId } : {}), + }; + const delivered = + control?.sendEvent?.( + 'session.event', + payload, + identity, + options?.retained ? { preserveConnectionOnFailure: true } : undefined + ) === true; + if (!delivered) startPublicationFailure(identity, 'Session event delivery failed'); + return delivered; + }, retireRuntime: reason => shutdown(1, reason, heartbeatReasonFrom(reason)), onShutdown: () => shutdown(0, 'Sandbox shutting down'), }); + function beginPublicationFailure( + runtime: WorktreeKiloRuntime, + identity: SessionEventIdentity, + reason: string + ): PublicationFailureAttempt | undefined { + const root = identity.rootKiloSessionId ?? identity.kiloSessionId; + if (!root || (identity.nativeRuntimeId && identity.nativeRuntimeId !== runtime.runtimeId)) + return undefined; + const nativeRuntimeId = identity.nativeRuntimeId ?? runtime.runtimeId; + const target = { runtimeId: runtime.runtimeId, client: runtime.kiloClient }; + const deadlineAt = Date.now() + KILO_CONTROL_REQUEST_TIMEOUT_MS; + return deps.operations.escalateRootPublication({ + directory: identity.directory, + root, + nativeRuntimeId, + target, + reason, + deadlineAt, + }); + } + + async function retirePublicationFailure( + runtime: WorktreeKiloRuntime, + identity: SessionEventIdentity, + reason: string + ): Promise { + const attempt = beginPublicationFailure(runtime, identity, reason); + if (!attempt) return 'stale'; + return attempt.physical; + } + + function startPublicationFailure(identity: SessionEventIdentity, reason: string): void { + const runtime = kiloRuntimes.get(identity.directory); + if (!runtime) return; + const attempt = beginPublicationFailure(runtime, identity, reason); + if (!attempt) return; + void attempt.physical.catch(() => { + diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'failed' }); + }); + } + const mutationNotifications = createWorktreeMutationNotifications({ sessions: deps.sessions, kiloRuntimes, @@ -338,20 +404,39 @@ function main( getRuntime: directory => kiloRuntimes.get(directory), onFailure: (failure, runtime) => { reportOutboxRetirement(failure, runtime.runtimeId, 'started'); - void deps.operations - .retireDirectory( - failure.publication.session.directory, - `Session event delivery ${failure.reason}`, - Date.now() + KILO_CONTROL_REQUEST_TIMEOUT_MS, - { runtimeId: runtime.runtimeId, client: runtime.kiloClient } - ) - .then(() => { - reportOutboxRetirement(failure, runtime.runtimeId, 'retired', true); - }) - .catch(() => { + const attempt = beginPublicationFailure( + runtime, + failure.publication.session, + `Session event delivery ${failure.reason}` + ); + if (!attempt) return; + return attempt.cleanup.then( + cleanup => { + if (cleanup === 'confirmed') { + reportOutboxRetirement(failure, runtime.runtimeId, 'retired', true); + return; + } + void attempt.physical.then( + result => { + if (result === 'retired' || result === 'stale') + reportOutboxRetirement(failure, runtime.runtimeId, 'retired', true); + else { + reportOutboxRetirement(failure, runtime.runtimeId, 'failed', false); + diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'failed' }); + } + }, + () => { + reportOutboxRetirement(failure, runtime.runtimeId, 'failed', false); + diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'failed' }); + } + ); + }, + () => { + void attempt.physical.catch(() => undefined); reportOutboxRetirement(failure, runtime.runtimeId, 'failed', false); diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'failed' }); - }); + } + ); }, }), onDisconnected: () => shutdown(1, 'Sandbox control connection lost', 'control_disconnected'), diff --git a/services/cloud-agent-next/wrapper/src/control/operation-registry.test.ts b/services/cloud-agent-next/wrapper/src/control/operation-registry.test.ts index e92f9a8638..225087fcc1 100644 --- a/services/cloud-agent-next/wrapper/src/control/operation-registry.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/operation-registry.test.ts @@ -1,4 +1,4 @@ -import { afterEach, beforeEach, describe, expect, it, setSystemTime, spyOn } from 'bun:test'; +import { afterEach, beforeEach, describe, expect, it, mock, setSystemTime, spyOn } from 'bun:test'; import fs from 'node:fs'; import os from 'node:os'; import path from 'node:path'; @@ -28,7 +28,12 @@ import { type Completion, } from './control-test-fixtures'; import { operationIntent } from './operation-intent'; -import { rememberAttachedRoot, resetSessionDirectoryState } from './session-directories'; +import { + rememberAttachedRoot, + rememberChildSession, + resetSessionDirectoryState, +} from './session-directories'; +import { createOperationRegistry } from './operation-registry'; import { resetDirectoryOperationState } from './worktree-operations'; let homeRoot: string; @@ -58,6 +63,575 @@ function onlyOperation(handlerDeps: HandlerDeps) { } describe('operation admission and lookup', () => { + it('scopes a confirmed publication failure to A and keeps B plus fresh A work live', async () => { + const pendingA = Promise.withResolvers(); + const pendingB = Promise.withResolvers(); + let statusCalls = 0; + let abortCalls = 0; + const retired = mock(); + const client = fakeKilo({ + sendPrompt: async options => + options.messageId === 'message_a' ? pendingA.promise : pendingB.promise, + abortSession: async () => { + abortCalls += 1; + return true; + }, + getSessionStatuses: async () => { + statusCalls += 1; + return { kilo_a: { type: 'idle' } }; + }, + }); + const handlerDeps = deps({ kiloClient: client, retireRuntime: retired }); + const sessionA = { ...session, sessionId: 'ses_a', kiloSessionId: 'kilo_a' }; + const sessionB = { ...session, sessionId: 'ses_b', kiloSessionId: 'kilo_b' }; + rememberAttachedRoot(sessionA.kiloSessionId, sessionA.directory); + rememberAttachedRoot(sessionB.kiloSessionId, sessionB.directory); + const authorizationA = operationAuthorization('session.prompt', 'message_a', sessionA); + const requestA = handleControlRequest( + 'session.prompt', + sessionA, + { ...promptPayload, messageId: 'message_a' }, + handlerDeps, + authorizationA + ); + const requestB = handleControlRequest( + 'session.prompt', + sessionB, + { ...promptPayload, messageId: 'message_b' }, + handlerDeps + ); + try { + for (let index = 0; index < 10; index += 1) { + if ( + handlerDeps.operations.active(sessionA.kiloSessionId)?.snapshot().native.state === + 'pending' + ) + break; + await Promise.resolve(); + } + for (let index = 0; index < 10; index += 1) { + if ( + handlerDeps.operations.active(sessionB.kiloSessionId)?.snapshot().native.state === + 'pending' + ) + break; + await Promise.resolve(); + } + const runtime = handlerDeps.kiloRuntimes?.get(sessionA.directory); + if (!runtime) throw new Error('Missing native runtime'); + const publicationCleanup = await handlerDeps.operations.retireRootPublication({ + directory: sessionA.directory, + root: sessionA.kiloSessionId, + nativeRuntimeId: runtime.runtimeId, + target: { runtimeId: runtime.runtimeId, client: runtime.kiloClient }, + reason: 'event rejected', + deadlineAt: Date.now() + 1_000, + }); + expect(publicationCleanup).toBe('confirmed'); + expect(abortCalls).toBe(1); + expect(statusCalls).toBe(1); + expect(handlerDeps.operations.active(sessionB.kiloSessionId)?.signal.aborted).toBe(false); + expect(retired).not.toHaveBeenCalled(); + + pendingA.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + await requestA; + for ( + let index = 0; + index < 10 && handlerDeps.operations.active(sessionA.kiloSessionId); + index += 1 + ) + await Promise.resolve(); + expect(handlerDeps.operations.active(sessionB.kiloSessionId)?.signal.aborted).toBe(false); + + pendingB.resolve(completion()); + expect((await requestB).ok).toBe(true); + expect(await handleControlRequest('session.detach', sessionB, {}, handlerDeps)).toMatchObject( + { + ok: true, + result: { detached: true }, + } + ); + const fresh = await handleControlRequest( + 'session.prompt', + sessionA, + { ...promptPayload, messageId: 'message_fresh' }, + handlerDeps, + operationAuthorization('session.prompt', 'message_fresh', sessionA) + ); + expect(fresh.ok).toBe(true); + } finally { + pendingA.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + pendingB.resolve(completion()); + } + }); + + it('retains an unconfirmed root/incarnation failure and gates only fresh A work', async () => { + const handlerDeps = deps({ + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + const sessionA = { ...session, sessionId: 'ses_a', kiloSessionId: 'kilo_a' }; + const sessionB = { ...session, sessionId: 'ses_b', kiloSessionId: 'kilo_b' }; + rememberAttachedRoot(sessionA.kiloSessionId, sessionA.directory); + rememberAttachedRoot(sessionB.kiloSessionId, sessionB.directory); + const runtime = handlerDeps.kiloRuntimes?.get(sessionA.directory); + if (!runtime) throw new Error('Missing native runtime'); + const input = { + directory: sessionA.directory, + root: sessionA.kiloSessionId, + nativeRuntimeId: runtime.runtimeId, + target: { runtimeId: runtime.runtimeId, client: runtime.kiloClient }, + reason: 'event rejected', + deadlineAt: Date.now() + 1_000, + }; + expect(await handlerDeps.operations.retireRootPublication(input)).toBe('unconfirmed'); + expect(await handlerDeps.operations.retireRootPublication(input)).toBe('unconfirmed'); + + const authorizationA = operationAuthorization('session.prompt', 'message_a', sessionA); + expect( + await handleControlRequest( + 'session.prompt', + sessionA, + { ...promptPayload, messageId: 'message_a' }, + handlerDeps, + authorizationA + ) + ).toMatchObject({ ok: false, error: { code: 'not_ready', retryable: true } }); + expect( + await handleControlRequest( + 'session.prompt', + sessionA, + { ...promptPayload, messageId: 'message_b' }, + handlerDeps + ) + ).toMatchObject({ ok: false, error: { code: 'not_ready', retryable: true } }); + expect( + await handleControlRequest('session.operation.get', sessionA, authorizationA, handlerDeps) + ).toMatchObject({ ok: true, result: { state: 'missing' } }); + expect( + await handleControlRequest('session.operation.ack', sessionA, {}, handlerDeps) + ).toMatchObject({ ok: false, error: { code: 'unauthorized' } }); + + const authorizationB = operationAuthorization('session.prompt', 'message_b', sessionB); + expect( + await handleControlRequest( + 'session.prompt', + sessionB, + { ...promptPayload, messageId: 'message_b' }, + handlerDeps, + authorizationB + ) + ).toMatchObject({ ok: true, result: { status: 'accepted' } }); + const bRecord = handlerDeps.operations + .retained() + .find(record => record.messageId === 'message_b'); + if (!bRecord) throw new Error('Missing B operation record'); + await bRecord.done; + await bRecord.waitForDelivery(); + expect( + await handleControlRequest('session.operation.get', sessionB, authorizationB, handlerDeps) + ).toMatchObject({ ok: true, result: { state: 'completed' } }); + const bDelivery = bRecord.deliveryResult(); + if (!bDelivery) throw new Error('Missing B delivery'); + expect( + await handleControlRequest( + 'session.operation.ack', + sessionB, + await acknowledgeOperation(bDelivery), + handlerDeps + ) + ).toMatchObject({ ok: true, result: { acknowledged: true } }); + }); + + it('retains terminal unconfirmed state across routing loss until explicit root notification', async () => { + const handlerDeps = deps(); + const sessionA = { ...session, sessionId: 'ses_a', kiloSessionId: 'kilo_a' }; + rememberAttachedRoot(sessionA.kiloSessionId, sessionA.directory); + const runtime = handlerDeps.kiloRuntimes?.get(sessionA.directory); + if (!runtime) throw new Error('Missing native runtime'); + const input = { + directory: sessionA.directory, + root: sessionA.kiloSessionId, + nativeRuntimeId: runtime.runtimeId, + target: { runtimeId: runtime.runtimeId, client: runtime.kiloClient }, + reason: 'event rejected', + deadlineAt: Date.now() + 1_000, + }; + expect(await handlerDeps.operations.retireRootPublication(input)).toBe('unconfirmed'); + handlerDeps.operations.settleRootPublication({ ...input, result: 'unconfirmed' }); + + expect(await handleControlRequest('session.detach', sessionA, {}, handlerDeps)).toMatchObject({ + ok: true, + result: { detached: true }, + }); + handlerDeps.operations.prune(); + rememberAttachedRoot(sessionA.kiloSessionId, sessionA.directory); + expect(await handlerDeps.operations.retireRootPublication(input)).toBe('unconfirmed'); + expect(handlerDeps.operations.admission('session.prompt', sessionA, undefined).kind).toBe( + 'reply' + ); + + handlerDeps.operations.notifyRootDisappeared({ + directory: sessionA.directory, + root: sessionA.kiloSessionId, + nativeRuntimeId: runtime.runtimeId, + }); + expect(handlerDeps.operations.admission('session.prompt', sessionA, undefined).kind).toBe( + 'continue' + ); + }); + + it('retains an unconfirmed record after physical unregistration and later no-op failures', async () => { + const fixture = deps(); + const runtime = fixture.kiloRuntimes?.get(session.directory); + if (!runtime) throw new Error('Missing native runtime'); + const nativeRetirement = mock(async () => 'unconfirmed' as const); + const operations = createOperationRegistry({ + native: { + get: () => runtime, + getRetained: () => runtime, + retireRuntime: nativeRetirement, + verifyQuiescence: async () => false, + }, + onStarted: () => {}, + onCompleted: () => {}, + retireRuntime: () => {}, + }); + const sessionA = { ...session, kiloSessionId: 'kilo_a', sessionId: 'ses_a' }; + rememberAttachedRoot(sessionA.kiloSessionId, sessionA.directory); + const input = { + directory: sessionA.directory, + root: sessionA.kiloSessionId, + nativeRuntimeId: runtime.runtimeId, + target: { runtimeId: runtime.runtimeId, client: runtime.kiloClient }, + reason: 'event rejected', + deadlineAt: Date.now() + 1_000, + }; + expect(await operations.retireRootPublication(input)).toBe('unconfirmed'); + expect( + await operations.retireDirectory( + sessionA.directory, + 'physical', + input.deadlineAt, + input.target + ) + ).toBe('unconfirmed'); + expect(nativeRetirement).toHaveBeenCalledTimes(1); + + expect(await handleControlRequest('session.detach', sessionA, {}, fixture)).toMatchObject({ + ok: true, + result: { detached: true }, + }); + operations.prune(); + rememberAttachedRoot(sessionA.kiloSessionId, sessionA.directory); + expect(await operations.retireRootPublication(input)).toBe('unconfirmed'); + expect(operations.admission('session.prompt', sessionA, undefined).kind).toBe('reply'); + + operations.notifyRootDisappeared({ + directory: sessionA.directory, + root: sessionA.kiloSessionId, + nativeRuntimeId: runtime.runtimeId, + }); + expect(operations.admission('session.prompt', sessionA, undefined).kind).toBe('continue'); + }); + + it('clears a scoped record after successful targeted retirement without a deferred intent', async () => { + const fixture = deps(); + const runtime = fixture.kiloRuntimes?.get(session.directory); + if (!runtime) throw new Error('Missing native runtime'); + const nativeRetirement = mock(async () => 'retired' as const); + const operations = createOperationRegistry({ + native: { + get: () => runtime, + getRetained: () => runtime, + retireRuntime: nativeRetirement, + verifyQuiescence: async () => true, + }, + onStarted: () => {}, + onCompleted: () => {}, + retireRuntime: () => {}, + }); + const sessionA = { ...session, kiloSessionId: 'kilo_a', sessionId: 'ses_a' }; + rememberAttachedRoot(sessionA.kiloSessionId, sessionA.directory); + const target = { runtimeId: runtime.runtimeId, client: runtime.kiloClient }; + const input = { + directory: sessionA.directory, + root: sessionA.kiloSessionId, + nativeRuntimeId: runtime.runtimeId, + target, + reason: 'targeted retirement', + deadlineAt: Date.now() + 1_000, + }; + expect(await operations.retireRootPublication(input)).toBe('unconfirmed'); + expect( + await operations.retireDirectory(sessionA.directory, 'targeted', input.deadlineAt, target) + ).toBe('retired'); + + expect(operations.admission('session.prompt', sessionA, undefined).kind).toBe('continue'); + }); + + it('invalidates a pending escalation before an unchanged same-root reattach can join the runtime', async () => { + const handlerDeps = deps(); + const sessionA = { ...session, kiloSessionId: 'kilo_a', sessionId: 'ses_a' }; + const sessionB = { ...session, kiloSessionId: 'kilo_b', sessionId: 'ses_b' }; + rememberAttachedRoot(sessionA.kiloSessionId, sessionA.directory); + rememberAttachedRoot(sessionB.kiloSessionId, sessionB.directory); + const runtime = handlerDeps.kiloRuntimes?.get(sessionA.directory); + if (!runtime) throw new Error('Missing native runtime'); + let gateCalls = 0; + const runtimes = handlerDeps.kiloRuntimes; + if (!runtimes) throw new Error('Missing worktree runtimes'); + runtimes.retireRuntimeIfUnshared = async () => { + gateCalls += 1; + return 'shared'; + }; + const input = { + directory: sessionA.directory, + root: sessionA.kiloSessionId, + nativeRuntimeId: runtime.runtimeId, + target: { runtimeId: runtime.runtimeId, client: runtime.kiloClient }, + reason: 'event rejected', + deadlineAt: Date.now() + 1_000, + }; + const escalation = handlerDeps.operations.escalateRootPublication(input); + handlerDeps.operations.notifyRootDisappeared(input); + rememberAttachedRoot(sessionA.kiloSessionId, sessionA.directory); + expect(await escalation.physical).toBe('stale'); + expect(gateCalls).toBe(0); + expect(handlerDeps.operations.admission('session.prompt', sessionA, undefined).kind).toBe( + 'continue' + ); + + handlerDeps.operations.notifyRootDisappeared({ + directory: sessionB.directory, + root: sessionB.kiloSessionId, + nativeRuntimeId: runtime.runtimeId, + }); + expect(gateCalls).toBe(0); + }); + + it('does not let a retained A1 Stop reselect fresh active A2 after root invalidation', async () => { + const runningA1 = Promise.withResolvers(); + const runningA2 = Promise.withResolvers(); + const runningB = Promise.withResolvers(); + const startedA1 = Promise.withResolvers(); + const startedA2 = Promise.withResolvers(); + const startedB = Promise.withResolvers(); + let status: 'busy' | 'idle' = 'busy'; + let gateCalls = 0; + const client = fakeKilo({ + sendPrompt: async options => { + if (options.messageId === 'message_a1') { + startedA1.resolve(); + return runningA1.promise; + } + if (options.messageId === 'message_a2') { + startedA2.resolve(); + return runningA2.promise; + } + startedB.resolve(); + return runningB.promise; + }, + getSessionStatuses: async () => ({ kilo_a: { type: status } }), + abortSession: async () => true, + }); + const handlerDeps = deps({ + kiloClient: client, + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + const sessionA = { ...session, sessionId: 'ses_a', kiloSessionId: 'kilo_a' }; + const sessionB = { ...session, sessionId: 'ses_b', kiloSessionId: 'kilo_b' }; + rememberAttachedRoot(sessionA.kiloSessionId, sessionA.directory); + rememberAttachedRoot(sessionB.kiloSessionId, sessionB.directory); + const runtimes = handlerDeps.kiloRuntimes; + if (!runtimes) throw new Error('Missing worktree runtimes'); + runtimes.retireRuntimeIfUnshared = async () => { + gateCalls += 1; + return 'shared'; + }; + const authorizationA1 = operationAuthorization('session.prompt', 'message_a1', sessionA); + const requestA1 = handleControlRequest( + 'session.prompt', + sessionA, + { ...promptPayload, messageId: 'message_a1' }, + handlerDeps, + authorizationA1 + ); + try { + await startedA1.promise; + const target = handlerDeps.kiloRuntimes?.get(sessionA.directory); + if (!target) throw new Error('Missing native runtime'); + expect( + await handlerDeps.operations.retireRootPublication({ + directory: sessionA.directory, + root: sessionA.kiloSessionId, + nativeRuntimeId: target.runtimeId, + target: { runtimeId: target.runtimeId, client: target.kiloClient }, + reason: 'event rejected', + deadlineAt: Date.now() + 20, + }) + ).toBe('unconfirmed'); + runningA1.resolve( + completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } }) + ); + await requestA1; + const retainedA1 = handlerDeps.operations + .retained() + .find(record => record.messageId === 'message_a1'); + if (!retainedA1) throw new Error('Missing retained A1 operation'); + await retainedA1.done; + await retainedA1.waitForDelivery(); + + expect(await handleControlRequest('session.detach', sessionA, {}, handlerDeps)).toMatchObject( + { + ok: true, + result: { detached: true }, + } + ); + handlerDeps.operations.notifyRootDisappeared({ + directory: sessionA.directory, + root: sessionA.kiloSessionId, + nativeRuntimeId: target.runtimeId, + target: { runtimeId: target.runtimeId, client: target.kiloClient }, + }); + const attachAuthorization = operationAuthorization('session.attach', undefined, sessionA); + expect( + await handleControlRequest( + 'session.attach', + sessionA, + { kilo }, + handlerDeps, + attachAuthorization + ) + ).toMatchObject({ ok: true }); + + const requestB = handleControlRequest( + 'session.prompt', + sessionB, + { ...promptPayload, messageId: 'message_b' }, + handlerDeps + ); + const requestA2 = handleControlRequest( + 'session.prompt', + sessionA, + { ...promptPayload, messageId: 'message_a2' }, + handlerDeps + ); + await startedB.promise; + await startedA2.promise; + const a2 = handlerDeps.operations.active(sessionA.kiloSessionId); + const b = handlerDeps.operations.active(sessionB.kiloSessionId); + if (!a2 || !b) throw new Error('Missing fresh active operations'); + const stopped = await handleControlRequest( + 'session.abort', + sessionA, + { messageId: 'message_a1', operationId: '11111111-1111-4111-8111-111111111111' }, + handlerDeps + ); + expect(stopped).toMatchObject({ ok: true, result: { quiescent: false } }); + expect(a2.signal.aborted).toBe(false); + expect(b.signal.aborted).toBe(false); + expect(gateCalls).toBe(0); + + status = 'idle'; + runningB.resolve(completion({ name: 'MessageAbortedError', data: { message: 'detached' } })); + await requestB; + expect(await handleControlRequest('session.detach', sessionB, {}, handlerDeps)).toMatchObject( + { + ok: true, + result: { detached: true }, + } + ); + expect(a2.signal.aborted).toBe(false); + expect(gateCalls).toBe(0); + runningA2.resolve( + completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } }) + ); + await requestA2; + } finally { + status = 'idle'; + runningA1.resolve( + completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } }) + ); + runningA2.resolve( + completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } }) + ); + runningB.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + await Promise.allSettled([requestA1]); + for (const task of handlerDeps.operations.activeOperations()) await task.done; + } + }); + + it('coalesces authorized and unauthorized active operations under one root claim', async () => { + const pendingAuthorized = Promise.withResolvers(); + const pendingUnauthorized = Promise.withResolvers(); + const client = fakeKilo({ + sendPrompt: async options => + options.messageId === 'message_authorized' + ? pendingAuthorized.promise + : pendingUnauthorized.promise, + abortSession: async () => true, + getSessionStatuses: async () => ({ kilo_a: { type: 'idle' } }), + }); + const handlerDeps = deps({ kiloClient: client }); + const sessionA = { ...session, sessionId: 'ses_a', kiloSessionId: 'kilo_a' }; + const child = { ...session, sessionId: 'ses_child', kiloSessionId: 'child_a' }; + rememberAttachedRoot(sessionA.kiloSessionId, sessionA.directory); + rememberChildSession({ childId: child.kiloSessionId, parentId: sessionA.kiloSessionId }); + const runtime = handlerDeps.kiloRuntimes?.get(sessionA.directory); + if (!runtime) throw new Error('Missing native runtime'); + const authorized = handlerDeps.operations.start( + sessionA, + operationAuthorization('session.prompt', 'message_authorized', sessionA), + { + operation: 'session.prompt', + payload: { ...promptPayload, messageId: 'message_authorized' }, + runtime, + }, + { emitSessionEvent: () => true } + ); + const unauthorized = handlerDeps.operations.start( + child, + undefined, + { + operation: 'session.prompt', + payload: { ...promptPayload, messageId: 'message_unauthorized' }, + runtime, + }, + { emitSessionEvent: () => true } + ); + try { + for (let index = 0; index < 10; index += 1) { + if ( + authorized.snapshot().native.state === 'pending' && + unauthorized.snapshot().native.state === 'pending' + ) + break; + await Promise.resolve(); + } + const input = { + directory: sessionA.directory, + root: sessionA.kiloSessionId, + nativeRuntimeId: runtime.runtimeId, + target: { runtimeId: runtime.runtimeId, client: runtime.kiloClient }, + reason: 'coalesced publication failure', + deadlineAt: Date.now() + 1_000, + }; + const first = handlerDeps.operations.retireRootPublication(input); + const duplicate = handlerDeps.operations.retireRootPublication(input); + expect(duplicate).toBe(first); + expect(await first).toBe('confirmed'); + expect(authorized.publicationScope()?.claim).toBe(unauthorized.publicationScope()?.claim); + } finally { + pendingAuthorized.resolve( + completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } }) + ); + pendingUnauthorized.resolve( + completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } }) + ); + await Promise.all([authorized.done, unauthorized.done]); + } + }); + it('looks up the same operation during and after completion and rejects changed intent', async () => { const running = Promise.withResolvers(); const started = Promise.withResolvers(); diff --git a/services/cloud-agent-next/wrapper/src/control/operation-registry.ts b/services/cloud-agent-next/wrapper/src/control/operation-registry.ts index f0df4fe7ad..52e406ce7e 100644 --- a/services/cloud-agent-next/wrapper/src/control/operation-registry.ts +++ b/services/cloud-agent-next/wrapper/src/control/operation-registry.ts @@ -8,8 +8,13 @@ import { type SessionRequestIdentity, } from '../../../src/shared/sandbox-control-protocol.js'; import { rejectBeforeAdmission } from './control-handler-result.js'; +import { rootForSession } from './session-directories.js'; import type { WorktreeKiloRuntimes } from './worktree-runtime.js'; -import type { NativeOperationTarget, NativeRetirement } from './session-operation-cleanup.js'; +import type { + NativeOperationTarget, + NativeRetirement, + RootScopedCleanupResult, +} from './session-operation-cleanup.js'; import { SessionOperation, type ControlHandlerResult, @@ -31,6 +36,13 @@ type OperationRegistryDependencies = { deadlineAt: number, target?: NativeOperationTarget ): Promise; + retireRuntimeIfUnshared?( + directory: string, + target: NativeOperationTarget, + retiringRoot: string, + deadlineAt: number, + reason?: string + ): Promise; verifyQuiescence( directory: string, target: NativeOperationTarget, @@ -51,6 +63,38 @@ type Admission = | { kind: 'continue' } | { kind: 'reply'; result: ControlHandlerResult | Promise }; +type ScopedFailure = { + root: string; + nativeRuntimeId: string; + directory: string; + target?: NativeOperationTarget; + deadlineAt: number; + cleanup: Promise; + physical?: Promise; + result?: RootScopedCleanupResult; + claim: symbol; +}; + +type RootPublicationInput = { + directory: string; + root: string; + nativeRuntimeId: string; + target?: NativeOperationTarget; + reason: string; + deadlineAt: number; + expectedClaim?: symbol; +}; + +type RootPublicationSettlement = { + directory: string; + root: string; + nativeRuntimeId: string; + target?: NativeOperationTarget; + result: NativeRetirement; +}; + +type RootPublicationDisappearance = Omit; + function key(authorization: SessionOperationAuthorization): string { return JSON.stringify([ authorization.session.sessionId, @@ -67,9 +111,196 @@ function fail(code: string, message: string, retryable: boolean): ControlHandler return { ok: false, error: { code, message, retryable } }; } +const ROOT_SCOPED_WORK = new Set(['session.attach', 'session.prompt', 'session.terminal.create']); + export function createOperationRegistry(deps: OperationRegistryDependencies) { const active = new Map(); const retained = new Map(); + const scopedFailures = new Map(); + + function scopedFailureKey(root: string, nativeRuntimeId: string): string { + return JSON.stringify([root, nativeRuntimeId]); + } + + function currentRuntime(directory: string) { + return deps.native.getRetained(directory) ?? deps.native.get(directory); + } + + function compatibleTarget( + left: NativeOperationTarget | undefined, + right: NativeOperationTarget | undefined + ): boolean { + return ( + left === undefined || + right === undefined || + (left.runtimeId === right.runtimeId && + (left.client === undefined || right.client === undefined || left.client === right.client)) + ); + } + + function clearStaleScopedFailures(): void { + for (const [id, failure] of scopedFailures) { + const runtime = currentRuntime(failure.directory); + if (runtime && runtime.runtimeId !== failure.nativeRuntimeId) scopedFailures.delete(id); + } + } + + function matchingPublicationOperations(input: RootPublicationInput): SessionOperation[] { + return [...active.values()].filter(operation => { + if (operation.session.directory !== input.directory) return false; + if (rootForSession(operation.session.kiloSessionId, input.directory) !== input.root) + return false; + const operationTarget = operation.nativeTarget(); + return ( + operationTarget?.runtimeId === input.nativeRuntimeId && + compatibleTarget(operationTarget, input.target) + ); + }); + } + + function createScopedFailure( + input: RootPublicationInput, + matching: SessionOperation[] + ): Promise { + const claim = Symbol('publication-scoped-failure'); + for (const operation of matching) + operation.markPublicationScoped(input.reason, input.deadlineAt, claim); + const cleanup = (async (): Promise => { + if (matching.length === 0) return 'unconfirmed'; + const results = await Promise.all( + matching.map(operation => operation.runRootScopedCleanup()) + ); + return results.every(result => result === 'confirmed') ? 'confirmed' : 'unconfirmed'; + })().catch(() => 'unconfirmed' as const); + const failure: ScopedFailure = { + root: input.root, + nativeRuntimeId: input.nativeRuntimeId, + directory: input.directory, + cleanup, + deadlineAt: input.deadlineAt, + claim, + }; + failure.target = input.target; + failure.deadlineAt = input.deadlineAt; + failure.cleanup = cleanup; + scopedFailures.set(scopedFailureKey(input.root, input.nativeRuntimeId), failure); + void cleanup.then(result => { + if (scopedFailures.get(scopedFailureKey(input.root, input.nativeRuntimeId)) === failure) + failure.result = result; + }); + return cleanup; + } + + function retireRootPublication(input: RootPublicationInput): Promise { + const id = scopedFailureKey(input.root, input.nativeRuntimeId); + clearStaleScopedFailures(); + let existing = scopedFailures.get(id); + if (input.expectedClaim !== undefined) { + if ( + !existing || + existing.claim !== input.expectedClaim || + !compatibleTarget(existing.target, input.target) + ) + return Promise.resolve('unconfirmed'); + return existing.cleanup; + } + if (existing && !compatibleTarget(existing.target, input.target)) { + scopedFailures.delete(id); + existing = undefined; + } else if (existing && existing.result !== 'confirmed') { + return existing.cleanup; + } + const matching = matchingPublicationOperations(input); + return createScopedFailure(input, matching); + } + + function escalateRootPublication(input: RootPublicationInput): { + cleanup: Promise; + physical: Promise; + } { + const id = scopedFailureKey(input.root, input.nativeRuntimeId); + const current = scopedFailures.get(id); + if (input.expectedClaim !== undefined && current?.claim !== input.expectedClaim) + return { cleanup: Promise.resolve('unconfirmed'), physical: Promise.resolve('stale') }; + const cleanup = retireRootPublication(input); + const failure = scopedFailures.get(id); + if (!failure) return { cleanup, physical: Promise.resolve('unconfirmed') }; + if (failure.physical) return { cleanup, physical: failure.physical }; + failure.physical = cleanup.then(async cleanupResult => { + if (scopedFailures.get(id) !== failure) return 'stale'; + if (cleanupResult === 'confirmed') return cleanupResult; + const retirement = + (await deps.native.retireRuntimeIfUnshared?.( + input.directory, + input.target ?? { runtimeId: input.nativeRuntimeId }, + input.root, + input.deadlineAt, + input.reason + )) ?? 'unconfirmed'; + if (retirement === 'retired' || retirement === 'stale') { + if (scopedFailures.get(id) !== failure) return retirement; + settleRootPublication({ + directory: input.directory, + root: input.root, + nativeRuntimeId: input.nativeRuntimeId, + target: input.target, + result: retirement, + }); + } + return retirement; + }); + return { cleanup, physical: failure.physical }; + } + + function settleRootPublication(input: RootPublicationSettlement): void { + const id = scopedFailureKey(input.root, input.nativeRuntimeId); + const failure = scopedFailures.get(id); + if (!failure || failure.directory !== input.directory) return; + if (!compatibleTarget(failure.target, input.target)) return; + if (input.result === 'retired' || input.result === 'stale') scopedFailures.delete(id); + else failure.result = 'unconfirmed'; + } + + function settleScopedFailuresForIncarnation( + directory: string, + target: NativeOperationTarget, + result: NativeRetirement + ): void { + if (result !== 'retired' && result !== 'stale') return; + for (const [id, failure] of scopedFailures) { + if ( + failure.directory === directory && + failure.nativeRuntimeId === target.runtimeId && + compatibleTarget(failure.target, target) + ) + scopedFailures.delete(id); + } + } + + function notifyRootDisappeared(input: RootPublicationDisappearance): void { + const id = scopedFailureKey(input.root, input.nativeRuntimeId); + const failure = scopedFailures.get(id); + if (failure?.directory === input.directory && compatibleTarget(failure.target, input.target)) + scopedFailures.delete(id); + } + + function publicationFailureBlocks(operation: string, session: SessionRequestIdentity): boolean { + if (!ROOT_SCOPED_WORK.has(operation)) return false; + clearStaleScopedFailures(); + const root = rootForSession(session.kiloSessionId, session.directory); + if (!root) return false; + const runtime = currentRuntime(session.directory); + for (const failure of scopedFailures.values()) { + if ( + failure.directory === session.directory && + failure.root === root && + failure.result !== 'confirmed' && + (runtime === undefined || runtime.runtimeId === failure.nativeRuntimeId) + ) + return true; + } + return false; + } function prune(now = Date.now()): void { for (const [id, operation] of retained) { @@ -88,6 +319,7 @@ export function createOperationRegistry(deps: OperationRegistryDependencies) { }); } } + clearStaleScopedFailures(); } function admission( @@ -96,7 +328,15 @@ export function createOperationRegistry(deps: OperationRegistryDependencies) { payload: unknown, authorization?: SessionOperationAuthorization ): Admission { - if (operation !== 'session.operation.get' && !authorization) return { kind: 'continue' }; + clearStaleScopedFailures(); + if (operation !== 'session.operation.get' && !authorization) { + if (publicationFailureBlocks(operation, session)) + return { + kind: 'reply', + result: rejectBeforeAdmission('not_ready', 'Native runtime cleanup is unconfirmed', true), + }; + return { kind: 'continue' }; + } const reply = (result: ControlHandlerResult | Promise): Admission => ({ kind: 'reply', result, @@ -159,6 +399,10 @@ export function createOperationRegistry(deps: OperationRegistryDependencies) { ); } if (operation === 'session.operation.get') return reply(ok({ state: 'missing' })); + if (publicationFailureBlocks(operation, session)) + return reply( + rejectBeforeAdmission('not_ready', 'Native runtime cleanup is unconfirmed', true) + ); if (Date.now() >= target.dispatchDeadlineAt) return reply( rejectBeforeAdmission('not_ready', 'Operation dispatch authorization expired', false) @@ -206,6 +450,7 @@ export function createOperationRegistry(deps: OperationRegistryDependencies) { const retirement = await deps.native.retireRuntime(directory, deadlineAt, target); for (const operation of matching) operation.confirmCleanup(retirement === 'retired' || retirement === 'stale', deadlineAt); + if (target) settleScopedFailuresForIncarnation(directory, target, retirement); return retirement; } @@ -251,6 +496,10 @@ export function createOperationRegistry(deps: OperationRegistryDependencies) { return { admission, acknowledge, + retireRootPublication, + escalateRootPublication, + settleRootPublication, + notifyRootDisappeared, start, prune, active: (rootKiloSessionId: string) => active.get(rootKiloSessionId), diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts index 5255bb7028..eda33d770d 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts @@ -2456,6 +2456,289 @@ describe('owned control execution', () => { } }); + it('routes an early publication-scoped failed Stop through shared-root retirement without cancelling B', async () => { + const runningA = Promise.withResolvers(); + const runningB = Promise.withResolvers(); + const startedA = Promise.withResolvers(); + const startedB = Promise.withResolvers(); + const sibling = { ...session, sessionId: 'ses_b', kiloSessionId: 'kilo_b' }; + rememberAttachedRoot(sibling.kiloSessionId, sibling.directory); + const { handlerDeps, retired } = runtimeDeps( + fakeKilo({ + sendPrompt: options => { + if (options.messageId === 'message_a') { + startedA.resolve(); + return runningA.promise; + } + startedB.resolve(); + return runningB.promise; + }, + getSessionStatuses: async () => ({ [session.kiloSessionId]: { type: 'idle' } }), + abortSession: async () => false, + }) + ); + const runtimes = handlerDeps.kiloRuntimes; + if (!runtimes) throw new Error('Missing worktree runtimes'); + let scopedRetirementCalls = 0; + runtimes.retireRuntimeIfUnshared = async () => { + scopedRetirementCalls += 1; + return 'shared'; + }; + const directoryRetirement = spyOn(handlerDeps.operations, 'retireDirectory'); + const operationId = '11111111-1111-4111-8111-111111111111'; + const requestA = handleControlRequest( + 'session.prompt', + session, + { ...promptPayload, messageId: 'message_a' }, + handlerDeps + ); + const requestB = handleControlRequest( + 'session.prompt', + sibling, + { ...promptPayload, messageId: 'message_b' }, + handlerDeps + ); + try { + await startedA.promise; + await startedB.promise; + const taskA = handlerDeps.operations.active(session.kiloSessionId); + if (!taskA) throw new Error('Missing A operation'); + taskA.markPublicationScoped('publication failure', Date.now() + 1_000); + + const stopped = await handleControlRequest( + 'session.abort', + session, + { messageId: 'message_a', operationId, cleanupDeadlineAt: Date.now() + 1_000 }, + handlerDeps + ); + expect(stopped).toMatchObject({ + ok: true, + result: { status: 'unconfirmed', quiescent: false }, + }); + expect(scopedRetirementCalls).toBe(1); + expect(directoryRetirement).not.toHaveBeenCalled(); + expect(retired).toEqual([]); + expect(handlerDeps.operations.active(sibling.kiloSessionId)?.signal.aborted).toBe(false); + } finally { + runningA.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + runningB.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + await Promise.allSettled([requestA, requestB]); + await waitForTasks(handlerDeps); + directoryRetirement.mockRestore(); + } + }); + + it('rechecks a publication claim that arrives while Stop cleanup is awaiting native abort', async () => { + const runningA = Promise.withResolvers(); + const runningB = Promise.withResolvers(); + const abortPending = Promise.withResolvers(); + const abortStarted = Promise.withResolvers(); + const startedA = Promise.withResolvers(); + const startedB = Promise.withResolvers(); + let abortCalls = 0; + const sibling = { ...session, sessionId: 'ses_b', kiloSessionId: 'kilo_b' }; + rememberAttachedRoot(sibling.kiloSessionId, sibling.directory); + const { handlerDeps, retired } = runtimeDeps( + fakeKilo({ + sendPrompt: options => { + if (options.messageId === 'message_a') { + startedA.resolve(); + return runningA.promise; + } + startedB.resolve(); + return runningB.promise; + }, + getSessionStatuses: async () => ({ [session.kiloSessionId]: { type: 'idle' } }), + abortSession: async () => { + abortCalls += 1; + if (abortCalls === 1) { + abortStarted.resolve(); + return abortPending.promise; + } + return false; + }, + }) + ); + const runtimes = handlerDeps.kiloRuntimes; + if (!runtimes) throw new Error('Missing worktree runtimes'); + let scopedRetirementCalls = 0; + runtimes.retireRuntimeIfUnshared = async () => { + scopedRetirementCalls += 1; + return 'shared'; + }; + const directoryRetirement = spyOn(handlerDeps.operations, 'retireDirectory'); + const requestA = handleControlRequest( + 'session.prompt', + session, + { ...promptPayload, messageId: 'message_a' }, + handlerDeps + ); + const requestB = handleControlRequest( + 'session.prompt', + sibling, + { ...promptPayload, messageId: 'message_b' }, + handlerDeps + ); + try { + await startedA.promise; + await startedB.promise; + const taskA = handlerDeps.operations.active(session.kiloSessionId); + if (!taskA) throw new Error('Missing A operation'); + const stopping = handleControlRequest( + 'session.abort', + session, + { messageId: 'message_a', operationId: '11111111-1111-4111-8111-111111111111' }, + handlerDeps + ); + await abortStarted.promise; + taskA.markPublicationScoped('publication failure during Stop', Date.now() + 1_000); + expect(taskA.publicationScope()).toBeDefined(); + abortPending.resolve(false); + expect(await stopping).toMatchObject({ + ok: true, + result: { status: 'unconfirmed', quiescent: false }, + }); + expect(scopedRetirementCalls).toBe(1); + expect(directoryRetirement).not.toHaveBeenCalled(); + expect(retired).toEqual([]); + expect(handlerDeps.operations.active(sibling.kiloSessionId)?.signal.aborted).toBe(false); + } finally { + abortPending.resolve(false); + runningA.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + runningB.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + await Promise.allSettled([requestA, requestB]); + await waitForTasks(handlerDeps); + directoryRetirement.mockRestore(); + } + }); + + it('does not install a shared-root retirement after scoped cleanup is confirmed', async () => { + const running = Promise.withResolvers(); + const started = Promise.withResolvers(); + const abortCalled = Promise.withResolvers(); + const { handlerDeps, retired } = runtimeDeps( + fakeKilo({ + sendPrompt: () => { + started.resolve(); + return running.promise; + }, + getSessionStatuses: async () => ({ [session.kiloSessionId]: { type: 'idle' } }), + abortSession: async () => { + abortCalled.resolve(); + return true; + }, + }) + ); + const runtimes = handlerDeps.kiloRuntimes; + if (!runtimes) throw new Error('Missing worktree runtimes'); + let scopedRetirementCalls = 0; + runtimes.retireRuntimeIfUnshared = async () => { + scopedRetirementCalls += 1; + return 'shared'; + }; + const directoryRetirement = spyOn(handlerDeps.operations, 'retireDirectory'); + const promptRequest = handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps + ); + try { + await started.promise; + const task = handlerDeps.operations.active(session.kiloSessionId); + if (!task) throw new Error('Missing operation record'); + task.markPublicationScoped('publication cleanup confirmed', Date.now() + 1_000); + const stopping = handleControlRequest( + 'session.abort', + session, + { + messageId: promptPayload.messageId, + operationId: '11111111-1111-4111-8111-111111111111', + cleanupDeadlineAt: Date.now() + 1_000, + }, + handlerDeps + ); + await abortCalled.promise; + await Promise.resolve(); + expect(scopedRetirementCalls).toBe(0); + expect(directoryRetirement).not.toHaveBeenCalled(); + running.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + expect(await stopping).toMatchObject({ + ok: true, + result: { status: 'unconfirmed', quiescent: false }, + }); + expect(retired).toEqual([]); + } finally { + running.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + await Promise.allSettled([promptRequest]); + await waitForTasks(handlerDeps); + directoryRetirement.mockRestore(); + } + }); + + it('reports quiescence only after claimed Stop completes sole-root physical retirement', async () => { + const running = Promise.withResolvers(); + const started = Promise.withResolvers(); + const { handlerDeps, retired } = runtimeDeps( + fakeKilo({ + sendPrompt: () => { + started.resolve(); + return running.promise; + }, + getSessionStatuses: async () => ({ [session.kiloSessionId]: { type: 'idle' } }), + abortSession: async () => false, + }) + ); + const runtimes = handlerDeps.kiloRuntimes; + if (!runtimes) throw new Error('Missing worktree runtimes'); + let scopedRetirementCalls = 0; + runtimes.retireRuntimeIfUnshared = async () => { + scopedRetirementCalls += 1; + return 'retired'; + }; + const directoryRetirement = spyOn(handlerDeps.operations, 'retireDirectory'); + const promptRequest = handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps + ); + try { + await started.promise; + const task = handlerDeps.operations.active(session.kiloSessionId); + if (!task) throw new Error('Missing operation record'); + task.markPublicationScoped('publication cleanup failed', Date.now() + 1_000); + const stopping = handleControlRequest( + 'session.abort', + session, + { + messageId: promptPayload.messageId, + operationId: '11111111-1111-4111-8111-111111111111', + cleanupDeadlineAt: Date.now() + 1_000, + }, + handlerDeps + ); + const stopped = await stopping; + expect(stopped).toMatchObject({ + ok: true, + result: { + status: 'aborted', + quiescent: true, + runtimeRetired: true, + nativeRuntimeId: 'native_1', + }, + }); + expect(scopedRetirementCalls).toBe(1); + expect(directoryRetirement).not.toHaveBeenCalled(); + expect(retired).toEqual([]); + } finally { + running.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + await Promise.allSettled([promptRequest]); + await waitForTasks(handlerDeps); + directoryRetirement.mockRestore(); + } + }); + it('keeps a scoped Stop unconfirmed when preparation fails after cancellation', async () => { const started = Promise.withResolvers(); const handlerDeps = deps({ diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts index 475e3faa56..ddb22fe3d7 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts @@ -340,6 +340,14 @@ export function createControlHandlerDeps(input: Omit) retireRuntime: (directory, deadlineAt, target) => input.kiloRuntimes?.retireRuntime?.(directory, deadlineAt, target) ?? Promise.resolve('unconfirmed'), + retireRuntimeIfUnshared: (directory, target, retiringRoot, deadlineAt, reason) => + input.kiloRuntimes?.retireRuntimeIfUnshared?.( + directory, + target, + retiringRoot, + deadlineAt, + reason + ) ?? Promise.resolve('unconfirmed'), verifyQuiescence: (directory, target, deadlineAt) => input.kiloRuntimes?.verifyQuiescence?.(directory, target, deadlineAt) ?? Promise.resolve(false), @@ -989,10 +997,29 @@ async function handleAbort( parsed.data.cleanupDeadlineAt ?? Date.now() + SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS ); let quiescent = await task.cleanupOwnedWork(deadlineAt); + const publicationScope = task.publicationScope(); + const scopedCleanup = publicationScope ? await task.runRootScopedCleanup() : undefined; + if (publicationScope) quiescent = false; let runtimeRetired = false; let nativeRuntimeId: string | undefined; const target = task.nativeTarget(); - if (!quiescent && target && Date.now() < deadlineAt) { + if (!quiescent && target && publicationScope && scopedCleanup === 'unconfirmed') { + const retiringRoot = + rootForSession(session.kiloSessionId, session.directory) ?? session.kiloSessionId; + const escalation = deps.operations.escalateRootPublication({ + directory: session.directory, + root: retiringRoot, + nativeRuntimeId: target.runtimeId, + target, + reason: publicationScope.reason, + deadlineAt: publicationScope.deadlineAt, + ...(publicationScope.claim === undefined ? {} : { expectedClaim: publicationScope.claim }), + }); + const retirement = await escalation.physical; + runtimeRetired = retirement === 'retired'; + if (runtimeRetired) nativeRuntimeId = target.runtimeId; + quiescent = retirement === 'retired'; + } else if (!quiescent && !publicationScope && target && Date.now() < deadlineAt) { const retirementReason = 'Native cancellation did not settle'; const retirement = await deps.operations.retireDirectory( session.directory, @@ -1004,7 +1031,7 @@ async function handleAbort( if (runtimeRetired) nativeRuntimeId = target.runtimeId; quiescent = task.confirmCleanup(retirement !== 'unconfirmed', deadlineAt); if (retirement === 'unconfirmed') deps.retireRuntime(retirementReason); - } else if (!quiescent) { + } else if (!quiescent && !publicationScope) { task.requestRetirement('Kilo cancellation failed', deadlineAt); } const result = await task.done; diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup-proof.test.ts b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup-proof.test.ts index bab7a71348..5986887027 100644 --- a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup-proof.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup-proof.test.ts @@ -56,10 +56,33 @@ function fixture(overrides: Partial = {}) { completionEvidence: 'unconfirmed', cancel: () => {}, }), + runRoot: (deadlineAt = Date.now() + 150) => + cleanup.cleanupRootScoped({ + deadlineAt, + target, + completionEvidence: 'unconfirmed', + cancel: () => {}, + }), }; } describe('native cancellation proof', () => { + it('keeps normal cleanup runtime-wide while explicit root cleanup skips verification', async () => { + const f = fixture({ abortSession: async () => true }); + expect(await f.run()).toBe(true); + expect(f.verify).toHaveBeenCalledTimes(1); + f.verify.mockClear(); + expect(await f.runRoot()).toBe('confirmed'); + expect(f.verify).not.toHaveBeenCalled(); + }); + + it('reports root-scoped cleanup as unconfirmed when owned process stop fails', async () => { + const f = fixture({ abortSession: async () => true }); + f.stop.mockResolvedValue(false); + expect(await f.runRoot()).toBe('unconfirmed'); + expect(f.verify).not.toHaveBeenCalled(); + }); + it('accepts native empty-map idle only with exact process proof after acknowledged abort', async () => { const abort = jest.fn(async () => true); const f = fixture({ abortSession: abort }); diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts index b3b9578f83..4e9ea10dd6 100644 --- a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts @@ -83,6 +83,92 @@ describe('SessionOperation cleanup', () => { await operation.done; }); + it('confirms root-scoped cleanup from owned processes and root idleness without physical verification', async () => { + const pending = Promise.withResolvers>(); + const abort = jest.fn(async () => true); + const verify = jest.fn(async () => true); + const client = fakeKilo({ + sendPrompt: () => pending.promise, + abortSession: abort, + getSessionStatuses: async () => ({ [session.kiloSessionId]: { type: 'idle' } }), + }); + const runtime: WorktreeKiloRuntime = { + scopeId: 'worktree_1', + runtimeId: 'native_1', + directory: session.directory, + env: {}, + kiloClient: client, + signal: new AbortController().signal, + }; + const operation = new SessionOperation( + session, + undefined, + { operation: 'session.prompt', payload: promptPayload, runtime }, + { + isCurrent: () => true, + getRuntime: () => runtime, + verifyQuiescence: verify, + retireRuntime: () => {}, + emitSessionEvent: () => {}, + onLocalCompletion: () => {}, + onCleanupConfirmed: () => {}, + } + ); + for (let index = 0; index < 5 && operation.snapshot().native.state !== 'pending'; index += 1) + await Promise.resolve(); + + operation.markPublicationScoped('event rejected', Date.now() + 1_000); + const first = operation.runRootScopedCleanup(); + expect(operation.runRootScopedCleanup()).toBe(first); + expect(await first).toBe('confirmed'); + expect(abort).toHaveBeenCalledTimes(1); + expect(verify).not.toHaveBeenCalled(); + + pending.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + await operation.done; + }); + + it('keeps root-scoped cleanup unconfirmed when the root remains busy', async () => { + const pending = Promise.withResolvers>(); + const verify = jest.fn(async () => true); + const client = fakeKilo({ + sendPrompt: () => pending.promise, + abortSession: async () => true, + getSessionStatuses: async () => ({ [session.kiloSessionId]: { type: 'busy' } }), + }); + const runtime: WorktreeKiloRuntime = { + scopeId: 'worktree_1', + runtimeId: 'native_1', + directory: session.directory, + env: {}, + kiloClient: client, + signal: new AbortController().signal, + }; + const operation = new SessionOperation( + session, + undefined, + { operation: 'session.prompt', payload: promptPayload, runtime }, + { + isCurrent: () => true, + getRuntime: () => runtime, + verifyQuiescence: verify, + retireRuntime: () => {}, + emitSessionEvent: () => {}, + onLocalCompletion: () => {}, + onCleanupConfirmed: () => {}, + } + ); + for (let index = 0; index < 5 && operation.snapshot().native.state !== 'pending'; index += 1) + await Promise.resolve(); + + operation.markPublicationScoped('event rejected', Date.now() + 20); + expect(await operation.runRootScopedCleanup()).toBe('unconfirmed'); + expect(verify).not.toHaveBeenCalled(); + + pending.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + await operation.done; + }); + it('does not accept an idle root while a scoped native child is active', async () => { const pending = Promise.withResolvers>(); let verified = 0; diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.ts b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.ts index 03e3bb201f..ac2cba146f 100644 --- a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.ts +++ b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.ts @@ -16,6 +16,7 @@ export type NativeOperationTarget = Readonly<{ export type NativeRetirement = 'retired' | 'stale' | 'unconfirmed'; export type NativeCleanupEvidence = 'not_issued' | 'finished' | 'unconfirmed'; +export type RootScopedCleanupResult = 'confirmed' | 'unconfirmed'; type CleanupState = 'not_requested' | 'acknowledged' | 'confirmed' | 'unconfirmed'; @@ -23,6 +24,8 @@ export class SessionOperationCleanup { private deadlineAt?: number; private state: CleanupState = 'not_requested'; private pending?: Promise; + private rootScopedPending?: Promise; + private rootScopedResult?: RootScopedCleanupResult; private processStop?: Promise; private nativeAbort?: Promise; @@ -95,6 +98,40 @@ export class SessionOperationCleanup { return this.pending; } + cleanupRootScoped(input: { + deadlineAt: number; + target?: NativeOperationTarget; + completionEvidence: NativeCleanupEvidence; + cancel: () => void; + }): Promise { + if (this.rootScopedResult) return Promise.resolve(this.rootScopedResult); + if (this.rootScopedPending) return this.rootScopedPending; + + const deadlineAt = this.captureDeadline(input.deadlineAt); + const processes = this.stopProcesses(deadlineAt); + if (input.completionEvidence === 'unconfirmed') input.cancel(); + this.rootScopedPending = (async () => { + if (!(await processes)) return 'unconfirmed'; + const target = input.target; + const client = target?.client; + if (!target || !client) + return input.completionEvidence === 'unconfirmed' ? 'unconfirmed' : 'confirmed'; + if (!this.rootCurrent(target, deadlineAt)) return 'unconfirmed'; + if ( + input.completionEvidence === 'unconfirmed' && + !(await this.abortNative(target, deadlineAt)) + ) + return 'unconfirmed'; + return (await this.observeRootIdle(target, client, deadlineAt)) ? 'confirmed' : 'unconfirmed'; + })() + .catch(() => 'unconfirmed' as const) + .then(result => { + this.rootScopedResult = result; + return result; + }); + return this.rootScopedPending; + } + confirm(confirmed: boolean, deadlineAt: number): boolean { if (this.state === 'confirmed') return true; const quiescent = @@ -151,6 +188,71 @@ export class SessionOperationCleanup { } } + private rootCurrent(target: NativeOperationTarget, deadlineAt: number): boolean { + const { kiloSessionId, directory } = this.session; + const root = rootForSession(kiloSessionId, directory); + const attachment = root === undefined ? undefined : rootAttachmentId(root); + return ( + root !== undefined && + attachment !== undefined && + rootAttachmentId(root) === attachment && + rootForSession(kiloSessionId, directory) === root && + directoriesForRoot(root, directory).every(value => value === directory) && + this.isCurrent(target) && + Date.now() < Math.min(deadlineAt, this.deadlineAt ?? Infinity) + ); + } + + private async observeRootIdle( + target: NativeOperationTarget, + client: WrapperKiloClient, + deadlineAt: number + ): Promise { + const { kiloSessionId, directory } = this.session; + const root = rootForSession(kiloSessionId, directory); + const attachment = root === undefined ? undefined : rootAttachmentId(root); + const current = () => + root !== undefined && + attachment !== undefined && + rootAttachmentId(root) === attachment && + rootForSession(kiloSessionId, directory) === root && + directoriesForRoot(root, directory).every(value => value === directory) && + this.isCurrent(target) && + Date.now() < Math.min(deadlineAt, this.deadlineAt ?? Infinity); + if (!current()) return false; + const controller = new AbortController(); + try { + return await withTimeoutAndAbort( + withKiloRequestDeadline(async signal => { + if (root === undefined) return false; + const session = await client.getSessionDetails(root, directory, signal); + if (session.id !== root || session.directory !== directory) return false; + while (current()) { + const statuses = await client.getSessionStatuses(directory, signal); + if (!current()) return false; + let idle = true; + for (const [id, status] of Object.entries(statuses)) { + if (status.type === 'idle') continue; + const statusRoot = rootForSession(id, directory); + if (!statusRoot) return false; + if (statusRoot === root) idle = false; + } + if (idle) return current(); + await delay(Math.min(25, Math.max(1, deadlineAt - Date.now())), undefined, { signal }); + } + return false; + }, controller.signal), + { + timeoutMs: Math.max(1, deadlineAt - Date.now()), + timeoutMessage: 'Kilo root cleanup status probe timed out', + abortMessage: 'Kilo root cleanup status probe cancelled', + } + ); + } finally { + controller.abort(); + } + } + private abortNative(target: NativeOperationTarget, deadlineAt: number): Promise { if (this.nativeAbort) return this.nativeAbort; const client = target.client; diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation.test.ts b/services/cloud-agent-next/wrapper/src/control/session-operation.test.ts index 92c3e54cc2..8d75adeded 100644 --- a/services/cloud-agent-next/wrapper/src/control/session-operation.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/session-operation.test.ts @@ -25,8 +25,10 @@ import { session, type Completion, } from './control-test-fixtures'; +import { SessionOperation } from './session-operation'; import { rememberAttachedRoot, resetSessionDirectoryState } from './session-directories'; import { resetDirectoryOperationState } from './worktree-operations'; +import type { WorktreeKiloRuntime } from './worktree-runtime'; let homeRoot: string; @@ -181,6 +183,84 @@ describe('operation results and delivery', () => { expect(record.snapshot().outcome?.status).toBe('cancelled'); }); + it('does not retire the wrapper when a claimed operation publishes a failed outcome', async () => { + const prompt = Promise.withResolvers(); + let retirementCalls = 0; + let outcomeEvents = 0; + const client = fakeKilo({ sendPrompt: () => prompt.promise }); + const runtime: WorktreeKiloRuntime = { + scopeId: 'worktree_1', + runtimeId: 'native_1', + directory: session.directory, + env: {}, + kiloClient: client, + signal: new AbortController().signal, + }; + const operation = new SessionOperation( + session, + undefined, + { operation: 'session.prompt', payload: promptPayload, runtime }, + { + isCurrent: () => true, + getRuntime: () => runtime, + verifyQuiescence: async () => true, + retireRuntime: () => { + retirementCalls += 1; + }, + emitSessionEvent: () => { + outcomeEvents += 1; + return false; + }, + onLocalCompletion: () => {}, + onCleanupConfirmed: () => {}, + } + ); + operation.markPublicationScoped('outcome publication failed', Date.now() + 1_000); + prompt.resolve(completion()); + await operation.done; + expect(retirementCalls).toBe(0); + expect(outcomeEvents).toBe(1); + }); + + it('does not retire the wrapper when execution expiry follows a publication claim', async () => { + const prompt = Promise.withResolvers(); + let retirementCalls = 0; + const client = fakeKilo({ + sendPrompt: () => prompt.promise, + abortSession: async () => true, + }); + const runtime: WorktreeKiloRuntime = { + scopeId: 'worktree_1', + runtimeId: 'native_1', + directory: session.directory, + env: {}, + kiloClient: client, + signal: new AbortController().signal, + }; + const operation = new SessionOperation( + session, + undefined, + { operation: 'session.prompt', payload: promptPayload, runtime }, + { + isCurrent: () => true, + getRuntime: () => runtime, + verifyQuiescence: async () => true, + retireRuntime: () => { + retirementCalls += 1; + }, + emitSessionEvent: () => true, + onLocalCompletion: () => {}, + onCleanupConfirmed: () => {}, + } + ); + operation.markPublicationScoped('execution publication failed', Date.now() + 1_000); + (operation as unknown as { expire(): void }).expire(); + expect(retirementCalls).toBe(0); + prompt.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + await operation.done; + expect(retirementCalls).toBe(0); + }); + it('aborts an already-awaiting finalizer and retains its late original result without false quiescence', async () => { const entered = Promise.withResolvers(); const finalized = Promise.withResolvers(); diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation.ts b/services/cloud-agent-next/wrapper/src/control/session-operation.ts index a520493903..5ead5fb5f7 100644 --- a/services/cloud-agent-next/wrapper/src/control/session-operation.ts +++ b/services/cloud-agent-next/wrapper/src/control/session-operation.ts @@ -37,6 +37,7 @@ import { type NativeCleanupEvidence, type NativeOperationTarget, type NativeRetirement, + type RootScopedCleanupResult, } from './session-operation-cleanup.js'; import { operationIntent } from './operation-intent.js'; import { @@ -160,6 +161,11 @@ export class SessionOperation { private delivery?: OperationResultDelivery; private readonly cleanupOwner: SessionOperationCleanup; private deadlineCleanup?: Promise; + private publicationScoped?: Readonly<{ + reason: string; + deadlineAt: number; + claim?: symbol; + }>; constructor( session: SessionRequestIdentity, @@ -263,6 +269,10 @@ export class SessionOperation { reason = 'Session aborted', status: 'failed' | 'cancelled' = 'cancelled' ): Promise { + if (this.publicationScoped) { + await this.runRootScopedCleanup(); + return false; + } return this.cleanupOwner.cleanup({ deadlineAt, target: this.target, @@ -272,6 +282,38 @@ export class SessionOperation { }); } + markPublicationScoped(reason: string, deadlineAt: number, claim?: symbol): void { + const captured = this.captureCleanupDeadline(deadlineAt); + this.publicationScoped = this.publicationScoped + ? { + reason: this.publicationScoped.reason, + deadlineAt: Math.min(this.publicationScoped.deadlineAt, captured), + ...(claim === undefined && this.publicationScoped.claim === undefined + ? {} + : { claim: claim ?? this.publicationScoped.claim }), + } + : { reason, deadlineAt: captured, ...(claim === undefined ? {} : { claim }) }; + } + + publicationScope(): Readonly<{ reason: string; deadlineAt: number; claim?: symbol }> | undefined { + return this.publicationScoped; + } + + runRootScopedCleanup( + reason = 'Session event delivery failed', + deadlineAt = Date.now() + SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS + ): Promise { + if (!this.publicationScoped) this.markPublicationScoped(reason, deadlineAt); + const scoped = this.publicationScoped; + if (!scoped) return Promise.resolve('unconfirmed'); + return this.cleanupOwner.cleanupRootScoped({ + deadlineAt: scoped.deadlineAt, + target: this.target, + completionEvidence: this.cleanupEvidence(), + cancel: () => this.cancel(scoped.reason, 'failed', scoped.deadlineAt), + }); + } + snapshot() { const retained = this.retainedNotifications.snapshot(); return { @@ -327,12 +369,19 @@ export class SessionOperation { } cancel(reason: string, status: 'failed' | 'cancelled', cleanupDeadlineAt?: number): void { - if (cleanupDeadlineAt !== undefined) this.captureCleanupDeadline(cleanupDeadlineAt); + if (cleanupDeadlineAt !== undefined) { + const captured = this.captureCleanupDeadline(cleanupDeadlineAt); + if (this.publicationScoped) + this.publicationScoped = { + ...this.publicationScoped, + deadlineAt: Math.min(this.publicationScoped.deadlineAt, captured), + }; + } if (!this.local) this.controller.abort(new ControlTaskCancellation(status, reason)); } requestRetirement(reason: string, deadlineAt: number): void { - if (this.cleanupOwner.cleanupState === 'confirmed') return; + if (this.publicationScoped || this.cleanupOwner.cleanupState === 'confirmed') return; this.deps.retireRuntime(reason, this.captureCleanupDeadline(deadlineAt), this.nativeTarget()); } diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts b/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts index a75feda600..cef90cd44b 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts @@ -6,12 +6,15 @@ import { buildWorktreeKiloEnvironment, createWorktreeKiloRuntimes, startWorktreeKiloServer, + type RootRuntimeRetirement, type WorktreeKiloAuth, type WorktreeKiloRuntimes, } from './worktree-runtime'; +import type { NativeRetirement } from './session-operation-cleanup'; import type { OwnedProcessScope } from './owned-processes'; import { SANDBOX_CONTROL_RECOVERY_MAX_ATTEMPTS, + SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS, SANDBOX_CONTROL_EXECUTION_TIMEOUT_MS, sessionMessageOutcomeSchema, type SessionEventIdentity, @@ -34,6 +37,7 @@ import { childFromSessionCreated, eventKiloSessionId, sessionEventIdentity } fro import { directoryForSession, rememberChildSession, + rememberAttachedRoot, resetSessionDirectoryState, rootForSession, } from './session-directories'; @@ -315,6 +319,25 @@ function createHandlerDeps(registry: WorktreeKiloRuntimes): HandlerDeps { }); } +function createIntegratedRegistry( + overrides: Partial[0]> = {} +) { + const context: { handlerDeps?: HandlerDeps } = {}; + const settlements: RootRuntimeRetirement[] = []; + const harness = createRegistry({ + onRootRetirement: settlement => { + settlements.push(settlement); + context.handlerDeps?.operations.settleRootPublication(settlement); + }, + onRootDisappeared: disappearance => + context.handlerDeps?.operations.notifyRootDisappeared(disappearance), + ...overrides, + }); + const dependencies = createHandlerDeps(harness.registry); + context.handlerDeps = dependencies; + return { ...harness, handlerDeps: dependencies, settlements }; +} + beforeEach(() => { resetSessionDirectoryState(); tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'worktree-kilo-test-')); @@ -2403,4 +2426,429 @@ setInterval(() => {}, 1000); expect(attachment.signal.aborted).toBe(true); stopped.resolve(); }); + + it('defers failed-root retirement behind a pending sibling and uses a fresh deadline when it becomes sole', async () => { + const stopDeadlines: number[] = []; + const rootRetirements: string[] = []; + const registry = createWorktreeKiloRuntimes({ + homeRoot: path.join(tmpDir, 'homes'), + inheritedEnv: inherited, + startServer: async options => { + const server = createKiloStub(); + servers.push(server); + options.onProcessScope?.({ + stop: async (deadlineAt: number) => { + stopDeadlines.push(deadlineAt); + return true; + }, + } as unknown as OwnedProcessScope); + return { url: server.url, close: () => {} }; + }, + onRootRetirement: retirement => { + if (retirement.result === 'retired') rootRetirements.push(retirement.root); + }, + onUnexpectedClose: () => {}, + }); + registries.push(registry); + const directory = path.join(tmpDir, 'shared'); + const first = registry.attach(rootIdentity(directory, 'first'), auth); + const runtime = await first.ready; + first.commit(); + first.release(); + const pending = registry.attach(rootIdentity(directory, 'pending'), auth); + await pending.ready; + + const originalDeadline = Date.now() - 1; + expect( + await registry.retireRuntimeIfUnshared?.( + directory, + { runtimeId: runtime.runtimeId, client: runtime.kiloClient }, + 'root_first', + originalDeadline, + 'event rejected' + ) + ).toBe('shared'); + expect(registry.get(directory)).toBe(runtime); + expect(stopDeadlines).toEqual([]); + + pending.release(); + await waitUntil(() => rootRetirements.includes('root_first')); + expect(stopDeadlines[0]).toBeGreaterThan(originalDeadline); + expect(stopDeadlines[0]).toBeLessThanOrEqual(Date.now() + SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS); + expect(registry.get(directory)).toBeUndefined(); + }); + + it('rejects a stale target without retiring the current runtime', async () => { + const harness = createRegistry(); + const { registry } = harness; + const directory = path.join(tmpDir, 'stale-target'); + const attachment = registry.attach(rootIdentity(directory), auth); + const runtime = await attachment.ready; + attachment.commit(); + attachment.release(); + + expect( + await registry.retireRuntimeIfUnshared?.( + directory, + { runtimeId: crypto.randomUUID(), client: runtime.kiloClient }, + 'root_stale-target', + Date.now() + 1_000, + 'stale event' + ) + ).toBe('stale'); + expect(registry.get(directory)).toBe(runtime); + expect(harness.closes).toBe(0); + }); + + it('reports immediate root retirement with its captured incarnation even without a deferred intent', async () => { + const reports: string[] = []; + const harness = createRegistry({ + onRootRetirement: retirement => + reports.push(`${retirement.root}:${retirement.nativeRuntimeId}:${retirement.result}`), + }); + const directory = path.join(tmpDir, 'immediate-root-retirement'); + const attachment = harness.registry.attach(rootIdentity(directory), auth); + const runtime = await attachment.ready; + attachment.commit(); + attachment.release(); + const target = { runtimeId: runtime.runtimeId, client: runtime.kiloClient }; + + expect( + await harness.registry.retireRuntimeIfUnshared?.( + directory, + target, + 'root_immediate-root-retirement', + Date.now() + 1_000, + 'event rejected' + ) + ).toBe('retired'); + expect(reports).toEqual([`root_immediate-root-retirement:${runtime.runtimeId}:retired`]); + expect(harness.closes).toBe(1); + }); + + it('does not settle a current N2 deferred failure from a stale N1 cleanup callback', async () => { + const harness = createRegistry(); + const { registry } = harness; + const directory = path.join(tmpDir, 'stale-deferred-target'); + const first = registry.attach(rootIdentity(directory, 'first'), auth); + const sibling = registry.attach(rootIdentity(directory, 'sibling'), auth); + const runtime = await first.ready; + await sibling.ready; + first.commit(); + sibling.commit(); + first.release(); + sibling.release(); + const currentTarget = { runtimeId: runtime.runtimeId, client: runtime.kiloClient }; + expect( + await registry.retireRuntimeIfUnshared?.( + directory, + currentTarget, + 'root_first', + Date.now() + 1_000, + 'N2 publication failure' + ) + ).toBe('shared'); + if (!registry.retireRuntime) throw new Error('Missing runtime retirement API'); + expect( + await registry.retireRuntime(directory, Date.now() + 1_000, { + runtimeId: 'N1', + client: runtime.kiloClient, + }) + ).toBe('stale'); + expect(registry.get(directory)).toBe(runtime); + + registry.detach(rootIdentity(directory, 'sibling')); + await waitUntil(() => registry.get(directory) === undefined); + expect(harness.closes).toBe(1); + }); + + it('does not report physical unconfirmed unregistration as explicit root disappearance', async () => { + const rootRetirements: NativeRetirement[] = []; + const disappearedRoots: string[] = []; + const server = createKiloStub(); + servers.push(server); + const harness = createRegistry({ + startServer: async options => { + options.onProcessScope?.({ + stop: async (_deadlineAt: number) => false, + } as unknown as OwnedProcessScope); + return { url: server.url, close: () => {} }; + }, + onRootRetirement: retirement => rootRetirements.push(retirement.result), + onRootDisappeared: disappearance => disappearedRoots.push(disappearance.root), + }); + const directory = path.join(tmpDir, 'unconfirmed-unregistration'); + const first = harness.registry.attach(rootIdentity(directory, 'first'), auth); + const sibling = harness.registry.attach(rootIdentity(directory, 'sibling'), auth); + const runtime = await first.ready; + await sibling.ready; + first.commit(); + sibling.commit(); + first.release(); + sibling.release(); + expect( + await harness.registry.retireRuntimeIfUnshared?.( + directory, + { runtimeId: runtime.runtimeId, client: runtime.kiloClient }, + 'root_first', + Date.now() + 1_000, + 'event rejected' + ) + ).toBe('shared'); + harness.registry.detach(rootIdentity(directory, 'sibling')); + await waitUntil(() => rootRetirements.length > 0); + expect(rootRetirements).toEqual(['unconfirmed']); + expect(disappearedRoots).toEqual(['root_sibling']); + }); + + it('does not retire a healthy survivor when multiple failed roots disappear', async () => { + const harness = createRegistry(); + const { registry } = harness; + const directory = path.join(tmpDir, 'multiple-failures'); + const roots = ['first', 'second', 'healthy'].map(name => rootIdentity(directory, name)); + const attachments = roots.map(root => registry.attach(root, auth)); + const runtime = await attachments[0]?.ready; + if (!runtime) throw new Error('Missing native runtime'); + for (const attachment of attachments) { + attachment.commit(); + attachment.release(); + } + const target = { runtimeId: runtime.runtimeId, client: runtime.kiloClient }; + expect( + await registry.retireRuntimeIfUnshared?.( + directory, + target, + roots[0]?.kiloSessionId ?? 'root_first', + Date.now() + 1_000, + 'first event rejected' + ) + ).toBe('shared'); + expect( + await registry.retireRuntimeIfUnshared?.( + directory, + target, + roots[1]?.kiloSessionId ?? 'root_second', + Date.now() + 1_000, + 'second event rejected' + ) + ).toBe('shared'); + + registry.detach(roots[0]!); + await Promise.resolve(); + expect(registry.get(directory)).toBe(runtime); + registry.detach(roots[1]!); + await Promise.resolve(); + expect(registry.get(directory)).toBe(runtime); + expect(harness.closes).toBe(0); + }); + + it.each([{ order: ['first', 'second'] as const }, { order: ['second', 'first'] as const }])( + 'retires either failed root when the other failed root detaches first', + async ({ order }) => { + const harness = createRegistry(); + const { registry } = harness; + const directory = path.join(tmpDir, `failed-order-${order[0]}`); + const first = registry.attach(rootIdentity(directory, 'first'), auth); + const second = registry.attach(rootIdentity(directory, 'second'), auth); + const runtime = await first.ready; + await second.ready; + first.commit(); + second.commit(); + first.release(); + second.release(); + const target = { runtimeId: runtime.runtimeId, client: runtime.kiloClient }; + expect( + await registry.retireRuntimeIfUnshared?.( + directory, + target, + 'root_first', + Date.now() + 1_000, + 'first event rejected' + ) + ).toBe('shared'); + expect( + await registry.retireRuntimeIfUnshared?.( + directory, + target, + 'root_second', + Date.now() + 1_000, + 'second event rejected' + ) + ).toBe('shared'); + + expect(registry.detach(rootIdentity(directory, order[0]))).toBe(true); + await waitUntil(() => registry.get(directory) === undefined); + expect(harness.closes).toBe(1); + } + ); +}); + +describe('runtime-to-registry root settlement', () => { + it('keeps a shared runtime alive after confirmed A cleanup, fresh A work, and B detach', async () => { + const integrated = createIntegratedRegistry(); + const directory = path.join(tmpDir, 'confirmed-shared-runtime'); + const identityA = rootIdentity(directory, 'a'); + const identityB = rootIdentity(directory, 'b'); + expect( + await handleControlRequest( + 'session.attach', + identityA, + { kilo: auth }, + integrated.handlerDeps + ) + ).toMatchObject({ + ok: true, + }); + expect( + await handleControlRequest( + 'session.attach', + identityB, + { kilo: auth }, + integrated.handlerDeps + ) + ).toMatchObject({ + ok: true, + }); + const runtime = integrated.registry.get(directory); + const server = servers.at(-1); + if (!runtime || !server) throw new Error('Missing shared runtime'); + server.holdPrompts(); + const prompt = { + messageId: 'confirmed_a1', + turn: { type: 'prompt' as const, prompt: 'confirm root cleanup' }, + agent: { mode: 'code', model: 'test' }, + }; + const promptRequest = handleControlRequest( + 'session.prompt', + identityA, + prompt, + integrated.handlerDeps + ); + const waitForTasks = () => + Promise.all(integrated.handlerDeps.operations.activeOperations().map(task => task.done)); + try { + await waitUntil( + () => + integrated.handlerDeps.operations.active(identityA.kiloSessionId)?.snapshot().native + .state === 'pending' + ); + const escalation = integrated.handlerDeps.operations.escalateRootPublication({ + directory, + root: identityA.kiloSessionId, + nativeRuntimeId: runtime.runtimeId, + target: { runtimeId: runtime.runtimeId, client: runtime.kiloClient }, + reason: 'confirmed publication cleanup', + deadlineAt: Date.now() + 1_000, + }); + expect(await escalation.physical).toBe('confirmed'); + expect(integrated.registry.get(directory)).toBe(runtime); + await promptRequest; + await waitForTasks(); + + const freshPrompt = { + ...prompt, + messageId: 'confirmed_a2', + turn: { type: 'prompt' as const, prompt: 'fresh A work' }, + }; + expect( + await handleControlRequest('session.prompt', identityA, freshPrompt, integrated.handlerDeps) + ).toMatchObject({ ok: true, result: { status: 'accepted' } }); + await waitForTasks(); + expect(integrated.registry.get(directory)).toBe(runtime); + + expect(integrated.registry.detach(identityB)).toBe(true); + expect(integrated.registry.get(directory)).toBe(runtime); + expect(runtime.signal.aborted).toBe(false); + expect(integrated.closes).toBe(0); + } finally { + server.releasePrompts(); + await Promise.allSettled([promptRequest]); + await waitForTasks(); + } + }); + + it('settles a non-deferred scoped record when failRuntime retires the runtime', async () => { + const exited = Promise.withResolvers(); + const server = createKiloStub(); + servers.push(server); + const integrated = createIntegratedRegistry({ + startServer: async options => { + options.onProcessScope?.({ stop: async () => true } as unknown as OwnedProcessScope); + return { url: server.url, close: () => {}, exited: exited.promise }; + }, + }); + const directory = path.join(tmpDir, 'settled-failure-record'); + const attachment = integrated.registry.attach(rootIdentity(directory), auth); + const runtime = await attachment.ready; + attachment.commit(); + attachment.release(); + const input = { + directory, + root: `root_${path.basename(directory)}`, + nativeRuntimeId: runtime.runtimeId, + target: { runtimeId: runtime.runtimeId, client: runtime.kiloClient }, + reason: 'event rejected', + deadlineAt: Date.now() + 1_000, + }; + expect(await integrated.handlerDeps.operations.retireRootPublication(input)).toBe( + 'unconfirmed' + ); + + exited.resolve(); + await waitUntil(() => + integrated.settlements.some( + settlement => settlement.root === input.root && settlement.result === 'retired' + ) + ); + expect(integrated.registry.get(directory)).toBeUndefined(); + rememberAttachedRoot(input.root, directory); + expect( + integrated.handlerDeps.operations.admission( + 'session.prompt', + rootIdentity(directory), + undefined + ).kind + ).toBe('continue'); + }); + + it('preserves a non-deferred scoped record when failRuntime retirement is unconfirmed', async () => { + const exited = Promise.withResolvers(); + const server = createKiloStub(); + servers.push(server); + const integrated = createIntegratedRegistry({ + startServer: async options => { + options.onProcessScope?.({ stop: async () => false } as unknown as OwnedProcessScope); + return { url: server.url, close: () => {}, exited: exited.promise }; + }, + }); + const directory = path.join(tmpDir, 'unconfirmed-failure-record'); + const attachment = integrated.registry.attach(rootIdentity(directory), auth); + const runtime = await attachment.ready; + attachment.commit(); + attachment.release(); + const sessionA = rootIdentity(directory); + const input = { + directory, + root: sessionA.kiloSessionId, + nativeRuntimeId: runtime.runtimeId, + target: { runtimeId: runtime.runtimeId, client: runtime.kiloClient }, + reason: 'event rejected', + deadlineAt: Date.now() + 1_000, + }; + expect(await integrated.handlerDeps.operations.retireRootPublication(input)).toBe( + 'unconfirmed' + ); + + exited.resolve(); + await waitUntil(() => + integrated.settlements.some( + settlement => settlement.root === input.root && settlement.result === 'unconfirmed' + ) + ); + integrated.handlerDeps.operations.prune(); + rememberAttachedRoot(sessionA.kiloSessionId, directory); + expect( + integrated.handlerDeps.operations.admission('session.prompt', sessionA, undefined).kind + ).toBe('reply'); + }); }); diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-runtime.ts b/services/cloud-agent-next/wrapper/src/control/worktree-runtime.ts index 8b1a821fb9..48cbab7e9c 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-runtime.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-runtime.ts @@ -56,6 +56,16 @@ export type WorktreeKiloAttachment = { release(): void; }; +export type RootRuntimeRetirement = { + directory: string; + root: string; + nativeRuntimeId: string; + target: NativeOperationTarget; + result: NativeRetirement; +}; + +export type RootRuntimeDisappearance = Omit; + export type WorktreeKiloRuntimes = { readonly kiloCliVersion?: string | null; attach( @@ -73,6 +83,13 @@ export type WorktreeKiloRuntimes = { deadlineAt: number, target?: NativeOperationTarget ): Promise; + retireRuntimeIfUnshared?( + directory: string, + target: NativeOperationTarget, + retiringRoot: string, + deadlineAt: number, + reason?: string + ): Promise; verifyQuiescence?( directory: string, target: NativeOperationTarget, @@ -331,6 +348,8 @@ export function createWorktreeKiloRuntimes(options: { inheritedEnv?: NodeJS.ProcessEnv; startServer?: (options: ServerOptions) => Promise; onEvent?: (runtime: WorktreeKiloRuntime, event: WorktreeKiloEvent) => unknown; + onRootRetirement?: (retirement: RootRuntimeRetirement) => void; + onRootDisappeared?: (disappearance: RootRuntimeDisappearance) => void; onDiagnostic?: ControlDiagnosticReporter; onUnexpectedClose: (failure: WorktreeKiloFailure) => void; }): WorktreeKiloRuntimes { @@ -340,6 +359,19 @@ export function createWorktreeKiloRuntimes(options: { const roots = new Map(); const homesByDirectory = new Map>(); const deletedDirectories = new Set(); + const deferredRetirements = new Map< + string, + { + directory: string; + root: string; + nativeRuntimeId: string; + entry: RuntimeEntry; + target: NativeOperationTarget; + reason: string; + deadlineAt: number; + } + >(); + const evaluatingDeferredRetirements = new Set(); let observedVersion: string | null | undefined; let closed = false; @@ -377,6 +409,116 @@ export function createWorktreeKiloRuntimes(options: { return entry.cleanupDeadlineAt; } + function deferredRetirementKey(root: string, nativeRuntimeId: string): string { + return JSON.stringify([root, nativeRuntimeId]); + } + + function runtimeTargetMatches(entry: RuntimeEntry, target: NativeOperationTarget): boolean { + return ( + entry.runtimeId === target.runtimeId && + (target.client === undefined || target.client === entry.kiloClient) + ); + } + + function reportRootRetirement( + intent: { + directory: string; + root: string; + nativeRuntimeId: string; + target: NativeOperationTarget; + }, + result: NativeRetirement + ): void { + options.onRootRetirement?.({ ...intent, result }); + } + + function reportRuntimeRetirement( + intent: { + directory: string; + root: string; + nativeRuntimeId: string; + target: NativeOperationTarget; + }, + retirement: Promise + ): Promise { + void retirement.then(result => reportRootRetirement(intent, result)); + return retirement; + } + + function settleDeferredRetirement( + intent: { + directory: string; + root: string; + nativeRuntimeId: string; + }, + result: NativeRetirement + ): void { + const id = deferredRetirementKey(intent.root, intent.nativeRuntimeId); + const current = deferredRetirements.get(id); + if (!current || current.directory !== intent.directory) return; + deferredRetirements.delete(id); + reportRootRetirement({ ...intent, target: current.target }, result); + } + + function settleEntryDeferredRetirements( + entry: RuntimeEntry, + target: NativeOperationTarget | undefined, + result: NativeRetirement + ): Set { + const settled = new Set(); + for (const intent of [...deferredRetirements.values()]) { + if (intent.entry !== entry) continue; + if ( + target && + (target.runtimeId !== intent.nativeRuntimeId || + (target.client !== undefined && + intent.target.client !== undefined && + target.client !== intent.target.client)) + ) + continue; + settled.add(deferredRetirementKey(intent.root, intent.nativeRuntimeId)); + settleDeferredRetirement(intent, result); + } + return settled; + } + + function liveRoots(entry: RuntimeEntry): RootAttachment[] { + return [...entry.roots].filter(root => root.attached || root.pending.size > 0); + } + + function evaluateDeferredRetirements(entry: RuntimeEntry): void { + if (entry.retiring || evaluatingDeferredRetirements.has(entry)) return; + const live = liveRoots(entry); + evaluatingDeferredRetirements.add(entry); + try { + for (const intent of [...deferredRetirements.values()]) { + if (intent.entry !== entry) continue; + const failedRoot = [...entry.roots].find( + root => root.identity.kiloSessionId === intent.root + ); + if (!failedRoot || !live.includes(failedRoot)) { + settleDeferredRetirement(intent, 'stale'); + continue; + } + if (!runtimeTargetMatches(entry, intent.target)) { + settleDeferredRetirement(intent, 'stale'); + continue; + } + if (live.some(root => root !== failedRoot)) continue; + + deferredRetirements.delete(deferredRetirementKey(intent.root, intent.nativeRuntimeId)); + const now = Date.now(); + const physicalDeadlineAt = + now >= intent.deadlineAt + ? now + SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS + : Math.min(intent.deadlineAt, now + SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS); + void retire(entry, physicalDeadlineAt, intent.target); + } + } finally { + evaluatingDeferredRetirements.delete(entry); + } + } + function unregisterRoot(root: RootAttachment): void { root.entry.roots.delete(root); root.attached = false; @@ -384,7 +526,15 @@ export function createWorktreeKiloRuntimes(options: { if (roots.get(root.identity.kiloSessionId) === root) { roots.delete(root.identity.kiloSessionId); forgetAttachedRoot(root.identity.kiloSessionId, root.identity.directory); + if (!root.entry.retiring) + options.onRootDisappeared?.({ + directory: root.entry.directory, + root: root.identity.kiloSessionId, + nativeRuntimeId: root.entry.runtimeId, + target: { runtimeId: root.entry.runtimeId, client: root.entry.kiloClient }, + }); } + evaluateDeferredRetirements(root.entry); } function retire( @@ -392,7 +542,9 @@ export function createWorktreeKiloRuntimes(options: { requested?: number, target?: NativeOperationTarget ): Promise { - return retireWorktreeRuntime(entry, requested, target, { + const settlementTarget = target ?? { runtimeId: entry.runtimeId, client: entry.kiloClient }; + const affectedRoots = [...entry.roots].map(root => root.identity.kiloSessionId); + const retirement = retireWorktreeRuntime(entry, requested, target, { cleanupDeadline, unregisterRoot, removeEntry: retiring => { @@ -403,6 +555,71 @@ export function createWorktreeKiloRuntimes(options: { } }, }); + void retirement.then(result => { + const settled = settleEntryDeferredRetirements(entry, target, result); + for (const root of affectedRoots) { + const id = deferredRetirementKey(root, settlementTarget.runtimeId); + if (!settled.has(id)) + reportRootRetirement( + { + directory: entry.directory, + root, + nativeRuntimeId: settlementTarget.runtimeId, + target: settlementTarget, + }, + result + ); + } + }); + return retirement; + } + + function retireRuntimeIfUnshared( + directory: string, + target: NativeOperationTarget, + retiringRoot: string, + deadlineAt: number, + reason = 'Native runtime retirement requested' + ): Promise { + const entry = entries.get(directory); + if (!entry || !runtimeTargetMatches(entry, target)) { + const intent = { + directory, + root: retiringRoot, + nativeRuntimeId: target.runtimeId, + target, + }; + settleDeferredRetirement(intent, 'stale'); + return reportRuntimeRetirement(intent, Promise.resolve('stale')); + } + const live = liveRoots(entry); + const intent = { + directory, + root: retiringRoot, + nativeRuntimeId: target.runtimeId, + target, + }; + if (entry.retiring) return entry.retiring; + const failedRoot = [...entry.roots].find(root => root.identity.kiloSessionId === retiringRoot); + if (!failedRoot || !(failedRoot.attached || failedRoot.pending.size > 0)) { + settleDeferredRetirement(intent, 'stale'); + return reportRuntimeRetirement(intent, Promise.resolve('stale')); + } + if (live.some(root => root !== failedRoot)) { + const id = deferredRetirementKey(retiringRoot, target.runtimeId); + deferredRetirements.set(id, { + directory, + root: retiringRoot, + nativeRuntimeId: target.runtimeId, + entry, + target, + reason, + deadlineAt, + }); + return Promise.resolve('shared'); + } + deferredRetirements.delete(deferredRetirementKey(retiringRoot, target.runtimeId)); + return retire(entry, Date.now() + SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS, target); } function removeRoot(root: RootAttachment): void { @@ -638,6 +855,7 @@ export function createWorktreeKiloRuntimes(options: { if (!(await stopped) || Date.now() >= deadlineAt) { throw new Error('Original native execution is not contained'); } + settleEntryDeferredRetirements(entry, undefined, 'stale'); entry.kilo = { ...kilo, targets: { ...kilo.targets } }; entry.env = env; entry.cleanupDeadlineAt = undefined; @@ -747,6 +965,7 @@ export function createWorktreeKiloRuntimes(options: { } } if (!entry) { + if (previous) settleEntryDeferredRetirements(previous, undefined, 'stale'); const homeId = createHash('sha256') .update(kilo.scopeId) .update('\0') @@ -853,6 +1072,9 @@ export function createWorktreeKiloRuntimes(options: { if (!entry) return 'stale'; return retire(entry, deadlineAt, target); }, + async retireRuntimeIfUnshared(directory, target, retiringRoot, deadlineAt, reason) { + return retireRuntimeIfUnshared(directory, target, retiringRoot, deadlineAt, reason); + }, async verifyQuiescence(directory, target, deadlineAt) { const entry = entries.get(directory); if ( @@ -902,6 +1124,8 @@ export function createWorktreeKiloRuntimes(options: { closed = true; for (const root of roots.values()) removeRoot(root); for (const entry of entries.values()) void retire(entry); + for (const intent of [...deferredRetirements.values()]) + settleDeferredRetirement(intent, 'stale'); directoriesByScope.clear(); }, };