1use std::collections::BTreeSet;
30use std::path::{Path, PathBuf};
31use std::time::{Duration, Instant};
32
33use axum::{Json, extract::State, http::StatusCode, response::IntoResponse};
34use serde::Serialize;
35use serde_json::json;
36
37use super::team::TeamAppState;
38
39pub(super) const STORAGE_CACHE_TTL: Duration = Duration::from_mins(1);
41
42pub(super) const QUOTA_ENV: &str = "LEANCTX_TEAM_STORAGE_QUOTA_BYTES";
47
48pub(super) const DEFAULT_TEAM_STORAGE_QUOTA_BYTES: u64 = 5 * 1024 * 1024 * 1024;
55
56#[derive(Debug, Clone, Serialize)]
58#[serde(rename_all = "camelCase")]
59pub struct StorageComponent {
60 pub id: String,
62 pub bytes: u64,
63}
64
65#[derive(Debug, Clone)]
67pub struct StorageReport {
68 pub used_bytes: u64,
69 pub components: Vec<StorageComponent>,
70}
71
72#[derive(Default)]
74pub struct StorageCache(Option<(Instant, StorageReport)>);
75
76#[derive(Clone)]
78pub struct StorageRoots {
79 pub data_root: PathBuf,
82 pub workspaces: Vec<(String, PathBuf)>,
85 pub quota_bytes: u64,
88}
89
90impl StorageRoots {
91 fn measure(&self) -> StorageReport {
93 let mut components = Vec::new();
94 let data_bytes = dir_allocated_bytes(&self.data_root);
95 components.push(StorageComponent {
96 id: "server-data".into(),
97 bytes: data_bytes,
98 });
99
100 for (id, state_dir) in &self.workspaces {
101 if state_dir.starts_with(&self.data_root) {
102 continue; }
104 components.push(StorageComponent {
105 id: format!("workspace:{id}"),
106 bytes: dir_allocated_bytes(state_dir),
107 });
108 }
109
110 let used_bytes = components
111 .iter()
112 .map(|c| c.bytes)
113 .fold(0u64, u64::saturating_add);
114 StorageReport {
115 used_bytes,
116 components,
117 }
118 }
119}
120
121pub async fn v1_storage(State(state): State<TeamAppState>) -> impl IntoResponse {
123 let (report, age) = cached_report(&state).await;
124 let body = json!({
125 "schemaVersion": 1,
126 "measuredAt": chrono::Utc::now().to_rfc3339(),
127 "usedBytes": report.used_bytes,
128 "quotaBytes": state.team.storage_roots.quota_bytes,
129 "components": report.components,
130 "cacheAgeSeconds": age.as_secs(),
131 });
132 (StatusCode::OK, Json(body))
133}
134
135pub async fn v1_usage(State(state): State<TeamAppState>) -> impl IntoResponse {
137 let dir = state.team.savings_store_dir.lock().await.clone();
138 let summary = tokio::task::spawn_blocking(move || super::savings_summary::aggregate(&dir))
139 .await
140 .unwrap_or_default();
141 let (report, _) = cached_report(&state).await;
142
143 let storage = json!({
144 "used_bytes": report.used_bytes,
145 "quota_bytes": state.team.storage_roots.quota_bytes,
146 });
147
148 let connectors = state.team.connectors.clone();
152 let connectors_dir = state.team.connectors_state_dir.as_ref().clone();
153 let connectors_usage = tokio::task::spawn_blocking(move || {
154 super::team::connectors::usage_rollup(&connectors, &connectors_dir)
155 })
156 .await
157 .unwrap_or_else(|_| json!({}));
158
159 let body = json!({
160 "schemaVersion": 1,
161 "generatedAt": chrono::Utc::now().to_rfc3339(),
162 "savings": {
163 "memberCount": summary.member_count,
164 "savedTokens": summary.totals.saved_tokens,
165 "netSavedTokens": summary.totals.net_saved_tokens,
166 "savedUsd": summary.totals.saved_usd,
167 },
168 "toolCalls": summary.totals.total_events,
171 "storage": storage,
172 "connectors": connectors_usage,
173 });
174 (StatusCode::OK, Json(body))
175}
176
177async fn cached_report(state: &TeamAppState) -> (StorageReport, Duration) {
179 {
180 let cache = state.team.storage_cache.lock().await;
181 if let Some((at, report)) = cache.0.as_ref() {
182 let age = at.elapsed();
183 if age < STORAGE_CACHE_TTL {
184 return (report.clone(), age);
185 }
186 }
187 }
188
189 let roots = state.team.storage_roots.clone();
190 let report = tokio::task::spawn_blocking(move || roots.measure())
191 .await
192 .unwrap_or(StorageReport {
193 used_bytes: 0,
194 components: Vec::new(),
195 });
196
197 let mut cache = state.team.storage_cache.lock().await;
198 cache.0 = Some((Instant::now(), report.clone()));
199 (report, Duration::ZERO)
200}
201
202fn quota_bytes_from_env() -> Option<u64> {
203 std::env::var(QUOTA_ENV).ok()?.trim().parse::<u64>().ok()
204}
205
206pub(super) fn resolve_quota_bytes(env_override: Option<u64>, config_quota: Option<u64>) -> u64 {
210 env_override
211 .or(config_quota)
212 .unwrap_or(DEFAULT_TEAM_STORAGE_QUOTA_BYTES)
213}
214
215pub(super) fn storage_roots_from_config(
219 audit_log_path: &Path,
220 workspaces: &[(String, PathBuf)],
221 config_quota: Option<u64>,
222) -> StorageRoots {
223 let data_root = audit_log_path
224 .parent()
225 .unwrap_or_else(|| Path::new("."))
226 .to_path_buf();
227 let workspaces = workspaces
228 .iter()
229 .map(|(id, root)| (id.clone(), root.join(".lean-ctx")))
230 .collect();
231 StorageRoots {
232 data_root,
233 workspaces,
234 quota_bytes: resolve_quota_bytes(quota_bytes_from_env(), config_quota),
235 }
236}
237
238fn dir_allocated_bytes(path: &Path) -> u64 {
243 let mut seen: BTreeSet<(u64, u64)> = BTreeSet::new();
244 let mut total = 0u64;
245 let mut stack = vec![path.to_path_buf()];
246 while let Some(p) = stack.pop() {
247 let Ok(meta) = std::fs::symlink_metadata(&p) else {
248 continue;
249 };
250 if meta.is_symlink() {
251 continue;
252 }
253 if meta.is_dir() {
254 if let Ok(entries) = std::fs::read_dir(&p) {
255 stack.extend(entries.flatten().map(|e| e.path()));
256 }
257 continue;
258 }
259 #[cfg(unix)]
260 {
261 use std::os::unix::fs::MetadataExt;
262 if !seen.insert((meta.dev(), meta.ino())) {
263 continue; }
265 total = total.saturating_add(meta.blocks().saturating_mul(512));
266 }
267 #[cfg(not(unix))]
268 {
269 let _ = &mut seen;
270 total = total.saturating_add(meta.len());
271 }
272 }
273 total
274}
275
276#[must_use]
283pub(crate) fn is_over_quota(data_root: &Path, quota_bytes: u64) -> bool {
284 quota_bytes > 0 && dir_allocated_bytes(data_root) >= quota_bytes
285}
286
287#[cfg(test)]
288mod tests {
289 use super::*;
290
291 fn temp_dir(tag: &str) -> PathBuf {
292 let d =
293 std::env::temp_dir().join(format!("leanctx_team_billing_{tag}_{}", std::process::id()));
294 let _ = std::fs::remove_dir_all(&d);
295 std::fs::create_dir_all(&d).unwrap();
296 d
297 }
298
299 #[test]
300 fn missing_dir_measures_zero() {
301 let missing = std::env::temp_dir().join("leanctx_team_billing_does_not_exist_xyz");
302 let _ = std::fs::remove_dir_all(&missing);
303 assert_eq!(dir_allocated_bytes(&missing), 0);
304 }
305
306 #[test]
311 fn quota_resolution_precedence() {
312 assert_eq!(resolve_quota_bytes(Some(7), Some(9)), 7);
313 assert_eq!(resolve_quota_bytes(None, Some(9)), 9);
314 assert_eq!(
315 resolve_quota_bytes(None, None),
316 DEFAULT_TEAM_STORAGE_QUOTA_BYTES
317 );
318 assert_eq!(DEFAULT_TEAM_STORAGE_QUOTA_BYTES, 5_368_709_120);
319 }
320
321 #[test]
325 fn over_quota_only_trips_with_a_positive_quota() {
326 let d = temp_dir("overquota");
327 std::fs::write(d.join("a.bin"), vec![b'x'; 20_000]).unwrap();
328 assert!(!is_over_quota(&d, 0), "0 quota must never trip");
329 assert!(is_over_quota(&d, 1_000), "20 KiB must exceed a 1 KiB quota");
330 assert!(
331 !is_over_quota(&d, 10 * 1024 * 1024),
332 "20 KiB must not exceed a 10 MiB quota"
333 );
334 let missing = std::env::temp_dir().join("leanctx_team_billing_overquota_missing_xyz");
335 let _ = std::fs::remove_dir_all(&missing);
336 assert!(
337 !is_over_quota(&missing, 1),
338 "missing dir is never over quota"
339 );
340 let _ = std::fs::remove_dir_all(&d);
341 }
342
343 #[test]
344 fn allocated_bytes_cover_written_content() {
345 let d = temp_dir("alloc");
346 std::fs::write(d.join("a.jsonl"), vec![b'x'; 10_000]).unwrap();
347 std::fs::create_dir_all(d.join("nested")).unwrap();
348 std::fs::write(d.join("nested/b.bin"), vec![b'y'; 5_000]).unwrap();
349 let measured = dir_allocated_bytes(&d);
350 assert!(measured >= 15_000, "measured {measured} < logical 15000");
354 assert!(
355 measured < 15_000 * 64,
356 "measured {measured} implausibly large"
357 );
358 let _ = std::fs::remove_dir_all(&d);
359 }
360
361 #[cfg(unix)]
362 #[test]
363 fn hard_links_count_once_and_symlinks_do_not_escape() {
364 let d = temp_dir("links");
365 std::fs::write(d.join("real.bin"), vec![b'z'; 8_192]).unwrap();
366 std::fs::hard_link(d.join("real.bin"), d.join("hard.bin")).unwrap();
367 let outside = temp_dir("links_outside");
368 std::fs::write(outside.join("big.bin"), vec![b'w'; 100_000]).unwrap();
369 std::os::unix::fs::symlink(outside.join("big.bin"), d.join("escape.bin")).unwrap();
370
371 let measured = dir_allocated_bytes(&d);
372 let single = dir_allocated_bytes(&outside); assert!(measured < single, "symlink target must not be billed");
374 assert!(
376 (8_192..16_384).contains(&measured),
377 "hard link double-counted: {measured}"
378 );
379
380 let _ = std::fs::remove_dir_all(&d);
381 let _ = std::fs::remove_dir_all(&outside);
382 }
383
384 #[test]
385 fn workspace_state_dirs_under_data_root_are_not_double_counted() {
386 let d = temp_dir("dedupe");
387 let audit = d.join("audit.jsonl");
388 std::fs::write(&audit, "x").unwrap();
389 let ws_root = d.join("ws1");
392 std::fs::create_dir_all(ws_root.join(".lean-ctx")).unwrap();
393 std::fs::write(ws_root.join(".lean-ctx/events.jsonl"), vec![b'e'; 4_096]).unwrap();
394 let ext = temp_dir("dedupe_ext");
396 std::fs::create_dir_all(ext.join(".lean-ctx")).unwrap();
397 std::fs::write(ext.join(".lean-ctx/k.jsonl"), vec![b'k'; 4_096]).unwrap();
398
399 let roots = storage_roots_from_config(
400 &audit,
401 &[
402 ("inside".into(), ws_root.clone()),
403 ("outside".into(), ext.clone()),
404 ],
405 None,
406 );
407 let report = roots.measure();
408 let ids: Vec<&str> = report.components.iter().map(|c| c.id.as_str()).collect();
409 assert!(ids.contains(&"server-data"));
410 assert!(
411 !ids.contains(&"workspace:inside"),
412 "nested workspace double-counted"
413 );
414 assert!(ids.contains(&"workspace:outside"));
415 let sum: u64 = report.components.iter().map(|c| c.bytes).sum();
416 assert_eq!(report.used_bytes, sum);
417
418 let _ = std::fs::remove_dir_all(&d);
419 let _ = std::fs::remove_dir_all(&ext);
420 }
421}