Skip to main content

tiny_agent/messages/
mod.rs

1//! Messages:框架对外的**实时消息通道**。
2//!
3//! 它是推给**前端实时渲染**的内容流,承载流式增量(`Partial*`)与成型消息。
4//! 观测/监控则走标准 `tracing` span/event,由宿主程序决定是否导出到 OpenTelemetry。
5//!
6//! 故意做成一个**具体 struct**([`MessageChannel`])而非 trait —— 内部就是一个 mpsc,
7//! 使用方拿 `receiver` 消费即可,不需要实现任何 sink。未配置时用 [`MessageChannel::disconnected`],
8//! `send` 静默丢弃。
9
10use crate::shared::ContentBlock;
11use serde::{Deserialize, Serialize};
12use serde_json::Value;
13use std::time::{SystemTime, UNIX_EPOCH};
14use tokio::sync::mpsc;
15
16/// 一条对前端的消息:会话 id + 发生时间 + 具体内容。
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct AgentMessage {
19    pub session_id: String,
20    pub at: u64,
21    #[serde(flatten)]
22    pub kind: MessageKind,
23}
24
25impl AgentMessage {
26    pub fn new(session_id: &str, kind: MessageKind) -> Self {
27        let at = SystemTime::now()
28            .duration_since(UNIX_EPOCH)
29            .map(|d| d.as_millis() as u64)
30            .unwrap_or(0);
31        Self {
32            session_id: session_id.to_string(),
33            at,
34            kind,
35        }
36    }
37}
38
39/// 消息内容。`Partial*` 是流式增量,与之配对的成型变体在内容完整时发出。
40///
41/// 序列化为 `{ "type": "tool_call", ... }`,方便前端按 tag 分发。
42#[derive(Debug, Clone, Serialize, Deserialize)]
43#[serde(tag = "type", rename_all = "snake_case")]
44pub enum MessageKind {
45    /// 用户提交的一条消息。
46    UserMessage { text: String },
47    /// 一次推理开始(provider 即将被调用)。
48    ReasoningStarted,
49    /// assistant 文本的流式增量。(预留:需 provider 支持 token 流式后才发出。)
50    PartialAssistantMessage { delta: String },
51    /// 一条成型的 assistant 消息(含 thinking / text / tool_use 块)。
52    AssistantMessage { blocks: Vec<ContentBlock> },
53    /// 工具调用参数的流式拼装增量。(预留:同上。)
54    PartialToolCall {
55        call_id: String,
56        name: String,
57        delta: String,
58    },
59    /// 模型决定调用某工具(参数已成型)。
60    ToolCall {
61        call_id: String,
62        name: String,
63        arguments: Value,
64    },
65    /// 工具执行期间的实时原始输出(如 bash 边跑边吐)。**纯前端展示用**:
66    /// 不进 transcript、模型也看不到——它不是 `ToolResponse` 的一部分。
67    /// 模型看到的真正结果始终是下面那条完整的 `ToolResponse`(后台工具任务则是 `{id, running}` 信封)。
68    /// 不带 call_id —— 归属"当前正在执行的工具"(非只读工具独占串行,任一时刻至多一个)。
69    ToolOutput { chunk: String },
70    /// 工具执行结束的**完整**结果(模型看到的、也是会落进 transcript 的那条消息)。
71    /// 后台工具任务超时让步时,这里就是 `{id, status:"running", …}` 那个异步信封。
72    ToolResponse {
73        call_id: String,
74        output: Vec<ContentBlock>,
75        is_error: bool,
76    },
77}
78
79/// 实时消息通道。`Clone` 即克隆底层 sender,可多生产者并发推送。
80#[derive(Clone)]
81pub struct MessageChannel {
82    tx: mpsc::UnboundedSender<AgentMessage>,
83}
84
85impl MessageChannel {
86    /// 创建通道,返回 `(channel, receiver)`。所有 `channel` 副本 drop 后,`receiver` 收到 `None`。
87    pub fn new() -> (Self, mpsc::UnboundedReceiver<AgentMessage>) {
88        let (tx, rx) = mpsc::unbounded_channel();
89        (Self { tx }, rx)
90    }
91
92    /// 一个没有消费者的通道:`send` 全部静默丢弃。未配置 messages 时的默认值。
93    pub fn disconnected() -> Self {
94        let (tx, _rx) = mpsc::unbounded_channel();
95        Self { tx }
96    }
97
98    /// 投递一条消息。fire-and-forget:消费者已断开则静默丢弃,不阻断主流程。
99    pub fn send(&self, session_id: &str, kind: MessageKind) {
100        let _ = self.tx.send(AgentMessage::new(session_id, kind));
101    }
102}
103
104impl Default for MessageChannel {
105    fn default() -> Self {
106        Self::disconnected()
107    }
108}