edgecrab_types/
trajectory.rs

1//! Trajectory types for session recording and RL training.
2//!
3//! Trajectories capture the full conversation history including tool
4//! calls, token usage, and metadata for offline analysis and training.
5
6use serde::{Deserialize, Serialize};
7
8use crate::Message;
9
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct Trajectory {
12    pub session_id: String,
13    pub model: String,
14    pub timestamp: String,
15    pub messages: Vec<Message>,
16    pub metadata: TrajectoryMetadata,
17}
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct TrajectoryMetadata {
21    pub task_id: Option<String>,
22    pub total_tokens: u64,
23    pub total_cost: f64,
24    pub api_calls: u32,
25    pub tools_used: Vec<String>,
26    pub completed: bool,
27    pub duration_seconds: f64,
28}
29
30/// Extract thinking blocks from assistant content.
31///
32/// Returns (cleaned_content, optional_reasoning).
33/// Handles `<think>…</think>` tags used by reasoning models.
34pub fn extract_reasoning(content: &str) -> (String, Option<String>) {
35    let re = regex::Regex::new(r"(?s)<think>(.*?)</think>").expect("valid regex");
36    let reasoning = re.captures(content).map(|c| c[1].trim().to_string());
37    let cleaned = re.replace_all(content, "").trim().to_string();
38    (cleaned, reasoning)
39}
40
41/// Check if content after think block is empty (thinking exhaustion).
42pub fn has_content_after_think(content: &str) -> bool {
43    let re = regex::Regex::new(r"(?s)<think>.*?</think>").expect("valid regex");
44    let after = re.replace_all(content, "").trim().to_string();
45    !after.is_empty()
46}
47
48/// Convert legacy `<REASONING_SCRATCHPAD>` tags to `<think>` tags.
49///
50/// Both tag formats are supported by reasoning models.
51pub fn convert_scratchpad_to_think(content: &str) -> String {
52    content
53        .replace("<REASONING_SCRATCHPAD>", "<think>")
54        .replace("</REASONING_SCRATCHPAD>", "</think>")
55}
56
57/// Check if content has an opening scratchpad tag without a closing one.
58pub fn has_incomplete_scratchpad(content: &str) -> bool {
59    content.contains("<REASONING_SCRATCHPAD>") && !content.contains("</REASONING_SCRATCHPAD>")
60}
61
62/// Save trajectory as JSONL (append).
63pub fn save_trajectory(path: &std::path::Path, trajectory: &Trajectory) -> std::io::Result<()> {
64    use std::io::Write;
65    let json = serde_json::to_string(trajectory)
66        .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
67    let mut file = std::fs::OpenOptions::new()
68        .create(true)
69        .append(true)
70        .open(path)?;
71    writeln!(file, "{json}")?;
72    Ok(())
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    #[test]
80    fn extract_reasoning_with_think_tags() {
81        let content = "<think>Let me figure this out...</think>The answer is 42.";
82        let (cleaned, reasoning) = extract_reasoning(content);
83        assert_eq!(cleaned, "The answer is 42.");
84        assert_eq!(reasoning.as_deref(), Some("Let me figure this out..."));
85    }
86
87    #[test]
88    fn extract_reasoning_no_tags() {
89        let content = "Just a normal response.";
90        let (cleaned, reasoning) = extract_reasoning(content);
91        assert_eq!(cleaned, "Just a normal response.");
92        assert!(reasoning.is_none());
93    }
94
95    #[test]
96    fn has_content_after_think_true() {
97        assert!(has_content_after_think(
98            "<think>thinking...</think>Real content here."
99        ));
100    }
101
102    #[test]
103    fn has_content_after_think_false() {
104        assert!(!has_content_after_think("<think>only thinking</think>"));
105    }
106
107    #[test]
108    fn convert_scratchpad() {
109        let input = "<REASONING_SCRATCHPAD>stuff</REASONING_SCRATCHPAD>";
110        assert_eq!(convert_scratchpad_to_think(input), "<think>stuff</think>");
111    }
112
113    #[test]
114    fn incomplete_scratchpad() {
115        assert!(has_incomplete_scratchpad(
116            "<REASONING_SCRATCHPAD>partial reasoning"
117        ));
118        assert!(!has_incomplete_scratchpad(
119            "<REASONING_SCRATCHPAD>complete</REASONING_SCRATCHPAD>"
120        ));
121    }
122
123    #[test]
124    fn trajectory_roundtrip() {
125        let traj = Trajectory {
126            session_id: "s1".into(),
127            model: "test-model".into(),
128            timestamp: "2026-03-28T00:00:00Z".into(),
129            messages: vec![Message::user("hello"), Message::assistant("hi")],
130            metadata: TrajectoryMetadata {
131                task_id: Some("t1".into()),
132                total_tokens: 100,
133                total_cost: 0.001,
134                api_calls: 1,
135                tools_used: vec!["read_file".into()],
136                completed: true,
137                duration_seconds: 2.5,
138            },
139        };
140        let json = serde_json::to_string(&traj).expect("serialize");
141        let deser: Trajectory = serde_json::from_str(&json).expect("deserialize");
142        assert_eq!(deser.session_id, "s1");
143        assert_eq!(deser.messages.len(), 2);
144    }
145}
edgecrab_types/trajectory.rs

edgecrab_types/
trajectory.rs
