diff --git a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts index 60d1449656a40..7cb00e2ee1ef8 100644 --- a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts +++ b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts @@ -7650,6 +7650,12 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( function* (unloadInput) { const nativeThreadId = yield* getNativeThreadId(unloadInput.providerThread); sessionApprovals.delete(nativeThreadId); + // Archive also unloads, inline and without thread/closed; resumeNativeThread + // unarchives it on the next turn. + if (unloadInput.archive === true) { + yield* client.request("thread/archive", { threadId: nativeThreadId }); + return; + } // A resume may already wait on the pending entry; replacing it would strand that wait. if (!threadClosings.has(nativeThreadId)) { threadClosings.set(nativeThreadId, yield* Deferred.make()); diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts index 5e494b8b31904..740143c1500b2 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts @@ -448,6 +448,8 @@ const makeProviderAdapter = Effect.fnUntraced(function* ( readonly startTurn?: Effect.Effect; readonly interruptTurn?: Effect.Effect; readonly beforeUnload?: Effect.Effect; + readonly unloads?: Queue.Queue; + readonly archive?: Effect.Effect; readonly spawnBeforeOpen?: boolean; readonly scopeCloseReached?: Deferred.Deferred; } = {}, @@ -457,7 +459,11 @@ const makeProviderAdapter = Effect.fnUntraced(function* ( Ref.update(state, (current) => ({ ...current, closeCount: current.closeCount + 1, - })), + })).pipe( + Effect.andThen( + options.unloads === undefined ? Effect.void : Queue.offer(options.unloads, "close"), + ), + ), ); // Register after countClose to wedge teardown before its finalizer, matching a provider stream that never yields. const hangClose = Effect.addFinalizer(() => @@ -563,7 +569,7 @@ const makeProviderAdapter = Effect.fnUntraced(function* ( ...current, interruptCount: current.interruptCount + 1, })).pipe(Effect.andThen(options.interruptTurn ?? Effect.void)), - unloadThread: ({ providerThread }) => + unloadThread: ({ providerThread, archive }) => (options.beforeUnload ?? Effect.void).pipe( Effect.andThen( Ref.update(state, (current) => ({ @@ -574,6 +580,15 @@ const makeProviderAdapter = Effect.fnUntraced(function* ( ], })), ), + Effect.andThen( + options.unloads === undefined + ? Effect.void + : Queue.offer( + options.unloads, + `${archive === true ? "archive" : "unload"}:${providerThread.nativeThreadRef?.nativeId}`, + ), + ), + Effect.andThen(archive === true ? (options.archive ?? Effect.void) : Effect.void), ), respondToRuntimeRequest: () => Effect.void, readThreadSnapshot: () => unimplemented("readThreadSnapshot unused in test"), @@ -614,6 +629,8 @@ function layerTest(input: { readonly startTurn?: Effect.Effect; readonly interruptTurn?: Effect.Effect; readonly beforeUnload?: Effect.Effect; + readonly unloads?: Queue.Queue; + readonly archive?: Effect.Effect; readonly spawnBeforeOpen?: boolean; readonly scopeCloseReached?: Deferred.Deferred; readonly serverSettingsLayer?: ReturnType; @@ -657,6 +674,8 @@ function layerTest(input: { ...(input.startTurn === undefined ? {} : { startTurn: input.startTurn }), ...(input.interruptTurn === undefined ? {} : { interruptTurn: input.interruptTurn }), ...(input.beforeUnload === undefined ? {} : { beforeUnload: input.beforeUnload }), + ...(input.unloads === undefined ? {} : { unloads: input.unloads }), + ...(input.archive === undefined ? {} : { archive: input.archive }), ...(input.spawnBeforeOpen === undefined ? {} : { spawnBeforeOpen: input.spawnBeforeOpen }), ...(input.scopeCloseReached === undefined ? {} @@ -4706,6 +4725,169 @@ it.effect( }), ); +it.effect.each([ + { keptOpen: true, archive: "succeeds" }, + { keptOpen: false, archive: "succeeds" }, + { keptOpen: false, archive: "fails" }, + { keptOpen: false, archive: "hangs" }, +] as const)( + "ProviderSessionManagerV2 archives an idle delegated child once its parent has the result (session kept open: $keptOpen, archive $archive)", + ({ keptOpen, archive }) => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const unloads = yield* Queue.unbounded(); + const effect = 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 parentThreadId = ThreadId.make("thread-provider-session-manager-archive-parent"); + const parentRunId = idAllocator.derive.run({ threadId: parentThreadId, ordinal: 1 }); + const providerSessionId = idAllocator.derive.providerSession({ + providerInstanceId: modelSelection.instanceId, + }); + yield* eventSink.write({ + events: [yield* makeThreadCreatedEvent({ idAllocator, threadId: parentThreadId, now })], + }); + let deliverPending: OrchestrationV2DomainEvent | undefined; + for (const delivery of ["pending", "delivered"] as const) { + const threadId = ThreadId.make(`thread-provider-session-manager-archive-${delivery}`); + const runId = idAllocator.derive.run({ threadId, ordinal: 1 }); + const taskId = idAllocator.derive.rootNode({ runId }); + const providerThread = makeProviderThread({ + idAllocator, + threadId, + providerSessionId, + now, + nativeThreadId: `native-${delivery}`, + }); + const created = yield* makeThreadCreatedEvent({ + idAllocator, + threadId, + now, + delegatedBy: parentThreadId, + }); + const taskEvent = (state: "pending" | "delivered") => + Effect.gen(function* () { + return { + id: yield* idAllocator.allocate.event({ threadId: parentThreadId }), + type: "subagent.updated", + threadId: parentThreadId, + nodeId: taskId, + providerInstanceId: modelSelection.instanceId, + occurredAt: now, + payload: { + id: taskId, + threadId: parentThreadId, + runId: parentRunId, + parentNodeId: idAllocator.derive.rootNode({ runId: parentRunId }), + origin: "app_owned", + createdBy: "agent", + driver: CODEX_DRIVER, + providerInstanceId: modelSelection.instanceId, + providerThreadId: null, + childThreadId: threadId, + nativeTaskRef: null, + prompt: "Delegated work", + title: null, + model: null, + completionDelivery: { state, observedByRunId: null }, + status: "completed", + result: "Done.", + startedAt: now, + completedAt: now, + updatedAt: now, + }, + } satisfies OrchestrationV2DomainEvent; + }); + if (delivery === "pending") deliverPending = yield* taskEvent("delivered"); + yield* eventSink.write({ + events: [ + { + ...created, + payload: { ...created.payload, forkedFrom: { type: "node", nodeId: taskId } }, + }, + yield* taskEvent(delivery), + ], + }); + const runtime = yield* manager.open({ + threadId, + providerSessionId, + modelSelection, + runtimePolicy, + }); + yield* runtime.startTurn({ + appThread: (yield* projectionStore.getThreadProjection(threadId)).thread, + threadId, + runId, + runOrdinal: 1, + providerTurnOrdinal: 1, + attemptId: idAllocator.derive.runAttempt({ runId, attemptOrdinal: 1 }), + rootNodeId: idAllocator.derive.rootNode({ runId }), + providerThread, + message: { + createdBy: "agent", + creationSource: "server", + messageId: yield* idAllocator.allocate.message({ threadId, ordinal: 1 }), + text: "Delegated work", + attachments: [], + }, + modelSelection, + runtimePolicy, + }); + // The pump marks the thread idle before subscribers see its terminal. + const subscription = yield* runtime.subscribeEvents!; + yield* Queue.offer((yield* Ref.get(state)).eventQueues.get(String(providerSessionId))!, { + type: "turn.terminal", + driver: CODEX_DRIVER, + providerThreadId: providerThread.id, + providerTurnId: idAllocator.derive.providerTurn({ + driver: CODEX_DRIVER, + nativeTurnId: `native-turn-${delivery}`, + }), + runOrdinal: 1, + status: "completed", + failure: null, + threadDisposition: "reusable", + }); + yield* subscription.events.pipe(Stream.runHead); + } + + yield* TestClock.adjust("1 second"); + assert.equal(yield* Queue.take(unloads), "archive:native-delivered"); + if (archive !== "succeeds") { + if (archive === "hangs") yield* TestClock.adjust("10 seconds"); + // An archive that never finished may have left the thread loaded. + assert.equal(yield* Queue.take(unloads), "unload:native-delivered"); + return; + } + // The pending child stays loaded, and keeps the session open, until its parent takes the result. + yield* eventSink.write({ events: [deliverPending!] }); + yield* TestClock.adjust("1 second"); + assert.equal(yield* Queue.take(unloads), "archive:native-pending"); + if (!keptOpen) assert.equal(yield* Queue.take(unloads), "close"); + }); + + yield* effect.pipe( + Effect.provide( + layerTest({ + state, + idleTimeoutMs: 1000, + unloads, + ...(archive === "fails" + ? { archive: unimplemented("archive refused") } + : archive === "hangs" + ? { archive: Effect.never } + : {}), + // Another thread's work keeps the shared session open past the children's idle unloads. + hasPendingBackgroundWork: Effect.succeed(keptOpen), + }), + ), + ); + }), +); + it.effect( "ProviderSessionManagerV2 re-attaching a thread waits for its in-flight unload, then reloads it", () => diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.ts index 33088785615a4..0c259bcd1361d 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.ts @@ -1216,11 +1216,19 @@ export const layerWithOptions = ( // Capture runtime identity before yielding; replacements can reuse the session id. const probedRuntime = entry.runtime; const hasPendingWork = - probedRuntime.hasPendingBackgroundWork === undefined + (probedRuntime.hasPendingBackgroundWork === undefined ? false : yield* probedRuntime.hasPendingBackgroundWork.pipe( Effect.catchCause(() => Effect.succeed(false)), - ); + )) || + // Closing now would skip the archive of a child whose parent has not taken its result. + Option.isSome( + yield* Effect.findFirst(entry.idleThreadUnloads.keys(), (threadId) => + delegatedChildArchive(threadId).pipe( + Effect.map((child) => child === "awaiting_delivery"), + ), + ), + ); if (hasPendingWork) { const now = yield* Clock.currentTimeMillis; const pinnedSinceMs = entry.pinnedSinceMs ?? now; @@ -1255,6 +1263,18 @@ export const layerWithOptions = ( pinnedForMs: now - pinnedSinceMs, }); } + // Closing skips the threads' own idle unloads, so archive delivered children first. + yield* Effect.forEach( + entry.idleThreadUnloads, + ([threadId, pending]) => + unloadIdleThread({ + providerSessionId: input.providerSessionId, + threadId, + generation: pending.generation, + archiveOnly: true, + }), + { discard: true }, + ); // Adapter probes yield; guard busyTurns and idleGeneration atomically when removing the entry. yield* releaseEntry({ providerSessionId: input.providerSessionId, @@ -1672,11 +1692,44 @@ export const layerWithOptions = ( }, ); + // Provider clients list a delegated child as an ordinary conversation, so a + // finished one is archived once its parent has the result. + const delegatedChildArchive = Effect.fn("ProviderSessionManager.delegatedChildArchive")( + function* (threadId: ThreadId) { + const { thread, runs } = yield* projectionStore.getThreadRecords(threadId, ["runs"], { + runStatuses: [], + includeLatestRun: true, + }); + // Follow-up tasks name their run; the initial task is the node the child forked from. + const taskId = + runs.at(-1)?.delegatedTaskId ?? + (thread.forkedFrom?.type === "node" ? thread.forkedFrom.nodeId : undefined); + const parentThreadId = thread.lineage.parentThreadId; + if (parentThreadId === null || taskId === undefined) return undefined; + const task = (yield* projectionStore.getThreadRecords(parentThreadId, ["subagents"], { + subagentIds: [taskId], + })).subagents[0]; + if ( + task?.origin !== "app_owned" || + task.childThreadId !== threadId || + task.status !== "completed" + ) { + return undefined; + } + const delivery = task.completionDelivery?.state; + if (delivery === "pending" || delivery === "claimed") return "awaiting_delivery"; + return delivery === "delivered" || delivery === "acknowledged" ? "ready" : undefined; + }, + (effect) => effect.pipe(Effect.orElseSucceed(() => undefined)), + ); + const unloadIdleThread = Effect.fn("ProviderSessionManager.unloadIdleThread")( function* (input: { readonly providerSessionId: ProviderSessionId; readonly threadId: ThreadId; readonly generation: number; + /** Session release unloads only the threads it would archive, and never waits. */ + readonly archiveOnly?: boolean; }): Effect.fn.Return { const outcome = yield* threadAttachment.withLock( threadAttachmentKey(input), @@ -1703,6 +1756,11 @@ export const layerWithOptions = ( .hasPendingBackgroundWorkForThread(pending.providerThread) .pipe(Effect.catchCause(() => Effect.succeed(false))); if (hasPendingWork) return "deferred" as const; + const child = yield* delegatedChildArchive(input.threadId); + // Unloading now would miss the archive once the parent takes the result. + if (child === "awaiting_delivery") return "deferred" as const; + const archive = child === "ready"; + if (input.archiveOnly === true && !archive) return "skipped" as const; const unloading = yield* Ref.modify(sessions, (current) => { const latest = current.get(key); if ( @@ -1726,8 +1784,23 @@ export const layerWithOptions = ( return [true, updated] as const; }); if (!unloading) return "skipped" as const; - yield* unloadThread({ providerThread: pending.providerThread }).pipe( - Effect.timeout(UNLOAD_THREAD_TIMEOUT_MS), + const unload = (archiveThread: boolean) => + unloadThread({ + providerThread: pending.providerThread, + archive: archiveThread, + }).pipe(Effect.timeout(UNLOAD_THREAD_TIMEOUT_MS)); + yield* unload(archive).pipe( + // A failed or timed-out archive may have left the thread loaded. + Effect.catch((cause) => + archive + ? Effect.logWarning("orchestration-v2.driver-session.idle-archive-failed", { + providerSessionId: input.providerSessionId, + threadId: input.threadId, + providerThreadId: pending.providerThread.id, + cause, + }).pipe(Effect.andThen(unload(false))) + : Effect.fail(cause), + ), Effect.tap(() => Ref.update(sessions, (current) => { const latest = current.get(key); @@ -1750,7 +1823,7 @@ export const layerWithOptions = ( }), ); // Wait outside the lock so the next attach is not blocked. - if (outcome === "deferred") { + if (outcome === "deferred" && input.archiveOnly !== true) { yield* Effect.sleep(Duration.millis(idleTimeoutMs)); return yield* unloadIdleThread(input); } diff --git a/packages/provider-core/src/server/ProviderAdapter.ts b/packages/provider-core/src/server/ProviderAdapter.ts index 7ae8c0c43f003..047eb8aae1005 100644 --- a/packages/provider-core/src/server/ProviderAdapter.ts +++ b/packages/provider-core/src/server/ProviderAdapter.ts @@ -724,6 +724,8 @@ export interface ProviderAdapterV2SessionRuntime { }) => Effect.Effect; readonly unloadThread?: (input: { readonly providerThread: OrchestrationV2ProviderThread; + /** Also hide it from the provider's own clients where it can; resuming it must restore it. */ + readonly archive?: boolean; }) => Effect.Effect; readonly respondToRuntimeRequest: ( input: ProviderAdapterV2RuntimeRequestResponseInput,