Skip to content

fix(cloud-agent-next): isolate control-event retirement #6002

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
eshurakov merged 1 commit into from
Sep 10, 2026
Merged

fix(cloud-agent-next): isolate control-event retirement #6002

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 49 additions & 4 deletions services/cloud-agent-next/wrapper/src/control/control-event-transport.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<void>();
const handleFailure = createControlEventFailureHandler({
getRuntime: () => current,
onFailure: retired,
onFailure: (...args) => {
retired(...args);
return cleanup.promise;
},
});
const failures: ControlEventOutboxFailure[] = [];
const published: ControlEventPublication[] = [];
Expand Down Expand Up @@ -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, {
Expand All @@ -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);
Expand All @@ -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<void>();
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 = {
Expand Down
35 changes: 29 additions & 6 deletions services/cloud-agent-next/wrapper/src/control/control-event-transport.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,18 +9,41 @@ type EventKind = 'session.event' | 'session.preparing';

export function createControlEventFailureHandler<Runtime extends { runtimeId: string }>(options: {
getRuntime: (directory: string) => Runtime | undefined;
onFailure: (failure: ControlEventOutboxFailure, runtime: Runtime) => void;
onFailure: (failure: ControlEventOutboxFailure, runtime: Runtime) => unknown;
}) {
const failedRuntimes = new WeakSet<Runtime>();
const inFlight = new WeakMap<Runtime, Set<string>>();
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<string>();
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);
}
);
};
}

Expand Down
161 changes: 123 additions & 38 deletions services/cloud-agent-next/wrapper/src/control/main.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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 {
Expand All @@ -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<RootScopedCleanupResult>;
physical: Promise<PublicationRetirementResult>;
};
import type { ControlEventOutboxFailure } from './control-event-outbox';

function main(
Expand All @@ -49,8 +64,16 @@ function main(
let control: ReturnType<typeof maybeStartSandboxControlClient> = 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({
Expand All @@ -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' });
}
Expand Down Expand Up @@ -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<PublicationRetirementResult> {
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,
Expand Down Expand Up @@ -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'),
Expand Down
Loading