Add explicit Runner session release and resume support (#19683)

This commit is contained in:
Nicky Leach authored and GitHub committed 2026-11-06 18:31:23 -08:00
1 parent f09b4704b6
commit 03498f90ee
15 files changed
+1131 -8

No files matched your search

@@ -366,6 +366,67 @@ describe("runner provider-ingress WebSocket wire", () => {
}
});
it.each(["readiness", "authenticated observer"])("does not start liveness when quiesce wins %s", async mode => {
vi.useFakeTimers();
const create = sockets.create, socket = new FakeSocket();
sockets.create = () => openingSocket(socket);
const endpoint = { kind: "authenticated_websocket" as const,
websocketUrl: "wss://fixture.example.test/connect", secretHeaders: [], generation: "fixture",
refresh: vi.fn(async () => endpoint), close: vi.fn(async () => {}) };
const handle = connectRunnerPrpIngress({ authority: {
attachWireConnection: () => ({ isAuthenticated: () => {
if (mode === "readiness") handle.quiesce();
return true;
} }),
} as unknown as DurablePrpControlPlane, endpoint,
onStateChange: state => { if (state === "authenticated" && mode === "authenticated observer") handle.quiesce(); } });
try {
await vi.advanceTimersByTimeAsync(100_000);
expect(socket.ping).not.toHaveBeenCalled();
expect(socket.terminate).not.toHaveBeenCalled();
expect(socket.close).not.toHaveBeenCalled();
expect(endpoint.refresh).not.toHaveBeenCalled();
} finally { await handle.close(); sockets.create = create; vi.useRealTimers(); }
});
it("quiets liveness and reconnect while preserving the final receive/ACK path", async () => {
vi.useFakeTimers();
const create = sockets.create;
const socket = new FakeSocket();
let opens = 0;
sockets.create = () => { opens += 1; return openingSocket(socket); };
const endpoint = { kind: "authenticated_websocket" as const,
websocketUrl: "wss://fixture.example.test/connect", secretHeaders: [], generation: "fixture",
refresh: vi.fn(async () => endpoint), close: vi.fn(async () => {}) };
const received = vi.fn();
const attach = vi.fn((wire: WsJsonWireConnection) => {
wire.onJson(value => { received(value); wire.sendJson({ kind: "ack" }); });
return { isAuthenticated: () => true };
});
const handle = connectRunnerPrpIngress({ authority: { attachWireConnection: attach } as unknown as DurablePrpControlPlane,
endpoint, livenessIntervalMs: 10, livenessTimeoutMs: 30 });
const failure = vi.fn();
void handle.failure.catch(failure);
try {
await vi.advanceTimersByTimeAsync(0);
await handle.ready;
handle.quiesce();
socket.ping.mockClear();
socket.emit("message", Buffer.from(JSON.stringify({ kind: "last_event" })), false);
expect(received).toHaveBeenCalledWith({ kind: "last_event" });
expect(socket.send).toHaveBeenCalledWith(JSON.stringify({ kind: "ack" }));
await vi.advanceTimersByTimeAsync(100_000);
expect(socket.ping).not.toHaveBeenCalled();
expect(socket.terminate).not.toHaveBeenCalled();
expect(socket.close).not.toHaveBeenCalled();
socket.close();
await vi.advanceTimersByTimeAsync(100_000);
expect(opens).toBe(1);
expect(endpoint.refresh).not.toHaveBeenCalled();
expect(failure).not.toHaveBeenCalled();
} finally { await handle.close(); sockets.create = create; vi.useRealTimers(); }
});
it("does not report intentional close as an active-run ingress failure", async () => {
const endpoint = {
kind: "authenticated_websocket" as const,
+13 -1
View File
@@ -305,6 +305,8 @@ export interface RunnerPrpOutboundHandle {
readonly ready: Promise<void>;
/** Rejects only when active-run recovery has exhausted its fixed budget. */
readonly failure: Promise<never>;
/** Stop liveness and reconnect work, preserving an established receive/ACK path. */
quiesce(): void;
close(): Promise<void>;
}
@@ -337,6 +339,7 @@ export function connectRunnerPrpIngress(input: {
let readyResolve!: () => void;
let readyReject!: (error: unknown) => void;
let readySettled = false;
let stopLiveness = () => {};
const ready = new Promise<void>((resolve, reject) => {
readyResolve = resolve;
readyReject = reject;
@@ -386,22 +389,26 @@ export function connectRunnerPrpIngress(input: {
"Runner ingress did not complete PRP authentication before the deadline.",
);
}
if (abort.signal.aborted) return;
if (!readySettled) {
readySettled = true;
readyResolve();
}
input.onStateChange?.("authenticated");
if (abort.signal.aborted) return;
attempt = 0;
leaseFailureCode = null;
const stopLiveness = opened.wire.monitorLiveness({
stopLiveness = opened.wire.monitorLiveness({
intervalMs: input.livenessIntervalMs ?? DEFAULT_LIVENESS_INTERVAL_MS,
timeoutMs: input.livenessTimeoutMs ?? DEFAULT_LIVENESS_TIMEOUT_MS,
});
if (abort.signal.aborted) { stopLiveness(); return; }
try {
await socketClosed(opened.wire);
} finally {
stopLiveness();
}
if (abort.signal.aborted) return;
recoveryDeadline = Date.now() + recoveryGraceMs;
const leaseFailure = input.authority.connectionLeaseFailure ?? null;
if (leaseFailure) {
@@ -480,6 +487,11 @@ export function connectRunnerPrpIngress(input: {
return {
ready,
failure,
quiesce() {
stopLiveness();
// Established wires removed their handshake abort listener. Keep their ACK path.
abort.abort();
},
async close() {
abort.abort();
await Promise.all([...ownedWires].map(closeOwnedWire));
@@ -538,6 +538,62 @@ describe("plugin Runner release/resume facade", () => {
? { kind: "released", taskId: "task" }
: { kind: "resume", taskId: "task", outcome: "resumed" });
it.each(["codex", "claude"])("quiesces %s before release and closes only after its ACK", async harness => {
const { task } = provider();
const perform = task.getMockImplementation()!;
const order: string[] = [];
let acknowledge!: (value: unknown) => void;
task.mockImplementation(op => {
if (op.kind !== "release") return perform(op);
order.push("release");
return new Promise(resolve => { acknowledge = resolve; });
});
const quiesce = vi.fn(() => { order.push("quiesce"); });
const close = vi.fn(async () => { order.push("close"); });
ingress.connect.mockReset().mockReturnValue({ ready: Promise.resolve(), quiesce, close });
const owner = createPluginRunnerTaskOwner({ ...input(harness), task });
const registration = await owner.registration({} as never, "run");
registration.activate(); await registration.ready();
try {
const releasing = owner.driver.releaseResume!.release(requirements.identity);
expect(order).toEqual(["quiesce", "release"]);
expect(close).not.toHaveBeenCalled();
acknowledge({ kind: "released", taskId: "task" });
await releasing;
expect(order).toEqual(["quiesce", "release", "close"]);
expect(task.mock.calls.at(-1)![0]).toEqual({ kind: "release", identity: requirements.identity });
expect(ingress.connect).toHaveBeenCalledOnce();
} finally { await registration.release("detach"); }
});
it.each(["codex", "claude"])("keeps an uncertain %s release quiet for an exact retry", async harness => {
const { task } = provider();
const perform = task.getMockImplementation()!;
const lost = new Error("fixture release reply lost");
let attempts = 0;
task.mockImplementation(async op => {
if (op.kind !== "release") return perform(op);
if (attempts++ === 0) throw lost;
return { kind: "released", taskId: "task" };
});
const quiesce = vi.fn(), close = vi.fn(async () => {});
ingress.connect.mockReset().mockReturnValue({ ready: Promise.resolve(), quiesce, close });
const owner = createPluginRunnerTaskOwner({ ...input(harness), task });
const registration = await owner.registration({} as never, "run");
registration.activate(); await registration.ready();
try {
await expect(owner.driver.releaseResume!.release(requirements.identity)).rejects.toBe(lost);
expect(quiesce).toHaveBeenCalledOnce();
expect(close).not.toHaveBeenCalled();
await owner.driver.releaseResume!.release(requirements.identity);
expect(quiesce).toHaveBeenCalledTimes(2);
expect(close).toHaveBeenCalledOnce();
expect(task.mock.calls.filter(([op]) => op.kind === "release").map(([op]) => op))
.toEqual([{ kind: "release", identity: requirements.identity }, { kind: "release", identity: requirements.identity }]);
expect(ingress.connect).toHaveBeenCalledOnce();
} finally { await registration.release("detach"); }
});
it.each(["codex", "claude"])("forwards both %s operations through the current checked owner", async harness => {
const firstTask = lifecycleTask();
const currentTask = lifecycleTask();
@@ -741,3 +797,44 @@ it.each(["codex", "claude"])("refuses warm %s account, harness, project, or cred
await expect(prepare(owner)).rejects.toThrow("credential broker v1");
expect(task.mock.calls.map(([op])=>op.kind)).toEqual(["describe"]);
});
it.each((["codex", "claude"] as const).flatMap(harness => ["prepared", "unprepared", "foreign target"].map(mode => [harness, mode] as const)))(
"authenticates adopted %s source before activating only its exact %s successor", async (harness, mode) => {
const { task } = provider(), perform = task.getMockImplementation()!;
const source = launchRequirements(harness), target = { ...source, identity: { ...source.identity, runId: "successor", turnId: "next-turn", itemId: "next-item" } };
let nativeIdentity = source.identity;
task.mockImplementation(async operation => operation.kind === "read_state"
? { kind: "state", taskId: "task", state: { schema: "paperclip.runner.durable.state.v1", lifecycle: "ready", ...nativeIdentity } }
: perform(operation));
const connect = vi.fn(() => ({ ready: Promise.resolve(), close: async () => {} }));
ingress.connect.mockImplementation(connect);
const owner = createPluginRunnerTaskOwner({ taskId: "task", task, projectIds: [], harness, authorityRunId: target.identity.runId,
projectAuthCheckVersion: 1, accountIdentity: "account", adoptIdentity: source.identity });
const initial = await owner.driver.controlPlaneRegistration!({} as never, source.identity);
await initial.activate!(); await initial.ready!();
expect(task.mock.calls.some(([operation]) => operation.kind === "read_state" && operation.identity.runId === source.identity.runId)).toBe(true);
const requested = mode === "foreign target" ? { ...target.identity, itemId: "unprepared-item" } : target.identity;
const next = await owner.driver.controlPlaneRegistration!({} as never, requested);
if (mode !== "unprepared") await prepare(owner, target);
nativeIdentity = target.identity;
try {
await next.activate!();
if (mode === "prepared") { await expect(next.ready!()).resolves.toBeUndefined(); expect(connect).toHaveBeenCalledTimes(2); }
else { await expect(next.ready!()).rejects.toThrow("adoption identity mismatch"); expect(connect).toHaveBeenCalledTimes(1); }
} finally { await next.release!("detach"); await initial.release!(); }
},
);
it.each(["codex", "claude"] as const)("rejects a foreign %s source tuple before opening adopted ingress", async harness => {
const { task } = provider();
const connect = vi.fn(() => ({ ready: Promise.resolve(), close: async () => {} }));
ingress.connect.mockImplementation(connect);
const source = launchRequirements(harness).identity;
const owner = createPluginRunnerTaskOwner({ taskId: "task", task, projectIds: [], harness, adoptIdentity: source });
const registration = await owner.driver.controlPlaneRegistration!({} as never, { ...source, itemId: "foreign-item" });
try {
await registration.activate!();
await expect(registration.ready!()).rejects.toThrow("adoption identity mismatch");
expect(task).not.toHaveBeenCalled(); expect(connect).not.toHaveBeenCalled();
} finally { await registration.release!("detach"); }
});
@@ -48,6 +48,9 @@ export function createPluginRunnerTaskOwner(input: {
pollMs?: number;
warmTaskVersion?: 1;
releaseResumeVersion?: 1;
releasedSource?: DurableRecoveryIdentity;
onResumed?: (identity: DurableRecoveryIdentity) => Promise<void>;
onRetired?: (identity: DurableRecoveryIdentity) => Promise<void>;
githubBroker?: { url: string; token: string };
adoptIdentity?: DurableRecoveryIdentity;
}) {
@@ -268,7 +271,8 @@ export function createPluginRunnerTaskOwner(input: {
});
return { child, completion, restart: null };
};
const registration = async (authority: DurablePrpControlPlane, runId: string, signal?: AbortSignal) => {
const registration = async (authority: DurablePrpControlPlane, runId: string, signal?: AbortSignal, expectedIdentity?: DurableRecoveryIdentity) => {
const desired = expectedIdentity ? structuredClone(expectedIdentity) : undefined;
const generation = ++registrationGeneration;
let registrationReleased = false;
let outbound: ReturnType<typeof connectRunnerPrpIngress> | null = null;
@@ -284,7 +288,15 @@ export function createPluginRunnerTaskOwner(input: {
activate: () => {
activation = (async () => {
if (input.adoptIdentity) {
await verifyPluginRunnerTaskReattachment(input.task, input.adoptIdentity);
if (!desired || runId === input.adoptIdentity.runId) {
if (desired && Object.entries(input.adoptIdentity).some(([key, value]) => desired[key as keyof DurableRecoveryIdentity] !== value)) {
throw new Error("Plugin Runner adoption identity mismatch");
}
await verifyPluginRunnerTaskReattachment(input.task, input.adoptIdentity);
} else if (!preparedIdentity || Object.entries(preparedIdentity).some(([key, value]) => desired[key as keyof DurableRecoveryIdentity] !== value)
|| ["runnerInstanceId", "environmentLeaseId", "normalizedSessionId"].some(key => desired[key as keyof DurableRecoveryIdentity] !== input.adoptIdentity![key as keyof DurableRecoveryIdentity])) {
throw new Error("Plugin Runner adoption identity mismatch");
}
} else await running;
if (registrationReleased || signal?.aborted) throw new Error("Plugin Runner startup cancelled");
const ingress = await endpoint();
@@ -318,6 +330,8 @@ export function createPluginRunnerTaskOwner(input: {
if (submitted) {
await call({ kind: "complete" });
completed = true;
const identity = preparedIdentity ?? launchIdentity ?? input.priorIdentity;
if (identity) await input.onRetired?.(identity);
}
},
release: async (reason?: "detach") => {
@@ -330,6 +344,25 @@ export function createPluginRunnerTaskOwner(input: {
};
};
const driver: RunnerExecutionDriver = {
get releaseResume() {
if (input.releaseResumeVersion !== 1 || input.warmTaskVersion !== 1 || !input.accountIdentity) return undefined;
return {
releasedSource: input.releasedSource,
release: async (identity: DurableRecoveryIdentity) => {
requireReleaseResume();
const outbound = activeOutbound;
outbound?.quiesce();
await release(identity);
await outbound?.close();
if (activeOutbound === outbound) activeOutbound = null;
},
resume: async (identity: DurableRecoveryIdentity, ticket: string) => {
const outcome = await resume(identity, ticket);
if (outcome === "resumed") await input.onResumed?.(identity);
return outcome;
},
};
},
get authorityRunId() { return input.authorityRunId ?? input.executionContext?.execution.binding.runId; },
get codexSourceRecoveryVersion() { return input.codexSourceRecoveryVersion; },
prepareLaunch, authorizeRunnerAdoption, readRunnerState, runnerProcessLauncher: launcher,
@@ -341,7 +374,7 @@ export function createPluginRunnerTaskOwner(input: {
if (result.kind !== "state") throw new Error("Plugin Runner archived state unavailable");
return result.state;
},
controlPlaneRegistration: (authority, identity, signal) => registration(authority, identity?.runId ?? "", signal),
controlPlaneRegistration: (authority, identity, signal) => registration(authority, identity?.runId ?? "", signal, identity),
...(input.adoptIdentity ? { adoptExistingRunner: {
// Provider ownership is attested by Launcher, never by a controller-local PID.
pid: null, processGroupId: null, startedAt: null,