1use crate::diagnostic::{
2 ProviderTerminalDiagnostic, decode_terminal_diagnostic, encode_terminal_diagnostic,
3 is_terminal_diagnostic_content,
4};
5use crate::turn_outcome::{
6 TurnTranscriptOutcomeRecord, decode_turn_transcript_outcome, encode_turn_transcript_outcome,
7 is_turn_transcript_outcome_content,
8};
9use crate::{Session, SessionEntry, SessionError, SessionMetadata};
10use chrono::Utc;
11use std::collections::HashSet;
12use talos_core::message::{AgentEvent, Message};
13use uuid::Uuid;
14
15impl Session {
16 pub fn append(&self, message: &Message) -> Result<(), SessionError> {
17 self.append_with_metadata(message, SessionMetadata::default())
18 }
19
20 pub fn append_with_metadata(
21 &self,
22 message: &Message,
23 mut metadata: SessionMetadata,
24 ) -> Result<(), SessionError> {
25 let (role, content) = message_parts(message);
26 if let Message::Assistant {
27 reasoning: Some(r), ..
28 } = message
29 {
30 metadata.reasoning = Some(r.clone());
31 }
32 let entry = self.build_entry(&role, &content, metadata)?;
33 self.append_entry_locked(&entry)
34 }
35
36 pub fn append_terminal_diagnostic(
37 &self,
38 diagnostic: &ProviderTerminalDiagnostic,
39 ) -> Result<(), SessionError> {
40 let content = encode_terminal_diagnostic(diagnostic)
41 .map_err(|error| SessionError::InvalidJson(error.to_string()))?;
42 let entry = self.build_entry(
43 "system",
44 &content,
45 SessionMetadata {
46 provider: diagnostic.provider.clone(),
47 model: diagnostic.model.clone(),
48 ..SessionMetadata::default()
49 },
50 )?;
51 self.append_entry_locked(&entry)
52 }
53
54 pub fn read_terminal_diagnostics(
55 &self,
56 ) -> Result<Vec<ProviderTerminalDiagnostic>, SessionError> {
57 Ok(self
58 .read_entries()?
59 .iter()
60 .filter_map(|entry| decode_terminal_diagnostic(&entry.content))
61 .collect())
62 }
63
64 pub fn append_turn_transcript_outcome(
70 &self,
71 outcome: &TurnTranscriptOutcomeRecord,
72 ) -> Result<(), SessionError> {
73 let content = encode_turn_transcript_outcome(outcome)
74 .map_err(|error| SessionError::InvalidJson(error.to_string()))?;
75 let entry = self.build_entry(
76 "system",
77 &content,
78 SessionMetadata {
79 turn_id: (outcome.outcome == crate::TurnTranscriptOutcome::Success)
80 .then(|| outcome.turn_id.clone()),
81 ..SessionMetadata::default()
82 },
83 )?;
84 self.append_entry_locked(&entry)
85 }
86
87 pub fn read_turn_transcript_outcomes(
88 &self,
89 ) -> Result<Vec<TurnTranscriptOutcomeRecord>, SessionError> {
90 Ok(self
91 .read_entries()?
92 .iter()
93 .filter_map(|entry| decode_turn_transcript_outcome(&entry.content))
94 .collect())
95 }
96
97 pub fn append_event(&self, event: &AgentEvent) -> Result<(), SessionError> {
98 if matches!(
99 event,
100 AgentEvent::ThinkingDelta { .. } | AgentEvent::ReasoningComplete { .. }
101 ) {
102 return Ok(());
103 }
104 let content =
105 serde_json::to_string(event).map_err(|e| SessionError::InvalidJson(e.to_string()))?;
106 let entry = self.build_entry("system", &content, SessionMetadata::default())?;
107 self.append_entry_locked(&entry)
108 }
109
110 fn build_entry(
111 &self,
112 role: &str,
113 content: &str,
114 metadata: SessionMetadata,
115 ) -> Result<SessionEntry, SessionError> {
116 let parent_id = {
117 let guard = self
118 .last_entry_id
119 .lock()
120 .expect("last_entry_id mutex poisoned");
121 if guard.is_none() {
122 drop(guard);
123 let id = self.store.read_last_entry_id(&self.file_path);
124 *self
125 .last_entry_id
126 .lock()
127 .expect("last_entry_id mutex poisoned") = id.clone();
128 id
129 } else {
130 guard.clone()
131 }
132 };
133
134 Ok(SessionEntry {
135 id: Uuid::new_v4().to_string(),
136 parent_id,
137 timestamp: Utc::now(),
138 role: role.to_string(),
139 content: content.to_string(),
140 metadata,
141 })
142 }
143
144 fn append_entry_locked(&self, entry: &SessionEntry) -> Result<(), SessionError> {
145 let _lock = self.write_lock.lock().expect("write_lock mutex poisoned");
146 self.store.append_entry(&self.file_path, entry)
147 }
148
149 pub fn read_entries(&self) -> Result<Vec<SessionEntry>, SessionError> {
150 self.store.read_entries(&self.file_path)
151 }
152
153 pub fn read_messages(&self) -> Result<Vec<Message>, SessionError> {
154 let entries = self.read_entries()?;
155 let mut messages = Vec::new();
156 let mut pending_tool_call_ids = HashSet::new();
157
158 for entry in entries {
159 let msg = match entry.role.as_str() {
160 "user" => {
161 pending_tool_call_ids.clear();
162 Some(Message::User {
163 content: entry.content,
164 })
165 }
166 "assistant" => {
167 let tool_calls =
168 talos_core::message::extract_tool_calls_from_text(&entry.content);
169 let cleaned = talos_core::message::strip_tool_syntax(&entry.content);
170 pending_tool_call_ids.clear();
171 pending_tool_call_ids.extend(tool_calls.iter().map(|call| call.id.clone()));
172 Message::Assistant {
173 content: cleaned,
174 tool_calls,
175 reasoning: entry.metadata.reasoning,
176 }
177 .into()
178 }
179 "system" => {
180 if is_terminal_diagnostic_content(&entry.content)
181 || is_turn_transcript_outcome_content(&entry.content)
182 {
183 None
184 } else if let Some(sys_content) = entry.content.strip_prefix("__SYSTEM__:") {
185 Some(Message::System {
186 content: sys_content.to_string(),
187 cache_markers: Vec::new(),
188 })
189 } else if serde_json::from_str::<AgentEvent>(&entry.content).is_ok() {
190 None
191 } else {
192 let (is_error, tool_use_id, content) = parse_tool_result(&entry.content);
193 if pending_tool_call_ids.remove(&tool_use_id) {
194 Some(Message::Tool {
195 result: talos_core::message::MessageToolResult {
196 tool_use_id,
197 content,
198 is_error,
199 },
200 })
201 } else {
202 None
203 }
204 }
205 }
206 _ => None,
207 };
208
209 if let Some(msg) = msg {
210 messages.push(msg);
211 }
212 }
213
214 Ok(messages)
215 }
216
217 pub fn read_events(&self) -> Result<Vec<AgentEvent>, SessionError> {
218 let entries = self.read_entries()?;
219 let mut events = Vec::new();
220
221 for entry in entries {
222 if entry.role == "system"
223 && let Ok(event) = serde_json::from_str::<AgentEvent>(&entry.content)
224 {
225 events.push(event);
226 }
227 }
228
229 Ok(events)
230 }
231}
232
233pub(crate) fn parse_tool_result(content: &str) -> (bool, String, String) {
234 if let Some(rest) = content.strip_prefix("__ERROR__:")
235 && let Some((id, body)) = rest.split_once("__\n")
236 {
237 return (true, id.to_string(), body.to_string());
238 }
239 if let Some(rest) = content.strip_prefix("__OK__:")
240 && let Some((id, body)) = rest.split_once("__\n")
241 {
242 return (false, id.to_string(), body.to_string());
243 }
244 (false, "unknown".to_string(), content.to_string())
245}
246
247pub(crate) fn message_parts(message: &Message) -> (String, String) {
248 match message {
249 Message::User { content } => ("user".to_string(), content.clone()),
250 Message::Assistant {
251 content,
252 tool_calls,
253 ..
254 } => {
255 if tool_calls.is_empty() {
256 return ("assistant".to_string(), content.clone());
257 }
258 let mut full = content.clone();
259 for tc in tool_calls {
260 let block = serde_json::json!({
261 "id": tc.id,
262 "name": tc.name,
263 "args": tc.input,
264 });
265 full.push_str(&format!("\n```json-tool\n{block}\n```"));
266 }
267 ("assistant".to_string(), full)
268 }
269 Message::Tool { result } => {
270 let prefix = if result.is_error {
271 format!("__ERROR__:{}__\n", result.tool_use_id)
272 } else {
273 format!("__OK__:{}__\n", result.tool_use_id)
274 };
275 ("system".to_string(), format!("{prefix}{}", result.content))
276 }
277 Message::System { content, .. } => ("system".to_string(), format!("__SYSTEM__:{content}")),
278 Message::Context { content } => ("user".to_string(), content.clone()),
279 Message::Multimodal { parts } => {
280 let mut content = String::new();
281 for part in parts {
282 match part {
283 talos_core::message::ContentPart::Text { text } => content.push_str(text),
284 talos_core::message::ContentPart::Image {
285 path,
286 mime,
287 byte_count,
288 content_digest: _,
289 } => {
290 let filename = path
291 .file_name()
292 .and_then(|n| n.to_str())
293 .unwrap_or("(unknown)");
294 content.push_str(&format!(
295 " [Image: {} ({} bytes, {})]",
296 filename, byte_count, mime
297 ));
298 }
299 }
300 }
301 ("user".to_string(), content)
302 }
303 }
304}
305
306pub(crate) fn preview_text(content: &str) -> String {
307 const MAX_PREVIEW_CHARS: usize = 100;
308 let mut chars = content.chars();
309 let preview: String = chars.by_ref().take(MAX_PREVIEW_CHARS).collect();
310 if chars.next().is_some() {
311 format!("{preview}...")
312 } else {
313 preview
314 }
315}