diff --git a/README.md b/README.md index 6587f7e..1059407 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,8 @@ If you're building an [Agent](https://agentclientprotocol.com/protocol/overview# If you're building a [Client](https://agentclientprotocol.com/protocol/overview#client), start with `client({ name })`, register client-side handlers such as `requestPermission(...)` and `sessionUpdate(...)`, then run your agent workflow with `connectWith(stream, async (ctx) => ...)`. +If you're building something that sits between the two — a logger, an authorization gate, a message transformer — start with `proxy()`. Register typed handlers with `onRequestFromClient(...)` / `onNotificationFromAgent(...)` (the method names say which direction they intercept) just like the agent/client builders, then call `connect({ client, agent })`. Anything you don't claim is forwarded untouched in both directions; handlers can rewrite, answer, or drop messages before they cross. + ### Study a Production Implementation For a complete, production-ready implementation, check out the [Gemini CLI Agent](https://github.com/google-gemini/gemini-cli/blob/main/packages/cli/src/zed-integration/zedIntegration.ts). diff --git a/src/acp.ts b/src/acp.ts index 0fb42c4..29ac3bc 100644 --- a/src/acp.ts +++ b/src/acp.ts @@ -21,6 +21,19 @@ export { PROTOCOL_VERSION, } from "./schema/index.js"; export * from "./stream.js"; +export { proxy } from "./proxy.js"; +// ProxyBuilder is type-only on purpose: proxy() is the sole factory, so its +// constructor is not part of the public contract. +export type { + ProxyBuilder, + ProxyHandle, + ProxyNotificationContext, + ProxyNotificationHandler, + ProxyRequestContext, + ProxyRequestHandler, + ProxySideConnection, + ProxyStreams, +} from "./proxy.js"; export { RequestError } from "./jsonrpc.js"; export type { AnyMessage, diff --git a/src/examples/README.md b/src/examples/README.md index f00c3cd..58c5254 100644 --- a/src/examples/README.md +++ b/src/examples/README.md @@ -4,6 +4,7 @@ This directory contains examples using the [ACP](https://agentclientprotocol.com - [`agent.ts`](./agent.ts) - A minimal agent implementation that simulates LLM interaction - [`client.ts`](./client.ts) - A minimal client implementation that spawns the [`agent.ts`](./agent.ts) as a subprocess +- [`proxy.ts`](./proxy.ts) - A pass-through proxy that wraps any agent command and logs the messages crossing it in both directions - [`http-server.ts`](./http-server.ts) - A minimal ACP Streamable HTTP server with WebSocket upgrade support - [`http-client.ts`](./http-client.ts) - A minimal client using `createHttpStream` - [`ws-client.ts`](./ws-client.ts) - A minimal client using `createWebSocketStream` diff --git a/src/examples/proxy.ts b/src/examples/proxy.ts new file mode 100644 index 0000000..f703088 --- /dev/null +++ b/src/examples/proxy.ts @@ -0,0 +1,70 @@ +#!/usr/bin/env node + +import { spawn } from "node:child_process"; +import { dirname, join } from "node:path"; +import { Readable, Writable } from "node:stream"; +import { fileURLToPath } from "node:url"; + +import * as acp from "../acp.js"; + +// A minimal ACP proxy: it presents as an agent on its own stdio, spawns the +// real agent as a subprocess, and forwards every message between the two +// while logging traffic to stderr. Point an ACP client (like Zed) at this +// script exactly as it would run the agent directly — neither side needs to +// know the proxy is there. +// +// Usage: proxy.ts [command args...] +// Runs the example agent from this directory when no command is given. + +function logAndForward( + direction: string, +): acp.ProxyRequestHandler { + return ({ method, params, forward }) => { + console.error(`[proxy] ${direction} request: ${method}`); + return forward(params); + }; +} + +function logAndForwardNotification( + direction: string, +): acp.ProxyNotificationHandler { + return async ({ method, params, forward }) => { + console.error(`[proxy] ${direction} notification: ${method}`); + await forward(params); + }; +} + +// Spawn the wrapped agent: the command from argv, or the example agent. +const npxCmd = process.platform === "win32" ? "npx.cmd" : "npx"; +const exampleAgent = join(dirname(fileURLToPath(import.meta.url)), "agent.ts"); +const [command, ...args] = + process.argv.length > 2 + ? process.argv.slice(2) + : [npxCmd, "tsx", exampleAgent]; +const agentProcess = spawn(command, args, { + stdio: ["pipe", "pipe", "inherit"], +}); + +// "*" catches anything without an exact registration; typed interception is +// also available, e.g. +// .onRequestFromClient("session/prompt", async ({ params, forward }) => ...). +const handle = acp + .proxy() + .onRequestFromClient("*", logAndForward("client → agent")) + .onNotificationFromClient("*", logAndForwardNotification("client → agent")) + .onRequestFromAgent("*", logAndForward("agent → client")) + .onNotificationFromAgent("*", logAndForwardNotification("agent → client")) + .connect({ + client: acp.ndJsonStream( + Writable.toWeb(process.stdout), + Readable.toWeb(process.stdin) as ReadableStream, + ), + agent: acp.ndJsonStream( + Writable.toWeb(agentProcess.stdin!), + Readable.toWeb(agentProcess.stdout!) as ReadableStream, + ), + }); + +agentProcess.once("exit", () => handle.close()); +await handle.closed; +agentProcess.kill(); diff --git a/src/jsonrpc.ts b/src/jsonrpc.ts index 504e0b1..536cd74 100644 --- a/src/jsonrpc.ts +++ b/src/jsonrpc.ts @@ -426,7 +426,12 @@ function requestCancelledError(reason?: unknown): RequestError { return RequestError.requestCancelled(reason); } -function errorToRequestResult( +/** + * Maps an error thrown while handling a request to a JSON-RPC result: + * `RequestError` keeps its code/message/data, an abort after cancellation + * maps to request-cancelled, anything else becomes a generic internal error. + */ +export function errorToRequestResult( error: unknown, signal: AbortSignal, ): Result { diff --git a/src/proxy.test.ts b/src/proxy.test.ts new file mode 100644 index 0000000..e806edb --- /dev/null +++ b/src/proxy.test.ts @@ -0,0 +1,952 @@ +import { PassThrough, Readable, Writable } from "node:stream"; + +import { describe, expect, it, vi } from "vitest"; + +import { agent, client, ndJsonStream } from "./acp.js"; +import { Connection, Handled, RequestError } from "./jsonrpc.js"; +import type { RequestResponder } from "./jsonrpc.js"; +import { proxy } from "./proxy.js"; +import type { ProxyBuilder, ProxyHandle } from "./proxy.js"; +import type { AnyMessage } from "./jsonrpc.js"; +import type { Stream } from "./stream.js"; + +function inMemoryStreamPair(): [Stream, Stream] { + const leftToRight = new TransformStream(); + const rightToLeft = new TransformStream(); + return [ + { readable: rightToLeft.readable, writable: leftToRight.writable }, + { readable: leftToRight.readable, writable: rightToLeft.writable }, + ]; +} + +type ProxySetup = { + clientStream: Stream; + agentStream: Stream; + builder: ProxyBuilder; + handle: ProxyHandle; +}; + +function setupProxy(configure?: (p: ProxyBuilder) => void): ProxySetup { + const [clientStream, proxyClientSide] = inMemoryStreamPair(); + const [proxyAgentSide, agentStream] = inMemoryStreamPair(); + const builder = proxy(); + configure?.(builder); + const handle = builder.connect({ + client: proxyClientSide, + agent: proxyAgentSide, + }); + return { clientStream, agentStream, builder, handle }; +} + +describe("proxy forwarding", () => { + it("forwards client requests to the agent and responses back", async () => { + const { clientStream, agentStream } = setupProxy(); + Connection.builder() + .onReceiveRequest( + "echo", + (params) => params, + (request, responder) => responder.respond({ echoed: request }), + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + const response = await clientEnd.sendRequest("echo", { value: 42 }); + + expect(response).toEqual({ echoed: { value: 42 } }); + }); + + it("forwards agent-initiated requests to the client", async () => { + const { clientStream, agentStream } = setupProxy(); + Connection.builder() + .onReceiveRequest( + "confirm", + (params) => params, + (_request, responder) => responder.respond({ approved: true }), + ) + .connect(clientStream); + const agentEnd = new Connection(agentStream, []); + + const response = await agentEnd.sendRequest("confirm", { action: "run" }); + + expect(response).toEqual({ approved: true }); + }); + + it("forwards notifications", async () => { + const { clientStream, agentStream } = setupProxy(); + const received = vi.fn(); + const seen = Promise.withResolvers(); + Connection.builder() + .onReceiveNotification( + "log", + (params) => params, + (notification) => { + received(notification); + seen.resolve(); + }, + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + await clientEnd.sendNotification("log", { line: "hello" }); + await seen.promise; + + expect(received).toHaveBeenCalledWith({ line: "hello" }); + }); + + it("preserves error responses across the proxy", async () => { + const { clientStream, agentStream } = setupProxy(); + Connection.builder() + .onReceiveRequest( + "always-fails", + (params) => params, + () => { + throw new RequestError(1234, "nope", { reason: "test" }); + }, + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + const failure = clientEnd.sendRequest("always-fails", {}); + + await expect(failure).rejects.toMatchObject({ + code: 1234, + message: "nope", + data: { reason: "test" }, + }); + }); + + it("responds method-not-found from the far side, not the proxy", async () => { + const { clientStream, agentStream } = setupProxy(); + Connection.builder().connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + const failure = clientEnd.sendRequest("no-such-method", {}); + + await expect(failure).rejects.toMatchObject({ code: -32601 }); + }); + + it("propagates cancellation to the side handling the request", async () => { + const { clientStream, agentStream } = setupProxy(); + Connection.builder() + .onReceiveRequest( + "slow", + (params) => params, + (_request, responder) => { + // Never respond on our own; settle only when the caller cancels. + responder.signal.addEventListener("abort", () => { + void responder.respondWithError(RequestError.requestCancelled()); + }); + }, + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + const canceller = new AbortController(); + const failure = clientEnd.sendRequest("slow", {}, undefined, { + cancellationSignal: canceller.signal, + }); + canceller.abort(); + + await expect(failure).rejects.toMatchObject({ code: -32800 }); + }); + + it("preserves notification order when an earlier handler is slow", async () => { + const { clientStream, agentStream } = setupProxy((p) => { + p.onNotificationFromAgent( + "update", + (params) => params as { tag: string; delay: number }, + async ({ params, forward }) => { + await new Promise((resolve) => setTimeout(resolve, params.delay)); + await forward(params); + }, + ); + }); + const received: string[] = []; + const gotBoth = Promise.withResolvers(); + Connection.builder() + .onReceiveNotification( + "update", + (params) => params as { tag: string }, + (notification) => { + received.push(notification.tag); + if (received.length === 2) { + gotBoth.resolve(); + } + }, + ) + .connect(clientStream); + const agentEnd = new Connection(agentStream, []); + + // Without serialized dispatch the second (instant) chain would overtake + // the first (delayed) one. + await agentEnd.sendNotification("update", { tag: "first", delay: 25 }); + await agentEnd.sendNotification("update", { tag: "second", delay: 0 }); + await gotBoth.promise; + + expect(received).toEqual(["first", "second"]); + }); + + it("does not block later messages behind a pending request round trip", async () => { + // `slow` goes through a typed handler that awaits forward(...), which + // must release the dispatch loop at the send, not the response. + const { clientStream, agentStream } = setupProxy((p) => { + p.onRequestFromClient( + "slow", + (params) => params, + async ({ params, forward }) => forward(params), + ); + }); + const slowArrived = Promise.withResolvers>(); + Connection.builder() + .onReceiveRequest( + "slow", + (params) => params, + (_request, responder) => { + // Hold the request open; it is answered after `ping` completes. + slowArrived.resolve(responder); + }, + ) + .onReceiveRequest( + "ping", + (params) => params, + (_request, responder) => responder.respond({ pong: true }), + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + const slow = clientEnd.sendRequest("slow", {}); + const ping = await clientEnd.sendRequest("ping", {}); + expect(ping).toEqual({ pong: true }); + + const responder = await slowArrived.promise; + await responder.respond({ done: true }); + await expect(slow).resolves.toEqual({ done: true }); + }); + + it("closes the other side when one side of the proxy closes", async () => { + const { handle } = setupProxy(); + + handle.client.close(new Error("client side went away")); + + await expect(handle.closed).resolves.toBeUndefined(); + expect(handle.agent.signal.aborted).toBe(true); + // The close reason crosses to the other side's pending requests. + expect(handle.agent.signal.reason).toMatchObject({ + message: "client side went away", + }); + }); + + it("closes both sides when the client transport reaches EOF", async () => { + // Over a real byte transport (stdio, sockets) a disconnect arrives as + // EOF on the read loop, not as a close() call. Byte pipes reproduce + // that, unlike the in-memory object streams used elsewhere. + const clientToProxy = new PassThrough(); + const proxyToClient = new PassThrough(); + const [proxyAgentSide, agentStream] = inMemoryStreamPair(); + const handle = proxy().connect({ + client: ndJsonStream( + Writable.toWeb(proxyToClient), + Readable.toWeb(clientToProxy) as ReadableStream, + ), + agent: proxyAgentSide, + }); + const agentEnd = new Connection(agentStream, []); + // In-flight traffic when the transport dies; over the in-memory agent + // pair this pending request can never settle, so it is not awaited. + void agentEnd.sendRequest("confirm", {}).catch(() => {}); + + // The client goes away mid-request. + clientToProxy.end(); + + await expect(handle.closed).resolves.toBeUndefined(); + expect(handle.client.signal.aborted).toBe(true); + expect(handle.agent.signal.aborted).toBe(true); + }); + + it("closes both sides through the handle", async () => { + const { handle } = setupProxy(); + + handle.close(); + + await expect(handle.closed).resolves.toBeUndefined(); + expect(handle.client.signal.aborted).toBe(true); + expect(handle.agent.signal.aborted).toBe(true); + }); +}); + +describe("proxy error and cancellation edges", () => { + it("propagates cancellation through a registered handler's forward", async () => { + const { clientStream, agentStream } = setupProxy((p) => { + p.onRequestFromClient( + "slow", + (params) => params, + async ({ params, forward }) => forward(params), + ); + }); + Connection.builder() + .onReceiveRequest( + "slow", + (params) => params, + (_request, responder) => { + responder.signal.addEventListener("abort", () => { + void responder.respondWithError(RequestError.requestCancelled()); + }); + }, + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + const canceller = new AbortController(); + const failure = clientEnd.sendRequest("slow", {}, undefined, { + cancellationSignal: canceller.signal, + }); + canceller.abort(); + + await expect(failure).rejects.toMatchObject({ code: -32800 }); + }); + + it("maps a handler's abort after cancellation to request-cancelled", async () => { + const { clientStream } = setupProxy((p) => { + p.onRequestFromClient( + "watch", + (params) => params, + ({ signal }) => + new Promise((_resolve, reject) => { + signal.addEventListener("abort", () => { + const abortError = new Error("The operation was aborted"); + abortError.name = "AbortError"; + reject(abortError); + }); + }), + ); + }); + const clientEnd = new Connection(clientStream, []); + + const canceller = new AbortController(); + const failure = clientEnd.sendRequest("watch", {}, undefined, { + cancellationSignal: canceller.signal, + }); + canceller.abort(); + + await expect(failure).rejects.toMatchObject({ code: -32800 }); + }); + + it("answers a generic handler throw as an internal error", async () => { + const { clientStream } = setupProxy((p) => { + p.onRequestFromClient( + "boom", + (params) => params, + () => { + throw new Error("kaput"); + }, + ); + }); + const clientEnd = new Connection(clientStream, []); + + await expect(clientEnd.sendRequest("boom", {})).rejects.toMatchObject({ + code: -32603, + }); + }); + + it("answers a parser failure as invalid params", async () => { + class FakeZodError extends Error { + name = "ZodError"; + issues: unknown[] = []; + format(): unknown { + return { _errors: ["expected a number"] }; + } + } + const { clientStream } = setupProxy((p) => { + p.onRequestFromClient( + "strict", + () => { + throw new FakeZodError("invalid"); + }, + async ({ params, forward }) => forward(params), + ); + }); + const clientEnd = new Connection(clientStream, []); + + await expect( + clientEnd.sendRequest("strict", { n: "not-a-number" }), + ).rejects.toMatchObject({ code: -32602 }); + }); + + it("drops a throwing notification handler's message and keeps dispatching", async () => { + const consoleError = vi + .spyOn(console, "error") + .mockImplementation(() => {}); + const { clientStream, agentStream } = setupProxy((p) => { + p.onNotificationFromClient( + "log", + (params) => params as { fail: boolean }, + async ({ params, forward }) => { + if (params.fail) { + throw new Error("handler exploded"); + } + await forward(params); + }, + ); + }); + const received = vi.fn(); + const gotSecond = Promise.withResolvers(); + Connection.builder() + .onReceiveNotification( + "log", + (params) => params, + (notification) => { + received(notification); + gotSecond.resolve(); + }, + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + await clientEnd.sendNotification("log", { fail: true }); + await clientEnd.sendNotification("log", { fail: false }); + await gotSecond.promise; + + expect(received).toHaveBeenCalledTimes(1); + expect(received).toHaveBeenCalledWith({ fail: false }); + consoleError.mockRestore(); + }); +}); + +describe("proxy ordering and correlation", () => { + it("preserves request send order through an async registered handler", async () => { + const { clientStream, agentStream } = setupProxy((p) => { + p.onRequestFromClient( + "echo", + (params) => params as { n: number }, + async ({ params, forward }) => { + // Delay before forwarding; serialization must keep send order. + await new Promise((resolve) => setTimeout(resolve, 15 - params.n)); + return forward(params); + }, + ); + }); + const arrivals: number[] = []; + Connection.builder() + .onReceiveRequest( + "echo", + (params) => params as { n: number }, + (request, responder) => { + arrivals.push(request.n); + return responder.respond({ echoed: request.n }); + }, + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + const [first, second] = await Promise.all([ + clientEnd.sendRequest("echo", { n: 1 }), + clientEnd.sendRequest("echo", { n: 2 }), + ]); + + expect(arrivals).toEqual([1, 2]); + expect(first).toEqual({ echoed: 1 }); + expect(second).toEqual({ echoed: 2 }); + }); + + it("correlates concurrent round trips answered out of order", async () => { + const { clientStream, agentStream } = setupProxy(); + const held: Array<{ n: number; responder: RequestResponder }> = []; + const gotBoth = Promise.withResolvers(); + Connection.builder() + .onReceiveRequest( + "hold", + (params) => params as { n: number }, + (request, responder) => { + held.push({ n: request.n, responder }); + if (held.length === 2) { + gotBoth.resolve(); + } + }, + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + const first = clientEnd.sendRequest("hold", { n: 1 }); + const second = clientEnd.sendRequest("hold", { n: 2 }); + await gotBoth.promise; + + // Answer in reverse order; each caller must still get its own response. + await held[1]!.responder.respond({ answered: 2 }); + await held[0]!.responder.respond({ answered: 1 }); + + await expect(second).resolves.toEqual({ answered: 2 }); + await expect(first).resolves.toEqual({ answered: 1 }); + }); + + it("keeps pass-through traffic ordered behind a busy registered handler", async () => { + const { clientStream, agentStream } = setupProxy((p) => { + p.onNotificationFromAgent( + "update", + (params) => params, + async ({ params, forward }) => { + await new Promise((resolve) => setTimeout(resolve, 20)); + await forward(params); + }, + ); + }); + const received: string[] = []; + const gotBoth = Promise.withResolvers(); + Connection.builder() + .onReceiveMessage((message) => { + received.push(message.method); + if (received.length === 2) { + gotBoth.resolve(); + } + return Handled.yes(); + }) + .connect(clientStream); + const agentEnd = new Connection(agentStream, []); + + // "update" is registered (slow); "other" is unregistered pass-through. + // The fast path must not let "other" overtake the busy queue. + await agentEnd.sendNotification("update", {}); + await agentEnd.sendNotification("other", {}); + await gotBoth.promise; + + expect(received).toEqual(["update", "other"]); + }); +}); + +describe("proxy agent-side interception", () => { + it("intercepts agent-initiated requests with onRequestFromAgent", async () => { + const clientSaw = vi.fn(); + const { clientStream, agentStream } = setupProxy((p) => { + p.onRequestFromAgent( + "confirm", + (params) => params as { action: string }, + // Answer in the proxy; the client must never be consulted. + ({ params }) => ({ approved: params.action === "safe" }), + ); + }); + Connection.builder() + .onReceiveMessage((message) => { + clientSaw(message.method); + return Handled.no(message); + }) + .connect(clientStream); + const agentEnd = new Connection(agentStream, []); + + await expect( + agentEnd.sendRequest("confirm", { action: "safe" }), + ).resolves.toEqual({ + approved: true, + }); + await expect( + agentEnd.sendRequest("confirm", { action: "risky" }), + ).resolves.toEqual({ approved: false }); + + expect(clientSaw).not.toHaveBeenCalled(); + }); + + it("parses params with an object-form (schema-style) parser", async () => { + const parsed: unknown[] = []; + const { clientStream, agentStream } = setupProxy((p) => { + p.onRequestFromClient( + "strict", + // Object form, like a zod schema: { parse(input) { ... } }. + { + parse: (input: unknown) => { + const record = input as { n: number }; + return { n: record.n * 10 }; + }, + }, + async ({ params, forward }) => { + parsed.push(params); + return forward(params); + }, + ); + }); + Connection.builder() + .onReceiveRequest( + "strict", + (params) => params, + (request, responder) => responder.respond({ echoed: request }), + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + const response = await clientEnd.sendRequest("strict", { n: 4 }); + + expect(parsed).toEqual([{ n: 40 }]); + expect(response).toEqual({ echoed: { n: 40 } }); + }); +}); + +describe("proxy composition", () => { + it("chains two proxies, each rewriting in flight", async () => { + const [clientStream, firstClientSide] = inMemoryStreamPair(); + const [firstAgentSide, secondClientSide] = inMemoryStreamPair(); + const [secondAgentSide, agentStream] = inMemoryStreamPair(); + proxy() + .onRequestFromClient( + "echo", + (params) => params as Record, + async ({ params, forward }) => forward({ ...params, hop1: true }), + ) + .connect({ client: firstClientSide, agent: firstAgentSide }); + proxy() + .onRequestFromClient( + "echo", + (params) => params as Record, + async ({ params, forward }) => forward({ ...params, hop2: true }), + ) + .connect({ client: secondClientSide, agent: secondAgentSide }); + Connection.builder() + .onReceiveRequest( + "echo", + (params) => params, + (request, responder) => responder.respond({ echoed: request }), + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + const response = await clientEnd.sendRequest("echo", { original: true }); + + expect(response).toEqual({ + echoed: { original: true, hop1: true, hop2: true }, + }); + }); + + it("connects one configured builder to multiple stream pairs independently", async () => { + const builder = proxy().onRequestFromClient( + "echo", + (params) => params as Record, + async ({ params, forward }) => forward({ ...params, stamped: true }), + ); + + for (const label of ["a", "b"]) { + const [clientStream, proxyClientSide] = inMemoryStreamPair(); + const [proxyAgentSide, agentStream] = inMemoryStreamPair(); + builder.connect({ client: proxyClientSide, agent: proxyAgentSide }); + Connection.builder() + .onReceiveRequest( + "echo", + (params) => params, + (request, responder) => responder.respond({ echoed: request }), + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + await expect( + clientEnd.sendRequest("echo", { via: label }), + ).resolves.toEqual({ echoed: { via: label, stamped: true } }); + } + }); +}); + +describe("proxy typed registration", () => { + it("rewrites request params before forwarding", async () => { + const { clientStream, agentStream } = setupProxy((p) => { + p.onRequestFromClient( + "echo", + (params) => params as Record, + async ({ forward }) => forward({ rewritten: true }), + ); + }); + Connection.builder() + .onReceiveRequest( + "echo", + (params) => params, + (request, responder) => responder.respond({ echoed: request }), + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + const response = await clientEnd.sendRequest("echo", { original: true }); + + expect(response).toEqual({ echoed: { rewritten: true } }); + }); + + it("rewrites responses on the way back", async () => { + const { clientStream, agentStream } = setupProxy((p) => { + p.onRequestFromClient( + "echo", + (params) => params, + async ({ params, forward }) => { + const response = (await forward(params)) as Record; + return { ...response, stamped: true }; + }, + ); + }); + Connection.builder() + .onReceiveRequest( + "echo", + (params) => params, + (request, responder) => responder.respond({ echoed: request }), + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + const response = await clientEnd.sendRequest("echo", { value: 1 }); + + expect(response).toEqual({ echoed: { value: 1 }, stamped: true }); + }); + + it("answers intercepted requests without the agent seeing them", async () => { + const agentSaw = vi.fn(); + const { clientStream, agentStream } = setupProxy((p) => { + p.onRequestFromClient( + "denied", + (params) => params, + () => { + throw new RequestError(-32001, "not allowed"); + }, + ); + }); + Connection.builder() + .onReceiveMessage((message) => { + agentSaw(message.method); + return Handled.no(message); + }) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + const denied = clientEnd.sendRequest("denied", {}); + await expect(denied).rejects.toMatchObject({ code: -32001 }); + + // A permitted request afterwards proves the denied one never crossed. + await clientEnd.sendRequest("allowed", {}).catch(() => {}); + expect(agentSaw).toHaveBeenCalledWith("allowed"); + expect(agentSaw).not.toHaveBeenCalledWith("denied"); + }); + + it("answers without forwarding when the handler returns its own response", async () => { + const { clientStream } = setupProxy((p) => { + p.onRequestFromClient( + "cached", + (params) => params, + () => ({ fromProxy: true }), + ); + }); + const clientEnd = new Connection(clientStream, []); + + const response = await clientEnd.sendRequest("cached", {}); + + expect(response).toEqual({ fromProxy: true }); + }); + + it("drops notifications when the handler skips forward", async () => { + const { clientStream, agentStream } = setupProxy((p) => { + p.onNotificationFromAgent( + "update", + (params) => params as { secret: boolean }, + async ({ params, forward }) => { + if (!params.secret) { + await forward(params); + } + }, + ); + }); + const received = vi.fn(); + const gotPublic = Promise.withResolvers(); + Connection.builder() + .onReceiveNotification( + "update", + (params) => params, + (notification) => { + received(notification); + gotPublic.resolve(); + }, + ) + .connect(clientStream); + const agentEnd = new Connection(agentStream, []); + + await agentEnd.sendNotification("update", { secret: true }); + await agentEnd.sendNotification("update", { secret: false }); + await gotPublic.promise; + + expect(received).toHaveBeenCalledTimes(1); + expect(received).toHaveBeenCalledWith({ secret: false }); + }); + + it("routes unclaimed traffic to '*' with exact registrations winning", async () => { + const wildcardSaw: string[] = []; + const { clientStream, agentStream } = setupProxy((p) => { + p.onRequestFromClient( + "specific", + (params) => params, + () => ({ via: "specific" }), + ).onRequestFromClient("*", ({ method, params, forward }) => { + wildcardSaw.push(method); + return forward(params); + }); + }); + Connection.builder() + .onReceiveRequest( + "other", + (params) => params, + (_request, responder) => responder.respond({ via: "agent" }), + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + await expect(clientEnd.sendRequest("specific", {})).resolves.toEqual({ + via: "specific", + }); + await expect(clientEnd.sendRequest("other", {})).resolves.toEqual({ + via: "agent", + }); + + expect(wildcardSaw).toEqual(["other"]); + }); + + it("snapshots registrations at connect", async () => { + const { clientStream, agentStream, builder } = setupProxy(); + Connection.builder() + .onReceiveRequest( + "echo", + (params) => params, + (request, responder) => responder.respond({ echoed: request }), + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + // Registering after connect is allowed but must not affect the + // already-connected proxy — same semantics as the fluent app builders. + builder.onRequestFromClient( + "echo", + (params) => params as Record, + async ({ forward }) => forward({ late: true }), + ); + + await expect(clientEnd.sendRequest("echo", { n: 1 })).resolves.toEqual({ + echoed: { n: 1 }, + }); + }); + + it("rejects duplicate registrations for the same method", () => { + const p = proxy(); + p.onRequestFromClient( + "echo", + (params) => params, + async ({ params, forward }) => forward(params), + ); + + expect(() => + p.onRequestFromClient( + "echo", + (params) => params, + async ({ params, forward }) => forward(params), + ), + ).toThrow("already registered"); + }); + + it("routes unclaimed notifications to '*'", async () => { + const wildcardSaw: string[] = []; + const { clientStream, agentStream } = setupProxy((p) => { + p.onNotificationFromClient("*", async ({ method, params, forward }) => { + wildcardSaw.push(method); + await forward(params); + }); + }); + const seen = Promise.withResolvers(); + Connection.builder() + .onReceiveNotification( + "log", + (params) => params, + () => seen.resolve(), + ) + .connect(agentStream); + const clientEnd = new Connection(clientStream, []); + + await clientEnd.sendNotification("log", { line: "hi" }); + await seen.promise; + + expect(wildcardSaw).toEqual(["log"]); + }); +}); + +describe("proxy with fluent apps", () => { + it("relays a full initialize/new/prompt flow with a typed interceptor", async () => { + const promptTexts: string[] = []; + const { clientStream, agentStream } = setupProxy((p) => { + p.onRequestFromClient("session/prompt", async ({ params, forward }) => { + // Typed literal: params is PromptRequest, response is PromptResponse. + for (const block of params.prompt) { + if (block.type === "text") { + promptTexts.push(block.text); + } + } + const response = await forward(params); + return { ...response, _meta: { ...response._meta, proxied: true } }; + }); + }); + + agent({ name: "e2e-agent" }) + .onRequest("initialize", ({ params }) => ({ + protocolVersion: params.protocolVersion, + agentCapabilities: {}, + })) + .onRequest("session/new", () => ({ sessionId: "s-1" })) + .onRequest("session/prompt", ({ params }) => ({ + stopReason: "end_turn" as const, + _meta: { sawSession: params.sessionId }, + })) + .connect(agentStream); + + const connection = client({ name: "e2e-client" }).connect(clientStream); + try { + const init = await connection.agent.request("initialize", { + protocolVersion: 1, + clientCapabilities: {}, + }); + expect(init.protocolVersion).toBe(1); + + const session = await connection.agent.request("session/new", { + cwd: "/tmp", + mcpServers: [], + }); + expect(session.sessionId).toBe("s-1"); + + const result = await connection.agent.request("session/prompt", { + sessionId: session.sessionId, + prompt: [{ type: "text", text: "hello through the proxy" }], + }); + + expect(result.stopReason).toBe("end_turn"); + expect(result._meta).toEqual({ sawSession: "s-1", proxied: true }); + expect(promptTexts).toEqual(["hello through the proxy"]); + } finally { + connection.close(); + } + }); + + it("connects a client app to an agent app through the proxy", async () => { + const promptsSeen: string[] = []; + const { clientStream, agentStream } = setupProxy((p) => { + p.onRequestFromClient( + "_test/echo", + (params) => params as { value: number }, + async ({ params, forward }) => { + promptsSeen.push(`echo:${params.value}`); + return forward(params); + }, + ); + }); + agent({ name: "test-agent" }) + .onRequest( + "_test/echo", + (params) => params as { value: number }, + ({ params }) => ({ doubled: params.value * 2 }), + ) + .connect(agentStream); + + const connection = client({ name: "test-client" }).connect(clientStream); + try { + const response = await connection.agent.request<{ doubled: number }>( + "_test/echo", + { value: 21 }, + ); + + expect(response).toEqual({ doubled: 42 }); + expect(promptsSeen).toEqual(["echo:21"]); + } finally { + connection.close(); + } + }); +}); diff --git a/src/proxy.ts b/src/proxy.ts new file mode 100644 index 0000000..76a3ead --- /dev/null +++ b/src/proxy.ts @@ -0,0 +1,596 @@ +import { Connection, Handled, errorToRequestResult } from "./jsonrpc.js"; +import type { + HandleResult, + IncomingMessage, + IncomingNotification, + IncomingRequest, + JsonRpcHandler, + MaybePromise, +} from "./jsonrpc.js"; +import type { Stream } from "./stream.js"; +import type { + AgentNotificationMethod, + AgentNotificationParamsByMethod, + AgentRequestMethod, + AgentRequestParamsByMethod, + AgentRequestResponsesByMethod, + ClientNotificationMethod, + ClientNotificationParamsByMethod, + ClientRequestMethod, + ClientRequestParamsByMethod, + ClientRequestResponsesByMethod, + ParamsParser, +} from "./acp.js"; + +/** + * Context passed to a typed proxy request handler. + */ +export type ProxyRequestContext = { + /** + * Method name of the intercepted request. + */ + method: string; + /** + * Request params as received on the wire. + * + * Built-in method literals type the params for convenience, but the proxy + * does not validate or normalize them — schema validation happens at the + * destination app, and a proxy that re-parsed params would strip unknown + * extension fields from traffic it relays. Register with a parser when you + * want runtime validation. + */ + params: Params; + /** + * Aborts when the caller cancels this request or the connection closes. + */ + signal: AbortSignal; + /** + * Forwards the request to the other side and resolves with its response. + * + * Pass the received params to forward unchanged, or a modified copy to + * rewrite the request in flight. Cancellation by the original caller is + * propagated automatically. + */ + forward(params: Params): Promise; +}; + +/** + * Typed proxy request handler. + * + * The caller is waiting on a response, so the handler must produce one: + * return `forward(params)`'s result (unchanged or modified) to relay the + * request, return your own response without calling `forward` to answer in + * the proxy, or throw a `RequestError` to reject with a specific JSON-RPC + * error. Unlike notifications, a request can never be silently dropped. + */ +export type ProxyRequestHandler = ( + context: ProxyRequestContext, +) => MaybePromise; + +/** + * Context passed to a typed proxy notification handler. + */ +export type ProxyNotificationContext = { + /** + * Method name of the intercepted notification. + */ + method: string; + /** + * Notification params as received on the wire (see + * {@link ProxyRequestContext.params} on validation). + */ + params: Params; + /** + * Forwards the notification to the other side. + */ + forward(params: Params): Promise; +}; + +/** + * Typed proxy notification handler. + * + * Call `forward(params)` to deliver the notification to the other side, + * unchanged or rewritten. Returning without calling `forward` drops the + * notification: it is never delivered, and — as with any JSON-RPC + * notification — neither side is told, which makes skipping `forward` the + * intentional way to filter traffic. + */ +export type ProxyNotificationHandler = ( + context: ProxyNotificationContext, +) => MaybePromise; + +/** + * Handler with its context type erased for storage; dispatch restores the + * matching context shape per registration kind. + */ +type ErasedHandler = (context: never) => MaybePromise; + +type Registration = { + parse?: (params: unknown) => unknown; + handler: ErasedHandler; +}; + +/** + * Streams accepted by `ProxyBuilder.connect(...)`. + */ +export type ProxyStreams = { + /** + * Stream connected to the client side. The proxy presents as an agent on + * this stream. + */ + client: Stream; + /** + * Stream connected to the agent side. The proxy presents as a client on + * this stream. + */ + agent: Stream; +}; + +/** + * One side of a running proxy. Matches the shape of the fluent + * `AcpConnection` interface. + */ +export type ProxySideConnection = { + /** + * Aborts when this side's connection closes; `signal.reason` carries the + * close reason. + */ + readonly signal: AbortSignal; + /** + * Promise that resolves when this side's connection closes. + */ + readonly closed: Promise; + /** + * Closes this side. The other side is closed with the same reason. + */ + close(error?: unknown): void; +}; + +/** + * Handle to a running proxy returned by `ProxyBuilder.connect(...)`. + */ +export type ProxyHandle = { + /** + * The side facing the client stream. + */ + readonly client: ProxySideConnection; + /** + * The side facing the agent stream. + */ + readonly agent: ProxySideConnection; + /** + * Promise that resolves once both sides of the proxy have closed. Either + * side closing (for example the client stream ending) closes the other, + * propagating the close reason to its pending requests. + */ + readonly closed: Promise; + /** + * Closes both sides of the proxy and rejects pending requests. + */ + close(error?: unknown): void; +}; + +/** + * Builder for an ACP proxy between a client stream and an agent stream. + * + * The four registration methods name the traffic they intercept: requests + * and notifications arriving **from the client** (agent-bound methods such + * as `session/prompt`) or **from the agent** (client-bound methods such as + * `session/request_permission`). Built-in method literals infer their + * params and response types from the same ByMethod maps the + * `agent(...)`/`client(...)` builders use. + * + * Handlers claim their method: at most one runs per message. Register `"*"` + * to catch traffic no exact registration claims (most-specific wins); + * anything still unclaimed is forwarded untouched. + * + * Each `connect(...)` snapshots the registrations, exactly like the + * `agent(...)`/`client(...)` builders: registering afterwards is allowed but + * applies only to subsequent connects, never to already-connected proxies. + * For behavior that changes mid-session, keep the changing state in your + * handler's closure instead of changing the registrations. + */ +export class ProxyBuilder { + private readonly clientRequests = new Map(); + private readonly clientNotifications = new Map(); + private readonly agentRequests = new Map(); + private readonly agentNotifications = new Map(); + + /** + * Registers a typed handler for requests arriving from the client. + * + * Built-in method literals infer their params and response types from + * `method`. Pass a parser as the second argument for custom extension + * methods or to opt into runtime validation. Register `"*"` to catch any + * request no exact registration claims. + */ + onRequestFromClient( + method: Method, + handler: ProxyRequestHandler< + AgentRequestParamsByMethod[Method], + AgentRequestResponsesByMethod[Method] + >, + ): this; + onRequestFromClient( + method: "*", + handler: ProxyRequestHandler, + ): this; + onRequestFromClient( + method: string, + params: ParamsParser, + handler: ProxyRequestHandler, + ): this; + onRequestFromClient( + method: string, + handlerOrParams: ErasedHandler | ParamsParser, + handler?: ErasedHandler, + ): this { + register(this.clientRequests, method, handlerOrParams, handler); + return this; + } + + /** + * Registers a typed handler for notifications arriving from the client. + * Same registration forms as `onRequestFromClient`. + */ + onNotificationFromClient( + method: Method, + handler: ProxyNotificationHandler, + ): this; + onNotificationFromClient( + method: "*", + handler: ProxyNotificationHandler, + ): this; + onNotificationFromClient( + method: string, + params: ParamsParser, + handler: ProxyNotificationHandler, + ): this; + onNotificationFromClient( + method: string, + handlerOrParams: ErasedHandler | ParamsParser, + handler?: ErasedHandler, + ): this { + register(this.clientNotifications, method, handlerOrParams, handler); + return this; + } + + /** + * Registers a typed handler for requests arriving from the agent. + * Same registration forms as `onRequestFromClient`. + */ + onRequestFromAgent( + method: Method, + handler: ProxyRequestHandler< + ClientRequestParamsByMethod[Method], + ClientRequestResponsesByMethod[Method] + >, + ): this; + onRequestFromAgent( + method: "*", + handler: ProxyRequestHandler, + ): this; + onRequestFromAgent( + method: string, + params: ParamsParser, + handler: ProxyRequestHandler, + ): this; + onRequestFromAgent( + method: string, + handlerOrParams: ErasedHandler | ParamsParser, + handler?: ErasedHandler, + ): this { + register(this.agentRequests, method, handlerOrParams, handler); + return this; + } + + /** + * Registers a typed handler for notifications arriving from the agent. + * Same registration forms as `onRequestFromClient`. + */ + onNotificationFromAgent( + method: Method, + handler: ProxyNotificationHandler, + ): this; + onNotificationFromAgent( + method: "*", + handler: ProxyNotificationHandler, + ): this; + onNotificationFromAgent( + method: string, + params: ParamsParser, + handler: ProxyNotificationHandler, + ): this; + onNotificationFromAgent( + method: string, + handlerOrParams: ErasedHandler | ParamsParser, + handler?: ErasedHandler, + ): this { + register(this.agentNotifications, method, handlerOrParams, handler); + return this; + } + + /** + * Connects the proxy between the two streams and starts relaying. + */ + connect(streams: ProxyStreams): ProxyHandle { + // Snapshot so registrations made after connect(...) apply only to + // subsequent connects — the same semantics as the fluent app builders. + const client: Connection = new Connection(streams.client, [ + serialize( + dispatcher( + new Map(this.clientRequests), + new Map(this.clientNotifications), + () => agent, + ), + ), + ]); + const agent: Connection = new Connection(streams.agent, [ + serialize( + dispatcher( + new Map(this.agentRequests), + new Map(this.agentNotifications), + () => client, + ), + ), + ]); + // When either side closes, close the other with the same reason so its + // pending requests reject with the true cause. + void client.closed.then(() => agent.close(client.signal.reason)); + void agent.closed.then(() => client.close(agent.signal.reason)); + + return { + client, + agent, + closed: Promise.all([client.closed, agent.closed]).then(() => {}), + close(error?: unknown): void { + client.close(error); + agent.close(error); + }, + }; + } +} + +function register( + table: Map, + method: string, + handlerOrParams: ErasedHandler | ParamsParser, + handler?: ErasedHandler, +): void { + if (table.has(method)) { + throw new Error(`Proxy handler already registered: ${method}`); + } + + if (handler) { + const parser = handlerOrParams as ParamsParser; + table.set(method, { + parse: typeof parser === "function" ? parser : (p) => parser.parse(p), + handler, + }); + return; + } + + table.set(method, { handler: handlerOrParams as ErasedHandler }); +} + +/** + * Creates an ACP proxy builder. + * + * A proxy sits between a client and an agent, intercepting messages in both + * directions. The proxy terminates the protocol on both sides: each side is + * a full JSON-RPC connection, so forwarded requests are re-issued with the + * proxy's own request ids and their responses are correlated back to the + * original caller automatically. Client-initiated requests such as + * `session/prompt` flow toward the agent, and agent-initiated requests such + * as `session/request_permission` flow toward the client. The design + * matches the `Proxy` role in ACP's other SDKs. + * + * Messages that no handler claims are forwarded untouched, preserving the + * observable protocol behavior of a direct connection: + * + * - Error responses pass through with their original code, message, and data. + * - `$/cancel_request` from the original caller is propagated to the side + * handling the request, and the eventual response (which may be a normal + * result) still settles the original request. + * - When either side closes, the other side is closed and its pending + * requests are rejected. + * + * Each side processes messages one at a time in arrival order: the next + * message is not dispatched until the previous handler completes. Request + * handlers release the loop as soon as they forward (or settle without + * forwarding) rather than holding it across the round trip, so a pending + * request never blocks later messages such as `session/cancel`. + * + * Proxies chain by connecting one proxy's agent stream to the next proxy's + * client stream. + * + * @example + * ```ts + * const handle = proxy() + * .onRequestFromClient("session/prompt", async ({ params, forward }) => { + * audit(params); + * return forward(params); + * }) + * .onNotificationFromAgent("session/update", async ({ params, forward }) => { + * if (!redacted(params)) await forward(params); + * }) + * .connect({ client: clientStream, agent: agentStream }); + * ``` + */ +export function proxy(): ProxyBuilder { + return new ProxyBuilder(); +} + +/** + * Builds the dispatch function for one proxy side. + * + * Each message runs the most specific registration — exact method, then + * `"*"` — or, unclaimed, is forwarded untouched on a fully synchronous fast + * path (the send is enqueued in arrival order and the loop is never held). + * + * Request registrations run detached from the dispatch loop once they + * forward or settle: the loop resumes when the request has been sent (or + * answered), not when the handler finishes waiting on the round trip, so a + * pending request never blocks later messages. Notification registrations + * hold the loop until the handler settles, which is what preserves + * delivery ordering across async handlers. + */ +function dispatcher( + requests: Map, + notifications: Map, + target: () => Connection, +): (message: IncomingMessage) => MaybePromise { + const runRequest = ( + message: IncomingRequest, + registration: Registration, + ): MaybePromise => { + const { responder } = message; + let released = false; + let resolveReleased: (() => void) | undefined; + const release = () => { + released = true; + resolveReleased?.(); + }; + const run = registration.handler as ( + context: ProxyRequestContext, + ) => MaybePromise; + + void (async () => { + try { + const response = await run({ + method: message.method, + params: registration.parse + ? registration.parse(message.params) + : message.params, + signal: message.signal, + forward: (params: unknown) => { + const sent = target().sendRequest( + message.method, + params, + undefined, + { cancellationSignal: message.signal }, + ); + release(); + return sent; + }, + }); + await responder.respond(response ?? null); + } catch (error) { + await responder + .respondWithResult(errorToRequestResult(error, message.signal)) + .catch(() => {}); + } finally { + release(); + } + })(); + + // A handler that forwards before its first await has already released; + // skip the promise entirely so the loop continues on the same tick. + if (released) { + return Handled.yes(); + } + + return new Promise((resolve) => { + resolveReleased = () => resolve(Handled.yes()); + }); + }; + + const runNotification = async ( + message: IncomingNotification, + registration: Registration, + ): Promise => { + const run = registration.handler as ( + context: ProxyNotificationContext, + ) => MaybePromise; + await run({ + method: message.method, + params: registration.parse + ? registration.parse(message.params) + : message.params, + forward: (params: unknown) => + target().sendNotification(message.method, params), + }); + return Handled.yes(); + }; + + return (message) => { + const table = message.kind === "request" ? requests : notifications; + // Most-specific wins: an exact method registration beats "*", and "*" + // catches only otherwise-unclaimed traffic. + const registration = table.get(message.method) ?? table.get("*"); + + if (!registration) { + // Pass-through: the send is enqueued synchronously (so send order is + // arrival order) and the loop is never held. + if (message.kind === "request") { + const { responder } = message; + target() + .sendRequest(message.method, message.params, undefined, { + cancellationSignal: message.signal, + }) + .then( + (result) => responder.respond(result), + (error) => + responder.respondWithResult( + errorToRequestResult(error, message.signal), + ), + ) + // The response cannot be delivered when the caller's side is + // already closed; there is nowhere left to report it. + .catch(() => {}); + } else { + // Write failures close the connection via the write queue; there is + // no per-notification error to report. + void target() + .sendNotification(message.method, message.params) + .catch(() => {}); + } + + return Handled.yes(); + } + + return message.kind === "request" + ? runRequest(message, registration) + : runNotification(message, registration); + }; +} + +/** + * Serializes one side's dispatch: messages run one at a time, in arrival + * order. The underlying `Connection` dispatches each message as its own + * async task, which would let slow handlers be overtaken. Synchronous + * dispatches (pass-through traffic, handlers that forward immediately) + * bypass the queue entirely; a rejected dispatch fails only its own + * message. + */ +function serialize( + dispatch: (message: IncomingMessage) => MaybePromise, +): JsonRpcHandler { + let pending = 0; + let tail: Promise = Promise.resolve(); + + const track = (result: Promise): Promise => { + pending++; + tail = result.then( + () => { + pending--; + }, + () => { + pending--; + }, + ); + return result; + }; + + return { + handleMessage(message) { + if (pending === 0) { + const result = dispatch(message); + return result instanceof Promise ? track(result) : result; + } + + return track(tail.then(() => dispatch(message))); + }, + describe: () => "proxy:dispatch", + }; +}