Skip to main content

codei_session/
model.rs

1use std::path::PathBuf;
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5use uuid::Uuid;
6
7use codei_llm::{Message, Role as LlmRole, ToolCall};
8
9pub type SessionId = String;
10pub type MessageId = String;
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct Session {
14    pub id: SessionId,
15    pub title: Option<String>,
16    pub cwd: PathBuf,
17    pub created_at: DateTime<Utc>,
18    pub updated_at: DateTime<Utc>,
19    pub messages: Vec<StoredMessage>,
20}
21
22#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
23#[serde(rename_all = "lowercase")]
24pub enum Role {
25    System,
26    User,
27    Assistant,
28    Tool,
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct StoredMessage {
33    pub id: MessageId,
34    pub role: Role,
35    pub content: MessageContent,
36    pub tool_calls: Option<Vec<ToolCallRecord>>,
37    pub tool_call_id: Option<String>,
38    pub created_at: DateTime<Utc>,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub enum MessageContent {
43    Text(String),
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct ToolCallRecord {
48    pub id: String,
49    pub name: String,
50    pub arguments: String,
51}
52
53impl Session {
54    pub fn new(cwd: PathBuf) -> Self {
55        let now = Utc::now();
56        Self {
57            id: Uuid::new_v4().to_string(),
58            title: None,
59            cwd,
60            created_at: now,
61            updated_at: now,
62            messages: Vec::new(),
63        }
64    }
65
66    pub fn touch(&mut self) {
67        self.updated_at = Utc::now();
68    }
69
70    pub fn push_user(&mut self, content: impl Into<String>) -> &StoredMessage {
71        self.push_message(Role::User, MessageContent::Text(content.into()), None, None)
72    }
73
74    pub fn push_assistant(
75        &mut self,
76        content: String,
77        tool_calls: Option<Vec<ToolCallRecord>>,
78    ) -> &StoredMessage {
79        self.push_message(
80            Role::Assistant,
81            MessageContent::Text(content),
82            tool_calls,
83            None,
84        )
85    }
86
87    pub fn push_tool(
88        &mut self,
89        tool_call_id: impl Into<String>,
90        content: impl Into<String>,
91    ) -> &StoredMessage {
92        self.push_message(
93            Role::Tool,
94            MessageContent::Text(content.into()),
95            None,
96            Some(tool_call_id.into()),
97        )
98    }
99
100    pub fn clear_messages(&mut self) {
101        self.messages.clear();
102        self.touch();
103    }
104
105    /// Drop older messages, keeping the most recent `keep_recent` entries.
106    pub fn compact(&mut self, keep_recent: usize) {
107        if self.messages.len() <= keep_recent {
108            return;
109        }
110        let remove = self.messages.len() - keep_recent;
111        self.messages.drain(0..remove);
112        self.touch();
113    }
114
115    /// Replace older messages with a summary, keeping the most recent `keep_recent` entries.
116    pub fn compact_with_summary(&mut self, keep_recent: usize, summary: String) {
117        if self.messages.len() <= keep_recent {
118            return;
119        }
120        let remove = self.messages.len() - keep_recent;
121        self.messages.drain(0..remove);
122        let now = Utc::now();
123        self.messages.insert(
124            0,
125            StoredMessage {
126                id: Uuid::new_v4().to_string(),
127                role: Role::User,
128                content: MessageContent::Text(format!(
129                    "[Context summary — earlier conversation compressed]\n\n{summary}"
130                )),
131                tool_calls: None,
132                tool_call_id: None,
133                created_at: now,
134            },
135        );
136        self.touch();
137    }
138
139    fn push_message(
140        &mut self,
141        role: Role,
142        content: MessageContent,
143        tool_calls: Option<Vec<ToolCallRecord>>,
144        tool_call_id: Option<String>,
145    ) -> &StoredMessage {
146        let msg = StoredMessage {
147            id: Uuid::new_v4().to_string(),
148            role,
149            content,
150            tool_calls,
151            tool_call_id,
152            created_at: Utc::now(),
153        };
154        self.messages.push(msg);
155        self.touch();
156        self.messages.last().expect("message just pushed")
157    }
158}
159
160impl From<ToolCall> for ToolCallRecord {
161    fn from(value: ToolCall) -> Self {
162        Self {
163            id: value.id,
164            name: value.name,
165            arguments: value.arguments,
166        }
167    }
168}
169
170impl StoredMessage {
171    pub fn text(&self) -> Option<&str> {
172        match &self.content {
173            MessageContent::Text(s) => Some(s),
174        }
175    }
176}
177
178pub fn to_llm_messages(session: &Session, system_prompt: &str) -> Vec<Message> {
179    let mut messages = vec![Message::system(system_prompt)];
180    for msg in &session.messages {
181        match msg.role {
182            Role::User => {
183                if let Some(text) = msg.text() {
184                    messages.push(Message::user(text));
185                }
186            }
187            Role::Assistant => {
188                let tool_calls: Option<Vec<ToolCall>> = msg.tool_calls.as_ref().map(|calls| {
189                    calls
190                        .iter()
191                        .map(|c| ToolCall {
192                            id: c.id.clone(),
193                            name: c.name.clone(),
194                            arguments: c.arguments.clone(),
195                        })
196                        .collect()
197                });
198                let text = msg.text().unwrap_or("").to_string();
199                if let Some(calls) = &tool_calls {
200                    if !calls.is_empty() {
201                        messages.push(Message {
202                            role: LlmRole::Assistant,
203                            content: if text.is_empty() { None } else { Some(text) },
204                            tool_calls,
205                            tool_call_id: None,
206                        });
207                    } else {
208                        messages.push(Message::assistant(text, None));
209                    }
210                } else {
211                    messages.push(Message::assistant(text, None));
212                }
213            }
214            Role::Tool => {
215                if let (Some(id), Some(text)) = (&msg.tool_call_id, msg.text()) {
216                    messages.push(Message::tool(id.clone(), text));
217                }
218            }
219            Role::System => {}
220        }
221    }
222    messages
223}