diff --git a/src/agents/main-session-recovery-store.test.ts b/src/agents/main-session-recovery-store.test.ts index 9260b7cd6b1d..0f93ff5d1588 100644 --- a/src/agents/main-session-recovery-store.test.ts +++ b/src/agents/main-session-recovery-store.test.ts @@ -403,7 +403,10 @@ describe("main session recovery store", () => { }, ); - const immediateRelease = releaseMainSessionRecoveryOwner(claim.lease); + const onDeferredSuccess = vi.fn(); + const immediateRelease = releaseMainSessionRecoveryOwner(claim.lease, { + onDeferredSuccess, + }); const immediateReleaseRejected = expect(immediateRelease).rejects.toThrow( "transient session-store failure", ); @@ -415,6 +418,11 @@ describe("main session recovery store", () => { await vi.waitFor(() => { expect(read().mainRestartRecovery?.foregroundClaims).toBeUndefined(); }); + expect(onDeferredSuccess).toHaveBeenCalledWith({ + sessionId: "session-1", + sessionKey, + storePath, + }); } finally { vi.useRealTimers(); } diff --git a/src/agents/main-session-recovery-store.ts b/src/agents/main-session-recovery-store.ts index 350463345982..2124eb4252ce 100644 --- a/src/agents/main-session-recovery-store.ts +++ b/src/agents/main-session-recovery-store.ts @@ -339,23 +339,35 @@ async function releaseMainSessionRecoveryOwnerWithRetries( return { sessionId: entry.sessionId, sessionKey, storePath: lease.storePath }; } -function scheduleMainSessionRecoveryOwnerRelease(lease: MainSessionRecoveryOwnerLease): void { +function scheduleMainSessionRecoveryOwnerRelease( + lease: MainSessionRecoveryOwnerLease, + onDeferredSuccess?: ( + pending: MainSessionRecoveryPendingTarget | undefined, + ) => void | Promise, +): void { // A token is process-owned but durably blocks recovery. Keep exact-token // cleanup alive through transient writer outages until release or restart. scheduleMainSessionRecoveryMutation({ mutation: () => releaseMainSessionRecoveryOwnerWithRetries(lease), - onSuccess: async (pending) => { - if (pending) { - const { scheduleMainSessionRecoveryPendingTarget } = - await import("./main-session-recovery-owner-release.js"); - scheduleMainSessionRecoveryPendingTarget(pending); - } - }, + onSuccess: + onDeferredSuccess ?? + (async (pending) => { + if (pending) { + const { scheduleMainSessionRecoveryPendingTarget } = + await import("./main-session-recovery-owner-release.js"); + scheduleMainSessionRecoveryPendingTarget(pending); + } + }), }); } export async function releaseMainSessionRecoveryOwner( lease: MainSessionRecoveryOwnerLease | undefined, + options?: { + onDeferredSuccess?: ( + pending: MainSessionRecoveryPendingTarget | undefined, + ) => void | Promise; + }, ): Promise { if (!lease) { return undefined; @@ -363,7 +375,7 @@ export async function releaseMainSessionRecoveryOwner( try { return await releaseMainSessionRecoveryOwnerWithRetries(lease); } catch (error) { - scheduleMainSessionRecoveryOwnerRelease(lease); + scheduleMainSessionRecoveryOwnerRelease(lease, options?.onDeferredSuccess); throw error; } } diff --git a/src/auto-reply/reply/reply-run-registry.test.ts b/src/auto-reply/reply/reply-run-registry.test.ts index 250478d7ffb3..a2766740b1d9 100644 --- a/src/auto-reply/reply/reply-run-registry.test.ts +++ b/src/auto-reply/reply/reply-run-registry.test.ts @@ -28,7 +28,9 @@ import { clearReplyRunForResetBySessionId, REPLY_RUN_IDLE_SETTLE_TIMEOUT_MS, REPLY_RUN_TERMINAL_SETTLE_TIMEOUT_MS, + registerReplyOperationSuccessorBarrier, ReplyRunAlreadyActiveError, + ReplyRunSuccessorAdmissionBlockedError, replyRunRegistry, markReplyOperationGlobalLaneWaitProgress, runAfterReplyOperationClear, @@ -37,6 +39,7 @@ import { resolveReplyRunPhaseForSessionId, waitForReplyOperationOwnerSettlement, waitForReplyRunEndBySessionId, + waitForReplyRunSuccessorAdmission, } from "./reply-run-registry.js"; import { testing } from "./reply-run-registry.test-support.js"; import { admitReplyTurn } from "./reply-turn-admission.js"; @@ -851,6 +854,92 @@ describe("reply run registry", () => { } }); + it("fences every durable alias until successor handoff settles", async () => { + vi.useFakeTimers(); + try { + const requestKey = "agent:main:telegram:alias:request"; + const canonicalKey = "agent:main:telegram:alias:canonical"; + const adoptedKey = "agent:main:telegram:alias:adopted"; + const operation = createTestReplyOperation({ + sessionKey: requestKey, + sessionId: "alias-session", + }); + let releaseFirstBarrier = () => {}; + const firstBarrier = new Promise((resolve) => { + releaseFirstBarrier = resolve; + }); + registerReplyOperationSuccessorBarrier({ + operation, + sessionId: "alias-session", + sessionKeys: [requestKey, canonicalKey], + start: () => firstBarrier, + }); + let releaseSecondBarrier = () => {}; + const secondBarrier = new Promise((resolve) => { + releaseSecondBarrier = resolve; + }); + registerReplyOperationSuccessorBarrier({ + operation, + sessionId: "alias-session", + sessionKeys: [adoptedKey], + start: () => secondBarrier, + }); + + operation.updateSessionId("rotated-alias-session"); + operation.complete(); + for (const sessionKey of [requestKey, canonicalKey, adoptedKey]) { + expect(() => createTestReplyOperation({ sessionKey })).toThrow( + ReplyRunSuccessorAdmissionBlockedError, + ); + } + const timedWait = waitForReplyRunSuccessorAdmission(canonicalKey, 100); + await vi.advanceTimersByTimeAsync(100); + await expect(timedWait).resolves.toEqual({ settled: false }); + + const requestWait = waitForReplyRunSuccessorAdmission(requestKey, 100); + const canonicalWait = waitForReplyRunSuccessorAdmission(canonicalKey, 100); + releaseFirstBarrier(); + for (const wait of [requestWait, canonicalWait]) { + await expect(wait).resolves.toEqual({ + settled: true, + sessionId: "rotated-alias-session", + }); + } + expect(() => createTestReplyOperation({ sessionKey: adoptedKey })).toThrow( + ReplyRunSuccessorAdmissionBlockedError, + ); + releaseSecondBarrier(); + await expect(waitForReplyRunSuccessorAdmission(adoptedKey, 100)).resolves.toEqual({ + settled: true, + sessionId: "rotated-alias-session", + }); + const successor = createTestReplyOperation({ sessionKey: canonicalKey }); + successor.complete(); + } finally { + await vi.runOnlyPendingTimersAsync(); + vi.useRealTimers(); + } + }); + + it("stops a successor wait when its signal aborts", async () => { + const operation = createTestReplyOperation(); + registerReplyOperationSuccessorBarrier({ + operation, + sessionId: operation.sessionId, + sessionKeys: [operation.key], + start: () => new Promise(() => {}), + }); + operation.complete(); + const controller = new AbortController(); + const wait = waitForReplyRunSuccessorAdmission(operation.key, null, { + signal: controller.signal, + }); + + controller.abort(); + + await expect(wait).resolves.toEqual({ settled: false }); + }); + it("extends a hung delivery barrier only while bounded owner work remains active", async () => { vi.useFakeTimers(); try { diff --git a/src/auto-reply/reply/reply-run-registry.ts b/src/auto-reply/reply/reply-run-registry.ts index 7cafc42461de..100cec5d8fcf 100644 --- a/src/auto-reply/reply/reply-run-registry.ts +++ b/src/auto-reply/reply/reply-run-registry.ts @@ -311,7 +311,7 @@ type ReplyRunWaiter = { timer?: NodeJS.Timeout; }; -type ReplyRunFollowupAdmissionBarrier = { +type ReplyRunAdmissionBarrier = { settled: Promise; sessionId: string; }; @@ -322,7 +322,8 @@ type ReplyRunState = { activeKeysBySessionId: Map; waitKeysBySessionId: Map; waitersByKey: Map>; - followupAdmissionBarriersByKey: Map; + followupAdmissionBarriersByKey: Map; + successorAdmissionBarriersByKey: Map; evictOperationByOperation?: WeakMap void>; }; @@ -334,10 +335,12 @@ const replyRunState = resolveGlobalSingleton(REPLY_RUN_STATE_KEY, activeKeysBySessionId: new Map(), waitKeysBySessionId: new Map(), waitersByKey: new Map>(), - followupAdmissionBarriersByKey: new Map(), + followupAdmissionBarriersByKey: new Map(), + successorAdmissionBarriersByKey: new Map(), evictOperationByOperation: new WeakMap void>(), })); replyRunState.followupAdmissionBarriersByKey ??= new Map(); +replyRunState.successorAdmissionBarriersByKey ??= new Map(); const evictReplyOperationByOperation = replyRunState.evictOperationByOperation ?? (replyRunState.evictOperationByOperation = new WeakMap void>()); @@ -363,6 +366,13 @@ export class ReplyRunFollowupAdmissionBlockedError extends Error { } } +export class ReplyRunSuccessorAdmissionBlockedError extends Error { + constructor(sessionKey: string) { + super(`Reply successor admission is blocked for ${sessionKey}`); + this.name = "ReplyRunSuccessorAdmissionBlockedError"; + } +} + function createUserAbortError(): Error { return createAbortError("Reply operation aborted by user"); } @@ -440,6 +450,18 @@ const afterClearCallbacksByOperation = new WeakMap< ReplyOperation, Set<(sessionId: string) => void> >(); +const successorBarrierStartsByOperation = new WeakMap void>>(); +type ReplyOperationSuccessorBarrierGroup = { + registrationKey: string; + barriers: Set; +}; +// Alias-keyed fences registered for one lane rotate together. Rekeyed command +// operations retain prior-lane identities so source successors do not adopt +// the target session. +const successorBarrierGroupsByOperation = new WeakMap< + ReplyOperation, + Set +>(); type ReplyOperationStaleExpiryOptions = { afterClearBarrier?: PromiseLike; followupAdmissionBarrierTimeout?: number | ReplyFollowupAdmissionBarrierTimeoutPolicy; @@ -566,6 +588,104 @@ export function runAfterReplyOperationClear( afterClearCallbacksByOperation.set(operation, callbacks); } +function registerSuccessorAdmissionBarrier( + sessionKey: string, + sessionId: string, + barrier: Promise, +): ReplyRunAdmissionBarrier { + const barriersByKey = replyRunState.successorAdmissionBarriersByKey; + const previous = barriersByKey.get(sessionKey)?.settled; + const settled = previous ? Promise.all([previous, barrier]).then(() => undefined) : barrier; + const entry = { settled, sessionId }; + barriersByKey.set(sessionKey, entry); + void settled.then(() => { + if (barriersByKey.get(sessionKey) === entry) { + barriersByKey.delete(sessionKey); + } + }); + return entry; +} + +/** Fence successor admission until owner handoff started at slot clear settles. */ +export function registerReplyOperationSuccessorBarrier(params: { + operation: ReplyOperation; + sessionId: string; + sessionKeys: readonly string[]; + start: () => PromiseLike; +}): void { + const settlement = createDeferred(); + const barriers = new Set(); + for (const sessionKey of new Set(params.sessionKeys.map(normalizeOptionalString))) { + if (sessionKey) { + barriers.add( + registerSuccessorAdmissionBarrier(sessionKey, params.sessionId, settlement.promise), + ); + } + } + let started = false; + const start = () => { + if (started) { + return; + } + started = true; + try { + void Promise.resolve(params.start()).then( + () => settlement.resolve(undefined), + () => {}, + ); + } catch { + // A failed handoff leaves the fence closed. Visible callers stay + // abortably blocked; bounded queued callers cannot observe a partial release. + } + }; + if (replyRunState.activeRunsByKey.get(params.operation.key) !== params.operation) { + start(); + return; + } + const groups = + successorBarrierGroupsByOperation.get(params.operation) ?? + new Set(); + groups.add({ registrationKey: params.operation.key, barriers }); + successorBarrierGroupsByOperation.set(params.operation, groups); + const starts = successorBarrierStartsByOperation.get(params.operation) ?? new Set<() => void>(); + starts.add(start); + successorBarrierStartsByOperation.set(params.operation, starts); +} + +function startReplyOperationSuccessorBarriers(operation: ReplyOperation): void { + const starts = successorBarrierStartsByOperation.get(operation); + // These maps are operation-owned lifecycle metadata, not identity indexes. + // Clear drops both before handoff starts so adoption cannot retain stale groups. + successorBarrierStartsByOperation.delete(operation); + successorBarrierGroupsByOperation.delete(operation); + if (!starts) { + return; + } + for (const start of starts) { + start(); + } +} + +function updateSuccessorAdmissionSessionId(operation: ReplyOperation, sessionId: string): void { + for (const group of successorBarrierGroupsByOperation.get(operation) ?? []) { + if (group.registrationKey !== operation.key) { + continue; + } + for (const barrier of group.barriers) { + barrier.sessionId = sessionId; + } + } +} + +export function isReplyRunSuccessorAdmissionBlocked(sessionKey: string): boolean { + const normalizedSessionKey = normalizeOptionalString(sessionKey); + return Boolean( + normalizedSessionKey && + !replyRunState.activeRunsByKey.has(normalizedSessionKey) && + replyRunState.successorAdmissionBarriersByKey.has(normalizedSessionKey), + ); +} + function flushReplyOperationAfterClear(operation: ReplyOperation, sessionId: string): void { const callbacks = afterClearCallbacksByOperation.get(operation); if (!callbacks) { @@ -635,7 +755,7 @@ function registerFollowupAdmissionBarrier( sessionId: string, barrier: PromiseLike, timeout: number | ReplyFollowupAdmissionBarrierTimeoutPolicy = REPLY_RUN_IDLE_SETTLE_TIMEOUT_MS, -): ReplyRunFollowupAdmissionBarrier { +): ReplyRunAdmissionBarrier { const barriersByKey = replyRunState.followupAdmissionBarriersByKey; const previous = barriersByKey.get(sessionKey)?.settled; const current = waitForReplyBarrierSettlement(barrier, timeout); @@ -718,6 +838,9 @@ export function createReplyOperation(params: { if (replyRunState.activeRunsByKey.has(sessionKey)) { throw new ReplyRunAlreadyActiveError(sessionKey); } + if (replyRunState.successorAdmissionBarriersByKey.has(sessionKey)) { + throw new ReplyRunSuccessorAdmissionBlockedError(sessionKey); + } const controller = new AbortController(); // Mutable so updateSessionKey can move the run slot (command-turn continuation @@ -730,7 +853,7 @@ export function createReplyOperation(params: { let result: ReplyOperationResult | null = null; let stateCleared = false; let clearBarrierSettlement: Promise | undefined; - let pendingClearBarrier: ReplyRunFollowupAdmissionBarrier | undefined; + let pendingClearBarrier: ReplyRunAdmissionBarrier | undefined; let retainFailureUntilComplete = false; let terminalRecovery = false; let acceptedSteeredInboundAudio = false; @@ -787,6 +910,9 @@ export function createReplyOperation(params: { : pendingClearBarrier; pendingClearBarrier = undefined; updateFollowupAdmissionSessionId(currentSessionKey, currentSessionId); + // Recovery-owner handoff must begin before the old slot wakes a successor; + // otherwise that successor can snapshot durable state the handoff then mutates. + startReplyOperationSuccessorBarriers(operation); markReplyRunDiagnosticProgress({ sessionKey: currentSessionKey, sessionId: currentSessionId, @@ -965,6 +1091,7 @@ export function createReplyOperation(params: { currentSessionId = normalizedNextSessionId; ownedSessionIds.add(currentSessionId); updateFollowupAdmissionSessionId(currentSessionKey, currentSessionId); + updateSuccessorAdmissionSessionId(operation, currentSessionId); replyRunState.activeSessionIdsByKey.set(currentSessionKey, currentSessionId); replyRunState.activeKeysBySessionId.set(currentSessionId, currentSessionKey); registerWaitSessionId(currentSessionKey, currentSessionId); @@ -990,6 +1117,9 @@ export function createReplyOperation(params: { if (replyRunState.activeRunsByKey.has(normalizedNextKey)) { throw new ReplyRunAlreadyActiveError(normalizedNextKey); } + if (replyRunState.successorAdmissionBarriersByKey.has(normalizedNextKey)) { + throw new ReplyRunSuccessorAdmissionBlockedError(normalizedNextKey); + } recordActivity(); const previousKey = currentSessionKey; replyRunState.activeRunsByKey.delete(previousKey); @@ -1613,43 +1743,51 @@ export function waitForReplyRunEndBySessionId( return replyRunRegistry.waitForIdle(waitKey, timeoutMs); } -export async function waitForReplyRunFollowupAdmission( - sessionKey: string, - timeoutMs: number, - opts?: { signal?: AbortSignal }, -): Promise<{ settled: boolean; sessionId?: string }> { - const normalizedSessionKey = normalizeOptionalString(sessionKey); - if (!normalizedSessionKey) { - return { settled: true }; - } - const resolvedTimeoutMs = resolveTimerTimeoutMs(timeoutMs, 100, 100); - const deadline = Date.now() + resolvedTimeoutMs; +async function waitForReplyRunAdmissionBarrier(params: { + barriersByKey: Map; + minimumTimeoutMs: number; + sessionKey: string; + signal?: AbortSignal; + timeoutMs?: number | null; +}): Promise<{ settled: boolean; sessionId?: string }> { + const deadline = + typeof params.timeoutMs === "number" + ? Date.now() + + resolveTimerTimeoutMs(params.timeoutMs, params.minimumTimeoutMs, params.minimumTimeoutMs) + : undefined; let sessionId: string | undefined; while (true) { - if (opts?.signal?.aborted) { + if (params.signal?.aborted) { return { settled: false }; } - const barrier = replyRunState.followupAdmissionBarriersByKey.get(normalizedSessionKey); + const barrier = params.barriersByKey.get(params.sessionKey); if (!barrier) { return { settled: true, sessionId }; } - const remainingMs = deadline - Date.now(); - if (remainingMs <= 0) { + const remainingMs = deadline === undefined ? undefined : deadline - Date.now(); + if (remainingMs !== undefined && remainingMs <= 0) { return { settled: false }; } let timer: NodeJS.Timeout | undefined; let abortHandler: (() => void) | undefined; const outcome = await Promise.race([ barrier.settled.then(() => true), - new Promise((resolve) => { - timer = setTimeout(() => resolve(false), remainingMs); - timer.unref?.(); - }), - ...(opts?.signal + ...(remainingMs !== undefined + ? [ + new Promise((resolve) => { + timer = setTimeout(() => resolve(false), Math.max(1, remainingMs)); + timer.unref?.(); + }), + ] + : []), + ...(params.signal ? [ new Promise((resolve) => { abortHandler = () => resolve(false); - opts.signal?.addEventListener("abort", abortHandler, { once: true }); + params.signal?.addEventListener("abort", abortHandler, { once: true }); + if (params.signal?.aborted) { + abortHandler(); + } }), ] : []), @@ -1658,7 +1796,7 @@ export async function waitForReplyRunFollowupAdmission( clearTimeout(timer); } if (abortHandler) { - opts?.signal?.removeEventListener("abort", abortHandler); + params.signal?.removeEventListener("abort", abortHandler); } if (!outcome) { return { settled: false }; @@ -1667,6 +1805,40 @@ export async function waitForReplyRunFollowupAdmission( } } +export async function waitForReplyRunFollowupAdmission( + sessionKey: string, + timeoutMs: number, + opts?: { signal?: AbortSignal }, +): Promise<{ settled: boolean; sessionId?: string }> { + const normalizedSessionKey = normalizeOptionalString(sessionKey); + return normalizedSessionKey + ? await waitForReplyRunAdmissionBarrier({ + barriersByKey: replyRunState.followupAdmissionBarriersByKey, + minimumTimeoutMs: 100, + sessionKey: normalizedSessionKey, + signal: opts?.signal, + timeoutMs, + }) + : { settled: true }; +} + +export async function waitForReplyRunSuccessorAdmission( + sessionKey: string, + timeoutMs?: number | null, + opts?: { signal?: AbortSignal }, +): Promise<{ settled: boolean; sessionId?: string }> { + const normalizedSessionKey = normalizeOptionalString(sessionKey); + return normalizedSessionKey + ? await waitForReplyRunAdmissionBarrier({ + barriersByKey: replyRunState.successorAdmissionBarriersByKey, + minimumTimeoutMs: 0, + sessionKey: normalizedSessionKey, + signal: opts?.signal, + timeoutMs, + }) + : { settled: true }; +} + export function abortActiveReplyRuns(opts: { mode: "all" | "compacting"; onAbortError?: (sessionId: string, error: unknown) => void; @@ -1784,6 +1956,7 @@ const replyRunRegistryTestApi = { } replyRunState.waitersByKey.clear(); replyRunState.followupAdmissionBarriersByKey.clear(); + replyRunState.successorAdmissionBarriersByKey.clear(); }, }; diff --git a/src/auto-reply/reply/reply-turn-admission.test.ts b/src/auto-reply/reply/reply-turn-admission.test.ts index 18d6bc00877a..3a73496b9c4d 100644 --- a/src/auto-reply/reply/reply-turn-admission.test.ts +++ b/src/auto-reply/reply/reply-turn-admission.test.ts @@ -2,6 +2,7 @@ import path from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; import { useAutoCleanupTempDirTracker } from "../../../test/helpers/temp-dir.js"; +import * as sessionAccessor from "../../config/sessions/session-accessor.js"; import { deleteSessionEntryLifecycle, loadSessionEntry, @@ -30,9 +31,25 @@ import { testing } from "./reply-run-registry.test-support.js"; import { admitReplyTurn, runWithReplyOperationLifecycleAdmission } from "./reply-turn-admission.js"; const recoveryOwnerReleaseMocks = vi.hoisted(() => ({ + beforeRelease: vi.fn(async () => {}), schedulePendingTarget: vi.fn(), })); +vi.mock("../../agents/main-session-recovery-store.js", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + releaseMainSessionRecoveryOwner: async ( + lease: Parameters[0], + options: Parameters[1], + ) => { + await recoveryOwnerReleaseMocks.beforeRelease(); + return await actual.releaseMainSessionRecoveryOwner(lease, options); + }, + }; +}); + vi.mock("../../agents/main-session-recovery-owner-release.js", async (importOriginal) => ({ ...(await importOriginal()), scheduleMainSessionRecoveryPendingTarget: recoveryOwnerReleaseMocks.schedulePendingTarget, @@ -84,6 +101,7 @@ describe("reply turn admission", () => { afterEach(() => { testing.resetReplyRunRegistry(); resetDiagnosticRunActivityForTest(); + recoveryOwnerReleaseMocks.beforeRelease.mockClear(); recoveryOwnerReleaseMocks.schedulePendingTarget.mockClear(); }); @@ -407,6 +425,229 @@ describe("reply turn admission", () => { }, ); + it.each(["visible", "queued_followup"] as const)( + "waits for restart-recovery owner release before %s successor admission", + async (kind) => { + const sessionKey = `agent:main:telegram:topic:recovery-successor:${kind}`; + const sessionId = "interrupted-session"; + const storePath = createSessionStore({ + [sessionKey]: { + sessionId, + updatedAt: 100, + status: "running", + abortedLastRun: true, + mainRestartRecovery: { + cycleId: "cycle-1", + revision: 1, + chargedAttempts: 0, + }, + }, + }); + const owner = await admitTestReplyTurn({ + sessionKey, + sessionId, + expectedSessionId: sessionId, + storePath, + }); + expect(owner.status).toBe("owned"); + if (owner.status !== "owned") { + return; + } + + const releaseStarted = createDeferred(); + const allowRelease = createDeferred(); + recoveryOwnerReleaseMocks.beforeRelease.mockImplementationOnce(async () => { + releaseStarted.resolve(); + await allowRelease.promise; + }); + owner.operation.complete(); + await releaseStarted.promise; + + const successor = admitTestReplyTurn({ + sessionKey, + sessionId, + expectedSessionId: sessionId, + storePath, + kind, + }); + let successorSettled = false; + void successor.then(() => { + successorSettled = true; + }); + await Promise.resolve(); + expect(successorSettled).toBe(false); + await expect( + admitTestReplyTurn({ + sessionKey, + sessionId, + expectedSessionId: sessionId, + storePath, + kind: "heartbeat", + }), + ).resolves.toEqual({ status: "skipped", reason: "active-run" }); + + allowRelease.resolve(); + const admitted = await successor; + expect(admitted.status).toBe("owned"); + if (admitted.status === "owned") { + admitted.operation.complete(); + } + await vi.waitFor(async () => { + const entry = await readSessionEntry(storePath, sessionKey); + expect(entry?.mainRestartRecovery?.foregroundClaims).toBeUndefined(); + }); + }, + ); + + it("waits through deferred owner release retries beyond one settle slice", async () => { + vi.useFakeTimers(); + try { + const sessionKey = "agent:main:telegram:topic:deferred-recovery-release"; + const sessionId = "interrupted-session"; + const storePath = createSessionStore({ + [sessionKey]: { + sessionId, + updatedAt: 100, + status: "running", + abortedLastRun: true, + mainRestartRecovery: { + cycleId: "cycle-1", + revision: 1, + chargedAttempts: 0, + }, + }, + }); + const owner = await admitTestReplyTurn({ + sessionKey, + sessionId, + expectedSessionId: sessionId, + storePath, + }); + expect(owner.status).toBe("owned"); + if (owner.status !== "owned") { + return; + } + const applySessionEntryReplacements = sessionAccessor.applySessionEntryReplacements; + let failures = 0; + vi.spyOn(sessionAccessor, "applySessionEntryReplacements").mockImplementation( + async (params) => { + if (failures < 15) { + failures += 1; + throw new Error("transient session-store failure"); + } + return await applySessionEntryReplacements(params); + }, + ); + + owner.operation.complete(); + const successor = admitTestReplyTurn({ + sessionKey, + sessionId, + expectedSessionId: sessionId, + storePath, + }); + let successorSettled = false; + void successor.then(() => { + successorSettled = true; + }); + await vi.advanceTimersByTimeAsync(REPLY_RUN_IDLE_SETTLE_TIMEOUT_MS + 1); + expect(successorSettled).toBe(false); + + await vi.advanceTimersByTimeAsync(20_000); + const admitted = await successor; + expect(admitted.status).toBe("owned"); + if (admitted.status === "owned") { + admitted.operation.complete(); + } + } finally { + await vi.runOnlyPendingTimersAsync(); + vi.useRealTimers(); + } + }); + + it("preserves a source recovery identity after adopting a distinct target session", async () => { + const sourceSessionKey = "agent:main:telegram:slash:recovery-source"; + const sourceSessionId = "recovery-source-session"; + const targetSessionKey = "agent:main:telegram:group:recovery-target"; + const targetSessionId = "recovery-target-session"; + const storePath = createSessionStore({ + [sourceSessionKey]: { + sessionId: sourceSessionId, + updatedAt: 100, + status: "running", + abortedLastRun: true, + mainRestartRecovery: { + cycleId: "cycle-1", + revision: 1, + chargedAttempts: 0, + }, + }, + [targetSessionKey]: { sessionId: targetSessionId, updatedAt: 100 }, + }); + const source = await admitTestReplyTurn({ + sessionKey: sourceSessionKey, + sessionId: sourceSessionId, + expectedSessionId: sourceSessionId, + storePath, + }); + expect(source.status).toBe("owned"); + if (source.status !== "owned") { + return; + } + + const adoption = await admitTestReplyTurn({ + sessionKey: targetSessionKey, + sessionId: source.operation.sessionId, + expectedSessionId: targetSessionId, + storePath, + waitForActive: false, + adoptOperation: source.operation, + }); + expect(adoption.status).toBe("owned"); + if (adoption.status !== "owned") { + source.operation.complete(); + return; + } + adoption.operation.updateSessionId(targetSessionId); + expect(adoption.operation).toBe(source.operation); + expect(adoption.operation.key).toBe(targetSessionKey); + expect(adoption.operation.sessionId).toBe(targetSessionId); + + const releaseStarted = createDeferred(); + const allowRelease = createDeferred(); + recoveryOwnerReleaseMocks.beforeRelease.mockImplementationOnce(async () => { + releaseStarted.resolve(); + await allowRelease.promise; + }); + adoption.operation.complete(); + await releaseStarted.promise; + + const successor = admitTestReplyTurn({ + sessionKey: sourceSessionKey, + sessionId: sourceSessionId, + expectedSessionId: sourceSessionId, + storePath, + }); + let successorSettled = false; + void successor.then(() => { + successorSettled = true; + }); + await Promise.resolve(); + expect(successorSettled).toBe(false); + + allowRelease.resolve(); + const admitted = await successor; + expect(admitted.status).toBe("owned"); + if (admitted.status === "owned") { + expect(admitted.operation.sessionId).toBe(sourceSessionId); + admitted.operation.complete(); + } + await vi.waitFor(async () => { + const entry = await readSessionEntry(storePath, sourceSessionKey); + expect(entry?.mainRestartRecovery?.foregroundClaims).toBeUndefined(); + }); + }); + it.each(["visible", "heartbeat"] as const)( "rejects %s reply admission for a tombstoned recovery session", async (kind) => { diff --git a/src/auto-reply/reply/reply-turn-admission.ts b/src/auto-reply/reply/reply-turn-admission.ts index 77f8a4f855bb..e1d5ecda7aab 100644 --- a/src/auto-reply/reply/reply-turn-admission.ts +++ b/src/auto-reply/reply/reply-turn-admission.ts @@ -24,16 +24,20 @@ import { import { createReplyOperation, expireStaleReplyOperation, + isReplyRunSuccessorAdmissionBlocked, isReplyRunEvidenceStale, REPLY_RUN_IDLE_SETTLE_TIMEOUT_MS, REPLY_RUN_TERMINAL_SETTLE_TIMEOUT_MS, replyRunRegistry, ReplyRunAlreadyActiveError, ReplyRunFollowupAdmissionBlockedError, + ReplyRunSuccessorAdmissionBlockedError, + registerReplyOperationSuccessorBarrier, retainReplyOperationUntilComplete, runAfterReplyOperationClear, type ReplyOperation, waitForReplyRunFollowupAdmission, + waitForReplyRunSuccessorAdmission, } from "./reply-run-registry.js"; /** Kinds of turns that compete for one reply run slot per session. */ @@ -57,11 +61,22 @@ const lifecycleAdmissionByOperation = new WeakMap { + if (!lease) { + return undefined; + } + let settleDeferredRelease: ( + pending: MainSessionRecoveryPendingTarget | undefined, + ) => void = () => {}; + const deferredRelease = new Promise((resolve) => { + settleDeferredRelease = resolve; + }); try { - return await releaseMainSessionRecoveryOwner(lease); + return await releaseMainSessionRecoveryOwner(lease, { + onDeferredSuccess: settleDeferredRelease, + }); } catch (error) { log.warn(`failed to release main-session recovery reply owner: ${formatErrorMessage(error)}`); - return undefined; + return await deferredRelease; } } @@ -175,6 +190,29 @@ async function admitReplyTurnWithWaitSignal( if (isAbortSignalAborted(params.upstreamAbortSignal)) { return { status: "skipped", reason: "aborted" }; } + if (isReplyRunSuccessorAdmissionBlocked(params.sessionKey)) { + if (params.kind === "heartbeat") { + return { status: "skipped", reason: "active-run" }; + } + const successorAdmission = await waitForAdmission(() => + waitForReplyRunSuccessorAdmission( + params.sessionKey, + params.kind === "visible" ? null : waitTimeoutMs, + { signal: params.upstreamAbortSignal }, + ), + ); + if (!successorAdmission.settled) { + return { + status: "skipped", + reason: isAbortSignalAborted(params.upstreamAbortSignal) ? "aborted" : "active-run", + }; + } + sessionId = successorAdmission.sessionId ?? sessionId; + if (expectedSessionId && successorAdmission.sessionId) { + expectedSessionId = successorAdmission.sessionId; + } + continue; + } try { const storePath = params.storePath; let operation: ReplyOperation | undefined; @@ -255,6 +293,9 @@ async function admitReplyTurnWithWaitSignal( }) : undefined; try { + if (isReplyRunSuccessorAdmissionBlocked(params.sessionKey)) { + throw new ReplyRunSuccessorAdmissionBlockedError(params.sessionKey); + } if ( storePath && !params.resetTriggered && @@ -335,10 +376,21 @@ async function admitReplyTurnWithWaitSignal( // idempotent), so both identities free on operation clear. retainReplyOperationUntilComplete(operation); lifecycleAdmissionByOperation.set(operation, admission); + let recoveryOwnerRelease: Promise | undefined; + const releaseRecoveryOwner = () => + (recoveryOwnerRelease ??= releaseReplyRecoveryOwner(recoveryOwnerLease)); + if (recoveryOwnerLease) { + registerReplyOperationSuccessorBarrier({ + operation, + sessionId: recoveryOwnerLease.sessionId, + sessionKeys: [params.sessionKey, recoveryOwnerLease.sessionKey], + start: releaseRecoveryOwner, + }); + } runAfterReplyOperationClear(operation, () => { lifecycleAdmissionByOperation.delete(operation); // Keep reset/delete behind durable owner release and its writer lock. - void releaseReplyRecoveryOwner(recoveryOwnerLease).then((pendingTarget) => { + void releaseRecoveryOwner().then((pendingTarget) => { admission.release(); scheduleMainSessionRecoveryPendingTarget(pendingTarget); }); @@ -356,6 +408,12 @@ async function admitReplyTurnWithWaitSignal( if (error instanceof QueuedFollowupLifecycleInvalidatedError) { return { status: "skipped", reason: "lifecycle-invalidated" }; } + if (error instanceof ReplyRunSuccessorAdmissionBlockedError) { + if (params.kind === "heartbeat") { + return { status: "skipped", reason: "active-run" }; + } + continue; + } if (error instanceof ReplyRunFollowupAdmissionBlockedError) { if (params.kind === "heartbeat") { return { status: "skipped", reason: "active-run" }; diff --git a/src/auto-reply/reply/restart-recovery-claim.test.ts b/src/auto-reply/reply/restart-recovery-claim.test.ts index ebef7af6d8b0..1481e728a8df 100644 --- a/src/auto-reply/reply/restart-recovery-claim.test.ts +++ b/src/auto-reply/reply/restart-recovery-claim.test.ts @@ -2,6 +2,10 @@ import path from "node:path"; import { DatabaseSync } from "node:sqlite"; import { afterEach, describe, expect, it, vi } from "vitest"; import { useAutoCleanupTempDirTracker } from "../../../test/helpers/temp-dir.js"; +import { + claimMainSessionRecoveryOwner, + releaseMainSessionRecoveryOwner, +} from "../../agents/main-session-recovery-store.js"; import { loadSessionEntry, replaceSessionEntry, @@ -9,6 +13,7 @@ import { } from "../../config/sessions/session-accessor.js"; import { resolveSqliteReadScope } from "../../config/sessions/session-accessor.sqlite-scope.js"; import type { InternalSessionEntry, SessionEntry } from "../../config/sessions/types.js"; +import { getAgentEventLifecycleGeneration } from "../../infra/agent-events.js"; import { createUserTurnTranscriptRecorder } from "../../sessions/user-turn-transcript.js"; import type { UserTurnTranscriptRecorder, @@ -356,6 +361,104 @@ describe("createReplyRestartRecoveryClaimController", () => { }); }); + it("rejects durable admission when the captured recovery owner releases", async () => { + const root = tempDirs.make("openclaw-reply-admission-owner-release-"); + const storePath = path.join(root, "sessions.json"); + const sessionKey = "agent:main:telegram:group:chat:topic:owner-release"; + const sessionId = "channel-session-id"; + const sourceTurnId = "telegram-update-new"; + const deliveryContext = { + channel: "telegram", + to: "chat", + accountId: "default", + threadId: "thread", + }; + let entry: InternalSessionEntry = { + sessionId, + updatedAt: 10, + abortedLastRun: true, + status: "running", + mainRestartRecovery: { + cycleId: "cycle-1", + revision: 1, + chargedAttempts: 0, + }, + }; + await replaceSessionEntry({ storePath, sessionKey }, entry); + const owner = await claimMainSessionRecoveryOwner({ + lifecycleGeneration: getAgentEventLifecycleGeneration(), + sessionId, + target: { sessionKey, storePath }, + }); + expect(owner.kind).toBe("claimed"); + if (owner.kind !== "claimed") { + return; + } + entry = (await updateSessionEntry({ storePath, sessionKey }, () => ({ + abortedLastRun: false, + restartRecoveryDeliveryContext: deliveryContext, + restartRecoveryDeliveryRunId: "orphaned-run", + restartRecoveryDeliverySourceRunId: "telegram-update-old", + status: "done", + }))) as InternalSessionEntry; + const sourceMessage = { + role: "user" as const, + content: "continue", + idempotencyKey: sourceTurnId, + timestamp: Date.now(), + }; + const admission = createTestAdmission({ + entryId: sourceTurnId, + sessionId, + sessionKey, + storePath, + }); + const delegate = createUserTurnTranscriptRecorder({ + message: sourceMessage, + target: { + agentId: "main", + sessionEntry: entry, + sessionId, + sessionKey, + storePath, + }, + updateMode: "none", + }); + const recorder = { + ...delegate, + getAdmissionReceipt: () => admission, + persistApproved: async ( + options?: Parameters[0], + ) => { + await releaseMainSessionRecoveryOwner(owner.lease); + return await delegate.persistApproved(options); + }, + } satisfies UserTurnTranscriptRecorder; + const controller = createReplyRestartRecoveryClaimController({ + getEntry: () => entry, + getSessionId: () => sessionId, + isRestartAbort: () => false, + resolveDeliveryContext: () => deliveryContext, + sessionKey, + setEntry: (next) => { + entry = next; + }, + sourceTurnId, + storePath, + }); + + await expect(controller.admitUserTurn(recorder)).rejects.toThrow( + "session changed before durable user-turn admission", + ); + const persisted = loadSessionEntry({ storePath, sessionKey }); + expect(persisted).not.toHaveProperty("mainRestartRecovery"); + expect(persisted).toMatchObject({ + restartRecoveryDeliveryRunId: "orphaned-run", + restartRecoveryDeliverySourceRunId: "telegram-update-old", + status: "done", + }); + }); + it("confirms an active replacement claim after SQLite lease loss", async () => { const root = tempDirs.make("openclaw-reply-restart-handoff-"); const storePath = path.join(root, "sessions.json");