Skip to main content

talos_session/
jsonl.rs

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    /// Appends the authoritative terminal transcript outcome for one Turn.
65    ///
66    /// Callers must write this marker only after every transcript message for
67    /// the outcome has been durably appended. Recovery never infers Success
68    /// from ordinary or partial message entries alone.
69    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}