pi-codemcp 1.3.1 → 1.3.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-codemcp",
-  "version": "1.3.1",
+  "version": "1.3.3",
   "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

@@ -1511,6 +1511,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",

package/sidecar/mcp_config.py

@@ -4,21 +4,34 @@ import hashlib
 import json
 import os
 import re
-from typing import TYPE_CHECKING, override
+import time
+from contextlib import suppress
+from typing import TYPE_CHECKING, Any, override
 from urllib.parse import urlsplit
 
+import httpx
 from fastmcp.client.auth import OAuth
+from fastmcp.client.auth.oauth import TokenStorageAdapter
 from fastmcp.mcp_config import (
     MCPConfig,
     RemoteMCPServer,
     StdioMCPServer,
     infer_transport_type_from_url,
 )
+from key_value.aio.adapters.pydantic import PydanticAdapter
 from key_value.aio.stores.filetree import FileTreeStore
 from key_value.aio.stores.filetree.store import (
     FileTreeV1CollectionSanitizationStrategy,
     FileTreeV1KeySanitizationStrategy,
 )
+from mcp.client.auth.utils import (
+    build_oauth_authorization_server_metadata_discovery_urls,
+    build_protected_resource_metadata_discovery_urls,
+    create_oauth_metadata_request,
+    handle_auth_metadata_response,
+    handle_protected_resource_response,
+)
+from mcp.shared.auth import OAuthMetadata, ProtectedResourceMetadata
 from pydantic import BaseModel, ConfigDict
 
 from .json_types import JSON_VALUE_ADAPTER, JsonObject, JsonValue
@@ -27,7 +40,9 @@ from .models import NormalizedServerInfo, ServerAuth
 if TYPE_CHECKING:
     from pathlib import Path
 
-    import httpx
+    from key_value.aio.protocols import AsyncKeyValue
+    from mcp.shared.auth import OAuthClientInformationFull
+    from pydantic import AnyUrl
 
 PI_ONLY_FIELDS = {"directTools", "lifecycle", "idleTimeout", "disabled", "enabled"}
 REMOTE_TRANSPORTS = {"http", "streamable-http", "sse"}
@@ -59,16 +74,165 @@ class NormalizedConfig(BaseModel):
     servers: list[NormalizedServerInfo]
 
 
+class CodemcpTokenStorage(TokenStorageAdapter):
+    """Token storage that keeps OAuth state usable across sidecar restarts.
+
+    The upstream adapter evicts the client registration once the server-announced
+    client_secret_expires_at passes (Linear announces 24 hours), which silently
+    forces a full browser re-login after the next token expiry. A genuinely dead
+    secret still surfaces as invalid_client and re-registers, so persisting the
+    registration is strictly better. This adapter also persists the discovered
+    authorization-server metadata so token refresh hits the real token endpoint
+    instead of the SDK's "<origin>/token" fallback (a 404 for e.g. Outline).
+    """
+
+    def __init__(self, async_key_value: AsyncKeyValue, server_url: str) -> None:
+        super().__init__(async_key_value, server_url)
+        self._storage_oauth_metadata = PydanticAdapter[OAuthMetadata](
+            default_collection="mcp-oauth-metadata",
+            key_value=async_key_value,
+            pydantic_model=OAuthMetadata,
+            raise_on_validation_error=True,
+        )
+        self._storage_protected_resource = PydanticAdapter[ProtectedResourceMetadata](
+            default_collection="mcp-oauth-protected-resource",
+            key_value=async_key_value,
+            pydantic_model=ProtectedResourceMetadata,
+            raise_on_validation_error=True,
+        )
+
+    @override
+    async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
+        await self._storage_client_info.put(
+            key=self._get_client_info_cache_key(),
+            value=client_info,
+        )
+
+    @override
+    async def clear(self) -> None:
+        await super().clear()
+        await self._storage_oauth_metadata.delete(key=self._oauth_metadata_cache_key())
+        await self._storage_protected_resource.delete(key=self._protected_resource_cache_key())
+
+    async def get_oauth_metadata(self) -> OAuthMetadata | None:
+        result: OAuthMetadata | None = await self._storage_oauth_metadata.get(
+            key=self._oauth_metadata_cache_key()
+        )
+        return result
+
+    async def set_oauth_metadata(self, metadata: OAuthMetadata) -> None:
+        await self._storage_oauth_metadata.put(
+            key=self._oauth_metadata_cache_key(),
+            value=metadata,
+        )
+
+    async def get_protected_resource_metadata(self) -> ProtectedResourceMetadata | None:
+        result: ProtectedResourceMetadata | None = await self._storage_protected_resource.get(
+            key=self._protected_resource_cache_key()
+        )
+        return result
+
+    async def set_protected_resource_metadata(self, metadata: ProtectedResourceMetadata) -> None:
+        await self._storage_protected_resource.put(
+            key=self._protected_resource_cache_key(),
+            value=metadata,
+        )
+
+    def _oauth_metadata_cache_key(self) -> str:
+        return f"{self._server_url}/oauth_metadata"
+
+    def _protected_resource_cache_key(self) -> str:
+        return f"{self._server_url}/protected_resource"
+
+
 class PersistentCallbackOAuth(OAuth):
-    """Reuse the callback registered with a persisted dynamic OAuth client."""
+    """OAuth provider hardened for long-lived shared file token storage.
+
+    On top of reusing the callback registered with a persisted dynamic client:
+    - discovered authorization-server metadata is persisted and restored so token
+      refresh works in fresh sidecar processes (the SDK otherwise falls back to
+      "<origin>/token", which 404s for servers like Outline and turns every
+      access-token expiry into a forced interactive re-login);
+    - a refresh response without refresh_token keeps the previous one (RFC 6749
+      section 6 allows omission when the refresh token does not rotate);
+    - a failed refresh adopts fresher tokens another sidecar process may have
+      stored instead of dropping straight into the browser flow.
+    """
+
+    def __init__(
+        self,
+        *,
+        mcp_url: str,
+        client_name: str,
+        token_storage: AsyncKeyValue,
+        additional_client_metadata: dict[str, Any] | None = None,
+    ) -> None:
+        self._codemcp_token_store = token_storage
+        self._persisted_oauth_metadata: OAuthMetadata | None = None
+        self._persisted_protected_resource: ProtectedResourceMetadata | None = None
+        super().__init__(
+            mcp_url=mcp_url,
+            client_name=client_name,
+            token_storage=token_storage,
+            additional_client_metadata=additional_client_metadata,
+        )
+
+    @override
+    def _bind(self, mcp_url: str) -> None:
+        super()._bind(mcp_url)
+        if isinstance(self.token_storage_adapter, CodemcpTokenStorage):
+            return
+        storage = CodemcpTokenStorage(self._codemcp_token_store, self.mcp_url)
+        self.token_storage_adapter = storage
+        self.context.storage = storage
 
     @override
     async def _initialize(self) -> None:
         await super()._initialize()
         client_info = self.context.client_info
-        if client_info is None or not client_info.redirect_uris:
+        if client_info is not None and client_info.redirect_uris:
+            self._reuse_registered_callback(client_info.redirect_uris[0])
+        storage = self.token_storage_adapter
+        if not isinstance(storage, CodemcpTokenStorage):
             return
-        redirect_uri = client_info.redirect_uris[0]
+        if self.context.oauth_metadata is None:
+            self.context.oauth_metadata = await storage.get_oauth_metadata()
+            self._persisted_oauth_metadata = self.context.oauth_metadata
+        if self.context.protected_resource_metadata is None:
+            self.context.protected_resource_metadata = (
+                await storage.get_protected_resource_metadata()
+            )
+            self._persisted_protected_resource = self.context.protected_resource_metadata
+        tokens = self.context.current_tokens
+        if tokens is not None and tokens.expires_in and await storage.get_token_expiry() is None:
+            # Without the absolute expiry record the upstream fallback re-applies the
+            # stale relative expires_in from now; the expired access token then looks
+            # valid, gets rejected with a 401, and the flow goes interactive.
+            self.context.token_expiry_time = time.time() - 1
+
+    @override
+    async def _refresh_token(self) -> httpx.Request:
+        if self.context.oauth_metadata is None:
+            await self._discover_server_metadata()
+        return await super()._refresh_token()
+
+    @override
+    async def _handle_token_response(self, response: httpx.Response) -> None:
+        await super()._handle_token_response(response)
+        await self._persist_discovered_metadata()
+
+    @override
+    async def _handle_refresh_response(self, response: httpx.Response) -> bool:
+        previous_tokens = self.context.current_tokens
+        previous_access_token = previous_tokens.access_token if previous_tokens else None
+        previous_refresh_token = previous_tokens.refresh_token if previous_tokens else None
+        if await super()._handle_refresh_response(response):
+            await self._restore_unrotated_refresh_token(previous_refresh_token)
+            await self._persist_discovered_metadata()
+            return True
+        return await self._adopt_tokens_refreshed_elsewhere(previous_access_token)
+
+    def _reuse_registered_callback(self, redirect_uri: AnyUrl) -> None:
         parsed = urlsplit(str(redirect_uri))
         if (
             parsed.scheme != "http"
@@ -82,6 +246,90 @@ class PersistentCallbackOAuth(OAuth):
         self._callback_host = parsed.hostname
         self.context.client_metadata.redirect_uris = [redirect_uri]
 
+    async def _discover_server_metadata(self) -> None:
+        """Best-effort OAuth discovery so refresh uses the real token endpoint."""
+        with suppress(httpx.HTTPError, ValueError):
+            async with self.httpx_client_factory() as client:
+                await self._discover_protected_resource(client)
+                await self._discover_authorization_server(client)
+        await self._persist_discovered_metadata()
+
+    async def _discover_protected_resource(self, client: httpx.AsyncClient) -> None:
+        if self.context.protected_resource_metadata is not None:
+            return
+        for url in build_protected_resource_metadata_discovery_urls(
+            None,
+            self.context.server_url,
+        ):
+            response = await client.send(create_oauth_metadata_request(url))
+            prm = await handle_protected_resource_response(response)
+            if prm is not None:
+                self.context.protected_resource_metadata = prm
+                self.context.auth_server_url = str(prm.authorization_servers[0])
+                return
+
+    async def _discover_authorization_server(self, client: httpx.AsyncClient) -> None:
+        if self.context.oauth_metadata is not None:
+            return
+        for url in build_oauth_authorization_server_metadata_discovery_urls(
+            self.context.auth_server_url,
+            self.context.server_url,
+        ):
+            response = await client.send(create_oauth_metadata_request(url))
+            ok, metadata = await handle_auth_metadata_response(response)
+            if not ok:
+                return
+            if metadata is not None:
+                self.context.oauth_metadata = metadata
+                return
+
+    async def _persist_discovered_metadata(self) -> None:
+        storage = self.token_storage_adapter
+        if not isinstance(storage, CodemcpTokenStorage):
+            return
+        metadata = self.context.oauth_metadata
+        if metadata is not None and metadata != self._persisted_oauth_metadata:
+            await storage.set_oauth_metadata(metadata)
+            self._persisted_oauth_metadata = metadata
+        resource = self.context.protected_resource_metadata
+        if resource is not None and resource != self._persisted_protected_resource:
+            await storage.set_protected_resource_metadata(resource)
+            self._persisted_protected_resource = resource
+
+    async def _restore_unrotated_refresh_token(self, previous_refresh_token: str | None) -> None:
+        tokens = self.context.current_tokens
+        if tokens is None or tokens.refresh_token is not None or previous_refresh_token is None:
+            return
+        # RFC 6749 section 6: the server may omit refresh_token when it does not
+        # rotate; the SDK overwrites the stored token set and would lose it.
+        tokens.refresh_token = previous_refresh_token
+        await self.context.storage.set_tokens(tokens)
+
+    async def _adopt_tokens_refreshed_elsewhere(self, previous_access_token: str | None) -> bool:
+        storage = self.token_storage_adapter
+        if not isinstance(storage, CodemcpTokenStorage):
+            return False
+        stored = await storage.get_tokens()
+        if (
+            stored is None
+            or not stored.access_token
+            or stored.access_token == previous_access_token
+        ):
+            return False
+        expiry = await storage.get_token_expiry()
+        if expiry is not None and time.time() > expiry:
+            return False
+        # Another sidecar process rotated the refresh token first and stored the
+        # result; adopt it instead of dropping into the interactive flow.
+        self.context.current_tokens = stored
+        if expiry is not None:
+            self.context.token_expiry_time = expiry
+        elif stored.expires_in is not None:
+            self.context.token_expiry_time = time.time() + stored.expires_in
+        else:
+            self.context.token_expiry_time = None
+        return True
+
 
 def load_mcp_json(path: Path) -> JsonObject:
     if not path.exists():

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}`)}`;
   }