1use std::collections::BTreeMap;
20use std::sync::Mutex;
21use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
22use std::time::Duration;
23
24use serde::{Deserialize, Serialize};
25
26const MAX_TRACKED_PENDING: usize = 4_096;
29const MAX_METHOD_BYTES: usize = 64;
31const PUBLISH_INTERVAL: Duration = Duration::from_secs(2);
33pub const HEARTBEAT: Duration = Duration::from_secs(10);
36pub const STALE_AFTER: Duration = Duration::from_secs(120);
38
39#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
41#[serde(rename_all = "snake_case")]
42pub enum Stage {
43 PreDispatch,
45 Handler,
47 Finalizing,
49}
50
51impl Stage {
52 pub fn label(self) -> &'static str {
53 match self {
54 Self::PreDispatch => "pre-dispatch",
55 Self::Handler => "handler",
56 Self::Finalizing => "finalizing",
57 }
58 }
59}
60
61#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
63#[serde(rename_all = "snake_case")]
64pub enum Outcome {
65 Completed,
66 Errored,
67 Abandoned,
69}
70
71#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
72pub struct PendingRequest {
73 pub method: String,
74 pub stage: Stage,
75 pub admitted_at_ms: u64,
76 #[serde(default, skip_serializing_if = "Option::is_none")]
77 pub age_ms: Option<u64>,
78}
79
80#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
81pub struct CompletionRecord {
82 pub method: String,
83 pub outcome: Outcome,
84 pub finished_at_ms: u64,
85 pub duration_ms: u64,
86 #[serde(default, skip_serializing_if = "Option::is_none")]
87 pub age_ms: Option<u64>,
88}
89
90#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
91pub struct PhaseRecord {
92 pub name: String,
93 pub running_ms: u64,
94}
95
96#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
99pub struct RequestLedgerSnapshot {
100 pub pid: u32,
101 pub kind: String,
103 pub generation: String,
105 pub started_at_ms: u64,
106 pub snapshot_at_ms: u64,
107 pub admitted: u64,
108 pub completed: u64,
109 pub errored: u64,
110 pub refused: u64,
111 pub abandoned: u64,
112 pub pending: u64,
113 pub untracked_pending: u64,
115 pub oldest_pending: Option<PendingRequest>,
116 pub last_tool_completion: Option<CompletionRecord>,
117 pub last_initialize: Option<CompletionRecord>,
118 pub current_phase: Option<PhaseRecord>,
119 pub observations: crate::core::events::DetachedObservationStats,
120}
121
122struct Tracked {
123 method: String,
124 stage: Stage,
125 admitted_at_ms: u64,
126}
127
128struct Ledger {
129 pending: BTreeMap<u64, Tracked>,
130 last_tool_completion: Option<CompletionRecord>,
131 last_initialize: Option<CompletionRecord>,
132}
133
134static LEDGER: Mutex<Ledger> = Mutex::new(Ledger {
135 pending: BTreeMap::new(),
136 last_tool_completion: None,
137 last_initialize: None,
138});
139static NEXT_ID: AtomicU64 = AtomicU64::new(1);
140static ADMITTED: AtomicU64 = AtomicU64::new(0);
141static COMPLETED: AtomicU64 = AtomicU64::new(0);
142static ERRORED: AtomicU64 = AtomicU64::new(0);
143static REFUSED: AtomicU64 = AtomicU64::new(0);
144static ABANDONED: AtomicU64 = AtomicU64::new(0);
145static STARTED_AT_MS: AtomicU64 = AtomicU64::new(0);
146static SERVING: AtomicBool = AtomicBool::new(false);
147static FINISHED: AtomicBool = AtomicBool::new(false);
148static DIRTY: AtomicBool = AtomicBool::new(false);
149static KIND: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
150
151fn ledger() -> std::sync::MutexGuard<'static, Ledger> {
152 LEDGER
153 .lock()
154 .unwrap_or_else(std::sync::PoisonError::into_inner)
155}
156
157fn now_ms() -> u64 {
158 std::time::SystemTime::now()
159 .duration_since(std::time::UNIX_EPOCH)
160 .map(|d| u64::try_from(d.as_millis()).unwrap_or(u64::MAX))
161 .unwrap_or_default()
162}
163
164fn sanitize_method(method: &str) -> String {
166 method
167 .chars()
168 .filter(char::is_ascii_graphic)
169 .take(MAX_METHOD_BYTES)
170 .collect()
171}
172
173pub fn start_serving(kind: &'static str) {
176 crate::core::events::enable_detached_observations();
177 if SERVING.swap(true, Ordering::AcqRel) {
178 return;
179 }
180 let _ = KIND.set(kind);
181 STARTED_AT_MS.store(now_ms(), Ordering::Relaxed);
182 DIRTY.store(true, Ordering::Release);
183 if let Some(dir) = snapshot_dir() {
184 let _ = std::thread::Builder::new()
185 .name("lean-ctx-request-ledger".to_string())
186 .spawn(move || publish_loop(&dir));
187 }
188}
189
190pub fn serving() -> bool {
191 SERVING.load(Ordering::Acquire)
192}
193
194pub fn finish_serving(timeout: Duration) {
200 if FINISHED.swap(true, Ordering::AcqRel) {
201 return;
202 }
203 crate::core::events::flush_detached_observations(timeout);
204 if SERVING.load(Ordering::Acquire)
205 && let Some(dir) = snapshot_dir()
206 {
207 let _ = std::fs::remove_file(dir.join(format!("{}.json", std::process::id())));
208 }
209}
210
211pub fn refused() {
213 REFUSED.fetch_add(1, Ordering::Relaxed);
214 DIRTY.store(true, Ordering::Release);
215}
216
217#[must_use = "the request is abandoned when the guard drops unfinished"]
221pub fn begin(method: &str) -> RequestGuard {
222 let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
223 let admitted_at_ms = now_ms();
224 let method = sanitize_method(method);
225 ADMITTED.fetch_add(1, Ordering::Relaxed);
226 {
227 let mut ledger = ledger();
228 if ledger.pending.len() < MAX_TRACKED_PENDING {
229 ledger.pending.insert(
230 id,
231 Tracked {
232 method: method.clone(),
233 stage: Stage::PreDispatch,
234 admitted_at_ms,
235 },
236 );
237 }
238 }
239 DIRTY.store(true, Ordering::Release);
240 RequestGuard {
241 id,
242 method,
243 admitted_at_ms,
244 finished: false,
245 }
246}
247
248pub struct RequestGuard {
249 id: u64,
250 method: String,
251 admitted_at_ms: u64,
252 finished: bool,
253}
254
255impl RequestGuard {
256 pub fn stage(&self, stage: Stage) {
257 if let Some(tracked) = ledger().pending.get_mut(&self.id) {
258 tracked.stage = stage;
259 }
260 DIRTY.store(true, Ordering::Release);
261 }
262
263 pub fn finish(mut self, outcome: Outcome) {
264 self.record(outcome);
265 }
266
267 fn record(&mut self, outcome: Outcome) {
268 if std::mem::replace(&mut self.finished, true) {
269 return;
270 }
271 match outcome {
272 Outcome::Completed => &COMPLETED,
273 Outcome::Errored => &ERRORED,
274 Outcome::Abandoned => &ABANDONED,
275 }
276 .fetch_add(1, Ordering::Relaxed);
277 let finished_at_ms = now_ms();
278 let record = CompletionRecord {
279 method: std::mem::take(&mut self.method),
280 outcome,
281 finished_at_ms,
282 duration_ms: finished_at_ms.saturating_sub(self.admitted_at_ms),
283 age_ms: None,
284 };
285 {
286 let mut ledger = ledger();
287 ledger.pending.remove(&self.id);
288 if record.method == "initialize" {
289 ledger.last_initialize = Some(record);
290 } else if record.method.starts_with("tools/call") {
291 ledger.last_tool_completion = Some(record);
292 }
293 }
294 DIRTY.store(true, Ordering::Release);
295 }
296}
297
298impl Drop for RequestGuard {
299 fn drop(&mut self) {
300 self.record(Outcome::Abandoned);
301 }
302}
303
304pub fn snapshot() -> RequestLedgerSnapshot {
307 let (oldest_pending, tracked, last_tool_completion, last_initialize) = {
308 let ledger = ledger();
309 let oldest = ledger
312 .pending
313 .first_key_value()
314 .map(|(_, t)| PendingRequest {
315 method: t.method.clone(),
316 stage: t.stage,
317 admitted_at_ms: t.admitted_at_ms,
318 age_ms: None,
319 });
320 (
321 oldest,
322 ledger.pending.len() as u64,
323 ledger.last_tool_completion.clone(),
324 ledger.last_initialize.clone(),
325 )
326 };
327 let completed = COMPLETED.load(Ordering::Relaxed);
330 let errored = ERRORED.load(Ordering::Relaxed);
331 let abandoned = ABANDONED.load(Ordering::Relaxed);
332 let admitted = ADMITTED.load(Ordering::Relaxed);
333 let pending = admitted.saturating_sub(completed + errored + abandoned);
334 RequestLedgerSnapshot {
335 pid: std::process::id(),
336 kind: KIND.get().copied().unwrap_or("unknown").to_string(),
337 generation: crate::core::context_os::observation::producer_instance().to_string(),
338 started_at_ms: STARTED_AT_MS.load(Ordering::Relaxed),
339 snapshot_at_ms: now_ms(),
340 admitted,
341 completed,
342 errored,
343 refused: REFUSED.load(Ordering::Relaxed),
344 abandoned,
345 pending,
346 untracked_pending: if tracked < MAX_TRACKED_PENDING as u64 {
350 0
351 } else {
352 pending.saturating_sub(tracked)
353 },
354 oldest_pending,
355 last_tool_completion,
356 last_initialize,
357 current_phase: crate::core::phase_timing::current().map(|(name, running)| PhaseRecord {
358 name: name.to_string(),
359 running_ms: u64::try_from(running.as_millis()).unwrap_or(u64::MAX),
360 }),
361 observations: crate::core::events::detached_observation_stats(),
362 }
363}
364
365impl RequestLedgerSnapshot {
366 #[must_use]
368 pub fn with_ages(mut self, now_ms: u64) -> Self {
369 if let Some(pending) = self.oldest_pending.as_mut() {
370 pending.age_ms = Some(now_ms.saturating_sub(pending.admitted_at_ms));
371 }
372 for record in [
373 self.last_tool_completion.as_mut(),
374 self.last_initialize.as_mut(),
375 ]
376 .into_iter()
377 .flatten()
378 {
379 record.age_ms = Some(now_ms.saturating_sub(record.finished_at_ms));
380 }
381 self
382 }
383
384 pub fn render(&self, now_ms: u64) -> String {
386 let age = |at: u64| format_age(now_ms.saturating_sub(at));
387 let oldest = if self.pending == 0 {
388 "none".to_string()
389 } else {
390 match &self.oldest_pending {
391 Some(p) => format!(
392 "{} {} in {}",
393 p.method,
394 age(p.admitted_at_ms),
395 p.stage.label()
396 ),
397 None if self.untracked_pending > 0 => "unknown (untracked)".to_string(),
398 None => "none".to_string(),
399 }
400 };
401 let last = |record: &Option<CompletionRecord>| match record {
402 Some(r) => format!(
403 "{} ago ({}, {}ms)",
404 age(r.finished_at_ms),
405 outcome_label(r.outcome),
406 r.duration_ms
407 ),
408 None => "none yet".to_string(),
409 };
410 let mut line = format!(
411 "pid {} ({}), up {}: admitted {} · completed {} · errored {} · refused {} · abandoned {} · pending {}; oldest pending {}; last tool completion {}; last initialize {}",
412 self.pid,
413 self.kind,
414 age(self.started_at_ms),
415 self.admitted,
416 self.completed,
417 self.errored,
418 self.refused,
419 self.abandoned,
420 self.pending,
421 oldest,
422 last(&self.last_tool_completion),
423 last(&self.last_initialize),
424 );
425 if let Some(phase) = &self.current_phase {
426 line.push_str(&format!(
427 "; phase {} {}",
428 phase.name,
429 format_age(phase.running_ms)
430 ));
431 }
432 if self.observations.dropped > 0 || self.observations.failed > 0 {
433 line.push_str(&format!(
434 "; observations dropped {} · failed {}",
435 self.observations.dropped, self.observations.failed
436 ));
437 }
438 let snapshot_age = now_ms.saturating_sub(self.snapshot_at_ms);
439 if snapshot_age > HEARTBEAT.as_millis() as u64 * 2 {
440 line.push_str(&format!("; snapshot {} old", format_age(snapshot_age)));
441 }
442 line
443 }
444
445 pub fn oldest_pending_age(&self, now_ms: u64) -> Option<Duration> {
447 self.oldest_pending
448 .as_ref()
449 .map(|p| Duration::from_millis(now_ms.saturating_sub(p.admitted_at_ms)))
450 }
451}
452
453fn outcome_label(outcome: Outcome) -> &'static str {
454 match outcome {
455 Outcome::Completed => "completed",
456 Outcome::Errored => "errored",
457 Outcome::Abandoned => "abandoned",
458 }
459}
460
461fn format_age(ms: u64) -> String {
462 match ms {
463 0..1_000 => format!("{ms}ms"),
464 1_000..120_000 => format!("{}s", ms / 1_000),
465 120_000..7_200_000 => format!("{}m", ms / 60_000),
466 _ => format!("{}h", ms / 3_600_000),
467 }
468}
469
470pub fn epoch_ms() -> u64 {
472 now_ms()
473}
474
475fn snapshot_dir() -> Option<std::path::PathBuf> {
476 crate::core::data_dir::lean_ctx_data_dir()
477 .ok()
478 .map(|dir| dir.join("requests"))
479}
480
481fn publish_loop(dir: &std::path::Path) {
482 if std::fs::create_dir_all(dir).is_err() {
483 return;
484 }
485 sweep_stale(dir, now_ms());
486 let path = dir.join(format!("{}.json", std::process::id()));
487 let mut last_write_ms = 0u64;
488 loop {
489 if FINISHED.load(Ordering::Acquire) {
491 return;
492 }
493 let now = now_ms();
494 let heartbeat_due = now.saturating_sub(last_write_ms) >= HEARTBEAT.as_millis() as u64;
495 if DIRTY.swap(false, Ordering::AcqRel) || heartbeat_due {
496 if let Ok(bytes) = serde_json::to_vec(&snapshot()) {
497 let _ = crate::core::atomic_fs::try_atomic_write(&path, &bytes, None);
498 }
499 if FINISHED.load(Ordering::Acquire) {
502 let _ = std::fs::remove_file(&path);
503 return;
504 }
505 last_write_ms = now;
506 }
507 std::thread::sleep(PUBLISH_INTERVAL);
508 }
509}
510
511fn sweep_stale(dir: &std::path::Path, now: u64) {
514 for snapshot in read_dir_snapshots(dir) {
515 if now.saturating_sub(snapshot.1.snapshot_at_ms) > STALE_AFTER.as_millis() as u64 * 5 {
516 let _ = std::fs::remove_file(&snapshot.0);
517 }
518 }
519}
520
521fn read_dir_snapshots(dir: &std::path::Path) -> Vec<(std::path::PathBuf, RequestLedgerSnapshot)> {
522 const MAX_BYTES: u64 = 64 * 1024;
523 let Ok(entries) = std::fs::read_dir(dir) else {
524 return Vec::new();
525 };
526 let mut snapshots: Vec<_> = entries
527 .flatten()
528 .map(|entry| entry.path())
529 .filter(|path| path.extension().is_some_and(|ext| ext == "json"))
530 .filter(|path| std::fs::metadata(path).is_ok_and(|m| m.len() <= MAX_BYTES))
531 .filter_map(|path| {
532 let text = std::fs::read_to_string(&path).ok()?;
533 let snapshot = serde_json::from_str(&text).ok()?;
534 Some((path, snapshot))
535 })
536 .collect();
537 snapshots.sort_by_key(|(_, s): &(_, RequestLedgerSnapshot)| s.pid);
538 snapshots
539}
540
541pub fn published_snapshots() -> Vec<RequestLedgerSnapshot> {
543 let Some(dir) = snapshot_dir() else {
544 return Vec::new();
545 };
546 let now = now_ms();
547 read_dir_snapshots(&dir)
548 .into_iter()
549 .map(|(_, snapshot)| snapshot)
550 .filter(|s| now.saturating_sub(s.snapshot_at_ms) <= STALE_AFTER.as_millis() as u64)
551 .collect()
552}
553
554#[cfg(test)]
555mod tests {
556 use super::*;
557
558 fn counts() -> (u64, u64, u64, u64) {
559 (
560 ADMITTED.load(Ordering::Relaxed),
561 COMPLETED.load(Ordering::Relaxed),
562 ERRORED.load(Ordering::Relaxed),
563 ABANDONED.load(Ordering::Relaxed),
564 )
565 }
566
567 #[test]
572 fn each_request_path_counts_its_outcome_exactly_once() {
573 if crate::test_env::run_in_own_process(
574 "core::request_ledger::tests::each_request_path_counts_its_outcome_exactly_once",
575 &[],
576 ) {
577 return;
578 }
579 let init = begin("initialize");
580 init.finish(Outcome::Completed);
581
582 let stuck = begin("tools/call:ctx_read");
583 stuck.stage(Stage::Handler);
584 std::thread::sleep(Duration::from_millis(5));
585 let completed = begin("tools/call:ctx_search");
586 completed.stage(Stage::Handler);
587 completed.stage(Stage::Finalizing);
588 completed.finish(Outcome::Completed);
589 begin("tools/call:ctx_shell").finish(Outcome::Errored);
590 drop(begin("tools/call:ctx_tree"));
592 refused();
593
594 assert_eq!(counts(), (5, 2, 1, 1));
595 let snap = snapshot();
596 assert_eq!(snap.pending, 1);
597 assert_eq!(snap.refused, 1);
598 assert_eq!(snap.untracked_pending, 0);
599 let oldest = snap.oldest_pending.clone().expect("pending tracked");
600 assert_eq!(
601 (oldest.method.as_str(), oldest.stage),
602 ("tools/call:ctx_read", Stage::Handler)
603 );
604 let tool = snap.last_tool_completion.clone().expect("tool completion");
605 assert_eq!(tool.method, "tools/call:ctx_tree");
607 assert_eq!(tool.outcome, Outcome::Abandoned);
608 assert_eq!(
609 snap.last_initialize.as_ref().map(|r| r.outcome),
610 Some(Outcome::Completed)
611 );
612
613 let now = oldest.admitted_at_ms + 74_000;
614 let line = snap.clone().render(now);
615 assert!(
616 line.contains("oldest pending tools/call:ctx_read 74s in handler"),
617 "{line}"
618 );
619 assert!(
620 line.contains(
621 "admitted 5 · completed 2 · errored 1 · refused 1 · abandoned 1 · pending 1"
622 ),
623 "{line}"
624 );
625 assert_eq!(
626 snap.with_ages(now).oldest_pending.and_then(|p| p.age_ms),
627 Some(74_000)
628 );
629
630 stuck.finish(Outcome::Completed);
631 assert_eq!(counts(), (5, 3, 1, 1));
632 assert_eq!(snapshot().pending, 0);
633 }
634
635 #[test]
638 fn finish_serving_withdraws_the_published_snapshot() {
639 let data = tempfile::tempdir().expect("tempdir");
640 let data_dir = data.path().to_str().expect("utf-8 temp path");
641 if crate::test_env::run_in_own_process(
642 "core::request_ledger::tests::finish_serving_withdraws_the_published_snapshot",
643 &[("LEAN_CTX_DATA_DIR", data_dir)],
644 ) {
645 return;
646 }
647 let path = snapshot_dir()
648 .expect("data dir")
649 .join(format!("{}.json", std::process::id()));
650 start_serving("test");
651 let deadline = std::time::Instant::now() + PUBLISH_INTERVAL * 3;
652 while !path.exists() && std::time::Instant::now() < deadline {
653 std::thread::sleep(Duration::from_millis(50));
654 }
655 assert!(path.exists(), "snapshot published at {}", path.display());
656
657 finish_serving(Duration::from_secs(5));
658 assert!(!path.exists(), "snapshot withdrawn");
659 refused();
660 std::thread::sleep(PUBLISH_INTERVAL * 2);
661 assert!(!path.exists(), "publisher wrote the snapshot back");
662 }
663
664 #[test]
665 fn method_names_are_bounded_printable_identifiers() {
666 let long = format!("tools/call:{}\n{}", "x".repeat(200), "secret path");
667 let kept = sanitize_method(&long);
668 assert!(kept.len() <= MAX_METHOD_BYTES);
669 assert!(kept.chars().all(|c| c.is_ascii_graphic()));
670 }
671}