Skip to main content

ytil_agents/agent/
session.rs

1use std::fmt::Display;
2use std::fmt::Formatter;
3use std::path::PathBuf;
4use std::str::FromStr;
5
6use chrono::DateTime;
7use chrono::Utc;
8use rootcause::option_ext::OptionExt;
9use rootcause::report;
10
11use crate::agent::Agent;
12
13const SEARCH_TEXT_MAX_BYTES: usize = 32 * 1024;
14
15#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
16pub struct SessionKey {
17    agent: Agent,
18    id: String,
19}
20
21impl SessionKey {
22    pub fn new(agent: Agent, id: impl Into<String>) -> Self {
23        Self { agent, id: id.into() }
24    }
25
26    pub const fn agent(&self) -> Agent {
27        self.agent
28    }
29
30    pub fn id(&self) -> &str {
31        &self.id
32    }
33}
34
35impl Display for SessionKey {
36    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
37        write!(f, "{}:{}", self.agent.name(), self.id)
38    }
39}
40
41impl FromStr for SessionKey {
42    type Err = rootcause::Report;
43
44    fn from_str(value: &str) -> Result<Self, Self::Err> {
45        let Some((agent, id)) = value.split_once(':') else {
46            return Err(report!("invalid session key").attach(format!("value={value}")));
47        };
48        let agent =
49            Agent::from_name(agent).map_err(|err| report!("invalid session key agent").attach(err.to_string()))?;
50        if id.is_empty() {
51            return Err(report!("invalid session key").attach(format!("value={value}")));
52        }
53        Ok(Self::new(agent, id))
54    }
55}
56
57#[derive(Clone, Debug, Eq, PartialEq)]
58pub struct Session {
59    pub id: String,
60    pub agent: Agent,
61    pub name: String,
62    pub search_text: String,
63    pub workspace: PathBuf,
64    pub path: PathBuf,
65    pub created_at: DateTime<Utc>,
66    pub updated_at: DateTime<Utc>,
67}
68
69impl Session {
70    pub fn new(
71        agent: Agent,
72        session_id: String,
73        workspace_dir: PathBuf,
74        path: PathBuf,
75        name: Option<String>,
76        created_at: DateTime<Utc>,
77    ) -> Self {
78        let name = name.filter(|name| !name.trim().is_empty()).unwrap_or_else(|| {
79            workspace_dir
80                .file_name()
81                .and_then(|name| name.to_str())
82                .filter(|name| !name.is_empty())
83                .map_or_else(|| session_id.clone(), str::to_owned)
84        });
85
86        Self {
87            id: session_id,
88            agent,
89            search_text: name.clone(),
90            name,
91            workspace: workspace_dir,
92            path,
93            created_at,
94            updated_at: created_at,
95        }
96    }
97
98    /// Build the argv required to resume this session with its owning agent CLI.
99    ///
100    /// # Errors
101    /// Returns an error when the workspace path is not UTF-8 or the agent has no
102    /// supported resume command.
103    pub fn build_resume_command(&self) -> rootcause::Result<(&'static str, Vec<String>)> {
104        let workspace = self.workspace.to_str().context("non-utf8 workspace dir".to_owned())?;
105        match self.agent {
106            Agent::Claude => Ok(("claude", vec!["--resume".into(), self.id.clone()])),
107            Agent::Codex => Ok((
108                "codex",
109                self.build_codex_resume_args(workspace, std::env::var_os("ZELLIJ").is_some()),
110            )),
111            Agent::Cursor => Ok((
112                "cursor-agent",
113                vec![
114                    "--resume".into(),
115                    self.id.clone(),
116                    "--workspace".into(),
117                    workspace.into(),
118                ],
119            )),
120            Agent::Gemini | Agent::Opencode => {
121                Err(report!("resume is not supported for this agent").attach(format!("agent={}", self.agent)))
122            }
123        }
124    }
125
126    fn build_codex_resume_args(&self, workspace: &str, is_zellij: bool) -> Vec<String> {
127        let mut args = vec!["resume".into(), self.id.clone()];
128        // In Zellij, Codex's mouse-aware TUI captures wheel events before
129        // Zellij can use them for inline scrollback. Keep inline mode only
130        // outside Zellij, where terminal scrollback works as intended.
131        if !is_zellij {
132            args.push("--no-alt-screen".into());
133        }
134        args.extend(["--cd".into(), workspace.into()]);
135        args
136    }
137}
138
139#[derive(Debug, Default)]
140pub struct SearchTextBuilder {
141    snippets_text: String,
142    first_snippet: Option<String>,
143    last_snippet: Option<String>,
144    reached_limit: bool,
145}
146
147impl SearchTextBuilder {
148    pub fn push(&mut self, raw: &str) {
149        if self.reached_limit {
150            return;
151        }
152
153        let snippet = raw.split_whitespace().collect::<Vec<_>>().join(" ");
154        let Some(snippet) = (!snippet.is_empty()).then_some(snippet) else {
155            return;
156        };
157        if self.last_snippet.as_ref().is_some_and(|last| last == &snippet) {
158            return;
159        }
160        if self.first_snippet.is_none() {
161            self.first_snippet = Some(snippet.clone());
162        }
163
164        self.reached_limit = !push_normalized_snippet(&mut self.snippets_text, &mut self.last_snippet, &snippet);
165    }
166
167    pub fn build(self, fallback: &str) -> String {
168        let fallback = fallback.split_whitespace().collect::<Vec<_>>().join(" ");
169        let Some(fallback) = (!fallback.is_empty()).then_some(fallback) else {
170            return self.snippets_text;
171        };
172
173        if self.first_snippet.as_ref().is_some_and(|first| first == &fallback) {
174            return self.snippets_text;
175        }
176
177        let mut search_text = String::new();
178        let mut last_snippet = None::<String>;
179        if !push_normalized_snippet(&mut search_text, &mut last_snippet, &fallback) {
180            return search_text;
181        }
182        if self.snippets_text.is_empty() {
183            return search_text;
184        }
185
186        let separator_len = usize::from(!search_text.is_empty());
187        if search_text.len().saturating_add(separator_len) >= SEARCH_TEXT_MAX_BYTES {
188            return search_text;
189        }
190        if !search_text.is_empty() {
191            search_text.push(' ');
192        }
193
194        let remaining = SEARCH_TEXT_MAX_BYTES.saturating_sub(search_text.len());
195        if let Some(truncated) = truncate_to_boundary(&self.snippets_text, remaining) {
196            search_text.push_str(truncated);
197        }
198
199        search_text
200    }
201}
202
203fn push_normalized_snippet(search_text: &mut String, last_snippet: &mut Option<String>, snippet: &str) -> bool {
204    let separator_len = usize::from(!search_text.is_empty());
205    if search_text.len().saturating_add(separator_len) >= SEARCH_TEXT_MAX_BYTES {
206        return false;
207    }
208    if !search_text.is_empty() {
209        search_text.push(' ');
210    }
211
212    let remaining = SEARCH_TEXT_MAX_BYTES.saturating_sub(search_text.len());
213    if remaining == 0 {
214        return false;
215    }
216
217    let snippet_len = snippet.len();
218    truncate_to_boundary(snippet, remaining).is_some_and(|truncated| {
219        let is_full_snippet = truncated.len() == snippet_len;
220        search_text.push_str(truncated);
221        *last_snippet = Some(snippet.to_owned());
222        is_full_snippet
223    })
224}
225
226fn truncate_to_boundary(text: &str, max_bytes: usize) -> Option<&str> {
227    if max_bytes == 0 {
228        return None;
229    }
230    if text.len() <= max_bytes {
231        return Some(text);
232    }
233
234    let mut end = 0;
235    for (idx, ch) in text.char_indices() {
236        let next = idx.saturating_add(ch.len_utf8());
237        if next > max_bytes {
238            break;
239        }
240        end = next;
241    }
242
243    (end > 0).then(|| text.get(..end)).flatten()
244}
245
246#[cfg(test)]
247mod tests {
248    use chrono::DateTime;
249    use tempfile::tempdir;
250
251    use super::*;
252
253    #[test]
254    fn test_session_key_string_round_trip_uses_agent_session_format() {
255        assert2::assert!(let Ok(key) = "codex:session-id".parse::<SessionKey>());
256
257        pretty_assertions::assert_eq!(key, SessionKey::new(Agent::Codex, "session-id"));
258        pretty_assertions::assert_eq!(key.to_string(), "codex:session-id");
259    }
260
261    #[test]
262    fn test_build_resume_command_matches_agent() {
263        let tempdir = tempdir().expect("tempdir should be created");
264        let workspace = tempdir.path().join("workspace");
265        let path = tempdir.path().join("session.jsonl");
266        std::fs::create_dir_all(&workspace).expect("workspace should be created");
267        let created_at = DateTime::from_timestamp_millis(1).expect("test timestamp should be valid");
268
269        let claude = Session {
270            agent: Agent::Claude,
271            id: "session-id".into(),
272            workspace: workspace.clone(),
273            name: "session-name".into(),
274            search_text: "session-name".into(),
275            path,
276            created_at: created_at.to_utc(),
277            updated_at: created_at.to_utc(),
278        };
279        let codex = Session {
280            agent: Agent::Codex,
281            ..claude.clone()
282        };
283        let cursor = Session {
284            agent: Agent::Cursor,
285            ..claude.clone()
286        };
287
288        assert2::assert!(let Ok((_, claude_args)) = claude.build_resume_command());
289        pretty_assertions::assert_eq!(claude_args, vec!["--resume".to_owned(), "session-id".to_owned()]);
290        let workspace_str = workspace.to_str().expect("workspace test path should be utf8");
291        pretty_assertions::assert_eq!(
292            codex.build_codex_resume_args(workspace_str, false),
293            vec![
294                "resume".to_owned(),
295                "session-id".to_owned(),
296                "--no-alt-screen".to_owned(),
297                "--cd".to_owned(),
298                workspace_str.to_owned(),
299            ]
300        );
301        pretty_assertions::assert_eq!(
302            codex.build_codex_resume_args(workspace_str, true),
303            vec![
304                "resume".to_owned(),
305                "session-id".to_owned(),
306                "--cd".to_owned(),
307                workspace_str.to_owned(),
308            ]
309        );
310        assert2::assert!(let Ok((_, cursor_args)) = cursor.build_resume_command());
311        pretty_assertions::assert_eq!(
312            cursor_args,
313            vec![
314                "--resume".to_owned(),
315                "session-id".to_owned(),
316                "--workspace".to_owned(),
317                workspace_str.to_owned(),
318            ]
319        );
320    }
321
322    #[test]
323    fn test_session_new_sets_search_text_from_resolved_name() {
324        let tempdir = tempdir().expect("tempdir should be created");
325        let workspace = tempdir.path().join("workspace");
326        std::fs::create_dir_all(&workspace).expect("workspace should be created");
327        let created_at = DateTime::from_timestamp_millis(1).expect("test timestamp should be valid");
328
329        let session = Session::new(
330            Agent::Codex,
331            "session-id".into(),
332            workspace,
333            PathBuf::from("session.jsonl"),
334            Some("hello world".into()),
335            created_at.to_utc(),
336        );
337
338        pretty_assertions::assert_eq!(session.name, "hello world");
339        pretty_assertions::assert_eq!(session.search_text, "hello world");
340    }
341
342    #[test]
343    fn test_search_text_builder_normalizes_dedupes_and_falls_back() {
344        let mut builder = SearchTextBuilder::default();
345        for snippet in ["  fallback  ", "first\nline", "", "first line", "second\tline"] {
346            builder.push(snippet);
347        }
348        let search_text = builder.build("fallback");
349
350        pretty_assertions::assert_eq!(search_text, "fallback first line second line");
351    }
352}