Skip to main content

tiny_agent/trajectory/
mod.rs

1//! Trajectory:一次 agent 运行过程中产生的**只读事件流**。
2//!
3//! 它回答的是"刚刚发生了什么",区别于 `Transcript`(喂给模型的消息历史)和
4//! `Checkpoint`(可恢复的状态快照)。典型消费者:
5//! - `examples/websocket`:run 在一个 task 里跑,事件经 channel 实时推给客户端;
6//! - 日志 / 回放 / 评测。
7//!
8//! 运行时在关键节点同步 `emit`,因此 sink 收到事件的顺序天然等于发生顺序;
9//! `at`(毫秒时间戳)只是附带的排序/展示元数据。
10
11use crate::shared::{ContentBlock, UserInteraction};
12use async_trait::async_trait;
13use serde::{Deserialize, Serialize};
14use serde_json::Value;
15use std::sync::Arc;
16use std::time::{SystemTime, UNIX_EPOCH};
17
18/// 单条轨迹事件:会话 id + 发生时间 + 具体事件。
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct TrajectoryEvent {
21    pub session_id: String,
22    pub at: u64,
23    #[serde(flatten)]
24    pub kind: TrajectoryEventKind,
25}
26
27impl TrajectoryEvent {
28    pub fn new(session_id: &str, kind: TrajectoryEventKind) -> Self {
29        let at = SystemTime::now()
30            .duration_since(UNIX_EPOCH)
31            .map(|d| d.as_millis() as u64)
32            .unwrap_or(0);
33        Self {
34            session_id: session_id.to_string(),
35            at,
36            kind,
37        }
38    }
39}
40
41/// 事件种类。序列化为 `{ "type": "user_message", ... }`,方便前端按 tag 分发。
42#[derive(Debug, Clone, Serialize, Deserialize)]
43#[serde(tag = "type", rename_all = "snake_case")]
44pub enum TrajectoryEventKind {
45    /// 用户向会话提交了一条消息(`submit` 入口)。轨迹的对话起点。
46    UserMessage { text: String },
47    /// 一次推理开始(provider 即将被调用)。
48    /// `request` 为 composer + middleware 之后,由 provider 投影出的请求快照。
49    ReasoningStarted { request: Value },
50    /// 一条 assistant 消息已成型并落库(含 thinking / text / tool_use 块)。
51    AssistantMessage { blocks: Vec<ContentBlock> },
52    /// 某个工具开始执行。
53    ToolStarted {
54        call_id: String,
55        name: String,
56        arguments: Value,
57    },
58    /// 某个工具执行结束。
59    ToolCompleted {
60        call_id: String,
61        name: String,
62        output: Vec<ContentBlock>,
63        is_error: bool,
64    },
65    /// 运行被取消并已存档 checkpoint。
66    Interrupted,
67    /// 工具已返回合法结果,但该结果要求用户交互后再继续会话。
68    WaitingForUser { interaction: UserInteraction },
69    /// 进入失败(可能仍会重试,见 max_consecutive_fail_count)。
70    Failed { kind: String, message: String },
71    /// 推理未再产生工具调用,运行成功结束。
72    Succeeded { final_resp: Option<String> },
73}
74
75impl TrajectoryEventKind {
76    pub fn kind_type(&self) -> &'static str {
77        match self {
78            Self::UserMessage { .. } => "user_message",
79            Self::ReasoningStarted { .. } => "reasoning_started",
80            Self::AssistantMessage { .. } => "assistant_message",
81            Self::ToolStarted { .. } => "tool_started",
82            Self::ToolCompleted { .. } => "tool_completed",
83            Self::Interrupted => "interrupted",
84            Self::WaitingForUser { .. } => "waiting_for_user",
85            Self::Failed { .. } => "failed",
86            Self::Succeeded { .. } => "succeeded",
87        }
88    }
89}
90
91/// 轨迹事件的下游。实现者负责把事件送到日志 / channel / websocket 等。
92///
93/// `emit` 不返回错误:可观测性不应阻断主流程,投递失败(如 channel 关闭)由实现内部吞掉。
94#[async_trait]
95pub trait TrajectorySink: Send + Sync {
96    async fn emit(&self, event: TrajectoryEvent);
97}
98
99/// 默认 sink:丢弃所有事件。未配置 trajectory 时使用。
100pub struct NoopSink;
101
102#[async_trait]
103impl TrajectorySink for NoopSink {
104    async fn emit(&self, _event: TrajectoryEvent) {}
105}
106
107/// 基于无界 mpsc 的 sink,为 websocket 场景准备:
108/// 一个 task 跑 `AgentRunTime::run`,另一个 task 从 `receiver` 拉事件推给客户端。
109pub struct ChannelSink {
110    tx: tokio::sync::mpsc::UnboundedSender<TrajectoryEvent>,
111}
112
113impl ChannelSink {
114    /// 返回 sink 与配对的 receiver。run 结束、sink 被 drop 后,receiver 收到 `None`。
115    pub fn new() -> (
116        Arc<Self>,
117        tokio::sync::mpsc::UnboundedReceiver<TrajectoryEvent>,
118    ) {
119        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
120        (Arc::new(Self { tx }), rx)
121    }
122}
123
124#[async_trait]
125impl TrajectorySink for ChannelSink {
126    async fn emit(&self, event: TrajectoryEvent) {
127        // 客户端已断开 → receiver 被 drop,这里静默丢弃。
128        let _ = self.tx.send(event);
129    }
130}