pi-codemcp 1.2.1 → 1.3.0

package/sidecar/mcp_config.py

@@ -4,7 +4,8 @@ import hashlib
 import json
 import os
 import re
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, override
+from urllib.parse import urlsplit
 
 from fastmcp.client.auth import OAuth
 from fastmcp.mcp_config import (
@@ -58,6 +59,30 @@ class NormalizedConfig(BaseModel):
     servers: list[NormalizedServerInfo]
 
 
+class PersistentCallbackOAuth(OAuth):
+    """Reuse the callback registered with a persisted dynamic OAuth client."""
+
+    @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:
+            return
+        redirect_uri = client_info.redirect_uris[0]
+        parsed = urlsplit(str(redirect_uri))
+        if (
+            parsed.scheme != "http"
+            or parsed.hostname not in {"localhost", "127.0.0.1", "::1"}
+            or parsed.port is None
+        ):
+            return
+        if parsed.path != "/callback" or parsed.query or parsed.fragment:
+            return
+        self.redirect_port = parsed.port
+        self._callback_host = parsed.hostname
+        self.context.client_metadata.redirect_uris = [redirect_uri]
+
+
 def load_mcp_json(path: Path) -> JsonObject:
     if not path.exists():
         raise FileNotFoundError(f"MCP config not found: {path}")
@@ -231,7 +256,7 @@ def normalize_mcp_config(
             auth: str | httpx.Auth | None
             auth_kind: ServerAuth | None = None
             if raw_auth == "oauth":
-                auth = OAuth(
+                auth = PersistentCallbackOAuth(
                     mcp_url=url,
                     client_name=oauth_client_name,
                     token_storage=oauth_storage,

package/sidecar/refinement_cache.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import math
+import time
+import uuid
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal
+
+from .json_types import JSON_VALUE_ADAPTER, JsonValue
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+DEFAULT_REFINEMENT_ENTRY_LIMIT = 8
+DEFAULT_REFINEMENT_ENTRY_BYTES = 1024 * 1024
+DEFAULT_REFINEMENT_TOTAL_BYTES = 4 * 1024 * 1024
+DEFAULT_REFINEMENT_TTL_SECONDS = 300.0
+
+
+type ReferenceFailureReason = Literal[
+    "expired",
+    "evicted",
+    "cross_sidecar",
+    "unknown",
+]
+
+
+@dataclass(frozen=True, slots=True)
+class RetainedResult:
+    reference: str
+    expires_in_seconds: int
+
+
+@dataclass(frozen=True, slots=True)
+class _CacheEntry:
+    payload: bytes
+    expires_at: float
+
+
+class ResultReferenceError(ValueError):
+    def __init__(self, reason: ReferenceFailureReason, reference: str) -> None:
+        self.reason = reason
+        self.reference = reference
+        messages = {
+            "expired": "Result reference has expired",
+            "evicted": "Result reference was evicted",
+            "cross_sidecar": "Result reference belongs to another sidecar",
+            "unknown": "Unknown result reference",
+        }
+        super().__init__(f"{messages[reason]}: {reference}")
+
+
+class RefinementCache:
+    def __init__(
+        self,
+        *,
+        entry_limit: int = DEFAULT_REFINEMENT_ENTRY_LIMIT,
+        entry_byte_limit: int = DEFAULT_REFINEMENT_ENTRY_BYTES,
+        total_byte_limit: int = DEFAULT_REFINEMENT_TOTAL_BYTES,
+        ttl_seconds: float = DEFAULT_REFINEMENT_TTL_SECONDS,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        if entry_limit < 1:
+            raise ValueError("entry_limit must be positive")
+        if entry_byte_limit < 1:
+            raise ValueError("entry_byte_limit must be positive")
+        if total_byte_limit < entry_byte_limit:
+            raise ValueError("total_byte_limit must be at least entry_byte_limit")
+        if ttl_seconds <= 0:
+            raise ValueError("ttl_seconds must be positive")
+        self.entry_limit = entry_limit
+        self.entry_byte_limit = entry_byte_limit
+        self.total_byte_limit = total_byte_limit
+        self.ttl_seconds = ttl_seconds
+        self._clock = clock
+        self._instance_id = uuid.uuid4().hex[:12]
+        self._prefix = f"result_{self._instance_id}_"
+        self._entries: OrderedDict[str, _CacheEntry] = OrderedDict()
+        self._invalid: OrderedDict[str, ReferenceFailureReason] = OrderedDict()
+        self._total_bytes = 0
+
+    @property
+    def entry_count(self) -> int:
+        return len(self._entries)
+
+    @property
+    def total_bytes(self) -> int:
+        return self._total_bytes
+
+    def retain(self, value: JsonValue) -> RetainedResult | None:
+        payload = JSON_VALUE_ADAPTER.dump_json(value)
+        payload_bytes = len(payload)
+        if payload_bytes > self.entry_byte_limit or payload_bytes > self.total_byte_limit:
+            return None
+        now = self._clock()
+        self._prune_expired(now)
+        while (
+            len(self._entries) >= self.entry_limit
+            or self._total_bytes + payload_bytes > self.total_byte_limit
+        ):
+            self._evict_oldest()
+        reference = f"{self._prefix}{uuid.uuid4().hex[:16]}"
+        self._entries[reference] = _CacheEntry(
+            payload=payload,
+            expires_at=now + self.ttl_seconds,
+        )
+        self._total_bytes += payload_bytes
+        return RetainedResult(
+            reference=reference,
+            expires_in_seconds=math.ceil(self.ttl_seconds),
+        )
+
+    def resolve(self, reference: str) -> JsonValue:
+        now = self._clock()
+        self._prune_expired(now)
+        entry = self._entries.get(reference)
+        if entry is not None:
+            self._entries.move_to_end(reference)
+            return JSON_VALUE_ADAPTER.validate_json(entry.payload)
+        reason = self._invalid.get(reference)
+        if reason is not None:
+            raise ResultReferenceError(reason, reference)
+        if not reference.startswith(self._prefix):
+            raise ResultReferenceError("cross_sidecar", reference)
+        raise ResultReferenceError("unknown", reference)
+
+    def clear(self) -> None:
+        self._entries.clear()
+        self._invalid.clear()
+        self._total_bytes = 0
+
+    def _prune_expired(self, now: float) -> None:
+        expired = [
+            reference for reference, entry in self._entries.items() if entry.expires_at <= now
+        ]
+        for reference in expired:
+            self._remove(reference, "expired")
+
+    def _evict_oldest(self) -> None:
+        reference = next(iter(self._entries))
+        self._remove(reference, "evicted")
+
+    def _remove(
+        self,
+        reference: str,
+        reason: Literal["expired", "evicted"],
+    ) -> None:
+        entry = self._entries.pop(reference)
+        self._total_bytes -= len(entry.payload)
+        self._invalid[reference] = reason
+        tombstone_limit = self.entry_limit * 4
+        while len(self._invalid) > tombstone_limit:
+            self._invalid.popitem(last=False)

package/sidecar/sandbox_api.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+INSPECT_JSON_NAME = "inspect_json"
+EXPECT_OBJECT_NAME = "expect_object"
+EXPECT_LIST_NAME = "expect_list"
+EXPECT_STRING_NAME = "expect_string"
+EXPECT_INTEGER_NAME = "expect_integer"
+
+SANDBOX_FUNCTION_EXTERNALS = {
+    INSPECT_JSON_NAME: "__codemcp_inspect_json",
+    EXPECT_OBJECT_NAME: "__codemcp_expect_object",
+    EXPECT_LIST_NAME: "__codemcp_expect_list",
+    EXPECT_STRING_NAME: "__codemcp_expect_string",
+    EXPECT_INTEGER_NAME: "__codemcp_expect_integer",
+}
+
+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"]',
+)
+SANDBOX_CAPABILITY_STUB = (
+    "# Prebound SDK/helpers. Supported: import asyncio; asyncio.gather.\n"
+    "# Unavailable: collections.Counter, base64, gzip, asyncio.create_task,\n"
+    "# __import__. Use dict counts."
+)
+INSPECT_JSON_STUB = (
+    "def inspect_json(value: JsonValue, *, "
+    "samples: Literal[1, 2, 3] = 2, "
+    "max_depth: Literal[1, 2, 3, 4, 5, 6] = 3) -> JsonValue: ..."
+)
+NARROWING_HELPER_STUBS = (
+    "def expect_object(value: JsonValue) -> dict[str, JsonValue]: ...",
+    "def expect_list(value: JsonValue) -> list[JsonValue]: ...",
+    "def expect_string(value: JsonValue) -> str: ...",
+    "def expect_integer(value: JsonValue) -> int: ...",
+)
+STUB_PRELUDE = "\n\n".join((
+    STUB_IMPORTS,
+    *JSON_TYPE_STUBS,
+    SANDBOX_CAPABILITY_STUB,
+    INSPECT_JSON_STUB,
+    *NARROWING_HELPER_STUBS,
+))