pi-codemcp 1.2.2 → 1.3.1

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):
@@ -1055,42 +1047,67 @@ def _resolve_ref(ref: str, root: JsonSchema) -> JsonObject:
 def _merge_all_of(
     schema: JsonObject,
     root: JsonSchema,
-) -> JsonObject | None:
-    merged_properties: JsonObject = {}
-    merged_required: list[JsonValue] = []
-    additional_properties: JsonValue = True
-    has_additional_properties = False
+) -> JsonSchema | None:
     raw_members = schema.get("allOf")
     if not isinstance(raw_members, list):
         return None
+    # Sibling constraints (type, properties, ...) apply on top of allOf members.
+    parts: list[JsonObject] = [{key: value for key, value in schema.items() if key != "allOf"}]
     for member in raw_members:
         if not isinstance(member, dict):
             return None
         ref = member.get("$ref")
-        resolved = _resolve_ref(ref, root) if isinstance(ref, str) else member
-        if resolved.get("type") not in {None, "object"}:
-            return None
-        raw_properties = resolved.get("properties")
+        parts.append(_resolve_ref(ref, root) if isinstance(ref, str) else member)
+
+    declared_types: set[str] = set()
+    for part in parts:
+        part_type = part.get("type")
+        if isinstance(part_type, str):
+            declared_types.add(part_type)
+    if len(declared_types) > 1:
+        return None
+    merged_type = declared_types.pop() if declared_types else None
+
+    if merged_type == "object" or (
+        merged_type is None and any("properties" in part for part in parts)
+    ):
+        return _merge_object_parts(parts)
+
+    if merged_type is None:
+        # Annotation-only allOf (descriptions, patterns, ...): no type information.
+        return True
+    merged: JsonObject = {"type": merged_type}
+    for part in parts:
+        if "items" in part and "items" not in merged:
+            merged["items"] = part["items"]
+    return merged
+
+
+def _merge_object_parts(parts: list[JsonObject]) -> JsonObject:
+    merged_properties: JsonObject = {}
+    merged_required: list[JsonValue] = []
+    additional_properties: JsonValue = True
+    has_additional_properties = False
+    for part in parts:
+        raw_properties = part.get("properties")
         if isinstance(raw_properties, dict):
             merged_properties.update(raw_properties)
-        raw_required = resolved.get("required")
+        raw_required = part.get("required")
         if isinstance(raw_required, list):
             for required in raw_required:
                 if isinstance(required, str) and required not in merged_required:
                     merged_required.append(required)
-        if "additionalProperties" in resolved:
-            additional_properties = JSON_VALUE_ADAPTER.validate_python(
-                resolved["additionalProperties"]
-            )
+        if "additionalProperties" in part:
+            additional_properties = JSON_VALUE_ADAPTER.validate_python(part["additionalProperties"])
             has_additional_properties = True
-    merged: JsonObject = {
+    merged_object: JsonObject = {
         "type": "object",
         "properties": merged_properties,
         "required": merged_required,
     }
     if has_additional_properties:
-        merged["additionalProperties"] = additional_properties
-    return merged
+        merged_object["additionalProperties"] = additional_properties
+    return merged_object
 
 
 def _dedupe(blocks: Iterable[str]) -> list[str]:

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";
@@ -67,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>();
@@ -94,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(
@@ -106,6 +115,7 @@ export class SavedChainManager {
         code: input.code,
         input_schema: input.inputSchema,
         output_schema: input.outputSchema,
+        trace_id: traceId,
       },
       signal,
     );
@@ -116,8 +126,8 @@ 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;
@@ -127,11 +137,12 @@ export class SavedChainManager {
     name: string,
     scope: ChainScope,
     enabled: boolean,
+    traceId: string,
     signal?: AbortSignal,
   ): Promise<SavedChainView> {
     const result = await this.lifecycle.request(
       "set_chain_enabled",
-      { name, scope, enabled },
+      { name, scope, enabled, trace_id: traceId },
       signal,
     );
     const view = parseSavedChainView(result, "set_chain_enabled");
@@ -140,16 +151,34 @@ export class SavedChainManager {
     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;
@@ -163,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,
@@ -171,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/mcp-client.ts

@@ -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 }),

package/src/modal.ts

@@ -72,6 +72,7 @@ export interface StatsModalState {
     p95Ms: number;
     maxMs: number;
   }>;
+  outcomes: Array<{ name: string; count: number }>;
   failures: Array<{ stage: string; count: number }>;
   upstreamOutputBytes: number;
   cacheHits: number;
@@ -258,6 +259,10 @@ export function statsStateFromSnapshot(snapshot: Record<string, unknown>): Stats
       },
     ];
   });
+  const rawOutcomes = isRecord(snapshot.outcomes) ? snapshot.outcomes : {};
+  const outcomes = Object.entries(rawOutcomes)
+    .flatMap(([name, count]) => (typeof count === "number" && count >= 0 ? [{ name, count }] : []))
+    .sort((left, right) => right.count - left.count || left.name.localeCompare(right.name));
   const rawFailures = isRecord(snapshot.failures) ? snapshot.failures : {};
   const failures = Object.entries(rawFailures)
     .flatMap(([stage, count]) =>
@@ -278,6 +283,7 @@ export function statsStateFromSnapshot(snapshot: Record<string, unknown>): Stats
     recent,
     operations,
     phases,
+    outcomes,
     failures,
     upstreamOutputBytes,
     cacheHits: numberField(cache, "hits"),
@@ -755,6 +761,11 @@ class ServerManagerModal implements Component, Focusable {
         ),
       );
     }
+    lines.push("", this.theme.fg("dim", this.theme.bold("OUTCOMES")));
+    if (stats.outcomes.length === 0) lines.push(this.theme.fg("muted", "No outcomes yet"));
+    for (const outcome of stats.outcomes) {
+      lines.push(`${outcome.name.padEnd(22)} ${outcome.count.toLocaleString().padStart(8)}`);
+    }
     lines.push("", this.theme.fg("dim", this.theme.bold("FAILURE STAGES")));
     if (stats.failures.length === 0) lines.push(this.theme.fg("muted", "No failures yet"));
     for (const failure of stats.failures) {

package/src/prompts.ts

@@ -8,10 +8,10 @@ export const INSPECT_PROMPT_GUIDELINES = [
 ] as const;
 
 export const EXECUTE_PROMPT_GUIDELINES = [
-  "Use programmatic execution for a bounded workflow when code can deterministically filter, join, aggregate, deduplicate, validate, or reduce intermediate results.",
+  "Use bounded execution to deterministically filter, aggregate, sample, join, or reduce upstream data in the sandbox; do not return raw payloads.",
   "Keep a model turn between calls when an intermediate result changes the semantic decision or user approval is required.",
-  "Return the smallest result that answers the request; oversized results fail explicitly with bounded structural inspection data.",
-  "SDK facades are prebound globals and must not be imported. Import supported stdlib normally (for example, `import asyncio`); class declarations and `__import__` are unsupported.",
+  "Return the smallest answer. If an oversized result provides result_ref, refine it via inputRef instead of repeating upstream calls.",
+  "SDK facades are prebound globals and must not be imported. Use `import asyncio` for `asyncio.gather`; other imports, classes, `asyncio.create_task`, and `__import__` are unsupported.",
 ] as const;
 
 export const SAVE_CHAIN_PROMPT_GUIDELINES = [

package/src/tools.ts

@@ -144,6 +144,12 @@ const ExecuteParameters = Type.Object({
     description:
       "Sandboxed Python body. Call typed SDK methods such as await linear.list_issues(arguments) and return a compact final value.",
   }),
+  inputRef: Type.Optional(
+    Type.String({
+      minLength: 1,
+      description: "Opaque retained-result reference exposed to sandbox code as input",
+    }),
+  ),
 });
 
 export function registerCodeMcpTools(
@@ -159,7 +165,7 @@ export function registerCodeMcpTools(
     promptSnippet: "Discover compact MCP capabilities or inventory",
     promptGuidelines: [...SEARCH_PROMPT_GUIDELINES],
     parameters: SearchParameters,
-    async execute(_toolCallId, params, signal, onUpdate) {
+    async execute(toolCallId, params, signal, onUpdate) {
       onUpdate?.({
         content: [{ type: "text", text: "Searching MCP tools..." }],
         details: undefined,
@@ -173,6 +179,7 @@ export function registerCodeMcpTools(
           limit: params.limit ?? 5,
           cursor: params.cursor ?? 0,
           ...(params.server === undefined ? {} : { server: params.server }),
+          trace_id: toolCallId,
         },
         signal,
       );
@@ -247,12 +254,16 @@ export function registerCodeMcpTools(
     promptSnippet: "Load exact typed contracts for selected MCP calls",
     promptGuidelines: [...INSPECT_PROMPT_GUIDELINES],
     parameters: InspectParameters,
-    async execute(_toolCallId, params, signal, onUpdate) {
+    async execute(toolCallId, params, signal, onUpdate) {
       onUpdate?.({
         content: [{ type: "text", text: "Inspecting MCP tool contracts..." }],
         details: undefined,
       });
-      const result = await lifecycle.request("inspect", { calls: params.calls }, signal);
+      const result = await lifecycle.request(
+        "inspect",
+        { calls: params.calls, trace_id: toolCallId },
+        signal,
+      );
       const output = formatCodeMcpOutput(result, outputLimits(lifecycle));
       const results = Array.isArray(result.results) ? result.results : [];
       return {
@@ -292,23 +303,34 @@ export function registerCodeMcpTools(
     name: "codemcp_execute",
     label: "MCP Execute",
     description:
-      "Type-check and execute one bounded sandboxed Python MCP call graph. Use it when code can deterministically process intermediate results; preserve a model turn for semantic decisions or approvals. The sandbox has no host filesystem, environment, network, or subprocess access. Oversized results fail with bounded shape, size, and sample diagnostics.",
+      "Type-check and execute one bounded sandboxed Python MCP call graph. Filter, aggregate, or sample upstream data inside the sandbox instead of returning raw payloads; an optional inputRef exposes a retained oversized value as input without repeating upstream calls. Preserve a model turn for semantic decisions or approvals. The sandbox has no host filesystem, environment, network, or subprocess access. Oversized results fail with bounded shape, size, and sample diagnostics.",
     promptSnippet: "Run a bounded typed MCP workflow and return a compact result",
     promptGuidelines: [...EXECUTE_PROMPT_GUIDELINES],
     parameters: ExecuteParameters,
-    async execute(_toolCallId, params, signal, onUpdate) {
+    async execute(toolCallId, params, signal, onUpdate) {
       onUpdate?.({
         content: [{ type: "text", text: "Type-checking MCP chain..." }],
         details: undefined,
       });
-      const result = await lifecycle.request("execute", { code: params.code }, signal);
+      const result = await lifecycle.request(
+        "execute",
+        {
+          code: params.code,
+          trace_id: toolCallId,
+          ...(params.inputRef === undefined ? {} : { input_ref: params.inputRef }),
+        },
+        signal,
+      );
       const ok = result.ok === true;
       const modelValue = ok
         ? result.result
         : {
             failure_stage: result.failure_stage,
             error: result.error,
+            failure: result.failure,
             shape: result.shape,
+            result_ref: result.result_ref,
+            expires_in_seconds: result.expires_in_seconds,
             calls_made: result.calls_made,
             chain_calls: result.chain_calls,
           };
@@ -365,7 +387,7 @@ export function registerCodeMcpTools(
     promptSnippet: "Save a repeated MCP execution as a typed reusable native tool",
     promptGuidelines: [...SAVE_CHAIN_PROMPT_GUIDELINES],
     parameters: SaveChainParameters,
-    async execute(_toolCallId, params, signal, onUpdate) {
+    async execute(toolCallId, params, signal, onUpdate) {
       const scope = requireChainScope(params.scope ?? "project");
       if (scope === "project" && lifecycle.projectChainsPath === undefined) {
         throw new Error(
@@ -385,6 +407,7 @@ export function registerCodeMcpTools(
           inputSchema: params.inputSchema,
           outputSchema: params.outputSchema,
         },
+        toolCallId,
         signal,
       );
       const result = {
@@ -455,11 +478,11 @@ export function registerCodeMcpTools(
     promptSnippet: "List or explicitly manage saved MCP chains",
     promptGuidelines: [...MANAGE_CHAIN_PROMPT_GUIDELINES],
     parameters: ManageChainsParameters,
-    async execute(_toolCallId, params, signal, onUpdate) {
+    async execute(toolCallId, params, signal, onUpdate) {
       const action = params.action;
       let views: SavedChainView[];
       if (action === "list") {
-        views = await chains.list(signal);
+        views = await chains.list(toolCallId, signal);
       } else {
         if (params.confirmedByUser !== true) {
           throw new Error(`${action} requires confirmedByUser=true after explicit user approval`);
@@ -476,13 +499,13 @@ export function registerCodeMcpTools(
           details: undefined,
         });
         if (action === "enable" || action === "disable") {
-          await chains.setEnabled(params.name, scope, action === "enable", signal);
-          views = await chains.list(signal);
+          await chains.setEnabled(params.name, scope, action === "enable", toolCallId, signal);
+          views = await chains.list(toolCallId, signal);
         } else if (action === "revalidate") {
-          await chains.revalidate(params.name, scope, signal);
-          views = await chains.list(signal);
+          await chains.revalidate(params.name, scope, toolCallId, signal);
+          views = await chains.list(toolCallId, signal);
         } else {
-          views = await chains.delete(params.name, scope, signal);
+          views = await chains.delete(params.name, scope, toolCallId, signal);
         }
       }
       const result = {