diff --git a/apps/server/src/peer/PeerLinks.ts b/apps/server/src/peer/PeerLinks.ts index 48f61e6554a17..550d3ed757f07 100644 --- a/apps/server/src/peer/PeerLinks.ts +++ b/apps/server/src/peer/PeerLinks.ts @@ -47,10 +47,10 @@ const PeerHttpApi = HttpApi.make(EnvironmentHttpApi.identifier) /** The redirect registered for a link. Nothing ever listens there: the code comes back in the response. */ const LINK_REDIRECT_URI = "http://127.0.0.1/t3code-peer-link"; /** - * How long a peer's descriptor may take before a route counts as unreachable. - * A tailnet or relayed route to a machine that is waking up takes longer. + * How long a peer's descriptor or MCP handshake may take before a route counts + * as unreachable. A tailnet or relayed route to a machine that is waking up takes longer. */ -const probeTimeout = (origin: string) => +export const probeTimeout = (origin: string) => isLocalLoopbackHost(new URL(origin).hostname) ? Duration.seconds(3) : Duration.seconds(10); /** How long a listing waits on a link's addresses; a real call still gives a waking peer `probeTimeout`. */ const LIST_PROBE_TIMEOUT = Duration.seconds(3); diff --git a/apps/server/src/peer/PeerMcpClient.test.ts b/apps/server/src/peer/PeerMcpClient.test.ts index ef8a3aeff6a65..e8703f794f365 100644 --- a/apps/server/src/peer/PeerMcpClient.test.ts +++ b/apps/server/src/peer/PeerMcpClient.test.ts @@ -9,12 +9,15 @@ import { import * as Context from "effect/Context"; import * as DateTime from "effect/DateTime"; import * as Deferred from "effect/Deferred"; +import * as Duration from "effect/Duration"; import * as Effect from "effect/Effect"; import * as Fiber from "effect/Fiber"; import * as Layer from "effect/Layer"; +import * as Queue from "effect/Queue"; import * as Ref from "effect/Ref"; import * as Schema from "effect/Schema"; import { Tool, Toolkit } from "effect/ai"; +import * as TestClock from "effect/testing/TestClock"; import { HttpClient, HttpClientResponse } from "effect/http"; import * as McpHttpServer from "../mcp/McpHttpServer.ts"; @@ -44,6 +47,49 @@ const toolkit = Toolkit.make( }), ); +/** One `probe` call to a linked peer at a remote address, answered by `http`. */ +const probeThrough = Effect.fn(function* (http: HttpClient.HttpClient) { + 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, + }), + ), + ), + ), + ); + return Context.get(context, PeerMcpClient.PeerMcpClient).call({ + environmentId: box.environmentId, + tool: "probe", + arguments: {}, + limits: { runtimeMode: "approval-required", interactionMode: "plan" }, + success: Result, + }); +}); + +const probeResult = JSON.stringify({ + result: { content: [{ type: "text", text: '{"modes":"approval-required/plan"}' }] }, +}); + const refusal = new OrchestratorMcpFailure({ code: "thread_not_found", message: "Thread thread-gone was not found.", @@ -295,3 +341,45 @@ it.effect("renews a missing peer session once and reuses its replacement", () => ]); }).pipe(Effect.scoped), ); + +it.effect("bounds the session handshake, not a tool call that waits longer", () => + Effect.gen(function* () { + const arrived = yield* Queue.unbounded(); + const answer = yield* Deferred.make(); + let requests = 0; + const http = HttpClient.make((request) => + Effect.gen(function* () { + const index = requests++; + yield* Queue.offer(arrived, index); + // The first initialize, then the next initialized notification, never answer; + // the third session's tool call answers late. + if (index === 0 || index === 2) return yield* Effect.never; + if (index === 5) yield* Deferred.await(answer); + return HttpClientResponse.fromWeb( + request, + new Response(index === 5 ? probeResult : "", { + status: index === 4 ? 202 : 200, + headers: { "mcp-session-id": "session" }, + }), + ); + }), + ); + const probe = yield* probeThrough(http); + + for (const requestsBeforeStall of [1, 2]) { + const stalled = yield* probe.pipe(Effect.flip, Effect.forkScoped); + yield* Queue.takeN(arrived, requestsBeforeStall); + yield* TestClock.adjust(Duration.seconds(10)); + expect(yield* Fiber.join(stalled)).toMatchObject({ + code: "orchestration_error", + message: "The linked environment stopped answering.", + }); + } + + const waiting = yield* Effect.forkScoped(probe); + yield* Queue.takeN(arrived, 3); + yield* TestClock.adjust(Duration.minutes(1)); + yield* Deferred.succeed(answer, undefined); + expect(yield* Fiber.join(waiting)).toEqual({ modes: "approval-required/plan" }); + }).pipe(Effect.scoped), +); diff --git a/apps/server/src/peer/PeerMcpClient.ts b/apps/server/src/peer/PeerMcpClient.ts index e800333fb1e08..8c438ad702b53 100644 --- a/apps/server/src/peer/PeerMcpClient.ts +++ b/apps/server/src/peer/PeerMcpClient.ts @@ -90,6 +90,11 @@ const linkedEnvironmentsOffThere = new OrchestratorMcpFailure({ "The linked environment has Linked environments turned off. Turn it on in its Settings → Labs.", }); +const stoppedAnswering = new OrchestratorMcpFailure({ + code: "orchestration_error", + message: "The linked environment stopped answering.", +}); + /** The body of an MCP response: JSON, or a single server-sent event carrying JSON. */ const responseJson = (contentType: string | undefined, text: string) => contentType?.includes("text/event-stream") === true @@ -119,45 +124,45 @@ const make = Effect.gen(function* () { ), ); - const openSession = Effect.fn("PeerMcpClient.openSession")(function* ( - url: string, - token: string, - limitHeader: Record, - ) { - const initialized = yield* post( - url, - token, - { - jsonrpc: "2.0", - id: 0, - method: "initialize", - params: { - protocolVersion: PROTOCOL_VERSION, - capabilities: {}, - clientInfo: { name: "t3code-peer", version: "1" }, + const openSession = Effect.fn("PeerMcpClient.openSession")( + function* (url: string, token: string, limitHeader: Record) { + const initialized = yield* post( + url, + token, + { + jsonrpc: "2.0", + id: 0, + method: "initialize", + params: { + protocolVersion: PROTOCOL_VERSION, + capabilities: {}, + clientInfo: { name: "t3code-peer", version: "1" }, + }, }, - }, - limitHeader, - ); - const sessionId = initialized.headers["mcp-session-id"]; - if (initialized.status === 403) return yield* linkedEnvironmentsOffThere; - if (initialized.status !== 200 || sessionId === undefined) { - return yield* new OrchestratorMcpFailure({ - code: initialized.status === 401 ? "capability_denied" : "orchestration_error", - message: - initialized.status === 401 - ? "The linked environment no longer accepts this link. It may have been revoked there; link it again." - : `The linked environment refused to start an MCP session (HTTP ${initialized.status}).`, - }); - } - yield* post( - url, - token, - { jsonrpc: "2.0", method: "notifications/initialized" }, - { ...limitHeader, "mcp-session-id": sessionId, "mcp-protocol-version": PROTOCOL_VERSION }, - ); - return sessionId; - }); + limitHeader, + ); + const sessionId = initialized.headers["mcp-session-id"]; + if (initialized.status === 403) return yield* linkedEnvironmentsOffThere; + if (initialized.status !== 200 || sessionId === undefined) { + return yield* new OrchestratorMcpFailure({ + code: initialized.status === 401 ? "capability_denied" : "orchestration_error", + message: + initialized.status === 401 + ? "The linked environment no longer accepts this link. It may have been revoked there; link it again." + : `The linked environment refused to start an MCP session (HTTP ${initialized.status}).`, + }); + } + yield* post( + url, + token, + { jsonrpc: "2.0", method: "notifications/initialized" }, + { ...limitHeader, "mcp-session-id": sessionId, "mcp-protocol-version": PROTOCOL_VERSION }, + ); + return sessionId; + }, + // Only the handshake is bounded: a forwarded tool call may wait for minutes. + (effect, url) => Effect.timeout(effect, PeerLinks.probeTimeout(url)), + ); const call: PeerMcpClient["Service"]["call"] = Effect.fn("PeerMcpClient.call")( function* (input) { @@ -285,13 +290,8 @@ const make = Effect.gen(function* () { Effect.tap(() => links.recordOutcome(input.environmentId, { error: null })), // Recorded above with the transport's own message; the agent hears it plainly. Effect.catchTags({ - HttpClientError: () => - Effect.fail( - new OrchestratorMcpFailure({ - code: "orchestration_error", - message: "The linked environment stopped answering.", - }), - ), + HttpClientError: () => Effect.fail(stoppedAnswering), + TimeoutError: () => Effect.fail(stoppedAnswering), }), Effect.flatMap(Effect.fromResult), ),