code-muse 0.0.1__py3-none-any.whl

code_muse/stream_parser/utf8_stream_parser.pyx

@@ -0,0 +1,206 @@
+# cython: language_level=3
+"""UTF-8 byte-stream adapter that wraps a StreamTextParser.
+
+Buffers partial code points across chunk boundaries and rolls back entire
+chunks on invalid UTF-8 so the wrapped parser never sees malformed text.
+"""
+
+from typing import TypeVar
+
+from code_muse.stream_parser.stream_text_chunk import StreamTextChunk
+from code_muse.stream_parser.stream_text_parser import StreamTextParser
+
+T = TypeVar("T")
+
+
+class Utf8StreamParserError(Exception):
+    """Base error for Utf8StreamParser."""
+
+    pass
+
+
+class InvalidUtf8(Utf8StreamParserError):
+    """Raised when a pushed byte chunk contains an invalid UTF-8 sequence.
+
+    The chunk is rolled back; the parser's pending buffer is restored to its
+    state before the chunk was pushed.
+    """
+
+    def __init__(self, valid_up_to: int, error_len: int) -> None:
+        self.valid_up_to = valid_up_to
+        self.error_len = error_len
+        super().__init__(
+            f"invalid UTF-8 in streamed bytes at offset {valid_up_to} "
+            f"(error length {error_len})"
+        )
+
+
+class IncompleteUtf8AtEof(Utf8StreamParserError):
+    """Raised at finish() or into_inner() when an incomplete code point remains."""
+
+    def __init__(self) -> None:
+        super().__init__("incomplete UTF-8 code point at end of stream")
+
+
+class Utf8StreamParser:
+    """Wraps a StreamTextParser and accepts raw bytes.
+
+    * Buffers incomplete UTF-8 code points across chunk boundaries.
+    * On invalid UTF-8, rolls back the **entire** chunk and raises.
+    * On incomplete-but-valid UTF-8, processes the valid prefix and keeps
+      the incomplete tail buffered.
+    * ``finish()`` flushes pending bytes; raises if a code point is incomplete.
+    * ``into_inner()`` returns the wrapped parser only when no bytes are pending.
+    * ``into_inner_lossy()`` returns the wrapped parser, dropping pending bytes.
+
+    The wrapped parser never receives malformed text.
+    """
+
+    def __init__(self, inner: StreamTextParser[T]) -> None:
+        self.inner = inner
+        self._pending_utf8: bytearray = bytearray()
+
+    def push_bytes(self, const unsigned char[:] chunk) -> StreamTextChunk[T]:
+        """Feed a new byte chunk.
+
+        Args:
+            chunk: Raw bytes to decode and forward to the inner parser.
+
+        Returns:
+            Visible text and extracted payloads from the inner parser.
+
+        Raises:
+            InvalidUtf8: If the chunk (combined with any previously buffered
+                partial code point) contains an invalid UTF-8 sequence. The
+                entire chunk is rolled back so the buffer is restored to its
+                pre-call state.
+        """
+        cdef int old_len
+        cdef int valid_up_to
+        cdef bint is_incomplete
+
+        # Fast path: no pending bytes and chunk decodes cleanly.
+        if not self._pending_utf8:
+            try:
+                text = bytes(chunk).decode("utf-8")
+                return self.inner.push_str(text)
+            except UnicodeDecodeError as err:
+                is_incomplete = (
+                    err.reason == "unexpected end of data" and err.end == len(chunk)
+                )
+                if not is_incomplete:
+                    raise InvalidUtf8(err.start, err.end - err.start) from err
+                # Buffer the chunk; valid prefix will be consumed below.
+                self._pending_utf8.extend(chunk)
+                valid_up_to = err.start
+                if valid_up_to == 0:
+                    return StreamTextChunk()
+                text = bytes(chunk[:valid_up_to]).decode("utf-8")
+                out = self.inner.push_str(text)
+                del self._pending_utf8[:valid_up_to]
+                return out
+
+        # Slow path: append to existing pending buffer.
+        old_len = len(self._pending_utf8)
+        self._pending_utf8.extend(chunk)
+
+        try:
+            text = self._pending_utf8.decode("utf-8")
+        except UnicodeDecodeError as err:
+            # Distinguish "invalid byte sequence" from "incomplete at end".
+            # In Python an incomplete sequence at the very end reports
+            # ``reason == "unexpected end of data"`` and ``err.end`` equals
+            # ``len(self._pending_utf8)``.  Anything else is an actual error.
+            is_incomplete = err.reason == "unexpected end of data" and err.end == len(
+                self._pending_utf8
+            )
+
+            if not is_incomplete:
+                # Invalid sequence somewhere in the buffer.  Roll back the
+                # entire chunk so the inner parser never sees malformed data.
+                del self._pending_utf8[old_len:]
+                raise InvalidUtf8(err.start, err.end - err.start) from err
+
+            # Incomplete code point at the end of the buffer.
+            valid_up_to = err.start
+            if valid_up_to == 0:
+                # Nothing valid to forward yet.
+                return StreamTextChunk()
+
+            # Process the valid prefix.  Defend against the edge case where
+            # the prefix itself fails to decode (nested error).
+            try:
+                text = bytes(self._pending_utf8[:valid_up_to]).decode("utf-8")
+            except UnicodeDecodeError as nested_err:
+                del self._pending_utf8[old_len:]
+                raise InvalidUtf8(
+                    nested_err.start, nested_err.end - nested_err.start
+                ) from nested_err
+
+            out = self.inner.push_str(text)
+            del self._pending_utf8[:valid_up_to]
+            return out
+
+        # Full buffer decoded cleanly.
+        out = self.inner.push_str(text)
+        self._pending_utf8.clear()
+        return out
+
+    def finish(self) -> StreamTextChunk[T]:
+        """Flush any buffered bytes at end-of-stream.
+
+        Returns:
+            Any visible text and extracted payloads produced by the inner
+            parser for the final pending bytes.
+
+        Raises:
+            IncompleteUtf8AtEof: If an incomplete UTF-8 code point remains
+                in the buffer.  The bytes are **not** consumed; they stay
+                buffered so the caller can decide what to do next.
+        """
+        if not self._pending_utf8:
+            return self.inner.finish()
+
+        # Attempt to decode whatever is left.  If it is incomplete, raise.
+        try:
+            text = self._pending_utf8.decode("utf-8")
+        except UnicodeDecodeError as err:
+            if err.reason == "unexpected end of data" and err.end == len(
+                self._pending_utf8
+            ):
+                raise IncompleteUtf8AtEof() from err
+            # Should be unreachable for a previously-validated buffer, but
+            # treat it as invalid UTF-8.
+            raise InvalidUtf8(err.start, err.end - err.start) from err
+
+        out = self.inner.push_str(text)
+        self._pending_utf8.clear()
+        tail = self.inner.finish()
+        return StreamTextChunk(
+            visible_text=out.visible_text + tail.visible_text,
+            extracted=out.extracted + tail.extracted,
+        )
+
+    def into_inner(self) -> StreamTextParser[T]:
+        """Return the wrapped parser only if no bytes are pending.
+
+        Returns:
+            The inner StreamTextParser instance.
+
+        Raises:
+            IncompleteUtf8AtEof: If there are pending bytes (incomplete code
+                point) in the buffer.
+        """
+        if self._pending_utf8:
+            raise IncompleteUtf8AtEof()
+        return self.inner
+
+    def into_inner_lossy(self) -> StreamTextParser[T]:
+        """Return the wrapped parser, discarding any pending bytes.
+
+        Returns:
+            The inner StreamTextParser instance.  Any incomplete UTF-8 code
+            point buffered across previous chunks is silently dropped.
+        """
+        self._pending_utf8.clear()
+        return self.inner

code_muse/summarization_agent.py

@@ -0,0 +1,308 @@
+import asyncio
+import atexit
+import hashlib
+import logging
+import pathlib
+import threading
+
+# FREE-THREADED: ThreadPoolExecutor is compatible with free-threaded Python 3.14 —
+# no GIL contention for I/O-bound summarization work.
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+
+from pydantic_ai import Agent
+
+from code_muse.config import get_summarization_model_name
+from code_muse.model_factory import ModelFactory, make_model_settings
+
+logger = logging.getLogger(__name__)
+
+# Keep a module-level agent reference to avoid rebuilding per call
+_summarization_agent = None
+# FREE-THREADED: _agent_lock guards sync-only agent cache access.
+_agent_lock = threading.Lock()
+
+# P2-05/PERF-05: track the model name the cached agent was built for
+_cached_model_name: str | None = None
+
+# Safe sync runner for async agent.run calls
+# Avoids "event loop is already running" by offloading to a separate thread loop when needed
+_thread_pool: ThreadPoolExecutor | None = None
+
+# Reload counter
+_reload_count = 0
+
+
+# ---------------------------------------------------------------------------
+# P2-05/PERF-05: Model config cache with mtime invalidation
+# ---------------------------------------------------------------------------
+
+
+def _models_config_fingerprint() -> tuple[float, str]:
+    """Compute a lightweight fingerprint of all model config sources.
+
+    Returns (max_mtime, content_hash) — if either changes, the cached
+    config is stale and must be reloaded.
+    """
+    source_paths: list[pathlib.Path] = []
+
+    # Bundled models.json is always loaded
+    bundled = pathlib.Path(__file__).parent / "models.json"
+    source_paths.append(bundled)
+
+    # Extra model sources (mirrors ModelFactory.load_config)
+    try:
+        from code_muse.config import (
+            CHATGPT_MODELS_FILE,
+            CLAUDE_MODELS_FILE,
+            COPILOT_MODELS_FILE,
+            EXTRA_MODELS_FILE,
+            GEMINI_MODELS_FILE,
+            MODELS_FILE,
+        )
+
+        for p in (
+            MODELS_FILE,
+            EXTRA_MODELS_FILE,
+            CHATGPT_MODELS_FILE,
+            CLAUDE_MODELS_FILE,
+            GEMINI_MODELS_FILE,
+            COPILOT_MODELS_FILE,
+        ):
+            source_paths.append(pathlib.Path(p))
+    except Exception:
+        pass
+
+    max_mtime = 0.0
+    hasher = hashlib.md5(usedforsecurity=False)
+    for sp in source_paths:
+        try:
+            if sp.exists():
+                stat = sp.stat()
+                max_mtime = max(max_mtime, stat.st_mtime)
+                # Hash file contents (or just size+mtime as a cheap proxy)
+                hasher.update(f"{sp}:{stat.st_size}:{stat.st_mtime}".encode())
+        except OSError:
+            pass
+
+    return max_mtime, hasher.hexdigest()
+
+
+# Module-level model config cache: (config_dict, fingerprint)
+_models_config_cache: tuple[dict[str, Any | None], tuple[float, str | None]] = (
+    None,
+    None,
+)
+# FREE-THREADED: _models_config_lock guards sync-only config cache access.
+_models_config_lock = threading.Lock()
+
+
+def get_cached_models_config() -> dict[str, Any]:
+    """Return the models config, using a cache invalidated by mtime/hash changes.
+
+    This avoids re-reading ``models.json`` and extra model files on every call
+    to ``ModelFactory.load_config()`` when nothing has changed. The cache is
+    invalidated when any source file's mtime changes.
+
+    Falls back to ``ModelFactory.load_config()`` on any error.
+    """
+    global _models_config_cache
+
+    fingerprint = _models_config_fingerprint()
+
+    with _models_config_lock:
+        cached_config, cached_fp = _models_config_cache
+        if cached_config is not None and cached_fp == fingerprint:
+            return cached_config
+
+        # Cache miss — reload. Let exceptions propagate so callers
+        # (including reload_summarization_agent) see the same errors they
+        # would have seen without the cache.
+        config = ModelFactory.load_config()
+        _models_config_cache = (config, fingerprint)
+        return config
+
+
+def invalidate_models_config_cache() -> None:
+    """Force the next ``get_cached_models_config()`` call to reload.
+
+    Call this when settings or model files are known to have changed
+    (e.g. after a ``/set`` command that modifies model config).
+    """
+    global _models_config_cache
+    with _models_config_lock:
+        _models_config_cache = (None, None)
+
+
+def _ensure_thread_pool():
+    global _thread_pool
+    # Check if pool is None OR if it's been shutdown
+    if _thread_pool is None or _thread_pool._shutdown:
+        # FREE-THREADED: ThreadPoolExecutor is compatible with free-threaded Python 3.14.
+        _thread_pool = ThreadPoolExecutor(
+            max_workers=1, thread_name_prefix="summarizer-loop"
+        )
+    return _thread_pool
+
+
+def _shutdown_thread_pool():
+    global _thread_pool
+    if _thread_pool is not None:
+        _thread_pool.shutdown(wait=False)
+        _thread_pool = None
+
+
+atexit.register(_shutdown_thread_pool)
+
+
+async def _run_agent_async(agent: Agent, prompt: str, message_history: list):
+    return await agent.run(prompt, message_history=message_history)
+
+
+class SummarizationError(Exception):
+    """Raised when summarization fails with details about the failure."""
+
+    def __init__(self, message: str, original_error: Exception | None = None):
+        self.original_error = original_error
+        super().__init__(message)
+
+
+def run_summarization_sync(prompt: str, message_history: list) -> list:
+    """Run the summarization agent synchronously.
+
+    Raises:
+        SummarizationError: If summarization fails for any reason.
+    """
+    try:
+        agent = get_summarization_agent()
+    except Exception as e:
+        raise SummarizationError(
+            f"Failed to initialize summarization agent: {type(e).__name__}: {e}",
+            original_error=e,
+        ) from e
+
+    # Handle claude-code models: prepend system prompt to user prompt
+    from code_muse.model_utils import prepare_prompt_for_model
+
+    model_name = get_summarization_model_name()
+    prepared = prepare_prompt_for_model(
+        model_name, _get_summarization_instructions(), prompt
+    )
+    prompt = prepared.user_prompt
+
+    def _run_in_thread():
+        """
+        Run the async agent in a dedicated thread with its own event loop.
+        Uses run_until_complete instead of asyncio.run to avoid shutting down
+        the default executor (which may break plugins in the main thread).
+        Does NOT touch global event loop state.
+        """
+        loop = asyncio.new_event_loop()
+        try:
+            coro = agent.run(prompt, message_history=message_history)
+            return loop.run_until_complete(coro)
+        finally:
+            # Clean up without shutting down the default executor
+            try:
+                # Cancel pending tasks
+                pending = asyncio.all_tasks(loop)
+                for task in pending:
+                    task.cancel()
+                if pending:
+                    # NOTE: asyncio.gather for already-running tasks; TaskGroup cannot adopt existing tasks.
+                    loop.run_until_complete(
+                        asyncio.gather(*pending, return_exceptions=True)
+                    )
+                loop.run_until_complete(loop.shutdown_asyncgens())
+            finally:
+                loop.close()
+
+    try:
+        # Always use thread pool since we're likely in an existing event loop
+        pool = _ensure_thread_pool()
+        result = pool.submit(_run_in_thread).result()
+        return result.new_messages()
+    except Exception as e:
+        error_type = type(e).__name__
+        error_msg = str(e) if str(e) else "(no details available)"
+        raise SummarizationError(
+            f"LLM call failed during summarization: [{error_type}] {error_msg}",
+            original_error=e,
+        ) from e
+
+
+def _get_summarization_instructions() -> str:
+    """Get the system instructions for the summarization agent."""
+    return """You are a message summarization expert. Your task is to summarize conversation messages
+while preserving important context and information. The summaries should be concise but capture the essential content
+and intent of the original messages. This is to help manage token usage in a conversation history
+while maintaining context for the AI to continue the conversation effectively.
+
+When summarizing:
+1. Keep summary concise but informative
+2. Preserve important context and key information and decisions
+3. Keep any important technical details
+4. Don't summarize the system message
+5. Make sure all tool calls and responses are summarized, as they are vital
+6. Focus on token usage efficiency and system message preservation"""
+
+
+def reload_summarization_agent():
+    """Create a specialized agent for summarizing messages when context limit is reached."""
+    from code_muse.model_utils import prepare_prompt_for_model
+
+    # Always bust the cache on explicit reload — the caller expects fresh config
+    invalidate_models_config_cache()
+    models_config = get_cached_models_config()
+    model_name = get_summarization_model_name()
+    model = ModelFactory.get_model(model_name, models_config)
+
+    # Handle claude-code models: swap instructions (prompt prepending happens in run_summarization_sync)
+    instructions = _get_summarization_instructions()
+    prepared = prepare_prompt_for_model(
+        model_name, instructions, "", prepend_system_to_user=False
+    )
+    instructions = prepared.instructions
+
+    model_settings = make_model_settings(model_name)
+
+    agent = Agent(
+        model=model,
+        instructions=instructions,
+        output_type=str,
+        retries=1,  # Fewer retries for summarization
+        model_settings=model_settings,
+    )
+    # NOTE: We intentionally don't wrap the summarization agent.
+    # Summarization is a simple one-shot call that doesn't need durable execution,
+    # and wrapping can cause async event loop conflicts with run_sync().
+    return agent
+
+
+def get_summarization_agent(force_reload=False):
+    """Retrieve the summarization agent, caching across calls.
+
+    P2-05/PERF-05: The default is now ``force_reload=False``. The agent is
+    rebuilt only when:
+    - ``force_reload=True`` is explicitly passed
+    - The summarization model name has changed since the last build
+    - No agent has been built yet (first call)
+
+    Args:
+        force_reload: When True, unconditionally rebuild the agent.
+
+    Returns:
+        A ``pydantic_ai.Agent`` configured for summarization.
+    """
+    global _summarization_agent, _cached_model_name
+    current_model = get_summarization_model_name()
+    with _agent_lock:
+        needs_reload = (
+            force_reload
+            or _summarization_agent is None
+            or _cached_model_name != current_model
+        )
+        if needs_reload:
+            _summarization_agent = reload_summarization_agent()
+            _cached_model_name = current_model
+        return _summarization_agent