diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts index 40c2d28a590a..de020ca91f11 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts @@ -251,6 +251,8 @@ function makeThreadCreatedEvent(input: { readonly now: DateTime.Utc; readonly projectId?: ProjectId; readonly delegatedBy?: ThreadId; + readonly parentRelationship?: "fork" | "subagent"; + readonly remoteDelegation?: boolean; }) { return Effect.gen(function* () { const projectId = @@ -277,13 +279,27 @@ function makeThreadCreatedEvent(input: { activeProviderThreadId: providerThreadId, lineage: input.delegatedBy === undefined - ? { parentThreadId: null, relationshipToParent: null, rootThreadId: input.threadId } + ? { + parentThreadId: + input.parentRelationship === undefined ? null : ThreadId.make("parent"), + relationshipToParent: input.parentRelationship ?? null, + rootThreadId: input.threadId, + } : { parentThreadId: input.delegatedBy, relationshipToParent: "subagent", rootThreadId: input.delegatedBy, }, forkedFrom: null, + ...(input.remoteDelegation + ? { + delegatedFrom: { + environmentId: EnvironmentId.make("remote-parent"), + threadId: ThreadId.make("remote-parent-thread"), + title: "Remote parent", + }, + } + : {}), createdAt: input.now, updatedAt: input.now, archivedAt: null, @@ -615,7 +631,7 @@ const makeProviderAdapter = Effect.fnUntraced(function* ( function layerTest(input: { readonly state: Ref.Ref; - readonly idleTimeoutMs: number; + readonly idleTimeoutMs?: number; readonly maxIdlePinMs?: number; readonly failEventStream?: boolean; readonly normalizeThreadLoadModelSelection?: (selection: ModelSelection) => ModelSelection; @@ -727,7 +743,7 @@ function layerTest(input: { layerConfiguredMcpRegistry, McpProviderSessions.layer, ProviderSessionManager.layerWithOptions({ - idleTimeoutMs: input.idleTimeoutMs, + ...(input.idleTimeoutMs === undefined ? {} : { idleTimeoutMs: input.idleTimeoutMs }), ...(input.maxIdlePinMs === undefined ? {} : { maxIdlePinMs: input.maxIdlePinMs }), }).pipe( Layer.provide( @@ -3492,6 +3508,150 @@ it.effect("ProviderSessionManagerV2 keeps provider updates that land while relea }), ); +it.effect.each([ + { parentRelationship: undefined, remoteDelegation: false }, + { parentRelationship: "fork", remoteDelegation: false }, + { parentRelationship: "subagent", remoteDelegation: false }, + { parentRelationship: undefined, remoteDelegation: true }, +] as const)( + "ProviderSessionManagerV2 uses the default idle window (relationship: $parentRelationship, remote: $remoteDelegation)", + ({ parentRelationship, remoteDelegation }) => + 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 threadId = ThreadId.make("idle-window"); + const providerSessionId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + yield* eventSink.write({ + events: [ + yield* makeThreadCreatedEvent({ + idAllocator, + threadId, + now, + remoteDelegation, + ...(parentRelationship === undefined ? {} : { parentRelationship }), + }), + ], + }); + yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); + yield* TestClock.adjust("5 minutes"); + const delegated = parentRelationship === "subagent" || remoteDelegation; + assert.equal(yield* manager.isLive(providerSessionId), !delegated); + if (!delegated) yield* TestClock.adjust("25 minutes"); + assert.isFalse(yield* manager.isLive(providerSessionId)); + assert.equal((yield* Ref.get(state)).closeCount, 1); + + const projection = yield* projectionStore.getThreadProjection(threadId); + const releasedSession = projection.providerSessions.at(-1); + assert.equal(releasedSession?.status, "stopped"); + const runtime = yield* manager.open({ + threadId, + providerSessionId, + modelSelection, + runtimePolicy, + resumeFromSession: releasedSession!, + }); + yield* runtime.resumeThread({ + threadId, + providerThread: makeProviderThread({ idAllocator, threadId, providerSessionId, now }), + modelSelection, + runtimePolicy, + }); + assert.equal((yield* Ref.get(state)).openCount, 2); + assert.equal((yield* Ref.get(state)).resumeCount, 1); + }).pipe(Effect.provide(layerTest({ state }))); + }), +); + +it.effect.each([true, false])( + "ProviderSessionManagerV2 preserves the user idle window in a mixed session (child first: %s)", + (childFirst) => + 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 now = yield* DateTime.now; + const childId = ThreadId.make("mixed-child"); + const userId = ThreadId.make("mixed-user"); + const providerSessionId = idAllocator.derive.providerSession({ + providerInstanceId: modelSelection.instanceId, + }); + yield* eventSink.write({ + events: [ + yield* makeThreadCreatedEvent({ + idAllocator, + threadId: childId, + now, + parentRelationship: "subagent", + }), + yield* makeThreadCreatedEvent({ idAllocator, threadId: userId, now }), + ], + }); + for (const threadId of childFirst ? [childId, userId] : [userId, childId]) { + yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); + } + yield* TestClock.adjust("5 minutes"); + assert.isTrue(yield* manager.isLive(providerSessionId)); + yield* manager.detach({ providerSessionId, threadId: userId }); + yield* TestClock.adjust("5 minutes"); + assert.isFalse(yield* manager.isLive(providerSessionId)); + assert.equal((yield* Ref.get(state)).closeCount, 1); + }).pipe(Effect.provide(layerTest({ state }))); + }), +); + +it.effect("ProviderSessionManagerV2 pins delegated idle sessions for pending background work", () => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const checks = yield* Ref.make(0); + yield* Effect.gen(function* () { + const eventSink = yield* EventSink.EventSinkV2; + const idAllocator = yield* IdAllocator.IdAllocatorV2; + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const threadId = ThreadId.make("delegated-idle-pin"); + const now = yield* DateTime.now; + const providerSessionId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + yield* eventSink.write({ + events: [ + yield* makeThreadCreatedEvent({ + idAllocator, + threadId, + now, + parentRelationship: "subagent", + }), + ], + }); + yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); + yield* TestClock.adjust("5 minutes"); + assert.equal(yield* Ref.get(checks), 1); + assert.isTrue(yield* manager.isLive(providerSessionId)); + assert.equal((yield* Ref.get(state)).closeCount, 0); + yield* TestClock.adjust("4 hours"); + assert.isFalse(yield* manager.isLive(providerSessionId)); + assert.equal((yield* Ref.get(state)).closeCount, 1); + }).pipe( + Effect.provide( + layerTest({ + state, + hasPendingBackgroundWork: Ref.update(checks, (count) => count + 1).pipe(Effect.as(true)), + }), + ), + ); + }), +); + it.effect( "ProviderSessionManagerV2 does not idle-release a session that turns busy during the pending-work check", () => diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.ts index ade466854136..0c62c4dd98c0 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.ts @@ -63,6 +63,7 @@ import * as ProviderAdapterRegistry from "./ProviderAdapterRegistry.ts"; import * as ProjectionStore from "./ProjectionStore.ts"; const DEFAULT_IDLE_TIMEOUT_MS = 30 * 60 * 1000; +const DELEGATED_IDLE_TIMEOUT_MS = 5 * 60 * 1000; const DEFAULT_MAX_IDLE_PIN_MS = 4 * 60 * 60 * 1000; const RELEASE_SCOPE_CLOSE_TIMEOUT_MS = 30 * 1000; /** Past adapters' own 10 s stop fallbacks, so they settle first when the provider answers at all. */ @@ -274,6 +275,8 @@ export class ProviderSessionManagerV2 extends Context.Service< interface LiveSessionEntry { readonly attachedThreadIds: ReadonlySet; + readonly idleTimeoutMsByThread: ReadonlyMap; + readonly idleTimeoutMs: number; readonly loadedProviderThreadKeyByThread: ReadonlyMap; /** null marks an unknown catalog; undefined means MCP refresh is not configured. */ readonly openedWithLinkedEnvironments: boolean | null | undefined; @@ -505,6 +508,22 @@ export const layerWithOptions = ( }) => `${input.providerSessionId}\u0000${input.threadId}`; const idleTimeoutMs = Math.max(1, options.idleTimeoutMs ?? DEFAULT_IDLE_TIMEOUT_MS); const maxIdlePinMs = Math.max(0, options.maxIdlePinMs ?? DEFAULT_MAX_IDLE_PIN_MS); + const resolveThreadIdleTimeout = Effect.fn("ProviderSessionManager.resolveThreadIdleTimeout")( + function* (threadId: ThreadId) { + if (options.idleTimeoutMs !== undefined) return idleTimeoutMs; + const thread = yield* projectionStore.getThread(threadId); + // Peer-delegated children have a remote parent instead of local lineage. + return thread.lineage.relationshipToParent === "subagent" || + thread.delegatedFrom !== undefined + ? DELEGATED_IDLE_TIMEOUT_MS + : idleTimeoutMs; + }, + Effect.orElseSucceed(() => idleTimeoutMs), + ); + // Shared runtimes must preserve the user window while any user thread is attached. + // Recompute only on attachment changes; provider events read the cached maximum. + const attachedIdleTimeout = (timeouts: ReadonlyMap) => + timeouts.size === 0 ? idleTimeoutMs : Math.max(...timeouts.values()); interface PreparedMcpCredential { readonly mcpCredentialId: string | undefined; readonly issued: boolean; @@ -1389,8 +1408,8 @@ export const layerWithOptions = ( providerSessionId: input.providerSessionId, pinnedForMs: now - pinnedSinceMs, }); - // Recheck on this fiber; scheduleIdleReleaseInternal would interrupt it and self-deadlock. - yield* Effect.sleep(Duration.millis(idleTimeoutMs)); + // Recheck on this fiber; scheduling again would interrupt this fiber. + yield* Effect.sleep(Duration.millis(entry.idleTimeoutMs)); return yield* releaseIfStillIdle(input); } yield* Effect.logWarning("orchestration-v2.driver-session.idle-release-pin-expired", { @@ -1459,7 +1478,7 @@ export const layerWithOptions = ( yield* cancelIdleFiber(entry.idleFiber); const generation = entry.idleGeneration + 1; - const idleFiber = yield* Effect.sleep(Duration.millis(idleTimeoutMs)).pipe( + const idleFiber = yield* Effect.sleep(Duration.millis(entry.idleTimeoutMs)).pipe( Effect.andThen(releaseIfStillIdle({ providerSessionId, generation })), Effect.forkIn(layerScope), ); @@ -1497,17 +1516,26 @@ export const layerWithOptions = ( }) => withActivityError( input.providerSessionId, - Ref.modify(sessions, (current) => { - const entry = current.get(sessionKey(input.providerSessionId)); - if (entry === undefined || entry.attachedThreadIds.has(input.threadId)) { - return [undefined, current] as const; - } - const updated = new Map(current); - updated.set(sessionKey(input.providerSessionId), { - ...entry, - attachedThreadIds: new Set([...entry.attachedThreadIds, input.threadId]), + Effect.gen(function* () { + const existing = (yield* Ref.get(sessions)).get(sessionKey(input.providerSessionId)); + if (existing === undefined || existing.attachedThreadIds.has(input.threadId)) return; + const threadIdleTimeoutMs = yield* resolveThreadIdleTimeout(input.threadId); + return yield* Ref.modify(sessions, (current) => { + const entry = current.get(sessionKey(input.providerSessionId)); + if (entry === undefined || entry.attachedThreadIds.has(input.threadId)) { + return [undefined, current] as const; + } + const updated = new Map(current); + const idleTimeoutMsByThread = new Map(entry.idleTimeoutMsByThread); + idleTimeoutMsByThread.set(input.threadId, threadIdleTimeoutMs); + updated.set(sessionKey(input.providerSessionId), { + ...entry, + attachedThreadIds: new Set([...entry.attachedThreadIds, input.threadId]), + idleTimeoutMsByThread, + idleTimeoutMs: attachedIdleTimeout(idleTimeoutMsByThread), + }); + return [entry.runtime, updated] as const; }); - return [entry.runtime, updated] as const; }), ); @@ -1528,6 +1556,8 @@ export const layerWithOptions = ( } const attachedThreadIds = new Set(entry.attachedThreadIds); attachedThreadIds.delete(input.threadId); + const idleTimeoutMsByThread = new Map(entry.idleTimeoutMsByThread); + idleTimeoutMsByThread.delete(input.threadId); const loadedProviderThreadKeyByThread = new Map(entry.loadedProviderThreadKeyByThread); loadedProviderThreadKeyByThread.delete(input.threadId); const linkedEnvironmentsByThread = new Map(entry.linkedEnvironmentsByThread); @@ -1536,6 +1566,8 @@ export const layerWithOptions = ( updated.set(key, { ...entry, attachedThreadIds, + idleTimeoutMsByThread, + idleTimeoutMs: attachedIdleTimeout(idleTimeoutMsByThread), loadedProviderThreadKeyByThread, linkedEnvironmentsByThread, }); @@ -2858,6 +2890,7 @@ export const layerWithOptions = ( return existing.exposedRuntime; } + const threadIdleTimeoutMs = yield* resolveThreadIdleTimeout(input.threadId); const adapter = yield* registry.get(input.modelSelection.instanceId).pipe( Effect.mapError( (cause) => @@ -2943,6 +2976,8 @@ export const layerWithOptions = ( const exposedRuntime = decorateRuntime(runtime, eventSubscribers); const entry: LiveSessionEntry = { attachedThreadIds: new Set([input.threadId]), + idleTimeoutMsByThread: new Map([[input.threadId, threadIdleTimeoutMs]]), + idleTimeoutMs: threadIdleTimeoutMs, loadedProviderThreadKeyByThread: new Map(), openedWithLinkedEnvironments, linkedEnvironmentsByThread: new Map(), @@ -3163,6 +3198,8 @@ export const layerWithOptions = ( } const attachedThreadIds = new Set(entry.attachedThreadIds); attachedThreadIds.delete(input.threadId); + const idleTimeoutMsByThread = new Map(entry.idleTimeoutMsByThread); + idleTimeoutMsByThread.delete(input.threadId); const loadedProviderThreadKeyByThread = new Map( entry.loadedProviderThreadKeyByThread, ); @@ -3182,6 +3219,8 @@ export const layerWithOptions = ( const updatedEntry = { ...entry, attachedThreadIds, + idleTimeoutMsByThread, + idleTimeoutMs: attachedIdleTimeout(idleTimeoutMsByThread), loadedProviderThreadKeyByThread, linkedEnvironmentsByThread, mcpCredentialIdByThread,