code-puppy 0.0.170__py3-none-any.whl → 0.0.172__py3-none-any.whl

code_puppy/message_history_processor.py

@@ -3,7 +3,15 @@ import queue
 from typing import Any, List, Set, Tuple
 
 import pydantic
-from pydantic_ai.messages import ModelMessage, ModelRequest, TextPart, ToolCallPart
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    TextPart,
+    ToolCallPart,
+    ToolCallPartDelta,
+    ToolReturn,
+    ToolReturnPart,
+)
 
 from code_puppy.config import (
     get_model_name,
@@ -82,9 +90,46 @@ def estimate_tokens_for_message(message: ModelMessage) -> int:
 
 
 def filter_huge_messages(messages: List[ModelMessage]) -> List[ModelMessage]:
-    filtered = [m for m in messages if estimate_tokens_for_message(m) < 50000]
-    pruned = prune_interrupted_tool_calls(filtered)
-    return pruned
+    if not messages:
+        return []
+
+    # Never drop the system prompt, even if it is extremely large.
+    system_message, *rest = messages
+    filtered_rest = [
+        m for m in rest if estimate_tokens_for_message(m) < 50000
+    ]
+    return [system_message] + filtered_rest
+
+
+def _is_tool_call_part(part: Any) -> bool:
+    if isinstance(part, (ToolCallPart, ToolCallPartDelta)):
+        return True
+
+    part_kind = (getattr(part, "part_kind", "") or "").replace("_", "-")
+    if part_kind == "tool-call":
+        return True
+
+    has_tool_name = getattr(part, "tool_name", None) is not None
+    has_args = getattr(part, "args", None) is not None
+    has_args_delta = getattr(part, "args_delta", None) is not None
+
+    return bool(has_tool_name and (has_args or has_args_delta))
+
+
+def _is_tool_return_part(part: Any) -> bool:
+    if isinstance(part, (ToolReturnPart, ToolReturn)):
+        return True
+
+    part_kind = (getattr(part, "part_kind", "") or "").replace("_", "-")
+    if part_kind in {"tool-return", "tool-result"}:
+        return True
+
+    if getattr(part, "tool_call_id", None) is None:
+        return False
+
+    has_content = getattr(part, "content", None) is not None
+    has_content_delta = getattr(part, "content_delta", None) is not None
+    return bool(has_content or has_content_delta)
 
 
 def split_messages_for_protected_summarization(
@@ -126,19 +171,18 @@ def split_messages_for_protected_summarization(
         if protected_token_count + message_tokens > protected_tokens_limit:
             break
 
-        protected_messages.insert(0, message)  # Insert at beginning to maintain order
+        protected_messages.append(message)
         protected_token_count += message_tokens
 
-    # Add system message at the beginning of protected messages
+    # Messages that were added while scanning backwards are currently in reverse order.
+    # Reverse them to restore chronological ordering, then prepend the system prompt.
+    protected_messages.reverse()
     protected_messages.insert(0, system_message)
 
-    # Messages to summarize are everything between system message and protected zone
-    protected_start_idx = (
-        len(messages) - len(protected_messages) + 1
-    )  # +1 because system message is protected
-    messages_to_summarize = messages[
-        1:protected_start_idx
-    ]  # Start from 1 to skip system message
+    # Messages to summarize are everything between the system message and the
+    # protected tail zone we just constructed.
+    protected_start_idx = max(1, len(messages) - (len(protected_messages) - 1))
+    messages_to_summarize = messages[1:protected_start_idx]
 
     emit_info(
         f"🔒 Protecting {len(protected_messages)} recent messages ({protected_token_count} tokens, limit: {protected_tokens_limit})"
@@ -148,99 +192,36 @@ def split_messages_for_protected_summarization(
     return messages_to_summarize, protected_messages
 
 
-def deduplicate_tool_returns(messages: List[ModelMessage]) -> List[ModelMessage]:
-    """
-    Remove duplicate tool returns while preserving the first occurrence for each tool_call_id.
-
-    This function identifies tool-return parts that share the same tool_call_id and
-    removes duplicates, keeping only the first return for each id. This prevents
-    conversation corruption from duplicate tool_result blocks.
-    """
-    if not messages:
-        return messages
-
-    seen_tool_returns: Set[str] = set()
-    deduplicated: List[ModelMessage] = []
-    removed_count = 0
-
-    for msg in messages:
-        # Check if this message has any parts we need to filter
-        if not hasattr(msg, "parts") or not msg.parts:
-            deduplicated.append(msg)
-            continue
-
-        # Filter parts within this message
-        filtered_parts = []
-        msg_had_duplicates = False
-
-        for part in msg.parts:
-            tool_call_id = getattr(part, "tool_call_id", None)
-            part_kind = getattr(part, "part_kind", None)
-
-            # Check if this is a tool-return part
-            if tool_call_id and part_kind in {
-                "tool-return",
-                "tool-result",
-                "tool_result",
-            }:
-                if tool_call_id in seen_tool_returns:
-                    # This is a duplicate return, skip it
-                    msg_had_duplicates = True
-                    removed_count += 1
-                    continue
-                else:
-                    # First occurrence of this return, keep it
-                    seen_tool_returns.add(tool_call_id)
-                    filtered_parts.append(part)
-            else:
-                # Not a tool return, always keep
-                filtered_parts.append(part)
-
-        # If we filtered out parts, create a new message with filtered parts
-        if msg_had_duplicates and filtered_parts:
-            # Create a new message with the same attributes but filtered parts
-            new_msg = type(msg)(parts=filtered_parts)
-            # Copy over other attributes if they exist
-            for attr_name in dir(msg):
-                if (
-                    not attr_name.startswith("_")
-                    and attr_name != "parts"
-                    and hasattr(msg, attr_name)
-                ):
-                    try:
-                        setattr(new_msg, attr_name, getattr(msg, attr_name))
-                    except (AttributeError, TypeError):
-                        # Skip attributes that can't be set
-                        pass
-            deduplicated.append(new_msg)
-        elif filtered_parts:  # No duplicates but has parts
-            deduplicated.append(msg)
-        # If no parts remain after filtering, drop the entire message
-
-    if removed_count > 0:
-        emit_warning(f"Removed {removed_count} duplicate tool-return part(s)")
-
-    return deduplicated
-
-
 def summarize_messages(
-    messages: List[ModelMessage], with_protection=True
+    messages: List[ModelMessage], with_protection: bool = True
 ) -> Tuple[List[ModelMessage], List[ModelMessage]]:
     """
     Summarize messages while protecting recent messages up to PROTECTED_TOKENS.
 
     Returns:
-        List of messages: [system_message, summary_of_old_messages, ...protected_recent_messages]
+        Tuple of (compacted_messages, summarized_source_messages)
+        where compacted_messages always preserves the original system message
+        as the first entry.
     """
-    messages_to_summarize, protected_messages = messages, []
+    messages_to_summarize: List[ModelMessage]
+    protected_messages: List[ModelMessage]
+
     if with_protection:
         messages_to_summarize, protected_messages = (
             split_messages_for_protected_summarization(messages)
         )
+    else:
+        messages_to_summarize = messages[1:] if messages else []
+        protected_messages = messages[:1]
+
+    if not messages:
+        return [], []
+
+    system_message = messages[0]
 
     if not messages_to_summarize:
-        # Nothing to summarize, return protected messages as-is
-        return protected_messages, messages_to_summarize
+        # Nothing to summarize, so just return the original sequence
+        return prune_interrupted_tool_calls(messages), []
 
     instructions = (
         "The input will be a log of Agentic AI steps that have been taken"
@@ -257,12 +238,24 @@ def summarize_messages(
         new_messages = run_summarization_sync(
             instructions, message_history=messages_to_summarize
         )
-        # Return: [system_message, summary, ...protected_recent_messages]
-        result = new_messages + protected_messages[1:]
-        return prune_interrupted_tool_calls(result), messages_to_summarize
+
+        if not isinstance(new_messages, list):
+            emit_warning(
+                "Summarization agent returned non-list output; wrapping into message request"
+            )
+            new_messages = [ModelRequest([TextPart(str(new_messages))])]
+
+        compacted: List[ModelMessage] = [system_message] + list(new_messages)
+
+        # Drop the system message from protected_messages because we already included it
+        protected_tail = [msg for msg in protected_messages if msg is not system_message]
+
+        compacted.extend(protected_tail)
+
+        return prune_interrupted_tool_calls(compacted), messages_to_summarize
     except Exception as e:
         emit_error(f"Summarization failed during compaction: {e}")
-        return messages, messages_to_summarize  # Return original messages on failure
+        return messages, []  # Return original messages on failure
 
 
 def summarize_message(message: ModelMessage) -> ModelMessage:
@@ -329,11 +322,10 @@ def prune_interrupted_tool_calls(messages: List[ModelMessage]) -> List[ModelMess
             tool_call_id = getattr(part, "tool_call_id", None)
             if not tool_call_id:
                 continue
-            # Heuristic: if it's an explicit ToolCallPart or has a tool_name/args,
-            # consider it a call; otherwise it's a return/result.
-            if part.part_kind == "tool-call":
+
+            if _is_tool_call_part(part) and not _is_tool_return_part(part):
                 tool_call_ids.add(tool_call_id)
-            else:
+            elif _is_tool_return_part(part):
                 tool_return_ids.add(tool_call_id)
 
     mismatched: Set[str] = tool_call_ids.symmetric_difference(tool_return_ids)
@@ -362,12 +354,15 @@ def prune_interrupted_tool_calls(messages: List[ModelMessage]) -> List[ModelMess
 
 
 def message_history_processor(messages: List[ModelMessage]) -> List[ModelMessage]:
-    # First, prune any interrupted/mismatched tool-call conversations
-    total_current_tokens = sum(estimate_tokens_for_message(msg) for msg in messages)
+    cleaned_history = prune_interrupted_tool_calls(messages)
+
+    total_current_tokens = sum(
+        estimate_tokens_for_message(msg) for msg in cleaned_history
+    )
 
     model_max = get_model_context_length()
 
-    proportion_used = total_current_tokens / model_max
+    proportion_used = total_current_tokens / model_max if model_max else 0
 
     # Check if we're in TUI mode and can update the status bar
     from code_puppy.state_management import get_tui_app_instance, is_tui_mode
@@ -406,17 +401,15 @@ def message_history_processor(messages: List[ModelMessage]) -> List[ModelMessage
     compaction_strategy = get_compaction_strategy()
 
     if proportion_used > compaction_threshold:
+        filtered_history = filter_huge_messages(cleaned_history)
+
         if compaction_strategy == "truncation":
-            # Use truncation instead of summarization
             protected_tokens = get_protected_token_count()
-            result_messages = truncation(
-                filter_huge_messages(messages), protected_tokens
-            )
-            summarized_messages = []  # No summarization in truncation mode
+            result_messages = truncation(filtered_history, protected_tokens)
+            summarized_messages: List[ModelMessage] = []
         else:
-            # Default to summarization
             result_messages, summarized_messages = summarize_messages(
-                filter_huge_messages(messages)
+                filtered_history
             )
 
         final_token_count = sum(
@@ -447,7 +440,9 @@ def message_history_processor(messages: List[ModelMessage]) -> List[ModelMessage
         for m in summarized_messages:
             add_compacted_message_hash(hash_message(m))
         return result_messages
-    return messages
+
+    set_message_history(cleaned_history)
+    return cleaned_history
 
 
 def truncation(
@@ -475,16 +470,17 @@ def truncation(
 
 
 def message_history_accumulator(messages: List[Any]):
-    _message_history = get_message_history()
-    message_history_hashes = set([hash_message(m) for m in _message_history])
-    for msg in messages:
-        if (
-            hash_message(msg) not in message_history_hashes
-            and hash_message(msg) not in get_compacted_message_hashes()
-        ):
-            _message_history.append(msg)
-
-    # Apply message history trimming using the main processor
-    # This ensures we maintain global state while still managing context limits
-    message_history_processor(_message_history)
-    return get_message_history()
+    existing_history = list(get_message_history())
+    seen_hashes = {hash_message(message) for message in existing_history}
+    compacted_hashes = get_compacted_message_hashes()
+
+    for message in messages:
+        message_hash = hash_message(message)
+        if message_hash in seen_hashes or message_hash in compacted_hashes:
+            continue
+        existing_history.append(message)
+        seen_hashes.add(message_hash)
+
+    updated_history = message_history_processor(existing_history)
+    set_message_history(updated_history)
+    return updated_history

code_puppy/state_management.py

@@ -1,44 +1,34 @@
-from typing import Any, List
+import json
+from types import ModuleType
+from typing import Any, List, Set
 
-# Legacy global state - maintained for backward compatibility
-_message_history: List[Any] = []
-_compacted_message_hashes = set()
+import pydantic
+
+from code_puppy.messaging import emit_info
 
-# Flag to control whether to use agent-specific history (True) or global history (False)
-_use_agent_specific_history = True
 _tui_mode: bool = False
 _tui_app_instance: Any = None
 
 
+def _require_agent_manager() -> ModuleType:
+    """Import the agent manager module, raising if it is unavailable."""
+    try:
+        from code_puppy.agents import agent_manager
+    except Exception as error:  # pragma: no cover - import errors surface immediately
+        raise RuntimeError("Agent manager module unavailable") from error
+    return agent_manager
+
+
 def add_compacted_message_hash(message_hash: str) -> None:
     """Add a message hash to the set of compacted message hashes."""
-    if _use_agent_specific_history:
-        try:
-            from code_puppy.agents.agent_manager import (
-                add_current_agent_compacted_message_hash,
-            )
+    manager = _require_agent_manager()
+    manager.add_current_agent_compacted_message_hash(message_hash)
 
-            add_current_agent_compacted_message_hash(message_hash)
-            return
-        except Exception:
-            # Fallback to global if agent system fails
-            pass
-    _compacted_message_hashes.add(message_hash)
 
-
-def get_compacted_message_hashes():
+def get_compacted_message_hashes() -> Set[str]:
     """Get the set of compacted message hashes."""
-    if _use_agent_specific_history:
-        try:
-            from code_puppy.agents.agent_manager import (
-                get_current_agent_compacted_message_hashes,
-            )
-
-            return get_current_agent_compacted_message_hashes()
-        except Exception:
-            # Fallback to global if agent system fails
-            pass
-    return _compacted_message_hashes
+    manager = _require_agent_manager()
+    return manager.get_current_agent_compacted_message_hashes()
 
 
 def set_tui_mode(enabled: bool) -> None:
@@ -89,112 +79,81 @@ def get_tui_mode() -> bool:
 
 
 def get_message_history() -> List[Any]:
-    """Get message history - uses agent-specific history if enabled, otherwise global."""
-    if _use_agent_specific_history:
-        try:
-            from code_puppy.agents.agent_manager import (
-                get_current_agent_message_history,
-            )
-
-            return get_current_agent_message_history()
-        except Exception:
-            # Fallback to global if agent system fails
-            return _message_history
-    return _message_history
+    """Get message history for the active agent."""
+    manager = _require_agent_manager()
+    return manager.get_current_agent_message_history()
 
 
 def set_message_history(history: List[Any]) -> None:
-    """Set message history - uses agent-specific history if enabled, otherwise global."""
-    if _use_agent_specific_history:
-        try:
-            from code_puppy.agents.agent_manager import (
-                set_current_agent_message_history,
-            )
-
-            set_current_agent_message_history(history)
-            return
-        except Exception:
-            # Fallback to global if agent system fails
-            pass
-    global _message_history
-    _message_history = history
+    """Replace the message history for the active agent."""
+    manager = _require_agent_manager()
+    manager.set_current_agent_message_history(history)
 
 
 def clear_message_history() -> None:
-    """Clear message history - uses agent-specific history if enabled, otherwise global."""
-    if _use_agent_specific_history:
-        try:
-            from code_puppy.agents.agent_manager import (
-                clear_current_agent_message_history,
-            )
-
-            clear_current_agent_message_history()
-            return
-        except Exception:
-            # Fallback to global if agent system fails
-            pass
-    global _message_history
-    _message_history = []
+    """Clear message history for the active agent."""
+    manager = _require_agent_manager()
+    manager.clear_current_agent_message_history()
 
 
 def append_to_message_history(message: Any) -> None:
-    """Append to message history - uses agent-specific history if enabled, otherwise global."""
-    if _use_agent_specific_history:
-        try:
-            from code_puppy.agents.agent_manager import (
-                append_to_current_agent_message_history,
-            )
-
-            append_to_current_agent_message_history(message)
-            return
-        except Exception:
-            # Fallback to global if agent system fails
-            pass
-    _message_history.append(message)
+    """Append a message to the active agent's history."""
+    manager = _require_agent_manager()
+    manager.append_to_current_agent_message_history(message)
 
 
 def extend_message_history(history: List[Any]) -> None:
-    """Extend message history - uses agent-specific history if enabled, otherwise global."""
-    if _use_agent_specific_history:
-        try:
-            from code_puppy.agents.agent_manager import (
-                extend_current_agent_message_history,
-            )
-
-            extend_current_agent_message_history(history)
-            return
-        except Exception:
-            # Fallback to global if agent system fails
-            pass
-    _message_history.extend(history)
-
-
-def set_use_agent_specific_history(enabled: bool) -> None:
-    """Enable or disable agent-specific message history.
-
-    Args:
-        enabled: True to use per-agent history, False to use global history.
-    """
-    global _use_agent_specific_history
-    _use_agent_specific_history = enabled
-
-
-def is_using_agent_specific_history() -> bool:
-    """Check if agent-specific message history is enabled.
-
-    Returns:
-        True if using per-agent history, False if using global history.
-    """
-    return _use_agent_specific_history
-
-
-def hash_message(message):
-    hashable_entities = []
-    for part in message.parts:
-        if hasattr(part, "timestamp"):
-            hashable_entities.append(part.timestamp.isoformat())
-        elif hasattr(part, "tool_call_id"):
-            hashable_entities.append(part.tool_call_id)
-        else:
-            hashable_entities.append(part.content)
-    return hash(",".join(hashable_entities))
+    """Extend the active agent's message history."""
+    manager = _require_agent_manager()
+    manager.extend_current_agent_message_history(history)
+
+
+def _stringify_part(part: Any) -> str:
+    """Create a stable string representation for a message part.
+
+    We deliberately ignore timestamps so identical content hashes the same even when
+    emitted at different times. This prevents status updates from blowing up the
+    history when they are repeated with new timestamps."""
+
+    attributes: List[str] = [part.__class__.__name__]
+
+    # Role/instructions help disambiguate parts that otherwise share content
+    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):
+        attributes.append(f"content={json.dumps(content.model_dump(), sort_keys=True)}")
+    elif isinstance(content, dict):
+        attributes.append(f"content={json.dumps(content, sort_keys=True)}")
+    else:
+        attributes.append(f"content={repr(content)}")
+    result = "|".join(attributes)
+    return result
+
+
+def hash_message(message: Any) -> int:
+    """Create a stable hash for a model message that ignores timestamps."""
+    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)
+    return hash(canonical)