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