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 obj
82 } else {
83 let qs: Vec<_> = queries
84 .iter()
85 .map(|q| {
86 let mut obj = json!({
87 "id": q.id.as_str(),
88 "prompt": q.prompt,
89 "system": q.system,
90 "max_tokens": q.max_tokens,
91 });
92 if q.grounded {
93 obj["grounded"] = json!(true);
94 }
95 if q.underspecified {
96 obj["underspecified"] = json!(true);
97 }
98 if let Some(cb) = &q.cache_breakpoint {
99 obj["cache_breakpoint"] = json!(cb);
100 }
101 obj
102 })
103 .collect();
104 json!({
105 "status": "needs_response",
106 "session_id": session_id,
107 "queries": qs,
108 })
109 }
110 }
111 Self::Finished(result) => match &result.state {
112 TerminalState::Completed { result: val } => json!({
113 "status": "completed",
114 "result": val,
115 "stats": result.metrics.to_json(),
116 }),
117 TerminalState::Failed { error } => json!({
118 "status": "error",
119 "error": error,
120 }),
121 TerminalState::Cancelled => json!({
122 "status": "cancelled",
123 "stats": result.metrics.to_json(),
124 }),
125 },
126 }
127 }
128}
129
130pub const DEFAULT_PROMPT_PREVIEW_CHARS: usize = 200;
136
137#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
144pub struct PendingFilter {
145 #[serde(default)]
146 pub query_id: bool,
147 #[serde(default)]
148 pub max_tokens: bool,
149 #[serde(default)]
150 pub system: bool,
151 #[serde(default)]
152 pub grounded: bool,
153 #[serde(default)]
154 pub underspecified: bool,
155 #[serde(default)]
156 pub cache_breakpoint: bool,
157 #[serde(default)]
158 pub prompt: PromptProjection,
159}
160
161#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
167#[serde(tag = "mode", rename_all = "snake_case")]
168pub enum PromptProjection {
169 #[default]
170 Off,
171 Preview {
172 chars: usize,
173 },
174 Full,
175}
176
177impl PendingFilter {
178 pub fn preset_meta() -> Self {
180 Self {
181 query_id: true,
182 max_tokens: true,
183 ..Self::default()
184 }
185 }
186
187 pub fn preset_preview() -> Self {
190 Self::preset_preview_with(DEFAULT_PROMPT_PREVIEW_CHARS)
191 }
192
193 pub fn preset_preview_with(chars: usize) -> Self {
196 Self {
197 query_id: true,
198 max_tokens: true,
199 prompt: PromptProjection::Preview { chars },
200 ..Self::default()
201 }
202 }
203
204 pub fn preset_full() -> Self {
206 Self {
207 query_id: true,
208 max_tokens: true,
209 system: true,
210 grounded: true,
211 underspecified: true,
212 cache_breakpoint: true,
213 prompt: PromptProjection::Full,
214 }
215 }
216
217 pub fn from_preset(name: &str) -> Option<Self> {
221 match name {
222 "meta" => Some(Self::preset_meta()),
223 "preview" => Some(Self::preset_preview()),
224 "full" => Some(Self::preset_full()),
225 _ => None,
226 }
227 }
228
229 pub fn from_preset_with(name: &str, preview_chars: usize) -> Option<Self> {
232 match name {
233 "meta" => Some(Self::preset_meta()),
234 "preview" => Some(Self::preset_preview_with(preview_chars)),
235 "full" => Some(Self::preset_full()),
236 _ => None,
237 }
238 }
239}
240
241fn project_query(q: &LlmQuery, f: &PendingFilter) -> serde_json::Value {
246 let mut obj = serde_json::Map::new();
247 if f.query_id {
248 obj.insert("query_id".into(), q.id.as_str().into());
249 }
250 if f.max_tokens {
251 obj.insert("max_tokens".into(), q.max_tokens.into());
252 }
253 if f.system {
254 obj.insert(
255 "system".into(),
256 match &q.system {
257 Some(s) => serde_json::Value::String(s.clone()),
258 None => serde_json::Value::Null,
259 },
260 );
261 }
262 if f.grounded {
263 obj.insert("grounded".into(), q.grounded.into());
264 }
265 if f.underspecified {
266 obj.insert("underspecified".into(), q.underspecified.into());
267 }
268 if f.cache_breakpoint {
269 obj.insert(
270 "cache_breakpoint".into(),
271 match &q.cache_breakpoint {
272 Some(s) => serde_json::Value::String(s.clone()),
273 None => serde_json::Value::Null,
274 },
275 );
276 }
277 match &f.prompt {
278 PromptProjection::Off => {}
279 PromptProjection::Full => {
280 obj.insert("prompt".into(), q.prompt.clone().into());
281 }
282 PromptProjection::Preview { chars } => {
283 let preview: String = q.prompt.chars().take(*chars).collect();
284 obj.insert("prompt_preview".into(), preview.into());
285 }
286 }
287 serde_json::Value::Object(obj)
288}
289
290pub struct Session {
298 state: ExecutionState,
299 metrics: ExecutionMetrics,
300 observer: MetricsObserver,
301 llm_rx: tokio::sync::mpsc::Receiver<LlmRequest>,
302 exec_task: AsyncTask,
303 resp_txs: HashMap<QueryId, tokio::sync::oneshot::Sender<Result<String, String>>>,
305 _vm_driver: AsyncIsleDriver,
308 last_active: std::time::Instant,
311 started_at_ms: i64,
313 last_activity_ms: Arc<AtomicI64>,
316}
317
318impl Session {
319 pub fn new(
335 llm_rx: tokio::sync::mpsc::Receiver<LlmRequest>,
336 exec_task: AsyncTask,
337 metrics: ExecutionMetrics,
338 vm_driver: AsyncIsleDriver,
339 ) -> Self {
340 let observer = metrics.create_observer();
341 let started_at_ms = SystemTime::now()
344 .duration_since(UNIX_EPOCH)
345 .unwrap_or_default()
346 .as_millis() as i64;
347 Self {
348 state: ExecutionState::Running,
349 metrics,
350 observer,
351 llm_rx,
352 exec_task,
353 resp_txs: HashMap::new(),
354 _vm_driver: vm_driver,
355 last_active: std::time::Instant::now(),
356 started_at_ms,
357 last_activity_ms: Arc::new(AtomicI64::new(started_at_ms)),
358 }
359 }
360
361 async fn wait_event(&mut self) -> Result<FeedResult, SessionError> {
366 tokio::select! {
367 result = &mut self.exec_task => {
368 match result {
369 Ok(json_str) => match serde_json::from_str::<serde_json::Value>(&json_str) {
370 Ok(v) => {
371 self.state.complete(v.clone()).map_err(|e| {
372 SessionError::InvalidTransition(e.to_string())
373 })?;
374 self.observer.on_completed(&v);
375 Ok(FeedResult::Finished(ExecutionResult {
376 state: TerminalState::Completed { result: v },
377 metrics: self.take_metrics(),
378 }))
379 }
380 Err(e) => self.fail_with(format!("JSON parse: {e}")),
381 },
382 Err(e) => self.fail_with(e.to_string()),
383 }
384 }
385 Some(req) = self.llm_rx.recv() => {
386 let queries: Vec<LlmQuery> = req.queries.iter().map(|qr| LlmQuery {
387 id: qr.id.clone(),
388 prompt: qr.prompt.clone(),
389 system: qr.system.clone(),
390 max_tokens: qr.max_tokens,
391 grounded: qr.grounded,
392 underspecified: qr.underspecified,
393 cache_breakpoint: qr.cache_breakpoint.clone(),
394 }).collect();
395
396 for qr in req.queries {
397 self.resp_txs.insert(qr.id, qr.resp_tx);
398 }
399
400 self.state.pause(queries.clone()).map_err(|e| {
401 SessionError::InvalidTransition(e.to_string())
402 })?;
403 self.observer.on_paused(&queries);
404 Ok(FeedResult::Paused { queries })
405 }
406 }
407 }
408
409 fn feed_one(
425 &mut self,
426 query_id: &QueryId,
427 response: String,
428 usage: Option<&algocline_core::TokenUsage>,
429 ) -> Result<bool, SessionError> {
430 self.last_active = std::time::Instant::now();
432 let now_ms = SystemTime::now()
435 .duration_since(UNIX_EPOCH)
436 .unwrap_or_default()
437 .as_millis() as i64;
438 self.last_activity_ms.store(now_ms, Ordering::Relaxed);
439
440 self.observer.on_response_fed(query_id, &response, usage);
442
443 if let Some(tx) = self.resp_txs.remove(query_id) {
445 let _ = tx.send(Ok(response.clone()));
446 }
447
448 let complete = self
450 .state
451 .feed(query_id, response)
452 .map_err(SessionError::Feed)?;
453
454 if complete {
455 self.state
457 .take_responses()
458 .map_err(|e| SessionError::InvalidTransition(e.to_string()))?;
459 self.observer.on_resumed();
460 } else {
461 self.observer
462 .on_partial_feed(query_id, self.state.remaining());
463 }
464
465 Ok(complete)
466 }
467
468 fn fail_with(&mut self, msg: String) -> Result<FeedResult, SessionError> {
469 self.state
470 .fail(msg.clone())
471 .map_err(|e| SessionError::InvalidTransition(e.to_string()))?;
472 self.observer.on_failed(&msg);
473 Ok(FeedResult::Finished(ExecutionResult {
474 state: TerminalState::Failed { error: msg },
475 metrics: self.take_metrics(),
476 }))
477 }
478
479 fn take_metrics(&mut self) -> ExecutionMetrics {
480 std::mem::take(&mut self.metrics)
481 }
482
483 pub fn snapshot(
510 &self,
511 pending_filter: Option<&PendingFilter>,
512 include_history: bool,
513 ) -> serde_json::Value {
514 let state_label = match &self.state {
515 ExecutionState::Running => "running",
516 ExecutionState::Paused(_) => "paused",
517 _ => "terminal",
518 };
519
520 let phase = match &self.state {
521 ExecutionState::Running => "running",
522 ExecutionState::Paused(_) => "paused",
523 ExecutionState::Completed { .. } => "completed",
524 ExecutionState::Failed { .. } => "failed",
525 ExecutionState::Cancelled => "cancelled",
526 };
527
528 let mut json = serde_json::json!({
529 "state": state_label,
530 "phase": phase,
531 "started_at": self.started_at_ms,
532 "last_activity_at": self.last_activity_ms.load(Ordering::Relaxed),
533 });
534
535 let metrics = self.metrics.snapshot(include_history);
536 if !metrics.is_null() {
537 json["metrics"] = metrics;
538 }
539
540 if let ExecutionState::Paused(pending) = &self.state {
542 json["pending_queries"] = pending.remaining().into();
543
544 if let Some(filter) = pending_filter {
545 let items: Vec<serde_json::Value> = pending
546 .pending_queries()
547 .iter()
548 .map(|q| project_query(q, filter))
549 .collect();
550 json["pending"] = serde_json::Value::Array(items);
551 }
552 }
553
554 json
555 }
556
557 pub fn is_expired(&self, ttl: Duration) -> bool {
562 is_expired_impl(self.last_active, ttl)
563 }
564
565 pub(crate) fn into_driver_parts(
573 self,
574 ) -> (
575 AsyncTask,
576 tokio::sync::mpsc::Receiver<LlmRequest>,
577 AsyncIsleDriver,
578 ExecutionMetrics,
579 ) {
580 (self.exec_task, self.llm_rx, self._vm_driver, self.metrics)
581 }
582}
583
584fn is_expired_impl(last_active: std::time::Instant, ttl: Duration) -> bool {
586 std::time::Instant::now().saturating_duration_since(last_active) >= ttl
587}
588
589pub struct SessionRegistry {
624 sessions: Arc<Mutex<HashMap<String, Session>>>,
625}
626
627impl Default for SessionRegistry {
628 fn default() -> Self {
629 Self::new()
630 }
631}
632
633impl SessionRegistry {
634 pub fn new() -> Self {
635 Self {
636 sessions: Arc::new(Mutex::new(HashMap::new())),
637 }
638 }
639
640 pub async fn start_execution(
642 &self,
643 mut session: Session,
644 ) -> Result<(String, FeedResult), SessionError> {
645 let session_id = gen_session_id();
646 let result = session.wait_event().await?;
647
648 if matches!(result, FeedResult::Paused { .. }) {
649 self.sessions
650 .lock()
651 .await
652 .insert(session_id.clone(), session);
653 }
654
655 Ok((session_id, result))
656 }
657
658 pub async fn feed_response(
664 &self,
665 session_id: &str,
666 query_id: &QueryId,
667 response: String,
668 usage: Option<&algocline_core::TokenUsage>,
669 ) -> Result<FeedResult, SessionError> {
670 let complete = {
672 let mut map = self.sessions.lock().await;
673 let session = map
674 .get_mut(session_id)
675 .ok_or_else(|| SessionError::NotFound(session_id.into()))?;
676
677 let complete = session.feed_one(query_id, response, usage)?;
678
679 if !complete {
680 return Ok(FeedResult::Accepted {
681 remaining: session.state.remaining(),
682 });
683 }
684
685 complete
686 };
687
688 debug_assert!(complete);
690 let mut session = {
691 let mut map = self.sessions.lock().await;
692 map.remove(session_id)
693 .ok_or_else(|| SessionError::NotFound(session_id.into()))?
694 };
695
696 let result = session.wait_event().await?;
697
698 if matches!(result, FeedResult::Paused { .. }) {
699 self.sessions
700 .lock()
701 .await
702 .insert(session_id.into(), session);
703 }
704
705 Ok(result)
706 }
707
708 pub async fn resolve_sole_pending_id(&self, session_id: &str) -> Result<QueryId, SessionError> {
714 let map = self.sessions.lock().await;
715 let session = map
716 .get(session_id)
717 .ok_or_else(|| SessionError::NotFound(session_id.into()))?;
718 let keys: Vec<QueryId> = session.resp_txs.keys().cloned().collect();
719 match keys.len() {
720 0 => Err(SessionError::InvalidTransition("no pending queries".into())),
721 1 => keys
722 .into_iter()
723 .next()
724 .ok_or_else(|| SessionError::InvalidTransition("unexpected empty keys".into())),
725 n => Err(SessionError::InvalidTransition(format!(
726 "{n} queries pending; specify query_id explicitly"
727 ))),
728 }
729 }
730
731 pub async fn list_snapshots(
747 &self,
748 pending_filter: Option<&PendingFilter>,
749 include_history: bool,
750 ) -> HashMap<String, serde_json::Value> {
751 let map = self.sessions.lock().await;
752 map.iter()
753 .map(|(id, session)| {
754 (
755 id.clone(),
756 session.snapshot(pending_filter, include_history),
757 )
758 })
759 .collect()
760 }
761
762 pub fn spawn_gc_task(&self, ttl: Duration) {
768 let sessions = Arc::clone(&self.sessions);
769 tokio::spawn(async move {
770 let mut interval = tokio::time::interval(Duration::from_secs(60));
771 loop {
772 interval.tick().await;
773 let mut map = sessions.lock().await;
774 let expired: Vec<String> = map
775 .iter()
776 .filter(|(_, s)| s.is_expired(ttl))
777 .map(|(id, _)| id.clone())
778 .collect();
779 for id in &expired {
780 tracing::info!(session_id = %id, "GC: reaping expired session");
781 map.remove(id);
782 }
783 }
784 });
785 }
786}
787
788fn gen_session_id() -> String {
807 use rand::RngExt;
808 use std::time::{SystemTime, UNIX_EPOCH};
809 let ts = SystemTime::now()
810 .duration_since(UNIX_EPOCH)
811 .unwrap_or_default()
812 .as_nanos();
813 let random: u64 = rand::rng().random();
814 format!("s-{ts:x}-{random:016x}")
815}
816
817#[cfg(test)]
818mod tests {
819 use super::*;
820 use algocline_core::{ExecutionMetrics, LlmQuery, QueryId};
821 use serde_json::json;
822
823 fn make_query(index: usize) -> LlmQuery {
824 LlmQuery {
825 id: QueryId::batch(index),
826 prompt: format!("prompt-{index}"),
827 system: None,
828 max_tokens: 100,
829 grounded: false,
830 underspecified: false,
831 cache_breakpoint: None,
832 }
833 }
834
835 #[test]
838 fn to_json_accepted() {
839 let result = FeedResult::Accepted { remaining: 3 };
840 let json = result.to_json("s-123");
841 assert_eq!(json["status"], "accepted");
842 assert_eq!(json["remaining"], 3);
843 }
844
845 #[test]
846 fn to_json_paused_single_query() {
847 let query = LlmQuery {
848 id: QueryId::single(),
849 prompt: "What is 2+2?".into(),
850 system: Some("You are a calculator.".into()),
851 max_tokens: 50,
852 grounded: false,
853 underspecified: false,
854 cache_breakpoint: None,
855 };
856 let result = FeedResult::Paused {
857 queries: vec![query],
858 };
859 let json = result.to_json("s-abc");
860
861 assert_eq!(json["status"], "needs_response");
862 assert_eq!(json["session_id"], "s-abc");
863 assert_eq!(json["prompt"], "What is 2+2?");
864 assert_eq!(json["system"], "You are a calculator.");
865 assert_eq!(json["max_tokens"], 50);
866 assert!(json.get("queries").is_none());
868 assert!(
870 json.get("grounded").is_none(),
871 "grounded key must be absent when false"
872 );
873 assert!(
875 json.get("underspecified").is_none(),
876 "underspecified key must be absent when false"
877 );
878 }
879
880 #[test]
881 fn to_json_paused_single_query_grounded() {
882 let query = LlmQuery {
883 id: QueryId::single(),
884 prompt: "verify this claim".into(),
885 system: None,
886 max_tokens: 200,
887 grounded: true,
888 underspecified: false,
889 cache_breakpoint: None,
890 };
891 let result = FeedResult::Paused {
892 queries: vec![query],
893 };
894 let json = result.to_json("s-grounded");
895
896 assert_eq!(json["status"], "needs_response");
897 assert_eq!(
898 json["grounded"], true,
899 "grounded must appear in single-query MCP JSON"
900 );
901 }
902
903 #[test]
904 fn to_json_paused_single_query_underspecified() {
905 let query = LlmQuery {
906 id: QueryId::single(),
907 prompt: "what output format do you need?".into(),
908 system: None,
909 max_tokens: 200,
910 grounded: false,
911 underspecified: true,
912 cache_breakpoint: None,
913 };
914 let result = FeedResult::Paused {
915 queries: vec![query],
916 };
917 let json = result.to_json("s-underspec");
918
919 assert_eq!(json["status"], "needs_response");
920 assert_eq!(
921 json["underspecified"], true,
922 "underspecified must appear in single-query MCP JSON"
923 );
924 assert!(
925 json.get("grounded").is_none(),
926 "grounded must be absent when false"
927 );
928 }
929
930 #[test]
931 fn to_json_paused_multiple_queries_mixed_grounded() {
932 let grounded_query = LlmQuery {
933 id: QueryId::batch(0),
934 prompt: "verify".into(),
935 system: None,
936 max_tokens: 100,
937 grounded: true,
938 underspecified: false,
939 cache_breakpoint: None,
940 };
941 let normal_query = LlmQuery {
942 id: QueryId::batch(1),
943 prompt: "generate".into(),
944 system: None,
945 max_tokens: 100,
946 grounded: false,
947 underspecified: false,
948 cache_breakpoint: None,
949 };
950 let result = FeedResult::Paused {
951 queries: vec![grounded_query, normal_query],
952 };
953 let json = result.to_json("s-batch");
954
955 let qs = json["queries"].as_array().expect("queries should be array");
956 assert_eq!(
957 qs[0]["grounded"], true,
958 "grounded query must have grounded=true"
959 );
960 assert!(
961 qs[1].get("grounded").is_none(),
962 "non-grounded query must omit grounded key"
963 );
964 }
965
966 #[test]
967 fn to_json_paused_single_query_cache_breakpoint() {
968 let query = LlmQuery {
969 id: QueryId::single(),
970 prompt: "cached user turn".into(),
971 system: Some("shared system prompt".into()),
972 max_tokens: 512,
973 grounded: false,
974 underspecified: false,
975 cache_breakpoint: Some("context".into()),
976 };
977 let result = FeedResult::Paused {
978 queries: vec![query],
979 };
980 let json = result.to_json("s-cache");
981
982 assert_eq!(json["status"], "needs_response");
983 assert_eq!(
984 json["cache_breakpoint"], "context",
985 "cache_breakpoint must appear verbatim in single-query MCP JSON"
986 );
987 assert!(
988 json.get("grounded").is_none(),
989 "grounded must be absent when false"
990 );
991 assert!(
992 json.get("underspecified").is_none(),
993 "underspecified must be absent when false"
994 );
995 }
996
997 #[test]
998 fn to_json_paused_single_query_no_cache_breakpoint_absent() {
999 let query = LlmQuery {
1000 id: QueryId::single(),
1001 prompt: "uncached".into(),
1002 system: None,
1003 max_tokens: 128,
1004 grounded: false,
1005 underspecified: false,
1006 cache_breakpoint: None,
1007 };
1008 let result = FeedResult::Paused {
1009 queries: vec![query],
1010 };
1011 let json = result.to_json("s-nocache");
1012
1013 assert!(
1014 json.get("cache_breakpoint").is_none(),
1015 "cache_breakpoint key must be absent when None"
1016 );
1017 }
1018
1019 #[test]
1020 fn to_json_paused_multiple_queries_mixed_cache_breakpoint() {
1021 let cached_query = LlmQuery {
1022 id: QueryId::batch(0),
1023 prompt: "cached".into(),
1024 system: None,
1025 max_tokens: 100,
1026 grounded: false,
1027 underspecified: false,
1028 cache_breakpoint: Some("prompt".into()),
1029 };
1030 let normal_query = LlmQuery {
1031 id: QueryId::batch(1),
1032 prompt: "uncached".into(),
1033 system: None,
1034 max_tokens: 100,
1035 grounded: false,
1036 underspecified: false,
1037 cache_breakpoint: None,
1038 };
1039 let result = FeedResult::Paused {
1040 queries: vec![cached_query, normal_query],
1041 };
1042 let json = result.to_json("s-batch-cache");
1043
1044 let qs = json["queries"].as_array().expect("queries should be array");
1045 assert_eq!(
1046 qs[0]["cache_breakpoint"], "prompt",
1047 "cached query must forward cache_breakpoint verbatim"
1048 );
1049 assert!(
1050 qs[1].get("cache_breakpoint").is_none(),
1051 "uncached query must omit cache_breakpoint key"
1052 );
1053 }
1054
1055 #[test]
1056 fn to_json_paused_multiple_queries_mixed_underspecified() {
1057 let underspec_query = LlmQuery {
1058 id: QueryId::batch(0),
1059 prompt: "clarify intent".into(),
1060 system: None,
1061 max_tokens: 100,
1062 grounded: false,
1063 underspecified: true,
1064 cache_breakpoint: None,
1065 };
1066 let normal_query = LlmQuery {
1067 id: QueryId::batch(1),
1068 prompt: "generate".into(),
1069 system: None,
1070 max_tokens: 100,
1071 grounded: false,
1072 underspecified: false,
1073 cache_breakpoint: None,
1074 };
1075 let result = FeedResult::Paused {
1076 queries: vec![underspec_query, normal_query],
1077 };
1078 let json = result.to_json("s-batch-us");
1079
1080 let qs = json["queries"].as_array().expect("queries should be array");
1081 assert_eq!(
1082 qs[0]["underspecified"], true,
1083 "underspecified query must have underspecified=true"
1084 );
1085 assert!(
1086 qs[1].get("underspecified").is_none(),
1087 "non-underspecified query must omit underspecified key"
1088 );
1089 }
1090
1091 #[test]
1092 fn to_json_paused_single_query_no_system() {
1093 let query = LlmQuery {
1094 id: QueryId::single(),
1095 prompt: "hello".into(),
1096 system: None,
1097 max_tokens: 1024,
1098 grounded: false,
1099 underspecified: false,
1100 cache_breakpoint: None,
1101 };
1102 let result = FeedResult::Paused {
1103 queries: vec![query],
1104 };
1105 let json = result.to_json("s-x");
1106
1107 assert_eq!(json["status"], "needs_response");
1108 assert!(json["system"].is_null());
1109 }
1110
1111 #[test]
1112 fn to_json_paused_multiple_queries() {
1113 let queries = vec![make_query(0), make_query(1), make_query(2)];
1114 let result = FeedResult::Paused { queries };
1115 let json = result.to_json("s-multi");
1116
1117 assert_eq!(json["status"], "needs_response");
1118 assert_eq!(json["session_id"], "s-multi");
1119
1120 let qs = json["queries"].as_array().expect("queries should be array");
1121 assert_eq!(qs.len(), 3);
1122 assert_eq!(qs[0]["id"], "q-0");
1123 assert_eq!(qs[0]["prompt"], "prompt-0");
1124 assert_eq!(qs[1]["id"], "q-1");
1125 assert_eq!(qs[2]["id"], "q-2");
1126 }
1127
1128 #[test]
1129 fn to_json_finished_completed() {
1130 let result = FeedResult::Finished(ExecutionResult {
1131 state: TerminalState::Completed {
1132 result: json!({"answer": 42}),
1133 },
1134 metrics: ExecutionMetrics::new(),
1135 });
1136 let json = result.to_json("s-done");
1137
1138 assert_eq!(json["status"], "completed");
1139 assert_eq!(json["result"]["answer"], 42);
1140 assert!(json.get("stats").is_some());
1141 }
1142
1143 #[test]
1144 fn to_json_finished_failed() {
1145 let result = FeedResult::Finished(ExecutionResult {
1146 state: TerminalState::Failed {
1147 error: "lua error: bad argument".into(),
1148 },
1149 metrics: ExecutionMetrics::new(),
1150 });
1151 let json = result.to_json("s-err");
1152
1153 assert_eq!(json["status"], "error");
1154 assert_eq!(json["error"], "lua error: bad argument");
1155 }
1156
1157 #[test]
1158 fn to_json_finished_cancelled() {
1159 let result = FeedResult::Finished(ExecutionResult {
1160 state: TerminalState::Cancelled,
1161 metrics: ExecutionMetrics::new(),
1162 });
1163 let json = result.to_json("s-cancel");
1164
1165 assert_eq!(json["status"], "cancelled");
1166 assert!(json.get("stats").is_some());
1167 }
1168
1169 #[test]
1172 fn session_id_starts_with_prefix() {
1173 let id = gen_session_id();
1174 assert!(id.starts_with("s-"), "id should start with 's-': {id}");
1175 }
1176
1177 #[test]
1178 fn session_id_uniqueness() {
1179 let ids: Vec<String> = (0..10).map(|_| gen_session_id()).collect();
1180 let set: std::collections::HashSet<&String> = ids.iter().collect();
1181 assert_eq!(set.len(), 10, "10 IDs should all be unique");
1182 }
1183
1184 #[test]
1191 fn is_expired_impl_fresh_instant_not_expired() {
1192 let now = std::time::Instant::now();
1194 assert!(!is_expired_impl(now, Duration::from_secs(1)));
1195 }
1196
1197 #[test]
1198 fn is_expired_impl_old_instant_expired() {
1199 let two_hours_ago = std::time::Instant::now()
1201 .checked_sub(Duration::from_secs(7200))
1202 .expect("checked_sub should succeed with sane duration");
1203 assert!(is_expired_impl(two_hours_ago, Duration::from_secs(3600)));
1205 }
1206
1207 #[test]
1208 fn is_expired_impl_not_yet_expired() {
1209 let one_hour_ago = std::time::Instant::now()
1211 .checked_sub(Duration::from_secs(3600))
1212 .expect("checked_sub should succeed with sane duration");
1213 assert!(!is_expired_impl(one_hour_ago, Duration::from_secs(10800)));
1215 }
1216
1217 #[test]
1218 fn is_expired_impl_zero_ttl_always_expired() {
1219 let now = std::time::Instant::now();
1221 assert!(is_expired_impl(now, Duration::ZERO));
1222 }
1223
1224 #[test]
1227 fn pending_filter_default_is_all_off() {
1228 let f = PendingFilter::default();
1229 assert!(!f.query_id);
1230 assert!(!f.max_tokens);
1231 assert!(!f.system);
1232 assert!(!f.grounded);
1233 assert!(!f.underspecified);
1234 assert!(matches!(f.prompt, PromptProjection::Off));
1235 }
1236
1237 #[test]
1238 fn pending_filter_preset_meta_flags() {
1239 let f = PendingFilter::preset_meta();
1240 assert!(f.query_id);
1241 assert!(f.max_tokens);
1242 assert!(!f.system);
1243 assert!(!f.grounded);
1244 assert!(!f.underspecified);
1245 assert!(
1246 matches!(f.prompt, PromptProjection::Off),
1247 "meta preset must not project prompt content"
1248 );
1249 }
1250
1251 #[test]
1252 fn pending_filter_preset_preview_uses_default_chars() {
1253 let f = PendingFilter::preset_preview();
1254 assert!(f.query_id);
1255 assert!(f.max_tokens);
1256 match f.prompt {
1257 PromptProjection::Preview { chars } => {
1258 assert_eq!(chars, DEFAULT_PROMPT_PREVIEW_CHARS);
1259 }
1260 other => panic!("expected Preview, got {other:?}"),
1261 }
1262 }
1263
1264 #[test]
1265 fn pending_filter_preset_preview_with_custom_chars() {
1266 let f = PendingFilter::preset_preview_with(42);
1267 match f.prompt {
1268 PromptProjection::Preview { chars } => assert_eq!(chars, 42),
1269 other => panic!("expected Preview {{chars: 42}}, got {other:?}"),
1270 }
1271 }
1272
1273 #[test]
1274 fn pending_filter_preset_full_flags_all_on() {
1275 let f = PendingFilter::preset_full();
1276 assert!(f.query_id);
1277 assert!(f.max_tokens);
1278 assert!(f.system);
1279 assert!(f.grounded);
1280 assert!(f.underspecified);
1281 assert!(matches!(f.prompt, PromptProjection::Full));
1282 }
1283
1284 #[test]
1285 fn pending_filter_from_preset_known_names() {
1286 assert!(PendingFilter::from_preset("meta").is_some());
1287 assert!(PendingFilter::from_preset("preview").is_some());
1288 assert!(PendingFilter::from_preset("full").is_some());
1289 }
1290
1291 #[test]
1292 fn pending_filter_from_preset_unknown_returns_none() {
1293 assert!(PendingFilter::from_preset("").is_none());
1296 assert!(PendingFilter::from_preset("META").is_none());
1297 assert!(PendingFilter::from_preset("bogus").is_none());
1298 }
1299
1300 #[test]
1301 fn pending_filter_from_preset_with_overrides_preview_chars() {
1302 let f = PendingFilter::from_preset_with("preview", 73).unwrap();
1305 match f.prompt {
1306 PromptProjection::Preview { chars } => assert_eq!(chars, 73),
1307 other => panic!("expected Preview {{chars: 73}}, got {other:?}"),
1308 }
1309
1310 let f_meta = PendingFilter::from_preset_with("meta", 73).unwrap();
1311 assert!(matches!(f_meta.prompt, PromptProjection::Off));
1312
1313 let f_full = PendingFilter::from_preset_with("full", 73).unwrap();
1314 assert!(matches!(f_full.prompt, PromptProjection::Full));
1315 }
1316
1317 #[test]
1320 fn project_query_default_filter_produces_empty_object() {
1321 let q = make_query(0);
1322 let v = project_query(&q, &PendingFilter::default());
1323 let obj = v.as_object().expect("object");
1324 assert!(obj.is_empty(), "default filter should project nothing");
1325 }
1326
1327 #[test]
1328 fn project_query_meta_preset_has_id_and_max_tokens_only() {
1329 let q = make_query(0);
1330 let v = project_query(&q, &PendingFilter::preset_meta());
1331 let obj = v.as_object().expect("object");
1332 assert_eq!(obj.len(), 2);
1333 assert_eq!(v["query_id"], "q-0");
1334 assert_eq!(v["max_tokens"], 100);
1335 assert!(obj.get("prompt").is_none());
1336 assert!(obj.get("prompt_preview").is_none());
1337 assert!(obj.get("system").is_none());
1338 assert!(obj.get("grounded").is_none());
1339 assert!(obj.get("underspecified").is_none());
1340 }
1341
1342 #[test]
1343 fn project_query_full_preset_has_all_fields() {
1344 let q = LlmQuery {
1345 id: QueryId::batch(0),
1346 prompt: "hi".into(),
1347 system: Some("sys".into()),
1348 max_tokens: 100,
1349 grounded: true,
1350 underspecified: true,
1351 cache_breakpoint: None,
1352 };
1353 let v = project_query(&q, &PendingFilter::preset_full());
1354 assert_eq!(v["query_id"], "q-0");
1355 assert_eq!(v["max_tokens"], 100);
1356 assert_eq!(v["system"], "sys");
1357 assert_eq!(v["grounded"], true);
1358 assert_eq!(v["underspecified"], true);
1359 assert_eq!(v["prompt"], "hi");
1360 assert!(v.get("prompt_preview").is_none());
1361 }
1362
1363 #[test]
1364 fn project_query_preview_truncates_at_char_count() {
1365 let q = LlmQuery {
1366 id: QueryId::batch(0),
1367 prompt: "abcdefghij".into(),
1368 system: None,
1369 max_tokens: 10,
1370 grounded: false,
1371 underspecified: false,
1372 cache_breakpoint: None,
1373 };
1374 let v = project_query(&q, &PendingFilter::preset_preview_with(5));
1375 assert_eq!(v["prompt_preview"], "abcde");
1376 assert!(v.get("prompt").is_none());
1377 }
1378
1379 #[test]
1380 fn project_query_preview_utf8_multibyte_safe() {
1381 let prompt = "あいうえお";
1385 let q = LlmQuery {
1386 id: QueryId::batch(0),
1387 prompt: prompt.to_string(),
1388 system: None,
1389 max_tokens: 10,
1390 grounded: false,
1391 underspecified: false,
1392 cache_breakpoint: None,
1393 };
1394 let v = project_query(&q, &PendingFilter::preset_preview_with(3));
1395 let preview = v["prompt_preview"].as_str().expect("str");
1396 assert_eq!(preview, "あいう");
1397 assert_eq!(preview.chars().count(), 3);
1398 }
1399
1400 #[test]
1401 fn project_query_preview_chars_over_length_returns_whole_prompt() {
1402 let q = LlmQuery {
1403 id: QueryId::batch(0),
1404 prompt: "abc".into(),
1405 system: None,
1406 max_tokens: 10,
1407 grounded: false,
1408 underspecified: false,
1409 cache_breakpoint: None,
1410 };
1411 let v = project_query(&q, &PendingFilter::preset_preview_with(100));
1412 assert_eq!(v["prompt_preview"], "abc");
1413 }
1414
1415 #[test]
1416 fn project_query_system_field_null_when_absent() {
1417 let q = LlmQuery {
1418 id: QueryId::batch(0),
1419 prompt: "p".into(),
1420 system: None,
1421 max_tokens: 10,
1422 grounded: false,
1423 underspecified: false,
1424 cache_breakpoint: None,
1425 };
1426 let filter = PendingFilter {
1427 system: true,
1428 ..Default::default()
1429 };
1430 let v = project_query(&q, &filter);
1431 assert!(
1432 v["system"].is_null(),
1433 "absent system must serialize as null"
1434 );
1435 }
1436
1437 #[test]
1440 fn pending_filter_deserialize_custom_object_preview() {
1441 let raw = serde_json::json!({
1443 "query_id": true,
1444 "prompt": { "mode": "preview", "chars": 50 }
1445 });
1446 let f: PendingFilter = serde_json::from_value(raw).expect("deserialize");
1447 assert!(f.query_id);
1448 match f.prompt {
1449 PromptProjection::Preview { chars } => assert_eq!(chars, 50),
1450 other => panic!("expected Preview, got {other:?}"),
1451 }
1452 }
1453
1454 #[test]
1455 fn pending_filter_deserialize_partial_object_uses_field_defaults() {
1456 let raw = serde_json::json!({});
1459 let f: PendingFilter = serde_json::from_value(raw).expect("deserialize");
1460 assert!(!f.query_id);
1461 assert!(matches!(f.prompt, PromptProjection::Off));
1462 }
1463
1464 #[test]
1465 fn pending_filter_deserialize_prompt_full_tag() {
1466 let raw = serde_json::json!({ "prompt": { "mode": "full" } });
1467 let f: PendingFilter = serde_json::from_value(raw).expect("deserialize");
1468 assert!(matches!(f.prompt, PromptProjection::Full));
1469 }
1470
1471 fn tmp_dirs() -> (
1479 std::sync::Arc<crate::state::JsonFileStore>,
1480 std::sync::Arc<crate::card::FileCardStore>,
1481 std::path::PathBuf,
1482 ) {
1483 let tmp = tempfile::tempdir().expect("test tempdir");
1484 let root = tmp.path().to_path_buf();
1485 std::mem::forget(tmp);
1486 (
1487 std::sync::Arc::new(crate::state::JsonFileStore::new(root.join("state"))),
1488 std::sync::Arc::new(crate::card::FileCardStore::new(root.join("cards"))),
1489 root.join("scenarios"),
1490 )
1491 }
1492
1493 #[tokio::test]
1500 async fn snapshot_v2_contains_phase_and_timestamps() {
1501 let executor = crate::executor::Executor::new(vec![]).await.unwrap();
1502 let (state_store, card_store, scenarios_dir) = tmp_dirs();
1503
1504 let code = r#"
1506 local response = alc.llm("what is 2+2?")
1507 return response
1508 "#
1509 .to_string();
1510
1511 let session = executor
1512 .start_session(
1513 code,
1514 serde_json::json!({}),
1515 vec![],
1516 vec![],
1517 state_store,
1518 card_store,
1519 scenarios_dir,
1520 )
1521 .await
1522 .unwrap();
1523
1524 let snap = session.snapshot(None, false);
1527
1528 assert!(
1530 snap.get("phase").is_some(),
1531 "snapshot must have 'phase' field"
1532 );
1533 assert_eq!(snap["phase"], "running", "initial state must be running");
1534
1535 assert_eq!(snap["state"], "running");
1537
1538 let started_at = snap["started_at"].as_i64().expect("started_at must be i64");
1540 assert!(started_at > 0, "started_at must be > 0 (unix ms)");
1541
1542 let last_activity = snap["last_activity_at"]
1544 .as_i64()
1545 .expect("last_activity_at must be i64");
1546 assert_eq!(
1547 started_at, last_activity,
1548 "last_activity_at should equal started_at before any feed"
1549 );
1550 }
1551
1552 #[test]
1556 fn snapshot_phase_running_state_label() {
1557 let cases: &[(&str, &str)] = &[
1561 ("running", "running"),
1562 ("paused", "paused"),
1563 ("completed", "completed"),
1564 ("failed", "failed"),
1565 ("cancelled", "cancelled"),
1566 ];
1567 for (state_str, expected_phase) in cases {
1568 let three_value_state = match *state_str {
1572 "running" => "running",
1573 "paused" => "paused",
1574 _ => "terminal",
1575 };
1576 assert_eq!(
1578 *expected_phase, *state_str,
1579 "phase for {state_str} must be the same string"
1580 );
1581 if *state_str != "running" && *state_str != "paused" {
1582 assert_eq!(
1583 three_value_state, "terminal",
1584 "{state_str} must map to 'terminal' in 3-value state"
1585 );
1586 }
1587 }
1588 }
1589
1590 #[tokio::test]
1592 async fn snapshot_conversation_history_opt_in() {
1593 let executor = crate::executor::Executor::new(vec![]).await.unwrap();
1594 let (state_store, card_store, scenarios_dir) = tmp_dirs();
1595
1596 let code = r#"
1597 local response = alc.llm("explain recursion")
1598 return response
1599 "#
1600 .to_string();
1601
1602 let session = executor
1603 .start_session(
1604 code,
1605 serde_json::json!({}),
1606 vec![],
1607 vec![],
1608 state_store,
1609 card_store,
1610 scenarios_dir,
1611 )
1612 .await
1613 .unwrap();
1614
1615 let snap_false = session.snapshot(None, false);
1617 assert!(
1618 snap_false
1619 .get("metrics")
1620 .and_then(|m| m.get("conversation_history"))
1621 .is_none(),
1622 "conversation_history must be absent with include_history=false"
1623 );
1624
1625 let snap_true = session.snapshot(None, true);
1627 if let Some(metrics) = snap_true.get("metrics") {
1630 let _ = metrics.get("conversation_history");
1634 }
1635 }
1636
1637 #[tokio::test]
1639 async fn snapshot_last_activity_at_starts_equal_to_started_at() {
1640 let executor = crate::executor::Executor::new(vec![]).await.unwrap();
1641 let (state_store, card_store, scenarios_dir) = tmp_dirs();
1642
1643 let code = r#"
1644 local response = alc.llm("test query")
1645 return response
1646 "#
1647 .to_string();
1648
1649 let session = executor
1650 .start_session(
1651 code,
1652 serde_json::json!({}),
1653 vec![],
1654 vec![],
1655 state_store,
1656 card_store,
1657 scenarios_dir,
1658 )
1659 .await
1660 .unwrap();
1661
1662 let snap = session.snapshot(None, false);
1663 let started_at = snap["started_at"].as_i64().unwrap_or(-1);
1664 let last_activity = snap["last_activity_at"].as_i64().unwrap_or(-2);
1665
1666 assert_eq!(
1667 started_at, last_activity,
1668 "last_activity_at must equal started_at before any feed_one"
1669 );
1670 assert!(started_at > 0, "started_at must be positive unix ms");
1671 }
1672}