     1→/**
     2→ * WebSocket Hook for Real-time Updates
     3→ */
     4→
     5→import { useEffect, useRef, useState, useCallback } from 'react'
     6→import type { WSMessage, AgentStatus } from '../lib/types'
     7→
     8→interface WebSocketState {
     9→  progress: {
    10→    passing: number
    11→    in_progress: number
    12→    total: number
    13→    percentage: number
    14→  }
    15→  agentStatus: AgentStatus
    16→  logs: Array<{ line: string; timestamp: string }>
    17→  isConnected: boolean
    18→}
    19→
    20→const MAX_LOGS = 100 // Keep last 100 log lines
    21→
    22→export function useProjectWebSocket(projectName: string | null) {
    23→  const [state, setState] = useState<WebSocketState>({
    24→    progress: { passing: 0, in_progress: 0, total: 0, percentage: 0 },
    25→    agentStatus: 'stopped',
    26→    logs: [],
    27→    isConnected: false,
    28→  })
    29→
    30→  const wsRef = useRef<WebSocket | null>(null)
    31→  const reconnectTimeoutRef = useRef<number | null>(null)
    32→  const reconnectAttempts = useRef(0)
    33→
    34→  const connect = useCallback(() => {
    35→    if (!projectName) return
    36→
    37→    // Build WebSocket URL
    38→    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
    39→    const host = window.location.host
    40→    const wsUrl = `${protocol}//${host}/ws/projects/${encodeURIComponent(projectName)}`
    41→
    42→    try {
    43→      const ws = new WebSocket(wsUrl)
    44→      wsRef.current = ws
    45→
    46→      ws.onopen = () => {
    47→        setState(prev => ({ ...prev, isConnected: true }))
    48→        reconnectAttempts.current = 0
    49→      }
    50→
    51→      ws.onmessage = (event) => {
    52→        try {
    53→          const message: WSMessage = JSON.parse(event.data)
    54→
    55→          switch (message.type) {
    56→            case 'progress':
    57→              setState(prev => ({
    58→                ...prev,
    59→                progress: {
    60→                  passing: message.passing,
    61→                  in_progress: message.in_progress,
    62→                  total: message.total,
    63→                  percentage: message.percentage,
    64→                },
    65→              }))
    66→              break
    67→
    68→            case 'agent_status':
    69→              setState(prev => ({
    70→                ...prev,
    71→                agentStatus: message.status,
    72→              }))
    73→              break
    74→
    75→            case 'log':
    76→              setState(prev => ({
    77→                ...prev,
    78→                logs: [
    79→                  ...prev.logs.slice(-MAX_LOGS + 1),
    80→                  { line: message.line, timestamp: message.timestamp },
    81→                ],
    82→              }))
    83→              break
    84→
    85→            case 'feature_update':
    86→              // Feature updates will trigger a refetch via React Query
    87→              break
    88→
    89→            case 'pong':
    90→              // Heartbeat response
    91→              break
    92→          }
    93→        } catch {
    94→          console.error('Failed to parse WebSocket message')
    95→        }
    96→      }
    97→
    98→      ws.onclose = () => {
    99→        setState(prev => ({ ...prev, isConnected: false }))
   100→        wsRef.current = null
   101→
   102→        // Exponential backoff reconnection
   103→        const delay = Math.min(1000 * Math.pow(2, reconnectAttempts.current), 30000)
   104→        reconnectAttempts.current++
   105→
   106→        reconnectTimeoutRef.current = window.setTimeout(() => {
   107→          connect()
   108→        }, delay)
   109→      }
   110→
   111→      ws.onerror = () => {
   112→        ws.close()
   113→      }
   114→    } catch {
   115→      // Failed to connect, will retry via onclose
   116→    }
   117→  }, [projectName])
   118→
   119→  // Send ping to keep connection alive
   120→  const sendPing = useCallback(() => {
   121→    if (wsRef.current?.readyState === WebSocket.OPEN) {
   122→      wsRef.current.send(JSON.stringify({ type: 'ping' }))
   123→    }
   124→  }, [])
   125→
   126→  // Connect when project changes
   127→  useEffect(() => {
   128→    if (!projectName) {
   129→      // Disconnect if no project
   130→      if (wsRef.current) {
   131→        wsRef.current.close()
   132→        wsRef.current = null
   133→      }
   134→      return
   135→    }
   136→
   137→    connect()
   138→
   139→    // Ping every 30 seconds
   140→    const pingInterval = setInterval(sendPing, 30000)
   141→
   142→    return () => {
   143→      clearInterval(pingInterval)
   144→      if (reconnectTimeoutRef.current) {
   145→        clearTimeout(reconnectTimeoutRef.current)
   146→      }
   147→      if (wsRef.current) {
   148→        wsRef.current.close()
   149→        wsRef.current = null
   150→      }
   151→    }
   152→  }, [projectName, connect, sendPing])
   153→
   154→  // Clear logs function
   155→  const clearLogs = useCallback(() => {
   156→    setState(prev => ({ ...prev, logs: [] }))
   157→  }, [])
   158→
   159→  return {
   160→    ...state,
   161→    clearLogs,
   162→  }
   163→}
   164→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

/**
 * WebSocket Hook for Real-time Updates
 */

import { useEffect, useRef, useState, useCallback } from 'react'
import type { WSMessage, AgentStatus } from '../lib/types'

interface WebSocketState {
  progress: {
    passing: number
    in_progress: number
    total: number
    percentage: number
  }
  agentStatus: AgentStatus
  logs: Array<{ line: string; timestamp: string }>
  isConnected: boolean
}

const MAX_LOGS = 100 // Keep last 100 log lines

export function useProjectWebSocket(projectName: string | null) {
  const [state, setState] = useState<WebSocketState>({
    progress: { passing: 0, in_progress: 0, total: 0, percentage: 0 },
    agentStatus: 'stopped',
    logs: [],
    isConnected: false,
  })

  const wsRef = useRef<WebSocket | null>(null)
  const reconnectTimeoutRef = useRef<number | null>(null)
  const reconnectAttempts = useRef(0)

  const connect = useCallback(() => {
    if (!projectName) return

    // Build WebSocket URL
    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
    const host = window.location.host
    const wsUrl = `${protocol}//${host}/ws/projects/${encodeURIComponent(projectName)}`

    try {
      const ws = new WebSocket(wsUrl)
      wsRef.current = ws

      ws.onopen = () => {
        setState(prev => ({ ...prev, isConnected: true }))
        reconnectAttempts.current = 0
      }

      ws.onmessage = (event) => {
        try {
          const message: WSMessage = JSON.parse(event.data)

          switch (message.type) {
            case 'progress':
              setState(prev => ({
                ...prev,
                progress: {
                  passing: message.passing,
                  in_progress: message.in_progress,
                  total: message.total,
                  percentage: message.percentage,
                },
              }))
              break

            case 'agent_status':
              setState(prev => ({
                ...prev,
                agentStatus: message.status,
              }))
              break

            case 'log':
              setState(prev => ({
                ...prev,
                logs: [
                  ...prev.logs.slice(-MAX_LOGS + 1),
                  { line: message.line, timestamp: message.timestamp },
                ],
              }))
              break

            case 'feature_update':
              // Feature updates will trigger a refetch via React Query
              break

            case 'pong':
              // Heartbeat response
              break
          }
        } catch {
          console.error('Failed to parse WebSocket message')
        }
      }

      ws.onclose = () => {
        setState(prev => ({ ...prev, isConnected: false }))
        wsRef.current = null

        // Exponential backoff reconnection
        const delay = Math.min(1000 * Math.pow(2, reconnectAttempts.current), 30000)
        reconnectAttempts.current++

        reconnectTimeoutRef.current = window.setTimeout(() => {
          connect()
        }, delay)
      }

      ws.onerror = () => {
        ws.close()
      }
    } catch {
      // Failed to connect, will retry via onclose
    }
  }, [projectName])

  // Send ping to keep connection alive
  const sendPing = useCallback(() => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify({ type: 'ping' }))
    }
  }, [])

  // Connect when project changes
  useEffect(() => {
    if (!projectName) {
      // Disconnect if no project
      if (wsRef.current) {
        wsRef.current.close()
        wsRef.current = null
      }
      return
    }

    connect()

    // Ping every 30 seconds
    const pingInterval = setInterval(sendPing, 30000)

    return () => {
      clearInterval(pingInterval)
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current)
      }
      if (wsRef.current) {
        wsRef.current.close()
        wsRef.current = null
      }
    }
  }, [projectName, connect, sendPing])

  // Clear logs function
  const clearLogs = useCallback(() => {
    setState(prev => ({ ...prev, logs: [] }))
  }, [])

  return {
    ...state,
    clearLogs,
  }
}


Perfect! Now I have comprehensive understanding of the codebase. Let me create a detailed analysis document for the user.

## Comprehensive Autocoder 2.0 Token & Cost Tracking Analysis

Based on my exploration of the Autocoder 2.0 codebase, here's a detailed breakdown of integration points and patterns for implementing token and cost tracking:

---

### 1. CURRENT CLAUDE SDK USAGE & RESPONSE METADATA

**Files:** `client.py`, `agent.py`, `server/services/assistant_chat_session.py`

**Current Implementation:**
- Claude SDK client initialized via `ClaudeSDKClient` with `ClaudeAgentOptions` (lines 183-201 in `client.py`)
- API calls made through `client.query(message)` and `client.receive_response()` (agent.py:56, 60)
- Streaming responses via async generator: `async for msg in client.receive_response()`

**Token/Cost Data Available:**
According to Claude SDK documentation, the `receive_response()` stream yields message objects with:
- `AssistantMessage` objects containing content blocks
- Each response should have metadata with token usage (input/output tokens)
- **Key integration point**: Messages are streamed, so token metadata needs to be extracted per-response cycle

**Current Response Handling Pattern** (agent.py:60-98):
```python
async for msg in client.receive_response():
    msg_type = type(msg).__name__
    if msg_type == "AssistantMessage" and hasattr(msg, "content"):
        for block in msg.content:
            # Process TextBlock and ToolUseBlock
    elif msg_type == "UserMessage" and hasattr(msg, "content"):
        # Process ToolResultBlock
```

**What's Missing:**
- No extraction of `usage` metadata from response messages
- No tracking of `stop_reason` (e.g., "end_turn", "tool_use", "max_tokens")
- SDK hook for final response metrics not being used

**Integration Points for Token Tracking:**
1. After `client.receive_response()` completes → extract final usage stats
2. Hook into message objects' `.usage` attribute (if available on final message)
3. Alternative: Use `PostToolUse` or custom hooks (client.py:192-196 shows hook pattern)

---

### 2. DATABASE PATTERNS FOR NEW TABLES

**File:** `api/database.py`

**Current SQLAlchemy Model Pattern:**
```python
class Feature(Base):
    __tablename__ = "features"

    id = Column(Integer, primary_key=True, index=True)
    priority = Column(Integer, nullable=False, default=999, index=True)
    category = Column(String(100), nullable=False)
    # ... more columns

    def to_dict(self) -> dict:
        return {...}
```

**Database Setup Functions:**
- `create_database(project_dir)` → returns `(engine, SessionLocal)` tuple
- `get_database_url(project_dir)` → constructs SQLite URL with POSIX paths
- `_migrate_add_in_progress_column(engine)` → shows migration pattern using `ALTER TABLE`
- Global session maker pattern: `set_session_maker()` and `get_db()` dependency

**Recommended Tables for Token Tracking:**

**Table 1: `usage_sessions` (per agent run)**
```python
class UsageSession(Base):
    __tablename__ = "usage_sessions"

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

    # Cumulative totals
    total_input_tokens = Column(Integer, default=0)
    total_output_tokens = Column(Integer, default=0)
    total_calls = Column(Integer, default=0)

    model = Column(String(50))  # e.g., "claude-opus-4-5"
    status = Column(String(20))  # "active", "completed", "failed"
```

**Table 2: `api_calls` (per Claude API call)**
```python
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
    call_type = Column(String(50))  # "query", "message", "tool_result"
    prompt_length = Column(Integer)

    # Output
    input_tokens = Column(Integer)
    output_tokens = Column(Integer)
    stop_reason = Column(String(50))  # "end_turn", "tool_use", "max_tokens"
    response_time_ms = Column(Integer)

    # Cost calculation
    cost_cents = Column(Integer)  # Store as cents to avoid float precision issues
```

**Table 3: `usage_summary` (cached daily/session stats)**
```python
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)
    total_calls = Column(Integer, default=0)

    # For display
    cost_display = Column(String(20))  # "$12.34"
    tokens_display = Column(String(50))  # "45k input / 12k output"
```

**Migration Pattern:**
```python
def _migrate_add_usage_tables(engine) -> None:
    """Add usage tracking tables to existing databases."""
    from sqlalchemy import text
    with engine.connect() as conn:
        # Check if table exists
        result = conn.execute(text("SELECT name FROM sqlite_master WHERE type='table' AND name='api_calls'"))
        if not result.fetchone():
            # Table doesn't exist, create via Base.metadata.create_all()
            UsageSession.__table__.create(engine, checkfirst=True)
            APICall.__table__.create(engine, checkfirst=True)
            conn.commit()
```

---

### 3. WEBSOCKET REAL-TIME UPDATES

**File:** `server/websocket.py`

**Current WebSocket Message Types:**
```python
{
    "type": "progress",
    "passing": int,
    "in_progress": int,
    "total": int,
    "percentage": float
}

{
    "type": "agent_status",
    "status": str
}

{
    "type": "log",
    "line": str,
    "timestamp": datetime
}
```

**How WebSocket Architecture Works:**
- `ConnectionManager` class manages per-project WebSocket connections (lines 50-98)
- `broadcast_to_project(project_name, message)` sends to all connections (lines 75-93)
- `poll_progress()` polls database every 2 seconds for progress updates (lines 112-143)
- Callbacks system for agent output: `agent_manager.add_output_callback()` (lines 195-196)

**Integration Points for Token Updates:**

**1. Add new WebSocket message type:**
```python
{
    "type": "usage_update",
    "session_id": str,
    "total_input_tokens": int,
    "total_output_tokens": int,
    "total_calls": int,
    "estimated_cost": "$X.XX",
    "timestamp": datetime
}
```

**2. Polling approach (similar to progress):**
```python
async def poll_usage(websocket: WebSocket, project_name: str, session_id: str):
    """Poll usage database for cost/token updates."""
    last_total_calls = -1

    while True:
        try:
            usage = get_usage_summary(session_id)  # Query API_CALLS table

            if usage['total_calls'] != last_total_calls:
                last_total_calls = usage['total_calls']
                await websocket.send_json({
                    "type": "usage_update",
                    "total_input_tokens": usage['input_tokens'],
                    "total_output_tokens": usage['output_tokens'],
                    "total_calls": usage['total_calls'],
                    "estimated_cost": format_cost(usage['cost_cents']),
                    "timestamp": datetime.now().isoformat(),
                })

            await asyncio.sleep(1)  # Poll every second
        except asyncio.CancelledError:
            raise
        except Exception as e:
            logger.warning(f"Usage polling error: {e}")
            break
```

**3. Callback approach (real-time on each API call):**
```python
# In agent process manager (process_manager.py)
def _notify_usage_update(self, usage_data: dict) -> None:
    """Notify all registered callbacks of usage update."""
    with self._callbacks_lock:
        callbacks = list(self._usage_callbacks)

    for callback in callbacks:
        try:
            loop = asyncio.get_running_loop()
            loop.create_task(self._safe_callback(callback, usage_data))
        except RuntimeError:
            pass

# In WebSocket handler
async def on_usage_update(data: dict):
    await websocket.send_json({
        "type": "usage_update",
        **data,
        "timestamp": datetime.now().isoformat(),
    })

agent_manager.add_usage_callback(on_usage_update)
```

---

### 4. SESSION MANAGEMENT & TOKEN INJECTION POINTS

**Files:** `agent.py`, `server/services/assistant_chat_session.py`, `server/services/process_manager.py`

**Current Session Architecture:**

1. **Agent Sessions** (autonomous_agent_demo.py):
   - One agent process per project
   - Continuous loop with multiple iterations (agent.py:164-209)
   - Each iteration creates fresh client context (agent.py:177)

2. **Assistant Chat Sessions** (assistant_chat_session.py):
   - Per-project conversational assistant sessions
   - Uses `AssistantChatSession` class with session registry (lines 99-373)
   - Persists to `assistant.db` (ConversationMessage table)

**Where to Inject Token Counting:**

**Option A: Agent Loop Level** (agent.py:34-105)
```python
# Before: await client.query(message)
session_id = generate_session_id()
start_time = time.time()

await client.query(message)

# After: async for msg in client.receive_response()
usage_data = extract_usage_from_final_message()  # NEW
track_api_call(
    session_id=session_id,
    call_type="query",
    input_tokens=usage_data['input_tokens'],
    output_tokens=usage_data['output_tokens'],
    duration_ms=int((time.time() - start_time) * 1000),
    cost_cents=calculate_cost(usage_data)
)
```

**Option B: Client Hook Level** (client.py:192-196)
```python
# Add PostToolUse hook to capture token data after each tool execution
hooks={
    "PostToolUse": [
        HookMatcher(matcher="*", hooks=[token_tracking_hook]),
    ]
}

def token_tracking_hook(hook_input) -> None:
    """Hook called after tool execution to capture token usage."""
    # Extract usage metadata from hook_input
    if hasattr(hook_input, 'usage'):
        track_api_call(hook_input.usage)
```

**Option C: Assistant Chat Session Level** (assistant_chat_session.py:255-294)
```python
# In _query_claude() method, after response stream completes
async def _query_claude(self, message: str):
    start_time = time.time()

    await self.client.query(message)
    full_response = ""

    async for msg in self.client.receive_response():
        # ... existing code ...

    # NEW: Extract and track usage
    if hasattr(msg, 'usage'):
        track_conversation_usage(
            conversation_id=self.conversation_id,
            input_tokens=msg.usage.input_tokens,
            output_tokens=msg.usage.output_tokens,
            duration_ms=int((time.time() - start_time) * 1000),
        )
```

**Session Tracking Context:**

The assistant_database pattern shows good session persistence:
```python
class Conversation(Base):
    __tablename__ = "conversations"
    id = Column(Integer, primary_key=True)
    project_name = Column(String(100), index=True)
    created_at = Column(DateTime)
    # Add usage tracking
    # total_tokens_used = Column(Integer, default=0)
    # total_cost_cents = Column(Integer, default=0)
```

---

### 5. UI INTEGRATION POINTS

**Files:** `ui/src/App.tsx`, `ui/src/components/ProgressDashboard.tsx`, `ui/src/hooks/useWebSocket.ts`

**Current UI Message Flow:**

1. **WebSocket Types** (useWebSocket.ts:8-18):
```typescript
interface WebSocketState {
  progress: { passing, in_progress, total, percentage }
  agentStatus: AgentStatus
  logs: Array<{ line, timestamp }>
  isConnected: boolean
}
```

2. **Message Handler** (useWebSocket.ts:51-92):
```typescript
switch (message.type) {
    case 'progress': setState(...)
    case 'agent_status': setState(...)
    case 'log': setState(...)
    case 'feature_update': // triggers refetch
    case 'pong': // heartbeat
}
```

**Integration Points for Cost Display:**

**1. Extend WebSocket State** (useWebSocket.ts):
```typescript
interface WebSocketState {
  progress: { ... }
  agentStatus: AgentStatus
  logs: Array<{ line, timestamp }>
  isConnected: boolean

  // NEW: Usage tracking
  usage: {
    totalInputTokens: number
    totalOutputTokens: number
    totalCalls: number
    estimatedCost: string  // "$12.34"
    lastUpdated: string
  }
}
```

**2. Add message handler** (useWebSocket.ts:51-92):
```typescript
case 'usage_update':
    setState(prev => ({
        ...prev,
        usage: {
            totalInputTokens: message.total_input_tokens,
            totalOutputTokens: message.total_output_tokens,
            totalCalls: message.total_calls,
            estimatedCost: message.estimated_cost,
            lastUpdated: message.timestamp,
        }
    }))
    break
```

**3. Display in Header or Dashboard:**

**Option A: Cost Badge in Header** (App.tsx:160-239)
```typescript
// Add next to AgentControl component
{selectedProject && (
    <div className="flex items-center gap-2 px-3 py-2 bg-black/20 rounded border border-white/20">
        <DollarSign size={16} />
        <span className="font-mono text-sm">{wsState.usage.estimatedCost}</span>
        <span className="text-xs text-white/60">{wsState.usage.totalCalls} calls</span>
    </div>
)}
```

**Option B: Extended Progress Dashboard** (ProgressDashboard.tsx)
```typescript
// Below percentage display, add usage section
<div className="border-t border-[var(--color-neo-border)] pt-4 mt-4">
    <h3 className="font-display text-sm uppercase mb-3">Usage</h3>
    <div className="grid grid-cols-3 gap-4 text-center">
        <div>
            <span className="font-mono text-lg">{usage.totalInputTokens.toLocaleString()}</span>
            <span className="block text-xs text-[var(--color-neo-text-secondary)]">Input Tokens</span>
        </div>
        <div>
            <span className="font-mono text-lg">{usage.totalOutputTokens.toLocaleString()}</span>
            <span className="block text-xs text-[var(--color-neo-text-secondary)]">Output Tokens</span>
        </div>
        <div>
            <span className="font-mono text-lg text-[var(--color-neo-done)]">{usage.estimatedCost}</span>
            <span className="block text-xs text-[var(--color-neo-text-secondary)]">Est. Cost</span>
        </div>
    </div>
</div>
```

**Option C: Separate Usage Panel** (new component)
```typescript
// UsagePanel.tsx
export function UsagePanel({ usage }: { usage: UsageData }) {
    return (
        <div className="neo-card p-4">
            <h2 className="font-display font-bold mb-4">API Usage</h2>
            <div className="space-y-3">
                <div className="flex justify-between">
                    <span>Tokens</span>
                    <span className="font-mono">{total} tokens</span>
                </div>
                <div className="flex justify-between">
                    <span>API Calls</span>
                    <span className="font-mono">{usage.totalCalls}</span>
                </div>
                <div className="flex justify-between border-t pt-3">
                    <span className="font-bold">Estimated Cost</span>
                    <span className="font-mono font-bold text-[var(--color-neo-done)]">
                        {usage.estimatedCost}
                    </span>
                </div>
            </div>
        </div>
    )
}
```

**API Endpoints for Cost Data:**

Could add REST endpoint for historical data:
```python
# server/routers/usage.py (NEW)
@router.get("/api/projects/{project_name}/usage/summary")
async def get_usage_summary(project_name: str):
    """Get usage summary for a project."""
    return {
        "total_input_tokens": int,
        "total_output_tokens": int,
        "total_calls": int,
        "estimated_cost": str,
        "calls_by_day": [...],
    }

@router.get("/api/projects/{project_name}/usage/timeline")
async def get_usage_timeline(project_name: str, days: int = 7):
    """Get hourly usage timeline."""
    return [
        {
            "timestamp": datetime,
            "input_tokens": int,
            "output_tokens": int,
            "calls": int,
            "cost_cents": int,
        }
    ]
```

---

### 6. COST CALCULATION UTILITIES

**Pricing Constants** (should be in `config.py` or new `pricing.py`):
```python
# Current Claude API pricing (Opus 4.5)
PRICING = {
    "claude-opus-4-5": {
        "input": 0.00015,  # $0.15 per 1M input tokens
        "output": 0.0006,   # $0.60 per 1M output tokens
    },
    "claude-opus-4-5-sonnet": {
        "input": 0.000003,
        "output": 0.000015,
    },
    "claude-3-5-haiku": {
        "input": 0.000008,
        "output": 0.00004,
    },
}

def calculate_cost(input_tokens: int, output_tokens: int, model: str) -> int:
    """Calculate cost in cents."""
    rates = PRICING.get(model, PRICING["claude-opus-4-5"])
    cost_dollars = (input_tokens * rates["input"]) + (output_tokens * rates["output"])
    return int(round(cost_dollars * 100))  # Convert to cents

def format_cost(cents: int) -> str:
    """Format cost as currency string."""
    return f"${cents / 100:.2f}"
```

---

### 7. IMPLEMENTATION SUMMARY & MINIMAL BREAKING CHANGES

**Phase 1: Database Setup (No Breaking Changes)**
1. Add migration in `api/migration.py` or `api/database.py`
2. Create `Usage`, `APICall`, `UsageSummary` models
3. Initialize tables on server startup

**Phase 2: Token Extraction (Minimal Changes)**
1. Modify `agent.py` to capture usage after `receive_response()`
2. Modify `assistant_chat_session.py` to track tokens per message
3. Store in database via new functions

**Phase 3: WebSocket Broadcasting (Isolated)**
1. Add `usage_callbacks` to `AgentProcessManager` (follows existing pattern)
2. Add new message type to `websocket.py` poll function
3. Keep existing message types unchanged

**Phase 4: UI Enhancement (Non-Breaking)**
1. Extend `WebSocketState` type in `useWebSocket.ts`
2. Add new case in message handler switch statement
3. Add usage display components (new files, no changes to existing components)

**Phase 5: REST API (Optional)**
1. Create new router `server/routers/usage.py`
2. Query usage tables for historical data
3. Independent from existing API

---

### 8. KEY FILES TO MODIFY (In Priority Order)

| File | Change Type | Purpose |
|------|-------------|---------|
| `api/database.py` | Add models | `Usage`, `APICall`, `UsageSummary` tables |
| `agent.py` | Insert tracking | Extract tokens after each response |
| `server/services/assistant_chat_session.py` | Insert tracking | Extract tokens in `_query_claude()` |
| `server/websocket.py` | Add polling | `poll_usage()` function, new message type |
| `ui/src/hooks/useWebSocket.ts` | Extend state | Add `usage` field to `WebSocketState` |
| `ui/src/components/ProgressDashboard.tsx` | Add display | Usage stats section (or new component) |
| `config.py` | Add pricing | Pricing constants and calculator |

---

### 9. EXISTING PATTERNS YOU CAN REUSE

✅ **Message Type System**: WebSocket already uses `type` field for routing (useWebSocket.ts:55-92)

✅ **Polling Pattern**: `poll_progress()` model can be copied for `poll_usage()` (websocket.py:112-143)

✅ **Callback System**: `AgentProcessManager` already has output/status callbacks (process_manager.py:118-136)

✅ **Database Migration**: `_migrate_add_in_progress_column()` shows safe migration pattern (database.py:61-73)

✅ **Session Management**: Both agent and assistant already track sessions (registry.py, assistant_database.py)

✅ **TypeScript Message Types**: `WSMessage` union type can be extended (useWebSocket.ts:6)

✅ **SQLAlchemy ORM**: Established models with `to_dict()` serialization (database.py:33-44, assistant_database.py)

---

This exploration provides everything you need to extend Autocoder 2.0 with token and cost tracking while maintaining backward compatibility and leveraging existing architectural patterns. The integration points are well-defined and isolated, making implementation straightforward without breaking current functionality.