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                Value::Str(s) => {
57                    let turn_id = ctx
58                        .turn_id
59                        .clone()
60                        .unwrap_or_else(crate::event::TurnId::now);
61                    vec![Message::assistant_text(turn_id, s)]
62                }
63                other => {
64                    return Err(RuntimeError::TypeMismatch {
65                        expected: "message, list of message, or string".into(),
66                        actual: other.kind_name().into(),
67                    });
68                }
69            };
70            let Some(handle) = &ctx.session_messages_handle else {
71                return Err(RuntimeError::ToolFailed(
72                    "session.push: no session messages handle available".into(),
73                ));
74            };
75            let _compact_guard = match &ctx.compact_lock_handle {
76                Some(lock) => Some(lock.lock().await),
77                None => None,
78            };
79            for msg in msgs {
80                let msg = crate::tools::tool_output::maybe_truncate_tool_message(
81                    &msg,
82                    ctx.session_dir.as_deref(),
83                );
84                emit_message_event(ctx, &msg);
85                if let Some(tx) = &ctx.stream_tx {
86                    let _ = tx.send(crate::stream::StreamFrame::ToolResultMsg {
87                        flow_run_id: ctx.flow_run_id.as_ref().map(|r| r.0.to_string()),
88                        message: msg.clone(),
89                    });
90                }
91                handle.lock().unwrap().push(msg);
92            }
93            Ok(Value::Unit)
94        })
95    }
96}
97
98fn emit_message_event(ctx: &ToolCtx, msg: &Message) {
99    use crate::event::{Event, TurnId};
100    let Some(sink) = &ctx.events else {
101        return;
102    };
103    let msg =
104        crate::tools::tool_output::maybe_truncate_tool_message(msg, ctx.session_dir.as_deref());
105    let turn_id = ctx.turn_id.clone().unwrap_or_else(TurnId::now);
106    let event = match msg.role {
107        MessageRole::User => Event::UserMsg {
108            turn_id,
109            message: msg.clone(),
110        },
111        MessageRole::Assistant => Event::AssistantMsg {
112            turn_id,
113            flow_run_id: ctx.flow_run_id.clone(),
114            message: msg.clone(),
115        },
116        MessageRole::Tool => Event::ToolResultMsg {
117            turn_id,
118            flow_run_id: ctx.flow_run_id.clone(),
119            message: msg.clone(),
120        },
121        MessageRole::System => Event::SystemMsg {
122            turn_id,
123            message: msg.clone(),
124        },
125    };
126    sink.emit(event);
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132
133    #[test]
134    fn session_push_name_and_tier() {
135        let tool = SessionPush;
136        assert_eq!(tool.name(), "session.push");
137        assert_eq!(tool.tier(), Tier::Zero);
138        assert!(tool.description().is_some());
139    }
140}