Skip to content

Commit ce2d234

Browse files
maria-rcksmaria-rcks
committed
fix(server): refresh reused provider MCP clients and reclaim safely
1 parent 7df56fd commit ce2d234

7 files changed

Lines changed: 204 additions & 33 deletions

File tree

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

Lines changed: 72 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -481,6 +481,7 @@ function layerTest(input: {
481481
readonly armed: Ref.Ref<boolean>;
482482
readonly paused: Deferred.Deferred<void>;
483483
};
484+
readonly afterMcpSet?: Ref.Ref<Effect.Effect<void>>;
484485
readonly hasPendingBackgroundWork?: Effect.Effect<boolean>;
485486
readonly hasPendingBackgroundWorkForThread?: Effect.Effect<boolean>;
486487
readonly backgroundTasksEndWithProcess?: boolean;
@@ -534,6 +535,22 @@ function layerTest(input: {
534535
input.pauseResolve === undefined
535536
? layerTestMcpRegistry
536537
: layerPausingMcpRegistry(input.pauseResolve);
538+
const layerMcpSessions =
539+
input.afterMcpSet === undefined
540+
? McpProviderSessions.layer
541+
: Layer.effect(
542+
McpProviderSessions.McpProviderSessions,
543+
Effect.gen(function* () {
544+
const delegate = yield* McpProviderSessions.McpProviderSessions;
545+
return McpProviderSessions.McpProviderSessions.of({
546+
...delegate,
547+
set: (config) =>
548+
delegate
549+
.set(config)
550+
.pipe(Effect.andThen(Ref.get(input.afterMcpSet!)), Effect.flatten),
551+
});
552+
}),
553+
).pipe(Layer.provide(McpProviderSessions.layer));
537554
const layerProviderEventIngestorTest = ProviderEventIngestor.layer.pipe(
538555
Layer.provide(
539556
Layer.mergeAll(
@@ -556,14 +573,14 @@ function layerTest(input: {
556573
layerConfiguredEventSink,
557574
IdAllocator.layer,
558575
layerConfiguredMcpRegistry,
559-
McpProviderSessions.layer,
576+
layerMcpSessions,
560577
ProviderSessionManager.layerWithOptions({
561578
idleTimeoutMs: input.idleTimeoutMs,
562579
...(input.maxIdlePinMs === undefined ? {} : { maxIdlePinMs: input.maxIdlePinMs }),
563580
}).pipe(
564581
Layer.provide(
565582
Layer.mergeAll(
566-
layerRegistry.pipe(Layer.provideMerge(McpProviderSessions.layer)),
583+
layerRegistry.pipe(Layer.provideMerge(layerMcpSessions)),
567584
CommandOutputHub.layer.pipe(
568585
Layer.provide(Layer.merge(layerConfiguredEventSink, layerTestStores)),
569586
),
@@ -2281,11 +2298,12 @@ it.effect(
22812298
}),
22822299
);
22832300

2284-
it.effect(
2285-
"ProviderSessionManagerV2 takes the thread's MCP credential back when a run returns to a still-attached session",
2286-
() =>
2301+
it.effect.each([false, true])(
2302+
"ProviderSessionManagerV2 takes the thread's MCP credential back when a run returns to a still-attached session, interrupted=%s",
2303+
(interrupted) =>
22872304
Effect.gen(function* () {
22882305
const state = yield* Ref.make(emptyState);
2306+
const afterMcpSet = yield* Ref.make<Effect.Effect<void>>(Effect.void);
22892307
const effect = Effect.gen(function* () {
22902308
const eventSink = yield* EventSink.EventSinkV2;
22912309
const idAllocator = yield* IdAllocator.IdAllocatorV2;
@@ -2304,7 +2322,31 @@ it.effect(
23042322
yield* eventSink.write({
23052323
events: [yield* makeThreadCreatedEvent({ idAllocator, threadId, now })],
23062324
});
2307-
yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy });
2325+
const runtime = yield* manager.open({
2326+
threadId,
2327+
providerSessionId,
2328+
modelSelection,
2329+
runtimePolicy,
2330+
});
2331+
const providerThread = makeProviderThread({
2332+
idAllocator,
2333+
threadId,
2334+
providerSessionId,
2335+
now,
2336+
});
2337+
yield* runtime.resumeThread({ providerThread, modelSelection, runtimePolicy });
2338+
yield* runtime.resumeThread({ providerThread, modelSelection, runtimePolicy });
2339+
assert.equal((yield* Ref.get(state)).resumeCount, 1);
2340+
const staleSessionId = yield* idAllocator.allocate.providerSession({
2341+
providerInstanceId: modelSelection.instanceId,
2342+
threadId,
2343+
});
2344+
yield* manager.open({
2345+
threadId,
2346+
providerSessionId: staleSessionId,
2347+
modelSelection,
2348+
runtimePolicy,
2349+
});
23082350

23092351
// The thread switches to another provider, whose attach replaces the
23102352
// thread's credential while this session stays attached.
@@ -2315,9 +2357,31 @@ it.effect(
23152357
});
23162358
yield* mcpSessions.set(otherProvider.config);
23172359

2360+
if (interrupted) {
2361+
const published = yield* Deferred.make<void>();
2362+
const resume = yield* Deferred.make<void>();
2363+
yield* Ref.set(
2364+
afterMcpSet,
2365+
Deferred.succeed(published, undefined).pipe(Effect.andThen(Deferred.await(resume))),
2366+
);
2367+
const reclaim = yield* manager
2368+
.open({ threadId, providerSessionId, modelSelection, runtimePolicy })
2369+
.pipe(Effect.forkChild({ startImmediately: true }));
2370+
yield* Deferred.await(published);
2371+
const issued = yield* mcpSessions.read(threadId);
2372+
reclaim.interruptUnsafe();
2373+
yield* Deferred.succeed(resume, undefined);
2374+
yield* Fiber.await(reclaim);
2375+
assert.isUndefined(yield* registry.resolve(tokenOf(issued)));
2376+
assert.isUndefined(yield* mcpSessions.read(threadId));
2377+
return;
2378+
}
2379+
23182380
// The thread switches back and its next turn reuses this session.
23192381
yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy });
2320-
assert.equal((yield* Ref.get(state)).openCount, 1);
2382+
assert.equal((yield* Ref.get(state)).openCount, 2);
2383+
yield* runtime.resumeThread({ providerThread, modelSelection, runtimePolicy });
2384+
assert.equal((yield* Ref.get(state)).resumeCount, 2);
23212385

23222386
const reclaimed = yield* mcpSessions.read(threadId);
23232387
assert.equal(reclaimed?.providerInstanceId, modelSelection.instanceId);
@@ -2333,7 +2397,7 @@ it.effect(
23332397
assert.isUndefined(yield* mcpSessions.read(threadId));
23342398
});
23352399

2336-
yield* effect.pipe(Effect.provide(layerTest({ state, idleTimeoutMs: 1_000 })));
2400+
yield* effect.pipe(Effect.provide(layerTest({ state, idleTimeoutMs: 1_000, afterMcpSet })));
23372401
}),
23382402
);
23392403

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

Lines changed: 41 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -441,6 +441,7 @@ export const layerWithOptions = (
441441
const prepareMcpSession = (
442442
threadId: ThreadId,
443443
providerInstanceId: ProviderInstanceId,
444+
onPrepared?: (prepared: PreparedMcpCredential) => void,
444445
): Effect.Effect<PreparedMcpCredential> =>
445446
options.configureMcp === false
446447
? mcpSessions
@@ -485,21 +486,31 @@ export const layerWithOptions = (
485486
// A credential issued while the lineage read failed is replaced once it works.
486487
(resolved.thread.subagent === true) === subagent
487488
) {
488-
return { mcpCredentialId: existing.providerSessionId, issued: false };
489+
const prepared = { mcpCredentialId: existing.providerSessionId, issued: false };
490+
onPrepared?.(prepared);
491+
return prepared;
489492
}
490493
dropMcpCredentialReservation(threadId, existing.providerSessionId);
491494
}
492-
yield* mcpSessionRegistry.revokeThread(threadId);
493-
const credential = yield* mcpSessionRegistry.issue({
494-
threadId,
495-
providerInstanceId,
496-
browserToolsAvailable,
497-
capabilities,
498-
subagent,
495+
const issue = Effect.gen(function* () {
496+
yield* mcpSessionRegistry.revokeThread(threadId);
497+
const credential = yield* mcpSessionRegistry.issue({
498+
threadId,
499+
providerInstanceId,
500+
browserToolsAvailable,
501+
capabilities,
502+
subagent,
503+
});
504+
yield* mcpSessions.set(credential.config);
505+
reserveMcpCredential(threadId, credential.config.providerSessionId);
506+
const prepared = {
507+
mcpCredentialId: credential.config.providerSessionId,
508+
issued: true,
509+
};
510+
onPrepared?.(prepared);
511+
return prepared;
499512
});
500-
yield* mcpSessions.set(credential.config);
501-
reserveMcpCredential(threadId, credential.config.providerSessionId);
502-
return { mcpCredentialId: credential.config.providerSessionId, issued: true };
513+
return yield* onPrepared === undefined ? issue : Effect.uninterruptible(issue);
503514
}),
504515
);
505516
const clearMcpSession = (threadId: ThreadId, mcpCredentialId?: string) =>
@@ -1087,8 +1098,7 @@ export const layerWithOptions = (
10871098
Array.from(current.values()).some(
10881099
(other) =>
10891100
other !== entry &&
1090-
(other.attachedThreadIds.has(threadId) ||
1091-
other.mcpCredentialIdByThread.get(threadId) === mcpCredentialId),
1101+
other.mcpCredentialIdByThread.get(threadId) === mcpCredentialId,
10921102
);
10931103
return heldElsewhere
10941104
? Effect.void
@@ -1363,8 +1373,13 @@ export const layerWithOptions = (
13631373
.get(sessionKey(input.providerSessionId))
13641374
?.attachedThreadIds.has(input.threadId) === true;
13651375
if (attached === undefined && !reclaimsCredential) return;
1366-
const prepared = yield* prepareMcpSession(input.threadId, input.providerInstanceId);
1367-
preparedForCleanup = prepared;
1376+
const prepared = yield* prepareMcpSession(
1377+
input.threadId,
1378+
input.providerInstanceId,
1379+
(prepared) => {
1380+
preparedForCleanup = prepared;
1381+
},
1382+
);
13681383
if (prepared.mcpCredentialId !== undefined) {
13691384
const mcpCredentialId = prepared.mcpCredentialId;
13701385
yield* Ref.update(sessions, (current) => {
@@ -1373,8 +1388,16 @@ export const layerWithOptions = (
13731388
if (entry === undefined) return current;
13741389
const mcpCredentialIdByThread = new Map(entry.mcpCredentialIdByThread);
13751390
mcpCredentialIdByThread.set(input.threadId, mcpCredentialId);
1391+
const loadedProviderThreadKeyByThread = new Map(
1392+
entry.loadedProviderThreadKeyByThread,
1393+
);
1394+
if (reclaimsCredential) loadedProviderThreadKeyByThread.delete(input.threadId);
13761395
const updated = new Map(current);
1377-
updated.set(key, { ...entry, mcpCredentialIdByThread });
1396+
updated.set(key, {
1397+
...entry,
1398+
mcpCredentialIdByThread,
1399+
loadedProviderThreadKeyByThread,
1400+
});
13781401
return updated;
13791402
});
13801403
}
@@ -1422,10 +1445,8 @@ export const layerWithOptions = (
14221445
Array.from(current.values()).some(
14231446
(other) =>
14241447
other.runtime !== attachedRuntime &&
1425-
(other.attachedThreadIds.has(input.threadId) ||
1426-
(mcpCredentialId !== undefined &&
1427-
other.mcpCredentialIdByThread.get(input.threadId) ===
1428-
mcpCredentialId)),
1448+
mcpCredentialId !== undefined &&
1449+
other.mcpCredentialIdByThread.get(input.threadId) === mcpCredentialId,
14291450
);
14301451
return held
14311452
? Effect.void

‎packages/provider-pi/src/server/adapter.test.ts‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1574,7 +1574,7 @@ describe("PiAdapterV2", () => {
15741574
browserToolsAvailable: true,
15751575
});
15761576
const fake = yield* makeFakePi;
1577-
yield* openRuntime(fake);
1577+
const { runtime } = yield* openRuntime(fake);
15781578
const spawn = fake.lastSpawn();
15791579
assert.isTrue(spawn.args.includes("--extension"));
15801580
const extensions = spawn.args.flatMap((arg, index) =>
@@ -1585,6 +1585,21 @@ describe("PiAdapterV2", () => {
15851585
assert.equal(spawn.env.T3_MCP_URL, "http://127.0.0.1:43123/mcp");
15861586
assert.equal(spawn.env.T3_MCP_BEARER_TOKEN, "secret-pi-token");
15871587
assert.equal(spawn.env.T3_PI_RUNTIME_MODE, "full-access");
1588+
const credentialPath = spawn.env.T3_PI_MCP_CREDENTIAL_PATH;
1589+
assert.isDefined(credentialPath);
1590+
const fileSystem = yield* FileSystem.FileSystem;
1591+
assert.equal(yield* fileSystem.readFileString(credentialPath!), "Bearer secret-pi-token");
1592+
const mcpSessions = yield* McpProviderSessions.McpProviderSessions;
1593+
const config = yield* mcpSessions.read(THREAD_ID);
1594+
yield* mcpSessions.set({ ...config!, authorizationHeader: "Bearer reclaimed-pi-token" });
1595+
const providerThread = yield* runtime.ensureThread({
1596+
threadId: THREAD_ID,
1597+
modelSelection: modelSelection("default"),
1598+
runtimePolicy,
1599+
});
1600+
yield* startTurn(runtime, providerThread);
1601+
yield* fake.takeRequest("prompt");
1602+
assert.equal(yield* fileSystem.readFileString(credentialPath!), "Bearer reclaimed-pi-token");
15881603
}).pipe(Effect.scoped, Effect.provide(layerTest)),
15891604
);
15901605

‎packages/provider-pi/src/server/adapter.ts‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ import {
7070
materializePiT3McpExtension,
7171
resolvePiLaunchArgs,
7272
} from "./mcpInjection.ts";
73-
import { PI_FILE_CHANGE_TOOLS } from "./mcpExtensionSource.ts";
73+
import { PI_FILE_CHANGE_TOOLS, T3_PI_MCP_CREDENTIAL_PATH_ENV } from "./mcpExtensionSource.ts";
7474

7575
export const PI_PROVIDER = ProviderDriverKind.make("pi");
7676
const PI_DRIVER_KIND = PI_PROVIDER;
@@ -441,7 +441,7 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* (
441441
const scope = yield* Effect.scope;
442442
const cwd = input.runtimePolicy.cwd ?? host.paths.cwd;
443443
const mcpSession = yield* mcpSessions.read(input.threadId);
444-
const provideCacheFs = <A, E>(effect: Effect.Effect<A, E, FileSystem.FileSystem>) =>
444+
const provideCacheFs = <A, E, R>(effect: Effect.Effect<A, E, R | FileSystem.FileSystem>) =>
445445
effect.pipe(
446446
Effect.provideService(FileSystem.FileSystem, fileSystem),
447447
Effect.mapError(
@@ -468,6 +468,19 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* (
468468
extensionPath,
469469
runtimeMode: input.runtimePolicy.runtimeMode,
470470
});
471+
const credentialPath =
472+
mcpSession === undefined
473+
? undefined
474+
: `${yield* provideCacheFs(fileSystem.makeTempDirectoryScoped({ prefix: "t3-pi-mcp-" }))}/credential`;
475+
let authorizationHeader = mcpSession?.authorizationHeader;
476+
if (credentialPath !== undefined && authorizationHeader !== undefined) {
477+
yield* provideCacheFs(
478+
fileSystem.writeFileString(credentialPath, authorizationHeader, { mode: 0o600 }),
479+
);
480+
launch.env[T3_PI_MCP_CREDENTIAL_PATH_ENV] = credentialPath;
481+
} else {
482+
delete launch.env[T3_PI_MCP_CREDENTIAL_PATH_ENV];
483+
}
471484
const scopedLaunch = yield* agentScope.wrap({
472485
command: options.settings.binaryPath || "pi",
473486
args: launch.args,
@@ -2462,6 +2475,25 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* (
24622475
) {
24632476
return yield* protocolError("Pi turn requested for a different native session");
24642477
}
2478+
if (credentialPath !== undefined) {
2479+
const current = yield* mcpSessions.read(turnInput.threadId);
2480+
if (current === undefined || current.providerInstanceId !== options.instanceId) {
2481+
return yield* protocolError("Pi turn has no current MCP credential");
2482+
}
2483+
if (current.authorizationHeader !== authorizationHeader) {
2484+
yield* fileSystem.writeFileString(credentialPath, current.authorizationHeader).pipe(
2485+
Effect.mapError(
2486+
(cause) =>
2487+
new ProviderAdapter.ProviderAdapterProtocolError({
2488+
driver: PI_PROVIDER,
2489+
cause,
2490+
detail: "Could not refresh Pi MCP credential",
2491+
}),
2492+
),
2493+
);
2494+
authorizationHeader = current.authorizationHeader;
2495+
}
2496+
}
24652497
// Fork adoption uses the orchestrator-allocated row; preserve that identity in later updates.
24662498
const continuation =
24672499
turnInput.message.createdBy === "agent" &&

‎packages/provider-pi/src/server/mcpBridge.testkit.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// @effect-diagnostics nodeBuiltinImport:off -- Executes the shipped Pi extension at its native JavaScript boundary.
22
import * as NodeModule from "node:module";
33
import * as NodePath from "node:path";
4+
import * as NodeFSP from "node:fs/promises";
45
import * as NodeVM from "node:vm";
56
import { PI_T3_MCP_EXTENSION_SOURCE } from "./mcpExtensionSource.ts";
67

@@ -48,6 +49,7 @@ export async function loadMcpBridge(
4849
readonly allowsTool?: (name: string) => boolean;
4950
readonly runtimeMode?: string;
5051
readonly inputSchema?: Record<string, unknown>;
52+
readonly credentialPath?: string;
5153
} = {},
5254
) {
5355
const handlers = new Map<string, AgentStartHook>();
@@ -90,12 +92,14 @@ export async function loadMcpBridge(
9092
env: {
9193
T3_MCP_URL: "http://fixture.invalid/mcp",
9294
T3_MCP_BEARER_TOKEN: "fixture-token",
95+
T3_PI_MCP_CREDENTIAL_PATH: options.credentialPath,
9396
T3_PI_MCP_EXTENSION_PATH: "/fixture/pi-t3-extension.ts",
9497
T3_PI_RUNTIME_MODE: options.runtimeMode,
9598
},
9699
},
97100
AbortSignal,
98101
NodePath,
102+
NodeFSP,
99103
Type: { Unsafe: (schema: unknown) => schema },
100104
fetch: async (
101105
url: string,

0 commit comments

Comments
 (0)