1use std::collections::HashMap;
7use std::sync::atomic::{AtomicI64, Ordering};
8use std::sync::Arc;
9use std::time::{Duration, SystemTime, UNIX_EPOCH};
10
11use algocline_core::{
12 ExecutionMetrics, ExecutionObserver, ExecutionState, LlmQuery, MetricsObserver, QueryId,
13 TerminalState,
14};
15use mlua_isle::{AsyncIsleDriver, AsyncTask};
16use serde_json::json;
17use tokio::sync::Mutex;
18
19use crate::llm_bridge::LlmRequest;
20
21#[derive(Debug, thiserror::Error)]
24pub enum SessionError {
25 #[error("session '{0}' not found")]
26 NotFound(String),
27 #[error(transparent)]
28 Feed(#[from] algocline_core::FeedError),
29 #[error("invalid transition: {0}")]
30 InvalidTransition(String),
31}
32
33#[derive(serde::Serialize)]
37pub struct ExecutionResult {
38 pub state: TerminalState,
39 pub metrics: ExecutionMetrics,
40}
41
42#[derive(serde::Serialize)]
44pub enum FeedResult {
45 Accepted { remaining: usize },
47 Paused { queries: Vec<LlmQuery> },
49 Finished(ExecutionResult),
51}
52
53impl FeedResult {
54 pub fn to_json(&self, session_id: &str) -> serde_json::Value {
56 match self {
57 Self::Accepted { remaining } => json!({
58 "status": "accepted",
59 "remaining": remaining,
60 }),
61 Self::Paused { queries } => {
62 if queries.len() == 1 {
63 let q = &queries[0];
64 let mut obj = json!({
65 "status": "needs_response",
66 "session_id": session_id,
67 "query_id": q.id.as_str(),
68 "prompt": q.prompt,
69 "system": q.system,
70 "max_tokens": q.max_tokens,
71 });
72 if q.grounded {
73 obj["grounded"] = json!(true);
74 }
75 if q.underspecified {
76 obj["underspecified"] = json!(true);
77 }
78 if let Some(cb) = &q.cache_breakpoint {
79 obj["cache_breakpoint"] = json!(cb);
80 }
81 if let Some(role) = &q.role {
82 obj["role"] = json!(role);
83 }
84 obj
85 } else {
86 let qs: Vec<_> = queries
87 .iter()
88 .map(|q| {
89 let mut obj = json!({
90 "id": q.id.as_str(),
91 "prompt": q.prompt,
92 "system": q.system,
93 "max_tokens": q.max_tokens,
94 });
95 if q.grounded {
96 obj["grounded"] = json!(true);
97 }
98 if q.underspecified {
99 obj["underspecified"] = json!(true);
100 }
101 if let Some(cb) = &q.cache_breakpoint {
102 obj["cache_breakpoint"] = json!(cb);
103 }
104 if let Some(role) = &q.role {
105 obj["role"] = json!(role);
106 }
107 obj
108 })
109 .collect();
110 json!({
111 "status": "needs_response",
112 "session_id": session_id,
113 "queries": qs,
114 })
115 }
116 }
117 Self::Finished(result) => match &result.state {
118 TerminalState::Completed { result: val } => json!({
119 "status": "completed",
120 "result": val,
121 "stats": result.metrics.to_json(),
122 }),
123 TerminalState::Failed { error } => json!({
124 "status": "error",
125 "error": error,
126 }),
127 TerminalState::Cancelled => json!({
128 "status": "cancelled",
129 "stats": result.metrics.to_json(),
130 }),
131 },
132 }
133 }
134}
135
136pub const DEFAULT_PROMPT_PREVIEW_CHARS: usize = 200;
142
143#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
150pub struct PendingFilter {
151 #[serde(default)]
152 pub query_id: bool,
153 #[serde(default)]
154 pub max_tokens: bool,
155 #[serde(default)]
156 pub system: bool,
157 #[serde(default)]
158 pub grounded: bool,
159 #[serde(default)]
160 pub underspecified: bool,
161 #[serde(default)]
162 pub cache_breakpoint: bool,
163 #[serde(default)]
169 pub role: bool,
170 #[serde(default)]
171 pub prompt: PromptProjection,
172}
173
174#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
180#[serde(tag = "mode", rename_all = "snake_case")]
181pub enum PromptProjection {
182 #[default]
183 Off,
184 Preview {
185 chars: usize,
186 },
187 Full,
188}
189
190impl PendingFilter {
191 pub fn preset_meta() -> Self {
193 Self {
194 query_id: true,
195 max_tokens: true,
196 ..Self::default()
197 }
198 }
199
200 pub fn preset_preview() -> Self {
203 Self::preset_preview_with(DEFAULT_PROMPT_PREVIEW_CHARS)
204 }
205
206 pub fn preset_preview_with(chars: usize) -> Self {
213 Self {
214 query_id: true,
215 max_tokens: true,
216 role: true,
217 prompt: PromptProjection::Preview { chars },
218 ..Self::default()
219 }
220 }
221
222 pub fn preset_full() -> Self {
224 Self {
225 query_id: true,
226 max_tokens: true,
227 system: true,
228 grounded: true,
229 underspecified: true,
230 cache_breakpoint: true,
231 role: true,
232 prompt: PromptProjection::Full,
233 }
234 }
235
236 pub fn from_preset(name: &str) -> Option<Self> {
240 match name {
241 "meta" => Some(Self::preset_meta()),
242 "preview" => Some(Self::preset_preview()),
243 "full" => Some(Self::preset_full()),
244 _ => None,
245 }
246 }
247
248 pub fn from_preset_with(name: &str, preview_chars: usize) -> Option<Self> {
251 match name {
252 "meta" => Some(Self::preset_meta()),
253 "preview" => Some(Self::preset_preview_with(preview_chars)),
254 "full" => Some(Self::preset_full()),
255 _ => None,
256 }
257 }
258}
259
260fn project_query(q: &LlmQuery, f: &PendingFilter) -> serde_json::Value {
265 let mut obj = serde_json::Map::new();
266 if f.query_id {
267 obj.insert("query_id".into(), q.id.as_str().into());
268 }
269 if f.max_tokens {
270 obj.insert("max_tokens".into(), q.max_tokens.into());
271 }
272 if f.system {
273 obj.insert(
274 "system".into(),
275 match &q.system {
276 Some(s) => serde_json::Value::String(s.clone()),
277 None => serde_json::Value::Null,
278 },
279 );
280 }
281 if f.grounded {
282 obj.insert("grounded".into(), q.grounded.into());
283 }
284 if f.underspecified {
285 obj.insert("underspecified".into(), q.underspecified.into());
286 }
287 if f.cache_breakpoint {
288 obj.insert(
289 "cache_breakpoint".into(),
290 match &q.cache_breakpoint {
291 Some(s) => serde_json::Value::String(s.clone()),
292 None => serde_json::Value::Null,
293 },
294 );
295 }
296 if f.role {
297 obj.insert(
298 "role".into(),
299 match &q.role {
300 Some(s) => serde_json::Value::String(s.clone()),
301 None => serde_json::Value::Null,
302 },
303 );
304 }
305 match &f.prompt {
306 PromptProjection::Off => {}
307 PromptProjection::Full => {
308 obj.insert("prompt".into(), q.prompt.clone().into());
309 }
310 PromptProjection::Preview { chars } => {
311 let preview: String = q.prompt.chars().take(*chars).collect();
312 obj.insert("prompt_preview".into(), preview.into());
313 }
314 }
315 serde_json::Value::Object(obj)
316}
317
318pub struct Session {
326 state: ExecutionState,
327 metrics: ExecutionMetrics,
328 observer: MetricsObserver,
329 llm_rx: tokio::sync::mpsc::Receiver<LlmRequest>,
330 exec_task: AsyncTask,
331 resp_txs: HashMap<QueryId, tokio::sync::oneshot::Sender<Result<String, String>>>,
333 _vm_driver: AsyncIsleDriver,
336 last_active: std::time::Instant,
339 started_at_ms: i64,
341 last_activity_ms: Arc<AtomicI64>,
344 _log_sink_registration: Option<crate::card::LogSinkRegistration>,
353}
354
355impl Session {
356 pub fn new(
372 llm_rx: tokio::sync::mpsc::Receiver<LlmRequest>,
373 exec_task: AsyncTask,
374 metrics: ExecutionMetrics,
375 vm_driver: AsyncIsleDriver,
376 ) -> Self {
377 let observer = metrics.create_observer();
378 let started_at_ms = SystemTime::now()
381 .duration_since(UNIX_EPOCH)
382 .unwrap_or_default()
383 .as_millis() as i64;
384 Self {
385 state: ExecutionState::Running,
386 metrics,
387 observer,
388 llm_rx,
389 exec_task,
390 resp_txs: HashMap::new(),
391 _vm_driver: vm_driver,
392 last_active: std::time::Instant::now(),
393 started_at_ms,
394 last_activity_ms: Arc::new(AtomicI64::new(started_at_ms)),
395 _log_sink_registration: None,
396 }
397 }
398
399 pub(crate) fn attach_log_sink_registration(
409 &mut self,
410 guard: Option<crate::card::LogSinkRegistration>,
411 ) {
412 self._log_sink_registration = guard;
413 }
414
415 async fn wait_event(&mut self) -> Result<FeedResult, SessionError> {
420 tokio::select! {
421 result = &mut self.exec_task => {
422 match result {
423 Ok(json_str) => match serde_json::from_str::<serde_json::Value>(&json_str) {
424 Ok(v) => {
425 self.state.complete(v.clone()).map_err(|e| {
426 SessionError::InvalidTransition(e.to_string())
427 })?;
428 self.observer.on_completed(&v);
429 Ok(FeedResult::Finished(ExecutionResult {
430 state: TerminalState::Completed { result: v },
431 metrics: self.take_metrics(),
432 }))
433 }
434 Err(e) => self.fail_with(format!("JSON parse: {e}")),
435 },
436 Err(e) => self.fail_with(e.to_string()),
437 }
438 }
439 Some(req) = self.llm_rx.recv() => {
440 let queries: Vec<LlmQuery> = req.queries.iter().map(|qr| LlmQuery {
441 id: qr.id.clone(),
442 prompt: qr.prompt.clone(),
443 system: qr.system.clone(),
444 max_tokens: qr.max_tokens,
445 grounded: qr.grounded,
446 underspecified: qr.underspecified,
447 cache_breakpoint: qr.cache_breakpoint.clone(),
448 role: qr.role.clone(),
449 }).collect();
450
451 for qr in req.queries {
452 self.resp_txs.insert(qr.id, qr.resp_tx);
453 }
454
455 self.state.pause(queries.clone()).map_err(|e| {
456 SessionError::InvalidTransition(e.to_string())
457 })?;
458 self.observer.on_paused(&queries);
459 Ok(FeedResult::Paused { queries })
460 }
461 }
462 }
463
464 fn feed_one(
480 &mut self,
481 query_id: &QueryId,
482 response: String,
483 usage: Option<&algocline_core::TokenUsage>,
484 ) -> Result<bool, SessionError> {
485 self.last_active = std::time::Instant::now();
487 let now_ms = SystemTime::now()
490 .duration_since(UNIX_EPOCH)
491 .unwrap_or_default()
492 .as_millis() as i64;
493 self.last_activity_ms.store(now_ms, Ordering::Relaxed);
494
495 self.observer.on_response_fed(query_id, &response, usage);
497
498 if let Some(tx) = self.resp_txs.remove(query_id) {
500 let _ = tx.send(Ok(response.clone()));
501 }
502
503 let complete = self
505 .state
506 .feed(query_id, response)
507 .map_err(SessionError::Feed)?;
508
509 if complete {
510 self.state
512 .take_responses()
513 .map_err(|e| SessionError::InvalidTransition(e.to_string()))?;
514 self.observer.on_resumed();
515 } else {
516 self.observer
517 .on_partial_feed(query_id, self.state.remaining());
518 }
519
520 Ok(complete)
521 }
522
523 fn fail_with(&mut self, msg: String) -> Result<FeedResult, SessionError> {
524 self.state
525 .fail(msg.clone())
526 .map_err(|e| SessionError::InvalidTransition(e.to_string()))?;
527 self.observer.on_failed(&msg);
528 Ok(FeedResult::Finished(ExecutionResult {
529 state: TerminalState::Failed { error: msg },
530 metrics: self.take_metrics(),
531 }))
532 }
533
534 fn take_metrics(&mut self) -> ExecutionMetrics {
535 std::mem::take(&mut self.metrics)
536 }
537
538 pub fn snapshot(
565 &self,
566 pending_filter: Option<&PendingFilter>,
567 include_history: bool,
568 ) -> serde_json::Value {
569 let state_label = match &self.state {
570 ExecutionState::Running => "running",
571 ExecutionState::Paused(_) => "paused",
572 _ => "terminal",
573 };
574
575 let phase = match &self.state {
576 ExecutionState::Running => "running",
577 ExecutionState::Paused(_) => "paused",
578 ExecutionState::Completed { .. } => "completed",
579 ExecutionState::Failed { .. } => "failed",
580 ExecutionState::Cancelled => "cancelled",
581 };
582
583 let mut json = serde_json::json!({
584 "state": state_label,
585 "phase": phase,
586 "started_at": self.started_at_ms,
587 "last_activity_at": self.last_activity_ms.load(Ordering::Relaxed),
588 });
589
590 let metrics = self.metrics.snapshot(include_history);
591 if !metrics.is_null() {
592 json["metrics"] = metrics;
593 }
594
595 if let ExecutionState::Paused(pending) = &self.state {
597 json["pending_queries"] = pending.remaining().into();
598
599 if let Some(filter) = pending_filter {
600 let items: Vec<serde_json::Value> = pending
601 .pending_queries()
602 .iter()
603 .map(|q| project_query(q, filter))
604 .collect();
605 json["pending"] = serde_json::Value::Array(items);
606 }
607 }
608
609 json
610 }
611
612 pub fn is_expired(&self, ttl: Duration) -> bool {
617 is_expired_impl(self.last_active, ttl)
618 }
619
620 pub(crate) fn into_driver_parts(
628 self,
629 ) -> (
630 AsyncTask,
631 tokio::sync::mpsc::Receiver<LlmRequest>,
632 AsyncIsleDriver,
633 ExecutionMetrics,
634 ) {
635 (self.exec_task, self.llm_rx, self._vm_driver, self.metrics)
636 }
637}
638
639fn is_expired_impl(last_active: std::time::Instant, ttl: Duration) -> bool {
641 std::time::Instant::now().saturating_duration_since(last_active) >= ttl
642}
643
644pub struct SessionRegistry {
679 sessions: Arc<Mutex<HashMap<String, Session>>>,
680}
681
682impl Default for SessionRegistry {
683 fn default() -> Self {
684 Self::new()
685 }
686}
687
688impl SessionRegistry {
689 pub fn new() -> Self {
690 Self {
691 sessions: Arc::new(Mutex::new(HashMap::new())),
692 }
693 }
694
695 pub async fn start_execution(
697 &self,
698 mut session: Session,
699 ) -> Result<(String, FeedResult), SessionError> {
700 let session_id = gen_session_id();
701 let result = session.wait_event().await?;
702
703 if matches!(result, FeedResult::Paused { .. }) {
704 self.sessions
705 .lock()
706 .await
707 .insert(session_id.clone(), session);
708 }
709
710 Ok((session_id, result))
711 }
712
713 pub async fn feed_response(
719 &self,
720 session_id: &str,
721 query_id: &QueryId,
722 response: String,
723 usage: Option<&algocline_core::TokenUsage>,
724 ) -> Result<FeedResult, SessionError> {
725 let complete = {
727 let mut map = self.sessions.lock().await;
728 let session = map
729 .get_mut(session_id)
730 .ok_or_else(|| SessionError::NotFound(session_id.into()))?;
731
732 let complete = session.feed_one(query_id, response, usage)?;
733
734 if !complete {
735 return Ok(FeedResult::Accepted {
736 remaining: session.state.remaining(),
737 });
738 }
739
740 complete
741 };
742
743 debug_assert!(complete);
745 let mut session = {
746 let mut map = self.sessions.lock().await;
747 map.remove(session_id)
748 .ok_or_else(|| SessionError::NotFound(session_id.into()))?
749 };
750
751 let result = session.wait_event().await?;
752
753 if matches!(result, FeedResult::Paused { .. }) {
754 self.sessions
755 .lock()
756 .await
757 .insert(session_id.into(), session);
758 }
759
760 Ok(result)
761 }
762
763 pub async fn resolve_sole_pending_id(&self, session_id: &str) -> Result<QueryId, SessionError> {
769 let map = self.sessions.lock().await;
770 let session = map
771 .get(session_id)
772 .ok_or_else(|| SessionError::NotFound(session_id.into()))?;
773 let keys: Vec<QueryId> = session.resp_txs.keys().cloned().collect();
774 match keys.len() {
775 0 => Err(SessionError::InvalidTransition("no pending queries".into())),
776 1 => keys
777 .into_iter()
778 .next()
779 .ok_or_else(|| SessionError::InvalidTransition("unexpected empty keys".into())),
780 n => Err(SessionError::InvalidTransition(format!(
781 "{n} queries pending; specify query_id explicitly"
782 ))),
783 }
784 }
785
786 pub async fn list_snapshots(
802 &self,
803 pending_filter: Option<&PendingFilter>,
804 include_history: bool,
805 ) -> HashMap<String, serde_json::Value> {
806 let map = self.sessions.lock().await;
807 map.iter()
808 .map(|(id, session)| {
809 (
810 id.clone(),
811 session.snapshot(pending_filter, include_history),
812 )
813 })
814 .collect()
815 }
816
817 pub fn spawn_gc_task(&self, ttl: Duration) {
823 let sessions = Arc::clone(&self.sessions);
824 tokio::spawn(async move {
825 let mut interval = tokio::time::interval(Duration::from_secs(60));
826 loop {
827 interval.tick().await;
828 let mut map = sessions.lock().await;
829 let expired: Vec<String> = map
830 .iter()
831 .filter(|(_, s)| s.is_expired(ttl))
832 .map(|(id, _)| id.clone())
833 .collect();
834 for id in &expired {
835 tracing::info!(session_id = %id, "GC: reaping expired session");
836 map.remove(id);
837 }
838 }
839 });
840 }
841}
842
843fn gen_session_id() -> String {
862 use rand::RngExt;
863 use std::time::{SystemTime, UNIX_EPOCH};
864 let ts = SystemTime::now()
865 .duration_since(UNIX_EPOCH)
866 .unwrap_or_default()
867 .as_nanos();
868 let random: u64 = rand::rng().random();
869 format!("s-{ts:x}-{random:016x}")
870}
871
872#[cfg(test)]
873mod tests {
874 use super::*;
875 use algocline_core::{ExecutionMetrics, LlmQuery, QueryId};
876 use serde_json::json;
877
878 fn make_query(index: usize) -> LlmQuery {
879 LlmQuery {
880 id: QueryId::batch(index),
881 prompt: format!("prompt-{index}"),
882 system: None,
883 max_tokens: 100,
884 grounded: false,
885 underspecified: false,
886 cache_breakpoint: None,
887 role: None,
888 }
889 }
890
891 #[test]
894 fn to_json_accepted() {
895 let result = FeedResult::Accepted { remaining: 3 };
896 let json = result.to_json("s-123");
897 assert_eq!(json["status"], "accepted");
898 assert_eq!(json["remaining"], 3);
899 }
900
901 #[test]
902 fn to_json_paused_single_query() {
903 let query = LlmQuery {
904 id: QueryId::single(),
905 prompt: "What is 2+2?".into(),
906 system: Some("You are a calculator.".into()),
907 max_tokens: 50,
908 grounded: false,
909 underspecified: false,
910 cache_breakpoint: None,
911 role: None,
912 };
913 let result = FeedResult::Paused {
914 queries: vec![query],
915 };
916 let json = result.to_json("s-abc");
917
918 assert_eq!(json["status"], "needs_response");
919 assert_eq!(json["session_id"], "s-abc");
920 assert_eq!(json["prompt"], "What is 2+2?");
921 assert_eq!(json["system"], "You are a calculator.");
922 assert_eq!(json["max_tokens"], 50);
923 assert!(json.get("queries").is_none());
925 assert!(
927 json.get("grounded").is_none(),
928 "grounded key must be absent when false"
929 );
930 assert!(
932 json.get("underspecified").is_none(),
933 "underspecified key must be absent when false"
934 );
935 }
936
937 #[test]
938 fn to_json_paused_single_query_grounded() {
939 let query = LlmQuery {
940 id: QueryId::single(),
941 prompt: "verify this claim".into(),
942 system: None,
943 max_tokens: 200,
944 grounded: true,
945 underspecified: false,
946 cache_breakpoint: None,
947 role: None,
948 };
949 let result = FeedResult::Paused {
950 queries: vec![query],
951 };
952 let json = result.to_json("s-grounded");
953
954 assert_eq!(json["status"], "needs_response");
955 assert_eq!(
956 json["grounded"], true,
957 "grounded must appear in single-query MCP JSON"
958 );
959 }
960
961 #[test]
962 fn to_json_paused_single_query_underspecified() {
963 let query = LlmQuery {
964 id: QueryId::single(),
965 prompt: "what output format do you need?".into(),
966 system: None,
967 max_tokens: 200,
968 grounded: false,
969 underspecified: true,
970 cache_breakpoint: None,
971 role: None,
972 };
973 let result = FeedResult::Paused {
974 queries: vec![query],
975 };
976 let json = result.to_json("s-underspec");
977
978 assert_eq!(json["status"], "needs_response");
979 assert_eq!(
980 json["underspecified"], true,
981 "underspecified must appear in single-query MCP JSON"
982 );
983 assert!(
984 json.get("grounded").is_none(),
985 "grounded must be absent when false"
986 );
987 }
988
989 #[test]
990 fn to_json_paused_multiple_queries_mixed_grounded() {
991 let grounded_query = LlmQuery {
992 id: QueryId::batch(0),
993 prompt: "verify".into(),
994 system: None,
995 max_tokens: 100,
996 grounded: true,
997 underspecified: false,
998 cache_breakpoint: None,
999 role: None,
1000 };
1001 let normal_query = LlmQuery {
1002 id: QueryId::batch(1),
1003 prompt: "generate".into(),
1004 system: None,
1005 max_tokens: 100,
1006 grounded: false,
1007 underspecified: false,
1008 cache_breakpoint: None,
1009 role: None,
1010 };
1011 let result = FeedResult::Paused {
1012 queries: vec![grounded_query, normal_query],
1013 };
1014 let json = result.to_json("s-batch");
1015
1016 let qs = json["queries"].as_array().expect("queries should be array");
1017 assert_eq!(
1018 qs[0]["grounded"], true,
1019 "grounded query must have grounded=true"
1020 );
1021 assert!(
1022 qs[1].get("grounded").is_none(),
1023 "non-grounded query must omit grounded key"
1024 );
1025 }
1026
1027 #[test]
1028 fn to_json_paused_single_query_cache_breakpoint() {
1029 let query = LlmQuery {
1030 id: QueryId::single(),
1031 prompt: "cached user turn".into(),
1032 system: Some("shared system prompt".into()),
1033 max_tokens: 512,
1034 grounded: false,
1035 underspecified: false,
1036 cache_breakpoint: Some("context".into()),
1037 role: None,
1038 };
1039 let result = FeedResult::Paused {
1040 queries: vec![query],
1041 };
1042 let json = result.to_json("s-cache");
1043
1044 assert_eq!(json["status"], "needs_response");
1045 assert_eq!(
1046 json["cache_breakpoint"], "context",
1047 "cache_breakpoint must appear verbatim in single-query MCP JSON"
1048 );
1049 assert!(
1050 json.get("grounded").is_none(),
1051 "grounded must be absent when false"
1052 );
1053 assert!(
1054 json.get("underspecified").is_none(),
1055 "underspecified must be absent when false"
1056 );
1057 }
1058
1059 #[test]
1060 fn to_json_paused_single_query_no_cache_breakpoint_absent() {
1061 let query = LlmQuery {
1062 id: QueryId::single(),
1063 prompt: "uncached".into(),
1064 system: None,
1065 max_tokens: 128,
1066 grounded: false,
1067 underspecified: false,
1068 cache_breakpoint: None,
1069 role: None,
1070 };
1071 let result = FeedResult::Paused {
1072 queries: vec![query],
1073 };
1074 let json = result.to_json("s-nocache");
1075
1076 assert!(
1077 json.get("cache_breakpoint").is_none(),
1078 "cache_breakpoint key must be absent when None"
1079 );
1080 }
1081
1082 #[test]
1083 fn to_json_paused_single_query_role() {
1084 let query = LlmQuery {
1085 id: QueryId::single(),
1086 prompt: "grade this answer".into(),
1087 system: None,
1088 max_tokens: 256,
1089 grounded: false,
1090 underspecified: false,
1091 cache_breakpoint: None,
1092 role: Some("grader".into()),
1093 };
1094 let result = FeedResult::Paused {
1095 queries: vec![query],
1096 };
1097 let json = result.to_json("s-role");
1098
1099 assert_eq!(json["status"], "needs_response");
1100 assert_eq!(
1101 json["role"], "grader",
1102 "role must appear verbatim in single-query MCP JSON"
1103 );
1104 }
1105
1106 #[test]
1107 fn to_json_paused_single_query_no_role_absent() {
1108 let query = LlmQuery {
1109 id: QueryId::single(),
1110 prompt: "no role".into(),
1111 system: None,
1112 max_tokens: 64,
1113 grounded: false,
1114 underspecified: false,
1115 cache_breakpoint: None,
1116 role: None,
1117 };
1118 let result = FeedResult::Paused {
1119 queries: vec![query],
1120 };
1121 let json = result.to_json("s-norole");
1122
1123 assert!(
1124 json.get("role").is_none(),
1125 "role key must be absent when None"
1126 );
1127 }
1128
1129 #[test]
1130 fn to_json_paused_multiple_queries_mixed_role() {
1131 let judge_query = LlmQuery {
1132 id: QueryId::batch(0),
1133 prompt: "grade".into(),
1134 system: None,
1135 max_tokens: 100,
1136 grounded: false,
1137 underspecified: false,
1138 cache_breakpoint: None,
1139 role: Some("grader".into()),
1140 };
1141 let normal_query = LlmQuery {
1142 id: QueryId::batch(1),
1143 prompt: "generate".into(),
1144 system: None,
1145 max_tokens: 100,
1146 grounded: false,
1147 underspecified: false,
1148 cache_breakpoint: None,
1149 role: None,
1150 };
1151 let result = FeedResult::Paused {
1152 queries: vec![judge_query, normal_query],
1153 };
1154 let json = result.to_json("s-batch-role");
1155
1156 let qs = json["queries"].as_array().expect("queries should be array");
1157 assert_eq!(
1158 qs[0]["role"], "grader",
1159 "grader query must forward role verbatim"
1160 );
1161 assert!(
1162 qs[1].get("role").is_none(),
1163 "non-role query must omit role key"
1164 );
1165 }
1166
1167 #[test]
1168 fn to_json_paused_multiple_queries_mixed_cache_breakpoint() {
1169 let cached_query = LlmQuery {
1170 id: QueryId::batch(0),
1171 prompt: "cached".into(),
1172 system: None,
1173 max_tokens: 100,
1174 grounded: false,
1175 underspecified: false,
1176 cache_breakpoint: Some("prompt".into()),
1177 role: None,
1178 };
1179 let normal_query = LlmQuery {
1180 id: QueryId::batch(1),
1181 prompt: "uncached".into(),
1182 system: None,
1183 max_tokens: 100,
1184 grounded: false,
1185 underspecified: false,
1186 cache_breakpoint: None,
1187 role: None,
1188 };
1189 let result = FeedResult::Paused {
1190 queries: vec![cached_query, normal_query],
1191 };
1192 let json = result.to_json("s-batch-cache");
1193
1194 let qs = json["queries"].as_array().expect("queries should be array");
1195 assert_eq!(
1196 qs[0]["cache_breakpoint"], "prompt",
1197 "cached query must forward cache_breakpoint verbatim"
1198 );
1199 assert!(
1200 qs[1].get("cache_breakpoint").is_none(),
1201 "uncached query must omit cache_breakpoint key"
1202 );
1203 }
1204
1205 #[test]
1206 fn to_json_paused_multiple_queries_mixed_underspecified() {
1207 let underspec_query = LlmQuery {
1208 id: QueryId::batch(0),
1209 prompt: "clarify intent".into(),
1210 system: None,
1211 max_tokens: 100,
1212 grounded: false,
1213 underspecified: true,
1214 cache_breakpoint: None,
1215 role: None,
1216 };
1217 let normal_query = LlmQuery {
1218 id: QueryId::batch(1),
1219 prompt: "generate".into(),
1220 system: None,
1221 max_tokens: 100,
1222 grounded: false,
1223 underspecified: false,
1224 cache_breakpoint: None,
1225 role: None,
1226 };
1227 let result = FeedResult::Paused {
1228 queries: vec![underspec_query, normal_query],
1229 };
1230 let json = result.to_json("s-batch-us");
1231
1232 let qs = json["queries"].as_array().expect("queries should be array");
1233 assert_eq!(
1234 qs[0]["underspecified"], true,
1235 "underspecified query must have underspecified=true"
1236 );
1237 assert!(
1238 qs[1].get("underspecified").is_none(),
1239 "non-underspecified query must omit underspecified key"
1240 );
1241 }
1242
1243 #[test]
1244 fn to_json_paused_single_query_no_system() {
1245 let query = LlmQuery {
1246 id: QueryId::single(),
1247 prompt: "hello".into(),
1248 system: None,
1249 max_tokens: 1024,
1250 grounded: false,
1251 underspecified: false,
1252 cache_breakpoint: None,
1253 role: None,
1254 };
1255 let result = FeedResult::Paused {
1256 queries: vec![query],
1257 };
1258 let json = result.to_json("s-x");
1259
1260 assert_eq!(json["status"], "needs_response");
1261 assert!(json["system"].is_null());
1262 }
1263
1264 #[test]
1265 fn to_json_paused_multiple_queries() {
1266 let queries = vec![make_query(0), make_query(1), make_query(2)];
1267 let result = FeedResult::Paused { queries };
1268 let json = result.to_json("s-multi");
1269
1270 assert_eq!(json["status"], "needs_response");
1271 assert_eq!(json["session_id"], "s-multi");
1272
1273 let qs = json["queries"].as_array().expect("queries should be array");
1274 assert_eq!(qs.len(), 3);
1275 assert_eq!(qs[0]["id"], "q-0");
1276 assert_eq!(qs[0]["prompt"], "prompt-0");
1277 assert_eq!(qs[1]["id"], "q-1");
1278 assert_eq!(qs[2]["id"], "q-2");
1279 }
1280
1281 #[test]
1282 fn to_json_finished_completed() {
1283 let result = FeedResult::Finished(ExecutionResult {
1284 state: TerminalState::Completed {
1285 result: json!({"answer": 42}),
1286 },
1287 metrics: ExecutionMetrics::new(),
1288 });
1289 let json = result.to_json("s-done");
1290
1291 assert_eq!(json["status"], "completed");
1292 assert_eq!(json["result"]["answer"], 42);
1293 assert!(json.get("stats").is_some());
1294 }
1295
1296 #[test]
1297 fn to_json_finished_failed() {
1298 let result = FeedResult::Finished(ExecutionResult {
1299 state: TerminalState::Failed {
1300 error: "lua error: bad argument".into(),
1301 },
1302 metrics: ExecutionMetrics::new(),
1303 });
1304 let json = result.to_json("s-err");
1305
1306 assert_eq!(json["status"], "error");
1307 assert_eq!(json["error"], "lua error: bad argument");
1308 }
1309
1310 #[test]
1311 fn to_json_finished_cancelled() {
1312 let result = FeedResult::Finished(ExecutionResult {
1313 state: TerminalState::Cancelled,
1314 metrics: ExecutionMetrics::new(),
1315 });
1316 let json = result.to_json("s-cancel");
1317
1318 assert_eq!(json["status"], "cancelled");
1319 assert!(json.get("stats").is_some());
1320 }
1321
1322 #[test]
1325 fn session_id_starts_with_prefix() {
1326 let id = gen_session_id();
1327 assert!(id.starts_with("s-"), "id should start with 's-': {id}");
1328 }
1329
1330 #[test]
1331 fn session_id_uniqueness() {
1332 let ids: Vec<String> = (0..10).map(|_| gen_session_id()).collect();
1333 let set: std::collections::HashSet<&String> = ids.iter().collect();
1334 assert_eq!(set.len(), 10, "10 IDs should all be unique");
1335 }
1336
1337 #[test]
1344 fn is_expired_impl_fresh_instant_not_expired() {
1345 let now = std::time::Instant::now();
1347 assert!(!is_expired_impl(now, Duration::from_secs(1)));
1348 }
1349
1350 #[test]
1351 fn is_expired_impl_old_instant_expired() {
1352 let two_hours_ago = std::time::Instant::now()
1354 .checked_sub(Duration::from_secs(7200))
1355 .expect("checked_sub should succeed with sane duration");
1356 assert!(is_expired_impl(two_hours_ago, Duration::from_secs(3600)));
1358 }
1359
1360 #[test]
1361 fn is_expired_impl_not_yet_expired() {
1362 let one_hour_ago = std::time::Instant::now()
1364 .checked_sub(Duration::from_secs(3600))
1365 .expect("checked_sub should succeed with sane duration");
1366 assert!(!is_expired_impl(one_hour_ago, Duration::from_secs(10800)));
1368 }
1369
1370 #[test]
1371 fn is_expired_impl_zero_ttl_always_expired() {
1372 let now = std::time::Instant::now();
1374 assert!(is_expired_impl(now, Duration::ZERO));
1375 }
1376
1377 #[test]
1380 fn pending_filter_default_is_all_off() {
1381 let f = PendingFilter::default();
1382 assert!(!f.query_id);
1383 assert!(!f.max_tokens);
1384 assert!(!f.system);
1385 assert!(!f.grounded);
1386 assert!(!f.underspecified);
1387 assert!(matches!(f.prompt, PromptProjection::Off));
1388 }
1389
1390 #[test]
1391 fn pending_filter_preset_meta_flags() {
1392 let f = PendingFilter::preset_meta();
1393 assert!(f.query_id);
1394 assert!(f.max_tokens);
1395 assert!(!f.system);
1396 assert!(!f.grounded);
1397 assert!(!f.underspecified);
1398 assert!(
1399 !f.role,
1400 "meta preset must not project role (query_id + max_tokens only)"
1401 );
1402 assert!(
1403 matches!(f.prompt, PromptProjection::Off),
1404 "meta preset must not project prompt content"
1405 );
1406 }
1407
1408 #[test]
1409 fn pending_filter_preset_preview_uses_default_chars() {
1410 let f = PendingFilter::preset_preview();
1411 assert!(f.query_id);
1412 assert!(f.max_tokens);
1413 assert!(
1414 f.role,
1415 "preview preset must project role so judge pauses are distinguishable from strategy pauses"
1416 );
1417 match f.prompt {
1418 PromptProjection::Preview { chars } => {
1419 assert_eq!(chars, DEFAULT_PROMPT_PREVIEW_CHARS);
1420 }
1421 other => panic!("expected Preview, got {other:?}"),
1422 }
1423 }
1424
1425 #[test]
1426 fn pending_filter_preset_preview_with_custom_chars() {
1427 let f = PendingFilter::preset_preview_with(42);
1428 match f.prompt {
1429 PromptProjection::Preview { chars } => assert_eq!(chars, 42),
1430 other => panic!("expected Preview {{chars: 42}}, got {other:?}"),
1431 }
1432 }
1433
1434 #[test]
1435 fn pending_filter_preset_full_flags_all_on() {
1436 let f = PendingFilter::preset_full();
1437 assert!(f.query_id);
1438 assert!(f.max_tokens);
1439 assert!(f.system);
1440 assert!(f.grounded);
1441 assert!(f.underspecified);
1442 assert!(f.role);
1443 assert!(matches!(f.prompt, PromptProjection::Full));
1444 }
1445
1446 #[test]
1447 fn project_query_role_projection_grader() {
1448 let mut q = make_query(0);
1449 q.role = Some("grader".into());
1450 let f = PendingFilter {
1451 role: true,
1452 ..PendingFilter::default()
1453 };
1454 let v = project_query(&q, &f);
1455 assert_eq!(v["role"], "grader");
1456 }
1457
1458 #[test]
1459 fn project_query_role_projection_absent_when_flag_off() {
1460 let mut q = make_query(0);
1461 q.role = Some("grader".into());
1462 let f = PendingFilter::preset_meta();
1463 let v = project_query(&q, &f);
1464 let obj = v.as_object().expect("object");
1465 assert!(
1466 obj.get("role").is_none(),
1467 "meta preset must not project role even when the query carries one"
1468 );
1469 }
1470
1471 #[test]
1472 fn project_query_role_projection_null_when_query_role_absent() {
1473 let q = make_query(0);
1474 let f = PendingFilter {
1476 role: true,
1477 ..PendingFilter::default()
1478 };
1479 let v = project_query(&q, &f);
1480 assert!(
1481 v["role"].is_null(),
1482 "role field must be null when the query carries no role tag"
1483 );
1484 }
1485
1486 #[test]
1487 fn pending_filter_from_preset_known_names() {
1488 assert!(PendingFilter::from_preset("meta").is_some());
1489 assert!(PendingFilter::from_preset("preview").is_some());
1490 assert!(PendingFilter::from_preset("full").is_some());
1491 }
1492
1493 #[test]
1494 fn pending_filter_from_preset_unknown_returns_none() {
1495 assert!(PendingFilter::from_preset("").is_none());
1498 assert!(PendingFilter::from_preset("META").is_none());
1499 assert!(PendingFilter::from_preset("bogus").is_none());
1500 }
1501
1502 #[test]
1503 fn pending_filter_from_preset_with_overrides_preview_chars() {
1504 let f = PendingFilter::from_preset_with("preview", 73).unwrap();
1507 match f.prompt {
1508 PromptProjection::Preview { chars } => assert_eq!(chars, 73),
1509 other => panic!("expected Preview {{chars: 73}}, got {other:?}"),
1510 }
1511
1512 let f_meta = PendingFilter::from_preset_with("meta", 73).unwrap();
1513 assert!(matches!(f_meta.prompt, PromptProjection::Off));
1514
1515 let f_full = PendingFilter::from_preset_with("full", 73).unwrap();
1516 assert!(matches!(f_full.prompt, PromptProjection::Full));
1517 }
1518
1519 #[test]
1522 fn project_query_default_filter_produces_empty_object() {
1523 let q = make_query(0);
1524 let v = project_query(&q, &PendingFilter::default());
1525 let obj = v.as_object().expect("object");
1526 assert!(obj.is_empty(), "default filter should project nothing");
1527 }
1528
1529 #[test]
1530 fn project_query_meta_preset_has_id_and_max_tokens_only() {
1531 let q = make_query(0);
1532 let v = project_query(&q, &PendingFilter::preset_meta());
1533 let obj = v.as_object().expect("object");
1534 assert_eq!(obj.len(), 2);
1535 assert_eq!(v["query_id"], "q-0");
1536 assert_eq!(v["max_tokens"], 100);
1537 assert!(obj.get("prompt").is_none());
1538 assert!(obj.get("prompt_preview").is_none());
1539 assert!(obj.get("system").is_none());
1540 assert!(obj.get("grounded").is_none());
1541 assert!(obj.get("underspecified").is_none());
1542 }
1543
1544 #[test]
1545 fn project_query_full_preset_has_all_fields() {
1546 let q = LlmQuery {
1547 id: QueryId::batch(0),
1548 prompt: "hi".into(),
1549 system: Some("sys".into()),
1550 max_tokens: 100,
1551 grounded: true,
1552 underspecified: true,
1553 cache_breakpoint: None,
1554 role: None,
1555 };
1556 let v = project_query(&q, &PendingFilter::preset_full());
1557 assert_eq!(v["query_id"], "q-0");
1558 assert_eq!(v["max_tokens"], 100);
1559 assert_eq!(v["system"], "sys");
1560 assert_eq!(v["grounded"], true);
1561 assert_eq!(v["underspecified"], true);
1562 assert_eq!(v["prompt"], "hi");
1563 assert!(v.get("prompt_preview").is_none());
1564 }
1565
1566 #[test]
1567 fn project_query_preview_truncates_at_char_count() {
1568 let q = LlmQuery {
1569 id: QueryId::batch(0),
1570 prompt: "abcdefghij".into(),
1571 system: None,
1572 max_tokens: 10,
1573 grounded: false,
1574 underspecified: false,
1575 cache_breakpoint: None,
1576 role: None,
1577 };
1578 let v = project_query(&q, &PendingFilter::preset_preview_with(5));
1579 assert_eq!(v["prompt_preview"], "abcde");
1580 assert!(v.get("prompt").is_none());
1581 }
1582
1583 #[test]
1584 fn project_query_preview_utf8_multibyte_safe() {
1585 let prompt = "あいうえお";
1589 let q = LlmQuery {
1590 id: QueryId::batch(0),
1591 prompt: prompt.to_string(),
1592 system: None,
1593 max_tokens: 10,
1594 grounded: false,
1595 underspecified: false,
1596 cache_breakpoint: None,
1597 role: None,
1598 };
1599 let v = project_query(&q, &PendingFilter::preset_preview_with(3));
1600 let preview = v["prompt_preview"].as_str().expect("str");
1601 assert_eq!(preview, "あいう");
1602 assert_eq!(preview.chars().count(), 3);
1603 }
1604
1605 #[test]
1606 fn project_query_preview_chars_over_length_returns_whole_prompt() {
1607 let q = LlmQuery {
1608 id: QueryId::batch(0),
1609 prompt: "abc".into(),
1610 system: None,
1611 max_tokens: 10,
1612 grounded: false,
1613 underspecified: false,
1614 cache_breakpoint: None,
1615 role: None,
1616 };
1617 let v = project_query(&q, &PendingFilter::preset_preview_with(100));
1618 assert_eq!(v["prompt_preview"], "abc");
1619 }
1620
1621 #[test]
1622 fn project_query_system_field_null_when_absent() {
1623 let q = LlmQuery {
1624 id: QueryId::batch(0),
1625 prompt: "p".into(),
1626 system: None,
1627 max_tokens: 10,
1628 grounded: false,
1629 underspecified: false,
1630 cache_breakpoint: None,
1631 role: None,
1632 };
1633 let filter = PendingFilter {
1634 system: true,
1635 ..Default::default()
1636 };
1637 let v = project_query(&q, &filter);
1638 assert!(
1639 v["system"].is_null(),
1640 "absent system must serialize as null"
1641 );
1642 }
1643
1644 #[test]
1647 fn pending_filter_deserialize_custom_object_preview() {
1648 let raw = serde_json::json!({
1650 "query_id": true,
1651 "prompt": { "mode": "preview", "chars": 50 }
1652 });
1653 let f: PendingFilter = serde_json::from_value(raw).expect("deserialize");
1654 assert!(f.query_id);
1655 match f.prompt {
1656 PromptProjection::Preview { chars } => assert_eq!(chars, 50),
1657 other => panic!("expected Preview, got {other:?}"),
1658 }
1659 }
1660
1661 #[test]
1662 fn pending_filter_deserialize_partial_object_uses_field_defaults() {
1663 let raw = serde_json::json!({});
1666 let f: PendingFilter = serde_json::from_value(raw).expect("deserialize");
1667 assert!(!f.query_id);
1668 assert!(matches!(f.prompt, PromptProjection::Off));
1669 }
1670
1671 #[test]
1672 fn pending_filter_deserialize_prompt_full_tag() {
1673 let raw = serde_json::json!({ "prompt": { "mode": "full" } });
1674 let f: PendingFilter = serde_json::from_value(raw).expect("deserialize");
1675 assert!(matches!(f.prompt, PromptProjection::Full));
1676 }
1677
1678 fn tmp_dirs() -> crate::executor::SessionDirs {
1686 let tmp = tempfile::tempdir().expect("test tempdir");
1687 let root = tmp.path().to_path_buf();
1688 std::mem::forget(tmp);
1689 crate::executor::SessionDirs {
1690 state_store: std::sync::Arc::new(crate::state::JsonFileStore::new(root.join("state"))),
1691 card_store: std::sync::Arc::new(crate::card::FileCardStore::new(root.join("cards"))),
1692 scenarios_dir: root.join("scenarios"),
1693 nn_dir: root.join("nn"),
1694 }
1695 }
1696
1697 #[tokio::test]
1704 async fn snapshot_v2_contains_phase_and_timestamps() {
1705 let executor = crate::executor::Executor::new(vec![]).await.unwrap();
1706 let dirs = tmp_dirs();
1707
1708 let code = r#"
1710 local response = alc.llm("what is 2+2?")
1711 return response
1712 "#
1713 .to_string();
1714
1715 let session = executor
1716 .start_session(code, serde_json::json!({}), vec![], vec![], dirs, false)
1717 .await
1718 .unwrap();
1719
1720 let snap = session.snapshot(None, false);
1723
1724 assert!(
1726 snap.get("phase").is_some(),
1727 "snapshot must have 'phase' field"
1728 );
1729 assert_eq!(snap["phase"], "running", "initial state must be running");
1730
1731 assert_eq!(snap["state"], "running");
1733
1734 let started_at = snap["started_at"].as_i64().expect("started_at must be i64");
1736 assert!(started_at > 0, "started_at must be > 0 (unix ms)");
1737
1738 let last_activity = snap["last_activity_at"]
1740 .as_i64()
1741 .expect("last_activity_at must be i64");
1742 assert_eq!(
1743 started_at, last_activity,
1744 "last_activity_at should equal started_at before any feed"
1745 );
1746 }
1747
1748 #[test]
1752 fn snapshot_phase_running_state_label() {
1753 let cases: &[(&str, &str)] = &[
1757 ("running", "running"),
1758 ("paused", "paused"),
1759 ("completed", "completed"),
1760 ("failed", "failed"),
1761 ("cancelled", "cancelled"),
1762 ];
1763 for (state_str, expected_phase) in cases {
1764 let three_value_state = match *state_str {
1768 "running" => "running",
1769 "paused" => "paused",
1770 _ => "terminal",
1771 };
1772 assert_eq!(
1774 *expected_phase, *state_str,
1775 "phase for {state_str} must be the same string"
1776 );
1777 if *state_str != "running" && *state_str != "paused" {
1778 assert_eq!(
1779 three_value_state, "terminal",
1780 "{state_str} must map to 'terminal' in 3-value state"
1781 );
1782 }
1783 }
1784 }
1785
1786 #[tokio::test]
1788 async fn snapshot_conversation_history_opt_in() {
1789 let executor = crate::executor::Executor::new(vec![]).await.unwrap();
1790 let dirs = tmp_dirs();
1791
1792 let code = r#"
1793 local response = alc.llm("explain recursion")
1794 return response
1795 "#
1796 .to_string();
1797
1798 let session = executor
1799 .start_session(code, serde_json::json!({}), vec![], vec![], dirs, false)
1800 .await
1801 .unwrap();
1802
1803 let snap_false = session.snapshot(None, false);
1805 assert!(
1806 snap_false
1807 .get("metrics")
1808 .and_then(|m| m.get("conversation_history"))
1809 .is_none(),
1810 "conversation_history must be absent with include_history=false"
1811 );
1812
1813 let snap_true = session.snapshot(None, true);
1815 if let Some(metrics) = snap_true.get("metrics") {
1818 let _ = metrics.get("conversation_history");
1822 }
1823 }
1824
1825 #[tokio::test]
1827 async fn snapshot_last_activity_at_starts_equal_to_started_at() {
1828 let executor = crate::executor::Executor::new(vec![]).await.unwrap();
1829 let dirs = tmp_dirs();
1830
1831 let code = r#"
1832 local response = alc.llm("test query")
1833 return response
1834 "#
1835 .to_string();
1836
1837 let session = executor
1838 .start_session(code, serde_json::json!({}), vec![], vec![], dirs, false)
1839 .await
1840 .unwrap();
1841
1842 let snap = session.snapshot(None, false);
1843 let started_at = snap["started_at"].as_i64().unwrap_or(-1);
1844 let last_activity = snap["last_activity_at"].as_i64().unwrap_or(-2);
1845
1846 assert_eq!(
1847 started_at, last_activity,
1848 "last_activity_at must equal started_at before any feed_one"
1849 );
1850 assert!(started_at > 0, "started_at must be positive unix ms");
1851 }
1852}