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;
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                append_message_to_context(ctx, msg)?;
79            }
80            Ok(Value::Unit)
81        })
82    }
83}
84
85pub(crate) fn append_message_to_context(ctx: &ToolCtx, msg: Message) -> Result<(), RuntimeError> {
86    if matches!(ctx.history_segment, crate::tool::HistorySegment::Root)
87        && let Some(session) = ctx.session_runtime.as_ref()
88    {
89        let stream_flow_run_id = match msg.role {
90            MessageRole::Assistant | MessageRole::Tool => ctx.flow_run_id.clone(),
91            MessageRole::User | MessageRole::System => None,
92        };
93        session.append_message_with_stream_scope(
94            msg,
95            ctx.message_flow_run_id(),
96            stream_flow_run_id,
97        );
98        return Ok(());
99    }
100    let Some(handle) = &ctx.session_messages_handle else {
101        return Err(RuntimeError::ToolFailed(
102            "session message context is unavailable".into(),
103        ));
104    };
105    emit_message_event(ctx, &msg);
106    let flow_run_id = match msg.role {
107        MessageRole::Assistant | MessageRole::Tool => {
108            ctx.flow_run_id.as_ref().map(|run_id| run_id.0.to_string())
109        }
110        MessageRole::User => ctx.message_flow_run_id().map(|run_id| run_id.0.to_string()),
111        MessageRole::System => None,
112    };
113    if msg.origin != crate::message::MessageOrigin::Internal
114        && msg.role != MessageRole::System
115        && let Some(tx) = &ctx.stream_tx
116    {
117        let frame = match msg.role {
118            MessageRole::Assistant => crate::stream::StreamFrame::AssistantMsg {
119                flow_run_id,
120                message: msg.clone(),
121            },
122            MessageRole::User | MessageRole::Tool => crate::stream::StreamFrame::ToolResultMsg {
123                flow_run_id,
124                message: msg.clone(),
125            },
126            MessageRole::System => unreachable!(),
127        };
128        let _ = tx.send(frame);
129    }
130    let mut messages = handle.lock().unwrap();
131    messages.push(msg);
132    Ok(())
133}
134
135fn emit_message_event(ctx: &ToolCtx, msg: &Message) {
136    use crate::event::{Event, TurnId};
137    let Some(sink) = &ctx.events else {
138        return;
139    };
140    let turn_id = ctx.turn_id.clone().unwrap_or_else(TurnId::now);
141    let flow_run_id = ctx.message_flow_run_id();
142    let event = match msg.role {
143        MessageRole::User => Event::UserMsg {
144            turn_id,
145            flow_run_id,
146            message: msg.clone(),
147        },
148        MessageRole::Assistant => Event::AssistantMsg {
149            turn_id,
150            flow_run_id,
151            message: msg.clone(),
152        },
153        MessageRole::Tool => Event::ToolResultMsg {
154            turn_id,
155            flow_run_id,
156            message: msg.clone(),
157        },
158        MessageRole::System => Event::SystemMsg {
159            turn_id,
160            flow_run_id: ctx.message_flow_run_id(),
161            message: msg.clone(),
162        },
163    };
164    sink.emit(event);
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170
171    #[test]
172    fn spawned_context_does_not_rewrite_messages_outside_compaction() {
173        let messages = std::sync::Arc::new(std::sync::Mutex::new(
174            (0..100)
175                .map(|index| {
176                    Message::assistant_text(crate::event::TurnId::now(), format!("old-{index}"))
177                })
178                .collect(),
179        ));
180        let mut ctx = ToolCtx::new()
181            .with_history_segment(crate::tool::HistorySegment::Spawned)
182            .with_session_messages_handle(std::sync::Arc::clone(&messages));
183        ctx.context_epoch_handle = Some(std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)));
184
185        append_message_to_context(
186            &ctx,
187            Message::assistant_text(crate::event::TurnId::now(), "new"),
188        )
189        .unwrap();
190
191        assert_eq!(messages.lock().unwrap().len(), 101);
192        assert_eq!(messages.lock().unwrap()[0].text_concat(), "old-0");
193        assert_eq!(ctx.context_epoch_seed().as_deref(), Some("generation:0"));
194    }
195
196    #[test]
197    fn session_push_name_and_tier() {
198        let tool = SessionPush;
199        assert_eq!(tool.name(), "session.push");
200        assert_eq!(tool.tier(), Tier::Zero);
201        assert!(tool.description().is_some());
202    }
203
204    #[tokio::test]
205    async fn session_push_scopes_live_tool_result_without_scoping_session_history() {
206        use crate::event::{Event, FlowRunId, TurnId};
207        use crate::message::{MessageOrigin, MessagePart};
208
209        let session = std::sync::Arc::new(crate::session::Session::open_ephemeral());
210        let run_id = FlowRunId::now();
211        let mut stream_rx = session.stream_subscribe();
212        let ctx = ToolCtx::new()
213            .with_anchors(Some(TurnId::now()), Some(run_id.clone()), None)
214            .with_events(session.sink().clone())
215            .with_session_messages_handle(session.messages_handle())
216            .with_session_runtime(session.clone());
217        let message = Message {
218            role: MessageRole::Tool,
219            parts: vec![MessagePart::ToolResult {
220                tool_use_id: "tu_1".into(),
221                content: "done".into(),
222                is_error: false,
223            }],
224            turn_id: TurnId::now(),
225            origin: MessageOrigin::User,
226        };
227
228        SessionPush
229            .call(
230                ToolArgs {
231                    positional: vec![Value::Message(message)],
232                    named: Vec::new(),
233                },
234                &ctx,
235            )
236            .await
237            .unwrap();
238
239        let crate::stream::StreamFrame::ToolResultMsg { flow_run_id, .. } =
240            stream_rx.recv().await.unwrap()
241        else {
242            panic!("tool result frame");
243        };
244        assert_eq!(flow_run_id.as_deref(), Some(run_id.0.to_string().as_str()));
245        assert!(session.sink().snapshot().iter().any(|event| {
246            matches!(
247                event,
248                Event::ToolResultMsg {
249                    flow_run_id: None,
250                    ..
251                }
252            )
253        }));
254        assert!(
255            session
256                .messages()
257                .iter()
258                .any(|message| message.role == MessageRole::Tool)
259        );
260    }
261
262    #[test]
263    fn internal_child_system_message_is_audited_without_becoming_a_live_transcript_frame() {
264        use crate::context_plan::{
265            ContextRecord, ContextRecordAuthority, ContextRecordBody, ContextRecordRetention,
266        };
267        use crate::event::{Event, FlowRunId, TurnId};
268
269        let session = std::sync::Arc::new(crate::session::Session::open_ephemeral());
270        let run_id = FlowRunId::now();
271        let (stream_tx, mut stream_rx) = tokio::sync::broadcast::channel(8);
272        let messages = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
273        let ctx = ToolCtx::new()
274            .with_anchors(Some(TurnId::now()), Some(run_id.clone()), None)
275            .with_history_segment(crate::tool::HistorySegment::Spawned)
276            .with_events(session.sink().clone())
277            .with_session_messages_handle(messages)
278            .with_stream_tx(stream_tx);
279        let message = Message::context_record(
280            TurnId::now(),
281            ContextRecord::new(
282                "handoff.parent",
283                1,
284                ContextRecordAuthority::Runtime,
285                ContextRecordRetention::Latest,
286                ContextRecordBody::text("delegated"),
287            ),
288        );
289
290        append_message_to_context(&ctx, message).unwrap();
291
292        assert!(matches!(
293            stream_rx.try_recv(),
294            Err(tokio::sync::broadcast::error::TryRecvError::Empty)
295        ));
296        assert!(session.sink().snapshot().iter().any(|event| {
297            matches!(
298                event,
299                Event::SystemMsg {
300                    flow_run_id: Some(owner),
301                    ..
302                } if owner == &run_id
303            )
304        }));
305    }
306}