From 378e40f919af15d1ec5f3b3bf99a2b7c8dc7a9a4 Mon Sep 17 00:00:00 2001 From: jieyuexing Date: Thu, 8 Oct 2026 14:03:21 +0800 Subject: [PATCH 1/3] fix(server): T3 tools work again after a thread switches back to a reused provider session A thread keeps its old provider session attached after switching providers, and the new provider's attach rotates the thread's single MCP credential slot to its own scope. When the thread switched back, the reused session skipped credential preparation, so the returning provider ran with a credential scoped to the other provider and every live-caller check failed with parent_not_active (#16565). ensureThreadAttached now also prepares the credential when the thread is already attached but the slot belongs to another provider instance, and records the new credential on the session so release revokes the right one. Co-Authored-By: Claude Opus 5.5 (1M context) Rebased onto current main (history rewritten since the PR opened): the slot read now goes through the McpProviderSessions service, which replaced the readMcpProviderSession/setMcpProviderSession helpers. Co-Authored-By: Claude Opus 5.5 --- .../ProviderSessionManager.test.ts | 56 +++++++++++++++++++ .../ProviderSessionManager.ts | 42 +++++++++----- 2 files changed, 83 insertions(+), 15 deletions(-) diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts index 3c635be6753df..391e9d911a0b5 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts @@ -2107,6 +2107,62 @@ it.effect( }), ); +it.effect( + "ProviderSessionManagerV2 takes the thread's MCP credential back when a run returns to a still-attached session", + () => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const effect = Effect.gen(function* () { + const eventSink = yield* EventSink.EventSinkV2; + const idAllocator = yield* IdAllocator.IdAllocatorV2; + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const registry = yield* McpSessionRegistry.McpSessionRegistry; + const mcpSessions = yield* McpProviderSessions.McpProviderSessions; + const now = yield* DateTime.now; + const threadId = ThreadId.make("thread-provider-session-manager-returning-provider"); + const providerSessionId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + const tokenOf = (config: McpProviderSession.McpProviderSessionConfig | undefined) => + config?.authorizationHeader.replace(/^Bearer\s+/, "") ?? ""; + + yield* eventSink.write({ + events: [yield* makeThreadCreatedEvent({ idAllocator, threadId, now })], + }); + yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); + + // The thread switches to another provider, whose attach replaces the + // thread's credential while this session stays attached. + yield* registry.revokeThread(threadId); + const otherProvider = yield* registry.issue({ + threadId, + providerInstanceId: ProviderInstanceId.make("claudeAgent"), + }); + yield* mcpSessions.set(otherProvider.config); + + // The thread switches back and its next turn reuses this session. + yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); + assert.equal((yield* Ref.get(state)).openCount, 1); + + const reclaimed = yield* mcpSessions.read(threadId); + assert.equal(reclaimed?.providerInstanceId, modelSelection.instanceId); + assert.equal( + (yield* registry.resolve(tokenOf(reclaimed)))?.thread.providerInstanceId, + modelSelection.instanceId, + ); + assert.isUndefined(yield* registry.resolve(tokenOf(otherProvider.config))); + + // The session's record follows the new credential, so release revokes it. + yield* manager.close(providerSessionId); + assert.isUndefined(yield* registry.resolve(tokenOf(reclaimed))); + assert.isUndefined(yield* mcpSessions.read(threadId)); + }); + + yield* effect.pipe(Effect.provide(layerTest({ state, idleTimeoutMs: 1_000 }))); + }), +); + it.effect( "ProviderSessionManagerV2 protects a reused credential from a predecessor release during open", () => diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.ts index 00b42d510bd91..f83c318626775 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.ts @@ -1328,22 +1328,34 @@ export const layerWithOptions = ( Effect.tap((runtime) => Effect.sync(() => (attachedTo = runtime))), Effect.uninterruptible, ); + // A thread stays attached to the session it ran on before switching + // providers, and the other provider's attach took over the thread's + // credential slot. A run returning to this session must take it + // back, or the process calls T3 tools scoped to the other provider. + const reclaimsCredential = + attached === undefined && + (yield* mcpSessions.read(input.threadId))?.providerInstanceId !== + input.providerInstanceId && + (yield* Ref.get(sessions)) + .get(sessionKey(input.providerSessionId)) + ?.attachedThreadIds.has(input.threadId) === true; + if (attached === undefined && !reclaimsCredential) return; + const prepared = yield* prepareMcpSession(input.threadId, input.providerInstanceId); + preparedForCleanup = prepared; + if (prepared.mcpCredentialId !== undefined) { + const mcpCredentialId = prepared.mcpCredentialId; + yield* Ref.update(sessions, (current) => { + const key = sessionKey(input.providerSessionId); + const entry = current.get(key); + if (entry === undefined) return current; + const mcpCredentialIdByThread = new Map(entry.mcpCredentialIdByThread); + mcpCredentialIdByThread.set(input.threadId, mcpCredentialId); + const updated = new Map(current); + updated.set(key, { ...entry, mcpCredentialIdByThread }); + return updated; + }); + } if (attached !== undefined) { - const prepared = yield* prepareMcpSession(input.threadId, input.providerInstanceId); - preparedForCleanup = prepared; - if (prepared.mcpCredentialId !== undefined) { - const mcpCredentialId = prepared.mcpCredentialId; - yield* Ref.update(sessions, (current) => { - const key = sessionKey(input.providerSessionId); - const entry = current.get(key); - if (entry === undefined) return current; - const mcpCredentialIdByThread = new Map(entry.mcpCredentialIdByThread); - mcpCredentialIdByThread.set(input.threadId, mcpCredentialId); - const updated = new Map(current); - updated.set(key, { ...entry, mcpCredentialIdByThread }); - return updated; - }); - } const entry = (yield* Ref.get(sessions)).get(sessionKey(input.providerSessionId)); if (entry !== undefined) { yield* withActivityError( From 7c69e728e5919801fc83a2076421862ffcdbd001 Mon Sep 17 00:00:00 2001 From: "Abhi (onFire)" Date: Mon, 26 Oct 2026 14:57:39 +0545 Subject: [PATCH 2/3] fix(server): revoke MCP credential issued during interrupted reclaim Handle failed reclaims in ensureThreadAttached separately from failed new attachments: revoke the credential issued by the current call when it has no other holder, instead of skipping cleanup because attachedTo is undefined. Co-Authored-By: Claude Haiku 5.5 --- .../ProviderSessionManager.ts | 72 +++++++++++-------- 1 file changed, 42 insertions(+), 30 deletions(-) diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.ts index f83c318626775..9f86fca43fbdc 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.ts @@ -1370,37 +1370,49 @@ export const layerWithOptions = ( } } }).pipe( - Effect.onError(() => - attachedTo === undefined - ? Effect.void - : removeThreadAttachment({ ...input, runtime: attachedTo }).pipe( - Effect.andThen( - Effect.suspend(() => { - dropReservation(); - // Other live processes hold reused MCP credentials; revocation would break their clients. - if (preparedForCleanup?.issued !== true) return Effect.void; - const mcpCredentialId = preparedForCleanup.mcpCredentialId; - const attachedRuntime = attachedTo; - return Ref.get(sessions).pipe( - Effect.flatMap((current) => - (mcpCredentialId !== undefined && - isMcpCredentialReserved(input.threadId, mcpCredentialId)) || - Array.from(current.values()).some( - (other) => - other.runtime !== attachedRuntime && - (other.attachedThreadIds.has(input.threadId) || - (mcpCredentialId !== undefined && - other.mcpCredentialIdByThread.get(input.threadId) === - mcpCredentialId)), - ) - ? Effect.void - : clearMcpSession(input.threadId, mcpCredentialId), - ), - ); + Effect.onError(() => { + // A reclaim (thread already attached, credential rotated) has + // attachedTo === undefined but still issued a credential that must + // be cleaned up on failure. Only skip when nothing was prepared. + const isReclaim = attachedTo === undefined && preparedForCleanup !== undefined; + if (attachedTo === undefined && !isReclaim) return Effect.void; + const removeAttachment = + attachedTo === undefined + ? Effect.void + : removeThreadAttachment({ ...input, runtime: attachedTo }); + return removeAttachment.pipe( + Effect.andThen( + Effect.suspend(() => { + dropReservation(); + // Other live processes hold reused MCP credentials; revocation would break their clients. + if (preparedForCleanup?.issued !== true) return Effect.void; + const mcpCredentialId = preparedForCleanup.mcpCredentialId; + // For a reclaim, exclude the session being reused (not torn + // down) from the holder check, mirroring new attachments. + return Ref.get(sessions).pipe( + Effect.flatMap((current) => { + const attachedRuntime = + attachedTo ?? current.get(sessionKey(input.providerSessionId))?.runtime; + const held = + (mcpCredentialId !== undefined && + isMcpCredentialReserved(input.threadId, mcpCredentialId)) || + Array.from(current.values()).some( + (other) => + other.runtime !== attachedRuntime && + (other.attachedThreadIds.has(input.threadId) || + (mcpCredentialId !== undefined && + other.mcpCredentialIdByThread.get(input.threadId) === + mcpCredentialId)), + ); + return held + ? Effect.void + : clearMcpSession(input.threadId, mcpCredentialId); }), - ), - ), - ), + ); + }), + ), + ); + }), ); return threadAttachment .withLock(threadAttachmentKey(input), attach) From ce2d234054843def86a5d1ead05aece08fb3cfb7 Mon Sep 17 00:00:00 2001 From: maria-rcks Date: Sun, 1 Nov 2026 02:47:03 +0000 Subject: [PATCH 3/3] fix(server): refresh reused provider MCP clients and reclaim safely --- .../ProviderSessionManager.test.ts | 80 +++++++++++++++++-- .../ProviderSessionManager.ts | 61 +++++++++----- .../provider-pi/src/server/adapter.test.ts | 17 +++- packages/provider-pi/src/server/adapter.ts | 36 ++++++++- .../src/server/mcpBridge.testkit.ts | 4 + .../src/server/mcpExtensionSource.test.ts | 22 +++++ .../src/server/mcpExtensionSource.ts | 17 +++- 7 files changed, 204 insertions(+), 33 deletions(-) diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts index 83f210b836922..c02ac71f3d993 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts @@ -481,6 +481,7 @@ function layerTest(input: { readonly armed: Ref.Ref; readonly paused: Deferred.Deferred; }; + readonly afterMcpSet?: Ref.Ref>; readonly hasPendingBackgroundWork?: Effect.Effect; readonly hasPendingBackgroundWorkForThread?: Effect.Effect; readonly backgroundTasksEndWithProcess?: boolean; @@ -534,6 +535,22 @@ function layerTest(input: { input.pauseResolve === undefined ? layerTestMcpRegistry : layerPausingMcpRegistry(input.pauseResolve); + const layerMcpSessions = + input.afterMcpSet === undefined + ? McpProviderSessions.layer + : Layer.effect( + McpProviderSessions.McpProviderSessions, + Effect.gen(function* () { + const delegate = yield* McpProviderSessions.McpProviderSessions; + return McpProviderSessions.McpProviderSessions.of({ + ...delegate, + set: (config) => + delegate + .set(config) + .pipe(Effect.andThen(Ref.get(input.afterMcpSet!)), Effect.flatten), + }); + }), + ).pipe(Layer.provide(McpProviderSessions.layer)); const layerProviderEventIngestorTest = ProviderEventIngestor.layer.pipe( Layer.provide( Layer.mergeAll( @@ -556,14 +573,14 @@ function layerTest(input: { layerConfiguredEventSink, IdAllocator.layer, layerConfiguredMcpRegistry, - McpProviderSessions.layer, + layerMcpSessions, ProviderSessionManager.layerWithOptions({ idleTimeoutMs: input.idleTimeoutMs, ...(input.maxIdlePinMs === undefined ? {} : { maxIdlePinMs: input.maxIdlePinMs }), }).pipe( Layer.provide( Layer.mergeAll( - layerRegistry.pipe(Layer.provideMerge(McpProviderSessions.layer)), + layerRegistry.pipe(Layer.provideMerge(layerMcpSessions)), CommandOutputHub.layer.pipe( Layer.provide(Layer.merge(layerConfiguredEventSink, layerTestStores)), ), @@ -2281,11 +2298,12 @@ it.effect( }), ); -it.effect( - "ProviderSessionManagerV2 takes the thread's MCP credential back when a run returns to a still-attached session", - () => +it.effect.each([false, true])( + "ProviderSessionManagerV2 takes the thread's MCP credential back when a run returns to a still-attached session, interrupted=%s", + (interrupted) => Effect.gen(function* () { const state = yield* Ref.make(emptyState); + const afterMcpSet = yield* Ref.make>(Effect.void); const effect = Effect.gen(function* () { const eventSink = yield* EventSink.EventSinkV2; const idAllocator = yield* IdAllocator.IdAllocatorV2; @@ -2304,7 +2322,31 @@ it.effect( yield* eventSink.write({ events: [yield* makeThreadCreatedEvent({ idAllocator, threadId, now })], }); - yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); + const runtime = yield* manager.open({ + threadId, + providerSessionId, + modelSelection, + runtimePolicy, + }); + const providerThread = makeProviderThread({ + idAllocator, + threadId, + providerSessionId, + now, + }); + yield* runtime.resumeThread({ providerThread, modelSelection, runtimePolicy }); + yield* runtime.resumeThread({ providerThread, modelSelection, runtimePolicy }); + assert.equal((yield* Ref.get(state)).resumeCount, 1); + const staleSessionId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + yield* manager.open({ + threadId, + providerSessionId: staleSessionId, + modelSelection, + runtimePolicy, + }); // The thread switches to another provider, whose attach replaces the // thread's credential while this session stays attached. @@ -2315,9 +2357,31 @@ it.effect( }); yield* mcpSessions.set(otherProvider.config); + if (interrupted) { + const published = yield* Deferred.make(); + const resume = yield* Deferred.make(); + yield* Ref.set( + afterMcpSet, + Deferred.succeed(published, undefined).pipe(Effect.andThen(Deferred.await(resume))), + ); + const reclaim = yield* manager + .open({ threadId, providerSessionId, modelSelection, runtimePolicy }) + .pipe(Effect.forkChild({ startImmediately: true })); + yield* Deferred.await(published); + const issued = yield* mcpSessions.read(threadId); + reclaim.interruptUnsafe(); + yield* Deferred.succeed(resume, undefined); + yield* Fiber.await(reclaim); + assert.isUndefined(yield* registry.resolve(tokenOf(issued))); + assert.isUndefined(yield* mcpSessions.read(threadId)); + return; + } + // The thread switches back and its next turn reuses this session. yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); - assert.equal((yield* Ref.get(state)).openCount, 1); + assert.equal((yield* Ref.get(state)).openCount, 2); + yield* runtime.resumeThread({ providerThread, modelSelection, runtimePolicy }); + assert.equal((yield* Ref.get(state)).resumeCount, 2); const reclaimed = yield* mcpSessions.read(threadId); assert.equal(reclaimed?.providerInstanceId, modelSelection.instanceId); @@ -2333,7 +2397,7 @@ it.effect( assert.isUndefined(yield* mcpSessions.read(threadId)); }); - yield* effect.pipe(Effect.provide(layerTest({ state, idleTimeoutMs: 1_000 }))); + yield* effect.pipe(Effect.provide(layerTest({ state, idleTimeoutMs: 1_000, afterMcpSet }))); }), ); diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.ts index e357579dd9052..6a117fd0ca6b5 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.ts @@ -441,6 +441,7 @@ export const layerWithOptions = ( const prepareMcpSession = ( threadId: ThreadId, providerInstanceId: ProviderInstanceId, + onPrepared?: (prepared: PreparedMcpCredential) => void, ): Effect.Effect => options.configureMcp === false ? mcpSessions @@ -485,21 +486,31 @@ export const layerWithOptions = ( // A credential issued while the lineage read failed is replaced once it works. (resolved.thread.subagent === true) === subagent ) { - return { mcpCredentialId: existing.providerSessionId, issued: false }; + const prepared = { mcpCredentialId: existing.providerSessionId, issued: false }; + onPrepared?.(prepared); + return prepared; } dropMcpCredentialReservation(threadId, existing.providerSessionId); } - yield* mcpSessionRegistry.revokeThread(threadId); - const credential = yield* mcpSessionRegistry.issue({ - threadId, - providerInstanceId, - browserToolsAvailable, - capabilities, - subagent, + const issue = Effect.gen(function* () { + yield* mcpSessionRegistry.revokeThread(threadId); + const credential = yield* mcpSessionRegistry.issue({ + threadId, + providerInstanceId, + browserToolsAvailable, + capabilities, + subagent, + }); + yield* mcpSessions.set(credential.config); + reserveMcpCredential(threadId, credential.config.providerSessionId); + const prepared = { + mcpCredentialId: credential.config.providerSessionId, + issued: true, + }; + onPrepared?.(prepared); + return prepared; }); - yield* mcpSessions.set(credential.config); - reserveMcpCredential(threadId, credential.config.providerSessionId); - return { mcpCredentialId: credential.config.providerSessionId, issued: true }; + return yield* onPrepared === undefined ? issue : Effect.uninterruptible(issue); }), ); const clearMcpSession = (threadId: ThreadId, mcpCredentialId?: string) => @@ -1087,8 +1098,7 @@ export const layerWithOptions = ( Array.from(current.values()).some( (other) => other !== entry && - (other.attachedThreadIds.has(threadId) || - other.mcpCredentialIdByThread.get(threadId) === mcpCredentialId), + other.mcpCredentialIdByThread.get(threadId) === mcpCredentialId, ); return heldElsewhere ? Effect.void @@ -1363,8 +1373,13 @@ export const layerWithOptions = ( .get(sessionKey(input.providerSessionId)) ?.attachedThreadIds.has(input.threadId) === true; if (attached === undefined && !reclaimsCredential) return; - const prepared = yield* prepareMcpSession(input.threadId, input.providerInstanceId); - preparedForCleanup = prepared; + const prepared = yield* prepareMcpSession( + input.threadId, + input.providerInstanceId, + (prepared) => { + preparedForCleanup = prepared; + }, + ); if (prepared.mcpCredentialId !== undefined) { const mcpCredentialId = prepared.mcpCredentialId; yield* Ref.update(sessions, (current) => { @@ -1373,8 +1388,16 @@ export const layerWithOptions = ( if (entry === undefined) return current; const mcpCredentialIdByThread = new Map(entry.mcpCredentialIdByThread); mcpCredentialIdByThread.set(input.threadId, mcpCredentialId); + const loadedProviderThreadKeyByThread = new Map( + entry.loadedProviderThreadKeyByThread, + ); + if (reclaimsCredential) loadedProviderThreadKeyByThread.delete(input.threadId); const updated = new Map(current); - updated.set(key, { ...entry, mcpCredentialIdByThread }); + updated.set(key, { + ...entry, + mcpCredentialIdByThread, + loadedProviderThreadKeyByThread, + }); return updated; }); } @@ -1422,10 +1445,8 @@ export const layerWithOptions = ( Array.from(current.values()).some( (other) => other.runtime !== attachedRuntime && - (other.attachedThreadIds.has(input.threadId) || - (mcpCredentialId !== undefined && - other.mcpCredentialIdByThread.get(input.threadId) === - mcpCredentialId)), + mcpCredentialId !== undefined && + other.mcpCredentialIdByThread.get(input.threadId) === mcpCredentialId, ); return held ? Effect.void diff --git a/packages/provider-pi/src/server/adapter.test.ts b/packages/provider-pi/src/server/adapter.test.ts index 5c94de3040209..dfada6cbede68 100644 --- a/packages/provider-pi/src/server/adapter.test.ts +++ b/packages/provider-pi/src/server/adapter.test.ts @@ -1574,7 +1574,7 @@ describe("PiAdapterV2", () => { browserToolsAvailable: true, }); const fake = yield* makeFakePi; - yield* openRuntime(fake); + const { runtime } = yield* openRuntime(fake); const spawn = fake.lastSpawn(); assert.isTrue(spawn.args.includes("--extension")); const extensions = spawn.args.flatMap((arg, index) => @@ -1585,6 +1585,21 @@ describe("PiAdapterV2", () => { assert.equal(spawn.env.T3_MCP_URL, "http://127.0.0.1:43123/mcp"); assert.equal(spawn.env.T3_MCP_BEARER_TOKEN, "secret-pi-token"); assert.equal(spawn.env.T3_PI_RUNTIME_MODE, "full-access"); + const credentialPath = spawn.env.T3_PI_MCP_CREDENTIAL_PATH; + assert.isDefined(credentialPath); + const fileSystem = yield* FileSystem.FileSystem; + assert.equal(yield* fileSystem.readFileString(credentialPath!), "Bearer secret-pi-token"); + const mcpSessions = yield* McpProviderSessions.McpProviderSessions; + const config = yield* mcpSessions.read(THREAD_ID); + yield* mcpSessions.set({ ...config!, authorizationHeader: "Bearer reclaimed-pi-token" }); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + yield* startTurn(runtime, providerThread); + yield* fake.takeRequest("prompt"); + assert.equal(yield* fileSystem.readFileString(credentialPath!), "Bearer reclaimed-pi-token"); }).pipe(Effect.scoped, Effect.provide(layerTest)), ); diff --git a/packages/provider-pi/src/server/adapter.ts b/packages/provider-pi/src/server/adapter.ts index a32a5759ba707..9d4d93b02acc2 100644 --- a/packages/provider-pi/src/server/adapter.ts +++ b/packages/provider-pi/src/server/adapter.ts @@ -70,7 +70,7 @@ import { materializePiT3McpExtension, resolvePiLaunchArgs, } from "./mcpInjection.ts"; -import { PI_FILE_CHANGE_TOOLS } from "./mcpExtensionSource.ts"; +import { PI_FILE_CHANGE_TOOLS, T3_PI_MCP_CREDENTIAL_PATH_ENV } from "./mcpExtensionSource.ts"; export const PI_PROVIDER = ProviderDriverKind.make("pi"); const PI_DRIVER_KIND = PI_PROVIDER; @@ -441,7 +441,7 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* ( const scope = yield* Effect.scope; const cwd = input.runtimePolicy.cwd ?? host.paths.cwd; const mcpSession = yield* mcpSessions.read(input.threadId); - const provideCacheFs = (effect: Effect.Effect) => + const provideCacheFs = (effect: Effect.Effect) => effect.pipe( Effect.provideService(FileSystem.FileSystem, fileSystem), Effect.mapError( @@ -468,6 +468,19 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* ( extensionPath, runtimeMode: input.runtimePolicy.runtimeMode, }); + const credentialPath = + mcpSession === undefined + ? undefined + : `${yield* provideCacheFs(fileSystem.makeTempDirectoryScoped({ prefix: "t3-pi-mcp-" }))}/credential`; + let authorizationHeader = mcpSession?.authorizationHeader; + if (credentialPath !== undefined && authorizationHeader !== undefined) { + yield* provideCacheFs( + fileSystem.writeFileString(credentialPath, authorizationHeader, { mode: 0o600 }), + ); + launch.env[T3_PI_MCP_CREDENTIAL_PATH_ENV] = credentialPath; + } else { + delete launch.env[T3_PI_MCP_CREDENTIAL_PATH_ENV]; + } const scopedLaunch = yield* agentScope.wrap({ command: options.settings.binaryPath || "pi", args: launch.args, @@ -2462,6 +2475,25 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* ( ) { return yield* protocolError("Pi turn requested for a different native session"); } + if (credentialPath !== undefined) { + const current = yield* mcpSessions.read(turnInput.threadId); + if (current === undefined || current.providerInstanceId !== options.instanceId) { + return yield* protocolError("Pi turn has no current MCP credential"); + } + if (current.authorizationHeader !== authorizationHeader) { + yield* fileSystem.writeFileString(credentialPath, current.authorizationHeader).pipe( + Effect.mapError( + (cause) => + new ProviderAdapter.ProviderAdapterProtocolError({ + driver: PI_PROVIDER, + cause, + detail: "Could not refresh Pi MCP credential", + }), + ), + ); + authorizationHeader = current.authorizationHeader; + } + } // Fork adoption uses the orchestrator-allocated row; preserve that identity in later updates. const continuation = turnInput.message.createdBy === "agent" && diff --git a/packages/provider-pi/src/server/mcpBridge.testkit.ts b/packages/provider-pi/src/server/mcpBridge.testkit.ts index 133a4b2981ac1..36aa39a6a3f06 100644 --- a/packages/provider-pi/src/server/mcpBridge.testkit.ts +++ b/packages/provider-pi/src/server/mcpBridge.testkit.ts @@ -1,6 +1,7 @@ // @effect-diagnostics nodeBuiltinImport:off -- Executes the shipped Pi extension at its native JavaScript boundary. import * as NodeModule from "node:module"; import * as NodePath from "node:path"; +import * as NodeFSP from "node:fs/promises"; import * as NodeVM from "node:vm"; import { PI_T3_MCP_EXTENSION_SOURCE } from "./mcpExtensionSource.ts"; @@ -48,6 +49,7 @@ export async function loadMcpBridge( readonly allowsTool?: (name: string) => boolean; readonly runtimeMode?: string; readonly inputSchema?: Record; + readonly credentialPath?: string; } = {}, ) { const handlers = new Map(); @@ -90,12 +92,14 @@ export async function loadMcpBridge( env: { T3_MCP_URL: "http://fixture.invalid/mcp", T3_MCP_BEARER_TOKEN: "fixture-token", + T3_PI_MCP_CREDENTIAL_PATH: options.credentialPath, T3_PI_MCP_EXTENSION_PATH: "/fixture/pi-t3-extension.ts", T3_PI_RUNTIME_MODE: options.runtimeMode, }, }, AbortSignal, NodePath, + NodeFSP, Type: { Unsafe: (schema: unknown) => schema }, fetch: async ( url: string, diff --git a/packages/provider-pi/src/server/mcpExtensionSource.test.ts b/packages/provider-pi/src/server/mcpExtensionSource.test.ts index cedf3974f535d..3ed0ba92d2da8 100644 --- a/packages/provider-pi/src/server/mcpExtensionSource.test.ts +++ b/packages/provider-pi/src/server/mcpExtensionSource.test.ts @@ -107,6 +107,28 @@ describe("Pi project trust", () => { }); describe("Pi MCP tool results", () => { + it("refreshes a reused client's credential before the next turn", async () => { + const directory = await NodeFSP.mkdtemp(NodePath.join(NodeOS.tmpdir(), "pi-mcp-reclaim-")); + try { + const credentialPath = NodePath.join(directory, "credential"); + await NodeFSP.writeFile(credentialPath, "Bearer fixture-token"); + const bridge = await loadMcpBridge({ credentialPath }); + const tool = bridge.tools.find((tool) => tool.name === "mcp__t3-code__delegate_task")!; + await tool.execute("first", {}); + assert.equal(bridge.transports.at(-1)?.authorization, "Bearer fixture-token"); + await NodeFSP.writeFile(credentialPath, "Bearer reclaimed-token"); + await bridge.handlers.get("before_agent_start")!( + { systemPrompt: "base" }, + { ui: { notify: assert.fail } }, + ); + await tool.execute("second", {}); + assert.equal(bridge.transports.at(-1)?.authorization, "Bearer reclaimed-token"); + assert.equal(bridge.requests.filter((request) => request.method === "initialize").length, 2); + } finally { + await NodeFSP.rm(directory, { recursive: true, force: true }); + } + }); + it("returns mixed screenshot blocks and typed script output without repeating structured text", async () => { const content = [ { type: "text", text: "Browser screenshot" }, diff --git a/packages/provider-pi/src/server/mcpExtensionSource.ts b/packages/provider-pi/src/server/mcpExtensionSource.ts index ee72c38f75966..3849e1b6573f1 100644 --- a/packages/provider-pi/src/server/mcpExtensionSource.ts +++ b/packages/provider-pi/src/server/mcpExtensionSource.ts @@ -7,6 +7,7 @@ export const T3_MCP_URL_ENV = "T3_MCP_URL"; export const T3_MCP_BEARER_ENV = "T3_MCP_BEARER_TOKEN"; export const T3_PI_RUNTIME_MODE_ENV = "T3_PI_RUNTIME_MODE"; export const T3_PI_MCP_EXTENSION_PATH_ENV = "T3_PI_MCP_EXTENSION_PATH"; +export const T3_PI_MCP_CREDENTIAL_PATH_ENV = "T3_PI_MCP_CREDENTIAL_PATH"; // These tools trigger file-change approvals, bypassed by Auto-accept edits. export const PI_FILE_CHANGE_TOOLS = ["edit", "write"] as const; @@ -21,6 +22,7 @@ const URL_ENV = ${JSON.stringify(T3_MCP_URL_ENV)}; const TOKEN_ENV = ${JSON.stringify(T3_MCP_BEARER_ENV)}; const RUNTIME_MODE_ENV = ${JSON.stringify(T3_PI_RUNTIME_MODE_ENV)}; const EXTENSION_PATH_ENV = ${JSON.stringify(T3_PI_MCP_EXTENSION_PATH_ENV)}; +const CREDENTIAL_PATH_ENV = ${JSON.stringify(T3_PI_MCP_CREDENTIAL_PATH_ENV)}; const ORCHESTRATION_INSTRUCTIONS = ${JSON.stringify(T3_CODE_ORCHESTRATION_INSTRUCTIONS.trim())}; const PROTOCOL = "2025-06-18"; const READ_ONLY_TOOLS = new Set(["read", "grep", "find", "ls"]); @@ -366,7 +368,8 @@ export default async function t3McpExtension(pi: ExtensionAPI) { let deferOptionalTools = supportsExposure; let catalog: ReadonlyArray = []; - const client = createMcpClient(endpoint, token); + let currentToken = token; + let client = createMcpClient(endpoint, currentToken); let started: Promise | undefined; const registerTools = () => { @@ -463,7 +466,17 @@ export default async function t3McpExtension(pi: ExtensionAPI) { }); // Keep guidance in the system prompt so leading slash commands still expand. - pi.on("before_agent_start", (event, ctx) => { + pi.on("before_agent_start", async (event, ctx) => { + const credentialPath = env(CREDENTIAL_PATH_ENV); + if (credentialPath !== undefined) { + const nextToken = (await NodeFSP.readFile(credentialPath, "utf8")).replace(/^Bearer\\s+/, ""); + if (nextToken !== currentToken) { + currentToken = nextToken; + client = createMcpClient(endpoint, currentToken); + started = undefined; + } + await ensureStarted(); + } reconcileDiscovery(ctx); const sections = event.systemPromptOptions?.sections; if (sections && typeof sections === "object" && event.systemPromptOptions?.forceSystemPrompt === undefined) {