Skip to content
Source
Sign inSign up

mcpExtensionSource.ts

507 lines20.5 KB

1// No T3 imports inside the string; resolve Pi and typebox from the user's Pi install.
2import { T3_CODE_ORCHESTRATION_INSTRUCTIONS } from "@t3tools/provider-core/server/orchestrationInstructions";
3​
4export const PI_T3_MCP_EXTENSION_FILENAME = "pi-t3-mcp-extension.ts";
5​
6export const T3_MCP_URL_ENV = "T3_MCP_URL";
7export const T3_MCP_BEARER_ENV = "T3_MCP_BEARER_TOKEN";
8export const T3_PI_RUNTIME_MODE_ENV = "T3_PI_RUNTIME_MODE";
9export const T3_PI_MCP_EXTENSION_PATH_ENV = "T3_PI_MCP_EXTENSION_PATH";
10export const PI_T3_SERVICE_TIER_COMMAND = "t3-service-tier";
11​
12// These tools trigger file-change approvals, bypassed by Auto-accept edits.
13export const PI_FILE_CHANGE_TOOLS = ["edit", "write"] as const;
14​
15export const PI_T3_MCP_EXTENSION_SOURCE = `\
16import { stripFrontmatter, type ExtensionAPI, type ExtensionContext } from "@earendil-works/pi-coding-agent";
17import * as NodeCrypto from "node:crypto";
18import * as NodeFSP from "node:fs/promises";
19import * as NodePath from "node:path";
20import { Type } from "typebox";
21​
22const URL_ENV = ${JSON.stringify(T3_MCP_URL_ENV)};
23const TOKEN_ENV = ${JSON.stringify(T3_MCP_BEARER_ENV)};
24const RUNTIME_MODE_ENV = ${JSON.stringify(T3_PI_RUNTIME_MODE_ENV)};
25const EXTENSION_PATH_ENV = ${JSON.stringify(T3_PI_MCP_EXTENSION_PATH_ENV)};
26const ORCHESTRATION_INSTRUCTIONS = ${JSON.stringify(T3_CODE_ORCHESTRATION_INSTRUCTIONS.trim())};
27const SERVICE_TIER_COMMAND = ${JSON.stringify(PI_T3_SERVICE_TIER_COMMAND)};
28const PROTOCOL = "2025-06-18";
29const READ_ONLY_TOOLS = new Set(["read", "grep", "find", "ls"]);
30const FILE_CHANGE_TOOLS = new Set(${JSON.stringify(PI_FILE_CHANGE_TOOLS)});
31​
32type RuntimeMode = "approval-required" | "auto-accept-edits" | "auto" | "full-access";
33​
34type JsonValue = string | number | boolean | null | JsonValue[] | { [key: string]: JsonValue };
35​
36type JsonRpcResponse = {
37 readonly id?: number | string;
38 readonly result?: JsonValue;
39 readonly error?: { readonly message?: string };
40};
41​
42type McpTool = {
43 readonly name: string;
44 readonly description?: string;
45 readonly inputSchema?: Record<string, unknown>;
46 readonly annotations?: { readonly readOnlyHint?: boolean };
47 readonly outputSchema?: Record<string, unknown>;
48};
49​
50function env(name: string): string | undefined {
51 const value = process.env[name];
52 return value && value.length > 0 ? value : undefined;
54​
55function runtimeMode(): RuntimeMode {
56 const value = env(RUNTIME_MODE_ENV);
57 return value === "approval-required" ||
58 value === "auto-accept-edits" ||
59 value === "auto" ||
60 value === "full-access"
61 ? value
62 : "full-access";
64​
65// T3 keys "Always allow this session" on this text, so a cut summary ends with the
66// full input's hash: calls that differ only past the cut get separate grants.
67function toolInputSummary(input: unknown): string {
68 let text: string;
69 try {
70 text = JSON.stringify(input, null, 2) ?? String(input);
71 } catch {
72 text = String(input);
73 }
74 if (text.length <= 4_000) return text;
75 const hash = NodeCrypto.createHash("sha256").update(text, "utf16le").digest("hex");
76 return text.slice(0, 4_000) + "\\n... " + (text.length - 4_000) +
77 " more characters not shown (sha256 " + hash + ")";
79​
80function parseSseOrJson(body: string, contentType: string): JsonRpcResponse {
81 if (contentType.includes("text/event-stream")) {
82 for (const line of body.split("\\n")) {
83 const trimmed = line.startsWith("data:") ? line.slice(5).trim() : "";
84 if (trimmed.length === 0) continue;
85 const parsed = JSON.parse(trimmed) as JsonRpcResponse;
86 if (parsed.id !== undefined || parsed.result !== undefined || parsed.error !== undefined) {
87 return parsed;
88 }
89 }
90 throw new Error("MCP SSE response had no JSON-RPC payload.");
91 }
92 return JSON.parse(body) as JsonRpcResponse;
94​
95function jsonSchemaToTypebox(schema: Record<string, unknown> | undefined) {
96 const unsafe = (Type as { Unsafe?: (value: unknown) => unknown }).Unsafe;
97 if (typeof unsafe === "function" && schema !== undefined) {
98 // Zero-argument MCP tools may omit "properties", which some OpenAI-compatible
99 // servers reject. Declare it, as Pi does for its own MCP tools.
100 return unsafe({
101 ...schema,
102 type: schema.type ?? "object",
103 ...(schema.properties === undefined ? { properties: {} } : {}),
104 });
105 }
106 return Type.Object({}, { additionalProperties: true });
108​
109type ModelContent =
110 | { type: "text"; text: string }
111 | { type: "image"; data: string; mimeType: string };
112​
113function mcpModelContent(result: unknown): ModelContent[] {
114 if (result === null || result === undefined) return [];
115 if (typeof result !== "object") return [{ type: "text", text: String(result) }];
116 const record = result as {
117 readonly content?: ReadonlyArray<{
118 readonly type?: string;
119 readonly text?: string;
120 readonly data?: string;
121 readonly mimeType?: string;
122 }>;
123 readonly structuredContent?: unknown;
124 };
125 const content: ModelContent[] = [];
126 if (Array.isArray(record.content)) {
127 for (const part of record.content) {
128 if (part?.type === "text" && typeof part.text === "string") {
129 content.push({ type: "text", text: part.text });
130 } else if (part?.type === "image" && typeof part.data === "string" && typeof part.mimeType === "string") {
131 content.push({ type: "image", data: part.data, mimeType: part.mimeType });
132 }
133 }
134 }
135 // Structured output belongs to scripts. Only use it as model text when
136 // the server supplied no model-facing blocks, matching Pi's native MCP tools.
137 if (content.length === 0) {
138 content.push({ type: "text", text: JSON.stringify(record.structuredContent ?? result) });
139 }
140 return content;
142​
143function mcpScriptResult(result: JsonValue | undefined) {
144 if (typeof result !== "object" || result === null || Array.isArray(result)) return result;
145 const { _meta: _ignored, ...scriptResult } = result;
146 return scriptResult;
148​
149function mcpOutputSchema(tool: McpTool) {
150 return jsonSchemaToTypebox({
151 type: "object",
152 properties: {
153 content: { type: "array", items: { type: "object" } },
154 ...(tool.outputSchema === undefined ? {} : { structuredContent: tool.outputSchema }),
155 isError: { type: "boolean" },
156 _meta: { type: "object" },
157 },
158 required: ["content"],
159 });
161​
162function isMcpToolError(result: unknown): boolean {
163 return (
164 typeof result === "object" &&
165 result !== null &&
166 "isError" in result &&
167 result.isError === true
168 );
170​
171function createMcpClient(endpoint: string, token: string) {
172 let nextId = 1;
173 let sessionId: string | undefined;
174​
175 const headers = (): Record<string, string> => {
176 const next: Record<string, string> = {
177 accept: "application/json, text/event-stream",
178 authorization: token.startsWith("Bearer ") ? token : \`Bearer \${token}\`,
179 "content-type": "application/json",
180 // Effect's HTTP MCP rejects post-initialize requests without this
181 // (400). The worktree client in McpHttpServer tests sends the same
182 // header; initialize itself does not require it.
183 "mcp-protocol-version": PROTOCOL,
184 };
185 if (sessionId !== undefined) next["mcp-session-id"] = sessionId;
186 return next;
187 };
188​
189 const request = async (method: string, params?: unknown, signal?: AbortSignal) => {
190 const id = nextId++;
191 const response = await fetch(endpoint, {
192 method: "POST",
193 headers: headers(),
194 body: JSON.stringify({ jsonrpc: "2.0", id, method, params }),
195 signal,
196 });
197 const nextSession = response.headers.get("mcp-session-id");
198 if (nextSession) sessionId = nextSession;
199 const body = await response.text();
200 if (!response.ok) {
201 throw new Error(\`MCP \${method} failed (\${response.status}): \${body.slice(0, 400)}\`);
202 }
203 if (body.length === 0) return undefined;
204 const parsed = parseSseOrJson(body, response.headers.get("content-type") ?? "");
205 if (parsed.error) {
206 throw new Error(parsed.error.message ?? \`MCP \${method} returned an error\`);
207 }
208 return parsed.result;
209 };
210​
211 const notify = async (method: string, params?: unknown, signal?: AbortSignal) => {
212 await fetch(endpoint, {
213 method: "POST",
214 headers: headers(),
215 body: JSON.stringify({ jsonrpc: "2.0", method, params }),
216 signal,
217 });
218 };
219​
220 return {
221 async connect(signal?: AbortSignal) {
222 await request(
223 "initialize",
224 {
225 protocolVersion: PROTOCOL,
226 capabilities: {},
227 clientInfo: { name: "t3-pi-mcp", version: "1.0.0" },
228 },
229 signal,
230 );
231 await notify("notifications/initialized", {}, signal).catch(() => undefined);
232 },
233 async listTools(signal?: AbortSignal) {
234 const tools: McpTool[] = [];
235 let cursor: string | undefined;
236 do {
237 const result = (await request(
238 "tools/list",
239 cursor === undefined ? {} : { cursor },
240 signal,
241 )) as { tools?: McpTool[]; nextCursor?: string } | undefined;
242 tools.push(...(result?.tools ?? []));
243 cursor = result?.nextCursor;
244 } while (cursor);
245 return tools;
246 },
247 async callTool(name: string, args: Record<string, unknown>, signal?: AbortSignal) {
248 return request("tools/call", { name, arguments: args }, signal);
249 },
250 };
252​
253export default async function t3McpExtension(pi: ExtensionAPI) {
254 // RPC has no trust prompt, so Pi skips an untrusted folder's protected resources. Session
255 // replacement reruns this factory without project_trust, and T3 drops notices outside a turn,
256 // so check each session's first turn. An explicit --no-approve is deliberate and stays quiet.
257 let projectTrustChecked = false;
258 pi.on("before_agent_start", (_event, ctx) => {
259 if (projectTrustChecked) return;
260 projectTrustChecked = true;
261 if (ctx.isProjectTrusted() || process.argv.some((arg) => arg === "--no-approve" || arg === "-na")) return;
262 ctx.ui.notify(
263 \`Pi did not load project settings, extensions, skills, prompts, MCP servers or system prompt files for \${ctx.cwd} because the folder is not trusted. To load them, run the configured Pi binary there with the same PI_CODING_AGENT_DIR (if set), trust the folder or its parent with /trust, then start a new thread in that same worktree. To trust every project, add --approve to Pi's launch arguments in T3 Code settings; to keep them off without this notice, add --no-approve.\`,
264 "warning",
265 );
266 });
267​
268 // Keep skill chips in the request. Pi expands only one leading /skill:
269 // command, so use its own frontmatter loader for every selected chip here.
270 pi.on("input", async (event, ctx) => {
271 if (!event.text.includes("$")) return;
272 const commands = pi.getCommands().filter((command) => command.source === "skill");
273 const nativeSkill = /^\\/skill:([^ ]+)(?: |$)/.exec(event.text)?.[1];
274 const seen = new Set<string>(nativeSkill === undefined ? [] : [nativeSkill]);
275 const blocks: string[] = [];
276 for (const match of event.text.matchAll(/(^|\\s)\\$([^\\s]+)(?=\\s|$)/g)) {
277 const name = match[2];
278 if (name === undefined || seen.has(name)) continue;
279 const command = commands.find((candidate) => candidate.name === "skill:" + name);
280 if (command === undefined) continue;
281 seen.add(name);
282 const path = command.sourceInfo.path;
283 try {
284 const body = stripFrontmatter(await NodeFSP.readFile(path, "utf8")).trim();
285 blocks.push(
286 \`<skill name="\${name}" location="\${path}">\\nReferences are relative to \${NodePath.dirname(path)}.\\n\\n\${body}\\n</skill>\`,
287 );
288 } catch {
289 ctx.ui.notify(\`Could not load skill \${name} from \${path}.\`, "error");
290 }
291 }
292 if (blocks.length === 0) return;
293 return {
294 action: "transform",
295 text: event.text + "\\n\\n" + blocks.join("\\n\\n"),
296 images: event.images,
297 };
298 });
299​
300 // Pi has no service tier setting, so T3 sets one through this command before each turn.
301 // Only a tier the active model advertises is sent; Standard sends none.
302 let serviceTier: string | undefined;
303 pi.registerCommand(SERVICE_TIER_COMMAND, {
304 description: "Set the service tier T3 Code sends with provider requests",
305 handler: async (args) => {
306 serviceTier = args.trim() || undefined;
307 },
308 });
309​
310 pi.on("before_provider_request", (event, ctx) => {
311 const payload = event.payload;
312 if (typeof payload !== "object" || payload === null || Array.isArray(payload)) return;
313 const replacement = { ...payload } as Record<string, unknown>;
314 let changed = false;
315 // Workaround for an upstream Pi context-budgeting bug: pi-ai reuses the
316 // previous response's usage even when a fork's instructions/tools differ,
317 // then reserves almost all remaining context for output. OpenRouter can
318 // reject even a short conversation. Remove this cap when Pi accounts for
319 // the current request prefix reliably (api/simple-options + utils/estimate).
320 if (ctx.model?.provider === "openrouter") {
321 for (const key of ["max_tokens", "max_completion_tokens"]) {
322 const limit = replacement[key];
323 if (typeof limit === "number" && Number.isFinite(limit) && limit > 32_768) {
324 replacement[key] = 32_768;
325 changed = true;
326 }
327 }
328 }
329 const tiers = (ctx.model as { serviceTiers?: unknown } | undefined)?.serviceTiers;
330 if (
331 serviceTier !== undefined &&
332 Array.isArray(tiers) &&
333 tiers.some((tier) => tier?.id === serviceTier)
334 ) {
335 replacement.service_tier = serviceTier;
336 changed = true;
337 }
338 if (changed) return replacement;
339 });
340​
341 // Pi deliberately leaves permission policy to extensions. T3's injected
342 // bridge uses Pi's public blocking tool hook so the shared runtime modes
343 // keep their normal meaning without replacing or shadowing Pi's runtime.
344 // Only Pi's own search is known to be read-only. An extension that replaces
345 // it keeps the name, and cannot discover deferred tools either.
346 const hasBuiltinToolSearch = () =>
347 typeof pi.getAllTools === "function" &&
348 pi.getAllTools().some((tool) => tool.name === "tool_search" && tool.sourceInfo?.path === "builtin:tool-search");
349​
350 const readOnlyMcpTools = new Set<string>();
351 const isReadOnlyMcpTool = (name: string) => {
352 const extensionPath = env(EXTENSION_PATH_ENV);
353 if (extensionPath === undefined || !readOnlyMcpTools.has(name) || typeof pi.getAllTools !== "function") return false;
354 // A user extension may own the same name. Only our registered HTTP bridge
355 // may inherit the canonical T3 server's read-only annotation.
356 return pi.getAllTools().some((tool) => tool.name === name &&
357 typeof tool.sourceInfo?.path === "string" &&
358 NodePath.resolve(tool.sourceInfo.path) === NodePath.resolve(extensionPath));
359 };
360​
361 pi.on("tool_call", async (event, ctx) => {
362 const mode = runtimeMode();
363 if (mode === "full-access") return;
364 if ((event.toolName === "tool_search" ? hasBuiltinToolSearch() : READ_ONLY_TOOLS.has(event.toolName)) || isReadOnlyMcpTool(event.toolName)) {
365 return;
366 }
367 if (mode === "auto-accept-edits" && FILE_CHANGE_TOOLS.has(event.toolName)) {
368 return;
369 }
370 const approved = await ctx.ui.confirm(
371 \`Allow \${event.toolName}?\`,
372 toolInputSummary(event.input),
373 { signal: ctx.signal },
374 );
375 if (!approved) {
376 return { block: true, reason: \`\${event.toolName} was declined in T3 Code.\` };
377 }
378 });
379​
380 const endpoint = env(URL_ENV);
381 const token = env(TOKEN_ENV);
382 if (endpoint === undefined || token === undefined) {
383 pi.on("session_start", async (_event, ctx) => {
384 ctx.ui.notify(
385 "t3-code MCP unavailable: T3_MCP_URL or T3_MCP_BEARER_TOKEN is missing.",
386 "warning",
387 );
388 });
389 return;
390 }
391​
392 // Tool exposure arrived with registerMcpServer in Pi 0.99. Keep the HTTP
393 // bridge as the credential owner: mcp.json overrides native registrations.
394 const supportsExposure = "registerMcpServer" in pi && typeof pi.registerMcpServer === "function";
395 const directTools = new Set(["orchestrator_capabilities", "delegate_task", "task_status"]);
396 let deferOptionalTools = supportsExposure;
397 let catalog: ReadonlyArray<McpTool> = [];
398​
399 const client = createMcpClient(endpoint, token);
400 let started: Promise<void> | undefined;
401​
402 const registerTools = () => {
403 // Preserve public names for saved loadouts and tool selectors. Hidden
404 // canonical names reserve ownership against Pi's configured MCP servers.
405 const prefixes = supportsExposure ? ["mcp__t3-code__", "mcp__t3_code__"] : ["mcp__t3-code__"];
406 readOnlyMcpTools.clear();
407 for (const tool of catalog) {
408 const name = tool.name;
409 for (const prefix of prefixes) {
410 if (tool.annotations?.readOnlyHint === true) readOnlyMcpTools.add(\`\${prefix}\${name}\`);
411 const exposure = prefix === "mcp__t3_code__" ? "hidden" :
412 deferOptionalTools && !directTools.has(name) ? "deferred" : "direct";
413 pi.registerTool({
414 name: \`\${prefix}\${name}\`,
415 label: name,
416 description: tool.description ?? name,
417 parameters: jsonSchemaToTypebox(tool.inputSchema),
418 ...(supportsExposure ? { outputSchema: mcpOutputSchema(tool) } : {}),
419 ...(supportsExposure ? { exposure } : {}),
420 async execute(_toolCallId, params, signal) {
421 const result = await client.callTool(
422 name,
423 (params ?? {}) as Record<string, unknown>,
424 signal,
425 );
426 const content = mcpModelContent(result);
427 if (isMcpToolError(result) && !content.some((part) => part.type === "text" && part.text.length > 0)) {
428 content.push({ type: "text", text: "MCP tool t3-code/" + name + " returned an error" });
429 }
430 return {
431 content,
432 structuredContent: mcpScriptResult(result),
433 details: { server: "t3-code", tool: name },
434 ...(isMcpToolError(result) ? { isError: true } : {}),
435 };
436 },
437 });
438 }
439 }
440 };
441​
442 const ensureStarted = () => {
443 if (started !== undefined) return started;
444 const attempt = (async () => {
445 const signal = AbortSignal.timeout(10_000);
446 await client.connect(signal);
447 catalog = await client.listTools(signal);
448 registerTools();
449 })();
450 started = attempt;
451 void attempt.catch(() => {
452 if (started === attempt) started = undefined;
453 });
454 return attempt;
455 };
456​
457 // CLI extensions load before builtins, and Pi keeps the first registration
458 // of a tool name. Register now so the bridge owns the T3 namespace even when
459 // mcp.json configures it; retry a failed connection at session_start.
460 await ensureStarted().catch(() => undefined);
461​
462 const reconcileDiscovery = (ctx: Pick<ExtensionContext, "model">) => {
463 if (!supportsExposure) return;
464 // A disabled or replaced search builtin cannot discover deferred tools.
465 const hasToolSearch = hasBuiltinToolSearch();
466 const exposureChanged = deferOptionalTools !== hasToolSearch;
467 deferOptionalTools = hasToolSearch;
468 if (exposureChanged) registerTools();
469 const active = pi.getActiveTools();
470 // Completion gateways can silently drop calls to undeclared tools named
471 // in history or guidance. Eagerly load T3 schemas, including after tree
472 // restores, while retaining deferred exposure for Pi's allowlist rules.
473 // Pi still applies --tools/--exclude-tools to this active loadout.
474 const required = [
475 ...(hasToolSearch ? ["tool_search"] : []),
476 ...(ctx.model?.api === "openai-completions" ? catalog.map((tool) => "mcp__t3-code__" + tool.name) : []),
477 ];
478 const missing = required.filter((name) => !active.includes(name));
479 if (missing.length > 0) pi.setActiveTools([...active, ...missing]);
480 };
481​
482 // Tree navigation restores its saved loadout after session_start.
483 pi.on("session_tree", (_event, ctx) => reconcileDiscovery(ctx));
484 pi.on("model_select", (_event, ctx) => reconcileDiscovery(ctx));
485 pi.on("session_start", async (_event, ctx) => {
486 reconcileDiscovery(ctx);
487 try {
488 await ensureStarted();
489 } catch (error) {
490 const message = error instanceof Error ? error.message : String(error);
491 ctx.ui.notify(\`t3-code MCP unavailable: \${message}\`, "warning");
492 }
493 });
494​
495 // Keep guidance in the system prompt so leading slash commands still expand.
496 pi.on("before_agent_start", (event, ctx) => {
497 reconcileDiscovery(ctx);
498 const sections = event.systemPromptOptions?.sections;
499 if (sections && typeof sections === "object" && event.systemPromptOptions?.forceSystemPrompt === undefined) {
500 // A full prompt override bypasses Pi's later section rendering.
501 sections.t3_code_orchestration = ORCHESTRATION_INSTRUCTIONS;
502 return undefined;
503 }
504 return { systemPrompt: event.systemPrompt + "\\n\\n" + ORCHESTRATION_INSTRUCTIONS };
505 });
507`;