diff --git a/apps/server/src/mcp/McpSessionRegistry.test.ts b/apps/server/src/mcp/McpSessionRegistry.test.ts index 9fc482d366d3..a33819400c02 100644 --- a/apps/server/src/mcp/McpSessionRegistry.test.ts +++ b/apps/server/src/mcp/McpSessionRegistry.test.ts @@ -2,6 +2,7 @@ import * as NodeServices from "@effect/platform-node/NodeServices"; import { expect, it } from "@effect/vitest"; import { EnvironmentId, ProviderInstanceId, ThreadId } from "@t3tools/contracts"; import * as Effect from "effect/Effect"; +import * as Layer from "effect/Layer"; import { HttpServer } from "effect/unstable/http"; import * as NetAddress from "effect/unstable/net/NetAddress"; @@ -163,3 +164,47 @@ it.effect("does not keep credentials of other threads alive", () => expect(yield* registry.resolve(token)).toBeUndefined(); }), ); + +/** + * A replacement credential can be issued without revoking the token a live + * provider query is still presenting. + */ +function retainsExistingThreadCredentialUntilThatSessionIsRevoked() { + const layer = McpSessionRegistry.layer.pipe( + Layer.provide(Layer.succeed(HttpServer.HttpServer, fakeHttpServer)), + Layer.provide(Layer.succeed(ServerEnvironment.ServerEnvironment, fakeEnvironment)), + Layer.provide(NodeServices.layer), + ); + return Effect.gen(function* () { + const registry = yield* McpSessionRegistry.McpSessionRegistry; + const threadId = ThreadId.make("thread-retain"); + const first = yield* McpSessionRegistry.issueActiveMcpCredential({ + threadId, + providerInstanceId: ProviderInstanceId.make("claudeAgent"), + capabilities: new Set(["preview"]), + }); + const second = yield* McpSessionRegistry.issueActiveMcpCredential({ + threadId, + providerInstanceId: ProviderInstanceId.make("claudeAgent"), + capabilities: new Set(["preview"]), + retainExisting: true, + }); + if (first === undefined || second === undefined) { + return yield* Effect.die("expected MCP credentials"); + } + const firstToken = first.config.authorizationHeader.replace(/^Bearer\s+/, ""); + const secondToken = second.config.authorizationHeader.replace(/^Bearer\s+/, ""); + + expect((yield* registry.resolve(firstToken))?.threadId).toBe(threadId); + expect((yield* registry.resolve(secondToken))?.threadId).toBe(threadId); + + yield* McpSessionRegistry.revokeActiveMcpProviderSession(second.config.providerSessionId); + expect((yield* registry.resolve(firstToken))?.threadId).toBe(threadId); + expect(yield* registry.resolve(secondToken)).toBeUndefined(); + }).pipe(Effect.provide(layer)); +} + +it.effect( + "retains an existing thread credential when the replacement asks to keep it", + retainsExistingThreadCredentialUntilThatSessionIsRevoked, +); diff --git a/apps/server/src/mcp/McpSessionRegistry.ts b/apps/server/src/mcp/McpSessionRegistry.ts index fa0a17d826f9..628402bc03b8 100644 --- a/apps/server/src/mcp/McpSessionRegistry.ts +++ b/apps/server/src/mcp/McpSessionRegistry.ts @@ -16,6 +16,11 @@ export interface McpCredentialRequest { readonly threadId: ThreadId; readonly providerInstanceId: ProviderInstanceId; readonly capabilities: ReadonlySet; + /** + * Keep credentials already issued for this thread. A refused session + * replacement still needs the token the running query baked in. + */ + readonly retainExisting?: boolean; } export interface McpIssuedCredential { @@ -223,14 +228,31 @@ const make = Effect.acquireRelease( export const layer = Layer.effect(McpSessionRegistry, make); +/** + * Issue the MCP credential a provider session presents to `/mcp`. + * Replaces any credential already issued for the thread unless + * `retainExisting` is set, which keeps the live query's token valid until + * the caller revokes it. + */ export const issueActiveMcpCredential = ( request: McpCredentialRequest, ): Effect.Effect => activeMcpSessionRegistry - ? activeMcpSessionRegistry - .revokeThread(request.threadId) - .pipe(Effect.andThen(activeMcpSessionRegistry.issue(request))) - : Effect.sync((): McpIssuedCredential | undefined => undefined); + ? (request.retainExisting === true + ? Effect.void + : activeMcpSessionRegistry.revokeThread(request.threadId) + ).pipe(Effect.andThen(activeMcpSessionRegistry.issue(request))) + : Effect.succeed(undefined); + +/** + * Revoke one issued MCP credential without touching the thread's other tokens. + * Drops a rejected replacement credential, or the credential a stopped query + * was using, while a live query keeps its own. + */ +export const revokeActiveMcpProviderSession = (providerSessionId: string): Effect.Effect => + activeMcpSessionRegistry + ? activeMcpSessionRegistry.revokeProviderSession(providerSessionId) + : Effect.void; /** * Refreshes the liveness of a thread's MCP credential. Called on every provider diff --git a/apps/server/src/provider/Layers/ClaudeAdapter.test.ts b/apps/server/src/provider/Layers/ClaudeAdapter.test.ts index ca596e6501cf..283af610fa57 100644 --- a/apps/server/src/provider/Layers/ClaudeAdapter.test.ts +++ b/apps/server/src/provider/Layers/ClaudeAdapter.test.ts @@ -313,6 +313,178 @@ async function readPromptMessages( } const THREAD_ID = ThreadId.make("thread-claude-1"); + +/** + * Claude adapter layer that records every query `startSession` creates. + */ +function makeQueryRecordingClaudeLayer() { + const queries: Array = []; + const layer = Layer.effect( + ClaudeAdapter, + Effect.gen(function* () { + const claudeConfig = decodeClaudeSettings({}); + return yield* makeClaudeAdapter(claudeConfig, { + createQuery: () => { + const query = new FakeClaudeQuery(); + queries.push(query); + return query; + }, + }); + }), + ).pipe( + Layer.provideMerge(ServerConfig.layerTest("/tmp/claude-adapter-test", "/tmp")), + Layer.provideMerge(ServerSettingsService.layerTest()), + Layer.provideMerge(NodeServices.layer), + ); + return { queries, layer }; +} + +/** + * Refuses Claude session replacement while a child task is still running. + */ +function refusesClaudeSessionReplacementWhileLiveTasks() { + const harness = makeQueryRecordingClaudeLayer(); + return Effect.gen(function* () { + const adapter = yield* ClaudeAdapter; + const taskStartedFiber = yield* adapter.streamEvents.pipe( + Stream.filter((event) => event.type === "task.started"), + Stream.take(1), + Stream.runCollect, + Effect.forkChild, + ); + const firstSession = yield* adapter.startSession({ + threadId: THREAD_ID, + provider: ProviderDriverKind.make("claudeAgent"), + runtimeMode: "full-access", + }); + yield* adapter.sendTurn({ + threadId: firstSession.threadId, + input: "spawn agents", + attachments: [], + }); + harness.queries[0]?.emit({ + type: "system", + subtype: "task_started", + task_id: "task-live", + description: "Agent A", + task_type: "local_agent", + uuid: "task-live-uuid", + session_id: "sdk-session", + } as unknown as SDKMessage); + yield* Fiber.join(taskStartedFiber); + + const warningFiber = yield* adapter.streamEvents.pipe( + Stream.filter((event) => event.type === "runtime.warning"), + Stream.take(1), + Stream.runCollect, + Effect.forkChild, + ); + const replaceError = yield* adapter + .startSession({ + threadId: THREAD_ID, + provider: ProviderDriverKind.make("claudeAgent"), + runtimeMode: "approval-required", + resumeCursor: firstSession.resumeCursor, + }) + .pipe(Effect.flip); + const warnings = Array.from(yield* Fiber.join(warningFiber)); + const activeSessions = yield* adapter.listSessions(); + + assert.equal(replaceError._tag, "ProviderAdapterRequestError"); + if (replaceError._tag !== "ProviderAdapterRequestError") { + return; + } + assert.equal(replaceError.method, "startSession"); + assert.match(replaceError.detail, /live task/); + assert.equal(harness.queries.length, 1); + assert.equal(harness.queries[0]?.closeCalls, 0); + assert.equal(yield* adapter.hasSession(THREAD_ID), true); + assert.equal(activeSessions.length, 1); + assert.equal(activeSessions[0]?.threadId, firstSession.threadId); + assert.equal(activeSessions[0]?.runtimeMode, "full-access"); + assert.equal(warnings.length, 1); + assert.equal(warnings[0]?.type, "runtime.warning"); + if (warnings[0]?.type === "runtime.warning") { + assert.match(warnings[0].payload.message, /live task/); + } + }).pipe( + Effect.provideService(Random.Random, makeDeterministicRandomService()), + Effect.provide(harness.layer), + ); +} + +/** + * Replaces a Claude session once its child tasks have finished. + */ +function replacesClaudeSessionAfterLiveTasksFinish() { + const harness = makeQueryRecordingClaudeLayer(); + return Effect.gen(function* () { + const adapter = yield* ClaudeAdapter; + const taskStartedFiber = yield* adapter.streamEvents.pipe( + Stream.filter((event) => event.type === "task.started"), + Stream.take(1), + Stream.runCollect, + Effect.forkChild, + ); + const firstSession = yield* adapter.startSession({ + threadId: THREAD_ID, + provider: ProviderDriverKind.make("claudeAgent"), + runtimeMode: "full-access", + }); + yield* adapter.sendTurn({ + threadId: firstSession.threadId, + input: "spawn agents", + attachments: [], + }); + harness.queries[0]?.emit({ + type: "system", + subtype: "task_started", + task_id: "task-done", + description: "Agent A", + task_type: "local_agent", + uuid: "task-done-uuid", + session_id: "sdk-session", + } as unknown as SDKMessage); + yield* Fiber.join(taskStartedFiber); + + const taskCompletedFiber = yield* adapter.streamEvents.pipe( + Stream.filter((event) => event.type === "task.completed"), + Stream.take(1), + Stream.runCollect, + Effect.forkChild, + ); + harness.queries[0]?.emit({ + type: "system", + subtype: "task_notification", + task_id: "task-done", + status: "completed", + summary: "done", + uuid: "task-done-notify", + session_id: "sdk-session", + } as unknown as SDKMessage); + yield* Fiber.join(taskCompletedFiber); + + const secondSession = yield* adapter.startSession({ + threadId: THREAD_ID, + provider: ProviderDriverKind.make("claudeAgent"), + runtimeMode: "approval-required", + resumeCursor: firstSession.resumeCursor, + }); + const activeSessions = yield* adapter.listSessions(); + + assert.equal(harness.queries.length, 2); + assert.equal(harness.queries[0]?.closeCalls, 1); + assert.equal(harness.queries[1]?.closeCalls, 0); + assert.equal(yield* adapter.hasSession(THREAD_ID), true); + assert.equal(activeSessions.length, 1); + assert.equal(activeSessions[0]?.runtimeMode, "approval-required"); + assert.deepEqual(activeSessions[0]?.resumeCursor, secondSession.resumeCursor); + }).pipe( + Effect.provideService(Random.Random, makeDeterministicRandomService()), + Effect.provide(harness.layer), + ); +} + const RESUME_THREAD_ID = ThreadId.make("thread-claude-resume"); const SYNTHETIC_SUBAGENT_MODEL = "claude-synthetic-subagent[expanded]"; const CLAUDE_ORIGINAL_SESSION_ID = "550e8400-e29b-41d4-a716-446655440010"; @@ -4186,6 +4358,16 @@ describe("ClaudeAdapterLive", () => { ); }); + it.effect( + "does not replace a session while live tasks are still running", + refusesClaudeSessionReplacementWhileLiveTasks, + ); + + it.effect( + "replaces a session after its live tasks have finished", + replacesClaudeSessionAfterLiveTasksFinish, + ); + it.effect("stopSession does not throw into the SDK prompt consumer", () => { // The SDK consumes user messages via `for await (... of prompt)`. // Stopping a session must end that loop cleanly — not throw an error. diff --git a/apps/server/src/provider/Layers/ClaudeAdapter.ts b/apps/server/src/provider/Layers/ClaudeAdapter.ts index 13645a33a05b..ae2e245fb517 100644 --- a/apps/server/src/provider/Layers/ClaudeAdapter.ts +++ b/apps/server/src/provider/Layers/ClaudeAdapter.ts @@ -2064,6 +2064,17 @@ function sdkNativeItemId(message: SDKMessage): string | undefined { return undefined; } +/** + * True when replacing the Claude query would stop child tasks that are still + * running. That stop records them completed and suppresses `session.exited`. + */ +function shouldBlockClaudeSessionReplacementForLiveTasks(existingContext: { + readonly stopped: boolean; + readonly liveTaskIds: ReadonlySet; +}): boolean { + return !existingContext.stopped && existingContext.liveTaskIds.size > 0; +} + export const makeClaudeAdapter = Effect.fn("makeClaudeAdapter")(function* ( claudeSettings: ClaudeSettings, options?: ClaudeAdapterLiveOptions, @@ -4391,7 +4402,17 @@ export const makeClaudeAdapter = Effect.fn("makeClaudeAdapter")(function* ( return Effect.succeed(context); }; + /** + * Start or replace the Claude query for a thread. + * Refuses replacement while child tasks are still running so the open + * query and those tasks stay in place. + */ const startSession: ClaudeAdapterShape["startSession"] = Effect.fn("startSession")( + /** + * Start or replace the Claude query for a thread. + * Refuses replacement while child tasks are still running so the open + * query and those tasks stay in place. + */ function* (input) { const modelCatalog = yield* modelCatalogEffect; if (input.provider !== undefined && input.provider !== PROVIDER) { @@ -4404,6 +4425,25 @@ export const makeClaudeAdapter = Effect.fn("makeClaudeAdapter")(function* ( const existingContext = sessions.get(input.threadId); if (existingContext) { + // Fail instead of handing back the current session: callers bind the + // requested runtime mode and send the turn on whatever comes back. + // Stopping here would mark live tasks completed without session.exited. + if (shouldBlockClaudeSessionReplacementForLiveTasks(existingContext)) { + const liveTaskCount = existingContext.liveTaskIds.size; + yield* emitRuntimeWarning( + existingContext, + `Claude session replacement blocked: ${liveTaskCount} live task(s) are still running.`, + { + existingSessionStatus: existingContext.session.status, + liveTaskCount, + }, + ); + return yield* new ProviderAdapterRequestError({ + provider: PROVIDER, + method: "startSession", + detail: `Claude session replacement is blocked: ${liveTaskCount} live task(s) are still running. Retry after background work finishes.`, + }); + } yield* Effect.logWarning("claude.session.replacing", { threadId: input.threadId, existingSessionStatus: existingContext.session.status, diff --git a/apps/server/src/provider/Layers/ProviderService.test.ts b/apps/server/src/provider/Layers/ProviderService.test.ts index b9997e1df312..b9241cb5ae06 100644 --- a/apps/server/src/provider/Layers/ProviderService.test.ts +++ b/apps/server/src/provider/Layers/ProviderService.test.ts @@ -50,6 +50,8 @@ import * as Schema from "effect/Schema"; import * as Scope from "effect/Scope"; import * as Stream from "effect/Stream"; import * as TestClock from "effect/testing/TestClock"; +import { HttpServer } from "effect/unstable/http"; +import * as NetAddress from "effect/unstable/net/NetAddress"; import * as SqlClient from "effect/unstable/sql/SqlClient"; import { @@ -74,6 +76,9 @@ import { SqlitePersistenceMemory, } from "../../persistence/Layers/Sqlite.ts"; import * as ServerConfig from "../../config.ts"; +import * as ServerEnvironment from "../../environment/ServerEnvironment.ts"; +import * as McpProviderSession from "../../mcp/McpProviderSession.ts"; +import * as McpSessionRegistry from "../../mcp/McpSessionRegistry.ts"; import * as ServerSettings from "../../serverSettings.ts"; import * as AnalyticsService from "../../telemetry/AnalyticsService.ts"; import { makeAdapterRegistryMock } from "../testUtils/providerAdapterRegistryMock.ts"; @@ -5150,3 +5155,239 @@ describe("agent browser access", () => { }).pipe(Effect.provide(NodeServices.layer)), ); }); + +/** + * Provider service wired to the real MCP registry and one Claude adapter. + * `startSession` records the authorization header the adapter would bake in. + */ +function makeClaudeMcpProviderHarness() { + const claude = makeFakeCodexAdapter(CLAUDE_AGENT_DRIVER); + const seenAuthorizationHeaders: Array = []; + const control = { rejectNextStart: false, failNextStart: false }; + const originalStart = claude.startSession.getMockImplementation() as + | (( + input: ProviderSessionStartInput, + ) => Effect.Effect) + | undefined; + /** + * Records the MCP header a start would bake in, and can fail while leaving + * the existing session in place. + */ + const implementStart = (input: ProviderSessionStartInput) => + Effect.gen(function* () { + seenAuthorizationHeaders.push( + McpProviderSession.readMcpProviderSession(input.threadId)?.authorizationHeader, + ); + if (control.failNextStart || control.rejectNextStart) { + control.failNextStart = false; + control.rejectNextStart = false; + return yield* new ProviderAdapterRequestError({ + provider: CLAUDE_AGENT_DRIVER, + method: "startSession", + detail: "Claude session replacement is blocked: 1 live task(s) are still running.", + }); + } + if (originalStart === undefined) { + return yield* Effect.die("Claude startSession mock is missing its original implementation"); + } + return yield* originalStart(input); + }); + ( + claude.startSession as unknown as { + mockImplementation(fn: typeof implementStart): void; + } + ).mockImplementation(implementStart); + + const httpServer = HttpServer.HttpServer.of({ + address: NetAddress.inetAddressFromIpStringUnsafe("127.0.0.1", 43123), + serve: (() => Effect.void) as HttpServer.HttpServer["Service"]["serve"], + }); + const mcpRegistryLayer = McpSessionRegistry.layer.pipe( + Layer.provide(Layer.succeed(HttpServer.HttpServer, httpServer)), + Layer.provide( + Layer.succeed( + ServerEnvironment.ServerEnvironment, + ServerEnvironment.ServerEnvironment.of({ + getEnvironmentId: Effect.succeed(EnvironmentId.make("environment-mcp-retain")), + getDescriptor: Effect.die("unused"), + }), + ), + ), + Layer.provide(NodeServices.layer), + ); + const providerLayer = makeProviderServiceLive().pipe( + Layer.provide(NodeServices.layer), + Layer.provide( + Layer.succeed( + ProviderAdapterRegistry.ProviderAdapterRegistry, + makeAdapterRegistryMock({ [CLAUDE_AGENT_DRIVER]: claude.adapter }), + ), + ), + Layer.provide( + ProviderSessionDirectoryLive.pipe( + Layer.provide(ProviderSessionRuntime.layer.pipe(Layer.provide(SqlitePersistenceMemory))), + ), + ), + Layer.provide(defaultServerSettingsLayer), + Layer.provide(serverConfigTestLayer), + Layer.provide(AnalyticsService.layerTest), + Layer.provide( + Layer.succeed( + ProviderEventLoggers.ProviderEventLoggers, + ProviderEventLoggers.NoOpProviderEventLoggers, + ), + ), + ); + + return { + claude, + seenAuthorizationHeaders, + rejectNextStart: () => { + control.rejectNextStart = true; + }, + failNextStart: () => { + control.failNextStart = true; + }, + layer: Layer.mergeAll(providerLayer, mcpRegistryLayer), + }; +} + +/** Bearer token carried by an MCP authorization header. */ +function mcpBearerToken(authorizationHeader: string | undefined): string { + return authorizationHeader?.replace(/^Bearer\s+/, "") ?? ""; +} + +/** + * A refused restart leaves the credential the live query already presented. + */ +function keepsLiveMcpCredentialWhenSessionReplacementIsRefused() { + const harness = makeClaudeMcpProviderHarness(); + const threadId = asThreadId("thread-mcp-refuse"); + const cwd = fixtureCwd("mcp-refuse"); + return Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const registry = yield* McpSessionRegistry.McpSessionRegistry; + yield* provider.startSession(threadId, { + provider: CLAUDE_AGENT_DRIVER, + providerInstanceId: claudeAgentInstanceId, + threadId, + cwd, + runtimeMode: "full-access", + }); + const liveHeader = harness.seenAuthorizationHeaders[0]; + const liveToken = mcpBearerToken(liveHeader); + assert.equal(liveToken.length > 20, true); + assert.equal((yield* registry.resolve(liveToken))?.threadId, threadId); + + harness.rejectNextStart(); + const failure = yield* provider + .startSession(threadId, { + provider: CLAUDE_AGENT_DRIVER, + providerInstanceId: claudeAgentInstanceId, + threadId, + cwd, + runtimeMode: "approval-required", + }) + .pipe(Effect.flip); + const replacementToken = mcpBearerToken(harness.seenAuthorizationHeaders[1]); + + assert.equal(failure._tag, "ProviderAdapterRequestError"); + assert.notEqual(harness.seenAuthorizationHeaders[1], liveHeader); + assert.equal( + McpProviderSession.readMcpProviderSession(threadId)?.authorizationHeader, + liveHeader, + ); + assert.equal((yield* registry.resolve(liveToken))?.threadId, threadId); + assert.equal(yield* registry.resolve(replacementToken), undefined); + const sessions = yield* harness.claude.adapter.listSessions(); + assert.equal(sessions.length, 1); + assert.equal(sessions[0]?.runtimeMode, "full-access"); + }).pipe(Effect.provide(harness.layer)); +} + +/** + * An accepted replacement rotates the MCP credential off the stopped query. + */ +function rotatesMcpCredentialWhenSessionReplacementIsAccepted() { + const harness = makeClaudeMcpProviderHarness(); + const threadId = asThreadId("thread-mcp-replace"); + const cwd = fixtureCwd("mcp-replace"); + return Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const registry = yield* McpSessionRegistry.McpSessionRegistry; + yield* provider.startSession(threadId, { + provider: CLAUDE_AGENT_DRIVER, + providerInstanceId: claudeAgentInstanceId, + threadId, + cwd, + runtimeMode: "full-access", + }); + const liveHeader = harness.seenAuthorizationHeaders[0]; + const liveToken = mcpBearerToken(liveHeader); + + const replaced = yield* provider.startSession(threadId, { + provider: CLAUDE_AGENT_DRIVER, + providerInstanceId: claudeAgentInstanceId, + threadId, + cwd, + runtimeMode: "approval-required", + }); + const nextHeader = harness.seenAuthorizationHeaders[1]; + const nextToken = mcpBearerToken(nextHeader); + + assert.equal(replaced.runtimeMode, "approval-required"); + assert.notEqual(nextHeader, liveHeader); + assert.equal( + McpProviderSession.readMcpProviderSession(threadId)?.authorizationHeader, + nextHeader, + ); + assert.equal(yield* registry.resolve(liveToken), undefined); + assert.equal((yield* registry.resolve(nextToken))?.threadId, threadId); + }).pipe(Effect.provide(harness.layer)); +} + +/** + * A failed first start still revokes the credential it just issued. + */ +function clearsMcpCredentialWhenInitialStartFails() { + const harness = makeClaudeMcpProviderHarness(); + const threadId = asThreadId("thread-mcp-initial-fail"); + const cwd = fixtureCwd("mcp-initial-fail"); + return Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const registry = yield* McpSessionRegistry.McpSessionRegistry; + harness.failNextStart(); + const failure = yield* provider + .startSession(threadId, { + provider: CLAUDE_AGENT_DRIVER, + providerInstanceId: claudeAgentInstanceId, + threadId, + cwd, + runtimeMode: "full-access", + }) + .pipe(Effect.flip); + const issuedToken = mcpBearerToken(harness.seenAuthorizationHeaders[0]); + + assert.equal(failure._tag, "ProviderAdapterRequestError"); + assert.equal(yield* registry.resolve(issuedToken), undefined); + assert.equal(McpProviderSession.readMcpProviderSession(threadId), undefined); + assert.equal(yield* harness.claude.adapter.hasSession(threadId), false); + }).pipe(Effect.provide(harness.layer)); +} + +describe("MCP credential across session replacement", () => { + it.effect( + "keeps the live query credential when replacement is refused", + keepsLiveMcpCredentialWhenSessionReplacementIsRefused, + ); + + it.effect( + "rotates the MCP credential when replacement is accepted", + rotatesMcpCredentialWhenSessionReplacementIsAccepted, + ); + + it.effect( + "clears the MCP credential when the first start fails", + clearsMcpCredentialWhenInitialStartFails, + ); +}); diff --git a/apps/server/src/provider/Layers/ProviderService.ts b/apps/server/src/provider/Layers/ProviderService.ts index cdac979c4dfd..5eed20de9675 100644 --- a/apps/server/src/provider/Layers/ProviderService.ts +++ b/apps/server/src/provider/Layers/ProviderService.ts @@ -940,10 +940,31 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( } satisfies Record; }); - const prepareMcpSession = (threadId: ThreadId, providerInstanceId: ProviderInstanceId) => - Effect.gen(function* () { + /** + * Issue the MCP credential for a provider start and publish it for the + * adapter to bake into the new process. `retainExisting` leaves the thread's + * current token registered so a rejected restart can keep the live query's + * credential. + */ + const prepareMcpSession = Effect.fn("ProviderService.prepareMcpSession")( + /** + * Issue the MCP credential for a provider start and publish it for the + * adapter to bake into the new process. `retainExisting` leaves the thread's + * current token registered so a rejected restart can keep the live query's + * credential. + */ + function* ( + threadId: ThreadId, + providerInstanceId: ProviderInstanceId, + options?: { readonly retainExisting?: boolean }, + ) { const capabilities = yield* agentAccessCapabilities(threadId); - const credential = yield* issueMcpCredential({ threadId, providerInstanceId, capabilities }); + const credential = yield* issueMcpCredential({ + threadId, + providerInstanceId, + capabilities, + ...(options?.retainExisting === true ? { retainExisting: true } : {}), + }); if (credential) { const deviceEnvironment = capabilities.has("device") ? yield* agentDeviceEnvironment @@ -956,12 +977,67 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( ); } return credential; - }); + }, + ); const clearMcpSession = (threadId: ThreadId) => McpSessionRegistry.revokeActiveMcpThread(threadId).pipe( Effect.tap(() => Effect.sync(() => McpProviderSession.clearMcpProviderSession(threadId))), ); + /** + * Settle MCP credentials after a provider start attempt. + * A rejection that leaves the existing session up restores the credential + * that query already has. An accepted replacement revokes the credential the + * stopped query was using. Any other failure revokes the thread. + */ + const settleMcpCredentialAfterStart = Effect.fn("ProviderService.settleMcpCredentialAfterStart")( + /** + * Settle MCP credentials after a provider start attempt. + * A rejection that leaves the existing session up restores the credential + * that query already has. An accepted replacement revokes the credential the + * stopped query was using. Any other failure revokes the thread. + */ + function* (input: { + readonly adapter: ProviderAdapterShape; + readonly threadId: ThreadId; + readonly hadActiveSession: boolean; + readonly previous: McpProviderSession.McpProviderSessionConfig | undefined; + readonly issuedProviderSessionId: string | undefined; + readonly accepted: boolean; + }) { + if (!input.accepted) { + // A probe failure must not revoke the token a live query may still hold. + const sessionStillActive = input.hadActiveSession + ? yield* input.adapter.hasSession(input.threadId).pipe(Effect.orElseSucceed(() => true)) + : false; + if (sessionStillActive) { + if (input.issuedProviderSessionId !== undefined) { + yield* McpSessionRegistry.revokeActiveMcpProviderSession(input.issuedProviderSessionId); + } + yield* Effect.sync(() => { + if (input.previous) { + McpProviderSession.setMcpProviderSession(input.previous); + return; + } + McpProviderSession.clearMcpProviderSession(input.threadId); + }); + return; + } + yield* clearMcpSession(input.threadId); + return; + } + + if ( + input.hadActiveSession && + input.previous !== undefined && + input.issuedProviderSessionId !== undefined && + input.previous.providerSessionId !== input.issuedProviderSessionId + ) { + yield* McpSessionRegistry.revokeActiveMcpProviderSession(input.previous.providerSessionId); + } + }, + ); + const publishRuntimeEvent = (event: ProviderRuntimeEvent): Effect.Effect => Effect.succeed(event).pipe( Effect.tap((canonicalEvent) => @@ -1394,7 +1470,19 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( ); }); + /** + * Start a provider session for a thread. + * When a session is already running, its MCP credential stays valid until + * this start is accepted. A rejection that leaves that session up restores + * the credential the live query is using. + */ const startSession: ProviderServiceMethod<"startSession"> = Effect.fn("startSession")( + /** + * Start a provider session for a thread. + * When a session is already running, its MCP credential stays valid until + * this start is accepted. A rejection that leaves that session up restores + * the credential the live query is using. + */ function* (threadId, rawInput) { const parsed = yield* decodeInputOrValidationError({ operation: "ProviderService.startSession", @@ -1499,8 +1587,16 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( } } const adapter = yield* registry.getByInstance(resolvedInstanceId); + // Hold the running query's MCP token until this start is accepted. + const hadActiveSession = yield* adapter.hasSession(threadId); + const previousMcpSession = hadActiveSession + ? McpProviderSession.readMcpProviderSession(threadId) + : undefined; yield* clearTurnAnalyticsSession(resolvedInstanceId, threadId); - yield* prepareMcpSession(threadId, resolvedInstanceId); + const issuedCredential = yield* prepareMcpSession(threadId, resolvedInstanceId, { + retainExisting: hadActiveSession, + }); + const issuedProviderSessionId = issuedCredential?.config.providerSessionId; const session = yield* adapter .startSession({ ...input, @@ -1508,7 +1604,18 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( ...(effectiveCwd !== undefined ? { cwd: effectiveCwd } : {}), ...(effectiveResumeCursor !== undefined ? { resumeCursor: effectiveResumeCursor } : {}), }) - .pipe(Effect.onError(() => clearMcpSession(threadId))); + .pipe( + Effect.onError(() => + settleMcpCredentialAfterStart({ + adapter, + threadId, + hadActiveSession, + previous: previousMcpSession, + issuedProviderSessionId, + accepted: false, + }), + ), + ); if (session.provider !== adapter.provider) { yield* clearMcpSession(threadId); @@ -1517,6 +1624,14 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( `Adapter/provider mismatch: requested '${adapter.provider}', received '${session.provider}'.`, ); } + yield* settleMcpCredentialAfterStart({ + adapter, + threadId, + hadActiveSession, + previous: previousMcpSession, + issuedProviderSessionId, + accepted: true, + }); const sessionWithInstance = { ...session, providerInstanceId: resolvedInstanceId,