pi-codemcp 1.2.1 → 1.3.0

package/sidecar/tool_catalog.py

@@ -24,6 +24,7 @@ from .json_types import (
     JsonValue,
 )
 from .models import SearchDetail, ToolSchemaView
+from .sandbox_api import STUB_PRELUDE
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -32,15 +33,6 @@ if TYPE_CHECKING:
 
 SEARCH_SCORE_CUTOFF = 20.0
 MIN_PLURAL_TOKEN_LENGTH = 4
-STUB_IMPORTS = "from typing import Literal, Never, NotRequired, TypeAlias, TypedDict"
-JSON_TYPE_STUBS = (
-    "JsonScalar: TypeAlias = bool | int | float | str | None",
-    'JsonValue: TypeAlias = JsonScalar | list["JsonValue"] | dict[str, "JsonValue"]',
-)
-INSPECT_JSON_STUB = (
-    "def inspect_json(value: JsonValue, *, samples: int = 2, max_depth: int = 3) -> JsonValue: ..."
-)
-STUB_PRELUDE = "\n\n".join([STUB_IMPORTS, *JSON_TYPE_STUBS, INSPECT_JSON_STUB])
 
 
 class ToolSpec(BaseModel):

package/src/chains.ts

@@ -1,3 +1,4 @@
+import { randomUUID } from "node:crypto";
 import { existsSync, readdirSync, readFileSync } from "node:fs";
 import { join } from "node:path";
 import { type ExtensionAPI, highlightCode, type Theme } from "@earendil-works/pi-coding-agent";
@@ -46,17 +47,6 @@ export interface SavedChainView {
   calledBy: string[];
 }
 
-export interface ChainEnabledChange {
-  name: string;
-  scope: ChainScope;
-  enabled: boolean;
-}
-
-export interface ManagerApplyResult {
-  chains: SavedChainView[];
-  status: Record<string, unknown>;
-}
-
 export interface SaveChainInput {
   scope: ChainScope;
   name: string;
@@ -78,6 +68,10 @@ interface LoadedChains {
 
 const CHAIN_NAME = /^[a-z][a-z0-9_]{0,63}$/;
 
+export function newCodeMcpTraceId(source: string): string {
+  return `${source}:${randomUUID()}`;
+}
+
 export class SavedChainManager {
   readonly startupErrors: string[] = [];
   private readonly manifests = new Map<string, ScopedSavedChain>();
@@ -105,7 +99,11 @@ export class SavedChainManager {
     }
   }
 
-  async save(input: SaveChainInput, signal?: AbortSignal): Promise<SavedChainView> {
+  async save(
+    input: SaveChainInput,
+    traceId: string,
+    signal?: AbortSignal,
+  ): Promise<SavedChainView> {
     assertChainName(input.name);
     this.assertToolNameAvailable(input.name);
     const result = await this.lifecycle.request(
@@ -117,6 +115,7 @@ export class SavedChainManager {
         code: input.code,
         input_schema: input.inputSchema,
         output_schema: input.outputSchema,
+        trace_id: traceId,
       },
       signal,
     );
@@ -127,40 +126,59 @@ export class SavedChainManager {
     return view;
   }
 
-  async list(signal?: AbortSignal): Promise<SavedChainView[]> {
-    const result = await this.lifecycle.request("list_chains", {}, signal);
+  async list(traceId: string, signal?: AbortSignal): Promise<SavedChainView[]> {
+    const result = await this.lifecycle.request("list_chains", { trace_id: traceId }, signal);
     const views = parseViewList(result.chains, "list_chains.chains");
     this.synchronizeViews(views);
     return views;
   }
 
-  async applyEnabled(
-    changes: readonly ChainEnabledChange[],
+  async setEnabled(
+    name: string,
+    scope: ChainScope,
+    enabled: boolean,
+    traceId: string,
     signal?: AbortSignal,
-  ): Promise<ManagerApplyResult> {
+  ): Promise<SavedChainView> {
     const result = await this.lifecycle.request(
-      "apply_manager_changes",
-      {
-        changes: changes.map((change) => ({ ...change })),
-      },
+      "set_chain_enabled",
+      { name, scope, enabled, trace_id: traceId },
       signal,
     );
-    const views = parseViewList(result.chains, "apply_manager_changes.chains");
-    const status = requireRecord(result.status, "apply_manager_changes.status");
-    this.synchronizeViews(views);
-    return { chains: views, status };
+    const view = parseSavedChainView(result, "set_chain_enabled");
+    this.upsertView(view);
+    this.refreshNativeTools();
+    return view;
   }
 
-  async revalidate(name: string, scope: ChainScope, signal?: AbortSignal): Promise<SavedChainView> {
-    const result = await this.lifecycle.request("revalidate_chain", { name, scope }, signal);
+  async revalidate(
+    name: string,
+    scope: ChainScope,
+    traceId: string,
+    signal?: AbortSignal,
+  ): Promise<SavedChainView> {
+    const result = await this.lifecycle.request(
+      "revalidate_chain",
+      { name, scope, trace_id: traceId },
+      signal,
+    );
     const view = parseSavedChainView(result, "revalidate_chain");
     this.upsertView(view);
     this.refreshNativeTools();
     return view;
   }
 
-  async delete(name: string, scope: ChainScope, signal?: AbortSignal): Promise<SavedChainView[]> {
-    const result = await this.lifecycle.request("delete_chain", { name, scope }, signal);
+  async delete(
+    name: string,
+    scope: ChainScope,
+    traceId: string,
+    signal?: AbortSignal,
+  ): Promise<SavedChainView[]> {
+    const result = await this.lifecycle.request(
+      "delete_chain",
+      { name, scope, trace_id: traceId },
+      signal,
+    );
     const views = parseViewList(result.chains, "delete_chain.chains");
     this.synchronizeViews(views);
     return views;
@@ -174,7 +192,7 @@ export class SavedChainManager {
       label: chain.name,
       description: chain.description,
       parameters: chain.inputSchema,
-      async execute(_toolCallId, params, signal, onUpdate) {
+      async execute(toolCallId, params, signal, onUpdate) {
         onUpdate?.({
           content: [{ type: "text", text: `Running saved MCP chain ${chain.name}...` }],
           details: undefined,
@@ -182,13 +200,24 @@ export class SavedChainManager {
         const arguments_ = requireRecord(params, `${chain.name} arguments`);
         const result = await manager.lifecycle.request(
           "execute_chain",
-          { name: chain.name, arguments: arguments_ },
+          { name: chain.name, arguments: arguments_, trace_id: toolCallId },
           signal,
         );
         if (result.ok !== true) {
-          const error =
-            typeof result.error === "string" ? result.error : `Saved chain ${chain.name} failed`;
-          throw new Error(error);
+          throw new Error(
+            JSON.stringify({
+              failure_stage: result.failure_stage,
+              error:
+                typeof result.error === "string"
+                  ? result.error
+                  : `Saved chain ${chain.name} failed`,
+              failure: result.failure,
+              result_ref: result.result_ref,
+              expires_in_seconds: result.expires_in_seconds,
+              calls_made: result.calls_made,
+              chain_calls: result.chain_calls,
+            }),
+          );
         }
         const settings = manager.lifecycle.loadSettings();
         const output = formatCodeMcpOutput(result.result, {

package/src/config.ts

@@ -5,54 +5,23 @@ import {
   writeJsonObjectAtomically,
 } from "./json-file.js";
 
-export interface McpServerEnabledChange {
-  name: string;
-  enabled: boolean;
-}
-
 export function setMcpServerEnabled(configPath: string, name: string, enabled: boolean): void {
-  setMcpServersEnabled(configPath, [{ name, enabled }]);
-}
-
-export function setMcpServersEnabled(
-  configPath: string,
-  changes: readonly McpServerEnabledChange[],
-): void {
-  if (changes.length === 0) return;
-  const duplicate = duplicateName(changes.map((change) => change.name));
-  if (duplicate) throw new Error(`Duplicate MCP server change: ${JSON.stringify(duplicate)}`);
-
   const root = readJsonObject(configPath, "mcp.json root");
   const hasServerBlock = Object.hasOwn(root, "mcpServers");
   const servers = hasServerBlock ? requireJsonObject(root.mcpServers, "mcp.json mcpServers") : root;
-  const updatedServers: JsonRecord = { ...servers };
+  const server = requireJsonObject(servers[name], `MCP server ${JSON.stringify(name)}`);
+  const updatedServer: JsonRecord = { ...server };
 
-  for (const change of changes) {
-    const server = requireJsonObject(
-      servers[change.name],
-      `MCP server ${JSON.stringify(change.name)}`,
-    );
-    const updatedServer: JsonRecord = { ...server };
-    if (typeof server.enabled === "boolean") {
-      updatedServer.enabled = change.enabled;
-      delete updatedServer.disabled;
-    } else {
-      updatedServer.disabled = !change.enabled;
-    }
-    updatedServers[change.name] = updatedServer;
+  if (typeof server.enabled === "boolean") {
+    updatedServer.enabled = enabled;
+    delete updatedServer.disabled;
+  } else {
+    updatedServer.disabled = !enabled;
   }
 
+  const updatedServers: JsonRecord = { ...servers, [name]: updatedServer };
   const updatedRoot: JsonRecord = hasServerBlock
     ? { ...root, mcpServers: updatedServers }
     : updatedServers;
   writeJsonObjectAtomically(configPath, updatedRoot);
 }
-
-function duplicateName(names: readonly string[]): string | undefined {
-  const seen = new Set<string>();
-  for (const name of names) {
-    if (seen.has(name)) return name;
-    seen.add(name);
-  }
-  return undefined;
-}

package/src/mcp-client.ts

@@ -22,11 +22,11 @@ export type SidecarToolName =
   | "inspect"
   | "discover"
   | "reload_settings"
-  | "apply_manager_changes"
   | "execute"
   | "save_chain"
   | "list_chains"
   | "execute_chain"
+  | "set_chain_enabled"
   | "revalidate_chain"
   | "delete_chain"
   | "stats"
@@ -99,6 +99,7 @@ export class SidecarClient {
       ...definedProcessEnvironment(),
       ...(options.environment ?? {}),
       PI_CODEMCP_AGENT_DIR: this.agentDir,
+      PI_CODEMCP_PACKAGE_VERSION: this.packageVersion,
       ...(this.projectChainsDirectory === undefined
         ? {}
         : { PI_CODEMCP_PROJECT_CHAINS_DIR: this.projectChainsDirectory }),