feat: add MCP bridge exposing pi tools to the CLI
A localhost-only MCP HTTP server, guarded by a per-turn bearer token, lists pi's tools and holds each call until pi has run it and returned the result.
This commit is contained in:
@@ -0,0 +1,193 @@
|
|||||||
|
import { randomBytes, randomUUID, timingSafeEqual } from "node:crypto";
|
||||||
|
import * as http from "node:http";
|
||||||
|
import type { AddressInfo } from "node:net";
|
||||||
|
|
||||||
|
import type { ImageContent, JsonObject, TextContent, Tool, ToolResultMessage } from "@earendil-works/pi-ai";
|
||||||
|
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
|
||||||
|
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
|
||||||
|
import {
|
||||||
|
CallToolRequestSchema,
|
||||||
|
ListToolsRequestSchema,
|
||||||
|
type CallToolResult,
|
||||||
|
type Tool as McpTool,
|
||||||
|
} from "@modelcontextprotocol/sdk/types.js";
|
||||||
|
|
||||||
|
export const BRIDGE_SERVER_NAME = "learnharness";
|
||||||
|
|
||||||
|
const MCP_TOOL_PREFIX = `mcp__${BRIDGE_SERVER_NAME}__`;
|
||||||
|
const LOOPBACK_HOST = "127.0.0.1";
|
||||||
|
const TOKEN_BYTES = 32;
|
||||||
|
|
||||||
|
export type PendingToolCall = { name: string; arguments: JsonObject };
|
||||||
|
|
||||||
|
export type ToolBridge = {
|
||||||
|
mcpConfig: string;
|
||||||
|
allowedTools: string[];
|
||||||
|
waitForToolCalls(count: number): Promise<PendingToolCall[]>;
|
||||||
|
resolveToolCall(call: PendingToolCall, result: ToolResultMessage): void;
|
||||||
|
getPendingToolCalls(): PendingToolCall[];
|
||||||
|
close(): Promise<void>;
|
||||||
|
};
|
||||||
|
|
||||||
|
type ToolCallWaiter = {
|
||||||
|
count: number;
|
||||||
|
resolve: (calls: PendingToolCall[]) => void;
|
||||||
|
reject: (error: Error) => void;
|
||||||
|
};
|
||||||
|
|
||||||
|
export function toolNameFromMcp(name: string): string {
|
||||||
|
return name.startsWith(MCP_TOOL_PREFIX) ? name.slice(MCP_TOOL_PREFIX.length) : name;
|
||||||
|
}
|
||||||
|
|
||||||
|
function toMcpContent(block: TextContent | ImageContent): CallToolResult["content"][number] {
|
||||||
|
switch (block.type) {
|
||||||
|
case "text":
|
||||||
|
return { type: "text", text: block.text };
|
||||||
|
|
||||||
|
case "image":
|
||||||
|
return { type: "image", data: block.data, mimeType: block.mimeType };
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function toCallToolResult(result: ToolResultMessage): CallToolResult {
|
||||||
|
return { content: result.content.map(toMcpContent), isError: result.isError };
|
||||||
|
}
|
||||||
|
|
||||||
|
function toMcpTool(tool: Tool): McpTool {
|
||||||
|
return { name: tool.name, description: tool.description, inputSchema: tool.parameters as McpTool["inputSchema"] };
|
||||||
|
}
|
||||||
|
|
||||||
|
class ToolCallQueue {
|
||||||
|
private readonly results = new Map<PendingToolCall, (result: CallToolResult) => void>();
|
||||||
|
private waiters: ToolCallWaiter[] = [];
|
||||||
|
|
||||||
|
add(call: PendingToolCall): Promise<CallToolResult> {
|
||||||
|
const result = new Promise<CallToolResult>((resolve) => this.results.set(call, resolve));
|
||||||
|
|
||||||
|
this.wakeWaiters();
|
||||||
|
|
||||||
|
return result;
|
||||||
|
}
|
||||||
|
|
||||||
|
waitFor(count: number): Promise<PendingToolCall[]> {
|
||||||
|
return new Promise((resolve, reject) => {
|
||||||
|
this.waiters.push({ count, resolve, reject });
|
||||||
|
this.wakeWaiters();
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
resolve(call: PendingToolCall, result: ToolResultMessage): void {
|
||||||
|
this.results.get(call)?.(toCallToolResult(result));
|
||||||
|
this.results.delete(call);
|
||||||
|
}
|
||||||
|
|
||||||
|
list(): PendingToolCall[] {
|
||||||
|
return [...this.results.keys()];
|
||||||
|
}
|
||||||
|
|
||||||
|
clear(): void {
|
||||||
|
this.results.clear();
|
||||||
|
|
||||||
|
for (const waiter of this.waiters) waiter.reject(new Error("claude-code tool bridge closed"));
|
||||||
|
this.waiters = [];
|
||||||
|
}
|
||||||
|
|
||||||
|
private wakeWaiters(): void {
|
||||||
|
const calls = this.list();
|
||||||
|
|
||||||
|
const ready = this.waiters.filter((waiter) => calls.length >= waiter.count);
|
||||||
|
this.waiters = this.waiters.filter((waiter) => calls.length < waiter.count);
|
||||||
|
|
||||||
|
for (const waiter of ready) waiter.resolve(calls);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Low-level Server on purpose: McpServer.registerTool only takes zod schemas, and pi tools carry
|
||||||
|
// JSON Schema that must reach the CLI unchanged.
|
||||||
|
function createMcpServer(tools: Tool[], queue: ToolCallQueue): Server {
|
||||||
|
const server = new Server({ name: BRIDGE_SERVER_NAME, version: "1.0.0" }, { capabilities: { tools: {} } });
|
||||||
|
|
||||||
|
server.setRequestHandler(ListToolsRequestSchema, () => ({ tools: tools.map(toMcpTool) }));
|
||||||
|
|
||||||
|
server.setRequestHandler(CallToolRequestSchema, (request) => {
|
||||||
|
const call = { name: request.params.name, arguments: (request.params.arguments ?? {}) as JsonObject };
|
||||||
|
return queue.add(call);
|
||||||
|
});
|
||||||
|
|
||||||
|
return server;
|
||||||
|
}
|
||||||
|
|
||||||
|
function isAuthorized(header: string | undefined, expected: Buffer): boolean {
|
||||||
|
if (header === undefined) return false;
|
||||||
|
|
||||||
|
const actual = Buffer.from(header);
|
||||||
|
return actual.length === expected.length && timingSafeEqual(actual, expected);
|
||||||
|
}
|
||||||
|
|
||||||
|
function createHttpServer(transport: StreamableHTTPServerTransport, authorization: string): http.Server {
|
||||||
|
const expected = Buffer.from(authorization);
|
||||||
|
|
||||||
|
return http.createServer((req, res) => {
|
||||||
|
if (!isAuthorized(req.headers.authorization, expected)) {
|
||||||
|
res.writeHead(401).end();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
transport.handleRequest(req, res).catch((error: unknown) => {
|
||||||
|
console.error("claude-code tool bridge request failed", error);
|
||||||
|
|
||||||
|
if (!res.headersSent) res.writeHead(500);
|
||||||
|
res.end();
|
||||||
|
});
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
function listen(server: http.Server): Promise<AddressInfo> {
|
||||||
|
return new Promise((resolve, reject) => {
|
||||||
|
server.once("error", reject);
|
||||||
|
|
||||||
|
server.listen(0, LOOPBACK_HOST, () => resolve(server.address() as AddressInfo));
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
function closeHttpServer(server: http.Server): Promise<void> {
|
||||||
|
server.closeAllConnections();
|
||||||
|
|
||||||
|
return new Promise((resolve, reject) => {
|
||||||
|
server.close((error) => {
|
||||||
|
if (error) reject(error);
|
||||||
|
else resolve();
|
||||||
|
});
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
function toMcpConfig(port: number, authorization: string): string {
|
||||||
|
const server = { type: "http", url: `http://${LOOPBACK_HOST}:${port}/mcp`, headers: { Authorization: authorization } };
|
||||||
|
|
||||||
|
return JSON.stringify({ mcpServers: { [BRIDGE_SERVER_NAME]: server } });
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function createToolBridge(tools: Tool[]): Promise<ToolBridge> {
|
||||||
|
const authorization = `Bearer ${randomBytes(TOKEN_BYTES).toString("hex")}`;
|
||||||
|
|
||||||
|
const queue = new ToolCallQueue();
|
||||||
|
const mcpServer = createMcpServer(tools, queue);
|
||||||
|
const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => randomUUID() });
|
||||||
|
await mcpServer.connect(transport);
|
||||||
|
|
||||||
|
const httpServer = createHttpServer(transport, authorization);
|
||||||
|
const { port } = await listen(httpServer);
|
||||||
|
|
||||||
|
return {
|
||||||
|
mcpConfig: toMcpConfig(port, authorization),
|
||||||
|
allowedTools: tools.map((tool) => `${MCP_TOOL_PREFIX}${tool.name}`),
|
||||||
|
waitForToolCalls: (count) => queue.waitFor(count),
|
||||||
|
resolveToolCall: (call, result) => queue.resolve(call, result),
|
||||||
|
getPendingToolCalls: () => queue.list(),
|
||||||
|
close: async () => {
|
||||||
|
queue.clear();
|
||||||
|
await mcpServer.close();
|
||||||
|
await closeHttpServer(httpServer);
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
@@ -0,0 +1,225 @@
|
|||||||
|
/**
|
||||||
|
* Claude Code MCP tool bridge (src/agent/claude-code/tool-bridge.ts).
|
||||||
|
*
|
||||||
|
* A real MCP SDK client talks to the bridge over loopback HTTP, the same way
|
||||||
|
* the claude CLI does, so listing, auth, pending calls and their results
|
||||||
|
* are all exercised end to end.
|
||||||
|
*/
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import * as net from "node:net";
|
||||||
|
import * as os from "node:os";
|
||||||
|
import test from "node:test";
|
||||||
|
|
||||||
|
import type { Tool, ToolResultMessage } from "@earendil-works/pi-ai";
|
||||||
|
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
|
||||||
|
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
|
||||||
|
import { Type } from "typebox";
|
||||||
|
|
||||||
|
import { BRIDGE_SERVER_NAME, createToolBridge, toolNameFromMcp, type ToolBridge } from "../src/agent/claude-code/tool-bridge.ts";
|
||||||
|
|
||||||
|
const READ_TOOL: Tool = { name: "read", description: "Read a file", parameters: Type.Object({ path: Type.String() }) };
|
||||||
|
const SCREENSHOT_TOOL: Tool = { name: "screenshot", description: "Capture the screen", parameters: Type.Object({}) };
|
||||||
|
|
||||||
|
type BridgeServerConfig = { type: string; url: string; headers: { Authorization: string } };
|
||||||
|
|
||||||
|
function getServerConfig(bridge: ToolBridge): BridgeServerConfig {
|
||||||
|
const config = JSON.parse(bridge.mcpConfig) as { mcpServers: Record<string, BridgeServerConfig> };
|
||||||
|
const server = config.mcpServers[BRIDGE_SERVER_NAME];
|
||||||
|
assert.ok(server);
|
||||||
|
return server;
|
||||||
|
}
|
||||||
|
|
||||||
|
async function connectClient(bridge: ToolBridge, authorization = getServerConfig(bridge).headers.Authorization): Promise<Client> {
|
||||||
|
const transport = new StreamableHTTPClientTransport(new URL(getServerConfig(bridge).url), {
|
||||||
|
requestInit: { headers: { Authorization: authorization } },
|
||||||
|
});
|
||||||
|
const client = new Client({ name: "test", version: "0" });
|
||||||
|
await client.connect(transport);
|
||||||
|
|
||||||
|
return client;
|
||||||
|
}
|
||||||
|
|
||||||
|
function toolResult(content: ToolResultMessage["content"], isError = false): ToolResultMessage {
|
||||||
|
return { role: "toolResult", toolCallId: "call-1", toolName: "read", content, isError, timestamp: 0 };
|
||||||
|
}
|
||||||
|
|
||||||
|
test("lists exactly the given tools with their JSON schemas", async (t) => {
|
||||||
|
const bridge = await createToolBridge([READ_TOOL, SCREENSHOT_TOOL]);
|
||||||
|
t.after(() => bridge.close());
|
||||||
|
const client = await connectClient(bridge);
|
||||||
|
t.after(() => client.close());
|
||||||
|
|
||||||
|
const { tools } = await client.listTools();
|
||||||
|
|
||||||
|
assert.deepEqual(
|
||||||
|
tools.map((tool) => ({ name: tool.name, description: tool.description, inputSchema: tool.inputSchema })),
|
||||||
|
[
|
||||||
|
{ name: "read", description: "Read a file", inputSchema: JSON.parse(JSON.stringify(READ_TOOL.parameters)) },
|
||||||
|
{ name: "screenshot", description: "Capture the screen", inputSchema: JSON.parse(JSON.stringify(SCREENSHOT_TOOL.parameters)) },
|
||||||
|
],
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.deepEqual(bridge.allowedTools, ["mcp__learnharness__read", "mcp__learnharness__screenshot"]);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("rejects requests without the bearer token", async (t) => {
|
||||||
|
const bridge = await createToolBridge([READ_TOOL]);
|
||||||
|
t.after(() => bridge.close());
|
||||||
|
|
||||||
|
await assert.rejects(connectClient(bridge, "Bearer wrong"), (error: { code?: number }) => error.code === 401);
|
||||||
|
const response = await fetch(getServerConfig(bridge).url, {
|
||||||
|
method: "POST",
|
||||||
|
headers: { "Content-Type": "application/json" },
|
||||||
|
body: "{}",
|
||||||
|
});
|
||||||
|
assert.equal(response.status, 401);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("binds 127.0.0.1", async (t) => {
|
||||||
|
const bridge = await createToolBridge([READ_TOOL]);
|
||||||
|
t.after(() => bridge.close());
|
||||||
|
|
||||||
|
const url = new URL(getServerConfig(bridge).url);
|
||||||
|
assert.equal(url.hostname, "127.0.0.1");
|
||||||
|
|
||||||
|
const externalAddress = Object.values(os.networkInterfaces())
|
||||||
|
.flat()
|
||||||
|
.find((address) => address && !address.internal && address.family === "IPv4");
|
||||||
|
|
||||||
|
if (!externalAddress) {
|
||||||
|
t.diagnostic("no non-loopback IPv4 interface; skipped external connect check");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
await assert.rejects(
|
||||||
|
new Promise((resolve, reject) => {
|
||||||
|
const socket = net.connect(Number(url.port), externalAddress.address, () => resolve(socket.end()));
|
||||||
|
socket.once("error", reject);
|
||||||
|
}),
|
||||||
|
{ code: "ECONNREFUSED" },
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("tools/call stays pending until resolved and returns its text and image content", async (t) => {
|
||||||
|
const bridge = await createToolBridge([READ_TOOL]);
|
||||||
|
t.after(() => bridge.close());
|
||||||
|
const client = await connectClient(bridge);
|
||||||
|
t.after(() => client.close());
|
||||||
|
|
||||||
|
let settled = false;
|
||||||
|
const pending = client.callTool({ name: "read", arguments: { path: "notes.md" } }).finally(() => {
|
||||||
|
settled = true;
|
||||||
|
});
|
||||||
|
|
||||||
|
const [call] = await bridge.waitForToolCalls(1);
|
||||||
|
assert.ok(call);
|
||||||
|
assert.deepEqual(call, { name: "read", arguments: { path: "notes.md" } });
|
||||||
|
assert.deepEqual(bridge.getPendingToolCalls(), [call]);
|
||||||
|
|
||||||
|
await new Promise((resolve) => setTimeout(resolve, 20));
|
||||||
|
assert.equal(settled, false);
|
||||||
|
|
||||||
|
bridge.resolveToolCall(
|
||||||
|
call,
|
||||||
|
toolResult([
|
||||||
|
{ type: "text", text: "# Notes" },
|
||||||
|
{ type: "image", data: "AAAA", mimeType: "image/png" },
|
||||||
|
]),
|
||||||
|
);
|
||||||
|
|
||||||
|
const result = await pending;
|
||||||
|
assert.deepEqual(result.content, [
|
||||||
|
{ type: "text", text: "# Notes" },
|
||||||
|
{ type: "image", data: "AAAA", mimeType: "image/png" },
|
||||||
|
]);
|
||||||
|
assert.notEqual(result.isError, true);
|
||||||
|
assert.deepEqual(bridge.getPendingToolCalls(), []);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("waitForToolCalls resolves once that many calls are pending, in arrival order", async (t) => {
|
||||||
|
const bridge = await createToolBridge([READ_TOOL, SCREENSHOT_TOOL]);
|
||||||
|
t.after(() => bridge.close());
|
||||||
|
const client = await connectClient(bridge);
|
||||||
|
t.after(() => client.close());
|
||||||
|
|
||||||
|
let waited = false;
|
||||||
|
const waiting = bridge.waitForToolCalls(2).finally(() => {
|
||||||
|
waited = true;
|
||||||
|
});
|
||||||
|
|
||||||
|
void client.callTool({ name: "read", arguments: { path: "a.md" } }).catch(() => {});
|
||||||
|
await new Promise((resolve) => setTimeout(resolve, 20));
|
||||||
|
assert.equal(waited, false);
|
||||||
|
|
||||||
|
void client.callTool({ name: "screenshot", arguments: {} }).catch(() => {});
|
||||||
|
|
||||||
|
assert.deepEqual(await waiting, [
|
||||||
|
{ name: "read", arguments: { path: "a.md" } },
|
||||||
|
{ name: "screenshot", arguments: {} },
|
||||||
|
]);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("isError result becomes MCP isError true", async (t) => {
|
||||||
|
const bridge = await createToolBridge([READ_TOOL]);
|
||||||
|
t.after(() => bridge.close());
|
||||||
|
const client = await connectClient(bridge);
|
||||||
|
t.after(() => client.close());
|
||||||
|
|
||||||
|
const pending = client.callTool({ name: "read", arguments: { path: "missing.md" } });
|
||||||
|
const [call] = await bridge.waitForToolCalls(1);
|
||||||
|
assert.ok(call);
|
||||||
|
|
||||||
|
bridge.resolveToolCall(call, toolResult([{ type: "text", text: "ENOENT" }], true));
|
||||||
|
|
||||||
|
const result = await pending;
|
||||||
|
assert.equal(result.isError, true);
|
||||||
|
assert.deepEqual(result.content, [{ type: "text", text: "ENOENT" }]);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("close fails pending calls and frees the port", async () => {
|
||||||
|
const bridge = await createToolBridge([READ_TOOL]);
|
||||||
|
const port = Number(new URL(getServerConfig(bridge).url).port);
|
||||||
|
const client = await connectClient(bridge);
|
||||||
|
const pending = client.callTool({ name: "read", arguments: { path: "notes.md" } });
|
||||||
|
await bridge.waitForToolCalls(1);
|
||||||
|
|
||||||
|
await bridge.close();
|
||||||
|
|
||||||
|
await assert.rejects(pending);
|
||||||
|
assert.deepEqual(bridge.getPendingToolCalls(), []);
|
||||||
|
await client.close();
|
||||||
|
|
||||||
|
const probe = net.createServer();
|
||||||
|
await new Promise<void>((resolve, reject) => {
|
||||||
|
probe.once("error", reject);
|
||||||
|
probe.listen(port, "127.0.0.1", resolve);
|
||||||
|
});
|
||||||
|
await new Promise((resolve) => probe.close(resolve));
|
||||||
|
});
|
||||||
|
|
||||||
|
test("close rejects consumers still waiting for tool calls", { timeout: 1000 }, async () => {
|
||||||
|
const bridge = await createToolBridge([READ_TOOL]);
|
||||||
|
const waiting = bridge.waitForToolCalls(1);
|
||||||
|
|
||||||
|
await bridge.close();
|
||||||
|
await assert.rejects(waiting, /tool bridge closed/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test('mcpConfig is {"mcpServers":{"learnharness":{"type":"http","url":...,"headers":{"Authorization":"Bearer ..."}}}}', async (t) => {
|
||||||
|
const bridge = await createToolBridge([READ_TOOL]);
|
||||||
|
t.after(() => bridge.close());
|
||||||
|
|
||||||
|
const config = JSON.parse(bridge.mcpConfig) as { mcpServers: Record<string, unknown> };
|
||||||
|
assert.deepEqual(Object.keys(config.mcpServers), ["learnharness"]);
|
||||||
|
|
||||||
|
const server = getServerConfig(bridge);
|
||||||
|
assert.deepEqual(Object.keys(server).sort(), ["headers", "type", "url"]);
|
||||||
|
assert.equal(server.type, "http");
|
||||||
|
assert.match(server.url, /^http:\/\/127\.0\.0\.1:\d+\/mcp$/);
|
||||||
|
assert.match(server.headers.Authorization, /^Bearer [0-9a-f]{64}$/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("toolNameFromMcp strips the bridge prefix", () => {
|
||||||
|
assert.equal(toolNameFromMcp("mcp__learnharness__read"), "read");
|
||||||
|
assert.equal(toolNameFromMcp("read"), "read");
|
||||||
|
});
|
||||||
Reference in New Issue
Block a user