diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts index 696a40d8d176..e3d94a0914f8 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts @@ -27,6 +27,7 @@ import * as Layer from "effect/Layer"; import * as Option from "effect/Option"; import * as Queue from "effect/Queue"; import * as Ref from "effect/Ref"; +import * as Scheduler from "effect/Scheduler"; import * as Stream from "effect/Stream"; import { TestClock } from "effect/testing"; import { HttpServer } from "effect/unstable/http"; @@ -81,6 +82,65 @@ const FailingReleaseEventSinkLayer = Layer.effect( }), ).pipe(Layer.provide(TestEventSinkLayer)); +/** + * Pauses the write that attaches a just-opened session to its thread: + * signals `started`, then waits for `resume`. Lets a test interrupt `open()` + * while its entry is already recorded in `sessions` but setup has not + * finished, the window after `Ref.update(sessions, ...)` and before + * `startEventPump` / `scheduleIdleRelease`. + */ +const makeAttachWritePauseLayer = (gate: { + readonly started: Deferred.Deferred; + readonly resume: Deferred.Deferred; +}) => + Layer.effect( + EventSink.EventSinkV2, + Effect.gen(function* () { + const delegate = yield* EventSink.EventSinkV2; + return EventSink.EventSinkV2.of({ + ...delegate, + write: (input) => + Effect.gen(function* () { + if (input.events.some((event) => event.type === "provider-session.attached")) { + yield* Deferred.succeed(gate.started, undefined); + yield* Deferred.await(gate.resume); + } + return yield* delegate.write(input); + }), + }); + }), + ).pipe(Layer.provide(TestEventSinkLayer)); + +/** + * Pauses the first `revokeProviderSession` call: signals `started`, then + * waits for `resume`. Lets a test hold a credential revoke in flight (and, + * with it, any per-thread lock the caller acquired around the revoke) while + * a concurrent prepare for the same thread is attempted. + */ +const makeRevokePauseRegistryLayer = (gate: { + readonly started: Deferred.Deferred; + readonly resume: Deferred.Deferred; +}) => + Layer.effect( + McpSessionRegistry.McpSessionRegistry, + Effect.gen(function* () { + const delegate = yield* McpSessionRegistry.McpSessionRegistry; + let paused = false; + return McpSessionRegistry.McpSessionRegistry.of({ + ...delegate, + revokeProviderSession: (providerSessionId) => + Effect.gen(function* () { + if (!paused) { + paused = true; + yield* Deferred.succeed(gate.started, undefined); + yield* Deferred.await(gate.resume); + } + yield* delegate.revokeProviderSession(providerSessionId); + }), + }); + }), + ).pipe(Layer.provide(TestMcpRegistryLayer)); + interface FlakyReleaseWrites { /** Which release writes fail right now. */ readonly failing: Ref.Ref<"none" | "session" | "session-and-requests">; @@ -293,6 +353,8 @@ function makeProviderAdapter( }) => Effect.Effect; readonly hasPendingBackgroundWork?: Effect.Effect; readonly hangSessionScopeClose?: boolean; + /** Runs after the session's resources are registered, before open returns. */ + readonly afterOpen?: Effect.Effect; readonly beforeUnload?: Effect.Effect; } = {}, ): ProviderAdapterV2Shape { @@ -340,6 +402,9 @@ function makeProviderAdapter( // that never yields its message stream. yield* Effect.addFinalizer(() => Effect.never); } + if (options.afterOpen !== undefined) { + yield* options.afterOpen; + } return { instanceId: ProviderInstanceId.make("codex"), @@ -407,8 +472,17 @@ function makeTestLayer(input: { }) => Effect.Effect; readonly failReleaseEventWrites?: boolean; readonly flakyReleaseWrites?: FlakyReleaseWrites; + readonly attachWritePause?: { + readonly started: Deferred.Deferred; + readonly resume: Deferred.Deferred; + }; + readonly revokePause?: { + readonly started: Deferred.Deferred; + readonly resume: Deferred.Deferred; + }; readonly hasPendingBackgroundWork?: Effect.Effect; readonly hangSessionScopeClose?: boolean; + readonly afterOpen?: Effect.Effect; readonly beforeUnload?: Effect.Effect; readonly serverSettingsLayer?: ReturnType; readonly projectServiceLayer?: Layer.Layer; @@ -418,7 +492,9 @@ function makeTestLayer(input: { ? makeFlakyReleaseEventSinkLayer(input.flakyReleaseWrites) : input.failReleaseEventWrites ? FailingReleaseEventSinkLayer - : TestEventSinkLayer; + : input.attachWritePause !== undefined + ? makeAttachWritePauseLayer(input.attachWritePause) + : TestEventSinkLayer; const registryLayer = ProviderAdapterRegistry.makeSingleLayer( makeProviderAdapter(input.state, { failEventStream: input.failEventStream ?? false, @@ -431,6 +507,7 @@ function makeTestLayer(input: { ...(input.hangSessionScopeClose === undefined ? {} : { hangSessionScopeClose: input.hangSessionScopeClose }), + ...(input.afterOpen === undefined ? {} : { afterOpen: input.afterOpen }), ...(input.beforeUnload === undefined ? {} : { beforeUnload: input.beforeUnload }), }), ); @@ -444,11 +521,15 @@ function makeTestLayer(input: { ), ), ); + const configuredMcpRegistryLayer = + input.revokePause !== undefined + ? makeRevokePauseRegistryLayer(input.revokePause) + : TestMcpRegistryLayer; return Layer.mergeAll( TestStoresLayer, configuredEventSinkLayer, IdAllocator.layer, - TestMcpRegistryLayer, + configuredMcpRegistryLayer, ProviderSessionManager.layerWithOptions({ idleTimeoutMs: input.idleTimeoutMs, ...(input.maxIdlePinMs === undefined ? {} : { maxIdlePinMs: input.maxIdlePinMs }), @@ -459,7 +540,7 @@ function makeTestLayer(input: { configuredEventSinkLayer, IdAllocator.layer, providerEventIngestorTestLayer, - TestMcpRegistryLayer, + configuredMcpRegistryLayer, TestStoresLayer, ...(input.serverSettingsLayer === undefined ? [] : [input.serverSettingsLayer]), ...(input.projectServiceLayer === undefined ? [] : [input.projectServiceLayer]), @@ -776,6 +857,598 @@ it.effect("ProviderSessionManagerV2 opens independent sessions concurrently", () }), ); +it.effect("ProviderSessionManagerV2 closes the session scope when an open is interrupted", () => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const openBlocked = yield* Deferred.make(); + const effect = Effect.gen(function* () { + const eventSink = yield* EventSink.EventSinkV2; + const idAllocator = yield* IdAllocator.IdAllocatorV2; + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const now = yield* DateTime.now; + const threadId = ThreadId.make("thread-provider-session-manager-interrupted-open"); + const providerSessionId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + yield* eventSink.write({ + events: [yield* makeThreadCreatedEvent({ idAllocator, threadId, now })], + }); + + // Like a start timeout or a Stop interrupting an open that spawned its + // provider process and then never heard back from it. + const fiber = yield* manager + .open({ threadId, providerSessionId, modelSelection, runtimePolicy }) + .pipe(Effect.forkScoped); + yield* Deferred.await(openBlocked); + yield* Fiber.interrupt(fiber); + + assert.equal((yield* Ref.get(state)).openCount, 1); + assert.equal((yield* Ref.get(state)).closeCount, 1); + assert.isTrue(Option.isNone(yield* manager.get(providerSessionId))); + }); + + yield* effect.pipe( + Effect.provide( + makeTestLayer({ + state, + idleTimeoutMs: 60_000, + afterOpen: Deferred.succeed(openBlocked, undefined).pipe(Effect.andThen(Effect.never)), + }), + ), + ); + }), +); + +it.effect("ProviderSessionManagerV2 finishes an interrupted open when its scope close hangs", () => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const openBlocked = yield* Deferred.make(); + const effect = Effect.gen(function* () { + const eventSink = yield* EventSink.EventSinkV2; + const idAllocator = yield* IdAllocator.IdAllocatorV2; + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const now = yield* DateTime.now; + const threadId = ThreadId.make("thread-provider-session-manager-interrupted-hung-open"); + const providerSessionId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + yield* eventSink.write({ + events: [yield* makeThreadCreatedEvent({ idAllocator, threadId, now })], + }); + + const fiber = yield* manager + .open({ threadId, providerSessionId, modelSelection, runtimePolicy }) + .pipe(Effect.forkScoped); + yield* Deferred.await(openBlocked); + // A start timeout interrupting an open whose provider process never + // yields its message stream must still get its thread lane back. + const interrupter = yield* Fiber.interrupt(fiber).pipe(Effect.forkScoped); + yield* TestClock.adjust("29 seconds"); + assert.isUndefined(interrupter.pollUnsafe()); + + yield* TestClock.adjust("1 second"); + yield* Fiber.join(interrupter); + assert.equal((yield* Ref.get(state)).closeCount, 0); + assert.isTrue(Option.isNone(yield* manager.get(providerSessionId))); + }); + + yield* effect.pipe( + Effect.provide( + makeTestLayer({ + state, + idleTimeoutMs: 60_000, + hangSessionScopeClose: true, + afterOpen: Deferred.succeed(openBlocked, undefined).pipe(Effect.andThen(Effect.never)), + }), + ), + ); + }), +); + +it.effect( + "ProviderSessionManagerV2 keeps a credential a concurrent open still holds when the first open is interrupted", + () => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const openStartedCount = yield* Ref.make(0); + const firstStarted = yield* Deferred.make(); + const secondStarted = yield* Deferred.make(); + const releaseSecond = yield* Deferred.make(); + // sessionOpen serializes by providerSessionId, not threadId, so the + // first and second open below run concurrently. Each call's own + // prepareMcpSession (serialized per thread) finishes, and only then + // does it reach this hook, so call order here matches prepare order: + // the first call mints the credential, the second reuses it. + const beforeOpen = () => + Effect.gen(function* () { + const callNumber = yield* Ref.modify(openStartedCount, (count) => [count + 1, count + 1]); + if (callNumber === 1) { + yield* Deferred.succeed(firstStarted, undefined); + return yield* Effect.never; + } else { + yield* Deferred.succeed(secondStarted, undefined); + yield* Deferred.await(releaseSecond); + } + }); + + const effect = Effect.gen(function* () { + const eventSink = yield* EventSink.EventSinkV2; + const idAllocator = yield* IdAllocator.IdAllocatorV2; + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const registry = yield* McpSessionRegistry.McpSessionRegistry; + const now = yield* DateTime.now; + const threadId = ThreadId.make("thread-provider-session-manager-open-credential-race"); + const s1 = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + const s2 = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + yield* eventSink.write({ + events: [yield* makeThreadCreatedEvent({ idAllocator, threadId, now })], + }); + + const firstFiber = yield* manager + .open({ threadId, providerSessionId: s1, modelSelection, runtimePolicy }) + .pipe(Effect.forkScoped); + yield* Deferred.await(firstStarted); + const originalToken = McpProviderSession.readMcpProviderSession( + threadId, + )?.authorizationHeader.replace(/^Bearer\s+/, ""); + assert.isDefined(originalToken); + + const secondFiber = yield* manager + .open({ threadId, providerSessionId: s2, modelSelection, runtimePolicy }) + .pipe(Effect.forkScoped); + yield* Deferred.await(secondStarted); + + yield* Fiber.interrupt(firstFiber); + assert.isDefined( + yield* registry.resolve(originalToken!), + "interrupting s1 must not revoke the credential s2 is still configuring its process with", + ); + + yield* Deferred.succeed(releaseSecond, undefined); + yield* Fiber.join(secondFiber); + assert.isDefined(yield* registry.resolve(originalToken!)); + yield* manager.close(s2); + assert.isUndefined( + yield* registry.resolve(originalToken!), + "closing the session that actually holds the credential still revokes it", + ); + }); + + yield* effect.pipe( + Effect.provide( + makeTestLayer({ + state, + idleTimeoutMs: 60_000, + beforeOpen, + }), + ), + ); + }), +); + +it.effect( + "ProviderSessionManagerV2 releases an entry when open is interrupted after it is recorded", + () => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const gate = { + started: yield* Deferred.make(), + resume: yield* Deferred.make(), + }; + const effect = Effect.gen(function* () { + const eventSink = yield* EventSink.EventSinkV2; + const idAllocator = yield* IdAllocator.IdAllocatorV2; + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const now = yield* DateTime.now; + const threadId = ThreadId.make("thread-provider-session-manager-interrupted-attach"); + const providerSessionId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + yield* eventSink.write({ + events: [yield* makeThreadCreatedEvent({ idAllocator, threadId, now })], + }); + + // Like a Stop landing after the session is recorded in `sessions` + // (so manager.get sees it) but before setup finishes persisting the + // attachment, starting the event pump, and scheduling idle release. + const fiber = yield* manager + .open({ threadId, providerSessionId, modelSelection, runtimePolicy }) + .pipe(Effect.forkScoped); + yield* Deferred.await(gate.started); + assert.isTrue( + Option.isSome(yield* manager.get(providerSessionId)), + "the entry must already be recorded while the attach write is in flight", + ); + + yield* Fiber.interrupt(fiber); + + assert.isTrue( + Option.isNone(yield* manager.get(providerSessionId)), + "an open interrupted after recording its entry must release it", + ); + assert.equal((yield* Ref.get(state)).closeCount, 1); + }); + + yield* effect.pipe( + Effect.provide( + makeTestLayer({ + state, + idleTimeoutMs: 60_000, + attachWritePause: gate, + }), + ), + ); + }), +); + +it.effect( + "ProviderSessionManagerV2 blocks a concurrent prepare for the same thread while revoking an abandoned open's credential", + () => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const openStartedCount = yield* Ref.make(0); + const firstStarted = yield* Deferred.make(); + const secondAdapterStarted = yield* Deferred.make(); + const revokeGate = { + started: yield* Deferred.make(), + resume: yield* Deferred.make(), + }; + // Call 1 (s1) mints the credential and parks, holding it open forever + // until interrupted. Call 2 (s2) only signals once it reaches the + // adapter, which happens strictly after its own prepareMcpSession (and + // therefore the per-thread lock it needs) has returned. + const beforeOpen = () => + Effect.gen(function* () { + const callNumber = yield* Ref.modify(openStartedCount, (count) => [count + 1, count + 1]); + if (callNumber === 1) { + yield* Deferred.succeed(firstStarted, undefined); + return yield* Effect.never; + } + yield* Deferred.succeed(secondAdapterStarted, undefined); + }); + + const effect = Effect.gen(function* () { + const eventSink = yield* EventSink.EventSinkV2; + const idAllocator = yield* IdAllocator.IdAllocatorV2; + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const registry = yield* McpSessionRegistry.McpSessionRegistry; + const now = yield* DateTime.now; + const threadId = ThreadId.make("thread-provider-session-manager-revoke-lock-race"); + const s1 = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + const s2 = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + yield* eventSink.write({ + events: [yield* makeThreadCreatedEvent({ idAllocator, threadId, now })], + }); + + const first = yield* manager + .open({ threadId, providerSessionId: s1, modelSelection, runtimePolicy }) + .pipe(Effect.forkScoped); + yield* Deferred.await(firstStarted); + + // Interrupt s1: abandonOpen's revokeIfUnheld finds nothing else + // holds the credential yet, decides to revoke, and (via the paused + // registry) is now parked mid-revoke, still holding the per-thread + // prepare lock. + const stopping = yield* Fiber.interrupt(first).pipe(Effect.forkScoped); + yield* Deferred.await(revokeGate.started); + + // s2 starts for the same thread while the revoke is in flight. If + // the revoke and the lock are not serialized, s2's prepare can run + // right now and adopt the credential s1 is about to revoke anyway. + const second = yield* manager + .open({ threadId, providerSessionId: s2, modelSelection, runtimePolicy }) + .pipe(Effect.forkScoped); + const secondAdapterStartedWaiter = yield* Deferred.await(secondAdapterStarted).pipe( + Effect.forkScoped, + ); + // Give the scheduler every chance to run s2 as far as it can go + // without the lock before checking whether it is still blocked: a + // bare pollUnsafe right after forking can pass even on unfixed code + // simply because s2 has not been scheduled yet. + for (let i = 0; i < 200; i++) { + yield* Effect.yieldNow; + } + // Capture the result instead of asserting here: without the lock s2 + // is not actually blocked, so an immediate assert would throw before + // the gate below ever releases s1's parked, uninterruptible revoke + // handler, and the whole fiber tree (and this test) would hang + // instead of failing cleanly. + const secondStillWaitingOnLock = secondAdapterStartedWaiter.pollUnsafe() === undefined; + + yield* Deferred.succeed(revokeGate.resume, undefined); + yield* Fiber.join(stopping); + yield* Fiber.join(secondAdapterStartedWaiter); + yield* Fiber.join(second); + + assert.isTrue( + secondStillWaitingOnLock, + "s2's prepare must still be waiting on the per-thread lock while s1's revoke is in flight", + ); + + const current = McpProviderSession.readMcpProviderSession(threadId); + assert.isDefined(current, "s2 must end up with its own valid credential"); + const token = current!.authorizationHeader.replace(/^Bearer\s+/, ""); + assert.isDefined( + yield* registry.resolve(token), + "s2's credential must not be the one s1 revoked out from under it", + ); + }); + + yield* effect.pipe( + Effect.provide( + makeTestLayer({ + state, + idleTimeoutMs: 60_000, + beforeOpen, + revokePause: revokeGate, + }), + ), + ); + }), +); + +it.effect( + "ProviderSessionManagerV2 revokes an abandoned open's credential despite an unrelated older entry attached to the same thread", + () => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const openBlocked = yield* Deferred.make(); + // Only the interrupted open (the second adapter call) blocks; the + // older entry's two opens (minting its own credential, then attaching + // to the contested thread) both complete normally. + const afterOpen = Ref.get(state).pipe( + Effect.flatMap((s) => + s.openCount === 2 + ? Deferred.succeed(openBlocked, undefined).pipe(Effect.andThen(Effect.never)) + : Effect.void, + ), + ); + const effect = Effect.gen(function* () { + const eventSink = yield* EventSink.EventSinkV2; + const idAllocator = yield* IdAllocator.IdAllocatorV2; + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const registry = yield* McpSessionRegistry.McpSessionRegistry; + const now = yield* DateTime.now; + const ownThreadId = ThreadId.make("thread-provider-session-manager-holder-own"); + const sharedThreadId = ThreadId.make("thread-provider-session-manager-holder-shared"); + const olderEntryId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId: ownThreadId, + }); + const abandonedId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId: sharedThreadId, + }); + yield* eventSink.write({ + events: [ + yield* makeThreadCreatedEvent({ idAllocator, threadId: ownThreadId, now }), + yield* makeThreadCreatedEvent({ idAllocator, threadId: sharedThreadId, now }), + ], + }); + + // The older entry mints its own credential X on its own thread, then + // also attaches to the shared thread (recording a credential of its + // own there too, immaterial to the test other than that it is not + // the credential the abandoned open later mints). + yield* manager.open({ + threadId: ownThreadId, + providerSessionId: olderEntryId, + modelSelection, + runtimePolicy, + }); + const ownToken = McpProviderSession.readMcpProviderSession( + ownThreadId, + )?.authorizationHeader.replace(/^Bearer\s+/, ""); + assert.isDefined( + ownToken, + "the older entry must hold its own credential on its own thread", + ); + yield* manager.open({ + threadId: sharedThreadId, + providerSessionId: olderEntryId, + modelSelection, + runtimePolicy, + }); + + // The credential the older entry recorded for the shared thread dies + // externally (e.g. expiry), so the next open for that thread rotates + // to a fresh credential instead of reusing it. The older entry stays + // attached to the shared thread with its now-stale record. + yield* registry.revokeThread(sharedThreadId); + + const fiber = yield* manager + .open({ + threadId: sharedThreadId, + providerSessionId: abandonedId, + modelSelection, + runtimePolicy, + }) + .pipe(Effect.forkScoped); + yield* Deferred.await(openBlocked); + const abandonedToken = McpProviderSession.readMcpProviderSession( + sharedThreadId, + )?.authorizationHeader.replace(/^Bearer\s+/, ""); + assert.isDefined(abandonedToken, "the abandoned open must have minted its own credential"); + assert.notEqual(abandonedToken, ownToken); + + yield* Fiber.interrupt(fiber); + + assert.isUndefined( + yield* registry.resolve(abandonedToken!), + "nothing holds the abandoned open's credential on the shared thread; it must be revoked despite the older entry's unrelated attachment there", + ); + assert.isDefined( + yield* registry.resolve(ownToken!), + "the older entry's own credential on its own thread must survive", + ); + }); + + yield* effect.pipe( + Effect.provide( + makeTestLayer({ + state, + idleTimeoutMs: 60_000, + afterOpen, + }), + ), + ); + }), +); + +// Forces an interrupt to land at a chosen scheduler check for open()'s own +// fiber, counted from when this test's afterOpen hook runs (so armed starts +// after the MCP credential is minted and the adapter call has returned, but +// before setup finishes). A plain Deferred-signaled race cannot reliably +// land an interrupt in a window with no real suspension in it; forcing +// shouldYield to fire can, because it makes the runtime actually stop there +// before deciding what runs next. delay is not meant to target a specific +// line: it is swept across a range by the caller so the case is exercised +// at every scheduler check in and around the publish, not just whichever +// one happened to reproduce the bug first. +function runPublishInterruptCase(delay: number) { + return Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const mcpConfigs = yield* Ref.make< + ReadonlyArray + >([]); + const parentScheduler = yield* Scheduler.Scheduler; + let armed = false; + let targetFiber: Fiber.Fiber | undefined; + let sinceArmed = 0; + let fired = false; + let pendingInterrupt: Fiber.Fiber | undefined; + const scheduler: Scheduler.Scheduler = { + executionMode: parentScheduler.executionMode, + shouldYield(fiber) { + if (armed && !fired && fiber === targetFiber) { + sinceArmed++; + if (sinceArmed > delay) { + fired = true; + pendingInterrupt = fiber; + return true; + } + } + return parentScheduler.shouldYield(fiber); + }, + makeDispatcher() { + const delegate = parentScheduler.makeDispatcher(); + return { + scheduleTask(task, priority) { + const target = pendingInterrupt; + pendingInterrupt = undefined; + delegate.scheduleTask(() => { + if (target !== undefined) target.interruptUnsafe(); + task(); + }, priority); + }, + flush() { + delegate.flush(); + }, + }; + }, + }; + + const effect = Effect.gen(function* () { + const eventSink = yield* EventSink.EventSinkV2; + const idAllocator = yield* IdAllocator.IdAllocatorV2; + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const registry = yield* McpSessionRegistry.McpSessionRegistry; + const now = yield* DateTime.now; + const threadId = ThreadId.make(`thread-provider-session-manager-publish-interrupt-${delay}`); + const providerSessionId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + yield* eventSink.write({ + events: [yield* makeThreadCreatedEvent({ idAllocator, threadId, now })], + }); + + const fiber = yield* manager + .open({ threadId, providerSessionId, modelSelection, runtimePolicy }) + .pipe(Effect.provideService(Scheduler.Scheduler, scheduler), Effect.forkScoped); + targetFiber = fiber; + const exit = yield* Fiber.await(fiber); + + assert.isTrue(fired, `delay=${delay}: the probe must have armed and forced the interrupt`); + assert.isTrue(Exit.isFailure(exit), `delay=${delay}`); + assert.equal( + (yield* Ref.get(state)).closeCount, + 1, + `delay=${delay}: interrupt must close the scope before leaving open`, + ); + assert.isTrue( + Option.isNone(yield* manager.get(providerSessionId)), + `delay=${delay}: interrupt must not leave a published entry with no cleanup`, + ); + + // The adapter captures the credential the moment openSession is + // called, before this test's afterOpen hook (and so before armed is + // even set): whichever scheduler check the interrupt lands on, the + // credential was already minted and reserved by then. + const captured = (yield* Ref.get(mcpConfigs))[0]; + assert.isDefined(captured, `delay=${delay}: the open must have minted a credential`); + const token = captured?.authorizationHeader.replace(/^Bearer\s+/, ""); + // clearMcpSession's credential-id branch, which performs both of + // these, only runs once the reservation count for this thread and + // credential has dropped back to 0 (releaseEntry's holder check + // otherwise treats the reservation itself as a holder and skips + // revoking). Observing both confirms the reservation was actually + // dropped, not just that cleanup ran some other way. + assert.isUndefined( + McpProviderSession.readMcpProviderSession(threadId), + `delay=${delay}: an abandoned open must not leave its credential in the thread's slot`, + ); + assert.isUndefined( + yield* registry.resolve(token!), + `delay=${delay}: an abandoned open must revoke the credential it reserved, not just publish-and-leak it`, + ); + }); + + yield* effect.pipe( + Effect.provide( + makeTestLayer({ + state, + idleTimeoutMs: 60_000, + mcpConfigs, + afterOpen: Effect.sync(() => { + armed = true; + }), + }), + ), + ); + }); +} + +it.effect( + "ProviderSessionManagerV2 releases the entry, and drops its MCP credential reservation, when interrupted exactly as it publishes", + () => + Effect.gen(function* () { + // Sweeping every delay instead of a single fixed one demonstrates the + // fix holds at whichever scheduler check the interrupt lands on, not + // just the one value (7) that reproduced the bug before the fix. Each + // iteration is a few milliseconds of virtual scheduling, so the full + // sweep is cheap enough to keep as a committed regression test. + for (let delay = 0; delay <= 20; delay++) { + yield* runPublishInterruptCase(delay); + } + }), +); + it.effect("ProviderSessionManagerV2 closes every live session for a provider instance", () => Effect.gen(function* () { const state = yield* Ref.make(emptyState); diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.ts index 20e240e79c58..d42b4aadc879 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.ts @@ -62,6 +62,51 @@ export const ProviderSessionReleaseReason = Schema.Literals([ ]); export type ProviderSessionReleaseReason = typeof ProviderSessionReleaseReason.Type; +/** + * Closes a provider session scope, waiting at most + * `RELEASE_SCOPE_CLOSE_TIMEOUT_MS`. Scope close can wedge on a misbehaving + * adapter finalizer (e.g. a provider process that never yields its message + * stream). Past the bound, the close keeps running detached and its outcome is + * logged, so the caller gets its thread lane back instead of waiting on a + * process that may never let go. + */ +const closeSessionScopeBounded = ( + scope: Scope.Closeable, + annotations: { + readonly providerSessionId: ProviderSessionId; + readonly reason: ProviderSessionReleaseReason | "open_abandoned"; + }, +) => + Effect.gen(function* () { + const closeFiber = yield* Scope.close(scope, Exit.void).pipe( + Effect.exit, + Effect.forkDetach({ startImmediately: true }), + ); + const closeExit = yield* Fiber.join(closeFiber).pipe( + Effect.timeoutOption(RELEASE_SCOPE_CLOSE_TIMEOUT_MS), + ); + if (Option.isSome(closeExit)) return closeExit; + yield* Effect.logWarning("orchestration-v2.provider-session-scope-close-timeout", { + ...annotations, + timeoutMs: RELEASE_SCOPE_CLOSE_TIMEOUT_MS, + }); + yield* Fiber.join(closeFiber).pipe( + Effect.flatMap((exit) => + Exit.isFailure(exit) + ? Effect.logWarning("orchestration-v2.provider-session-scope-close-failed", { + ...annotations, + cause: exit.cause, + }) + : Effect.logInfo( + "orchestration-v2.provider-session-scope-close-completed-late", + annotations, + ), + ), + Effect.forkDetach, + ); + return closeExit; + }); + /** * ProviderSessionManager owns live session residency: open sessions, idle release, * explicit shutdown, and release-on-runtime-failure. @@ -882,49 +927,10 @@ export const layerWithOptions = ( input.detail ?? `Provider session released: ${input.reason}.`, ); } - // Scope close can wedge on a misbehaving adapter finalizer - // (e.g. a provider process that never yields its message - // stream). Time-box it so release still persists released - // events and leaves a diagnosable trail instead of silently - // parking the session as "ready" forever. - const closeFiber = yield* Scope.close(entry.scope, Exit.void).pipe( - Effect.exit, - Effect.forkDetach({ startImmediately: true }), - ); - const closeExit = yield* Fiber.join(closeFiber).pipe( - Effect.timeoutOption(RELEASE_SCOPE_CLOSE_TIMEOUT_MS), - ); - if (Option.isNone(closeExit)) { - yield* Effect.logWarning( - "orchestration-v2.provider-session-scope-close-timeout", - { - providerSessionId: input.providerSessionId, - reason: input.reason, - timeoutMs: RELEASE_SCOPE_CLOSE_TIMEOUT_MS, - }, - ); - yield* Fiber.join(closeFiber).pipe( - Effect.flatMap((exit) => - Exit.isFailure(exit) - ? Effect.logWarning( - "orchestration-v2.provider-session-scope-close-failed", - { - providerSessionId: input.providerSessionId, - reason: input.reason, - cause: exit.cause, - }, - ) - : Effect.logInfo( - "orchestration-v2.provider-session-scope-close-completed-late", - { - providerSessionId: input.providerSessionId, - reason: input.reason, - }, - ), - ), - Effect.forkDetach, - ); - } + const closeExit = yield* closeSessionScopeBounded(entry.scope, { + providerSessionId: input.providerSessionId, + reason: input.reason, + }); const records = { entry, reason: input.reason, @@ -1750,105 +1756,184 @@ export const layerWithOptions = ( } }); const sessionScope = yield* Scope.make(); - const runtime = yield* adapter - .openSession({ - threadId: input.threadId, + // sessionOpen serializes by providerSessionId, not threadId, so a + // second open for the same thread (a different providerSessionId) + // can run prepareMcpSession concurrently with this one and reuse + // the credential this open mints, holding its own reservation. + // Revoke it only if nothing else still depends on it: the same + // check releaseEntry uses before revoking a session's credential. + // The check and the revoke run under the same per-thread lock + // prepareMcpSession uses, so a concurrent prepare for this + // thread cannot adopt the credential in the gap between this + // open's holder check and its revoke: either prepare runs + // first and the check sees its reservation, or this finishes + // and releases the credential before prepare can reuse it. + // Compares only the credential id, not thread attachment: an + // older entry can be attached to this thread while recording a + // different credential (a stale record, or its own credential + // for a thread it also serves), and that must not veto + // revoking the credential this open minted and nothing holds. + // releaseEntry's own holder check also accepts bare thread + // attachment, but for a different reason: it runs without this + // lock, so it still needs to catch an in-flight attach whose + // credential isn't recorded yet. That gap can't happen here, + // since the lock excludes any concurrent prepare for this + // thread for the entire span of this check and this revoke. + const revokeIfUnheld = mcpPrepareLock.withLock( + input.threadId, + Effect.gen(function* () { + if (!prepared.issued || mcpCredentialId === undefined) return; + const current = yield* Ref.get(sessions); + const heldElsewhere = + isMcpCredentialReserved(input.threadId, mcpCredentialId) || + Array.from(current.values()).some( + (other) => + other.mcpCredentialIdByThread.get(input.threadId) === mcpCredentialId, + ); + if (!heldElsewhere) { + yield* clearMcpSession(input.threadId, mcpCredentialId); + } + }), + ); + // A failure or an interrupt anywhere from here through the end + // of setup, opening the provider session, constructing the + // entry, publishing it, dropping the reservation, persisting + // the attachment, must not leave the provider process this + // open spawned behind. Splitting that span across more than + // one handler-wrapped effect does not close this: the fiber + // is unprotected in the gap between one piece's wrapper + // ending and the next piece's wrapper attaching, and a Stop + // landing exactly there is not caught by either. One handler + // wrapping the whole span as a single effect removes the gap + // instead of racing to close it: abandonSetup adapts to + // whatever progress was made (the adapter call never + // finished, the entry was built but not published, or it was + // published) and cleans up accordingly, so it is correct no + // matter where in the span the interrupt lands. + // + // Identifies "still ours" by the adapter runtime reference, + // not the entry object: manager.get and idle scheduling + // legitimately replace the map's entry object (a spread with + // an updated lastActivityAtMs/idleFiber) on every touch, so + // comparing entry objects by identity would treat routine + // activity as "replaced by someone else" and skip releasing + // a session that is still this open's own. + let openedRuntime: ProviderAdapterV2SessionRuntime | undefined; + const abandonSetup = Effect.gen(function* () { + // Unconditional and first: the success path drops the + // reservation right after publishing the entry, but an + // interrupt can land in the gap between the entry going into + // the map (openedRuntime recorded) and that drop actually + // running. Landing there takes the releaseEntry branch below, + // which revokes the credential only if nothing holds it; the + // reservation this open never got to drop would make that + // check see the credential as still in use and skip revoking + // it, leaking it (and the reservation count) forever. Calling + // it here first, before either branch, closes that gap no + // matter where in this span the interrupt landed. + yield* dropReservation; + if (openedRuntime !== undefined) { + const current = yield* Ref.get(sessions); + if (current.get(key)?.runtime === openedRuntime) { + yield* releaseEntry({ + providerSessionId: input.providerSessionId, + reason: "runtime_error", + detail: + "Opening the provider session was abandoned before it finished setup.", + }).pipe(logReleaseFailure(input.providerSessionId)); + return; + } + } + // Bound the close the same way release does: a misbehaving + // adapter finalizer must not turn the abandonment itself + // into a new hang. + yield* closeSessionScopeBounded(sessionScope, { providerSessionId: input.providerSessionId, - modelSelection: input.modelSelection, - runtimePolicy: input.runtimePolicy, - ...(input.resumeFromSession === undefined - ? {} - : { resumeFromSession: input.resumeFromSession }), - ...(input.initialNativeThreadId === undefined - ? {} - : { initialNativeThreadId: input.initialNativeThreadId }), - ...(input.initialProviderItemIdentityVersion === undefined - ? {} - : { - initialProviderItemIdentityVersion: - input.initialProviderItemIdentityVersion, - }), - }) - .pipe( - Effect.provideService(Scope.Scope, sessionScope), - Effect.tapError(() => - Scope.close(sessionScope, Exit.void).pipe( - Effect.ignore, - Effect.andThen(dropReservation), - // Revoke only a credential this open freshly minted: a - // reused credential is held by another live provider - // process and must survive this open's failure. - Effect.andThen( - prepared.issued - ? clearMcpSession(input.threadId, mcpCredentialId) - : Effect.void, - ), - ), - ), - Effect.onInterrupt(() => dropReservation), - Effect.mapError( - (cause) => - new ProviderSessionOpenError({ - instanceId: input.modelSelection.instanceId, - providerSessionId: input.providerSessionId, - cause, - }), - ), - ); - const eventSubscribers = yield* Ref.make< - ReadonlyMap> - >(new Map()); - const exposedRuntime = decorateRuntime(runtime, eventSubscribers); - const now = yield* Clock.currentTimeMillis; - const entry: LiveSessionEntry = { - attachedThreadIds: new Set([input.threadId]), - loadedProviderThreadKeyByThread: new Map(), - mcpCredentialIdByThread: - mcpCredentialId === undefined - ? new Map() - : new Map([[input.threadId, mcpCredentialId]]), - supportsMultipleProviderThreads: - runtime.providerSession.capabilities.sessions - .supportsMultipleProviderThreadsPerSession, - runtime, - exposedRuntime, - eventSubscribers, - requestEventPermit: yield* Semaphore.make(1), - scope: sessionScope, - idleGeneration: 0, - busyCount: 0, - lastActivityAtMs: now, - idleFiber: null, - pinnedSinceMs: null, - }; - yield* Ref.update(sessions, (current) => { - const updated = new Map(current); - updated.set(key, entry); - return updated; + reason: "open_abandoned", + }).pipe(Effect.andThen(revokeIfUnheld)); }); - // The entry now guards the credential via its recorded id, so - // the pre-open reservation can be dropped. - yield* dropReservation; - yield* withActivityError( - input.providerSessionId, - writeProviderSessionEvents({ - runtime, - threadIds: [input.threadId], - type: "provider-session.attached", - payload: runtime.providerSession, - }), - ).pipe( - Effect.tapError(() => - releaseEntry({ + return yield* Effect.gen(function* () { + const runtime = yield* adapter + .openSession({ + threadId: input.threadId, providerSessionId: input.providerSessionId, - reason: "runtime_error", - detail: "Failed to persist the provider-session attachment.", - }).pipe(logReleaseFailure(input.providerSessionId)), - ), + modelSelection: input.modelSelection, + runtimePolicy: input.runtimePolicy, + ...(input.resumeFromSession === undefined + ? {} + : { resumeFromSession: input.resumeFromSession }), + ...(input.initialNativeThreadId === undefined + ? {} + : { initialNativeThreadId: input.initialNativeThreadId }), + ...(input.initialProviderItemIdentityVersion === undefined + ? {} + : { + initialProviderItemIdentityVersion: + input.initialProviderItemIdentityVersion, + }), + }) + .pipe( + Effect.provideService(Scope.Scope, sessionScope), + Effect.mapError( + (cause) => + new ProviderSessionOpenError({ + instanceId: input.modelSelection.instanceId, + providerSessionId: input.providerSessionId, + cause, + }), + ), + ); + const eventSubscribers = yield* Ref.make< + ReadonlyMap> + >(new Map()); + const exposedRuntime = decorateRuntime(runtime, eventSubscribers); + const now = yield* Clock.currentTimeMillis; + const entry: LiveSessionEntry = { + attachedThreadIds: new Set([input.threadId]), + loadedProviderThreadKeyByThread: new Map(), + mcpCredentialIdByThread: + mcpCredentialId === undefined + ? new Map() + : new Map([[input.threadId, mcpCredentialId]]), + supportsMultipleProviderThreads: + runtime.providerSession.capabilities.sessions + .supportsMultipleProviderThreadsPerSession, + runtime, + exposedRuntime, + eventSubscribers, + requestEventPermit: yield* Semaphore.make(1), + scope: sessionScope, + idleGeneration: 0, + busyCount: 0, + lastActivityAtMs: now, + idleFiber: null, + pinnedSinceMs: null, + }; + yield* Ref.update(sessions, (current) => { + const updated = new Map(current); + updated.set(key, entry); + return updated; + }); + openedRuntime = runtime; + // The entry now guards the credential via its recorded id, + // so the pre-open reservation can be dropped. + yield* dropReservation; + yield* withActivityError( + input.providerSessionId, + writeProviderSessionEvents({ + runtime, + threadIds: [input.threadId], + type: "provider-session.attached", + payload: runtime.providerSession, + }), + ); + yield* startEventPump(entry); + yield* scheduleIdleRelease(input.providerSessionId); + return exposedRuntime; + }).pipe( + Effect.tapError(() => abandonSetup), + Effect.onInterrupt(() => abandonSetup), ); - yield* startEventPump(entry); - yield* scheduleIdleRelease(input.providerSessionId); - return exposedRuntime; }), ), get: (providerSessionId) =>