fix(server): T3 tools work again after a thread switches back to a reused provider session

apps/server/src/orchestration-v2/ProviderSessionManager.test.ts

@@ -481,6 +481,7 @@ function layerTest(input: {
     readonly armed: Ref.Ref<boolean>;
     readonly paused: Deferred.Deferred<void>;
   };
+  readonly afterMcpSet?: Ref.Ref<Effect.Effect<void>>;
   readonly hasPendingBackgroundWork?: Effect.Effect<boolean>;
   readonly hasPendingBackgroundWorkForThread?: Effect.Effect<boolean>;
   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,6 +2298,109 @@ it.effect(
     }),
 );
 
+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.Effect<void>>(Effect.void);
+      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 })],
+        });
+        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.
+        yield* registry.revokeThread(threadId);
+        const otherProvider = yield* registry.issue({
+          threadId,
+          providerInstanceId: ProviderInstanceId.make("claudeAgent"),
+        });
+        yield* mcpSessions.set(otherProvider.config);
+
+        if (interrupted) {
+          const published = yield* Deferred.make<void>();
+          const resume = yield* Deferred.make<void>();
+          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, 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);
+        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, afterMcpSet })));
+    }),
+);
+
 it.effect(
   "ProviderSessionManagerV2 protects a reused credential from a predecessor release during open",
   () =>

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<PreparedMcpCredential> =>
         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
@@ -1351,22 +1361,47 @@ export const layerWithOptions = (
               Effect.tap((runtime) => Effect.sync(() => (attachedTo = runtime))),
               Effect.uninterruptible,
             );
-            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;
+            // 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,
+              (prepared) => {
+                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 loadedProviderThreadKeyByThread = new Map(
+                  entry.loadedProviderThreadKeyByThread,
+                );
+                if (reclaimsCredential) loadedProviderThreadKeyByThread.delete(input.threadId);
+                const updated = new Map(current);
+                updated.set(key, {
+                  ...entry,
+                  mcpCredentialIdByThread,
+                  loadedProviderThreadKeyByThread,
                 });
-              }
+                return updated;
+              });
+            }
+            if (attached !== undefined) {
               const entry = (yield* Ref.get(sessions)).get(sessionKey(input.providerSessionId));
               if (entry !== undefined) {
                 yield* withActivityError(
@@ -1381,37 +1416,47 @@ 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 &&
+                              mcpCredentialId !== undefined &&
+                              other.mcpCredentialIdByThread.get(input.threadId) === mcpCredentialId,
+                          );
+                        return held
+                          ? Effect.void
+                          : clearMcpSession(input.threadId, mcpCredentialId);
                       }),
-                    ),
-                  ),
-            ),
+                    );
+                  }),
+                ),
+              );
+            }),
           );
           return threadAttachment
             .withLock(threadAttachmentKey(input), attach)

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)),
   );