tunacode-cli 0.1.21__py3-none-any.whl

tunacode/core/agents/agent_components/state_transition.py

@@ -0,0 +1,112 @@
+"""State transition management for agent response processing."""
+
+import threading
+from dataclasses import dataclass
+from enum import Enum
+
+from tunacode.types import AgentState
+
+
+class InvalidStateTransitionError(Exception):
+    """Raised when an invalid state transition is attempted."""
+
+    def __init__(self, from_state: Enum, to_state: Enum, message: str | None = 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,271 @@
+"""Streaming instrumentation and handling for agent model request nodes.
+
+This module encapsulates verbose streaming 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 collections.abc import Awaitable, Callable
+
+from pydantic_ai.messages import PartDeltaEvent, TextPartDelta
+
+from tunacode.core.state import StateManager
+
+
+def _find_overlap_length(pre_text: str, delta_text: str) -> int:
+    """Find length of longest pre_text suffix that equals delta_text prefix.
+
+    This detects when delta_text starts with content already in pre_text,
+    so we can avoid emitting duplicate text.
+
+    Returns:
+        Number of characters that overlap (0 if no overlap).
+    """
+    if not pre_text or not delta_text:
+        return 0
+
+    max_check = min(len(pre_text), len(delta_text))
+    for overlap_len in range(max_check, 0, -1):
+        if delta_text[:overlap_len] == pre_text[-overlap_len:]:
+            return overlap_len
+    return 0
+
+
+async def stream_model_request_node(
+    node,
+    agent_run_ctx,
+    state_manager: StateManager,
+    streaming_callback: Callable[[str], Awaitable[None]] | 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: str | None = None
+
+            # Helper to extract text from a possible final-result object
+            def _extract_text(obj) -> str | None:
+                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
+                                    ):
+                                        # Find overlap: longest suffix of pre_first_delta_text
+                                        # that matches a prefix of delta_text
+                                        overlap_len = _find_overlap_length(
+                                            pre_first_delta_text, delta_text
+                                        )
+                                        # Emit the non-overlapping prefix
+                                        prefix_len = len(pre_first_delta_text) - overlap_len
+                                        prefix_to_emit = pre_first_delta_text[:prefix_len]
+
+                                        if prefix_to_emit.strip():
+                                            await streaming_callback(prefix_to_emit)
+                                            seeded_prefix_sent = True
+                                            preview_msg = (
+                                                f"[src] seeded_prefix overlap={overlap_len} "
+                                                f"len={len(prefix_to_emit)} "
+                                                f"preview={repr(prefix_to_emit[:20])}"
+                                            )
+                                            state_manager.session._debug_events.append(preview_msg)
+                                        else:
+                                            skip_msg = (
+                                                f"[src] seed_skip overlap={overlap_len} "
+                                                f"delta_len={len(delta_text)} "
+                                                f"pre_len={len(pre_first_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:
+        # Reset node state to allow graceful degradation to non-streaming mode
+        try:
+            if hasattr(node, "_did_stream"):
+                node._did_stream = False
+        except Exception:
+            pass

tunacode/core/agents/agent_components/task_completion.py

@@ -0,0 +1,40 @@
+"""Task completion detection utilities."""
+
+import re
+
+_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]:
+    """
+    Check if the content indicates task completion.
+
+    Args:
+        content: The text content to check
+
+    Returns:
+        Tuple of (is_complete, cleaned_content)
+        - is_complete: True if task completion marker found
+        - cleaned_content: Content with marker removed
+    """
+    if not content:
+        return False, 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

@@ -0,0 +1,44 @@
+"""Tool buffer for managing parallel execution of read-only tools."""
+
+from typing import Any
+
+
+class ToolBuffer:
+    """Buffer for collecting read-only tool calls to execute in parallel."""
+
+    def __init__(self):
+        self.read_only_tasks: list[tuple[Any, Any]] = []
+
+    def add(self, part: Any, node: Any) -> None:
+        """Add a read-only tool call to the buffer."""
+        self.read_only_tasks.append((part, node))
+
+    def flush(self) -> list[tuple[Any, Any]]:
+        """Return buffered tasks and clear the buffer."""
+        tasks = self.read_only_tasks
+        self.read_only_tasks = []
+        return tasks
+
+    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

@@ -0,0 +1,101 @@
+"""Tool execution and parallelization functionality with automatic retry."""
+
+import asyncio
+import os
+import random
+from typing import Any
+
+from pydantic_ai import ModelRetry
+
+from tunacode.constants import (
+    TOOL_MAX_RETRIES,
+    TOOL_RETRY_BASE_DELAY,
+    TOOL_RETRY_MAX_DELAY,
+)
+from tunacode.exceptions import (
+    ConfigurationError,
+    FileOperationError,
+    ToolExecutionError,
+    UserAbortError,
+    ValidationError,
+)
+from tunacode.types import ToolCallback
+
+# Errors that should NOT be retried - they represent user intent or unrecoverable states
+NON_RETRYABLE_ERRORS = (
+    UserAbortError,
+    ModelRetry,
+    KeyboardInterrupt,
+    SystemExit,
+    ValidationError,
+    ConfigurationError,
+    ToolExecutionError,
+    FileOperationError,
+)
+
+
+def _calculate_backoff(attempt: int) -> float:
+    """Exponential backoff with jitter."""
+    delay = min(TOOL_RETRY_BASE_DELAY * (2 ** (attempt - 1)), TOOL_RETRY_MAX_DELAY)
+    jitter = random.uniform(0, delay * 0.1)  # nosec B311 - not for crypto
+    return delay + jitter
+
+
+async def execute_tools_parallel(
+    tool_calls: list[tuple[Any, Any]], callback: ToolCallback
+) -> list[Any]:
+    """
+    Execute multiple tool calls in parallel using asyncio with automatic retry.
+
+    Each tool gets up to TOOL_MAX_RETRIES attempts before failing.
+    Non-retryable errors (user abort, validation, etc.) propagate immediately.
+
+    Args:
+        tool_calls: List of (part, node) tuples
+        callback: The tool callback function to execute
+
+    Returns:
+        List of results in the same order as input
+
+    Raises:
+        Exception: Re-raises after all retry attempts exhausted
+    """
+    max_parallel = int(os.environ.get("TUNACODE_MAX_PARALLEL", os.cpu_count() or 4))
+
+    async def execute_with_retry(part, node):
+        for attempt in range(1, TOOL_MAX_RETRIES + 1):
+            try:
+                result = await callback(part, node)
+
+                return result
+            except NON_RETRYABLE_ERRORS:
+                raise
+            except Exception:
+                if attempt == TOOL_MAX_RETRIES:
+                    raise
+                backoff = _calculate_backoff(attempt)
+                await asyncio.sleep(backoff)
+
+        raise AssertionError("unreachable")
+
+    # Execute in batches if we have more tools than max_parallel
+    if len(tool_calls) > max_parallel:
+        results = []
+        for i in range(0, len(tool_calls), max_parallel):
+            batch = tool_calls[i : i + max_parallel]
+            batch_tasks = [execute_with_retry(part, node) for part, node in batch]
+            batch_results = await asyncio.gather(*batch_tasks, return_exceptions=True)
+            results.extend(batch_results)
+            # Check for errors after each batch
+            for result in batch_results:
+                if isinstance(result, Exception):
+                    raise result
+        return results
+    else:
+        tasks = [execute_with_retry(part, node) for part, node in tool_calls]
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        # Check for errors - raise the first one
+        for result in results:
+            if isinstance(result, Exception):
+                raise result
+        return results

tunacode/core/agents/agent_components/truncation_checker.py

@@ -0,0 +1,37 @@
+"""Truncation detection utilities for agent responses."""
+
+
+def check_for_truncation(combined_content: str) -> bool:
+    """Check if content appears to be truncated.
+
+    Args:
+        combined_content: The text content to check for truncation
+
+    Returns:
+        bool: True if the content appears truncated, False otherwise
+    """
+    if not combined_content:
+        return False
+
+    # Truncation indicators:
+    # 1. Ends with "..." or "…" (but not part of a complete sentence)
+    # 2. Contains incomplete markdown/code blocks
+    # 3. Ends with incomplete parentheses/brackets
+
+    # Check for ellipsis at end suggesting truncation
+    if combined_content.endswith(("...", "…")) and not combined_content.endswith(("....", "….")):
+        return True
+
+    # Check for unclosed markdown code blocks
+    code_block_count = combined_content.count("```")
+    if code_block_count % 2 != 0:
+        return True
+
+    # Check for unclosed brackets/parentheses (more opens than closes)
+    open_brackets = (
+        combined_content.count("[") + combined_content.count("(") + combined_content.count("{")
+    )
+    close_brackets = (
+        combined_content.count("]") + combined_content.count(")") + combined_content.count("}")
+    )
+    return open_brackets > close_brackets