pi-codemcp 1.3.2 → 1.4.0

package/README.md

@@ -7,6 +7,7 @@ Instead of putting every upstream MCP tool definition into the model context, pi
 - `codemcp_search` ranks capabilities or pages through a compact inventory without loading full schemas.
 - `codemcp_inspect` returns exact typed SDK stubs only for selected calls.
 - `codemcp_execute` runs one sandboxed Python call graph across one or many MCP servers.
+- `codemcp_edit` applies one exact replacement to the previous execution and reruns it without resending the full code.
 - `codemcp_save_chain` turns a repeated call graph into a reusable native Pi tool.
 - `codemcp_manage_chains` lists chains or performs an explicitly confirmed enable, disable, revalidate, or delete.
 
@@ -177,6 +178,8 @@ Incomplete upstream schemas become recursive `JsonValue`, not `Any`; use the pre
 
 When an oversized value fits the bounded in-memory refinement cache, the failure also returns an opaque `result_ref` and expiry. Pass that reference back as `inputRef` on one follow-up `codemcp_execute`; the retained JSON is exposed as `input`, so code can filter or aggregate it without repeating upstream calls. References expire after five minutes, are valid only in the originating sidecar, and are never persisted.
 
+For a small correction, `codemcp_edit` replaces one uniquely matching `oldText` with `newText` in the most recent execution and reruns it with the same `inputRef`. The state is in-memory and disappears when the sidecar restarts. The full call graph runs again, including upstream MCP calls.
+
 ## Saved-chain CLI flow
 
 Saved chains are JSON manifests with sandboxed code plus explicit input/output JSON Schemas. Project-scoped chains live under `<project>/.pi/pi-codemcp/chains`; global chains live under `<agent-dir>/pi-codemcp/chains`.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-codemcp",
-  "version": "1.3.2",
+  "version": "1.4.0",
   "description": "Typed, sandboxed Code Mode access to configured MCP servers for Pi",
   "type": "module",
   "packageManager": "bun@1.3.10",

package/sidecar/executor.py

@@ -50,9 +50,10 @@ INSPECT_BYTE_LIMIT = 8 * 1024
 CHAIN_INPUT_EXTERNAL = "__codemcp_saved_chain_input"
 
 
-type FailureStage = Literal["preflight", "runtime", "timeout", "cancelled", "result"]
+type FailureStage = Literal["preflight", "arguments", "runtime", "timeout", "cancelled", "result"]
 type FailureKind = Literal[
     "preflight",
+    "argument_validation",
     "result",
     "result_reference",
     "sandbox_runtime",
@@ -487,7 +488,18 @@ class MontyExecutor:
                 raise RuntimeError(
                     f"Call limit exceeded: maximum {context.settings.max_calls} total calls"
                 )
-            validated = catalog.validate_arguments(name, arguments)
+            try:
+                validated = catalog.validate_arguments(name, arguments)
+            except (TypeError, ValidationError, ValueError) as error:
+                message = f"{spec.call}: invalid arguments: {error}"
+                context.failure = ExecutionFailureInfo(
+                    kind="argument_validation",
+                    server=spec.server,
+                    tool=spec.backend_name,
+                    retryable=False,
+                    message=message,
+                )
+                raise ValueError(message) from error
             if spec.kind == "saved_chain":
                 context.chain_calls += 1
                 return await context.call_tool(name, validated, context)
@@ -624,8 +636,10 @@ class MontyExecutor:
             context.metrics.runtime_ms += _elapsed_ms(runtime_started)
             message = error.display("type-msg").strip()
             lowered = message.lower()
-            stage: Literal["runtime", "timeout"] = (
-                "timeout"
+            stage: Literal["arguments", "runtime", "timeout"] = (
+                "arguments"
+                if context.failure is not None and context.failure.kind == "argument_validation"
+                else "timeout"
                 if (context.failure is not None and context.failure.kind == "upstream_timeout")
                 or "duration" in lowered
                 or "timed out" in lowered
@@ -728,7 +742,7 @@ class MontyExecutor:
     @staticmethod
     def _failure(
         context: ExecutionContext,
-        stage: Literal["preflight", "runtime", "timeout", "cancelled", "result"],
+        stage: FailureStage,
         error: str,
     ) -> ExecutionResponse:
         failure = context.failure or _execution_failure_info(stage, error)
@@ -749,6 +763,8 @@ def _execution_failure_info(
 ) -> ExecutionFailureInfo:
     if stage == "preflight":
         kind: FailureKind = "preflight"
+    elif stage == "arguments":
+        kind = "argument_validation"
     elif stage == "result":
         kind = "result"
     elif stage == "timeout":

package/sidecar/gateway.py

@@ -206,6 +206,11 @@ class SavedChainHandlers(NamedTuple):
     delete: Callable[[str, ChainScope, str], Awaitable[ChainListResponse]]
 
 
+class _LastExecution(NamedTuple):
+    code: str
+    input_ref: str | None
+
+
 class SavedChainRuntime:
     def __init__(self, handlers: SavedChainHandlers) -> None:
         self.handlers = handlers
@@ -307,6 +312,8 @@ class GatewayRuntime:
             )
         )
         self._catalog_lock = asyncio.Lock()
+        self._execute_lock = asyncio.Lock()
+        self._last_execution: _LastExecution | None = None
 
     @classmethod
     def create(
@@ -595,6 +602,38 @@ class GatewayRuntime:
         trace_id: str,
         input_ref: str | None = None,
     ) -> ExecutionResponse:
+        async with self._execute_lock:
+            return await self._execute(code, trace_id, input_ref)
+
+    async def edit_execute(
+        self,
+        old_text: str,
+        new_text: str,
+        trace_id: str,
+    ) -> ExecutionResponse:
+        async with self._execute_lock:
+            previous = self._last_execution
+            if previous is None:
+                raise ValueError("No previous CodeMCP execution is available in this sidecar")
+            if not old_text:
+                raise ValueError("old_text must not be empty")
+            if old_text == new_text:
+                raise ValueError("old_text and new_text must differ")
+            matches = previous.code.count(old_text)
+            if matches != 1:
+                raise ValueError(
+                    f"old_text must match the previous code exactly once; found {matches} matches"
+                )
+            code = previous.code.replace(old_text, new_text, 1)
+            return await self._execute(code, trace_id, previous.input_ref)
+
+    async def _execute(
+        self,
+        code: str,
+        trace_id: str,
+        input_ref: str | None,
+    ) -> ExecutionResponse:
+        self._last_execution = None
         started = time.perf_counter()
         input_bytes = len(code.encode()) + len((input_ref or "").encode())
         input_value: JsonValue = None
@@ -619,6 +658,7 @@ class GatewayRuntime:
                     response,
                     trace_id,
                 )
+                self._last_execution = _LastExecution(code, input_ref)
                 return response
         discovery_started = time.perf_counter()
         try:
@@ -655,6 +695,7 @@ class GatewayRuntime:
             )
             raise
         self._record_execution("execute", started, input_bytes, response, trace_id)
+        self._last_execution = _LastExecution(code, input_ref)
         return response
 
     async def _execute_chain(
@@ -1511,6 +1552,7 @@ def _execution_failure_subtype(response: ExecutionResponse) -> str:
     failure = response.failure
     if failure is not None and failure.kind in {
         "result_reference",
+        "argument_validation",
         "sandbox_timeout",
         "upstream",
         "upstream_transport",
@@ -1678,6 +1720,16 @@ async def execute(
     return await _require_runtime().execute(code, trace_id, input_ref)
 
 
+@mcp.tool
+async def edit_execute(
+    trace_id: str,
+    old_text: str,
+    new_text: str,
+) -> ExecutionResponse:
+    """Patch and rerun the last sandbox execution."""
+    return await _require_runtime().edit_execute(old_text, new_text, trace_id)
+
+
 @mcp.tool
 async def save_chain(
     trace_id: str,

package/sidecar/sandbox_api.py

@@ -14,7 +14,7 @@ SANDBOX_FUNCTION_EXTERNALS = {
     EXPECT_INTEGER_NAME: "__codemcp_expect_integer",
 }
 
-STUB_IMPORTS = "from typing import Literal, Never, NotRequired, TypeAlias, TypedDict"
+STUB_IMPORTS = "from typing import Literal, Mapping, Never, NotRequired, TypeAlias, TypedDict"
 JSON_TYPE_STUBS = (
     "JsonScalar: TypeAlias = bool | int | float | str | None",
     'JsonValue: TypeAlias = JsonScalar | list["JsonValue"] | dict[str, "JsonValue"]',

package/sidecar/stats.py

@@ -98,6 +98,7 @@ DISTINCT_NAME_QUERIES = {
 FailureOutcome = Literal[
     "success",
     "preflight_rejection",
+    "argument_rejection",
     "result_refinement",
     "upstream_failure",
     "transport_failure",
@@ -996,6 +997,8 @@ def _operation_outcome(
         return "internal_error"
     if failure.stage == "preflight":
         return "preflight_rejection"
+    if failure.stage == "arguments":
+        return "argument_rejection"
     if failure.stage == "result":
         return "result_refinement"
     if failure.stage == "cancelled":

package/sidecar/tool_catalog.py

@@ -426,7 +426,8 @@ class ToolCatalog(BaseModel):
         for spec in specs:
             definitions.extend(spec.stub.split("\n\n") if spec.stub else [])
             facade_methods.setdefault(spec.namespace, []).append(
-                f"    async def {spec.method}(self, arguments: {spec.input_type_name}) "
+                f"    async def {spec.method}(self, arguments: "
+                f"{spec.input_type_name} | Mapping[str, JsonValue]) "
                 f"-> {spec.output_type_name}: ..."
             )
         facades: list[str] = []

package/src/execution-rendering.ts

@@ -114,6 +114,9 @@ function renderCompactFailure(
   if (stage === "preflight") {
     return theme.fg("warning", `✗ Preflight · code not run · ${summary}`);
   }
+  if (stage === "arguments") {
+    return theme.fg("warning", `✗ Argument validation · code run · ${summary}`);
+  }
   if (stage === "timeout") {
     return theme.fg("error", `✗ Timeout · stopped after ${summary}`);
   }
@@ -131,6 +134,9 @@ function failureHeading(stage: string, calls: number, chainCalls: number, theme:
   if (stage === "preflight") {
     return `${theme.fg("warning", theme.bold("Preflight failed"))}\n${theme.fg("muted", "Code was not executed; no upstream side effects")}`;
   }
+  if (stage === "arguments") {
+    return `${theme.fg("warning", theme.bold("Argument validation failed"))}\n${theme.fg("muted", "Code ran; no invalid MCP call was sent")}`;
+  }
   if (stage === "timeout") {
     return `${theme.fg("error", theme.bold("Execution timed out"))}\n${theme.fg("muted", `Stopped after ${summary}`)}`;
   }

package/src/mcp-client.ts

@@ -23,6 +23,7 @@ export type SidecarToolName =
   | "discover"
   | "reload_settings"
   | "execute"
+  | "edit_execute"
   | "save_chain"
   | "list_chains"
   | "execute_chain"
@@ -34,6 +35,7 @@ export type SidecarToolName =
 
 const LONG_RUNNING_TOOLS = new Set<SidecarToolName>([
   "execute",
+  "edit_execute",
   "save_chain",
   "execute_chain",
   "revalidate_chain",

package/src/tools.ts

@@ -152,6 +152,16 @@ const ExecuteParameters = Type.Object({
   ),
 });
 
+const EditExecuteParameters = Type.Object({
+  oldText: Type.String({
+    minLength: 1,
+    description: "Exact text that must occur once in the previous execution code",
+  }),
+  newText: Type.String({
+    description: "Replacement text; may be empty",
+  }),
+});
+
 export function registerCodeMcpTools(
   pi: ExtensionAPI,
   lifecycle: CodeMcpLifecycle,
@@ -321,32 +331,7 @@ export function registerCodeMcpTools(
         },
         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,
-          };
-      const output = formatCodeMcpOutput(modelValue, outputLimits(lifecycle));
-      return {
-        content: [{ type: "text", text: output.text }],
-        details: {
-          ...output.details,
-          ok,
-          failureStage: typeof result.failure_stage === "string" ? result.failure_stage : undefined,
-          callsMade: Number(result.calls_made ?? 0),
-          chainCalls: Number(result.chain_calls ?? 0),
-          ...(isRecord(result.timings) ? { timings: result.timings } : {}),
-          preview: previewExecutionValue(ok ? result.result : result.error),
-        },
-      };
+      return formatExecutionResponse(result, lifecycle);
     },
     renderCall(args, theme, context) {
       const code = args.code.trim();
@@ -379,6 +364,46 @@ export function registerCodeMcpTools(
     },
   });
 
+  pi.registerTool({
+    name: "codemcp_edit",
+    label: "MCP Edit",
+    description:
+      "Patch and rerun the most recent codemcp_execute in the current sidecar without resending its full code. oldText must match exactly once. The original inputRef is reused and the entire call graph, including upstream MCP calls, runs again.",
+    promptSnippet: "Patch and rerun the previous CodeMCP execution",
+    parameters: EditExecuteParameters,
+    async execute(toolCallId, params, signal, onUpdate) {
+      onUpdate?.({
+        content: [{ type: "text", text: "Patching and type-checking MCP chain..." }],
+        details: undefined,
+      });
+      const result = await lifecycle.request(
+        "edit_execute",
+        {
+          old_text: params.oldText,
+          new_text: params.newText,
+          trace_id: toolCallId,
+        },
+        signal,
+      );
+      return formatExecutionResponse(result, lifecycle);
+    },
+    renderCall(args, theme) {
+      const lines = args.oldText ? args.oldText.split("\n").length : 0;
+      const title = theme.fg("toolTitle", theme.bold("MCP Edit"));
+      const patchLabel = theme.fg(
+        "accent",
+        theme.bold(`Edited ${lines} ${lines === 1 ? "line" : "lines"}`),
+      );
+      return new Text(`${title} ${theme.fg("muted", "·")} ${patchLabel}`, 0, 0);
+    },
+    renderResult(result, state, theme) {
+      return renderExecutionResult(result, state, theme, {
+        partialText: "Applying edit, then execution...",
+        expandDescription: "full output",
+      });
+    },
+  });
+
   pi.registerTool({
     name: "codemcp_save_chain",
     label: "Save MCP Chain",
@@ -552,6 +577,35 @@ export function registerCodeMcpTools(
   });
 }
 
+function formatExecutionResponse(result: Record<string, unknown>, lifecycle: CodeMcpLifecycle) {
+  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,
+      };
+  const output = formatCodeMcpOutput(modelValue, outputLimits(lifecycle));
+  return {
+    content: [{ type: "text" as const, text: output.text }],
+    details: {
+      ...output.details,
+      ok,
+      failureStage: typeof result.failure_stage === "string" ? result.failure_stage : undefined,
+      callsMade: Number(result.calls_made ?? 0),
+      chainCalls: Number(result.chain_calls ?? 0),
+      ...(isRecord(result.timings) ? { timings: result.timings } : {}),
+      preview: previewExecutionValue(ok ? result.result : result.error),
+    },
+  };
+}
+
 function compactChainView(view: SavedChainView) {
   return {
     name: view.chain.name,