# FWAuto 用量追蹤實作方案

## 概述

透過 Claude Agent SDK 的 `ResultMessage` 回報用量到雲端資料庫，實現計費追蹤。

## 架構

```
┌─────────────────────────────────────────────────────────────┐
│  用戶執行 fwauto feat / fwauto build（失敗自動修復）         │
└─────────────────────────┬───────────────────────────────────┘
                          │
                          ▼
┌─────────────────────────────────────────────────────────────┐
│  ai_brain.py                                                │
│                                                             │
│  claude_agent_sdk.query()                                   │
│    • Read / Write / Edit 檔案                               │
│    • 執行 Bash 指令                                         │
│    • 自動修復 build errors                                  │
│                                                             │
│  執行完成後收到 ResultMessage:                               │
│  {                                                          │
│    "usage": {                                               │
│      "input_tokens": 5000,                                  │
│      "output_tokens": 2000                                  │
│    },                                                       │
│    "total_cost_usd": 0.035,                                 │
│    "duration_ms": 45000,                                    │
│    "num_turns": 8                                           │
│  }                                                          │
│                                                             │
│  ──► report_usage() ────────────────────────────────────────┼──►
└─────────────────────────────────────────────────────────────┘   │
                                                                  │
                          ┌───────────────────────────────────────┘
                          │
                          ▼
┌─────────────────────────────────────────────────────────────┐
│  Flask Server (35.229.221.64.nip.io:5000)                   │
│                                                             │
│  POST /api/report-usage                                     │
│  Authorization: Bearer {user_token}                         │
│                                                             │
│  ──► INSERT INTO usage_log (...)                            │
└─────────────────────────────────────────────────────────────┘
                          │
                          ▼
┌─────────────────────────────────────────────────────────────┐
│  Database                                                   │
│                                                             │
│  usage_log:                                                 │
│  | user_id | task_type | input_tokens | output_tokens | ... │
│  |    1    | fix_build |     5000     |     2000      | ... │
└─────────────────────────────────────────────────────────────┘
```

---

## 第一部分：ai_brain.py 修改

### 新增用量回報函數

在 `src/fwauto/nodes/ai_brain.py` 加入以下程式碼：

```python
# ============================================================
# 用量回報（新增）
# ============================================================

import requests
from pathlib import Path
import json
from claude_agent_sdk.types import ResultMessage

# 雲端 Server URL
USAGE_REPORT_URL = "http://35.229.221.64.nip.io:5000/api/report-usage"


def _get_user_token() -> str | None:
    """
    從本地配置讀取 user token

    配置檔案位置: ~/.claude_client_config.json
    格式: {"user_token": "ut-xxx", "user_email": "user@example.com"}
    """
    config_file = Path.home() / '.claude_client_config.json'
    if config_file.exists():
        try:
            with open(config_file, 'r', encoding='utf-8') as f:
                config = json.load(f)
                return config.get('user_token')
        except Exception:
            pass
    return None


def _report_usage_to_server(result_message: ResultMessage, task_type: str) -> None:
    """
    回報用量到雲端 Server

    Args:
        result_message: Claude Agent SDK 執行完成後的 ResultMessage
        task_type: 任務類型 (fix_build_errors / implement_feature)
    """
    logger = get_logger("nodes.ai_brain")

    # 取得 user token
    token = _get_user_token()
    if not token:
        logger.debug("No user token found, skipping usage report")
        return

    # 檢查是否有用量資訊
    if not result_message.usage:
        logger.debug("No usage info in ResultMessage, skipping report")
        return

    try:
        response = requests.post(
            USAGE_REPORT_URL,
            headers={
                "Authorization": f"Bearer {token}",
                "Content-Type": "application/json"
            },
            json={
                "task_type": task_type,
                "model": "claude-sonnet-4",
                "input_tokens": result_message.usage.get("input_tokens", 0),
                "output_tokens": result_message.usage.get("output_tokens", 0),
                "total_cost_usd": result_message.total_cost_usd,
                "duration_ms": result_message.duration_ms,
                "num_turns": result_message.num_turns,
            },
            timeout=5
        )

        if response.status_code == 200:
            logger.debug(f"Usage reported: {result_message.usage}")
        else:
            logger.debug(f"Usage report failed: {response.status_code}")

    except Exception as e:
        # 回報失敗不影響主流程
        logger.debug(f"Usage report error: {e}")
```

### 修改 run_ai_brain 函數

```python
async def run_ai_brain():
    """Execute AI brain processing with prompt templates."""
    try:
        # 1. Analyze and determine task type
        task_type = analyze_task_type(state)

        # 2. Generate prompt from template (prompts/{task_type}.md)
        prompt = prompt_manager.get_prompt(task_type, dict(state))

        # 3. Configure AI options with STM32 project access
        options = get_ai_options(state)

        # 4. Execute AI query
        responses = []
        result_message = None  # ← 新增：儲存 ResultMessage

        async for message in query(prompt=prompt, options=options):
            logger.debug(f"🧠 AI Brain: {message}")
            print_ai_message(message, prefix="🧠")
            responses.append(str(message))

            # ← 新增：擷取 ResultMessage
            if isinstance(message, ResultMessage):
                result_message = message

        # ← 新增：回報用量到雲端
        if result_message:
            _report_usage_to_server(result_message, task_type)

        return {"success": True, "responses": responses, "task_type": task_type}

    except Exception as e:
        logger.debug(f"AI Brain exception: {str(e)}")
        return {"success": False, "error": str(e)}
```

### 完整修改後的 ai_brain.py

```python
"""AI Brain - Unified intelligent decision maker with prompt templates."""

import asyncio
import json
import requests
from pathlib import Path
from typing import Any

from claude_agent_sdk import ClaudeAgentOptions, query
from claude_agent_sdk.types import ResultMessage

from ..ai_utils import print_ai_message
from ..logging_config import (
    get_logger,
    log_node_error,
    log_node_input_state,
    log_node_output_state,
    log_node_start,
    log_node_success,
)
from ..prompt_manager import PromptManager
from ..state import FirmwareState

# Global prompt manager
prompt_manager = PromptManager()

# 用量回報設定
USAGE_REPORT_URL = "http://35.229.221.64.nip.io:5000/api/report-usage"


def _get_user_token() -> str | None:
    """從本地配置讀取 user token"""
    config_file = Path.home() / '.claude_client_config.json'
    if config_file.exists():
        try:
            with open(config_file, 'r', encoding='utf-8') as f:
                config = json.load(f)
                return config.get('user_token')
        except Exception:
            pass
    return None


def _report_usage_to_server(
    result_message: ResultMessage,
    task_type: str,
    logger
) -> None:
    """回報用量到雲端 Server"""
    token = _get_user_token()
    if not token:
        logger.debug("No user token found, skipping usage report")
        return

    if not result_message.usage:
        logger.debug("No usage info in ResultMessage, skipping report")
        return

    try:
        response = requests.post(
            USAGE_REPORT_URL,
            headers={
                "Authorization": f"Bearer {token}",
                "Content-Type": "application/json"
            },
            json={
                "task_type": task_type,
                "model": "claude-sonnet-4",
                "input_tokens": result_message.usage.get("input_tokens", 0),
                "output_tokens": result_message.usage.get("output_tokens", 0),
                "total_cost_usd": result_message.total_cost_usd,
                "duration_ms": result_message.duration_ms,
                "num_turns": result_message.num_turns,
            },
            timeout=5
        )

        if response.status_code == 200:
            usage = result_message.usage
            cost = result_message.total_cost_usd or 0
            logger.debug(
                f"Usage reported: {usage.get('input_tokens', 0)} in + "
                f"{usage.get('output_tokens', 0)} out = ${cost:.4f}"
            )
        else:
            logger.debug(f"Usage report failed: {response.status_code}")

    except Exception as e:
        logger.debug(f"Usage report error: {e}")


def ai_brain_node(state: FirmwareState) -> dict[str, Any]:
    """
    AI Brain - Unified intelligent processor using prompt templates.
    """
    logger = get_logger("nodes.ai_brain")

    log_node_input_state("ai_brain", state)
    log_node_start("ai_brain")

    async def run_ai_brain():
        """Execute AI brain processing with prompt templates."""
        try:
            task_type = analyze_task_type(state)
            prompt = prompt_manager.get_prompt(task_type, dict(state))
            options = get_ai_options(state)

            responses = []
            result_message = None

            async for message in query(prompt=prompt, options=options):
                logger.debug(f"🧠 AI Brain: {message}")
                print_ai_message(message, prefix="🧠")
                responses.append(str(message))

                if isinstance(message, ResultMessage):
                    result_message = message

            # 回報用量
            if result_message:
                _report_usage_to_server(result_message, task_type, logger)

            return {"success": True, "responses": responses, "task_type": task_type}

        except Exception as e:
            logger.debug(f"AI Brain exception: {str(e)}")
            return {"success": False, "error": str(e)}

    result = asyncio.run(run_ai_brain())
    output_state = update_state_from_result(state, result)

    if result["success"]:
        log_node_success("ai_brain", f"Task completed: {result.get('task_type', 'unknown')}")
    else:
        log_node_error("ai_brain", result["error"])

    log_node_output_state("ai_brain", output_state)
    return output_state


def analyze_task_type(state: FirmwareState) -> str:
    """Analyze current state to determine task type for prompt selection."""
    logger = get_logger("nodes.ai_brain")
    logger.debug("🧠 AI Brain: Analyzing task type...")

    if state.get("build_errors") and state.get("build_status") == "error":
        logger.debug("  ✓ Build errors detected → fix_build_errors")
        return "fix_build_errors"
    elif state.get("execution_mode") == "feat":
        logger.debug("  ✓ Feature mode detected → implement_feature")
        return "implement_feature"
    else:
        logger.error(f"  ✗ Unexpected state: execution_mode={state.get('execution_mode')}")
        raise ValueError(
            f"Invalid execution_mode for AI Brain: {state.get('execution_mode')}. "
            "AI Brain only supports 'feat' mode or build error auto-fix."
        )


def get_ai_options(state: FirmwareState) -> ClaudeAgentOptions:
    """Configure AI options with project directory access."""
    logger = get_logger("nodes.ai_brain")
    project_root = state.get("project_root")

    options = {
        "allowed_tools": ["Read", "Write", "Edit", "MultiEdit", "Glob", "Grep", "Bash"],
        "permission_mode": "bypassPermissions",
    }

    if project_root:
        options["add_dirs"] = [project_root]
        options["cwd"] = project_root
        logger.debug(f"🧠 AI Brain: Working in project: {project_root}")
    else:
        logger.debug("🧠 AI Brain: No project path found, using default directory")

    return ClaudeAgentOptions(**options)


def update_state_from_result(state: FirmwareState, result: dict[str, Any]) -> dict[str, Any]:
    """Update firmware state based on AI processing result."""
    logger = get_logger("nodes.ai_brain")

    if result["success"]:
        logger.debug(f"AI Brain result: Success - {result.get('task_type', 'unknown')}")
        return {
            **state,
            "build_status": "success",
            "build_errors": [],
            "build_log": f"🧠 AI Brain completed: {result.get('task_type', 'unknown')}",
        }
    else:
        logger.debug(f"AI Brain result: Failed - {result['error']}")
        return {
            **state,
            "build_status": "error",
            "build_errors": [result["error"]],
            "build_log": f"🧠 AI Brain failed: {result['error']}",
            "iteration": state.get("iteration", 0) + 1,
        }
```

---

## 第二部分：Flask Server 端

### 新增 API 端點

```python
# app.py

from flask import Flask, request, jsonify, g
import sqlite3
from functools import wraps

app = Flask(__name__)
DATABASE_PATH = "fwauto_usage.db"


def get_db():
    if 'db' not in g:
        g.db = sqlite3.connect(DATABASE_PATH)
        g.db.row_factory = sqlite3.Row
    return g.db


@app.teardown_appcontext
def close_db(error):
    db = g.pop('db', None)
    if db is not None:
        db.close()


def require_auth(f):
    """認證裝飾器"""
    @wraps(f)
    def decorated(*args, **kwargs):
        auth_header = request.headers.get('Authorization', '')

        if not auth_header.startswith('Bearer '):
            return jsonify({"error": "Unauthorized"}), 401

        token = auth_header[7:]

        db = get_db()
        user = db.execute(
            'SELECT * FROM users WHERE user_token = ? AND is_active = 1',
            (token,)
        ).fetchone()

        if not user:
            return jsonify({"error": "Invalid token"}), 401

        g.current_user = dict(user)
        return f(*args, **kwargs)

    return decorated


@app.route('/api/report-usage', methods=['POST'])
@require_auth
def report_usage():
    """
    接收 Claude Agent SDK 用量回報

    Request:
    {
        "task_type": "fix_build_errors",
        "model": "claude-sonnet-4",
        "input_tokens": 5000,
        "output_tokens": 2000,
        "total_cost_usd": 0.035,
        "duration_ms": 45000,
        "num_turns": 8
    }
    """
    user = g.current_user
    data = request.json

    db = get_db()
    db.execute('''
        INSERT INTO usage_log
        (user_id, task_type, model, input_tokens, output_tokens, cost_usd, duration_ms, num_turns)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
    ''', (
        user['id'],
        data.get('task_type', 'unknown'),
        data.get('model', 'claude-sonnet-4'),
        data.get('input_tokens', 0),
        data.get('output_tokens', 0),
        data.get('total_cost_usd'),
        data.get('duration_ms'),
        data.get('num_turns'),
    ))
    db.commit()

    return jsonify({"status": "ok"})


@app.route('/api/usage', methods=['GET'])
@require_auth
def get_usage():
    """查詢用戶用量統計"""
    user = g.current_user
    db = get_db()

    result = db.execute('''
        SELECT
            COALESCE(SUM(input_tokens), 0) as total_input,
            COALESCE(SUM(output_tokens), 0) as total_output,
            COALESCE(SUM(cost_usd), 0) as total_cost,
            COUNT(*) as total_requests
        FROM usage_log
        WHERE user_id = ?
    ''', (user['id'],)).fetchone()

    return jsonify({
        "user": user['email'],
        "total_input_tokens": result['total_input'],
        "total_output_tokens": result['total_output'],
        "total_cost_usd": result['total_cost'],
        "total_requests": result['total_requests']
    })
```

### 資料庫 Schema

```sql
-- schema.sql

CREATE TABLE IF NOT EXISTS users (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    email TEXT UNIQUE NOT NULL,
    google_id TEXT,
    user_token TEXT UNIQUE NOT NULL,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    is_active BOOLEAN DEFAULT 1
);

CREATE TABLE IF NOT EXISTS usage_log (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id INTEGER NOT NULL,
    task_type TEXT,                    -- fix_build_errors / implement_feature
    model TEXT,                        -- claude-sonnet-4
    input_tokens INTEGER DEFAULT 0,
    output_tokens INTEGER DEFAULT 0,
    cost_usd REAL,                     -- 0.035
    duration_ms INTEGER,               -- 45000
    num_turns INTEGER,                 -- 8
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    FOREIGN KEY (user_id) REFERENCES users(id)
);

-- 索引
CREATE INDEX IF NOT EXISTS idx_usage_user_id ON usage_log(user_id);
CREATE INDEX IF NOT EXISTS idx_usage_created_at ON usage_log(created_at);
```

### 初始化資料庫

```python
def init_db():
    conn = sqlite3.connect(DATABASE_PATH)
    conn.executescript(open('schema.sql').read())
    conn.commit()
    conn.close()
    print("✅ Database initialized")

if __name__ == '__main__':
    init_db()
    app.run(host='0.0.0.0', port=5000)
```

---

## 第三部分：用戶配置

### 登入後儲存 user_token

Google SSO 登入成功後，在 `~/.claude_client_config.json` 儲存：

```json
{
    "user_token": "ut-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
    "user_email": "user@example.com"
}
```

### 現有 kenttest.py 修改

Google 登入成功後，回傳 `user_token` 而不是 `api_key`：

```python
# 原本
CallbackHandler.auth_code = f"api_key={api_key}&email={email}"

# 改成
CallbackHandler.auth_code = f"user_token={user_token}&email={email}"
```

---

## 第四部分：移除 LiteLLM

### 不再需要的東西

1. **LiteLLM Server (port 4000)** - 直接關掉
2. **kenttest.py 的 litellm_url** - 移除
3. **LiteLLM API Key** - 不需要了

### 保留的東西

1. **ANTHROPIC_API_KEY** - Claude Agent SDK 還是需要（環境變數）
2. **Flask Server (port 5000)** - 繼續用於認證 + 用量追蹤
3. **Google SSO** - 登入流程不變

---

## 改動摘要

| 檔案 | 改動內容 | 行數 |
|------|---------|------|
| `ai_brain.py` | 新增 `_report_usage_to_server()` | +40 行 |
| Flask `app.py` | 新增 `/api/report-usage` 端點 | +30 行 |
| `schema.sql` | 新增 `usage_log` 表 | +15 行 |
| LiteLLM | **刪除** | -∞ |

## 優點

1. **極簡改動** - 只加約 70 行程式碼
2. **不影響主流程** - 用量回報失敗不會中斷 AI 執行
3. **完整追蹤** - tokens、費用、執行時間、任務類型都記錄
4. **移除依賴** - 不需要維護 LiteLLM Server

## API 文件

### POST /api/report-usage

**用途**：接收 Claude Agent SDK 用量回報

**Headers**：
```
Authorization: Bearer {user_token}
Content-Type: application/json
```

**Request Body**：
```json
{
    "task_type": "fix_build_errors",
    "model": "claude-sonnet-4",
    "input_tokens": 5000,
    "output_tokens": 2000,
    "total_cost_usd": 0.035,
    "duration_ms": 45000,
    "num_turns": 8
}
```

**Response**：
```json
{"status": "ok"}
```

### GET /api/usage

**用途**：查詢用戶用量統計

**Headers**：
```
Authorization: Bearer {user_token}
```

**Response**：
```json
{
    "user": "user@example.com",
    "total_input_tokens": 50000,
    "total_output_tokens": 20000,
    "total_cost_usd": 0.35,
    "total_requests": 15
}
```
