perf(server): linked calls reuse sessions while a tool waits

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<void>();
+        const release = yield* Deferred.make<void>();
+        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<string | undefined> = [];
+    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),
+);

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({