code-muse 0.0.1__py3-none-any.whl

code_muse/pydantic_patches.py

@@ -0,0 +1,410 @@
+"""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_muse.pydantic_patches import apply_all_patches
+    apply_all_patches()
+"""
+
+import contextlib
+import importlib.metadata
+from typing import Any
+
+
+def _get_muse_version() -> str:
+    """Get the current Muse version."""
+    try:
+        return importlib.metadata.version("code-muse")
+    except Exception:
+        return "0.0.0-dev"
+
+
+def patch_user_agent() -> None:
+    """Patch pydantic-ai's User-Agent to use Muse's 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
+    - 'Muse/{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_muse_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_muse.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"Muse/{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.
+
+    In newer pydantic-ai versions the function was removed, so this becomes a no-op.
+    """
+    try:
+        from pydantic_ai import _agent_graph
+
+        if not hasattr(_agent_graph, "_process_message_history"):
+            return
+
+        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 (
+                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:
+                        messages = await 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 tool handling to support callbacks and Claude Code tool names.
+
+    Claude Code OAuth prefixes tool names with ``cp_`` on the wire.  pydantic-ai
+    classifies tool calls *before* ``_call_tool`` runs, so unprefixing only in
+    ``_call_tool`` is too late: prefixed tools get marked as ``unknown`` and can
+    burn through result retries, eventually raising ``UnexpectedModelBehavior``.
+
+    This patch normalizes Claude Code tool names early (during lookup/dispatch)
+    and wraps ``_call_tool`` so every tool invocation also triggers the
+    ``pre_tool_call`` and ``post_tool_call`` callbacks defined in
+    ``code_muse.callbacks``.
+    """
+    import time
+
+    try:
+        from pydantic_ai._tool_manager import ToolManager
+
+        _original_call_tool = ToolManager._call_tool
+        _original_get_tool_def = ToolManager.get_tool_def
+        _original_handle_call = ToolManager.handle_call
+
+        # 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_"
+
+        def _normalize_tool_name(name: Any) -> Any:
+            """Strip the ``cp_`` prefix if present."""
+            if isinstance(name, str) and name.startswith(TOOL_PREFIX):
+                return name[len(TOOL_PREFIX) :]
+            return name
+
+        def _normalize_call_tool_name(call: Any) -> tuple[Any, Any]:
+            """Normalize the tool_name on a call object in-place."""
+            tool_name = getattr(call, "tool_name", None)
+            normalized_name = _normalize_tool_name(tool_name)
+            if normalized_name != tool_name:
+                with contextlib.suppress(AttributeError, TypeError):
+                    call.tool_name = normalized_name
+            return normalized_name, call
+
+        # -- Early normalization patches -----------------------------------------
+        # These run *before* pydantic-ai classifies the tool as function/output/
+        # unknown, so prefixed names resolve correctly.
+
+        def _patched_get_tool_def(self, name: str):
+            return _original_get_tool_def(self, _normalize_tool_name(name))
+
+        async def _patched_handle_call(
+            self,
+            call,
+            allow_partial: bool = False,
+            wrap_validation_errors: bool = True,
+            *,
+            approved: bool = False,
+            metadata: Any = None,
+        ):
+            _normalize_call_tool_name(call)
+            return await _original_handle_call(
+                self,
+                call,
+                allow_partial=allow_partial,
+                wrap_validation_errors=wrap_validation_errors,
+                approved=approved,
+                metadata=metadata,
+            )
+
+        # -- _call_tool wrapper with callbacks -----------------------------------
+
+        async def _patched_call_tool(
+            self,
+            call,
+            *,
+            allow_partial: bool,
+            wrap_validation_errors: bool,
+            approved: bool,
+            metadata: Any = None,
+        ):
+            tool_name, call = _normalize_call_tool_name(call)
+
+            # 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_muse import callbacks
+                from code_muse.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_muse import callbacks
+
+                    await callbacks.on_post_tool_call(
+                        tool_name, tool_args, final_result, duration_ms
+                    )
+                except Exception:
+                    pass  # never block tool execution
+
+        ToolManager.get_tool_def = _patched_get_tool_def
+        ToolManager.handle_call = _patched_handle_call
+        ToolManager._call_tool = _patched_call_tool
+
+    except ImportError:
+        pass
+    except Exception:
+        pass
+
+
+def patch_prompt_toolkit_emoji_width() -> None:
+    """Patch prompt_toolkit's character width calculation for emojis.
+
+    Modern terminals render most emojis as 2 cells wide, but wcwidth often
+    returns 1 for many emoji codepoints. This causes cursor misalignment.
+
+    This patch:
+    1. Returns 0 for variation selectors (zero-width modifiers)
+    2. Returns 2 for emoji codepoints (terminals render them wide)
+    3. Falls back to wcwidth for non-emoji characters
+    """
+    try:
+        import wcwidth
+        from prompt_toolkit import utils as pt_utils
+
+        _original_get_cwidth = pt_utils.get_cwidth
+
+        def _patched_get_cwidth(char: str) -> int:
+            """Get character width with better emoji support."""
+            code = ord(char)
+
+            # Variation selectors are zero-width
+            if 0xFE00 <= code <= 0xFE0F:  # VS1-VS16
+                return 0
+
+            # Emoji codepoints - terminals render these as 2 cells wide
+            # even when wcwidth says 1
+            if (
+                0x1F300 <= code <= 0x1F9FF  # Misc Symbols/Pictographs, Emoticons
+                or 0x1F600 <= code <= 0x1F64F  # Emoticons
+                or 0x1F680 <= code <= 0x1F6FF  # Transport/Map symbols
+                or 0x1FA00 <= code <= 0x1FAFF  # Symbols/Pictographs Extended-A
+                or 0x2600 <= code <= 0x26FF  # Misc Symbols (☀️, ⚡, etc)
+                or 0x2700 <= code <= 0x27BF  # Dingbats (✂️, ✈️, etc)
+                or 0x1F1E0 <= code <= 0x1F1FF  # Regional indicators (flags)
+            ):
+                return 2
+
+            # Use wcwidth for non-emoji
+            w = wcwidth.wcwidth(char)
+            if w >= 0:
+                return w
+
+            return _original_get_cwidth(char)
+
+        pt_utils.get_cwidth = _patched_get_cwidth
+
+    except ImportError:
+        pass  # wcwidth or prompt_toolkit not available
+    except Exception:
+        pass  # Don't crash on patch failure
+
+
+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()
+    patch_prompt_toolkit_emoji_width()

code_muse/reopenable_async_client.py

@@ -0,0 +1,233 @@
+"""
+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
+
+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: 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: httpx.AsyncClient | None = None
+        self._is_closed = True
+        self._lock = asyncio.Lock()
+        # FREE-THREADED: _sync_lock guards build_request() which is a sync method
+        # that creates a temporary sync httpx.Client. It must remain threading.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: 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: 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: 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: 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: 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: 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: 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: 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: 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: 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) -> httpx.Timeout | None:
+        """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}]>"