diff --git a/apps/server/src/mcp/McpSessionRegistry.ts b/apps/server/src/mcp/McpSessionRegistry.ts index ebc8ed74b9385..bbd5394746cc2 100644 --- a/apps/server/src/mcp/McpSessionRegistry.ts +++ b/apps/server/src/mcp/McpSessionRegistry.ts @@ -33,7 +33,8 @@ export interface McpSessionRegistryShape { ) => Effect.Effect; readonly touch: (threadId: ThreadId) => Effect.Effect; readonly revokeProviderSession: (providerSessionId: string) => Effect.Effect; - readonly revokeThread: (threadId: ThreadId) => Effect.Effect; + /** `keep` names provider sessions whose credentials a live process still holds. */ + readonly revokeThread: (threadId: ThreadId, keep?: ReadonlySet) => Effect.Effect; readonly revokeAll: Effect.Effect; } @@ -212,8 +213,12 @@ const makeWithOptions = Effect.fn("McpSessionRegistry.make")(function* ( yield* revokeWhere((record) => record.scope.thread.providerSessionId === providerSessionId); }, ), - revokeThread: Effect.fn("McpSessionRegistry.revokeThread")(function* (threadId) { - yield* revokeWhere((record) => record.scope.thread.threadId === threadId); + revokeThread: Effect.fn("McpSessionRegistry.revokeThread")(function* (threadId, keep) { + yield* revokeWhere( + (record) => + record.scope.thread.threadId === threadId && + keep?.has(record.scope.thread.providerSessionId) !== true, + ); }), revokeAll: SynchronizedRef.set(state, { records: new Map() }), }); diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts index 09441e2659166..6e249e518cdd5 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts @@ -466,8 +466,10 @@ const makeProviderAdapter = Effect.fnUntraced(function* ( readonly archive?: Effect.Effect; readonly spawnBeforeOpen?: boolean; readonly scopeCloseReached?: Deferred.Deferred; + readonly instanceId?: ProviderInstanceId; } = {}, ) { + const instanceId = options.instanceId ?? modelSelection.instanceId; const mcpSessions = yield* McpProviderSessions.McpProviderSessions; const countClose = Effect.addFinalizer(() => Ref.update(state, (current) => ({ @@ -487,7 +489,7 @@ const makeProviderAdapter = Effect.fnUntraced(function* ( ).pipe(Effect.andThen(Effect.never)), ); return ProviderAdapter.ProviderAdapterV2.of({ - instanceId: ProviderInstanceId.make("codex"), + instanceId, driver: CODEX_DRIVER, getCapabilities: () => Effect.succeed(options.capabilities ?? CodexCapabilities), planSelectionTransition: () => Effect.succeed({ type: "apply_on_next_turn" }), @@ -535,7 +537,7 @@ const makeProviderAdapter = Effect.fnUntraced(function* ( } return { - instanceId: ProviderInstanceId.make("codex"), + instanceId, driver: CODEX_DRIVER, providerSessionId: input.providerSessionId, ...(options.normalizeThreadLoadModelSelection === undefined @@ -651,6 +653,7 @@ function layerTest(input: { readonly serverSettingsLayer?: ReturnType; readonly projectServiceLayer?: Layer.Layer; readonly installationGateLayer?: Layer.Layer; + readonly secondInstanceId?: ProviderInstanceId; }) { const layerConfiguredEventSink = input.flakyReleaseWrites !== undefined @@ -698,7 +701,17 @@ function layerTest(input: { ...(input.scopeCloseReached === undefined ? {} : { scopeCloseReached: input.scopeCloseReached }), - }).pipe(Effect.map(ProviderAdapterRegistry.layerSingle)), + }).pipe( + Effect.flatMap((adapter) => + input.secondInstanceId === undefined + ? Effect.succeed([adapter]) + : makeProviderAdapter(input.state, { + ...(input.mcpConfigs === undefined ? {} : { mcpConfigs: input.mcpConfigs }), + instanceId: input.secondInstanceId, + }).pipe(Effect.map((second) => [adapter, second])), + ), + Effect.map(ProviderAdapterRegistry.layerFromAdapters), + ), ); const layerConfiguredMcpRegistry = input.pauseResolve === undefined @@ -2958,6 +2971,101 @@ it.effect("ProviderSessionManagerV2 duplicate detach preserves replacement MCP c }), ); +// regression: #37000 a queued hand-off keeps a busy session attached while another instance opens +it.effect( + "ProviderSessionManagerV2 keeps a retained session's MCP credential when a hand-off opens another instance", + () => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const settings = yield* ServerSettings.ServerSettingsService.pipe( + Effect.provide(ServerSettings.layerTest()), + ); + const otherInstanceId = ProviderInstanceId.make("codex-other"); + 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 mcpSessions = yield* McpProviderSessions.McpProviderSessions; + const now = yield* DateTime.now; + const threadId = ThreadId.make("thread-provider-session-manager-retained-handoff"); + const retainedSessionId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + const handoffSessionId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: otherInstanceId, + threadId, + }); + yield* eventSink.write({ + events: [yield* makeThreadCreatedEvent({ idAllocator, threadId, now })], + }); + + yield* manager.open({ + threadId, + providerSessionId: retainedSessionId, + modelSelection, + runtimePolicy, + }); + const retainedToken = (yield* mcpSessions.read(threadId))?.authorizationHeader.replace( + /^Bearer\s+/, + "", + ); + assert.isDefined(retainedToken); + + // The busy session stays attached; the hand-off opens its own session. + yield* manager.open({ + threadId, + providerSessionId: handoffSessionId, + modelSelection: { instanceId: otherInstanceId, model: "gpt-5.4" }, + runtimePolicy, + }); + assert.equal( + (yield* registry.resolve(retainedToken!))?.thread.providerInstanceId, + modelSelection.instanceId, + "the retained process still holds its credential", + ); + + // Switching back reuses the retained, still-attached session. + yield* manager.open({ + threadId, + providerSessionId: retainedSessionId, + modelSelection, + runtimePolicy, + }); + const current = yield* mcpSessions.read(threadId); + const resolved = yield* registry.resolve( + current!.authorizationHeader.replace(/^Bearer\s+/, ""), + ); + assert.equal(resolved?.thread.providerInstanceId, modelSelection.instanceId); + + // Narrowed agent access is not kept alive by the retained session. + yield* settings.updateSettings({ enableAgentBrowserAccess: false }); + yield* manager.open({ + threadId, + providerSessionId: yield* idAllocator.allocate.providerSession({ + providerInstanceId: otherInstanceId, + threadId, + }), + modelSelection: { instanceId: otherInstanceId, model: "gpt-5.4" }, + runtimePolicy, + }); + assert.isUndefined(yield* registry.resolve(retainedToken!)); + }); + + yield* effect.pipe( + Effect.provide( + layerTest({ + state, + idleTimeoutMs: 1_000, + secondInstanceId: otherInstanceId, + serverSettingsLayer: Layer.succeed(ServerSettings.ServerSettingsService, settings), + }), + ), + ); + }), +); + it.effect( "ProviderSessionManagerV2 keeps a thread's MCP credential stable across detach and re-attach on a shared session", () => diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.ts index ca54c14b2cb95..e32d8e1162898 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.ts @@ -547,6 +547,8 @@ export const layerWithOptions = ( const isMcpCredentialReserved = (threadId: ThreadId, mcpCredentialId: string) => (mcpCredentialReservations.get(mcpReservationKey(threadId, mcpCredentialId)) ?? 0) > 0; const mcpPrepareLock = yield* KeyedLock.make(); + /** The thread's MCP slot holds one config; sessions still attached get theirs back on reuse. */ + const mcpConfigByCredentialId = new Map[0]>(); /** Serialize prepares per thread so concurrent rotations cannot revoke each other's credentials. */ const prepareMcpSession = ( threadId: ThreadId, @@ -599,7 +601,30 @@ export const layerWithOptions = ( } dropMcpCredentialReservation(threadId, existing.providerSessionId); } - yield* mcpSessionRegistry.revokeThread(threadId); + // A hand-off can keep a busy session attached; its process still uses its token + // unless agent access was narrowed since it was issued. + const held = new Set(); + for (const entry of (yield* Ref.get(sessions)).values()) { + const heldId = entry.mcpCredentialIdByThread.get(threadId); + const heldCapabilities = + heldId === undefined + ? undefined + : mcpConfigByCredentialId.get(heldId)?.capabilities; + if ( + heldId !== undefined && + entry.attachedThreadIds.has(threadId) && + heldCapabilities?.has("preview") === browserToolsAvailable && + heldCapabilities.has("device") === deviceToolsAvailable + ) { + held.add(heldId); + } + } + yield* mcpSessionRegistry.revokeThread(threadId, held); + for (const [id, config] of mcpConfigByCredentialId) { + if (config.threadId === threadId && !held.has(id)) { + mcpConfigByCredentialId.delete(id); + } + } const credential = yield* mcpSessionRegistry.issue({ threadId, providerInstanceId, @@ -607,7 +632,7 @@ export const layerWithOptions = ( capabilities, subagent, }); - yield* mcpSessions.set({ + const config = { ...credential.config, hasPendingWake: projectionStore .getThreadRecords(threadId, ["runs", "subagents"]) @@ -615,19 +640,29 @@ export const layerWithOptions = ( Effect.map(delegatedTaskWillWake), Effect.orElseSucceed(() => false), ), - }); + }; + mcpConfigByCredentialId.set(credential.config.providerSessionId, config); + yield* mcpSessions.set(config); reserveMcpCredential(threadId, credential.config.providerSessionId); return { mcpCredentialId: credential.config.providerSessionId, issued: true }; }), ); const clearMcpSession = (threadId: ThreadId, mcpCredentialId?: string) => mcpCredentialId === undefined - ? mcpSessionRegistry - .revokeThread(threadId) - .pipe(Effect.tap(() => mcpSessions.clear(threadId))) + ? mcpSessionRegistry.revokeThread(threadId).pipe( + Effect.tap(() => + Effect.sync(() => { + for (const [id, config] of mcpConfigByCredentialId) { + if (config.threadId === threadId) mcpConfigByCredentialId.delete(id); + } + }), + ), + Effect.tap(() => mcpSessions.clear(threadId)), + ) : mcpSessionRegistry.revokeProviderSession(mcpCredentialId).pipe( Effect.tap(() => Effect.gen(function* () { + mcpConfigByCredentialId.delete(mcpCredentialId); if ((yield* mcpSessions.read(threadId))?.providerSessionId === mcpCredentialId) { yield* mcpSessions.clear(threadId); } @@ -635,6 +670,27 @@ export const layerWithOptions = ( ), ); + const restoreHeldMcpSession = (providerSessionId: ProviderSessionId, threadId: ThreadId) => + mcpPrepareLock.withLock( + threadId, + Effect.gen(function* () { + const mcpCredentialId = (yield* Ref.get(sessions)) + .get(sessionKey(providerSessionId)) + ?.mcpCredentialIdByThread.get(threadId); + const config = + mcpCredentialId === undefined + ? undefined + : mcpConfigByCredentialId.get(mcpCredentialId); + if ( + config === undefined || + (yield* mcpSessions.read(threadId))?.providerSessionId === mcpCredentialId + ) { + return; + } + yield* mcpSessions.set(config); + }), + ); + const subscriberReleaseCleanup = new WeakMap< Queue.Queue, Effect.Effect @@ -1767,6 +1823,8 @@ export const layerWithOptions = ( }), ); } + } else { + yield* restoreHeldMcpSession(input.providerSessionId, input.threadId); } }).pipe( Effect.onError(() =>