agtrace_runtime/ops/
pack.rs

1use agtrace_engine::{SessionDigest, analyze_and_select_sessions, assemble_session};
2use agtrace_index::{Database, SessionSummary};
3use anyhow::Result;
4use std::collections::HashMap;
5
6use crate::storage::{LoadOptions, SessionRepository};
7
8pub struct PackResult {
9    pub selections: Vec<SessionDigest>,
10    pub balanced_count: usize,
11    pub raw_count: usize,
12}
13
14pub struct PackService<'a> {
15    db: &'a Database,
16}
17
18impl<'a> PackService<'a> {
19    pub fn new(db: &'a Database) -> Self {
20        Self { db }
21    }
22
23    pub fn select_sessions(&self, project_hash: Option<&str>, limit: usize) -> Result<PackResult> {
24        let raw_sessions = self.db.list_sessions(project_hash, 1000)?;
25        let balanced_sessions = balance_sessions_by_provider(&raw_sessions, 200);
26
27        let mut digests = Vec::new();
28        let loader = SessionRepository::new(self.db);
29        let options = LoadOptions::default();
30
31        for (i, session) in balanced_sessions.iter().enumerate() {
32            if let Ok(events) = loader.load_events(&session.id, &options)
33                && let Some(agent_session) = assemble_session(&events)
34            {
35                let recency_boost = (balanced_sessions.len() - i) as u32;
36                let digest = SessionDigest::new(
37                    &session.id,
38                    &session.provider,
39                    agent_session,
40                    recency_boost,
41                );
42                digests.push(digest);
43            }
44        }
45
46        let selections = analyze_and_select_sessions(digests, limit);
47
48        Ok(PackResult {
49            selections,
50            balanced_count: balanced_sessions.len(),
51            raw_count: raw_sessions.len(),
52        })
53    }
54}
55
56fn balance_sessions_by_provider(
57    sessions: &[SessionSummary],
58    target_per_provider: usize,
59) -> Vec<SessionSummary> {
60    let mut by_provider: HashMap<String, Vec<SessionSummary>> = HashMap::new();
61    for session in sessions {
62        by_provider
63            .entry(session.provider.clone())
64            .or_default()
65            .push(session.clone());
66    }
67
68    let mut balanced = Vec::new();
69    for (_, mut list) in by_provider {
70        if list.len() > target_per_provider {
71            list.truncate(target_per_provider);
72        }
73        balanced.extend(list);
74    }
75
76    balanced.sort_by(|a, b| b.start_ts.cmp(&a.start_ts));
77    balanced
78}
agtrace_runtime/ops/pack.rs

agtrace_runtime/ops/
pack.rs
