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;
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_with_budget(
74                    &msg,
75                    ctx.output_store.as_deref(),
76                    ctx.tool_output_budget,
77                );
78                emit_message_event(ctx, &msg);
79                let flow_run_id = match msg.role {
80                    MessageRole::Assistant | MessageRole::Tool => {
81                        ctx.flow_run_id.as_ref().map(|r| r.0.to_string())
82                    }
83                    MessageRole::User | MessageRole::System => None,
84                };
85                if let Some(tx) = &ctx.stream_tx {
86                    let _ = tx.send(crate::stream::StreamFrame::ToolResultMsg {
87                        flow_run_id,
88                        message: msg.clone(),
89                    });
90                }
91                handle.lock().unwrap().push(msg.clone());
92                // Cap ephemeral sub-agent segments; root persists via event sink.
93                if ctx.session_runtime.is_none() && handle.lock().unwrap().len() > 100 {
94                    let mut h = handle.lock().unwrap();
95                    let start = h.len() - 100;
96                    h.drain(..start);
97                }
98            }
99            Ok(Value::Unit)
100        })
101    }
102}
103
104fn emit_message_event(ctx: &ToolCtx, msg: &Message) {
105    use crate::event::{Event, TurnId};
106    let Some(sink) = &ctx.events else {
107        return;
108    };
109    let turn_id = ctx.turn_id.clone().unwrap_or_else(TurnId::now);
110    let flow_run_id = ctx.message_flow_run_id();
111    let event = match msg.role {
112        MessageRole::User => Event::UserMsg {
113            turn_id,
114            flow_run_id,
115            message: msg.clone(),
116        },
117        MessageRole::Assistant => Event::AssistantMsg {
118            turn_id,
119            flow_run_id,
120            message: msg.clone(),
121        },
122        MessageRole::Tool => Event::ToolResultMsg {
123            turn_id,
124            flow_run_id,
125            message: msg.clone(),
126        },
127        MessageRole::System => Event::SystemMsg {
128            turn_id,
129            message: msg.clone(),
130        },
131    };
132    sink.emit(event);
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    #[test]
140    fn session_push_name_and_tier() {
141        let tool = SessionPush;
142        assert_eq!(tool.name(), "session.push");
143        assert_eq!(tool.tier(), Tier::Zero);
144        assert!(tool.description().is_some());
145    }
146
147    #[tokio::test]
148    async fn session_push_scopes_live_tool_result_without_scoping_session_history() {
149        use crate::event::{Event, FlowRunId, TurnId};
150        use crate::message::{MessageOrigin, MessagePart};
151
152        let session = std::sync::Arc::new(crate::session::Session::open_ephemeral());
153        let run_id = FlowRunId::now();
154        let (stream_tx, mut stream_rx) = tokio::sync::broadcast::channel(8);
155        let ctx = ToolCtx::new()
156            .with_anchors(Some(TurnId::now()), Some(run_id.clone()), None)
157            .with_events(session.sink().clone())
158            .with_session_messages_handle(session.messages_handle())
159            .with_session_runtime(session.clone())
160            .with_stream_tx(stream_tx);
161        let message = Message {
162            role: MessageRole::Tool,
163            parts: vec![MessagePart::ToolResult {
164                tool_use_id: "tu_1".into(),
165                content: "done".into(),
166                is_error: false,
167            }],
168            turn_id: TurnId::now(),
169            origin: MessageOrigin::User,
170        };
171
172        SessionPush
173            .call(
174                ToolArgs {
175                    positional: vec![Value::Message(message)],
176                    named: Vec::new(),
177                },
178                &ctx,
179            )
180            .await
181            .unwrap();
182
183        let crate::stream::StreamFrame::ToolResultMsg { flow_run_id, .. } =
184            stream_rx.recv().await.unwrap()
185        else {
186            panic!("tool result frame");
187        };
188        assert_eq!(flow_run_id.as_deref(), Some(run_id.0.to_string().as_str()));
189        assert!(session.sink().snapshot().iter().any(|event| {
190            matches!(
191                event,
192                Event::ToolResultMsg {
193                    flow_run_id: None,
194                    ..
195                }
196            )
197        }));
198        assert!(
199            session
200                .messages()
201                .iter()
202                .any(|message| message.role == MessageRole::Tool)
203        );
204    }
205}
atman_runtime/tools/session.rs

atman_runtime/tools/
session.rs
