diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts index 40c2d28a590a9..3dff759550c78 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts @@ -1626,6 +1626,73 @@ it.effect("ProviderSessionManagerV2 gives up on a provider that never confirms a }), ); +it.effect("ProviderSessionManagerV2 still idles out a session releaseUnlessBusy kept", () => + 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("thread-provider-session-manager-kept-idles"); + const providerSessionId = idAllocator.derive.providerSession({ + providerInstanceId: modelSelection.instanceId, + }); + yield* eventSink.write({ + events: [yield* makeThreadCreatedEvent({ idAllocator, threadId, now })], + }); + const runtime = yield* manager.open({ + threadId, + providerSessionId, + modelSelection, + runtimePolicy, + }); + const providerThread = makeProviderThread({ + idAllocator, + threadId, + providerSessionId, + now, + nativeThreadId: `native-${threadId}`, + }); + const runId = idAllocator.derive.run({ threadId, ordinal: 1 }); + 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: "user", + creationSource: "web", + messageId: yield* idAllocator.allocate.message({ threadId, ordinal: 1 }), + text: "turn", + attachments: [], + }, + modelSelection, + runtimePolicy, + }); + + assert.isFalse( + yield* manager.releaseUnlessBusy({ + providerSessionId, + ownProviderThreadIds: [providerThread.id], + throughRunOrdinal: 1, + keepIf: Effect.succeed(true), + reason: "manual_shutdown", + }), + ); + assert.isTrue(yield* manager.isLive(providerSessionId)); + // The stopped run 1 never reports its end. + yield* TestClock.adjust("2 seconds"); + assert.isFalse(yield* manager.isLive(providerSessionId)); + }).pipe(Effect.provide(layerTest({ state, idleTimeoutMs: 1_000 }))); + }), +); + it.effect( "ProviderSessionManagerV2 keeps a shared session a sibling is busy on when releasing", () => @@ -1692,6 +1759,7 @@ it.effect( manager.releaseUnlessBusy({ providerSessionId, ownProviderThreadIds: [providerThreadOf(own).id], + throughRunOrdinal: 1, reason: "manual_shutdown", }); diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.ts index ade4668541362..c2f7eb767945d 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.ts @@ -71,6 +71,17 @@ export const PROVIDER_INTERRUPT_TIMEOUT_MS = 15 * 1000; const busyTurnPrefix = (providerThreadId: ProviderThreadId) => `${providerThreadId}#`; const busyTurnKey = (providerThreadId: ProviderThreadId, runOrdinal: number) => `${busyTurnPrefix(providerThreadId)}${runOrdinal}`; +interface OwnBusyTurns { + readonly providerThreadIds: ReadonlyArray; + readonly throughRunOrdinal: number; +} +const isOwnBusyTurn = (turnKey: string, own: OwnBusyTurns) => + own.providerThreadIds.some((id) => { + const prefix = busyTurnPrefix(id); + return ( + turnKey.startsWith(prefix) && Number(turnKey.slice(prefix.length)) <= own.throughRunOrdinal + ); + }); const UNLOAD_THREAD_TIMEOUT_MS = 10 * 1000; // Providers can announce a retry and never resume. Longer than any adapter's own stall handling. const PROVIDER_RETRY_STALL_TIMEOUT = Duration.minutes(15); @@ -243,13 +254,16 @@ export interface ProviderSessionManagerV2Shape { readonly revokeMcpCredential?: boolean; }) => Effect.Effect; /** - * Releases the session unless, at the moment of removal, a turn of a provider thread outside - * `ownProviderThreadIds` is busy on it. Checked atomically with turn starts. Succeeds with - * whether it was released. + * Releases the session unless, at the moment of removal, a turn is busy on it other than one + * of `ownProviderThreadIds` from a run up to `throughRunOrdinal`. Checked atomically with turn + * starts. Succeeds with whether it was released. */ readonly releaseUnlessBusy: (input: { readonly providerSessionId: ProviderSessionId; readonly ownProviderThreadIds: ReadonlyArray; + readonly throughRunOrdinal: number; + /** Keeps the session when it succeeds with true, checked atomically with the removal. */ + readonly keepIf?: Effect.Effect; readonly reason: ProviderSessionReleaseReason; readonly detail?: string; }) => Effect.Effect; @@ -1125,7 +1139,7 @@ export const layerWithOptions = ( function* (input: { readonly providerSessionId: ProviderSessionId; readonly onlyIfIdleGeneration?: number; - readonly keepIfBusyOtherThan?: ReadonlyArray; + readonly keepIfBusyOtherThan?: OwnBusyTurns; readonly keepIf?: Effect.Effect; }): Effect.fn.Return, DateTime.Utc]> { const key = sessionKey(input.providerSessionId); @@ -1133,40 +1147,48 @@ export const layerWithOptions = ( if (candidate === undefined) { return [Option.none(), yield* DateTime.now] as const; } - const remove = Effect.zip( - Ref.modify(sessions, (current) => { - const existing = current.get(key); - if (existing !== candidate) { - return [existing === undefined ? "gone" : "changed", current] as const; - } - if ( - input.onlyIfIdleGeneration !== undefined && - (existing.busyTurns.size > 0 || - existing.idleGeneration !== input.onlyIfIdleGeneration) - ) { - return ["kept", current] as const; - } - const own = input.keepIfBusyOtherThan; - if ( - own !== undefined && - [...existing.busyTurns].some( - (turnKey) => !own.some((id) => turnKey.startsWith(busyTurnPrefix(id))), - ) - ) { - return ["kept", current] as const; - } - const updated = new Map(current); - updated.delete(key); - return ["removed", updated] as const; - }), - DateTime.now, + const remove = (keepRequested: boolean) => + Effect.zip( + Ref.modify(sessions, (current) => { + const existing = current.get(key); + if (existing !== candidate) { + return [existing === undefined ? "gone" : "changed", current] as const; + } + if ( + input.onlyIfIdleGeneration !== undefined && + (existing.busyTurns.size > 0 || + existing.idleGeneration !== input.onlyIfIdleGeneration) + ) { + return ["kept", current] as const; + } + const own = input.keepIfBusyOtherThan; + if ( + keepRequested || + (own !== undefined && + [...existing.busyTurns].some((turnKey) => !isOwnBusyTurn(turnKey, own))) + ) { + if (own === undefined) return ["kept", current] as const; + // The caller settled its own turns; one whose end never arrives must not keep + // the session from going idle once the turn it was kept for ends. + const kept = new Map(current); + kept.set(key, { + ...existing, + busyTurns: new Set( + [...existing.busyTurns].filter((turnKey) => !isOwnBusyTurn(turnKey, own)), + ), + }); + return ["kept", kept] as const; + } + const updated = new Map(current); + updated.delete(key); + return ["removed", updated] as const; + }), + DateTime.now, + ); + const [outcome, releasedAt] = yield* (input.keepIf ?? Effect.succeed(false)).pipe( + Effect.flatMap(remove), + candidate.requestEventPermit.withPermits(1), ); - const [outcome, releasedAt] = yield* Effect.gen(function* () { - if (input.keepIf !== undefined && (yield* input.keepIf)) { - return ["kept" as const, yield* DateTime.now] as const; - } - return yield* remove; - }).pipe(candidate.requestEventPermit.withPermits(1)); if (outcome === "changed") return yield* removeLiveEntry(input); return [ outcome === "removed" ? Option.some(candidate) : Option.none(), @@ -1218,7 +1240,7 @@ export const layerWithOptions = ( readonly detachedThreadId?: ThreadId; readonly cancelIdleFiber?: boolean; readonly onlyIfIdleGeneration?: number; - readonly keepIfBusyOtherThan?: ReadonlyArray; + readonly keepIfBusyOtherThan?: OwnBusyTurns; readonly keepIf?: Effect.Effect; readonly gracefulSubscribers?: boolean; readonly eventStreamCause?: Cause.Cause; @@ -3076,10 +3098,19 @@ export const layerWithOptions = ( providerSessionId: input.providerSessionId, reason: input.reason, ...(input.detail === undefined ? {} : { detail: input.detail }), - keepIfBusyOtherThan: input.ownProviderThreadIds, + keepIfBusyOtherThan: { + providerThreadIds: input.ownProviderThreadIds, + throughRunOrdinal: input.throughRunOrdinal, + }, + ...(input.keepIf === undefined ? {} : { keepIf: input.keepIf }), }); - // Identity, not presence: open can install a replacement under the same id meanwhile. - return before !== undefined && (yield* Ref.get(sessions)).get(key) !== before; + // Scope, not presence: open can install a replacement under the same id meanwhile, and + // a kept entry is rewritten with the same scope. + const released = + before !== undefined && (yield* Ref.get(sessions)).get(key)?.scope !== before.scope; + // The stopped turn's markBusy cancelled the idle timer and its markIdle may never come. + if (!released) yield* scheduleIdleReleaseInternal(input.providerSessionId); + return released; }), otherWorkingThreads: Effect.fn("ProviderSessionManager.otherWorkingThreads")( function* (input) { diff --git a/apps/server/src/orchestration-v2/ThreadManagementService.ts b/apps/server/src/orchestration-v2/ThreadManagementService.ts index 1dd01ae504ca5..a69d308884c72 100644 --- a/apps/server/src/orchestration-v2/ThreadManagementService.ts +++ b/apps/server/src/orchestration-v2/ThreadManagementService.ts @@ -1303,6 +1303,28 @@ const make = Effect.gen(function* () { const ownProviderThreadIds = records.providerThreads.flatMap((providerThread) => providerThread.appThreadId === input.threadId ? [providerThread.id] : [], ); + // Settling the target can promote a queued follow-up onto the same session; that run is + // not Force stop's to end. Only a newer run bound to a provider thread on this session + // counts. An unreadable thread counts as one, so its sessions are kept. + const throughRunOrdinal = target?.ordinal ?? 0; + const laterRunActive = (providerSessionId: ProviderSessionId) => + orchestrator.getThreadRecords(input.threadId, ["runs", "providerThreads"]).pipe( + Effect.map(({ runs, providerThreads }) => { + const onSession = new Set( + providerThreads.flatMap((providerThread) => + providerThread.providerSessionId === providerSessionId ? [providerThread.id] : [], + ), + ); + return runs.some( + (run) => + run.ordinal > throughRunOrdinal && + isActiveRun(run) && + run.providerThreadId !== null && + onSession.has(run.providerThreadId), + ); + }), + Effect.orElseSucceed(() => true), + ); const detach = (providerSessionId: ProviderSessionId) => sessions .detach({ providerSessionId, threadId: input.threadId, detail: FORCE_STOP_REASON }) @@ -1340,6 +1362,9 @@ const make = Effect.gen(function* () { .releaseUnlessBusy({ providerSessionId, ownProviderThreadIds, + throughRunOrdinal, + // A promoted follow-up may not have marked its turn busy yet. + keepIf: laterRunActive(providerSessionId), reason: "manual_shutdown", detail: FORCE_STOP_REASON, }) @@ -1354,7 +1379,7 @@ const make = Effect.gen(function* () { ); if (Option.isNone(released)) { releasing = true; - } else if (!released.value) { + } else if (!released.value && !(yield* laterRunActive(release.providerSessionId))) { shared = true; yield* detach(release.providerSessionId); } diff --git a/apps/server/src/orchestration-v2/ThreadStop.test.ts b/apps/server/src/orchestration-v2/ThreadStop.test.ts index 7c6bf40044537..d07cbc7867718 100644 --- a/apps/server/src/orchestration-v2/ThreadStop.test.ts +++ b/apps/server/src/orchestration-v2/ThreadStop.test.ts @@ -20,7 +20,9 @@ import { import * as DateTime from "effect/DateTime"; import * as Effect from "effect/Effect"; import * as Exit from "effect/Exit"; +import * as Fiber from "effect/Fiber"; import * as Layer from "effect/Layer"; +import * as Queue from "effect/Queue"; import * as Schema from "effect/Schema"; import * as SqlClient from "effect/sql/SqlClient"; import * as Stream from "effect/Stream"; @@ -38,8 +40,10 @@ import * as EffectWorker from "./EffectWorker.ts"; import * as EffectOutbox from "./EffectOutbox.ts"; import * as ProjectionStore from "./ProjectionStore.ts"; import * as ProviderAdapterRegistry from "./ProviderAdapterRegistry.ts"; +import * as ProviderSessionManager from "./ProviderSessionManager.ts"; import * as ThreadManagementService from "./ThreadManagementService.ts"; import * as ProviderReplayHarness from "./testkit/ProviderReplayHarness.ts"; +import { checkpointWorkspace } from "@t3tools/provider-testing/replayWorkspace"; import type * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter"; const instanceId = ProviderInstanceId.make("codex"); @@ -1834,3 +1838,261 @@ it.effect.each(["Stop", "Force stop"] as const)( ); }).pipe(Effect.provide(layerTest)), ); + +it.effect.each([ + { followUp: "started its turn", startedBeforeRelease: true }, + { followUp: "was promoted", startedBeforeRelease: false }, + { followUp: "switched provider instance", startedBeforeRelease: true, switchesInstance: true }, + { followUp: "has fresh context", startedBeforeRelease: true, freshContext: true }, +])( + "Force stop keeps only the session of a queued follow-up that $followUp", + ({ startedBeforeRelease, switchesInstance = false, freshContext = false }) => + Effect.scoped( + Effect.gen(function* () { + const cwd = yield* checkpointWorkspace("thread-stop-force-stop-follow-up"); + const driver = adapter.driver; + const otherInstanceId = ProviderInstanceId.make("codex-other"); + const eventQueues = new Map< + ProviderSessionId, + Queue.Queue + >(); + const started: Array = []; + const interrupted: Array = []; + let opened = 0; + // Fresh context on a one-thread-per-session provider puts the follow-up on its own session. + // The session itself reports shared threads so fresh-context cleanup leaves the stopped + // run's session to Force stop, as pending background work would. + const adapterCapabilities = freshContext + ? { + ...CodexProviderCapabilitiesV2, + sessions: { + ...CodexProviderCapabilitiesV2.sessions, + supportsMultipleProviderThreadsPerSession: false, + }, + } + : CodexProviderCapabilitiesV2; + const liveAdapter = ( + instanceId: ProviderInstanceId, + ): ProviderAdapter.ProviderAdapterV2["Service"] => ({ + ...adapter, + instanceId, + getCapabilities: () => Effect.succeed(adapterCapabilities), + openSession: (input) => + Effect.gen(function* () { + opened += 1; + const now = yield* DateTime.now; + const events = yield* Queue.unbounded(); + eventQueues.set(input.providerSessionId, events); + return { + instanceId, + driver, + providerSessionId: input.providerSessionId, + providerSession: { + id: input.providerSessionId, + driver, + providerInstanceId: instanceId, + status: "ready", + cwd, + model: modelSelection.model, + capabilities: CodexProviderCapabilitiesV2, + createdAt: now, + updatedAt: now, + lastError: null, + }, + events: Stream.fromQueue(events), + ensureThread: ({ threadId }) => + Effect.succeed({ + id: ProviderThreadId.make( + `provider-thread:${input.providerSessionId}:${threadId}`, + ), + driver, + providerInstanceId: instanceId, + providerSessionId: input.providerSessionId, + appThreadId: threadId, + ownerNodeId: null, + nativeThreadRef: { driver, nativeId: "native-thread", strength: "strong" }, + nativeConversationHeadRef: null, + status: "idle", + firstRunOrdinal: null, + lastRunOrdinal: null, + handoffIds: [], + forkedFrom: null, + createdAt: now, + updatedAt: now, + }), + resumeThread: ({ providerThread }) => Effect.succeed(providerThread), + startTurn: (turn) => + Effect.gen(function* () { + started.push(turn); + yield* Queue.offer(events, { + type: "provider_turn.updated", + driver, + providerTurn: { + id: ProviderTurnId.make(`provider-turn:${turn.attemptId}`), + providerThreadId: turn.providerThread.id, + nodeId: turn.rootNodeId, + runAttemptId: turn.attemptId, + nativeTurnRef: { + driver, + nativeId: `native:${turn.attemptId}`, + strength: "strong", + }, + ordinal: turn.providerTurnOrdinal, + status: "running", + startedAt: now, + completedAt: null, + }, + }); + }), + steerTurn: () => Effect.die("unused"), + interruptTurn: (interrupt) => Effect.sync(() => interrupted.push(interrupt)), + respondToRuntimeRequest: () => Effect.die("unused"), + readThreadSnapshot: () => Effect.die("unused"), + rollbackThread: () => Effect.die("unused"), + forkThread: () => Effect.die("unused"), + }; + }), + }); + // The settlement promotes the follow-up; optionally its turn starts before Force stop releases. + const layerSessionsStartingFollowUpFirst = Layer.effect( + ProviderSessionManager.ProviderSessionManagerV2, + Effect.gen(function* () { + const sessions = yield* ProviderSessionManager.ProviderSessionManagerV2; + const worker = yield* EffectWorker.OrchestrationEffectWorkerV2; + return ProviderSessionManager.ProviderSessionManagerV2.of({ + ...sessions, + releaseUnlessBusy: (input) => + (startedBeforeRelease ? worker.drain().pipe(Effect.orDie) : Effect.void).pipe( + Effect.andThen(sessions.releaseUnlessBusy(input)), + ), + }); + }), + ); + const layerLive = ThreadManagementService.layer.pipe( + Layer.provide(layerSessionsStartingFollowUpFirst), + Layer.provideMerge( + Layer.mergeAll( + layerDatabase, + ProjectionStore.layer.pipe(Layer.provide(layerDatabase)), + EventStore.layer.pipe(Layer.provide(layerDatabase)), + EffectOutbox.layer.pipe(Layer.provide(layerDatabase)), + ProviderReplayHarness.layerWithRegistry( + { name: "thread-stop-force-stop-follow-up" }, + ProviderAdapterRegistry.layerFromAdapters([ + liveAdapter(instanceId), + liveAdapter(otherInstanceId), + ]), + { databaseLayer: layerDatabase, runEffectWorker: false }, + ), + ), + ), + Layer.provideMerge(ServerSettings.layerTest()), + Layer.provideMerge(CommandReceiptStore.layer.pipe(Layer.provide(layerDatabase))), + ); + + yield* Effect.gen(function* () { + const orchestrator = yield* Orchestrator.OrchestratorV2; + const worker = yield* EffectWorker.OrchestrationEffectWorkerV2; + const sessions = yield* ProviderSessionManager.ProviderSessionManagerV2; + const threads = yield* ThreadManagementService.ThreadManagementService; + const threadId = ThreadId.make("thread:force-stop-follow-up"); + const runningTurn = orchestrator.streamDomainEvents.pipe( + Stream.filter( + (event) => + event.type === "provider-turn.updated" && event.payload.status === "running", + ), + Stream.take(1), + Stream.runDrain, + Effect.forkScoped, + ); + yield* orchestrator.dispatch({ + type: "thread.create", + commandId: CommandId.make("create:force-stop-follow-up"), + threadId, + projectId: ProjectId.make("project:thread-stop"), + title: "Force stop follow-up", + modelSelection, + runtimeMode: "full-access", + interactionMode: "default", + branch: null, + worktreePath: cwd, + createdBy: "user", + creationSource: "web", + }); + const firstRunning = yield* runningTurn; + yield* send(threadId, "first", "start_immediately"); + yield* worker.drain(); + yield* Fiber.join(firstRunning); + yield* orchestrator.dispatch({ + type: "message.dispatch", + commandId: CommandId.make("send:force-stop-follow-up"), + threadId, + messageId: MessageId.make("message:force-stop-follow-up"), + text: "follow-up", + attachments: [], + dispatchMode: { type: "queue_after_active" }, + ...(switchesInstance + ? { modelSelection: { ...modelSelection, instanceId: otherInstanceId } } + : {}), + ...(freshContext ? { freshContext } : {}), + createdBy: "user", + creationSource: "web", + }); + + const result = yield* threads.forceStopThread({ threadId }); + yield* worker.drain(); + + assert.deepEqual( + started.map((turn) => turn.runOrdinal), + [1, 2], + "the queued follow-up did not start its turn", + ); + assert.equal( + opened, + switchesInstance || freshContext ? 2 : 1, + "Force stop tore down the session the follow-up was promoted onto", + ); + const [stopped, followUp] = [started[0]!, started[1]!]; + const providerSessionId = followUp.providerThread.providerSessionId!; + assert.equal( + yield* sessions.isLive(stopped.providerThread.providerSessionId!), + !switchesInstance && !freshContext, + "Force stop kept the stopped run's session for a follow-up on another session", + ); + assert.isFalse( + interrupted.some((interrupt) => + String(interrupt.providerTurnId).includes(followUp.attemptId), + ), + "Force stop of the first run interrupted the follow-up's turn", + ); + assert.isTrue( + yield* sessions.isLive(providerSessionId), + "Force stop of the first run released the session the follow-up's turn is running on", + ); + assert.deepEqual( + result, + { outcome: "released", sharedWithThreadIds: [] }, + "Force stop detached the thread the follow-up runs on as if its session were shared", + ); + + // The first run's provider never confirms the stop; the follow-up's turn ends normally. + yield* Queue.offer(eventQueues.get(providerSessionId)!, { + type: "turn.terminal", + driver, + providerThreadId: followUp.providerThread.id, + providerTurnId: ProviderTurnId.make(`provider-turn:${followUp.attemptId}`), + runOrdinal: followUp.runOrdinal, + status: "completed", + failure: null, + threadDisposition: "reusable", + }); + yield* worker.drain(); + yield* TestClock.adjust("31 minutes"); + assert.isFalse( + yield* sessions.isLive(providerSessionId), + "the stopped run's unconfirmed turn kept the session from going idle", + ); + }).pipe(Effect.provide(layerLive)); + }), + ), +); diff --git a/packages/contracts/src/orchestrationV2.ts b/packages/contracts/src/orchestrationV2.ts index 24e3a48ba76dc..35ac16fc063a4 100644 --- a/packages/contracts/src/orchestrationV2.ts +++ b/packages/contracts/src/orchestrationV2.ts @@ -4180,10 +4180,10 @@ export type OrchestrationV2ForceStopThreadInput = typeof OrchestrationV2ForceSto export const OrchestrationV2ForceStopThreadResult = Schema.Struct({ /** - * `killed`: the thread's own agent processes were killed. `released`: its provider session - * was torn down, which kills the agent's process group. `releasing`: that teardown has not - * finished yet and continues on the server. `shared`: the agent process also works for - * `sharedWithThreadIds`, so only this thread's turn was ended. + * `killed`: its own agent processes were killed. `released`: its provider session was torn + * down (killing the agent's process group) or handed to its newer run. `releasing`: teardown + * continues on the server. `shared`: the agent process also works for `sharedWithThreadIds`, + * so only this thread's turn was ended. */ outcome: Schema.Literals(["killed", "released", "releasing", "shared"]), sharedWithThreadIds: Schema.Array(ThreadId),