trustee-core 0.4.0

Core types, session state, and workflow logic for Trustee agent
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
//! Session discovery for the API/Web layer.
//!
//! Wraps `abk::checkpoint` to provide serializable session listing, detail,
//! and resume-info creation — used by the Trustee REST API so users can
//! browse and resume checkpoint sessions from the Web UI.

use serde::{Deserialize, Serialize};

use abk::checkpoint::{
    get_storage_manager,
    models::{CheckpointMetadata, ChatMessage, SessionMetadata as AbkSessionMetadata},
    storage::ProjectMetadata as AbkProjectMetadata,
    CheckpointStorageManager,
};
use abk::cli::ResumeInfo;

/// Compact session info suitable for JSON API responses.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionSummary {
    pub session_id: String,
    pub project_name: String,
    pub project_path: String,
    /// Project hash (storage partition key) — used for per-user namespace filtering.
    pub project_id: String,
    pub checkpoint_count: usize,
    pub created_at: chrono::DateTime<chrono::Utc>,
    pub last_accessed: chrono::DateTime<chrono::Utc>,
    pub description: Option<String>,
    pub is_current_project: bool,
}

/// Compact checkpoint info for the session detail endpoint.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointSummary {
    pub checkpoint_id: String,
    pub session_id: String,
    pub iteration: u32,
    pub workflow_step: String,
    pub created_at: chrono::DateTime<chrono::Utc>,
}

// -----------------------------------------------------------------------
// Conversion helpers
// -----------------------------------------------------------------------

fn session_to_summary(
    session: &AbkSessionMetadata,
    project: &AbkProjectMetadata,
    is_current: bool,
) -> SessionSummary {
    SessionSummary {
        session_id: session.session_id.clone(),
        project_name: project.name.clone(),
        project_path: project.project_path.to_string_lossy().to_string(),
        project_id: project.project_hash.clone(),
        checkpoint_count: session.checkpoint_count as usize,
        created_at: session.created_at,
        last_accessed: session.last_accessed,
        description: session.description.clone(),
        is_current_project: is_current,
    }
}

fn checkpoint_to_summary(cp: &CheckpointMetadata) -> CheckpointSummary {
    CheckpointSummary {
        checkpoint_id: cp.checkpoint_id.clone(),
        session_id: cp.session_id.clone(),
        iteration: cp.iteration,
        workflow_step: format!("{:?}", cp.workflow_step),
        created_at: cp.created_at,
    }
}

/// Derive the working directory from a trustee config TOML string.
///
/// Looks for `[agent] working_dir`. Falls back to the current directory
/// of the process if not specified or unparseable.
fn config_working_dir(config_toml: &str) -> std::path::PathBuf {
    if let Ok(value) = toml::from_str::<toml::Value>(config_toml) {
        if let Some(agent) = value.get("agent") {
            if let Some(wd) = agent.get("working_dir").and_then(|v| v.as_str()) {
                return std::path::PathBuf::from(wd);
            }
        }
    }
    std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."))
}

/// Check whether two paths refer to the same project by canonicalising.
fn paths_match(a: &std::path::Path, b: &std::path::Path) -> bool {
    let ca = a.canonicalize().unwrap_or_else(|_| a.to_path_buf());
    let cb = b.canonicalize().unwrap_or_else(|_| b.to_path_buf());
    ca == cb
}

/// Create a storage manager from an optional home_dir override.
///
/// When `home_dir` is `Some`, creates a per-user storage manager.
/// When `None`, falls back to the global `get_storage_manager()`.
fn make_storage_manager(
    home_dir: Option<&std::path::Path>,
    agent_name: &str,
) -> anyhow::Result<CheckpointStorageManager> {
    if let Some(dir) = home_dir {
        CheckpointStorageManager::with_home_dir(dir.to_path_buf(), agent_name)
            .map_err(|e| anyhow::anyhow!("Failed to create storage manager: {}", e))
    } else {
        get_storage_manager()
            .map_err(|e| anyhow::anyhow!("Failed to get storage manager: {}", e))
    }
}

// -----------------------------------------------------------------------
// Public API
// -----------------------------------------------------------------------

/// List all sessions across all projects that have at least one checkpoint.
///
/// Sessions from the current project (derived from `config_toml`) are listed
/// first, then everything else sorted by `last_accessed` descending.
///
/// When `home_dir` is `Some`, queries that specific user's checkpoint storage
/// instead of the global default.
pub async fn list_all_sessions(
    config_toml: &str,
    home_dir: Option<&std::path::Path>,
) -> anyhow::Result<Vec<SessionSummary>> {
    let current_dir = config_working_dir(config_toml);
    let manager = make_storage_manager(home_dir, "trustee")?;

    let projects = manager
        .list_projects()
        .await
        .map_err(|e| anyhow::anyhow!("Failed to list projects: {}", e))?;

    let mut summaries = Vec::new();

    for project in &projects {
        // Skip projects whose path can't be resolved (e.g. deleted dirs,
        // permission denied) — don't let one bad project kill the whole list.
        let project_storage = match manager.get_project_storage(&project.project_path).await {
            Ok(ps) => ps,
            Err(e) => {
                tracing::debug!("Skipping project {}: {}", project.project_path.display(), e);
                continue;
            }
        };

        let sessions = match project_storage.list_sessions().await {
            Ok(s) => s,
            Err(e) => {
                tracing::debug!("Failed to list sessions for {}: {}", project.project_path.display(), e);
                continue;
            }
        };

        let is_current = paths_match(&project.project_path, &current_dir);

        for session in sessions {
            // Only include sessions that have checkpoints (resumable)
            if session.checkpoint_count > 0 {
                summaries.push(session_to_summary(&session, project, is_current));
            }
        }
    }

    // Sort: current project first, then by last_accessed descending
    summaries.sort_by(|a, b| match (a.is_current_project, b.is_current_project) {
        (true, false) => std::cmp::Ordering::Less,
        (false, true) => std::cmp::Ordering::Greater,
        _ => b.last_accessed.cmp(&a.last_accessed),
    });

    Ok(summaries)
}

/// Get detailed information about a specific session, including its checkpoints.
///
/// Searches all projects for the given `session_id`.
/// Returns `None` if the session is not found.
///
/// When `home_dir` is `Some`, searches that specific user's checkpoint storage.
pub async fn get_session_detail(
    _config_toml: &str,
    session_id: &str,
    home_dir: Option<&std::path::Path>,
) -> anyhow::Result<Option<(SessionSummary, Vec<CheckpointSummary>)>> {
    let manager = make_storage_manager(home_dir, "trustee")?;

    let projects = manager
        .list_projects()
        .await
        .map_err(|e| anyhow::anyhow!("Failed to list projects: {}", e))?;

    for project in &projects {
        let project_storage = match manager.get_project_storage(&project.project_path).await {
            Ok(ps) => ps,
            Err(e) => {
                tracing::debug!("Skipping project {}: {}", project.project_path.display(), e);
                continue;
            }
        };

        let sessions = match project_storage.list_sessions().await {
            Ok(s) => s,
            Err(e) => {
                tracing::debug!("Failed to list sessions for {}: {}", project.project_path.display(), e);
                continue;
            }
        };

        if let Some(session) = sessions.iter().find(|s| s.session_id == session_id) {
            let session_storage = project_storage
                .create_session(session_id)
                .await
                .map_err(|e| anyhow::anyhow!("Failed to get session storage: {}", e))?;

            let checkpoints = session_storage
                .list_checkpoints()
                .await
                .map_err(|e| anyhow::anyhow!("Failed to list checkpoints: {}", e))?;

            let summary = session_to_summary(session, project, false);
            let cp_summaries: Vec<CheckpointSummary> =
                checkpoints.iter().map(checkpoint_to_summary).collect();

            return Ok(Some((summary, cp_summaries)));
        }
    }

    Ok(None)
}

/// Create a `ResumeInfo` from the latest checkpoint of a given session.
///
/// Searches all projects for `session_id`, finds the most recent checkpoint,
/// and returns a `ResumeInfo` suitable for setting on `Session::resume_info`.
/// Returns `None` if the session or its checkpoints are not found.
///
/// When `home_dir` is `Some`, searches that specific user's checkpoint storage.
pub async fn create_resume_info(
    _config_toml: &str,
    session_id: &str,
    home_dir: Option<&std::path::Path>,
) -> anyhow::Result<Option<ResumeInfo>> {
    let manager = make_storage_manager(home_dir, "trustee")?;

    let projects = manager
        .list_projects()
        .await
        .map_err(|e| anyhow::anyhow!("Failed to list projects: {}", e))?;

    for project in &projects {
        let project_storage = match manager.get_project_storage(&project.project_path).await {
            Ok(ps) => ps,
            Err(e) => {
                tracing::debug!("Skipping project {}: {}", project.project_path.display(), e);
                continue;
            }
        };

        let sessions = match project_storage.list_sessions().await {
            Ok(s) => s,
            Err(e) => {
                tracing::debug!("Failed to list sessions for {}: {}", project.project_path.display(), e);
                continue;
            }
        };

        if sessions.iter().any(|s| s.session_id == session_id) {
            let session_storage = project_storage
                .create_session(session_id)
                .await
                .map_err(|e| anyhow::anyhow!("Failed to get session storage: {}", e))?;

            let checkpoints = session_storage
                .list_checkpoints()
                .await
                .map_err(|e| anyhow::anyhow!("Failed to list checkpoints: {}", e))?;

            if let Some(latest) = checkpoints.iter().max_by_key(|cp| cp.created_at) {
                return Ok(Some(ResumeInfo {
                    session_id: session_id.to_string(),
                    checkpoint_id: latest.checkpoint_id.clone(),
                    iteration: latest.iteration,
                    project_path: Some(project.project_path.clone()),
                }));
            }

            // Session found but no checkpoints
            return Ok(None);
        }
    }

    Ok(None)
}

// -----------------------------------------------------------------------
// Conversation history loading
// -----------------------------------------------------------------------

/// A single chat message rendered in the Web UI conversation history.
///
/// Each message maps to how the TUI/Web already renders things:
/// - `user` messages → right-aligned chat bubble
/// - `assistant` messages → left-aligned agent bubble (markdown)
/// - `assistant` with `tool_calls` → tool-pending/tool-done lines
/// - `tool` messages → hidden (tool results are shown via tool_calls)
/// - `reasoning` → collapsible reasoning section
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HistoryMessage {
    /// "user", "assistant", or "tool"
    pub role: String,
    /// Main message text
    pub content: String,
    /// Reasoning/thinking content (if any)
    pub reasoning: Option<String>,
    /// Tool calls (assistant messages that invoke tools)
    pub tool_calls: Option<Vec<HistoryToolCall>>,
    /// Tool name (for tool-role messages)
    pub name: Option<String>,
}

/// A tool call within a message.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HistoryToolCall {
    pub name: String,
    /// Short description of what the tool was called with (for display)
    pub hint: String,
}

/// Metadata about the session/task for the history header.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionHistory {
    pub session_id: String,
    pub checkpoint_id: String,
    pub task_description: String,
    pub iteration: u32,
    pub total_messages: usize,
    pub messages: Vec<HistoryMessage>,
}

/// Load conversation history from a session's latest checkpoint.
///
/// Returns messages suitable for rendering in the Web UI. System messages
/// are filtered out, tool results are summarized, and reasoning is
/// preserved separately.
///
/// When `home_dir` is `Some`, loads from that specific user's checkpoint storage.
pub async fn load_session_history(
    session_id: &str,
    home_dir: Option<&std::path::Path>,
) -> anyhow::Result<Option<SessionHistory>> {
    let manager = make_storage_manager(home_dir, "trustee")?;;

    let projects = manager
        .list_projects()
        .await
        .map_err(|e| anyhow::anyhow!("Failed to list projects: {}", e))?;

    for project in &projects {
        let project_storage = match manager.get_project_storage(&project.project_path).await {
            Ok(ps) => ps,
            Err(e) => {
                tracing::debug!("Skipping project {}: {}", project.project_path.display(), e);
                continue;
            }
        };

        let sessions = match project_storage.list_sessions().await {
            Ok(s) => s,
            Err(e) => {
                tracing::debug!("Failed to list sessions for {}: {}", project.project_path.display(), e);
                continue;
            }
        };

        if !sessions.iter().any(|s| s.session_id == session_id) {
            continue;
        }

        let session_storage = project_storage
            .create_session(session_id)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to get session storage: {}", e))?;

        let checkpoints = session_storage
            .list_checkpoints()
            .await
            .map_err(|e| anyhow::anyhow!("Failed to list checkpoints: {}", e))?;

        let latest = match checkpoints.iter().max_by_key(|cp| cp.created_at) {
            Some(cp) => cp,
            None => return Ok(None),
        };

        let checkpoint_id = latest.checkpoint_id.clone();
        let iteration = latest.iteration;

        let checkpoint = session_storage
            .load_checkpoint(&checkpoint_id)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to load checkpoint: {}", e))?;

        let task_description = checkpoint.agent_state.task_description.clone();
        let total_messages = checkpoint.conversation_state.messages.len();

        let messages = convert_messages(&checkpoint.conversation_state.messages);

        return Ok(Some(SessionHistory {
            session_id: session_id.to_string(),
            checkpoint_id,
            task_description,
            iteration,
            total_messages,
            messages,
        }));
    }

    Ok(None)
}

/// Convert ABK `ChatMessage`s to Web UI-friendly `HistoryMessage`s.
///
/// - System messages are dropped (not useful in UI)
/// - Tool-role messages are dropped (tool results shown via assistant's tool_calls)
/// - Very long content is truncated to avoid massive payloads
fn convert_messages(messages: &[ChatMessage]) -> Vec<HistoryMessage> {
    const MAX_CONTENT_LEN: usize = 10_000;

    let mut result = Vec::new();

    for msg in messages {
        // Skip system messages — not useful in the conversation view
        if msg.role == "system" {
            continue;
        }

        // Skip tool-role messages — tool results are shown via the
        // assistant's tool_calls and a compact tool-done line
        if msg.role == "tool" {
            continue;
        }

        let content = if msg.content.len() > MAX_CONTENT_LEN {
            format!("{}...\n[truncated]", &msg.content[..MAX_CONTENT_LEN])
        } else {
            msg.content.clone()
        };

        // Convert tool calls if present
        let tool_calls = msg.tool_calls.as_ref().map(|calls| {
            calls
                .iter()
                .map(|tc| {
                    let hint = summarize_tool_args(&tc.function.name, &tc.function.arguments);
                    HistoryToolCall {
                        name: tc.function.name.clone(),
                        hint,
                    }
                })
                .collect::<Vec<_>>()
        });

        // For assistant messages that only contain tool calls (empty content),
        // we still emit them so tool call lines render
        if msg.role == "assistant" && content.is_empty() && tool_calls.is_some() {
            result.push(HistoryMessage {
                role: "assistant".to_string(),
                content: String::new(),
                reasoning: msg.reasoning.clone(),
                tool_calls,
                name: None,
            });
            continue;
        }

        result.push(HistoryMessage {
            role: msg.role.clone(),
            content,
            reasoning: msg.reasoning.clone(),
            tool_calls,
            name: msg.name.clone(),
        });
    }

    result
}

/// Create a short human-readable hint from tool call arguments.
///
/// e.g. `{"command": "ls -la /tmp"}` → `ls -la /tmp`
///      `{"file_path": "/foo/bar.rs"}` → `/foo/bar.rs`
fn summarize_tool_args(name: &str, args: &str) -> String {
    let parsed: serde_json::Value = match serde_json::from_str(args) {
        Ok(v) => v,
        Err(_) => return args.chars().take(200).collect(),
    };

    let obj = match parsed.as_object() {
        Some(o) => o,
        None => return args.chars().take(200).collect(),
    };

    match name {
        "bash" | "execute_command" => {
            obj.get("command").and_then(|v| v.as_str()).map(|s| s.to_string())
        }
        "read" | "read_file" => {
            obj.get("file_path").and_then(|v| v.as_str()).map(|s| s.to_string())
        }
        "write" | "write_file" | "edit" => {
            obj.get("file_path").and_then(|v| v.as_str()).map(|s| s.to_string())
        }
        "grep" | "search" => {
            obj.get("pattern").and_then(|v| v.as_str()).map(|s| s.to_string())
        }
        "glob" => {
            obj.get("pattern").and_then(|v| v.as_str()).map(|s| s.to_string())
        }
        "todowrite" => {
            let count = obj.get("todos").and_then(|v| v.as_array()).map(|a| a.len()).unwrap_or(0);
            Some(format!("{} items", count))
        }
        "websearch" | "webfetch" => {
            obj.get("query").or_else(|| obj.get("url")).and_then(|v| v.as_str()).map(|s| s.to_string())
        }
        _ => {
            obj.values().next().and_then(|v| match v {
                serde_json::Value::String(s) => Some(s.clone()),
                _ => Some(v.to_string()),
            })
        }
    }
    .unwrap_or_default()
    .chars()
    .take(200)
    .collect()
}