agent_base/engine/
session.rs

1use std::collections::HashSet;
2
3use serde::{Deserialize, Serialize};
4
5use crate::types::{ChatMessage, ImageAttachment, Message, MessageRole, ToolCallMessage};
6
7use crate::types::SessionId;
8
9#[derive(Clone, Debug, Default, Serialize, Deserialize)]
10pub struct AgentSession {
11    id: Option<SessionId>,
12    messages: Vec<Message>,
13    chat_messages: Vec<ChatMessage>,
14    always_allowed_actions: HashSet<String>,
15}
16
17impl AgentSession {
18    pub fn new(id: SessionId) -> Self {
19        Self {
20            id: Some(id),
21            messages: Vec::new(),
22            chat_messages: Vec::new(),
23            always_allowed_actions: HashSet::new(),
24        }
25    }
26
27    pub fn id(&self) -> Option<SessionId> {
28        self.id.clone()
29    }
30
31    pub fn messages(&self) -> &[Message] {
32        &self.messages
33    }
34
35    pub fn chat_messages(&self) -> &[ChatMessage] {
36        &self.chat_messages
37    }
38
39    pub fn is_action_allowed(&self, action_key: &str) -> bool {
40        self.always_allowed_actions.contains(action_key)
41    }
42
43    pub fn allow_action(&mut self, action_key: impl Into<String>) {
44        self.always_allowed_actions.insert(action_key.into());
45    }
46
47    pub fn push_message(&mut self, role: MessageRole, content: impl Into<String>) {
48        let content = content.into();
49        self.messages.push(Message {
50            role: role.clone(),
51            content: content.clone(),
52        });
53        let chat_msg = match role {
54            MessageRole::System => ChatMessage::system(content),
55            MessageRole::User => ChatMessage::user(content),
56            MessageRole::Assistant => ChatMessage::assistant(content),
57            MessageRole::Tool => ChatMessage::tool(String::new(), content),
58        };
59        self.chat_messages.push(chat_msg);
60    }
61
62    pub fn push_user_message_with_images(
63        &mut self,
64        content: impl Into<String>,
65        images: Vec<ImageAttachment>,
66    ) {
67        let content = content.into();
68        self.messages.push(Message {
69            role: MessageRole::User,
70            content: content.clone(),
71        });
72        self.chat_messages
73            .push(ChatMessage::user_with_images(content, images));
74    }
75
76    pub fn push_assistant_tool_call(
77        &mut self,
78        tool_call_id: &str,
79        tool_name: &str,
80        arguments_json: &str,
81    ) {
82        self.chat_messages.push(ChatMessage::assistant_tool_call(
83            tool_call_id,
84            tool_name,
85            arguments_json,
86        ));
87    }
88
89    pub fn push_assistant_tool_calls(&mut self, tool_calls: &[(String, String, String)]) {
90        let calls: Vec<ToolCallMessage> = tool_calls
91            .iter()
92            .map(|(id, name, args)| ToolCallMessage {
93                id: id.clone(),
94                name: name.clone(),
95                arguments: args.clone(),
96            })
97            .collect();
98        self.chat_messages.push(ChatMessage::Assistant {
99            content: None,
100            reasoning_content: None,
101            tool_calls: Some(calls),
102        });
103    }
104
105    pub fn push_tool_result(&mut self, tool_call_id: &str, content: impl Into<String>) {
106        let content = content.into();
107        self.messages.push(Message {
108            role: MessageRole::Tool,
109            content: content.clone(),
110        });
111        self.chat_messages.push(ChatMessage::tool(
112            tool_call_id,
113            content,
114        ));
115    }
116}
agent_base/engine/session.rs

agent_base/engine/
session.rs
