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 { 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 = if ctx.session_runtime.is_some() {
79                    None
80                } else {
81                    ctx.flow_run_id.as_ref().map(|r| r.0.to_string())
82                };
83                if let Some(tx) = &ctx.stream_tx {
84                    let _ = tx.send(crate::stream::StreamFrame::ToolResultMsg {
85                        flow_run_id,
86                        message: msg.clone(),
87                    });
88                }
89                handle.lock().unwrap().push(msg.clone());
90                // Cap ephemeral sub-agent segments; root persists via event sink.
91                if ctx.session_runtime.is_none() && handle.lock().unwrap().len() > 100 {
92                    let mut h = handle.lock().unwrap();
93                    let start = h.len() - 100;
94                    h.drain(..start);
95                }
96            }
97            Ok(Value::Unit)
98        })
99    }
100}
101
102fn emit_message_event(ctx: &ToolCtx, msg: &Message) {
103    use crate::event::{Event, TurnId};
104    let Some(sink) = &ctx.events else {
105        return;
106    };
107    let msg =
108        crate::tools::tool_output::maybe_truncate_tool_message(msg, ctx.session_dir.as_deref());
109    let turn_id = ctx.turn_id.clone().unwrap_or_else(TurnId::now);
110    let flow_run_id = if ctx.session_runtime.is_some() {
111        None
112    } else {
113        ctx.flow_run_id.clone()
114    };
115    let event = match msg.role {
116        MessageRole::User => Event::UserMsg {
117            turn_id,
118            flow_run_id,
119            message: msg.clone(),
120        },
121        MessageRole::Assistant => Event::AssistantMsg {
122            turn_id,
123            flow_run_id,
124            message: msg.clone(),
125        },
126        MessageRole::Tool => Event::ToolResultMsg {
127            turn_id,
128            flow_run_id,
129            message: msg.clone(),
130        },
131        MessageRole::System => Event::SystemMsg {
132            turn_id,
133            message: msg.clone(),
134        },
135    };
136    sink.emit(event);
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn session_push_name_and_tier() {
145        let tool = SessionPush;
146        assert_eq!(tool.name(), "session.push");
147        assert_eq!(tool.tier(), Tier::Zero);
148        assert!(tool.description().is_some());
149    }
150}