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                emit_message_event(ctx, &msg);
81                if let Some(tx) = &ctx.stream_tx {
82                    let _ = tx.send(crate::stream::StreamFrame::ToolResultMsg {
83                        flow_run_id: ctx.flow_run_id.as_ref().map(|r| r.0.to_string()),
84                        message: msg.clone(),
85                    });
86                }
87                handle.lock().unwrap().push(msg);
88            }
89            Ok(Value::Unit)
90        })
91    }
92}
93
94fn emit_message_event(ctx: &ToolCtx, msg: &Message) {
95    use crate::event::{Event, TurnId};
96    let Some(sink) = &ctx.events else {
97        return;
98    };
99    let turn_id = ctx.turn_id.clone().unwrap_or_else(TurnId::now);
100    let event = match msg.role {
101        MessageRole::User => Event::UserMsg {
102            turn_id,
103            message: msg.clone(),
104        },
105        MessageRole::Assistant => Event::AssistantMsg {
106            turn_id,
107            flow_run_id: ctx.flow_run_id.clone(),
108            message: msg.clone(),
109        },
110        MessageRole::Tool => Event::ToolResultMsg {
111            turn_id,
112            flow_run_id: ctx.flow_run_id.clone(),
113            message: msg.clone(),
114        },
115        MessageRole::System => Event::SystemMsg {
116            turn_id,
117            message: msg.clone(),
118        },
119    };
120    sink.emit(event);
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126
127    #[test]
128    fn session_push_name_and_tier() {
129        let tool = SessionPush;
130        assert_eq!(tool.name(), "session.push");
131        assert_eq!(tool.tier(), Tier::Zero);
132        assert!(tool.description().is_some());
133    }
134}