From eba0521563946a0834af881465d3e1ff7a4ca2f1 Mon Sep 17 00:00:00 2001 From: Rhys Sullivan <39114868+RhysSullivan@users.noreply.github.com> Date: Wed, 7 Oct 2026 13:54:17 -0700 Subject: [PATCH] fix(server): unload Codex threads left idle on the shared app-server (#16917) Co-authored-by: Julius Marminge <51714798+juliusmarminge@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) --- .../ProviderSessionManager.test.ts | 207 +++++++++++++++ .../ProviderSessionManager.ts | 244 +++++++++++++++++- 2 files changed, 440 insertions(+), 11 deletions(-) diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts index bfb076ee3bf3b6d3f3048efa917574e38f2fb6bd..2f2b0402f6205624354a6d02d56ff60476505da9 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts @@ -319,6 +319,7 @@ function makeProviderAdapter( readonly initialProviderItemIdentityVersion?: 2; }) => Effect.Effect; readonly hasPendingBackgroundWork?: Effect.Effect; + readonly hasPendingBackgroundWorkForThread?: Effect.Effect; readonly hangSessionScopeClose?: boolean; readonly startTurn?: Effect.Effect; readonly beforeUnload?: Effect.Effect; @@ -407,6 +408,11 @@ function makeProviderAdapter( ...(options.hasPendingBackgroundWork === undefined ? {} : { hasPendingBackgroundWork: options.hasPendingBackgroundWork }), + ...(options.hasPendingBackgroundWorkForThread === undefined + ? {} + : { + hasPendingBackgroundWorkForThread: () => options.hasPendingBackgroundWorkForThread!, + }), ensureThread: () => unimplemented("ensureThread unused in test"), resumeThread: (threadInput) => Ref.update(state, (current) => ({ @@ -463,6 +469,7 @@ function layerTest(input: { readonly paused: Deferred.Deferred; }; readonly hasPendingBackgroundWork?: Effect.Effect; + readonly hasPendingBackgroundWorkForThread?: Effect.Effect; readonly hangSessionScopeClose?: boolean; readonly startTurn?: Effect.Effect; readonly beforeUnload?: Effect.Effect; @@ -488,6 +495,9 @@ function layerTest(input: { ...(input.hasPendingBackgroundWork === undefined ? {} : { hasPendingBackgroundWork: input.hasPendingBackgroundWork }), + ...(input.hasPendingBackgroundWorkForThread === undefined + ? {} + : { hasPendingBackgroundWorkForThread: input.hasPendingBackgroundWorkForThread }), ...(input.hangSessionScopeClose === undefined ? {} : { hangSessionScopeClose: input.hangSessionScopeClose }), @@ -4210,6 +4220,203 @@ it.effect( }), ); +function runIdleThreadUnloadScenario( + name: string, + scenario: (input: { + readonly state: Ref.Ref; + readonly manager: ProviderSessionManager.ProviderSessionManagerV2Shape; + readonly providerSessionId: ProviderSessionId; + readonly threadA: ThreadId; + /** Starts run `ordinal` on a thread and returns once the provider accepted it. */ + readonly startTurn: (threadId: ThreadId, ordinal: number) => Effect.Effect; + /** Ends run `ordinal` on a thread and waits for the session to process it. */ + readonly endTurn: (threadId: ThreadId, ordinal: number) => Effect.Effect; + readonly resume: (threadId: ThreadId) => Effect.Effect; + }) => Effect.Effect, + options: { readonly hasPendingBackgroundWorkForThread?: Effect.Effect } = {}, +) { + return Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + yield* Effect.gen(function* () { + const eventSink = yield* EventSink.EventSinkV2; + const idAllocator = yield* IdAllocator.IdAllocatorV2; + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const projectionStore = yield* ProjectionStore.ProjectionStoreV2; + const now = yield* DateTime.now; + const threadA = ThreadId.make(`thread-provider-session-manager-${name}-a`); + const threadB = ThreadId.make(`thread-provider-session-manager-${name}-b`); + const providerSessionId = idAllocator.derive.providerSession({ + providerInstanceId: modelSelection.instanceId, + }); + yield* eventSink.write({ + events: [ + yield* makeThreadCreatedEvent({ idAllocator, threadId: threadA, now }), + yield* makeThreadCreatedEvent({ idAllocator, threadId: threadB, now }), + ], + }); + const runtime = yield* manager.open({ + threadId: threadA, + providerSessionId, + modelSelection, + runtimePolicy, + }); + yield* manager.open({ threadId: threadB, providerSessionId, modelSelection, runtimePolicy }); + const providerThreadOf = (threadId: ThreadId) => + makeProviderThread({ + idAllocator, + threadId, + providerSessionId, + now, + nativeThreadId: `native-${threadId}`, + }); + const resume = (threadId: ThreadId) => + runtime + .resumeThread({ + providerThread: providerThreadOf(threadId), + threadId, + modelSelection, + runtimePolicy, + }) + .pipe(Effect.asVoid, Effect.orDie); + const startTurn = (threadId: ThreadId, ordinal: number) => + Effect.gen(function* () { + yield* resume(threadId); + const runId = idAllocator.derive.run({ threadId, ordinal }); + yield* runtime.startTurn({ + appThread: (yield* projectionStore.getThreadProjection(threadId)).thread, + threadId, + runId, + runOrdinal: ordinal, + providerTurnOrdinal: ordinal, + attemptId: idAllocator.derive.runAttempt({ runId, attemptOrdinal: 1 }), + rootNodeId: idAllocator.derive.rootNode({ runId }), + providerThread: providerThreadOf(threadId), + message: { + createdBy: "user", + creationSource: "web", + messageId: yield* idAllocator.allocate.message({ threadId, ordinal }), + text: "turn", + attachments: [], + }, + modelSelection, + runtimePolicy, + }); + }).pipe(Effect.orDie); + const subscribe = runtime.subscribeEvents; + assert.isDefined(subscribe); + const endTurn = (threadId: ThreadId, ordinal: number) => + Effect.scoped( + Effect.gen(function* () { + // The pump hands an event to subscribers only after the session + // has processed it, so receiving it is the receipt. + const subscription = yield* Effect.acquireRelease(subscribe!, (sub) => sub.close); + const received = yield* subscription.events.pipe( + Stream.filter((event) => event.type === "turn.terminal"), + Stream.runHead, + Effect.forkScoped, + ); + const queue = (yield* Ref.get(state)).eventQueues.get(String(providerSessionId)); + assert.isDefined(queue); + yield* Queue.offer(queue!, { + type: "turn.terminal", + driver: CODEX_DRIVER, + providerThreadId: providerThreadOf(threadId).id, + providerTurnId: idAllocator.derive.providerTurn({ + driver: CODEX_DRIVER, + nativeTurnId: `native-turn-${threadId}-${ordinal}`, + }), + runOrdinal: ordinal, + status: "completed", + failure: null, + threadDisposition: "reusable", + }); + yield* Fiber.join(received); + }), + ).pipe(Effect.orDie); + + // B's running turn keeps the shared runtime itself busy throughout. + yield* startTurn(threadB, 1); + yield* scenario({ + state, + manager, + providerSessionId, + threadA, + startTurn, + endTurn, + resume, + }); + }).pipe( + Effect.provide( + layerTest({ + state, + idleTimeoutMs: 1000, + ...(options.hasPendingBackgroundWorkForThread === undefined + ? {} + : { hasPendingBackgroundWorkForThread: options.hasPendingBackgroundWorkForThread }), + }), + ), + ); + }); +} + +it.effect( + "ProviderSessionManagerV2 unloads a shared-runtime thread left idle, and reloads it on its next turn", + () => + runIdleThreadUnloadScenario( + "idle-unload", + ({ state, manager, providerSessionId, threadA, startTurn, endTurn, resume }) => + Effect.gen(function* () { + yield* startTurn(threadA, 1); + assert.equal((yield* Ref.get(state)).resumeCount, 2); + yield* endTurn(threadA, 1); + + // A follow-up before the timeout keeps the thread loaded. + yield* TestClock.adjust("500 millis"); + yield* startTurn(threadA, 2); + yield* TestClock.adjust("1 second"); + assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, []); + + // Idle for the full timeout after its last turn, A is unloaded while + // the runtime stays up for B. + yield* endTurn(threadA, 2); + yield* TestClock.adjust("1 second"); + assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, [`native-${threadA}`]); + assert.isTrue(Option.isSome(yield* manager.get(providerSessionId))); + assert.equal((yield* Ref.get(state)).closeCount, 0); + + // The next turn's resume reaches the provider and reloads A. + const resumes = (yield* Ref.get(state)).resumeCount; + yield* resume(threadA); + assert.equal((yield* Ref.get(state)).resumeCount, resumes + 1); + }), + ), +); + +it.effect( + "ProviderSessionManagerV2 keeps an idle shared-runtime thread loaded while its background work runs", + () => + Effect.gen(function* () { + const pendingWork = yield* Ref.make(true); + yield* runIdleThreadUnloadScenario( + "idle-unload-pinned", + ({ state, threadA, startTurn, endTurn }) => + Effect.gen(function* () { + yield* startTurn(threadA, 1); + yield* endTurn(threadA, 1); + yield* TestClock.adjust("3 seconds"); + assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, []); + + yield* Ref.set(pendingWork, false); + yield* TestClock.adjust("1 second"); + assert.deepEqual((yield* Ref.get(state)).unloadedNativeThreadIds, [ + `native-${threadA}`, + ]); + }), + { hasPendingBackgroundWorkForThread: Ref.get(pendingWork) }, + ); + }), +); + it.effect( "ProviderSessionManagerV2 rejects a second thread when the provider runtime is exclusive", () => diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.ts index f84441129cc3ae10ea738bd3ebce245ea877ec1d..d181497ec2754f6b2a6641cca95989aee3dd9012 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.ts @@ -61,9 +61,10 @@ const DEFAULT_IDLE_TIMEOUT_MS = 30 * 60 * 1000; const DEFAULT_MAX_IDLE_PIN_MS = 4 * 60 * 60 * 1000; const RELEASE_SCOPE_CLOSE_TIMEOUT_MS = 30 * 1000; +const busyTurnPrefix = (providerThreadId: ProviderThreadId) => `${providerThreadId}#`; /** The identity a turn's start and its `turn.terminal` share. */ const busyTurnKey = (providerThreadId: ProviderThreadId, runOrdinal: number) => - `${providerThreadId}#${runOrdinal}`; + `${busyTurnPrefix(providerThreadId)}${runOrdinal}`; const UNLOAD_THREAD_TIMEOUT_MS = 10 * 1000; export const ProviderSessionReleaseReason = Schema.Literals([ @@ -228,6 +229,18 @@ interface LiveSessionEntry { readonly idleFiber: Fiber.Fiber | null; /** Set when idle release is deferred for pending background work; bounds total deferral. */ readonly pinnedSinceMs: number | null; + /** + * Shared runtimes only: the provider thread each attached app thread last + * started a turn on, and the timer that unloads it once it has been idle + * for `idleTimeoutMs`. + */ + readonly idleThreadUnloads: ReadonlyMap; +} + +interface IdleThreadUnload { + readonly providerThread: OrchestrationV2ProviderThread; + readonly generation: number; + readonly fiber: Fiber.Fiber | null; } type ProviderSessionEventSignal = @@ -1400,7 +1413,11 @@ export const layerWithOptions = ( // Clearing a turn that is not marked busy (one whose failed start already // cleared it, or a subagent turn the manager never started) only // records activity. - const markIdle = (providerSessionId: ProviderSessionId, turnKey: string) => + const markIdle = ( + providerSessionId: ProviderSessionId, + providerThreadId: ProviderThreadId, + runOrdinal: number, + ) => withActivityError( providerSessionId, Effect.gen(function* () { @@ -1412,7 +1429,7 @@ export const layerWithOptions = ( return current; } const busyTurns = new Set(entry.busyTurns); - busyTurns.delete(turnKey); + busyTurns.delete(busyTurnKey(providerThreadId, runOrdinal)); const updated = new Map(current); updated.set(key, { ...entry, @@ -1422,9 +1439,181 @@ export const layerWithOptions = ( return updated; }); yield* scheduleIdleReleaseInternal(providerSessionId); + yield* scheduleThreadUnload(providerSessionId, providerThreadId); }), ); + const hasBusyTurn = (entry: LiveSessionEntry, providerThreadId: ProviderThreadId) => { + const prefix = busyTurnPrefix(providerThreadId); + for (const turnKey of entry.busyTurns) { + if (turnKey.startsWith(prefix)) return true; + } + return false; + }; + + const updateIdleThreadUnload = ( + providerSessionId: ProviderSessionId, + threadId: ThreadId, + update: (current: IdleThreadUnload | undefined) => IdleThreadUnload | undefined, + ) => + Ref.modify(sessions, (current) => { + const key = sessionKey(providerSessionId); + const entry = current.get(key); + if (entry === undefined) return [undefined, current] as const; + const previous = entry.idleThreadUnloads.get(threadId); + const next = update(previous); + const idleThreadUnloads = new Map(entry.idleThreadUnloads); + if (next === undefined) idleThreadUnloads.delete(threadId); + else idleThreadUnloads.set(threadId, next); + const updated = new Map(current); + updated.set(key, { ...entry, idleThreadUnloads }); + return [previous, updated] as const; + }); + + /** + * Starts tracking the provider thread an app thread runs its turns on, and + * stops any unload pending for it. Called before the thread is resumed or + * given a turn, so an unload cannot land between a resume that found the + * thread loaded and the turn that relies on it. + */ + const holdThreadLoaded = (input: { + readonly providerSessionId: ProviderSessionId; + readonly threadId: ThreadId; + readonly providerThread: OrchestrationV2ProviderThread; + }) => + Effect.gen(function* () { + const entry = (yield* Ref.get(sessions)).get(sessionKey(input.providerSessionId)); + if ( + entry === undefined || + !entry.supportsMultipleProviderThreads || + entry.exposedRuntime.unloadThread === undefined || + input.providerThread.nativeThreadRef === null + ) { + return; + } + const previous = yield* updateIdleThreadUnload( + input.providerSessionId, + input.threadId, + (current) => ({ + providerThread: input.providerThread, + generation: (current?.generation ?? 0) + 1, + fiber: null, + }), + ); + yield* cancelIdleFiber(previous?.fiber ?? null); + }); + + /** + * A shared runtime never goes idle while any of its threads is in use, so + * its threads get the idle timeout one by one: a thread with no turn for + * `idleTimeoutMs` is unloaded from the runtime, along with the native MCP + * servers it started. Its next turn's resume loads it again. + */ + const scheduleThreadUnload = ( + providerSessionId: ProviderSessionId, + providerThreadId: ProviderThreadId, + ) => + Effect.gen(function* () { + const entry = (yield* Ref.get(sessions)).get(sessionKey(providerSessionId)); + if (entry === undefined || hasBusyTurn(entry, providerThreadId)) return; + const tracked = Array.from(entry.idleThreadUnloads).find( + ([, pending]) => pending.providerThread.id === providerThreadId, + ); + if (tracked === undefined) return; + const [threadId, pending] = tracked; + const generation = pending.generation + 1; + const fiber = yield* Effect.sleep(Duration.millis(idleTimeoutMs)).pipe( + Effect.andThen(unloadIdleThread({ providerSessionId, threadId, generation })), + Effect.forkIn(layerScope), + ); + // A turn that started meanwhile already moved the generation on. + const previous = yield* updateIdleThreadUnload(providerSessionId, threadId, (current) => + current?.generation === pending.generation + ? { ...current, generation, fiber } + : current, + ); + yield* cancelIdleFiber( + previous?.generation === pending.generation ? previous.fiber : fiber, + ); + }); + + const unloadIdleThread = (input: { + readonly providerSessionId: ProviderSessionId; + readonly threadId: ThreadId; + readonly generation: number; + }): Effect.Effect => + Effect.gen(function* () { + const outcome = yield* threadAttachment.withLock( + threadAttachmentKey(input), + Effect.gen(function* () { + const key = sessionKey(input.providerSessionId); + const entry = (yield* Ref.get(sessions)).get(key); + const pending = entry?.idleThreadUnloads.get(input.threadId); + const unloadThread = entry?.exposedRuntime.unloadThread; + if ( + entry === undefined || + pending === undefined || + pending.generation !== input.generation || + unloadThread === undefined || + !entry.attachedThreadIds.has(input.threadId) || + hasBusyTurn(entry, pending.providerThread.id) + ) { + return "skipped" as const; + } + // Unloading stops the native thread's background terminals, so a + // thread still running background work stays loaded. + const hasPendingWork = + entry.runtime.hasPendingBackgroundWorkForThread === undefined + ? false + : yield* entry.runtime + .hasPendingBackgroundWorkForThread(pending.providerThread) + .pipe(Effect.catchCause(() => Effect.succeed(false))); + if (hasPendingWork) return "deferred" as const; + const unloading = yield* Ref.modify(sessions, (current) => { + const latest = current.get(key); + if ( + latest?.runtime !== entry.runtime || + latest.idleThreadUnloads.get(input.threadId)?.generation !== input.generation + ) { + 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 updated = new Map(current); + updated.set(key, { + ...latest, + loadedProviderThreadKeyByThread, + idleThreadUnloads, + }); + return [true, updated] as const; + }); + if (!unloading) return "skipped" as const; + yield* unloadThread({ providerThread: pending.providerThread }).pipe( + Effect.timeout(UNLOAD_THREAD_TIMEOUT_MS), + Effect.catchCause((cause) => + Effect.logWarning("orchestration-v2.driver-session.idle-unload-failed", { + providerSessionId: input.providerSessionId, + threadId: input.threadId, + providerThreadId: pending.providerThread.id, + cause, + }), + ), + ); + return "unloaded" as const; + }), + ); + // Re-check on this fiber after another idle window, outside the + // lock so the thread's next attach is not held up meanwhile. + if (outcome === "deferred") { + yield* Effect.sleep(Duration.millis(idleTimeoutMs)); + return yield* unloadIdleThread(input); + } + }); + const observeActivity = ( providerSessionId: ProviderSessionId, activity: Effect.Effect, @@ -1544,6 +1733,13 @@ export const layerWithOptions = ( providerInstanceId: runtime.instanceId, }), ).pipe( + Effect.andThen( + holdThreadLoaded({ + providerSessionId, + threadId, + providerThread: input.providerThread, + }), + ), Effect.andThen( isProviderThreadLoaded({ providerSessionId, threadId, providerThreadKey }), ), @@ -1607,6 +1803,13 @@ export const layerWithOptions = ( // the adapter emits the terminal anyway, clearing the same turn // again changes nothing, so another thread's turn on a shared // session stays busy either way. + Effect.andThen( + holdThreadLoaded({ + providerSessionId, + threadId: input.threadId, + providerThread: input.providerThread, + }), + ), Effect.andThen( Effect.acquireUseRelease( observeActivity( @@ -1622,10 +1825,7 @@ export const layerWithOptions = ( Exit.isFailure(exit) ? observeActivity( providerSessionId, - markIdle( - providerSessionId, - busyTurnKey(input.providerThread.id, input.runOrdinal), - ), + markIdle(providerSessionId, input.providerThread.id, input.runOrdinal), ) : Effect.void, ), @@ -1692,7 +1892,8 @@ export const layerWithOptions = ( event.type === "turn.terminal" ? markIdle( entry.runtime.providerSessionId, - busyTurnKey(event.providerThreadId, event.runOrdinal), + event.providerThreadId, + event.runOrdinal, ) : touchActivity(entry.runtime.providerSessionId), ).pipe( @@ -1965,6 +2166,7 @@ export const layerWithOptions = ( lastActivityAtMs: now, idleFiber: null, pinnedSinceMs: null, + idleThreadUnloads: new Map(), }; yield* Ref.update(sessions, (current) => { const updated = new Map(current); @@ -2111,10 +2313,16 @@ export const layerWithOptions = ( ); } } - const detached = yield* Ref.modify(sessions, (current) => { + const detachResult = yield* Ref.modify(sessions, (current) => { const entry = current.get(key); if (entry === undefined || !entry.attachedThreadIds.has(input.threadId)) { - return [Option.none(), current] as const; + return [ + Option.none<{ + readonly entry: LiveSessionEntry; + readonly idleUnloadFiber: Fiber.Fiber | null; + }>(), + current, + ] as const; } const attachedThreadIds = new Set(entry.attachedThreadIds); attachedThreadIds.delete(input.threadId); @@ -2135,16 +2343,30 @@ export const layerWithOptions = ( return pruned; })() : entry.mcpCredentialIdByThread; + // The detach unloads the thread itself below. + const idleThreadUnloads = new Map(entry.idleThreadUnloads); + idleThreadUnloads.delete(input.threadId); const updatedEntry = { ...entry, attachedThreadIds, loadedProviderThreadKeyByThread, mcpCredentialIdByThread, + idleThreadUnloads, }; const updated = new Map(current); updated.set(key, updatedEntry); - return [Option.some(updatedEntry), updated] as const; + return [ + Option.some({ + entry: updatedEntry, + idleUnloadFiber: entry.idleThreadUnloads.get(input.threadId)?.fiber ?? null, + }), + updated, + ] as const; }); + if (Option.isSome(detachResult)) { + yield* cancelIdleFiber(detachResult.value.idleUnloadFiber); + } + const detached = Option.map(detachResult, (result) => result.entry); // Plain detaches deliberately do not revoke: a detached thread's // provider process may still be alive (shared multi-thread codex // session across a workspace handoff) and holds its MCP client's -- 2.54.0