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(
24        &self,
25        project_hash: Option<&agtrace_types::ProjectHash>,
26        limit: usize,
27    ) -> Result<PackResult> {
28        let raw_sessions = self.db.list_sessions(project_hash, 1000)?;
29        let balanced_sessions = balance_sessions_by_provider(&raw_sessions, 200);
30
31        let mut digests = Vec::new();
32        let loader = SessionRepository::new(self.db);
33        let options = LoadOptions::default();
34
35        for (i, session) in balanced_sessions.iter().enumerate() {
36            if let Ok(events) = loader.load_events(&session.id, &options)
37                && let Some(agent_session) = assemble_session(&events)
38            {
39                let recency_boost = (balanced_sessions.len() - i) as u32;
40                let digest = SessionDigest::new(
41                    &session.id,
42                    &session.provider,
43                    agent_session,
44                    recency_boost,
45                );
46                digests.push(digest);
47            }
48        }
49
50        let selections = analyze_and_select_sessions(digests, limit);
51
52        Ok(PackResult {
53            selections,
54            balanced_count: balanced_sessions.len(),
55            raw_count: raw_sessions.len(),
56        })
57    }
58}
59
60fn balance_sessions_by_provider(
61    sessions: &[SessionSummary],
62    target_per_provider: usize,
63) -> Vec<SessionSummary> {
64    let mut by_provider: HashMap<String, Vec<SessionSummary>> = HashMap::new();
65    for session in sessions {
66        by_provider
67            .entry(session.provider.clone())
68            .or_default()
69            .push(session.clone());
70    }
71
72    let mut balanced = Vec::new();
73    for (_, mut list) in by_provider {
74        if list.len() > target_per_provider {
75            list.truncate(target_per_provider);
76        }
77        balanced.extend(list);
78    }
79
80    balanced.sort_by(|a, b| b.start_ts.cmp(&a.start_ts));
81    balanced
82}
agtrace_runtime/ops/pack.rs

agtrace_runtime/ops/
pack.rs
