[P1] Add token tracking after LLM calls in WorkerAgent

[frankbria]

Overview

After implementing actual LLM calls in execute_task(), we need to track token usage for cost monitoring and metrics.

Current State

File: codeframe/agents/worker_agent.py line 98

# TODO: Add token tracking after LLM call (see docstring example)

Expected Behavior

After each LLM API call:

• Extract input_tokens and output_tokens from response
• Record token usage via MetricsTracker.record_token_usage()
• Store in database with task_id, agent_id, model_name, timestamp
• Calculate cost using MODEL_PRICING constants

Implementation Requirements

1. Use the pattern shown in the docstring (lines 82-93)
2. Call MetricsTracker.record_token_usage() after each API call
3. Handle cases where usage info is missing (graceful degradation)
4. Support multiple models (Sonnet 4.5, Opus 4, Haiku 4)

Acceptance Criteria

• Token usage recorded after every LLM call
• Data includes task_id, agent_id, project_id, model, tokens
• Cost calculated correctly using MODEL_PRICING
• Dashboard displays real-time token usage
• Historical metrics queryable via API

Priority

P1 - Critical: Required for cost monitoring and preventing budget overruns.

Dependencies

• Requires: [P1] Implement WorkerAgent.execute_task() with Anthropic API integration #98 (Implement execute_task with API calls)

References

• Docstring example: worker_agent.py lines 82-93
• MetricsTracker: codeframe/lib/metrics_tracker.py
• User Documentation: CODEFRAME_USER_DOCUMENTATION.md lines 299-304

[traycerai]

Plan

Observations

The WorkerAgent's execute_task() method currently returns a mock response (line 99) without making actual LLM calls. The docstring (lines 82-93) provides a clear example of how token tracking should be implemented. The HybridWorkerAgent demonstrates a working implementation pattern with the _record_token_usage() helper method (lines 308-344). The MetricsTracker is fully implemented with support for all three Claude models (Sonnet 4.5, Opus 4, Haiku 4) and includes cost calculation using MODEL_PRICING constants. The issue depends on #98 for actual LLM API integration.

Approach

The implementation will follow the pattern established in HybridWorkerAgent by creating a helper method _record_token_usage() that extracts token information from LLM responses and calls MetricsTracker. This approach provides graceful degradation when usage information is missing and maintains consistency across the codebase. The implementation will be added to the base WorkerAgent class so all specialized workers (Backend, Frontend, Test, Review) inherit the functionality automatically. Token tracking will be integrated into the execute_task() method after LLM calls are implemented, using CallType.TASK_EXECUTION as the default call type.

Implementation Steps

1. Add Token Tracking Helper Method to WorkerAgent

File: file:codeframe/agents/worker_agent.py

Add a new private method _record_token_usage() after the execute_task() method (around line 100):

async def _record_token_usage(self, task: Task, response: Dict[str, Any]) -> None:
    """Record token usage metrics for this task.
    
    Extracts token usage from LLM response and records it via MetricsTracker.
    Handles graceful degradation if usage information is missing.
    
    Args:
        task: Task that was executed
        response: LLM provider response with usage info
    """

Implementation details:

• Import MetricsTracker at the top of the method: from codeframe.lib.metrics_tracker import MetricsTracker
• Extract usage from response: usage = response.get("usage", {})
• Get input_tokens and output_tokens with defaults: input_tokens = usage.get("input_tokens", 0) and output_tokens = usage.get("output_tokens", 0)
• Return early if both are 0 (no tokens to record)
• Get project_id from task: project_id = task.project_id if hasattr(task, "project_id") else None
• Log warning and return if project_id is None
• Determine model_name from self.provider or use a default (e.g., "claude-sonnet-4-5")
• Create MetricsTracker instance: tracker = MetricsTracker(db=self.db)
• Call await tracker.record_token_usage() with all required parameters
• Wrap in try-except to catch any exceptions and log warnings without failing task execution
• Use CallType.TASK_EXECUTION as the call_type parameter
• Pass session_id=None (or extract from response if available)

Reference pattern: Follow the implementation in file:codeframe/agents/hybrid_worker.py lines 308-344

2. Integrate Token Tracking into execute_task()

File: file:codeframe/agents/worker_agent.py

Modify the execute_task() method (lines 67-99) to call the token tracking helper:

Location: After line 98 where the TODO comment exists

Steps:

• Remove the TODO comment at line 98
• After the LLM API call is made (when [P1] Implement WorkerAgent.execute_task() with Anthropic API integration #98 is implemented), add: await self._record_token_usage(task, response)
• Ensure this is called before returning the result
• The call should be placed after the LLM response is received but before any error handling that might return early

Note: Since actual LLM integration depends on issue #98, the token tracking call should be added in the location where the LLM response will be available. The current mock response at line 99 should be replaced with actual LLM call + token tracking.

3. Handle Missing Usage Information

Implementation considerations:

The _record_token_usage() helper method should implement graceful degradation:

• Check if usage dictionary exists in response
• Check if input_tokens and output_tokens are present and non-zero
• Log appropriate warnings when usage info is missing
• Never raise exceptions that would fail task execution
• Return silently if no tokens to record

Error scenarios to handle:

• Response has no "usage" key
• Usage dictionary is empty
• Token counts are 0 or negative
• Model name is unknown (MetricsTracker will raise ValueError)
• Database connection issues
• Project context is missing

4. Model Name Resolution

File: file:codeframe/agents/worker_agent.py

The model name needs to be determined for cost calculation:

Options:

1. Use self.provider if it contains model information
2. Add a model_name attribute to WorkerAgent.init()
3. Extract from LLM response if available
4. Use a default based on agent maturity level

Recommended approach:

• Add optional model_name parameter to __init__() with default "claude-sonnet-4-5"
• Store as self.model_name instance attribute
• Use this in _record_token_usage() for MetricsTracker calls
• Allow specialized workers to override with their preferred models

5. Update Specialized Worker Agents

Files to verify:

• file:codeframe/agents/backend_worker_agent.py
• file:codeframe/agents/frontend_worker_agent.py
• file:codeframe/agents/test_worker_agent.py
• file:codeframe/agents/review_worker_agent.py

Verification steps:

• Ensure they inherit from WorkerAgent and will automatically get token tracking
• Check if they override execute_task() - if so, they need to call the parent method or implement token tracking themselves
• Backend and Frontend workers have their own execute_task() implementations that will need token tracking added
• Review and Test workers should be checked for similar patterns

6. Dashboard Integration Verification

Files to check:

• file:web-ui/src/components/metrics/TokenUsageChart.tsx
• file:web-ui/src/components/metrics/CostDashboard.tsx
• file:codeframe/ui/routers/metrics.py

Verification:

• Confirm that existing dashboard components can display the recorded token usage
• Check that API endpoints in file:codeframe/ui/routers/metrics.py are properly wired
• Verify WebSocket broadcasts for real-time updates (if implemented)
• No changes should be needed if MetricsTracker is working correctly

7. Testing Considerations

Test file: file:tests/lib/test_metrics_tracker.py

Verification:

• Existing tests cover MetricsTracker functionality comprehensively
• Tests verify cost calculation for all three models
• Tests verify aggregation and timeline queries
• No new tests needed for MetricsTracker itself

New tests needed:

• Test WorkerAgent._record_token_usage() with valid response
• Test graceful degradation with missing usage info
• Test with different model names
• Test error handling scenarios
• Mock MetricsTracker to verify it's called with correct parameters

Implementation Flow

sequenceDiagram
    participant Task
    participant WorkerAgent
    participant LLMProvider
    participant MetricsTracker
    participant Database
    participant Dashboard

    Task->>WorkerAgent: execute_task(task)
    WorkerAgent->>WorkerAgent: Set current_task
    WorkerAgent->>LLMProvider: send_message(prompt)
    LLMProvider-->>WorkerAgent: response {content, usage}
    WorkerAgent->>WorkerAgent: _record_token_usage(task, response)
    WorkerAgent->>MetricsTracker: record_token_usage(...)
    MetricsTracker->>MetricsTracker: calculate_cost(model, tokens)
    MetricsTracker->>Database: save_token_usage(TokenUsage)
    Database-->>MetricsTracker: usage_id
    MetricsTracker-->>WorkerAgent: usage_id
    WorkerAgent-->>Task: {status, output, usage}
    Dashboard->>Database: Query token usage
    Database-->>Dashboard: Display metrics

Loading

Key Files Reference

File	Purpose	Lines of Interest
file:codeframe/agents/worker_agent.py	Base worker agent - add token tracking here	67-99 (execute_task), 78-93 (docstring example)
file:codeframe/agents/hybrid_worker.py	Reference implementation	308-344 (_record_token_usage)
file:codeframe/lib/metrics_tracker.py	Token tracking service	126-203 (record_token_usage), 47-51 (MODEL_PRICING)
file:codeframe/core/models.py	Data models	102-108 (CallType), 701-747 (TokenUsage)
file:codeframe/providers/anthropic.py	LLM response structure	55-115 (send_message)
file:tests/lib/test_metrics_tracker.py	Test examples	All tests show expected behavior

Dependencies

• Prerequisite: Issue [P1] Implement WorkerAgent.execute_task() with Anthropic API integration #98 must be completed first (Implement execute_task with API calls)
• Database: Requires migration_007 applied (token_usage table exists)
• MetricsTracker: Already fully implemented and tested
• MODEL_PRICING: Already defined with current Anthropic pricing

Acceptance Criteria Mapping

Criterion	Implementation Location
Token usage recorded after every LLM call	WorkerAgent._record_token_usage() called in execute_task()
Data includes task_id, agent_id, project_id, model, tokens	All parameters passed to MetricsTracker.record_token_usage()
Cost calculated correctly using MODEL_PRICING	MetricsTracker.calculate_cost() uses MODEL_PRICING constants
Dashboard displays real-time token usage	Existing dashboard components query database via metrics API
Historical metrics queryable via API	Existing endpoints in file:codeframe/ui/routers/metrics.py

Import In IDE

🤖 Prompt for AI Agents

## Observations

The WorkerAgent's `execute_task()` method currently returns a mock response (line 99) without making actual LLM calls. The docstring (lines 82-93) provides a clear example of how token tracking should be implemented. The HybridWorkerAgent demonstrates a working implementation pattern with the `_record_token_usage()` helper method (lines 308-344). The MetricsTracker is fully implemented with support for all three Claude models (Sonnet 4.5, Opus 4, Haiku 4) and includes cost calculation using MODEL_PRICING constants. The issue depends on #98 for actual LLM API integration.

## Approach

The implementation will follow the pattern established in HybridWorkerAgent by creating a helper method `_record_token_usage()` that extracts token information from LLM responses and calls MetricsTracker. This approach provides graceful degradation when usage information is missing and maintains consistency across the codebase. The implementation will be added to the base WorkerAgent class so all specialized workers (Backend, Frontend, Test, Review) inherit the functionality automatically. Token tracking will be integrated into the execute_task() method after LLM calls are implemented, using CallType.TASK_EXECUTION as the default call type.

## Implementation Steps

### 1. Add Token Tracking Helper Method to WorkerAgent

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

Add a new private method `_record_token_usage()` after the `execute_task()` method (around line 100):

```python
async def _record_token_usage(self, task: Task, response: Dict[str, Any]) -> None:
    """Record token usage metrics for this task.
    
    Extracts token usage from LLM response and records it via MetricsTracker.
    Handles graceful degradation if usage information is missing.
    
    Args:
        task: Task that was executed
        response: LLM provider response with usage info
    """
```

**Implementation details:**
- Import MetricsTracker at the top of the method: `from codeframe.lib.metrics_tracker import MetricsTracker`
- Extract usage from response: `usage = response.get("usage", {})`
- Get input_tokens and output_tokens with defaults: `input_tokens = usage.get("input_tokens", 0)` and `output_tokens = usage.get("output_tokens", 0)`
- Return early if both are 0 (no tokens to record)
- Get project_id from task: `project_id = task.project_id if hasattr(task, "project_id") else None`
- Log warning and return if project_id is None
- Determine model_name from self.provider or use a default (e.g., "claude-sonnet-4-5")
- Create MetricsTracker instance: `tracker = MetricsTracker(db=self.db)`
- Call `await tracker.record_token_usage()` with all required parameters
- Wrap in try-except to catch any exceptions and log warnings without failing task execution
- Use `CallType.TASK_EXECUTION` as the call_type parameter
- Pass `session_id=None` (or extract from response if available)

**Reference pattern:** Follow the implementation in `file:codeframe/agents/hybrid_worker.py` lines 308-344

### 2. Integrate Token Tracking into execute_task()

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

Modify the `execute_task()` method (lines 67-99) to call the token tracking helper:

**Location:** After line 98 where the TODO comment exists

**Steps:**
- Remove the TODO comment at line 98
- After the LLM API call is made (when #98 is implemented), add: `await self._record_token_usage(task, response)`
- Ensure this is called before returning the result
- The call should be placed after the LLM response is received but before any error handling that might return early

**Note:** Since actual LLM integration depends on issue #98, the token tracking call should be added in the location where the LLM response will be available. The current mock response at line 99 should be replaced with actual LLM call + token tracking.

### 3. Handle Missing Usage Information

**Implementation considerations:**

The `_record_token_usage()` helper method should implement graceful degradation:
- Check if usage dictionary exists in response
- Check if input_tokens and output_tokens are present and non-zero
- Log appropriate warnings when usage info is missing
- Never raise exceptions that would fail task execution
- Return silently if no tokens to record

**Error scenarios to handle:**
- Response has no "usage" key
- Usage dictionary is empty
- Token counts are 0 or negative
- Model name is unknown (MetricsTracker will raise ValueError)
- Database connection issues
- Project context is missing

### 4. Model Name Resolution

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

The model name needs to be determined for cost calculation:

**Options:**
1. Use `self.provider` if it contains model information
2. Add a `model_name` attribute to WorkerAgent.__init__()
3. Extract from LLM response if available
4. Use a default based on agent maturity level

**Recommended approach:**
- Add optional `model_name` parameter to `__init__()` with default "claude-sonnet-4-5"
- Store as `self.model_name` instance attribute
- Use this in `_record_token_usage()` for MetricsTracker calls
- Allow specialized workers to override with their preferred models

### 5. Update Specialized Worker Agents

**Files to verify:**
- `file:codeframe/agents/backend_worker_agent.py`
- `file:codeframe/agents/frontend_worker_agent.py`
- `file:codeframe/agents/test_worker_agent.py`
- `file:codeframe/agents/review_worker_agent.py`

**Verification steps:**
- Ensure they inherit from WorkerAgent and will automatically get token tracking
- Check if they override execute_task() - if so, they need to call the parent method or implement token tracking themselves
- Backend and Frontend workers have their own execute_task() implementations that will need token tracking added
- Review and Test workers should be checked for similar patterns

### 6. Dashboard Integration Verification

**Files to check:**
- `file:web-ui/src/components/metrics/TokenUsageChart.tsx`
- `file:web-ui/src/components/metrics/CostDashboard.tsx`
- `file:codeframe/ui/routers/metrics.py`

**Verification:**
- Confirm that existing dashboard components can display the recorded token usage
- Check that API endpoints in `file:codeframe/ui/routers/metrics.py` are properly wired
- Verify WebSocket broadcasts for real-time updates (if implemented)
- No changes should be needed if MetricsTracker is working correctly

### 7. Testing Considerations

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

**Verification:**
- Existing tests cover MetricsTracker functionality comprehensively
- Tests verify cost calculation for all three models
- Tests verify aggregation and timeline queries
- No new tests needed for MetricsTracker itself

**New tests needed:**
- Test WorkerAgent._record_token_usage() with valid response
- Test graceful degradation with missing usage info
- Test with different model names
- Test error handling scenarios
- Mock MetricsTracker to verify it's called with correct parameters

## Implementation Flow

```mermaid
sequenceDiagram
    participant Task
    participant WorkerAgent
    participant LLMProvider
    participant MetricsTracker
    participant Database
    participant Dashboard

    Task->>WorkerAgent: execute_task(task)
    WorkerAgent->>WorkerAgent: Set current_task
    WorkerAgent->>LLMProvider: send_message(prompt)
    LLMProvider-->>WorkerAgent: response {content, usage}
    WorkerAgent->>WorkerAgent: _record_token_usage(task, response)
    WorkerAgent->>MetricsTracker: record_token_usage(...)
    MetricsTracker->>MetricsTracker: calculate_cost(model, tokens)
    MetricsTracker->>Database: save_token_usage(TokenUsage)
    Database-->>MetricsTracker: usage_id
    MetricsTracker-->>WorkerAgent: usage_id
    WorkerAgent-->>Task: {status, output, usage}
    Dashboard->>Database: Query token usage
    Database-->>Dashboard: Display metrics
```

## Key Files Reference

| File | Purpose | Lines of Interest |
|------|---------|-------------------|
| `file:codeframe/agents/worker_agent.py` | Base worker agent - add token tracking here | 67-99 (execute_task), 78-93 (docstring example) |
| `file:codeframe/agents/hybrid_worker.py` | Reference implementation | 308-344 (_record_token_usage) |
| `file:codeframe/lib/metrics_tracker.py` | Token tracking service | 126-203 (record_token_usage), 47-51 (MODEL_PRICING) |
| `file:codeframe/core/models.py` | Data models | 102-108 (CallType), 701-747 (TokenUsage) |
| `file:codeframe/providers/anthropic.py` | LLM response structure | 55-115 (send_message) |
| `file:tests/lib/test_metrics_tracker.py` | Test examples | All tests show expected behavior |

## Dependencies

- **Prerequisite:** Issue #98 must be completed first (Implement execute_task with API calls)
- **Database:** Requires migration_007 applied (token_usage table exists)
- **MetricsTracker:** Already fully implemented and tested
- **MODEL_PRICING:** Already defined with current Anthropic pricing

## Acceptance Criteria Mapping

| Criterion | Implementation Location |
|-----------|------------------------|
| Token usage recorded after every LLM call | WorkerAgent._record_token_usage() called in execute_task() |
| Data includes task_id, agent_id, project_id, model, tokens | All parameters passed to MetricsTracker.record_token_usage() |
| Cost calculated correctly using MODEL_PRICING | MetricsTracker.calculate_cost() uses MODEL_PRICING constants |
| Dashboard displays real-time token usage | Existing dashboard components query database via metrics API |
| Historical metrics queryable via API | Existing endpoints in file:codeframe/ui/routers/metrics.py |

---

Execution Information

Branch: main
Commit: 1a87e4e

---

💡 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

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

🔗 Similar Issues

Possible Duplicates

• [P1] Implement WorkerAgent.execute_task() with Anthropic API integration #98

Related Issues

• Implement advanced Metrics UI features (token stats, filtering, export, trends) #46
• Optimize Context Assembly Order for LLM Cache Hits #65
• Optimize Context Assembly Order for LLM Cache Hits #66
• Optimize Context Assembly Order for LLM Cache Hits #64

👤 Suggested Assignees

• frankbria

---

🧪 Issue enrichment is currently in open beta.

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 or schedule a call!

[frankbria]

Closed with PR #126