From f13ca86b8b1338552bf3875dc6216c2e12255c61 Mon Sep 17 00:00:00 2001 From: Adamulek123 Date: Wed, 28 Oct 2026 17:36:30 +0100 Subject: [PATCH] perf(server): reuse initialized peer sessions during pending calls --- apps/server/src/peer/PeerMcpClient.test.ts | 177 +++++++++++++++++++++ apps/server/src/peer/PeerMcpClient.ts | 6 +- 2 files changed, 181 insertions(+), 2 deletions(-) create mode 100644 apps/server/src/peer/PeerMcpClient.test.ts diff --git a/apps/server/src/peer/PeerMcpClient.test.ts b/apps/server/src/peer/PeerMcpClient.test.ts new file mode 100644 index 0000000000000..7d19f4e856bf7 --- /dev/null +++ b/apps/server/src/peer/PeerMcpClient.test.ts @@ -0,0 +1,177 @@ +import { NodeHttpServer } from "@effect/platform-node"; +import { expect, it } from "@effect/vitest"; +import * as Context from "effect/Context"; +import * as DateTime from "effect/DateTime"; +import * as Deferred from "effect/Deferred"; +import * as Effect from "effect/Effect"; +import * as Fiber from "effect/Fiber"; +import * as Layer from "effect/Layer"; +import * as Ref from "effect/Ref"; +import * as Schema from "effect/Schema"; +import { Tool, Toolkit } from "effect/ai"; +import { HttpClient, HttpClientResponse } from "effect/http"; + +import * as McpHttpServer from "../mcp/McpHttpServer.ts"; +import * as McpInvocationContext from "../mcp/McpInvocationContext.ts"; +import * as McpToolAccess from "../mcp/McpToolAccess.ts"; +import * as PeerLinks from "./PeerLinks.ts"; +import * as PeerMcpClient from "./PeerMcpClient.ts"; +import { descriptorOf, layerLinkingEnvironment, linkTo, servePeer } from "./PeerLinks.testkit.ts"; + +const Result = Schema.Struct({ modes: Schema.String }); +const toolkit = Toolkit.make( + Tool.make("probe", { + description: "Returns this request's mode limits.", + success: Result, + dependencies: [McpInvocationContext.McpInvocationContext], + }), + Tool.make("hold", { + description: "Holds a tool request open.", + success: Result, + dependencies: [McpInvocationContext.McpInvocationContext], + }), +); + +it.effect.each([false, true])( + "reuses an initialized session before its first call finishes (cancelled: %s)", + (cancelFirst) => + Effect.scoped( + Effect.gen(function* () { + const entered = yield* Deferred.make(); + const release = yield* Deferred.make(); + const modes = McpInvocationContext.McpInvocationContext.pipe( + Effect.map((scope) => { + const ceiling = McpInvocationContext.clientModeCeiling(scope.client); + return { modes: `${ceiling.runtimeMode}/${ceiling.interactionMode}` }; + }), + ); + const handlers = McpToolAccess.toLayer(toolkit, { + probe: McpToolAccess.reads(() => modes), + hold: McpToolAccess.reads(() => + Deferred.succeed(entered, undefined).pipe( + Effect.andThen(Deferred.await(release)), + Effect.andThen(modes), + ), + ), + }); + const box = descriptorOf("environment-box", "Box"); + const b = yield* servePeer(box, McpHttpServer.toolkitRegistration(toolkit, handlers)); + const context = yield* Layer.build( + layerLinkingEnvironment(descriptorOf("environment-a", "A")), + ); + const links = Context.get(context, PeerLinks.PeerLinks); + const client = Context.get(context, PeerMcpClient.PeerMcpClient); + yield* linkTo(links, b); + const held = yield* client + .call({ + environmentId: box.environmentId, + tool: "hold", + arguments: {}, + limits: { runtimeMode: "full-access", interactionMode: "default" }, + success: Result, + }) + .pipe(Effect.forkScoped); + yield* Deferred.await(entered); + if (cancelFirst) yield* Fiber.interrupt(held); + const before = (yield* Ref.get(b.bearers)).length; + expect( + yield* client.call({ + environmentId: box.environmentId, + tool: "probe", + arguments: {}, + limits: { runtimeMode: "approval-required", interactionMode: "plan" }, + success: Result, + }), + ).toEqual({ modes: "approval-required/plan" }); + // One tool POST; no second initialize or initialized notification. + expect((yield* Ref.get(b.bearers)).length - before).toBe(1); + yield* Deferred.succeed(release, undefined); + if (!cancelFirst) expect(yield* Fiber.join(held)).toEqual({ modes: "auto/default" }); + }), + ).pipe(Effect.provide(NodeHttpServer.layerTest)), +); + +it.effect("renews a missing peer session once and reuses its replacement", () => + Effect.gen(function* () { + const sessionHeaders: Array = []; + const statuses = [200, 202, 404, 200, 202, 200, 200]; + const http = HttpClient.make((request) => + Effect.sync(() => { + const index = sessionHeaders.length; + sessionHeaders.push(request.headers["mcp-session-id"]); + const status = statuses[index]; + if (status === undefined) throw new Error("Unexpected retry"); + return HttpClientResponse.fromWeb( + request, + new Response( + index >= 5 + ? JSON.stringify({ + result: { + content: [{ type: "text", text: '{"modes":"approval-required/plan"}' }], + }, + }) + : "", + { + status, + headers: + index === 0 + ? { "mcp-session-id": "old" } + : index === 3 + ? { "mcp-session-id": "renewed" } + : {}, + }, + ), + ); + }), + ); + const box = descriptorOf("environment-box", "Box"); + const context = yield* Layer.build( + PeerMcpClient.layer.pipe( + Layer.provide( + Layer.mergeAll( + Layer.succeed(HttpClient.HttpClient, http), + Layer.mock(PeerLinks.PeerLinks)({ + resolve: () => + Effect.succeed({ + link: { + environmentId: box.environmentId, + label: box.label, + urls: ["https://peer.example"], + access: "auto", + linkedAt: DateTime.makeUnsafe(0), + expiresAt: DateTime.makeUnsafe(1), + lastReachedAt: null, + lastError: null, + }, + url: "https://peer.example", + token: "fixture", + }), + recordOutcome: () => Effect.void, + }), + ), + ), + ), + ); + const client = Context.get(context, PeerMcpClient.PeerMcpClient); + for (let call = 0; call < 2; call++) { + expect( + yield* client.call({ + environmentId: box.environmentId, + tool: "probe", + arguments: {}, + limits: { runtimeMode: "approval-required", interactionMode: "plan" }, + success: Result, + }), + ).toEqual({ modes: "approval-required/plan" }); + } + expect(sessionHeaders).toEqual([ + undefined, + "old", + "old", + undefined, + "renewed", + "renewed", + "renewed", + ]); + }).pipe(Effect.scoped), +); diff --git a/apps/server/src/peer/PeerMcpClient.ts b/apps/server/src/peer/PeerMcpClient.ts index b953d458f8e27..27397d54ba7de 100644 --- a/apps/server/src/peer/PeerMcpClient.ts +++ b/apps/server/src/peer/PeerMcpClient.ts @@ -175,13 +175,15 @@ const make = Effect.gen(function* () { ); const existing = (yield* Ref.get(sessions)).get(sessionKey); const sessionId = existing ?? (yield* openSession(resolved.url, resolved.token, limitHeader)); + // The initialized session outlives a tool request, including a long wait or cancellation. + if (existing === undefined) { + yield* Ref.update(sessions, (current) => new Map(current).set(sessionKey, sessionId)); + } let response = yield* callOnce(sessionId); if (response.status === 404) { const renewed = yield* openSession(resolved.url, resolved.token, limitHeader); yield* Ref.update(sessions, (current) => new Map(current).set(sessionKey, renewed)); response = yield* callOnce(renewed); - } else if (existing === undefined) { - yield* Ref.update(sessions, (current) => new Map(current).set(sessionKey, sessionId)); } if (response.status === 401) { return yield* new OrchestratorMcpFailure({