diff --git a/pi/files.linux/agent/settings.json b/pi/files.linux/agent/settings.json index 8f24fad..4947790 100644 --- a/pi/files.linux/agent/settings.json +++ b/pi/files.linux/agent/settings.json @@ -1,7 +1,7 @@ { "lastChangelogVersion": "0.67.3", - "defaultProvider": "cursor-acp", - "defaultModel": "auto", + "defaultProvider": "openai-codex", + "defaultModel": "gpt-5.3-codex", "defaultThinkingLevel": "medium", "theme": "matugen", "lsp": { diff --git a/pi/files/agent/extensions/package.json b/pi/files/agent/extensions/package.json index 0143cc8..85b56bb 100644 --- a/pi/files/agent/extensions/package.json +++ b/pi/files/agent/extensions/package.json @@ -4,6 +4,7 @@ "type": "module", "dependencies": { "@anthropic-ai/sdk": "^0.52.0", + "@mariozechner/jiti": "^2.6.5", "@mozilla/readability": "^0.5.0", "@sinclair/typebox": "^0.34.0", "linkedom": "^0.16.0", diff --git a/pi/files/agent/extensions/pi-claude-code-use/LICENSE b/pi/files/agent/extensions/pi-claude-code-use/LICENSE new file mode 100644 index 0000000..0308ed5 --- /dev/null +++ b/pi/files/agent/extensions/pi-claude-code-use/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Ben Vargas + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/pi/files/agent/extensions/pi-claude-code-use/README.md b/pi/files/agent/extensions/pi-claude-code-use/README.md new file mode 100644 index 0000000..05449e5 --- /dev/null +++ b/pi/files/agent/extensions/pi-claude-code-use/README.md @@ -0,0 +1,115 @@ +# @benvargas/pi-claude-code-use + +`pi-claude-code-use` keeps Pi's built-in `anthropic` provider intact and applies the smallest payload changes needed for Anthropic OAuth subscription use in Pi. + +It does not register a new provider or replace Pi's Anthropic request transport. Pi core remains in charge of OAuth transport, headers, model definitions, and streaming. + +## What It Changes + +When Pi is using Anthropic OAuth, this extension intercepts outbound API requests via the `before_provider_request` hook and: + +- **System prompt rewrite** -- rewrites a small set of Pi-identifying prompt phrases in system prompt text: + - `pi itself` → `the cli itself` + - `pi .md files` → `cli .md files` + - `pi packages` → `cli packages` + Preserves Pi's original `system[]` structure, `cache_control` metadata, and non-text blocks. +- **Tool filtering** -- passes through core Claude Code tools, Anthropic-native typed tools (e.g. `web_search`), and any tool prefixed with `mcp__`. Unknown flat-named tools are filtered out. +- **Companion tool remapping** -- renames known companion extension tools from their flat names to MCP-style aliases (e.g. `web_search_exa` becomes `mcp__exa__web_search`). Duplicate flat entries are removed after remapping. +- **tool_choice remapping** -- if `tool_choice` references a flat companion name that was remapped, the reference is updated to the MCP alias. If it references a tool that was filtered out, `tool_choice` is removed from the payload. +- **Message history rewriting** -- `tool_use` blocks in conversation history that reference flat companion names are rewritten to their MCP aliases so the model sees consistent tool names across the conversation. +- **Companion alias registration** -- at session start and before each agent turn, discovers loaded companion extensions, captures their tool definitions via a jiti-based shim, and registers MCP-alias copies so the model can invoke them under Claude Code-compatible names. +- **Alias activation tracking** -- auto-activates MCP aliases when their flat counterpart is active under Anthropic OAuth. Tracks provenance (auto-managed vs user-selected) so that disabling OAuth only removes auto-activated aliases, preserving any the user explicitly enabled. + +Non-OAuth Anthropic usage and non-Anthropic providers are left completely unchanged. + +## Install + +```bash +pi install npm:@benvargas/pi-claude-code-use +``` + +Or load it directly without installing: + +```bash +pi -e /path/to/pi-packages/packages/pi-claude-code-use/extensions/index.ts +``` + +## Usage + +Install the package and continue using the normal `anthropic` provider with Anthropic OAuth login: + +```bash +/login anthropic +/model anthropic/claude-opus-4-6 +``` + +No extra configuration is required. + +## Environment Variables + +| Variable | Description | +|---|---| +| `PI_CLAUDE_CODE_USE_DEBUG_LOG` | Set to a file path to enable debug logging. Writes two JSON entries per Anthropic OAuth request: one with `"stage": "before"` (the original payload from Pi) and one with `"stage": "after"` (the transformed payload sent to Anthropic). | +| `PI_CLAUDE_CODE_USE_DISABLE_TOOL_FILTER` | Set to `1` to disable tool filtering. System prompt rewriting still applies, but all tools pass through unchanged. Useful for debugging whether a tool-filtering issue is causing a problem. | + +Example: + +```bash +PI_CLAUDE_CODE_USE_DEBUG_LOG=/tmp/pi-claude-debug.log pi -e /path/to/extensions/index.ts --model anthropic/claude-sonnet-4-20250514 +``` + +## Companion Tool Aliases + +When these companion extensions from this monorepo are loaded alongside `pi-claude-code-use`, MCP aliases are automatically registered and remapped: + +| Flat name | MCP alias | +|---|---| +| `web_search_exa` | `mcp__exa__web_search` | +| `get_code_context_exa` | `mcp__exa__get_code_context` | +| `firecrawl_scrape` | `mcp__firecrawl__scrape` | +| `firecrawl_map` | `mcp__firecrawl__map` | +| `firecrawl_search` | `mcp__firecrawl__search` | +| `generate_image` | `mcp__antigravity__generate_image` | +| `image_quota` | `mcp__antigravity__image_quota` | + +### How companion discovery works + +The extension identifies companion tools by matching `sourceInfo` metadata that Pi attaches to each registered tool: + +1. **baseDir match** -- if the tool's `sourceInfo.baseDir` directory name matches the companion's directory name (e.g. `pi-exa-mcp`). +2. **Path match** -- if the tool's `sourceInfo.path` contains the companion's scoped package name (e.g. `@benvargas/pi-exa-mcp`) or directory name as a path segment. This handles npm installs, git clones, and monorepo layouts where `baseDir` points to the repo root rather than the individual package. + +Once a companion tool is identified, its extension factory is loaded via jiti into a capture shim to obtain the full tool definition, which is then re-registered under the MCP alias name. + +## Core Tools Allowlist + +The following tool names always pass through filtering (case-insensitive). This list mirrors Pi core's `claudeCodeTools` in `packages/ai/src/providers/anthropic.ts`: + +`Read`, `Write`, `Edit`, `Bash`, `Grep`, `Glob`, `AskUserQuestion`, `EnterPlanMode`, `ExitPlanMode`, `KillShell`, `NotebookEdit`, `Skill`, `Task`, `TaskOutput`, `TodoWrite`, `WebFetch`, `WebSearch` + +Additionally, any tool with a `type` field (Anthropic-native tools like `web_search`) and any tool prefixed with `mcp__` always passes through. + +## Guidance For Extension Authors + +Anthropic's OAuth subscription path appears to fingerprint tool names. Flat extension tool names such as `web_search_exa` were rejected in live testing, while MCP-style names such as `mcp__exa__web_search` were accepted. + +If you want a custom tool to survive Anthropic OAuth filtering cleanly, prefer registering it directly under an MCP-style name: + +```text +mcp____ +``` + +Examples: + +- `mcp__exa__web_search` +- `mcp__firecrawl__scrape` +- `mcp__mytools__lookup_customer` + +If an extension keeps a flat legacy name for non-Anthropic use, it can also register an MCP-style alias alongside it. `pi-claude-code-use` already does this centrally for the known companion tools in this repo, but unknown non-MCP tool names will still be filtered out on Anthropic OAuth requests. + +## Notes + +- The extension activates for all Anthropic OAuth requests regardless of model, rather than using a fixed model allowlist. +- Non-OAuth Anthropic usage (API key auth) is left unchanged. +- In practice, unknown non-MCP extension tools were the remaining trigger for Anthropic's extra-usage classification, so this package keeps core tools, keeps MCP-style tools, auto-aliases the known companion tools above, and filters the rest. +- Pi may show its built-in OAuth subscription warning banner even when the request path works correctly. That banner is UI logic in Pi, not a signal that the upstream request is being billed as extra usage. diff --git a/pi/files/agent/extensions/pi-claude-code-use/index.ts b/pi/files/agent/extensions/pi-claude-code-use/index.ts new file mode 100644 index 0000000..367b603 --- /dev/null +++ b/pi/files/agent/extensions/pi-claude-code-use/index.ts @@ -0,0 +1,641 @@ +import { appendFileSync } from "node:fs"; +import { basename, dirname } from "node:path"; +import { createJiti } from "@mariozechner/jiti"; +import * as piAiModule from "@mariozechner/pi-ai"; +import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; +import * as piCodingAgentModule from "@mariozechner/pi-coding-agent"; +import * as typeboxModule from "@sinclair/typebox"; + +// ============================================================================ +// Types +// ============================================================================ + +interface CompanionSpec { + dirName: string; + packageName: string; + aliases: ReadonlyArray; +} + +type ToolRegistration = Parameters[0]; +type ToolInfo = ReturnType[number]; + +// ============================================================================ +// Constants +// ============================================================================ + +/** + * Core Claude Code tool names that always pass through Anthropic OAuth filtering. + * Stored lowercase for case-insensitive matching. + * Mirrors Pi core's claudeCodeTools list in packages/ai/src/providers/anthropic.ts + */ +const CORE_TOOL_NAMES = new Set([ + "read", + "write", + "edit", + "bash", + "grep", + "glob", + "askuserquestion", + "enterplanmode", + "exitplanmode", + "killshell", + "notebookedit", + "skill", + "task", + "taskoutput", + "todowrite", + "webfetch", + "websearch", +]); + +/** Flat companion tool name → MCP-style alias. */ +const FLAT_TO_MCP = new Map([ + ["web_search_exa", "mcp__exa__web_search"], + ["get_code_context_exa", "mcp__exa__get_code_context"], + ["firecrawl_scrape", "mcp__firecrawl__scrape"], + ["firecrawl_map", "mcp__firecrawl__map"], + ["firecrawl_search", "mcp__firecrawl__search"], + ["generate_image", "mcp__antigravity__generate_image"], + ["image_quota", "mcp__antigravity__image_quota"], +]); + +/** Known companion extensions and the tools they provide. */ +const COMPANIONS: CompanionSpec[] = [ + { + dirName: "pi-exa-mcp", + packageName: "@benvargas/pi-exa-mcp", + aliases: [ + ["web_search_exa", "mcp__exa__web_search"], + ["get_code_context_exa", "mcp__exa__get_code_context"], + ], + }, + { + dirName: "pi-firecrawl", + packageName: "@benvargas/pi-firecrawl", + aliases: [ + ["firecrawl_scrape", "mcp__firecrawl__scrape"], + ["firecrawl_map", "mcp__firecrawl__map"], + ["firecrawl_search", "mcp__firecrawl__search"], + ], + }, + { + dirName: "pi-antigravity-image-gen", + packageName: "@benvargas/pi-antigravity-image-gen", + aliases: [ + ["generate_image", "mcp__antigravity__generate_image"], + ["image_quota", "mcp__antigravity__image_quota"], + ], + }, +]; + +/** Reverse lookup: flat tool name → its companion spec. */ +const TOOL_TO_COMPANION = new Map( + COMPANIONS.flatMap((spec) => spec.aliases.map(([flat]) => [flat, spec] as const)), +); + +// ============================================================================ +// Helpers +// ============================================================================ + +function isPlainObject(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function lower(name: string | undefined): string { + return (name ?? "").trim().toLowerCase(); +} + +// ============================================================================ +// System prompt rewrite (PRD §1.1) +// +// Replace "pi itself" → "the cli itself" in system prompt text. +// Preserves cache_control, non-text blocks, and payload shape. +// ============================================================================ + +function rewritePromptText(text: string): string { + return text + .replaceAll("pi itself", "the cli itself") + .replaceAll("pi .md files", "cli .md files") + .replaceAll("pi packages", "cli packages"); +} + +function rewriteSystemField(system: unknown): unknown { + if (typeof system === "string") { + return rewritePromptText(system); + } + if (!Array.isArray(system)) { + return system; + } + return system.map((block) => { + if (!isPlainObject(block) || block.type !== "text" || typeof block.text !== "string") { + return block; + } + const rewritten = rewritePromptText(block.text); + return rewritten === block.text ? block : { ...block, text: rewritten }; + }); +} + +// ============================================================================ +// Tool filtering and MCP alias remapping (PRD §1.2) +// +// Rules applied per tool: +// 1. Anthropic-native typed tools (have a `type` field) → pass through +// 2. Core Claude Code tool names → pass through +// 3. Tools already prefixed with mcp__ → pass through +// 4. Known companion tools whose MCP alias is also advertised → rename to alias +// 5. Known companion tools without an advertised alias → filtered out +// 6. Unknown flat-named tools → filtered out (unless disableFilter) +// ============================================================================ + +function collectToolNames(tools: unknown[]): Set { + const names = new Set(); + for (const tool of tools) { + if (isPlainObject(tool) && typeof tool.name === "string") { + names.add(lower(tool.name)); + } + } + return names; +} + +function filterAndRemapTools(tools: unknown[] | undefined, disableFilter: boolean): unknown[] | undefined { + if (!Array.isArray(tools)) return tools; + + const advertised = collectToolNames(tools); + const emitted = new Set(); + const result: unknown[] = []; + + for (const tool of tools) { + if (!isPlainObject(tool)) continue; + + // Rule 1: native typed tools always pass through + if (typeof tool.type === "string" && tool.type.trim().length > 0) { + result.push(tool); + continue; + } + + const name = typeof tool.name === "string" ? tool.name : ""; + if (!name) continue; + const nameLc = lower(name); + + // Rules 2 & 3: core tools and mcp__-prefixed pass through (with dedup) + if (CORE_TOOL_NAMES.has(nameLc) || nameLc.startsWith("mcp__")) { + if (!emitted.has(nameLc)) { + emitted.add(nameLc); + result.push(tool); + } + continue; + } + + // Rules 4 & 5: known companion tool + const mcpAlias = FLAT_TO_MCP.get(nameLc); + if (mcpAlias) { + const aliasLc = lower(mcpAlias); + if (advertised.has(aliasLc) && !emitted.has(aliasLc)) { + // Alias exists in tool list → rename flat to alias, dedup + emitted.add(aliasLc); + result.push({ ...tool, name: mcpAlias }); + } else if (disableFilter && !emitted.has(nameLc)) { + // Filter disabled: keep flat name if not yet emitted + emitted.add(nameLc); + result.push(tool); + } + continue; + } + + // Rule 6: unknown flat-named tool + if (disableFilter && !emitted.has(nameLc)) { + emitted.add(nameLc); + result.push(tool); + } + } + + return result; +} + +function remapToolChoice( + toolChoice: Record, + survivingNames: Map, +): Record | undefined { + if (toolChoice.type !== "tool" || typeof toolChoice.name !== "string") { + return toolChoice; + } + + const nameLc = lower(toolChoice.name); + const actualName = survivingNames.get(nameLc); + if (actualName) { + return actualName === toolChoice.name ? toolChoice : { ...toolChoice, name: actualName }; + } + + const mcpAlias = FLAT_TO_MCP.get(nameLc); + if (mcpAlias && survivingNames.has(lower(mcpAlias))) { + return { ...toolChoice, name: mcpAlias }; + } + + return undefined; +} + +function remapMessageToolNames(messages: unknown[], survivingNames: Map): unknown[] { + let anyChanged = false; + const result = messages.map((msg) => { + if (!isPlainObject(msg) || !Array.isArray(msg.content)) return msg; + + let msgChanged = false; + const content = (msg.content as unknown[]).map((block) => { + if (!isPlainObject(block) || block.type !== "tool_use" || typeof block.name !== "string") { + return block; + } + const mcpAlias = FLAT_TO_MCP.get(lower(block.name)); + if (mcpAlias && survivingNames.has(lower(mcpAlias))) { + msgChanged = true; + return { ...block, name: mcpAlias }; + } + return block; + }); + + if (msgChanged) { + anyChanged = true; + return { ...msg, content }; + } + return msg; + }); + + return anyChanged ? result : messages; +} + +// ============================================================================ +// Full payload transform +// ============================================================================ + +function transformPayload(raw: Record, disableFilter: boolean): Record { + // Deep clone to avoid mutating the original + const payload = JSON.parse(JSON.stringify(raw)) as Record; + + // 1. System prompt rewrite (always applies) + if (payload.system !== undefined) { + payload.system = rewriteSystemField(payload.system); + } + + // When escape hatch is active, skip all tool filtering/remapping + if (disableFilter) { + return payload; + } + + // 2. Tool filtering and alias remapping + payload.tools = filterAndRemapTools(payload.tools as unknown[] | undefined, false); + + // 3. Build map of tool names that survived filtering (lowercase → actual name) + const survivingNames = new Map(); + if (Array.isArray(payload.tools)) { + for (const tool of payload.tools) { + if (isPlainObject(tool) && typeof tool.name === "string") { + survivingNames.set(lower(tool.name), tool.name as string); + } + } + } + + // 4. Remap tool_choice if it references a renamed or filtered tool + if (isPlainObject(payload.tool_choice)) { + const remapped = remapToolChoice(payload.tool_choice, survivingNames); + if (remapped === undefined) { + delete payload.tool_choice; + } else { + payload.tool_choice = remapped; + } + } + + // 5. Rewrite historical tool_use blocks in message history + if (Array.isArray(payload.messages)) { + payload.messages = remapMessageToolNames(payload.messages, survivingNames); + } + + return payload; +} + +// ============================================================================ +// Debug logging (PRD §1.4) +// ============================================================================ + +const debugLogPath = process.env.PI_CLAUDE_CODE_USE_DEBUG_LOG; + +function writeDebugLog(payload: unknown): void { + if (!debugLogPath) return; + try { + appendFileSync(debugLogPath, `${new Date().toISOString()}\n${JSON.stringify(payload, null, 2)}\n---\n`, "utf-8"); + } catch { + // Debug logging must never break actual requests + } +} + +// ============================================================================ +// Companion alias registration (PRD §1.3) +// +// Discovers loaded companion extensions, captures their tool definitions via +// a shim ExtensionAPI, and registers MCP-alias versions so the model can +// invoke them under Claude Code-compatible names. +// ============================================================================ + +const registeredMcpAliases = new Set(); +const autoActivatedAliases = new Set(); +let lastManagedToolList: string[] | undefined; + +const captureCache = new Map>>(); +let jitiLoader: { import(path: string, opts?: { default?: boolean }): Promise } | undefined; + +function getJitiLoader() { + if (!jitiLoader) { + jitiLoader = createJiti(import.meta.url, { + moduleCache: false, + tryNative: false, + virtualModules: { + "@mariozechner/pi-ai": piAiModule, + "@mariozechner/pi-coding-agent": piCodingAgentModule, + "@sinclair/typebox": typeboxModule, + }, + }); + } + return jitiLoader; +} + +async function loadFactory(baseDir: string): Promise<((pi: ExtensionAPI) => void | Promise) | undefined> { + const dir = baseDir.replace(/\/$/, ""); + const candidates = [`${dir}/index.ts`, `${dir}/index.js`, `${dir}/extensions/index.ts`, `${dir}/extensions/index.js`]; + + const loader = getJitiLoader(); + for (const path of candidates) { + try { + const mod = await loader.import(path, { default: true }); + if (typeof mod === "function") return mod as (pi: ExtensionAPI) => void | Promise; + } catch { + // Try next candidate + } + } + return undefined; +} + +function isCompanionSource(tool: ToolInfo | undefined, spec: CompanionSpec): boolean { + if (!tool?.sourceInfo) return false; + + const baseDir = tool.sourceInfo.baseDir; + if (baseDir) { + const dirName = basename(baseDir); + if (dirName === spec.dirName) return true; + if (dirName === "extensions" && basename(dirname(baseDir)) === spec.dirName) return true; + } + + const fullPath = tool.sourceInfo.path; + if (typeof fullPath !== "string") return false; + // Normalize backslashes for Windows paths before segment-bounded check + const normalized = fullPath.replaceAll("\\", "/"); + // Check for scoped package name (npm install) or directory name (git/monorepo) + return normalized.includes(`/${spec.packageName}/`) || normalized.includes(`/${spec.dirName}/`); +} + +function buildCaptureShim(realPi: ExtensionAPI, captured: Map): ExtensionAPI { + const shimFlags = new Set(); + return { + registerTool(def) { + captured.set(def.name, def as unknown as ToolRegistration); + }, + registerFlag(name, _options) { + shimFlags.add(name); + }, + getFlag(name) { + return shimFlags.has(name) ? realPi.getFlag(name) : undefined; + }, + on() {}, + registerCommand() {}, + registerShortcut() {}, + registerMessageRenderer() {}, + registerProvider() {}, + unregisterProvider() {}, + sendMessage() {}, + sendUserMessage() {}, + appendEntry() {}, + setSessionName() {}, + getSessionName() { + return undefined; + }, + setLabel() {}, + exec(command, args, options) { + return realPi.exec(command, args, options); + }, + getActiveTools() { + return realPi.getActiveTools(); + }, + getAllTools() { + return realPi.getAllTools(); + }, + setActiveTools(names) { + realPi.setActiveTools(names); + }, + getCommands() { + return realPi.getCommands(); + }, + setModel(model) { + return realPi.setModel(model); + }, + getThinkingLevel() { + return realPi.getThinkingLevel(); + }, + setThinkingLevel(level) { + realPi.setThinkingLevel(level); + }, + events: realPi.events, + } as ExtensionAPI; +} + +async function captureCompanionTools(baseDir: string, realPi: ExtensionAPI): Promise> { + let pending = captureCache.get(baseDir); + if (!pending) { + pending = (async () => { + const factory = await loadFactory(baseDir); + if (!factory) return new Map(); + const tools = new Map(); + await factory(buildCaptureShim(realPi, tools)); + return tools; + })(); + captureCache.set(baseDir, pending); + } + return pending; +} + +async function registerAliasesForLoadedCompanions(pi: ExtensionAPI): Promise { + // Clear capture cache so flag/config changes since last call take effect + captureCache.clear(); + + const allTools = pi.getAllTools(); + const toolIndex = new Map(); + const knownNames = new Set(); + for (const tool of allTools) { + toolIndex.set(lower(tool.name), tool); + knownNames.add(lower(tool.name)); + } + + for (const spec of COMPANIONS) { + for (const [flatName, mcpName] of spec.aliases) { + if (registeredMcpAliases.has(mcpName) || knownNames.has(lower(mcpName))) continue; + + const tool = toolIndex.get(lower(flatName)); + if (!tool || !isCompanionSource(tool, spec)) continue; + + // Prefer the extension file's directory for loading (sourceInfo.path is the actual + // entry point). Fall back to baseDir only if path is unavailable. baseDir can be + // the monorepo root which doesn't contain the extension entry point directly. + const loadDir = tool.sourceInfo?.path ? dirname(tool.sourceInfo.path) : tool.sourceInfo?.baseDir; + if (!loadDir) continue; + + const captured = await captureCompanionTools(loadDir, pi); + const def = captured.get(flatName); + if (!def) continue; + + pi.registerTool({ + ...def, + name: mcpName, + label: def.label?.startsWith("MCP ") ? def.label : `MCP ${def.label ?? mcpName}`, + }); + registeredMcpAliases.add(mcpName); + knownNames.add(lower(mcpName)); + } + } +} + +/** + * Synchronize MCP alias tool activation with the current model state. + * When OAuth is active, auto-activate aliases for any active companion tools. + * When OAuth is inactive, remove auto-activated aliases (but preserve user-selected ones). + */ +function syncAliasActivation(pi: ExtensionAPI, enableAliases: boolean): void { + const activeNames = pi.getActiveTools(); + const allNames = new Set(pi.getAllTools().map((t) => t.name)); + + if (enableAliases) { + // Determine which aliases should be active based on their flat counterpart being active + const activeLc = new Set(activeNames.map(lower)); + const desiredAliases: string[] = []; + for (const [flat, mcp] of FLAT_TO_MCP) { + if (activeLc.has(flat) && allNames.has(mcp) && registeredMcpAliases.has(mcp)) { + desiredAliases.push(mcp); + } + } + const desiredSet = new Set(desiredAliases); + + // Promote auto-activated aliases to user-selected when the user explicitly kept + // the alias while removing its flat counterpart from the tool picker. + // We detect this by checking: (a) user changed the tool list since our last sync, + // (b) the flat tool was previously managed but is no longer active, and + // (c) the alias is still active. This means the user deliberately kept the alias. + if (lastManagedToolList !== undefined) { + const activeSet = new Set(activeNames); + const lastManaged = new Set(lastManagedToolList); + for (const alias of autoActivatedAliases) { + if (!activeSet.has(alias) || desiredSet.has(alias)) continue; + // Find the flat name for this alias + const flatName = [...FLAT_TO_MCP.entries()].find(([, mcp]) => mcp === alias)?.[0]; + if (flatName && lastManaged.has(flatName) && !activeSet.has(flatName)) { + // User removed the flat tool but kept the alias → promote to user-selected + autoActivatedAliases.delete(alias); + } + } + } + + // Find registered aliases currently in the active list + const activeRegistered = activeNames.filter((n) => registeredMcpAliases.has(n) && allNames.has(n)); + + // Per-alias provenance: an alias is "user-selected" if it's active and was NOT + // auto-activated by us. Only preserve those; auto-activated aliases get re-derived + // from the desired set each sync. + const preserved = activeRegistered.filter((n) => !autoActivatedAliases.has(n)); + + // Build result: non-alias tools + preserved user aliases + desired aliases + const nonAlias = activeNames.filter((n) => !registeredMcpAliases.has(n)); + const next = Array.from(new Set([...nonAlias, ...preserved, ...desiredAliases])); + + // Update auto-activation tracking: aliases we added this sync that weren't user-preserved + const preservedSet = new Set(preserved); + autoActivatedAliases.clear(); + for (const name of desiredAliases) { + if (!preservedSet.has(name)) { + autoActivatedAliases.add(name); + } + } + + if (next.length !== activeNames.length || next.some((n, i) => n !== activeNames[i])) { + pi.setActiveTools(next); + lastManagedToolList = [...next]; + } + } else { + // Remove only auto-activated aliases; user-selected ones are preserved + const next = activeNames.filter((n) => !autoActivatedAliases.has(n)); + autoActivatedAliases.clear(); + + if (next.length !== activeNames.length || next.some((n, i) => n !== activeNames[i])) { + pi.setActiveTools(next); + lastManagedToolList = [...next]; + } else { + lastManagedToolList = undefined; + } + } +} + +// ============================================================================ +// Extension entry point +// ============================================================================ + +export default async function piClaudeCodeUse(pi: ExtensionAPI): Promise { + pi.on("session_start", async () => { + await registerAliasesForLoadedCompanions(pi); + }); + + pi.on("before_agent_start", async (_event, ctx) => { + await registerAliasesForLoadedCompanions(pi); + const model = ctx.model; + const isOAuth = model?.provider === "anthropic" && ctx.modelRegistry.isUsingOAuth(model); + syncAliasActivation(pi, isOAuth); + }); + + pi.on("before_provider_request", (event, ctx) => { + const model = ctx.model; + if (!model || model.provider !== "anthropic" || !ctx.modelRegistry.isUsingOAuth(model)) { + return undefined; + } + if (!isPlainObject(event.payload)) { + return undefined; + } + + writeDebugLog({ stage: "before", payload: event.payload }); + const disableFilter = process.env.PI_CLAUDE_CODE_USE_DISABLE_TOOL_FILTER === "1"; + const transformed = transformPayload(event.payload as Record, disableFilter); + writeDebugLog({ stage: "after", payload: transformed }); + return transformed; + }); +} + +// ============================================================================ +// Test exports +// ============================================================================ + +export const _test = { + CORE_TOOL_NAMES, + FLAT_TO_MCP, + COMPANIONS, + TOOL_TO_COMPANION, + autoActivatedAliases, + buildCaptureShim, + collectToolNames, + filterAndRemapTools, + getLastManagedToolList: () => lastManagedToolList, + isCompanionSource, + isPlainObject, + lower, + registerAliasesForLoadedCompanions, + registeredMcpAliases, + remapMessageToolNames, + remapToolChoice, + rewritePromptText, + rewriteSystemField, + setLastManagedToolList: (v: string[] | undefined) => { + lastManagedToolList = v; + }, + syncAliasActivation, + transformPayload, +}; diff --git a/pi/files/agent/extensions/pnpm-lock.yaml b/pi/files/agent/extensions/pnpm-lock.yaml index 8274043..af4def9 100644 --- a/pi/files/agent/extensions/pnpm-lock.yaml +++ b/pi/files/agent/extensions/pnpm-lock.yaml @@ -11,6 +11,9 @@ importers: '@anthropic-ai/sdk': specifier: ^0.52.0 version: 0.52.0 + '@mariozechner/jiti': + specifier: ^2.6.5 + version: 2.6.5 '@mozilla/readability': specifier: ^0.5.0 version: 0.5.0