From 0b31e51ec76398418be309148ec0de22d9efcef8 Mon Sep 17 00:00:00 2001 From: Igor Bedesqui Date: Mon, 20 Jul 2026 14:07:06 -0300 Subject: [PATCH 1/4] fix(mcp): revoke credentials by provider session --- .../server/src/mcp/McpSessionRegistry.test.ts | 25 +++++++++++++++++++ apps/server/src/mcp/McpSessionRegistry.ts | 12 +++++++++ 2 files changed, 37 insertions(+) diff --git a/apps/server/src/mcp/McpSessionRegistry.test.ts b/apps/server/src/mcp/McpSessionRegistry.test.ts index a91d98febd8..fd6da9c389b 100644 --- a/apps/server/src/mcp/McpSessionRegistry.test.ts +++ b/apps/server/src/mcp/McpSessionRegistry.test.ts @@ -55,6 +55,31 @@ it.effect("stores only a token hash, resolves the bearer token, and revokes by t }), ); +it.effect("revokes one provider session without revoking another credential for the thread", () => + Effect.gen(function* () { + const registry = yield* makeRegistry(() => 1_000); + const threadId = ThreadId.make("thread-selective-revocation"); + const first = yield* registry.issue({ + threadId, + providerInstanceId: ProviderInstanceId.make("codex"), + }); + const second = yield* registry.issue({ + threadId, + providerInstanceId: ProviderInstanceId.make("claudeAgent"), + }); + const firstToken = first.config.authorizationHeader.replace(/^Bearer\s+/, ""); + const secondToken = second.config.authorizationHeader.replace(/^Bearer\s+/, ""); + + yield* registry.revokeProviderSession(second.config.providerSessionId); + + expect(yield* registry.resolve(firstToken)).toMatchObject({ + threadId, + providerSessionId: first.config.providerSessionId, + }); + expect(yield* registry.resolve(secondToken)).toBeUndefined(); + }), +); + it.effect("builds MCP endpoints from the bound server host", () => Effect.gen(function* () { const cases = [ diff --git a/apps/server/src/mcp/McpSessionRegistry.ts b/apps/server/src/mcp/McpSessionRegistry.ts index 67c4f2f0ff0..1e6f85b5fef 100644 --- a/apps/server/src/mcp/McpSessionRegistry.ts +++ b/apps/server/src/mcp/McpSessionRegistry.ts @@ -202,6 +202,18 @@ export const issueActiveMcpCredential = ( .pipe(Effect.andThen(activeMcpSessionRegistry.issue(request))) : Effect.sync((): McpIssuedCredential | undefined => undefined); +export const issueUncommittedMcpCredential = ( + request: McpCredentialRequest, +): Effect.Effect => + activeMcpSessionRegistry + ? activeMcpSessionRegistry.issue(request) + : Effect.sync((): McpIssuedCredential | undefined => undefined); + +export const revokeActiveMcpProviderSession = (providerSessionId: string): Effect.Effect => + activeMcpSessionRegistry + ? activeMcpSessionRegistry.revokeProviderSession(providerSessionId) + : Effect.void; + export const revokeActiveMcpThread = (threadId: ThreadId): Effect.Effect => activeMcpSessionRegistry ? activeMcpSessionRegistry.revokeThread(threadId) : Effect.void; From d3b944e5faff0ee23199107e2d94de4d183bf9a2 Mon Sep 17 00:00:00 2001 From: Igor Bedesqui Date: Mon, 20 Jul 2026 14:09:33 -0300 Subject: [PATCH 2/4] fix(provider): make session ownership transitions atomic --- .../Layers/ProviderCommandReactor.ts | 25 +- .../provider/Layers/ProviderService.test.ts | 999 ++++++++++++++++++ .../src/provider/Layers/ProviderService.ts | 569 ++++++---- .../src/provider/Services/ProviderService.ts | 1 + 4 files changed, 1407 insertions(+), 187 deletions(-) diff --git a/apps/server/src/orchestration/Layers/ProviderCommandReactor.ts b/apps/server/src/orchestration/Layers/ProviderCommandReactor.ts index 9c7a7c94bb1..3b1f2ad62bb 100644 --- a/apps/server/src/orchestration/Layers/ProviderCommandReactor.ts +++ b/apps/server/src/orchestration/Layers/ProviderCommandReactor.ts @@ -474,16 +474,21 @@ const make = Effect.gen(function* () { const startProviderSession = (input?: { readonly resumeCursor?: unknown; readonly provider?: ProviderDriverKind; + readonly activeSession?: "reuse" | "replace"; }) => - providerService.startSession(threadId, { + providerService.startSession( threadId, - ...(preferredProvider ? { provider: preferredProvider } : {}), - providerInstanceId: desiredInstanceId, - ...(effectiveCwd ? { cwd: effectiveCwd } : {}), - modelSelection: desiredModelSelection, - ...(input?.resumeCursor !== undefined ? { resumeCursor: input.resumeCursor } : {}), - runtimeMode: desiredRuntimeMode, - }); + { + threadId, + ...(preferredProvider ? { provider: preferredProvider } : {}), + providerInstanceId: desiredInstanceId, + ...(effectiveCwd ? { cwd: effectiveCwd } : {}), + modelSelection: desiredModelSelection, + ...(input?.resumeCursor !== undefined ? { resumeCursor: input.resumeCursor } : {}), + runtimeMode: desiredRuntimeMode, + }, + { activeSession: input?.activeSession ?? "reuse" }, + ); const bindSessionToThread = (session: ProviderSession) => Effect.gen(function* () { @@ -564,7 +569,9 @@ const make = Effect.gen(function* () { hasResumeCursor: resumeCursor !== undefined, }); const restartedSession = yield* startProviderSession( - resumeCursor !== undefined ? { resumeCursor } : undefined, + resumeCursor !== undefined + ? { resumeCursor, activeSession: "replace" } + : { activeSession: "replace" }, ); yield* Effect.logInfo("provider command reactor restarted provider session", { threadId, diff --git a/apps/server/src/provider/Layers/ProviderService.test.ts b/apps/server/src/provider/Layers/ProviderService.test.ts index ccbbce1759f..8166d7d47de 100644 --- a/apps/server/src/provider/Layers/ProviderService.test.ts +++ b/apps/server/src/provider/Layers/ProviderService.test.ts @@ -12,6 +12,7 @@ import type { } from "@t3tools/contracts"; import { ApprovalRequestId, + EnvironmentId, EventId, ProviderDriverKind, ProviderInstanceId, @@ -22,6 +23,8 @@ import { import { createModelSelection } from "@t3tools/shared/model"; import { it, assert, vi } from "@effect/vitest"; +import * as Cause from "effect/Cause"; +import * as Deferred from "effect/Deferred"; import * as Effect from "effect/Effect"; import * as Exit from "effect/Exit"; import * as Fiber from "effect/Fiber"; @@ -33,11 +36,13 @@ import * as Ref from "effect/Ref"; 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 SqlClient from "effect/unstable/sql/SqlClient"; import { ProviderAdapterRequestError, ProviderAdapterSessionNotFoundError, + ProviderSessionDirectoryPersistenceError, ProviderUnsupportedError, ProviderValidationError, type ProviderAdapterError, @@ -51,10 +56,13 @@ import * as ProviderEventLoggers from "./ProviderEventLoggers.ts"; import { ProviderSessionDirectoryLive } from "./ProviderSessionDirectory.ts"; import * as NodeServices from "@effect/platform-node/NodeServices"; import * as ProviderSessionRuntime from "../../persistence/ProviderSessionRuntime.ts"; +import * as McpProviderSession from "../../mcp/McpProviderSession.ts"; +import * as McpSessionRegistry from "../../mcp/McpSessionRegistry.ts"; import { makeSqlitePersistenceLive, SqlitePersistenceMemory, } from "../../persistence/Layers/Sqlite.ts"; +import * as ServerEnvironment from "../../environment/ServerEnvironment.ts"; import * as ServerSettings from "../../serverSettings.ts"; import * as AnalyticsService from "../../telemetry/AnalyticsService.ts"; import { makeAdapterRegistryMock } from "../testUtils/providerAdapterRegistryMock.ts"; @@ -71,6 +79,52 @@ const CODEX_DRIVER = ProviderDriverKind.make("codex"); const CLAUDE_AGENT_DRIVER = ProviderDriverKind.make("claudeAgent"); const CURSOR_DRIVER = ProviderDriverKind.make("cursor"); +const makeStartInput = ( + threadId: ThreadId, + overrides: Partial> = {}, +): ProviderSessionStartInput => ({ + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + runtimeMode: "full-access", + ...overrides, + threadId, +}); + +const makeServiceLayer = ( + registry: ProviderAdapterRegistry.ProviderAdapterRegistryShape, + directory: ProviderSessionDirectory.ProviderSessionDirectoryShape, +) => + makeProviderServiceLive().pipe( + Layer.provide(Layer.succeed(ProviderAdapterRegistry.ProviderAdapterRegistry, registry)), + Layer.provide(Layer.succeed(ProviderSessionDirectory.ProviderSessionDirectory, directory)), + Layer.provide(defaultServerSettingsLayer), + Layer.provide(AnalyticsService.layerTest), + Layer.provide( + Layer.succeed( + ProviderEventLoggers.ProviderEventLoggers, + ProviderEventLoggers.NoOpProviderEventLoggers, + ), + ), + ); + +const mcpTestEnvironment = ServerEnvironment.ServerEnvironment.of({ + getEnvironmentId: Effect.succeed(EnvironmentId.make("environment-provider-service-test")), + getDescriptor: Effect.die("unused test environment method"), +}); +const mcpTestHttpServer = HttpServer.HttpServer.of({ + address: { _tag: "TcpAddress", hostname: "127.0.0.1", port: 43123 }, + serve: (() => Effect.void) as HttpServer.HttpServer["Service"]["serve"], +}); +const mcpTestRegistryLayer = McpSessionRegistry.layer.pipe( + Layer.provide( + Layer.mergeAll( + Layer.succeed(HttpServer.HttpServer, mcpTestHttpServer), + Layer.succeed(ServerEnvironment.ServerEnvironment, mcpTestEnvironment), + NodeServices.layer, + ), + ), +); + type LegacyProviderRuntimeEvent = { readonly type: string; readonly eventId: EventId; @@ -234,6 +288,9 @@ function makeFakeCodexAdapter(provider: ProviderDriverKind = CODEX_DRIVER) { } sessions.set(threadId, update(existing)); }; + const setSession = (session: ProviderSession): void => { + sessions.set(session.threadId, session); + }; return { adapter, @@ -250,6 +307,7 @@ function makeFakeCodexAdapter(provider: ProviderDriverKind = CODEX_DRIVER) { readThread, rollbackThread, stopAll, + setSession, }; } @@ -315,6 +373,901 @@ function makeProviderServiceLayer() { }; } +function makeBindingFailureHarness( + initialBindings: ReadonlyArray = [], + listBindingsFailure?: ProviderSessionDirectoryPersistenceError, +) { + const codex = makeFakeCodexAdapter(); + const claude = makeFakeCodexAdapter(CLAUDE_AGENT_DRIVER); + const bindings = new Map(initialBindings.map((binding) => [binding.threadId, binding])); + const persistenceFailure = new ProviderSessionDirectoryPersistenceError({ + operation: "upsert", + detail: "injected binding failure", + }); + const upsert = vi.fn(() => + Effect.fail(persistenceFailure), + ); + const directory = ProviderSessionDirectory.ProviderSessionDirectory.of({ + upsert, + getProvider: (threadId) => + Option.match(bindings.get(threadId) ? Option.some(bindings.get(threadId)!) : Option.none(), { + onNone: () => Effect.die("missing test binding"), + onSome: (binding) => Effect.succeed(binding.provider), + }), + getBinding: (threadId) => { + const binding = bindings.get(threadId); + return Effect.succeed(binding === undefined ? Option.none() : Option.some(binding)); + }, + listThreadIds: () => Effect.succeed(Array.from(bindings.keys())), + listBindings: () => + listBindingsFailure + ? Effect.fail(listBindingsFailure) + : Effect.succeed( + Array.from(bindings.values(), (binding) => ({ + ...binding, + lastSeenAt: "2026-01-01T00:00:00.000Z", + })), + ), + }); + const registry = makeAdapterRegistryMock({ + [CODEX_DRIVER]: codex.adapter, + [CLAUDE_AGENT_DRIVER]: claude.adapter, + }); + const layer = makeServiceLayer(registry, directory); + return { claude, codex, layer, persistenceFailure, upsert }; +} + +function makeStopGenerationHarness( + threadId: ThreadId, + upsert: ProviderSessionDirectory.ProviderSessionDirectoryShape["upsert"], +) { + const codex = makeFakeCodexAdapter(); + const binding: ProviderSessionDirectory.ProviderRuntimeBinding = { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: "full-access", + resumeCursor: { opaque: "persisted-resume" }, + }; + const directory = ProviderSessionDirectory.ProviderSessionDirectory.of({ + upsert, + getProvider: () => Effect.succeed(binding.provider), + getBinding: () => Effect.succeed(Option.some(binding)), + listThreadIds: () => Effect.succeed([threadId]), + listBindings: () => Effect.succeed([{ ...binding, lastSeenAt: "2026-01-01T00:00:00.000Z" }]), + }); + const registry = makeAdapterRegistryMock({ [CODEX_DRIVER]: codex.adapter }); + const layer = makeServiceLayer(registry, directory); + return { codex, layer }; +} + +it.effect("ProviderServiceLive compensates a newly started session when binding fails", () => + Effect.gen(function* () { + const harness = makeBindingFailureHarness(); + const threadId = asThreadId("thread-binding-failure-new"); + const exit = yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + return yield* Effect.exit(provider.startSession(threadId, makeStartInput(threadId))); + }).pipe(Effect.provide(harness.layer)); + + assert.equal(Exit.isFailure(exit), true); + if (Exit.isFailure(exit)) { + assert.equal( + Option.getOrUndefined(Cause.findErrorOption(exit.cause)), + harness.persistenceFailure, + ); + } + assert.deepEqual(harness.codex.stopSession.mock.calls, [[threadId]]); + }), +); + +it.effect("ProviderServiceLive does not compensate a preexisting adapter session", () => + Effect.gen(function* () { + const threadId = asThreadId("thread-binding-failure-existing"); + const harness = makeBindingFailureHarness([ + { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: "full-access", + resumeCursor: { opaque: "existing-resume" }, + }, + ]); + const previousMcpSession = { + environmentId: EnvironmentId.make("environment-test"), + threadId, + providerSessionId: "provider-session-existing", + providerInstanceId: codexInstanceId, + endpoint: "http://127.0.0.1/mcp", + authorizationHeader: "Bearer existing", + }; + McpProviderSession.setMcpProviderSession(previousMcpSession); + yield* harness.codex.adapter.startSession(makeStartInput(threadId)); + + const exit = yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + return yield* Effect.exit( + provider.startSession(threadId, makeStartInput(threadId), { activeSession: "reuse" }), + ); + }).pipe(Effect.provide(harness.layer)); + + assert.equal(Exit.isFailure(exit), true); + if (Exit.isFailure(exit)) { + assert.equal( + Option.getOrUndefined(Cause.findErrorOption(exit.cause)), + harness.persistenceFailure, + ); + } + assert.equal(harness.codex.stopSession.mock.calls.length, 0); + assert.equal(yield* harness.codex.adapter.hasSession(threadId), true); + assert.equal(McpProviderSession.readMcpProviderSession(threadId), previousMcpSession); + McpProviderSession.clearMcpProviderSession(threadId); + }), +); + +it.effect("ProviderServiceLive owns and compensates an explicit replacement", () => + Effect.gen(function* () { + const harness = makeBindingFailureHarness(); + const threadId = asThreadId("thread-binding-failure-replacement"); + yield* harness.codex.adapter.startSession(makeStartInput(threadId)); + + const exit = yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + return yield* Effect.exit( + provider.startSession(threadId, makeStartInput(threadId), { activeSession: "replace" }), + ); + }).pipe(Effect.provide(harness.layer)); + + assert.equal(Exit.isFailure(exit), true); + if (Exit.isFailure(exit)) { + assert.equal( + Option.getOrUndefined(Cause.findErrorOption(exit.cause)), + harness.persistenceFailure, + ); + } + assert.equal(harness.codex.startSession.mock.calls.length, 2); + assert.deepEqual(harness.codex.stopSession.mock.calls, [[threadId], [threadId]]); + assert.equal(yield* harness.codex.adapter.hasSession(threadId), false); + }), +); + +it.effect("ProviderServiceLive preserves binding failures when compensation cleanup fails", () => + Effect.gen(function* () { + const harness = makeBindingFailureHarness(); + const threadId = asThreadId("thread-binding-failure-cleanup-failure"); + harness.codex.stopSession.mockImplementation(() => + Effect.fail( + new ProviderAdapterRequestError({ + provider: String(CODEX_DRIVER), + method: "stopSession", + detail: "injected cleanup failure", + }), + ), + ); + + const exit = yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + return yield* Effect.exit(provider.startSession(threadId, makeStartInput(threadId))); + }).pipe(Effect.provide(harness.layer)); + + assert.equal(Exit.isFailure(exit), true); + if (Exit.isFailure(exit)) { + assert.equal( + Option.getOrUndefined(Cause.findErrorOption(exit.cause)), + harness.persistenceFailure, + ); + } + assert.deepEqual(harness.codex.stopSession.mock.calls, [[threadId]]); + }), +); + +it.effect("ProviderServiceLive serializes competing starts through persistence and cleanup", () => + Effect.gen(function* () { + const firstUpsertStarted = yield* Deferred.make(); + const releaseFirstUpsert = yield* Deferred.make(); + const codex = makeFakeCodexAdapter(); + const persistenceFailure = new ProviderSessionDirectoryPersistenceError({ + operation: "upsert", + detail: "injected first binding failure", + }); + let upsertCount = 0; + const directory = ProviderSessionDirectory.ProviderSessionDirectory.of({ + upsert: () => { + upsertCount += 1; + return upsertCount === 1 + ? Deferred.succeed(firstUpsertStarted, undefined).pipe( + Effect.andThen(Deferred.await(releaseFirstUpsert)), + Effect.andThen(Effect.fail(persistenceFailure)), + ) + : Effect.void; + }, + getProvider: () => Effect.die("unused test directory method"), + getBinding: () => Effect.succeed(Option.none()), + listThreadIds: () => Effect.succeed([]), + listBindings: () => Effect.succeed([]), + }); + const registry = makeAdapterRegistryMock({ [CODEX_DRIVER]: codex.adapter }); + const layer = makeServiceLayer(registry, directory); + const threadId = asThreadId("thread-binding-failure-concurrent"); + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const start = () => provider.startSession(threadId, makeStartInput(threadId)); + const first = yield* Effect.forkChild(start()); + yield* Deferred.await(firstUpsertStarted); + const second = yield* Effect.forkChild(start()); + yield* Effect.yieldNow; + assert.equal(codex.startSession.mock.calls.length, 1); + + yield* Deferred.succeed(releaseFirstUpsert, undefined); + const firstExit = yield* Fiber.await(first); + assert.equal(Exit.isFailure(firstExit), true); + if (Exit.isFailure(firstExit)) { + assert.equal( + Option.getOrUndefined(Cause.findErrorOption(firstExit.cause)), + persistenceFailure, + ); + } + yield* Fiber.join(second); + assert.equal(yield* codex.adapter.hasSession(threadId), true); + }).pipe(Effect.provide(layer)); + + assert.equal(codex.startSession.mock.calls.length, 2); + assert.deepEqual(codex.stopSession.mock.calls, [[threadId]]); + }), +); + +it.effect("ProviderServiceLive rejects stale updates after same-instance replacement", () => + Effect.gen(function* () { + const replacementUpsertStarted = yield* Deferred.make(); + const releaseReplacementUpsert = yield* Deferred.make(); + const staleTurnStarted = yield* Deferred.make(); + const releaseStaleTurn = yield* Deferred.make(); + const threadId = asThreadId("thread-binding-recovery-race"); + const codex = makeFakeCodexAdapter(); + codex.sendTurn.mockImplementation((input) => + Deferred.succeed(staleTurnStarted, undefined).pipe( + Effect.andThen(Deferred.await(releaseStaleTurn)), + Effect.as({ + threadId: input.threadId, + turnId: TurnId.make("turn-stale-provider"), + }), + ), + ); + let binding: ProviderSessionDirectory.ProviderRuntimeBinding = { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: "full-access", + resumeCursor: { opaque: "codex-resume" }, + }; + let upsertCount = 0; + let staleTurnWrite: ProviderSessionDirectory.ProviderRuntimeBinding | undefined; + const directory = ProviderSessionDirectory.ProviderSessionDirectory.of({ + upsert: (nextBinding) => { + upsertCount += 1; + const commit = Effect.sync(() => { + binding = { ...binding, ...nextBinding }; + if ( + nextBinding.runtimePayload && + typeof nextBinding.runtimePayload === "object" && + "activeTurnId" in nextBinding.runtimePayload && + nextBinding.runtimePayload.activeTurnId === "turn-stale-provider" + ) { + staleTurnWrite = nextBinding; + } + }); + return upsertCount === 1 + ? Deferred.succeed(replacementUpsertStarted, undefined).pipe( + Effect.andThen(Deferred.await(releaseReplacementUpsert)), + Effect.andThen(commit), + ) + : commit; + }, + getProvider: () => Effect.succeed(binding.provider), + getBinding: () => Effect.succeed(Option.some(binding)), + listThreadIds: () => Effect.succeed([threadId]), + listBindings: () => Effect.succeed([{ ...binding, lastSeenAt: "2026-01-01T00:00:00.000Z" }]), + }); + const registry = makeAdapterRegistryMock({ [CODEX_DRIVER]: codex.adapter }); + const layer = makeServiceLayer(registry, directory); + yield* codex.adapter.startSession(makeStartInput(threadId)); + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const turn = yield* Effect.forkChild( + provider.sendTurn({ + threadId, + input: "route before replacement", + attachments: [], + }), + ); + yield* Deferred.await(staleTurnStarted); + const replacement = yield* Effect.forkChild( + provider.startSession(threadId, makeStartInput(threadId), { activeSession: "replace" }), + ); + yield* Deferred.await(replacementUpsertStarted); + yield* Deferred.succeed(releaseStaleTurn, undefined); + yield* Effect.yieldNow; + yield* Deferred.succeed(releaseReplacementUpsert, undefined); + yield* Fiber.join(replacement); + yield* Fiber.join(turn); + }).pipe(Effect.provide(layer)); + + assert.equal(binding.providerInstanceId, codexInstanceId); + assert.equal(codex.startSession.mock.calls.length, 2); + assert.equal(codex.stopSession.mock.calls.length, 1); + assert.equal(codex.sendTurn.mock.calls.length, 1); + assert.equal(staleTurnWrite, undefined); + }), +); + +it.effect("ProviderServiceLive does not persist a stale turn after failed replacement", () => + Effect.gen(function* () { + const staleTurnStarted = yield* Deferred.make(); + const releaseStaleTurn = yield* Deferred.make(); + const threadId = asThreadId("thread-binding-failed-replacement-race"); + const harness = makeBindingFailureHarness([ + { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: "full-access", + resumeCursor: { opaque: "old-session" }, + }, + ]); + let replacementWriteSeen = false; + let delayedWrite: ProviderSessionDirectory.ProviderRuntimeBinding | undefined; + harness.upsert.mockImplementation((binding) => { + if (!replacementWriteSeen) { + replacementWriteSeen = true; + return Effect.fail(harness.persistenceFailure); + } + return Effect.sync(() => { + delayedWrite = binding; + }); + }); + yield* harness.codex.adapter.startSession(makeStartInput(threadId)); + harness.codex.sendTurn.mockImplementation((input) => + Deferred.succeed(staleTurnStarted, undefined).pipe( + Effect.andThen(Deferred.await(releaseStaleTurn)), + Effect.as({ + threadId: input.threadId, + turnId: TurnId.make("turn-stale-failed-replacement"), + }), + ), + ); + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const turn = yield* Effect.forkChild( + provider.sendTurn({ + threadId, + input: "started before failed replacement", + attachments: [], + }), + ); + yield* Deferred.await(staleTurnStarted); + const replacementExit = yield* Effect.exit( + provider.startSession(threadId, makeStartInput(threadId), { activeSession: "replace" }), + ); + assert.equal(Exit.isFailure(replacementExit), true); + yield* Deferred.succeed(releaseStaleTurn, undefined); + yield* Fiber.join(turn); + assert.equal(delayedWrite, undefined); + }).pipe(Effect.provide(harness.layer)); + }), +); + +it.effect("ProviderServiceLive rejects delayed turn updates when stop persistence fails", () => + Effect.gen(function* () { + const threadId = asThreadId("thread-stop-persistence-failure-race"); + const turnStarted = yield* Deferred.make(); + const releaseTurn = yield* Deferred.make(); + const persistenceFailure = new ProviderSessionDirectoryPersistenceError({ + operation: "upsert", + detail: "injected stop binding failure", + }); + const upsert = vi.fn(() => + Effect.fail(persistenceFailure), + ); + const harness = makeStopGenerationHarness(threadId, upsert); + yield* harness.codex.adapter.startSession(makeStartInput(threadId)); + harness.codex.sendTurn.mockImplementation((input) => + Deferred.succeed(turnStarted, undefined).pipe( + Effect.andThen(Deferred.await(releaseTurn)), + Effect.as({ threadId: input.threadId, turnId: TurnId.make("turn-before-failed-stop") }), + ), + ); + let upsertCountBeforeStop = 0; + let operationWrites: ReadonlyArray = []; + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + upsertCountBeforeStop = upsert.mock.calls.length; + const turn = yield* Effect.forkChild( + provider.sendTurn({ threadId, input: "before failed stop", attachments: [] }), + ); + yield* Deferred.await(turnStarted); + const stopExit = yield* Effect.exit(provider.stopSession({ threadId })); + assert.equal(Exit.isFailure(stopExit), true); + if (Exit.isFailure(stopExit)) { + assert.equal( + Option.getOrUndefined(Cause.findErrorOption(stopExit.cause)), + persistenceFailure, + ); + } + yield* Deferred.succeed(releaseTurn, undefined); + yield* Fiber.join(turn); + operationWrites = upsert.mock.calls.slice(upsertCountBeforeStop).map(([binding]) => binding); + }).pipe(Effect.provide(harness.layer)); + + assert.deepEqual( + operationWrites.map((binding) => binding.status), + ["stopped"], + ); + }), +); + +it.effect("ProviderServiceLive rejects delayed turn updates when stop is interrupted", () => + Effect.gen(function* () { + const threadId = asThreadId("thread-stop-interruption-race"); + const turnStarted = yield* Deferred.make(); + const releaseTurn = yield* Deferred.make(); + const stopStarted = yield* Deferred.make(); + const upsert = vi.fn( + () => Effect.void, + ); + const harness = makeStopGenerationHarness(threadId, upsert); + yield* harness.codex.adapter.startSession(makeStartInput(threadId)); + harness.codex.sendTurn.mockImplementation((input) => + Deferred.succeed(turnStarted, undefined).pipe( + Effect.andThen(Deferred.await(releaseTurn)), + Effect.as({ + threadId: input.threadId, + turnId: TurnId.make("turn-before-interrupted-stop"), + }), + ), + ); + harness.codex.stopSession.mockImplementation(() => + Deferred.succeed(stopStarted, undefined).pipe(Effect.andThen(Effect.never)), + ); + let upsertCountBeforeStop = 0; + let operationWrites: ReadonlyArray = []; + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + upsertCountBeforeStop = upsert.mock.calls.length; + const turn = yield* Effect.forkChild( + provider.sendTurn({ threadId, input: "before interrupted stop", attachments: [] }), + ); + yield* Deferred.await(turnStarted); + const stop = yield* Effect.forkChild(provider.stopSession({ threadId })); + yield* Deferred.await(stopStarted); + yield* Fiber.interrupt(stop); + yield* Deferred.succeed(releaseTurn, undefined); + yield* Fiber.join(turn); + operationWrites = upsert.mock.calls.slice(upsertCountBeforeStop).map(([binding]) => binding); + }).pipe(Effect.provide(harness.layer)); + + assert.deepEqual(operationWrites, []); + }), +); + +it.effect( + "ProviderServiceLive rejects delayed turn updates when replacement stop is interrupted", + () => + Effect.gen(function* () { + const threadId = asThreadId("thread-replacement-stop-interruption-race"); + const turnStarted = yield* Deferred.make(); + const releaseTurn = yield* Deferred.make(); + const stopStarted = yield* Deferred.make(); + const upsert = vi.fn( + () => Effect.void, + ); + const harness = makeStopGenerationHarness(threadId, upsert); + yield* harness.codex.adapter.startSession(makeStartInput(threadId)); + harness.codex.sendTurn.mockImplementation((input) => + Deferred.succeed(turnStarted, undefined).pipe( + Effect.andThen(Deferred.await(releaseTurn)), + Effect.as({ + threadId: input.threadId, + turnId: TurnId.make("turn-before-interrupted-replacement"), + }), + ), + ); + harness.codex.stopSession.mockImplementation(() => + Deferred.succeed(stopStarted, undefined).pipe(Effect.andThen(Effect.never)), + ); + let operationWrites: ReadonlyArray = []; + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const upsertCountBeforeReplacement = upsert.mock.calls.length; + const turn = yield* Effect.forkChild( + provider.sendTurn({ threadId, input: "before interrupted replacement", attachments: [] }), + ); + yield* Deferred.await(turnStarted); + const replacement = yield* Effect.forkChild( + provider.startSession(threadId, makeStartInput(threadId), { activeSession: "replace" }), + ); + yield* Deferred.await(stopStarted); + yield* Fiber.interrupt(replacement); + yield* Deferred.succeed(releaseTurn, undefined); + yield* Fiber.join(turn); + operationWrites = upsert.mock.calls + .slice(upsertCountBeforeReplacement) + .map(([binding]) => binding); + }).pipe(Effect.provide(harness.layer)); + + assert.deepEqual(operationWrites, []); + }), +); + +it.effect("ProviderServiceLive compensates an interrupted start after adapter startup", () => + Effect.gen(function* () { + const upsertStarted = yield* Deferred.make(); + const codex = makeFakeCodexAdapter(); + const directory = ProviderSessionDirectory.ProviderSessionDirectory.of({ + upsert: () => Deferred.succeed(upsertStarted, undefined).pipe(Effect.andThen(Effect.never)), + getProvider: () => Effect.die("unused test directory method"), + getBinding: () => Effect.succeed(Option.none()), + listThreadIds: () => Effect.succeed([]), + listBindings: () => Effect.succeed([]), + }); + const registry = makeAdapterRegistryMock({ [CODEX_DRIVER]: codex.adapter }); + const layer = makeServiceLayer(registry, directory); + const threadId = asThreadId("thread-binding-failure-interrupted"); + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const start = yield* Effect.forkChild( + provider.startSession(threadId, makeStartInput(threadId)), + ); + yield* Deferred.await(upsertStarted); + yield* Fiber.interrupt(start); + }).pipe(Effect.provide(layer)); + + assert.deepEqual(codex.stopSession.mock.calls, [[threadId]]); + assert.equal(yield* codex.adapter.hasSession(threadId), false); + }), +); + +it.effect("ProviderServiceLive compensates a resumed session when binding fails", () => + Effect.gen(function* () { + const threadId = asThreadId("thread-binding-failure-recovery"); + const harness = makeBindingFailureHarness([ + { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: "full-access", + resumeCursor: { opaque: "persisted-resume" }, + }, + ]); + const exit = yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + return yield* Effect.exit(provider.rollbackConversation({ threadId, numTurns: 1 })); + }).pipe(Effect.provide(harness.layer)); + + assert.equal(Exit.isFailure(exit), true); + if (Exit.isFailure(exit)) { + assert.equal( + Option.getOrUndefined(Cause.findErrorOption(exit.cause)), + harness.persistenceFailure, + ); + } + assert.equal(harness.codex.startSession.mock.calls.length, 1); + assert.deepEqual(harness.codex.stopSession.mock.calls, [[threadId]]); + }), +); + +it.effect("ProviderServiceLive compensates a started session with a mismatched provider", () => + Effect.gen(function* () { + const registry = yield* McpSessionRegistry.McpSessionRegistry; + const harness = makeBindingFailureHarness(); + const threadId = asThreadId("thread-provider-mismatch-start"); + let attemptCredential: McpProviderSession.McpProviderSessionConfig | undefined; + harness.codex.startSession.mockImplementation((input) => + Effect.sync(() => { + attemptCredential = McpProviderSession.readMcpProviderSession(threadId); + const session: ProviderSession = { + provider: CLAUDE_AGENT_DRIVER, + providerInstanceId: codexInstanceId, + status: "ready", + runtimeMode: input.runtimeMode, + threadId: input.threadId, + cwd: input.cwd ?? process.cwd(), + createdAt: "2026-01-01T00:00:00.000Z", + updatedAt: "2026-01-01T00:00:00.000Z", + }; + harness.codex.setSession(session); + return session; + }), + ); + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const exit = yield* Effect.exit(provider.startSession(threadId, makeStartInput(threadId))); + assert.equal(Exit.isFailure(exit), true); + if (Exit.isFailure(exit)) { + const error = Option.getOrUndefined(Cause.findErrorOption(exit.cause)); + assert.equal(error?._tag, "ProviderValidationError"); + if (error?._tag === "ProviderValidationError") { + assert.equal( + error.issue, + "Adapter/provider mismatch: requested 'codex', received 'claudeAgent'.", + ); + } + } + assert.deepEqual(harness.codex.stopSession.mock.calls, [[threadId]]); + assert.equal(yield* harness.codex.adapter.hasSession(threadId), false); + assert.equal(McpProviderSession.readMcpProviderSession(threadId), undefined); + assert.equal(attemptCredential !== undefined, true); + const attemptToken = attemptCredential?.authorizationHeader.replace(/^Bearer\s+/, "") ?? ""; + assert.equal(yield* registry.resolve(attemptToken), undefined); + }).pipe(Effect.provide(harness.layer)); + }).pipe(Effect.provide(mcpTestRegistryLayer)), +); + +it.effect("ProviderServiceLive restores MCP state when recovery mismatch cleanup fails", () => + Effect.gen(function* () { + const registry = yield* McpSessionRegistry.McpSessionRegistry; + const threadId = asThreadId("thread-provider-mismatch-recovery"); + const harness = makeBindingFailureHarness([ + { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: "full-access", + resumeCursor: { opaque: "persisted-resume" }, + }, + ]); + let attemptCredential: McpProviderSession.McpProviderSessionConfig | undefined; + harness.codex.startSession.mockImplementation((input) => + Effect.sync(() => { + attemptCredential = McpProviderSession.readMcpProviderSession(threadId); + const session: ProviderSession = { + provider: CLAUDE_AGENT_DRIVER, + providerInstanceId: codexInstanceId, + status: "ready", + runtimeMode: input.runtimeMode, + threadId: input.threadId, + cwd: input.cwd ?? process.cwd(), + createdAt: "2026-01-01T00:00:00.000Z", + updatedAt: "2026-01-01T00:00:00.000Z", + }; + harness.codex.setSession(session); + return session; + }), + ); + const cleanupFailure = new ProviderAdapterRequestError({ + provider: String(CODEX_DRIVER), + method: "stopSession", + detail: "injected mismatch cleanup failure", + }); + harness.codex.stopSession.mockImplementation(() => Effect.fail(cleanupFailure)); + const oldCredential = yield* registry.issue({ + threadId, + providerInstanceId: codexInstanceId, + }); + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + McpProviderSession.setMcpProviderSession(oldCredential.config); + const exit = yield* Effect.exit(provider.rollbackConversation({ threadId, numTurns: 1 })); + + assert.equal(Exit.isFailure(exit), true); + if (Exit.isFailure(exit)) { + const error = Option.getOrUndefined(Cause.findErrorOption(exit.cause)); + assert.equal(error?._tag, "ProviderValidationError"); + if (error?._tag === "ProviderValidationError") { + assert.equal( + error.issue, + `Adapter/provider mismatch while recovering thread '${threadId}'. Expected 'codex', received 'claudeAgent'.`, + ); + } + } + assert.deepEqual(harness.codex.stopSession.mock.calls, [[threadId]]); + assert.equal(McpProviderSession.readMcpProviderSession(threadId), oldCredential.config); + assert.equal(attemptCredential !== undefined, true); + const oldToken = oldCredential.config.authorizationHeader.replace(/^Bearer\s+/, ""); + const attemptToken = attemptCredential?.authorizationHeader.replace(/^Bearer\s+/, "") ?? ""; + assert.equal( + (yield* registry.resolve(oldToken))?.providerSessionId, + oldCredential.config.providerSessionId, + ); + assert.equal(yield* registry.resolve(attemptToken), undefined); + McpProviderSession.clearMcpProviderSession(threadId); + }).pipe(Effect.provide(harness.layer)); + }).pipe(Effect.provide(mcpTestRegistryLayer)), +); + +it.effect( + "ProviderServiceLive replaces stale target sessions without revoking the old provider", + () => + Effect.gen(function* () { + const threadId = asThreadId("thread-binding-failure-provider-switch"); + const harness = makeBindingFailureHarness([ + { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: "full-access", + resumeCursor: { opaque: "codex-resume" }, + }, + ]); + const previousMcpSession = { + environmentId: EnvironmentId.make("environment-test"), + threadId, + providerSessionId: "provider-session-old-provider", + providerInstanceId: codexInstanceId, + endpoint: "http://127.0.0.1/mcp", + authorizationHeader: "Bearer old-provider", + }; + McpProviderSession.setMcpProviderSession(previousMcpSession); + yield* harness.codex.adapter.startSession(makeStartInput(threadId)); + yield* harness.claude.adapter.startSession( + makeStartInput(threadId, { + provider: CLAUDE_AGENT_DRIVER, + providerInstanceId: claudeAgentInstanceId, + }), + ); + + const exit = yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + return yield* Effect.exit( + provider.startSession( + threadId, + makeStartInput(threadId, { + provider: CLAUDE_AGENT_DRIVER, + providerInstanceId: claudeAgentInstanceId, + }), + { activeSession: "reuse" }, + ), + ); + }).pipe(Effect.provide(harness.layer)); + + assert.equal(Exit.isFailure(exit), true); + if (Exit.isFailure(exit)) { + assert.equal( + Option.getOrUndefined(Cause.findErrorOption(exit.cause)), + harness.persistenceFailure, + ); + } + assert.equal(yield* harness.codex.adapter.hasSession(threadId), true); + assert.equal(harness.codex.stopSession.mock.calls.length, 0); + assert.deepEqual(harness.claude.stopSession.mock.calls, [[threadId], [threadId]]); + assert.equal(yield* harness.claude.adapter.hasSession(threadId), false); + assert.equal(McpProviderSession.readMcpProviderSession(threadId), previousMcpSession); + McpProviderSession.clearMcpProviderSession(threadId); + }), +); + +it.effect( + "ProviderServiceLive retains the old MCP credential when replacement persistence fails", + () => + Effect.gen(function* () { + const registry = yield* McpSessionRegistry.McpSessionRegistry; + const threadId = asThreadId("thread-mcp-failed-replacement"); + const harness = makeBindingFailureHarness([ + { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: "full-access", + resumeCursor: { opaque: "codex-resume" }, + }, + ]); + yield* harness.codex.adapter.startSession(makeStartInput(threadId)); + const oldCredential = yield* registry.issue({ + threadId, + providerInstanceId: codexInstanceId, + }); + let attemptCredential: McpProviderSession.McpProviderSessionConfig | undefined; + harness.upsert.mockImplementation(() => + Effect.sync(() => { + const current = McpProviderSession.readMcpProviderSession(threadId); + if (current?.providerInstanceId === claudeAgentInstanceId) { + attemptCredential = current; + } + return current; + }).pipe( + Effect.flatMap((current) => + current?.providerInstanceId === claudeAgentInstanceId + ? Effect.fail(harness.persistenceFailure) + : Effect.void, + ), + ), + ); + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + McpProviderSession.setMcpProviderSession(oldCredential.config); + const exit = yield* Effect.exit( + provider.startSession( + threadId, + makeStartInput(threadId, { + provider: CLAUDE_AGENT_DRIVER, + providerInstanceId: claudeAgentInstanceId, + }), + ), + ); + assert.equal(Exit.isFailure(exit), true); + assert.equal(McpProviderSession.readMcpProviderSession(threadId), oldCredential.config); + assert.equal(attemptCredential !== undefined, true); + const oldToken = oldCredential.config.authorizationHeader.replace(/^Bearer\s+/, ""); + const attemptToken = attemptCredential?.authorizationHeader.replace(/^Bearer\s+/, "") ?? ""; + assert.equal( + (yield* registry.resolve(oldToken))?.providerSessionId, + oldCredential.config.providerSessionId, + ); + assert.equal(yield* registry.resolve(attemptToken), undefined); + }).pipe(Effect.provide(harness.layer)); + + McpProviderSession.clearMcpProviderSession(threadId); + }).pipe(Effect.provide(mcpTestRegistryLayer)), +); + +it.effect("ProviderServiceLive commits the new MCP credential after successful replacement", () => + Effect.gen(function* () { + const registry = yield* McpSessionRegistry.McpSessionRegistry; + const codex = makeFakeCodexAdapter(); + const claude = makeFakeCodexAdapter(CLAUDE_AGENT_DRIVER); + const bindings = new Map(); + const directory = ProviderSessionDirectory.ProviderSessionDirectory.of({ + upsert: (binding) => + Effect.sync(() => { + bindings.set(binding.threadId, binding); + }), + getProvider: (threadId) => Effect.succeed(bindings.get(threadId)?.provider ?? CODEX_DRIVER), + getBinding: (threadId) => { + const binding = bindings.get(threadId); + return Effect.succeed(binding ? Option.some(binding) : Option.none()); + }, + listThreadIds: () => Effect.succeed(Array.from(bindings.keys())), + listBindings: () => + Effect.succeed( + Array.from(bindings.values(), (binding) => ({ + ...binding, + lastSeenAt: "2026-01-01T00:00:00.000Z", + })), + ), + }); + const adapterRegistry = makeAdapterRegistryMock({ + [CODEX_DRIVER]: codex.adapter, + [CLAUDE_AGENT_DRIVER]: claude.adapter, + }); + const providerLayer = makeServiceLayer(adapterRegistry, directory); + const threadId = asThreadId("thread-mcp-successful-replacement"); + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + yield* provider.startSession(threadId, makeStartInput(threadId)); + const oldCredential = McpProviderSession.readMcpProviderSession(threadId); + assert.equal(oldCredential !== undefined, true); + + yield* provider.startSession( + threadId, + makeStartInput(threadId, { + provider: CLAUDE_AGENT_DRIVER, + providerInstanceId: claudeAgentInstanceId, + }), + ); + const currentCredential = McpProviderSession.readMcpProviderSession(threadId); + assert.equal(currentCredential !== undefined, true); + const oldToken = oldCredential?.authorizationHeader.replace(/^Bearer\s+/, "") ?? ""; + const currentToken = currentCredential?.authorizationHeader.replace(/^Bearer\s+/, "") ?? ""; + assert.equal(yield* registry.resolve(oldToken), undefined); + assert.equal( + (yield* registry.resolve(currentToken))?.providerSessionId, + currentCredential?.providerSessionId, + ); + McpProviderSession.clearMcpProviderSession(threadId); + }).pipe(Effect.provide(providerLayer)); + }).pipe(Effect.provide(mcpTestRegistryLayer)), +); + it.effect("ProviderServiceLive catches stopAll failures during shutdown", () => Effect.gen(function* () { const codex = makeFakeCodexAdapter(); @@ -1062,6 +2015,7 @@ routing.layer("ProviderServiceLive routing", (it) => { cwd: "/tmp/project-binding-mismatch", runtimeMode: "full-access", }); + yield* directory.upsert({ threadId, provider: ProviderDriverKind.make("claudeAgent"), @@ -1119,6 +2073,51 @@ routing.layer("ProviderServiceLive routing", (it) => { }), ); + it.effect("finishes stale-session cleanup when interrupted after replacement persistence", () => + Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const threadId = asThreadId("thread-provider-replacement-interrupted"); + const staleStopStarted = yield* Deferred.make(); + const releaseStaleStop = yield* Deferred.make(); + + yield* provider.startSession(threadId, { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: "full-access", + }); + const originalStopSession = routing.codex.stopSession.getMockImplementation(); + assert.equal(originalStopSession !== undefined, true); + routing.codex.stopSession.mockImplementation((stoppedThreadId) => + Deferred.succeed(staleStopStarted, undefined).pipe( + Effect.andThen(Deferred.await(releaseStaleStop)), + Effect.andThen(originalStopSession!(stoppedThreadId)), + ), + ); + + const replacement = yield* Effect.forkChild( + provider.startSession(threadId, { + provider: CLAUDE_AGENT_DRIVER, + providerInstanceId: claudeAgentInstanceId, + threadId, + runtimeMode: "full-access", + }), + ); + yield* Deferred.await(staleStopStarted); + const interruption = yield* Effect.forkChild(Fiber.interrupt(replacement)); + yield* Effect.yieldNow; + yield* Deferred.succeed(releaseStaleStop, undefined); + yield* Fiber.join(interruption); + routing.codex.stopSession.mockImplementation(originalStopSession!); + + assert.equal(yield* routing.codex.adapter.hasSession(threadId), false); + assert.equal(yield* routing.claude.adapter.hasSession(threadId), true); + const directory = yield* ProviderSessionDirectory.ProviderSessionDirectory; + const binding = Option.getOrUndefined(yield* directory.getBinding(threadId)); + assert.equal(binding?.providerInstanceId, claudeAgentInstanceId); + }), + ); + it.effect("recovers stale sessions for sendTurn using persisted cwd", () => Effect.gen(function* () { const provider = yield* ProviderService.ProviderService; diff --git a/apps/server/src/provider/Layers/ProviderService.ts b/apps/server/src/provider/Layers/ProviderService.ts index 2eaaeb8ce3c..53c6a05d027 100644 --- a/apps/server/src/provider/Layers/ProviderService.ts +++ b/apps/server/src/provider/Layers/ProviderService.ts @@ -33,7 +33,9 @@ import * as PubSub from "effect/PubSub"; import * as Ref from "effect/Ref"; import * as Schema from "effect/Schema"; import * as SchemaIssue from "effect/SchemaIssue"; +import * as Semaphore from "effect/Semaphore"; import * as Stream from "effect/Stream"; +import * as SynchronizedRef from "effect/SynchronizedRef"; import { increment, @@ -214,18 +216,103 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( const directory = yield* ProviderSessionDirectory.ProviderSessionDirectory; const runtimeEventPubSub = yield* PubSub.unbounded(); const nowIso = Effect.map(DateTime.now, DateTime.formatIso); - const prepareMcpSession = (threadId: ThreadId, providerInstanceId: ProviderInstanceId) => - McpSessionRegistry.issueActiveMcpCredential({ threadId, providerInstanceId }).pipe( - Effect.tap((credential) => - credential - ? Effect.sync(() => McpProviderSession.setMcpProviderSession(credential.config)) - : Effect.void, - ), - ); + const threadLocks = yield* SynchronizedRef.make(new Map()); + const threadGenerations = yield* Ref.make(new Map()); + const getThreadLock = (threadId: ThreadId) => + SynchronizedRef.modifyEffect(threadLocks, (current) => { + const existing = current.get(threadId); + if (existing) return Effect.succeed([existing, current] as const); + return Semaphore.make(1).pipe( + Effect.map((created) => [created, new Map(current).set(threadId, created)] as const), + ); + }); + const withThreadLock = (threadId: ThreadId, effect: Effect.Effect) => + Effect.flatMap(getThreadLock(threadId), (lock) => lock.withPermit(effect)); + const getThreadGeneration = (threadId: ThreadId) => + Ref.get(threadGenerations).pipe(Effect.map((generations) => generations.get(threadId) ?? 0)); + const advanceThreadGeneration = (threadId: ThreadId) => + Ref.update(threadGenerations, (generations) => { + const next = new Map(generations); + next.set(threadId, (next.get(threadId) ?? 0) + 1); + return next; + }); + + interface PreparedMcpSession { + readonly previous: McpProviderSession.McpProviderSessionConfig | undefined; + readonly current: McpProviderSession.McpProviderSessionConfig | undefined; + } + const prepareMcpSession = Effect.fn("ProviderService.prepareMcpSession")(function* ( + threadId: ThreadId, + providerInstanceId: ProviderInstanceId, + ): Effect.fn.Return { + const previous = McpProviderSession.readMcpProviderSession(threadId); + const credential = yield* McpSessionRegistry.issueUncommittedMcpCredential({ + threadId, + providerInstanceId, + }); + const current = credential?.config; + if (current) McpProviderSession.setMcpProviderSession(current); + return { previous, current }; + }); + const commitMcpSession = Effect.fn("ProviderService.commitMcpSession")(function* ( + prepared: PreparedMcpSession, + ) { + if ( + prepared.previous && + prepared.previous.providerSessionId !== prepared.current?.providerSessionId + ) { + yield* McpSessionRegistry.revokeActiveMcpProviderSession(prepared.previous.providerSessionId); + } + }); + const rollbackMcpSession = Effect.fn("ProviderService.rollbackMcpSession")(function* ( + threadId: ThreadId, + prepared: PreparedMcpSession, + ) { + if (prepared.current) { + yield* McpSessionRegistry.revokeActiveMcpProviderSession(prepared.current.providerSessionId); + } + if (prepared.previous) McpProviderSession.setMcpProviderSession(prepared.previous); + else McpProviderSession.clearMcpProviderSession(threadId); + }); const clearMcpSession = (threadId: ThreadId) => McpSessionRegistry.revokeActiveMcpThread(threadId).pipe( Effect.tap(() => Effect.sync(() => McpProviderSession.clearMcpProviderSession(threadId))), ); + const clearMcpSessionForInstance = Effect.fn("ProviderService.clearMcpSessionForInstance")( + function* (threadId: ThreadId, providerInstanceId: ProviderInstanceId) { + const current = McpProviderSession.readMcpProviderSession(threadId); + if (!current || current.providerInstanceId !== providerInstanceId) return; + yield* McpSessionRegistry.revokeActiveMcpProviderSession(current.providerSessionId); + McpProviderSession.clearMcpProviderSession(threadId); + }, + ); + + const compensateFailedSessionBinding = Effect.fn( + "ProviderService.compensateFailedSessionBinding", + )(function* (input: { + readonly adapter: ProviderAdapterShape; + readonly threadId: ThreadId; + readonly mcp: PreparedMcpSession; + }) { + yield* rollbackMcpSession(input.threadId, input.mcp).pipe( + Effect.catchCause((cause) => + Effect.logWarning("provider.session.binding-compensation-mcp-clear-failed", { + threadId: input.threadId, + provider: input.adapter.provider, + cause, + }), + ), + ); + yield* input.adapter.stopSession(input.threadId).pipe( + Effect.catchCause((cause) => + Effect.logWarning("provider.session.binding-compensation-stop-failed", { + threadId: input.threadId, + provider: input.adapter.provider, + cause, + }), + ), + ); + }); const publishRuntimeEvent = (event: ProviderRuntimeEvent): Effect.Effect => Effect.succeed(event).pipe( @@ -281,6 +368,25 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( }); }); + const upsertSessionBindingIfCurrent = Effect.fn("ProviderService.upsertSessionBindingIfCurrent")( + function* (input: { + readonly threadId: ThreadId; + readonly instanceId: ProviderInstanceId; + readonly generation: number; + readonly binding: ProviderSessionDirectory.ProviderRuntimeBinding; + }) { + yield* withThreadLock( + input.threadId, + Effect.gen(function* () { + const current = Option.getOrUndefined(yield* directory.getBinding(input.threadId)); + if (current?.providerInstanceId !== input.instanceId) return; + if ((yield* getThreadGeneration(input.threadId)) !== input.generation) return; + yield* directory.upsert(input.binding); + }), + ); + }, + ); + const processRuntimeEvent = ( source: { readonly instanceId: ProviderInstanceId; @@ -353,87 +459,113 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( ).pipe(Effect.forkScoped); const recoverSessionForThread = Effect.fn("recoverSessionForThread")(function* (input: { - readonly binding: ProviderSessionDirectory.ProviderRuntimeBinding; + readonly threadId: ThreadId; readonly operation: string; }) { - const bindingInstanceId = yield* requireBindingInstanceId(input.operation, input.binding); - yield* Effect.annotateCurrentSpan({ - "provider.operation": "recover-session", - "provider.kind": input.binding.provider, - "provider.instance_id": bindingInstanceId, - "provider.thread_id": input.binding.threadId, - }); + let metricProvider = "unknown"; return yield* Effect.gen(function* () { + const binding = Option.getOrUndefined(yield* directory.getBinding(input.threadId)); + if (!binding) { + return yield* toValidationError( + input.operation, + `Cannot recover thread '${input.threadId}' because no persisted provider binding exists.`, + ); + } + const bindingInstanceId = yield* requireBindingInstanceId(input.operation, binding); + metricProvider = binding.provider; + yield* Effect.annotateCurrentSpan({ + "provider.operation": "recover-session", + "provider.kind": binding.provider, + "provider.instance_id": bindingInstanceId, + "provider.thread_id": binding.threadId, + }); const adapter = yield* registry.getByInstance(bindingInstanceId); - const hasResumeCursor = - input.binding.resumeCursor !== null && input.binding.resumeCursor !== undefined; - const hasActiveSession = yield* adapter.hasSession(input.binding.threadId); + const hasResumeCursor = binding.resumeCursor !== null && binding.resumeCursor !== undefined; + const hasActiveSession = yield* adapter.hasSession(binding.threadId); if (hasActiveSession) { const activeSessions = yield* adapter.listSessions(); - const existing = activeSessions.find( - (session) => session.threadId === input.binding.threadId, - ); + const existing = activeSessions.find((session) => session.threadId === binding.threadId); if (existing) { - yield* upsertSessionBinding( - { ...existing, providerInstanceId: bindingInstanceId }, - input.binding.threadId, - ); + yield* persistCommittedSession({ + threadId: binding.threadId, + currentInstanceId: bindingInstanceId, + persistence: upsertSessionBinding( + { ...existing, providerInstanceId: bindingInstanceId }, + binding.threadId, + ), + }); yield* analytics.record("provider.session.recovered", { provider: existing.provider, strategy: "adopt-existing", hasResumeCursor: existing.resumeCursor !== undefined, }); - return { adapter, session: existing } as const; + return { adapter, instanceId: bindingInstanceId, session: existing } as const; } } if (!hasResumeCursor) { return yield* toValidationError( input.operation, - `Cannot recover thread '${input.binding.threadId}' because no provider resume state is persisted.`, + `Cannot recover thread '${binding.threadId}' because no provider resume state is persisted.`, ); } - const persistedCwd = readPersistedCwd(input.binding.runtimePayload); - const persistedModelSelection = readPersistedModelSelection(input.binding.runtimePayload); + const persistedCwd = readPersistedCwd(binding.runtimePayload); + const persistedModelSelection = readPersistedModelSelection(binding.runtimePayload); - yield* prepareMcpSession(input.binding.threadId, bindingInstanceId); + const mcp = yield* prepareMcpSession(binding.threadId, bindingInstanceId); const resumed = yield* adapter .startSession({ - threadId: input.binding.threadId, - provider: input.binding.provider, + threadId: binding.threadId, + provider: binding.provider, providerInstanceId: bindingInstanceId, ...(persistedCwd ? { cwd: persistedCwd } : {}), ...(persistedModelSelection ? { modelSelection: persistedModelSelection } : {}), - ...(hasResumeCursor ? { resumeCursor: input.binding.resumeCursor } : {}), - runtimeMode: input.binding.runtimeMode ?? "full-access", + ...(hasResumeCursor ? { resumeCursor: binding.resumeCursor } : {}), + runtimeMode: binding.runtimeMode ?? "full-access", }) - .pipe(Effect.onError(() => clearMcpSession(input.binding.threadId))); + .pipe(Effect.onError(() => rollbackMcpSession(binding.threadId, mcp))); if (resumed.provider !== adapter.provider) { - yield* clearMcpSession(input.binding.threadId); + yield* compensateFailedSessionBinding({ + adapter, + threadId: binding.threadId, + mcp, + }); return yield* toValidationError( input.operation, - `Adapter/provider mismatch while recovering thread '${input.binding.threadId}'. Expected '${adapter.provider}', received '${resumed.provider}'.`, + `Adapter/provider mismatch while recovering thread '${binding.threadId}'. Expected '${adapter.provider}', received '${resumed.provider}'.`, ); } - yield* upsertSessionBinding( - { ...resumed, providerInstanceId: bindingInstanceId }, - input.binding.threadId, - ); + yield* persistCommittedSession({ + threadId: binding.threadId, + currentInstanceId: bindingInstanceId, + mcp, + persistence: upsertSessionBinding( + { ...resumed, providerInstanceId: bindingInstanceId }, + binding.threadId, + ).pipe( + Effect.onError(() => + compensateFailedSessionBinding({ + adapter, + threadId: binding.threadId, + mcp, + }), + ), + ), + }); yield* analytics.record("provider.session.recovered", { provider: resumed.provider, strategy: "resume-thread", hasResumeCursor: resumed.resumeCursor !== undefined, }); - return { adapter, session: resumed } as const; + return { adapter, instanceId: bindingInstanceId, session: resumed } as const; }).pipe( withMetrics({ counter: providerSessionsTotal, - attributes: providerMetricAttributes(input.binding.provider, { - operation: "recover", - }), + attributes: () => providerMetricAttributes(metricProvider, { operation: "recover" }), }), + (effect) => withThreadLock(input.threadId, effect), ); }); @@ -473,12 +605,12 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( } const recovered = yield* recoverSessionForThread({ - binding, + threadId: input.threadId, operation: input.operation, }); return { adapter: recovered.adapter, - instanceId, + instanceId: recovered.instanceId, threadId: input.threadId, isActive: true, } as const; @@ -519,128 +651,197 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( ); }); - const startSession: ProviderServiceMethod<"startSession"> = Effect.fn("startSession")( - function* (threadId, rawInput) { - const parsed = yield* decodeInputOrValidationError({ - operation: "ProviderService.startSession", - schema: ProviderSessionStartInput, - payload: rawInput, - }); + const finalizeCommittedSession = Effect.fn("ProviderService.finalizeCommittedSession")( + function* (input: { + readonly threadId: ThreadId; + readonly currentInstanceId: ProviderInstanceId; + readonly mcp?: PreparedMcpSession; + }) { + if (input.mcp) yield* commitMcpSession(input.mcp); + yield* stopStaleSessionsForThread(input); + }, + Effect.uninterruptible, + ); + const persistCommittedSession = (input: { + readonly threadId: ThreadId; + readonly currentInstanceId: ProviderInstanceId; + readonly persistence: Effect.Effect; + readonly mcp?: PreparedMcpSession; + }) => + Effect.uninterruptibleMask((restore) => + restore(input.persistence).pipe( + Effect.andThen(advanceThreadGeneration(input.threadId)), + Effect.andThen(finalizeCommittedSession(input)), + ), + ); - const resolvedInstanceId = yield* requireBindingInstanceId( - "ProviderService.startSession", - parsed, - ); - let metricProvider = parsed.provider ?? String(resolvedInstanceId); + const startSession: ProviderServiceMethod<"startSession"> = Effect.fn("startSession")(function* ( + threadId, + rawInput, + options = { activeSession: "replace" as const }, + ) { + const parsed = yield* decodeInputOrValidationError({ + operation: "ProviderService.startSession", + schema: ProviderSessionStartInput, + payload: rawInput, + }); + + const resolvedInstanceId = yield* requireBindingInstanceId( + "ProviderService.startSession", + parsed, + ); + let metricProvider = parsed.provider ?? String(resolvedInstanceId); + yield* Effect.annotateCurrentSpan({ + "provider.operation": "start-session", + "provider.instance_id": resolvedInstanceId, + "provider.thread_id": threadId, + "provider.runtime_mode": parsed.runtimeMode, + }); + return yield* Effect.gen(function* () { + const instanceInfo = yield* registry.getInstanceInfo(resolvedInstanceId); + const resolvedProvider = instanceInfo.driverKind; + metricProvider = resolvedProvider; + if (parsed.provider !== undefined && parsed.provider !== resolvedProvider) { + return yield* toValidationError( + "ProviderService.startSession", + `Provider instance '${resolvedInstanceId}' belongs to driver '${resolvedProvider}', not '${parsed.provider}'.`, + ); + } + const input = { + ...parsed, + threadId, + provider: resolvedProvider, + }; + if (!instanceInfo.enabled) { + return yield* toValidationError( + "ProviderService.startSession", + `Provider instance '${resolvedInstanceId}' is disabled in T3 Code settings.`, + ); + } + const persistedBinding = Option.getOrUndefined(yield* directory.getBinding(threadId)); + const effectiveResumeCursor = + input.resumeCursor ?? + (persistedBinding?.providerInstanceId === resolvedInstanceId + ? persistedBinding.resumeCursor + : undefined); + const effectiveCwd = + input.cwd ?? + (persistedBinding?.providerInstanceId === resolvedInstanceId + ? readPersistedCwd(persistedBinding.runtimePayload) + : undefined); yield* Effect.annotateCurrentSpan({ - "provider.operation": "start-session", - "provider.instance_id": resolvedInstanceId, - "provider.thread_id": threadId, - "provider.runtime_mode": parsed.runtimeMode, + "provider.kind": resolvedProvider, + "provider.resume_cursor.source": + input.resumeCursor !== undefined + ? "request" + : effectiveResumeCursor !== undefined && + persistedBinding?.providerInstanceId === resolvedInstanceId + ? "persisted" + : "none", + "provider.resume_cursor.present": effectiveResumeCursor !== undefined, + "provider.cwd.source": + input.cwd !== undefined + ? "request" + : effectiveCwd !== undefined && + persistedBinding?.providerInstanceId === resolvedInstanceId + ? "persisted" + : "none", + "provider.cwd.effective": effectiveCwd ?? "", }); - return yield* Effect.gen(function* () { - const instanceInfo = yield* registry.getInstanceInfo(resolvedInstanceId); - const resolvedProvider = instanceInfo.driverKind; - metricProvider = resolvedProvider; - if (parsed.provider !== undefined && parsed.provider !== resolvedProvider) { - return yield* toValidationError( - "ProviderService.startSession", - `Provider instance '${resolvedInstanceId}' belongs to driver '${resolvedProvider}', not '${parsed.provider}'.`, - ); - } - const input = { - ...parsed, - threadId, - provider: resolvedProvider, - }; - if (!instanceInfo.enabled) { - return yield* toValidationError( - "ProviderService.startSession", - `Provider instance '${resolvedInstanceId}' is disabled in T3 Code settings.`, - ); - } - const persistedBinding = Option.getOrUndefined(yield* directory.getBinding(threadId)); - const effectiveResumeCursor = - input.resumeCursor ?? - (persistedBinding?.providerInstanceId === resolvedInstanceId - ? persistedBinding.resumeCursor - : undefined); - const effectiveCwd = - input.cwd ?? - (persistedBinding?.providerInstanceId === resolvedInstanceId - ? readPersistedCwd(persistedBinding.runtimePayload) - : undefined); - yield* Effect.annotateCurrentSpan({ - "provider.kind": resolvedProvider, - "provider.resume_cursor.source": - input.resumeCursor !== undefined - ? "request" - : effectiveResumeCursor !== undefined && - persistedBinding?.providerInstanceId === resolvedInstanceId - ? "persisted" - : "none", - "provider.resume_cursor.present": effectiveResumeCursor !== undefined, - "provider.cwd.source": - input.cwd !== undefined - ? "request" - : effectiveCwd !== undefined && - persistedBinding?.providerInstanceId === resolvedInstanceId - ? "persisted" - : "none", - "provider.cwd.effective": effectiveCwd ?? "", - }); - const adapter = yield* registry.getByInstance(resolvedInstanceId); - yield* prepareMcpSession(threadId, resolvedInstanceId); - const session = yield* adapter - .startSession({ - ...input, + const adapter = yield* registry.getByInstance(resolvedInstanceId); + if (yield* adapter.hasSession(threadId)) { + const existing = (yield* adapter.listSessions()).find( + (session) => session.threadId === threadId, + ); + const currentMcpSession = McpProviderSession.readMcpProviderSession(threadId); + const canReuse = + existing !== undefined && + options.activeSession === "reuse" && + persistedBinding?.providerInstanceId === resolvedInstanceId && + (currentMcpSession === undefined || + currentMcpSession.providerInstanceId === resolvedInstanceId); + if (existing && canReuse) { + const existingWithInstance = { + ...existing, providerInstanceId: resolvedInstanceId, - ...(effectiveCwd !== undefined ? { cwd: effectiveCwd } : {}), - ...(effectiveResumeCursor !== undefined ? { resumeCursor: effectiveResumeCursor } : {}), - }) - .pipe(Effect.onError(() => clearMcpSession(threadId))); - - if (session.provider !== adapter.provider) { - yield* clearMcpSession(threadId); - return yield* toValidationError( - "ProviderService.startSession", - `Adapter/provider mismatch: requested '${adapter.provider}', received '${session.provider}'.`, - ); + }; + yield* persistCommittedSession({ + threadId, + currentInstanceId: resolvedInstanceId, + persistence: upsertSessionBinding(existingWithInstance, threadId, { + modelSelection: input.modelSelection, + }), + }); + return existingWithInstance; } - const sessionWithInstance = { - ...session, + yield* advanceThreadGeneration(threadId); + yield* adapter.stopSession(threadId); + yield* clearMcpSessionForInstance(threadId, resolvedInstanceId); + } + const mcp = yield* prepareMcpSession(threadId, resolvedInstanceId); + const session = yield* adapter + .startSession({ + ...input, providerInstanceId: resolvedInstanceId, - }; + ...(effectiveCwd !== undefined ? { cwd: effectiveCwd } : {}), + ...(effectiveResumeCursor !== undefined ? { resumeCursor: effectiveResumeCursor } : {}), + }) + .pipe(Effect.onError(() => rollbackMcpSession(threadId, mcp))); - yield* stopStaleSessionsForThread({ + if (session.provider !== adapter.provider) { + yield* compensateFailedSessionBinding({ + adapter, threadId, - currentInstanceId: resolvedInstanceId, - }); - yield* upsertSessionBinding(sessionWithInstance, threadId, { - modelSelection: input.modelSelection, - }); - yield* analytics.record("provider.session.started", { - provider: sessionWithInstance.provider, - runtimeMode: input.runtimeMode, - hasResumeCursor: sessionWithInstance.resumeCursor !== undefined, - hasCwd: typeof effectiveCwd === "string" && effectiveCwd.trim().length > 0, - hasModel: - typeof input.modelSelection?.model === "string" && - input.modelSelection.model.trim().length > 0, + mcp, }); + return yield* toValidationError( + "ProviderService.startSession", + `Adapter/provider mismatch: requested '${adapter.provider}', received '${session.provider}'.`, + ); + } + const sessionWithInstance = { + ...session, + providerInstanceId: resolvedInstanceId, + }; - return sessionWithInstance; - }).pipe( - withMetrics({ - counter: providerSessionsTotal, - attributes: () => - providerMetricAttributes(metricProvider, { - operation: "start", + yield* persistCommittedSession({ + threadId, + currentInstanceId: resolvedInstanceId, + mcp, + persistence: upsertSessionBinding(sessionWithInstance, threadId, { + modelSelection: input.modelSelection, + }).pipe( + Effect.onError(() => + compensateFailedSessionBinding({ + adapter, + threadId, + mcp, }), - }), - ); - }, - ); + ), + ), + }); + yield* analytics.record("provider.session.started", { + provider: sessionWithInstance.provider, + runtimeMode: input.runtimeMode, + hasResumeCursor: sessionWithInstance.resumeCursor !== undefined, + hasCwd: typeof effectiveCwd === "string" && effectiveCwd.trim().length > 0, + hasModel: + typeof input.modelSelection?.model === "string" && + input.modelSelection.model.trim().length > 0, + }); + + return sessionWithInstance; + }).pipe( + withMetrics({ + counter: providerSessionsTotal, + attributes: () => + providerMetricAttributes(metricProvider, { + operation: "start", + }), + }), + (effect) => withThreadLock(threadId, effect), + ); + }); const sendTurn: ProviderServiceMethod<"sendTurn"> = Effect.fn("sendTurn")(function* (rawInput) { const parsed = yield* decodeInputOrValidationError({ @@ -673,6 +874,7 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( operation: "ProviderService.sendTurn", allowRecovery: true, }); + const generation = yield* withThreadLock(input.threadId, getThreadGeneration(input.threadId)); metricProvider = routed.adapter.provider; metricModel = input.modelSelection?.model; yield* Effect.annotateCurrentSpan({ @@ -680,17 +882,22 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( ...(input.modelSelection?.model ? { "provider.model": input.modelSelection.model } : {}), }); const turn = yield* routed.adapter.sendTurn(input); - yield* directory.upsert({ + yield* upsertSessionBindingIfCurrent({ threadId: input.threadId, - provider: routed.adapter.provider, - providerInstanceId: routed.instanceId, - status: "running", - ...(turn.resumeCursor !== undefined ? { resumeCursor: turn.resumeCursor } : {}), - runtimePayload: { - ...(input.modelSelection !== undefined ? { modelSelection: input.modelSelection } : {}), - activeTurnId: turn.turnId, - lastRuntimeEvent: "provider.sendTurn", - lastRuntimeEventAt: yield* nowIso, + instanceId: routed.instanceId, + generation, + binding: { + threadId: input.threadId, + provider: routed.adapter.provider, + providerInstanceId: routed.instanceId, + status: "running", + ...(turn.resumeCursor !== undefined ? { resumeCursor: turn.resumeCursor } : {}), + runtimePayload: { + ...(input.modelSelection !== undefined ? { modelSelection: input.modelSelection } : {}), + activeTurnId: turn.turnId, + lastRuntimeEvent: "provider.sendTurn", + lastRuntimeEventAt: yield* nowIso, + }, }, }); yield* analytics.record("provider.turn.sent", { @@ -846,19 +1053,24 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( "provider.kind": routed.adapter.provider, "provider.thread_id": input.threadId, }); + yield* advanceThreadGeneration(input.threadId); if (routed.isActive) { yield* routed.adapter.stopSession(routed.threadId); } yield* clearMcpSession(input.threadId); - yield* directory.upsert({ - threadId: input.threadId, - provider: routed.adapter.provider, - providerInstanceId: routed.instanceId, - status: "stopped", - runtimePayload: { - activeTurnId: null, - }, - }); + yield* Effect.uninterruptible( + directory + .upsert({ + threadId: input.threadId, + provider: routed.adapter.provider, + providerInstanceId: routed.instanceId, + status: "stopped", + runtimePayload: { + activeTurnId: null, + }, + }) + .pipe(Effect.asVoid), + ); yield* analytics.record("provider.session.stopped", { provider: routed.adapter.provider, }); @@ -870,6 +1082,7 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( operation: "stop", }), }), + (effect) => withThreadLock(input.threadId, effect), ); }, ); diff --git a/apps/server/src/provider/Services/ProviderService.ts b/apps/server/src/provider/Services/ProviderService.ts index 4d4cb4fa01a..2559a8ab41b 100644 --- a/apps/server/src/provider/Services/ProviderService.ts +++ b/apps/server/src/provider/Services/ProviderService.ts @@ -42,6 +42,7 @@ export interface ProviderServiceShape { readonly startSession: ( threadId: ThreadId, input: ProviderSessionStartInput, + options?: { readonly activeSession: "reuse" | "replace" }, ) => Effect.Effect; /** From fd0b53d791523f9941e9205a452ed2d5c0bb6ee9 Mon Sep 17 00:00:00 2001 From: Igor Bedesqui Date: Mon, 20 Jul 2026 14:10:35 -0300 Subject: [PATCH 3/4] fix(provider): list only committed sessions --- .../provider/Layers/ProviderService.test.ts | 82 +++++++++++++++- .../src/provider/Layers/ProviderService.ts | 94 ++++++++----------- 2 files changed, 118 insertions(+), 58 deletions(-) diff --git a/apps/server/src/provider/Layers/ProviderService.test.ts b/apps/server/src/provider/Layers/ProviderService.test.ts index 8166d7d47de..0935991f2b2 100644 --- a/apps/server/src/provider/Layers/ProviderService.test.ts +++ b/apps/server/src/provider/Layers/ProviderService.test.ts @@ -598,6 +598,7 @@ it.effect("ProviderServiceLive serializes competing starts through persistence a const second = yield* Effect.forkChild(start()); yield* Effect.yieldNow; assert.equal(codex.startSession.mock.calls.length, 1); + assert.deepEqual(yield* provider.listSessions(), []); yield* Deferred.succeed(releaseFirstUpsert, undefined); const firstExit = yield* Fiber.await(first); @@ -617,6 +618,52 @@ it.effect("ProviderServiceLive serializes competing starts through persistence a }), ); +it.effect("ProviderServiceLive does not expose a same-instance replacement before commit", () => + Effect.gen(function* () { + const threadId = asThreadId("thread-list-same-instance-replacement"); + const upsertStarted = yield* Deferred.make(); + const releaseUpsert = yield* Deferred.make(); + const harness = makeBindingFailureHarness([ + { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: "full-access", + resumeCursor: { opaque: "old-session" }, + }, + ]); + yield* harness.codex.adapter.startSession(makeStartInput(threadId)); + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + harness.upsert.mockImplementation(() => + Deferred.succeed(upsertStarted, undefined).pipe( + Effect.andThen(Deferred.await(releaseUpsert)), + Effect.andThen(Effect.fail(harness.persistenceFailure)), + ), + ); + const replacement = yield* Effect.forkChild( + provider.startSession(threadId, makeStartInput(threadId), { activeSession: "replace" }), + ); + yield* Deferred.await(upsertStarted); + const listingCompleted = yield* Deferred.make(); + const listing = yield* Effect.forkChild( + provider + .listSessions() + .pipe(Effect.tap(() => Deferred.succeed(listingCompleted, undefined))), + ); + yield* Effect.yieldNow; + const listingCompletedBeforeCommit = yield* Deferred.isDone(listingCompleted); + + yield* Deferred.succeed(releaseUpsert, undefined); + const replacementExit = yield* Fiber.await(replacement); + assert.equal(Exit.isFailure(replacementExit), true); + assert.equal(listingCompletedBeforeCommit, false); + assert.deepEqual(yield* Fiber.join(listing), []); + }).pipe(Effect.provide(harness.layer)); + }), +); + it.effect("ProviderServiceLive rejects stale updates after same-instance replacement", () => Effect.gen(function* () { const replacementUpsertStarted = yield* Deferred.make(); @@ -1268,6 +1315,26 @@ it.effect("ProviderServiceLive commits the new MCP credential after successful r }).pipe(Effect.provide(mcpTestRegistryLayer)), ); +it.effect("ProviderServiceLive fails listSessions when committed bindings cannot be read", () => + Effect.gen(function* () { + const readFailure = new ProviderSessionDirectoryPersistenceError({ + operation: "listBindings", + detail: "injected binding read failure", + }); + const harness = makeBindingFailureHarness([], readFailure); + + const exit = yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + return yield* Effect.exit(provider.listSessions()); + }).pipe(Effect.provide(harness.layer)); + + assert.equal(Exit.isFailure(exit), true); + if (Exit.isFailure(exit)) { + assert.equal(Cause.squash(exit.cause), readFailure); + } + }), +); + it.effect("ProviderServiceLive catches stopAll failures during shutdown", () => Effect.gen(function* () { const codex = makeFakeCodexAdapter(); @@ -2002,7 +2069,7 @@ routing.layer("ProviderServiceLive routing", (it) => { }), ); - it.effect("dies when an active session conflicts with its persisted binding", () => + it.effect("lists only active sessions owned by their persisted binding", () => Effect.gen(function* () { const provider = yield* ProviderService.ProviderService; const directory = yield* ProviderSessionDirectory.ProviderSessionDirectory; @@ -2016,6 +2083,13 @@ routing.layer("ProviderServiceLive routing", (it) => { runtimeMode: "full-access", }); + assert.deepEqual( + (yield* provider.listSessions()) + .filter((session) => session.threadId === threadId) + .map((session) => session.providerInstanceId), + [codexInstanceId], + ); + yield* directory.upsert({ threadId, provider: ProviderDriverKind.make("claudeAgent"), @@ -2023,8 +2097,10 @@ routing.layer("ProviderServiceLive routing", (it) => { runtimeMode: "full-access", }); - const exit = yield* Effect.exit(provider.listSessions()); - assert.equal(Exit.hasDies(exit), true); + assert.deepEqual( + (yield* provider.listSessions()).filter((session) => session.threadId === threadId), + [], + ); yield* directory.upsert({ threadId, provider: ProviderDriverKind.make("codex"), diff --git a/apps/server/src/provider/Layers/ProviderService.ts b/apps/server/src/provider/Layers/ProviderService.ts index 53c6a05d027..c012453fb3d 100644 --- a/apps/server/src/provider/Layers/ProviderService.ts +++ b/apps/server/src/provider/Layers/ProviderService.ts @@ -1101,74 +1101,58 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( ), ); const activeSessions = sessionsByProvider.flatMap((sessions) => sessions); - const persistedBindings = yield* directory.listThreadIds().pipe( - Effect.flatMap((threadIds) => - Effect.forEach( - threadIds, - (threadId) => - directory - .getBinding(threadId) - .pipe( - Effect.orElseSucceed(() => - Option.none(), - ), - ), - { concurrency: "unbounded" }, - ), - ), - Effect.orElseSucceed( - () => [] as Array>, - ), - ); + const persistedBindings = yield* directory.listBindings().pipe(Effect.orDie); const bindingsByThreadId = new Map< ThreadId, ProviderSessionDirectory.ProviderRuntimeBinding >(); - for (const bindingOption of persistedBindings) { - const binding = Option.getOrUndefined(bindingOption); - if (binding) { - bindingsByThreadId.set(binding.threadId, binding); - } + for (const binding of persistedBindings) { + bindingsByThreadId.set(binding.threadId, binding); } const sessions: ProviderSession[] = []; - for (const session of activeSessions) { - const binding = bindingsByThreadId.get(session.threadId); + const activeSessionsByThread = Map.groupBy(activeSessions, (session) => session.threadId); + for (const [threadId] of activeSessionsByThread) { + const binding = bindingsByThreadId.get(threadId); if (!binding) { - sessions.push(session); continue; } - const overrides: { - resumeCursor?: ProviderSession["resumeCursor"]; - runtimeMode?: ProviderSession["runtimeMode"]; - providerInstanceId?: ProviderSession["providerInstanceId"]; - } = {}; - overrides.providerInstanceId = dieOnMissingBindingInstanceId( - "ProviderService.listSessions", - binding, + const refreshed = yield* withThreadLock( + threadId, + Effect.gen(function* () { + const currentBinding = Option.getOrUndefined(yield* directory.getBinding(threadId)); + if (!currentBinding) return undefined; + const bindingInstanceId = dieOnMissingBindingInstanceId( + "ProviderService.listSessions", + currentBinding, + ); + const currentAdapter = yield* registry.getByInstance(bindingInstanceId); + const committedSessions = (yield* currentAdapter.listSessions()) + .filter( + (session) => + session.threadId === threadId && session.provider === currentBinding.provider, + ) + .map((session) => ({ ...session, providerInstanceId: bindingInstanceId })); + return { binding: currentBinding, instanceId: bindingInstanceId, committedSessions }; + }).pipe(Effect.orDie), ); - if (binding.provider !== session.provider) { - return yield* Effect.die( - new Error( - `ProviderService.listSessions: thread '${session.threadId}' is active on provider '${session.provider}' but persisted binding names provider '${binding.provider}'.`, - ), - ); - } - if (overrides.providerInstanceId !== session.providerInstanceId) { - return yield* Effect.die( - new Error( - `ProviderService.listSessions: thread '${session.threadId}' is active on provider instance '${session.providerInstanceId}' but persisted binding names '${overrides.providerInstanceId}'.`, - ), - ); - } - if (session.resumeCursor === undefined && binding.resumeCursor !== undefined) { - overrides.resumeCursor = binding.resumeCursor; - } - if (binding.runtimeMode !== undefined) { - overrides.runtimeMode = binding.runtimeMode; + if (!refreshed) continue; + + for (const session of refreshed.committedSessions) { + const overrides: { + resumeCursor?: ProviderSession["resumeCursor"]; + runtimeMode?: ProviderSession["runtimeMode"]; + providerInstanceId?: ProviderSession["providerInstanceId"]; + } = { providerInstanceId: refreshed.instanceId }; + if (session.resumeCursor === undefined && refreshed.binding.resumeCursor !== undefined) { + overrides.resumeCursor = refreshed.binding.resumeCursor; + } + if (refreshed.binding.runtimeMode !== undefined) { + overrides.runtimeMode = refreshed.binding.runtimeMode; + } + sessions.push(Object.assign({}, session, overrides)); } - sessions.push(Object.assign({}, session, overrides)); } return sessions; }, From 91f004ef8247743a288328c595804d129e1e237d Mon Sep 17 00:00:00 2001 From: Igor Bedesqui Date: Mon, 20 Jul 2026 15:19:33 -0300 Subject: [PATCH 4/4] fix(provider): clear stale mcp config without registry --- .../provider/Layers/ProviderService.test.ts | 45 +++++++++++++++++++ .../src/provider/Layers/ProviderService.ts | 1 + 2 files changed, 46 insertions(+) diff --git a/apps/server/src/provider/Layers/ProviderService.test.ts b/apps/server/src/provider/Layers/ProviderService.test.ts index 0935991f2b2..09f0adf6c05 100644 --- a/apps/server/src/provider/Layers/ProviderService.test.ts +++ b/apps/server/src/provider/Layers/ProviderService.test.ts @@ -1009,6 +1009,51 @@ it.effect("ProviderServiceLive compensates a resumed session when binding fails" }), ); +it.effect( + "ProviderServiceLive clears stale MCP state when credential issuance is unavailable", + () => + Effect.gen(function* () { + const threadId = asThreadId("thread-mcp-registry-unavailable"); + const codex = makeFakeCodexAdapter(); + let observedMcpSession: McpProviderSession.McpProviderSessionConfig | undefined; + const startSession = codex.startSession.getMockImplementation()!; + codex.startSession.mockImplementation((input) => + Effect.sync(() => { + observedMcpSession = McpProviderSession.readMcpProviderSession(threadId); + }).pipe(Effect.andThen(startSession(input))), + ); + const directory = ProviderSessionDirectory.ProviderSessionDirectory.of({ + upsert: () => Effect.void, + getProvider: () => Effect.die("unused test directory method"), + getBinding: () => Effect.succeed(Option.none()), + listThreadIds: () => Effect.succeed([]), + listBindings: () => Effect.succeed([]), + }); + const registry = makeAdapterRegistryMock({ [CODEX_DRIVER]: codex.adapter }); + const layer = makeServiceLayer(registry, directory); + const previousMcpSession = { + environmentId: EnvironmentId.make("environment-test"), + threadId, + providerSessionId: "provider-session-stale", + providerInstanceId: codexInstanceId, + endpoint: "http://127.0.0.1/mcp", + authorizationHeader: "Bearer stale", + }; + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + McpProviderSession.setMcpProviderSession(previousMcpSession); + yield* provider.startSession(threadId, makeStartInput(threadId)); + + assert.equal(observedMcpSession, undefined); + assert.equal(McpProviderSession.readMcpProviderSession(threadId), undefined); + }).pipe( + Effect.provide(layer), + Effect.ensuring(Effect.sync(() => McpProviderSession.clearMcpProviderSession(threadId))), + ); + }), +); + it.effect("ProviderServiceLive compensates a started session with a mismatched provider", () => Effect.gen(function* () { const registry = yield* McpSessionRegistry.McpSessionRegistry; diff --git a/apps/server/src/provider/Layers/ProviderService.ts b/apps/server/src/provider/Layers/ProviderService.ts index c012453fb3d..e2c6fd08741 100644 --- a/apps/server/src/provider/Layers/ProviderService.ts +++ b/apps/server/src/provider/Layers/ProviderService.ts @@ -252,6 +252,7 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( }); const current = credential?.config; if (current) McpProviderSession.setMcpProviderSession(current); + else McpProviderSession.clearMcpProviderSession(threadId); return { previous, current }; }); const commitMcpSession = Effect.fn("ProviderService.commitMcpSession")(function* (