atman_runtime/tools/
session.rs1use crate::error::RuntimeError;
2use crate::message::{Message, MessageRole};
3use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
4use crate::value::Value;
5
6pub struct SessionPush;
7
8impl Tool for SessionPush {
9 fn name(&self) -> &str {
10 "session.push"
11 }
12
13 fn tier(&self) -> Tier {
14 Tier::Zero
15 }
16
17 fn description(&self) -> Option<&str> {
18 Some(
19 "Push a Message value into the current session's message history. \
20 Use after dispatch_all to persist tool results so the next \
21 llm.call(context: \"session\") call can see them. The message role \
22 (user/assistant/tool/system) is preserved. Returns unit.",
23 )
24 }
25
26 fn input_schema(&self) -> serde_json::Value {
27 serde_json::json!({
28 "type": "object",
29 "properties": {
30 "message": {
31 "type": "object",
32 "description": "The Message value to push (e.g. a tool_result from dispatch_all). Pass the value returned by dispatch_all directly."
33 }
34 },
35 "required": ["message"]
36 })
37 }
38
39 fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
40 Box::pin(async move {
41 let val = match args.named("message").or_else(|| args.positional(0).ok()) {
42 Some(v) => v.clone(),
43 None => {
44 return Err(RuntimeError::MissingArg("session.push: message".into()));
45 }
46 };
47 let msgs = match val {
48 Value::Message(m) => vec![m],
49 Value::List(items) => items
50 .into_iter()
51 .filter_map(|v| match v {
52 Value::Message(m) => Some(m),
53 _ => None,
54 })
55 .collect(),
56 other => {
57 return Err(RuntimeError::TypeMismatch {
58 expected: "message or list of message".into(),
59 actual: other.kind_name().into(),
60 });
61 }
62 };
63 let Some(handle) = &ctx.session_messages_handle else {
64 return Err(RuntimeError::ToolFailed(
65 "session.push: no session messages handle available".into(),
66 ));
67 };
68 let _compact_guard = match &ctx.compact_lock_handle {
69 Some(lock) => Some(lock.lock().await),
70 None => None,
71 };
72 for msg in msgs {
73 let msg = crate::tools::tool_output::maybe_truncate_tool_message(
74 &msg,
75 ctx.session_dir.as_deref(),
76 );
77 emit_message_event(ctx, &msg);
78 let flow_run_id = match msg.role {
79 MessageRole::Assistant | MessageRole::Tool => {
80 ctx.flow_run_id.as_ref().map(|r| r.0.to_string())
81 }
82 MessageRole::User | MessageRole::System => None,
83 };
84 if let Some(tx) = &ctx.stream_tx {
85 let _ = tx.send(crate::stream::StreamFrame::ToolResultMsg {
86 flow_run_id,
87 message: msg.clone(),
88 });
89 }
90 handle.lock().unwrap().push(msg.clone());
91 if ctx.session_runtime.is_none() && handle.lock().unwrap().len() > 100 {
93 let mut h = handle.lock().unwrap();
94 let start = h.len() - 100;
95 h.drain(..start);
96 }
97 }
98 Ok(Value::Unit)
99 })
100 }
101}
102
103fn emit_message_event(ctx: &ToolCtx, msg: &Message) {
104 use crate::event::{Event, TurnId};
105 let Some(sink) = &ctx.events else {
106 return;
107 };
108 let msg =
109 crate::tools::tool_output::maybe_truncate_tool_message(msg, ctx.session_dir.as_deref());
110 let turn_id = ctx.turn_id.clone().unwrap_or_else(TurnId::now);
111 let flow_run_id = if ctx.session_runtime.is_some() {
112 None
113 } else {
114 ctx.flow_run_id.clone()
115 };
116 let event = match msg.role {
117 MessageRole::User => Event::UserMsg {
118 turn_id,
119 flow_run_id,
120 message: msg.clone(),
121 },
122 MessageRole::Assistant => Event::AssistantMsg {
123 turn_id,
124 flow_run_id,
125 message: msg.clone(),
126 },
127 MessageRole::Tool => Event::ToolResultMsg {
128 turn_id,
129 flow_run_id,
130 message: msg.clone(),
131 },
132 MessageRole::System => Event::SystemMsg {
133 turn_id,
134 message: msg.clone(),
135 },
136 };
137 sink.emit(event);
138}
139
140#[cfg(test)]
141mod tests {
142 use super::*;
143
144 #[test]
145 fn session_push_name_and_tier() {
146 let tool = SessionPush;
147 assert_eq!(tool.name(), "session.push");
148 assert_eq!(tool.tier(), Tier::Zero);
149 assert!(tool.description().is_some());
150 }
151
152 #[tokio::test]
153 async fn session_push_scopes_live_tool_result_without_scoping_session_history() {
154 use crate::event::{Event, FlowRunId, TurnId};
155 use crate::message::{MessageOrigin, MessagePart};
156
157 let session = std::sync::Arc::new(crate::session::Session::open_ephemeral());
158 let run_id = FlowRunId::now();
159 let (stream_tx, mut stream_rx) = tokio::sync::broadcast::channel(8);
160 let ctx = ToolCtx::new()
161 .with_anchors(Some(TurnId::now()), Some(run_id.clone()), None)
162 .with_events(session.sink().clone())
163 .with_session_messages_handle(session.messages_handle())
164 .with_session_runtime(session.clone())
165 .with_stream_tx(stream_tx);
166 let message = Message {
167 role: MessageRole::Tool,
168 parts: vec![MessagePart::ToolResult {
169 tool_use_id: "tu_1".into(),
170 content: "done".into(),
171 is_error: false,
172 }],
173 turn_id: TurnId::now(),
174 origin: MessageOrigin::User,
175 };
176
177 SessionPush
178 .call(
179 ToolArgs {
180 positional: vec![Value::Message(message)],
181 named: Vec::new(),
182 },
183 &ctx,
184 )
185 .await
186 .unwrap();
187
188 let crate::stream::StreamFrame::ToolResultMsg { flow_run_id, .. } =
189 stream_rx.recv().await.unwrap()
190 else {
191 panic!("tool result frame");
192 };
193 assert_eq!(flow_run_id.as_deref(), Some(run_id.0.to_string().as_str()));
194 assert!(session.sink().snapshot().iter().any(|event| {
195 matches!(
196 event,
197 Event::ToolResultMsg {
198 flow_run_id: None,
199 ..
200 }
201 )
202 }));
203 assert!(
204 session
205 .messages()
206 .iter()
207 .any(|message| message.role == MessageRole::Tool)
208 );
209 }
210}