Skip to main content

atman_runtime/tools/
session.rs

1use crate::error::RuntimeError;
2use crate::message::{Message, MessageRole};
3use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
4use crate::value::Value;
5use chrono::Utc;
6
7pub struct SessionPush;
8
9impl Tool for SessionPush {
10    fn name(&self) -> &str {
11        "session.push"
12    }
13
14    fn tier(&self) -> Tier {
15        Tier::Zero
16    }
17
18    fn description(&self) -> Option<&str> {
19        Some(
20            "Push a Message value into the current session's message history. \
21             Use after dispatch_all to persist tool results so the next \
22             llm { context: session } call can see them. The message role \
23             (user/assistant/tool/system) is preserved. Returns unit.",
24        )
25    }
26
27    fn input_schema(&self) -> serde_json::Value {
28        serde_json::json!({
29            "type": "object",
30            "properties": {
31                "message": {
32                    "type": "object",
33                    "description": "The Message value to push (e.g. a tool_result from dispatch_all). Pass the value returned by dispatch_all directly."
34                }
35            },
36            "required": ["message"]
37        })
38    }
39
40    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
41        Box::pin(async move {
42            let val = match args.named("message").or_else(|| args.positional(0).ok()) {
43                Some(v) => v.clone(),
44                None => {
45                    return Err(RuntimeError::MissingArg("session.push: message".into()));
46                }
47            };
48            let msgs = match val {
49                Value::Message(m) => vec![m],
50                Value::List(items) => items
51                    .into_iter()
52                    .filter_map(|v| match v {
53                        Value::Message(m) => Some(m),
54                        _ => None,
55                    })
56                    .collect(),
57                Value::Str(s) => {
58                    let turn_id = ctx
59                        .turn_id
60                        .clone()
61                        .unwrap_or_else(crate::event::TurnId::now);
62                    vec![Message::assistant_text(turn_id, s)]
63                }
64                other => {
65                    return Err(RuntimeError::TypeMismatch {
66                        expected: "message, list of message, or string".into(),
67                        actual: other.kind_name().into(),
68                    });
69                }
70            };
71            let Some(handle) = &ctx.session_messages_handle else {
72                return Err(RuntimeError::ToolFailed(
73                    "session.push: no session messages handle available".into(),
74                ));
75            };
76            let _compact_guard = match &ctx.compact_lock_handle {
77                Some(lock) => Some(lock.lock().await),
78                None => None,
79            };
80            for msg in msgs {
81                emit_message_event(ctx, &msg);
82                if let Some(tx) = &ctx.stream_tx {
83                    let _ = tx.send(crate::stream::StreamFrame::ToolResultMsg {
84                        flow_run_id: ctx.flow_run_id.as_ref().map(|r| r.0.to_string()),
85                        message: msg.clone(),
86                    });
87                }
88                handle.lock().unwrap().push(msg);
89            }
90            Ok(Value::Unit)
91        })
92    }
93}
94
95fn emit_message_event(ctx: &ToolCtx, msg: &Message) {
96    use crate::event::{Event, TurnId};
97    let Some(sink) = &ctx.events else {
98        return;
99    };
100    let ts = Utc::now();
101    let turn_id = ctx.turn_id.clone().unwrap_or_else(TurnId::now);
102    let event = match msg.role {
103        MessageRole::User => Event::UserMsg {
104            seq: 0,
105            turn_id,
106            message: msg.clone(),
107            ts,
108        },
109        MessageRole::Assistant => Event::AssistantMsg {
110            seq: 0,
111            turn_id,
112            flow_run_id: ctx.flow_run_id.clone(),
113            message: msg.clone(),
114            ts,
115        },
116        MessageRole::Tool => Event::ToolResultMsg {
117            seq: 0,
118            turn_id,
119            flow_run_id: ctx.flow_run_id.clone(),
120            message: msg.clone(),
121            ts,
122        },
123        MessageRole::System => Event::SystemMsg {
124            seq: 0,
125            turn_id,
126            message: msg.clone(),
127            ts,
128        },
129    };
130    sink.emit(event);
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136
137    #[test]
138    fn session_push_name_and_tier() {
139        let tool = SessionPush;
140        assert_eq!(tool.name(), "session.push");
141        assert_eq!(tool.tier(), Tier::Zero);
142        assert!(tool.description().is_some());
143    }
144}