tunacode-cli 0.0.55__py3-none-any.whl → 0.0.78.6__py3-none-any.whl

tunacode/core/agents/agent_components/state_transition.py

@@ -0,0 +1,116 @@
+"""State transition management for agent response processing."""
+
+import threading
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING, Dict, Set
+
+from tunacode.types import AgentState
+
+if TYPE_CHECKING:
+    pass
+
+
+class InvalidStateTransitionError(Exception):
+    """Raised when an invalid state transition is attempted."""
+
+    def __init__(self, from_state: Enum, to_state: Enum, message: str = None):
+        self.from_state = from_state
+        self.to_state = to_state
+        self.message = message or f"Invalid state transition: {from_state.value} → {to_state.value}"
+        super().__init__(self.message)
+
+
+@dataclass
+class StateTransitionRules:
+    """Defines valid state transitions for the agent state machine."""
+
+    # Valid transitions for each state
+    valid_transitions: Dict[Enum, Set[Enum]]
+
+    def is_valid_transition(self, from_state: Enum, to_state: Enum) -> bool:
+        """Check if a transition between states is valid."""
+        return to_state in self.valid_transitions.get(from_state, set())
+
+    def get_valid_next_states(self, current_state: Enum) -> Set[Enum]:
+        """Get all valid next states from the current state."""
+        return self.valid_transitions.get(current_state, set())
+
+
+class AgentStateMachine:
+    """Thread-safe state machine for agent response processing."""
+
+    def __init__(self, initial_state: "AgentState", rules: StateTransitionRules):
+        """
+        Initialize the state machine.
+
+        Args:
+            initial_state: The starting state
+            rules: Transition rules defining valid state changes
+        """
+        self._state = initial_state
+        self._rules = rules
+        self._lock = threading.RLock()  # Reentrant lock for thread safety
+        self._completion_detected = False
+
+    @property
+    def current_state(self) -> "AgentState":
+        """Get the current state."""
+        with self._lock:
+            return self._state
+
+    def transition_to(self, new_state: "AgentState") -> None:
+        """
+        Transition to a new state.
+
+        Args:
+            new_state: The state to transition to
+
+        Raises:
+            InvalidStateTransitionError: If the transition is not valid
+        """
+        with self._lock:
+            if not self._rules.is_valid_transition(self._state, new_state):
+                raise InvalidStateTransitionError(
+                    self._state,
+                    new_state,
+                    f"Invalid state transition: {self._state.value} → {new_state.value}",
+                )
+
+            # Handle self-transitions as no-ops
+            if self._state == new_state:
+                return
+
+            self._state = new_state
+
+    def can_transition_to(self, target_state: "AgentState") -> bool:
+        """Check if a transition to the target state is allowed."""
+        with self._lock:
+            return self._rules.is_valid_transition(self._state, target_state)
+
+    def set_completion_detected(self, detected: bool = True) -> None:
+        """Mark that completion has been detected in the RESPONSE state."""
+        with self._lock:
+            self._completion_detected = detected
+
+    def is_completed(self) -> bool:
+        """Check if the task is completed (only valid in RESPONSE state)."""
+        with self._lock:
+            return self._state == AgentState.RESPONSE and self._completion_detected
+
+    def reset(self, initial_state: "AgentState" = None) -> None:
+        """Reset the state machine to initial state."""
+        with self._lock:
+            self._state = initial_state or AgentState.USER_INPUT
+            self._completion_detected = False
+
+
+# Define the transition rules for the agent state machine
+AGENT_TRANSITION_RULES = StateTransitionRules(
+    valid_transitions={
+        AgentState.USER_INPUT: {AgentState.ASSISTANT},
+        AgentState.ASSISTANT: {AgentState.TOOL_EXECUTION, AgentState.RESPONSE},
+        AgentState.TOOL_EXECUTION: {AgentState.RESPONSE},
+        AgentState.RESPONSE: {AgentState.ASSISTANT},  # Can transition back to continue
+    }
+)

tunacode/core/agents/agent_components/streaming.py

@@ -0,0 +1,296 @@
+"""Streaming instrumentation and handling for agent model request nodes.
+
+This module encapsulates verbose streaming + logging logic used during
+token-level streaming from the LLM provider. It updates session debug fields
+and streams deltas to the provided callback while being resilient to errors.
+"""
+
+from __future__ import annotations
+
+from typing import Awaitable, Callable, Optional
+
+from pydantic_ai.messages import PartDeltaEvent, TextPartDelta
+
+from tunacode.core.logging.logger import get_logger
+from tunacode.core.state import StateManager
+
+logger = get_logger(__name__)
+
+
+async def stream_model_request_node(
+    node,
+    agent_run_ctx,
+    state_manager: StateManager,
+    streaming_callback: Optional[Callable[[str], Awaitable[None]]],
+    request_id: str,
+    iteration_index: int,
+) -> None:
+    """Stream token deltas for a model request node with detailed instrumentation.
+
+    This function mirrors the prior inline logic in main.py but is extracted to
+    keep main.py lean. On streaming failure, it degrades gracefully to allow
+    non-streaming processing of the node.
+    """
+    if not streaming_callback:
+        return
+
+    # Gracefully handle streaming errors from LLM provider
+    try:
+        async with node.stream(agent_run_ctx) as request_stream:
+            # Initialize per-node debug accumulators
+            state_manager.session._debug_raw_stream_accum = ""
+            state_manager.session._debug_events = []
+            first_delta_logged = False
+            debug_event_count = 0
+            first_delta_seen = False
+            seeded_prefix_sent = False
+            pre_first_delta_text: Optional[str] = None
+
+            # Helper to extract text from a possible final-result object
+            def _extract_text(obj) -> Optional[str]:
+                try:
+                    if obj is None:
+                        return None
+                    if isinstance(obj, str):
+                        return obj
+                    # Common attributes that may hold text
+                    for attr in ("output", "text", "content", "message"):
+                        v = getattr(obj, attr, None)
+                        if isinstance(v, str) and v:
+                            return v
+                    # Parts-based result
+                    parts = getattr(obj, "parts", None)
+                    if isinstance(parts, (list, tuple)) and parts:
+                        texts: list[str] = []
+                        for p in parts:
+                            c = getattr(p, "content", None)
+                            if isinstance(c, str) and c:
+                                texts.append(c)
+                        if texts:
+                            return "".join(texts)
+                    # Nested .result or .response
+                    for attr in ("result", "response", "final"):
+                        v = getattr(obj, attr, None)
+                        t = _extract_text(v)
+                        if t:
+                            return t
+                except Exception:
+                    return None
+                return None
+
+            # Mark stream open
+            try:
+                import time as _t
+
+                state_manager.session._debug_events.append(
+                    f"[src] stream_opened ts_ns={_t.perf_counter_ns()}"
+                )
+            except Exception:
+                pass
+
+            async for event in request_stream:
+                debug_event_count += 1
+                # Log first few raw event types for diagnosis
+                if debug_event_count <= 5:
+                    try:
+                        etype = type(event).__name__
+                        d = getattr(event, "delta", None)
+                        dtype = type(d).__name__ if d is not None else None
+                        c = getattr(d, "content_delta", None) if d is not None else None
+                        clen = len(c) if isinstance(c, str) else None
+                        cpreview = repr(c[:5]) if isinstance(c, str) else None
+                        # Probe common fields on non-delta events to see if they contain text
+                        r = getattr(event, "result", None)
+                        rtype = type(r).__name__ if r is not None else None
+                        rpreview = None
+                        rplen = None
+                        # Also inspect event.part if present (e.g., PartStartEvent)
+                        p = getattr(event, "part", None)
+                        ptype = type(p).__name__ if p is not None else None
+                        pkind = getattr(p, "part_kind", None)
+                        pcontent = getattr(p, "content", None)
+                        ppreview = repr(pcontent[:20]) if isinstance(pcontent, str) else None
+                        pplen = len(pcontent) if isinstance(pcontent, str) else None
+                        try:
+                            if isinstance(r, str):
+                                rpreview = repr(r[:20])
+                                rplen = len(r)
+                            elif r is not None:
+                                # Try a few common shapes: .output, .text, .parts
+                                r_output = getattr(r, "output", None)
+                                r_text = getattr(r, "text", None)
+                                r_parts = getattr(r, "parts", None)
+                                if isinstance(r_output, str):
+                                    rpreview = repr(r_output[:20])
+                                    rplen = len(r_output)
+                                elif isinstance(r_text, str):
+                                    rpreview = repr(r_text[:20])
+                                    rplen = len(r_text)
+                                elif isinstance(r_parts, (list, tuple)) and r_parts:
+                                    # render a compact preview of first textual part
+                                    for _rp in r_parts:
+                                        rc = getattr(_rp, "content", None)
+                                        if isinstance(rc, str) and rc:
+                                            rpreview = repr(rc[:20])
+                                            rplen = len(rc)
+                                            break
+                        except Exception:
+                            pass
+                        event_info = (
+                            f"[src] event[{debug_event_count}] etype={etype} d={dtype} "
+                            f"clen={clen} cprev={cpreview} rtype={rtype} "
+                            f"rprev={rpreview} rlen={rplen} ptype={ptype} "
+                            f"pkind={pkind} pprev={ppreview} plen={pplen}"
+                        )
+                        state_manager.session._debug_events.append(event_info)
+                    except Exception:
+                        pass
+
+                # Attempt to capture pre-first-delta text from non-delta events
+                if not first_delta_seen:
+                    try:
+                        # event might be a PartStartEvent with .part.content
+                        if hasattr(event, "part") and hasattr(event.part, "content"):
+                            pc = event.part.content
+                            if isinstance(pc, str) and pc and not pc.lstrip().startswith("\n"):
+                                # capture a short potential prefix
+                                pre_first_delta_text = pc[:100] if len(pc) > 100 else pc
+                    except Exception:
+                        pass
+
+                # Handle delta events
+                if isinstance(event, PartDeltaEvent):
+                    if isinstance(event.delta, TextPartDelta):
+                        if event.delta.content_delta is not None and streaming_callback:
+                            # Seed prefix logic before the first true delta
+                            if not first_delta_seen:
+                                first_delta_seen = True
+                                try:
+                                    delta_text = event.delta.content_delta or ""
+                                    # Only seed when we have a short, safe candidate
+                                    if (
+                                        pre_first_delta_text
+                                        and len(pre_first_delta_text) <= 100
+                                        and not seeded_prefix_sent
+                                    ):
+                                        # If delta contains the candidate,
+                                        # emit the prefix up to that point
+                                        probe = pre_first_delta_text[:20]
+                                        idx = pre_first_delta_text.find(probe)
+                                        if idx > 0:
+                                            prefix = pre_first_delta_text[:idx]
+                                            if prefix:
+                                                await streaming_callback(prefix)
+                                                seeded_prefix_sent = True
+                                                preview_msg = (
+                                                    f"[src] seeded_prefix idx={idx} "
+                                                    f"len={len(prefix)} preview={repr(prefix)}"
+                                                )
+                                                state_manager.session._debug_events.append(
+                                                    preview_msg
+                                                )
+                                        elif idx == -1:
+                                            # Delta text does not appear in pre-text;
+                                            # emit the pre-text directly as a seed
+                                            # Safe for short pre-text (e.g., first word)
+                                            # to avoid duplication
+                                            if pre_first_delta_text.strip():
+                                                await streaming_callback(pre_first_delta_text)
+                                                seeded_prefix_sent = True
+                                                text_len = len(pre_first_delta_text)
+                                                preview_repr = repr(pre_first_delta_text)
+                                                direct_msg = (
+                                                    f"[src] seeded_prefix_direct "
+                                                    f"len={text_len} preview={preview_repr}"
+                                                )
+                                                state_manager.session._debug_events.append(
+                                                    direct_msg
+                                                )
+                                        else:
+                                            # idx == 0 means pre-text is already the
+                                            # start of delta; skip
+                                            skip_msg = (
+                                                f"[src] seed_skip idx={idx} "
+                                                f"delta_len={len(delta_text)}"
+                                            )
+                                            state_manager.session._debug_events.append(skip_msg)
+                                except Exception:
+                                    pass
+                                finally:
+                                    pre_first_delta_text = None
+
+                            # Record first-delta instrumentation
+                            if not first_delta_logged:
+                                try:
+                                    import time as _t
+
+                                    ts_ns = _t.perf_counter_ns()
+                                except Exception:
+                                    ts_ns = 0
+                                # Store debug event summary for later display
+                                chunk_preview = repr(
+                                    event.delta.content_delta[:5]
+                                    if event.delta.content_delta
+                                    else ""
+                                )
+                                chunk_len = len(event.delta.content_delta or "")
+                                delta_msg = (
+                                    f"[src] first_delta_received ts_ns={ts_ns} "
+                                    f"chunk_repr={chunk_preview} len={chunk_len}"
+                                )
+                                state_manager.session._debug_events.append(delta_msg)
+                                first_delta_logged = True
+
+                            # Accumulate full raw stream for comparison and forward delta
+                            delta_text = event.delta.content_delta or ""
+                            state_manager.session._debug_raw_stream_accum += delta_text
+                            await streaming_callback(delta_text)
+                    else:
+                        # Log empty or non-text deltas encountered
+                        state_manager.session._debug_events.append(
+                            "[src] empty_or_nontext_delta_skipped"
+                        )
+                else:
+                    # Capture any final result text for diagnostics
+                    try:
+                        final_text = _extract_text(getattr(event, "result", None))
+                        if final_text:
+                            final_msg = (
+                                f"[src] final_text_preview len={len(final_text)} "
+                                f"preview={repr(final_text[:20])}"
+                            )
+                            state_manager.session._debug_events.append(final_msg)
+                    except Exception:
+                        pass
+    except Exception as stream_err:
+        # Log with context and optionally notify UI, then degrade gracefully
+        logger.warning(
+            "Streaming error req=%s iter=%s: %s",
+            request_id,
+            iteration_index,
+            stream_err,
+            exc_info=True,
+        )
+
+        # Reset node state to allow graceful degradation to non-streaming mode
+        try:
+            if hasattr(node, "_did_stream"):
+                node._did_stream = False
+                logger.debug(
+                    "Reset node._did_stream after streaming error (req=%s iter=%s)",
+                    request_id,
+                    iteration_index,
+                )
+        except Exception as reset_err:
+            logger.debug(
+                "Failed to reset node._did_stream (req=%s iter=%s): %s",
+                request_id,
+                iteration_index,
+                reset_err,
+                exc_info=True,
+            )
+
+        if getattr(state_manager.session, "show_thoughts", False):
+            from tunacode.ui import console as ui
+
+            await ui.warning("Streaming failed; falling back to non-streaming mode")

tunacode/core/agents/agent_components/task_completion.py

@@ -1,7 +1,13 @@
 """Task completion detection utilities."""
 
+import re
 from typing import Tuple
 
+_COMPLETION_MARKERS = (
+    re.compile(r"^\s*TUNACODE\s+DONE:\s*", re.IGNORECASE),
+    re.compile(r"^\s*TUNACODE[_\s]+TASK_COMPLETE\s*:?[\s]*", re.IGNORECASE),
+)
+
 
 def check_task_completion(content: str) -> Tuple[bool, str]:
     """
@@ -18,11 +24,18 @@ def check_task_completion(content: str) -> Tuple[bool, str]:
     if not content:
         return False, content
 
-    lines = content.strip().split("\n")
-    if lines and lines[0].strip() == "TUNACODE_TASK_COMPLETE":
-        # Task is complete, return cleaned content
-        cleaned_lines = lines[1:] if len(lines) > 1 else []
-        cleaned_content = "\n".join(cleaned_lines).strip()
-        return True, cleaned_content
+    lines = content.split("\n")
+
+    for idx, line in enumerate(lines):
+        for pattern in _COMPLETION_MARKERS:
+            match = pattern.match(line)
+            if match:
+                remainder = line[match.end() :].strip()
+                cleaned_lines = lines[:idx]
+                if remainder:
+                    cleaned_lines.append(remainder)
+                cleaned_lines.extend(lines[idx + 1 :])
+                cleaned = "\n".join(cleaned_lines).strip()
+                return True, cleaned
 
     return False, content

tunacode/core/agents/agent_components/tool_buffer.py

@@ -1,6 +1,6 @@
 """Tool buffer for managing parallel execution of read-only tools."""
 
-from typing import Any, List, Tuple
+from typing import Any, Dict, List, Tuple
 
 
 class ToolBuffer:
@@ -22,3 +22,23 @@ class ToolBuffer:
     def has_tasks(self) -> bool:
         """Check if there are buffered tasks."""
         return len(self.read_only_tasks) > 0
+
+    def size(self) -> int:
+        """Return the number of buffered tasks."""
+        return len(self.read_only_tasks)
+
+    def peek(self) -> List[Tuple[Any, Any]]:
+        """Return buffered tasks without clearing the buffer."""
+        return self.read_only_tasks.copy()
+
+    def count_by_type(self) -> Dict[str, int]:
+        """Count buffered tools by type for metrics and debugging."""
+        counts: Dict[str, int] = {}
+        for part, _ in self.read_only_tasks:
+            tool_name = getattr(part, "tool_name", "unknown")
+            counts[tool_name] = counts.get(tool_name, 0) + 1
+        return counts
+
+    def clear(self) -> None:
+        """Clear all buffered tasks without executing them."""
+        self.read_only_tasks.clear()

tunacode/core/agents/agent_components/tool_executor.py

@@ -34,6 +34,16 @@ async def execute_tools_parallel(
         except Exception as e:
             logger.error(f"Error executing parallel tool: {e}", exc_info=True)
             return e
+        finally:
+            # Tool execution completed - resource cleanup handled by BaseTool.execute()
+            # Each tool's cleanup() method is called automatically in its execute()
+            # finally block. This ensures resources (file handles, connections,
+            # processes) are freed regardless of success or failure.
+            tool_name = getattr(part, "tool_name", "<unknown>")
+            logger.debug(
+                "Parallel tool execution completed (success or failure): tool=%s",
+                tool_name,
+            )
 
     # If we have more tools than max_parallel, execute in batches
     if len(tool_calls) > max_parallel: