| ... | @@ -9,6 +9,308 @@ type RequestHook = ( | ... | @@ -9,6 +9,308 @@ type RequestHook = ( |
| 9 | ctx: { model: { provider: string } }, | 9 | ctx: { model: { provider: string } }, |
| 10 | ) => Record<string, unknown> | undefined; | 10 | ) => Record<string, unknown> | undefined; |
| 11 | | 11 | |
| | 12 | interface RegisteredTool { |
| | 13 | readonly name: string; |
| | 14 | readonly description: string; |
| | 15 | readonly parameters: unknown; |
| | 16 | readonly exposure?: string; |
| | 17 | readonly promptSnippet?: string; |
| | 18 | readonly promptGuidelines?: ReadonlyArray<string>; |
| | 19 | readonly execute: ( |
| | 20 | id: string, |
| | 21 | args: Record<string, unknown>, |
| | 22 | signal?: AbortSignal, |
| | 23 | ) => Promise<{ |
| | 24 | readonly content: ReadonlyArray<{ readonly type: string; readonly text: string }>; |
| | 25 | }>; |
| | 26 | } |
| | 27 | |
| | 28 | type AgentStartHook = ( |
| | 29 | event: { systemPrompt: string }, |
| | 30 | ctx: { ui: { notify: (message: string, severity: string) => void } }, |
| | 31 | ) => Promise<{ systemPrompt: string }>; |
| | 32 | |
| | 33 | async function loadMcpBridge( |
| | 34 | options: { |
| | 35 | readonly modern?: boolean; |
| | 36 | readonly toolSearchAvailable?: boolean; |
| | 37 | readonly toolSearchDisabled?: boolean; |
| | 38 | readonly allowsTool?: (name: string) => boolean; |
| | 39 | } = {}, |
| | 40 | ) { |
| | 41 | const handlers = new Map<string, AgentStartHook>(); |
| | 42 | const tools: RegisteredTool[] = []; |
| | 43 | const requests: Array<{ readonly method: string; readonly params?: unknown }> = []; |
| | 44 | let activeTools = ["read"]; |
| | 45 | const transports: Array<{ |
| | 46 | readonly url: string; |
| | 47 | readonly authorization: string; |
| | 48 | readonly signal: AbortSignal | undefined; |
| | 49 | }> = []; |
| | 50 | const servers: Array<{ readonly name: string; readonly config: Record<string, unknown> }> = []; |
| | 51 | const catalog = [ |
| | 52 | { name: "orchestrator_capabilities", description: "Discover available providers and models." }, |
| | 53 | { name: "delegate_task", description: "Delegate work to another agent." }, |
| | 54 | { name: "task_status", description: "Check delegated work." }, |
| | 55 | { name: "preview_snapshot", description: "Inspect the collaborative browser." }, |
| | 56 | ].map((tool) => ({ ...tool, inputSchema: { type: "object", properties: {} } })); |
| | 57 | const source = NodeModule.stripTypeScriptTypes( |
| | 58 | PI_T3_MCP_EXTENSION_SOURCE.replace('import { Type } from "typebox";', "").replace( |
| | 59 | "export default async function", |
| | 60 | "async function", |
| | 61 | ), |
| | 62 | ); |
| | 63 | await NodeVM.runInNewContext(`${source}\nt3McpExtension(pi)`, { |
| | 64 | process: { |
| | 65 | env: { T3_MCP_URL: "http://fixture.invalid/mcp", T3_MCP_BEARER_TOKEN: "fixture-token" }, |
| | 66 | }, |
| | 67 | AbortSignal, |
| | 68 | Type: { Unsafe: (schema: unknown) => schema }, |
| | 69 | fetch: async ( |
| | 70 | url: string, |
| | 71 | options: { body: string; headers: Record<string, string>; signal?: AbortSignal }, |
| | 72 | ) => { |
| | 73 | transports.push({ |
| | 74 | url, |
| | 75 | authorization: options.headers.authorization!, |
| | 76 | signal: options.signal, |
| | 77 | }); |
| | 78 | const request = JSON.parse(options.body) as { id: number; method: string; params?: unknown }; |
| | 79 | requests.push(request); |
| | 80 | const result = |
| | 81 | request.method === "tools/list" |
| | 82 | ? { tools: catalog } |
| | 83 | : request.method === "tools/call" |
| | 84 | ? { content: [{ type: "text", text: "browser snapshot" }] } |
| | 85 | : {}; |
| | 86 | return new Response(JSON.stringify({ jsonrpc: "2.0", id: request.id, result }), { |
| | 87 | headers: { "content-type": "application/json" }, |
| | 88 | }); |
| | 89 | }, |
| | 90 | pi: { |
| | 91 | on: (name: string, handler: AgentStartHook) => handlers.set(name, handler), |
| | 92 | registerTool: (tool: RegisteredTool) => { |
| | 93 | if (options.allowsTool && !options.allowsTool(tool.name)) return; |
| | 94 | const index = tools.findIndex((current) => current.name === tool.name); |
| | 95 | if (index === -1) tools.push(tool); |
| | 96 | else tools[index] = tool; |
| | 97 | }, |
| | 98 | getActiveTools: () => activeTools, |
| | 99 | setActiveTools: (names: string[]) => { |
| | 100 | activeTools = names.filter((name) => options.allowsTool?.(name) ?? true); |
| | 101 | }, |
| | 102 | getAllTools: () => |
| | 103 | options.toolSearchAvailable && |
| | 104 | !options.toolSearchDisabled && |
| | 105 | (options.allowsTool?.("tool_search") ?? true) |
| | 106 | ? [{ name: "tool_search", sourceInfo: { path: "builtin:tool-search" } }] |
| | 107 | : [], |
| | 108 | ...(options.modern |
| | 109 | ? { |
| | 110 | registerMcpServer: (name: string, config: Record<string, unknown>) => |
| | 111 | servers.push({ name, config }), |
| | 112 | unregisterMcpServer: () => servers.splice(0), |
| | 113 | } |
| | 114 | : {}), |
| | 115 | }, |
| | 116 | }); |
| | 117 | return { |
| | 118 | handlers, |
| | 119 | tools, |
| | 120 | requests, |
| | 121 | servers, |
| | 122 | transports, |
| | 123 | getActiveTools: () => activeTools, |
| | 124 | restoreActiveTools: (names: string[]) => { |
| | 125 | activeTools = names; |
| | 126 | }, |
| | 127 | }; |
| | 128 | } |
| | 129 | |
| | 130 | describe("Pi MCP tool exposure", () => { |
| | 131 | it("keeps orchestration direct and optional bridge tools discoverable on modern Pi", async () => { |
| | 132 | const bridge = await loadMcpBridge({ modern: true, toolSearchAvailable: true }); |
| | 133 | await bridge.handlers.get("session_start")!( |
| | 134 | { systemPrompt: "" }, |
| | 135 | { ui: { notify: () => undefined } }, |
| | 136 | ); |
| | 137 | const start = bridge.handlers.get("before_agent_start"); |
| | 138 | assert.isDefined(start); |
| | 139 | const prompt = await start!( |
| | 140 | { systemPrompt: "Pi system prompt" }, |
| | 141 | { ui: { notify: () => undefined } }, |
| | 142 | ); |
| | 143 | assert.include(prompt.systemPrompt, "orchestrator_capabilities"); |
| | 144 | assert.equal(bridge.servers.length, 0); |
| | 145 | assert.equal(bridge.tools.length, 8); |
| | 146 | assert.deepEqual(bridge.getActiveTools(), ["read", "tool_search"]); |
| | 147 | assert.deepEqual( |
| | 148 | bridge.tools |
| | 149 | .filter((tool) => tool.exposure !== "hidden") |
| | 150 | .map((tool) => [tool.name, tool.exposure]), |
| | 151 | [ |
| | 152 | ["mcp__t3-code__orchestrator_capabilities", "direct"], |
| | 153 | ["mcp__t3-code__delegate_task", "direct"], |
| | 154 | ["mcp__t3-code__task_status", "direct"], |
| | 155 | ["mcp__t3-code__preview_snapshot", "deferred"], |
| | 156 | ], |
| | 157 | ); |
| | 158 | assert.deepEqual( |
| | 159 | bridge.tools.filter((tool) => tool.exposure === "hidden").map((tool) => tool.name), |
| | 160 | [ |
| | 161 | "mcp__t3_code__orchestrator_capabilities", |
| | 162 | "mcp__t3_code__delegate_task", |
| | 163 | "mcp__t3_code__task_status", |
| | 164 | "mcp__t3_code__preview_snapshot", |
| | 165 | ], |
| | 166 | ); |
| | 167 | const result = await bridge.tools |
| | 168 | .find((tool) => tool.name === "mcp__t3-code__preview_snapshot")! |
| | 169 | .execute("call-1", { depth: 2 }); |
| | 170 | assert.equal(result.content[0]?.text, "browser snapshot"); |
| | 171 | assert.equal(bridge.requests.at(-1)?.method, "tools/call"); |
| | 172 | }); |
| | 173 | |
| | 174 | it.each(["legacy Pi", "disabled tool search"])( |
| | 175 | "keeps tool execution available with %s", |
| | 176 | async (mode) => { |
| | 177 | const bridge = await loadMcpBridge({ |
| | 178 | modern: mode !== "legacy Pi", |
| | 179 | toolSearchAvailable: mode === "disabled tool search", |
| | 180 | toolSearchDisabled: mode === "disabled tool search", |
| | 181 | }); |
| | 182 | if (mode !== "legacy Pi") { |
| | 183 | const start = bridge.handlers.get("session_start"); |
| | 184 | await start!({ systemPrompt: "Pi system prompt" }, { ui: { notify: () => undefined } }); |
| | 185 | } |
| | 186 | assert.equal(bridge.tools.filter((tool) => tool.exposure !== "hidden").length, 4); |
| | 187 | assert.isTrue( |
| | 188 | bridge.tools |
| | 189 | .filter((tool) => tool.exposure !== "hidden") |
| | 190 | .every((tool) => tool.exposure === undefined || tool.exposure === "direct"), |
| | 191 | ); |
| | 192 | const tool = bridge.tools.find((tool) => tool.name === "mcp__t3-code__preview_snapshot"); |
| | 193 | assert.isDefined(tool); |
| | 194 | const controller = new AbortController(); |
| | 195 | const result = await tool!.execute("call-1", { depth: 2 }, controller.signal); |
| | 196 | assert.equal(result.content[0]?.text, "browser snapshot"); |
| | 197 | assert.strictEqual(bridge.transports.at(-1)?.signal, controller.signal); |
| | 198 | assert.equal(bridge.transports.at(-1)?.url, "http://fixture.invalid/mcp"); |
| | 199 | assert.equal(bridge.transports.at(-1)?.authorization, "Bearer fixture-token"); |
| | 200 | assert.deepEqual(JSON.parse(JSON.stringify(bridge.requests.at(-1))), { |
| | 201 | jsonrpc: "2.0", |
| | 202 | id: 3, |
| | 203 | method: "tools/call", |
| | 204 | params: { name: "preview_snapshot", arguments: { depth: 2 } }, |
| | 205 | }); |
| | 206 | assert.isUndefined(tool?.promptSnippet); |
| | 207 | assert.isUndefined(tool?.promptGuidelines); |
| | 208 | }, |
| | 209 | ); |
| | 210 | |
| | 211 | it("preserves legacy wildcard tool selection", async () => { |
| | 212 | const bridge = await loadMcpBridge({ |
| | 213 | modern: true, |
| | 214 | toolSearchAvailable: true, |
| | 215 | allowsTool: (name) => |
| | 216 | name === "read" || name === "tool_search" || name.startsWith("mcp__t3-code__"), |
| | 217 | }); |
| | 218 | await bridge.handlers.get("session_start")!( |
| | 219 | { systemPrompt: "" }, |
| | 220 | { ui: { notify: () => undefined } }, |
| | 221 | ); |
| | 222 | assert.equal(bridge.tools.length, 4); |
| | 223 | assert.equal(bridge.tools.filter((tool) => tool.exposure === "direct").length, 3); |
| | 224 | const tool = bridge.tools.find((tool) => tool.name === "mcp__t3-code__preview_snapshot"); |
| | 225 | assert.isDefined(tool); |
| | 226 | assert.equal((await tool!.execute("selected", {})).content[0]?.text, "browser snapshot"); |
| | 227 | }); |
| | 228 | |
| | 229 | it.each([false, true])( |
| | 230 | "reconciles tree loadouts while honoring search exclusion: %s", |
| | 231 | async (excludeSearch) => { |
| | 232 | const bridge = await loadMcpBridge({ |
| | 233 | modern: true, |
| | 234 | toolSearchAvailable: true, |
| | 235 | allowsTool: (name) => |
| | 236 | name !== "mcp__t3-code__delegate_task" && (!excludeSearch || name !== "tool_search"), |
| | 237 | }); |
| | 238 | await bridge.handlers.get("session_start")!( |
| | 239 | { systemPrompt: "" }, |
| | 240 | { ui: { notify: () => undefined } }, |
| | 241 | ); |
| | 242 | bridge.restoreActiveTools([ |
| | 243 | "read", |
| | 244 | "mcp__t3-code__task_status", |
| | 245 | "mcp__t3-code__preview_snapshot", |
| | 246 | ]); |
| | 247 | const tree = bridge.handlers.get("session_tree"); |
| | 248 | assert.isDefined(tree); |
| | 249 | await tree!({ systemPrompt: "" }, { ui: { notify: () => undefined } }); |
| | 250 | assert.deepEqual(bridge.getActiveTools(), [ |
| | 251 | "read", |
| | 252 | "mcp__t3-code__task_status", |
| | 253 | "mcp__t3-code__preview_snapshot", |
| | 254 | ...(!excludeSearch ? ["tool_search"] : []), |
| | 255 | ]); |
| | 256 | assert.isFalse( |
| | 257 | bridge.tools.some( |
| | 258 | (tool) => tool.exposure !== "hidden" && tool.name.endsWith("__delegate_task"), |
| | 259 | ), |
| | 260 | ); |
| | 261 | }, |
| | 262 | ); |
| | 263 | }); |
| | 264 | |
| | 265 | describe("Pi tool discovery permissions", () => { |
| | 266 | it("allows discovery without confirmation and still gates the discovered tool", async () => { |
| | 267 | type ToolCallHook = ( |
| | 268 | event: { toolName: string; input: unknown }, |
| | 269 | ctx: { ui: { confirm: (title: string, detail: string) => Promise<boolean> } }, |
| | 270 | ) => Promise<{ block: true; reason: string } | undefined>; |
| | 271 | let toolCall: ToolCallHook | undefined; |
| | 272 | let searchPath = "builtin:tool-search"; |
| | 273 | const source = NodeModule.stripTypeScriptTypes( |
| | 274 | PI_T3_MCP_EXTENSION_SOURCE.replace('import { Type } from "typebox";', "").replace( |
| | 275 | "export default async function", |
| | 276 | "async function", |
| | 277 | ), |
| | 278 | ); |
| | 279 | await NodeVM.runInNewContext(`${source}\nt3McpExtension(pi)`, { |
| | 280 | process: { env: { T3_PI_RUNTIME_MODE: "approval-required" } }, |
| | 281 | pi: { |
| | 282 | on: (name: string, handler: ToolCallHook) => { |
| | 283 | if (name === "tool_call") toolCall = handler; |
| | 284 | }, |
| | 285 | getAllTools: () => [{ name: "tool_search", sourceInfo: { path: searchPath } }], |
| | 286 | }, |
| | 287 | }); |
| | 288 | assert.isDefined(toolCall); |
| | 289 | const confirmations: string[] = []; |
| | 290 | const ctx = { |
| | 291 | ui: { |
| | 292 | confirm: async (title: string) => { |
| | 293 | confirmations.push(title); |
| | 294 | return false; |
| | 295 | }, |
| | 296 | }, |
| | 297 | }; |
| | 298 | assert.isUndefined( |
| | 299 | await toolCall!({ toolName: "tool_search", input: { query: "preview_snapshot" } }, ctx), |
| | 300 | ); |
| | 301 | assert.equal(confirmations.length, 0); |
| | 302 | const result = await toolCall!({ toolName: "mcp__t3-code__preview_snapshot", input: {} }, ctx); |
| | 303 | assert.equal(result?.block, true); |
| | 304 | assert.deepEqual(confirmations, ["Allow mcp__t3-code__preview_snapshot?"]); |
| | 305 | |
| | 306 | // An extension that replaces the search builtin is not known to be read-only. |
| | 307 | searchPath = "/extensions/custom-search.ts"; |
| | 308 | const replaced = await toolCall!({ toolName: "tool_search", input: {} }, ctx); |
| | 309 | assert.equal(replaced?.block, true); |
| | 310 | assert.equal(confirmations.at(-1), "Allow tool_search?"); |
| | 311 | }); |
| | 312 | }); |
| | 313 | |
| 12 | async function loadRequestHook(): Promise<RequestHook> { | 314 | async function loadRequestHook(): Promise<RequestHook> { |
| 13 | const handlers = new Map<string, RequestHook>(); | 315 | const handlers = new Map<string, RequestHook>(); |
| 14 | // Execute the shipped extension with MCP disabled; this path needs no Typebox. | 316 | // Execute the shipped extension with MCP disabled; this path needs no Typebox. |