diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts index 3c635be6753df..391e9d911a0b5 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts @@ -2107,6 +2107,62 @@ it.effect( }), ); +it.effect( + "ProviderSessionManagerV2 takes the thread's MCP credential back when a run returns to a still-attached session", + () => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + 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-returning-provider"); + const providerSessionId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + const tokenOf = (config: McpProviderSession.McpProviderSessionConfig | undefined) => + config?.authorizationHeader.replace(/^Bearer\s+/, "") ?? ""; + + yield* eventSink.write({ + events: [yield* makeThreadCreatedEvent({ idAllocator, threadId, now })], + }); + yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); + + // The thread switches to another provider, whose attach replaces the + // thread's credential while this session stays attached. + yield* registry.revokeThread(threadId); + const otherProvider = yield* registry.issue({ + threadId, + providerInstanceId: ProviderInstanceId.make("claudeAgent"), + }); + yield* mcpSessions.set(otherProvider.config); + + // The thread switches back and its next turn reuses this session. + yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); + assert.equal((yield* Ref.get(state)).openCount, 1); + + const reclaimed = yield* mcpSessions.read(threadId); + assert.equal(reclaimed?.providerInstanceId, modelSelection.instanceId); + assert.equal( + (yield* registry.resolve(tokenOf(reclaimed)))?.thread.providerInstanceId, + modelSelection.instanceId, + ); + assert.isUndefined(yield* registry.resolve(tokenOf(otherProvider.config))); + + // The session's record follows the new credential, so release revokes it. + yield* manager.close(providerSessionId); + assert.isUndefined(yield* registry.resolve(tokenOf(reclaimed))); + assert.isUndefined(yield* mcpSessions.read(threadId)); + }); + + yield* effect.pipe(Effect.provide(layerTest({ state, idleTimeoutMs: 1_000 }))); + }), +); + it.effect( "ProviderSessionManagerV2 protects a reused credential from a predecessor release during open", () => diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.ts index 00b42d510bd91..9f86fca43fbdc 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.ts @@ -1328,22 +1328,34 @@ export const layerWithOptions = ( Effect.tap((runtime) => Effect.sync(() => (attachedTo = runtime))), Effect.uninterruptible, ); + // A thread stays attached to the session it ran on before switching + // providers, and the other provider's attach took over the thread's + // credential slot. A run returning to this session must take it + // back, or the process calls T3 tools scoped to the other provider. + const reclaimsCredential = + attached === undefined && + (yield* mcpSessions.read(input.threadId))?.providerInstanceId !== + input.providerInstanceId && + (yield* Ref.get(sessions)) + .get(sessionKey(input.providerSessionId)) + ?.attachedThreadIds.has(input.threadId) === true; + if (attached === undefined && !reclaimsCredential) return; + const prepared = yield* prepareMcpSession(input.threadId, input.providerInstanceId); + preparedForCleanup = prepared; + if (prepared.mcpCredentialId !== undefined) { + const mcpCredentialId = prepared.mcpCredentialId; + yield* Ref.update(sessions, (current) => { + const key = sessionKey(input.providerSessionId); + const entry = current.get(key); + if (entry === undefined) return current; + const mcpCredentialIdByThread = new Map(entry.mcpCredentialIdByThread); + mcpCredentialIdByThread.set(input.threadId, mcpCredentialId); + const updated = new Map(current); + updated.set(key, { ...entry, mcpCredentialIdByThread }); + return updated; + }); + } if (attached !== undefined) { - const prepared = yield* prepareMcpSession(input.threadId, input.providerInstanceId); - preparedForCleanup = prepared; - if (prepared.mcpCredentialId !== undefined) { - const mcpCredentialId = prepared.mcpCredentialId; - yield* Ref.update(sessions, (current) => { - const key = sessionKey(input.providerSessionId); - const entry = current.get(key); - if (entry === undefined) return current; - const mcpCredentialIdByThread = new Map(entry.mcpCredentialIdByThread); - mcpCredentialIdByThread.set(input.threadId, mcpCredentialId); - const updated = new Map(current); - updated.set(key, { ...entry, mcpCredentialIdByThread }); - return updated; - }); - } const entry = (yield* Ref.get(sessions)).get(sessionKey(input.providerSessionId)); if (entry !== undefined) { yield* withActivityError( @@ -1358,37 +1370,49 @@ export const layerWithOptions = ( } } }).pipe( - Effect.onError(() => - attachedTo === undefined - ? Effect.void - : removeThreadAttachment({ ...input, runtime: attachedTo }).pipe( - Effect.andThen( - Effect.suspend(() => { - dropReservation(); - // Other live processes hold reused MCP credentials; revocation would break their clients. - if (preparedForCleanup?.issued !== true) return Effect.void; - const mcpCredentialId = preparedForCleanup.mcpCredentialId; - const attachedRuntime = attachedTo; - return Ref.get(sessions).pipe( - Effect.flatMap((current) => - (mcpCredentialId !== undefined && - isMcpCredentialReserved(input.threadId, mcpCredentialId)) || - Array.from(current.values()).some( - (other) => - other.runtime !== attachedRuntime && - (other.attachedThreadIds.has(input.threadId) || - (mcpCredentialId !== undefined && - other.mcpCredentialIdByThread.get(input.threadId) === - mcpCredentialId)), - ) - ? Effect.void - : clearMcpSession(input.threadId, mcpCredentialId), - ), - ); + Effect.onError(() => { + // A reclaim (thread already attached, credential rotated) has + // attachedTo === undefined but still issued a credential that must + // be cleaned up on failure. Only skip when nothing was prepared. + const isReclaim = attachedTo === undefined && preparedForCleanup !== undefined; + if (attachedTo === undefined && !isReclaim) return Effect.void; + const removeAttachment = + attachedTo === undefined + ? Effect.void + : removeThreadAttachment({ ...input, runtime: attachedTo }); + return removeAttachment.pipe( + Effect.andThen( + Effect.suspend(() => { + dropReservation(); + // Other live processes hold reused MCP credentials; revocation would break their clients. + if (preparedForCleanup?.issued !== true) return Effect.void; + const mcpCredentialId = preparedForCleanup.mcpCredentialId; + // For a reclaim, exclude the session being reused (not torn + // down) from the holder check, mirroring new attachments. + return Ref.get(sessions).pipe( + Effect.flatMap((current) => { + const attachedRuntime = + attachedTo ?? current.get(sessionKey(input.providerSessionId))?.runtime; + const held = + (mcpCredentialId !== undefined && + isMcpCredentialReserved(input.threadId, mcpCredentialId)) || + Array.from(current.values()).some( + (other) => + other.runtime !== attachedRuntime && + (other.attachedThreadIds.has(input.threadId) || + (mcpCredentialId !== undefined && + other.mcpCredentialIdByThread.get(input.threadId) === + mcpCredentialId)), + ); + return held + ? Effect.void + : clearMcpSession(input.threadId, mcpCredentialId); }), - ), - ), - ), + ); + }), + ), + ); + }), ); return threadAttachment .withLock(threadAttachmentKey(input), attach)