tunacode-cli 0.0.71__py3-none-any.whl → 0.0.73__py3-none-any.whl

tunacode/core/agents/agent_components/response_state.py

@@ -1,13 +1,116 @@
 """Response state management for tracking agent processing state."""
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from typing import Optional
+
+from tunacode.types import AgentState
+
+from .state_transition import AGENT_TRANSITION_RULES, AgentStateMachine
 
 
 @dataclass
 class ResponseState:
-    """Track state across agent response processing."""
+    """Enhanced response state using enum-based state machine."""
+
+    # Internal state machine
+    _state_machine: AgentStateMachine = field(
+        default_factory=lambda: AgentStateMachine(AgentState.USER_INPUT, AGENT_TRANSITION_RULES)
+    )
+
+    # Backward compatibility boolean flags (derived from enum state)
+    _has_user_response: bool = False
+    _task_completed: bool = False
+    _awaiting_user_guidance: bool = False
+    _has_final_synthesis: bool = False
+
+    def __post_init__(self):
+        """Initialize the state machine."""
+        if not hasattr(self, "_state_machine"):
+            self._state_machine = AgentStateMachine(AgentState.USER_INPUT, AGENT_TRANSITION_RULES)
+
+    @property
+    def current_state(self) -> AgentState:
+        """Get the current enum state."""
+        return self._state_machine.current_state
+
+    def transition_to(self, new_state: AgentState) -> None:
+        """Transition to a new state."""
+        self._state_machine.transition_to(new_state)
+
+    def can_transition_to(self, target_state: AgentState) -> bool:
+        """Check if a transition to the target state is allowed."""
+        return self._state_machine.can_transition_to(target_state)
+
+    # Backward compatibility properties
+    @property
+    def has_user_response(self) -> bool:
+        """Legacy boolean flag for user response detection."""
+        return self._has_user_response
+
+    @has_user_response.setter
+    def has_user_response(self, value: bool) -> None:
+        """Set the legacy has_user_response flag."""
+        self._has_user_response = value
+
+    @property
+    def task_completed(self) -> bool:
+        """Legacy boolean flag for task completion (derived from state machine)."""
+        # If explicitly set true, honor it; otherwise derive from state machine
+        return bool(self._task_completed or self._state_machine.is_completed())
+
+    @task_completed.setter
+    def task_completed(self, value: bool) -> None:
+        """Set the legacy task_completed flag and sync with state machine."""
+        self._task_completed = bool(value)
+        if value:
+            # Ensure state reflects completion in RESPONSE
+            try:
+                if (
+                    self._state_machine.current_state != AgentState.RESPONSE
+                    and self._state_machine.can_transition_to(AgentState.RESPONSE)
+                ):
+                    self._state_machine.transition_to(AgentState.RESPONSE)
+            except Exception:
+                # Best-effort: ignore invalid transition in legacy paths
+                pass
+            self._state_machine.set_completion_detected(True)
+        else:
+            self._state_machine.set_completion_detected(False)
+
+    @property
+    def awaiting_user_guidance(self) -> bool:
+        """Legacy boolean flag for awaiting user guidance."""
+        return self._awaiting_user_guidance
+
+    @awaiting_user_guidance.setter
+    def awaiting_user_guidance(self, value: bool) -> None:
+        """Set the legacy awaiting_user_guidance flag."""
+        self._awaiting_user_guidance = value
+
+    @property
+    def has_final_synthesis(self) -> bool:
+        """Legacy boolean flag for final synthesis."""
+        return self._has_final_synthesis
+
+    @has_final_synthesis.setter
+    def has_final_synthesis(self, value: bool) -> None:
+        """Set the legacy has_final_synthesis flag."""
+        self._has_final_synthesis = value
+
+    # Enhanced state management methods
+    def set_completion_detected(self, detected: bool = True) -> None:
+        """Mark that completion has been detected in the RESPONSE state."""
+        self._state_machine.set_completion_detected(detected)
+
+    def is_completed(self) -> bool:
+        """Check if the task is completed according to the state machine."""
+        return self._state_machine.is_completed()
 
-    has_user_response: bool = False
-    task_completed: bool = False
-    awaiting_user_guidance: bool = False
-    has_final_synthesis: bool = False
+    def reset_state(self, initial_state: Optional[AgentState] = None) -> None:
+        """Reset the state machine to initial state."""
+        self._state_machine.reset(initial_state)
+        # Reset legacy flags
+        self._has_user_response = False
+        self._task_completed = False
+        self._awaiting_user_guidance = False
+        self._has_final_synthesis = False

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/task_completion.py

@@ -1,5 +1,6 @@
 """Task completion detection utilities."""
 
+import re
 from typing import Tuple
 
 
@@ -18,11 +19,14 @@ 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")
+
+    # New marker: any line starting with "TUNACODE DONE:" (case-insensitive, allow leading whitespace)
+    done_pattern = re.compile(r"^\s*TUNACODE\s+DONE:\s*", re.IGNORECASE)
+    for idx, line in enumerate(lines):
+        if done_pattern.match(line):
+            # Remove the marker line and return remaining content
+            cleaned = "\n".join(lines[:idx] + lines[idx + 1 :]).strip()
+            return True, cleaned
 
     return False, content

tunacode/core/agents/main.py

@@ -97,7 +97,7 @@ async def check_query_satisfaction(
     state_manager: StateManager,
 ) -> bool:
     """Check if the response satisfies the original query."""
-    return True  # Agent uses TUNACODE_TASK_COMPLETE marker
+    return True  # Completion decided via DONE marker in RESPONSE
 
 
 async def process_request(
@@ -251,7 +251,7 @@ Task: {message[:200]}...
 
 You're describing actions but not executing them. You MUST:
 
-1. If task is COMPLETE: Start response with TUNACODE_TASK_COMPLETE
+1. If task is COMPLETE: Start response with TUNACODE DONE:
 2. If task needs work: Execute a tool RIGHT NOW (grep, read_file, bash, etc.)
 3. If stuck: Explain the specific blocker
 
@@ -269,7 +269,7 @@ NO MORE DESCRIPTIONS. Take ACTION or mark COMPLETE."""
                     unproductive_iterations = 0
 
                 # REMOVED: Recursive satisfaction check that caused empty responses
-                # The agent now decides completion using TUNACODE_TASK_COMPLETE marker
+                # The agent now decides completion using a DONE marker
                 # This eliminates recursive agent calls and gives control back to the agent
 
                 # Store original query for reference
@@ -302,7 +302,7 @@ Progress so far:
 - Iterations: {i}
 - Tools used: {tools_used_str}
 
-If the task is complete, I should respond with TUNACODE_TASK_COMPLETE.
+If the task is complete, I should respond with TUNACODE DONE:
 Otherwise, please provide specific guidance on what to do next."""
 
                     create_user_message(clarification_content, state_manager)

tunacode/prompts/system.md

@@ -10,12 +10,22 @@ CRITICAL BEHAVIOR RULES:
 1. ALWAYS ANNOUNCE YOUR INTENTIONS FIRST: Before executing any tools, briefly state what you're about to do (e.g., "I'll search for the main agent implementation" or "Let me examine the file structure")
 2. When you say "Let me..." or "I will..." you MUST execute the corresponding tool in THE SAME RESPONSE
 3. Never describe what you'll do without doing it  ALWAYS execute tools when discussing actions
-4. When a task is COMPLETE, start your response with: TUNACODE_TASK_COMPLETE
+4. When a task is COMPLETE, start your response with: TUNACODE DONE:
 5. If your response is cut off or truncated, you'll be prompted to continue  complete your action
 6. YOU MUST NOT USE ANY EMOJIS, YOU WILL BE PUNISHED FOR EMOJI USE
 
 You MUST follow these rules:
 
+### Completion Signaling
+
+When you have fully completed the user’s task:
+
+- Start your response with a single line: `TUNACODE DONE:` followed by a brief outcome summary.
+- Do not add explanations before the DONE line; keep it as the first line.
+- Do NOT mark DONE if you have queued tools in the same response — execute tools first, then mark DONE.
+- Example:
+  - `TUNACODE DONE: Implemented enum state machine and updated completion logic`
+
 ###Tool Access Rules###
 
 You have 9 powerful tools at your disposal. Understanding their categories is CRITICAL for performance:

tunacode/types.py

@@ -209,6 +209,15 @@ class PlanPhase(Enum):
     REVIEW_DECISION = "review"
 
 
+class AgentState(Enum):
+    """Agent loop states for enhanced completion detection."""
+
+    USER_INPUT = "user_input"  # Initial: user prompt received
+    ASSISTANT = "assistant"  # Reasoning/deciding phase
+    TOOL_EXECUTION = "tool_execution"  # Tool execution phase
+    RESPONSE = "response"  # Handling results, may complete or loop
+
+
 @dataclass
 class PlanDoc:
     """Structured plan document with all required sections."""

tunacode/ui/completers.py

@@ -1,7 +1,7 @@
 """Completers for file references and commands."""
 
 import os
-from typing import TYPE_CHECKING, Iterable, Optional
+from typing import TYPE_CHECKING, Iterable, List, Optional
 
 from prompt_toolkit.completion import (
     CompleteEvent,
@@ -13,6 +13,7 @@ from prompt_toolkit.document import Document
 
 if TYPE_CHECKING:
     from ..cli.commands import CommandRegistry
+    from ..utils.models_registry import ModelInfo, ModelsRegistry
 
 
 class CommandCompleter(Completer):
@@ -126,11 +127,212 @@ class FileReferenceCompleter(Completer):
             pass
 
 
-def create_completer(command_registry: Optional["CommandRegistry"] = None) -> Completer:
-    """Create a merged completer for both commands and file references."""
-    return merge_completers(
-        [
-            CommandCompleter(command_registry),
-            FileReferenceCompleter(),
-        ]
-    )
+class ModelCompleter(Completer):
+    """Completer for model names in /model command."""
+
+    def __init__(self, registry: Optional["ModelsRegistry"] = None):
+        """Initialize the model completer."""
+        self.registry = registry
+        self._models_cache: Optional[List[ModelInfo]] = None
+        self._registry_loaded = False
+
+    async def _ensure_registry_loaded(self):
+        """Ensure the models registry is loaded."""
+        if self.registry and not self._registry_loaded:
+            try:
+                # Try to load models (this will be fast if already loaded)
+                await self.registry.load()
+                self._registry_loaded = True
+                self._models_cache = (
+                    list(self.registry.models.values()) if self.registry.models else []
+                )
+            except Exception:
+                # If loading fails, use empty cache
+                self._models_cache = []
+                self._registry_loaded = True
+
+    def get_completions(
+        self, document: Document, _complete_event: CompleteEvent
+    ) -> Iterable[Completion]:
+        """Get completions for model names."""
+        if not self.registry:
+            return
+
+        text = document.text_before_cursor
+
+        # Check if we're in a /model command context
+        lines = text.split("\n")
+        current_line = lines[-1].strip()
+
+        # Must start with /model
+        if not current_line.startswith("/model"):
+            return
+
+        # Try to load registry synchronously if not loaded
+        # Note: This is a compromise - ideally we'd use async completion
+        if not self._registry_loaded:
+            try:
+                # Quick attempt to load cached data only
+                if self.registry._is_cache_valid() and self.registry._load_from_cache():
+                    self._registry_loaded = True
+                    self._models_cache = list(self.registry.models.values())
+                elif not self._models_cache:
+                    # Use fallback models for immediate completion
+                    self.registry._load_fallback_models()
+                    self._registry_loaded = True
+                    self._models_cache = list(self.registry.models.values())
+            except Exception:
+                return  # Skip completion if we can't load models
+
+        # Get the part after /model
+        parts = current_line.split()
+        if len(parts) < 2:
+            # Just "/model" - suggest popular searches and top models
+            popular_searches = ["claude", "gpt", "gemini", "openai", "anthropic"]
+            for search_term in popular_searches:
+                yield Completion(
+                    text=search_term, display=f"{search_term} (search)", display_meta="search term"
+                )
+
+            # Also show top 3 most popular models if we have them
+            if self._models_cache:
+                popular_models = []
+                # Look for common popular models
+                for model in self._models_cache:
+                    if any(pop in model.id.lower() for pop in ["gpt-4o", "claude-3", "gemini-2"]):
+                        popular_models.append(model)
+                        if len(popular_models) >= 3:
+                            break
+
+                for model in popular_models:
+                    display = f"{model.full_id} - {model.name}"
+                    if model.cost.input is not None:
+                        display += f" (${model.cost.input}/{model.cost.output})"
+
+                    yield Completion(
+                        text=model.full_id, display=display, display_meta=f"{model.provider} model"
+                    )
+            return
+
+        # Get the current word being typed
+        word_before_cursor = document.get_word_before_cursor(WORD=True)
+        if not word_before_cursor or not self._models_cache:
+            return
+
+        query = word_before_cursor.lower()
+
+        # Use the new grouped approach to find base models with variants
+        base_models = self.registry.find_base_models(query)
+
+        if not base_models:
+            return
+
+        results = []
+        shown_base_models = 0
+
+        # Sort base models by popularity/relevance
+        sorted_base_models = sorted(
+            base_models.items(),
+            key=lambda x: (
+                # Popular models first
+                -1
+                if any(
+                    pop in x[0] for pop in ["gpt-4o", "gpt-4", "claude-3", "gemini-2", "o3", "o1"]
+                )
+                else 0,
+                # Then by name
+                x[0],
+            ),
+        )
+
+        for base_model_name, variants in sorted_base_models:
+            if shown_base_models >= 5:  # Limit to top 5 base models
+                break
+
+            shown_variants = 0
+            for i, model in enumerate(variants):
+                if shown_variants >= 3:  # Show max 3 variants per base model
+                    break
+
+                # Calculate start position for replacement
+                start_pos = -len(word_before_cursor)
+
+                # Build display text with enhanced info
+                cost_str = ""
+                if model.cost.input is not None:
+                    if model.cost.input == 0:
+                        cost_str = " (FREE)"
+                    else:
+                        cost_str = f" (${model.cost.input}/{model.cost.output})"
+
+                # Format provider info
+                provider_display = self._get_provider_display_name(model.provider)
+
+                # Primary variant gets the bullet, others get indentation
+                if i == 0:
+                    # First variant - primary option with bullet
+                    display = f"● {model.full_id} - {model.name}{cost_str}"
+                    if model.cost.input == 0:
+                        display += " ⭐"  # Star for free models
+                else:
+                    # Additional variants - indented
+                    display = f"   {model.full_id} - {model.name}{cost_str}"
+                    if model.cost.input == 0:
+                        display += " ⭐"
+
+                meta_info = f"{provider_display}"
+                if len(variants) > 1:
+                    meta_info += f" ({len(variants)} sources)"
+
+                results.append(
+                    Completion(
+                        text=model.full_id,
+                        start_position=start_pos,
+                        display=display,
+                        display_meta=meta_info,
+                    )
+                )
+
+                shown_variants += 1
+
+            shown_base_models += 1
+
+        # Limit total results for readability
+        for completion in results[:20]:
+            yield completion
+
+    def _get_provider_display_name(self, provider: str) -> str:
+        """Get a user-friendly provider display name."""
+        provider_names = {
+            "openai": "OpenAI Direct",
+            "anthropic": "Anthropic Direct",
+            "google": "Google Direct",
+            "google-gla": "Google Labs",
+            "openrouter": "OpenRouter",
+            "github-models": "GitHub Models (FREE)",
+            "azure": "Azure OpenAI",
+            "fastrouter": "FastRouter",
+            "requesty": "Requesty",
+            "cloudflare-workers-ai": "Cloudflare",
+            "amazon-bedrock": "AWS Bedrock",
+            "chutes": "Chutes AI",
+            "deepinfra": "DeepInfra",
+            "venice": "Venice AI",
+        }
+        return provider_names.get(provider, provider.title())
+
+
+def create_completer(
+    command_registry: Optional["CommandRegistry"] = None,
+    models_registry: Optional["ModelsRegistry"] = None,
+) -> Completer:
+    """Create a merged completer for commands, file references, and models."""
+    completers = [
+        CommandCompleter(command_registry),
+        FileReferenceCompleter(),
+    ]
+
+    if models_registry:
+        completers.append(ModelCompleter(models_registry))
+
+    return merge_completers(completers)

tunacode/ui/input.py

@@ -95,6 +95,12 @@ async def multiline_input(
         )
     )
 
+    # Create models registry for auto-completion (lazy loaded)
+    from ..utils.models_registry import ModelsRegistry
+
+    models_registry = ModelsRegistry()
+    # Note: Registry will be loaded lazily by the completer when needed
+
     # Display input area (Plan Mode indicator is handled dynamically in prompt manager)
     result = await input(
         "multiline",
@@ -102,7 +108,7 @@ async def multiline_input(
         key_bindings=kb,
         multiline=True,
         placeholder=placeholder,
-        completer=create_completer(command_registry),
+        completer=create_completer(command_registry, models_registry),
         lexer=FileReferenceLexer(),
         state_manager=state_manager,
     )