Agent/cyber_agent/trace/trace_id.py

"""
Trace ID 生成与逐层解析工具。

Sub-Agent 创建 Trace 时生成嵌套 ID；查询和辅助显示只解析最后一层，
权威父子关系始终以 Trace.parent_trace_id 为准。

Trace ID 格式：
- 主 Trace: {uuid} (标准 UUID)
- Sub-Trace: {parent_id}@{mode}-{timestamp}-{seq}
  例如: 2f8d3a1c-4b6e-4f9a-8c2d-1e5b7a9f3c4d@explore-20260204220012-001
"""

import uuid
from datetime import datetime
from threading import Lock
from typing import Dict, Optional


# 全局计数器（线程安全）
_seq_lock = Lock()
_seq_counter: Dict[str, int] = {}  # key: "{parent_id}@{mode}-{timestamp}"


def generate_trace_id() -> str:
    """
    生成主 Trace ID

    Returns:
        标准 UUID 字符串
    """
    return str(uuid.uuid4())


def generate_sub_trace_id(parent_id: str, mode: str) -> str:
    """
    生成 Sub-Trace ID

    格式: {parent_id}@{mode}-{timestamp}-{seq}

    使用完整的 parent_id（不截断），避免 ID 冲突风险。
    同一秒内多次调用会递增序号。

    Args:
        parent_id: 父 Trace ID（完整 UUID）
        mode: 运行模式（explore, delegate, compaction 等）

    Returns:
        Sub-Trace ID

    Examples:
        >>> generate_sub_trace_id("2f8d3a1c-4b6e-4f9a-8c2d-1e5b7a9f3c4d", "explore")
        '2f8d3a1c-4b6e-4f9a-8c2d-1e5b7a9f3c4d@explore-20260204220012-001'

        >>> generate_sub_trace_id("2f8d3a1c-4b6e-4f9a-8c2d-1e5b7a9f3c4d", "delegate")
        '2f8d3a1c-4b6e-4f9a-8c2d-1e5b7a9f3c4d@delegate-20260204220030-001'
    """
    # 直接使用完整 UUID，不截断
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")

    # 生成序号（同一秒内递增）
    prefix = f"{parent_id}@{mode}-{timestamp}"
    with _seq_lock:
        seq = _seq_counter.get(prefix, 0) + 1
        _seq_counter[prefix] = seq

    return f"{prefix}-{seq:03d}"


def parse_parent_trace_id(trace_id: str) -> Optional[str]:
    """
    从嵌套 ID 最后一个 @ 之前解析直接父 Trace ID。

    该函数仅供追溯与辅助显示，不取代持久化的 parent_trace_id 安全校验。

    Args:
        trace_id: Trace ID

    Returns:
        父 Trace ID，如果是主 Trace 则返回 None

    Examples:
        >>> parse_parent_trace_id("2f8d3a1c-4b6e-4f9a-8c2d-1e5b7a9f3c4d@explore-20260204220012-001")
        '2f8d3a1c-4b6e-4f9a-8c2d-1e5b7a9f3c4d'

        >>> parse_parent_trace_id("2f8d3a1c-4b6e-4f9a-8c2d-1e5b7a9f3c4d")
        None
    """
    if '@' in trace_id:
        return trace_id.rsplit('@', 1)[0]
    return None


def is_sub_trace(trace_id: str) -> bool:
    """
    判断是否为 Sub-Trace

    Args:
        trace_id: Trace ID

    Returns:
        True 表示是 Sub-Trace，False 表示是主 Trace

    Examples:
        >>> is_sub_trace("2f8d3a1c-4b6e-4f9a-8c2d-1e5b7a9f3c4d@explore-20260204220012-001")
        True

        >>> is_sub_trace("2f8d3a1c-4b6e-4f9a-8c2d-1e5b7a9f3c4d")
        False
    """
    return '@' in trace_id


def extract_mode(trace_id: str) -> Optional[str]:
    """
    从最后一层 Sub-Trace ID 中提取当前子 Trace 的运行模式。

    递归 ID 可含多个 @，因此这里与 parse_parent_trace_id() 一样从末层解析。

    Args:
        trace_id: Trace ID

    Returns:
        运行模式（explore, delegate 等），如果是主 Trace 则返回 None

    Examples:
        >>> extract_mode("2f8d3a1c-4b6e-4f9a-8c2d-1e5b7a9f3c4d@explore-20260204220012-001")
        'explore'

        >>> extract_mode("2f8d3a1c-4b6e-4f9a-8c2d-1e5b7a9f3c4d@delegate-20260204220030-001")
        'delegate'

        >>> extract_mode("2f8d3a1c-4b6e-4f9a-8c2d-1e5b7a9f3c4d")
        None
    """
    if '@' not in trace_id:
        return None

    # 格式: parent@mode-timestamp-seq
    parts = trace_id.rsplit('@', 1)[1]  # "mode-timestamp-seq"
    mode = parts.split('-')[0]
    return mode


def reset_seq_counter():
    """
    重置序号计数器

    主要用于测试，生产环境不应调用此函数。
    """
    global _seq_counter
    with _seq_lock:
        _seq_counter.clear()