agentmux_srv\backend\blockcontroller/
session_stats.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Session metadata tracking for agent panes (Phase 1.4 — ultra-long-sessions).
5//!
6//! Tracks per-session stats as block metadata keys:
7//!   `session:start_ts_ms`    — Unix ms when the first output line arrived
8//!   `session:last_activity_ms` — Unix ms of most recent output line
9//!   `session:line_count`     — total output lines emitted this session
10//!   `session:token_estimate` — rough token count (chars / 4, cumulative)
11//!
12//! To avoid a `SetMeta` write on every output line (which can be very frequent),
13//! this module debounces flushes to at most once per second using a local
14//! `Instant`-based timestamp.  Accumulators live in `SessionStatsAccumulator`
15//! which each controller instance owns privately.
16
17use std::sync::Arc;
18use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
19
20use crate::backend::obj::MetaMapType;
21use crate::backend::storage::store::Store;
22
23/// Keys used for session stats in block metadata.
24pub const META_SESSION_START_TS_MS: &str = "session:start_ts_ms";
25pub const META_SESSION_LAST_ACTIVITY_MS: &str = "session:last_activity_ms";
26pub const META_SESSION_LINE_COUNT: &str = "session:line_count";
27pub const META_SESSION_TOKEN_ESTIMATE: &str = "session:token_estimate";
28
29/// Debounce interval: at most one Store write per second.
30const FLUSH_DEBOUNCE: Duration = Duration::from_secs(1);
31
32/// Returns the current Unix timestamp in milliseconds.
33fn now_ms() -> i64 {
34    SystemTime::now()
35        .duration_since(UNIX_EPOCH)
36        .unwrap_or_default()
37        .as_millis() as i64
38}
39
40/// In-memory accumulator for session stats.  One per controller instance.
41///
42/// All fields are plain integers — no locking needed because each controller
43/// calls `record_line` from its single async stdout-reader task only.
44pub struct SessionStatsAccumulator {
45    block_id: String,
46    /// Unix ms when the first line was seen; 0 = not yet set.
47    start_ts_ms: i64,
48    /// Unix ms of the most-recently flushed line.
49    last_activity_ms: i64,
50    /// Total lines seen since session start.
51    line_count: u64,
52    /// Cumulative token estimate (chars / 4).
53    token_estimate: u64,
54    /// Wall-clock instant of the last flush; `None` = never flushed.
55    last_flush: Option<Instant>,
56}
57
58impl SessionStatsAccumulator {
59    /// Create a new accumulator for `block_id`.
60    pub fn new(block_id: String) -> Self {
61        Self {
62            block_id,
63            start_ts_ms: 0,
64            last_activity_ms: 0,
65            line_count: 0,
66            token_estimate: 0,
67            last_flush: None,
68        }
69    }
70
71    /// Record one output line of `line_len` bytes.
72    ///
73    /// Updates in-memory counters.  Flushes to the Store if the debounce
74    /// interval has elapsed *or* if this is the very first line (so the
75    /// frontend sees `session:start_ts_ms` promptly).
76    pub fn record_line(&mut self, line_len: usize, wstore: &Option<Arc<Store>>) {
77        let ts = now_ms();
78        let is_first = self.start_ts_ms == 0;
79
80        if is_first {
81            self.start_ts_ms = ts;
82        }
83        self.last_activity_ms = ts;
84        self.line_count += 1;
85        self.token_estimate += (line_len / 4) as u64;
86
87        // Flush immediately on first line; otherwise debounce.
88        let should_flush = is_first || match self.last_flush {
89            None => true,
90            Some(last) => last.elapsed() >= FLUSH_DEBOUNCE,
91        };
92
93        if should_flush {
94            if let Some(ref store) = wstore {
95                self.flush(store);
96            }
97        }
98    }
99
100    /// Force-flush all accumulated stats to the Store right now.
101    ///
102    /// Called by `record_line` when the debounce window has elapsed.
103    fn flush(&mut self, wstore: &Arc<Store>) {
104        let oref_str = format!("block:{}", self.block_id);
105        let mut meta_update = MetaMapType::new();
106
107        if self.start_ts_ms != 0 {
108            meta_update.insert(
109                META_SESSION_START_TS_MS.to_string(),
110                serde_json::json!(self.start_ts_ms),
111            );
112        }
113        meta_update.insert(
114            META_SESSION_LAST_ACTIVITY_MS.to_string(),
115            serde_json::json!(self.last_activity_ms),
116        );
117        meta_update.insert(
118            META_SESSION_LINE_COUNT.to_string(),
119            serde_json::json!(self.line_count),
120        );
121        meta_update.insert(
122            META_SESSION_TOKEN_ESTIMATE.to_string(),
123            serde_json::json!(self.token_estimate),
124        );
125
126        match crate::server::service::update_object_meta(wstore, &oref_str, &meta_update) {
127            Ok(()) => {
128                tracing::trace!(
129                    block_id = %self.block_id,
130                    line_count = self.line_count,
131                    token_estimate = self.token_estimate,
132                    "session stats flushed"
133                );
134            }
135            Err(e) => {
136                tracing::warn!(
137                    block_id = %self.block_id,
138                    error = %e,
139                    "failed to flush session stats"
140                );
141            }
142        }
143
144        self.last_flush = Some(Instant::now());
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn test_accumulator_first_line_sets_start_ts() {
154        let mut acc = SessionStatsAccumulator::new("blk-1".to_string());
155        // No wstore — flush is skipped but counters still update.
156        acc.record_line(100, &None);
157        assert_ne!(acc.start_ts_ms, 0);
158        assert_eq!(acc.line_count, 1);
159        assert_eq!(acc.token_estimate, 25); // 100 / 4
160    }
161
162    #[test]
163    fn test_accumulator_multiple_lines() {
164        let mut acc = SessionStatsAccumulator::new("blk-2".to_string());
165        acc.record_line(40, &None);
166        acc.record_line(80, &None);
167        acc.record_line(120, &None);
168        assert_eq!(acc.line_count, 3);
169        // 40/4 + 80/4 + 120/4 = 10 + 20 + 30 = 60
170        assert_eq!(acc.token_estimate, 60);
171    }
172
173    #[test]
174    fn test_accumulator_start_ts_not_reset_on_second_line() {
175        let mut acc = SessionStatsAccumulator::new("blk-3".to_string());
176        acc.record_line(10, &None);
177        let first_ts = acc.start_ts_ms;
178        acc.record_line(10, &None);
179        assert_eq!(acc.start_ts_ms, first_ts, "start_ts must not change after first line");
180    }
181
182    #[test]
183    fn test_debounce_constants() {
184        assert_eq!(FLUSH_DEBOUNCE, Duration::from_secs(1));
185    }
186}