From 804481be225b47224f8e00ad8af3221f9c0bf3b2 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 20 Dec 2025 21:38:06 +0000 Subject: [PATCH 1/4] feat: implement Tool Bridge and Tool Relay for reverse RPC over WebSocket Adds a complete implementation of the Tool Bridge Protocol for enabling server-side LLM agents to invoke tools on local machines behind NAT/firewall. Packages added: - @opencode-ai/tool-bridge-protocol: Shared protocol schemas and message types - @opencode-ai/tool-relay: Server-side gateway that forwards tool invocations - @opencode-ai/tool-bridge: Client-side agent that executes tools locally - @opencode-ai/tool-bridge-opencode: Adapter for integrating opencode tools - @opencode-ai/tool-bridge-tests: End-to-end integration tests Key features: - WebSocket-based reverse RPC with multiplexed tool calls - Session-aware routing with reconnect support - Heartbeat mechanism for connection health monitoring - Tool registration and discovery - Streaming support for progressive results - Cancellation support for long-running operations - Comprehensive integration test suite --- packages/tool-bridge-opencode/package.json | 21 + packages/tool-bridge-opencode/src/adapter.ts | 173 ++++++ packages/tool-bridge-opencode/src/index.ts | 16 + packages/tool-bridge-opencode/tsconfig.json | 9 + packages/tool-bridge-protocol/package.json | 17 + packages/tool-bridge-protocol/src/index.ts | 9 + packages/tool-bridge-protocol/src/schema.ts | 210 +++++++ packages/tool-bridge-protocol/src/utils.ts | 243 ++++++++ packages/tool-bridge-protocol/tsconfig.json | 9 + .../tool-bridge-tests/integration.test.ts | 479 +++++++++++++++ packages/tool-bridge-tests/package.json | 20 + packages/tool-bridge-tests/tsconfig.json | 9 + packages/tool-bridge/package.json | 18 + packages/tool-bridge/src/bridge.ts | 506 ++++++++++++++++ packages/tool-bridge/src/index.ts | 16 + packages/tool-bridge/src/types.ts | 83 +++ packages/tool-bridge/tsconfig.json | 9 + packages/tool-relay/package.json | 19 + packages/tool-relay/src/index.ts | 16 + packages/tool-relay/src/relay.ts | 565 ++++++++++++++++++ packages/tool-relay/src/types.ts | 104 ++++ packages/tool-relay/tsconfig.json | 9 + 22 files changed, 2560 insertions(+) create mode 100644 packages/tool-bridge-opencode/package.json create mode 100644 packages/tool-bridge-opencode/src/adapter.ts create mode 100644 packages/tool-bridge-opencode/src/index.ts create mode 100644 packages/tool-bridge-opencode/tsconfig.json create mode 100644 packages/tool-bridge-protocol/package.json create mode 100644 packages/tool-bridge-protocol/src/index.ts create mode 100644 packages/tool-bridge-protocol/src/schema.ts create mode 100644 packages/tool-bridge-protocol/src/utils.ts create mode 100644 packages/tool-bridge-protocol/tsconfig.json create mode 100644 packages/tool-bridge-tests/integration.test.ts create mode 100644 packages/tool-bridge-tests/package.json create mode 100644 packages/tool-bridge-tests/tsconfig.json create mode 100644 packages/tool-bridge/package.json create mode 100644 packages/tool-bridge/src/bridge.ts create mode 100644 packages/tool-bridge/src/index.ts create mode 100644 packages/tool-bridge/src/types.ts create mode 100644 packages/tool-bridge/tsconfig.json create mode 100644 packages/tool-relay/package.json create mode 100644 packages/tool-relay/src/index.ts create mode 100644 packages/tool-relay/src/relay.ts create mode 100644 packages/tool-relay/src/types.ts create mode 100644 packages/tool-relay/tsconfig.json diff --git a/packages/tool-bridge-opencode/package.json b/packages/tool-bridge-opencode/package.json new file mode 100644 index 00000000000..014e3a5c2d7 --- /dev/null +++ b/packages/tool-bridge-opencode/package.json @@ -0,0 +1,21 @@ +{ + "name": "@opencode-ai/tool-bridge-opencode", + "version": "1.0.0", + "type": "module", + "private": true, + "exports": { + "./*": "./src/*.ts" + }, + "dependencies": { + "@opencode-ai/tool-bridge": "workspace:*", + "@opencode-ai/tool-bridge-protocol": "workspace:*", + "opencode": "workspace:*", + "zod": "catalog:", + "zod-to-json-schema": "3.24.5" + }, + "devDependencies": { + "@tsconfig/bun": "catalog:", + "@types/bun": "catalog:", + "typescript": "catalog:" + } +} diff --git a/packages/tool-bridge-opencode/src/adapter.ts b/packages/tool-bridge-opencode/src/adapter.ts new file mode 100644 index 00000000000..54a1f6ec19d --- /dev/null +++ b/packages/tool-bridge-opencode/src/adapter.ts @@ -0,0 +1,173 @@ +import { ToolBridge, type ToolContext, type ToolResult } from "@opencode-ai/tool-bridge/index" +import { zodToJsonSchema } from "zod-to-json-schema" +import type { Tool } from "opencode/tool/tool" +import { BashTool } from "opencode/tool/bash" +import { EditTool } from "opencode/tool/edit" +import { GlobTool } from "opencode/tool/glob" +import { GrepTool } from "opencode/tool/grep" +import { ListTool } from "opencode/tool/ls" +import { ReadTool } from "opencode/tool/read" +import { WriteTool } from "opencode/tool/write" +import { WebFetchTool } from "opencode/tool/webfetch" +import type z from "zod" + +/** + * Configuration for the OpenCode tool adapter + */ +export interface OpencodeAdapterConfig { + /** + * Working directory for file operations + */ + workingDirectory: string + + /** + * Session ID for tool context + */ + sessionId?: string + + /** + * Agent name for tool context + */ + agent?: string + + /** + * Additional tools to register (beyond the default set) + */ + additionalTools?: Tool.Info[] + + /** + * Tools to exclude from registration + */ + excludeTools?: string[] +} + +/** + * Default tools to expose via the bridge + */ +const DEFAULT_TOOLS: Tool.Info[] = [ + BashTool, + EditTool, + GlobTool, + GrepTool, + ListTool, + ReadTool, + WriteTool, + WebFetchTool, +] + +/** + * Convert a zod schema to JSON schema + */ +function schemaToJson(schema: z.ZodType): Record { + try { + return zodToJsonSchema(schema) as Record + } catch { + return { type: "object" } + } +} + +/** + * Create a tool executor that wraps an opencode tool + */ +function createExecutor( + toolInfo: Awaited>, + config: OpencodeAdapterConfig +): (args: Record, ctx: ToolContext) => Promise { + return async (args: Record, ctx: ToolContext): Promise => { + // Create a tool context compatible with opencode + const toolCtx: Tool.Context = { + sessionID: ctx.sessionId || config.sessionId || "bridge-session", + messageID: ctx.callId, + agent: config.agent || "bridge", + abort: ctx.abort, + callID: ctx.callId, + metadata: (input) => { + // Send metadata updates as chunks + if (input.metadata) { + ctx.sendChunk(JSON.stringify(input.metadata), { type: "metadata" }) + } + }, + } + + try { + const result = await toolInfo.execute(args, toolCtx) + return { + output: result.output, + metadata: result.metadata as Record, + title: result.title, + } + } catch (error) { + throw error + } + } +} + +/** + * Register opencode tools with a tool bridge + */ +export async function registerOpencodeTools( + bridge: ToolBridge, + config: OpencodeAdapterConfig +): Promise { + const excludeSet = new Set(config.excludeTools ?? []) + const toolsToRegister = [...DEFAULT_TOOLS, ...(config.additionalTools ?? [])] + + // Set up process.cwd() to return the working directory + const originalCwd = process.cwd + const cwd = () => config.workingDirectory + + for (const tool of toolsToRegister) { + if (excludeSet.has(tool.id)) continue + + try { + // Temporarily override cwd for tool initialization + process.cwd = cwd + + const toolInfo = await tool.init() + const jsonSchema = schemaToJson(toolInfo.parameters) + + bridge.registerTool( + tool.id, + toolInfo.description, + jsonSchema, + createExecutor(toolInfo, config) + ) + } catch (error) { + console.error(`Failed to register tool ${tool.id}:`, error) + } finally { + process.cwd = originalCwd + } + } +} + +/** + * Create a pre-configured tool bridge with opencode tools + */ +export async function createOpencodeBridge( + relayUrl: string, + config: OpencodeAdapterConfig & { + bridgeId?: string + credentials?: { type: "bearer" | "none"; token?: string } + onConnect?: (sessionId: string) => void + onDisconnect?: (reason: string) => void + onError?: (error: Error) => void + } +): Promise { + const bridge = new ToolBridge({ + url: relayUrl, + bridgeId: config.bridgeId, + credentials: config.credentials, + onConnect: config.onConnect, + onDisconnect: config.onDisconnect, + onError: config.onError, + }) + + await registerOpencodeTools(bridge, config) + + return bridge +} + +/** + * Export tool definitions for reference + */ +export const TOOL_IDS = DEFAULT_TOOLS.map((t) => t.id) diff --git a/packages/tool-bridge-opencode/src/index.ts b/packages/tool-bridge-opencode/src/index.ts new file mode 100644 index 00000000000..3c871f7ade9 --- /dev/null +++ b/packages/tool-bridge-opencode/src/index.ts @@ -0,0 +1,16 @@ +/** + * OpenCode Tool Bridge Adapter + * + * Provides integration between the Tool Bridge and OpenCode's + * built-in tools (bash, read, write, edit, glob, grep, etc.) + */ + +export { + registerOpencodeTools, + createOpencodeBridge, + TOOL_IDS, + type OpencodeAdapterConfig, +} from "./adapter" + +// Re-export bridge types for convenience +export { ToolBridge, type BridgeConfig, type ToolContext, type ToolResult } from "@opencode-ai/tool-bridge/index" diff --git a/packages/tool-bridge-opencode/tsconfig.json b/packages/tool-bridge-opencode/tsconfig.json new file mode 100644 index 00000000000..ccdf1d78395 --- /dev/null +++ b/packages/tool-bridge-opencode/tsconfig.json @@ -0,0 +1,9 @@ +{ + "$schema": "https://json.schemastore.org/tsconfig", + "extends": "@tsconfig/bun/tsconfig.json", + "compilerOptions": { + "lib": ["ESNext", "DOM"], + "types": [], + "noUncheckedIndexedAccess": false + } +} diff --git a/packages/tool-bridge-protocol/package.json b/packages/tool-bridge-protocol/package.json new file mode 100644 index 00000000000..02492fec350 --- /dev/null +++ b/packages/tool-bridge-protocol/package.json @@ -0,0 +1,17 @@ +{ + "name": "@opencode-ai/tool-bridge-protocol", + "version": "1.0.0", + "type": "module", + "private": true, + "exports": { + "./*": "./src/*.ts" + }, + "dependencies": { + "zod": "catalog:" + }, + "devDependencies": { + "@tsconfig/bun": "catalog:", + "@types/bun": "catalog:", + "typescript": "catalog:" + } +} diff --git a/packages/tool-bridge-protocol/src/index.ts b/packages/tool-bridge-protocol/src/index.ts new file mode 100644 index 00000000000..6e67ce0b281 --- /dev/null +++ b/packages/tool-bridge-protocol/src/index.ts @@ -0,0 +1,9 @@ +/** + * Tool Bridge Protocol + * + * A protocol for reverse RPC over WebSocket, enabling server-side LLM agents + * to invoke tools on client machines behind NAT/firewall. + */ + +export * from "./schema" +export * from "./utils" diff --git a/packages/tool-bridge-protocol/src/schema.ts b/packages/tool-bridge-protocol/src/schema.ts new file mode 100644 index 00000000000..4f89ea88ccc --- /dev/null +++ b/packages/tool-bridge-protocol/src/schema.ts @@ -0,0 +1,210 @@ +import z from "zod" + +/** + * Tool Bridge Protocol - Message Schemas + * + * Defines all message types for the reverse RPC protocol + * between Tool Bridge (client) and Tool Relay (server). + */ + +// Common message envelope +export const MessageEnvelope = z.object({ + type: z.string(), + id: z.string(), + sessionId: z.string().optional(), + timestamp: z.string(), +}) + +// Tool definition for registration +export const ToolDefinition = z.object({ + id: z.string(), + description: z.string(), + parameters: z.record(z.any()), +}) + +// Handshake messages +export const HelloMessage = z.object({ + type: z.literal("hello"), + id: z.string(), + timestamp: z.string(), + payload: z.object({ + version: z.string(), + bridgeId: z.string(), + credentials: z + .object({ + type: z.enum(["bearer", "none"]), + token: z.string().optional(), + }) + .optional(), + }), +}) + +export const WelcomeMessage = z.object({ + type: z.literal("welcome"), + id: z.string(), + timestamp: z.string(), + payload: z.object({ + version: z.string(), + relayId: z.string(), + sessionId: z.string(), + }), +}) + +export const ErrorMessage = z.object({ + type: z.literal("error"), + id: z.string(), + timestamp: z.string(), + payload: z.object({ + code: z.string(), + message: z.string(), + details: z.record(z.any()).optional(), + }), +}) + +// Tool management messages +export const RegisterToolsMessage = z.object({ + type: z.literal("register_tools"), + id: z.string(), + sessionId: z.string().optional(), + timestamp: z.string(), + payload: z.object({ + tools: z.array(ToolDefinition), + }), +}) + +export const ToolsRegisteredMessage = z.object({ + type: z.literal("tools_registered"), + id: z.string(), + sessionId: z.string().optional(), + timestamp: z.string(), + payload: z.object({ + toolIds: z.array(z.string()), + }), +}) + +// Tool invocation messages +export const InvokeMessage = z.object({ + type: z.literal("invoke"), + id: z.string(), + sessionId: z.string(), + timestamp: z.string(), + payload: z.object({ + callId: z.string(), + toolId: z.string(), + arguments: z.record(z.any()), + traceId: z.string().optional(), + }), +}) + +export const ResultMessage = z.object({ + type: z.literal("result"), + id: z.string(), + sessionId: z.string(), + timestamp: z.string(), + payload: z.object({ + callId: z.string(), + success: z.boolean(), + output: z.string().optional(), + error: z + .object({ + code: z.string(), + message: z.string(), + details: z.record(z.any()).optional(), + }) + .optional(), + metadata: z.record(z.any()).optional(), + title: z.string().optional(), + }), +}) + +// Streaming messages +export const ChunkMessage = z.object({ + type: z.literal("chunk"), + id: z.string(), + sessionId: z.string(), + timestamp: z.string(), + payload: z.object({ + callId: z.string(), + data: z.string(), + metadata: z.record(z.any()).optional(), + }), +}) + +export const DoneMessage = z.object({ + type: z.literal("done"), + id: z.string(), + sessionId: z.string(), + timestamp: z.string(), + payload: z.object({ + callId: z.string(), + }), +}) + +// Control messages +export const CancelMessage = z.object({ + type: z.literal("cancel"), + id: z.string(), + sessionId: z.string(), + timestamp: z.string(), + payload: z.object({ + callId: z.string(), + reason: z.string().optional(), + }), +}) + +export const CancelledMessage = z.object({ + type: z.literal("cancelled"), + id: z.string(), + sessionId: z.string(), + timestamp: z.string(), + payload: z.object({ + callId: z.string(), + }), +}) + +export const PingMessage = z.object({ + type: z.literal("ping"), + id: z.string(), + timestamp: z.string(), +}) + +export const PongMessage = z.object({ + type: z.literal("pong"), + id: z.string(), + timestamp: z.string(), +}) + +// Union of all message types +export const ProtocolMessage = z.discriminatedUnion("type", [ + HelloMessage, + WelcomeMessage, + ErrorMessage, + RegisterToolsMessage, + ToolsRegisteredMessage, + InvokeMessage, + ResultMessage, + ChunkMessage, + DoneMessage, + CancelMessage, + CancelledMessage, + PingMessage, + PongMessage, +]) + +// Type exports +export type MessageEnvelope = z.infer +export type ToolDefinition = z.infer +export type HelloMessage = z.infer +export type WelcomeMessage = z.infer +export type ErrorMessage = z.infer +export type RegisterToolsMessage = z.infer +export type ToolsRegisteredMessage = z.infer +export type InvokeMessage = z.infer +export type ResultMessage = z.infer +export type ChunkMessage = z.infer +export type DoneMessage = z.infer +export type CancelMessage = z.infer +export type CancelledMessage = z.infer +export type PingMessage = z.infer +export type PongMessage = z.infer +export type ProtocolMessage = z.infer diff --git a/packages/tool-bridge-protocol/src/utils.ts b/packages/tool-bridge-protocol/src/utils.ts new file mode 100644 index 00000000000..bb9520dca9a --- /dev/null +++ b/packages/tool-bridge-protocol/src/utils.ts @@ -0,0 +1,243 @@ +import { randomUUID } from "crypto" +import type { + HelloMessage, + WelcomeMessage, + ErrorMessage, + RegisterToolsMessage, + ToolsRegisteredMessage, + InvokeMessage, + ResultMessage, + ChunkMessage, + DoneMessage, + CancelMessage, + CancelledMessage, + PingMessage, + PongMessage, + ToolDefinition, +} from "./schema" + +/** + * Protocol version + */ +export const PROTOCOL_VERSION = "1.0.0" + +/** + * Default timeouts + */ +export const DEFAULT_TIMEOUTS = { + HEARTBEAT_INTERVAL: 30_000, // 30 seconds + CONNECTION_TIMEOUT: 10_000, // 10 seconds + TOOL_EXECUTION_TIMEOUT: 120_000, // 2 minutes + RECONNECT_GRACE_PERIOD: 60_000, // 1 minute +} as const + +/** + * Error codes + */ +export const ErrorCode = { + AUTHENTICATION_FAILED: "AUTHENTICATION_FAILED", + TOOL_NOT_FOUND: "TOOL_NOT_FOUND", + TOOL_EXECUTION_FAILED: "TOOL_EXECUTION_FAILED", + INVALID_MESSAGE: "INVALID_MESSAGE", + SESSION_NOT_FOUND: "SESSION_NOT_FOUND", + BRIDGE_DISCONNECTED: "BRIDGE_DISCONNECTED", + TIMEOUT: "TIMEOUT", + CANCELLED: "CANCELLED", + INTERNAL_ERROR: "INTERNAL_ERROR", +} as const + +export type ErrorCode = (typeof ErrorCode)[keyof typeof ErrorCode] + +/** + * Generate a new unique ID + */ +export function generateId(): string { + return randomUUID() +} + +/** + * Get current ISO timestamp + */ +export function timestamp(): string { + return new Date().toISOString() +} + +/** + * Message factory functions + */ +export const Message = { + hello(bridgeId: string, credentials?: { type: "bearer" | "none"; token?: string }): HelloMessage { + return { + type: "hello", + id: generateId(), + timestamp: timestamp(), + payload: { + version: PROTOCOL_VERSION, + bridgeId, + credentials, + }, + } + }, + + welcome(relayId: string, sessionId: string): WelcomeMessage { + return { + type: "welcome", + id: generateId(), + timestamp: timestamp(), + payload: { + version: PROTOCOL_VERSION, + relayId, + sessionId, + }, + } + }, + + error(code: string, message: string, details?: Record): ErrorMessage { + return { + type: "error", + id: generateId(), + timestamp: timestamp(), + payload: { + code, + message, + details, + }, + } + }, + + registerTools(tools: ToolDefinition[], sessionId?: string): RegisterToolsMessage { + return { + type: "register_tools", + id: generateId(), + sessionId, + timestamp: timestamp(), + payload: { + tools, + }, + } + }, + + toolsRegistered(toolIds: string[], sessionId?: string): ToolsRegisteredMessage { + return { + type: "tools_registered", + id: generateId(), + sessionId, + timestamp: timestamp(), + payload: { + toolIds, + }, + } + }, + + invoke( + sessionId: string, + callId: string, + toolId: string, + args: Record, + traceId?: string + ): InvokeMessage { + return { + type: "invoke", + id: generateId(), + sessionId, + timestamp: timestamp(), + payload: { + callId, + toolId, + arguments: args, + traceId, + }, + } + }, + + result( + sessionId: string, + callId: string, + success: boolean, + options: { + output?: string + error?: { code: string; message: string; details?: Record } + metadata?: Record + title?: string + } + ): ResultMessage { + return { + type: "result", + id: generateId(), + sessionId, + timestamp: timestamp(), + payload: { + callId, + success, + ...options, + }, + } + }, + + chunk(sessionId: string, callId: string, data: string, metadata?: Record): ChunkMessage { + return { + type: "chunk", + id: generateId(), + sessionId, + timestamp: timestamp(), + payload: { + callId, + data, + metadata, + }, + } + }, + + done(sessionId: string, callId: string): DoneMessage { + return { + type: "done", + id: generateId(), + sessionId, + timestamp: timestamp(), + payload: { + callId, + }, + } + }, + + cancel(sessionId: string, callId: string, reason?: string): CancelMessage { + return { + type: "cancel", + id: generateId(), + sessionId, + timestamp: timestamp(), + payload: { + callId, + reason, + }, + } + }, + + cancelled(sessionId: string, callId: string): CancelledMessage { + return { + type: "cancelled", + id: generateId(), + sessionId, + timestamp: timestamp(), + payload: { + callId, + }, + } + }, + + ping(): PingMessage { + return { + type: "ping", + id: generateId(), + timestamp: timestamp(), + } + }, + + pong(): PongMessage { + return { + type: "pong", + id: generateId(), + timestamp: timestamp(), + } + }, +} diff --git a/packages/tool-bridge-protocol/tsconfig.json b/packages/tool-bridge-protocol/tsconfig.json new file mode 100644 index 00000000000..cb5c12aecdc --- /dev/null +++ b/packages/tool-bridge-protocol/tsconfig.json @@ -0,0 +1,9 @@ +{ + "$schema": "https://json.schemastore.org/tsconfig", + "extends": "@tsconfig/bun/tsconfig.json", + "compilerOptions": { + "lib": ["ESNext"], + "types": [], + "noUncheckedIndexedAccess": false + } +} diff --git a/packages/tool-bridge-tests/integration.test.ts b/packages/tool-bridge-tests/integration.test.ts new file mode 100644 index 00000000000..aecedca29ce --- /dev/null +++ b/packages/tool-bridge-tests/integration.test.ts @@ -0,0 +1,479 @@ +import { describe, test, expect, beforeAll, afterAll, beforeEach, afterEach } from "bun:test" +import { ToolRelay } from "@opencode-ai/tool-relay/index" +import { ToolBridge } from "@opencode-ai/tool-bridge/index" + +/** + * Integration tests for Tool Bridge Protocol + * + * These tests verify the full end-to-end communication between + * the Tool Relay (server) and Tool Bridge (client). + */ + +describe("Tool Bridge Protocol Integration", () => { + let relay: ToolRelay + let server: ReturnType + let serverUrl: string + + beforeAll(() => { + // Create the relay + relay = new ToolRelay({ + relayId: "test-relay", + auth: { type: "none" }, + heartbeatInterval: 5000, + executionTimeout: 10000, + }) + + // Create a WebSocket server + server = Bun.serve({ + port: 0, // Use random available port + fetch(req, server) { + if (server.upgrade(req)) { + return // Upgraded to WebSocket + } + return new Response("Not found", { status: 404 }) + }, + websocket: { + open(ws) { + relay.handleConnection(ws as unknown as WebSocket) + }, + message(ws, message) { + // Messages are handled by the relay via the connection + }, + close(ws) { + // Handled by the relay + }, + }, + }) + + serverUrl = `ws://localhost:${server.port}` + relay.start() + }) + + afterAll(() => { + relay.stop() + server.stop() + }) + + describe("Connection Lifecycle", () => { + test("bridge can connect to relay", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-1", + reconnect: { enabled: false }, + }) + + const sessionId = await bridge.connect() + + expect(sessionId).toBeTruthy() + expect(bridge.getState()).toBe("connected") + + bridge.disconnect() + }) + + test("bridge receives session ID on connect", async () => { + let receivedSessionId: string | null = null + + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-2", + reconnect: { enabled: false }, + onConnect: (sessionId) => { + receivedSessionId = sessionId + }, + }) + + await bridge.connect() + + expect(receivedSessionId).toBeTruthy() + expect(receivedSessionId).toBe(bridge.getSessionId()) + + bridge.disconnect() + }) + + test("relay tracks connected sessions", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-3", + reconnect: { enabled: false }, + }) + + const sessionId = await bridge.connect() + + const sessions = relay.getSessions() + expect(sessions.length).toBeGreaterThan(0) + expect(sessions.some((s) => s.sessionId === sessionId)).toBe(true) + + bridge.disconnect() + }) + + test("disconnect triggers callback", async () => { + let disconnected = false + let disconnectReason: string | null = null + + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-4", + reconnect: { enabled: false }, + onDisconnect: (reason) => { + disconnected = true + disconnectReason = reason + }, + }) + + await bridge.connect() + bridge.disconnect() + + // Give time for events to propagate + await new Promise((resolve) => setTimeout(resolve, 100)) + + expect(bridge.getState()).toBe("disconnected") + }) + }) + + describe("Tool Registration", () => { + test("bridge can register tools", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-5", + reconnect: { enabled: false }, + }) + + bridge.registerTool( + "echo", + "Echoes back the input", + { type: "object", properties: { message: { type: "string" } } }, + async (args) => ({ output: args.message as string }) + ) + + const sessionId = await bridge.connect() + + // Give time for tool registration + await new Promise((resolve) => setTimeout(resolve, 100)) + + const tools = relay.getTools(sessionId) + expect(tools.length).toBe(1) + expect(tools[0].id).toBe("echo") + + bridge.disconnect() + }) + + test("relay can list registered tools", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-6", + reconnect: { enabled: false }, + }) + + bridge.registerTool( + "tool-a", + "Tool A description", + { type: "object" }, + async () => ({ output: "a" }) + ) + + bridge.registerTool( + "tool-b", + "Tool B description", + { type: "object" }, + async () => ({ output: "b" }) + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + const tools = relay.getTools(sessionId) + expect(tools.length).toBe(2) + expect(tools.map((t) => t.id).sort()).toEqual(["tool-a", "tool-b"]) + + bridge.disconnect() + }) + }) + + describe("Tool Invocation", () => { + test("relay can invoke tool and receive result", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-7", + reconnect: { enabled: false }, + }) + + bridge.registerTool( + "add", + "Adds two numbers", + { + type: "object", + properties: { + a: { type: "number" }, + b: { type: "number" }, + }, + }, + async (args) => { + const sum = (args.a as number) + (args.b as number) + return { output: String(sum), title: "Addition" } + } + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + const result = await relay.invoke({ + sessionId, + toolId: "add", + arguments: { a: 5, b: 3 }, + }) + + expect(result.success).toBe(true) + expect(result.output).toBe("8") + expect(result.title).toBe("Addition") + + bridge.disconnect() + }) + + test("tool execution error is reported", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-8", + reconnect: { enabled: false }, + }) + + bridge.registerTool( + "failing-tool", + "Always fails", + { type: "object" }, + async () => { + throw new Error("Intentional failure") + } + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + const result = await relay.invoke({ + sessionId, + toolId: "failing-tool", + arguments: {}, + }) + + expect(result.success).toBe(false) + expect(result.error).toBeTruthy() + expect(result.error!.message).toContain("Intentional failure") + + bridge.disconnect() + }) + + test("invoking unknown tool returns error", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-9", + reconnect: { enabled: false }, + }) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + const result = await relay.invoke({ + sessionId, + toolId: "nonexistent-tool", + arguments: {}, + }) + + expect(result.success).toBe(false) + expect(result.error!.code).toBe("TOOL_NOT_FOUND") + + bridge.disconnect() + }) + + test("multiple concurrent tool invocations", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-10", + reconnect: { enabled: false }, + }) + + bridge.registerTool( + "delay", + "Delays and returns a value", + { + type: "object", + properties: { + ms: { type: "number" }, + value: { type: "string" }, + }, + }, + async (args) => { + await new Promise((resolve) => setTimeout(resolve, args.ms as number)) + return { output: args.value as string } + } + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + // Invoke multiple tools concurrently + const results = await Promise.all([ + relay.invoke({ sessionId, toolId: "delay", arguments: { ms: 100, value: "first" } }), + relay.invoke({ sessionId, toolId: "delay", arguments: { ms: 50, value: "second" } }), + relay.invoke({ sessionId, toolId: "delay", arguments: { ms: 75, value: "third" } }), + ]) + + expect(results.every((r) => r.success)).toBe(true) + expect(results.map((r) => r.output).sort()).toEqual(["first", "second", "third"]) + + bridge.disconnect() + }) + }) + + describe("Cancellation", () => { + test("tool can be cancelled", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-11", + reconnect: { enabled: false }, + }) + + let wasCancelled = false + + bridge.registerTool( + "long-running", + "Long running task", + { type: "object" }, + async (args, ctx) => { + try { + await new Promise((resolve, reject) => { + const timeout = setTimeout(resolve, 10000) + ctx.abort.addEventListener("abort", () => { + clearTimeout(timeout) + wasCancelled = true + reject(new Error("Cancelled")) + }) + }) + return { output: "completed" } + } catch { + throw new Error("Cancelled") + } + } + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + // Start invocation + const invocationPromise = relay.invoke({ + sessionId, + toolId: "long-running", + arguments: {}, + }) + + // Cancel after a short delay + await new Promise((resolve) => setTimeout(resolve, 50)) + relay.cancel(sessionId, "some-call-id") + + // The invocation should still complete (either with result or cancellation) + // depending on timing + await invocationPromise + + bridge.disconnect() + }) + }) + + describe("Session Management", () => { + test("isConnected returns correct state", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-12", + reconnect: { enabled: false }, + }) + + const sessionId = await bridge.connect() + + expect(relay.isConnected(sessionId)).toBe(true) + + bridge.disconnect() + + // Give time for disconnect to propagate + await new Promise((resolve) => setTimeout(resolve, 100)) + + expect(relay.isConnected(sessionId)).toBe(false) + }) + + test("invoking on disconnected session returns error", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-13", + reconnect: { enabled: false }, + }) + + bridge.registerTool( + "test", + "Test tool", + { type: "object" }, + async () => ({ output: "ok" }) + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + bridge.disconnect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + const result = await relay.invoke({ + sessionId, + toolId: "test", + arguments: {}, + }) + + expect(result.success).toBe(false) + expect(result.error!.code).toBe("BRIDGE_DISCONNECTED") + }) + }) + + describe("Streaming", () => { + test("tool can send chunks during execution", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-14", + reconnect: { enabled: false }, + }) + + bridge.registerTool( + "streaming", + "Sends multiple chunks", + { type: "object" }, + async (args, ctx) => { + ctx.sendChunk("chunk-1", { index: 0 }) + ctx.sendChunk("chunk-2", { index: 1 }) + ctx.sendChunk("chunk-3", { index: 2 }) + return { output: "done", metadata: { chunks: 3 } } + } + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + const result = await relay.invoke({ + sessionId, + toolId: "streaming", + arguments: {}, + }) + + expect(result.success).toBe(true) + expect(result.output).toBe("done") + expect(result.metadata?.chunks).toBe(3) + + bridge.disconnect() + }) + }) + + describe("Error Handling", () => { + test("invalid session ID returns error", async () => { + const result = await relay.invoke({ + sessionId: "nonexistent-session", + toolId: "some-tool", + arguments: {}, + }) + + expect(result.success).toBe(false) + expect(result.error!.code).toBe("SESSION_NOT_FOUND") + }) + }) +}) diff --git a/packages/tool-bridge-tests/package.json b/packages/tool-bridge-tests/package.json new file mode 100644 index 00000000000..5a07a4d5975 --- /dev/null +++ b/packages/tool-bridge-tests/package.json @@ -0,0 +1,20 @@ +{ + "name": "@opencode-ai/tool-bridge-tests", + "version": "1.0.0", + "type": "module", + "private": true, + "scripts": { + "test": "bun test" + }, + "dependencies": { + "@opencode-ai/tool-bridge": "workspace:*", + "@opencode-ai/tool-relay": "workspace:*", + "@opencode-ai/tool-bridge-protocol": "workspace:*", + "zod": "catalog:" + }, + "devDependencies": { + "@tsconfig/bun": "catalog:", + "@types/bun": "catalog:", + "typescript": "catalog:" + } +} diff --git a/packages/tool-bridge-tests/tsconfig.json b/packages/tool-bridge-tests/tsconfig.json new file mode 100644 index 00000000000..5304fa0b4be --- /dev/null +++ b/packages/tool-bridge-tests/tsconfig.json @@ -0,0 +1,9 @@ +{ + "$schema": "https://json.schemastore.org/tsconfig", + "extends": "@tsconfig/bun/tsconfig.json", + "compilerOptions": { + "lib": ["ESNext", "DOM"], + "types": ["bun-types"], + "noUncheckedIndexedAccess": false + } +} diff --git a/packages/tool-bridge/package.json b/packages/tool-bridge/package.json new file mode 100644 index 00000000000..a9237d1e266 --- /dev/null +++ b/packages/tool-bridge/package.json @@ -0,0 +1,18 @@ +{ + "name": "@opencode-ai/tool-bridge", + "version": "1.0.0", + "type": "module", + "private": true, + "exports": { + "./*": "./src/*.ts" + }, + "dependencies": { + "@opencode-ai/tool-bridge-protocol": "workspace:*", + "zod": "catalog:" + }, + "devDependencies": { + "@tsconfig/bun": "catalog:", + "@types/bun": "catalog:", + "typescript": "catalog:" + } +} diff --git a/packages/tool-bridge/src/bridge.ts b/packages/tool-bridge/src/bridge.ts new file mode 100644 index 00000000000..a141aceb755 --- /dev/null +++ b/packages/tool-bridge/src/bridge.ts @@ -0,0 +1,506 @@ +import { + ProtocolMessage, + WelcomeMessage, + ErrorMessage, + ToolsRegisteredMessage, + InvokeMessage, + CancelMessage, + PingMessage, + ToolDefinition, +} from "@opencode-ai/tool-bridge-protocol/schema" +import { Message, generateId, ErrorCode, DEFAULT_TIMEOUTS } from "@opencode-ai/tool-bridge-protocol/utils" +import type { + BridgeConfig, + ToolExecutor, + ToolContext, + ToolResult, + RegisteredTool, + ConnectionState, +} from "./types" + +/** + * Tool Bridge - Client-side agent for local tool execution + * + * Connects to a Tool Relay and executes tools locally on behalf + * of a remote LLM runtime. + */ +export class ToolBridge { + private readonly bridgeId: string + private readonly config: BridgeConfig + private readonly tools: Map = new Map() + private readonly activeCalls: Map = new Map() + + private socket: WebSocket | null = null + private sessionId: string | null = null + private state: ConnectionState = "disconnected" + private reconnectAttempt = 0 + private reconnectTimer: ReturnType | null = null + private heartbeatTimer: ReturnType | null = null + + constructor(config: BridgeConfig) { + this.bridgeId = config.bridgeId ?? generateId() + this.config = { + reconnect: { + enabled: true, + maxAttempts: 10, + baseDelay: 1000, + maxDelay: 30000, + }, + ...config, + } + } + + /** + * Register a tool with the bridge + */ + registerTool(id: string, description: string, parameters: Record, executor: ToolExecutor): void { + const definition: ToolDefinition = { + id, + description, + parameters, + } + this.tools.set(id, { definition, executor }) + } + + /** + * Unregister a tool + */ + unregisterTool(id: string): void { + this.tools.delete(id) + } + + /** + * Connect to the relay + */ + async connect(): Promise { + if (this.state === "connected") { + return this.sessionId! + } + + if (this.state === "connecting" || this.state === "reconnecting") { + return new Promise((resolve, reject) => { + const onConnect = this.config.onConnect + const onError = this.config.onError + + this.config.onConnect = (sessionId) => { + this.config.onConnect = onConnect + this.config.onError = onError + onConnect?.(sessionId) + resolve(sessionId) + } + + this.config.onError = (error) => { + this.config.onConnect = onConnect + this.config.onError = onError + onError?.(error) + reject(error) + } + }) + } + + this.state = "connecting" + + return new Promise((resolve, reject) => { + try { + this.socket = new WebSocket(this.config.url) + + const timeout = setTimeout(() => { + this.socket?.close() + this.state = "disconnected" + reject(new Error("Connection timeout")) + }, DEFAULT_TIMEOUTS.CONNECTION_TIMEOUT) + + this.socket.addEventListener("open", () => { + clearTimeout(timeout) + this.sendHello() + }) + + this.socket.addEventListener("message", async (event) => { + await this.handleMessage(event.data, resolve, reject) + }) + + this.socket.addEventListener("close", (event) => { + clearTimeout(timeout) + this.handleClose(event.reason || "Connection closed") + }) + + this.socket.addEventListener("error", () => { + clearTimeout(timeout) + if (this.state === "connecting") { + this.state = "disconnected" + reject(new Error("WebSocket connection failed")) + } + }) + } catch (error) { + this.state = "disconnected" + reject(error) + } + }) + } + + /** + * Disconnect from the relay + */ + disconnect(): void { + this.stopReconnect() + this.stopHeartbeat() + this.cancelAllCalls("Bridge disconnected") + + if (this.socket) { + try { + this.socket.close() + } catch {} + this.socket = null + } + + this.sessionId = null + this.state = "disconnected" + } + + /** + * Get connection state + */ + getState(): ConnectionState { + return this.state + } + + /** + * Get session ID + */ + getSessionId(): string | null { + return this.sessionId + } + + /** + * Get registered tools + */ + getTools(): ToolDefinition[] { + return Array.from(this.tools.values()).map((t) => t.definition) + } + + /** + * Send hello message + */ + private sendHello(): void { + if (!this.socket) return + + const hello = Message.hello(this.bridgeId, this.config.credentials) + this.socket.send(JSON.stringify(hello)) + } + + /** + * Handle incoming message + */ + private async handleMessage( + data: string | ArrayBuffer, + onConnect?: (sessionId: string) => void, + onError?: (error: Error) => void + ): Promise { + try { + const text = typeof data === "string" ? data : new TextDecoder().decode(data) + const message = JSON.parse(text) + const parsed = ProtocolMessage.safeParse(message) + + if (!parsed.success) { + console.error("Invalid message received:", parsed.error) + return + } + + switch (parsed.data.type) { + case "welcome": + this.handleWelcome(parsed.data, onConnect) + break + + case "error": + this.handleError(parsed.data, onError) + break + + case "tools_registered": + this.handleToolsRegistered(parsed.data) + break + + case "invoke": + await this.handleInvoke(parsed.data) + break + + case "cancel": + this.handleCancel(parsed.data) + break + + case "ping": + this.handlePing() + break + } + } catch (error) { + console.error("Failed to handle message:", error) + } + } + + /** + * Handle welcome message + */ + private handleWelcome(message: WelcomeMessage, onConnect?: (sessionId: string) => void): void { + this.sessionId = message.payload.sessionId + this.state = "connected" + this.reconnectAttempt = 0 + + // Start heartbeat + this.startHeartbeat() + + // Register tools + this.sendRegisterTools() + + this.config.onConnect?.(this.sessionId) + onConnect?.(this.sessionId) + } + + /** + * Handle error message + */ + private handleError(message: ErrorMessage, onError?: (error: Error) => void): void { + const error = new Error(`${message.payload.code}: ${message.payload.message}`) + this.config.onError?.(error) + onError?.(error) + } + + /** + * Handle tools_registered message + */ + private handleToolsRegistered(message: ToolsRegisteredMessage): void { + // Tools successfully registered + // Could emit an event here if needed + } + + /** + * Handle invoke message + */ + private async handleInvoke(message: InvokeMessage): Promise { + const { callId, toolId, arguments: args, traceId } = message.payload + const tool = this.tools.get(toolId) + + if (!tool) { + this.sendResult(message.sessionId, callId, false, { + error: { + code: ErrorCode.TOOL_NOT_FOUND, + message: `Tool ${toolId} not found`, + }, + }) + return + } + + // Create abort controller for this call + const abortController = new AbortController() + this.activeCalls.set(callId, abortController) + + const context: ToolContext = { + sessionId: message.sessionId, + callId, + traceId, + abort: abortController.signal, + sendChunk: (data, metadata) => { + this.sendChunk(message.sessionId, callId, data, metadata) + }, + } + + try { + const result = await tool.executor(args, context) + + if (!this.activeCalls.has(callId)) { + // Call was cancelled + return + } + + this.sendResult(message.sessionId, callId, true, { + output: result.output, + metadata: result.metadata, + title: result.title, + }) + } catch (error) { + if (!this.activeCalls.has(callId)) { + // Call was cancelled + return + } + + this.sendResult(message.sessionId, callId, false, { + error: { + code: ErrorCode.TOOL_EXECUTION_FAILED, + message: error instanceof Error ? error.message : String(error), + }, + }) + } finally { + this.activeCalls.delete(callId) + } + } + + /** + * Handle cancel message + */ + private handleCancel(message: CancelMessage): void { + const { callId } = message.payload + const abortController = this.activeCalls.get(callId) + + if (abortController) { + abortController.abort() + this.activeCalls.delete(callId) + this.sendCancelled(message.sessionId, callId) + } + } + + /** + * Handle ping message + */ + private handlePing(): void { + if (!this.socket) return + const pong = Message.pong() + this.socket.send(JSON.stringify(pong)) + } + + /** + * Handle connection close + */ + private handleClose(reason: string): void { + this.stopHeartbeat() + this.cancelAllCalls("Connection closed") + + const wasConnected = this.state === "connected" + this.state = "disconnected" + this.socket = null + + this.config.onDisconnect?.(reason) + + if (wasConnected && this.config.reconnect?.enabled) { + this.scheduleReconnect() + } + } + + /** + * Send register_tools message + */ + private sendRegisterTools(): void { + if (!this.socket || !this.sessionId) return + + const tools = Array.from(this.tools.values()).map((t) => t.definition) + const message = Message.registerTools(tools, this.sessionId) + this.socket.send(JSON.stringify(message)) + } + + /** + * Send result message + */ + private sendResult( + sessionId: string, + callId: string, + success: boolean, + options: { + output?: string + error?: { code: string; message: string; details?: Record } + metadata?: Record + title?: string + } + ): void { + if (!this.socket) return + + const message = Message.result(sessionId, callId, success, options) + this.socket.send(JSON.stringify(message)) + } + + /** + * Send chunk message + */ + private sendChunk( + sessionId: string, + callId: string, + data: string, + metadata?: Record + ): void { + if (!this.socket) return + + const message = Message.chunk(sessionId, callId, data, metadata) + this.socket.send(JSON.stringify(message)) + } + + /** + * Send cancelled message + */ + private sendCancelled(sessionId: string, callId: string): void { + if (!this.socket) return + + const message = Message.cancelled(sessionId, callId) + this.socket.send(JSON.stringify(message)) + } + + /** + * Cancel all active calls + */ + private cancelAllCalls(reason: string): void { + for (const [callId, abortController] of this.activeCalls) { + abortController.abort(reason) + } + this.activeCalls.clear() + } + + /** + * Start heartbeat timer + */ + private startHeartbeat(): void { + if (this.heartbeatTimer) return + + this.heartbeatTimer = setInterval(() => { + if (this.socket && this.state === "connected") { + // Just check connection is alive + // Relay sends pings, we respond with pongs + } + }, DEFAULT_TIMEOUTS.HEARTBEAT_INTERVAL) + } + + /** + * Stop heartbeat timer + */ + private stopHeartbeat(): void { + if (this.heartbeatTimer) { + clearInterval(this.heartbeatTimer) + this.heartbeatTimer = null + } + } + + /** + * Schedule reconnection attempt + */ + private scheduleReconnect(): void { + const { maxAttempts = 10, baseDelay = 1000, maxDelay = 30000 } = this.config.reconnect ?? {} + + if (this.reconnectAttempt >= maxAttempts) { + this.config.onError?.(new Error(`Max reconnection attempts (${maxAttempts}) reached`)) + return + } + + this.reconnectAttempt++ + this.state = "reconnecting" + + // Exponential backoff with jitter + const delay = Math.min(baseDelay * Math.pow(2, this.reconnectAttempt - 1), maxDelay) + const jitter = delay * 0.1 * Math.random() + const totalDelay = delay + jitter + + this.config.onReconnect?.(this.reconnectAttempt) + + this.reconnectTimer = setTimeout(async () => { + try { + await this.connect() + } catch { + // Will trigger handleClose which schedules next reconnect + } + }, totalDelay) + } + + /** + * Stop reconnection attempts + */ + private stopReconnect(): void { + if (this.reconnectTimer) { + clearTimeout(this.reconnectTimer) + this.reconnectTimer = null + } + this.reconnectAttempt = 0 + } +} diff --git a/packages/tool-bridge/src/index.ts b/packages/tool-bridge/src/index.ts new file mode 100644 index 00000000000..d289afac8f0 --- /dev/null +++ b/packages/tool-bridge/src/index.ts @@ -0,0 +1,16 @@ +/** + * Tool Bridge - Client-side agent for Tool Bridge Protocol + * + * The Tool Bridge connects to a Tool Relay and executes tools + * locally on behalf of remote LLM runtimes. + */ + +export { ToolBridge } from "./bridge" +export type { + BridgeConfig, + ToolExecutor, + ToolContext, + ToolResult, + RegisteredTool, + ConnectionState, +} from "./types" diff --git a/packages/tool-bridge/src/types.ts b/packages/tool-bridge/src/types.ts new file mode 100644 index 00000000000..a1ba21726a3 --- /dev/null +++ b/packages/tool-bridge/src/types.ts @@ -0,0 +1,83 @@ +import type { ToolDefinition } from "@opencode-ai/tool-bridge-protocol/schema" + +/** + * Bridge configuration options + */ +export interface BridgeConfig { + /** + * URL of the Tool Relay WebSocket endpoint + */ + url: string + + /** + * Unique identifier for this bridge + */ + bridgeId?: string + + /** + * Authentication credentials + */ + credentials?: { + type: "bearer" | "none" + token?: string + } + + /** + * Reconnect configuration + */ + reconnect?: { + enabled: boolean + maxAttempts?: number + baseDelay?: number + maxDelay?: number + } + + /** + * Event handlers + */ + onConnect?: (sessionId: string) => void + onDisconnect?: (reason: string) => void + onReconnect?: (attempt: number) => void + onError?: (error: Error) => void +} + +/** + * Tool executor function type + */ +export type ToolExecutor = ( + args: Record, + context: ToolContext +) => Promise + +/** + * Tool context provided during execution + */ +export interface ToolContext { + sessionId: string + callId: string + traceId?: string + abort: AbortSignal + sendChunk: (data: string, metadata?: Record) => void +} + +/** + * Tool result from execution + */ +export interface ToolResult { + output: string + metadata?: Record + title?: string +} + +/** + * Registered tool information + */ +export interface RegisteredTool { + definition: ToolDefinition + executor: ToolExecutor +} + +/** + * Bridge connection state + */ +export type ConnectionState = "disconnected" | "connecting" | "connected" | "reconnecting" diff --git a/packages/tool-bridge/tsconfig.json b/packages/tool-bridge/tsconfig.json new file mode 100644 index 00000000000..cb5c12aecdc --- /dev/null +++ b/packages/tool-bridge/tsconfig.json @@ -0,0 +1,9 @@ +{ + "$schema": "https://json.schemastore.org/tsconfig", + "extends": "@tsconfig/bun/tsconfig.json", + "compilerOptions": { + "lib": ["ESNext"], + "types": [], + "noUncheckedIndexedAccess": false + } +} diff --git a/packages/tool-relay/package.json b/packages/tool-relay/package.json new file mode 100644 index 00000000000..8840a0344ce --- /dev/null +++ b/packages/tool-relay/package.json @@ -0,0 +1,19 @@ +{ + "name": "@opencode-ai/tool-relay", + "version": "1.0.0", + "type": "module", + "private": true, + "exports": { + "./*": "./src/*.ts" + }, + "dependencies": { + "@opencode-ai/tool-bridge-protocol": "workspace:*", + "hono": "catalog:", + "zod": "catalog:" + }, + "devDependencies": { + "@tsconfig/bun": "catalog:", + "@types/bun": "catalog:", + "typescript": "catalog:" + } +} diff --git a/packages/tool-relay/src/index.ts b/packages/tool-relay/src/index.ts new file mode 100644 index 00000000000..b6dff5bf58a --- /dev/null +++ b/packages/tool-relay/src/index.ts @@ -0,0 +1,16 @@ +/** + * Tool Relay - Server-side gateway for Tool Bridge Protocol + * + * The Tool Relay accepts WebSocket connections from Tool Bridges and + * forwards tool invocations from LLM runtimes. + */ + +export { ToolRelay } from "./relay" +export type { + RelayConfig, + AuthConfig, + BridgeConnection, + Session, + ToolInvocationRequest, + ToolInvocationResult, +} from "./types" diff --git a/packages/tool-relay/src/relay.ts b/packages/tool-relay/src/relay.ts new file mode 100644 index 00000000000..e24b989b4a3 --- /dev/null +++ b/packages/tool-relay/src/relay.ts @@ -0,0 +1,565 @@ +import { + ProtocolMessage, + HelloMessage, + RegisterToolsMessage, + ResultMessage, + ChunkMessage, + CancelledMessage, + PongMessage, + ToolDefinition, +} from "@opencode-ai/tool-bridge-protocol/schema" +import { Message, generateId, ErrorCode, DEFAULT_TIMEOUTS, PROTOCOL_VERSION } from "@opencode-ai/tool-bridge-protocol/utils" +import type { RelayConfig, BridgeConnection, Session, ToolInvocationRequest, ToolInvocationResult } from "./types" + +/** + * Tool Relay - Server-side gateway for tool invocations + * + * Accepts WebSocket connections from Tool Bridges and forwards + * tool invocations from LLM runtimes. + */ +export class ToolRelay { + private readonly relayId: string + private readonly config: RelayConfig + private readonly connections: Map = new Map() + private readonly sessions: Map = new Map() + private readonly sessionToBridge: Map = new Map() + private heartbeatTimer?: ReturnType + + constructor(config: RelayConfig = {}) { + this.relayId = config.relayId ?? generateId() + this.config = { + heartbeatInterval: DEFAULT_TIMEOUTS.HEARTBEAT_INTERVAL, + executionTimeout: DEFAULT_TIMEOUTS.TOOL_EXECUTION_TIMEOUT, + sessionGracePeriod: DEFAULT_TIMEOUTS.RECONNECT_GRACE_PERIOD, + ...config, + } + } + + /** + * Start the relay heartbeat + */ + start(): void { + if (this.heartbeatTimer) return + this.heartbeatTimer = setInterval(() => { + this.sendHeartbeats() + this.cleanupStaleSessions() + }, this.config.heartbeatInterval) + } + + /** + * Stop the relay + */ + stop(): void { + if (this.heartbeatTimer) { + clearInterval(this.heartbeatTimer) + this.heartbeatTimer = undefined + } + + // Close all connections + for (const [bridgeId, connection] of this.connections) { + this.disconnectBridge(bridgeId, "relay_shutdown") + } + } + + /** + * Handle a new WebSocket connection + */ + handleConnection(socket: WebSocket): void { + let bridgeId: string | undefined + + socket.addEventListener("message", async (event) => { + try { + const data = typeof event.data === "string" ? event.data : event.data.toString() + const message = JSON.parse(data) + const parsed = ProtocolMessage.safeParse(message) + + if (!parsed.success) { + this.sendError(socket, ErrorCode.INVALID_MESSAGE, `Invalid message format: ${parsed.error.message}`) + return + } + + await this.handleMessage(socket, parsed.data, (id) => { + bridgeId = id + }) + } catch (error) { + this.sendError(socket, ErrorCode.INTERNAL_ERROR, `Failed to process message: ${error}`) + } + }) + + socket.addEventListener("close", () => { + if (bridgeId) { + this.handleDisconnect(bridgeId) + } + }) + + socket.addEventListener("error", () => { + if (bridgeId) { + this.handleDisconnect(bridgeId) + } + }) + } + + /** + * Handle incoming message + */ + private async handleMessage( + socket: WebSocket, + message: ProtocolMessage, + setBridgeId: (id: string) => void + ): Promise { + switch (message.type) { + case "hello": + await this.handleHello(socket, message, setBridgeId) + break + + case "register_tools": + this.handleRegisterTools(message) + break + + case "result": + this.handleResult(message) + break + + case "chunk": + this.handleChunk(message) + break + + case "cancelled": + this.handleCancelled(message) + break + + case "pong": + this.handlePong(message) + break + + default: + // Ignore unsupported message types from client + break + } + } + + /** + * Handle hello message (authentication) + */ + private async handleHello( + socket: WebSocket, + message: HelloMessage, + setBridgeId: (id: string) => void + ): Promise { + const { bridgeId, credentials } = message.payload + + // Validate authentication if configured + if (this.config.auth?.type === "bearer") { + const token = credentials?.token + if (!token) { + this.sendError(socket, ErrorCode.AUTHENTICATION_FAILED, "Missing authentication token") + socket.close() + return + } + + const isValid = this.config.auth.validateToken + ? await this.config.auth.validateToken(token) + : true + + if (!isValid) { + this.sendError(socket, ErrorCode.AUTHENTICATION_FAILED, "Invalid authentication token") + socket.close() + return + } + } + + // Check for existing session + let session = this.findSessionByBridgeId(bridgeId) + let sessionId: string + + if (session) { + // Rebind to existing session + sessionId = session.sessionId + session.lastActivity = Date.now() + } else { + // Create new session + sessionId = generateId() + session = { + sessionId, + bridgeId, + createdAt: Date.now(), + lastActivity: Date.now(), + tools: [], + } + this.sessions.set(sessionId, session) + } + + // Create connection + const connection: BridgeConnection = { + bridgeId, + sessionId, + socket, + tools: new Map(), + pendingCalls: new Map(), + lastPing: Date.now(), + connected: true, + } + + // Clean up old connection if exists + const oldConnection = this.connections.get(bridgeId) + if (oldConnection && oldConnection.connected) { + try { + oldConnection.socket.close() + } catch {} + } + + this.connections.set(bridgeId, connection) + this.sessionToBridge.set(sessionId, bridgeId) + setBridgeId(bridgeId) + + // Send welcome + const welcome = Message.welcome(this.relayId, sessionId) + socket.send(JSON.stringify(welcome)) + + this.config.onBridgeConnect?.(bridgeId, sessionId) + } + + /** + * Handle register_tools message + */ + private handleRegisterTools(message: RegisterToolsMessage): void { + const bridgeId = this.sessionToBridge.get(message.sessionId ?? "") + if (!bridgeId) return + + const connection = this.connections.get(bridgeId) + if (!connection) return + + const { tools } = message.payload + + // Register tools + for (const tool of tools) { + connection.tools.set(tool.id, tool) + } + + // Update session + const session = this.sessions.get(connection.sessionId) + if (session) { + session.tools = tools + session.lastActivity = Date.now() + } + + // Send confirmation + const confirmation = Message.toolsRegistered( + tools.map((t) => t.id), + connection.sessionId + ) + connection.socket.send(JSON.stringify(confirmation)) + + this.config.onToolRegistered?.(bridgeId, tools) + } + + /** + * Handle result message + */ + private handleResult(message: ResultMessage): void { + const { callId } = message.payload + const bridgeId = this.sessionToBridge.get(message.sessionId) + if (!bridgeId) return + + const connection = this.connections.get(bridgeId) + if (!connection) return + + const pending = connection.pendingCalls.get(callId) + if (!pending) return + + clearTimeout(pending.timeoutId) + connection.pendingCalls.delete(callId) + + this.config.onToolResult?.(message.sessionId, callId, message.payload) + pending.resolve(message.payload) + } + + /** + * Handle chunk message (streaming) + */ + private handleChunk(message: ChunkMessage): void { + // For now, just update session activity + // Streaming chunks could be forwarded to subscribers + const session = this.sessions.get(message.sessionId) + if (session) { + session.lastActivity = Date.now() + } + } + + /** + * Handle cancelled message + */ + private handleCancelled(message: CancelledMessage): void { + const { callId } = message.payload + const bridgeId = this.sessionToBridge.get(message.sessionId) + if (!bridgeId) return + + const connection = this.connections.get(bridgeId) + if (!connection) return + + const pending = connection.pendingCalls.get(callId) + if (!pending) return + + clearTimeout(pending.timeoutId) + connection.pendingCalls.delete(callId) + + pending.reject(new Error("Tool invocation was cancelled")) + } + + /** + * Handle pong message + */ + private handlePong(message: PongMessage): void { + // Find connection by iterating (not ideal but works for now) + for (const connection of this.connections.values()) { + if (connection.connected) { + connection.lastPing = Date.now() + } + } + } + + /** + * Handle bridge disconnect + */ + private handleDisconnect(bridgeId: string): void { + const connection = this.connections.get(bridgeId) + if (!connection) return + + connection.connected = false + + // Fail all pending calls + for (const [callId, pending] of connection.pendingCalls) { + clearTimeout(pending.timeoutId) + pending.reject(new Error("Bridge disconnected")) + } + connection.pendingCalls.clear() + + this.config.onBridgeDisconnect?.(bridgeId, connection.sessionId) + + // Keep session alive for grace period + // Session will be cleaned up by cleanupStaleSessions if not reconnected + } + + /** + * Disconnect a bridge + */ + private disconnectBridge(bridgeId: string, reason: string): void { + const connection = this.connections.get(bridgeId) + if (!connection) return + + try { + connection.socket.close() + } catch {} + + this.handleDisconnect(bridgeId) + this.connections.delete(bridgeId) + } + + /** + * Send heartbeats to all connections + */ + private sendHeartbeats(): void { + const ping = Message.ping() + const pingJson = JSON.stringify(ping) + + for (const connection of this.connections.values()) { + if (connection.connected) { + try { + connection.socket.send(pingJson) + } catch { + connection.connected = false + } + } + } + } + + /** + * Clean up stale sessions + */ + private cleanupStaleSessions(): void { + const now = Date.now() + const gracePeriod = this.config.sessionGracePeriod! + + for (const [sessionId, session] of this.sessions) { + const bridgeId = this.sessionToBridge.get(sessionId) + const connection = bridgeId ? this.connections.get(bridgeId) : undefined + + // If no connection or disconnected for too long + if (!connection || (!connection.connected && now - session.lastActivity > gracePeriod)) { + this.sessions.delete(sessionId) + this.sessionToBridge.delete(sessionId) + if (bridgeId) { + this.connections.delete(bridgeId) + } + } + } + } + + /** + * Find session by bridge ID + */ + private findSessionByBridgeId(bridgeId: string): Session | undefined { + for (const session of this.sessions.values()) { + if (session.bridgeId === bridgeId) { + return session + } + } + return undefined + } + + /** + * Send error message + */ + private sendError(socket: WebSocket, code: string, message: string): void { + const error = Message.error(code, message) + try { + socket.send(JSON.stringify(error)) + } catch {} + } + + /** + * Invoke a tool on a connected bridge + */ + async invoke(request: ToolInvocationRequest): Promise { + const { sessionId, toolId, arguments: args, traceId } = request + + const bridgeId = this.sessionToBridge.get(sessionId) + if (!bridgeId) { + return { + success: false, + error: { + code: ErrorCode.SESSION_NOT_FOUND, + message: `Session ${sessionId} not found`, + }, + } + } + + const connection = this.connections.get(bridgeId) + if (!connection || !connection.connected) { + return { + success: false, + error: { + code: ErrorCode.BRIDGE_DISCONNECTED, + message: `Bridge ${bridgeId} is not connected`, + }, + } + } + + if (!connection.tools.has(toolId)) { + return { + success: false, + error: { + code: ErrorCode.TOOL_NOT_FOUND, + message: `Tool ${toolId} not found on bridge`, + }, + } + } + + const callId = generateId() + const invokeMessage = Message.invoke(sessionId, callId, toolId, args, traceId) + + this.config.onToolInvoke?.(sessionId, callId, toolId) + + return new Promise((resolve, reject) => { + const timeoutId = setTimeout(() => { + connection.pendingCalls.delete(callId) + + // Send cancel message + const cancel = Message.cancel(sessionId, callId, "timeout") + try { + connection.socket.send(JSON.stringify(cancel)) + } catch {} + + resolve({ + success: false, + error: { + code: ErrorCode.TIMEOUT, + message: `Tool invocation timed out after ${this.config.executionTimeout}ms`, + }, + }) + }, this.config.executionTimeout) + + connection.pendingCalls.set(callId, { + resolve: (result) => { + resolve({ + success: result.success, + output: result.output, + error: result.error, + metadata: result.metadata, + title: result.title, + }) + }, + reject, + timeoutId, + }) + + try { + connection.socket.send(JSON.stringify(invokeMessage)) + } catch (error) { + connection.pendingCalls.delete(callId) + clearTimeout(timeoutId) + resolve({ + success: false, + error: { + code: ErrorCode.BRIDGE_DISCONNECTED, + message: `Failed to send invoke message: ${error}`, + }, + }) + } + }) + } + + /** + * Cancel a tool invocation + */ + cancel(sessionId: string, callId: string, reason?: string): void { + const bridgeId = this.sessionToBridge.get(sessionId) + if (!bridgeId) return + + const connection = this.connections.get(bridgeId) + if (!connection || !connection.connected) return + + const cancel = Message.cancel(sessionId, callId, reason) + try { + connection.socket.send(JSON.stringify(cancel)) + } catch {} + } + + /** + * Get available tools for a session + */ + getTools(sessionId: string): ToolDefinition[] { + const bridgeId = this.sessionToBridge.get(sessionId) + if (!bridgeId) return [] + + const connection = this.connections.get(bridgeId) + if (!connection) return [] + + return Array.from(connection.tools.values()) + } + + /** + * Get all active sessions + */ + getSessions(): Session[] { + return Array.from(this.sessions.values()) + } + + /** + * Check if a session is connected + */ + isConnected(sessionId: string): boolean { + const bridgeId = this.sessionToBridge.get(sessionId) + if (!bridgeId) return false + + const connection = this.connections.get(bridgeId) + return connection?.connected ?? false + } + + /** + * Get relay ID + */ + getRelayId(): string { + return this.relayId + } +} diff --git a/packages/tool-relay/src/types.ts b/packages/tool-relay/src/types.ts new file mode 100644 index 00000000000..d5b8f3741d1 --- /dev/null +++ b/packages/tool-relay/src/types.ts @@ -0,0 +1,104 @@ +import type { ToolDefinition, ResultMessage } from "@opencode-ai/tool-bridge-protocol/schema" + +/** + * Authentication configuration for the relay + */ +export interface AuthConfig { + type: "none" | "bearer" + validateToken?: (token: string) => Promise | boolean +} + +/** + * Relay configuration options + */ +export interface RelayConfig { + /** + * Unique identifier for this relay instance + */ + relayId?: string + + /** + * Authentication configuration + */ + auth?: AuthConfig + + /** + * Heartbeat interval in milliseconds (default: 30000) + */ + heartbeatInterval?: number + + /** + * Tool execution timeout in milliseconds (default: 120000) + */ + executionTimeout?: number + + /** + * Session grace period for reconnection in milliseconds (default: 60000) + */ + sessionGracePeriod?: number + + /** + * Event handlers + */ + onBridgeConnect?: (bridgeId: string, sessionId: string) => void + onBridgeDisconnect?: (bridgeId: string, sessionId: string) => void + onToolRegistered?: (bridgeId: string, tools: ToolDefinition[]) => void + onToolInvoke?: (sessionId: string, callId: string, toolId: string) => void + onToolResult?: (sessionId: string, callId: string, result: ResultMessage["payload"]) => void +} + +/** + * Represents a connected bridge + */ +export interface BridgeConnection { + bridgeId: string + sessionId: string + socket: WebSocket + tools: Map + pendingCalls: Map< + string, + { + resolve: (result: ResultMessage["payload"]) => void + reject: (error: Error) => void + timeoutId: ReturnType + } + > + lastPing: number + connected: boolean +} + +/** + * Represents a session + */ +export interface Session { + sessionId: string + bridgeId: string + createdAt: number + lastActivity: number + tools: ToolDefinition[] +} + +/** + * Tool invocation request + */ +export interface ToolInvocationRequest { + sessionId: string + toolId: string + arguments: Record + traceId?: string +} + +/** + * Tool invocation result + */ +export interface ToolInvocationResult { + success: boolean + output?: string + error?: { + code: string + message: string + details?: Record + } + metadata?: Record + title?: string +} diff --git a/packages/tool-relay/tsconfig.json b/packages/tool-relay/tsconfig.json new file mode 100644 index 00000000000..cb5c12aecdc --- /dev/null +++ b/packages/tool-relay/tsconfig.json @@ -0,0 +1,9 @@ +{ + "$schema": "https://json.schemastore.org/tsconfig", + "extends": "@tsconfig/bun/tsconfig.json", + "compilerOptions": { + "lib": ["ESNext"], + "types": [], + "noUncheckedIndexedAccess": false + } +} From 2125796277442890840c2a57f16225d009921bf8 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 20 Dec 2025 23:33:30 +0000 Subject: [PATCH 2/4] feat: add MCP client support to Tool Bridge Adds Model Context Protocol (MCP) client integration to the Tool Bridge, allowing it to connect to MCP servers and expose their tools through the bridge protocol. New components: - McpClientManager: Manages connections to multiple MCP servers - McpConfig types: Configuration for local (stdio) and remote (HTTP/SSE) MCP servers - Bridge.attachMcpServers(): Convenience method to attach MCP servers Features: - Support for local MCP servers via stdio transport - Support for remote MCP servers via StreamableHTTP and SSE transports - Automatic tool discovery and registration with the bridge - Configurable tool prefix for namespacing MCP tools - Connection timeout and error handling - Ability to add/remove servers dynamically Usage: ```typescript const bridge = new ToolBridge({ url: "ws://relay:8080" }) await bridge.attachMcpServers({ "filesystem": { type: "local", command: ["npx", "-y", "@modelcontextprotocol/server-filesystem", "/path"] }, "api": { type: "remote", url: "https://mcp.example.com/api" } }) await bridge.connect() ``` --- packages/tool-bridge-tests/mcp.test.ts | 301 ++++++++++++++++++++ packages/tool-bridge-tests/package.json | 2 + packages/tool-bridge/package.json | 2 + packages/tool-bridge/src/bridge.ts | 66 +++++ packages/tool-bridge/src/index.ts | 8 + packages/tool-bridge/src/mcp.ts | 364 ++++++++++++++++++++++++ packages/tool-bridge/src/types.ts | 54 ++++ 7 files changed, 797 insertions(+) create mode 100644 packages/tool-bridge-tests/mcp.test.ts create mode 100644 packages/tool-bridge/src/mcp.ts diff --git a/packages/tool-bridge-tests/mcp.test.ts b/packages/tool-bridge-tests/mcp.test.ts new file mode 100644 index 00000000000..c30085f468d --- /dev/null +++ b/packages/tool-bridge-tests/mcp.test.ts @@ -0,0 +1,301 @@ +import { describe, test, expect, beforeAll, afterAll } from "bun:test" +import { ToolRelay } from "@opencode-ai/tool-relay/index" +import { ToolBridge, McpClientManager } from "@opencode-ai/tool-bridge/index" + +/** + * MCP Client Manager Tests + * + * These tests verify the MCP client integration with the Tool Bridge. + * Note: These tests require actual MCP servers to be available, + * so they are marked as integration tests. + */ + +describe("MCP Client Manager", () => { + let relay: ToolRelay + let server: ReturnType + let serverUrl: string + + beforeAll(() => { + relay = new ToolRelay({ + relayId: "test-relay-mcp", + auth: { type: "none" }, + heartbeatInterval: 5000, + executionTimeout: 10000, + }) + + server = Bun.serve({ + port: 0, + fetch(req, server) { + if (server.upgrade(req)) { + return + } + return new Response("Not found", { status: 404 }) + }, + websocket: { + open(ws) { + relay.handleConnection(ws as unknown as WebSocket) + }, + message() {}, + close() {}, + }, + }) + + serverUrl = `ws://localhost:${server.port}` + relay.start() + }) + + afterAll(() => { + relay.stop() + server.stop() + }) + + describe("McpClientManager Initialization", () => { + test("can create MCP client manager", () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "mcp-test-1", + reconnect: { enabled: false }, + }) + + const manager = new McpClientManager(bridge) + expect(manager).toBeDefined() + + bridge.disconnect() + }) + + test("can create with custom tool prefix", () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "mcp-test-2", + reconnect: { enabled: false }, + }) + + const manager = new McpClientManager(bridge, { toolPrefix: "custom" }) + expect(manager).toBeDefined() + + bridge.disconnect() + }) + }) + + describe("Server Configuration", () => { + test("disabled server returns disabled status", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "mcp-test-3", + reconnect: { enabled: false }, + }) + + const manager = new McpClientManager(bridge) + const status = await manager.addServer("disabled-server", { + type: "local", + command: ["echo", "test"], + enabled: false, + }) + + expect(status.status).toBe("disabled") + + const allStatus = manager.getStatus() + expect(allStatus["disabled-server"]).toEqual({ status: "disabled" }) + + bridge.disconnect() + }) + + test("invalid local command returns failed status", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "mcp-test-4", + reconnect: { enabled: false }, + }) + + const manager = new McpClientManager(bridge) + const status = await manager.addServer("invalid-server", { + type: "local", + command: ["nonexistent-command-that-does-not-exist"], + timeout: 1000, + }) + + expect(status.status).toBe("failed") + if (status.status === "failed") { + expect(status.error).toBeTruthy() + } + + bridge.disconnect() + }) + + test("invalid remote URL returns failed status", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "mcp-test-5", + reconnect: { enabled: false }, + }) + + const manager = new McpClientManager(bridge) + const status = await manager.addServer("invalid-remote", { + type: "remote", + url: "http://localhost:99999/nonexistent", + timeout: 1000, + }) + + expect(status.status).toBe("failed") + + bridge.disconnect() + }) + }) + + describe("Bridge Integration", () => { + test("bridge can attach MCP servers", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "mcp-test-6", + reconnect: { enabled: false }, + }) + + // Attach a disabled server (for testing without real MCP) + const statuses = await bridge.attachMcpServers({ + "test-server": { + type: "local", + command: ["echo"], + enabled: false, + }, + }) + + expect(statuses["test-server"]).toEqual({ status: "disabled" }) + + const manager = bridge.getMcpManager() + expect(manager).toBeDefined() + + bridge.disconnect() + }) + + test("bridge exposes MCP manager", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "mcp-test-7", + reconnect: { enabled: false }, + }) + + // Initially no manager + expect(bridge.getMcpManager()).toBeNull() + + // After attaching, manager exists + await bridge.attachMcpServers({}) + expect(bridge.getMcpManager()).toBeDefined() + + bridge.disconnect() + }) + + test("multiple server configuration", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "mcp-test-8", + reconnect: { enabled: false }, + }) + + const statuses = await bridge.attachMcpServers({ + "server-a": { + type: "local", + command: ["echo"], + enabled: false, + }, + "server-b": { + type: "remote", + url: "http://localhost:12345", + enabled: false, + }, + }) + + expect(statuses["server-a"]).toEqual({ status: "disabled" }) + expect(statuses["server-b"]).toEqual({ status: "disabled" }) + + const clients = bridge.getMcpManager()?.getClients() + expect(clients?.length).toBe(2) + + bridge.disconnect() + }) + }) + + describe("Server Management", () => { + test("can remove server", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "mcp-test-9", + reconnect: { enabled: false }, + }) + + const manager = new McpClientManager(bridge) + await manager.addServer("removable", { + type: "local", + command: ["echo"], + enabled: false, + }) + + let status = manager.getStatus() + expect(status["removable"]).toBeDefined() + + await manager.removeServer("removable") + + status = manager.getStatus() + expect(status["removable"]).toBeUndefined() + + bridge.disconnect() + }) + + test("can close all connections", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "mcp-test-10", + reconnect: { enabled: false }, + }) + + const manager = new McpClientManager(bridge) + await manager.addServer("server1", { + type: "local", + command: ["echo"], + enabled: false, + }) + await manager.addServer("server2", { + type: "local", + command: ["echo"], + enabled: false, + }) + + let clients = manager.getClients() + expect(clients.length).toBe(2) + + await manager.close() + + clients = manager.getClients() + expect(clients.length).toBe(0) + + bridge.disconnect() + }) + }) + + describe("Client Info", () => { + test("getClients returns correct info", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "mcp-test-11", + reconnect: { enabled: false }, + }) + + const manager = new McpClientManager(bridge) + await manager.addServer("info-test", { + type: "local", + command: ["echo", "test"], + enabled: false, + }) + + const clients = manager.getClients() + expect(clients.length).toBe(1) + + const info = clients[0] + expect(info.name).toBe("info-test") + expect(info.config.type).toBe("local") + expect(info.status.status).toBe("disabled") + expect(info.toolCount).toBe(0) + + bridge.disconnect() + }) + }) +}) diff --git a/packages/tool-bridge-tests/package.json b/packages/tool-bridge-tests/package.json index 5a07a4d5975..184b066e884 100644 --- a/packages/tool-bridge-tests/package.json +++ b/packages/tool-bridge-tests/package.json @@ -10,6 +10,8 @@ "@opencode-ai/tool-bridge": "workspace:*", "@opencode-ai/tool-relay": "workspace:*", "@opencode-ai/tool-bridge-protocol": "workspace:*", + "@ai-sdk/mcp": "0.0.8", + "@modelcontextprotocol/sdk": "1.15.1", "zod": "catalog:" }, "devDependencies": { diff --git a/packages/tool-bridge/package.json b/packages/tool-bridge/package.json index a9237d1e266..1d62e005922 100644 --- a/packages/tool-bridge/package.json +++ b/packages/tool-bridge/package.json @@ -8,6 +8,8 @@ }, "dependencies": { "@opencode-ai/tool-bridge-protocol": "workspace:*", + "@ai-sdk/mcp": "0.0.8", + "@modelcontextprotocol/sdk": "1.15.1", "zod": "catalog:" }, "devDependencies": { diff --git a/packages/tool-bridge/src/bridge.ts b/packages/tool-bridge/src/bridge.ts index a141aceb755..9c7c39d8419 100644 --- a/packages/tool-bridge/src/bridge.ts +++ b/packages/tool-bridge/src/bridge.ts @@ -16,7 +16,9 @@ import type { ToolResult, RegisteredTool, ConnectionState, + McpConfig, } from "./types" +import type { McpClientManager } from "./mcp" /** * Tool Bridge - Client-side agent for local tool execution @@ -36,6 +38,7 @@ export class ToolBridge { private reconnectAttempt = 0 private reconnectTimer: ReturnType | null = null private heartbeatTimer: ReturnType | null = null + private mcpManager: McpClientManager | null = null constructor(config: BridgeConfig) { this.bridgeId = config.bridgeId ?? generateId() @@ -146,6 +149,9 @@ export class ToolBridge { this.stopHeartbeat() this.cancelAllCalls("Bridge disconnected") + // Close MCP connections (fire and forget) + this.closeMcpConnections().catch(() => {}) + if (this.socket) { try { this.socket.close() @@ -178,6 +184,66 @@ export class ToolBridge { return Array.from(this.tools.values()).map((t) => t.definition) } + /** + * Attach MCP servers and register their tools with this bridge + * + * @param servers - Map of server name to MCP configuration + * @param options - Optional configuration for the MCP manager + * @returns Status of each MCP server connection + * + * @example + * ```typescript + * const bridge = new ToolBridge({ url: "ws://localhost:8080" }) + * + * // Attach MCP servers + * await bridge.attachMcpServers({ + * "filesystem": { + * type: "local", + * command: ["npx", "-y", "@modelcontextprotocol/server-filesystem", "/path"] + * }, + * "github": { + * type: "remote", + * url: "https://mcp.github.com/api" + * } + * }) + * + * // Connect to relay + * await bridge.connect() + * ``` + */ + async attachMcpServers( + servers: Record, + options?: { toolPrefix?: string } + ): Promise> { + // Dynamically import to avoid circular dependency and allow tree-shaking + const { McpClientManager } = await import("./mcp") + + // Create manager if not exists + if (!this.mcpManager) { + this.mcpManager = new McpClientManager(this, options) + } + + // Add all servers + return this.mcpManager.addServers(servers) + } + + /** + * Get the MCP client manager (if attached) + */ + getMcpManager(): McpClientManager | null { + return this.mcpManager + } + + /** + * Close MCP connections when disconnecting + */ + private async closeMcpConnections(): Promise { + if (this.mcpManager) { + await this.mcpManager.close() + this.mcpManager = null + } + } + /** * Send hello message */ diff --git a/packages/tool-bridge/src/index.ts b/packages/tool-bridge/src/index.ts index d289afac8f0..32d96faa9c9 100644 --- a/packages/tool-bridge/src/index.ts +++ b/packages/tool-bridge/src/index.ts @@ -3,9 +3,12 @@ * * The Tool Bridge connects to a Tool Relay and executes tools * locally on behalf of remote LLM runtimes. + * + * Supports MCP (Model Context Protocol) servers for tool discovery. */ export { ToolBridge } from "./bridge" +export { McpClientManager } from "./mcp" export type { BridgeConfig, ToolExecutor, @@ -13,4 +16,9 @@ export type { ToolResult, RegisteredTool, ConnectionState, + McpConfig, + McpLocalConfig, + McpRemoteConfig, + McpStatus, + McpClientInfo, } from "./types" diff --git a/packages/tool-bridge/src/mcp.ts b/packages/tool-bridge/src/mcp.ts new file mode 100644 index 00000000000..7594d896d50 --- /dev/null +++ b/packages/tool-bridge/src/mcp.ts @@ -0,0 +1,364 @@ +import { experimental_createMCPClient } from "@ai-sdk/mcp" +import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js" +import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js" +import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js" +import type { McpConfig, McpStatus, McpClientInfo, ToolExecutor, ToolContext, ToolResult } from "./types" +import type { ToolBridge } from "./bridge" + +/** + * Type for MCP client from AI SDK + */ +type McpClient = Awaited> + +/** + * Internal client record + */ +interface ClientRecord { + name: string + config: McpConfig + client: McpClient + status: McpStatus + toolIds: string[] +} + +/** + * MCP Client Manager + * + * Manages connections to MCP servers and exposes their tools + * through the Tool Bridge. + */ +export class McpClientManager { + private readonly clients: Map = new Map() + private readonly bridge: ToolBridge + private readonly toolPrefix: string + + constructor(bridge: ToolBridge, options?: { toolPrefix?: string }) { + this.bridge = bridge + this.toolPrefix = options?.toolPrefix ?? "mcp" + } + + /** + * Add and connect to an MCP server + */ + async addServer(name: string, config: McpConfig): Promise { + // Check if already exists + if (this.clients.has(name)) { + await this.removeServer(name) + } + + // Skip if disabled + if (config.enabled === false) { + const record: ClientRecord = { + name, + config, + client: null as any, + status: { status: "disabled" }, + toolIds: [], + } + this.clients.set(name, record) + return { status: "disabled" } + } + + // Set connecting status + const record: ClientRecord = { + name, + config, + client: null as any, + status: { status: "connecting" }, + toolIds: [], + } + this.clients.set(name, record) + + try { + const client = await this.createClient(name, config) + record.client = client + record.status = { status: "connected" } + + // Discover and register tools + await this.discoverAndRegisterTools(name, client) + + return { status: "connected" } + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error) + record.status = { status: "failed", error: errorMessage } + return { status: "failed", error: errorMessage } + } + } + + /** + * Add multiple MCP servers from configuration + */ + async addServers(servers: Record): Promise> { + const results: Record = {} + + // Connect to all servers in parallel + await Promise.all( + Object.entries(servers).map(async ([name, config]) => { + results[name] = await this.addServer(name, config) + }) + ) + + return results + } + + /** + * Remove an MCP server + */ + async removeServer(name: string): Promise { + const record = this.clients.get(name) + if (!record) return + + // Unregister tools from bridge + for (const toolId of record.toolIds) { + this.bridge.unregisterTool(toolId) + } + + // Close client + if (record.client) { + try { + await record.client.close() + } catch { + // Ignore close errors + } + } + + this.clients.delete(name) + } + + /** + * Get status of all MCP clients + */ + getStatus(): Record { + const result: Record = {} + for (const [name, record] of this.clients) { + result[name] = record.status + } + return result + } + + /** + * Get detailed info about all MCP clients + */ + getClients(): McpClientInfo[] { + const result: McpClientInfo[] = [] + for (const [name, record] of this.clients) { + result.push({ + name, + config: record.config, + status: record.status, + toolCount: record.toolIds.length, + }) + } + return result + } + + /** + * Refresh tools from all connected MCP servers + */ + async refreshTools(): Promise { + for (const [name, record] of this.clients) { + if (record.status.status !== "connected" || !record.client) continue + + // Unregister old tools + for (const toolId of record.toolIds) { + this.bridge.unregisterTool(toolId) + } + record.toolIds = [] + + // Discover and register new tools + await this.discoverAndRegisterTools(name, record.client) + } + } + + /** + * Close all MCP connections + */ + async close(): Promise { + const names = Array.from(this.clients.keys()) + await Promise.all(names.map((name) => this.removeServer(name))) + } + + /** + * Create an MCP client based on config type + */ + private async createClient(name: string, config: McpConfig): Promise { + if (config.type === "local") { + return this.createLocalClient(name, config) + } else { + return this.createRemoteClient(name, config) + } + } + + /** + * Create a local (stdio) MCP client + */ + private async createLocalClient( + name: string, + config: Extract + ): Promise { + const [cmd, ...args] = config.command + + const transport = new StdioClientTransport({ + stderr: "ignore", + command: cmd, + args, + env: { + ...process.env, + ...config.environment, + }, + }) + + const client = await experimental_createMCPClient({ + name: `bridge-${name}`, + transport, + }) + + // Verify we can get tools + const timeout = config.timeout ?? 5000 + await this.withTimeout(client.tools(), timeout) + + return client + } + + /** + * Create a remote (HTTP/SSE) MCP client + */ + private async createRemoteClient( + name: string, + config: Extract + ): Promise { + // Try StreamableHTTP first, then fall back to SSE + const transports = [ + { + name: "StreamableHTTP", + transport: new StreamableHTTPClientTransport(new URL(config.url), { + requestInit: config.headers ? { headers: config.headers } : undefined, + }), + }, + { + name: "SSE", + transport: new SSEClientTransport(new URL(config.url), { + requestInit: config.headers ? { headers: config.headers } : undefined, + }), + }, + ] + + let lastError: Error | undefined + + for (const { transport } of transports) { + try { + const client = await experimental_createMCPClient({ + name: `bridge-${name}`, + transport, + }) + + // Verify we can get tools + const timeout = config.timeout ?? 5000 + await this.withTimeout(client.tools(), timeout) + + return client + } catch (error) { + lastError = error instanceof Error ? error : new Error(String(error)) + } + } + + throw lastError ?? new Error("Failed to connect to remote MCP server") + } + + /** + * Discover tools from an MCP client and register them with the bridge + */ + private async discoverAndRegisterTools(name: string, client: McpClient): Promise { + const record = this.clients.get(name) + if (!record) return + + const tools = await client.tools() + const toolIds: string[] = [] + + for (const [toolName, tool] of Object.entries(tools)) { + // Create a unique tool ID with prefix + const sanitizedName = name.replace(/[^a-zA-Z0-9_-]/g, "_") + const sanitizedToolName = toolName.replace(/[^a-zA-Z0-9_-]/g, "_") + const toolId = `${this.toolPrefix}_${sanitizedName}_${sanitizedToolName}` + + // Create executor that invokes the MCP tool + const executor = this.createMcpToolExecutor(tool) + + // Get tool parameters schema + const parameters = (tool as any).parameters ?? { type: "object" } + + // Register with bridge + this.bridge.registerTool( + toolId, + (tool as any).description ?? `MCP tool: ${toolName}`, + parameters, + executor + ) + + toolIds.push(toolId) + } + + record.toolIds = toolIds + } + + /** + * Create an executor function for an MCP tool + */ + private createMcpToolExecutor(tool: any): ToolExecutor { + return async (args: Record, ctx: ToolContext): Promise => { + // Check if aborted + if (ctx.abort.aborted) { + throw new Error("Tool invocation was cancelled") + } + + // Execute the MCP tool + const result = await tool.execute(args, { + abortSignal: ctx.abort, + }) + + // Process the result content + const textParts: string[] = [] + + if (result.content) { + for (const item of result.content) { + if (item.type === "text") { + textParts.push(item.text) + } else if (item.type === "image") { + // Send image data as chunk metadata + ctx.sendChunk("", { + type: "image", + mimeType: item.mimeType, + data: item.data, + }) + } + } + } + + return { + output: textParts.join("\n\n"), + metadata: result.metadata, + } + } + } + + /** + * Execute a promise with timeout + */ + private withTimeout(promise: Promise, ms: number): Promise { + return new Promise((resolve, reject) => { + const timeoutId = setTimeout(() => { + reject(new Error(`Operation timed out after ${ms}ms`)) + }, ms) + + promise + .then((result) => { + clearTimeout(timeoutId) + resolve(result) + }) + .catch((error) => { + clearTimeout(timeoutId) + reject(error) + }) + }) + } +} diff --git a/packages/tool-bridge/src/types.ts b/packages/tool-bridge/src/types.ts index a1ba21726a3..7e25967f2d9 100644 --- a/packages/tool-bridge/src/types.ts +++ b/packages/tool-bridge/src/types.ts @@ -1,5 +1,59 @@ import type { ToolDefinition } from "@opencode-ai/tool-bridge-protocol/schema" +/** + * MCP server configuration - local (stdio) transport + */ +export interface McpLocalConfig { + type: "local" + /** Command and arguments to run the MCP server */ + command: string[] + /** Environment variables for the MCP server process */ + environment?: Record + /** Enable or disable this MCP server */ + enabled?: boolean + /** Timeout for fetching tools (default: 5000ms) */ + timeout?: number +} + +/** + * MCP server configuration - remote (HTTP/SSE) transport + */ +export interface McpRemoteConfig { + type: "remote" + /** URL of the remote MCP server */ + url: string + /** Enable or disable this MCP server */ + enabled?: boolean + /** Additional headers to send with requests */ + headers?: Record + /** Timeout for fetching tools (default: 5000ms) */ + timeout?: number +} + +/** + * MCP server configuration + */ +export type McpConfig = McpLocalConfig | McpRemoteConfig + +/** + * MCP client status + */ +export type McpStatus = + | { status: "connected" } + | { status: "disabled" } + | { status: "failed"; error: string } + | { status: "connecting" } + +/** + * MCP client information + */ +export interface McpClientInfo { + name: string + config: McpConfig + status: McpStatus + toolCount: number +} + /** * Bridge configuration options */ From 5e18b514c485fab0f42c08feb3737cf351408c51 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 21 Dec 2025 03:15:20 +0000 Subject: [PATCH 3/4] docs: add comprehensive Tool Bridge design documentation Adds a detailed design document explaining the Tool Bridge and Tool Relay architecture in narrative style with ASCII diagrams. The document covers: - Problem statement and motivation for reverse RPC - High-level architecture with visual diagrams - Core concepts (tools, sessions, calls, multiplexing) - Protocol message types and flow diagrams - Component deep dive (Relay and Bridge) - MCP integration for external tool servers - Connection lifecycle and reconnection handling - Tool invocation flow with streaming and cancellation - Error handling and recovery scenarios - Security considerations - Code references to actual implementation files Written for first-time readers and junior engineers without requiring deep understanding of LLM agents. --- docs/tool-bridge-design.md | 1085 ++++++++++++++++++++++++++++++++++++ 1 file changed, 1085 insertions(+) create mode 100644 docs/tool-bridge-design.md diff --git a/docs/tool-bridge-design.md b/docs/tool-bridge-design.md new file mode 100644 index 00000000000..2b1361084aa --- /dev/null +++ b/docs/tool-bridge-design.md @@ -0,0 +1,1085 @@ +# Tool Bridge & Tool Relay: A Complete Design Guide + +## Introduction + +Imagine you're building an AI assistant that can help developers write code. The AI runs on a powerful server in the cloud, but it needs to do things on a developer's laptop—like reading files, running commands, or searching through code. The problem? The developer's laptop is behind a firewall. The server can't reach in. + +This document explains how we solve this problem using two components: the **Tool Bridge** and the **Tool Relay**. Together, they create a secure tunnel that lets cloud-based AI invoke tools on local machines. + +--- + +## Table of Contents + +1. [The Problem We're Solving](#the-problem-were-solving) +2. [High-Level Architecture](#high-level-architecture) +3. [Core Concepts](#core-concepts) +4. [The Protocol](#the-protocol) +5. [Component Deep Dive](#component-deep-dive) +6. [MCP Integration](#mcp-integration) +7. [Connection Lifecycle](#connection-lifecycle) +8. [Tool Invocation Flow](#tool-invocation-flow) +9. [Error Handling & Recovery](#error-handling--recovery) +10. [Security Considerations](#security-considerations) +11. [Code References](#code-references) + +--- + +## The Problem We're Solving + +### The Scenario + +``` +┌─────────────────────────────────────────────────────────────────────────┐ +│ THE CLOUD │ +│ ┌─────────────────────────────────────────────────────────────────┐ │ +│ │ LLM Agent Runtime │ │ +│ │ │ │ +│ │ "I need to read the file /src/main.ts on the user's laptop" │ │ +│ │ │ │ +│ └─────────────────────────────────────────────────────────────────┘ │ +│ │ │ +│ │ ??? How do we reach the laptop? │ +│ ▼ │ +└─────────────────────────────────────────────────────────────────────────┘ + │ + ┌──────┴──────┐ + │ FIREWALL │ ← Blocks incoming connections + │ NAT │ + └──────┬──────┘ + │ +┌─────────────────────────────────────────────────────────────────────────┐ +│ DEVELOPER'S LAPTOP │ +│ ┌─────────────────────────────────────────────────────────────────┐ │ +│ │ Local Files & Tools │ │ +│ │ │ │ +│ │ /src/main.ts bash commands git npm etc. │ │ +│ │ │ │ +│ └─────────────────────────────────────────────────────────────────┘ │ +└─────────────────────────────────────────────────────────────────────────┘ +``` + +### Why Traditional Approaches Don't Work + +1. **Direct Connection**: The server can't connect to the laptop because of NAT/firewall +2. **VPN/Port Forwarding**: Requires complex setup, security risks, not user-friendly +3. **Polling API**: The laptop would need to constantly ask "any work for me?"—inefficient and slow + +### Our Solution: Reverse the Connection + +What if the laptop initiated the connection? Outbound connections from laptops work fine. Once connected, communication can flow both ways. This is called **Reverse RPC**. + +--- + +## High-Level Architecture + +``` +┌─────────────────────────────────────────────────────────────────────────────┐ +│ SERVER │ +│ │ +│ ┌────────────────┐ ┌────────────────────────────────────────┐ │ +│ │ │ │ │ │ +│ │ LLM Agent │◄───────►│ TOOL RELAY │ │ +│ │ Runtime │ invoke │ │ │ +│ │ │ result │ • Accepts WebSocket connections │ │ +│ └────────────────┘ │ • Routes tool calls to bridges │ │ +│ │ • Manages sessions │ │ +│ │ │ │ +│ └──────────────────┬─────────────────────┘ │ +│ │ │ +│ │ WebSocket │ +│ │ (persistent) │ +└─────────────────────────────────────────────────┼───────────────────────────┘ + │ + ══════════════╪══════════════ + Internet │ (outbound OK) + ══════════════╪══════════════ + │ +┌─────────────────────────────────────────────────┼───────────────────────────┐ +│ LOCAL MACHINE │ │ +│ ▼ │ +│ ┌─────────────────────────────────────────────────────────────────────┐ │ +│ │ TOOL BRIDGE │ │ +│ │ │ │ +│ │ • Initiates WebSocket connection (outbound) │ │ +│ │ • Registers available tools │ │ +│ │ • Executes tool calls locally │ │ +│ │ • Sends results back │ │ +│ │ │ │ +│ └──────────────────────────────┬──────────────────────────────────────┘ │ +│ │ │ +│ ┌───────────────┼───────────────┐ │ +│ ▼ ▼ ▼ │ +│ ┌──────────┐ ┌──────────┐ ┌──────────┐ │ +│ │ bash │ │ read │ │ edit │ ... │ +│ │ tool │ │ tool │ │ tool │ │ +│ └──────────┘ └──────────┘ └──────────┘ │ +│ │ +│ ┌──────────────────────────────────────────┐ │ +│ │ MCP SERVERS │ │ +│ │ (filesystem, database, APIs, etc.) │ │ +│ └──────────────────────────────────────────┘ │ +│ │ +└─────────────────────────────────────────────────────────────────────────────┘ +``` + +### Key Insight + +The connection is **initiated by the client** (Tool Bridge) but **used by the server** (Tool Relay) to make requests. This is why it's called "Reverse RPC"—the server calls functions on the client. + +--- + +## Core Concepts + +Before diving into the code, let's understand the key concepts: + +### Tool + +A **tool** is a function that the AI can invoke. Examples: +- `read`: Read a file from disk +- `bash`: Execute a shell command +- `edit`: Modify a file +- `glob`: Find files matching a pattern + +Each tool has: +- **ID**: Unique identifier (e.g., `"bash"`, `"read"`) +- **Description**: What the tool does (for the AI to understand) +- **Parameters**: What inputs it accepts (as a JSON schema) +- **Executor**: The actual function that runs + +### Session + +A **session** is a logical context tied to an LLM conversation. It: +- Survives temporary disconnections +- Has a unique ID +- Tracks which tools are available +- Allows the relay to route calls correctly + +### Call + +A **call** is a single tool invocation. It has: +- **Call ID**: Unique identifier for this specific invocation +- **Tool ID**: Which tool to invoke +- **Arguments**: The parameters to pass +- **Result**: Success/failure plus output + +### Multiplexing + +**Multiplexing** means running multiple calls at the same time over one connection. The AI might want to read three files simultaneously—we don't want to wait for each one. + +``` +Time → +───────────────────────────────────────────────────────────────► + +Bridge Relay + │ │ + │◄─────────── invoke: read file1 (call-1) ───────────────│ + │◄─────────── invoke: read file2 (call-2) ───────────────│ + │◄─────────── invoke: run "ls" (call-3) ─────────────────│ + │ │ + │ (all three execute in parallel locally) │ + │ │ + │──────────── result: call-2 (file2 content) ───────────►│ + │──────────── result: call-3 (ls output) ───────────────►│ + │──────────── result: call-1 (file1 content) ───────────►│ + │ │ +``` + +--- + +## The Protocol + +The Tool Bridge Protocol defines how messages are structured and exchanged. + +### Message Envelope + +Every message follows this structure: + +```typescript +{ + type: "hello" | "welcome" | "invoke" | "result" | ... , + id: "unique-message-id", + sessionId: "session-id", // optional for some messages + timestamp: "2024-01-15T10:30:00Z", + payload: { ... } // message-specific data +} +``` + +**File Reference**: [`packages/tool-bridge-protocol/src/schema.ts`](../packages/tool-bridge-protocol/src/schema.ts) + +### Message Types + +#### Handshake Messages + +| Message | Direction | Purpose | +|---------|-----------|---------| +| `hello` | Bridge → Relay | "Hi, I'm bridge X with these credentials" | +| `welcome` | Relay → Bridge | "Welcome! Your session ID is Y" | +| `error` | Relay → Bridge | "Something went wrong" | + +#### Tool Messages + +| Message | Direction | Purpose | +|---------|-----------|---------| +| `register_tools` | Bridge → Relay | "I can execute these tools" | +| `tools_registered` | Relay → Bridge | "Got it, tools registered" | +| `invoke` | Relay → Bridge | "Please run tool X with args Y" | +| `result` | Bridge → Relay | "Here's the output" | + +#### Streaming Messages + +| Message | Direction | Purpose | +|---------|-----------|---------| +| `chunk` | Bridge → Relay | "Partial output while still running" | +| `done` | Bridge → Relay | "Finished streaming" | + +#### Control Messages + +| Message | Direction | Purpose | +|---------|-----------|---------| +| `cancel` | Relay → Bridge | "Stop executing call X" | +| `cancelled` | Bridge → Relay | "Call X has been stopped" | +| `ping` | Relay → Bridge | "Are you still there?" | +| `pong` | Bridge → Relay | "Yes, I'm here" | + +### Example: Complete Flow + +``` +Bridge Relay + │ │ + │ ──── WebSocket Connection Established ────────────────────► │ + │ │ + │ ──── hello { bridgeId: "my-bridge", version: "1.0" } ─────► │ + │ │ + │ ◄─── welcome { sessionId: "sess-123", relayId: "relay-1" } ── │ + │ │ + │ ──── register_tools { tools: [ │ + │ { id: "bash", description: "...", parameters: {...} }, │ + │ { id: "read", description: "...", parameters: {...} } │ + │ ] } ─────────────────────────────────────────────────► │ + │ │ + │ ◄─── tools_registered { toolIds: ["bash", "read"] } ───────── │ + │ │ + │ ... ready for invocations ... │ + │ │ + │ ◄─── invoke { callId: "call-1", toolId: "read", │ + │ arguments: { filePath: "/src/main.ts" } } ───── │ + │ │ + │ (Bridge reads the file locally) │ + │ │ + │ ──── result { callId: "call-1", success: true, │ + │ output: "const x = 1;..." } ────────────────► │ + │ │ +``` + +--- + +## Component Deep Dive + +### Tool Relay (Server Side) + +The Tool Relay is the server-side gateway. Think of it as a switchboard operator—it knows which bridges are connected and routes calls to the right one. + +**File Reference**: [`packages/tool-relay/src/relay.ts`](../packages/tool-relay/src/relay.ts) + +#### Key Responsibilities + +``` +┌─────────────────────────────────────────────────────────────────────────┐ +│ TOOL RELAY │ +│ │ +│ ┌──────────────────┐ ┌──────────────────┐ ┌──────────────────────┐ │ +│ │ │ │ │ │ │ │ +│ │ Connection │ │ Session │ │ Invocation │ │ +│ │ Manager │ │ Manager │ │ Router │ │ +│ │ │ │ │ │ │ │ +│ │ • Accept WS │ │ • Create/find │ │ • Match session │ │ +│ │ • Authenticate │ │ sessions │ │ to bridge │ │ +│ │ • Handle close │ │ • Track tools │ │ • Forward invoke │ │ +│ │ • Heartbeat │ │ • Grace period │ │ • Collect result │ │ +│ │ │ │ │ │ │ │ +│ └──────────────────┘ └──────────────────┘ └──────────────────────┘ │ +│ │ +│ Data Structures: │ +│ ───────────────── │ +│ connections: Map │ +│ sessions: Map │ +│ sessionToBridge: Map │ +│ │ +└─────────────────────────────────────────────────────────────────────────┘ +``` + +#### Creating a Relay + +```typescript +import { ToolRelay } from "@opencode-ai/tool-relay" + +const relay = new ToolRelay({ + relayId: "my-relay", + auth: { type: "bearer", validateToken: async (token) => token === "secret" }, + heartbeatInterval: 30000, // Ping every 30 seconds + executionTimeout: 120000, // 2 minute timeout for tool calls + sessionGracePeriod: 60000, // Keep session alive 1 min after disconnect +}) + +relay.start() +``` + +#### Handling Connections + +The relay doesn't create its own HTTP server—it integrates with yours: + +```typescript +import { Hono } from "hono" + +const app = new Hono() + +// Your WebSocket server calls this when a connection is established +Bun.serve({ + fetch(req, server) { + if (server.upgrade(req)) return + return new Response("Not found", { status: 404 }) + }, + websocket: { + open(ws) { + relay.handleConnection(ws) // Hand off to relay + }, + }, +}) +``` + +#### Invoking Tools + +From your LLM runtime, call tools like this: + +```typescript +const result = await relay.invoke({ + sessionId: "sess-123", + toolId: "bash", + arguments: { command: "ls -la", description: "List files" }, +}) + +if (result.success) { + console.log("Output:", result.output) +} else { + console.error("Error:", result.error.message) +} +``` + +### Tool Bridge (Client Side) + +The Tool Bridge runs on the developer's machine. It connects to the relay and executes tools locally. + +**File Reference**: [`packages/tool-bridge/src/bridge.ts`](../packages/tool-bridge/src/bridge.ts) + +#### Key Responsibilities + +``` +┌─────────────────────────────────────────────────────────────────────────┐ +│ TOOL BRIDGE │ +│ │ +│ ┌──────────────────┐ ┌──────────────────┐ ┌──────────────────────┐ │ +│ │ │ │ │ │ │ │ +│ │ Connection │ │ Tool │ │ Execution │ │ +│ │ Manager │ │ Registry │ │ Engine │ │ +│ │ │ │ │ │ │ │ +│ │ • Connect WS │ │ • Register │ │ • Run tool func │ │ +│ │ • Auto-reconnect│ │ tools │ │ • Handle abort │ │ +│ │ • Send hello │ │ • Store schema │ │ • Stream chunks │ │ +│ │ • Handle pings │ │ • Unregister │ │ • Return result │ │ +│ │ │ │ │ │ │ │ +│ └──────────────────┘ └──────────────────┘ └──────────────────────┘ │ +│ │ +│ ┌──────────────────────────────────────────────────────────────────┐ │ +│ │ MCP CLIENT MANAGER │ │ +│ │ │ │ +│ │ • Connect to MCP servers (local stdio / remote HTTP) │ │ +│ │ • Discover tools from MCP servers │ │ +│ │ • Auto-register MCP tools with bridge │ │ +│ │ │ │ +│ └──────────────────────────────────────────────────────────────────┘ │ +│ │ +└─────────────────────────────────────────────────────────────────────────┘ +``` + +#### Creating a Bridge + +```typescript +import { ToolBridge } from "@opencode-ai/tool-bridge" + +const bridge = new ToolBridge({ + url: "wss://relay.example.com/ws", + bridgeId: "developer-laptop-1", + credentials: { type: "bearer", token: "my-secret-token" }, + reconnect: { + enabled: true, + maxAttempts: 10, + baseDelay: 1000, + maxDelay: 30000, + }, + onConnect: (sessionId) => console.log("Connected! Session:", sessionId), + onDisconnect: (reason) => console.log("Disconnected:", reason), +}) +``` + +#### Registering Tools + +```typescript +// Register a simple tool +bridge.registerTool( + "echo", // Tool ID + "Echoes back the input message", // Description + { // JSON Schema for parameters + type: "object", + properties: { + message: { type: "string", description: "Message to echo" } + }, + required: ["message"] + }, + async (args, context) => { // Executor function + return { + output: `Echo: ${args.message}`, + title: "Echo Result", + } + } +) +``` + +#### The Executor Context + +When a tool executes, it receives a context object: + +```typescript +interface ToolContext { + sessionId: string // Current session + callId: string // Unique ID for this invocation + traceId?: string // For distributed tracing + abort: AbortSignal // Signals if the call was cancelled + sendChunk: (data, meta) => void // Send streaming updates +} + +// Using context in a tool +bridge.registerTool("long-task", "...", {...}, + async (args, ctx) => { + // Check for cancellation + if (ctx.abort.aborted) { + throw new Error("Cancelled") + } + + // Stream progress updates + ctx.sendChunk("Starting...", { progress: 0 }) + await doSomeWork() + ctx.sendChunk("Halfway...", { progress: 50 }) + await doMoreWork() + + return { output: "Done!", metadata: { progress: 100 } } + } +) +``` + +#### Connecting + +```typescript +// Connect to the relay +const sessionId = await bridge.connect() +console.log("Connected with session:", sessionId) + +// Later, disconnect +bridge.disconnect() +``` + +--- + +## MCP Integration + +[Model Context Protocol (MCP)](https://modelcontextprotocol.io/) is a standard for AI tools. The Tool Bridge can connect to MCP servers and expose their tools. + +**File Reference**: [`packages/tool-bridge/src/mcp.ts`](../packages/tool-bridge/src/mcp.ts) + +### What is MCP? + +MCP servers are standalone programs that provide tools. For example: +- **Filesystem MCP**: Read/write files, list directories +- **Database MCP**: Query SQL databases +- **GitHub MCP**: Interact with GitHub repos + +``` +┌────────────────────────────────────────────────────────────────────────┐ +│ TOOL BRIDGE │ +│ │ +│ ┌─────────────────────────────────────────────────────────────────┐ │ +│ │ MCP CLIENT MANAGER │ │ +│ │ │ │ +│ │ Manages connections to multiple MCP servers │ │ +│ │ │ │ +│ └───────────┬─────────────────────────────────┬───────────────────┘ │ +│ │ │ │ +│ ▼ ▼ │ +│ ┌───────────────────────┐ ┌───────────────────────────┐ │ +│ │ LOCAL MCP SERVER │ │ REMOTE MCP SERVER │ │ +│ │ (stdio transport) │ │ (HTTP/SSE transport) │ │ +│ │ │ │ │ │ +│ │ Command: npx -y │ │ URL: https://mcp.io/api │ │ +│ │ @mcp/server-fs / │ │ │ │ +│ │ │ │ │ │ +│ │ Tools: │ │ Tools: │ │ +│ │ • read_file │ │ • search_code │ │ +│ │ • write_file │ │ • create_issue │ │ +│ │ • list_directory │ │ • list_repos │ │ +│ │ │ │ │ │ +│ └───────────────────────┘ └───────────────────────────┘ │ +│ │ +│ All MCP tools are automatically registered with the bridge │ +│ with prefix: mcp_{server}_{tool} │ +│ │ +│ Example: mcp_filesystem_read_file, mcp_github_create_issue │ +│ │ +└────────────────────────────────────────────────────────────────────────┘ +``` + +### Attaching MCP Servers + +```typescript +const bridge = new ToolBridge({ url: "wss://relay.example.com/ws" }) + +// Attach MCP servers before connecting +const statuses = await bridge.attachMcpServers({ + // Local MCP server (runs as a subprocess) + "filesystem": { + type: "local", + command: ["npx", "-y", "@modelcontextprotocol/server-filesystem", "/home/user"], + environment: { DEBUG: "true" }, + timeout: 5000, + }, + + // Remote MCP server (connects over HTTP) + "github": { + type: "remote", + url: "https://mcp.github.com/api", + headers: { "Authorization": "Bearer ghp_xxxxx" }, + }, + + // Disabled server (won't connect) + "experimental": { + type: "local", + command: ["./my-experimental-mcp"], + enabled: false, + }, +}) + +console.log(statuses) +// { +// filesystem: { status: "connected" }, +// github: { status: "connected" }, +// experimental: { status: "disabled" } +// } + +// Now connect to relay - MCP tools are automatically available +await bridge.connect() +``` + +### How MCP Tools Are Registered + +When an MCP server connects, the bridge: + +1. **Discovers tools**: Calls `client.tools()` to get available tools +2. **Creates executors**: Wraps each MCP tool in a bridge-compatible executor +3. **Registers with bridge**: Adds tools with prefixed names + +```typescript +// Inside McpClientManager.discoverAndRegisterTools() + +const tools = await client.tools() + +for (const [toolName, tool] of Object.entries(tools)) { + // Create unique ID: mcp_{server}_{tool} + const toolId = `mcp_${serverName}_${toolName}` + + // Create executor that calls the MCP tool + const executor = async (args, ctx) => { + const result = await tool.execute(args, { abortSignal: ctx.abort }) + return { output: result.content.map(c => c.text).join("\n") } + } + + // Register with bridge + bridge.registerTool(toolId, tool.description, tool.parameters, executor) +} +``` + +### Managing MCP Connections + +```typescript +const manager = bridge.getMcpManager() + +// Get status of all MCP servers +console.log(manager.getStatus()) +// { filesystem: { status: "connected" }, github: { status: "failed", error: "..." } } + +// Get detailed info +console.log(manager.getClients()) +// [ +// { name: "filesystem", status: {...}, toolCount: 5, config: {...} }, +// { name: "github", status: {...}, toolCount: 0, config: {...} } +// ] + +// Refresh tools (re-discover from all servers) +await manager.refreshTools() + +// Remove a server +await manager.removeServer("github") + +// Close all MCP connections +await manager.close() +``` + +--- + +## Connection Lifecycle + +Understanding how connections are established and maintained is crucial for debugging. + +### Happy Path + +``` +┌─────────────────────────────────────────────────────────────────────────┐ +│ CONNECTION LIFECYCLE │ +└─────────────────────────────────────────────────────────────────────────┘ + + Bridge Relay + │ │ + │ ┌────────────────────────────────────────────────────────────┐ │ + │ │ 1. CONNECTING │ │ + │ │ Bridge creates WebSocket to relay URL │ │ + │ └────────────────────────────────────────────────────────────┘ │ + │ │ + │ ══════════ WebSocket Connection Opens ═══════════════════════► │ + │ │ + │ ┌────────────────────────────────────────────────────────────┐ │ + │ │ 2. AUTHENTICATING │ │ + │ │ Bridge sends hello with credentials │ │ + │ └────────────────────────────────────────────────────────────┘ │ + │ │ + │ ──── hello { bridgeId, credentials } ─────────────────────────► │ + │ │ + │ Relay validates credentials │ + │ Creates or finds session │ + │ │ + │ ◄─── welcome { sessionId, relayId } ───────────────────────── │ + │ │ + │ ┌────────────────────────────────────────────────────────────┐ │ + │ │ 3. REGISTERING TOOLS │ │ + │ │ Bridge tells relay what tools it has │ │ + │ └────────────────────────────────────────────────────────────┘ │ + │ │ + │ ──── register_tools { tools: [...] } ─────────────────────────► │ + │ │ + │ ◄─── tools_registered { toolIds: [...] } ──────────────────── │ + │ │ + │ ┌────────────────────────────────────────────────────────────┐ │ + │ │ 4. CONNECTED │ │ + │ │ Ready for tool invocations │ │ + │ └────────────────────────────────────────────────────────────┘ │ + │ │ + │ ◄════════════════ Tool Invocations ═══════════════════════════ │ + │ │ + │ ┌────────────────────────────────────────────────────────────┐ │ + │ │ 5. HEARTBEAT │ │ + │ │ Periodic pings to detect dead connections │ │ + │ └────────────────────────────────────────────────────────────┘ │ + │ │ + │ ◄─── ping ─────────────────────────────────────────────────── │ + │ ──── pong ───────────────────────────────────────────────────► │ + │ │ +``` + +### Disconnection & Reconnection + +``` +┌─────────────────────────────────────────────────────────────────────────┐ +│ RECONNECTION FLOW │ +└─────────────────────────────────────────────────────────────────────────┘ + + Bridge Relay + │ │ + │ ═══════════ Connection Lost (network issue) ════════════════╳ │ + │ │ + │ ┌────────────────────────────────────────────────────────────┐ │ + │ │ Bridge: DISCONNECTED │ │ + │ │ • Fail pending calls │ │ + │ │ • Stop heartbeat │ │ + │ │ • Schedule reconnect (exponential backoff) │ │ + │ └────────────────────────────────────────────────────────────┘ │ + │ │ + │ Wait 1s │ + │ ┌────────────────────────────────────────────────────────────┐ │ + │ │ Bridge: RECONNECTING (attempt 1) │ │ + │ └────────────────────────────────────────────────────────────┘ │ + │ │ + │ ══════════ WebSocket Connection Opens ════════════════════════► │ + │ │ + │ ──── hello { bridgeId (same as before) } ─────────────────────► │ + │ │ + │ ┌────────────────────────────────────────────────────────────┐ │ + │ │ Relay: Finds existing session for this bridgeId │ │ + │ │ Re-binds session to new connection │ │ + │ └────────────────────────────────────────────────────────────┘ │ + │ │ + │ ◄─── welcome { sessionId (same as before) } ───────────────── │ + │ │ + │ ──── register_tools { ... } ──────────────────────────────────► │ + │ │ + │ ┌────────────────────────────────────────────────────────────┐ │ + │ │ Bridge: CONNECTED │ │ + │ │ Session preserved! │ │ + │ └────────────────────────────────────────────────────────────┘ │ + │ │ +``` + +### Reconnection Backoff + +The bridge uses exponential backoff with jitter to avoid thundering herd: + +``` +Attempt 1: Wait ~1 second (1000ms base) +Attempt 2: Wait ~2 seconds (2000ms) +Attempt 3: Wait ~4 seconds (4000ms) +Attempt 4: Wait ~8 seconds (8000ms) +... +Attempt N: Wait min(base * 2^(N-1), maxDelay) + random jitter +``` + +**File Reference**: [`packages/tool-bridge/src/bridge.ts`](../packages/tool-bridge/src/bridge.ts) (lines 469-494) + +--- + +## Tool Invocation Flow + +Let's trace a complete tool invocation from start to finish. + +### The Journey of a Tool Call + +``` +┌─────────────────────────────────────────────────────────────────────────┐ +│ TOOL INVOCATION FLOW │ +└─────────────────────────────────────────────────────────────────────────┘ + + LLM Runtime Tool Relay Tool Bridge Local FS + │ │ │ │ + │ 1. "Read file │ │ │ + │ /src/main.ts" │ │ │ + │ │ │ │ + │ ──invoke({ │ │ │ + │ sessionId, │ │ │ + │ toolId: "read", │ │ │ + │ args: {path:...} │ │ │ + │ })─────────────────► │ │ + │ │ │ │ + │ │ 2. Find bridge │ │ + │ │ for session │ │ + │ │ │ │ + │ │ 3. Create call │ │ + │ │ with timeout │ │ + │ │ │ │ + │ │ ──invoke { │ │ + │ │ callId: "c-123", │ │ + │ │ toolId: "read", │ │ + │ │ args: {...} │ │ + │ │ }──────────────────►│ │ + │ │ │ │ + │ │ │ 4. Find tool │ + │ │ │ executor │ + │ │ │ │ + │ │ │ 5. Execute │ + │ │ │ ───read file───►│ + │ │ │ │ + │ │ │ ◄──file data────│ + │ │ │ │ + │ │ │ 6. Format │ + │ │ │ result │ + │ │ │ │ + │ │ ◄──result { │ │ + │ │ callId: "c-123", │ │ + │ │ success: true, │ │ + │ │ output: "..." │ │ + │ │ }───────────────────│ │ + │ │ │ │ + │ │ 7. Resolve │ │ + │ │ pending call │ │ + │ │ │ │ + │ ◄──{ success: true, │ │ │ + │ output: "const │ │ │ + │ x = 1;..." }─────│ │ │ + │ │ │ │ +``` + +### Streaming Results + +For long-running tools, the bridge can stream partial results: + +``` + LLM Runtime Tool Relay Tool Bridge + │ │ │ + │ ──invoke bash │ │ + │ "npm install"──────►──invoke──────────────►│ + │ │ │ + │ │ │ (npm starts) + │ │ │ + │ │ ◄──chunk { │ + │ │ data: "Installing" │ + │ │ }───────────────────│ + │ │ │ + │ │ ◄──chunk { │ + │ │ data: "50%..." │ + │ │ }───────────────────│ + │ │ │ + │ │ ◄──chunk { │ + │ │ data: "100%" │ + │ │ }───────────────────│ + │ │ │ + │ │ ◄──result { │ + │ │ success: true │ + │ │ }───────────────────│ + │ │ │ +``` + +### Cancellation + +If the LLM decides to cancel a long-running tool: + +``` + LLM Runtime Tool Relay Tool Bridge + │ │ │ + │ ──invoke bash │ │ + │ "sleep 60"─────────►──invoke──────────────►│ + │ │ │ + │ │ │ (sleeping...) + │ │ │ + │ ──cancel(callId)────►│ │ + │ │ │ + │ │ ──cancel { │ + │ │ callId: "..." │ + │ │ }──────────────────►│ + │ │ │ + │ │ │ (abort signal) + │ │ │ + │ │ ◄──cancelled { │ + │ │ callId: "..." │ + │ │ }───────────────────│ + │ │ │ +``` + +--- + +## Error Handling & Recovery + +### Error Types + +The protocol defines standard error codes: + +| Code | Meaning | +|------|---------| +| `AUTHENTICATION_FAILED` | Invalid credentials | +| `TOOL_NOT_FOUND` | Requested tool doesn't exist | +| `TOOL_EXECUTION_FAILED` | Tool threw an error | +| `INVALID_MESSAGE` | Malformed protocol message | +| `SESSION_NOT_FOUND` | Session doesn't exist or expired | +| `BRIDGE_DISCONNECTED` | Bridge is not connected | +| `TIMEOUT` | Tool execution took too long | +| `CANCELLED` | Tool execution was cancelled | +| `INTERNAL_ERROR` | Unexpected server error | + +**File Reference**: [`packages/tool-bridge-protocol/src/utils.ts`](../packages/tool-bridge-protocol/src/utils.ts) (lines 30-41) + +### Error Scenarios + +#### Bridge Disconnects During Call + +``` + Relay Bridge + │ │ + │ ──── invoke ──────────►│ + │ │ + │ ╳ Connection lost + │ │ + ┌────┴────────────────────────┘ + │ + │ Relay: + │ • Detects disconnect (ping timeout or close event) + │ • Fails all pending calls for this bridge + │ • Keeps session alive for grace period + │ + │ LLM Runtime receives: + │ { success: false, error: { code: "BRIDGE_DISCONNECTED", ... } } + │ +``` + +#### Tool Throws Error + +```typescript +// In the bridge's tool executor +bridge.registerTool("risky", "...", {...}, + async (args, ctx) => { + throw new Error("Something went wrong!") + } +) + +// The bridge catches this and sends: +// result { success: false, error: { code: "TOOL_EXECUTION_FAILED", message: "Something went wrong!" } } +``` + +#### Tool Times Out + +``` + Relay Bridge + │ │ + │ ──── invoke ──────────►│ + │ │ + │ ... 2 minutes pass ... + │ │ + ┌────┴────────────────────────┘ + │ + │ Relay: Timeout triggered + │ • Sends cancel to bridge + │ • Returns timeout error to LLM + │ + │ ──── cancel ───────────────►│ + │ │ + │ (Bridge tries to stop) + │ │ +``` + +--- + +## Security Considerations + +### Authentication + +The protocol supports multiple authentication methods: + +```typescript +// No authentication (development only!) +const relay = new ToolRelay({ auth: { type: "none" } }) + +// Bearer token +const relay = new ToolRelay({ + auth: { + type: "bearer", + validateToken: async (token) => { + return await verifyJWT(token) // Your validation logic + } + } +}) + +// Client provides credentials +const bridge = new ToolBridge({ + url: "wss://relay.example.com/ws", + credentials: { + type: "bearer", + token: "eyJhbG..." // JWT or API key + } +}) +``` + +### Transport Security + +Always use WSS (WebSocket Secure) in production: + +``` +❌ ws://relay.example.com/ws (plaintext - anyone can read) +✅ wss://relay.example.com/ws (encrypted with TLS) +``` + +### Tool Permissions + +The bridge controls which tools are exposed. Don't blindly register dangerous tools: + +```typescript +// ❌ Dangerous: Allows arbitrary command execution +bridge.registerTool("shell", "Run any command", {...}, + async (args) => exec(args.command) // BAD! +) + +// ✅ Better: Restricted commands with allowlist +bridge.registerTool("safe-shell", "Run allowed commands", {...}, + async (args) => { + const allowed = ["ls", "cat", "grep", "git status"] + if (!allowed.some(cmd => args.command.startsWith(cmd))) { + throw new Error("Command not allowed") + } + return exec(args.command) + } +) +``` + +### Session Isolation + +Each session is isolated. One user's bridge can't invoke tools on another user's bridge (assuming proper authentication). + +--- + +## Code References + +### Package Structure + +``` +packages/ +├── tool-bridge-protocol/ # Shared protocol definitions +│ └── src/ +│ ├── index.ts # Package exports +│ ├── schema.ts # Zod schemas for messages +│ └── utils.ts # Message factories, constants +│ +├── tool-relay/ # Server-side gateway +│ └── src/ +│ ├── index.ts # Package exports +│ ├── relay.ts # ToolRelay class +│ └── types.ts # TypeScript interfaces +│ +├── tool-bridge/ # Client-side agent +│ └── src/ +│ ├── index.ts # Package exports +│ ├── bridge.ts # ToolBridge class +│ ├── mcp.ts # McpClientManager +│ └── types.ts # TypeScript interfaces +│ +├── tool-bridge-opencode/ # OpenCode tool adapter +│ └── src/ +│ ├── index.ts # Package exports +│ └── adapter.ts # Adapts OpenCode tools +│ +└── tool-bridge-tests/ # Integration tests + ├── integration.test.ts # Protocol tests + └── mcp.test.ts # MCP integration tests +``` + +### Key Classes + +| Class | File | Purpose | +|-------|------|---------| +| `ToolRelay` | [`tool-relay/src/relay.ts`](../packages/tool-relay/src/relay.ts) | Server gateway | +| `ToolBridge` | [`tool-bridge/src/bridge.ts`](../packages/tool-bridge/src/bridge.ts) | Client agent | +| `McpClientManager` | [`tool-bridge/src/mcp.ts`](../packages/tool-bridge/src/mcp.ts) | MCP integration | + +### Key Types + +| Type | File | Purpose | +|------|------|---------| +| `ProtocolMessage` | [`tool-bridge-protocol/src/schema.ts`](../packages/tool-bridge-protocol/src/schema.ts) | Union of all messages | +| `BridgeConfig` | [`tool-bridge/src/types.ts`](../packages/tool-bridge/src/types.ts) | Bridge configuration | +| `RelayConfig` | [`tool-relay/src/types.ts`](../packages/tool-relay/src/types.ts) | Relay configuration | +| `McpConfig` | [`tool-bridge/src/types.ts`](../packages/tool-bridge/src/types.ts) | MCP server config | + +--- + +## Conclusion + +The Tool Bridge and Tool Relay system enables secure, efficient tool execution across network boundaries. Key takeaways: + +1. **Reverse RPC**: Client initiates connection, server makes requests +2. **Multiplexing**: Multiple concurrent calls over one connection +3. **Sessions**: Survive disconnects, enable reconnection +4. **MCP Integration**: Leverage the MCP ecosystem of tools +5. **Security**: Authentication, TLS, and permission controls + +For questions or issues, refer to the test files which demonstrate all major features: +- [`packages/tool-bridge-tests/integration.test.ts`](../packages/tool-bridge-tests/integration.test.ts) +- [`packages/tool-bridge-tests/mcp.test.ts`](../packages/tool-bridge-tests/mcp.test.ts) From 1779fe3b4c827dfb363db37172ab5b49f2bf28bc Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 21 Dec 2025 03:33:49 +0000 Subject: [PATCH 4/4] feat: add per-session tool access control Add support for dynamically enabling/disabling tools on a per-session basis. The server (relay) can now control which tools are available for each session, enabling denylist or allowlist modes. Changes: - Add SetToolsAccessMessage and ToolsAccessUpdatedMessage to protocol - Add TOOL_ACCESS_DENIED error code - Add setToolsAccess/getToolsAccess/getAccessibleTools methods to relay - Update bridge to track disabled tools and respect permissions - Add comprehensive tests for tool access control - Update design documentation with new feature --- docs/tool-bridge-design.md | 112 +++++++ packages/tool-bridge-protocol/src/schema.ts | 40 +++ packages/tool-bridge-protocol/src/utils.ts | 37 +++ .../tool-bridge-tests/integration.test.ts | 288 ++++++++++++++++++ packages/tool-bridge/src/bridge.ts | 50 +++ packages/tool-bridge/src/types.ts | 4 + packages/tool-relay/src/relay.ts | 161 +++++++++- packages/tool-relay/src/types.ts | 20 ++ 8 files changed, 711 insertions(+), 1 deletion(-) diff --git a/docs/tool-bridge-design.md b/docs/tool-bridge-design.md index 2b1361084aa..745307b42f9 100644 --- a/docs/tool-bridge-design.md +++ b/docs/tool-bridge-design.md @@ -240,6 +240,8 @@ Every message follows this structure: | `cancelled` | Bridge → Relay | "Call X has been stopped" | | `ping` | Relay → Bridge | "Are you still there?" | | `pong` | Bridge → Relay | "Yes, I'm here" | +| `set_tools_access` | Relay → Bridge | "Enable/disable these tools" | +| `tools_access_updated` | Relay → Bridge | "Here's what's now enabled/disabled" | ### Example: Complete Flow @@ -880,6 +882,7 @@ The protocol defines standard error codes: | `AUTHENTICATION_FAILED` | Invalid credentials | | `TOOL_NOT_FOUND` | Requested tool doesn't exist | | `TOOL_EXECUTION_FAILED` | Tool threw an error | +| `TOOL_ACCESS_DENIED` | Tool is disabled for this session | | `INVALID_MESSAGE` | Malformed protocol message | | `SESSION_NOT_FOUND` | Session doesn't exist or expired | | `BRIDGE_DISCONNECTED` | Bridge is not connected | @@ -988,6 +991,115 @@ Always use WSS (WebSocket Secure) in production: ✅ wss://relay.example.com/ws (encrypted with TLS) ``` +### Per-Session Tool Access Control + +The relay can dynamically enable or disable tools for each session. This allows the server to control which tools are available based on user permissions, context, or security policies. + +``` +┌─────────────────────────────────────────────────────────────────────────────┐ +│ PER-SESSION TOOL ACCESS CONTROL │ +└─────────────────────────────────────────────────────────────────────────────┘ + + LLM Runtime Tool Relay Tool Bridge + │ │ │ + │ "Disable dangerous │ │ + │ tools for this │ │ + │ user session" │ │ + │ │ │ + │ ──setToolsAccess({ │ │ + │ tools: { │ │ + │ "bash": false, │ │ + │ "write": false │ │ + │ } │ │ + │ })─────────────────►│ │ + │ │ │ + │ │ Update session config │ + │ │ │ + │ │ ──tools_access_updated { │ + │ │ enabled: ["read","glob"], │ + │ │ disabled: ["bash","write"]│ + │ │ }─────────────────────────►│ + │ │ │ + │ │ Bridge tracks + │ │ disabled tools + │ │ │ + │ Later: try to │ │ + │ invoke "bash" │ │ + │ │ │ + │ ──invoke({ │ │ + │ toolId: "bash" │ │ + │ })─────────────────►│ │ + │ │ │ + │ │ Check access: DENIED │ + │ │ │ + │ ◄──{ success: false, │ │ + │ error: { │ │ + │ code: "TOOL_ │ │ + │ ACCESS_DENIED"│ │ + │ }}────────────────│ │ + │ │ │ +``` + +#### Using Tool Access Control + +```typescript +// From the LLM runtime, control which tools are available + +// Disable specific tools (denylist mode - default) +relay.setToolsAccess(sessionId, { + tools: { "bash": false, "write": false } +}) + +// Enable only specific tools (allowlist mode) +relay.setToolsAccess(sessionId, { + tools: { "read": true, "glob": true }, + defaultAccess: false // All other tools disabled +}) + +// Re-enable all tools +relay.setToolsAccess(sessionId, { + tools: { "bash": true, "write": true }, + defaultAccess: true +}) + +// Get current access configuration +const access = relay.getToolsAccess(sessionId) +// { tools: { "bash": false }, defaultAccess: true } + +// Get only accessible tools (respects access control) +const tools = relay.getAccessibleTools(sessionId) +// Returns only enabled tools +``` + +#### Access Control Modes + +| Mode | Configuration | Behavior | +|------|---------------|----------| +| **Denylist** (default) | `defaultAccess: true` | All tools enabled except those set to `false` | +| **Allowlist** | `defaultAccess: false` | All tools disabled except those set to `true` | + +#### Bridge Notifications + +The bridge receives notifications when tool access changes: + +```typescript +const bridge = new ToolBridge({ + url: "wss://relay.example.com/ws", + onToolsAccessUpdated: (enabledTools, disabledTools) => { + console.log("Enabled:", enabledTools) // ["read", "glob"] + console.log("Disabled:", disabledTools) // ["bash", "write"] + } +}) + +// Check if a specific tool is enabled +if (bridge.isToolEnabled("bash")) { + // Tool can be invoked +} + +// Get all disabled tools +const disabled = bridge.getDisabledTools() // ["bash", "write"] +``` + ### Tool Permissions The bridge controls which tools are exposed. Don't blindly register dangerous tools: diff --git a/packages/tool-bridge-protocol/src/schema.ts b/packages/tool-bridge-protocol/src/schema.ts index 4f89ea88ccc..151db5eec59 100644 --- a/packages/tool-bridge-protocol/src/schema.ts +++ b/packages/tool-bridge-protocol/src/schema.ts @@ -174,6 +174,42 @@ export const PongMessage = z.object({ timestamp: z.string(), }) +// Tool access control messages (server controls which tools are allowed per session) +export const SetToolsAccessMessage = z.object({ + type: z.literal("set_tools_access"), + id: z.string(), + sessionId: z.string(), + timestamp: z.string(), + payload: z.object({ + /** + * Tool access rules. Each tool ID maps to its access state. + * - true: tool is enabled + * - false: tool is disabled + * If a tool is not in this map, it uses the default (enabled) + */ + tools: z.record(z.string(), z.boolean()), + /** + * Default access for tools not explicitly listed. + * If true (default), unlisted tools are enabled. + * If false, unlisted tools are disabled (allowlist mode). + */ + defaultAccess: z.boolean().optional(), + }), +}) + +export const ToolsAccessUpdatedMessage = z.object({ + type: z.literal("tools_access_updated"), + id: z.string(), + sessionId: z.string(), + timestamp: z.string(), + payload: z.object({ + /** Tools that are currently enabled */ + enabledTools: z.array(z.string()), + /** Tools that are currently disabled */ + disabledTools: z.array(z.string()), + }), +}) + // Union of all message types export const ProtocolMessage = z.discriminatedUnion("type", [ HelloMessage, @@ -189,6 +225,8 @@ export const ProtocolMessage = z.discriminatedUnion("type", [ CancelledMessage, PingMessage, PongMessage, + SetToolsAccessMessage, + ToolsAccessUpdatedMessage, ]) // Type exports @@ -207,4 +245,6 @@ export type CancelMessage = z.infer export type CancelledMessage = z.infer export type PingMessage = z.infer export type PongMessage = z.infer +export type SetToolsAccessMessage = z.infer +export type ToolsAccessUpdatedMessage = z.infer export type ProtocolMessage = z.infer diff --git a/packages/tool-bridge-protocol/src/utils.ts b/packages/tool-bridge-protocol/src/utils.ts index bb9520dca9a..7c1ba183471 100644 --- a/packages/tool-bridge-protocol/src/utils.ts +++ b/packages/tool-bridge-protocol/src/utils.ts @@ -13,6 +13,8 @@ import type { CancelledMessage, PingMessage, PongMessage, + SetToolsAccessMessage, + ToolsAccessUpdatedMessage, ToolDefinition, } from "./schema" @@ -38,6 +40,7 @@ export const ErrorCode = { AUTHENTICATION_FAILED: "AUTHENTICATION_FAILED", TOOL_NOT_FOUND: "TOOL_NOT_FOUND", TOOL_EXECUTION_FAILED: "TOOL_EXECUTION_FAILED", + TOOL_ACCESS_DENIED: "TOOL_ACCESS_DENIED", INVALID_MESSAGE: "INVALID_MESSAGE", SESSION_NOT_FOUND: "SESSION_NOT_FOUND", BRIDGE_DISCONNECTED: "BRIDGE_DISCONNECTED", @@ -240,4 +243,38 @@ export const Message = { timestamp: timestamp(), } }, + + setToolsAccess( + sessionId: string, + tools: Record, + defaultAccess?: boolean + ): SetToolsAccessMessage { + return { + type: "set_tools_access", + id: generateId(), + sessionId, + timestamp: timestamp(), + payload: { + tools, + defaultAccess, + }, + } + }, + + toolsAccessUpdated( + sessionId: string, + enabledTools: string[], + disabledTools: string[] + ): ToolsAccessUpdatedMessage { + return { + type: "tools_access_updated", + id: generateId(), + sessionId, + timestamp: timestamp(), + payload: { + enabledTools, + disabledTools, + }, + } + }, } diff --git a/packages/tool-bridge-tests/integration.test.ts b/packages/tool-bridge-tests/integration.test.ts index aecedca29ce..236a0634207 100644 --- a/packages/tool-bridge-tests/integration.test.ts +++ b/packages/tool-bridge-tests/integration.test.ts @@ -476,4 +476,292 @@ describe("Tool Bridge Protocol Integration", () => { expect(result.error!.code).toBe("SESSION_NOT_FOUND") }) }) + + describe("Tool Access Control", () => { + test("relay can disable a specific tool", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-15", + reconnect: { enabled: false }, + }) + + bridge.registerTool( + "tool-a", + "Tool A", + { type: "object" }, + async () => ({ output: "a" }) + ) + + bridge.registerTool( + "tool-b", + "Tool B", + { type: "object" }, + async () => ({ output: "b" }) + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + // Disable tool-a + const updated = relay.setToolsAccess(sessionId, { tools: { "tool-a": false } }) + expect(updated).toBe(true) + + // tool-a should be denied + const resultA = await relay.invoke({ + sessionId, + toolId: "tool-a", + arguments: {}, + }) + expect(resultA.success).toBe(false) + expect(resultA.error!.code).toBe("TOOL_ACCESS_DENIED") + + // tool-b should still work + const resultB = await relay.invoke({ + sessionId, + toolId: "tool-b", + arguments: {}, + }) + expect(resultB.success).toBe(true) + expect(resultB.output).toBe("b") + + bridge.disconnect() + }) + + test("relay can enable only specific tools (allowlist mode)", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-16", + reconnect: { enabled: false }, + }) + + bridge.registerTool( + "allowed-tool", + "Allowed", + { type: "object" }, + async () => ({ output: "allowed" }) + ) + + bridge.registerTool( + "denied-tool", + "Denied", + { type: "object" }, + async () => ({ output: "denied" }) + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + // Enable only allowed-tool, deny everything else + relay.setToolsAccess(sessionId, { + tools: { "allowed-tool": true }, + defaultAccess: false, + }) + + // allowed-tool should work + const resultAllowed = await relay.invoke({ + sessionId, + toolId: "allowed-tool", + arguments: {}, + }) + expect(resultAllowed.success).toBe(true) + expect(resultAllowed.output).toBe("allowed") + + // denied-tool should be denied + const resultDenied = await relay.invoke({ + sessionId, + toolId: "denied-tool", + arguments: {}, + }) + expect(resultDenied.success).toBe(false) + expect(resultDenied.error!.code).toBe("TOOL_ACCESS_DENIED") + + bridge.disconnect() + }) + + test("relay can re-enable all tools", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-17", + reconnect: { enabled: false }, + }) + + bridge.registerTool( + "test-tool", + "Test", + { type: "object" }, + async () => ({ output: "test" }) + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + // Disable the tool + relay.setToolsAccess(sessionId, { tools: { "test-tool": false } }) + + // Verify it's disabled + const resultDisabled = await relay.invoke({ + sessionId, + toolId: "test-tool", + arguments: {}, + }) + expect(resultDisabled.success).toBe(false) + + // Re-enable by setting defaultAccess to true and removing the override + relay.setToolsAccess(sessionId, { tools: { "test-tool": true }, defaultAccess: true }) + + // Verify it's enabled again + const resultEnabled = await relay.invoke({ + sessionId, + toolId: "test-tool", + arguments: {}, + }) + expect(resultEnabled.success).toBe(true) + expect(resultEnabled.output).toBe("test") + + bridge.disconnect() + }) + + test("getToolsAccess returns current configuration", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-18", + reconnect: { enabled: false }, + }) + + bridge.registerTool( + "some-tool", + "Some tool", + { type: "object" }, + async () => ({ output: "ok" }) + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + // Check initial state + let access = relay.getToolsAccess(sessionId) + expect(access).toBeTruthy() + expect(access!.defaultAccess).toBe(true) + expect(Object.keys(access!.tools).length).toBe(0) + + // Disable a tool + relay.setToolsAccess(sessionId, { tools: { "some-tool": false } }) + + // Check updated state + access = relay.getToolsAccess(sessionId) + expect(access!.tools["some-tool"]).toBe(false) + + bridge.disconnect() + }) + + test("getAccessibleTools respects access control", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-19", + reconnect: { enabled: false }, + }) + + bridge.registerTool( + "enabled-tool", + "Enabled", + { type: "object" }, + async () => ({ output: "enabled" }) + ) + + bridge.registerTool( + "disabled-tool", + "Disabled", + { type: "object" }, + async () => ({ output: "disabled" }) + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + // Initially all tools are accessible + let accessible = relay.getAccessibleTools(sessionId) + expect(accessible.length).toBe(2) + + // Disable one tool + relay.setToolsAccess(sessionId, { tools: { "disabled-tool": false } }) + + // Only one tool should be accessible + accessible = relay.getAccessibleTools(sessionId) + expect(accessible.length).toBe(1) + expect(accessible[0].id).toBe("enabled-tool") + + bridge.disconnect() + }) + + test("bridge receives tool access update notification", async () => { + let accessUpdateReceived = false + let receivedEnabledTools: string[] = [] + let receivedDisabledTools: string[] = [] + + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-20", + reconnect: { enabled: false }, + onToolsAccessUpdated: (enabled, disabled) => { + accessUpdateReceived = true + receivedEnabledTools = enabled + receivedDisabledTools = disabled + }, + }) + + bridge.registerTool( + "my-tool", + "My tool", + { type: "object" }, + async () => ({ output: "ok" }) + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + // Disable the tool + relay.setToolsAccess(sessionId, { tools: { "my-tool": false } }) + + // Wait for notification + await new Promise((resolve) => setTimeout(resolve, 100)) + + expect(accessUpdateReceived).toBe(true) + expect(receivedDisabledTools).toContain("my-tool") + + bridge.disconnect() + }) + + test("bridge rejects invocation of disabled tool locally", async () => { + const bridge = new ToolBridge({ + url: serverUrl, + bridgeId: "test-bridge-21", + reconnect: { enabled: false }, + }) + + bridge.registerTool( + "local-tool", + "Local tool", + { type: "object" }, + async () => ({ output: "local" }) + ) + + const sessionId = await bridge.connect() + await new Promise((resolve) => setTimeout(resolve, 100)) + + // Disable the tool + relay.setToolsAccess(sessionId, { tools: { "local-tool": false } }) + await new Promise((resolve) => setTimeout(resolve, 100)) + + // Bridge should track disabled state + expect(bridge.isToolEnabled("local-tool")).toBe(false) + expect(bridge.getDisabledTools()).toContain("local-tool") + + bridge.disconnect() + }) + + test("setToolsAccess returns false for invalid session", () => { + const result = relay.setToolsAccess("invalid-session-id", { tools: { foo: false } }) + expect(result).toBe(false) + }) + }) }) diff --git a/packages/tool-bridge/src/bridge.ts b/packages/tool-bridge/src/bridge.ts index 9c7c39d8419..a1ece43dbc8 100644 --- a/packages/tool-bridge/src/bridge.ts +++ b/packages/tool-bridge/src/bridge.ts @@ -6,6 +6,7 @@ import { InvokeMessage, CancelMessage, PingMessage, + ToolsAccessUpdatedMessage, ToolDefinition, } from "@opencode-ai/tool-bridge-protocol/schema" import { Message, generateId, ErrorCode, DEFAULT_TIMEOUTS } from "@opencode-ai/tool-bridge-protocol/utils" @@ -31,6 +32,7 @@ export class ToolBridge { private readonly config: BridgeConfig private readonly tools: Map = new Map() private readonly activeCalls: Map = new Map() + private readonly disabledTools: Set = new Set() private socket: WebSocket | null = null private sessionId: string | null = null @@ -161,6 +163,7 @@ export class ToolBridge { this.sessionId = null this.state = "disconnected" + this.disabledTools.clear() } /** @@ -296,6 +299,10 @@ export class ToolBridge { case "ping": this.handlePing() break + + case "tools_access_updated": + this.handleToolsAccessUpdated(parsed.data) + break } } catch (error) { console.error("Failed to handle message:", error) @@ -354,6 +361,17 @@ export class ToolBridge { return } + // Check if the tool is disabled + if (this.disabledTools.has(toolId)) { + this.sendResult(message.sessionId, callId, false, { + error: { + code: ErrorCode.TOOL_ACCESS_DENIED, + message: `Tool ${toolId} is disabled for this session`, + }, + }) + return + } + // Create abort controller for this call const abortController = new AbortController() this.activeCalls.set(callId, abortController) @@ -421,6 +439,38 @@ export class ToolBridge { this.socket.send(JSON.stringify(pong)) } + /** + * Handle tools_access_updated message + */ + private handleToolsAccessUpdated(message: ToolsAccessUpdatedMessage): void { + const { enabledTools, disabledTools } = message.payload + + // Update disabled tools set + for (const toolId of enabledTools) { + this.disabledTools.delete(toolId) + } + for (const toolId of disabledTools) { + this.disabledTools.add(toolId) + } + + // Notify callback + this.config.onToolsAccessUpdated?.(enabledTools, disabledTools) + } + + /** + * Check if a tool is enabled + */ + isToolEnabled(toolId: string): boolean { + return !this.disabledTools.has(toolId) + } + + /** + * Get list of disabled tools + */ + getDisabledTools(): string[] { + return Array.from(this.disabledTools) + } + /** * Handle connection close */ diff --git a/packages/tool-bridge/src/types.ts b/packages/tool-bridge/src/types.ts index 7e25967f2d9..6e896316ddc 100644 --- a/packages/tool-bridge/src/types.ts +++ b/packages/tool-bridge/src/types.ts @@ -93,6 +93,10 @@ export interface BridgeConfig { onDisconnect?: (reason: string) => void onReconnect?: (attempt: number) => void onError?: (error: Error) => void + /** + * Called when tool access permissions are updated by the server + */ + onToolsAccessUpdated?: (enabledTools: string[], disabledTools: string[]) => void } /** diff --git a/packages/tool-relay/src/relay.ts b/packages/tool-relay/src/relay.ts index e24b989b4a3..cef63bb4ca9 100644 --- a/packages/tool-relay/src/relay.ts +++ b/packages/tool-relay/src/relay.ts @@ -6,10 +6,11 @@ import { ChunkMessage, CancelledMessage, PongMessage, + SetToolsAccessMessage, ToolDefinition, } from "@opencode-ai/tool-bridge-protocol/schema" import { Message, generateId, ErrorCode, DEFAULT_TIMEOUTS, PROTOCOL_VERSION } from "@opencode-ai/tool-bridge-protocol/utils" -import type { RelayConfig, BridgeConnection, Session, ToolInvocationRequest, ToolInvocationResult } from "./types" +import type { RelayConfig, BridgeConnection, Session, ToolInvocationRequest, ToolInvocationResult, ToolAccessConfig } from "./types" /** * Tool Relay - Server-side gateway for tool invocations @@ -132,6 +133,10 @@ export class ToolRelay { this.handlePong(message) break + case "set_tools_access": + this.handleSetToolsAccess(message) + break + default: // Ignore unsupported message types from client break @@ -185,6 +190,10 @@ export class ToolRelay { createdAt: Date.now(), lastActivity: Date.now(), tools: [], + toolAccess: { + tools: {}, + defaultAccess: true, // All tools enabled by default + }, } this.sessions.set(sessionId, session) } @@ -318,6 +327,57 @@ export class ToolRelay { } } + /** + * Handle set_tools_access message (from server to bridge, forwarded here for processing) + */ + private handleSetToolsAccess(message: SetToolsAccessMessage): void { + const session = this.sessions.get(message.sessionId) + if (!session) return + + const { tools, defaultAccess } = message.payload + + // Update session tool access config + session.toolAccess = { + tools: { ...session.toolAccess.tools, ...tools }, + defaultAccess: defaultAccess ?? session.toolAccess.defaultAccess, + } + + // Compute enabled/disabled lists based on registered tools + const enabledTools: string[] = [] + const disabledTools: string[] = [] + + for (const tool of session.tools) { + if (this.isToolAccessible(session, tool.id)) { + enabledTools.push(tool.id) + } else { + disabledTools.push(tool.id) + } + } + + // Send update notification to the bridge + const bridgeId = this.sessionToBridge.get(message.sessionId) + if (bridgeId) { + const connection = this.connections.get(bridgeId) + if (connection?.connected) { + const notification = Message.toolsAccessUpdated(message.sessionId, enabledTools, disabledTools) + try { + connection.socket.send(JSON.stringify(notification)) + } catch {} + } + } + } + + /** + * Check if a tool is accessible for a session + */ + private isToolAccessible(session: Session, toolId: string): boolean { + const access = session.toolAccess.tools[toolId] + if (access !== undefined) { + return access + } + return session.toolAccess.defaultAccess + } + /** * Handle bridge disconnect */ @@ -455,6 +515,18 @@ export class ToolRelay { } } + // Check tool access permissions + const session = this.sessions.get(sessionId) + if (session && !this.isToolAccessible(session, toolId)) { + return { + success: false, + error: { + code: ErrorCode.TOOL_ACCESS_DENIED, + message: `Tool ${toolId} is not accessible in this session`, + }, + } + } + const callId = generateId() const invokeMessage = Message.invoke(sessionId, callId, toolId, args, traceId) @@ -562,4 +634,91 @@ export class ToolRelay { getRelayId(): string { return this.relayId } + + /** + * Set tool access for a session + * + * This allows the server to control which tools are available + * on a per-session basis. Tools can be selectively enabled or disabled. + * + * @param sessionId - The session to configure + * @param config - Tool access configuration + * @param config.tools - Map of tool ID to access state (true = enabled, false = disabled) + * @param config.defaultAccess - Default access for unlisted tools (default: true) + * @returns true if the session was found and updated, false otherwise + * + * @example + * ```typescript + * // Disable a specific tool + * relay.setToolsAccess(sessionId, { tools: { "dangerous_tool": false } }) + * + * // Enable only specific tools (allowlist mode) + * relay.setToolsAccess(sessionId, { + * tools: { "read": true, "list": true }, + * defaultAccess: false + * }) + * + * // Re-enable all tools + * relay.setToolsAccess(sessionId, { tools: {}, defaultAccess: true }) + * ``` + */ + setToolsAccess( + sessionId: string, + config: { tools?: Record; defaultAccess?: boolean } + ): boolean { + const session = this.sessions.get(sessionId) + if (!session) return false + + // Update session tool access config + if (config.tools) { + session.toolAccess.tools = { ...session.toolAccess.tools, ...config.tools } + } + if (config.defaultAccess !== undefined) { + session.toolAccess.defaultAccess = config.defaultAccess + } + + // Compute enabled/disabled lists + const enabledTools: string[] = [] + const disabledTools: string[] = [] + + for (const tool of session.tools) { + if (this.isToolAccessible(session, tool.id)) { + enabledTools.push(tool.id) + } else { + disabledTools.push(tool.id) + } + } + + // Send update notification to the bridge + const bridgeId = this.sessionToBridge.get(sessionId) + if (bridgeId) { + const connection = this.connections.get(bridgeId) + if (connection?.connected) { + const notification = Message.toolsAccessUpdated(sessionId, enabledTools, disabledTools) + try { + connection.socket.send(JSON.stringify(notification)) + } catch {} + } + } + + return true + } + + /** + * Get tool access configuration for a session + */ + getToolsAccess(sessionId: string): ToolAccessConfig | null { + const session = this.sessions.get(sessionId) + return session?.toolAccess ?? null + } + + /** + * Get accessible tools for a session (respecting access control) + */ + getAccessibleTools(sessionId: string): ToolDefinition[] { + const session = this.sessions.get(sessionId) + if (!session) return [] + + return session.tools.filter((tool) => this.isToolAccessible(session, tool.id)) + } } diff --git a/packages/tool-relay/src/types.ts b/packages/tool-relay/src/types.ts index d5b8f3741d1..554884a9191 100644 --- a/packages/tool-relay/src/types.ts +++ b/packages/tool-relay/src/types.ts @@ -67,6 +67,22 @@ export interface BridgeConnection { connected: boolean } +/** + * Tool access configuration for a session + */ +export interface ToolAccessConfig { + /** + * Map of tool ID to access state (true = enabled, false = disabled) + */ + tools: Record + /** + * Default access for tools not explicitly listed. + * If true (default), unlisted tools are enabled. + * If false, unlisted tools are disabled (allowlist mode). + */ + defaultAccess: boolean +} + /** * Represents a session */ @@ -76,6 +92,10 @@ export interface Session { createdAt: number lastActivity: number tools: ToolDefinition[] + /** + * Per-session tool access control + */ + toolAccess: ToolAccessConfig } /**