agtrace_runtime/ops/
pack.rs

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

agtrace_runtime/ops/
pack.rs
