透传原始上下文¶

TFRobot 的遥测系统基于 OpenTelemetry 构建。在实际业务中，工程师经常需要在 Hook 回调中获取调用方的业务上下文——比如用户 ID、会话 ID、请求来源等信息。这些信息不属于 TFRobot 内部状态，但 Hook 处理函数需要它们来完成日志关联、审计、计费等工作。

本文介绍 TFRobot 如何利用 OpenTelemetry Baggage 机制解决这个问题，以及最终的效果。

要解决的问题¶

考虑以下场景：你注册了一个 after_chain_run 的 Hook，希望在 Chain 执行结束后把结果写入某个业务系统。但 Hook 的回调签名是固定的——你只能拿到 event_name 和 attributes（即 BaseEventContext），拿不到调用方的业务上下文。

如果不借助上下文透传，你只能依赖全局变量或闭包来传递这些信息，这在并发场景下会出错。

TFRobot 中的实现¶

`span_decorator` 自动采集 Baggage¶

TFRobot 的核心组件（Robot、Brain、Chain、LLM、Tool、Memory）都通过 span_decorator 包装了关键方法。在每次 Span 事件触发时，装饰器会自动读取当前线程的 Baggage 并写入事件上下文：

# tfrobot/telemetry/context_schema.py — span_decorator 内部逻辑（简化）
with tracer.start_as_current_span(func.__name__, attributes=attr) as span:
    span_context = span.get_span_context()
    bg = baggage.get_all()           # 读取调用方附着的所有 Baggage
    bg_str = json.dumps(dict(bg))

    before_event_attr.update(
        span_id=format(span_context.span_id, "016x"),
        trace_id=format(span_context.trace_id, "032x"),
        baggage=bg_str,              # 写入事件上下文
    )
    span.add_event(before_event_name.value, attributes=before_event_attr)

`BaseEventContext.baggage` 字段¶

所有事件上下文都继承自 BaseEventContext，其中包含 baggage 字段：

class BaseEventContext(BaseModel):
    scene: EventScene
    entity_id: str
    span_id: Optional[str] = NOT_SET
    trace_id: Optional[str] = NOT_SET
    desc: Optional[str] = None
    info: str
    baggage: Optional[str] = NOT_SET   # 调用方透传的上下文，JSON 字符串

这意味着：任何注册了 Hook 的监听者，都可以从事件属性中直接读取 baggage 字段，无需额外传参。

实际案例：TFRobotServer 中的会话 ID 透传¶

TFRobotServer 是 TFRobot 的 HTTP 服务层，它的核心链路是：FastAPI 接口 → Celery 异步任务 → TFRobot.run()。在这条链路中，conversation_id（会话 ID）需要从 HTTP 请求一直传递到 TFRobot 内部的 Hook，用于将 Agent 运行轨迹（Trajectory）关联到正确的会话。

以下是真实的实现过程。

第一步：FastAPI 接口附着 Baggage¶

用户发送消息时，add_message 接口将 conversation_id 写入 Baggage，然后触发 Celery 任务：

# tfrobotserver/api/v1/chat/chats.py — add_message 方法（简化）
bg_with_conversation_id = baggage.set_baggage(
    "conversation_info", json.dumps({"conversation_id": conversation_id})
)
token = attach(bg_with_conversation_id)
try:
    task_id = self.run_service.run(robot_id, tf_msg)  # 发送 Celery 任务
finally:
    detach(token)

Celery 的 CeleryInstrumentor 配置了 propagate_headers=["traceparent", "tracestate", "baggage"]，会自动将当前 OpenTelemetry Context（包括 Baggage）序列化到 Celery 消息头中。

第二步：Celery Worker 恢复 Baggage¶

Celery Worker 收到任务后，在 TFSRobotTask.__call__ 中从消息头提取并恢复上下文：

# tfrobotserver/robot_celery.py — TFSRobotTask（简化）
class TFSRobotTask(Task):
    def __call__(self, *args, **kwargs):
        headers = self.request.headers
        context = extract(headers) if headers else None

        tracer = get_tf_tracer_provider().get_tracer(self.name)
        ctx_token = None
        if context:
            # 必须显式 attach，仅在 start_as_current_span 中传 context 不会传播 Baggage
            ctx_token = attach(context)
        try:
            with tracer.start_as_current_span("TFRobotCelery::robot.run", context=context, ...):
                return super().__call__(*args, **kwargs)
        finally:
            if ctx_token:
                detach(ctx_token)

关键陷阱

仅在 tracer.start_as_current_span(context=context) 中传递 context，Baggage 不会被传播。必须显式调用 attach(context) 才能让后续的 baggage.get_all() 读取到值。这是 TFRobotServer 开发中踩过的坑。

第三步：TFRobot 内部自动采集¶

从这里开始就是 TFRobot 框架的工作了。robot.run() → Brain → Chain → LLM/Tool 的每个 span_decorator 都会自动调用 baggage.get_all() 并写入 EventContext 的 baggage 字段。调用方无需做任何额外操作。

第四步：Hook 中消费 Baggage¶

TFRobotServer 注册了一个 Hook 用于记录运行轨迹（Trajectory）。它从 EventContext 的 baggage 字段中解析出 conversation_id，将每个事件关联到对应的会话：

# tfrobotserver/telemetry/event_trajectory.py — record_trajectory_event（简化）
def record_trajectory_event(name: str, context: types.Attributes) -> Optional[Trajectory]:
    context = dict(context or {})
    if not context.get("conversation_id"):
        baggage_str = context.get("baggage", "{}")
        baggage_data = json.loads(baggage_str)
        conversation_info = json.loads(baggage_data.get("conversation_info", "{}"))
        conversation_id = conversation_info.get("conversation_id")
        if conversation_id is not None:
            context["conversation_id"] = conversation_id

    # 使用 conversation_id + trace_id + span_id 构建 Trajectory 记录并写入数据库
    ...

完整数据流¶

sequenceDiagram
    participant Client as 客户端
    participant API as FastAPI
    participant Celery as Celery Worker
    participant Robot as TFRobot
    participant Hook as Trajectory Hook

    Client->>API: POST /conversations/{id}/messages
    API->>API: baggage.set_baggage("conversation_info", ...)
    API->>API: attach(ctx)
    API->>Celery: send_task("robot.run")<br/>消息头携带 baggage

    Celery->>Celery: extract(headers) → context
    Celery->>Celery: attach(context)
    Celery->>Robot: robot.run(msg)

    Robot->>Robot: span_decorator 自动<br/>baggage.get_all() → EventContext
    Robot->>Hook: add_event 触发 HookManager
    Hook->>Hook: 从 context["baggage"]<br/>解析 conversation_id
    Hook->>Hook: 写入 Trajectory 数据库

效果与启示¶

维度	效果
解耦	FastAPI 层只管附着 Baggage，TFRobot 只管采集，Hook 只管消费——三方互不感知
跨进程透传	从 FastAPI 进程到 Celery Worker 进程，Baggage 通过消息头自动传播
并发安全	OpenTelemetry Context 是线程/协程隔离的，不同请求的 Baggage 互不干扰
零侵入	TFRobot 框架本身不需要知道 `conversation_id` 的存在，`span_decorator` 透明地传递

请求生命周期

Baggage 的生命周期与 attach/detach 的作用域一致。务必在 finally 块中调用 detach，避免上下文泄漏到后续请求。在异步场景中，确保 attach/detach 在同一个协程内配对使用。