tunacode-cli 0.0.39__py3-none-any.whl → 0.0.41__py3-none-any.whl

tunacode/context.py

@@ -1,3 +1,4 @@
+import logging
 import subprocess
 from pathlib import Path
 from typing import Dict, List
@@ -5,6 +6,8 @@ from typing import Dict, List
 from tunacode.utils.ripgrep import ripgrep
 from tunacode.utils.system import list_cwd
 
+logger = logging.getLogger(__name__)
+
 
 async def get_git_status() -> Dict[str, object]:
     """Return git branch and dirty state information."""
@@ -29,7 +32,8 @@ async def get_git_status() -> Dict[str, object]:
                     behind = int(part.split("behind")[1].strip().strip(" ]"))
         dirty = any(line for line in lines[1:])
         return {"branch": branch, "ahead": ahead, "behind": behind, "dirty": dirty}
-    except Exception:
+    except Exception as e:
+        logger.warning(f"Failed to get git status: {e}")
         return {}
 
 
@@ -54,8 +58,8 @@ async def get_code_style() -> str:
         if file.exists():
             try:
                 parts.append(file.read_text(encoding="utf-8"))
-            except Exception:
-                pass
+            except Exception as e:
+                logger.debug(f"Failed to read TUNACODE.md at {file}: {e}")
         if current == current.parent:
             break
         current = current.parent

tunacode/core/agents/main.py

@@ -6,6 +6,7 @@ Handles agent creation, configuration, and request processing.
 
 import asyncio
 import json
+import logging
 import os
 import re
 from datetime import datetime, timezone
@@ -30,6 +31,8 @@ except ImportError:
 
 from tunacode.constants import READ_ONLY_TOOLS
 from tunacode.core.state import StateManager
+from tunacode.core.token_usage.api_response_parser import ApiResponseParser
+from tunacode.core.token_usage.cost_calculator import CostCalculator
 from tunacode.services.mcp import get_mcp_servers
 from tunacode.tools.bash import bash
 from tunacode.tools.glob import glob
@@ -37,6 +40,7 @@ from tunacode.tools.grep import grep
 from tunacode.tools.list_dir import list_dir
 from tunacode.tools.read_file import read_file
 from tunacode.tools.run_command import run_command
+from tunacode.tools.todo import TodoTool
 from tunacode.tools.update_file import update_file
 from tunacode.tools.write_file import write_file
 from tunacode.types import (
@@ -50,8 +54,12 @@ from tunacode.types import (
     ToolCallback,
     ToolCallId,
     ToolName,
+    UsageTrackerProtocol,
 )
 
+# Configure logging
+logger = logging.getLogger(__name__)
+
 
 class ToolBuffer:
     """Buffer for collecting read-only tool calls to execute in parallel."""
@@ -110,6 +118,7 @@ async def execute_tools_parallel(
         try:
             return await callback(part, node)
         except Exception as e:
+            logger.error(f"Error executing parallel tool: {e}", exc_info=True)
             return e
 
     # If we have more tools than max_parallel, execute in batches
@@ -214,6 +223,7 @@ async def _process_node(
     state_manager: StateManager,
     tool_buffer: Optional[ToolBuffer] = None,
     streaming_callback: Optional[callable] = None,
+    usage_tracker: Optional[UsageTrackerProtocol] = None,
 ):
     from tunacode.ui import console as ui
     from tunacode.utils.token_counter import estimate_tokens
@@ -233,6 +243,9 @@ async def _process_node(
     if hasattr(node, "model_response"):
         state_manager.session.messages.append(node.model_response)
 
+        if usage_tracker:
+            await usage_tracker.track_and_display(node.model_response)
+
         # Stream content to callback if provided
         # Use this as fallback when true token streaming is not available
         if streaming_callback and not STREAMING_AVAILABLE:
@@ -313,8 +326,8 @@ async def _process_node(
                             thought_obj = json.loads(content)
                             if "thought" in thought_obj:
                                 await ui.muted(f"REASONING: {thought_obj['thought']}")
-                    except (json.JSONDecodeError, KeyError):
-                        pass
+                    except (json.JSONDecodeError, KeyError) as e:
+                        logger.debug(f"Failed to parse thought JSON: {e}")
 
                     # Pattern 3: Multi-line thoughts with context
                     multiline_pattern = r'\{"thought":\s*"([^"]+(?:\\.[^"]*)*?)"\}'
@@ -442,8 +455,9 @@ async def _process_node(
         # Handle tool returns
         for part in node.model_response.parts:
             if part.part_kind == "tool-return":
-                obs_msg = f"OBSERVATION[{part.tool_name}]: {part.content[:2_000]}"
-                state_manager.session.messages.append(obs_msg)
+                state_manager.session.messages.append(
+                    f"OBSERVATION[{part.tool_name}]: {part.content}"
+                )
 
                 # Display tool return when thoughts are enabled
                 if state_manager.session.show_thoughts:
@@ -510,9 +524,22 @@ def get_or_create_agent(model: ModelName, state_manager: StateManager) -> Pydant
             else:
                 # Log that TUNACODE.md was not found
                 print("📄 TUNACODE.md not found: Using default context")
-        except Exception:
-            # Ignore errors loading TUNACODE.md
-            pass
+        except Exception as e:
+            # Log errors loading TUNACODE.md at debug level
+            logger.debug(f"Error loading TUNACODE.md: {e}")
+
+        todo_tool = TodoTool(state_manager=state_manager)
+
+        try:
+            # Only add todo section if there are actual todos
+            current_todos = todo_tool.get_current_todos_sync()
+            if current_todos != "No todos found":
+                system_prompt += f'\n\n# Current Todo List\n\nYou have existing todos that need attention:\n\n{current_todos}\n\nRemember to check progress on these todos and update them as you work. Use todo("list") to see current status anytime.'
+        except Exception as e:
+            # Log error but don't fail agent creation
+            import sys
+
+            print(f"Warning: Failed to load todos: {e}", file=sys.stderr)
 
         state_manager.session.agents[model] = Agent(
             model=model,
@@ -524,6 +551,7 @@ def get_or_create_agent(model: ModelName, state_manager: StateManager) -> Pydant
                 Tool(list_dir, max_retries=max_retries),
                 Tool(read_file, max_retries=max_retries),
                 Tool(run_command, max_retries=max_retries),
+                Tool(todo_tool._execute, max_retries=max_retries),
                 Tool(update_file, max_retries=max_retries),
                 Tool(write_file, max_retries=max_retries),
             ],
@@ -622,7 +650,9 @@ async def parse_json_tool_calls(
                     if isinstance(parsed, dict) and "tool" in parsed and "args" in parsed:
                         potential_jsons.append((parsed["tool"], parsed["args"]))
                 except json.JSONDecodeError:
-                    pass
+                    logger.debug(
+                        f"Failed to parse potential JSON tool call: {potential_json[:50]}..."
+                    )
                 start_pos = -1
 
     matches = potential_jsons
@@ -719,7 +749,13 @@ async def process_request(
     fallback_enabled = state_manager.session.user_config.get("settings", {}).get(
         "fallback_response", True
     )
+    from tunacode.configuration.models import ModelRegistry
+    from tunacode.core.token_usage.usage_tracker import UsageTracker
 
+    parser = ApiResponseParser()
+    registry = ModelRegistry()
+    calculator = CostCalculator(registry)
+    usage_tracker = UsageTracker(parser, calculator, state_manager)
     response_state = ResponseState()
 
     # Reset iteration tracking for this request
@@ -763,7 +799,14 @@ async def process_request(
                             if event.delta.content_delta:
                                 await streaming_callback(event.delta.content_delta)
 
-            await _process_node(node, tool_callback, state_manager, tool_buffer, streaming_callback)
+            await _process_node(
+                node,
+                tool_callback,
+                state_manager,
+                tool_buffer,
+                streaming_callback,
+                usage_tracker,
+            )
             if hasattr(node, "result") and node.result and hasattr(node.result, "output"):
                 if node.result.output:
                     response_state.has_user_response = True

tunacode/core/setup/config_setup.py

@@ -318,6 +318,11 @@ class ConfigSetup(BaseSetup):
 
             self.state_manager.session.user_config["default_model"] = model
 
+        if self.cli_config.get("custom_context_window"):
+            self.state_manager.session.user_config["context_window_size"] = self.cli_config[
+                "custom_context_window"
+            ]
+
         # Set current model
         self.state_manager.session.current_model = self.state_manager.session.user_config[
             "default_model"

tunacode/core/state.py

@@ -14,9 +14,12 @@ from tunacode.types import (
     MessageHistory,
     ModelName,
     SessionId,
+    TodoItem,
     ToolName,
     UserConfig,
 )
+from tunacode.utils.message_utils import get_message_content
+from tunacode.utils.token_counter import estimate_tokens
 
 
 @dataclass
@@ -37,6 +40,7 @@ class SessionState:
     device_id: Optional[DeviceId] = None
     input_sessions: InputSessions = field(default_factory=dict)
     current_task: Optional[Any] = None
+    todos: list[TodoItem] = field(default_factory=list)
     # Enhanced tracking for thoughts display
     files_in_context: set[str] = field(default_factory=set)
     tool_calls: list[dict[str, Any]] = field(default_factory=list)
@@ -46,6 +50,31 @@ class SessionState:
     is_streaming_active: bool = False
     # Track streaming panel reference for tool handler access
     streaming_panel: Optional[Any] = None
+    # Context window tracking (estimation based)
+    total_tokens: int = 0
+    max_tokens: int = 0
+    # API usage tracking (actual from providers)
+    last_call_usage: dict = field(
+        default_factory=lambda: {
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "cost": 0.0,
+        }
+    )
+    session_total_usage: dict = field(
+        default_factory=lambda: {
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "cost": 0.0,
+        }
+    )
+
+    def update_token_count(self):
+        """Calculates the total token count from messages and files in context."""
+        message_contents = [get_message_content(msg) for msg in self.messages]
+        message_content = " ".join(c for c in message_contents if c)
+        file_content = " ".join(self.files_in_context)
+        self.total_tokens = estimate_tokens(message_content + file_content, self.current_model)
 
 
 class StateManager:
@@ -56,5 +85,25 @@ class StateManager:
     def session(self) -> SessionState:
         return self._session
 
-    def reset_session(self):
+    def add_todo(self, todo: TodoItem) -> None:
+        self._session.todos.append(todo)
+
+    def update_todo(self, todo_id: str, status: str) -> None:
+        from datetime import datetime
+
+        for todo in self._session.todos:
+            if todo.id == todo_id:
+                todo.status = status
+                if status == "completed" and not todo.completed_at:
+                    todo.completed_at = datetime.now()
+                break
+
+    def remove_todo(self, todo_id: str) -> None:
+        self._session.todos = [todo for todo in self._session.todos if todo.id != todo_id]
+
+    def clear_todos(self) -> None:
+        self._session.todos = []
+
+    def reset_session(self) -> None:
+        """Reset the session to a fresh state."""
         self._session = SessionState()

tunacode/core/token_usage/api_response_parser.py

@@ -0,0 +1,44 @@
+"""
+Module: tunacode.llm.api_response_parser
+Provides a parser to standardize token usage information from various LLM API responses.
+"""
+
+from typing import Any, Dict
+
+from tunacode.types import ModelName
+
+
+class ApiResponseParser:
+    """
+    Parses LLM API response objects to extract token usage and the actual model name used.
+    This version works directly with the pydantic-ai ModelResponse object.
+    """
+
+    def parse(self, model: ModelName, response_obj: Any) -> Dict[str, Any]:
+        """
+        Parses the standardized API response object.
+
+        Args:
+            model (ModelName): The model name that was requested. Used as a fallback.
+            response_obj (Any): The raw ModelResponse object from the agent.
+
+        Returns:
+            Dict[str, Any]: A standardized dictionary with 'prompt_tokens',
+                            'completion_tokens', and 'model_name'.
+        """
+        # --- FIX: Access attributes directly from the object ---
+        # Default to an empty object if usage is None
+        usage = getattr(response_obj, "usage", None) or {}
+
+        # Extract the actual model name, falling back to the requested model.
+        actual_model_name = getattr(response_obj, "model_name", model)
+
+        # The pydantic-ai Usage object standardizes keys to 'request_tokens'
+        # and 'response_tokens'. We access them as attributes.
+        parsed_data = {
+            "prompt_tokens": getattr(usage, "request_tokens", 0),
+            "completion_tokens": getattr(usage, "response_tokens", 0),
+            "model_name": actual_model_name,
+        }
+
+        return parsed_data

tunacode/core/token_usage/cost_calculator.py

@@ -0,0 +1,58 @@
+"""
+Module: tunacode.pricing.cost_calculator
+Provides a utility for calculating the cost of model usage based on token counts.
+"""
+
+from tunacode.configuration.models import ModelRegistry
+from tunacode.types import CostAmount, ModelName, TokenCount
+
+
+class CostCalculator:
+    """
+    Calculates the cost of a model interaction based on prompt and completion tokens.
+    """
+
+    def __init__(self, registry: ModelRegistry):
+        """
+        Initializes the CostCalculator with a model registry.
+
+        Args:
+            registry (ModelRegistry): An instance of ModelRegistry that contains
+                                      the pricing information for various models.
+        """
+        self._registry = registry
+
+    def calculate_cost(
+        self,
+        model_name: ModelName,
+        prompt_tokens: TokenCount,
+        completion_tokens: TokenCount,
+    ) -> CostAmount:
+        """
+        Calculates the total cost for a given model and token usage.
+
+        Args:
+            model_name (ModelName): The identifier for the model (e.g., "openai:gpt-4o").
+            prompt_tokens (TokenCount): The number of tokens in the input/prompt.
+            completion_tokens (TokenCount): The number of tokens in the output/completion.
+
+        Returns:
+            CostAmount: The calculated cost as a float. Returns 0.0 if the model
+                        is not found in the registry.
+        """
+        model_config = self._registry.get_model(model_name)
+
+        if not model_config:
+            return 0.0
+
+        TOKENS_PER_MILLION = 1_000_000
+
+        pricing = model_config.pricing
+
+        input_cost = (prompt_tokens / TOKENS_PER_MILLION) * pricing.input
+
+        output_cost = (completion_tokens / TOKENS_PER_MILLION) * pricing.output
+
+        total_cost = input_cost + output_cost
+
+        return total_cost

tunacode/core/token_usage/usage_tracker.py

@@ -0,0 +1,98 @@
+from typing import Any
+
+from tunacode.core.state import StateManager
+from tunacode.core.token_usage.api_response_parser import ApiResponseParser
+from tunacode.core.token_usage.cost_calculator import CostCalculator
+from tunacode.types import UsageTrackerProtocol
+from tunacode.ui import console as ui  # Import the ui console directly
+
+
+class UsageTracker(UsageTrackerProtocol):
+    """
+    Handles parsing, calculating, storing, and displaying token usage and cost.
+    """
+
+    def __init__(
+        self,
+        parser: ApiResponseParser,
+        calculator: CostCalculator,
+        state_manager: StateManager,
+    ):
+        self.parser = parser
+        self.calculator = calculator
+        self.state_manager = state_manager
+
+    async def track_and_display(self, response_obj: Any):
+        """
+        Main method to process a model response for usage tracking.
+        """
+        try:
+            # 1. Parse the response to get token data
+            requested_model = self.state_manager.session.current_model
+            parsed_data = self.parser.parse(model=requested_model, response_obj=response_obj)
+
+            if not parsed_data:
+                return
+
+            # 2. Calculate the cost
+            cost = self._calculate_cost(parsed_data)
+
+            # 3. Update the session state
+            self._update_state(parsed_data, cost)
+
+            # 4. Display the summary if enabled
+            if self.state_manager.session.show_thoughts:
+                await self._display_summary()
+
+        except Exception as e:
+            if self.state_manager.session.show_thoughts:
+                await ui.error(f"Error during cost calculation: {e}")
+
+    def _calculate_cost(self, parsed_data: dict) -> float:
+        """Calculates the cost for the given parsed data."""
+        requested_model = self.state_manager.session.current_model
+        api_model_name = parsed_data.get("model_name", requested_model)
+        final_model_name = api_model_name
+
+        # Logic to preserve the provider prefix
+        if ":" in requested_model:
+            provider_prefix = requested_model.split(":", 1)[0]
+            if not api_model_name.startswith(provider_prefix + ":"):
+                final_model_name = f"{provider_prefix}:{api_model_name}"
+
+        return self.calculator.calculate_cost(
+            prompt_tokens=parsed_data.get("prompt_tokens", 0),
+            completion_tokens=parsed_data.get("completion_tokens", 0),
+            model_name=final_model_name,
+        )
+
+    def _update_state(self, parsed_data: dict, cost: float):
+        """Updates the last_call and session_total usage in the state."""
+        session = self.state_manager.session
+        prompt_tokens = parsed_data.get("prompt_tokens", 0)
+        completion_tokens = parsed_data.get("completion_tokens", 0)
+
+        # Update last call usage
+        session.last_call_usage["prompt_tokens"] = prompt_tokens
+        session.last_call_usage["completion_tokens"] = completion_tokens
+        session.last_call_usage["cost"] = cost
+
+        # Accumulate session totals
+        session.session_total_usage["prompt_tokens"] += prompt_tokens
+        session.session_total_usage["completion_tokens"] += completion_tokens
+        session.session_total_usage["cost"] += cost
+
+    async def _display_summary(self):
+        """Formats and prints the usage summary to the console."""
+        session = self.state_manager.session
+        prompt = session.last_call_usage["prompt_tokens"]
+        completion = session.last_call_usage["completion_tokens"]
+        last_cost = session.last_call_usage["cost"]
+        session_cost = session.session_total_usage["cost"]
+
+        usage_summary = (
+            f"[ Tokens: {prompt + completion:,} (P: {prompt:,}, C: {completion:,}) | "
+            f"Cost: ${last_cost:.4f} | "
+            f"Session Total: ${session_cost:.4f} ]"
+        )
+        await ui.muted(usage_summary)

tunacode/prompts/system.md

@@ -12,7 +12,7 @@ You MUST follow these rules:
 
 \###Tool Access Rules###
 
-You have 8 powerful tools at your disposal. Understanding their categories is CRITICAL for performance:
+You have 9 powerful tools at your disposal. Understanding their categories is CRITICAL for performance:
 
 ** READ-ONLY TOOLS (Safe, Parallel-Executable)**
 These tools can and SHOULD be executed in parallel batches for 3x-10x performance gains:
@@ -30,19 +30,28 @@ These tools can and SHOULD be executed in parallel batches for 3x-10x performanc
    - Returns: Sorted list of matching file paths
    - Use for: Finding all \*.py files, configs, etc.
 
+** TASK MANAGEMENT TOOLS (Fast, Sequential)**
+These tools help organize and track complex multi-step tasks:
+
+5. `todo(action: str, content: str = None, todo_id: str = None, status: str = None, priority: str = None, todos: list = None)` — Manage task lists
+   - Actions: "add", "add_multiple", "update", "complete", "list", "remove"
+   - Use for: Breaking down complex tasks, tracking progress, organizing work
+   - **IMPORTANT**: Use this tool when tackling multi-step problems or complex implementations
+   - **Multiple todos**: Use `todo("add_multiple", todos=[{"content": "task1", "priority": "high"}, {"content": "task2", "priority": "medium"}])` to add many todos at once
+
 ** WRITE/EXECUTE TOOLS (Require Confirmation, Sequential)**
 These tools modify state and MUST run one at a time with user confirmation:
 
-5. `write_file(filepath: str, content: str)` — Create new files
+6. `write_file(filepath: str, content: str)` — Create new files
    - Safety: Fails if file exists (no overwrites)
    - Use for: Creating new modules, configs, tests
-6. `update_file(filepath: str, target: str, patch: str)` — Modify existing files
+7. `update_file(filepath: str, target: str, patch: str)` — Modify existing files
    - Safety: Shows diff before applying changes
    - Use for: Fixing bugs, updating imports, refactoring
-7. `run_command(command: str)` — Execute shell commands
+8. `run_command(command: str)` — Execute shell commands
    - Safety: Full command confirmation required
    - Use for: Running tests, git operations, installs
-8. `bash(command: str)` — Advanced shell with environment control
+9. `bash(command: str)` — Advanced shell with environment control
    - Safety: Enhanced security, output limits (5KB)
    - Use for: Complex scripts, interactive commands
 
@@ -85,12 +94,65 @@ These tools modify state and MUST run one at a time with user confirmation:
 - Need to see file content? → `read_file`
 - Need to find something? → `grep` (content) or `glob` (filenames)
 - Need to explore? → `list_dir`
+- Need to track tasks? → `todo` (for complex multi-step work)
 - Need to create? → `write_file`
 - Need to modify? → `update_file`
 - Need to run commands? → `run_command` (simple) or `bash` (complex)
 
 ---
 
+\###Task Management Best Practices###
+
+**IMPORTANT**: For complex, multi-step tasks, you MUST use the todo tool to break down work and track progress.
+
+**When to use the todo tool:**
+- User requests implementing new features (3+ steps involved)
+- Complex debugging that requires multiple investigation steps  
+- Refactoring that affects multiple files
+- Any task where you need to track progress across multiple tool executions
+
+**Todo workflow pattern:**
+1. **Break down complex requests**: `todo("add", "Analyze current authentication system", priority="high")`
+2. **Track progress**: `todo("update", todo_id="1", status="in_progress")`
+3. **Mark completion**: `todo("complete", todo_id="1")`
+4. **Show status**: `todo("list")` to display current work
+
+**Example multi-step task breakdown:**
+```
+User: "Add authentication to my Flask app"
+
+OPTIMAL approach (multiple individual adds):
+1. todo("add", "Analyze Flask app structure", priority="high")
+2. todo("add", "Create user model and database schema", priority="high") 
+3. todo("add", "Implement registration endpoint", priority="medium")
+4. todo("add", "Implement login endpoint", priority="medium")
+5. todo("add", "Add password hashing", priority="high")
+6. todo("add", "Create auth middleware", priority="medium")
+7. todo("add", "Write tests for auth system", priority="low")
+
+ALTERNATIVE (batch add for efficiency):
+todo("add_multiple", todos=[
+  {"content": "Analyze Flask app structure", "priority": "high"},
+  {"content": "Create user model and database schema", "priority": "high"}, 
+  {"content": "Implement registration endpoint", "priority": "medium"},
+  {"content": "Implement login endpoint", "priority": "medium"},
+  {"content": "Add password hashing", "priority": "high"},
+  {"content": "Create auth middleware", "priority": "medium"},
+  {"content": "Write tests for auth system", "priority": "low"}
+])
+
+Then work through each task systematically, marking progress as you go.
+```
+
+**Benefits of using todos:**
+- Helps users understand the full scope of work
+- Provides clear progress tracking  
+- Ensures no steps are forgotten
+- Makes complex tasks feel manageable
+- Shows professional project management approach
+
+---
+
 \###Working Directory Rules###
 
 **CRITICAL**: You MUST respect the user's current working directory:
@@ -371,11 +433,13 @@ RESPONSE TO USER: The main.py file contains a simple main function that prints '
 | **grep** | 🔍 Read | ✅ Yes | ❌ No | 4KB | Search text patterns |
 | **list_dir** | 🔍 Read | ✅ Yes | ❌ No | 200 entries | Browse directories |
 | **glob** | 🔍 Read | ✅ Yes | ❌ No | 1000 files | Find files by pattern |
+| **todo** | 📋 Task | ❌ No | ❌ No | - | Track multi-step tasks |
 | **write_file** | ⚡ Write | ❌ No | ✅ Yes | - | Create new files |
 | **update_file** | ⚡ Write | ❌ No | ✅ Yes | - | Modify existing files |
 | **run_command** | ⚡ Execute | ❌ No | ✅ Yes | 5KB | Simple shell commands |
 | **bash** | ⚡ Execute | ❌ No | ✅ Yes | 5KB | Complex shell scripts |
 
 **Remember**: ALWAYS batch 3-4 read-only tools together for optimal performance (3x faster)!
+**Remember**: Use the todo tool to break down and track complex multi-step tasks!
 
 ```