幫 Claude Agent SDK 接上 MLflow Tracing：官方 autolog + 30 行補丁層（cost、chat id），與全手動重建 span 樹的取捨

python scripts/sdk_surface.py
# claude-agent-sdk == 0.2.104
# ResultMessage -> ['subtype', 'duration_ms', 'duration_api_ms', 'is_error',
#   'num_turns', 'session_id', 'stop_reason', 'total_cost_usd', 'usage',
#   'result', 'structured_output', 'model_usage', ...]
# AssistantMessage -> ['content', 'model', 'parent_tool_use_id', 'error',
#   'usage', 'message_id', 'stop_reason', ...]
Code language: PHP (php)

{
  "mlflow-mcp": {
    "type": "stdio",
    "command": "uv",
    "args": ["run", "--with", "mlflow[mcp]>=3.5.1", "mlflow", "mcp", "run"],
    "env": { "MLFLOW_TRACKING_URI": "http://your-mlflow-server:5000" }
  }
}
Code language: JSON / JSON with Comments (json)

import mlflow
import mlflow.anthropic

mlflow.set_experiment("my_claude_app")
mlflow.anthropic.autolog()   # 必須在 ClaudeSDKClient 建立之前

client = ClaudeSDKClient(options=...)  # 之後照常用
Code language: PHP (php)

# mlflow/anthropic/__init__.py（節錄）
def autolog(...):
    from anthropic.resources import AsyncMessages, Messages   # ← 硬依賴 anthropic 套件！

    safe_patch(FLAVOR_NAME, Messages, "create", patched_class_call)
    safe_patch(FLAVOR_NAME, AsyncMessages, "create", async_patched_class_call)

    # Patch Claude Code SDK if available
    try:
        from claude_agent_sdk import ClaudeSDKClient
        safe_patch(FLAVOR_NAME, ClaudeSDKClient, "__init__", patched_claude_sdk_init)
    except ImportError:
        ...
Code language: PHP (php)

claude_code_conversation (AGENT)   ← root，usage 完整（來自 ResultMessage）
├── tool_Bash (TOOL)               duration = 0.1 ms（假的，實際跑了幾秒）
└── llm (LLM)                      duration = 0.1 ms（name 固定叫 "llm"）

metadata: mlflow.trace.session = CLI 內部 session id
          mlflow.trace.user    = OS 使用者名
token_usage: {'input_tokens': 18, 'output_tokens': 210,
              'cache_read_input_tokens': 18767, 'cache_creation_input_tokens': 18944}
Code language: JavaScript (javascript)

uv add mlflow-skinny anthropic python-dateutil

"""MLflow tracing via the official integration + a patch layer for what it misses."""

import asyncio
import json
import logging
import os

import mlflow
import mlflow.anthropic
from claude_agent_sdk import ResultMessage
from mlflow.tracing.constant import TraceMetadataKey

logger = logging.getLogger(__name__)

_enabled = False

def setup_tracing() -> bool:
    """App 啟動時（dotenv 之後、任何 ClaudeSDKClient 建立之前）呼叫一次。"""
    global _enabled
    uri = os.environ.get("MLFLOW_TRACKING_URI")
    if not uri:
        logger.info("MLflow tracing disabled (MLFLOW_TRACKING_URI not set)")
        return False
    try:
        mlflow.set_tracking_uri(uri)
        mlflow.set_experiment(os.environ.get("MLFLOW_EXPERIMENT_NAME", "chatapp"))
        mlflow.anthropic.autolog()   # patch ClaudeSDKClient.__init__
        _enabled = True
        logger.info("MLflow tracing enabled (official autolog) -> %s", uri)
    except Exception as error:
        logger.warning("MLflow tracing disabled (server unreachable?): %s", error)
        _enabled = False
    return _enabled
Code language: PHP (php)

async def finalize_turn_trace(
    chat_id: str,
    cli_session_id: str | None,
    result: ResultMessage | None,
    turn_cost: float,
    interrupted: bool = False,
) -> None:
    """幫 autolog 剛寫完的 trace 補上它不記的東西。best-effort、絕不拋錯。"""
    if not _enabled or result is None:
        return
    try:
        await asyncio.to_thread(   # 網路操作丟 thread，不卡 event loop
            _finalize_sync, chat_id, cli_session_id, result, turn_cost, interrupted
        )
    except Exception:
        logger.debug("trace finalize failed", exc_info=True)

def _finalize_sync(chat_id, cli_session_id, result, turn_cost, interrupted):
    mlflow.flush_trace_async_logging()          # 等 trace 真的上了 server
    trace_id = mlflow.get_last_active_trace_id()
    if not trace_id:
        return
    # 多聊天並行防搶錯：官方 trace 的 session = CLI session id
    metadata = mlflow.get_trace(trace_id).info.trace_metadata or {}
    trace_session = metadata.get(TraceMetadataKey.TRACE_SESSION)
    if cli_session_id and trace_session and trace_session != cli_session_id:
        return                                   # 不是我這輪的，放手
    tags = {"chat_id": chat_id, "turn_cost_usd": f"{turn_cost:.6f}"}
    if result.total_cost_usd:
        tags["cli_total_cost_usd"] = f"{result.total_cost_usd:.6f}"
    if result.num_turns is not None:
        tags["num_turns"] = str(result.num_turns)
    if result.model_usage:
        tags["model_usage"] = json.dumps(result.model_usage, ensure_ascii=False)[:4000]
    if interrupted:
        tags["interrupted"] = "true"
    for key, value in tags.items():
        mlflow.set_trace_tag(trace_id, key, value)
Code language: PHP (php)

async with self._turn_lock:
    self._last_result = None       # 迴圈裡收到 ResultMessage 時存起來
    self._last_turn_cost = 0.0
    try:
        await self.agent.client.query(prompt)
        async for message in self.agent.client.receive_response():
            await self._handle_sdk_message(message)   # 這裡面存 _last_result
        # ↑ 迭代結束的瞬間 autolog 建好 trace —— 現在才能補 tag
        await finalize_turn_trace(
            chat_id=self.chat_id,
            cli_session_id=self._session_id,
            result=self._last_result,
            turn_cost=self._last_turn_cost,
        )
    except Exception as error:
        ...   # 異常回合官方不記 trace，也就沒東西可補
Code language: PHP (php)

tags:
    chat_id            = 20260618-165650
    turn_cost_usd      = 0.145098
    cli_total_cost_usd = 0.145098
    num_turns          = 1
    model_usage        = {"claude-haiku-4-5-...": {...}, "claude-opus-4-8": {...}}
Code language: JavaScript (javascript)

chat_turn (AGENT)
├── claude-opus-4-8 (LLM)
├── Read (TOOL)
├── Agent (TOOL)                ← 子代理
│   ├── claude-opus-4-8 (LLM)
│   └── Bash (TOOL)
└── claude-opus-4-8 (LLM)

"""MLflow tracing: one chat turn = one trace, built from the SDK message stream."""

import logging
import os
from typing import Any

import mlflow
from claude_agent_sdk import (
    AssistantMessage,
    ResultMessage,
    TextBlock,
    ThinkingBlock,
    ToolResultBlock,
    ToolUseBlock,
    UserMessage,
)
from mlflow.entities import SpanType
from mlflow.tracing.constant import SpanAttributeKey, TraceMetadataKey

logger = logging.getLogger(__name__)

_TRUNCATE_CHARS = 4000  # 工具輸出可能是整包檔案，砍掉免得灌爆 trace store
_enabled = False

def setup_tracing() -> bool:
    """指向 tracking server。App 啟動時（dotenv 之後）呼叫一次。"""
    global _enabled
    uri = os.environ.get("MLFLOW_TRACKING_URI")
    if not uri:
        logger.info("MLflow tracing disabled (MLFLOW_TRACKING_URI not set)")
        return False
    try:
        mlflow.set_tracking_uri(uri)
        mlflow.set_experiment(os.environ.get("MLFLOW_EXPERIMENT_NAME", "chatapp"))
        # 3.14 起 trace 上傳預設就是 async，不會拖慢回合
        _enabled = True
        logger.info("MLflow tracing enabled -> %s", uri)
    except Exception as error:
        logger.warning("MLflow tracing disabled (server unreachable?): %s", error)
        _enabled = False
    return _enabled

def _truncate(value: Any) -> Any:
    text = value if isinstance(value, str) else repr(value)
    if len(text) > _TRUNCATE_CHARS:
        return text[:_TRUNCATE_CHARS] + f"… [+{len(text) - _TRUNCATE_CHARS} chars]"
    return value

def _usage_dict(usage: Any) -> dict | None:
    """把 SDK 的 Anthropic 風格 usage 映射到 MLflow 的 token-usage 格式。"""
    if not isinstance(usage, dict):
        return None
    keys = (
        "input_tokens",
        "output_tokens",
        "cache_read_input_tokens",
        "cache_creation_input_tokens",
    )
    out = {k: usage[k] for k in keys if isinstance(usage.get(k), int)}
    if not out:
        return None
    out["total_tokens"] = sum(out.values())
    return out

class TurnTrace:
    """每則 SDK 訊息餵給 observe()；回合結束呼叫 close()。"""

    def __init__(self, chat_id: str, prompt: str) -> None:
        self._root = None
        self._tools: dict[str, Any] = {}  # tool_use_id -> 開著的 TOOL span
        self._texts: list[str] = []
        self._result: ResultMessage | None = None
        if not _enabled:
            return
        try:
            self._root = mlflow.start_span_no_context(
                name="chat_turn",
                span_type=SpanType.AGENT,
                inputs={"prompt": prompt},
                # 標準 session 欄位：同一聊天的回合在 UI 會被歸成一組
                metadata={TraceMetadataKey.TRACE_SESSION: chat_id},
                tags={"chat_id": chat_id},
            )
        except Exception as error:
            logger.warning("MLflow trace not started: %s", error)
            self._root = None

    def observe(self, message: Any) -> None:
        if self._root is None:
            return
        try:
            if isinstance(message, AssistantMessage):
                self._on_assistant(message)
            elif isinstance(message, UserMessage):
                self._on_tool_results(message)
            elif isinstance(message, ResultMessage):
                self._result = message
        except Exception:
            logger.debug("trace observe failed", exc_info=True)

    def _parent_for(self, parent_tool_use_id: str | None) -> Any:
        # 子代理的訊息帶著 spawn 它的 Agent 工具 id，
        # 掛到那個（還開著的）TOOL span 底下 → 自然巢狀
        return self._tools.get(parent_tool_use_id) or self._root

    def _on_assistant(self, message: AssistantMessage) -> None:
        parent = self._parent_for(message.parent_tool_use_id)
        texts: list[str] = []
        thinking_chars = 0
        tool_uses: list[ToolUseBlock] = []
        for block in message.content:
            if isinstance(block, TextBlock):
                texts.append(block.text)
            elif isinstance(block, ThinkingBlock):
                thinking_chars += len(block.thinking)
            elif isinstance(block, ToolUseBlock):
                tool_uses.append(block)
        if message.parent_tool_use_id is None:
            self._texts.extend(texts)  # 只彙整主線回覆，子代理的不算

        # 訊息到達時已完整，LLM span 在這裡開了馬上關
        attributes: dict[str, Any] = {}
        if message.model:
            attributes[SpanAttributeKey.MODEL] = message.model
        usage = _usage_dict(message.usage)
        if usage:
            # 注意：放普通屬性，「不要」放標準 token-usage 鍵！
            # per-message usage 來自 stream 的 message_start，
            # output_tokens 是 stub —— 準確總量在 close() 放 root（坑 1）
            attributes["usage"] = usage
        span = mlflow.start_span_no_context(
            name=message.model or "assistant",
            span_type=SpanType.LLM,
            parent_span=parent,
        )
        outputs: dict[str, Any] = {}
        if texts:
            outputs["text"] = _truncate("\n".join(texts))
        if thinking_chars:
            outputs["thinking_chars"] = thinking_chars
        if tool_uses:
            outputs["tool_calls"] = [{"name": b.name, "id": b.id} for b in tool_uses]
        span.end(outputs=outputs or None, attributes=attributes)

        # 每個工具呼叫開一個 TOOL span，等配對的結果回來再關
        for block in tool_uses:
            self._tools[block.id] = mlflow.start_span_no_context(
                name=block.name,
                span_type=SpanType.TOOL,
                parent_span=parent,
                inputs=block.input,
            )

    def _on_tool_results(self, message: UserMessage) -> None:
        content = message.content
        if not isinstance(content, list):
            return
        for block in content:
            if isinstance(block, ToolResultBlock):
                span = self._tools.pop(block.tool_use_id, None)
                if span is not None:
                    span.end(
                        outputs=_truncate(block.content),
                        status="ERROR" if block.is_error else "OK",
                    )

    def close(self, error: BaseException | None = None) -> None:
        if self._root is None:
            return
        root, self._root = self._root, None  # 冪等
        try:
            # 回合中斷/失敗可能留下沒關的工具 span
            for span in self._tools.values():
                span.end(status="ERROR" if error else "OK")
            self._tools.clear()

            result = self._result
            attributes: dict[str, Any] = {}
            outputs: dict[str, Any] = {}
            if self._texts:
                outputs["reply"] = _truncate("\n".join(self._texts))
            if result is not None:
                outputs["subtype"] = result.subtype
                for key in ("total_cost_usd", "num_turns", "duration_ms", "duration_api_ms"):
                    value = getattr(result, key, None)
                    if value is not None:
                        attributes[key] = value
                if result.model_usage:
                    attributes["model_usage"] = result.model_usage
                # 本回合權威 token 總量：整條 trace「唯一」的標準
                # token-usage 屬性，所以 trace 層聚合 = 它（坑 1 的解法）
                usage = _usage_dict(result.usage)
                if usage:
                    attributes[SpanAttributeKey.CHAT_USAGE] = usage
            if error is not None:
                outputs["error"] = str(error)
            failed = error is not None or (result is not None and result.is_error)
            root.end(
                outputs=outputs or None,
                attributes=attributes or None,
                status="ERROR" if failed else "OK",
            )
        except Exception:
            logger.debug("trace close failed", exc_info=True)

from .tracing import TurnTrace

async with self._turn_lock:
    trace = TurnTrace(self.chat_id, display)          # 1. 開 trace
    try:
        await self.agent.client.query(prompt)
        async for message in self.agent.client.receive_response():
            trace.observe(message)                     # 2. 每則訊息餵進去
            await self._handle_sdk_message(message)
        trace.close()                                  # 3. 正常收尾
    except Exception as error:
        trace.close(error)                             # 4. 錯誤也收尾
        ...
Code language: PHP (php)

trace tokens: {'input_tokens': 7164, 'output_tokens': 1, ...}
Code language: JavaScript (javascript)

curl -X POST "$MLFLOW_TRACKING_URI/api/2.0/mlflow/experiments/delete" \
  -H "Content-Type: application/json" -d '{"experiment_id": "2"}'
Code language: JavaScript (javascript)

14:42 spans=15  duration=18ms    ← 15 個 span 擠在 18ms？
13:40 spans=48  duration=44ms    ← 48 個 span 擠在 44ms？？
12:54 spans=363 duration=22 分鐘  ← 這條才是正常的

pkill -f "npm run dev"; sleep 2; nohup npm run dev &
Code language: JavaScript (javascript)

trace: tr-0bc949373eef29a8... | state: OK
spans: claude_code_conversation (AGENT) └─ llm (LLM)
token_usage: {'input_tokens': 7411, 'output_tokens': 8,
              'cache_read_input_tokens': 17192, 'cache_creation_input_tokens': 9857}
tags: chat_id=20260618-165650  turn_cost_usd=0.145098
      num_turns=1  model_usage={...per-model...}
Code language: JavaScript (javascript)