@phi-code-admin/phi-code 0.86.0 → 0.87.0

package/extensions/phi/mcp/tool-bridge.ts

@@ -12,20 +12,19 @@
  *   - Image/audio/resource content → text description passthrough
  */
 
-import { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import type { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { CallToolResultSchema, ListToolsResultSchema } from "@modelcontextprotocol/sdk/types.js";
-import * as Type from "typebox";
-import type { TSchema } from "typebox";
 import type { ExtensionAPI } from "phi-code";
-import { McpError } from "./errors.js";
+import type { TSchema } from "typebox";
+import * as Type from "typebox";
 import type { Settings } from "./config.js";
+import { McpError } from "./errors.js";
 
 // ─── Types ────────────────────────────────────────────────────────────────────
 
 /** Subset of Pi's ExtensionAPI used by the bridge. */
 export type PiExtensionAPI = Pick<ExtensionAPI, "registerTool" | "getActiveTools" | "setActiveTools">;
 
-
 // ─── Schema Conversion ────────────────────────────────────────────────────────
 
 /**
@@ -40,162 +39,148 @@ export type PiExtensionAPI = Pick<ExtensionAPI, "registerTool" | "getActiveTools
  *   - allOf → TypeBox Intersect
  * Falls back to Type.Any() for unresolvable $ref or missing type.
  */
-export function convertJsonSchemaToTypebox(
-  schema: unknown,
-  depth = 0,
-  defs?: Record<string, unknown>,
-): TSchema {
-  // Guard against infinite recursion and malformed schemas
-  if (!schema || typeof schema !== "object" || Array.isArray(schema) || depth > 10) {
-    return Type.Any();
-  }
-
-  const s = schema as Record<string, unknown>;
-  const description = typeof s["description"] === "string" ? s["description"] : undefined;
-  const opts = description ? { description } : {};
-
-  // Extract $defs / definitions for $ref resolution (carried through recursive calls)
-  const resolvedDefs: Record<string, unknown> = {
-    ...((s["$defs"] ?? s["definitions"]) as Record<string, unknown> | undefined),
-    ...defs,
-  };
-
-  // ── Handle $ref ──────────────────────────────────────────────────────────
-  if (typeof s["$ref"] === "string") {
-    const ref = s["$ref"] as string;
-    let resolved: unknown;
-
-    // Local references: #/$defs/Foo, #/definitions/Foo
-    if (ref.startsWith("#/")) {
-      const parts = ref.slice(2).split("/");
-      if (parts[0] === "$defs" || parts[0] === "definitions") {
-        const key = parts.slice(1).join("/");
-        resolved = resolvedDefs[key];
-      } else {
-        // Fallback: try walking the defs map by the last part
-        const key = parts[parts.length - 1]!;
-        resolved = resolvedDefs[key];
-      }
-    } else {
-      // External $ref — cannot resolve, fall back
-      console.warn(
-        `[pi-mcp] Cannot resolve external $ref "${ref}", using Type.Any()`,
-      );
-      return Type.Any(opts);
-    }
-
-    if (!resolved) {
-      console.warn(
-        `[pi-mcp] Could not resolve $ref "${ref}", using Type.Any()`,
-      );
-      return Type.Any(opts);
-    }
-
-    // Merge description from referencing schema into resolved schema
-    const merged = { ...(resolved as Record<string, unknown>) };
-    if (description && !merged["description"]) {
-      merged["description"] = description;
-    }
-    return convertJsonSchemaToTypebox(merged, depth + 1, resolvedDefs);
-  }
-
-  // ── Handle oneOf / anyOf → TypeBox Union ─────────────────────────────────
-  if (Array.isArray(s["oneOf"])) {
-    const members = (s["oneOf"] as unknown[])
-      .map((sub) => convertJsonSchemaToTypebox(sub, depth + 1, resolvedDefs));
-    return members.length === 1 ? members[0]! : Type.Union(members, opts);
-  }
-
-  if (Array.isArray(s["anyOf"])) {
-    const members = (s["anyOf"] as unknown[])
-      .map((sub) => convertJsonSchemaToTypebox(sub, depth + 1, resolvedDefs));
-    return members.length === 1 ? members[0]! : Type.Union(members, opts);
-  }
-
-  // ── Handle allOf → TypeBox Intersect ─────────────────────────────────────
-  if (Array.isArray(s["allOf"])) {
-    const members = (s["allOf"] as unknown[])
-      .map((sub) => convertJsonSchemaToTypebox(sub, depth + 1, resolvedDefs));
-    return members.length === 1 ? members[0]! : Type.Intersect(members, opts);
-  }
-
-  // Handle nullable types: { "type": ["string", "null"] }
-  const rawType = s["type"];
-  const type = Array.isArray(rawType)
-    ? rawType.find((t) => t !== "null") as string | undefined
-    : typeof rawType === "string" ? rawType : undefined;
-
-  const isNullable = Array.isArray(rawType) && rawType.includes("null");
-
-  let base: TSchema;
-
-  switch (type) {
-    case "string": {
-      const enumVals = s["enum"];
-      if (Array.isArray(enumVals) && enumVals.every((v) => typeof v === "string")) {
-        // TypeBox doesn't have a built-in StringEnum — use Union of Literals
-        base = Type.Union(
-          (enumVals as string[]).map((v) => Type.Literal(v)),
-          opts,
-        );
-      } else {
-        base = Type.String(opts);
-      }
-      break;
-    }
-    case "number":
-    case "integer":
-      base = Type.Number(opts);
-      break;
-    case "boolean":
-      base = Type.Boolean(opts);
-      break;
-    case "null":
-      base = Type.Null(opts);
-      break;
-    case "array": {
-      const items = s["items"];
-      base = Type.Array(
-        items ? convertJsonSchemaToTypebox(items, depth + 1, resolvedDefs) : Type.Unknown(),
-        opts,
-      );
-      break;
-    }
-    case "object": {
-      const properties = s["properties"] as Record<string, unknown> | undefined;
-      const required = new Set<string>(
-        Array.isArray(s["required"]) ? (s["required"] as string[]) : [],
-      );
-      const additionalProperties = s["additionalProperties"];
-
-      if (!properties) {
-        // Open object — passthrough as Any to avoid over-constraining
-        base = Type.Record(Type.String(), Type.Unknown(), opts);
-        break;
-      }
-
-      const props: Record<string, TSchema> = {};
-      for (const [key, value] of Object.entries(properties)) {
-        const converted = convertJsonSchemaToTypebox(value, depth + 1, resolvedDefs);
-        props[key] = required.has(key) ? converted : Type.Optional(converted);
-      }
-
-      const objOpts: Record<string, unknown> = { ...opts };
-      if (additionalProperties === false) {
-        objOpts["additionalProperties"] = false;
-      }
-
-      base = Type.Object(props, objOpts as any);
-      break;
-    }
-    default: {
-      // Truly unsupported or missing type field
-      base = Type.Any(opts);
-      break;
-    }
-  }
-
-  return isNullable ? Type.Union([base, Type.Null()]) : base;
+export function convertJsonSchemaToTypebox(schema: unknown, depth = 0, defs?: Record<string, unknown>): TSchema {
+	// Guard against infinite recursion and malformed schemas
+	if (!schema || typeof schema !== "object" || Array.isArray(schema) || depth > 10) {
+		return Type.Any();
+	}
+
+	const s = schema as Record<string, unknown>;
+	const description = typeof s.description === "string" ? s.description : undefined;
+	const opts = description ? { description } : {};
+
+	// Extract $defs / definitions for $ref resolution (carried through recursive calls)
+	const resolvedDefs: Record<string, unknown> = {
+		...((s.$defs ?? s.definitions) as Record<string, unknown> | undefined),
+		...defs,
+	};
+
+	// ── Handle $ref ──────────────────────────────────────────────────────────
+	if (typeof s.$ref === "string") {
+		const ref = s.$ref as string;
+		let resolved: unknown;
+
+		// Local references: #/$defs/Foo, #/definitions/Foo
+		if (ref.startsWith("#/")) {
+			const parts = ref.slice(2).split("/");
+			if (parts[0] === "$defs" || parts[0] === "definitions") {
+				const key = parts.slice(1).join("/");
+				resolved = resolvedDefs[key];
+			} else {
+				// Fallback: try walking the defs map by the last part
+				const key = parts[parts.length - 1]!;
+				resolved = resolvedDefs[key];
+			}
+		} else {
+			// External $ref — cannot resolve, fall back
+			console.warn(`[pi-mcp] Cannot resolve external $ref "${ref}", using Type.Any()`);
+			return Type.Any(opts);
+		}
+
+		if (!resolved) {
+			console.warn(`[pi-mcp] Could not resolve $ref "${ref}", using Type.Any()`);
+			return Type.Any(opts);
+		}
+
+		// Merge description from referencing schema into resolved schema
+		const merged = { ...(resolved as Record<string, unknown>) };
+		if (description && !merged.description) {
+			merged.description = description;
+		}
+		return convertJsonSchemaToTypebox(merged, depth + 1, resolvedDefs);
+	}
+
+	// ── Handle oneOf / anyOf → TypeBox Union ─────────────────────────────────
+	if (Array.isArray(s.oneOf)) {
+		const members = (s.oneOf as unknown[]).map((sub) => convertJsonSchemaToTypebox(sub, depth + 1, resolvedDefs));
+		return members.length === 1 ? members[0]! : Type.Union(members, opts);
+	}
+
+	if (Array.isArray(s.anyOf)) {
+		const members = (s.anyOf as unknown[]).map((sub) => convertJsonSchemaToTypebox(sub, depth + 1, resolvedDefs));
+		return members.length === 1 ? members[0]! : Type.Union(members, opts);
+	}
+
+	// ── Handle allOf → TypeBox Intersect ─────────────────────────────────────
+	if (Array.isArray(s.allOf)) {
+		const members = (s.allOf as unknown[]).map((sub) => convertJsonSchemaToTypebox(sub, depth + 1, resolvedDefs));
+		return members.length === 1 ? members[0]! : Type.Intersect(members, opts);
+	}
+
+	// Handle nullable types: { "type": ["string", "null"] }
+	const rawType = s.type;
+	const type = Array.isArray(rawType)
+		? (rawType.find((t) => t !== "null") as string | undefined)
+		: typeof rawType === "string"
+			? rawType
+			: undefined;
+
+	const isNullable = Array.isArray(rawType) && rawType.includes("null");
+
+	let base: TSchema;
+
+	switch (type) {
+		case "string": {
+			const enumVals = s.enum;
+			if (Array.isArray(enumVals) && enumVals.every((v) => typeof v === "string")) {
+				// TypeBox doesn't have a built-in StringEnum — use Union of Literals
+				base = Type.Union(
+					(enumVals as string[]).map((v) => Type.Literal(v)),
+					opts,
+				);
+			} else {
+				base = Type.String(opts);
+			}
+			break;
+		}
+		case "number":
+		case "integer":
+			base = Type.Number(opts);
+			break;
+		case "boolean":
+			base = Type.Boolean(opts);
+			break;
+		case "null":
+			base = Type.Null(opts);
+			break;
+		case "array": {
+			const items = s.items;
+			base = Type.Array(items ? convertJsonSchemaToTypebox(items, depth + 1, resolvedDefs) : Type.Unknown(), opts);
+			break;
+		}
+		case "object": {
+			const properties = s.properties as Record<string, unknown> | undefined;
+			const required = new Set<string>(Array.isArray(s.required) ? (s.required as string[]) : []);
+			const additionalProperties = s.additionalProperties;
+
+			if (!properties) {
+				// Open object — passthrough as Any to avoid over-constraining
+				base = Type.Record(Type.String(), Type.Unknown(), opts);
+				break;
+			}
+
+			const props: Record<string, TSchema> = {};
+			for (const [key, value] of Object.entries(properties)) {
+				const converted = convertJsonSchemaToTypebox(value, depth + 1, resolvedDefs);
+				props[key] = required.has(key) ? converted : Type.Optional(converted);
+			}
+
+			const objOpts: Record<string, unknown> = { ...opts };
+			if (additionalProperties === false) {
+				objOpts.additionalProperties = false;
+			}
+
+			base = Type.Object(props, objOpts as any);
+			break;
+		}
+		default: {
+			// Truly unsupported or missing type field
+			base = Type.Any(opts);
+			break;
+		}
+	}
+
+	return isNullable ? Type.Union([base, Type.Null()]) : base;
 }
 
 // ─── Tool Name Sanitization ───────────────────────────────────────────────────
@@ -209,16 +194,14 @@ const MAX_TOOL_NAME_LEN = 64;
  * If truncation is needed, the last 8 chars are replaced with a hash to avoid collisions.
  */
 export function buildToolName(prefix: string, serverName: string, toolName: string): string {
-  const raw = `${prefix}_${serverName}_${toolName}`;
-  const safe = raw.replace(/[^a-zA-Z0-9_]/g, "_");
-  if (safe.length <= MAX_TOOL_NAME_LEN) return safe;
-  // Truncate with hash suffix to prevent collisions on long names
-  const hash = Math.abs(
-    safe.split("").reduce((acc, c) => ((acc << 5) - acc + c.charCodeAt(0)) | 0, 0),
-  )
-    .toString(36)
-    .slice(0, 8);
-  return safe.slice(0, MAX_TOOL_NAME_LEN - 9) + "_" + hash;
+	const raw = `${prefix}_${serverName}_${toolName}`;
+	const safe = raw.replace(/[^a-zA-Z0-9_]/g, "_");
+	if (safe.length <= MAX_TOOL_NAME_LEN) return safe;
+	// Truncate with hash suffix to prevent collisions on long names
+	const hash = Math.abs(safe.split("").reduce((acc, c) => ((acc << 5) - acc + c.charCodeAt(0)) | 0, 0))
+		.toString(36)
+		.slice(0, 8);
+	return `${safe.slice(0, MAX_TOOL_NAME_LEN - 9)}_${hash}`;
 }
 
 // ─── Content Conversion ───────────────────────────────────────────────────────
@@ -226,48 +209,48 @@ export function buildToolName(prefix: string, serverName: string, toolName: stri
 type PiTextContent = { type: "text"; text: string };
 
 function convertMcpContent(items: unknown[]): PiTextContent[] {
-  return items.map((item: any) => {
-    if (!item || typeof item !== "object") {
-      return { type: "text", text: String(item) };
-    }
-    switch (item.type) {
-      case "text":
-        return { type: "text", text: String(item.text ?? "") };
-      case "image":
-        return {
-          type: "text",
-          text: `[Image: ${item.mimeType ?? "unknown"}, base64 encoded]`,
-        };
-      case "audio":
-        return {
-          type: "text",
-          text: `[Audio: ${item.mimeType ?? "unknown"}, base64 encoded]`,
-        };
-      case "resource": {
-        const r = item.resource;
-        if (r?.text) return { type: "text", text: r.text };
-        if (r?.blob) return { type: "text", text: `[Resource blob: ${r.uri}]` };
-        return { type: "text", text: `[Resource: ${r?.uri ?? "unknown"}]` };
-      }
-      default:
-        return { type: "text", text: JSON.stringify(item) };
-    }
-  });
+	return items.map((item: any) => {
+		if (!item || typeof item !== "object") {
+			return { type: "text", text: String(item) };
+		}
+		switch (item.type) {
+			case "text":
+				return { type: "text", text: String(item.text ?? "") };
+			case "image":
+				return {
+					type: "text",
+					text: `[Image: ${item.mimeType ?? "unknown"}, base64 encoded]`,
+				};
+			case "audio":
+				return {
+					type: "text",
+					text: `[Audio: ${item.mimeType ?? "unknown"}, base64 encoded]`,
+				};
+			case "resource": {
+				const r = item.resource;
+				if (r?.text) return { type: "text", text: r.text };
+				if (r?.blob) return { type: "text", text: `[Resource blob: ${r.uri}]` };
+				return { type: "text", text: `[Resource: ${r?.uri ?? "unknown"}]` };
+			}
+			default:
+				return { type: "text", text: JSON.stringify(item) };
+		}
+	});
 }
 
 // ─── Tool Listing ─────────────────────────────────────────────────────────────
 
 export interface McpToolDefinition {
-  name: string;
-  description?: string;
-  inputSchema: Record<string, unknown>;
-  annotations?: {
-    readOnlyHint?: boolean;
-    destructiveHint?: boolean;
-    idempotentHint?: boolean;
-    openWorldHint?: boolean;
-    title?: string;
-  };
+	name: string;
+	description?: string;
+	inputSchema: Record<string, unknown>;
+	annotations?: {
+		readOnlyHint?: boolean;
+		destructiveHint?: boolean;
+		idempotentHint?: boolean;
+		openWorldHint?: boolean;
+		title?: string;
+	};
 }
 
 /**
@@ -275,33 +258,30 @@ export interface McpToolDefinition {
  * The MCP spec mandates clients follow nextCursor until exhausted.
  * Includes a max-page guard to prevent infinite loops from broken servers.
  */
-export async function listAllTools(
-  client: Client,
-  requestTimeoutMs: number,
-): Promise<McpToolDefinition[]> {
-  const tools: McpToolDefinition[] = [];
-  let cursor: string | undefined;
-  const MAX_PAGES = 100;
-  let pageCount = 0;
-
-  do {
-    if (pageCount >= MAX_PAGES) {
-      console.warn(
-        `[pi-mcp] tools/list pagination exceeded ${MAX_PAGES} pages, stopping. The server may be malfunctioning.`,
-      );
-      break;
-    }
-    const result = await client.request(
-      { method: "tools/list", params: cursor ? { cursor } : {} },
-      ListToolsResultSchema,
-      { timeout: requestTimeoutMs },
-    );
-    tools.push(...(result.tools as McpToolDefinition[]));
-    cursor = result.nextCursor;
-    pageCount++;
-  } while (cursor);
-
-  return tools;
+export async function listAllTools(client: Client, requestTimeoutMs: number): Promise<McpToolDefinition[]> {
+	const tools: McpToolDefinition[] = [];
+	let cursor: string | undefined;
+	const MAX_PAGES = 100;
+	let pageCount = 0;
+
+	do {
+		if (pageCount >= MAX_PAGES) {
+			console.warn(
+				`[pi-mcp] tools/list pagination exceeded ${MAX_PAGES} pages, stopping. The server may be malfunctioning.`,
+			);
+			break;
+		}
+		const result = await client.request(
+			{ method: "tools/list", params: cursor ? { cursor } : {} },
+			ListToolsResultSchema,
+			{ timeout: requestTimeoutMs },
+		);
+		tools.push(...(result.tools as McpToolDefinition[]));
+		cursor = result.nextCursor;
+		pageCount++;
+	} while (cursor);
+
+	return tools;
 }
 
 // ─── Tool Bridge ──────────────────────────────────────────────────────────────
@@ -311,184 +291,170 @@ export async function listAllTools(
  * Tools are registered once and activated/deactivated as servers connect/disconnect.
  */
 export class ToolBridge {
-  private readonly settings: Settings;
-  private readonly pi: PiExtensionAPI;
-  /** Tracks which Pi tool names belong to which MCP server. */
-  private readonly serverToolNames = new Map<string, Set<string>>();
-
-  constructor(settings: Settings, pi: PiExtensionAPI) {
-    this.settings = settings;
-    this.pi = pi;
-  }
-
-  /**
-   * Refresh tools for a server — called on initial connect and on list_changed.
-   * Always re-registers tools with the current client reference so that
-   * tool execute closures capture the latest client after reconnection.
-   * Deactivates tools that are no longer in the server's list.
-   * Note: Pi's registerTool() overwrites by name (Map.set), so re-registration is safe.
-   */
-  async refreshTools(serverName: string, client: Client): Promise<void> {
-    const timeoutMs = this.settings.requestTimeoutMs;
-
-    let tools: McpToolDefinition[];
-    try {
-      tools = await listAllTools(client, timeoutMs);
-    } catch (err) {
-      throw new McpError(
-        `Failed to list tools from ${serverName}: ${err instanceof Error ? err.message : String(err)}`,
-        serverName,
-        "protocol",
-        err,
-      );
-    }
-
-    const registeredForServer = this.serverToolNames.get(serverName) ?? new Set<string>();
-
-    // Build the set of currently valid Pi tool names for this server
-    const currentToolNames = new Set<string>();
-
-    for (const tool of tools) {
-      const piName = buildToolName(this.settings.toolPrefix, serverName, tool.name);
-      // Detect collision: two different MCP tools mapping to the same Pi name
-      // (e.g. "my-tool" and "my_tool" both sanitize to "my_tool")
-      if (currentToolNames.has(piName)) {
-        console.warn(
-          `[pi-mcp] Tool name collision: "${tool.name}" maps to "${piName}" which is already taken. ` +
-          `The later tool definition will overwrite the earlier one.`,
-        );
-      }
-      currentToolNames.add(piName);
-      // Always re-register — on reconnect the client reference changes and
-      // Pi's registerTool overwrites by name, so this is idempotent.
-      this._registerTool(piName, serverName, tool, client);
-    }
-
-    // Deactivate tools that were removed from the server (no longer in tools/list)
-    for (const existingName of registeredForServer) {
-      if (!currentToolNames.has(existingName)) {
-        this._deactivateServerTool(existingName);
-      }
-    }
-
-    this.serverToolNames.set(serverName, currentToolNames);
-
-    // Activate all current tools for this server
-    this._activateServerTools(serverName);
-  }
-
-  /** Deactivate all Pi tools belonging to a server (called on disconnect). */
-  deactivateServer(serverName: string): void {
-    this._deactivateServerTools(serverName);
-  }
-
-  /** Remove all tracking data for a server (called when config changes remove a server). */
-  removeServer(serverName: string): void {
-    this._deactivateServerTools(serverName);
-    this.serverToolNames.delete(serverName);
-  }
-
-  /** Re-activate all Pi tools belonging to a server (called on reconnect). */
-  activateServer(serverName: string): void {
-    this._activateServerTools(serverName);
-  }
-
-  // ─── Internal ───────────────────────────────────────────────────────────────
-
-  private _registerTool(
-    piName: string,
-    serverName: string,
-    tool: McpToolDefinition,
-    client: Client,
-  ): void {
-    // Build description with annotation hints for LLM guidance
-    let description = tool.description ?? `MCP tool: ${tool.name}`;
-    const ann = tool.annotations;
-    if (ann) {
-      const hints: string[] = [];
-      if (ann.readOnlyHint) hints.push("read-only");
-      if (ann.destructiveHint) hints.push("⚠️ destructive");
-      if (ann.idempotentHint) hints.push("idempotent");
-      if (ann.openWorldHint) hints.push("may have side effects");
-      if (hints.length > 0) description += ` [${hints.join(", ")}]`;
-    }
-
-    const schema = convertJsonSchemaToTypebox(tool.inputSchema);
-    const timeoutMs = this.settings.requestTimeoutMs;
-
-    this.pi.registerTool({
-      name: piName,
-      label: ann?.title ?? tool.name,
-      description,
-      promptSnippet: description.slice(0, 120),
-      parameters: schema,
-
-      async execute(_toolCallId, params, signal, _onUpdate, _ctx) {
-        if (signal?.aborted) {
-          return { content: [{ type: "text", text: "Cancelled" }], details: {} };
-        }
-
-        try {
-          const result = await client.request(
-            {
-              method: "tools/call",
-              params: { name: tool.name, arguments: params },
-            },
-            CallToolResultSchema,
-            // Pass AbortSignal to SDK — it will automatically send
-            // notifications/cancelled when the signal fires
-            { timeout: timeoutMs, ...(signal ? { signal } : {}) },
-          );
-
-          const content = convertMcpContent(result.content as unknown[]);
-
-          // Tool execution errors (isError: true) — distinct from protocol errors
-          if (result.isError) {
-            const errorText = content.map((c) => c.text).join("\n");
-            throw new McpError(
-              errorText || "Tool reported an error",
-              serverName,
-              "tool",
-            );
-          }
-
-          return { content, details: {} };
-        } catch (err) {
-          if (err instanceof McpError) throw err;
-          // Protocol-level errors (JSON-RPC error response, timeout, etc.)
-          throw new McpError(
-            err instanceof Error ? err.message : String(err),
-            serverName,
-            "protocol",
-            err,
-          );
-        }
-      },
-    });
-  }
-
-  private _activateServerTools(serverName: string): void {
-    const serverTools = this.serverToolNames.get(serverName);
-    if (!serverTools || serverTools.size === 0) return;
-
-    const currentActive = new Set(this.pi.getActiveTools());
-    for (const name of serverTools) currentActive.add(name);
-    this.pi.setActiveTools(Array.from(currentActive));
-  }
-
-  private _deactivateServerTools(serverName: string): void {
-    const serverTools = this.serverToolNames.get(serverName);
-    if (!serverTools || serverTools.size === 0) return;
-
-    const currentActive = this.pi.getActiveTools();
-    const remaining = currentActive.filter((n) => !serverTools.has(n));
-    this.pi.setActiveTools(remaining);
-  }
-
-  /** Deactivate a single tool by Pi name (used when a tool is removed on list_changed). */
-  private _deactivateServerTool(piName: string): void {
-    const currentActive = this.pi.getActiveTools();
-    const remaining = currentActive.filter((n) => n !== piName);
-    this.pi.setActiveTools(remaining);
-  }
+	private readonly settings: Settings;
+	private readonly pi: PiExtensionAPI;
+	/** Tracks which Pi tool names belong to which MCP server. */
+	private readonly serverToolNames = new Map<string, Set<string>>();
+
+	constructor(settings: Settings, pi: PiExtensionAPI) {
+		this.settings = settings;
+		this.pi = pi;
+	}
+
+	/**
+	 * Refresh tools for a server — called on initial connect and on list_changed.
+	 * Always re-registers tools with the current client reference so that
+	 * tool execute closures capture the latest client after reconnection.
+	 * Deactivates tools that are no longer in the server's list.
+	 * Note: Pi's registerTool() overwrites by name (Map.set), so re-registration is safe.
+	 */
+	async refreshTools(serverName: string, client: Client): Promise<void> {
+		const timeoutMs = this.settings.requestTimeoutMs;
+
+		let tools: McpToolDefinition[];
+		try {
+			tools = await listAllTools(client, timeoutMs);
+		} catch (err) {
+			throw new McpError(
+				`Failed to list tools from ${serverName}: ${err instanceof Error ? err.message : String(err)}`,
+				serverName,
+				"protocol",
+				err,
+			);
+		}
+
+		const registeredForServer = this.serverToolNames.get(serverName) ?? new Set<string>();
+
+		// Build the set of currently valid Pi tool names for this server
+		const currentToolNames = new Set<string>();
+
+		for (const tool of tools) {
+			const piName = buildToolName(this.settings.toolPrefix, serverName, tool.name);
+			// Detect collision: two different MCP tools mapping to the same Pi name
+			// (e.g. "my-tool" and "my_tool" both sanitize to "my_tool")
+			if (currentToolNames.has(piName)) {
+				console.warn(
+					`[pi-mcp] Tool name collision: "${tool.name}" maps to "${piName}" which is already taken. ` +
+						`The later tool definition will overwrite the earlier one.`,
+				);
+			}
+			currentToolNames.add(piName);
+			// Always re-register — on reconnect the client reference changes and
+			// Pi's registerTool overwrites by name, so this is idempotent.
+			this._registerTool(piName, serverName, tool, client);
+		}
+
+		// Deactivate tools that were removed from the server (no longer in tools/list)
+		for (const existingName of registeredForServer) {
+			if (!currentToolNames.has(existingName)) {
+				this._deactivateServerTool(existingName);
+			}
+		}
+
+		this.serverToolNames.set(serverName, currentToolNames);
+
+		// Activate all current tools for this server
+		this._activateServerTools(serverName);
+	}
+
+	/** Deactivate all Pi tools belonging to a server (called on disconnect). */
+	deactivateServer(serverName: string): void {
+		this._deactivateServerTools(serverName);
+	}
+
+	/** Remove all tracking data for a server (called when config changes remove a server). */
+	removeServer(serverName: string): void {
+		this._deactivateServerTools(serverName);
+		this.serverToolNames.delete(serverName);
+	}
+
+	/** Re-activate all Pi tools belonging to a server (called on reconnect). */
+	activateServer(serverName: string): void {
+		this._activateServerTools(serverName);
+	}
+
+	// ─── Internal ───────────────────────────────────────────────────────────────
+
+	private _registerTool(piName: string, serverName: string, tool: McpToolDefinition, client: Client): void {
+		// Build description with annotation hints for LLM guidance
+		let description = tool.description ?? `MCP tool: ${tool.name}`;
+		const ann = tool.annotations;
+		if (ann) {
+			const hints: string[] = [];
+			if (ann.readOnlyHint) hints.push("read-only");
+			if (ann.destructiveHint) hints.push("⚠️ destructive");
+			if (ann.idempotentHint) hints.push("idempotent");
+			if (ann.openWorldHint) hints.push("may have side effects");
+			if (hints.length > 0) description += ` [${hints.join(", ")}]`;
+		}
+
+		const schema = convertJsonSchemaToTypebox(tool.inputSchema);
+		const timeoutMs = this.settings.requestTimeoutMs;
+
+		this.pi.registerTool({
+			name: piName,
+			label: ann?.title ?? tool.name,
+			description,
+			promptSnippet: description.slice(0, 120),
+			parameters: schema,
+
+			async execute(_toolCallId, params, signal, _onUpdate, _ctx) {
+				if (signal?.aborted) {
+					return { content: [{ type: "text", text: "Cancelled" }], details: {} };
+				}
+
+				try {
+					const result = await client.request(
+						{
+							method: "tools/call",
+							params: { name: tool.name, arguments: params },
+						},
+						CallToolResultSchema,
+						// Pass AbortSignal to SDK — it will automatically send
+						// notifications/cancelled when the signal fires
+						{ timeout: timeoutMs, ...(signal ? { signal } : {}) },
+					);
+
+					const content = convertMcpContent(result.content as unknown[]);
+
+					// Tool execution errors (isError: true) — distinct from protocol errors
+					if (result.isError) {
+						const errorText = content.map((c) => c.text).join("\n");
+						throw new McpError(errorText || "Tool reported an error", serverName, "tool");
+					}
+
+					return { content, details: {} };
+				} catch (err) {
+					if (err instanceof McpError) throw err;
+					// Protocol-level errors (JSON-RPC error response, timeout, etc.)
+					throw new McpError(err instanceof Error ? err.message : String(err), serverName, "protocol", err);
+				}
+			},
+		});
+	}
+
+	private _activateServerTools(serverName: string): void {
+		const serverTools = this.serverToolNames.get(serverName);
+		if (!serverTools || serverTools.size === 0) return;
+
+		const currentActive = new Set(this.pi.getActiveTools());
+		for (const name of serverTools) currentActive.add(name);
+		this.pi.setActiveTools(Array.from(currentActive));
+	}
+
+	private _deactivateServerTools(serverName: string): void {
+		const serverTools = this.serverToolNames.get(serverName);
+		if (!serverTools || serverTools.size === 0) return;
+
+		const currentActive = this.pi.getActiveTools();
+		const remaining = currentActive.filter((n) => !serverTools.has(n));
+		this.pi.setActiveTools(remaining);
+	}
+
+	/** Deactivate a single tool by Pi name (used when a tool is removed on list_changed). */
+	private _deactivateServerTool(piName: string): void {
+		const currentActive = this.pi.getActiveTools();
+		const remaining = currentActive.filter((n) => n !== piName);
+		this.pi.setActiveTools(remaining);
+	}
 }