1use std::sync::{Arc, Mutex};
10use std::time::{Instant, SystemTime, UNIX_EPOCH};
11
12use crate::budget::Budget;
13use crate::observer::ExecutionObserver;
14use crate::progress::ProgressInfo;
15use crate::recent_log::{LogEntry, LogSink};
16use crate::tokens::{estimate_tokens, TokenCount, TokenSource};
17use crate::{BudgetHandle, CustomMetrics, CustomMetricsHandle, LlmQuery, ProgressHandle, QueryId};
18
19struct TranscriptEntry {
28 query_id: String,
29 prompt: String,
30 system: Option<String>,
31 response: Option<String>,
32 prompt_tokens: u64,
34 prompt_source: TokenSource,
35 response_tokens: u64,
38 response_source: TokenSource,
39 started_at_ms: i64,
41 completed_at_ms: Option<i64>,
44}
45
46impl TranscriptEntry {
47 fn to_json(&self) -> serde_json::Value {
48 serde_json::json!({
49 "query_id": self.query_id,
50 "prompt": self.prompt,
51 "system": self.system,
52 "response": self.response,
53 })
54 }
55
56 fn to_history_json(&self) -> serde_json::Value {
65 serde_json::json!({
66 "query_id": self.query_id,
67 "prompt": self.prompt,
68 "response": self.response,
69 "prompt_tokens": self.prompt_tokens,
70 "response_tokens": self.response_tokens,
71 "started_at": self.started_at_ms,
72 "completed_at": self.completed_at_ms,
73 })
74 }
75}
76
77pub(crate) struct SessionStatus {
131 started_at: Instant,
132 ended_at: Option<Instant>,
133 pub(crate) llm_calls: u64,
134 pauses: u64,
135 rounds: u64,
136 total_prompt_chars: u64,
137 total_response_chars: u64,
138 transcript: Vec<TranscriptEntry>,
139 pub(crate) budget: Option<Budget>,
140 pub(crate) progress: Option<ProgressInfo>,
141}
142
143impl SessionStatus {
144 fn new() -> Self {
145 Self {
146 started_at: Instant::now(),
147 ended_at: None,
148 llm_calls: 0,
149 pauses: 0,
150 rounds: 0,
151 total_prompt_chars: 0,
152 total_response_chars: 0,
153 transcript: Vec::new(),
154 budget: None,
155 progress: None,
156 }
157 }
158
159 fn prompt_token_count(&self) -> TokenCount {
161 let mut tc = TokenCount::new(TokenSource::Definite);
162 for e in &self.transcript {
163 tc.accumulate(e.prompt_tokens, e.prompt_source);
164 }
165 tc
166 }
167
168 fn response_token_count(&self) -> TokenCount {
170 let mut tc = TokenCount::new(TokenSource::Definite);
171 for e in &self.transcript {
172 tc.accumulate(e.response_tokens, e.response_source);
173 }
174 tc
175 }
176
177 fn total_tokens(&self) -> u64 {
179 self.transcript
180 .iter()
181 .map(|e| e.prompt_tokens + e.response_tokens)
182 .sum()
183 }
184
185 fn elapsed_ms(&self) -> u64 {
187 self.ended_at
188 .map(|end| end.duration_since(self.started_at).as_millis() as u64)
189 .unwrap_or_else(|| self.started_at.elapsed().as_millis() as u64)
190 }
191
192 fn to_json(&self) -> serde_json::Value {
193 let prompt_tc = self.prompt_token_count();
194 let response_tc = self.response_token_count();
195 let total_tc = TokenCount {
196 tokens: prompt_tc.tokens + response_tc.tokens,
197 source: prompt_tc.source.weaker(response_tc.source),
198 };
199 let mut json = serde_json::json!({
200 "elapsed_ms": self.elapsed_ms(),
201 "llm_calls": self.llm_calls,
202 "pauses": self.pauses,
203 "rounds": self.rounds,
204 "total_prompt_chars": self.total_prompt_chars,
205 "total_response_chars": self.total_response_chars,
206 "prompt_tokens": prompt_tc.to_json(),
207 "response_tokens": response_tc.to_json(),
208 "total_tokens": total_tc.to_json(),
209 });
210 if let Some(ref b) = self.budget {
211 json["budget"] = b.to_json();
212 }
213 json
214 }
215
216 pub(crate) fn check_budget(&self) -> Result<(), String> {
217 match self.budget {
218 Some(ref b) => b.check(self.llm_calls, self.elapsed_ms(), self.total_tokens()),
219 None => Ok(()),
220 }
221 }
222
223 fn snapshot(&self, include_history: bool, log_sink: &LogSink) -> serde_json::Value {
242 let prompt_tc = self.prompt_token_count();
244 let response_tc = self.response_token_count();
245 let total_tokens = prompt_tc.tokens + response_tc.tokens;
246
247 let current_query = self.transcript.last().and_then(|e| {
250 if e.response.is_none() {
251 Some(serde_json::json!({
252 "query_id": e.query_id,
253 "prompt_tokens": e.prompt_tokens,
254 "started_waiting_at": e.started_at_ms,
255 }))
256 } else {
257 None
258 }
259 });
260
261 let mut json = serde_json::json!({
262 "elapsed_ms": self.elapsed_ms(),
263 "llm_calls": self.llm_calls,
264 "rounds": self.rounds,
265 "tokens": {
266 "prompt_total": prompt_tc.tokens,
267 "response_total": response_tc.tokens,
268 "total": total_tokens,
269 "current_query": current_query,
270 },
271 "recent_logs": log_sink.to_json(),
272 });
273
274 if let Some(ref p) = self.progress {
275 json["progress"] = serde_json::json!({
276 "step": p.step,
277 "total": p.total,
278 "message": p.message,
279 });
280 }
281
282 if let Some(ref b) = self.budget {
283 json["budget_remaining"] =
284 b.remaining_json(self.llm_calls, self.elapsed_ms(), self.total_tokens());
285 }
286
287 if include_history {
288 let start = self.transcript.len().saturating_sub(10);
290 let history: Vec<serde_json::Value> = self.transcript[start..]
291 .iter()
292 .map(|e| e.to_history_json())
293 .collect();
294 json["conversation_history"] = serde_json::Value::Array(history);
295 }
296
297 json
298 }
299
300 pub(crate) fn budget_remaining(&self) -> serde_json::Value {
301 match self.budget {
302 None => serde_json::Value::Null,
303 Some(ref b) => b.remaining_json(self.llm_calls, self.elapsed_ms(), self.total_tokens()),
304 }
305 }
306}
307
308pub struct ExecutionMetrics {
321 auto: Arc<Mutex<SessionStatus>>,
322 custom: Arc<Mutex<CustomMetrics>>,
323 log_sink: LogSink,
324}
325
326impl ExecutionMetrics {
327 pub fn new() -> Self {
328 Self {
329 auto: Arc::new(Mutex::new(SessionStatus::new())),
330 custom: Arc::new(Mutex::new(CustomMetrics::new())),
331 log_sink: LogSink::new(),
332 }
333 }
334
335 pub fn to_json(&self) -> serde_json::Value {
337 let auto_json = self
338 .auto
339 .lock()
340 .map(|m| m.to_json())
341 .unwrap_or(serde_json::Value::Null);
342
343 let custom_json = self
344 .custom
345 .lock()
346 .map(|m| m.to_json())
347 .unwrap_or(serde_json::Value::Null);
348
349 serde_json::json!({
350 "auto": auto_json,
351 "custom": custom_json,
352 })
353 }
354
355 pub fn transcript_to_json(&self) -> Vec<serde_json::Value> {
357 self.auto
358 .lock()
359 .map(|m| m.transcript.iter().map(|e| e.to_json()).collect())
360 .unwrap_or_default()
361 }
362
363 pub fn custom_metrics_handle(&self) -> CustomMetricsHandle {
365 CustomMetricsHandle::new(Arc::clone(&self.custom))
366 }
367
368 pub fn set_budget(&self, budget: Budget) {
370 if let Ok(mut m) = self.auto.lock() {
371 m.budget = Some(budget);
372 }
373 }
374
375 pub fn budget_handle(&self) -> BudgetHandle {
377 BudgetHandle::new(Arc::clone(&self.auto))
378 }
379
380 pub fn progress_handle(&self) -> ProgressHandle {
382 ProgressHandle::new(Arc::clone(&self.auto))
383 }
384
385 pub fn snapshot(&self, include_history: bool) -> serde_json::Value {
400 self.auto
401 .lock()
402 .map(|m| m.snapshot(include_history, &self.log_sink))
403 .unwrap_or(serde_json::Value::Null)
404 }
405
406 pub fn create_observer(&self) -> MetricsObserver {
407 MetricsObserver::new(Arc::clone(&self.auto), self.log_sink.clone())
408 }
409
410 pub fn log_sink_handle(&self) -> LogSink {
420 self.log_sink.clone()
421 }
422
423 pub fn stats_handle(&self) -> StatsHandle {
426 StatsHandle::new(Arc::clone(&self.auto))
427 }
428
429 pub fn usage_aggregate(&self) -> Option<crate::TokenUsage> {
439 let m = self.auto.lock().ok()?;
440 if m.llm_calls == 0 {
441 return None;
442 }
443 Some(crate::TokenUsage {
444 prompt_tokens: Some(m.prompt_token_count().tokens),
445 completion_tokens: Some(m.response_token_count().tokens),
446 })
447 }
448}
449
450#[derive(Clone)]
463pub struct StatsHandle {
464 auto: Arc<Mutex<SessionStatus>>,
465}
466
467impl StatsHandle {
468 pub(crate) fn new(auto: Arc<Mutex<SessionStatus>>) -> Self {
469 Self { auto }
470 }
471
472 pub fn llm_calls(&self) -> u64 {
480 self.auto.lock().map(|m| m.llm_calls).unwrap_or(0)
481 }
482}
483
484impl Default for ExecutionMetrics {
485 fn default() -> Self {
486 Self::new()
487 }
488}
489
490impl serde::Serialize for ExecutionMetrics {
491 fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
492 self.to_json().serialize(serializer)
493 }
494}
495
496pub struct MetricsObserver {
498 auto: Arc<Mutex<SessionStatus>>,
499 log_sink: LogSink,
500}
501
502impl MetricsObserver {
503 pub(crate) fn new(auto: Arc<Mutex<SessionStatus>>, log_sink: LogSink) -> Self {
504 Self { auto, log_sink }
505 }
506}
507
508impl ExecutionObserver for MetricsObserver {
509 fn on_paused(&self, queries: &[LlmQuery]) {
510 let now_ms = SystemTime::now()
513 .duration_since(UNIX_EPOCH)
514 .unwrap_or_default()
515 .as_millis() as i64;
516 if let Ok(mut m) = self.auto.lock() {
517 m.pauses += 1;
518 m.llm_calls += queries.len() as u64;
519 for q in queries {
520 m.total_prompt_chars += q.prompt.len() as u64;
521 let mut est = estimate_tokens(&q.prompt);
522 if let Some(ref sys) = q.system {
523 m.total_prompt_chars += sys.len() as u64;
524 est += estimate_tokens(sys);
525 }
526 m.transcript.push(TranscriptEntry {
527 query_id: q.id.as_str().to_string(),
528 prompt: q.prompt.clone(),
529 system: q.system.clone(),
530 response: None,
531 prompt_tokens: est,
532 prompt_source: TokenSource::Estimated,
533 response_tokens: 0,
534 response_source: TokenSource::Estimated,
535 started_at_ms: now_ms,
536 completed_at_ms: None,
537 });
538 }
539 }
540 }
541
542 fn on_response_fed(
543 &self,
544 query_id: &QueryId,
545 response: &str,
546 usage: Option<&crate::TokenUsage>,
547 ) {
548 let now_ms = SystemTime::now()
551 .duration_since(UNIX_EPOCH)
552 .unwrap_or_default()
553 .as_millis() as i64;
554 if let Ok(mut m) = self.auto.lock() {
555 m.total_response_chars += response.len() as u64;
556
557 if let Some(entry) = m
558 .transcript
559 .iter_mut()
560 .rev()
561 .find(|e| e.query_id == query_id.as_str())
562 {
563 entry.response = Some(response.to_string());
564 entry.completed_at_ms = Some(now_ms);
565
566 if let Some(pt) = usage.and_then(|u| u.prompt_tokens) {
568 entry.prompt_tokens = pt;
569 entry.prompt_source = TokenSource::Provided;
570 }
571
572 match usage.and_then(|u| u.completion_tokens) {
574 Some(ct) => {
575 entry.response_tokens = ct;
576 entry.response_source = TokenSource::Provided;
577 }
578 None => {
579 entry.response_tokens = estimate_tokens(response);
580 entry.response_source = TokenSource::Estimated;
581 }
582 }
583 }
584 }
585 }
586
587 fn on_log(&self, entry: &LogEntry) {
588 self.log_sink.push(entry.clone());
589 }
590
591 fn on_resumed(&self) {
592 if let Ok(mut m) = self.auto.lock() {
593 m.rounds += 1;
594 }
595 }
596
597 fn on_completed(&self, _result: &serde_json::Value) {
598 if let Ok(mut m) = self.auto.lock() {
599 m.ended_at = Some(Instant::now());
600 }
601 }
602
603 fn on_failed(&self, _error: &str) {
604 if let Ok(mut m) = self.auto.lock() {
605 m.ended_at = Some(Instant::now());
606 }
607 }
608
609 fn on_cancelled(&self) {
610 if let Ok(mut m) = self.auto.lock() {
611 m.ended_at = Some(Instant::now());
612 }
613 }
614}
615
616#[cfg(test)]
617mod tests {
618 use super::*;
619 use crate::{LlmQuery, QueryId};
620
621 #[test]
622 fn metrics_to_json_has_auto_and_custom() {
623 let metrics = ExecutionMetrics::new();
624 let json = metrics.to_json();
625 assert!(json.get("auto").is_some());
626 assert!(json.get("custom").is_some());
627 }
628
629 #[test]
630 fn custom_handle_shares_state() {
631 let metrics = ExecutionMetrics::new();
632 let handle = metrics.custom_metrics_handle();
633
634 handle.record("key".into(), serde_json::json!("value"));
635
636 let json = metrics.to_json();
637 let custom = json.get("custom").unwrap();
638 assert_eq!(custom.get("key").unwrap(), "value");
639 }
640
641 #[test]
642 fn observer_updates_auto_metrics() {
643 let metrics = ExecutionMetrics::new();
644 let observer = metrics.create_observer();
645
646 let queries = vec![LlmQuery {
647 id: QueryId::batch(0),
648 prompt: "test".into(),
649 system: None,
650 max_tokens: 100,
651 grounded: false,
652 underspecified: false,
653 cache_breakpoint: None,
654 role: None,
655 }];
656
657 observer.on_paused(&queries);
658 observer.on_completed(&serde_json::json!(null));
659
660 let json = metrics.to_json();
661 let auto = json.get("auto").unwrap();
662 assert_eq!(auto.get("llm_calls").unwrap(), 1);
663 assert_eq!(auto.get("pauses").unwrap(), 1);
664 assert_eq!(auto.get("rounds").unwrap(), 0);
665 assert_eq!(auto.get("total_prompt_chars").unwrap(), 4); assert_eq!(auto.get("total_response_chars").unwrap(), 0);
667 }
668
669 #[test]
670 fn observer_tracks_prompt_and_response_chars() {
671 let metrics = ExecutionMetrics::new();
672 let observer = metrics.create_observer();
673
674 let queries = vec![
675 LlmQuery {
676 id: QueryId::batch(0),
677 prompt: "hello".into(), system: Some("sys".into()), max_tokens: 100,
680 grounded: false,
681 underspecified: false,
682 cache_breakpoint: None,
683 role: None,
684 },
685 LlmQuery {
686 id: QueryId::batch(1),
687 prompt: "world".into(), system: None,
689 max_tokens: 100,
690 grounded: false,
691 underspecified: false,
692 cache_breakpoint: None,
693 role: None,
694 },
695 ];
696
697 observer.on_paused(&queries);
698 observer.on_response_fed(&QueryId::batch(0), &"x".repeat(42), None);
699 observer.on_response_fed(&QueryId::batch(1), &"y".repeat(58), None);
700 observer.on_resumed();
701 observer.on_completed(&serde_json::json!(null));
702
703 let json = metrics.to_json();
704 let auto = json.get("auto").unwrap();
705 assert_eq!(auto.get("total_prompt_chars").unwrap(), 13); assert_eq!(auto.get("total_response_chars").unwrap(), 100); assert_eq!(auto.get("rounds").unwrap(), 1);
708 }
709
710 #[test]
711 fn observer_tracks_multiple_rounds() {
712 let metrics = ExecutionMetrics::new();
713 let observer = metrics.create_observer();
714
715 let q = vec![LlmQuery {
716 id: QueryId::single(),
717 prompt: "p".into(),
718 system: None,
719 max_tokens: 10,
720 grounded: false,
721 underspecified: false,
722 cache_breakpoint: None,
723 role: None,
724 }];
725
726 observer.on_paused(&q);
728 observer.on_response_fed(&QueryId::single(), &"x".repeat(10), None);
729 observer.on_resumed();
730 observer.on_paused(&q);
732 observer.on_response_fed(&QueryId::single(), &"y".repeat(20), None);
733 observer.on_resumed();
734 observer.on_paused(&q);
736 observer.on_response_fed(&QueryId::single(), &"z".repeat(30), None);
737 observer.on_resumed();
738
739 observer.on_completed(&serde_json::json!(null));
740
741 let json = metrics.to_json();
742 let auto = json.get("auto").unwrap();
743 assert_eq!(auto.get("rounds").unwrap(), 3);
744 assert_eq!(auto.get("pauses").unwrap(), 3);
745 assert_eq!(auto.get("llm_calls").unwrap(), 3);
746 assert_eq!(auto.get("total_prompt_chars").unwrap(), 3); assert_eq!(auto.get("total_response_chars").unwrap(), 60); }
749
750 #[test]
751 fn transcript_records_prompt_response_pairs() {
752 let metrics = ExecutionMetrics::new();
753 let observer = metrics.create_observer();
754
755 let queries = vec![LlmQuery {
756 id: QueryId::single(),
757 prompt: "What is 2+2?".into(),
758 system: Some("You are a calculator.".into()),
759 max_tokens: 50,
760 grounded: false,
761 underspecified: false,
762 cache_breakpoint: None,
763 role: None,
764 }];
765
766 observer.on_paused(&queries);
767 observer.on_response_fed(&QueryId::single(), "4", None);
768 observer.on_resumed();
769 observer.on_completed(&serde_json::json!(null));
770
771 let transcript = metrics.transcript_to_json();
772 assert_eq!(transcript.len(), 1);
773 assert_eq!(transcript[0]["query_id"], "q-0");
774 assert_eq!(transcript[0]["prompt"], "What is 2+2?");
775 assert_eq!(transcript[0]["system"], "You are a calculator.");
776 assert_eq!(transcript[0]["response"], "4");
777 }
778
779 #[test]
780 fn transcript_not_in_stats() {
781 let metrics = ExecutionMetrics::new();
782 let observer = metrics.create_observer();
783 observer.on_paused(&[LlmQuery {
784 id: QueryId::single(),
785 prompt: "p".into(),
786 system: None,
787 max_tokens: 10,
788 grounded: false,
789 underspecified: false,
790 cache_breakpoint: None,
791 role: None,
792 }]);
793 observer.on_response_fed(&QueryId::single(), "r", None);
794 observer.on_resumed();
795 observer.on_completed(&serde_json::json!(null));
796
797 let json = metrics.to_json();
798 assert!(json["auto"].get("transcript").is_none());
799 }
800
801 #[test]
802 fn transcript_multi_round() {
803 let metrics = ExecutionMetrics::new();
804 let observer = metrics.create_observer();
805
806 observer.on_paused(&[LlmQuery {
808 id: QueryId::single(),
809 prompt: "step1".into(),
810 system: None,
811 max_tokens: 100,
812 grounded: false,
813 underspecified: false,
814 cache_breakpoint: None,
815 role: None,
816 }]);
817 observer.on_response_fed(&QueryId::single(), "answer1", None);
818 observer.on_resumed();
819
820 observer.on_paused(&[LlmQuery {
822 id: QueryId::single(),
823 prompt: "step2".into(),
824 system: Some("expert".into()),
825 max_tokens: 100,
826 grounded: false,
827 underspecified: false,
828 cache_breakpoint: None,
829 role: None,
830 }]);
831 observer.on_response_fed(&QueryId::single(), "answer2", None);
832 observer.on_resumed();
833
834 observer.on_completed(&serde_json::json!(null));
835
836 let transcript = metrics.transcript_to_json();
837 assert_eq!(transcript.len(), 2);
838
839 assert_eq!(transcript[0]["prompt"], "step1");
840 assert!(transcript[0]["system"].is_null());
841 assert_eq!(transcript[0]["response"], "answer1");
842
843 assert_eq!(transcript[1]["prompt"], "step2");
844 assert_eq!(transcript[1]["system"], "expert");
845 assert_eq!(transcript[1]["response"], "answer2");
846 }
847
848 #[test]
849 fn transcript_batch_queries() {
850 let metrics = ExecutionMetrics::new();
851 let observer = metrics.create_observer();
852
853 let queries = vec![
854 LlmQuery {
855 id: QueryId::batch(0),
856 prompt: "q0".into(),
857 system: None,
858 max_tokens: 50,
859 grounded: false,
860 underspecified: false,
861 cache_breakpoint: None,
862 role: None,
863 },
864 LlmQuery {
865 id: QueryId::batch(1),
866 prompt: "q1".into(),
867 system: None,
868 max_tokens: 50,
869 grounded: false,
870 underspecified: false,
871 cache_breakpoint: None,
872 role: None,
873 },
874 ];
875
876 observer.on_paused(&queries);
877 observer.on_response_fed(&QueryId::batch(0), "r0", None);
878 observer.on_response_fed(&QueryId::batch(1), "r1", None);
879 observer.on_resumed();
880 observer.on_completed(&serde_json::json!(null));
881
882 let transcript = metrics.transcript_to_json();
883 assert_eq!(transcript.len(), 2);
884 assert_eq!(transcript[0]["query_id"], "q-0");
885 assert_eq!(transcript[0]["response"], "r0");
886 assert_eq!(transcript[1]["query_id"], "q-1");
887 assert_eq!(transcript[1]["response"], "r1");
888 }
889
890 #[test]
894 fn on_log_routes_to_log_sink() {
895 let metrics = ExecutionMetrics::new();
896 let observer = metrics.create_observer();
897
898 observer.on_log(&crate::LogEntry::new("info", "engine", "hello"));
899 observer.on_log(&crate::LogEntry::new("warn", "alc.log", "world"));
900
901 let sink = metrics.log_sink_handle();
902 let entries = sink.entries();
903 assert_eq!(entries.len(), 2);
904 assert_eq!(entries[0].level, "info");
905 assert_eq!(entries[0].source, "engine");
906 assert_eq!(entries[0].message, "hello");
907 assert_eq!(entries[1].level, "warn");
908 assert_eq!(entries[1].message, "world");
909 }
910
911 #[test]
913 fn on_log_cap_enforcement_via_observer() {
914 let metrics = ExecutionMetrics::new();
915 let observer = metrics.create_observer();
916
917 for i in 0..=20u32 {
918 observer.on_log(&crate::LogEntry::new("info", "engine", format!("msg-{i}")));
919 }
920
921 let sink = metrics.log_sink_handle();
922 let entries = sink.entries();
923 assert_eq!(entries.len(), crate::recent_log::LOG_SINK_CAP);
924 assert_eq!(entries[0].message, "msg-1");
925 assert_eq!(
926 entries[crate::recent_log::LOG_SINK_CAP - 1].message,
927 "msg-20"
928 );
929 }
930
931 #[test]
934 fn transcript_timestamps_recorded() {
935 let metrics = ExecutionMetrics::new();
936 let observer = metrics.create_observer();
937
938 let before = std::time::SystemTime::now()
939 .duration_since(std::time::UNIX_EPOCH)
940 .unwrap_or_default()
941 .as_millis() as i64;
942
943 observer.on_paused(&[LlmQuery {
944 id: QueryId::single(),
945 prompt: "ts-test".into(),
946 system: None,
947 max_tokens: 10,
948 grounded: false,
949 underspecified: false,
950 cache_breakpoint: None,
951 role: None,
952 }]);
953
954 observer.on_response_fed(&QueryId::single(), "response", None);
955
956 let after_fed = std::time::SystemTime::now()
957 .duration_since(std::time::UNIX_EPOCH)
958 .unwrap_or_default()
959 .as_millis() as i64;
960
961 let snap = metrics.snapshot(true);
963 let history = snap["conversation_history"]
964 .as_array()
965 .expect("conversation_history must be array");
966 assert_eq!(history.len(), 1);
967
968 let started_at = history[0]["started_at"]
969 .as_i64()
970 .expect("started_at must be i64");
971 let completed_at = history[0]["completed_at"]
972 .as_i64()
973 .expect("completed_at must be i64 (not null)");
974
975 assert!(
976 started_at >= before,
977 "started_at ({started_at}) should be >= before ({before})"
978 );
979 assert!(
980 completed_at >= started_at,
981 "completed_at ({completed_at}) should be >= started_at ({started_at})"
982 );
983 assert!(
984 completed_at <= after_fed,
985 "completed_at ({completed_at}) should be <= after_fed ({after_fed})"
986 );
987 }
988
989 #[test]
991 fn snapshot_current_query_while_paused() {
992 let metrics = ExecutionMetrics::new();
993 let observer = metrics.create_observer();
994
995 observer.on_paused(&[LlmQuery {
996 id: QueryId::single(),
997 prompt: "in-flight".into(),
998 system: None,
999 max_tokens: 10,
1000 grounded: false,
1001 underspecified: false,
1002 cache_breakpoint: None,
1003 role: None,
1004 }]);
1005
1006 let snap = metrics.snapshot(false);
1008
1009 let tokens = snap.get("tokens").expect("tokens field must be present");
1010 let current_query = tokens
1011 .get("current_query")
1012 .expect("current_query must be present");
1013 assert!(
1014 !current_query.is_null(),
1015 "current_query should be non-null while paused"
1016 );
1017 assert_eq!(current_query["query_id"], "q-0");
1018 assert!(
1020 snap.get("conversation_history").is_none(),
1021 "conversation_history must be absent when include_history=false"
1022 );
1023 }
1024
1025 #[test]
1027 fn snapshot_current_query_null_after_response() {
1028 let metrics = ExecutionMetrics::new();
1029 let observer = metrics.create_observer();
1030
1031 observer.on_paused(&[LlmQuery {
1032 id: QueryId::single(),
1033 prompt: "done".into(),
1034 system: None,
1035 max_tokens: 10,
1036 grounded: false,
1037 underspecified: false,
1038 cache_breakpoint: None,
1039 role: None,
1040 }]);
1041 observer.on_response_fed(&QueryId::single(), "answer", None);
1042
1043 let snap = metrics.snapshot(false);
1044 let tokens = snap.get("tokens").expect("tokens must be present");
1045 let current_query = &tokens["current_query"];
1046 assert!(
1047 current_query.is_null(),
1048 "current_query should be null after response is fed"
1049 );
1050 }
1051
1052 #[test]
1054 fn snapshot_conversation_history_opt_in() {
1055 let metrics = ExecutionMetrics::new();
1056 let observer = metrics.create_observer();
1057
1058 observer.on_paused(&[LlmQuery {
1059 id: QueryId::single(),
1060 prompt: "hello".into(),
1061 system: None,
1062 max_tokens: 50,
1063 grounded: false,
1064 underspecified: false,
1065 cache_breakpoint: None,
1066 role: None,
1067 }]);
1068 observer.on_response_fed(&QueryId::single(), "world", None);
1069 observer.on_resumed();
1070 observer.on_completed(&serde_json::json!(null));
1071
1072 let snap_false = metrics.snapshot(false);
1074 assert!(
1075 snap_false.get("conversation_history").is_none(),
1076 "conversation_history must be absent with include_history=false"
1077 );
1078
1079 let snap_true = metrics.snapshot(true);
1081 let history = snap_true
1082 .get("conversation_history")
1083 .expect("conversation_history must be present with include_history=true");
1084 let arr = history
1085 .as_array()
1086 .expect("conversation_history must be an array");
1087 assert_eq!(arr.len(), 1);
1088 assert_eq!(arr[0]["query_id"], "q-0");
1089 assert_eq!(arr[0]["prompt"], "hello");
1090 assert_eq!(arr[0]["response"], "world");
1091 assert!(arr[0].get("started_at").is_some());
1093 assert!(arr[0].get("completed_at").is_some());
1094 }
1095
1096 #[test]
1098 fn snapshot_conversation_history_capped_at_10() {
1099 let metrics = ExecutionMetrics::new();
1100 let observer = metrics.create_observer();
1101
1102 for i in 0..15u32 {
1103 observer.on_paused(&[LlmQuery {
1104 id: QueryId::single(),
1105 prompt: format!("prompt-{i}"),
1106 system: None,
1107 max_tokens: 10,
1108 grounded: false,
1109 underspecified: false,
1110 cache_breakpoint: None,
1111 role: None,
1112 }]);
1113 observer.on_response_fed(&QueryId::single(), &format!("resp-{i}"), None);
1114 observer.on_resumed();
1115 }
1116
1117 let snap = metrics.snapshot(true);
1118 let history = snap["conversation_history"]
1119 .as_array()
1120 .expect("must be array");
1121 assert_eq!(history.len(), 10, "capped at 10 entries");
1122 assert_eq!(history[0]["prompt"], "prompt-5");
1124 assert_eq!(history[9]["prompt"], "prompt-14");
1125 }
1126
1127 #[test]
1129 fn snapshot_includes_recent_logs() {
1130 let metrics = ExecutionMetrics::new();
1131 let observer = metrics.create_observer();
1132 observer.on_log(&crate::LogEntry::new("info", "engine", "test-log"));
1133
1134 let snap = metrics.snapshot(false);
1135 let logs = snap
1136 .get("recent_logs")
1137 .expect("recent_logs must be in snapshot");
1138 let arr = logs.as_array().expect("recent_logs must be array");
1139 assert_eq!(arr.len(), 1);
1140 assert_eq!(arr[0]["message"], "test-log");
1141 }
1142
1143 #[test]
1145 fn snapshot_tokens_aggregate() {
1146 let metrics = ExecutionMetrics::new();
1147 let observer = metrics.create_observer();
1148
1149 observer.on_paused(&[LlmQuery {
1150 id: QueryId::single(),
1151 prompt: "x".repeat(100),
1152 system: None,
1153 max_tokens: 50,
1154 grounded: false,
1155 underspecified: false,
1156 cache_breakpoint: None,
1157 role: None,
1158 }]);
1159 observer.on_response_fed(&QueryId::single(), &"y".repeat(50), None);
1160 observer.on_resumed();
1161
1162 let snap = metrics.snapshot(false);
1163 let tokens = snap.get("tokens").expect("tokens must be in snapshot");
1164 let prompt_total = tokens["prompt_total"]
1165 .as_u64()
1166 .expect("prompt_total must be u64");
1167 let response_total = tokens["response_total"]
1168 .as_u64()
1169 .expect("response_total must be u64");
1170 let total = tokens["total"].as_u64().expect("total must be u64");
1171 assert!(prompt_total > 0, "prompt_total must be positive");
1173 assert!(response_total > 0, "response_total must be positive");
1174 assert_eq!(total, prompt_total + response_total);
1175 }
1176
1177 #[test]
1182 fn usage_aggregate_none_when_no_llm_calls() {
1183 let metrics = ExecutionMetrics::new();
1184 assert!(
1185 metrics.usage_aggregate().is_none(),
1186 "fresh metrics with no LLM calls must return None"
1187 );
1188 }
1189
1190 #[test]
1191 fn usage_aggregate_some_when_llm_calls_recorded() {
1192 use crate::TokenUsage;
1193
1194 let metrics = ExecutionMetrics::new();
1195 let observer = metrics.create_observer();
1196
1197 let queries = vec![LlmQuery {
1198 id: QueryId::batch(0),
1199 prompt: "test".into(),
1200 system: None,
1201 max_tokens: 100,
1202 grounded: false,
1203 underspecified: false,
1204 cache_breakpoint: None,
1205 role: None,
1206 }];
1207 observer.on_paused(&queries);
1208 observer.on_response_fed(
1209 &QueryId::batch(0),
1210 "ans",
1211 Some(&TokenUsage {
1212 prompt_tokens: Some(10),
1213 completion_tokens: Some(5),
1214 }),
1215 );
1216
1217 let result = metrics.usage_aggregate();
1218 assert!(
1219 result.is_some(),
1220 "usage_aggregate must return Some after LLM call"
1221 );
1222 let usage = result.unwrap();
1223 assert_eq!(
1224 usage.prompt_tokens,
1225 Some(10),
1226 "prompt_tokens must match provided value"
1227 );
1228 assert_eq!(
1229 usage.completion_tokens,
1230 Some(5),
1231 "completion_tokens must match provided value"
1232 );
1233 }
1234
1235 #[test]
1236 fn usage_aggregate_some_with_estimated_path() {
1237 let metrics = ExecutionMetrics::new();
1241 let observer = metrics.create_observer();
1242
1243 let queries = vec![LlmQuery {
1244 id: QueryId::single(),
1245 prompt: "hello world".into(),
1246 system: None,
1247 max_tokens: 50,
1248 grounded: false,
1249 underspecified: false,
1250 cache_breakpoint: None,
1251 role: None,
1252 }];
1253 observer.on_paused(&queries);
1254 observer.on_response_fed(&QueryId::single(), "response text", None);
1256
1257 let result = metrics.usage_aggregate();
1258 assert!(
1259 result.is_some(),
1260 "usage_aggregate must return Some even when token counts are estimated"
1261 );
1262 let usage = result.unwrap();
1263 assert!(
1264 usage.prompt_tokens.is_some(),
1265 "prompt_tokens must be Some (estimated)"
1266 );
1267 assert!(
1268 usage.completion_tokens.is_some(),
1269 "completion_tokens must be Some (estimated)"
1270 );
1271 }
1272}