diff --git a/docs/tool-bridge-design.md b/docs/tool-bridge-design.md new file mode 100644 index 00000000000..745307b42f9 --- /dev/null +++ b/docs/tool-bridge-design.md @@ -0,0 +1,1197 @@ +# 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" | +| `set_tools_access` | Relay → Bridge | "Enable/disable these tools" | +| `tools_access_updated` | Relay → Bridge | "Here's what's now enabled/disabled" | + +### 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 | +| `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 | +| `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) +``` + +### 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: + +```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) 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..151db5eec59 --- /dev/null +++ b/packages/tool-bridge-protocol/src/schema.ts @@ -0,0 +1,250 @@ +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(), +}) + +// 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, + WelcomeMessage, + ErrorMessage, + RegisterToolsMessage, + ToolsRegisteredMessage, + InvokeMessage, + ResultMessage, + ChunkMessage, + DoneMessage, + CancelMessage, + CancelledMessage, + PingMessage, + PongMessage, + SetToolsAccessMessage, + ToolsAccessUpdatedMessage, +]) + +// 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 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 new file mode 100644 index 00000000000..7c1ba183471 --- /dev/null +++ b/packages/tool-bridge-protocol/src/utils.ts @@ -0,0 +1,280 @@ +import { randomUUID } from "crypto" +import type { + HelloMessage, + WelcomeMessage, + ErrorMessage, + RegisterToolsMessage, + ToolsRegisteredMessage, + InvokeMessage, + ResultMessage, + ChunkMessage, + DoneMessage, + CancelMessage, + CancelledMessage, + PingMessage, + PongMessage, + SetToolsAccessMessage, + ToolsAccessUpdatedMessage, + 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", + TOOL_ACCESS_DENIED: "TOOL_ACCESS_DENIED", + 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(), + } + }, + + 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-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..236a0634207 --- /dev/null +++ b/packages/tool-bridge-tests/integration.test.ts @@ -0,0 +1,767 @@ +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") + }) + }) + + 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-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 new file mode 100644 index 00000000000..184b066e884 --- /dev/null +++ b/packages/tool-bridge-tests/package.json @@ -0,0 +1,22 @@ +{ + "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:*", + "@ai-sdk/mcp": "0.0.8", + "@modelcontextprotocol/sdk": "1.15.1", + "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..1d62e005922 --- /dev/null +++ b/packages/tool-bridge/package.json @@ -0,0 +1,20 @@ +{ + "name": "@opencode-ai/tool-bridge", + "version": "1.0.0", + "type": "module", + "private": true, + "exports": { + "./*": "./src/*.ts" + }, + "dependencies": { + "@opencode-ai/tool-bridge-protocol": "workspace:*", + "@ai-sdk/mcp": "0.0.8", + "@modelcontextprotocol/sdk": "1.15.1", + "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..a1ece43dbc8 --- /dev/null +++ b/packages/tool-bridge/src/bridge.ts @@ -0,0 +1,622 @@ +import { + ProtocolMessage, + WelcomeMessage, + ErrorMessage, + ToolsRegisteredMessage, + 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" +import type { + BridgeConfig, + ToolExecutor, + ToolContext, + ToolResult, + RegisteredTool, + ConnectionState, + McpConfig, +} from "./types" +import type { McpClientManager } from "./mcp" + +/** + * 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 readonly disabledTools: Set = new Set() + + 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 + private mcpManager: McpClientManager | 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") + + // Close MCP connections (fire and forget) + this.closeMcpConnections().catch(() => {}) + + if (this.socket) { + try { + this.socket.close() + } catch {} + this.socket = null + } + + this.sessionId = null + this.state = "disconnected" + this.disabledTools.clear() + } + + /** + * 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) + } + + /** + * 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 + */ + 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 + + case "tools_access_updated": + this.handleToolsAccessUpdated(parsed.data) + 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 + } + + // 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) + + 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 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 + */ + 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..32d96faa9c9 --- /dev/null +++ b/packages/tool-bridge/src/index.ts @@ -0,0 +1,24 @@ +/** + * 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. + * + * Supports MCP (Model Context Protocol) servers for tool discovery. + */ + +export { ToolBridge } from "./bridge" +export { McpClientManager } from "./mcp" +export type { + BridgeConfig, + ToolExecutor, + ToolContext, + 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 new file mode 100644 index 00000000000..6e896316ddc --- /dev/null +++ b/packages/tool-bridge/src/types.ts @@ -0,0 +1,141 @@ +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 + */ +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 + /** + * Called when tool access permissions are updated by the server + */ + onToolsAccessUpdated?: (enabledTools: string[], disabledTools: string[]) => 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..cef63bb4ca9 --- /dev/null +++ b/packages/tool-relay/src/relay.ts @@ -0,0 +1,724 @@ +import { + ProtocolMessage, + HelloMessage, + RegisterToolsMessage, + ResultMessage, + 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, ToolAccessConfig } 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 + + case "set_tools_access": + this.handleSetToolsAccess(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: [], + toolAccess: { + tools: {}, + defaultAccess: true, // All tools enabled by default + }, + } + 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 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 + */ + 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`, + }, + } + } + + // 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) + + 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 + } + + /** + * 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 new file mode 100644 index 00000000000..554884a9191 --- /dev/null +++ b/packages/tool-relay/src/types.ts @@ -0,0 +1,124 @@ +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 +} + +/** + * 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 + */ +export interface Session { + sessionId: string + bridgeId: string + createdAt: number + lastActivity: number + tools: ToolDefinition[] + /** + * Per-session tool access control + */ + toolAccess: ToolAccessConfig +} + +/** + * 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 + } +}