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

agtrace_runtime/ops/
pack.rs
