From 94313d92894bc878b980f680ff58fef4257959a4 Mon Sep 17 00:00:00 2001 From: ADD-SP Date: Sat, 11 Apr 2026 18:42:49 -0700 Subject: [PATCH] feat: add `/admin/stats` endpoint with persistent runtime counters Exposes GET /admin/stats returning running totals for requests served, upstream prompt/completion/total tokens, emails hidden by the sender policy, and a per-address breakdown of filtered senders. --- README.md | 13 +- clawshell.example.toml | 7 + src/app.rs | 128 ++++-- src/app/tests.rs | 19 + src/config.rs | 38 +- src/main.rs | 49 ++- src/onboard/config_render.rs | 5 + src/stats.rs | 400 ++++++++++++++++++ .../config/invalid/dlp_invalid_regex.toml | 3 + .../config/invalid/key_missing_real_key.toml | 3 + .../config/invalid/missing_stats.toml | 2 + tests/fixtures/config/valid/all_fields.toml | 3 + .../config/valid/dlp_disabled_scan.toml | 3 + .../fixtures/config/valid/empty_base_url.toml | 3 + tests/fixtures/config/valid/empty_host.toml | 3 + tests/fixtures/config/valid/empty_keys.toml | 3 + tests/fixtures/config/valid/minimal.toml | 3 + tests/fixtures/config/valid/port_max.toml | 3 + tests/fixtures/config/valid/port_zero.toml | 3 + .../config_fixtures__all_fields.snap | 2 + .../config_fixtures__dlp_disabled_scan.snap | 4 +- .../config_fixtures__empty_base_url.snap | 4 +- .../config_fixtures__empty_host.snap | 4 +- .../config_fixtures__empty_keys.snap | 4 +- tests/snapshots/config_fixtures__minimal.snap | 4 +- .../config_fixtures__missing_stats.snap | 9 + .../snapshots/config_fixtures__port_max.snap | 4 +- .../snapshots/config_fixtures__port_zero.snap | 4 +- 28 files changed, 668 insertions(+), 62 deletions(-) create mode 100644 src/stats.rs create mode 100644 tests/fixtures/config/invalid/missing_stats.toml create mode 100644 tests/snapshots/config_fixtures__missing_stats.snap diff --git a/README.md b/README.md index bd99352..c6b1e4e 100644 --- a/README.md +++ b/README.md @@ -49,12 +49,21 @@ ClawShell supports OAuth-based authentication as an alternative to static API ke - **Automatic Token Refresh**: Access tokens are refreshed transparently before they expire. - **Request Translation**: Automatically translates OpenAI Chat Completions API requests to the ChatGPT Responses API format when using Codex OAuth. -### 5. Seamless Integration +### 5. Runtime Statistics + +ClawShell exposes running counters at `GET /admin/stats` so operators can audit proxy activity since startup and across restarts. + +- **What's Counted**: Total requests served, total upstream `prompt_tokens` / `completion_tokens` / `total_tokens` (parsed from non-streaming JSON responses — SSE streams are not counted), number of emails hidden by the sender policy, and a per-address breakdown of filtered senders. +- **Loopback-Only**: The endpoint is reachable without a virtual key but only from `127.0.0.1` / `::1` peers; non-loopback clients receive `403`. +- **Persistent**: Counters are flushed to disk every 30 seconds and on graceful shutdown. The location is a required config field — set `[stats] persist_path = "..."` in `clawshell.toml` (typically `/var/lib/clawshell/stats.json` under the hardened systemd unit, since `/etc/clawshell` is read-only there). +- **Bounded**: The filtered-address map is capped at 10,000 unique entries; further unique addresses are aggregated under an `` bucket so memory stays bounded. + +### 6. Seamless Integration - **Drop-in Sidecar**: The `clawshell onboard` wizard configures exactly one downstream LLM client per run — either OpenClaw or [Hermes Agent](https://github.com/NousResearch/hermes-agent) — to route all requests through ClawShell's proxy. See [Agent Target (pick one)](#agent-target-pick-one). - **No External Dependencies**: Uses Unix file system permissions to protect secrets. No IdP, Vault, or external key management service required. -### 6. Ultra Lightweight and Scalable +### 7. Ultra Lightweight and Scalable - Runs in under 10MB of memory. - Written in Rust with Tokio. diff --git a/clawshell.example.toml b/clawshell.example.toml index 0cccb3e..50947ee 100644 --- a/clawshell.example.toml +++ b/clawshell.example.toml @@ -53,6 +53,13 @@ patterns = [ { name = "amex_card", regex = '\b3[47][0-9]{13}\b', action = "redact" }, ] +# Runtime statistics persistence +# ClawShell counts total requests served, upstream prompt/completion/total +# tokens (from non-streaming responses), and per-sender email-filter +# activity, and exposes them at GET /admin/stats (loopback-only). +[stats] +persist_path = "/etc/clawshell/stats.json" + # Email read endpoint # If enabled, set exactly one mode: # - mode = "allowlist" with non-empty allow_senders and empty deny_senders diff --git a/src/app.rs b/src/app.rs index 53373ad..496f393 100644 --- a/src/app.rs +++ b/src/app.rs @@ -8,10 +8,11 @@ use crate::email::{ use crate::keys::{KeyManager, KeySource, ResolvedKey}; use crate::oauth::OAuthRegistry; use crate::proxy::ProxyClient; +use crate::stats::Stats; use axum::Router; use axum::body::Body; -use axum::extract::{DefaultBodyLimit, Path, Query, Request, State}; +use axum::extract::{ConnectInfo, DefaultBodyLimit, Path, Query, Request, State}; use axum::http::{HeaderMap, Method, StatusCode, Uri}; use axum::middleware::Next; use axum::response::{IntoResponse, Response}; @@ -20,6 +21,7 @@ use bytes::Bytes; use http_body_util::BodyExt; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; +use std::net::SocketAddr; use std::sync::Arc; use std::time::Instant; use tracing::{debug, error, info, trace, warn}; @@ -34,6 +36,7 @@ pub struct AppState { pub email_policy: Option, pub email_accounts: Arc>, pub email_service: Arc, + pub stats: Arc, } impl AppState { @@ -138,6 +141,7 @@ impl AppState { email_policy, email_accounts: Arc::new(email_accounts), email_service: Arc::new(EmailService::Imap(ImapEmailService::default())), + stats: Arc::new(Stats::new(Some(config.stats.persist_path.clone()))), }) } } @@ -149,20 +153,30 @@ pub fn build_router(state: AppState) -> Router { Router::new() .route("/v1/email/messages", get(handle_email_secure_messages)) .route("/v1/email/messages/{id}", get(handle_email_message_content)) + .route("/admin/stats", get(handle_stats)) .route("/", any(handle_request)) .route("/{*path}", any(handle_request)) .layer(DefaultBodyLimit::max(MAX_BODY_SIZE)) - .layer(axum::middleware::from_fn(log_request_completion)) + .layer(axum::middleware::from_fn_with_state( + state.clone(), + log_request_completion, + )) .with_state(state) } -async fn log_request_completion(request: Request, next: Next) -> Response { +async fn log_request_completion( + State(state): State, + request: Request, + next: Next, +) -> Response { let start = Instant::now(); let method = request.method().clone(); let path = request.uri().path().to_string(); let query = request.uri().query().map(|q| q.to_string()); let response = next.run(request).await; + state.stats.record_request(); + info!( method = %method, path = %path, @@ -317,6 +331,7 @@ async fn handle_email_secure_messages( continue; }; if !policy.sender_visible(from_header) { + state.stats.record_email_filtered(from_header); continue; } visible_messages.push(EmailSecureMessage { @@ -447,6 +462,9 @@ async fn handle_email_message_content( .as_deref() .or_else(|| content.headers.get("from").map(String::as_str)); if from_header.is_none_or(|from| !policy.sender_visible(from)) { + if let Some(from) = from_header { + state.stats.record_email_filtered(from); + } return Err(error_response( StatusCode::NOT_FOUND, "Email message not found", @@ -469,6 +487,27 @@ async fn handle_email_message_content( Ok(axum::Json(response).into_response()) } +/// Management endpoint that returns running counters for total requests, +/// upstream token usage, and per-sender email-filter activity. Only +/// reachable from loopback peers; streaming (SSE) responses are not +/// included in token totals (see `handle_request`). +async fn handle_stats( + State(state): State, + ConnectInfo(peer): ConnectInfo, +) -> Result { + if !peer.ip().is_loopback() { + warn!( + peer = %peer, + "Non-loopback client tried to hit /admin/stats" + ); + return Err(error_response( + StatusCode::FORBIDDEN, + "stats endpoint is loopback-only", + )); + } + Ok(axum::Json(state.stats.snapshot()).into_response()) +} + async fn handle_request( State(state): State, request: Request, @@ -630,31 +669,52 @@ async fn handle_request( })?, }; - // 5. DLP scan on response body (redact all PII before returning to client) - let response = if state.dlp_scanner.scan_responses() { - trace!("Response DLP scanning enabled, checking response body"); - let is_streaming = response - .headers() - .get("content-type") - .and_then(|v| v.to_str().ok()) - .is_some_and(|ct| ct.contains("text/event-stream")); + // 5. Response processing. + // Non-streaming responses are buffered so we can (a) record upstream + // token usage for stats and (b) optionally DLP-redact before sending. + // SSE streams are handled by the existing DLP SSE wrapper; we do not + // count tokens from SSE usage events. + let is_streaming = response + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()) + .is_some_and(|ct| ct.contains("text/event-stream")); - if !is_streaming { - trace!("Non-streaming response, scanning body for PII"); + let response = if is_streaming { + if state.dlp_scanner.scan_responses() { + debug!( + method = %method, + path = %path, + virtual_key = %virtual_key, + "Streaming response (SSE) — wrapping with DLP SSE scanner" + ); let (parts, body) = response.into_parts(); - let body = body - .collect() - .await - .map_err(|e| { - error!(error = %e, "Failed to read response body for DLP scan"); - error_response( - StatusCode::INTERNAL_SERVER_ERROR, - "Failed to process response", - ) - })? - .to_bytes(); - - let (redacted, redacted_names) = state.dlp_scanner.redact_all(&body); + let dlp_body = + crate::translate::wrap_body_with_dlp_sse_stream(body, state.dlp_scanner.clone()); + Response::from_parts(parts, dlp_body) + } else { + trace!("Streaming response, DLP response scanning disabled"); + response + } + } else { + let (parts, body) = response.into_parts(); + let body_bytes = body + .collect() + .await + .map_err(|e| { + error!(error = %e, "Failed to read response body"); + error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "Failed to process response", + ) + })? + .to_bytes(); + + // Count tokens from the upstream `usage` block before any DLP mutation. + state.stats.record_tokens_from_usage(&body_bytes); + + if state.dlp_scanner.scan_responses() { + let (redacted, redacted_names) = state.dlp_scanner.redact_all(&body_bytes); if !redacted_names.is_empty() { warn!( method = %method, @@ -669,23 +729,11 @@ async fn handle_request( parts.headers.remove("content-length"); Response::from_parts(parts, Body::from(redacted_bytes)) } else { - Response::from_parts(parts, Body::from(body)) + Response::from_parts(parts, Body::from(body_bytes)) } } else { - debug!( - method = %method, - path = %path, - virtual_key = %virtual_key, - "Streaming response (SSE) — wrapping with DLP SSE scanner" - ); - let (parts, body) = response.into_parts(); - let dlp_body = - crate::translate::wrap_body_with_dlp_sse_stream(body, state.dlp_scanner.clone()); - Response::from_parts(parts, dlp_body) + Response::from_parts(parts, Body::from(body_bytes)) } - } else { - trace!("Response DLP scanning disabled"); - response }; Ok(response) diff --git a/src/app/tests.rs b/src/app/tests.rs index 2a9fa26..e3520a1 100644 --- a/src/app/tests.rs +++ b/src/app/tests.rs @@ -19,6 +19,7 @@ use crate::email::{ use crate::keys::{KeyManager, KeySource, ResolvedKey}; use crate::oauth::OAuthRegistry; use crate::proxy::ProxyClient; +use crate::stats::Stats; fn make_app(upstream_url: &str) -> axum::Router { let mut key_map = BTreeMap::new(); @@ -75,6 +76,7 @@ fn make_app(upstream_url: &str) -> axum::Router { email_policy: None, email_accounts: Arc::new(BTreeMap::new()), email_service: Arc::new(EmailService::mock_disabled()), + stats: Arc::new(Stats::new(None)), }; build_router(state) @@ -117,6 +119,7 @@ fn make_app_with_anthropic(upstream_url: &str) -> axum::Router { email_policy: None, email_accounts: Arc::new(BTreeMap::new()), email_service: Arc::new(EmailService::mock_disabled()), + stats: Arc::new(Stats::new(None)), }; build_router(state) @@ -568,6 +571,8 @@ openai_base_url = "https://api.openai.com" [[keys]] virtual_key = "vk-1" real_key = "sk-real-1" +[stats] +persist_path = "/etc/clawshell/stats.json" "#; let config = Config::parse(toml_str).unwrap(); let state = AppState::from_config(&config).unwrap(); @@ -597,6 +602,8 @@ provider = "openai" virtual_key = "vk-ant" real_key = "sk-ant-key" provider = "anthropic" +[stats] +persist_path = "/etc/clawshell/stats.json" "#; let config = Config::parse(toml_str).unwrap(); let state = AppState::from_config(&config).unwrap(); @@ -635,6 +642,9 @@ email = "bot@gmail.com" app_password = "abcd efgh ijkl mnop" imap_host = "imap.gmail.com" imap_port = 993 + +[stats] +persist_path = "/etc/clawshell/stats.json" "#; let config = Config::parse(toml_str).unwrap(); let state = AppState::from_config(&config).unwrap(); @@ -665,6 +675,9 @@ allow_senders = ["alice@example.com"] virtual_key = "vk-email" email = "bot@gmail.com" app_password = "abcd efgh ijkl mnop" + +[stats] +persist_path = "/etc/clawshell/stats.json" "#; let config = Config::parse(toml_str).unwrap(); let state = AppState::from_config(&config).unwrap(); @@ -708,6 +721,7 @@ async fn test_proxy_error_on_unreachable_upstream() { email_policy: None, email_accounts: Arc::new(BTreeMap::new()), email_service: Arc::new(EmailService::mock_disabled()), + stats: Arc::new(Stats::new(None)), }; let app = build_router(state); @@ -855,6 +869,7 @@ async fn test_anthropic_dlp_blocks_sensitive_data() { email_policy: None, email_accounts: Arc::new(BTreeMap::new()), email_service: Arc::new(EmailService::mock_disabled()), + stats: Arc::new(Stats::new(None)), }; let app = build_router(state); @@ -957,6 +972,7 @@ async fn test_openai_and_openrouter_keys_map_to_distinct_real_keys() { email_policy: None, email_accounts: Arc::new(BTreeMap::new()), email_service: Arc::new(EmailService::mock_disabled()), + stats: Arc::new(Stats::new(None)), }); let openai_req = Request::builder() @@ -1028,6 +1044,7 @@ fn make_app_with_redact(upstream_url: &str) -> axum::Router { email_policy: None, email_accounts: Arc::new(BTreeMap::new()), email_service: Arc::new(EmailService::mock_disabled()), + stats: Arc::new(Stats::new(None)), }; build_router(state) @@ -1235,6 +1252,7 @@ async fn test_response_dlp_disabled() { email_policy: None, email_accounts: Arc::new(BTreeMap::new()), email_service: Arc::new(EmailService::mock_disabled()), + stats: Arc::new(Stats::new(None)), }; let app = build_router(state); @@ -1598,6 +1616,7 @@ fn make_email_app( email_policy: Some(policy), email_accounts: Arc::new(email_accounts), email_service, + stats: Arc::new(Stats::new(None)), }; build_router(state) diff --git a/src/config.rs b/src/config.rs index a303869..192fe8a 100644 --- a/src/config.rs +++ b/src/config.rs @@ -2,7 +2,7 @@ use regex::Regex; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; use std::env::VarError; -use std::path::Path; +use std::path::{Path, PathBuf}; #[derive(Debug, Default, Deserialize, Serialize, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] #[serde(rename_all = "lowercase")] @@ -38,6 +38,7 @@ pub struct Config { pub dlp: DlpConfig, #[serde(default, skip_serializing_if = "EmailConfig::is_default")] pub email: EmailConfig, + pub stats: StatsConfig, #[serde(default = "default_log_level")] pub log_level: String, #[serde(default, skip_serializing_if = "Vec::is_empty")] @@ -150,6 +151,14 @@ impl Default for DlpConfig { } } +#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct StatsConfig { + /// Where the running stats counters are persisted on disk so they + /// survive restarts. + pub persist_path: PathBuf, +} + #[derive(Debug, Deserialize, Serialize, Clone)] #[serde(deny_unknown_fields)] pub struct DlpPattern { @@ -718,6 +727,9 @@ allow_senders = ["alice@example.com", "@trusted.org"] virtual_key = "vk-email" email = "bot@gmail.com" app_password = "abcd efgh ijkl mnop" + +[stats] +persist_path = "/etc/clawshell/stats.json" "#; let parsed = Config::parse(cfg); assert!(parsed.is_ok()); @@ -743,6 +755,9 @@ deny_senders = ["bob@example.com"] virtual_key = "vk-email" email = "bot@gmail.com" app_password = "abcd efgh ijkl mnop" + +[stats] +persist_path = "/etc/clawshell/stats.json" "#; let err = Config::parse(cfg).unwrap_err(); assert!( @@ -769,6 +784,9 @@ mode = "denylist" virtual_key = "vk-email" email = "bot@gmail.com" app_password = "abcd efgh ijkl mnop" + +[stats] +persist_path = "/etc/clawshell/stats.json" "#; let err = Config::parse(cfg).unwrap_err(); assert!( @@ -796,6 +814,9 @@ allow_senders = ["not-an-email"] virtual_key = "vk-email" email = "bot@gmail.com" app_password = "abcd efgh ijkl mnop" + +[stats] +persist_path = "/etc/clawshell/stats.json" "#; let err = Config::parse(cfg).unwrap_err(); assert!(err.to_string().contains("invalid allow_senders entry")); @@ -821,6 +842,9 @@ default_max_results = 0 virtual_key = "vk-email" email = "bot@gmail.com" app_password = "abcd efgh ijkl mnop" + +[stats] +persist_path = "/etc/clawshell/stats.json" "#; let err = Config::parse(cfg).unwrap_err(); assert!( @@ -850,6 +874,9 @@ email = "bot@gmail.com" app_password = "abcd efgh ijkl mnop" imap_host = "imap.gmail.com" imap_port = 993 + +[stats] +persist_path = "/etc/clawshell/stats.json" "#; let parsed = Config::parse(cfg); assert!(parsed.is_ok()); @@ -897,6 +924,9 @@ allow_senders = ["alice@example.com"] virtual_key = "vk-email" email = "botgmail.com" app_password = "abcd efgh ijkl mnop" + +[stats] +persist_path = "/etc/clawshell/stats.json" "#; let err = Config::parse(cfg).unwrap_err(); assert!( @@ -998,6 +1028,9 @@ virtual_key = "vk-email" email = "bot@gmail.com" app_password = "abcd efgh ijkl mnop" imap_port = 0 + +[stats] +persist_path = "/etc/clawshell/stats.json" "#; let err = Config::parse(cfg).unwrap_err(); assert!( @@ -1015,6 +1048,9 @@ port = 3000 [upstream] openai_base_url = "https://api.openai.com" + +[stats] +persist_path = "/etc/clawshell/stats.json" "#; let parsed = Config::parse(cfg).expect("config should parse"); assert_eq!(parsed.listen_addr(), "127.0.0.1:3000"); diff --git a/src/main.rs b/src/main.rs index 3afec76..060ea5d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -17,6 +17,7 @@ mod openclaw_cli; mod platform; mod process; mod proxy; +mod stats; mod translate; mod tui; @@ -676,27 +677,51 @@ async fn cmd_start_inner(config_path: &str) -> Result<(), Box break, + _ = tokio::time::sleep(interval) => { + if let Err(err) = stats_for_task.persist() { + warn!(error = %err, "failed to persist stats"); + } + } + } + } + }); + let addr: SocketAddr = listen_addr.parse()?; + let stats_for_shutdown = app_state.stats.clone(); let app = build_router(app_state); let listener = tokio::net::TcpListener::bind(addr).await?; info!("Listening on {}", addr); process::drop_privileges()?; - axum::serve(listener, app) - .with_graceful_shutdown(async { - let ctrl_c = signal::ctrl_c(); - #[cfg(unix)] - let mut term = signal::unix::signal(signal::unix::SignalKind::terminate()).unwrap(); - #[cfg(unix)] - tokio::select! { _ = ctrl_c => {}, _ = term.recv() => {} }; - #[cfg(not(unix))] - ctrl_c.await.ok(); - info!("Shutdown signal received"); - }) - .await?; + axum::serve( + listener, + app.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(async { + let ctrl_c = signal::ctrl_c(); + #[cfg(unix)] + let mut term = signal::unix::signal(signal::unix::SignalKind::terminate()).unwrap(); + #[cfg(unix)] + tokio::select! { _ = ctrl_c => {}, _ = term.recv() => {} }; + #[cfg(not(unix))] + ctrl_c.await.ok(); + info!("Shutdown signal received"); + }) + .await?; cancel.cancel(); + if let Err(err) = stats_for_shutdown.persist() { + warn!(error = %err, "failed to persist stats on shutdown"); + } info!("ClawShell shut down"); Ok(()) } diff --git a/src/onboard/config_render.rs b/src/onboard/config_render.rs index 0dacdbf..1d63c2d 100644 --- a/src/onboard/config_render.rs +++ b/src/onboard/config_render.rs @@ -76,6 +76,9 @@ patterns = [ {{ name = "mastercard", regex = '\b5[1-5][0-9]{{14}}\b', action = "redact" }}, {{ name = "amex_card", regex = '\b3[47][0-9]{{13}}\b', action = "redact" }}, ] + +[stats] +persist_path = "/etc/clawshell/stats.json" {oauth_providers_section}"#, version = env!("CARGO_PKG_VERSION"), host = config.server_host, @@ -152,6 +155,8 @@ mod tests { assert!(toml_str.contains("log_level = \"info\"")); assert!(toml_str.contains(&format!("version = \"{}\"", env!("CARGO_PKG_VERSION")))); assert!(toml_str.contains("[dlp]")); + assert!(toml_str.contains("[stats]")); + assert!(toml_str.contains("persist_path =")); assert!(!toml_str.contains("[email]")); assert!(!toml_str.contains("[rate_limit]")); } diff --git a/src/stats.rs b/src/stats.rs new file mode 100644 index 0000000..f9c166b --- /dev/null +++ b/src/stats.rs @@ -0,0 +1,400 @@ +use crate::email::{extract_sender_address, normalize_sender_rule}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; +use std::sync::Mutex; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use tracing::{debug, warn}; + +/// Hard cap on the number of distinct sender addresses the stats map will +/// track. Further unique addresses are aggregated under [`OVERFLOW_KEY`] +/// so the map size is strictly bounded. +pub const MAX_TRACKED_ADDRESSES: usize = 10_000; + +/// Synthetic key used when the address map is full. Its count equals the +/// number of filtered messages from addresses that could not be tracked +/// individually. +pub const OVERFLOW_KEY: &str = ""; + +/// Max byte length of a single address we are willing to store. RFC 5321 +/// puts the hard limit for a path at 256 octets; 320 leaves slack for +/// display-name residue before we reject the entry. +const MAX_ADDRESS_LEN: usize = 320; + +pub struct Stats { + requests_total: AtomicU64, + prompt_tokens_total: AtomicU64, + completion_tokens_total: AtomicU64, + total_tokens_total: AtomicU64, + emails_filtered_total: AtomicU64, + /// Key: filtered sender address + /// Value: count of messages from that sender that were filtered + filtered_email_addresses: Mutex>, + persist_path: Option, + dirty: AtomicBool, + overflow_warned_this_cycle: AtomicBool, +} + +#[derive(Debug, Default, Clone, Serialize, Deserialize)] +pub struct StatsSnapshot { + pub requests_total: u64, + pub prompt_tokens_total: u64, + pub completion_tokens_total: u64, + pub total_tokens_total: u64, + pub emails_filtered_total: u64, + /// Key: filtered sender address + /// Value: count of messages from that sender that were filtered + pub filtered_email_addresses: BTreeMap, +} + +impl Stats { + pub fn new(persist_path: Option) -> Self { + let snapshot = match persist_path.as_deref() { + Some(path) if path.exists() => match Self::load_snapshot(path) { + Ok(snap) => snap, + Err(err) => { + warn!( + path = %path.display(), + error = %err, + "Failed to load stats from disk — starting with empty counters" + ); + StatsSnapshot::default() + } + }, + _ => StatsSnapshot::default(), + }; + + Self { + requests_total: AtomicU64::new(snapshot.requests_total), + prompt_tokens_total: AtomicU64::new(snapshot.prompt_tokens_total), + completion_tokens_total: AtomicU64::new(snapshot.completion_tokens_total), + total_tokens_total: AtomicU64::new(snapshot.total_tokens_total), + emails_filtered_total: AtomicU64::new(snapshot.emails_filtered_total), + filtered_email_addresses: Mutex::new(snapshot.filtered_email_addresses), + persist_path, + dirty: AtomicBool::new(false), + overflow_warned_this_cycle: AtomicBool::new(false), + } + } + + fn load_snapshot(path: &Path) -> std::io::Result { + let content = std::fs::read_to_string(path)?; + serde_json::from_str(&content) + .map_err(|e| std::io::Error::other(format!("failed to parse stats file: {e}"))) + } + + pub fn record_request(&self) { + self.requests_total.fetch_add(1, Ordering::Relaxed); + self.dirty.store(true, Ordering::Relaxed); + } + + /// Parse a response body for an LLM `usage` object and add it to the + /// running totals. Accepts both OpenAI-shaped + /// (`{prompt_tokens, completion_tokens, total_tokens}`) and + /// Anthropic-shaped (`{input_tokens, output_tokens}`) usage blocks. + /// Silently no-ops for non-JSON bodies or bodies without a usage object. + pub fn record_tokens_from_usage(&self, body: &[u8]) { + if body.is_empty() { + return; + } + let Ok(json) = serde_json::from_slice::(body) else { + return; + }; + let Some(usage) = json.get("usage") else { + return; + }; + + let prompt = usage + .get("prompt_tokens") + .or_else(|| usage.get("input_tokens")) + .and_then(Value::as_u64) + .unwrap_or(0); + let completion = usage + .get("completion_tokens") + .or_else(|| usage.get("output_tokens")) + .and_then(Value::as_u64) + .unwrap_or(0); + let total = usage + .get("total_tokens") + .and_then(Value::as_u64) + .unwrap_or_else(|| prompt.saturating_add(completion)); + + if prompt == 0 && completion == 0 && total == 0 { + return; + } + + self.prompt_tokens_total + .fetch_add(prompt, Ordering::Relaxed); + self.completion_tokens_total + .fetch_add(completion, Ordering::Relaxed); + self.total_tokens_total.fetch_add(total, Ordering::Relaxed); + self.dirty.store(true, Ordering::Relaxed); + } + + /// Record that a message from `from_header` was hidden by the email + /// policy. `from_header` is the raw `From:` value; we extract the bare + /// address when possible and fall back to a normalized form otherwise. + pub fn record_email_filtered(&self, from_header: &str) { + self.emails_filtered_total.fetch_add(1, Ordering::Relaxed); + self.dirty.store(true, Ordering::Relaxed); + + let key = extract_sender_address(from_header) + .unwrap_or_else(|| normalize_sender_rule(from_header)); + if key.is_empty() || key.len() > MAX_ADDRESS_LEN { + // Over-long or empty keys are aggregated into the overflow bucket + // so we still account for the filter event without letting a + // malformed From header explode the map. + self.bump_overflow(); + return; + } + + let mut map = match self.filtered_email_addresses.lock() { + Ok(guard) => guard, + Err(poisoned) => poisoned.into_inner(), + }; + if let Some(count) = map.get_mut(&key) { + *count = count.saturating_add(1); + return; + } + if map.len() < MAX_TRACKED_ADDRESSES { + map.insert(key, 1); + return; + } + // Cap reached: funnel into the overflow bucket without growing the map. + let overflow = map.entry(OVERFLOW_KEY.to_string()).or_insert(0); + *overflow = overflow.saturating_add(1); + drop(map); + if !self + .overflow_warned_this_cycle + .swap(true, Ordering::Relaxed) + { + warn!( + cap = MAX_TRACKED_ADDRESSES, + "filtered_email_addresses map hit its cap; additional unique senders are being aggregated under '{OVERFLOW_KEY}'" + ); + } + } + + fn bump_overflow(&self) { + let mut map = match self.filtered_email_addresses.lock() { + Ok(guard) => guard, + Err(poisoned) => poisoned.into_inner(), + }; + let overflow = map.entry(OVERFLOW_KEY.to_string()).or_insert(0); + *overflow = overflow.saturating_add(1); + } + + pub fn snapshot(&self) -> StatsSnapshot { + let map = match self.filtered_email_addresses.lock() { + Ok(guard) => guard.clone(), + Err(poisoned) => poisoned.into_inner().clone(), + }; + StatsSnapshot { + requests_total: self.requests_total.load(Ordering::Relaxed), + prompt_tokens_total: self.prompt_tokens_total.load(Ordering::Relaxed), + completion_tokens_total: self.completion_tokens_total.load(Ordering::Relaxed), + total_tokens_total: self.total_tokens_total.load(Ordering::Relaxed), + emails_filtered_total: self.emails_filtered_total.load(Ordering::Relaxed), + filtered_email_addresses: map, + } + } + + /// Write the current snapshot to disk if configured and there are + /// unsaved changes. Atomic: write to a sibling temp file, then rename. + pub fn persist(&self) -> std::io::Result<()> { + let Some(path) = self.persist_path.as_ref() else { + return Ok(()); + }; + if !self.dirty.swap(false, Ordering::Relaxed) { + return Ok(()); + } + // Reset the overflow-warned flag so the next cycle can log again if + // overflow is still happening. We do this regardless of whether the + // write succeeds; worst case is an extra log line per cycle. + self.overflow_warned_this_cycle + .store(false, Ordering::Relaxed); + + let snapshot = self.snapshot(); + let content = serde_json::to_string_pretty(&snapshot) + .map_err(|e| std::io::Error::other(format!("failed to serialize stats: {e}")))?; + + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + let tmp_path = path.with_extension("json.tmp"); + std::fs::write(&tmp_path, content)?; + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(&tmp_path, std::fs::Permissions::from_mode(0o600))?; + } + + std::fs::rename(&tmp_path, path)?; + debug!(path = %path.display(), "stats persisted to disk"); + Ok(()) + } +} + +impl std::fmt::Debug for Stats { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Stats") + .field("persist_path", &self.persist_path) + .finish_non_exhaustive() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn records_and_snapshots_requests() { + let stats = Stats::new(None); + stats.record_request(); + stats.record_request(); + stats.record_request(); + let snap = stats.snapshot(); + assert_eq!(snap.requests_total, 3); + } + + #[test] + fn parses_openai_shaped_usage() { + let stats = Stats::new(None); + let body = br#"{ + "id": "chatcmpl-1", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13} + }"#; + stats.record_tokens_from_usage(body); + let snap = stats.snapshot(); + assert_eq!(snap.prompt_tokens_total, 10); + assert_eq!(snap.completion_tokens_total, 3); + assert_eq!(snap.total_tokens_total, 13); + } + + #[test] + fn parses_anthropic_shaped_usage() { + let stats = Stats::new(None); + let body = br#"{"usage": {"input_tokens": 7, "output_tokens": 5}}"#; + stats.record_tokens_from_usage(body); + let snap = stats.snapshot(); + assert_eq!(snap.prompt_tokens_total, 7); + assert_eq!(snap.completion_tokens_total, 5); + assert_eq!(snap.total_tokens_total, 12); + } + + #[test] + fn token_parse_is_no_op_for_garbage_or_missing_usage() { + let stats = Stats::new(None); + stats.record_tokens_from_usage(b"not json at all"); + stats.record_tokens_from_usage(b"{\"id\": \"x\"}"); + stats.record_tokens_from_usage(b""); + let snap = stats.snapshot(); + assert_eq!(snap.prompt_tokens_total, 0); + assert_eq!(snap.completion_tokens_total, 0); + assert_eq!(snap.total_tokens_total, 0); + } + + #[test] + fn filtered_address_dedupes_and_counts() { + let stats = Stats::new(None); + stats.record_email_filtered("Spammer "); + stats.record_email_filtered("spam@example.com"); + stats.record_email_filtered("\"Bob\" "); + let snap = stats.snapshot(); + assert_eq!(snap.emails_filtered_total, 3); + assert_eq!( + snap.filtered_email_addresses.get("spam@example.com"), + Some(&2) + ); + assert_eq!( + snap.filtered_email_addresses.get("bob@example.com"), + Some(&1) + ); + } + + #[test] + fn filtered_address_overflow_bucket() { + let stats = Stats::new(None); + // Fill the map to exactly MAX_TRACKED_ADDRESSES unique addresses. + for i in 0..MAX_TRACKED_ADDRESSES { + stats.record_email_filtered(&format!("user{i}@example.com")); + } + // The next N unique addresses should all land in the overflow bucket. + for i in 0..5 { + stats.record_email_filtered(&format!("overflow{i}@example.com")); + } + // An already-tracked address should still increment its own counter. + stats.record_email_filtered("user0@example.com"); + + let snap = stats.snapshot(); + assert_eq!( + snap.emails_filtered_total, + (MAX_TRACKED_ADDRESSES as u64) + 5 + 1 + ); + // Map is capped at MAX + 1 (the overflow sentinel). + assert_eq!( + snap.filtered_email_addresses.len(), + MAX_TRACKED_ADDRESSES + 1 + ); + assert_eq!(snap.filtered_email_addresses.get(OVERFLOW_KEY), Some(&5)); + assert_eq!( + snap.filtered_email_addresses.get("user0@example.com"), + Some(&2) + ); + // Totals add up: sum of per-address counts equals emails_filtered_total. + let sum: u64 = snap.filtered_email_addresses.values().sum(); + assert_eq!(sum, snap.emails_filtered_total); + } + + #[test] + fn over_long_addresses_go_to_overflow() { + let stats = Stats::new(None); + let long = "a".repeat(MAX_ADDRESS_LEN + 1); + stats.record_email_filtered(&format!("<{long}@example.com>")); + let snap = stats.snapshot(); + assert_eq!(snap.emails_filtered_total, 1); + assert_eq!(snap.filtered_email_addresses.get(OVERFLOW_KEY), Some(&1)); + } + + #[test] + fn persist_round_trip() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("stats.json"); + + let stats = Stats::new(Some(path.clone())); + stats.record_request(); + stats.record_request(); + let body = br#"{"usage": {"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6}}"#; + stats.record_tokens_from_usage(body); + stats.record_email_filtered("spam@example.com"); + stats.persist().unwrap(); + + // Simulate restart. + let reloaded = Stats::new(Some(path)); + let snap = reloaded.snapshot(); + assert_eq!(snap.requests_total, 2); + assert_eq!(snap.prompt_tokens_total, 4); + assert_eq!(snap.completion_tokens_total, 2); + assert_eq!(snap.total_tokens_total, 6); + assert_eq!(snap.emails_filtered_total, 1); + assert_eq!( + snap.filtered_email_addresses.get("spam@example.com"), + Some(&1) + ); + } + + #[test] + fn persist_noop_when_clean() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("stats.json"); + let stats = Stats::new(Some(path.clone())); + stats.persist().unwrap(); + assert!( + !path.exists(), + "no write should happen when dirty flag is false" + ); + } +} diff --git a/tests/fixtures/config/invalid/dlp_invalid_regex.toml b/tests/fixtures/config/invalid/dlp_invalid_regex.toml index 494aa03..56e883a 100644 --- a/tests/fixtures/config/invalid/dlp_invalid_regex.toml +++ b/tests/fixtures/config/invalid/dlp_invalid_regex.toml @@ -5,3 +5,6 @@ patterns = [ { name = "bad", regex = '[invalid' }, ] + +[stats] +persist_path = "/etc/clawshell/stats.json" diff --git a/tests/fixtures/config/invalid/key_missing_real_key.toml b/tests/fixtures/config/invalid/key_missing_real_key.toml index 881e7aa..712ce57 100644 --- a/tests/fixtures/config/invalid/key_missing_real_key.toml +++ b/tests/fixtures/config/invalid/key_missing_real_key.toml @@ -3,3 +3,6 @@ [[keys]] virtual_key = "vk-1" + +[stats] +persist_path = "/etc/clawshell/stats.json" diff --git a/tests/fixtures/config/invalid/missing_stats.toml b/tests/fixtures/config/invalid/missing_stats.toml new file mode 100644 index 0000000..1901d97 --- /dev/null +++ b/tests/fixtures/config/invalid/missing_stats.toml @@ -0,0 +1,2 @@ +[server] +[upstream] diff --git a/tests/fixtures/config/valid/all_fields.toml b/tests/fixtures/config/valid/all_fields.toml index 939c765..25b107f 100644 --- a/tests/fixtures/config/valid/all_fields.toml +++ b/tests/fixtures/config/valid/all_fields.toml @@ -30,3 +30,6 @@ patterns = [ { name = "ssn", regex = '\b\d{3}-\d{2}-\d{4}\b', action = "block" }, { name = "email", regex = '\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b', action = "redact" }, ] + +[stats] +persist_path = "/var/lib/clawshell/stats.json" diff --git a/tests/fixtures/config/valid/dlp_disabled_scan.toml b/tests/fixtures/config/valid/dlp_disabled_scan.toml index 3920093..804a35a 100644 --- a/tests/fixtures/config/valid/dlp_disabled_scan.toml +++ b/tests/fixtures/config/valid/dlp_disabled_scan.toml @@ -3,3 +3,6 @@ [dlp] scan_responses = false + +[stats] +persist_path = "/etc/clawshell/stats.json" diff --git a/tests/fixtures/config/valid/empty_base_url.toml b/tests/fixtures/config/valid/empty_base_url.toml index d433fcc..665faa0 100644 --- a/tests/fixtures/config/valid/empty_base_url.toml +++ b/tests/fixtures/config/valid/empty_base_url.toml @@ -2,3 +2,6 @@ [upstream] openai_base_url = "" + +[stats] +persist_path = "/etc/clawshell/stats.json" diff --git a/tests/fixtures/config/valid/empty_host.toml b/tests/fixtures/config/valid/empty_host.toml index 8897b61..8969194 100644 --- a/tests/fixtures/config/valid/empty_host.toml +++ b/tests/fixtures/config/valid/empty_host.toml @@ -2,3 +2,6 @@ host = "" [upstream] + +[stats] +persist_path = "/etc/clawshell/stats.json" diff --git a/tests/fixtures/config/valid/empty_keys.toml b/tests/fixtures/config/valid/empty_keys.toml index 34017f1..a921882 100644 --- a/tests/fixtures/config/valid/empty_keys.toml +++ b/tests/fixtures/config/valid/empty_keys.toml @@ -5,3 +5,6 @@ [[keys]] virtual_key = "" real_key = "" + +[stats] +persist_path = "/etc/clawshell/stats.json" diff --git a/tests/fixtures/config/valid/minimal.toml b/tests/fixtures/config/valid/minimal.toml index 1901d97..ebea846 100644 --- a/tests/fixtures/config/valid/minimal.toml +++ b/tests/fixtures/config/valid/minimal.toml @@ -1,2 +1,5 @@ [server] [upstream] + +[stats] +persist_path = "/etc/clawshell/stats.json" diff --git a/tests/fixtures/config/valid/port_max.toml b/tests/fixtures/config/valid/port_max.toml index 2e18862..2c1dee4 100644 --- a/tests/fixtures/config/valid/port_max.toml +++ b/tests/fixtures/config/valid/port_max.toml @@ -2,3 +2,6 @@ port = 65535 [upstream] + +[stats] +persist_path = "/etc/clawshell/stats.json" diff --git a/tests/fixtures/config/valid/port_zero.toml b/tests/fixtures/config/valid/port_zero.toml index 7ed567f..0a8158e 100644 --- a/tests/fixtures/config/valid/port_zero.toml +++ b/tests/fixtures/config/valid/port_zero.toml @@ -2,3 +2,6 @@ port = 0 [upstream] + +[stats] +persist_path = "/etc/clawshell/stats.json" diff --git a/tests/snapshots/config_fixtures__all_fields.snap b/tests/snapshots/config_fixtures__all_fields.snap index b21fc2e..7ab05b4 100644 --- a/tests/snapshots/config_fixtures__all_fields.snap +++ b/tests/snapshots/config_fixtures__all_fields.snap @@ -33,6 +33,8 @@ dlp: regex: "\\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\\.[A-Za-z]{2,}\\b" action: redact scan_responses: true +stats: + persist_path: /var/lib/clawshell/stats.json log_level: debug derived: listen_addr: "0.0.0.0:8080" diff --git a/tests/snapshots/config_fixtures__dlp_disabled_scan.snap b/tests/snapshots/config_fixtures__dlp_disabled_scan.snap index 2953da4..f44345a 100644 --- a/tests/snapshots/config_fixtures__dlp_disabled_scan.snap +++ b/tests/snapshots/config_fixtures__dlp_disabled_scan.snap @@ -1,5 +1,5 @@ --- -source: tests/config_fixtures.rs +source: src/config.rs expression: snapshot --- server: @@ -13,6 +13,8 @@ keys: [] dlp: patterns: [] scan_responses: false +stats: + persist_path: /etc/clawshell/stats.json log_level: info derived: listen_addr: "127.0.0.1:18790" diff --git a/tests/snapshots/config_fixtures__empty_base_url.snap b/tests/snapshots/config_fixtures__empty_base_url.snap index 6ec784a..64a2aa3 100644 --- a/tests/snapshots/config_fixtures__empty_base_url.snap +++ b/tests/snapshots/config_fixtures__empty_base_url.snap @@ -1,5 +1,5 @@ --- -source: tests/config_fixtures.rs +source: src/config.rs expression: snapshot --- server: @@ -13,6 +13,8 @@ keys: [] dlp: patterns: [] scan_responses: true +stats: + persist_path: /etc/clawshell/stats.json log_level: info derived: listen_addr: "127.0.0.1:18790" diff --git a/tests/snapshots/config_fixtures__empty_host.snap b/tests/snapshots/config_fixtures__empty_host.snap index 120b945..5610f52 100644 --- a/tests/snapshots/config_fixtures__empty_host.snap +++ b/tests/snapshots/config_fixtures__empty_host.snap @@ -1,5 +1,5 @@ --- -source: tests/config_fixtures.rs +source: src/config.rs expression: snapshot --- server: @@ -13,6 +13,8 @@ keys: [] dlp: patterns: [] scan_responses: true +stats: + persist_path: /etc/clawshell/stats.json log_level: info derived: listen_addr: ":18790" diff --git a/tests/snapshots/config_fixtures__empty_keys.snap b/tests/snapshots/config_fixtures__empty_keys.snap index 2a114e0..2320eb2 100644 --- a/tests/snapshots/config_fixtures__empty_keys.snap +++ b/tests/snapshots/config_fixtures__empty_keys.snap @@ -1,5 +1,5 @@ --- -source: tests/config_fixtures.rs +source: src/config.rs expression: snapshot --- server: @@ -17,6 +17,8 @@ keys: dlp: patterns: [] scan_responses: true +stats: + persist_path: /etc/clawshell/stats.json log_level: info derived: listen_addr: "127.0.0.1:18790" diff --git a/tests/snapshots/config_fixtures__minimal.snap b/tests/snapshots/config_fixtures__minimal.snap index a82b5c6..2ae27d2 100644 --- a/tests/snapshots/config_fixtures__minimal.snap +++ b/tests/snapshots/config_fixtures__minimal.snap @@ -1,5 +1,5 @@ --- -source: tests/config_fixtures.rs +source: src/config.rs expression: snapshot --- server: @@ -13,6 +13,8 @@ keys: [] dlp: patterns: [] scan_responses: true +stats: + persist_path: /etc/clawshell/stats.json log_level: info derived: listen_addr: "127.0.0.1:18790" diff --git a/tests/snapshots/config_fixtures__missing_stats.snap b/tests/snapshots/config_fixtures__missing_stats.snap new file mode 100644 index 0000000..5138d8c --- /dev/null +++ b/tests/snapshots/config_fixtures__missing_stats.snap @@ -0,0 +1,9 @@ +--- +source: src/config.rs +expression: err.to_string() +--- +TOML parse error at line 1, column 1 + | +1 | [server] + | ^ +missing field `stats` diff --git a/tests/snapshots/config_fixtures__port_max.snap b/tests/snapshots/config_fixtures__port_max.snap index 5d3768f..4b7a692 100644 --- a/tests/snapshots/config_fixtures__port_max.snap +++ b/tests/snapshots/config_fixtures__port_max.snap @@ -1,5 +1,5 @@ --- -source: tests/config_fixtures.rs +source: src/config.rs expression: snapshot --- server: @@ -13,6 +13,8 @@ keys: [] dlp: patterns: [] scan_responses: true +stats: + persist_path: /etc/clawshell/stats.json log_level: info derived: listen_addr: "127.0.0.1:65535" diff --git a/tests/snapshots/config_fixtures__port_zero.snap b/tests/snapshots/config_fixtures__port_zero.snap index 8e8761a..3e1b7fc 100644 --- a/tests/snapshots/config_fixtures__port_zero.snap +++ b/tests/snapshots/config_fixtures__port_zero.snap @@ -1,5 +1,5 @@ --- -source: tests/config_fixtures.rs +source: src/config.rs expression: snapshot --- server: @@ -13,6 +13,8 @@ keys: [] dlp: patterns: [] scan_responses: true +stats: + persist_path: /etc/clawshell/stats.json log_level: info derived: listen_addr: "127.0.0.1:0"