diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts index 56db053920a2..608879e3f964 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts @@ -14,7 +14,7 @@ import { ProjectId, ProviderDriverKind, ProviderInstanceId, - type ProviderSessionId, + ProviderSessionId, ThreadId, } from "@t3tools/contracts"; import * as Cause from "effect/Cause"; @@ -334,7 +334,7 @@ const makeProviderAdapter = Effect.fnUntraced(function* ( readonly backgroundTasksEndWithProcess?: boolean; readonly hangSessionScopeClose?: boolean; readonly startTurn?: Effect.Effect; - readonly beforeUnload?: Effect.Effect; + readonly beforeUnload?: Effect.Effect; readonly spawnBeforeOpen?: boolean; readonly scopeCloseReached?: Deferred.Deferred; } = {}, @@ -483,7 +483,7 @@ function layerTest(input: { readonly backgroundTasksEndWithProcess?: boolean; readonly hangSessionScopeClose?: boolean; readonly startTurn?: Effect.Effect; - readonly beforeUnload?: Effect.Effect; + readonly beforeUnload?: Effect.Effect; readonly spawnBeforeOpen?: boolean; readonly scopeCloseReached?: Deferred.Deferred; readonly serverSettingsLayer?: ReturnType; @@ -4428,13 +4428,17 @@ it.effect( yield* resumeSecondThread; assert.equal((yield* Ref.get(state)).resumeCount, 4); - assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, []); + // The loaded snapshot owns cleanup even when its native row was never persisted. + assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, ["native-thread"]); yield* manager.detach({ providerSessionId, threadId: firstThreadId }); assert.isTrue(Option.isSome(yield* manager.get(providerSessionId))); assert.equal((yield* Ref.get(state)).closeCount, 0); assert.equal((yield* Ref.get(state)).interruptCount, 1); - assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, ["native-thread"]); + assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, [ + "native-thread", + "native-thread", + ]); yield* manager.detach({ providerSessionId, threadId: secondThreadId }); yield* TestClock.adjust("1 second"); @@ -4546,6 +4550,8 @@ function runIdleThreadUnloadScenario( name: string, scenario: (input: { readonly state: Ref.Ref; + readonly eventSink: EventSink.EventSinkV2["Service"]; + readonly idAllocator: IdAllocator.IdAllocatorV2["Service"]; readonly manager: ProviderSessionManager.ProviderSessionManagerV2Shape; readonly providerSessionId: ProviderSessionId; readonly threadA: ThreadId; @@ -4553,7 +4559,10 @@ function runIdleThreadUnloadScenario( readonly endTurn: (threadId: ThreadId, ordinal: number) => Effect.Effect; readonly resume: (threadId: ThreadId) => Effect.Effect; }) => Effect.Effect, - options: { readonly hasPendingBackgroundWorkForThread?: Effect.Effect } = {}, + options: { + readonly hasPendingBackgroundWorkForThread?: Effect.Effect; + readonly beforeUnload?: Effect.Effect; + } = {}, ) { return Effect.gen(function* () { const state = yield* Ref.make(emptyState); @@ -4655,6 +4664,8 @@ function runIdleThreadUnloadScenario( yield* startTurn(threadB, 1); yield* scenario({ state, + eventSink, + idAllocator, manager, providerSessionId, threadA, @@ -4667,6 +4678,7 @@ function runIdleThreadUnloadScenario( layerTest({ state, idleTimeoutMs: 1000, + ...(options.beforeUnload === undefined ? {} : { beforeUnload: options.beforeUnload }), ...(options.hasPendingBackgroundWorkForThread === undefined ? {} : { hasPendingBackgroundWorkForThread: options.hasPendingBackgroundWorkForThread }), @@ -4908,3 +4920,386 @@ it.effect( assert.isFalse(denied?.capabilities?.has("device")); }), ); + +it.effect("unloads the original native writer after its projection is rebound", () => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + yield* Effect.gen(function* () { + const eventSink = yield* EventSink.EventSinkV2; + const ids = yield* IdAllocator.IdAllocatorV2; + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const now = yield* DateTime.now; + const threadId = ThreadId.make("thread:rebound-native-writer"); + const sessionId = yield* ids.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + yield* eventSink.write({ + events: [yield* makeThreadCreatedEvent({ idAllocator: ids, threadId, now })], + }); + const runtime = yield* manager.open({ + threadId, + providerSessionId: sessionId, + modelSelection, + runtimePolicy, + }); + const loaded = yield* runtime.resumeThread({ + providerThread: makeProviderThread({ + idAllocator: ids, + threadId, + providerSessionId: sessionId, + now, + }), + threadId, + modelSelection, + runtimePolicy, + }); + const replacementSessionId = ProviderSessionId.make("session:replacement-writer"); + // The outbox runs after this rebind commits. It must unload the native + // reference in the old process, never the new row's replacement ref. + yield* eventSink.write({ + events: [ + { + id: yield* ids.allocate.event({ threadId }), + type: "provider-thread.updated", + threadId, + providerInstanceId: ProviderInstanceId.make("codex-alternate"), + occurredAt: now, + payload: { + ...loaded, + providerInstanceId: ProviderInstanceId.make("codex-alternate"), + providerSessionId: replacementSessionId, + nativeThreadRef: { + driver: CODEX_DRIVER, + nativeId: "replacement-native-ref", + strength: "strong", + }, + }, + }, + ], + }); + yield* manager.detach({ threadId, providerSessionId: sessionId }); + assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, [ + loaded.nativeThreadRef!.nativeId!, + ]); + yield* manager.detach({ threadId, providerSessionId: sessionId }); + assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, [ + loaded.nativeThreadRef!.nativeId!, + ]); + }).pipe( + Effect.provide(layerTest({ state, idleTimeoutMs: 1_000, capabilities: CodexCapabilities })), + ); + }), +); + +it.effect( + "blocks same-owner reattachment after an ambiguous idle release until closure is confirmed", + () => + Effect.gen(function* () { + const confirmed = yield* Ref.make(false); + const attempts = yield* Ref.make(0); + yield* runIdleThreadUnloadScenario( + "unconfirmed-idle-writer", + ({ state, threadA, startTurn, endTurn, resume }) => + Effect.gen(function* () { + yield* startTurn(threadA, 1); + yield* endTurn(threadA, 1); + yield* TestClock.adjust("1 second"); + assert.equal(yield* Ref.get(attempts), 1); + const resumes = (yield* Ref.get(state)).resumeCount; + const uncertain = yield* resume(threadA).pipe(Effect.exit); + assert.equal( + (yield* Ref.get(state)).resumeCount, + resumes, + "same account is not proof that release never began", + ); + assert.isTrue(Exit.isFailure(uncertain)); + assert.equal( + (yield* Ref.get(state)).closeCount, + 0, + "the busy pooled peer must survive", + ); + yield* Ref.set(confirmed, true); + yield* resume(threadA); + assert.equal((yield* Ref.get(state)).resumeCount, resumes + 1); + assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, [ + `native-${threadA}`, + ]); + const releaseAttempts = yield* Ref.get(attempts); + yield* startTurn(threadA, 2); + yield* TestClock.adjust("30 seconds"); + assert.equal( + yield* Ref.get(attempts), + releaseAttempts, + "stale idle cleanup cannot unload the reattached generation", + ); + assert.equal((yield* Ref.get(state)).closeCount, 0); + }), + { + beforeUnload: Ref.update(attempts, (n) => n + 1).pipe( + Effect.andThen( + Ref.get(confirmed).pipe( + Effect.flatMap((done) => + done + ? Effect.void + : Effect.fail( + new ProviderAdapter.ProviderAdapterProtocolError({ + driver: CODEX_DRIVER, + detail: "Native closure is unconfirmed", + }), + ), + ), + ), + ), + ), + }, + ); + }), +); + +it.effect("reuses an idle writer only after a proven refusal before native release", () => + Effect.gen(function* () { + const attempts = yield* Ref.make(0); + const refusedThread = yield* Ref.make(undefined); + yield* runIdleThreadUnloadScenario( + "proven-refusal", + ({ + state, + threadA, + providerSessionId, + startTurn, + endTurn, + resume, + eventSink: sink, + idAllocator: ids, + }) => + Effect.gen(function* () { + const now = yield* DateTime.now; + const providerThread = makeProviderThread({ + idAllocator: ids, + threadId: threadA, + providerSessionId, + now, + nativeThreadId: `native-${threadA}`, + }); + yield* Ref.set(refusedThread, providerThread); + yield* sink.write({ + events: [ + { + id: yield* ids.allocate.event({ threadId: threadA }), + type: "provider-thread.updated", + threadId: threadA, + driver: CODEX_DRIVER, + occurredAt: now, + payload: providerThread, + }, + ], + }); + yield* startTurn(threadA, 1); + yield* endTurn(threadA, 1); + yield* TestClock.adjust("1 second"); + assert.equal(yield* Ref.get(attempts), 1); + const resumes = (yield* Ref.get(state)).resumeCount; + yield* resume(threadA); + assert.equal((yield* Ref.get(state)).resumeCount, resumes); + yield* startTurn(threadA, 2); + yield* TestClock.adjust("30 seconds"); + assert.equal(yield* Ref.get(attempts), 1, "the refused generation must be cancelled"); + assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, []); + assert.equal((yield* Ref.get(state)).closeCount, 0, "the busy pooled peer must survive"); + }).pipe(Effect.orDie), + { + beforeUnload: Effect.gen(function* () { + yield* Ref.update(attempts, (n) => n + 1); + const thread = yield* Ref.get(refusedThread); + if (thread === undefined) return yield* Effect.die("Missing refused writer fixture"); + return yield* new ProviderAdapter.ProviderAdapterThreadUnloadRefusedError({ + driver: CODEX_DRIVER, + providerThreadId: thread.id, + }); + }), + }, + ); + }), +); + +it.effect("reuses the same writer after a plain detach refuses ongoing background work", () => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const background = yield* Ref.make(true); + const refusedThread = yield* Ref.make(undefined); + const attempts = yield* Ref.make(0); + yield* Effect.gen(function* () { + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const sink = yield* EventSink.EventSinkV2; + const ids = yield* IdAllocator.IdAllocatorV2; + const now = yield* DateTime.now; + const threadId = ThreadId.make("refused-detached-writer"); + const peer = ThreadId.make("refused-detach-peer"); + const providerSessionId = ids.derive.providerSession({ + providerInstanceId: modelSelection.instanceId, + }); + yield* sink.write({ + events: [ + yield* makeThreadCreatedEvent({ idAllocator: ids, threadId, now }), + yield* makeThreadCreatedEvent({ idAllocator: ids, threadId: peer, now }), + ], + }); + const input = { threadId, providerSessionId, modelSelection, runtimePolicy }; + const runtime = yield* manager.open(input); + yield* manager.open({ ...input, threadId: peer }); + const providerThread = makeProviderThread({ + idAllocator: ids, + threadId, + providerSessionId, + now, + }); + yield* Ref.set(refusedThread, providerThread); + yield* runtime.resumeThread({ ...input, providerThread }); + yield* sink.write({ + events: [ + { + id: yield* ids.allocate.event({ threadId }), + type: "provider-thread.updated", + threadId, + occurredAt: now, + payload: providerThread, + }, + ], + }); + // Detach may surface the refusal or retain the writer silently; either + // way the unchanged owner must be usable while the command keeps running. + yield* manager.detach({ threadId, providerSessionId }).pipe(Effect.exit); + assert.isAbove(yield* Ref.get(attempts), 0); + assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, []); + const reopened = yield* manager.open(input); + assert.strictEqual(reopened, runtime); + assert.isTrue(yield* reopened.hasPendingBackgroundWorkForThread!(providerThread)); + const resumed = yield* reopened.resumeThread({ ...input, providerThread }); + assert.deepEqual(resumed.nativeThreadRef, providerThread.nativeThreadRef); + assert.isTrue(yield* reopened.hasPendingBackgroundWorkForThread!(providerThread)); + const final = yield* Ref.get(state); + assert.equal(final.resumeCount, 1, "the native writer is already loaded"); + assert.deepEqual(final.unloadedNativeThreadIds, []); + assert.equal(final.openCount, 1); + assert.equal(final.closeCount, 0, "the pooled process and peer must survive"); + }).pipe( + Effect.provide( + layerTest({ + state, + idleTimeoutMs: 600_000, + capabilities: CodexProviderCapabilitiesV2, + hasPendingBackgroundWork: Ref.get(background), + hasPendingBackgroundWorkForThread: Ref.get(background), + beforeUnload: Effect.gen(function* () { + yield* Ref.update(attempts, (n) => n + 1); + const thread = yield* Ref.get(refusedThread); + if (thread === undefined) return yield* Effect.die("Missing loaded writer fixture"); + if (yield* Ref.get(background)) { + return yield* new ProviderAdapter.ProviderAdapterThreadUnloadRefusedError({ + driver: CODEX_DRIVER, + providerThreadId: thread.id, + }); + } + }), + }), + ), + ); + }), +); + +it.effect("retains a failed detached writer for retry while its pooled peer remains alive", () => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const confirmed = yield* Ref.make(false); + const attempts = yield* Ref.make(0); + yield* Effect.gen(function* () { + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const sink = yield* EventSink.EventSinkV2; + const ids = yield* IdAllocator.IdAllocatorV2; + const now = yield* DateTime.now; + const threadId = ThreadId.make("unconfirmed-detached-writer"), + peer = ThreadId.make("retained-peer"); + const providerSessionId = ids.derive.providerSession({ + providerInstanceId: modelSelection.instanceId, + }); + yield* sink.write({ + events: [ + yield* makeThreadCreatedEvent({ idAllocator: ids, threadId, now }), + yield* makeThreadCreatedEvent({ idAllocator: ids, threadId: peer, now }), + ], + }); + const runtime = yield* manager.open({ + threadId, + providerSessionId, + modelSelection, + runtimePolicy, + }); + yield* manager.open({ threadId: peer, providerSessionId, modelSelection, runtimePolicy }); + const providerThread = makeProviderThread({ + idAllocator: ids, + threadId, + providerSessionId, + now, + }); + const resume = runtime.resumeThread({ + providerThread, + threadId, + modelSelection, + runtimePolicy, + }); + yield* resume; + yield* sink.write({ + events: [ + { + id: yield* ids.allocate.event({ threadId }), + type: "provider-thread.updated", + threadId, + occurredAt: now, + payload: providerThread, + }, + ], + }); + const detached = yield* manager.detach({ threadId, providerSessionId }).pipe(Effect.exit); + assert.isTrue( + Exit.isFailure(detached), + "logical detach must not acknowledge an unconfirmed native release", + ); + const reattached = yield* resume.pipe(Effect.exit); + assert.isTrue(Exit.isFailure(reattached)); + assert.equal((yield* Ref.get(state)).resumeCount, 1); + assert.equal((yield* Ref.get(state)).closeCount, 0); + yield* Ref.set(confirmed, true); + yield* manager.detach({ threadId, providerSessionId }); + const releaseAttempts = yield* Ref.get(attempts); + yield* manager.detach({ threadId, providerSessionId }); + assert.equal(yield* Ref.get(attempts), releaseAttempts); + assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, ["native-thread"]); + yield* resume; + assert.equal((yield* Ref.get(state)).resumeCount, 2); + assert.equal((yield* Ref.get(state)).closeCount, 0); + assert.strictEqual(Option.getOrThrow(yield* manager.get(providerSessionId)), runtime); + }).pipe( + Effect.provide( + layerTest({ + state, + idleTimeoutMs: 600_000, + capabilities: CodexProviderCapabilitiesV2, + beforeUnload: Ref.update(attempts, (n) => n + 1).pipe( + Effect.andThen(Ref.get(confirmed)), + Effect.flatMap((done) => + done + ? Effect.void + : Effect.fail( + new ProviderAdapter.ProviderAdapterProtocolError({ + driver: CODEX_DRIVER, + detail: "Native closure is unconfirmed", + }), + ), + ), + ), + }), + ), + ); + }), +); diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.ts index edcaecfa8848..eed609ae80aa 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.ts @@ -120,6 +120,15 @@ export class ProviderSessionReleaseError extends Schema.TaggedError()( + "ProviderSessionBackgroundWorkError", + { providerSessionId: ProviderSessionId, threadId: ThreadId }, +) { + override get message(): string { + return "Provider account switch is waiting for background work. Let its tools and child agents finish, or stop them before retrying. The original native thread remains loaded."; + } +} + export class ProviderSessionActivityError extends Schema.TaggedError()( "ProviderSessionActivityError", { @@ -178,6 +187,7 @@ export interface ProviderSessionManagerV2Shape { readonly providerSessionId: ProviderSessionId; readonly threadId: ThreadId; readonly detail?: string; + readonly expectedTargetInstanceId?: ProviderInstanceId; readonly revokeMcpCredential?: boolean; }) => Effect.Effect; } @@ -189,7 +199,24 @@ export class ProviderSessionManagerV2 extends Context.Service< interface LiveSessionEntry { readonly attachedThreadIds: ReadonlySet; - readonly loadedProviderThreadKeyByThread: ReadonlyMap; + /** Logical detach is not proof of native closure. Retain retryable ownership. */ + readonly pendingThreadReleases: ReadonlyMap< + ThreadId, + { + readonly providerThreads: ReadonlyArray; + readonly idleUnload?: { + readonly generation: number; + readonly release: "unconfirmed" | "refused"; + }; + } + >; + readonly loadedProviderThreadByThread: ReadonlyMap< + ThreadId, + { + readonly key: string; + readonly providerThread: OrchestrationV2ProviderThread; + } + >; /** * Revocation on detach/release is scoped to these ids so tearing down a * superseded session cannot revoke a replacement session's credential for the @@ -1293,13 +1320,13 @@ export const layerWithOptions = ( } const attachedThreadIds = new Set(entry.attachedThreadIds); attachedThreadIds.delete(input.threadId); - const loadedProviderThreadKeyByThread = new Map(entry.loadedProviderThreadKeyByThread); - loadedProviderThreadKeyByThread.delete(input.threadId); + const loadedProviderThreadByThread = new Map(entry.loadedProviderThreadByThread); + loadedProviderThreadByThread.delete(input.threadId); const updated = new Map(current); updated.set(key, { ...entry, attachedThreadIds, - loadedProviderThreadKeyByThread, + loadedProviderThreadByThread, }); return updated; }); @@ -1310,37 +1337,162 @@ export const layerWithOptions = ( readonly providerThreadKey: string; }) => Ref.get(sessions).pipe( - Effect.map( - (current) => - current - .get(sessionKey(input.providerSessionId)) - ?.loadedProviderThreadKeyByThread.get(input.threadId) === input.providerThreadKey, - ), + Effect.map((current) => { + const entry = current.get(sessionKey(input.providerSessionId)); + return ( + !entry?.pendingThreadReleases.has(input.threadId) && + entry?.loadedProviderThreadByThread.get(input.threadId)?.key === + input.providerThreadKey + ); + }), ); const markProviderThreadLoaded = (input: { readonly providerSessionId: ProviderSessionId; readonly threadId: ThreadId; readonly providerThreadKey: string; + readonly providerThread: OrchestrationV2ProviderThread; + readonly runtime: ProviderAdapter.ProviderAdapterV2SessionRuntime; }) => Ref.update(sessions, (current) => { const key = sessionKey(input.providerSessionId); const entry = current.get(key); - if (entry === undefined) { + if ( + entry === undefined || + entry.runtime !== input.runtime || + entry.pendingThreadReleases.has(input.threadId) + ) { return current; } - const loadedProviderThreadKeyByThread = new Map(entry.loadedProviderThreadKeyByThread); - loadedProviderThreadKeyByThread.set(input.threadId, input.providerThreadKey); + const loadedProviderThreadByThread = new Map(entry.loadedProviderThreadByThread); + loadedProviderThreadByThread.set(input.threadId, { + key: input.providerThreadKey, + providerThread: input.providerThread, + }); const updated = new Map(current); - updated.set(key, { ...entry, loadedProviderThreadKeyByThread }); + updated.set(key, { ...entry, loadedProviderThreadByThread }); return updated; }); - const ensureThreadAttached = (input: { - readonly providerSessionId: ProviderSessionId; - readonly threadId: ThreadId; - readonly providerInstanceId: ProviderInstanceId; - }) => + // Runs under the attachment lock. Both detach retries and reattachment + // must settle the same native ownership record before discarding it. + const confirmPendingThreadRelease = ( + entry: LiveSessionEntry, + threadId: ThreadId, + waitForIdleTurn = true, + ) => + Effect.gen(function* () { + const pending = entry.pendingThreadReleases.get(threadId); + const unloadThread = entry.exposedRuntime.unloadThread; + if (pending === undefined || unloadThread === undefined) return; + yield* Effect.forEach( + pending.providerThreads.filter((thread) => thread.nativeThreadRef !== null), + (providerThread) => + Effect.gen(function* () { + if (waitForIdleTurn && hasBusyTurn(entry, providerThread.id)) + return yield* new ProviderSessionBackgroundWorkError({ + providerSessionId: entry.runtime.providerSessionId, + threadId, + }); + yield* unloadThread({ providerThread }).pipe( + Effect.timeout(UNLOAD_THREAD_TIMEOUT_MS), + ); + }), + { concurrency: 1, discard: true }, + ); + yield* Ref.update(sessions, (current) => { + const key = sessionKey(entry.runtime.providerSessionId); + const latest = current.get(key); + if ( + latest?.runtime !== entry.runtime || + latest.pendingThreadReleases.get(threadId) !== pending + ) + return current; + const pendingThreadReleases = new Map(latest.pendingThreadReleases); + pendingThreadReleases.delete(threadId); + const loadedProviderThreadByThread = new Map(latest.loadedProviderThreadByThread); + loadedProviderThreadByThread.delete(threadId); + const updated = new Map(current); + updated.set(key, { ...latest, pendingThreadReleases, loadedProviderThreadByThread }); + return updated; + }); + }); + + const reclaimRefusedThreadRelease = ( + entry: LiveSessionEntry, + threadId: ThreadId, + refusal?: ProviderAdapter.ProviderAdapterThreadUnloadRefusedError, + ) => + Effect.gen(function* () { + const pending = entry.pendingThreadReleases.get(threadId); + const idle = entry.idleThreadUnloads.get(threadId); + const loaded = entry.loadedProviderThreadByThread.get(threadId); + const original = pending?.providerThreads[0]; + if ( + pending === undefined || + pending.providerThreads.length !== 1 || + (pending.idleUnload === undefined + ? refusal?.driver !== entry.runtime.driver || + refusal.providerThreadId !== original?.id + : pending.idleUnload.release !== "refused" || + !entry.attachedThreadIds.has(threadId)) || + idle?.generation !== pending.idleUnload?.generation || + loaded === undefined || + original?.nativeThreadRef === null || + original === undefined || + original.nativeThreadRef.nativeId === null || + loaded.providerThread.id !== original.id || + loaded.providerThread.nativeThreadRef?.nativeId !== original.nativeThreadRef.nativeId || + original.providerInstanceId !== entry.runtime.instanceId || + original.providerSessionId !== entry.runtime.providerSessionId + ) + return false; + const generation = pending.idleUnload?.generation; + const projection = yield* projectionStore.getThreadRecords(threadId, ["providerThreads"]); + const binding = projection.providerThreads.find( + (thread) => thread.id === projection.thread.activeProviderThreadId, + ); + if ( + projection.thread.modelSelection.instanceId !== entry.runtime.instanceId || + binding?.id !== original.id || + binding.providerInstanceId !== entry.runtime.instanceId || + binding.providerSessionId !== entry.runtime.providerSessionId || + binding.nativeThreadRef?.nativeId !== original.nativeThreadRef.nativeId + ) + return false; + const reclaimed = yield* Ref.modify(sessions, (current) => { + const key = sessionKey(entry.runtime.providerSessionId); + const latest = current.get(key); + const latestIdle = latest?.idleThreadUnloads.get(threadId); + if ( + latest?.runtime !== entry.runtime || + latest.attachedThreadIds.has(threadId) !== entry.attachedThreadIds.has(threadId) || + latest.pendingThreadReleases.get(threadId) !== pending || + latest.loadedProviderThreadByThread.get(threadId) !== loaded || + latestIdle?.generation !== generation + ) + return [Option.none | null>(), current] as const; + const pendingThreadReleases = new Map(latest.pendingThreadReleases); + pendingThreadReleases.delete(threadId); + const idleThreadUnloads = new Map(latest.idleThreadUnloads); + idleThreadUnloads.delete(threadId); + const updated = new Map(current); + updated.set(key, { ...latest, pendingThreadReleases, idleThreadUnloads }); + return [Option.some(latestIdle?.fiber ?? null), updated] as const; + }); + if (Option.isNone(reclaimed)) return false; + yield* cancelIdleFiber(reclaimed.value); + return true; + }); + + const ensureThreadAttached = ( + input: { + readonly providerSessionId: ProviderSessionId; + readonly threadId: ThreadId; + readonly providerInstanceId: ProviderInstanceId; + }, + attachmentLockHeld = false, + ) => Effect.suspend(() => { let attachedTo: ProviderAdapter.ProviderAdapterV2SessionRuntime | undefined; let preparedForCleanup: PreparedMcpCredential | undefined; @@ -1426,9 +1578,11 @@ export const layerWithOptions = ( ), ), ); - return threadAttachment - .withLock(threadAttachmentKey(input), attach) - .pipe(Effect.ensuring(Effect.sync(dropReservation))); + return ( + attachmentLockHeld + ? attach + : threadAttachment.withLock(threadAttachmentKey(input), attach) + ).pipe(Effect.ensuring(Effect.sync(dropReservation))); }); const markBusy = (providerSessionId: ProviderSessionId, turnKey: string) => @@ -1602,8 +1756,12 @@ export const layerWithOptions = ( ? false : yield* entry.runtime .hasPendingBackgroundWorkForThread(pending.providerThread) - .pipe(Effect.catchCause(() => Effect.succeed(false))); + .pipe(Effect.catchCause(() => Effect.succeed(true))); if (hasPendingWork) return "deferred" as const; + const pendingRelease = { + providerThreads: [pending.providerThread], + idleUnload: { generation: input.generation, release: "unconfirmed" as const }, + }; const unloading = yield* Ref.modify(sessions, (current) => { const latest = current.get(key); if ( @@ -1612,32 +1770,75 @@ export const layerWithOptions = ( ) { return [false, current] as const; } - const loadedProviderThreadKeyByThread = new Map( - latest.loadedProviderThreadKeyByThread, - ); - loadedProviderThreadKeyByThread.delete(input.threadId); - const idleThreadUnloads = new Map(latest.idleThreadUnloads); - idleThreadUnloads.delete(input.threadId); + const pendingThreadReleases = new Map(latest.pendingThreadReleases); + pendingThreadReleases.set(input.threadId, pendingRelease); const updated = new Map(current); updated.set(key, { ...latest, - loadedProviderThreadKeyByThread, - idleThreadUnloads, + pendingThreadReleases, }); return [true, updated] as const; }); if (!unloading) return "skipped" as const; - yield* unloadThread({ providerThread: pending.providerThread }).pipe( + const unloaded = yield* unloadThread({ providerThread: pending.providerThread }).pipe( Effect.timeout(UNLOAD_THREAD_TIMEOUT_MS), + Effect.as(true), + Effect.catchTags({ + ProviderAdapterThreadUnloadRefusedError: (error) => + error.driver !== entry.runtime.driver || + error.providerThreadId !== pending.providerThread.id + ? Effect.fail(error) + : Ref.update(sessions, (current) => { + const latest = current.get(key); + if ( + latest?.runtime !== entry.runtime || + latest.pendingThreadReleases.get(input.threadId) !== pendingRelease || + latest.idleThreadUnloads.get(input.threadId)?.generation !== + input.generation + ) + return current; + const pendingThreadReleases = new Map(latest.pendingThreadReleases); + pendingThreadReleases.set(input.threadId, { + ...pendingRelease, + idleUnload: { generation: input.generation, release: "refused" }, + }); + const updated = new Map(current); + updated.set(key, { ...latest, pendingThreadReleases }); + return updated; + }).pipe(Effect.as(false)), + }), Effect.catchCause((cause) => Effect.logWarning("orchestration-v2.driver-session.idle-unload-failed", { providerSessionId: input.providerSessionId, threadId: input.threadId, providerThreadId: pending.providerThread.id, cause, - }), + }).pipe(Effect.as(false)), ), ); + if (!unloaded) return "deferred" as const; + yield* Ref.update(sessions, (current) => { + const latest = current.get(key); + if ( + latest?.runtime !== entry.runtime || + latest.idleThreadUnloads.get(input.threadId)?.generation !== input.generation + ) + return current; + const loadedProviderThreadByThread = new Map(latest.loadedProviderThreadByThread); + loadedProviderThreadByThread.delete(input.threadId); + const idleThreadUnloads = new Map(latest.idleThreadUnloads); + idleThreadUnloads.delete(input.threadId); + const pendingThreadReleases = new Map(latest.pendingThreadReleases); + pendingThreadReleases.delete(input.threadId); + const updated = new Map(current); + updated.set(key, { + ...latest, + loadedProviderThreadByThread, + idleThreadUnloads, + pendingThreadReleases, + }); + return updated; + }); return "unloaded" as const; }), ); @@ -1707,6 +1908,56 @@ export const layerWithOptions = ( >, ): ProviderAdapter.ProviderAdapterV2SessionRuntime => { const providerSessionId = runtime.providerSessionId; + const prepareThreadAttachment = (threadId: ThreadId) => { + const input = { providerSessionId, threadId, providerInstanceId: runtime.instanceId }; + return threadAttachment.withLock( + threadAttachmentKey(input), + Effect.gen(function* () { + const entry = (yield* Ref.get(sessions)).get(sessionKey(providerSessionId)); + if (entry?.runtime !== runtime) { + return yield* new ProviderAdapter.ProviderAdapterProtocolError({ + driver: runtime.driver, + detail: "Provider runtime changed before thread attachment; retry the operation.", + }); + } + yield* reclaimRefusedThreadRelease(entry, threadId).pipe( + Effect.flatMap((reclaimed) => + reclaimed + ? Effect.void + : confirmPendingThreadRelease(entry, threadId).pipe( + Effect.catchTags({ + ProviderAdapterThreadUnloadRefusedError: (refusal) => + reclaimRefusedThreadRelease(entry, threadId, refusal).pipe( + Effect.flatMap((reclaimed) => + reclaimed ? Effect.void : Effect.fail(refusal), + ), + ), + }), + ), + ), + Effect.mapError( + (cause) => + new ProviderAdapter.ProviderAdapterProtocolError({ + driver: runtime.driver, + detail: + "Native writer release must finish before reattaching this conversation.", + cause, + }), + ), + ); + // Activity persistence stays best effort; native release is strict. + yield* observeActivity(providerSessionId, ensureThreadAttached(input, true)); + if ( + (yield* Ref.get(sessions)).get(sessionKey(providerSessionId))?.runtime !== runtime + ) { + return yield* new ProviderAdapter.ProviderAdapterProtocolError({ + driver: runtime.driver, + detail: "Provider runtime changed during thread attachment; retry the operation.", + }); + } + }), + ); + }; const subscribeEvents = makeEventSubscription(eventSubscribers); const turnMetrics = (operation: string, model?: string) => withMetrics({ @@ -1725,19 +1976,14 @@ export const layerWithOptions = ( subscribeEvents.pipe(Effect.map((subscription) => subscription.events)), ), ensureThread: (input) => - observeActivity( - providerSessionId, - ensureThreadAttached({ - providerSessionId, - threadId: input.threadId, - providerInstanceId: runtime.instanceId, - }), - ).pipe( + prepareThreadAttachment(input.threadId).pipe( Effect.andThen(runtime.ensureThread(input)), Effect.tap((providerThread) => markProviderThreadLoaded({ providerSessionId, threadId: input.threadId, + providerThread, + runtime, providerThreadKey: providerThreadLoadKey({ providerThread, modelSelection: input.modelSelection, @@ -1758,14 +2004,7 @@ export const layerWithOptions = ( : { modelSelection: input.modelSelection }), ...(input.runtimePolicy === undefined ? {} : { runtimePolicy: input.runtimePolicy }), }); - return observeActivity( - providerSessionId, - ensureThreadAttached({ - providerSessionId, - threadId, - providerInstanceId: runtime.instanceId, - }), - ).pipe( + return prepareThreadAttachment(threadId).pipe( Effect.andThen( holdThreadLoaded({ providerSessionId, @@ -1783,6 +2022,8 @@ export const layerWithOptions = ( markProviderThreadLoaded({ providerSessionId, threadId, + providerThread, + runtime, providerThreadKey: providerThreadLoadKey({ providerThread, ...(input.modelSelection === undefined @@ -1797,19 +2038,14 @@ export const layerWithOptions = ( ); }, forkThread: (input) => - observeActivity( - providerSessionId, - ensureThreadAttached({ - providerSessionId, - threadId: input.targetThreadId, - providerInstanceId: runtime.instanceId, - }), - ).pipe( + prepareThreadAttachment(input.targetThreadId).pipe( Effect.andThen(runtime.forkThread(input)), Effect.tap((providerThread) => markProviderThreadLoaded({ providerSessionId, threadId: input.targetThreadId, + providerThread, + runtime, providerThreadKey: providerThreadLoadKey({ providerThread, ...(input.modelSelection === undefined @@ -1823,14 +2059,7 @@ export const layerWithOptions = ( ), ), startTurn: (input) => - observeActivity( - providerSessionId, - ensureThreadAttached({ - providerSessionId, - threadId: input.threadId, - providerInstanceId: runtime.instanceId, - }), - ).pipe( + prepareThreadAttachment(input.threadId).pipe( // A start that fails or is stopped may never emit turn.terminal, // so it clears its own turn or the session never goes idle. Effect.andThen( @@ -2062,7 +2291,7 @@ export const layerWithOptions = ( ), ); - return ProviderSessionManagerV2.of({ + const manager: ProviderSessionManagerV2Shape = ProviderSessionManagerV2.of({ shutdown, open: (input) => sessionOpen.withLock( @@ -2081,6 +2310,24 @@ export const layerWithOptions = ( }); } } + if (input.initialNativeThreadId !== undefined) { + const owners = [...(yield* Ref.get(sessions)).values()].filter( + (entry) => + entry.runtime.driver === "codex" && + entry.runtime.instanceId !== input.modelSelection.instanceId && + entry.loadedProviderThreadByThread.get(input.threadId)?.providerThread + .nativeThreadRef?.nativeId === input.initialNativeThreadId, + ); + // Projections can already name the target after a failed detach. The + // live owner must still confirm release before any target load. + for (const owner of owners) { + yield* manager.detach({ + providerSessionId: owner.runtime.providerSessionId, + threadId: input.threadId, + expectedTargetInstanceId: input.modelSelection.instanceId, + }); + } + } const key = sessionKey(input.providerSessionId); const existing = (yield* Ref.get(sessions)).get(key); if (existing !== undefined) { @@ -2186,7 +2433,8 @@ export const layerWithOptions = ( const now = yield* Clock.currentTimeMillis; const entry: LiveSessionEntry = { attachedThreadIds: new Set([input.threadId]), - loadedProviderThreadKeyByThread: new Map(), + loadedProviderThreadByThread: new Map(), + pendingThreadReleases: new Map(), mcpCredentialIdByThread: mcpCredentialId === undefined ? new Map() @@ -2310,21 +2558,93 @@ export const layerWithOptions = ( Effect.gen(function* () { const key = sessionKey(input.providerSessionId); const currentEntry = (yield* Ref.get(sessions)).get(key); - let detachedProviderThreads: ReadonlyArray = []; - if (currentEntry?.supportsMultipleProviderThreads === true) { + const pendingRelease = currentEntry?.pendingThreadReleases.get(input.threadId); + // The projection may already have rebound the provider-thread row + // before this detach effect executes. The live session still owns + // its loaded native thread and must unsubscribe that reference. + const loaded = currentEntry?.loadedProviderThreadByThread.get(input.threadId); + const loadedProviderThread = loaded?.providerThread; + let bindingChanged = false; + let nativeBindingTransferred = false; + let detachedProviderThreads: ReadonlyArray = + pendingRelease?.providerThreads ?? + (loadedProviderThread === undefined ? [] : [loadedProviderThread]); + if ( + currentEntry?.supportsMultipleProviderThreads === true && + (currentEntry.attachedThreadIds.has(input.threadId) || pendingRelease !== undefined) + ) { const projection = yield* Effect.option( projectionStore.getThreadRecords(input.threadId, [ "providerThreads", "providerTurns", + "runs", ]), ); if (Option.isSome(projection)) { + const activeThread = projection.value.providerThreads.find( + (thread) => thread.id === projection.value.thread.activeProviderThreadId, + ); + // A delayed retry belongs to the earlier account switch. A + // new run bound to the source must keep its process and tools. + if ( + input.expectedTargetInstanceId !== undefined && + activeThread?.providerInstanceId !== input.expectedTargetInstanceId && + (projection.value.runs.some((run) => + ["preparing", "starting", "running", "waiting"].includes(run.status), + ) || + projection.value.thread.modelSelection.instanceId !== + input.expectedTargetInstanceId) + ) + return; + bindingChanged = + projection.value.thread.modelSelection.instanceId !== + currentEntry.runtime.instanceId || + projection.value.providerThreads.some( + (thread) => + thread.id === loadedProviderThread?.id && + (thread.providerInstanceId !== currentEntry.runtime.instanceId || + thread.providerSessionId !== input.providerSessionId), + ); + nativeBindingTransferred = + bindingChanged && + activeThread?.driver === currentEntry.runtime.driver && + activeThread.nativeThreadRef?.nativeId != null && + activeThread.nativeThreadRef.nativeId === + loadedProviderThread?.nativeThreadRef?.nativeId; const providerThreads = new Map( projection.value.providerThreads .filter((thread) => thread.providerSessionId === input.providerSessionId) .map((thread) => [thread.id, thread] as const), ); - detachedProviderThreads = [...providerThreads.values()]; + const ownedProviderThreads = new Map(providerThreads); + if (loadedProviderThread !== undefined) { + ownedProviderThreads.set(loadedProviderThread.id, loadedProviderThread); + } + detachedProviderThreads = [...ownedProviderThreads.values()]; + // Account transfer must not terminate the source's background work. + if (currentEntry.runtime.unloadThread !== undefined && nativeBindingTransferred) { + const ownedThread = + loadedProviderThread ?? + detachedProviderThreads.find( + (thread) => thread.ownerNodeId === null && thread.nativeThreadRef !== null, + ); + const pendingWork = + (input.expectedTargetInstanceId !== undefined && + ownedThread !== undefined && + hasBusyTurn(currentEntry, ownedThread.id)) || + (currentEntry.runtime.hasPendingBackgroundWorkForThread !== undefined + ? ownedThread !== undefined && + (yield* currentEntry.runtime.hasPendingBackgroundWorkForThread(ownedThread)) + : currentEntry.runtime.hasPendingBackgroundWork === undefined + ? false + : yield* currentEntry.runtime.hasPendingBackgroundWork); + if (pendingWork) { + return yield* new ProviderSessionBackgroundWorkError({ + providerSessionId: input.providerSessionId, + threadId: input.threadId, + }); + } + } const activeTurns = projection.value.providerTurns.filter( (turn) => turn.status === "running" && providerThreads.has(turn.providerThreadId), ); @@ -2355,7 +2675,32 @@ export const layerWithOptions = ( } const detachResult = yield* Ref.modify(sessions, (current) => { const entry = current.get(key); - if (entry === undefined || !entry.attachedThreadIds.has(input.threadId)) { + if ( + entry === undefined || + entry.runtime !== currentEntry?.runtime || + (!entry.attachedThreadIds.has(input.threadId) && + entry.pendingThreadReleases.get(input.threadId) !== pendingRelease) + ) { + return [ + Option.none<{ + readonly entry: LiveSessionEntry; + readonly idleUnloadFiber: Fiber.Fiber | null; + }>(), + current, + ] as const; + } + if (!entry.attachedThreadIds.has(input.threadId)) { + return [ + pendingRelease === undefined + ? Option.none() + : Option.some({ entry, idleUnloadFiber: null }), + current, + ] as const; + } + if ( + input.expectedTargetInstanceId !== undefined && + entry.loadedProviderThreadByThread.get(input.threadId) !== loaded + ) { return [ Option.none<{ readonly entry: LiveSessionEntry; @@ -2366,10 +2711,10 @@ export const layerWithOptions = ( } const attachedThreadIds = new Set(entry.attachedThreadIds); attachedThreadIds.delete(input.threadId); - const loadedProviderThreadKeyByThread = new Map( - entry.loadedProviderThreadKeyByThread, - ); - loadedProviderThreadKeyByThread.delete(input.threadId); + const pendingThreadReleases = new Map(entry.pendingThreadReleases); + pendingThreadReleases.set(input.threadId, { + providerThreads: detachedProviderThreads, + }); const mcpCredentialIdByThread = input.revokeMcpCredential === true ? (() => { @@ -2383,7 +2728,7 @@ export const layerWithOptions = ( const updatedEntry = { ...entry, attachedThreadIds, - loadedProviderThreadKeyByThread, + pendingThreadReleases, mcpCredentialIdByThread, idleThreadUnloads, }; @@ -2423,39 +2768,38 @@ export const layerWithOptions = ( }); return; } - const unloadThread = detached.value.exposedRuntime.unloadThread; - if (detached.value.supportsMultipleProviderThreads && unloadThread !== undefined) { - // Serialized with re-attachment: a thread whose next turn - // attaches first stays loaded, and one that attaches during the - // unload waits for it, so its resume reloads the native thread. + if ( + detached.value.supportsMultipleProviderThreads && + detached.value.exposedRuntime.unloadThread !== undefined + ) { yield* threadAttachment.withLock( threadAttachmentKey(input), Effect.gen(function* () { const entry = (yield* Ref.get(sessions)).get(key); if ( entry?.runtime !== detached.value.runtime || - entry.attachedThreadIds.has(input.threadId) - ) { + entry.attachedThreadIds.has(input.threadId) || + entry.pendingThreadReleases.get(input.threadId) !== + detached.value.pendingThreadReleases.get(input.threadId) + ) return; - } - yield* Effect.forEach( - detachedProviderThreads.filter((thread) => thread.nativeThreadRef !== null), - (providerThread) => - unloadThread({ providerThread }).pipe( - Effect.timeout(UNLOAD_THREAD_TIMEOUT_MS), - Effect.catchCause((cause) => - Effect.logWarning( + yield* confirmPendingThreadRelease( + entry, + input.threadId, + input.expectedTargetInstanceId !== undefined && nativeBindingTransferred, + ).pipe( + Effect.catchCause((cause) => + bindingChanged && !nativeBindingTransferred + ? Effect.logWarning( "orchestration-v2.driver-session.detach-unload-failed", { providerSessionId: input.providerSessionId, threadId: input.threadId, - providerThreadId: providerThread.id, cause, }, - ), - ), - ), - { concurrency: 1, discard: true }, + ) + : Effect.failCause(cause), + ), ); }), ); @@ -2472,7 +2816,8 @@ export const layerWithOptions = ( ), ), ), - } satisfies ProviderSessionManagerV2Shape); + }); + return manager; }), );