Write

# Autocoder 2.0 Enhancement Implementation Plan

## Overview

This plan implements the priority enhancements for Autocoder 2.0 - an autonomous coding agent system. Based on comprehensive codebase exploration, I've identified the architecture patterns and integration points needed to add token tracking, session recovery, enhanced error handling, and multi-model support while preserving all existing functionality.

## Architecture Understanding

Autocoder 2.0 follows these patterns that I'll extend:
- **Singleton managers with thread safety** (process_manager.py)
- **SQLAlchemy models with migrations** (api/database.py)
- **WebSocket real-time updates** (server/websocket.py)
- **React Query hooks + WebSocket state** (ui/src/hooks/)
- **MCP server for agent tools** (mcp_server/feature_mcp.py)
- **Async subprocess management** (server/services/)

## Priority 1: Token and Cost Tracking

### Database Schema Extensions

Add three new tables to `api/database.py`:

```python
class UsageSession(Base):
    __tablename__ = "usage_sessions"

    id = Column(Integer, primary_key=True)
    session_id = Column(String(100), unique=True, index=True)
    started_at = Column(DateTime, default=datetime.utcnow, index=True)
    ended_at = Column(DateTime, nullable=True)

    total_input_tokens = Column(Integer, default=0)
    total_output_tokens = Column(Integer, default=0)
    total_calls = Column(Integer, default=0)
    total_cost_cents = Column(Integer, default=0)

    model = Column(String(50))
    status = Column(String(20))  # "active", "completed", "failed"

class APICall(Base):
    __tablename__ = "api_calls"

    id = Column(Integer, primary_key=True)
    session_id = Column(String(100), ForeignKey("usage_sessions.session_id"), index=True)
    timestamp = Column(DateTime, default=datetime.utcnow, index=True)

    input_tokens = Column(Integer)
    output_tokens = Column(Integer)
    stop_reason = Column(String(50))
    response_time_ms = Column(Integer)
    cost_cents = Column(Integer)

class UsageSummary(Base):
    __tablename__ = "usage_summary"

    id = Column(Integer, primary_key=True)
    session_id = Column(String(100), unique=True, index=True)
    total_cost_cents = Column(Integer, default=0)
    total_input_tokens = Column(Integer, default=0)
    total_output_tokens = Column(Integer, default=0)
```

**Migration strategy**: Extend `_migrate_add_in_progress_column()` pattern to add usage tables.

### Token Extraction Integration

Modify `agent.py` and `assistant_chat_session.py` to capture usage metadata:

**In `agent.py` line 34-105:**
```python
async def run_agent_session(client, message, project_dir):
    session_id = str(uuid.uuid4())
    start_usage_session(session_id, client.model)

    start_time = time.time()
    await client.query(message)

    async for msg in client.receive_response():
        # Existing response handling...
        pass

    # NEW: Extract usage from final message
    if hasattr(msg, 'usage'):
        track_api_call(
            session_id=session_id,
            input_tokens=msg.usage.input_tokens,
            output_tokens=msg.usage.output_tokens,
            duration_ms=int((time.time() - start_time) * 1000)
        )
```

**Integration points**:
- Hook into `client.receive_response()` completion
- Use Claude SDK's `.usage` attribute on response messages
- Store in database via new service functions

### Real-time Cost Updates

Extend WebSocket system in `server/websocket.py`:

**Add new message type:**
```python
{
    "type": "usage_update",
    "total_input_tokens": int,
    "total_output_tokens": int,
    "total_calls": int,
    "estimated_cost": "$X.XX",
    "timestamp": ISO_string
}
```

**Polling approach** (follows existing `poll_progress()` pattern):
```python
async def poll_usage(websocket, project_name, session_id):
    last_calls = -1
    while True:
        usage = get_usage_summary(session_id)
        if usage['total_calls'] != last_calls:
            await websocket.send_json({
                "type": "usage_update",
                **usage,
                "estimated_cost": format_cost(usage['cost_cents'])
            })
        await asyncio.sleep(1)
```

### UI Integration

**Extend WebSocket state** in `ui/src/hooks/useWebSocket.ts`:
```typescript
interface WebSocketState {
  progress: {...}
  agentStatus: AgentStatus
  logs: Array<{...}>
  isConnected: boolean

  // NEW:
  usage: {
    totalInputTokens: number
    totalOutputTokens: number
    totalCalls: number
    estimatedCost: string
    lastUpdated: string
  }
}
```

**Cost display options**:
1. **Header badge** - Real-time cost next to agent status
2. **Extended ProgressDashboard** - Usage section with token breakdown
3. **Separate UsagePanel** - Dedicated component for cost tracking

## Priority 2: Session Recovery and Checkpointing

### Session Journal System

Create `.session_journal/` directory in each project:

```
project_dir/.session_journal/
├── session_001.json
├── session_002.json
└── current_session.json
```

**Journal schema:**
```json
{
  "session_id": "uuid",
  "started_at": "ISO_timestamp",
  "ended_at": "ISO_timestamp",
  "feature_id": 42,
  "initial_stats": {"passing": 10, "total": 50},
  "final_stats": {"passing": 12, "total": 50},
  "messages": [
    {
      "role": "user|assistant",
      "content": "...",
      "tools_used": ["feature_mark_passing"],
      "timestamp": "ISO"
    }
  ],
  "blockers": ["Dependency on feature #5"],
  "decisions": ["Skipped feature #10 - unclear requirements"],
  "errors": ["Tool Write failed with error X"],
  "status": "success|incomplete|error"
}
```

### Feature Work Metadata

Extend Feature model with recovery columns:
```python
# Add to Feature class in api/database.py:
last_worked_on = Column(DateTime)
work_session_id = Column(String(100))
attempt_count = Column(Integer, default=0)
last_error = Column(Text)
recovery_notes = Column(Text)
```

### Recovery Context Injection

Modify `prompts.py` to include recovery context:
```python
def get_recovery_context(project_dir: Path) -> str:
    last_session = load_latest_session_journal(project_dir)
    if not last_session:
        return ""

    return f"""
## Recovery Context from Previous Session

**Previous Status**: {last_session.get('status')}
**Features Completed**: {last_session['final_stats']['passing']}
**Blockers**: {json.dumps(last_session.get('blockers', []))}
**Errors**: {json.dumps(last_session.get('errors', []))}
**Decisions**: {json.dumps(last_session.get('decisions', []))}
"""
```

### Stale Work Detection

Add recovery module to detect and handle stuck features:
```python
# NEW: server/services/session_recovery.py
def detect_stale_work(project_dir, timeout_minutes=60):
    """Find features stuck in in_progress state."""
    return session.query(Feature).filter(
        Feature.in_progress == True,
        Feature.last_worked_on < datetime.now() - timedelta(minutes=timeout_minutes)
    ).all()

def recover_stale_work(project_dir):
    """Reset stale work or add recovery notes."""
    for feature in detect_stale_work(project_dir):
        journal = load_session_journal(feature.work_session_id)
        if journal.get("status") == "error":
            feature.attempt_count += 1
            if feature.attempt_count > 3:
                feature_skip(feature.id)  # Move to end of queue
```

## Priority 3: Enhanced Error Recovery

### Stuck State Detection

Monitor for repeated failures on same feature:
- Track attempt count in Feature model
- Detect same error patterns via session journal
- Auto-escalate to Opus after 3 Sonnet failures
- Skip features after 5 total attempts

### Model Escalation Chain

**Add to config.py:**
```python
@dataclass
class ModelConfig:
    default_model: str = "claude-sonnet-4-20250514"
    escalation_enabled: bool = True
    escalation_threshold: int = 3  # failures before escalating
    stuck_timeout_minutes: int = 30

    # Model routing rules
    opus_keywords: list[str] = ["refactor", "architecture", "complex"]
    haiku_keywords: list[str] = ["simple", "test", "docs", "lint"]
```

**Escalation logic:**
```python
def should_escalate_model(feature_id: int, current_model: str) -> str:
    """Determine if model escalation is needed."""
    feature = get_feature(feature_id)

    if feature.attempt_count >= 3 and current_model != "claude-opus-4-5":
        return "claude-opus-4-5"  # Escalate to Opus

    # Check for complexity keywords
    if any(keyword in feature.description.lower() for keyword in OPUS_KEYWORDS):
        return "claude-opus-4-5"

    return current_model
```

### Skip and Retry Mechanisms

**UI Controls:**
- "Retry with Opus" button for stuck features
- "Skip for now" button to move to end of queue
- Bulk skip operations via checkbox selection

**Auto-skip criteria:**
- Features failed 5+ times across all models
- Features blocked by dependencies for >24 hours
- Features with unresolvable errors (missing requirements, etc.)

## Priority 4: Multi-Model Support

### Configuration-Driven Routing

**Extend client.py:**
```python
def get_model_for_task(task_description: str, feature_category: str, attempt_count: int) -> str:
    """Route task to appropriate model based on complexity."""

    # Escalation chain
    if attempt_count >= 3:
        return "claude-opus-4-5"

    # Keyword-based routing
    description_lower = task_description.lower()

    if any(kw in description_lower for kw in OPUS_KEYWORDS):
        return "claude-opus-4-5"
    elif any(kw in description_lower for kw in HAIKU_KEYWORDS):
        return "claude-3-5-haiku"
    else:
        return "claude-sonnet-4-20250514"  # Default
```

**Model specialization:**
- **Opus**: Complex architectural decisions, debugging difficult failures, large refactors
- **Sonnet**: Standard coding tasks (current default)
- **Haiku**: Simple validations, linting, documentation, quick tests

## Git Integration Assessment

**Discovery**: Autocoder 2.0 already has comprehensive git integration:
- `git_agent.py` - Full orchestration with auto-commits on feature completion
- `git_operations.py` - Async git wrapper with all operations
- `server/routers/git.py` - REST API endpoints
- Auto-commit on feature passing, session end
- Real-time git status via WebSocket
- Change tracking by agent and feature
- Conflict detection and analysis

**Recommendation**: Enhance existing system rather than rebuild. Add:
- PR creation via GitHub CLI (already allowed in security.py)
- Branch-per-feature workflows
- Better merge conflict resolution
- Code review integration

## Implementation Phases

### Phase 1: Database Foundation (3-4 days)
1. Add usage tables to `api/database.py`
2. Create migration functions
3. Add session journal file structure
4. Extend Feature model with recovery metadata

### Phase 2: Token Tracking (4-5 days)
1. Modify `agent.py` and `assistant_chat_session.py` for usage extraction
2. Add cost calculation utilities
3. Extend WebSocket system for real-time updates
4. Create UI components for cost display
5. Add REST endpoints for usage history

### Phase 3: Session Recovery (5-6 days)
1. Implement session journal logging
2. Add recovery context injection to prompts
3. Create stale work detection service
4. Build recovery workflow in agent startup
5. Add UI indicators for resumed sessions

### Phase 4: Error Recovery & Multi-Model (4-5 days)
1. Add model escalation logic
2. Implement stuck state detection
3. Create skip/retry UI controls
4. Add configuration for model routing rules
5. Extend process manager for model switching

### Phase 5: Testing & Integration (2-3 days)
1. Test all enhancements with existing functionality
2. Verify database migrations work correctly
3. Ensure WebSocket updates don't break existing clients
4. Test process cleanup with new components
5. Validate multi-model cost differences

## Key Files to Modify

| File | Purpose | Changes |
|------|---------|---------|
| `api/database.py` | Add usage/recovery tables | Add 3 new models, migration function |
| `agent.py` | Token tracking, recovery | Extract usage, load recovery context |
| `server/websocket.py` | Real-time updates | Add usage polling, new message types |
| `ui/src/hooks/useWebSocket.ts` | UI state management | Extend WebSocketState interface |
| `ui/src/components/ProgressDashboard.tsx` | Usage display | Add cost/token section |
| `config.py` | Model routing config | Add ModelConfig, escalation rules |
| `client.py` | Multi-model support | Model selection logic |

**New Files:**
- `server/services/session_recovery.py` - Recovery logic
- `server/services/usage_tracking.py` - Token/cost utilities
- `server/routers/usage.py` - Usage history API
- `ui/src/components/UsagePanel.tsx` - Cost display component

## Verification Strategy

Each enhancement verified by:
1. **Database migrations**: Run on test project, verify schema changes
2. **Token tracking**: Start session, verify real-time cost updates in UI
3. **Session recovery**: Kill agent mid-feature, restart, verify recovery context
4. **Error handling**: Create failing feature, verify escalation to Opus
5. **Multi-model**: Test routing based on feature keywords
6. **Process cleanup**: Ensure no zombie processes with new components

## Risk Mitigation

- **Backward compatibility**: All changes are additive, existing functionality preserved
- **Database safety**: Migrations use `checkfirst=True`, handle missing columns gracefully
- **Performance**: Usage polling is lightweight (1-second interval), follows existing patterns
- **Security**: No changes to bash allowlist or security model
- **Process isolation**: New components follow existing singleton + lock patterns

This plan leverages Autocoder 2.0's solid architecture to add sophisticated enhancements while maintaining the reliability and patterns that make the current system successful.


File created successfully at: C:\Users\MikeSwanson\.claude\plans\hazy-wiggling-octopus.md

# Autocoder 2.0 Enhancement Implementation Plan

## Overview

This plan implements the priority enhancements for Autocoder 2.0 - an autonomous coding agent system. Based on comprehensive codebase exploration, I've identified the architecture patterns and integration points needed to add token tracking, session recovery, enhanced error handling, and multi-model support while preserving all existing functionality.

## Architecture Understanding

Autocoder 2.0 follows these patterns that I'll extend:
- **Singleton managers with thread safety** (process_manager.py)
- **SQLAlchemy models with migrations** (api/database.py)
- **WebSocket real-time updates** (server/websocket.py)
- **React Query hooks + WebSocket state** (ui/src/hooks/)
- **MCP server for agent tools** (mcp_server/feature_mcp.py)
- **Async subprocess management** (server/services/)

## Priority 1: Token and Cost Tracking

### Database Schema Extensions

Add three new tables to `api/database.py`:

```python
class UsageSession(Base):
    __tablename__ = "usage_sessions"

    id = Column(Integer, primary_key=True)
    session_id = Column(String(100), unique=True, index=True)
    started_at = Column(DateTime, default=datetime.utcnow, index=True)
    ended_at = Column(DateTime, nullable=True)

    total_input_tokens = Column(Integer, default=0)
    total_output_tokens = Column(Integer, default=0)
    total_calls = Column(Integer, default=0)
    total_cost_cents = Column(Integer, default=0)

    model = Column(String(50))
    status = Column(String(20))  # "active", "completed", "failed"

class APICall(Base):
    __tablename__ = "api_calls"

    id = Column(Integer, primary_key=True)
    session_id = Column(String(100), ForeignKey("usage_sessions.session_id"), index=True)
    timestamp = Column(DateTime, default=datetime.utcnow, index=True)

    input_tokens = Column(Integer)
    output_tokens = Column(Integer)
    stop_reason = Column(String(50))
    response_time_ms = Column(Integer)
    cost_cents = Column(Integer)

class UsageSummary(Base):
    __tablename__ = "usage_summary"

    id = Column(Integer, primary_key=True)
    session_id = Column(String(100), unique=True, index=True)
    total_cost_cents = Column(Integer, default=0)
    total_input_tokens = Column(Integer, default=0)
    total_output_tokens = Column(Integer, default=0)
```

**Migration strategy**: Extend `_migrate_add_in_progress_column()` pattern to add usage tables.

### Token Extraction Integration

Modify `agent.py` and `assistant_chat_session.py` to capture usage metadata:

**In `agent.py` line 34-105:**
```python
async def run_agent_session(client, message, project_dir):
    session_id = str(uuid.uuid4())
    start_usage_session(session_id, client.model)

    start_time = time.time()
    await client.query(message)

    async for msg in client.receive_response():
        # Existing response handling...
        pass

    # NEW: Extract usage from final message
    if hasattr(msg, 'usage'):
        track_api_call(
            session_id=session_id,
            input_tokens=msg.usage.input_tokens,
            output_tokens=msg.usage.output_tokens,
            duration_ms=int((time.time() - start_time) * 1000)
        )
```

**Integration points**:
- Hook into `client.receive_response()` completion
- Use Claude SDK's `.usage` attribute on response messages
- Store in database via new service functions

### Real-time Cost Updates

Extend WebSocket system in `server/websocket.py`:

**Add new message type:**
```python
{
    "type": "usage_update",
    "total_input_tokens": int,
    "total_output_tokens": int,
    "total_calls": int,
    "estimated_cost": "$X.XX",
    "timestamp": ISO_string
}
```

**Polling approach** (follows existing `poll_progress()` pattern):
```python
async def poll_usage(websocket, project_name, session_id):
    last_calls = -1
    while True:
        usage = get_usage_summary(session_id)
        if usage['total_calls'] != last_calls:
            await websocket.send_json({
                "type": "usage_update",
                **usage,
                "estimated_cost": format_cost(usage['cost_cents'])
            })
        await asyncio.sleep(1)
```

### UI Integration

**Extend WebSocket state** in `ui/src/hooks/useWebSocket.ts`:
```typescript
interface WebSocketState {
  progress: {...}
  agentStatus: AgentStatus
  logs: Array<{...}>
  isConnected: boolean

  // NEW:
  usage: {
    totalInputTokens: number
    totalOutputTokens: number
    totalCalls: number
    estimatedCost: string
    lastUpdated: string
  }
}
```

**Cost display options**:
1. **Header badge** - Real-time cost next to agent status
2. **Extended ProgressDashboard** - Usage section with token breakdown
3. **Separate UsagePanel** - Dedicated component for cost tracking

## Priority 2: Session Recovery and Checkpointing

### Session Journal System

Create `.session_journal/` directory in each project:

```
project_dir/.session_journal/
├── session_001.json
├── session_002.json
└── current_session.json
```

**Journal schema:**
```json
{
  "session_id": "uuid",
  "started_at": "ISO_timestamp",
  "ended_at": "ISO_timestamp",
  "feature_id": 42,
  "initial_stats": {"passing": 10, "total": 50},
  "final_stats": {"passing": 12, "total": 50},
  "messages": [
    {
      "role": "user|assistant",
      "content": "...",
      "tools_used": ["feature_mark_passing"],
      "timestamp": "ISO"
    }
  ],
  "blockers": ["Dependency on feature #5"],
  "decisions": ["Skipped feature #10 - unclear requirements"],
  "errors": ["Tool Write failed with error X"],
  "status": "success|incomplete|error"
}
```

### Feature Work Metadata

Extend Feature model with recovery columns:
```python
# Add to Feature class in api/database.py:
last_worked_on = Column(DateTime)
work_session_id = Column(String(100))
attempt_count = Column(Integer, default=0)
last_error = Column(Text)
recovery_notes = Column(Text)
```

### Recovery Context Injection

Modify `prompts.py` to include recovery context:
```python
def get_recovery_context(project_dir: Path) -> str:
    last_session = load_latest_session_journal(project_dir)
    if not last_session:
        return ""

    return f"""
## Recovery Context from Previous Session

**Previous Status**: {last_session.get('status')}
**Features Completed**: {last_session['final_stats']['passing']}
**Blockers**: {json.dumps(last_session.get('blockers', []))}
**Errors**: {json.dumps(last_session.get('errors', []))}
**Decisions**: {json.dumps(last_session.get('decisions', []))}
"""
```

### Stale Work Detection

Add recovery module to detect and handle stuck features:
```python
# NEW: server/services/session_recovery.py
def detect_stale_work(project_dir, timeout_minutes=60):
    """Find features stuck in in_progress state."""
    return session.query(Feature).filter(
        Feature.in_progress == True,
        Feature.last_worked_on < datetime.now() - timedelta(minutes=timeout_minutes)
    ).all()

def recover_stale_work(project_dir):
    """Reset stale work or add recovery notes."""
    for feature in detect_stale_work(project_dir):
        journal = load_session_journal(feature.work_session_id)
        if journal.get("status") == "error":
            feature.attempt_count += 1
            if feature.attempt_count > 3:
                feature_skip(feature.id)  # Move to end of queue
```

## Priority 3: Enhanced Error Recovery

### Stuck State Detection

Monitor for repeated failures on same feature:
- Track attempt count in Feature model
- Detect same error patterns via session journal
- Auto-escalate to Opus after 3 Sonnet failures
- Skip features after 5 total attempts

### Model Escalation Chain

**Add to config.py:**
```python
@dataclass
class ModelConfig:
    default_model: str = "claude-sonnet-4-20250514"
    escalation_enabled: bool = True
    escalation_threshold: int = 3  # failures before escalating
    stuck_timeout_minutes: int = 30

    # Model routing rules
    opus_keywords: list[str] = ["refactor", "architecture", "complex"]
    haiku_keywords: list[str] = ["simple", "test", "docs", "lint"]
```

**Escalation logic:**
```python
def should_escalate_model(feature_id: int, current_model: str) -> str:
    """Determine if model escalation is needed."""
    feature = get_feature(feature_id)

    if feature.attempt_count >= 3 and current_model != "claude-opus-4-5":
        return "claude-opus-4-5"  # Escalate to Opus

    # Check for complexity keywords
    if any(keyword in feature.description.lower() for keyword in OPUS_KEYWORDS):
        return "claude-opus-4-5"

    return current_model
```

### Skip and Retry Mechanisms

**UI Controls:**
- "Retry with Opus" button for stuck features
- "Skip for now" button to move to end of queue
- Bulk skip operations via checkbox selection

**Auto-skip criteria:**
- Features failed 5+ times across all models
- Features blocked by dependencies for >24 hours
- Features with unresolvable errors (missing requirements, etc.)

## Priority 4: Multi-Model Support

### Configuration-Driven Routing

**Extend client.py:**
```python
def get_model_for_task(task_description: str, feature_category: str, attempt_count: int) -> str:
    """Route task to appropriate model based on complexity."""

    # Escalation chain
    if attempt_count >= 3:
        return "claude-opus-4-5"

    # Keyword-based routing
    description_lower = task_description.lower()

    if any(kw in description_lower for kw in OPUS_KEYWORDS):
        return "claude-opus-4-5"
    elif any(kw in description_lower for kw in HAIKU_KEYWORDS):
        return "claude-3-5-haiku"
    else:
        return "claude-sonnet-4-20250514"  # Default
```

**Model specialization:**
- **Opus**: Complex architectural decisions, debugging difficult failures, large refactors
- **Sonnet**: Standard coding tasks (current default)
- **Haiku**: Simple validations, linting, documentation, quick tests

## Git Integration Assessment

**Discovery**: Autocoder 2.0 already has comprehensive git integration:
- `git_agent.py` - Full orchestration with auto-commits on feature completion
- `git_operations.py` - Async git wrapper with all operations
- `server/routers/git.py` - REST API endpoints
- Auto-commit on feature passing, session end
- Real-time git status via WebSocket
- Change tracking by agent and feature
- Conflict detection and analysis

**Recommendation**: Enhance existing system rather than rebuild. Add:
- PR creation via GitHub CLI (already allowed in security.py)
- Branch-per-feature workflows
- Better merge conflict resolution
- Code review integration

## Implementation Phases

### Phase 1: Database Foundation (3-4 days)
1. Add usage tables to `api/database.py`
2. Create migration functions
3. Add session journal file structure
4. Extend Feature model with recovery metadata

### Phase 2: Token Tracking (4-5 days)
1. Modify `agent.py` and `assistant_chat_session.py` for usage extraction
2. Add cost calculation utilities
3. Extend WebSocket system for real-time updates
4. Create UI components for cost display
5. Add REST endpoints for usage history

### Phase 3: Session Recovery (5-6 days)
1. Implement session journal logging
2. Add recovery context injection to prompts
3. Create stale work detection service
4. Build recovery workflow in agent startup
5. Add UI indicators for resumed sessions

### Phase 4: Error Recovery & Multi-Model (4-5 days)
1. Add model escalation logic
2. Implement stuck state detection
3. Create skip/retry UI controls
4. Add configuration for model routing rules
5. Extend process manager for model switching

### Phase 5: Testing & Integration (2-3 days)
1. Test all enhancements with existing functionality
2. Verify database migrations work correctly
3. Ensure WebSocket updates don't break existing clients
4. Test process cleanup with new components
5. Validate multi-model cost differences

## Key Files to Modify

| File | Purpose | Changes |
|------|---------|---------|
| `api/database.py` | Add usage/recovery tables | Add 3 new models, migration function |
| `agent.py` | Token tracking, recovery | Extract usage, load recovery context |
| `server/websocket.py` | Real-time updates | Add usage polling, new message types |
| `ui/src/hooks/useWebSocket.ts` | UI state management | Extend WebSocketState interface |
| `ui/src/components/ProgressDashboard.tsx` | Usage display | Add cost/token section |
| `config.py` | Model routing config | Add ModelConfig, escalation rules |
| `client.py` | Multi-model support | Model selection logic |

**New Files:**
- `server/services/session_recovery.py` - Recovery logic
- `server/services/usage_tracking.py` - Token/cost utilities
- `server/routers/usage.py` - Usage history API
- `ui/src/components/UsagePanel.tsx` - Cost display component

## Verification Strategy

Each enhancement verified by:
1. **Database migrations**: Run on test project, verify schema changes
2. **Token tracking**: Start session, verify real-time cost updates in UI
3. **Session recovery**: Kill agent mid-feature, restart, verify recovery context
4. **Error handling**: Create failing feature, verify escalation to Opus
5. **Multi-model**: Test routing based on feature keywords
6. **Process cleanup**: Ensure no zombie processes with new components

## Risk Mitigation

- **Backward compatibility**: All changes are additive, existing functionality preserved
- **Database safety**: Migrations use `checkfirst=True`, handle missing columns gracefully
- **Performance**: Usage polling is lightweight (1-second interval), follows existing patterns
- **Security**: No changes to bash allowlist or security model
- **Process isolation**: New components follow existing singleton + lock patterns

This plan leverages Autocoder 2.0's solid architecture to add sophisticated enhancements while maintaining the reliability and patterns that make the current system successful.


Based on my comprehensive exploration of the Autocoder 2.0 codebase, I've created a detailed implementation plan that addresses all your priority enhancements. Here are the key discoveries and approach:

## Major Discovery: Git Integration Already Exists!

I found that Autocoder 2.0 already has sophisticated git integration:
- **Auto-commits** on feature completion via `git_agent.py`
- **Full async git operations** via `git_operations.py`
- **REST API endpoints** for all git operations
- **Real-time git status** via WebSocket
- **Change tracking** by agent and feature
- **Auto-push configuration** on feature/session completion

This means we can focus enhancement efforts on the other priorities rather than building git integration from scratch.

## Implementation Approach

The plan follows your existing architectural patterns:
- **Singleton managers with thread safety** (like process_manager.py)
- **SQLAlchemy models with migrations** (extending api/database.py)
- **WebSocket real-time updates** (extending server/websocket.py)
- **React Query + WebSocket state management** (extending UI hooks)

## Priority Implementation Order

1. **Token & Cost Tracking** (3-4 days)
   - New usage tables in database
   - Extract tokens from Claude SDK responses
   - Real-time cost updates via WebSocket
   - UI cost badge and usage dashboard

2. **Session Recovery** (4-5 days)
   - Session journal system with conversation history
   - Feature work metadata for stale detection
   - Recovery context injection into prompts
   - Intelligent resume after crashes

3. **Enhanced Error Recovery** (3-4 days)
   - Stuck state detection with retry limits
   - Model escalation (Haiku → Sonnet → Opus)
   - Skip mechanisms and bulk operations

4. **Multi-Model Support** (3-4 days)
   - Config-driven routing based on keywords/complexity
   - Automatic escalation chain on failures
   - Cost-aware model selection

Each phase is designed to be independently testable and non-breaking to existing functionality.


ExitPlanMode

# Autocoder 2.0 Enhancement Implementation Plan

## Overview

This plan implements the priority enhancements for Autocoder 2.0 - an autonomous coding agent system. Based on comprehensive codebase exploration, I've identified the architecture patterns and integration points needed to add token tracking, session recovery, enhanced error handling, and multi-model support while preserving all existing functionality.

## Architecture Understanding

Autocoder 2.0 follows these patterns that I'll extend:
- **Singleton managers with thread safety** (process_manager.py)
- **SQLAlchemy models with migrations** (api/database.py)
- **WebSocket real-time updates** (server/websocket.py)
- **React Query hooks + WebSocket state** (ui/src/hooks/)
- **MCP server for agent tools** (mcp_server/feature_mcp.py)
- **Async subprocess management** (server/services/)

## Priority 1: Token and Cost Tracking

### Database Schema Extensions

Add three new tables to `api/database.py`:

```python
class UsageSession(Base):
    __tablename__ = "usage_sessions"

    id = Column(Integer, primary_key=True)
    session_id = Column(String(100), unique=True, index=True)
    started_at = Column(DateTime, default=datetime.utcnow, index=True)
    ended_at = Column(DateTime, nullable=True)

    total_input_tokens = Column(Integer, default=0)
    total_output_tokens = Column(Integer, default=0)
    total_calls = Column(Integer, default=0)
    total_cost_cents = Column(Integer, default=0)

    model = Column(String(50))
    status = Column(String(20))  # "active", "completed", "failed"

class APICall(Base):
    __tablename__ = "api_calls"

    id = Column(Integer, primary_key=True)
    session_id = Column(String(100), ForeignKey("usage_sessions.session_id"), index=True)
    timestamp = Column(DateTime, default=datetime.utcnow, index=True)

    input_tokens = Column(Integer)
    output_tokens = Column(Integer)
    stop_reason = Column(String(50))
    response_time_ms = Column(Integer)
    cost_cents = Column(Integer)

class UsageSummary(Base):
    __tablename__ = "usage_summary"

    id = Column(Integer, primary_key=True)
    session_id = Column(String(100), unique=True, index=True)
    total_cost_cents = Column(Integer, default=0)
    total_input_tokens = Column(Integer, default=0)
    total_output_tokens = Column(Integer, default=0)
```

**Migration strategy**: Extend `_migrate_add_in_progress_column()` pattern to add usage tables.

### Token Extraction Integration

Modify `agent.py` and `assistant_chat_session.py` to capture usage metadata:

**In `agent.py` line 34-105:**
```python
async def run_agent_session(client, message, project_dir):
    session_id = str(uuid.uuid4())
    start_usage_session(session_id, client.model)

    start_time = time.time()
    await client.query(message)

    async for msg in client.receive_response():
        # Existing response handling...
        pass

    # NEW: Extract usage from final message
    if hasattr(msg, 'usage'):
        track_api_call(
            session_id=session_id,
            input_tokens=msg.usage.input_tokens,
            output_tokens=msg.usage.output_tokens,
            duration_ms=int((time.time() - start_time) * 1000)
        )
```

**Integration points**:
- Hook into `client.receive_response()` completion
- Use Claude SDK's `.usage` attribute on response messages
- Store in database via new service functions

### Real-time Cost Updates

Extend WebSocket system in `server/websocket.py`:

**Add new message type:**
```python
{
    "type": "usage_update",
    "total_input_tokens": int,
    "total_output_tokens": int,
    "total_calls": int,
    "estimated_cost": "$X.XX",
    "timestamp": ISO_string
}
```

**Polling approach** (follows existing `poll_progress()` pattern):
```python
async def poll_usage(websocket, project_name, session_id):
    last_calls = -1
    while True:
        usage = get_usage_summary(session_id)
        if usage['total_calls'] != last_calls:
            await websocket.send_json({
                "type": "usage_update",
                **usage,
                "estimated_cost": format_cost(usage['cost_cents'])
            })
        await asyncio.sleep(1)
```

### UI Integration

**Extend WebSocket state** in `ui/src/hooks/useWebSocket.ts`:
```typescript
interface WebSocketState {
  progress: {...}
  agentStatus: AgentStatus
  logs: Array<{...}>
  isConnected: boolean

  // NEW:
  usage: {
    totalInputTokens: number
    totalOutputTokens: number
    totalCalls: number
    estimatedCost: string
    lastUpdated: string
  }
}
```

**Cost display options**:
1. **Header badge** - Real-time cost next to agent status
2. **Extended ProgressDashboard** - Usage section with token breakdown
3. **Separate UsagePanel** - Dedicated component for cost tracking

## Priority 2: Session Recovery and Checkpointing

### Session Journal System

Create `.session_journal/` directory in each project:

```
project_dir/.session_journal/
├── session_001.json
├── session_002.json
└── current_session.json
```

**Journal schema:**
```json
{
  "session_id": "uuid",
  "started_at": "ISO_timestamp",
  "ended_at": "ISO_timestamp",
  "feature_id": 42,
  "initial_stats": {"passing": 10, "total": 50},
  "final_stats": {"passing": 12, "total": 50},
  "messages": [
    {
      "role": "user|assistant",
      "content": "...",
      "tools_used": ["feature_mark_passing"],
      "timestamp": "ISO"
    }
  ],
  "blockers": ["Dependency on feature #5"],
  "decisions": ["Skipped feature #10 - unclear requirements"],
  "errors": ["Tool Write failed with error X"],
  "status": "success|incomplete|error"
}
```

### Feature Work Metadata

Extend Feature model with recovery columns:
```python
# Add to Feature class in api/database.py:
last_worked_on = Column(DateTime)
work_session_id = Column(String(100))
attempt_count = Column(Integer, default=0)
last_error = Column(Text)
recovery_notes = Column(Text)
```

### Recovery Context Injection

Modify `prompts.py` to include recovery context:
```python
def get_recovery_context(project_dir: Path) -> str:
    last_session = load_latest_session_journal(project_dir)
    if not last_session:
        return ""

    return f"""
## Recovery Context from Previous Session

**Previous Status**: {last_session.get('status')}
**Features Completed**: {last_session['final_stats']['passing']}
**Blockers**: {json.dumps(last_session.get('blockers', []))}
**Errors**: {json.dumps(last_session.get('errors', []))}
**Decisions**: {json.dumps(last_session.get('decisions', []))}
"""
```

### Stale Work Detection

Add recovery module to detect and handle stuck features:
```python
# NEW: server/services/session_recovery.py
def detect_stale_work(project_dir, timeout_minutes=60):
    """Find features stuck in in_progress state."""
    return session.query(Feature).filter(
        Feature.in_progress == True,
        Feature.last_worked_on < datetime.now() - timedelta(minutes=timeout_minutes)
    ).all()

def recover_stale_work(project_dir):
    """Reset stale work or add recovery notes."""
    for feature in detect_stale_work(project_dir):
        journal = load_session_journal(feature.work_session_id)
        if journal.get("status") == "error":
            feature.attempt_count += 1
            if feature.attempt_count > 3:
                feature_skip(feature.id)  # Move to end of queue
```

## Priority 3: Enhanced Error Recovery

### Stuck State Detection

Monitor for repeated failures on same feature:
- Track attempt count in Feature model
- Detect same error patterns via session journal
- Auto-escalate to Opus after 3 Sonnet failures
- Skip features after 5 total attempts

### Model Escalation Chain

**Add to config.py:**
```python
@dataclass
class ModelConfig:
    default_model: str = "claude-sonnet-4-20250514"
    escalation_enabled: bool = True
    escalation_threshold: int = 3  # failures before escalating
    stuck_timeout_minutes: int = 30

    # Model routing rules
    opus_keywords: list[str] = ["refactor", "architecture", "complex"]
    haiku_keywords: list[str] = ["simple", "test", "docs", "lint"]
```

**Escalation logic:**
```python
def should_escalate_model(feature_id: int, current_model: str) -> str:
    """Determine if model escalation is needed."""
    feature = get_feature(feature_id)

    if feature.attempt_count >= 3 and current_model != "claude-opus-4-5":
        return "claude-opus-4-5"  # Escalate to Opus

    # Check for complexity keywords
    if any(keyword in feature.description.lower() for keyword in OPUS_KEYWORDS):
        return "claude-opus-4-5"

    return current_model
```

### Skip and Retry Mechanisms

**UI Controls:**
- "Retry with Opus" button for stuck features
- "Skip for now" button to move to end of queue
- Bulk skip operations via checkbox selection

**Auto-skip criteria:**
- Features failed 5+ times across all models
- Features blocked by dependencies for >24 hours
- Features with unresolvable errors (missing requirements, etc.)

## Priority 4: Multi-Model Support

### Configuration-Driven Routing

**Extend client.py:**
```python
def get_model_for_task(task_description: str, feature_category: str, attempt_count: int) -> str:
    """Route task to appropriate model based on complexity."""

    # Escalation chain
    if attempt_count >= 3:
        return "claude-opus-4-5"

    # Keyword-based routing
    description_lower = task_description.lower()

    if any(kw in description_lower for kw in OPUS_KEYWORDS):
        return "claude-opus-4-5"
    elif any(kw in description_lower for kw in HAIKU_KEYWORDS):
        return "claude-3-5-haiku"
    else:
        return "claude-sonnet-4-20250514"  # Default
```

**Model specialization:**
- **Opus**: Complex architectural decisions, debugging difficult failures, large refactors
- **Sonnet**: Standard coding tasks (current default)
- **Haiku**: Simple validations, linting, documentation, quick tests

## Git Integration Assessment

**Discovery**: Autocoder 2.0 already has comprehensive git integration:
- `git_agent.py` - Full orchestration with auto-commits on feature completion
- `git_operations.py` - Async git wrapper with all operations
- `server/routers/git.py` - REST API endpoints
- Auto-commit on feature passing, session end
- Real-time git status via WebSocket
- Change tracking by agent and feature
- Conflict detection and analysis

**Recommendation**: Enhance existing system rather than rebuild. Add:
- PR creation via GitHub CLI (already allowed in security.py)
- Branch-per-feature workflows
- Better merge conflict resolution
- Code review integration

## Implementation Phases

### Phase 1: Database Foundation (3-4 days)
1. Add usage tables to `api/database.py`
2. Create migration functions
3. Add session journal file structure
4. Extend Feature model with recovery metadata

### Phase 2: Token Tracking (4-5 days)
1. Modify `agent.py` and `assistant_chat_session.py` for usage extraction
2. Add cost calculation utilities
3. Extend WebSocket system for real-time updates
4. Create UI components for cost display
5. Add REST endpoints for usage history

### Phase 3: Session Recovery (5-6 days)
1. Implement session journal logging
2. Add recovery context injection to prompts
3. Create stale work detection service
4. Build recovery workflow in agent startup
5. Add UI indicators for resumed sessions

### Phase 4: Error Recovery & Multi-Model (4-5 days)
1. Add model escalation logic
2. Implement stuck state detection
3. Create skip/retry UI controls
4. Add configuration for model routing rules
5. Extend process manager for model switching

### Phase 5: Testing & Integration (2-3 days)
1. Test all enhancements with existing functionality
2. Verify database migrations work correctly
3. Ensure WebSocket updates don't break existing clients
4. Test process cleanup with new components
5. Validate multi-model cost differences

## Key Files to Modify

| File | Purpose | Changes |
|------|---------|---------|
| `api/database.py` | Add usage/recovery tables | Add 3 new models, migration function |
| `agent.py` | Token tracking, recovery | Extract usage, load recovery context |
| `server/websocket.py` | Real-time updates | Add usage polling, new message types |
| `ui/src/hooks/useWebSocket.ts` | UI state management | Extend WebSocketState interface |
| `ui/src/components/ProgressDashboard.tsx` | Usage display | Add cost/token section |
| `config.py` | Model routing config | Add ModelConfig, escalation rules |
| `client.py` | Multi-model support | Model selection logic |

**New Files:**
- `server/services/session_recovery.py` - Recovery logic
- `server/services/usage_tracking.py` - Token/cost utilities
- `server/routers/usage.py` - Usage history API
- `ui/src/components/UsagePanel.tsx` - Cost display component

## Verification Strategy

Each enhancement verified by:
1. **Database migrations**: Run on test project, verify schema changes
2. **Token tracking**: Start session, verify real-time cost updates in UI
3. **Session recovery**: Kill agent mid-feature, restart, verify recovery context
4. **Error handling**: Create failing feature, verify escalation to Opus
5. **Multi-model**: Test routing based on feature keywords
6. **Process cleanup**: Ensure no zombie processes with new components

## Risk Mitigation

- **Backward compatibility**: All changes are additive, existing functionality preserved
- **Database safety**: Migrations use `checkfirst=True`, handle missing columns gracefully
- **Performance**: Usage polling is lightweight (1-second interval), follows existing patterns
- **Security**: No changes to bash allowlist or security model
- **Process isolation**: New components follow existing singleton + lock patterns

This plan leverages Autocoder 2.0's solid architecture to add sophisticated enhancements while maintaining the reliability and patterns that make the current system successful.


Exit plan mode?