newcode 0.1.1__py3-none-any.whl

code_puppy/pydantic_patches.py

@@ -0,0 +1,317 @@
+"""Monkey patches for pydantic-ai.
+
+This module contains all monkey patches needed to customize pydantic-ai behavior.
+These patches MUST be applied before any other pydantic-ai imports to work correctly.
+
+Usage:
+    from code_puppy.pydantic_patches import apply_all_patches
+    apply_all_patches()
+"""
+
+import importlib.metadata
+from typing import Any
+
+
+def _get_code_puppy_version() -> str:
+    """Get the current package version."""
+    try:
+        return importlib.metadata.version("newcode")
+    except Exception:
+        return "0.0.0-dev"
+
+
+def patch_user_agent() -> None:
+    """Patch pydantic-ai's User-Agent to use our version.
+
+    pydantic-ai sets its own User-Agent ('pydantic-ai/x.x.x') via a @cache-decorated
+    function. We replace it with a dynamic function that returns:
+    - 'KimiCLI/0.63' for Kimi models
+    - 'NewCode/{version}' for all other models
+
+    This MUST be called before any pydantic-ai models are created.
+    """
+    try:
+        import pydantic_ai.models as pydantic_models
+
+        version = _get_code_puppy_version()
+
+        # Clear cache if already called
+        if hasattr(pydantic_models.get_user_agent, "cache_clear"):
+            pydantic_models.get_user_agent.cache_clear()
+
+        def _get_dynamic_user_agent() -> str:
+            """Return User-Agent based on current model selection."""
+            try:
+                from code_puppy.config import get_global_model_name
+
+                model_name = get_global_model_name()
+                if model_name and "kimi" in model_name.lower():
+                    return "KimiCLI/0.63"
+            except Exception:
+                pass
+            return f"NewCode/{version}"
+
+        pydantic_models.get_user_agent = _get_dynamic_user_agent
+    except Exception:
+        pass  # Don't crash on patch failure
+
+
+def patch_message_history_cleaning() -> None:
+    """Disable overly strict message history cleaning in pydantic-ai."""
+    try:
+        from pydantic_ai import _agent_graph
+
+        _agent_graph._clean_message_history = lambda messages: messages
+    except Exception:
+        pass
+
+
+def patch_process_message_history() -> None:
+    """Patch _process_message_history to skip strict ModelRequest validation.
+
+    Pydantic AI added a validation that history must end with ModelRequest,
+    but this breaks valid conversation flows. We patch it to skip that validation.
+    """
+    try:
+        from pydantic_ai import _agent_graph
+
+        async def _patched_process_message_history(messages, processors, run_context):
+            """Patched version that doesn't enforce ModelRequest at end."""
+            from pydantic_ai._agent_graph import (
+                _HistoryProcessorAsync,
+                _HistoryProcessorSync,
+                _HistoryProcessorSyncWithCtx,
+                cast,
+                exceptions,
+                is_async_callable,
+                is_takes_ctx,
+                run_in_executor,
+            )
+
+            for processor in processors:
+                takes_ctx = is_takes_ctx(processor)
+
+                if is_async_callable(processor):
+                    if takes_ctx:
+                        messages = await processor(run_context, messages)
+                    else:
+                        async_processor = cast(_HistoryProcessorAsync, processor)
+                        messages = await async_processor(messages)
+                else:
+                    if takes_ctx:
+                        sync_processor_with_ctx = cast(
+                            _HistoryProcessorSyncWithCtx, processor
+                        )
+                        messages = await run_in_executor(
+                            sync_processor_with_ctx, run_context, messages
+                        )
+                    else:
+                        sync_processor = cast(_HistoryProcessorSync, processor)
+                        messages = await run_in_executor(sync_processor, messages)
+
+            if len(messages) == 0:
+                raise exceptions.UserError("Processed history cannot be empty.")
+
+            # NOTE: We intentionally skip the "must end with ModelRequest" validation
+            # that was added in newer Pydantic AI versions.
+
+            return messages
+
+        _agent_graph._process_message_history = _patched_process_message_history
+    except Exception:
+        pass
+
+
+def patch_tool_call_json_repair() -> None:
+    """Patch pydantic-ai's _call_tool to auto-repair malformed JSON arguments.
+
+    LLMs sometimes produce slightly broken JSON in tool calls (trailing commas,
+    missing quotes, etc.). This patch intercepts tool calls and runs json_repair
+    on the arguments before validation, preventing unnecessary retries.
+    """
+    try:
+        import json_repair
+        from pydantic_ai._tool_manager import ToolManager
+
+        # Store the original method
+        _original_call_tool = ToolManager._call_tool
+
+        async def _patched_call_tool(
+            self,
+            call,
+            *,
+            allow_partial: bool,
+            wrap_validation_errors: bool,
+            approved: bool,
+            metadata: Any = None,
+        ):
+            """Patched _call_tool that repairs malformed JSON before validation."""
+            # Only attempt repair if args is a string (JSON)
+            if isinstance(call.args, str) and call.args:
+                try:
+                    repaired = json_repair.repair_json(call.args)
+                    if repaired != call.args:
+                        # Update the call args with repaired JSON
+                        call.args = repaired
+                except Exception:
+                    pass  # If repair fails, let original validation handle it
+
+            # Call the original method
+            return await _original_call_tool(
+                self,
+                call,
+                allow_partial=allow_partial,
+                wrap_validation_errors=wrap_validation_errors,
+                approved=approved,
+                metadata=metadata,
+            )
+
+        # Apply the patch
+        ToolManager._call_tool = _patched_call_tool
+
+    except ImportError:
+        pass  # json_repair or pydantic_ai not available
+    except Exception:
+        pass  # Don't crash on patch failure
+
+
+def patch_tool_call_callbacks() -> None:
+    """Patch pydantic-ai's _call_tool to fire pre/post tool callbacks.
+
+    This wraps ToolManager._call_tool so that every tool invocation
+    automatically triggers the ``pre_tool_call`` and ``post_tool_call``
+    callback hooks defined in ``code_puppy.callbacks``, without needing
+    to decorate each tool function individually.
+    """
+    import time
+
+    try:
+        from pydantic_ai._tool_manager import ToolManager
+
+        _original_call_tool = ToolManager._call_tool
+
+        # Tool name prefix used by Claude Code OAuth - tools are prefixed on
+        # outgoing requests, so we need to unprefix them when they come back
+        TOOL_PREFIX = "cp_"
+
+        async def _patched_call_tool(
+            self,
+            call,
+            *,
+            allow_partial: bool,
+            wrap_validation_errors: bool,
+            approved: bool,
+            metadata: Any = None,
+        ):
+            tool_name = call.tool_name
+
+            # Unprefix tool names from Claude Code OAuth responses
+            # The cp_ prefix is added for OAuth compatibility but needs to be
+            # stripped so pydantic-ai can find the actual tool
+            if tool_name and tool_name.startswith(TOOL_PREFIX):
+                unprefixed_name = tool_name[len(TOOL_PREFIX) :]
+                # Try to update the call object's tool_name
+                try:
+                    call.tool_name = unprefixed_name
+                    tool_name = unprefixed_name
+                except (AttributeError, TypeError):
+                    # If the object is immutable, we can't modify it directly
+                    # The tool lookup might still fail, but at least callbacks
+                    # will use the unprefixed name
+                    tool_name = unprefixed_name
+
+            # Normalise args to a dict for the callback contract
+            tool_args: dict = {}
+            if isinstance(call.args, dict):
+                tool_args = call.args
+            elif isinstance(call.args, str):
+                try:
+                    import json
+
+                    tool_args = json.loads(call.args)
+                except Exception:
+                    tool_args = {"raw": call.args}
+
+            # --- pre_tool_call (with blocking support) ---
+            # Returns a string tool-result on block so pydantic-ai sees a clean
+            # "BLOCKED: ..." message and the agent can react gracefully, without
+            # triggering UnexpectedModelBehavior crashes.
+            try:
+                from code_puppy import callbacks
+                from code_puppy.messaging import emit_warning
+
+                callback_results = await callbacks.on_pre_tool_call(
+                    tool_name, tool_args
+                )
+
+                for callback_result in callback_results:
+                    if (
+                        callback_result
+                        and isinstance(callback_result, dict)
+                        and callback_result.get("blocked")
+                    ):
+                        raw_reason = (
+                            callback_result.get("error_message")
+                            or callback_result.get("reason")
+                            or ""
+                        )
+                        if "[BLOCKED]" in raw_reason:
+                            clean_reason = raw_reason[
+                                raw_reason.index("[BLOCKED]") :
+                            ].strip()
+                        else:
+                            clean_reason = (
+                                raw_reason.strip() or "Tool execution blocked by hook"
+                            )
+                        block_msg = f"🚫 Hook blocked this tool call: {clean_reason}"
+                        emit_warning(block_msg)
+                        return f"ERROR: {block_msg}\n\nThe hook policy prevented this tool from running. Please inform the user and do not retry this specific command."
+            except Exception:
+                pass  # other errors don't block tool execution
+
+            start = time.perf_counter()
+            error: Exception | None = None
+            result = None
+            try:
+                result = await _original_call_tool(
+                    self,
+                    call,
+                    allow_partial=allow_partial,
+                    wrap_validation_errors=wrap_validation_errors,
+                    approved=approved,
+                    metadata=metadata,
+                )
+                return result
+            except Exception as exc:
+                error = exc
+                raise
+            finally:
+                duration_ms = (time.perf_counter() - start) * 1000
+                final_result = result if error is None else {"error": str(error)}
+                try:
+                    from code_puppy import callbacks
+
+                    await callbacks.on_post_tool_call(
+                        tool_name, tool_args, final_result, duration_ms
+                    )
+                except Exception:
+                    pass  # never block tool execution
+
+        ToolManager._call_tool = _patched_call_tool
+
+    except ImportError:
+        pass
+    except Exception:
+        pass
+
+
+def apply_all_patches() -> None:
+    """Apply all pydantic-ai monkey patches.
+
+    Call this at the very top of main.py, before any other imports.
+    """
+    patch_user_agent()
+    patch_message_history_cleaning()
+    patch_process_message_history()
+    patch_tool_call_json_repair()
+    patch_tool_call_callbacks()

code_puppy/reopenable_async_client.py

@@ -0,0 +1,232 @@
+"""
+ReopenableAsyncClient - A reopenable httpx.AsyncClient wrapper.
+
+This module provides a ReopenableAsyncClient class that extends httpx.AsyncClient
+to support reopening after being closed, which the standard httpx.AsyncClient
+doesn't support.
+"""
+
+import asyncio
+import threading
+from typing import Optional, Union
+
+import httpx
+
+
+class ReopenableAsyncClient:
+    """
+    A wrapper around httpx.AsyncClient that can be reopened after being closed.
+
+    Standard httpx.AsyncClient becomes unusable after calling aclose().
+    This class allows you to reopen the client and continue using it.
+
+    Example:
+        >>> client = ReopenableAsyncClient(timeout=30.0)
+        >>> await client.get("https://httpbin.org/get")
+        >>> await client.aclose()
+        >>> # Client is now closed, but can be reopened
+        >>> await client.reopen()
+        >>> await client.get("https://httpbin.org/get")  # Works!
+
+    The client preserves all original configuration when reopening.
+    """
+
+    class _StreamWrapper:
+        """Async context manager wrapper for streaming responses."""
+
+        def __init__(
+            self,
+            parent_client: "ReopenableAsyncClient",
+            method: str,
+            url: Union[str, httpx.URL],
+            **kwargs,
+        ):
+            self.parent_client = parent_client
+            self.method = method
+            self.url = url
+            self.kwargs = kwargs
+            self._stream_context = None
+
+        async def __aenter__(self):
+            client = await self.parent_client._ensure_client_open()
+            self._stream_context = client.stream(self.method, self.url, **self.kwargs)
+            return await self._stream_context.__aenter__()
+
+        async def __aexit__(self, exc_type, exc_val, exc_tb):
+            if self._stream_context:
+                return await self._stream_context.__aexit__(exc_type, exc_val, exc_tb)
+
+    def __init__(self, client_class=None, **kwargs):
+        """
+        Initialize the ReopenableAsyncClient.
+
+        Args:
+            client_class: Class to use for creating the internal client (defaults to httpx.AsyncClient)
+            **kwargs: All arguments that would be passed to the client constructor
+        """
+        self._client_class = client_class or httpx.AsyncClient
+        self._client_kwargs = kwargs.copy()
+        self._client: Optional[httpx.AsyncClient] = None
+        self._is_closed = True
+        self._lock = asyncio.Lock()
+        self._sync_lock = threading.Lock()
+
+    async def _ensure_client_open(self) -> httpx.AsyncClient:
+        """
+        Ensure the underlying client is open and ready to use.
+
+        Returns:
+            The active client instance
+
+        Raises:
+            RuntimeError: If client cannot be opened
+        """
+        async with self._lock:
+            if self._is_closed or self._client is None:
+                await self._create_client()
+            return self._client
+
+    async def _create_client(self) -> None:
+        """Create a new client with the stored configuration."""
+        if self._client is not None and not self._is_closed:
+            # Close existing client first
+            await self._client.aclose()
+
+        self._client = self._client_class(**self._client_kwargs)
+        self._is_closed = False
+
+    async def reopen(self) -> None:
+        """
+        Explicitly reopen the client after it has been closed.
+
+        This is useful when you want to reuse a client that was previously closed.
+        """
+        async with self._lock:
+            await self._create_client()
+
+    async def aclose(self) -> None:
+        """
+        Close the underlying httpx.AsyncClient.
+
+        After calling this, the client can be reopened using reopen() or
+        automatically when making the next request.
+        """
+        async with self._lock:
+            if self._client is not None and not self._is_closed:
+                await self._client.aclose()
+                self._is_closed = True
+
+    @property
+    def is_closed(self) -> bool:
+        """Check if the client is currently closed."""
+        return self._is_closed or self._client is None
+
+    # Delegate all httpx.AsyncClient methods to the underlying client
+
+    async def get(self, url: Union[str, httpx.URL], **kwargs) -> httpx.Response:
+        """Make a GET request."""
+        client = await self._ensure_client_open()
+        return await client.get(url, **kwargs)
+
+    async def post(self, url: Union[str, httpx.URL], **kwargs) -> httpx.Response:
+        """Make a POST request."""
+        client = await self._ensure_client_open()
+        return await client.post(url, **kwargs)
+
+    async def put(self, url: Union[str, httpx.URL], **kwargs) -> httpx.Response:
+        """Make a PUT request."""
+        client = await self._ensure_client_open()
+        return await client.put(url, **kwargs)
+
+    async def patch(self, url: Union[str, httpx.URL], **kwargs) -> httpx.Response:
+        """Make a PATCH request."""
+        client = await self._ensure_client_open()
+        return await client.patch(url, **kwargs)
+
+    async def delete(self, url: Union[str, httpx.URL], **kwargs) -> httpx.Response:
+        """Make a DELETE request."""
+        client = await self._ensure_client_open()
+        return await client.delete(url, **kwargs)
+
+    async def head(self, url: Union[str, httpx.URL], **kwargs) -> httpx.Response:
+        """Make a HEAD request."""
+        client = await self._ensure_client_open()
+        return await client.head(url, **kwargs)
+
+    async def options(self, url: Union[str, httpx.URL], **kwargs) -> httpx.Response:
+        """Make an OPTIONS request."""
+        client = await self._ensure_client_open()
+        return await client.options(url, **kwargs)
+
+    async def request(
+        self, method: str, url: Union[str, httpx.URL], **kwargs
+    ) -> httpx.Response:
+        """Make a request with the specified HTTP method."""
+        client = await self._ensure_client_open()
+        return await client.request(method, url, **kwargs)
+
+    async def send(self, request: httpx.Request, **kwargs) -> httpx.Response:
+        """Send a pre-built request."""
+        client = await self._ensure_client_open()
+        return await client.send(request, **kwargs)
+
+    def build_request(
+        self, method: str, url: Union[str, httpx.URL], **kwargs
+    ) -> httpx.Request:
+        """
+        Build a request without sending it.
+
+        Note: This creates a temporary client if none exists, but doesn't keep it open.
+        """
+        with self._sync_lock:
+            if self._client is None or self._is_closed:
+                # Create temporary sync client for building request only
+                # Use httpx.Client (sync) so we can properly close it
+                temp_client = httpx.Client(**self._client_kwargs)
+                try:
+                    return temp_client.build_request(method, url, **kwargs)
+                finally:
+                    temp_client.close()
+            return self._client.build_request(method, url, **kwargs)
+
+    def stream(self, method: str, url: Union[str, httpx.URL], **kwargs):
+        """Stream a request. Returns an async context manager."""
+        return self._StreamWrapper(self, method, url, **kwargs)
+
+    # Context manager support
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self._ensure_client_open()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.aclose()
+
+    # Properties that don't require an active client
+    @property
+    def timeout(self) -> Optional[httpx.Timeout]:
+        """Get the configured timeout."""
+        return self._client_kwargs.get("timeout")
+
+    @property
+    def headers(self) -> httpx.Headers:
+        """Get the configured headers."""
+        if self._client is not None:
+            return self._client.headers
+        # Return headers from kwargs if client doesn't exist
+        headers = self._client_kwargs.get("headers", {})
+        return httpx.Headers(headers)
+
+    @property
+    def cookies(self) -> httpx.Cookies:
+        """Get the current cookies."""
+        if self._client is not None and not self._is_closed:
+            return self._client.cookies
+        # Return empty cookies if client doesn't exist or is closed
+        return httpx.Cookies()
+
+    def __repr__(self) -> str:
+        """String representation of the client."""
+        status = "closed" if self.is_closed else "open"
+        return f"<ReopenableAsyncClient [{status}]>"

code_puppy/round_robin_model.py

@@ -0,0 +1,150 @@
+import threading
+from contextlib import asynccontextmanager, suppress
+from dataclasses import dataclass, field
+from typing import Any, AsyncIterator, List
+
+from pydantic_ai._run_context import RunContext
+from pydantic_ai.models import (
+    Model,
+    ModelMessage,
+    ModelRequestParameters,
+    ModelResponse,
+    ModelSettings,
+    StreamedResponse,
+)
+
+try:
+    from opentelemetry.context import get_current_span
+except ImportError:
+    # If opentelemetry is not installed, provide a dummy implementation
+    def get_current_span():
+        class DummySpan:
+            def is_recording(self):
+                return False
+
+            def set_attributes(self, attributes):
+                pass
+
+        return DummySpan()
+
+
+@dataclass(init=False)
+class RoundRobinModel(Model):
+    """A model that cycles through multiple models in a round-robin fashion.
+
+    This model distributes requests across multiple candidate models to help
+    overcome rate limits or distribute load.
+    """
+
+    models: List[Model]
+    _current_index: int = field(default=0, repr=False)
+    _model_name: str = field(repr=False)
+    _rotate_every: int = field(default=1, repr=False)
+    _request_count: int = field(default=0, repr=False)
+    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)
+
+    def __init__(
+        self,
+        *models: Model,
+        rotate_every: int = 1,
+        settings: ModelSettings | None = None,
+    ):
+        """Initialize a round-robin model instance.
+
+        Args:
+            models: The model instances to cycle through.
+            rotate_every: Number of requests before rotating to the next model (default: 1).
+            settings: Model settings that will be used as defaults for this model.
+        """
+        super().__init__(settings=settings)
+        if not models:
+            raise ValueError("At least one model must be provided")
+        if rotate_every < 1:
+            raise ValueError("rotate_every must be at least 1")
+        self.models = list(models)
+        self._current_index = 0
+        self._request_count = 0
+        self._rotate_every = rotate_every
+        self._lock = threading.Lock()
+
+    @property
+    def model_name(self) -> str:
+        """The model name showing this is a round-robin model with its candidates."""
+        base_name = f"round_robin:{','.join(model.model_name for model in self.models)}"
+        if self._rotate_every != 1:
+            return f"{base_name}:rotate_every={self._rotate_every}"
+        return base_name
+
+    @property
+    def system(self) -> str:
+        """System prompt from the current model."""
+        return self.models[self._current_index].system
+
+    @property
+    def base_url(self) -> str | None:
+        """Base URL from the current model."""
+        return self.models[self._current_index].base_url
+
+    def _get_next_model(self) -> Model:
+        """Get the next model in the round-robin sequence and update the index."""
+        with self._lock:
+            model = self.models[self._current_index]
+            self._request_count += 1
+            if self._request_count >= self._rotate_every:
+                self._current_index = (self._current_index + 1) % len(self.models)
+                self._request_count = 0
+            return model
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        """Make a request using the next model in the round-robin sequence."""
+        current_model = self._get_next_model()
+        # Use prepare_request to merge settings and customize parameters
+        merged_settings, prepared_params = current_model.prepare_request(
+            model_settings, model_request_parameters
+        )
+
+        try:
+            response = await current_model.request(
+                messages, merged_settings, prepared_params
+            )
+            self._set_span_attributes(current_model)
+            return response
+        except Exception:
+            # Unlike FallbackModel, we don't try other models here
+            # The round-robin strategy is about distribution, not failover
+            raise
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
+    ) -> AsyncIterator[StreamedResponse]:
+        """Make a streaming request using the next model in the round-robin sequence."""
+        current_model = self._get_next_model()
+        # Use prepare_request to merge settings and customize parameters
+        merged_settings, prepared_params = current_model.prepare_request(
+            model_settings, model_request_parameters
+        )
+
+        async with current_model.request_stream(
+            messages, merged_settings, prepared_params, run_context
+        ) as response:
+            self._set_span_attributes(current_model)
+            yield response
+
+    def _set_span_attributes(self, model: Model):
+        """Set span attributes for observability."""
+        with suppress(Exception):
+            span = get_current_span()
+            if span.is_recording():
+                attributes = getattr(span, "attributes", {})
+                if attributes.get("gen_ai.request.model") == self.model_name:
+                    span.set_attributes(model.model_attributes(model))