1use serde::{Deserialize, Serialize};
8
9use abk::checkpoint::{
10 get_storage_manager,
11 models::{ChatMessage, CheckpointMetadata, SessionMetadata as AbkSessionMetadata},
12 storage::ProjectMetadata as AbkProjectMetadata,
13 CheckpointStorageManager,
14};
15use abk::cli::ResumeInfo;
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct SessionSummary {
20 pub session_id: String,
21 pub project_name: String,
22 pub project_path: String,
23 pub project_id: String,
25 pub checkpoint_count: usize,
26 pub created_at: chrono::DateTime<chrono::Utc>,
27 pub last_accessed: chrono::DateTime<chrono::Utc>,
28 pub description: Option<String>,
29 pub is_current_project: bool,
30}
31
32#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct CheckpointSummary {
35 pub checkpoint_id: String,
36 pub session_id: String,
37 pub iteration: u32,
38 pub workflow_step: String,
39 pub created_at: chrono::DateTime<chrono::Utc>,
40}
41
42fn session_to_summary(
47 session: &AbkSessionMetadata,
48 project: &AbkProjectMetadata,
49 is_current: bool,
50) -> SessionSummary {
51 SessionSummary {
52 session_id: session.session_id.clone(),
53 project_name: project.name.clone(),
54 project_path: project.project_path.to_string_lossy().to_string(),
55 project_id: project.project_hash.clone(),
56 checkpoint_count: session.checkpoint_count as usize,
57 created_at: session.created_at,
58 last_accessed: session.last_accessed,
59 description: session.description.clone(),
60 is_current_project: is_current,
61 }
62}
63
64fn checkpoint_to_summary(cp: &CheckpointMetadata) -> CheckpointSummary {
65 CheckpointSummary {
66 checkpoint_id: cp.checkpoint_id.clone(),
67 session_id: cp.session_id.clone(),
68 iteration: cp.iteration,
69 workflow_step: format!("{:?}", cp.workflow_step),
70 created_at: cp.created_at,
71 }
72}
73
74fn config_working_dir(config_toml: &str) -> std::path::PathBuf {
79 if let Ok(value) = toml::from_str::<toml::Value>(config_toml) {
80 if let Some(agent) = value.get("agent") {
81 if let Some(wd) = agent.get("working_dir").and_then(|v| v.as_str()) {
82 return std::path::PathBuf::from(wd);
83 }
84 }
85 }
86 std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."))
87}
88
89fn paths_match(a: &std::path::Path, b: &std::path::Path) -> bool {
91 let ca = a.canonicalize().unwrap_or_else(|_| a.to_path_buf());
92 let cb = b.canonicalize().unwrap_or_else(|_| b.to_path_buf());
93 ca == cb
94}
95
96async fn make_storage_manager(
105 config_toml: &str,
106 home_dir: Option<&std::path::Path>,
107 agent_name: &str,
108) -> anyhow::Result<CheckpointStorageManager> {
109 let backend_config = parse_storage_backend_config(config_toml);
111
112 match (home_dir, backend_config) {
113 (Some(dir), Some(backend)) => CheckpointStorageManager::with_home_dir_and_backend(
115 dir.to_path_buf(),
116 agent_name,
117 backend,
118 )
119 .await
120 .map_err(|e| anyhow::anyhow!("Failed to create storage manager with backend: {}", e)),
121 (Some(dir), None) => CheckpointStorageManager::with_home_dir(dir.to_path_buf(), agent_name)
123 .map_err(|e| anyhow::anyhow!("Failed to create storage manager: {}", e)),
124 (None, _) => get_storage_manager()
126 .map_err(|e| anyhow::anyhow!("Failed to get storage manager: {}", e)),
127 }
128}
129
130fn parse_storage_backend_config(
134 config_toml: &str,
135) -> Option<abk::checkpoint::StorageBackendConfig> {
136 let value = toml::from_str::<toml::Value>(config_toml).ok()?;
137 let checkpointing = value.get("checkpointing")?;
138 let storage_backend = checkpointing.get("storage_backend")?;
139 let backend: abk::checkpoint::StorageBackendConfig = storage_backend.clone().try_into().ok()?;
140
141 match backend.backend_type {
143 abk::checkpoint::StorageBackendType::DocumentDB
144 | abk::checkpoint::StorageBackendType::MongoDB => Some(backend),
145 abk::checkpoint::StorageBackendType::File => None,
146 }
147}
148
149pub async fn list_all_sessions(
161 config_toml: &str,
162 home_dir: Option<&std::path::Path>,
163) -> anyhow::Result<Vec<SessionSummary>> {
164 let current_dir = config_working_dir(config_toml);
165 let manager = make_storage_manager(config_toml, home_dir, "trustee").await?;
166
167 let projects = manager
168 .list_projects()
169 .await
170 .map_err(|e| anyhow::anyhow!("Failed to list projects: {}", e))?;
171
172 let mut summaries = Vec::new();
173
174 for project in &projects {
175 let project_storage = match manager
178 .get_project_storage_with_id(&project.project_path, &project.project_hash)
179 .await
180 {
181 Ok(ps) => ps,
182 Err(e) => {
183 tracing::debug!("Skipping project {}: {}", project.project_path.display(), e);
184 continue;
185 }
186 };
187
188 let sessions = match project_storage.list_sessions().await {
189 Ok(s) => s,
190 Err(e) => {
191 tracing::debug!(
192 "Failed to list sessions for {}: {}",
193 project.project_path.display(),
194 e
195 );
196 continue;
197 }
198 };
199
200 let is_current = paths_match(&project.project_path, ¤t_dir);
201
202 for session in sessions {
203 if session.checkpoint_count > 0 {
205 summaries.push(session_to_summary(&session, project, is_current));
206 }
207 }
208 }
209
210 summaries.sort_by(|a, b| match (a.is_current_project, b.is_current_project) {
212 (true, false) => std::cmp::Ordering::Less,
213 (false, true) => std::cmp::Ordering::Greater,
214 _ => b.last_accessed.cmp(&a.last_accessed),
215 });
216
217 Ok(summaries)
218}
219
220pub async fn get_session_detail(
227 _config_toml: &str,
228 session_id: &str,
229 home_dir: Option<&std::path::Path>,
230) -> anyhow::Result<Option<(SessionSummary, Vec<CheckpointSummary>)>> {
231 let manager = make_storage_manager(_config_toml, home_dir, "trustee").await?;
232
233 let projects = manager
234 .list_projects()
235 .await
236 .map_err(|e| anyhow::anyhow!("Failed to list projects: {}", e))?;
237
238 for project in &projects {
239 let project_storage = match manager
240 .get_project_storage_with_id(&project.project_path, &project.project_hash)
241 .await
242 {
243 Ok(ps) => ps,
244 Err(e) => {
245 tracing::debug!("Skipping project {}: {}", project.project_path.display(), e);
246 continue;
247 }
248 };
249
250 let sessions = match project_storage.list_sessions().await {
251 Ok(s) => s,
252 Err(e) => {
253 tracing::debug!(
254 "Failed to list sessions for {}: {}",
255 project.project_path.display(),
256 e
257 );
258 continue;
259 }
260 };
261
262 if let Some(session) = sessions.iter().find(|s| s.session_id == session_id) {
263 let session_storage = project_storage
264 .create_session(session_id)
265 .await
266 .map_err(|e| anyhow::anyhow!("Failed to get session storage: {}", e))?;
267
268 let checkpoints = session_storage
269 .list_checkpoints()
270 .await
271 .map_err(|e| anyhow::anyhow!("Failed to list checkpoints: {}", e))?;
272
273 let summary = session_to_summary(session, project, false);
274 let cp_summaries: Vec<CheckpointSummary> =
275 checkpoints.iter().map(checkpoint_to_summary).collect();
276
277 return Ok(Some((summary, cp_summaries)));
278 }
279 }
280
281 Ok(None)
282}
283
284pub async fn create_resume_info(
292 _config_toml: &str,
293 session_id: &str,
294 home_dir: Option<&std::path::Path>,
295) -> anyhow::Result<Option<ResumeInfo>> {
296 let manager = make_storage_manager(_config_toml, home_dir, "trustee").await?;
297
298 let projects = manager
299 .list_projects()
300 .await
301 .map_err(|e| anyhow::anyhow!("Failed to list projects: {}", e))?;
302
303 for project in &projects {
304 let project_storage = match manager
305 .get_project_storage_with_id(&project.project_path, &project.project_hash)
306 .await
307 {
308 Ok(ps) => ps,
309 Err(e) => {
310 tracing::debug!("Skipping project {}: {}", project.project_path.display(), e);
311 continue;
312 }
313 };
314
315 let sessions = match project_storage.list_sessions().await {
316 Ok(s) => s,
317 Err(e) => {
318 tracing::debug!(
319 "Failed to list sessions for {}: {}",
320 project.project_path.display(),
321 e
322 );
323 continue;
324 }
325 };
326
327 if sessions.iter().any(|s| s.session_id == session_id) {
328 let session_storage = project_storage
329 .create_session(session_id)
330 .await
331 .map_err(|e| anyhow::anyhow!("Failed to get session storage: {}", e))?;
332
333 let checkpoints = session_storage
334 .list_checkpoints()
335 .await
336 .map_err(|e| anyhow::anyhow!("Failed to list checkpoints: {}", e))?;
337
338 if let Some(latest) = checkpoints.iter().max_by_key(|cp| cp.created_at) {
339 return Ok(Some(ResumeInfo {
340 session_id: session_id.to_string(),
341 checkpoint_id: latest.checkpoint_id.clone(),
342 iteration: latest.iteration,
343 project_path: Some(project.project_path.clone()),
344 }));
345 }
346
347 return Ok(None);
349 }
350 }
351
352 Ok(None)
353}
354
355pub async fn fork_session_chain(
379 config_toml: &str,
380 source_session_id: &str,
381 new_session_id: &str,
382 home_dir: Option<&std::path::Path>,
383) -> anyhow::Result<Option<ResumeInfo>> {
384 let manager = make_storage_manager(config_toml, home_dir, "trustee").await?;
385
386 let projects = manager
387 .list_projects()
388 .await
389 .map_err(|e| anyhow::anyhow!("Failed to list projects: {}", e))?;
390
391 for project in &projects {
392 let project_storage = match manager
393 .get_project_storage_with_id(&project.project_path, &project.project_hash)
394 .await
395 {
396 Ok(ps) => ps,
397 Err(e) => {
398 tracing::debug!("Skipping project {}: {}", project.project_path.display(), e);
399 continue;
400 }
401 };
402
403 let sessions = match project_storage.list_sessions().await {
404 Ok(s) => s,
405 Err(e) => {
406 tracing::debug!(
407 "Failed to list sessions for {}: {}",
408 project.project_path.display(),
409 e
410 );
411 continue;
412 }
413 };
414
415 if !sessions.iter().any(|s| s.session_id == source_session_id) {
416 continue;
417 }
418
419 let source_storage = project_storage
420 .create_session(source_session_id)
421 .await
422 .map_err(|e| anyhow::anyhow!("Failed to get source session storage: {}", e))?;
423
424 let checkpoints = source_storage
425 .list_checkpoints()
426 .await
427 .map_err(|e| anyhow::anyhow!("Failed to list source checkpoints: {}", e))?;
428
429 if checkpoints.is_empty() {
430 return Ok(None);
431 }
432
433 let mut fork_storage = project_storage
436 .create_session(new_session_id)
437 .await
438 .map_err(|e| anyhow::anyhow!("Failed to create fork session storage: {}", e))?;
439
440 let mut latest: Option<CheckpointMetadata> = None;
441 for cp in &checkpoints {
442 let mut checkpoint = source_storage
443 .load_checkpoint(&cp.checkpoint_id)
444 .await
445 .map_err(|e| {
446 anyhow::anyhow!("Failed to load checkpoint {}: {}", cp.checkpoint_id, e)
447 })?;
448 checkpoint.metadata.session_id = new_session_id.to_string();
449 fork_storage
450 .save_checkpoint(&checkpoint)
451 .await
452 .map_err(|e| {
453 anyhow::anyhow!(
454 "Failed to save forked checkpoint {}: {}",
455 cp.checkpoint_id,
456 e
457 )
458 })?;
459 if latest
460 .as_ref()
461 .map_or(true, |l| cp.created_at > l.created_at)
462 {
463 latest = Some(cp.clone());
464 }
465 }
466
467 let latest = latest.expect("checkpoints non-empty");
468 return Ok(Some(ResumeInfo {
469 session_id: new_session_id.to_string(),
470 checkpoint_id: latest.checkpoint_id.clone(),
471 iteration: latest.iteration,
472 project_path: Some(project.project_path.clone()),
473 }));
474 }
475
476 Ok(None)
477}
478
479#[derive(Debug, Clone, Serialize, Deserialize)]
492pub struct HistoryMessage {
493 pub role: String,
495 pub content: String,
497 pub reasoning: Option<String>,
499 pub tool_calls: Option<Vec<HistoryToolCall>>,
501 pub name: Option<String>,
503}
504
505#[derive(Debug, Clone, Serialize, Deserialize)]
507pub struct HistoryToolCall {
508 pub name: String,
509 pub hint: String,
511}
512
513#[derive(Debug, Clone, Serialize, Deserialize)]
515pub struct SessionHistory {
516 pub session_id: String,
517 pub checkpoint_id: String,
518 pub task_description: String,
519 pub iteration: u32,
520 pub total_messages: usize,
521 pub messages: Vec<HistoryMessage>,
522}
523
524pub async fn load_session_history(
532 session_id: &str,
533 home_dir: Option<&std::path::Path>,
534) -> anyhow::Result<Option<SessionHistory>> {
535 let manager = make_storage_manager("", home_dir, "trustee").await?;
536
537 let projects = manager
538 .list_projects()
539 .await
540 .map_err(|e| anyhow::anyhow!("Failed to list projects: {}", e))?;
541
542 for project in &projects {
543 let project_storage = match manager
544 .get_project_storage_with_id(&project.project_path, &project.project_hash)
545 .await
546 {
547 Ok(ps) => ps,
548 Err(e) => {
549 tracing::debug!("Skipping project {}: {}", project.project_path.display(), e);
550 continue;
551 }
552 };
553
554 let sessions = match project_storage.list_sessions().await {
555 Ok(s) => s,
556 Err(e) => {
557 tracing::debug!(
558 "Failed to list sessions for {}: {}",
559 project.project_path.display(),
560 e
561 );
562 continue;
563 }
564 };
565
566 if !sessions.iter().any(|s| s.session_id == session_id) {
567 continue;
568 }
569
570 let session_storage = project_storage
571 .create_session(session_id)
572 .await
573 .map_err(|e| anyhow::anyhow!("Failed to get session storage: {}", e))?;
574
575 let checkpoints = session_storage
576 .list_checkpoints()
577 .await
578 .map_err(|e| anyhow::anyhow!("Failed to list checkpoints: {}", e))?;
579
580 let latest = match checkpoints.iter().max_by_key(|cp| cp.created_at) {
581 Some(cp) => cp,
582 None => return Ok(None),
583 };
584
585 let checkpoint_id = latest.checkpoint_id.clone();
586 let iteration = latest.iteration;
587
588 let checkpoint = session_storage
589 .load_checkpoint(&checkpoint_id)
590 .await
591 .map_err(|e| anyhow::anyhow!("Failed to load checkpoint: {}", e))?;
592
593 let task_description = checkpoint.agent_state.task_description.clone();
594 let total_messages = checkpoint.conversation_state.messages.len();
595
596 let messages = convert_messages(&checkpoint.conversation_state.messages);
597
598 return Ok(Some(SessionHistory {
599 session_id: session_id.to_string(),
600 checkpoint_id,
601 task_description,
602 iteration,
603 total_messages,
604 messages,
605 }));
606 }
607
608 Ok(None)
609}
610
611fn convert_messages(messages: &[ChatMessage]) -> Vec<HistoryMessage> {
617 const MAX_CONTENT_LEN: usize = 10_000;
618
619 let mut result = Vec::new();
620
621 for msg in messages {
622 if msg.role == "system" {
624 continue;
625 }
626
627 if msg.role == "tool" {
630 continue;
631 }
632
633 let content = if msg.content.len() > MAX_CONTENT_LEN {
638 format!(
639 "{}...\n[truncated]",
640 abk::text::truncate_at_boundary(&msg.content, MAX_CONTENT_LEN)
641 )
642 } else {
643 msg.content.clone()
644 };
645
646 let tool_calls = msg.tool_calls.as_ref().map(|calls| {
648 calls
649 .iter()
650 .map(|tc| {
651 let hint = summarize_tool_args(&tc.function.name, &tc.function.arguments);
652 HistoryToolCall {
653 name: tc.function.name.clone(),
654 hint,
655 }
656 })
657 .collect::<Vec<_>>()
658 });
659
660 if msg.role == "assistant" && content.is_empty() && tool_calls.is_some() {
663 result.push(HistoryMessage {
664 role: "assistant".to_string(),
665 content: String::new(),
666 reasoning: msg.reasoning.clone(),
667 tool_calls,
668 name: None,
669 });
670 continue;
671 }
672
673 result.push(HistoryMessage {
674 role: msg.role.clone(),
675 content,
676 reasoning: msg.reasoning.clone(),
677 tool_calls,
678 name: msg.name.clone(),
679 });
680 }
681
682 result
683}
684
685fn summarize_tool_args(name: &str, args: &str) -> String {
690 let parsed: serde_json::Value = match serde_json::from_str(args) {
691 Ok(v) => v,
692 Err(_) => return args.chars().take(200).collect(),
693 };
694
695 let obj = match parsed.as_object() {
696 Some(o) => o,
697 None => return args.chars().take(200).collect(),
698 };
699
700 match name {
701 "bash" | "execute_command" => obj
702 .get("command")
703 .and_then(|v| v.as_str())
704 .map(|s| s.to_string()),
705 "read" | "read_file" => obj
706 .get("file_path")
707 .and_then(|v| v.as_str())
708 .map(|s| s.to_string()),
709 "write" | "write_file" | "edit" => obj
710 .get("file_path")
711 .and_then(|v| v.as_str())
712 .map(|s| s.to_string()),
713 "grep" | "search" => obj
714 .get("pattern")
715 .and_then(|v| v.as_str())
716 .map(|s| s.to_string()),
717 "glob" => obj
718 .get("pattern")
719 .and_then(|v| v.as_str())
720 .map(|s| s.to_string()),
721 "todowrite" => {
722 let count = obj
723 .get("todos")
724 .and_then(|v| v.as_array())
725 .map(|a| a.len())
726 .unwrap_or(0);
727 Some(format!("{} items", count))
728 }
729 "websearch" | "webfetch" => obj
730 .get("query")
731 .or_else(|| obj.get("url"))
732 .and_then(|v| v.as_str())
733 .map(|s| s.to_string()),
734 _ => obj.values().next().and_then(|v| match v {
735 serde_json::Value::String(s) => Some(s.clone()),
736 _ => Some(v.to_string()),
737 }),
738 }
739 .unwrap_or_default()
740 .chars()
741 .take(200)
742 .collect()
743}
744
745#[cfg(test)]
746mod history_truncation_tests {
747 use super::*;
748
749 fn msg(role: &str, content: String) -> ChatMessage {
750 ChatMessage {
751 role: role.to_string(),
752 content,
753 reasoning: None,
754 timestamp: chrono::Utc::now(),
755 token_count: None,
756 tool_calls: None,
757 tool_call_id: None,
758 name: None,
759 }
760 }
761
762 #[test]
766 fn persian_long_message_history_does_not_panic() {
767 let long = "ب".repeat(6_000); let out = convert_messages(&[msg("user", long.clone())]);
769 let first = &out[0].content;
770 assert!(
771 first.ends_with("...\n[truncated]"),
772 "marker missing: {:?}",
773 &first[first.len().saturating_sub(30)..]
774 );
775 let body = &first[..first.len() - "...\n[truncated]".len()];
776 assert!(body.len() <= 10_000, "body over budget: {}", body.len());
777 assert!(
778 long.starts_with(body),
779 "body must be a prefix of the original"
780 );
781 }
782
783 #[test]
785 fn ascii_long_message_truncates_identically_to_legacy() {
786 let out = convert_messages(&[msg("user", "A".repeat(20_000))]);
787 assert_eq!(
788 out[0].content,
789 format!("{}...\n[truncated]", "A".repeat(10_000))
790 );
791 }
792
793 #[test]
795 fn short_persian_message_unchanged() {
796 let out = convert_messages(&[msg("user", "سلام دنیا".to_string())]);
797 assert_eq!(out[0].content, "سلام دنیا");
798 }
799
800 #[test]
803 fn four_byte_char_straddling_cut_is_excluded() {
804 let long = "\u{1F680}".repeat(3_000); let out = convert_messages(&[msg("user", long.clone())]);
806 let first = &out[0].content;
807 let body = &first[..first.len() - "...\n[truncated]".len()];
808 assert_eq!(body.chars().count(), 2_500); assert!(long.starts_with(body));
810 }
811}