[Phase 5] Token and cost tracking per task

[frankbria]

Summary

Track LLM token usage and estimated costs per task and batch execution.

Problem

No visibility into LLM costs. Users can't optimize or budget for AI usage.

Proposed Solution

Tracking

• Input/output tokens per LLM call
• Aggregate by task, batch, and workspace
• Cost estimation based on model pricing

Storage

• Token counts in run records
• Cost summaries in workspace stats

CLI Commands

cf stats tokens                    # Workspace token summary
cf stats tokens --task <id>        # Per-task breakdown
cf stats costs --period month      # Cost report

Acceptance Criteria

• Token counts recorded for all LLM calls
• Per-task and per-batch aggregation
• Cost estimation with configurable pricing
• CLI commands to view usage
• Export to CSV/JSON

Related

• [Phase 5] Implement by_day token usage aggregation in MetricsTracker #120: Performance metrics
• Part of Phase 5 observability

[traycerai]

Plan

Observations

The codebase already has substantial token tracking infrastructure in place from Sprint 10. Key components exist:

• TokenCounter class for tiktoken-based token counting
• MetricsTracker class with cost calculation and aggregation
• TokenUsage Pydantic model and database schema
• TokenRepository for persistence
• CLI commands in file:codeframe/cli/metrics_commands.py (not yet registered)

The main gaps are: (1) integration of token tracking into all LLM call sites, (2) CLI command registration and enhancement, (3) batch-level aggregation, (4) workspace statistics, and (5) CSV/JSON export functionality.

Approach

The implementation will leverage existing infrastructure and add missing integration points. The strategy focuses on:

1. Instrumentation: Add token tracking to all LLM call sites (Executor, LeadAgent, worker agents)
2. CLI Enhancement: Register existing metrics commands and add new stats subcommands with export capabilities
3. Aggregation: Extend MetricsTracker with batch and workspace-level aggregation methods
4. Export: Add CSV/JSON export utilities for usage reports

This approach minimizes code duplication and builds on proven patterns already in the codebase.

Implementation Steps

1. Integrate Token Tracking into LLM Call Sites

File: file:codeframe/core/executor.py

Modify _generate_file_content and _generate_edit_content methods to record token usage after each LLM call:

• Import MetricsTracker and CallType from file:codeframe/lib/metrics_tracker.py and file:codeframe/core/models.py
• Add metrics_tracker: MetricsTracker parameter to Executor.__init__
• After self.llm.complete() calls (lines 627-633, 657-663), extract token counts from response.input_tokens and response.output_tokens
• Call await metrics_tracker.record_token_usage() with:
  • task_id from context
  • agent_id from context or "executor"
  • project_id from context
  • model_name from response.model
  • Token counts and CallType.TASK_EXECUTION

File: file:codeframe/agents/lead_agent.py

Integrate token tracking in the chat method and other LLM call sites:

• Add self.metrics_tracker = MetricsTracker(db=self.db) in __init__
• In chat method (around line 200-250), after self.provider.send_message(), record token usage from response
• In _generate_next_discovery_question, generate_prd, and other LLM-calling methods, add token tracking
• Use CallType.COORDINATION for discovery/planning calls, CallType.TASK_EXECUTION for task-related calls

File: file:codeframe/agents/worker_agent.py (if exists)

Similar integration pattern as LeadAgent:

• Initialize MetricsTracker in constructor
• Record token usage after each LLM call
• Use appropriate CallType based on operation

File: file:codeframe/adapters/llm/anthropic.py

Enhance send_message to return token usage in response:

• Ensure the response dictionary includes usage with input_tokens and output_tokens (already present at lines 104-107)
• No changes needed, just verify consistency

2. Enhance MetricsTracker with Batch and Workspace Aggregation

File: file:codeframe/lib/metrics_tracker.py

Add new methods for batch and workspace-level statistics:

Method: get_batch_costs

async def get_batch_costs(
    self,
    project_id: int,
    task_ids: List[int],
    start_date: Optional[datetime] = None,
    end_date: Optional[datetime] = None,
) -> Dict[str, Any]:

• Query token_usage table filtering by task_id IN task_ids
• Aggregate total cost, tokens, and calls
• Return breakdown by task with individual costs

Method: get_workspace_costs

async def get_workspace_costs(
    self,
    start_date: Optional[datetime] = None,
    end_date: Optional[datetime] = None,
) -> Dict[str, Any]:

• Query all projects (no project_id filter)
• Aggregate by project, agent, and model
• Return workspace-wide statistics

Method: export_usage_to_csv

def export_usage_to_csv(
    self,
    usage_records: List[Dict[str, Any]],
    output_path: Path,
) -> None:

• Use Python's csv module
• Write headers: timestamp, task_id, agent_id, model_name, input_tokens, output_tokens, cost_usd, call_type
• Write each record as a row

Method: export_usage_to_json

def export_usage_to_json(
    self,
    usage_records: List[Dict[str, Any]],
    output_path: Path,
) -> None:

• Use json.dump() with indent=2
• Include metadata: exported_at, total_records, total_cost

3. Register and Enhance CLI Commands

File: file:codeframe/cli/app.py

Register the metrics commands:

• Import metrics_app from file:codeframe/cli/metrics_commands.py
• Add app.add_typer(metrics_app, name="metrics") in the app registration section (around line 4153-4184)

File: file:codeframe/cli/metrics_commands.py

Enhance existing commands and add new ones:

Rename metrics_app to stats_app for consistency with spec:

• Change metrics_app = typer.Typer(name="metrics", ...) to stats_app = typer.Typer(name="stats", ...)

Enhance tokens command:

• Add --task <id> option for per-task breakdown
• Add --export <format> option (csv/json)
• Add --output <path> option for export file path
• When --task is provided, call MetricsTracker.get_token_usage(task_id=task_id)
• When --export is provided, call appropriate export method

Enhance costs command:

• Add --period <month|week|day> option for time-based filtering
• Calculate date range based on period (e.g., month = last 30 days)
• Add --export and --output options
• Display daily/weekly breakdown in table format

Add new batch command:

@stats_app.command()
def batch(
    batch_id: int,
    format: str = typer.Option("text", "--format", "-f"),
):

• Query batch execution record to get task IDs
• Call MetricsTracker.get_batch_costs(task_ids=task_ids)
• Display aggregated costs and per-task breakdown

Add new workspace command:

@stats_app.command()
def workspace(
    period: str = typer.Option("month", "--period", "-p"),
    format: str = typer.Option("text", "--format", "-f"),
):

• Call MetricsTracker.get_workspace_costs()
• Display workspace-wide statistics by project

4. Add Database Query Methods

File: file:codeframe/persistence/repositories/token_repository.py

Add helper methods for batch and workspace queries:

Method: get_batch_token_usage

def get_batch_token_usage(
    self,
    task_ids: List[int],
    start_date: Optional[datetime] = None,
    end_date: Optional[datetime] = None,
) -> List[Dict[str, Any]]:

• Query with WHERE task_id IN (?) using parameter list
• Apply date filters if provided
• Return list of usage records

Method: get_workspace_token_usage

def get_workspace_token_usage(
    self,
    start_date: Optional[datetime] = None,
    end_date: Optional[datetime] = None,
) -> List[Dict[str, Any]]:

• Query all records (no project filter)
• Apply date filters
• Order by timestamp DESC

Method: get_task_token_summary

def get_task_token_summary(
    self,
    task_id: int,
) -> Dict[str, Any]:

• Aggregate token usage for a single task
• Return total tokens, cost, and call count

5. Update Database Schema (if needed)

File: file:codeframe/persistence/schema_manager.py

Verify token_usage table has session_id column (already present at line 702):

• No schema changes needed
• Indexes already exist for efficient queries (lines 817-824)

6. Add Integration Tests

File: file:tests/lib/test_metrics_tracker.py

Add tests for new methods:

• test_get_batch_costs - Verify batch aggregation
• test_get_workspace_costs - Verify workspace-level stats
• test_export_to_csv - Verify CSV export format
• test_export_to_json - Verify JSON export format

File: file:tests/cli/test_metrics_commands.py

Add tests for enhanced CLI commands:

• test_stats_tokens_with_task_filter
• test_stats_costs_with_period
• test_stats_batch
• test_stats_workspace
• test_export_csv
• test_export_json

7. Update Documentation

File: file:README.md or File: file:docs/CLI_USAGE.md

Add documentation for new CLI commands:

## Token and Cost Tracking

### View token usage
```bash
cf stats tokens                    # Workspace token summary
cf stats tokens --task <id>        # Per-task breakdown
cf stats tokens --export csv       # Export to CSV

View costs

cf stats costs --period month      # Monthly cost report
cf stats costs --period week       # Weekly cost report
cf stats costs --export json       # Export to JSON

Batch execution costs

cf stats batch <batch_id>          # Batch execution costs

Workspace statistics

cf stats workspace                 # Workspace-wide stats

## Architecture Diagram

```mermaid
sequenceDiagram
    participant CLI as CLI (stats commands)
    participant MT as MetricsTracker
    participant TR as TokenRepository
    participant DB as Database (token_usage)
    participant Executor
    participant LLM as LLM Provider

    Note over Executor,LLM: Token Recording Flow
    Executor->>LLM: complete(messages, purpose)
    LLM-->>Executor: LLMResponse(tokens, content)
    Executor->>MT: record_token_usage(task_id, tokens, model)
    MT->>MT: calculate_cost(model, tokens)
    MT->>TR: save_token_usage(TokenUsage)
    TR->>DB: INSERT INTO token_usage

    Note over CLI,DB: Stats Retrieval Flow
    CLI->>MT: get_project_costs(project_id)
    MT->>TR: get_token_usage(project_id)
    TR->>DB: SELECT * FROM token_usage WHERE...
    DB-->>TR: usage_records[]
    TR-->>MT: usage_records[]
    MT->>MT: aggregate by agent, model
    MT-->>CLI: cost_breakdown{}
    CLI->>CLI: format_table() / export_csv()

Summary Table

Component	Action	Files Modified
Executor	Add token tracking	file:codeframe/core/executor.py
LeadAgent	Add token tracking	file:codeframe/agents/lead_agent.py
WorkerAgent	Add token tracking	file:codeframe/agents/worker_agent.py
MetricsTracker	Add batch/workspace methods	file:codeframe/lib/metrics_tracker.py
TokenRepository	Add query methods	file:codeframe/persistence/repositories/token_repository.py
CLI App	Register stats commands	file:codeframe/cli/app.py
CLI Commands	Enhance with export	file:codeframe/cli/metrics_commands.py
Tests	Add integration tests	file:tests/lib/test_metrics_tracker.py, file:tests/cli/test_metrics_commands.py
Docs	Document CLI usage	file:README.md or file:docs/CLI_USAGE.md

Import In IDE

🤖 Prompt for AI Agents

## Observations

The codebase already has substantial token tracking infrastructure in place from Sprint 10. Key components exist:
- `TokenCounter` class for tiktoken-based token counting
- `MetricsTracker` class with cost calculation and aggregation
- `TokenUsage` Pydantic model and database schema
- `TokenRepository` for persistence
- CLI commands in `file:codeframe/cli/metrics_commands.py` (not yet registered)

The main gaps are: (1) integration of token tracking into all LLM call sites, (2) CLI command registration and enhancement, (3) batch-level aggregation, (4) workspace statistics, and (5) CSV/JSON export functionality.

## Approach

The implementation will leverage existing infrastructure and add missing integration points. The strategy focuses on:
1. **Instrumentation**: Add token tracking to all LLM call sites (`Executor`, `LeadAgent`, worker agents)
2. **CLI Enhancement**: Register existing metrics commands and add new `stats` subcommands with export capabilities
3. **Aggregation**: Extend `MetricsTracker` with batch and workspace-level aggregation methods
4. **Export**: Add CSV/JSON export utilities for usage reports

This approach minimizes code duplication and builds on proven patterns already in the codebase.

## Implementation Steps

### 1. Integrate Token Tracking into LLM Call Sites

**File: `file:codeframe/core/executor.py`**

Modify `_generate_file_content` and `_generate_edit_content` methods to record token usage after each LLM call:

- Import `MetricsTracker` and `CallType` from `file:codeframe/lib/metrics_tracker.py` and `file:codeframe/core/models.py`
- Add `metrics_tracker: MetricsTracker` parameter to `Executor.__init__`
- After `self.llm.complete()` calls (lines 627-633, 657-663), extract token counts from `response.input_tokens` and `response.output_tokens`
- Call `await metrics_tracker.record_token_usage()` with:
  - `task_id` from context
  - `agent_id` from context or "executor"
  - `project_id` from context
  - `model_name` from `response.model`
  - Token counts and `CallType.TASK_EXECUTION`

**File: `file:codeframe/agents/lead_agent.py`**

Integrate token tracking in the `chat` method and other LLM call sites:

- Add `self.metrics_tracker = MetricsTracker(db=self.db)` in `__init__`
- In `chat` method (around line 200-250), after `self.provider.send_message()`, record token usage from response
- In `_generate_next_discovery_question`, `generate_prd`, and other LLM-calling methods, add token tracking
- Use `CallType.COORDINATION` for discovery/planning calls, `CallType.TASK_EXECUTION` for task-related calls

**File: `file:codeframe/agents/worker_agent.py`** (if exists)

Similar integration pattern as `LeadAgent`:
- Initialize `MetricsTracker` in constructor
- Record token usage after each LLM call
- Use appropriate `CallType` based on operation

**File: `file:codeframe/adapters/llm/anthropic.py`**

Enhance `send_message` to return token usage in response:

- Ensure the response dictionary includes `usage` with `input_tokens` and `output_tokens` (already present at lines 104-107)
- No changes needed, just verify consistency

### 2. Enhance MetricsTracker with Batch and Workspace Aggregation

**File: `file:codeframe/lib/metrics_tracker.py`**

Add new methods for batch and workspace-level statistics:

**Method: `get_batch_costs`**
```python
async def get_batch_costs(
    self,
    project_id: int,
    task_ids: List[int],
    start_date: Optional[datetime] = None,
    end_date: Optional[datetime] = None,
) -> Dict[str, Any]:
```
- Query `token_usage` table filtering by `task_id IN task_ids`
- Aggregate total cost, tokens, and calls
- Return breakdown by task with individual costs

**Method: `get_workspace_costs`**
```python
async def get_workspace_costs(
    self,
    start_date: Optional[datetime] = None,
    end_date: Optional[datetime] = None,
) -> Dict[str, Any]:
```
- Query all projects (no project_id filter)
- Aggregate by project, agent, and model
- Return workspace-wide statistics

**Method: `export_usage_to_csv`**
```python
def export_usage_to_csv(
    self,
    usage_records: List[Dict[str, Any]],
    output_path: Path,
) -> None:
```
- Use Python's `csv` module
- Write headers: `timestamp, task_id, agent_id, model_name, input_tokens, output_tokens, cost_usd, call_type`
- Write each record as a row

**Method: `export_usage_to_json`**
```python
def export_usage_to_json(
    self,
    usage_records: List[Dict[str, Any]],
    output_path: Path,
) -> None:
```
- Use `json.dump()` with `indent=2`
- Include metadata: `exported_at`, `total_records`, `total_cost`

### 3. Register and Enhance CLI Commands

**File: `file:codeframe/cli/app.py`**

Register the metrics commands:

- Import `metrics_app` from `file:codeframe/cli/metrics_commands.py`
- Add `app.add_typer(metrics_app, name="metrics")` in the app registration section (around line 4153-4184)

**File: `file:codeframe/cli/metrics_commands.py`**

Enhance existing commands and add new ones:

**Rename `metrics_app` to `stats_app`** for consistency with spec:
- Change `metrics_app = typer.Typer(name="metrics", ...)` to `stats_app = typer.Typer(name="stats", ...)`

**Enhance `tokens` command**:
- Add `--task <id>` option for per-task breakdown
- Add `--export <format>` option (csv/json)
- Add `--output <path>` option for export file path
- When `--task` is provided, call `MetricsTracker.get_token_usage(task_id=task_id)`
- When `--export` is provided, call appropriate export method

**Enhance `costs` command**:
- Add `--period <month|week|day>` option for time-based filtering
- Calculate date range based on period (e.g., `month` = last 30 days)
- Add `--export` and `--output` options
- Display daily/weekly breakdown in table format

**Add new `batch` command**:
```python
@stats_app.command()
def batch(
    batch_id: int,
    format: str = typer.Option("text", "--format", "-f"),
):
```
- Query batch execution record to get task IDs
- Call `MetricsTracker.get_batch_costs(task_ids=task_ids)`
- Display aggregated costs and per-task breakdown

**Add new `workspace` command**:
```python
@stats_app.command()
def workspace(
    period: str = typer.Option("month", "--period", "-p"),
    format: str = typer.Option("text", "--format", "-f"),
):
```
- Call `MetricsTracker.get_workspace_costs()`
- Display workspace-wide statistics by project

### 4. Add Database Query Methods

**File: `file:codeframe/persistence/repositories/token_repository.py`**

Add helper methods for batch and workspace queries:

**Method: `get_batch_token_usage`**
```python
def get_batch_token_usage(
    self,
    task_ids: List[int],
    start_date: Optional[datetime] = None,
    end_date: Optional[datetime] = None,
) -> List[Dict[str, Any]]:
```
- Query with `WHERE task_id IN (?)` using parameter list
- Apply date filters if provided
- Return list of usage records

**Method: `get_workspace_token_usage`**
```python
def get_workspace_token_usage(
    self,
    start_date: Optional[datetime] = None,
    end_date: Optional[datetime] = None,
) -> List[Dict[str, Any]]:
```
- Query all records (no project filter)
- Apply date filters
- Order by timestamp DESC

**Method: `get_task_token_summary`**
```python
def get_task_token_summary(
    self,
    task_id: int,
) -> Dict[str, Any]:
```
- Aggregate token usage for a single task
- Return total tokens, cost, and call count

### 5. Update Database Schema (if needed)

**File: `file:codeframe/persistence/schema_manager.py`**

Verify `token_usage` table has `session_id` column (already present at line 702):
- No schema changes needed
- Indexes already exist for efficient queries (lines 817-824)

### 6. Add Integration Tests

**File: `file:tests/lib/test_metrics_tracker.py`**

Add tests for new methods:
- `test_get_batch_costs` - Verify batch aggregation
- `test_get_workspace_costs` - Verify workspace-level stats
- `test_export_to_csv` - Verify CSV export format
- `test_export_to_json` - Verify JSON export format

**File: `file:tests/cli/test_metrics_commands.py`**

Add tests for enhanced CLI commands:
- `test_stats_tokens_with_task_filter`
- `test_stats_costs_with_period`
- `test_stats_batch`
- `test_stats_workspace`
- `test_export_csv`
- `test_export_json`

### 7. Update Documentation

**File: `file:README.md`** or **File: `file:docs/CLI_USAGE.md`**

Add documentation for new CLI commands:

```markdown
## Token and Cost Tracking

### View token usage
```bash
cf stats tokens                    # Workspace token summary
cf stats tokens --task <id>        # Per-task breakdown
cf stats tokens --export csv       # Export to CSV
```

### View costs
```bash
cf stats costs --period month      # Monthly cost report
cf stats costs --period week       # Weekly cost report
cf stats costs --export json       # Export to JSON
```

### Batch execution costs
```bash
cf stats batch <batch_id>          # Batch execution costs
```

### Workspace statistics
```bash
cf stats workspace                 # Workspace-wide stats
```
```

## Architecture Diagram

```mermaid
sequenceDiagram
    participant CLI as CLI (stats commands)
    participant MT as MetricsTracker
    participant TR as TokenRepository
    participant DB as Database (token_usage)
    participant Executor
    participant LLM as LLM Provider

    Note over Executor,LLM: Token Recording Flow
    Executor->>LLM: complete(messages, purpose)
    LLM-->>Executor: LLMResponse(tokens, content)
    Executor->>MT: record_token_usage(task_id, tokens, model)
    MT->>MT: calculate_cost(model, tokens)
    MT->>TR: save_token_usage(TokenUsage)
    TR->>DB: INSERT INTO token_usage

    Note over CLI,DB: Stats Retrieval Flow
    CLI->>MT: get_project_costs(project_id)
    MT->>TR: get_token_usage(project_id)
    TR->>DB: SELECT * FROM token_usage WHERE...
    DB-->>TR: usage_records[]
    TR-->>MT: usage_records[]
    MT->>MT: aggregate by agent, model
    MT-->>CLI: cost_breakdown{}
    CLI->>CLI: format_table() / export_csv()
```

## Summary Table

| Component | Action | Files Modified |
|-----------|--------|----------------|
| **Executor** | Add token tracking | `file:codeframe/core/executor.py` |
| **LeadAgent** | Add token tracking | `file:codeframe/agents/lead_agent.py` |
| **WorkerAgent** | Add token tracking | `file:codeframe/agents/worker_agent.py` |
| **MetricsTracker** | Add batch/workspace methods | `file:codeframe/lib/metrics_tracker.py` |
| **TokenRepository** | Add query methods | `file:codeframe/persistence/repositories/token_repository.py` |
| **CLI App** | Register stats commands | `file:codeframe/cli/app.py` |
| **CLI Commands** | Enhance with export | `file:codeframe/cli/metrics_commands.py` |
| **Tests** | Add integration tests | `file:tests/lib/test_metrics_tracker.py`, `file:tests/cli/test_metrics_commands.py` |
| **Docs** | Document CLI usage | `file:README.md` or `file:docs/CLI_USAGE.md` |

---

Execution Information

Branch: main
Commit: e768c1d

---

💡 Tips

Supported Commands (Inside Comments)

• Use @traycerai generate to iterate on the previous version of the implementation plan.

Supported Commands (Inside Description)

• Add @traycerai ignore anywhere in the ticket description to prevent this ticket from being processed.
• Add @traycerai branch:<branch-name> anywhere in the ticket description to specify the target branch for the implementation plan.

Community

• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

📝 CodeRabbit Plan Mode

Generate an implementation plan and prompts that you can use with your favorite coding agent.

• Create Plan

Examples

• Example 1
• Example 2

---

🔗 Related PRs

#126 - Add token tracking to WorkerAgent execute_task method [merged]

---

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord!