1use std::sync::{Arc, Mutex};
2use std::time::{Instant, SystemTime, UNIX_EPOCH};
3
4use crate::budget::Budget;
5use crate::observer::ExecutionObserver;
6use crate::progress::ProgressInfo;
7use crate::recent_log::{LogEntry, LogSink};
8use crate::tokens::{estimate_tokens, TokenCount, TokenSource};
9use crate::{BudgetHandle, CustomMetrics, CustomMetricsHandle, LlmQuery, ProgressHandle, QueryId};
10
11struct TranscriptEntry {
20 query_id: String,
21 prompt: String,
22 system: Option<String>,
23 response: Option<String>,
24 prompt_tokens: u64,
26 prompt_source: TokenSource,
27 response_tokens: u64,
30 response_source: TokenSource,
31 started_at_ms: i64,
33 completed_at_ms: Option<i64>,
36}
37
38impl TranscriptEntry {
39 fn to_json(&self) -> serde_json::Value {
40 serde_json::json!({
41 "query_id": self.query_id,
42 "prompt": self.prompt,
43 "system": self.system,
44 "response": self.response,
45 })
46 }
47
48 fn to_history_json(&self) -> serde_json::Value {
57 serde_json::json!({
58 "query_id": self.query_id,
59 "prompt": self.prompt,
60 "response": self.response,
61 "prompt_tokens": self.prompt_tokens,
62 "response_tokens": self.response_tokens,
63 "started_at": self.started_at_ms,
64 "completed_at": self.completed_at_ms,
65 })
66 }
67}
68
69pub(crate) struct SessionStatus {
123 started_at: Instant,
124 ended_at: Option<Instant>,
125 pub(crate) llm_calls: u64,
126 pauses: u64,
127 rounds: u64,
128 total_prompt_chars: u64,
129 total_response_chars: u64,
130 transcript: Vec<TranscriptEntry>,
131 pub(crate) budget: Option<Budget>,
132 pub(crate) progress: Option<ProgressInfo>,
133}
134
135impl SessionStatus {
136 fn new() -> Self {
137 Self {
138 started_at: Instant::now(),
139 ended_at: None,
140 llm_calls: 0,
141 pauses: 0,
142 rounds: 0,
143 total_prompt_chars: 0,
144 total_response_chars: 0,
145 transcript: Vec::new(),
146 budget: None,
147 progress: None,
148 }
149 }
150
151 fn prompt_token_count(&self) -> TokenCount {
153 let mut tc = TokenCount::new(TokenSource::Definite);
154 for e in &self.transcript {
155 tc.accumulate(e.prompt_tokens, e.prompt_source);
156 }
157 tc
158 }
159
160 fn response_token_count(&self) -> TokenCount {
162 let mut tc = TokenCount::new(TokenSource::Definite);
163 for e in &self.transcript {
164 tc.accumulate(e.response_tokens, e.response_source);
165 }
166 tc
167 }
168
169 fn total_tokens(&self) -> u64 {
171 self.transcript
172 .iter()
173 .map(|e| e.prompt_tokens + e.response_tokens)
174 .sum()
175 }
176
177 fn elapsed_ms(&self) -> u64 {
179 self.ended_at
180 .map(|end| end.duration_since(self.started_at).as_millis() as u64)
181 .unwrap_or_else(|| self.started_at.elapsed().as_millis() as u64)
182 }
183
184 fn to_json(&self) -> serde_json::Value {
185 let prompt_tc = self.prompt_token_count();
186 let response_tc = self.response_token_count();
187 let total_tc = TokenCount {
188 tokens: prompt_tc.tokens + response_tc.tokens,
189 source: prompt_tc.source.weaker(response_tc.source),
190 };
191 let mut json = serde_json::json!({
192 "elapsed_ms": self.elapsed_ms(),
193 "llm_calls": self.llm_calls,
194 "pauses": self.pauses,
195 "rounds": self.rounds,
196 "total_prompt_chars": self.total_prompt_chars,
197 "total_response_chars": self.total_response_chars,
198 "prompt_tokens": prompt_tc.to_json(),
199 "response_tokens": response_tc.to_json(),
200 "total_tokens": total_tc.to_json(),
201 });
202 if let Some(ref b) = self.budget {
203 json["budget"] = b.to_json();
204 }
205 json
206 }
207
208 pub(crate) fn check_budget(&self) -> Result<(), String> {
209 match self.budget {
210 Some(ref b) => b.check(self.llm_calls, self.elapsed_ms(), self.total_tokens()),
211 None => Ok(()),
212 }
213 }
214
215 fn snapshot(&self, include_history: bool, log_sink: &LogSink) -> serde_json::Value {
235 let prompt_tc = self.prompt_token_count();
237 let response_tc = self.response_token_count();
238 let total_tokens = prompt_tc.tokens + response_tc.tokens;
239
240 let current_query = self.transcript.last().and_then(|e| {
243 if e.response.is_none() {
244 Some(serde_json::json!({
245 "query_id": e.query_id,
246 "prompt_tokens": e.prompt_tokens,
247 "started_waiting_at": e.started_at_ms,
248 }))
249 } else {
250 None
251 }
252 });
253
254 let mut json = serde_json::json!({
255 "elapsed_ms": self.elapsed_ms(),
256 "llm_calls": self.llm_calls,
257 "rounds": self.rounds,
258 "tokens": {
259 "prompt_total": prompt_tc.tokens,
260 "response_total": response_tc.tokens,
261 "total": total_tokens,
262 "current_query": current_query,
263 },
264 "recent_logs": log_sink.to_json(),
265 });
266
267 if let Some(ref p) = self.progress {
268 json["progress"] = serde_json::json!({
269 "step": p.step,
270 "total": p.total,
271 "message": p.message,
272 });
273 }
274
275 if let Some(ref b) = self.budget {
276 json["budget_remaining"] =
277 b.remaining_json(self.llm_calls, self.elapsed_ms(), self.total_tokens());
278 }
279
280 if include_history {
281 let start = self.transcript.len().saturating_sub(10);
283 let history: Vec<serde_json::Value> = self.transcript[start..]
284 .iter()
285 .map(|e| e.to_history_json())
286 .collect();
287 json["conversation_history"] = serde_json::Value::Array(history);
288 }
289
290 json
291 }
292
293 pub(crate) fn budget_remaining(&self) -> serde_json::Value {
294 match self.budget {
295 None => serde_json::Value::Null,
296 Some(ref b) => b.remaining_json(self.llm_calls, self.elapsed_ms(), self.total_tokens()),
297 }
298 }
299}
300
301pub struct ExecutionMetrics {
314 auto: Arc<Mutex<SessionStatus>>,
315 custom: Arc<Mutex<CustomMetrics>>,
316 log_sink: LogSink,
317}
318
319impl ExecutionMetrics {
320 pub fn new() -> Self {
321 Self {
322 auto: Arc::new(Mutex::new(SessionStatus::new())),
323 custom: Arc::new(Mutex::new(CustomMetrics::new())),
324 log_sink: LogSink::new(),
325 }
326 }
327
328 pub fn to_json(&self) -> serde_json::Value {
330 let auto_json = self
331 .auto
332 .lock()
333 .map(|m| m.to_json())
334 .unwrap_or(serde_json::Value::Null);
335
336 let custom_json = self
337 .custom
338 .lock()
339 .map(|m| m.to_json())
340 .unwrap_or(serde_json::Value::Null);
341
342 serde_json::json!({
343 "auto": auto_json,
344 "custom": custom_json,
345 })
346 }
347
348 pub fn transcript_to_json(&self) -> Vec<serde_json::Value> {
350 self.auto
351 .lock()
352 .map(|m| m.transcript.iter().map(|e| e.to_json()).collect())
353 .unwrap_or_default()
354 }
355
356 pub fn custom_metrics_handle(&self) -> CustomMetricsHandle {
358 CustomMetricsHandle::new(Arc::clone(&self.custom))
359 }
360
361 pub fn set_budget(&self, budget: Budget) {
363 if let Ok(mut m) = self.auto.lock() {
364 m.budget = Some(budget);
365 }
366 }
367
368 pub fn budget_handle(&self) -> BudgetHandle {
370 BudgetHandle::new(Arc::clone(&self.auto))
371 }
372
373 pub fn progress_handle(&self) -> ProgressHandle {
375 ProgressHandle::new(Arc::clone(&self.auto))
376 }
377
378 pub fn snapshot(&self, include_history: bool) -> serde_json::Value {
393 self.auto
394 .lock()
395 .map(|m| m.snapshot(include_history, &self.log_sink))
396 .unwrap_or(serde_json::Value::Null)
397 }
398
399 pub fn create_observer(&self) -> MetricsObserver {
400 MetricsObserver::new(Arc::clone(&self.auto), self.log_sink.clone())
401 }
402
403 pub fn log_sink_handle(&self) -> LogSink {
413 self.log_sink.clone()
414 }
415
416 pub fn stats_handle(&self) -> StatsHandle {
419 StatsHandle::new(Arc::clone(&self.auto))
420 }
421}
422
423#[derive(Clone)]
436pub struct StatsHandle {
437 auto: Arc<Mutex<SessionStatus>>,
438}
439
440impl StatsHandle {
441 pub(crate) fn new(auto: Arc<Mutex<SessionStatus>>) -> Self {
442 Self { auto }
443 }
444
445 pub fn llm_calls(&self) -> u64 {
453 self.auto.lock().map(|m| m.llm_calls).unwrap_or(0)
454 }
455}
456
457impl Default for ExecutionMetrics {
458 fn default() -> Self {
459 Self::new()
460 }
461}
462
463impl serde::Serialize for ExecutionMetrics {
464 fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
465 self.to_json().serialize(serializer)
466 }
467}
468
469pub struct MetricsObserver {
471 auto: Arc<Mutex<SessionStatus>>,
472 log_sink: LogSink,
473}
474
475impl MetricsObserver {
476 pub(crate) fn new(auto: Arc<Mutex<SessionStatus>>, log_sink: LogSink) -> Self {
477 Self { auto, log_sink }
478 }
479}
480
481impl ExecutionObserver for MetricsObserver {
482 fn on_paused(&self, queries: &[LlmQuery]) {
483 let now_ms = SystemTime::now()
486 .duration_since(UNIX_EPOCH)
487 .unwrap_or_default()
488 .as_millis() as i64;
489 if let Ok(mut m) = self.auto.lock() {
490 m.pauses += 1;
491 m.llm_calls += queries.len() as u64;
492 for q in queries {
493 m.total_prompt_chars += q.prompt.len() as u64;
494 let mut est = estimate_tokens(&q.prompt);
495 if let Some(ref sys) = q.system {
496 m.total_prompt_chars += sys.len() as u64;
497 est += estimate_tokens(sys);
498 }
499 m.transcript.push(TranscriptEntry {
500 query_id: q.id.as_str().to_string(),
501 prompt: q.prompt.clone(),
502 system: q.system.clone(),
503 response: None,
504 prompt_tokens: est,
505 prompt_source: TokenSource::Estimated,
506 response_tokens: 0,
507 response_source: TokenSource::Estimated,
508 started_at_ms: now_ms,
509 completed_at_ms: None,
510 });
511 }
512 }
513 }
514
515 fn on_response_fed(
516 &self,
517 query_id: &QueryId,
518 response: &str,
519 usage: Option<&crate::TokenUsage>,
520 ) {
521 let now_ms = SystemTime::now()
524 .duration_since(UNIX_EPOCH)
525 .unwrap_or_default()
526 .as_millis() as i64;
527 if let Ok(mut m) = self.auto.lock() {
528 m.total_response_chars += response.len() as u64;
529
530 if let Some(entry) = m
531 .transcript
532 .iter_mut()
533 .rev()
534 .find(|e| e.query_id == query_id.as_str())
535 {
536 entry.response = Some(response.to_string());
537 entry.completed_at_ms = Some(now_ms);
538
539 if let Some(pt) = usage.and_then(|u| u.prompt_tokens) {
541 entry.prompt_tokens = pt;
542 entry.prompt_source = TokenSource::Provided;
543 }
544
545 match usage.and_then(|u| u.completion_tokens) {
547 Some(ct) => {
548 entry.response_tokens = ct;
549 entry.response_source = TokenSource::Provided;
550 }
551 None => {
552 entry.response_tokens = estimate_tokens(response);
553 entry.response_source = TokenSource::Estimated;
554 }
555 }
556 }
557 }
558 }
559
560 fn on_log(&self, entry: &LogEntry) {
561 self.log_sink.push(entry.clone());
562 }
563
564 fn on_resumed(&self) {
565 if let Ok(mut m) = self.auto.lock() {
566 m.rounds += 1;
567 }
568 }
569
570 fn on_completed(&self, _result: &serde_json::Value) {
571 if let Ok(mut m) = self.auto.lock() {
572 m.ended_at = Some(Instant::now());
573 }
574 }
575
576 fn on_failed(&self, _error: &str) {
577 if let Ok(mut m) = self.auto.lock() {
578 m.ended_at = Some(Instant::now());
579 }
580 }
581
582 fn on_cancelled(&self) {
583 if let Ok(mut m) = self.auto.lock() {
584 m.ended_at = Some(Instant::now());
585 }
586 }
587}
588
589#[cfg(test)]
590mod tests {
591 use super::*;
592 use crate::{LlmQuery, QueryId};
593
594 #[test]
595 fn metrics_to_json_has_auto_and_custom() {
596 let metrics = ExecutionMetrics::new();
597 let json = metrics.to_json();
598 assert!(json.get("auto").is_some());
599 assert!(json.get("custom").is_some());
600 }
601
602 #[test]
603 fn custom_handle_shares_state() {
604 let metrics = ExecutionMetrics::new();
605 let handle = metrics.custom_metrics_handle();
606
607 handle.record("key".into(), serde_json::json!("value"));
608
609 let json = metrics.to_json();
610 let custom = json.get("custom").unwrap();
611 assert_eq!(custom.get("key").unwrap(), "value");
612 }
613
614 #[test]
615 fn observer_updates_auto_metrics() {
616 let metrics = ExecutionMetrics::new();
617 let observer = metrics.create_observer();
618
619 let queries = vec![LlmQuery {
620 id: QueryId::batch(0),
621 prompt: "test".into(),
622 system: None,
623 max_tokens: 100,
624 grounded: false,
625 underspecified: false,
626 }];
627
628 observer.on_paused(&queries);
629 observer.on_completed(&serde_json::json!(null));
630
631 let json = metrics.to_json();
632 let auto = json.get("auto").unwrap();
633 assert_eq!(auto.get("llm_calls").unwrap(), 1);
634 assert_eq!(auto.get("pauses").unwrap(), 1);
635 assert_eq!(auto.get("rounds").unwrap(), 0);
636 assert_eq!(auto.get("total_prompt_chars").unwrap(), 4); assert_eq!(auto.get("total_response_chars").unwrap(), 0);
638 }
639
640 #[test]
641 fn observer_tracks_prompt_and_response_chars() {
642 let metrics = ExecutionMetrics::new();
643 let observer = metrics.create_observer();
644
645 let queries = vec![
646 LlmQuery {
647 id: QueryId::batch(0),
648 prompt: "hello".into(), system: Some("sys".into()), max_tokens: 100,
651 grounded: false,
652 underspecified: false,
653 },
654 LlmQuery {
655 id: QueryId::batch(1),
656 prompt: "world".into(), system: None,
658 max_tokens: 100,
659 grounded: false,
660 underspecified: false,
661 },
662 ];
663
664 observer.on_paused(&queries);
665 observer.on_response_fed(&QueryId::batch(0), &"x".repeat(42), None);
666 observer.on_response_fed(&QueryId::batch(1), &"y".repeat(58), None);
667 observer.on_resumed();
668 observer.on_completed(&serde_json::json!(null));
669
670 let json = metrics.to_json();
671 let auto = json.get("auto").unwrap();
672 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);
675 }
676
677 #[test]
678 fn observer_tracks_multiple_rounds() {
679 let metrics = ExecutionMetrics::new();
680 let observer = metrics.create_observer();
681
682 let q = vec![LlmQuery {
683 id: QueryId::single(),
684 prompt: "p".into(),
685 system: None,
686 max_tokens: 10,
687 grounded: false,
688 underspecified: false,
689 }];
690
691 observer.on_paused(&q);
693 observer.on_response_fed(&QueryId::single(), &"x".repeat(10), None);
694 observer.on_resumed();
695 observer.on_paused(&q);
697 observer.on_response_fed(&QueryId::single(), &"y".repeat(20), None);
698 observer.on_resumed();
699 observer.on_paused(&q);
701 observer.on_response_fed(&QueryId::single(), &"z".repeat(30), None);
702 observer.on_resumed();
703
704 observer.on_completed(&serde_json::json!(null));
705
706 let json = metrics.to_json();
707 let auto = json.get("auto").unwrap();
708 assert_eq!(auto.get("rounds").unwrap(), 3);
709 assert_eq!(auto.get("pauses").unwrap(), 3);
710 assert_eq!(auto.get("llm_calls").unwrap(), 3);
711 assert_eq!(auto.get("total_prompt_chars").unwrap(), 3); assert_eq!(auto.get("total_response_chars").unwrap(), 60); }
714
715 #[test]
716 fn transcript_records_prompt_response_pairs() {
717 let metrics = ExecutionMetrics::new();
718 let observer = metrics.create_observer();
719
720 let queries = vec![LlmQuery {
721 id: QueryId::single(),
722 prompt: "What is 2+2?".into(),
723 system: Some("You are a calculator.".into()),
724 max_tokens: 50,
725 grounded: false,
726 underspecified: false,
727 }];
728
729 observer.on_paused(&queries);
730 observer.on_response_fed(&QueryId::single(), "4", None);
731 observer.on_resumed();
732 observer.on_completed(&serde_json::json!(null));
733
734 let transcript = metrics.transcript_to_json();
735 assert_eq!(transcript.len(), 1);
736 assert_eq!(transcript[0]["query_id"], "q-0");
737 assert_eq!(transcript[0]["prompt"], "What is 2+2?");
738 assert_eq!(transcript[0]["system"], "You are a calculator.");
739 assert_eq!(transcript[0]["response"], "4");
740 }
741
742 #[test]
743 fn transcript_not_in_stats() {
744 let metrics = ExecutionMetrics::new();
745 let observer = metrics.create_observer();
746 observer.on_paused(&[LlmQuery {
747 id: QueryId::single(),
748 prompt: "p".into(),
749 system: None,
750 max_tokens: 10,
751 grounded: false,
752 underspecified: false,
753 }]);
754 observer.on_response_fed(&QueryId::single(), "r", None);
755 observer.on_resumed();
756 observer.on_completed(&serde_json::json!(null));
757
758 let json = metrics.to_json();
759 assert!(json["auto"].get("transcript").is_none());
760 }
761
762 #[test]
763 fn transcript_multi_round() {
764 let metrics = ExecutionMetrics::new();
765 let observer = metrics.create_observer();
766
767 observer.on_paused(&[LlmQuery {
769 id: QueryId::single(),
770 prompt: "step1".into(),
771 system: None,
772 max_tokens: 100,
773 grounded: false,
774 underspecified: false,
775 }]);
776 observer.on_response_fed(&QueryId::single(), "answer1", None);
777 observer.on_resumed();
778
779 observer.on_paused(&[LlmQuery {
781 id: QueryId::single(),
782 prompt: "step2".into(),
783 system: Some("expert".into()),
784 max_tokens: 100,
785 grounded: false,
786 underspecified: false,
787 }]);
788 observer.on_response_fed(&QueryId::single(), "answer2", None);
789 observer.on_resumed();
790
791 observer.on_completed(&serde_json::json!(null));
792
793 let transcript = metrics.transcript_to_json();
794 assert_eq!(transcript.len(), 2);
795
796 assert_eq!(transcript[0]["prompt"], "step1");
797 assert!(transcript[0]["system"].is_null());
798 assert_eq!(transcript[0]["response"], "answer1");
799
800 assert_eq!(transcript[1]["prompt"], "step2");
801 assert_eq!(transcript[1]["system"], "expert");
802 assert_eq!(transcript[1]["response"], "answer2");
803 }
804
805 #[test]
806 fn transcript_batch_queries() {
807 let metrics = ExecutionMetrics::new();
808 let observer = metrics.create_observer();
809
810 let queries = vec![
811 LlmQuery {
812 id: QueryId::batch(0),
813 prompt: "q0".into(),
814 system: None,
815 max_tokens: 50,
816 grounded: false,
817 underspecified: false,
818 },
819 LlmQuery {
820 id: QueryId::batch(1),
821 prompt: "q1".into(),
822 system: None,
823 max_tokens: 50,
824 grounded: false,
825 underspecified: false,
826 },
827 ];
828
829 observer.on_paused(&queries);
830 observer.on_response_fed(&QueryId::batch(0), "r0", None);
831 observer.on_response_fed(&QueryId::batch(1), "r1", None);
832 observer.on_resumed();
833 observer.on_completed(&serde_json::json!(null));
834
835 let transcript = metrics.transcript_to_json();
836 assert_eq!(transcript.len(), 2);
837 assert_eq!(transcript[0]["query_id"], "q-0");
838 assert_eq!(transcript[0]["response"], "r0");
839 assert_eq!(transcript[1]["query_id"], "q-1");
840 assert_eq!(transcript[1]["response"], "r1");
841 }
842
843 #[test]
847 fn on_log_routes_to_log_sink() {
848 let metrics = ExecutionMetrics::new();
849 let observer = metrics.create_observer();
850
851 observer.on_log(&crate::LogEntry::new("info", "engine", "hello"));
852 observer.on_log(&crate::LogEntry::new("warn", "alc.log", "world"));
853
854 let sink = metrics.log_sink_handle();
855 let entries = sink.entries();
856 assert_eq!(entries.len(), 2);
857 assert_eq!(entries[0].level, "info");
858 assert_eq!(entries[0].source, "engine");
859 assert_eq!(entries[0].message, "hello");
860 assert_eq!(entries[1].level, "warn");
861 assert_eq!(entries[1].message, "world");
862 }
863
864 #[test]
866 fn on_log_cap_enforcement_via_observer() {
867 let metrics = ExecutionMetrics::new();
868 let observer = metrics.create_observer();
869
870 for i in 0..=20u32 {
871 observer.on_log(&crate::LogEntry::new("info", "engine", format!("msg-{i}")));
872 }
873
874 let sink = metrics.log_sink_handle();
875 let entries = sink.entries();
876 assert_eq!(entries.len(), crate::recent_log::LOG_SINK_CAP);
877 assert_eq!(entries[0].message, "msg-1");
878 assert_eq!(
879 entries[crate::recent_log::LOG_SINK_CAP - 1].message,
880 "msg-20"
881 );
882 }
883
884 #[test]
887 fn transcript_timestamps_recorded() {
888 let metrics = ExecutionMetrics::new();
889 let observer = metrics.create_observer();
890
891 let before = std::time::SystemTime::now()
892 .duration_since(std::time::UNIX_EPOCH)
893 .unwrap_or_default()
894 .as_millis() as i64;
895
896 observer.on_paused(&[LlmQuery {
897 id: QueryId::single(),
898 prompt: "ts-test".into(),
899 system: None,
900 max_tokens: 10,
901 grounded: false,
902 underspecified: false,
903 }]);
904
905 observer.on_response_fed(&QueryId::single(), "response", None);
906
907 let after_fed = std::time::SystemTime::now()
908 .duration_since(std::time::UNIX_EPOCH)
909 .unwrap_or_default()
910 .as_millis() as i64;
911
912 let snap = metrics.snapshot(true);
914 let history = snap["conversation_history"]
915 .as_array()
916 .expect("conversation_history must be array");
917 assert_eq!(history.len(), 1);
918
919 let started_at = history[0]["started_at"]
920 .as_i64()
921 .expect("started_at must be i64");
922 let completed_at = history[0]["completed_at"]
923 .as_i64()
924 .expect("completed_at must be i64 (not null)");
925
926 assert!(
927 started_at >= before,
928 "started_at ({started_at}) should be >= before ({before})"
929 );
930 assert!(
931 completed_at >= started_at,
932 "completed_at ({completed_at}) should be >= started_at ({started_at})"
933 );
934 assert!(
935 completed_at <= after_fed,
936 "completed_at ({completed_at}) should be <= after_fed ({after_fed})"
937 );
938 }
939
940 #[test]
942 fn snapshot_current_query_while_paused() {
943 let metrics = ExecutionMetrics::new();
944 let observer = metrics.create_observer();
945
946 observer.on_paused(&[LlmQuery {
947 id: QueryId::single(),
948 prompt: "in-flight".into(),
949 system: None,
950 max_tokens: 10,
951 grounded: false,
952 underspecified: false,
953 }]);
954
955 let snap = metrics.snapshot(false);
957
958 let tokens = snap.get("tokens").expect("tokens field must be present");
959 let current_query = tokens
960 .get("current_query")
961 .expect("current_query must be present");
962 assert!(
963 !current_query.is_null(),
964 "current_query should be non-null while paused"
965 );
966 assert_eq!(current_query["query_id"], "q-0");
967 assert!(
969 snap.get("conversation_history").is_none(),
970 "conversation_history must be absent when include_history=false"
971 );
972 }
973
974 #[test]
976 fn snapshot_current_query_null_after_response() {
977 let metrics = ExecutionMetrics::new();
978 let observer = metrics.create_observer();
979
980 observer.on_paused(&[LlmQuery {
981 id: QueryId::single(),
982 prompt: "done".into(),
983 system: None,
984 max_tokens: 10,
985 grounded: false,
986 underspecified: false,
987 }]);
988 observer.on_response_fed(&QueryId::single(), "answer", None);
989
990 let snap = metrics.snapshot(false);
991 let tokens = snap.get("tokens").expect("tokens must be present");
992 let current_query = &tokens["current_query"];
993 assert!(
994 current_query.is_null(),
995 "current_query should be null after response is fed"
996 );
997 }
998
999 #[test]
1001 fn snapshot_conversation_history_opt_in() {
1002 let metrics = ExecutionMetrics::new();
1003 let observer = metrics.create_observer();
1004
1005 observer.on_paused(&[LlmQuery {
1006 id: QueryId::single(),
1007 prompt: "hello".into(),
1008 system: None,
1009 max_tokens: 50,
1010 grounded: false,
1011 underspecified: false,
1012 }]);
1013 observer.on_response_fed(&QueryId::single(), "world", None);
1014 observer.on_resumed();
1015 observer.on_completed(&serde_json::json!(null));
1016
1017 let snap_false = metrics.snapshot(false);
1019 assert!(
1020 snap_false.get("conversation_history").is_none(),
1021 "conversation_history must be absent with include_history=false"
1022 );
1023
1024 let snap_true = metrics.snapshot(true);
1026 let history = snap_true
1027 .get("conversation_history")
1028 .expect("conversation_history must be present with include_history=true");
1029 let arr = history
1030 .as_array()
1031 .expect("conversation_history must be an array");
1032 assert_eq!(arr.len(), 1);
1033 assert_eq!(arr[0]["query_id"], "q-0");
1034 assert_eq!(arr[0]["prompt"], "hello");
1035 assert_eq!(arr[0]["response"], "world");
1036 assert!(arr[0].get("started_at").is_some());
1038 assert!(arr[0].get("completed_at").is_some());
1039 }
1040
1041 #[test]
1043 fn snapshot_conversation_history_capped_at_10() {
1044 let metrics = ExecutionMetrics::new();
1045 let observer = metrics.create_observer();
1046
1047 for i in 0..15u32 {
1048 observer.on_paused(&[LlmQuery {
1049 id: QueryId::single(),
1050 prompt: format!("prompt-{i}"),
1051 system: None,
1052 max_tokens: 10,
1053 grounded: false,
1054 underspecified: false,
1055 }]);
1056 observer.on_response_fed(&QueryId::single(), &format!("resp-{i}"), None);
1057 observer.on_resumed();
1058 }
1059
1060 let snap = metrics.snapshot(true);
1061 let history = snap["conversation_history"]
1062 .as_array()
1063 .expect("must be array");
1064 assert_eq!(history.len(), 10, "capped at 10 entries");
1065 assert_eq!(history[0]["prompt"], "prompt-5");
1067 assert_eq!(history[9]["prompt"], "prompt-14");
1068 }
1069
1070 #[test]
1072 fn snapshot_includes_recent_logs() {
1073 let metrics = ExecutionMetrics::new();
1074 let observer = metrics.create_observer();
1075 observer.on_log(&crate::LogEntry::new("info", "engine", "test-log"));
1076
1077 let snap = metrics.snapshot(false);
1078 let logs = snap
1079 .get("recent_logs")
1080 .expect("recent_logs must be in snapshot");
1081 let arr = logs.as_array().expect("recent_logs must be array");
1082 assert_eq!(arr.len(), 1);
1083 assert_eq!(arr[0]["message"], "test-log");
1084 }
1085
1086 #[test]
1088 fn snapshot_tokens_aggregate() {
1089 let metrics = ExecutionMetrics::new();
1090 let observer = metrics.create_observer();
1091
1092 observer.on_paused(&[LlmQuery {
1093 id: QueryId::single(),
1094 prompt: "x".repeat(100),
1095 system: None,
1096 max_tokens: 50,
1097 grounded: false,
1098 underspecified: false,
1099 }]);
1100 observer.on_response_fed(&QueryId::single(), &"y".repeat(50), None);
1101 observer.on_resumed();
1102
1103 let snap = metrics.snapshot(false);
1104 let tokens = snap.get("tokens").expect("tokens must be in snapshot");
1105 let prompt_total = tokens["prompt_total"]
1106 .as_u64()
1107 .expect("prompt_total must be u64");
1108 let response_total = tokens["response_total"]
1109 .as_u64()
1110 .expect("response_total must be u64");
1111 let total = tokens["total"].as_u64().expect("total must be u64");
1112 assert!(prompt_total > 0, "prompt_total must be positive");
1114 assert!(response_total > 0, "response_total must be positive");
1115 assert_eq!(total, prompt_total + response_total);
1116 }
1117}