code-muse 0.0.1__py3-none-any.whl

code_muse/agents/_history.py

@@ -0,0 +1,382 @@
+"""Pure helpers for message history hashing, token estimation, and pruning.
+
+Extracted from the original ``BaseAgent`` god-class. Everything in here is a
+free function with no hidden state. Call sites pass messages (and, where
+needed, already-resolved strings / tool dicts) in explicitly.
+
+PERF-04 adds :class:`CompactionCache` — a per-compaction-run cache that avoids
+repeated ``hash_message`` / ``estimate_tokens_for_message`` calls on the same
+message objects within a single ``compact()`` invocation. No global caches
+that retain message objects.
+"""
+
+import dataclasses
+import json
+import math
+import pathlib
+import weakref
+from annotationlib import get_annotations
+from typing import Any
+
+import pydantic
+from pydantic_ai import BinaryContent
+from pydantic_ai.messages import ModelMessage
+
+
+def _json_safe(obj):
+    """JSON serializer for objects not supported by default json.dumps."""
+    if isinstance(obj, pathlib.PurePath):
+        return str(obj)
+    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
+
+
+def stringify_part(part: Any) -> str:
+    """Return a stable, timestamp-free string representation of a message part.
+
+    Used for both hashing and token estimation. Ignoring timestamps means two
+    otherwise-identical parts emitted at different times collapse to the same
+    string, which is exactly what we want for dedup.
+    """
+    attributes: list[str] = [part.__class__.__name__]
+
+    if hasattr(part, "role") and part.role:
+        attributes.append(f"role={part.role}")
+    if hasattr(part, "instructions") and part.instructions:
+        attributes.append(f"instructions={part.instructions}")
+
+    if hasattr(part, "tool_call_id") and part.tool_call_id:
+        attributes.append(f"tool_call_id={part.tool_call_id}")
+    if hasattr(part, "tool_name") and part.tool_name:
+        attributes.append(f"tool_name={part.tool_name}")
+
+    content = getattr(part, "content", None)
+    if content is None:
+        attributes.append("content=None")
+    elif isinstance(content, str):
+        attributes.append(f"content={content}")
+    elif isinstance(content, pydantic.BaseModel):
+        dumped = json.dumps(content.model_dump(), sort_keys=True, default=_json_safe)
+        attributes.append(f"content={dumped}")
+    elif isinstance(content, dict):
+        dumped = json.dumps(content, sort_keys=True, default=_json_safe)
+        attributes.append(f"content={dumped}")
+    elif isinstance(content, list):
+        for item in content:
+            if isinstance(item, str):
+                attributes.append(f"content={item}")
+            elif isinstance(item, BinaryContent):
+                attributes.append(f"BinaryContent={hash(item.data)}")
+    else:
+        attributes.append(f"content={repr(content)}")
+
+    return "|".join(attributes)
+
+
+# PERF-07: Cache hash results by object identity.  A message that appears
+# in multiple history lists (e.g. after copy-on-write) will hash
+# identically, so id()-keyed caching is correct as long as the message
+# object is alive — which it always is while it's in a history list.
+#
+# A ``weakref.finalize`` callback automatically drops the cached entry
+# when the message object is garbage-collected, preventing stale hits if the
+# same memory address is reused for a different message later (common in
+# tests that create many short-lived messages).
+_hash_cache: dict[int, int] = {}
+_HASH_CACHE_MAX = 8192
+
+
+def _evict_hash_cache(msg_id: int) -> None:
+    _hash_cache.pop(msg_id, None)
+
+
+def hash_message(message: Any) -> int:
+    """Stable hash for a ``ModelMessage`` that ignores timestamps."""
+    msg_id = id(message)
+    cached = _hash_cache.get(msg_id)
+    if cached is not None:
+        return cached
+
+    role = getattr(message, "role", None)
+    instructions = getattr(message, "instructions", None)
+    header_bits: list[str] = []
+    if role:
+        header_bits.append(f"role={role}")
+    if instructions:
+        header_bits.append(f"instructions={instructions}")
+
+    part_strings = [stringify_part(part) for part in getattr(message, "parts", [])]
+    canonical = "||".join(header_bits + part_strings)
+    result = hash(canonical)
+
+    # Bounded cache — evict oldest entries when full
+    if len(_hash_cache) >= _HASH_CACHE_MAX:
+        _hash_cache.clear()
+    _hash_cache[msg_id] = result
+    weakref.finalize(message, _evict_hash_cache, msg_id)
+    return result
+
+
+def estimate_tokens(text: str) -> int:
+    """Character-based token estimator tuned for code-heavy content.
+
+    Uses ``len / 3.0`` as the base divisor (code has higher entropy than
+    prose, and the old ``/2.5`` heuristic systematically undercounted for
+    most models).  Per-model calibration multipliers applied downstream
+    by ``estimate_tokens_for_message`` adjust for models with even denser
+    tokenizers.
+    """
+    return max(1, math.floor(len(text) / 3.0))
+
+
+# Models whose tokenizer the char/2.5 heuristic systematically *under*counts.
+# Bump these by a calibration factor so context-usage math stops lying to us.
+# Substring match is case-insensitive; both naming orders are accepted because
+# vendor naming is a coin flip.
+_TOKEN_MULTIPLIER_RULES: tuple[tuple[tuple[str, ...], float], ...] = (
+    (("opus-4-7", "4-7-opus"), 1.35),
+)
+
+
+def model_token_multiplier(model_name: str | None) -> float:
+    """Per-model fudge factor for our char-based token estimator.
+
+    Returns 1.0 when ``model_name`` is falsy or doesn't match any rule.
+    """
+    if not model_name:
+        return 1.0
+    lowered = model_name.lower()
+    for needles, factor in _TOKEN_MULTIPLIER_RULES:
+        if any(needle in lowered for needle in needles):
+            return factor
+    return 1.0
+
+
+def _apply_multiplier(raw_tokens: int, model_name: str | None) -> int:
+    multiplier = model_token_multiplier(model_name)
+    if multiplier == 1.0:
+        return raw_tokens
+    return max(1, math.floor(raw_tokens * multiplier))
+
+
+def estimate_tokens_for_message(
+    message: ModelMessage,
+    model_name: str | None = None,
+) -> int:
+    """Estimate the number of tokens in a single model message.
+
+    When ``model_name`` is provided, the raw count is scaled by
+    :func:`model_token_multiplier` to compensate for tokenizers that don't
+    play nicely with our char/2.5 heuristic.
+    """
+    total = 0
+    for part in getattr(message, "parts", []) or []:
+        part_str = stringify_part(part)
+        if part_str:
+            total += estimate_tokens(part_str)
+    return _apply_multiplier(max(1, total), model_name)
+
+
+def estimate_context_overhead(
+    system_prompt: str,
+    pydantic_tools: dict[str, Any | None],
+    model_name: str | None = None,
+) -> int:
+    """Estimate fixed token overhead for the system prompt + tool definitions.
+
+    The caller is responsible for resolving the system prompt for the active
+    model (e.g. via ``prepare_prompt_for_model``). external tool overhead is
+    deliberately ignored — it was guesswork anyway.
+
+    Args:
+        system_prompt: The already-resolved instruction/system prompt string.
+        pydantic_tools: The pydantic-ai agent's ``_tools`` dict, or ``None``.
+
+    Returns:
+        Estimated total token overhead.
+    """
+    total = 0
+    if system_prompt:
+        total += estimate_tokens(system_prompt)
+
+    if not pydantic_tools:
+        return _apply_multiplier(total, model_name)
+
+    for tool_name, tool_func in pydantic_tools.items():
+        total += estimate_tokens(tool_name)
+
+        description = getattr(tool_func, "__doc__", None) or ""
+        if description:
+            total += estimate_tokens(description)
+
+        schema = getattr(tool_func, "schema", None)
+        if schema is not None:
+            schema_str = json.dumps(schema) if isinstance(schema, dict) else str(schema)
+            total += estimate_tokens(schema_str)
+        else:
+            annotations = get_annotations(tool_func)
+            if annotations:
+                total += estimate_tokens(str(annotations))
+
+    return _apply_multiplier(total, model_name)
+
+
+# Pydantic-AI has FOUR part kinds that carry a tool_call_id:
+#   * tool-call            -> ToolCallPart            (regular tool call)
+#   * tool-return          -> ToolReturnPart          (regular tool response)
+#   * builtin-tool-call    -> BuiltinToolCallPart     (claude extended-thinking, etc.)
+#   * builtin-tool-return  -> BuiltinToolReturnPart   (builtin tool response)
+#   * retry-prompt         -> RetryPromptPart
+#         (assistant told to retry; acts as a response)
+#
+# Treating only `tool-call` / `tool-return` (and ignoring the others) caused
+# subtle bugs: e.g. builtin tool calls on Claude Opus were counted as pending
+# forever, deferring summarization on every turn.
+_TOOL_CALL_PART_KINDS: frozenset[str] = frozenset({"tool-call", "builtin-tool-call"})
+_TOOL_RETURN_PART_KINDS: frozenset[str] = frozenset(
+    {"tool-return", "builtin-tool-return", "retry-prompt"}
+)
+
+
+def _classify_tool_part(part: object) -> str | None:
+    """Return ``"call"``, ``"return"``, or ``None`` for a message part.
+
+    ``None`` means the part doesn't participate in tool_call_id pairing
+    (either no id, or an unrelated part kind).
+    """
+    if getattr(part, "tool_call_id", None) is None:
+        return None
+    pk = getattr(part, "part_kind", None)
+    if pk in _TOOL_CALL_PART_KINDS:
+        return "call"
+    if pk in _TOOL_RETURN_PART_KINDS:
+        return "return"
+    return None
+
+
+def prune_interrupted_tool_calls(
+    messages: list[ModelMessage],
+) -> list[ModelMessage]:
+    """Drop messages participating in mismatched tool_call/tool_return pairs.
+
+    A mismatched ``tool_call_id`` is one that appears only as a call or only
+    as a return. The model will reject such histories ("tool_use ids found
+    without tool_result blocks"), so we strip them out while preserving order.
+    """
+    if not messages:
+        return messages
+
+    tool_call_ids: set[str] = set()
+    tool_return_ids: set[str] = set()
+
+    for msg in messages:
+        for part in getattr(msg, "parts", []) or []:
+            kind = _classify_tool_part(part)
+            if kind == "call":
+                tool_call_ids.add(part.tool_call_id)
+            elif kind == "return":
+                tool_return_ids.add(part.tool_call_id)
+
+    mismatched = tool_call_ids.symmetric_difference(tool_return_ids)
+    if not mismatched:
+        return messages
+
+    pruned: list[ModelMessage] = []
+    for msg in messages:
+        if any(
+            getattr(part, "tool_call_id", None) in mismatched
+            for part in getattr(msg, "parts", []) or []
+        ):
+            continue
+        pruned.append(msg)
+    return pruned
+
+
+def has_pending_tool_calls(messages: list[ModelMessage]) -> bool:
+    """Return True if any tool call is still waiting for its response.
+
+    Recognizes both regular (``tool-call`` / ``tool-return``) and builtin
+    (``builtin-tool-call`` / ``builtin-tool-return``) pairings, plus
+    ``retry-prompt`` as a valid response form.
+    """
+    if not messages:
+        return False
+
+    tool_call_ids: set[str] = set()
+    tool_return_ids: set[str] = set()
+
+    for msg in messages:
+        for part in getattr(msg, "parts", []) or []:
+            kind = _classify_tool_part(part)
+            if kind == "call":
+                tool_call_ids.add(part.tool_call_id)
+            elif kind == "return":
+                tool_return_ids.add(part.tool_call_id)
+
+    return bool(tool_call_ids - tool_return_ids)
+
+
+def filter_huge_messages(
+    messages: list[ModelMessage],
+    model_name: str | None = None,
+    cache: Any | None = None,  # CompactionCache when available
+) -> list[ModelMessage]:
+    """Drop individual messages above a 50k-token budget, then prune orphans."""
+    filtered = [
+        m
+        for m in messages
+        if (
+            cache.estimate_tokens(m, model_name)
+            if cache
+            else estimate_tokens_for_message(m, model_name)
+        )
+        < 50000
+    ]
+    return prune_interrupted_tool_calls(filtered)
+
+
+# ---------------------------------------------------------------------------
+# PERF-04: Per-compaction-run cache
+# ---------------------------------------------------------------------------
+
+
+@dataclasses.dataclass
+class CompactionCache:
+    """Per-compaction-run cache for message hashes and token estimates.
+
+    Created fresh at the start of each ``compact()`` call. Avoids repeated
+    ``hash_message()`` and ``estimate_tokens_for_message()`` invocations on
+    the same ``ModelMessage`` objects within a single compaction cycle.
+
+    Importantly, this is *not* a global cache — it is scoped to one
+    ``compact()`` invocation and dropped afterwards so we never retain
+    message objects beyond their natural lifecycle.
+    """
+
+    # message id(id(m)) → int hash
+    _message_hashes: dict[int, int] = dataclasses.field(default_factory=dict)
+    # (id(m), model_name_or_None) → int token count
+    _token_counts: dict[tuple[int, str | None], int] = dataclasses.field(
+        default_factory=dict
+    )
+
+    def hash_message(self, message: ModelMessage) -> int:
+        """Cached wrapper around :func:`hash_message` (module-level)."""
+        key = id(message)
+        if key not in self._message_hashes:
+            self._message_hashes[key] = hash_message(message)
+        return self._message_hashes[key]
+
+    def estimate_tokens(
+        self, message: ModelMessage, model_name: str | None = None
+    ) -> int:
+        """Cached wrapper around :func:`estimate_tokens_for_message`."""
+        key = (id(message), model_name)
+        if key not in self._token_counts:
+            self._token_counts[key] = estimate_tokens_for_message(message, model_name)
+        return self._token_counts[key]
+
+    def sum_tokens(
+        self, messages: list[ModelMessage], model_name: str | None = None
+    ) -> int:
+        """Sum estimated tokens across all messages using the cache."""
+        return sum(self.estimate_tokens(m, model_name) for m in messages)

code_muse/agents/_key_listeners.py

@@ -0,0 +1,148 @@
+"""Keyboard listener thread helpers, extracted from ``BaseAgent``.
+
+These functions listen for Ctrl+X (shell cancel) and optionally the configured
+cancel-agent key (when it's not bound to a signal like SIGINT). Previously
+they lived as methods on ``BaseAgent`` but they never touched ``self``, so
+they're free functions now.
+"""
+
+import threading
+from collections.abc import Callable
+
+from code_muse.keymap import cancel_agent_uses_signal, get_cancel_agent_char_code
+from code_muse.messaging import emit_warning
+
+
+def spawn_key_listener(
+    stop_event: threading.Event,
+    on_escape: Callable[[], None],
+    on_cancel_agent: Callable[[], None] | None = None,
+) -> threading.Thread | None:
+    """Start a daemon thread that listens for Ctrl+X / cancel-agent keys.
+
+    Args:
+        stop_event: Signal the listener to stop.
+        on_escape: Callback for Ctrl+X (shell command cancel).
+        on_cancel_agent: Optional callback for the configured cancel-agent
+            key, only used when ``cancel_agent_uses_signal()`` is False.
+
+    Returns:
+        The started Thread, or None if stdin isn't a TTY (or otherwise
+        unusable, e.g. during tests).
+    """
+    try:
+        import sys
+    except ImportError:
+        return None
+
+    stdin = getattr(sys, "stdin", None)
+    if stdin is None or not hasattr(stdin, "isatty"):
+        return None
+    try:
+        if not stdin.isatty():
+            return None
+    except Exception:
+        return None
+
+    def listener() -> None:
+        try:
+            if sys.platform.startswith("win"):
+                _listen_windows(stop_event, on_escape, on_cancel_agent)
+            else:
+                _listen_posix(stop_event, on_escape, on_cancel_agent)
+        except Exception:
+            emit_warning("Key listener stopped unexpectedly; press Ctrl+C to cancel.")
+
+    thread = threading.Thread(
+        target=listener, name="muse-key-listener", daemon=True
+    )
+    thread.start()
+    return thread
+
+
+def _listen_windows(
+    stop_event: threading.Event,
+    on_escape: Callable[[], None],
+    on_cancel_agent: Callable[[], None] | None = None,
+) -> None:
+    import msvcrt
+    import time
+
+    cancel_agent_char: str | None = None
+    if on_cancel_agent is not None and not cancel_agent_uses_signal():
+        cancel_agent_char = get_cancel_agent_char_code()
+
+    while not stop_event.is_set():
+        try:
+            if msvcrt.kbhit():
+                key = msvcrt.getwch()
+                if key == "\x18":  # Ctrl+X
+                    try:
+                        on_escape()
+                    except Exception:
+                        emit_warning(
+                            "Ctrl+X handler raised unexpectedly; Ctrl+C still works."
+                        )
+                elif cancel_agent_char and on_cancel_agent and key == cancel_agent_char:
+                    try:
+                        on_cancel_agent()
+                    except Exception:
+                        emit_warning("Cancel agent handler raised unexpectedly.")
+        except Exception:
+            emit_warning(
+                "Windows key listener error; Ctrl+C is still available for cancel."
+            )
+            return
+        time.sleep(0.05)
+
+
+def _listen_posix(
+    stop_event: threading.Event,
+    on_escape: Callable[[], None],
+    on_cancel_agent: Callable[[], None] | None = None,
+) -> None:
+    import select
+    import sys
+    import termios
+    import tty
+
+    cancel_agent_char: str | None = None
+    if on_cancel_agent is not None and not cancel_agent_uses_signal():
+        cancel_agent_char = get_cancel_agent_char_code()
+
+    stdin = sys.stdin
+    try:
+        fd = stdin.fileno()
+    except AttributeError, ValueError, OSError:
+        return
+    try:
+        original_attrs = termios.tcgetattr(fd)
+    except Exception:
+        return
+
+    try:
+        tty.setcbreak(fd)
+        while not stop_event.is_set():
+            try:
+                read_ready, _, _ = select.select([stdin], [], [], 0.05)
+            except Exception:
+                break
+            if not read_ready:
+                continue
+            data = stdin.read(1)
+            if not data:
+                break
+            if data == "\x18":  # Ctrl+X
+                try:
+                    on_escape()
+                except Exception:
+                    emit_warning(
+                        "Ctrl+X handler raised unexpectedly; Ctrl+C still works."
+                    )
+            elif cancel_agent_char and on_cancel_agent and data == cancel_agent_char:
+                try:
+                    on_cancel_agent()
+                except Exception:
+                    emit_warning("Cancel agent handler raised unexpectedly.")
+    finally:
+        termios.tcsetattr(fd, termios.TCSADRAIN, original_attrs)

code_muse/agents/_non_streaming_render.py

@@ -0,0 +1,148 @@
+"""Fallback rendering for agent runs where streaming didn't emit text.
+
+Some backends buffer responses and never emit SSE text deltas. In that case
+``event_stream_handler`` runs the tool/thinking paths but never prints the
+final answer, so the user sees a silent agent. This module provides:
+
+* ``StreamingTextDetector`` — a thin wrapper around ``event_stream_handler``
+  that records whether a ``TextPart`` / ``TextPartDelta`` ever appeared.
+* ``render_result_without_streaming`` — a one-shot renderer that walks
+  ``result.all_messages()`` and prints thinking + final text via the
+  non-streaming markdown path.
+
+The detector is async-safe and forwards every event untouched; it only
+observes. The renderer is best-effort: a render failure must never kill the
+run because the caller still has the raw result.
+"""
+
+from collections.abc import Callable
+from typing import Any
+
+from pydantic_ai import PartDeltaEvent, PartStartEvent
+from pydantic_ai.messages import (
+    ModelResponse,
+    TextPart,
+    TextPartDelta,
+    ThinkingPart,
+)
+
+from code_muse.tools.display import display_non_streamed_result
+
+
+class StreamingTextDetector:
+    """Wraps an ``event_stream_handler`` and records TextPart activity."""
+
+    def __init__(self, inner: Callable[..., Any]) -> None:
+        self._inner = inner
+        self.streamed_text: bool = False
+
+    async def __call__(self, ctx: Any, events: Any) -> Any:
+        detector = self
+
+        async def _tee() -> Any:
+            async for event in events:
+                if isinstance(event, PartStartEvent) and isinstance(
+                    getattr(event, "part", None), TextPart
+                ):
+                    part = event.part
+                    content = getattr(part, "content", "") or ""
+                    if content.strip():
+                        detector.streamed_text = True
+                elif isinstance(event, PartDeltaEvent) and isinstance(
+                    getattr(event, "delta", None), TextPartDelta
+                ):
+                    delta = event.delta
+                    if getattr(delta, "content_delta", ""):
+                        detector.streamed_text = True
+                yield event
+
+        return await self._inner(ctx, _tee())
+
+
+def _collect_thinking_from_messages(result: Any) -> str:
+    """Concatenate ``ThinkingPart`` content from intermediate ModelResponses.
+
+    "Intermediate" means every ModelResponse except the final one — the final
+    response's text is the answer we render separately.
+    """
+    try:
+        messages = list(result.all_messages())
+    except Exception:
+        return ""
+
+    model_responses = [m for m in messages if isinstance(m, ModelResponse)]
+    if len(model_responses) <= 1:
+        return ""
+
+    chunks: list[str] = []
+    for response in model_responses[:-1]:
+        for part in getattr(response, "parts", []) or []:
+            if isinstance(part, ThinkingPart):
+                content = getattr(part, "content", "") or ""
+                if content.strip():
+                    chunks.append(content)
+    return "\n\n".join(chunks)
+
+
+def _collect_final_text_from_messages(result: Any) -> str:
+    """Concatenate ``TextPart`` content from the final ModelResponse only."""
+    try:
+        messages = list(result.all_messages())
+    except Exception:
+        return ""
+
+    for message in reversed(messages):
+        if isinstance(message, ModelResponse):
+            chunks = [
+                getattr(p, "content", "") or ""
+                for p in getattr(message, "parts", []) or []
+                if isinstance(p, TextPart)
+            ]
+            return "".join(chunks)
+    return ""
+
+
+def render_result_without_streaming(result: Any) -> None:
+    """Render ``result`` via the non-streaming markdown path.
+
+    Emits the thinking banner (if any intermediate thinking was captured) and
+    the final agent response. Guarded so a render failure can't kill the run.
+    """
+    try:
+        thinking = _collect_thinking_from_messages(result)
+        if thinking.strip():
+            display_non_streamed_result(
+                thinking,
+                banner_text="THINKING",
+                banner_name="thinking",
+            )
+
+        final_text = _collect_final_text_from_messages(result)
+        if final_text.strip():
+            display_non_streamed_result(final_text)
+    except Exception:
+        # Rendering is best-effort: the caller still gets the raw result.
+        pass
+
+
+def should_render_fallback(
+    detector: StreamingTextDetector | None,
+    *,
+    skip: bool,
+) -> bool:
+    """Return True if we should render the final result ourselves.
+
+    ``skip`` is honoured unconditionally (e.g. plugins may render their own output).
+    Otherwise: render if there was no detector (streaming disabled) or the
+    detector never saw a TextPart fire.
+    """
+    if skip:
+        return False
+    return detector is None or not detector.streamed_text
+
+
+__all__ = [
+    "StreamingTextDetector",
+    "render_result_without_streaming",
+    "should_render_fallback",
+]