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};
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.
33use agentmux_common::time::now_ms;
34
35/// In-memory accumulator for session stats.  One per controller instance.
36///
37/// All fields are plain integers — no locking needed because each controller
38/// calls `record_line` from its single async stdout-reader task only.
39pub struct SessionStatsAccumulator {
40    block_id: String,
41    /// Unix ms when the first line was seen; 0 = not yet set.
42    start_ts_ms: i64,
43    /// Unix ms of the most-recently flushed line.
44    last_activity_ms: i64,
45    /// Total lines seen since session start.
46    line_count: u64,
47    /// Cumulative token estimate (chars / 4).
48    token_estimate: u64,
49    /// Wall-clock instant of the last flush; `None` = never flushed.
50    last_flush: Option<Instant>,
51}
52
53impl SessionStatsAccumulator {
54    /// Create a new accumulator for `block_id`.
55    pub fn new(block_id: String) -> Self {
56        Self {
57            block_id,
58            start_ts_ms: 0,
59            last_activity_ms: 0,
60            line_count: 0,
61            token_estimate: 0,
62            last_flush: None,
63        }
64    }
65
66    /// Record one output line of `line_len` bytes.
67    ///
68    /// Updates in-memory counters.  Flushes to the Store if the debounce
69    /// interval has elapsed *or* if this is the very first line (so the
70    /// frontend sees `session:start_ts_ms` promptly).
71    pub fn record_line(&mut self, line_len: usize, mstore: &Option<Arc<Store>>) {
72        let ts = now_ms();
73        let is_first = self.start_ts_ms == 0;
74
75        if is_first {
76            self.start_ts_ms = ts;
77        }
78        self.last_activity_ms = ts;
79        self.line_count += 1;
80        self.token_estimate += (line_len / 4) as u64;
81
82        // Flush immediately on first line; otherwise debounce.
83        let should_flush = is_first || match self.last_flush {
84            None => true,
85            Some(last) => last.elapsed() >= FLUSH_DEBOUNCE,
86        };
87
88        if should_flush {
89            if let Some(ref store) = mstore {
90                self.flush(store);
91            }
92        }
93    }
94
95    /// Force-flush all accumulated stats to the Store right now.
96    ///
97    /// Called by `record_line` when the debounce window has elapsed.
98    fn flush(&mut self, mstore: &Arc<Store>) {
99        let oref_str = format!("block:{}", self.block_id);
100        let mut meta_update = MetaMapType::new();
101
102        if self.start_ts_ms != 0 {
103            meta_update.insert(
104                META_SESSION_START_TS_MS.to_string(),
105                serde_json::json!(self.start_ts_ms),
106            );
107        }
108        meta_update.insert(
109            META_SESSION_LAST_ACTIVITY_MS.to_string(),
110            serde_json::json!(self.last_activity_ms),
111        );
112        meta_update.insert(
113            META_SESSION_LINE_COUNT.to_string(),
114            serde_json::json!(self.line_count),
115        );
116        meta_update.insert(
117            META_SESSION_TOKEN_ESTIMATE.to_string(),
118            serde_json::json!(self.token_estimate),
119        );
120
121        match crate::server::service::update_object_meta(mstore, &oref_str, &meta_update) {
122            Ok(()) => {
123                tracing::trace!(
124                    block_id = %self.block_id,
125                    line_count = self.line_count,
126                    token_estimate = self.token_estimate,
127                    "session stats flushed"
128                );
129            }
130            Err(e) => {
131                tracing::warn!(
132                    block_id = %self.block_id,
133                    error = %e,
134                    "failed to flush session stats"
135                );
136            }
137        }
138
139        self.last_flush = Some(Instant::now());
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn test_accumulator_first_line_sets_start_ts() {
149        let mut acc = SessionStatsAccumulator::new("blk-1".to_string());
150        // No mstore — flush is skipped but counters still update.
151        acc.record_line(100, &None);
152        assert_ne!(acc.start_ts_ms, 0);
153        assert_eq!(acc.line_count, 1);
154        assert_eq!(acc.token_estimate, 25); // 100 / 4
155    }
156
157    #[test]
158    fn test_accumulator_multiple_lines() {
159        let mut acc = SessionStatsAccumulator::new("blk-2".to_string());
160        acc.record_line(40, &None);
161        acc.record_line(80, &None);
162        acc.record_line(120, &None);
163        assert_eq!(acc.line_count, 3);
164        // 40/4 + 80/4 + 120/4 = 10 + 20 + 30 = 60
165        assert_eq!(acc.token_estimate, 60);
166    }
167
168    #[test]
169    fn test_accumulator_start_ts_not_reset_on_second_line() {
170        let mut acc = SessionStatsAccumulator::new("blk-3".to_string());
171        acc.record_line(10, &None);
172        let first_ts = acc.start_ts_ms;
173        acc.record_line(10, &None);
174        assert_eq!(acc.start_ts_ms, first_ts, "start_ts must not change after first line");
175    }
176
177    #[test]
178    fn test_debounce_constants() {
179        assert_eq!(FLUSH_DEBOUNCE, Duration::from_secs(1));
180    }
181}
agentmux_srv/backend/blockcontroller/session_stats.rs

agentmux_srv/backend/blockcontroller/
session_stats.rs
