1use parking_lot::RwLock;
33use std::collections::HashSet;
34use std::sync::{
35 Arc,
36 atomic::{AtomicU32, Ordering},
37};
38use tokio::sync::Mutex;
39use tokio::task::JoinSet;
40
41use serde_json::Value as JsonValue;
42use tracing::{Instrument as _, info_span};
43use zeph_db::{DbPool, sql};
44use zeph_llm::LlmProvider;
45use zeph_llm::any::AnyProvider;
46use zeph_llm::provider::{Message, Role};
47
48use zeph_common::SessionId;
49
50use crate::agent::error::AgentError;
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
72#[non_exhaustive]
73pub enum ToolRiskCategory {
74 Shell,
76 FileWrite,
78 ExfilCapable,
80 McpUnclassified,
89 Low,
91}
92
93#[derive(Debug, Clone, PartialEq, Eq)]
97#[non_exhaustive]
98pub enum ProbeVerdict {
99 Allow,
101 Deny {
104 reason: String,
106 },
107 Skip,
110}
111
112#[derive(Debug, Clone)]
118pub struct SentinelEvent {
119 pub id: i64,
121 pub session_id: SessionId,
123 pub turn_number: u64,
125 pub event_type: String,
127 pub tool_id: Option<String>,
129 pub risk_signal: Option<String>,
131 pub risk_level: String,
133 pub probe_verdict: Option<String>,
135 pub context_summary: Option<String>,
137 pub created_at: i64,
139}
140
141pub trait SafetyProbe: Send + Sync {
157 fn evaluate<'a>(
165 &'a self,
166 tool_id: &'a str,
167 tool_args: &'a JsonValue,
168 trajectory: &'a [SentinelEvent],
169 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>>;
170}
171
172pub struct LlmSafetyProbe {
180 provider: Arc<AnyProvider>,
181 timeout_ms: u64,
182 deny_on_timeout: bool,
183}
184
185impl LlmSafetyProbe {
186 #[must_use]
194 pub fn new(provider: Arc<AnyProvider>, timeout_ms: u64, deny_on_timeout: bool) -> Self {
195 Self {
196 provider,
197 timeout_ms,
198 deny_on_timeout,
199 }
200 }
201
202 fn build_prompt(
203 tool_id: &str,
204 tool_args: &JsonValue,
205 trajectory: &[SentinelEvent],
206 ) -> Vec<Message> {
207 let context = if trajectory.is_empty() {
208 "No prior events in this session.".to_owned()
209 } else {
210 trajectory
211 .iter()
212 .map(|e| {
213 format!(
214 "[turn {}] {} | tool={} | {}",
215 e.turn_number,
216 e.event_type,
217 e.tool_id.as_deref().unwrap_or("-"),
218 e.context_summary.as_deref().unwrap_or(""),
219 )
220 })
221 .collect::<Vec<_>>()
222 .join("\n")
223 };
224
225 let args_preview = serde_json::to_string(tool_args)
226 .unwrap_or_default()
227 .chars()
228 .take(500)
229 .collect::<String>();
230
231 let system = "You are a safety evaluator for an AI agent. \
232 Evaluate whether the proposed tool call is safe to execute given the session context. \
233 Reply ONLY with a JSON object: {\"verdict\": \"allow\"} or {\"verdict\": \"deny\", \"reason\": \"...\"}. \
234 Deny only if there is clear evidence of harmful intent or a dangerous pattern. \
235 When uncertain, allow.";
236
237 let user =
238 format!("Tool: {tool_id}\nArgs: {args_preview}\n\nRecent session events:\n{context}");
239
240 vec![
241 Message::from_legacy(Role::System, system),
242 Message::from_legacy(Role::User, user),
243 ]
244 }
245
246 fn parse_verdict(response: &str) -> ProbeVerdict {
247 let start = response.find('{');
249 let end = response.rfind('}');
250 if let (Some(s), Some(e)) = (start, end)
251 && let Ok(v) = serde_json::from_str::<serde_json::Value>(&response[s..=e])
252 {
253 match v.get("verdict").and_then(|x| x.as_str()) {
254 Some("allow") => return ProbeVerdict::Allow,
255 Some("deny") => {
256 let reason = v
257 .get("reason")
258 .and_then(|r| r.as_str())
259 .unwrap_or("safety probe denied this tool call")
260 .to_owned();
261 return ProbeVerdict::Deny { reason };
262 }
263 _ => {}
264 }
265 }
266 tracing::warn!(
268 raw = %response,
269 "ShadowSentinel: probe response could not be parsed, defaulting to Allow"
270 );
271 ProbeVerdict::Allow
272 }
273}
274
275impl SafetyProbe for LlmSafetyProbe {
276 fn evaluate<'a>(
277 &'a self,
278 tool_id: &'a str,
279 tool_args: &'a JsonValue,
280 trajectory: &'a [SentinelEvent],
281 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>> {
282 let span = info_span!("security.shadow.probe", tool_id = %tool_id);
283 Box::pin(
284 async move {
285 let messages = Self::build_prompt(tool_id, tool_args, trajectory);
286 let timeout = std::time::Duration::from_millis(self.timeout_ms);
287
288 match tokio::time::timeout(timeout, self.provider.chat(&messages)).await {
289 Ok(Ok(response)) => Self::parse_verdict(&response),
290 Ok(Err(e)) => {
291 tracing::warn!(error = %e, "ShadowSentinel: probe LLM error");
292 if self.deny_on_timeout {
293 ProbeVerdict::Deny {
294 reason: format!("probe LLM error: {e}"),
295 }
296 } else {
297 ProbeVerdict::Allow
298 }
299 }
300 Err(_) => {
301 tracing::warn!(
302 timeout_ms = self.timeout_ms,
303 "ShadowSentinel: probe timed out"
304 );
305 if self.deny_on_timeout {
306 ProbeVerdict::Deny {
307 reason: "safety probe timed out".to_owned(),
308 }
309 } else {
310 ProbeVerdict::Allow
311 }
312 }
313 }
314 }
315 .instrument(span),
316 )
317 }
318}
319
320#[derive(Clone)]
327pub struct ShadowEventStore {
328 pool: DbPool,
329}
330
331impl ShadowEventStore {
332 #[must_use]
334 pub fn new(pool: DbPool) -> Self {
335 Self { pool }
336 }
337
338 #[tracing::instrument(name = "security.shadow.record", skip_all, fields(event_type = %event.event_type))]
346 pub async fn record(&self, event: &SentinelEvent) -> Result<(), AgentError> {
347 zeph_db::query(sql!(
348 "INSERT INTO safety_shadow_events \
349 (session_id, turn_number, event_type, tool_id, risk_signal, risk_level, \
350 probe_verdict, context_summary, created_at) \
351 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)"
352 ))
353 .bind(event.session_id.as_str())
354 .bind(i64::try_from(event.turn_number).unwrap_or(i64::MAX))
355 .bind(&event.event_type)
356 .bind(&event.tool_id)
357 .bind(&event.risk_signal)
358 .bind(&event.risk_level)
359 .bind(&event.probe_verdict)
360 .bind(&event.context_summary)
361 .bind(event.created_at)
362 .execute(&self.pool)
363 .await
364 .map_err(|e| AgentError::Db(e.into()))?;
365
366 Ok(())
367 }
368
369 #[tracing::instrument(name = "security.shadow.get_trajectory", skip(self), fields(session_id = %session_id))]
377 pub async fn get_trajectory(
378 &self,
379 session_id: &str,
380 limit: usize,
381 ) -> Result<Vec<SentinelEvent>, AgentError> {
382 let rows = zeph_db::query_as::<_, ShadowEventRow>(sql!(
383 "SELECT id, session_id, turn_number, event_type, tool_id, risk_signal, \
384 risk_level, probe_verdict, context_summary, created_at \
385 FROM safety_shadow_events \
386 WHERE session_id = ? \
387 ORDER BY created_at DESC \
388 LIMIT ?"
389 ))
390 .bind(session_id)
391 .bind(i64::try_from(limit).unwrap_or(i64::MAX))
392 .fetch_all(&self.pool)
393 .await
394 .map_err(|e| AgentError::Db(e.into()))?;
395
396 let mut events: Vec<SentinelEvent> = rows.into_iter().map(SentinelEvent::from).collect();
398 events.reverse();
399 Ok(events)
400 }
401
402 #[tracing::instrument(name = "security.shadow.get_tool_history", skip(self), fields(tool_id = %tool_id))]
414 pub async fn get_tool_history(
415 &self,
416 tool_id: &str,
417 exclude_session_id: &str,
418 limit: usize,
419 ) -> Result<Vec<SentinelEvent>, AgentError> {
420 let rows = zeph_db::query_as::<_, ShadowEventRow>(sql!(
421 "SELECT id, session_id, turn_number, event_type, tool_id, risk_signal, \
422 risk_level, probe_verdict, context_summary, created_at \
423 FROM safety_shadow_events \
424 WHERE tool_id = ? AND session_id != ? \
425 ORDER BY created_at DESC \
426 LIMIT ?"
427 ))
428 .bind(tool_id)
429 .bind(exclude_session_id)
430 .bind(i64::try_from(limit).unwrap_or(i64::MAX))
431 .fetch_all(&self.pool)
432 .await
433 .map_err(|e| AgentError::Db(e.into()))?;
434
435 Ok(rows.into_iter().map(SentinelEvent::from).collect())
436 }
437}
438
439#[derive(sqlx::FromRow)]
441struct ShadowEventRow {
442 id: i64,
443 session_id: String,
444 turn_number: i64,
445 event_type: String,
446 tool_id: Option<String>,
447 risk_signal: Option<String>,
448 risk_level: String,
449 probe_verdict: Option<String>,
450 context_summary: Option<String>,
451 created_at: i64,
452}
453
454impl From<ShadowEventRow> for SentinelEvent {
455 fn from(r: ShadowEventRow) -> Self {
456 Self {
457 id: r.id,
458 session_id: SessionId::new(r.session_id),
459 turn_number: u64::try_from(r.turn_number).unwrap_or(0),
460 event_type: r.event_type,
461 tool_id: r.tool_id,
462 risk_signal: r.risk_signal,
463 risk_level: r.risk_level,
464 probe_verdict: r.probe_verdict,
465 context_summary: r.context_summary,
466 created_at: r.created_at,
467 }
468 }
469}
470
471const MAX_PENDING_WRITES: usize = 32;
479
480pub struct ShadowSentinel {
498 store: ShadowEventStore,
499 probe: Box<dyn SafetyProbe>,
500 config: zeph_config::ShadowSentinelConfig,
501 probes_this_turn: AtomicU32,
510 exfil_probes_this_turn: AtomicU32,
520 session_id: SessionId,
521 pending_writes: Mutex<JoinSet<()>>,
524 mcp_tool_ids: Arc<RwLock<HashSet<String>>>,
545}
546
547impl ShadowSentinel {
548 #[must_use]
557 pub fn new(
558 store: ShadowEventStore,
559 probe: Box<dyn SafetyProbe>,
560 config: zeph_config::ShadowSentinelConfig,
561 session_id: impl Into<SessionId>,
562 ) -> Self {
563 Self {
564 store,
565 probe,
566 config,
567 probes_this_turn: AtomicU32::new(0),
568 exfil_probes_this_turn: AtomicU32::new(0),
569 session_id: session_id.into(),
570 pending_writes: Mutex::new(JoinSet::new()),
571 mcp_tool_ids: Arc::new(RwLock::new(HashSet::new())),
572 }
573 }
574
575 #[must_use]
578 pub fn mcp_tool_ids_handle(&self) -> Arc<RwLock<HashSet<String>>> {
579 Arc::clone(&self.mcp_tool_ids)
580 }
581
582 fn is_mcp_tool(&self, tool_id: &str) -> bool {
584 self.mcp_tool_ids.read().contains(tool_id)
585 }
586
587 #[must_use]
593 pub fn classify_tool(&self, qualified_tool_id: &str) -> ToolRiskCategory {
594 if qualified_tool_id == "builtin:shell"
596 || qualified_tool_id == "builtin:bash"
597 || qualified_tool_id.starts_with("builtin:shell")
598 || qualified_tool_id == "bash"
599 || qualified_tool_id == "shell"
600 || qualified_tool_id == "sh"
601 {
602 return ToolRiskCategory::Shell;
603 }
604 if qualified_tool_id == "builtin:write"
605 || qualified_tool_id == "builtin:edit"
606 || qualified_tool_id == "builtin:delete"
607 || qualified_tool_id == "write"
608 || qualified_tool_id == "edit"
609 || qualified_tool_id == "delete"
610 {
611 return ToolRiskCategory::FileWrite;
612 }
613
614 for pattern in &self.config.probe_patterns {
616 if glob_matches(pattern, qualified_tool_id) {
617 if pattern.contains("shell") || pattern.contains("exec") {
619 return ToolRiskCategory::Shell;
620 }
621 if pattern.contains("write")
622 || pattern.contains("edit")
623 || pattern.contains("delete")
624 || pattern.contains("file")
625 {
626 if self.is_mcp_tool(qualified_tool_id) {
627 return ToolRiskCategory::ExfilCapable;
628 }
629 return ToolRiskCategory::FileWrite;
630 }
631 return ToolRiskCategory::ExfilCapable;
632 }
633 }
634
635 if self.is_mcp_tool(qualified_tool_id) {
639 return ToolRiskCategory::McpUnclassified;
640 }
641
642 ToolRiskCategory::Low
643 }
644
645 fn probe_budget_exhausted(&self, category: ToolRiskCategory) -> bool {
661 let max_probes = u32::try_from(self.config.max_probes_per_turn).unwrap_or(u32::MAX);
662
663 if category == ToolRiskCategory::ExfilCapable {
664 let exfil_max = max_probes.saturating_mul(2);
665 let count = self.exfil_probes_this_turn.fetch_add(1, Ordering::Relaxed);
666 if count >= exfil_max {
667 self.exfil_probes_this_turn.fetch_sub(1, Ordering::Relaxed);
668 tracing::debug!(
669 max = exfil_max,
670 "ShadowSentinel: ExfilCapable probe budget exhausted for this turn, skipping"
671 );
672 return true;
673 }
674 return false;
675 }
676
677 let count = self.probes_this_turn.fetch_add(1, Ordering::Relaxed);
679 let effective_max = if category == ToolRiskCategory::McpUnclassified {
680 max_probes.saturating_sub(1)
681 } else {
682 max_probes
683 };
684
685 if count >= effective_max {
686 self.probes_this_turn.fetch_sub(1, Ordering::Relaxed);
688 tracing::debug!(
689 max = self.config.max_probes_per_turn,
690 ?category,
691 "ShadowSentinel: probe budget exhausted for this turn, skipping"
692 );
693 return true;
694 }
695 false
696 }
697
698 async fn load_probe_context(&self, qualified_tool_id: &str) -> Vec<SentinelEvent> {
714 let db_timeout_ms = self.config.probe_timeout_ms.min(2000);
715 let db_timeout = std::time::Duration::from_millis(db_timeout_ms);
716
717 let mut trajectory: Vec<SentinelEvent> = match tokio::time::timeout(
718 db_timeout,
719 self.store
720 .get_trajectory(&self.session_id, self.config.max_context_events),
721 )
722 .await
723 {
724 Ok(Ok(t)) => t
725 .into_iter()
726 .filter(|e| e.event_type != "probe_result")
727 .collect(),
728 Ok(Err(e)) => {
729 tracing::warn!(error = %e, "ShadowSentinel: failed to load trajectory, proceeding without context");
730 vec![]
731 }
732 Err(_) => {
733 tracing::warn!(
734 timeout_ms = db_timeout_ms,
735 "ShadowSentinel: trajectory load timed out, proceeding without context"
736 );
737 vec![]
738 }
739 };
740
741 let cross_session_budget = self.config.max_context_events / 2;
748 let session_budget = self.config.max_context_events - cross_session_budget;
749 if trajectory.len() > session_budget {
750 let excess = trajectory.len() - session_budget;
751 trajectory.drain(0..excess);
752 }
753
754 match tokio::time::timeout(
759 db_timeout,
760 self.store.get_tool_history(
761 qualified_tool_id,
762 self.session_id.as_str(),
763 self.config.max_context_events,
764 ),
765 )
766 .await
767 {
768 Ok(Ok(history)) => {
769 let mut cross_session: Vec<SentinelEvent> = history
772 .into_iter()
773 .filter(|e| e.event_type != "probe_result")
774 .rev()
775 .collect();
776 if cross_session.len() > cross_session_budget {
777 let excess = cross_session.len() - cross_session_budget;
778 cross_session.drain(0..excess);
779 }
780 trajectory.splice(0..0, cross_session);
781 }
782 Ok(Err(e)) => {
783 tracing::warn!(error = %e, "ShadowSentinel: failed to load cross-session tool history, proceeding without it");
784 }
785 Err(_) => {
786 tracing::warn!(
787 timeout_ms = db_timeout_ms,
788 "ShadowSentinel: cross-session tool history load timed out, proceeding without it"
789 );
790 }
791 }
792
793 trajectory
794 }
795
796 #[tracing::instrument(name = "security.shadow.check", skip(self, tool_args), fields(tool_id = %qualified_tool_id))]
818 pub async fn check_tool_call(
819 &self,
820 qualified_tool_id: &str,
821 tool_args: &JsonValue,
822 turn_number: u64,
823 current_risk_level: &str,
824 ) -> ProbeVerdict {
825 if !self.config.enabled {
826 return ProbeVerdict::Skip;
827 }
828
829 let category = self.classify_tool(qualified_tool_id);
830 if category == ToolRiskCategory::Low {
831 return ProbeVerdict::Skip;
832 }
833
834 if self.probe_budget_exhausted(category) {
835 return ProbeVerdict::Skip;
836 }
837
838 let trajectory = self.load_probe_context(qualified_tool_id).await;
839
840 let verdict = self
841 .probe
842 .evaluate(qualified_tool_id, tool_args, &trajectory)
843 .await;
844
845 let probe_verdict_str = match &verdict {
847 ProbeVerdict::Allow => "allow",
848 ProbeVerdict::Deny { .. } => "deny",
849 ProbeVerdict::Skip => "skip",
850 };
851 let summary = match &verdict {
852 ProbeVerdict::Deny { reason } => {
853 format!("probe denied: {}", &reason[..reason.len().min(120)])
854 }
855 ProbeVerdict::Allow => format!("probe allowed {qualified_tool_id}"),
856 ProbeVerdict::Skip => format!("probe skipped {qualified_tool_id}"),
857 };
858 let event = SentinelEvent {
859 id: 0,
860 session_id: self.session_id.clone(),
861 turn_number,
862 event_type: "probe_result".to_owned(),
863 tool_id: Some(qualified_tool_id.to_owned()),
864 risk_signal: None,
865 risk_level: current_risk_level.to_owned(),
866 probe_verdict: Some(probe_verdict_str.to_owned()),
867 context_summary: Some(summary),
868 created_at: unix_now(),
869 };
870 self.persist_event(event, "probe result").await;
871
872 verdict
873 }
874
875 pub async fn record_tool_event(
879 &self,
880 qualified_tool_id: &str,
881 turn_number: u64,
882 risk_level: &str,
883 context_summary: &str,
884 ) {
885 if !self.config.enabled {
886 return;
887 }
888 let event = SentinelEvent {
889 id: 0,
890 session_id: self.session_id.clone(),
891 turn_number,
892 event_type: "tool_call".to_owned(),
893 tool_id: Some(qualified_tool_id.to_owned()),
894 risk_signal: None,
895 risk_level: risk_level.to_owned(),
896 probe_verdict: None,
897 context_summary: Some(context_summary.chars().take(250).collect()),
898 created_at: unix_now(),
899 };
900 self.persist_event(event, "tool event").await;
901 }
902
903 pub async fn drain_pending(&self) {
908 let mut set = {
909 let mut guard = self.pending_writes.lock().await;
910 std::mem::take(&mut *guard)
911 };
912 while set.join_next().await.is_some() {}
913 }
914
915 async fn persist_event(&self, event: SentinelEvent, warn_context: &'static str) {
921 let store = self.store.clone();
922 self.spawn_persist(async move {
923 if let Err(e) = store.record(&event).await {
924 tracing::warn!(error = %e, "ShadowSentinel: failed to persist {warn_context}");
925 }
926 })
927 .await;
928 }
929
930 async fn spawn_persist<F>(&self, fut: F)
936 where
937 F: std::future::Future<Output = ()> + Send + 'static,
938 {
939 let mut set = self.pending_writes.lock().await;
940 while set.try_join_next().is_some() {}
943 if set.len() < MAX_PENDING_WRITES {
944 set.spawn(fut);
945 } else {
946 tracing::debug!(
947 max = MAX_PENDING_WRITES,
948 "ShadowSentinel: pending_writes at capacity, skipping persist"
949 );
950 }
951 }
952
953 pub fn advance_turn(&self) {
958 self.probes_this_turn.store(0, Ordering::Release);
959 self.exfil_probes_this_turn.store(0, Ordering::Release);
960 }
961}
962
963fn unix_now() -> i64 {
967 std::time::SystemTime::now()
968 .duration_since(std::time::UNIX_EPOCH)
969 .ok()
970 .and_then(|d| i64::try_from(d.as_secs()).ok())
971 .unwrap_or(0)
972}
973
974fn glob_matches(pattern: &str, value: &str) -> bool {
977 if pattern == "*" {
978 return true;
979 }
980 let parts: Vec<&str> = pattern.split('*').collect();
982 if parts.len() == 1 {
983 return pattern == value;
984 }
985 let mut remaining = value;
986 for (i, part) in parts.iter().enumerate() {
987 if part.is_empty() {
988 continue;
989 }
990 if i == 0 {
991 if !remaining.starts_with(part) {
992 return false;
993 }
994 remaining = &remaining[part.len()..];
995 } else if i == parts.len() - 1 {
996 return remaining.ends_with(part);
997 } else if let Some(pos) = remaining.find(part) {
998 remaining = &remaining[pos + part.len()..];
999 } else {
1000 return false;
1001 }
1002 }
1003 true
1004}
1005
1006#[cfg(test)]
1011mod tests {
1012 use super::*;
1013
1014 #[tokio::test]
1015 async fn classify_builtin_shell_is_shell_risk() {
1016 let config = zeph_config::ShadowSentinelConfig::default();
1017 let sentinel = make_test_sentinel(config).await;
1018 assert_eq!(
1019 sentinel.classify_tool("builtin:shell"),
1020 ToolRiskCategory::Shell
1021 );
1022 assert_eq!(
1023 sentinel.classify_tool("builtin:bash"),
1024 ToolRiskCategory::Shell
1025 );
1026 }
1027
1028 #[tokio::test]
1029 async fn classify_builtin_write_is_file_write_risk() {
1030 let config = zeph_config::ShadowSentinelConfig::default();
1031 let sentinel = make_test_sentinel(config).await;
1032 assert_eq!(
1033 sentinel.classify_tool("builtin:write"),
1034 ToolRiskCategory::FileWrite
1035 );
1036 assert_eq!(
1037 sentinel.classify_tool("builtin:edit"),
1038 ToolRiskCategory::FileWrite
1039 );
1040 }
1041
1042 #[tokio::test]
1043 async fn classify_low_risk_returns_low() {
1044 let config = zeph_config::ShadowSentinelConfig::default();
1045 let sentinel = make_test_sentinel(config).await;
1046 assert_eq!(
1047 sentinel.classify_tool("builtin:read"),
1048 ToolRiskCategory::Low
1049 );
1050 assert_eq!(
1051 sentinel.classify_tool("builtin:search"),
1052 ToolRiskCategory::Low
1053 );
1054 }
1055
1056 #[tokio::test]
1060 async fn classify_mcp_tool_with_no_keyword_match_is_mcp_unclassified() {
1061 let config = zeph_config::ShadowSentinelConfig::default();
1062 let sentinel = make_test_sentinel(config).await;
1063 sentinel
1064 .mcp_tool_ids_handle()
1065 .write()
1066 .insert("some-server_frobnicate".to_owned());
1067 assert_eq!(
1068 sentinel.classify_tool("some-server_frobnicate"),
1069 ToolRiskCategory::McpUnclassified
1070 );
1071 }
1072
1073 #[tokio::test]
1076 async fn classify_non_mcp_tool_with_no_keyword_match_stays_low() {
1077 let config = zeph_config::ShadowSentinelConfig::default();
1078 let sentinel = make_test_sentinel(config).await;
1079 assert_eq!(
1080 sentinel.classify_tool("some-server_frobnicate"),
1081 ToolRiskCategory::Low
1082 );
1083 }
1084
1085 #[tokio::test]
1086 async fn classify_bare_shell_names_are_shell_risk() {
1087 let config = zeph_config::ShadowSentinelConfig::default();
1088 let sentinel = make_test_sentinel(config).await;
1089 assert_eq!(sentinel.classify_tool("bash"), ToolRiskCategory::Shell);
1090 assert_eq!(sentinel.classify_tool("shell"), ToolRiskCategory::Shell);
1091 assert_eq!(sentinel.classify_tool("sh"), ToolRiskCategory::Shell);
1092 }
1093
1094 #[tokio::test]
1095 async fn classify_bare_file_write_names_are_file_write_risk() {
1096 let config = zeph_config::ShadowSentinelConfig::default();
1097 let sentinel = make_test_sentinel(config).await;
1098 assert_eq!(sentinel.classify_tool("write"), ToolRiskCategory::FileWrite);
1099 assert_eq!(sentinel.classify_tool("edit"), ToolRiskCategory::FileWrite);
1100 assert_eq!(
1101 sentinel.classify_tool("delete"),
1102 ToolRiskCategory::FileWrite
1103 );
1104 }
1105
1106 #[tokio::test]
1111 async fn classify_mcp_tool_write_pattern_escalates_to_exfil_capable() {
1112 let config = zeph_config::ShadowSentinelConfig {
1113 probe_patterns: vec!["*edit*".to_owned()],
1114 ..zeph_config::ShadowSentinelConfig::default()
1115 };
1116 let sentinel = make_test_sentinel(config).await;
1117 assert_eq!(
1119 sentinel.classify_tool("github_edit_file"),
1120 ToolRiskCategory::FileWrite
1121 );
1122 sentinel
1125 .mcp_tool_ids_handle()
1126 .write()
1127 .insert("github_edit_file".to_owned());
1128 assert_eq!(
1129 sentinel.classify_tool("github_edit_file"),
1130 ToolRiskCategory::ExfilCapable
1131 );
1132 }
1133
1134 #[tokio::test]
1143 async fn classify_mcp_tool_write_under_default_config_escalates_to_exfil_capable() {
1144 let config = zeph_config::ShadowSentinelConfig::default();
1145 let sentinel = make_test_sentinel(config).await;
1146 sentinel
1147 .mcp_tool_ids_handle()
1148 .write()
1149 .insert("fs-test_write_file".to_owned());
1150 assert_eq!(
1151 sentinel.classify_tool("fs-test_write_file"),
1152 ToolRiskCategory::ExfilCapable
1153 );
1154 }
1155
1156 #[tokio::test]
1157 async fn advance_turn_resets_counter() {
1158 let config = zeph_config::ShadowSentinelConfig::default();
1159 let sentinel = make_test_sentinel(config).await;
1160 sentinel.probes_this_turn.store(3, Ordering::Relaxed);
1161 sentinel.advance_turn();
1162 assert_eq!(sentinel.probes_this_turn.load(Ordering::Relaxed), 0);
1163 }
1164
1165 #[test]
1166 fn glob_matches_star_wildcard() {
1167 assert!(glob_matches("mcp:*/file_*", "mcp:myserver/file_read"));
1168 assert!(glob_matches("mcp:*/file_*", "mcp:other/file_write"));
1169 assert!(!glob_matches("mcp:*/file_*", "builtin:shell"));
1170 }
1171
1172 #[test]
1173 fn glob_matches_exact() {
1174 assert!(glob_matches("builtin:shell", "builtin:shell"));
1175 assert!(!glob_matches("builtin:shell", "builtin:write"));
1176 }
1177
1178 #[test]
1179 fn parse_verdict_allow() {
1180 let v = LlmSafetyProbe::parse_verdict(r#"{"verdict": "allow"}"#);
1181 assert_eq!(v, ProbeVerdict::Allow);
1182 }
1183
1184 #[test]
1185 fn parse_verdict_deny_with_reason() {
1186 let v =
1187 LlmSafetyProbe::parse_verdict(r#"{"verdict": "deny", "reason": "suspicious pattern"}"#);
1188 assert_eq!(
1189 v,
1190 ProbeVerdict::Deny {
1191 reason: "suspicious pattern".to_owned()
1192 }
1193 );
1194 }
1195
1196 #[test]
1197 fn parse_verdict_unparseable_allows() {
1198 let v = LlmSafetyProbe::parse_verdict("I think this is fine");
1199 assert_eq!(v, ProbeVerdict::Allow);
1200 }
1201
1202 #[tokio::test]
1203 async fn check_tool_call_skips_after_budget_exhausted() {
1204 let config = zeph_config::ShadowSentinelConfig {
1205 enabled: true,
1206 max_probes_per_turn: 2,
1207 ..zeph_config::ShadowSentinelConfig::default()
1208 };
1209 let sentinel = make_test_sentinel(config).await;
1210
1211 let args = serde_json::Value::Object(serde_json::Map::new());
1213 let v1 = sentinel
1214 .check_tool_call("builtin:shell", &args, 1, "calm")
1215 .await;
1216 let v2 = sentinel
1217 .check_tool_call("builtin:shell", &args, 1, "calm")
1218 .await;
1219 assert_ne!(v1, ProbeVerdict::Skip, "first call within budget");
1220 assert_ne!(v2, ProbeVerdict::Skip, "second call within budget");
1221
1222 let v3 = sentinel
1224 .check_tool_call("builtin:shell", &args, 1, "calm")
1225 .await;
1226 assert_eq!(
1227 v3,
1228 ProbeVerdict::Skip,
1229 "third call must be skipped (budget exhausted)"
1230 );
1231 }
1232
1233 #[tokio::test]
1237 async fn check_tool_call_exfil_capable_bypasses_shared_budget_exhaustion() {
1238 let config = zeph_config::ShadowSentinelConfig {
1239 enabled: true,
1240 max_probes_per_turn: 1,
1241 probe_patterns: vec!["*edit*".to_owned()],
1242 ..zeph_config::ShadowSentinelConfig::default()
1243 };
1244 let sentinel = make_test_sentinel(config).await;
1245 sentinel
1246 .mcp_tool_ids_handle()
1247 .write()
1248 .insert("server_edit_file".to_owned());
1249 assert_eq!(
1250 sentinel.classify_tool("server_edit_file"),
1251 ToolRiskCategory::ExfilCapable
1252 );
1253
1254 let args = serde_json::Value::Object(serde_json::Map::new());
1255
1256 let v1 = sentinel
1258 .check_tool_call("builtin:shell", &args, 1, "calm")
1259 .await;
1260 assert_ne!(v1, ProbeVerdict::Skip, "first Shell call within budget");
1261 let v2 = sentinel
1262 .check_tool_call("builtin:shell", &args, 1, "calm")
1263 .await;
1264 assert_eq!(
1265 v2,
1266 ProbeVerdict::Skip,
1267 "second Shell call must be skipped — budget exhausted"
1268 );
1269
1270 let v3 = sentinel
1272 .check_tool_call("server_edit_file", &args, 1, "calm")
1273 .await;
1274 assert_ne!(
1275 v3,
1276 ProbeVerdict::Skip,
1277 "ExfilCapable must not be starved by the shared per-turn budget"
1278 );
1279 }
1280
1281 #[tokio::test]
1285 async fn check_tool_call_exfil_capable_has_finite_cap() {
1286 let config = zeph_config::ShadowSentinelConfig {
1287 enabled: true,
1288 max_probes_per_turn: 1,
1289 probe_patterns: vec!["*edit*".to_owned()],
1290 ..zeph_config::ShadowSentinelConfig::default()
1291 };
1292 let sentinel = make_test_sentinel(config).await;
1293 sentinel
1294 .mcp_tool_ids_handle()
1295 .write()
1296 .insert("server_edit_file".to_owned());
1297
1298 let args = serde_json::Value::Object(serde_json::Map::new());
1299
1300 let v1 = sentinel
1302 .check_tool_call("server_edit_file", &args, 1, "calm")
1303 .await;
1304 let v2 = sentinel
1305 .check_tool_call("server_edit_file", &args, 1, "calm")
1306 .await;
1307 assert_ne!(
1308 v1,
1309 ProbeVerdict::Skip,
1310 "first ExfilCapable call within its own budget"
1311 );
1312 assert_ne!(
1313 v2,
1314 ProbeVerdict::Skip,
1315 "second ExfilCapable call within its own budget"
1316 );
1317
1318 let v3 = sentinel
1320 .check_tool_call("server_edit_file", &args, 1, "calm")
1321 .await;
1322 assert_eq!(
1323 v3,
1324 ProbeVerdict::Skip,
1325 "ExfilCapable's own budget must still be finite (2 * max_probes_per_turn)"
1326 );
1327 }
1328
1329 #[tokio::test]
1333 async fn check_tool_call_mcp_unclassified_reserves_budget_slot() {
1334 let config = zeph_config::ShadowSentinelConfig {
1335 enabled: true,
1336 max_probes_per_turn: 2,
1337 ..zeph_config::ShadowSentinelConfig::default()
1338 };
1339 let sentinel = make_test_sentinel(config).await;
1340 sentinel
1341 .mcp_tool_ids_handle()
1342 .write()
1343 .insert("some-server_frobnicate".to_owned());
1344 assert_eq!(
1345 sentinel.classify_tool("some-server_frobnicate"),
1346 ToolRiskCategory::McpUnclassified
1347 );
1348
1349 let args = serde_json::Value::Object(serde_json::Map::new());
1350
1351 let v1 = sentinel
1353 .check_tool_call("some-server_frobnicate", &args, 1, "calm")
1354 .await;
1355 assert_ne!(
1356 v1,
1357 ProbeVerdict::Skip,
1358 "first McpUnclassified call within reserved share"
1359 );
1360
1361 let v2 = sentinel
1364 .check_tool_call("some-server_frobnicate", &args, 1, "calm")
1365 .await;
1366 assert_eq!(
1367 v2,
1368 ProbeVerdict::Skip,
1369 "second McpUnclassified call must be skipped — reserved share exhausted"
1370 );
1371
1372 let v3 = sentinel
1374 .check_tool_call("builtin:shell", &args, 1, "calm")
1375 .await;
1376 assert_ne!(
1377 v3,
1378 ProbeVerdict::Skip,
1379 "Shell call must still probe using the slot reserved for non-McpUnclassified categories"
1380 );
1381 }
1382
1383 #[tokio::test]
1389 async fn check_tool_call_mcp_unclassified_fully_reserved_out_at_budget_one() {
1390 let config = zeph_config::ShadowSentinelConfig {
1391 enabled: true,
1392 max_probes_per_turn: 1,
1393 ..zeph_config::ShadowSentinelConfig::default()
1394 };
1395 let sentinel = make_test_sentinel(config).await;
1396 sentinel
1397 .mcp_tool_ids_handle()
1398 .write()
1399 .insert("some-server_frobnicate".to_owned());
1400
1401 let args = serde_json::Value::Object(serde_json::Map::new());
1402
1403 let v1 = sentinel
1406 .check_tool_call("some-server_frobnicate", &args, 1, "calm")
1407 .await;
1408 assert_eq!(
1409 v1,
1410 ProbeVerdict::Skip,
1411 "McpUnclassified must get zero share when max_probes_per_turn == 1"
1412 );
1413
1414 let v2 = sentinel
1417 .check_tool_call("builtin:shell", &args, 1, "calm")
1418 .await;
1419 assert_ne!(
1420 v2,
1421 ProbeVerdict::Skip,
1422 "Shell must not be starved by a prior McpUnclassified attempt at max_probes_per_turn == 1"
1423 );
1424 }
1425
1426 #[tokio::test]
1431 async fn check_tool_call_all_categories_skip_at_budget_zero() {
1432 let config = zeph_config::ShadowSentinelConfig {
1433 enabled: true,
1434 max_probes_per_turn: 0,
1435 probe_patterns: vec!["*edit*".to_owned()],
1436 ..zeph_config::ShadowSentinelConfig::default()
1437 };
1438 let sentinel = make_test_sentinel(config).await;
1439 sentinel
1440 .mcp_tool_ids_handle()
1441 .write()
1442 .insert("server_edit_file".to_owned());
1443 sentinel
1444 .mcp_tool_ids_handle()
1445 .write()
1446 .insert("some-server_frobnicate".to_owned());
1447 assert_eq!(
1448 sentinel.classify_tool("server_edit_file"),
1449 ToolRiskCategory::ExfilCapable
1450 );
1451 assert_eq!(
1452 sentinel.classify_tool("some-server_frobnicate"),
1453 ToolRiskCategory::McpUnclassified
1454 );
1455
1456 let args = serde_json::Value::Object(serde_json::Map::new());
1457 assert_eq!(
1458 sentinel
1459 .check_tool_call("builtin:shell", &args, 1, "calm")
1460 .await,
1461 ProbeVerdict::Skip,
1462 "Shell must skip when max_probes_per_turn == 0"
1463 );
1464 assert_eq!(
1465 sentinel
1466 .check_tool_call("some-server_frobnicate", &args, 1, "calm")
1467 .await,
1468 ProbeVerdict::Skip,
1469 "McpUnclassified must skip when max_probes_per_turn == 0"
1470 );
1471 assert_eq!(
1472 sentinel
1473 .check_tool_call("server_edit_file", &args, 1, "calm")
1474 .await,
1475 ProbeVerdict::Skip,
1476 "ExfilCapable's independent budget (2 * 0 == 0) must also skip, not run unbounded"
1477 );
1478 }
1479
1480 #[tokio::test]
1481 async fn check_tool_call_returns_skip_when_disabled() {
1482 let config = zeph_config::ShadowSentinelConfig {
1483 enabled: false,
1484 ..zeph_config::ShadowSentinelConfig::default()
1485 };
1486 let sentinel = make_test_sentinel(config).await;
1487 let args = serde_json::Value::Object(serde_json::Map::new());
1488 let verdict = sentinel
1489 .check_tool_call("builtin:shell", &args, 1, "calm")
1490 .await;
1491 assert_eq!(
1492 verdict,
1493 ProbeVerdict::Skip,
1494 "disabled sentinel must always return Skip without calling the probe"
1495 );
1496 }
1497
1498 #[tokio::test]
1502 async fn drain_pending_awaits_all_tasks() {
1503 use std::sync::atomic::{AtomicU32, Ordering};
1504
1505 let config = zeph_config::ShadowSentinelConfig::default();
1506 let sentinel = make_test_sentinel(config).await;
1507
1508 let counter = Arc::new(AtomicU32::new(0));
1509 for _ in 0..5 {
1510 let c = Arc::clone(&counter);
1511 sentinel
1512 .spawn_persist(async move {
1513 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1514 c.fetch_add(1, Ordering::Relaxed);
1515 })
1516 .await;
1517 }
1518
1519 sentinel.drain_pending().await;
1520
1521 assert_eq!(
1522 counter.load(Ordering::Relaxed),
1523 5,
1524 "drain_pending must join all 5 tasks before returning"
1525 );
1526 }
1527
1528 #[tokio::test]
1534 async fn spawn_persist_beyond_capacity_does_not_panic() {
1535 use std::sync::atomic::{AtomicU32, Ordering};
1536
1537 let config = zeph_config::ShadowSentinelConfig::default();
1538 let sentinel = make_test_sentinel(config).await;
1539 let counter = Arc::new(AtomicU32::new(0));
1540
1541 for _ in 0..(MAX_PENDING_WRITES * 2) {
1544 let c = Arc::clone(&counter);
1545 sentinel
1546 .spawn_persist(async move {
1547 c.fetch_add(1, Ordering::Relaxed);
1548 })
1549 .await;
1550 }
1551
1552 sentinel.drain_pending().await;
1553
1554 let ran = counter.load(Ordering::Relaxed);
1556 assert!(
1557 ran >= u32::try_from(MAX_PENDING_WRITES).unwrap(),
1558 "at least MAX_PENDING_WRITES tasks must complete; ran={ran}"
1559 );
1560 }
1561
1562 async fn make_test_sentinel(config: zeph_config::ShadowSentinelConfig) -> ShadowSentinel {
1567 struct NoopProbe;
1568 impl SafetyProbe for NoopProbe {
1569 fn evaluate<'a>(
1570 &'a self,
1571 _: &'a str,
1572 _: &'a JsonValue,
1573 _: &'a [SentinelEvent],
1574 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>>
1575 {
1576 Box::pin(async { ProbeVerdict::Allow })
1577 }
1578 }
1579 let pool = test_pool().await;
1580 let store = ShadowEventStore::new(pool);
1581 ShadowSentinel::new(store, Box::new(NoopProbe), config, "test-session")
1582 }
1583
1584 async fn test_pool() -> DbPool {
1587 zeph_db::DbConfig {
1588 url: ":memory:".to_owned(),
1589 ..Default::default()
1590 }
1591 .connect()
1592 .await
1593 .expect("connect + migrate in-memory sqlite pool")
1594 }
1595
1596 fn make_event(
1597 session_id: &str,
1598 turn_number: u64,
1599 tool_id: &str,
1600 summary: &str,
1601 ) -> SentinelEvent {
1602 SentinelEvent {
1603 id: 0,
1604 session_id: SessionId::new(session_id),
1605 turn_number,
1606 event_type: "tool_call".to_owned(),
1607 tool_id: Some(tool_id.to_owned()),
1608 risk_signal: None,
1609 risk_level: "elevated".to_owned(),
1610 probe_verdict: None,
1611 context_summary: Some(summary.to_owned()),
1612 created_at: unix_now(),
1613 }
1614 }
1615
1616 #[tokio::test]
1617 async fn get_tool_history_returns_events_across_sessions() {
1618 let store = ShadowEventStore::new(test_pool().await);
1619
1620 store
1621 .record(&make_event(
1622 "session-a",
1623 1,
1624 "builtin:shell",
1625 "session-a ran a command",
1626 ))
1627 .await
1628 .expect("record session-a event");
1629 store
1630 .record(&make_event(
1631 "session-b",
1632 1,
1633 "builtin:shell",
1634 "session-b ran a command",
1635 ))
1636 .await
1637 .expect("record session-b event");
1638 store
1639 .record(&make_event(
1640 "session-a",
1641 2,
1642 "builtin:write",
1643 "unrelated tool",
1644 ))
1645 .await
1646 .expect("record unrelated-tool event");
1647
1648 let history = store
1649 .get_tool_history("builtin:shell", "unrelated-session", 10)
1650 .await
1651 .expect("get_tool_history");
1652
1653 assert_eq!(
1654 history.len(),
1655 2,
1656 "must return events from both non-excluded sessions for the queried tool_id, \
1657 excluding other tools"
1658 );
1659 assert!(history.iter().any(|e| e.session_id.as_str() == "session-a"));
1660 assert!(history.iter().any(|e| e.session_id.as_str() == "session-b"));
1661
1662 let history_excluding_a = store
1663 .get_tool_history("builtin:shell", "session-a", 10)
1664 .await
1665 .expect("get_tool_history");
1666 assert_eq!(
1667 history_excluding_a.len(),
1668 1,
1669 "exclude_session_id must be applied in SQL, not just usable for client-side \
1670 filtering afterward"
1671 );
1672 assert!(
1673 history_excluding_a
1674 .iter()
1675 .all(|e| e.session_id.as_str() != "session-a")
1676 );
1677 }
1678
1679 #[tokio::test]
1682 async fn check_tool_call_incorporates_cross_session_tool_history() {
1683 struct CapturingProbe {
1684 captured: Arc<Mutex<Vec<SentinelEvent>>>,
1685 }
1686 impl SafetyProbe for CapturingProbe {
1687 fn evaluate<'a>(
1688 &'a self,
1689 _tool_id: &'a str,
1690 _tool_args: &'a JsonValue,
1691 trajectory: &'a [SentinelEvent],
1692 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>>
1693 {
1694 let captured = Arc::clone(&self.captured);
1695 let trajectory = trajectory.to_vec();
1696 Box::pin(async move {
1697 *captured.lock().await = trajectory;
1698 ProbeVerdict::Allow
1699 })
1700 }
1701 }
1702
1703 let store = ShadowEventStore::new(test_pool().await);
1704 let other_session = "other-session";
1705 store
1706 .record(&make_event(
1707 other_session,
1708 1,
1709 "builtin:shell",
1710 "other session ran rm -rf",
1711 ))
1712 .await
1713 .expect("record cross-session event");
1714
1715 let captured: Arc<Mutex<Vec<SentinelEvent>>> = Arc::new(Mutex::new(Vec::new()));
1716
1717 let config = zeph_config::ShadowSentinelConfig {
1718 enabled: true,
1719 ..zeph_config::ShadowSentinelConfig::default()
1720 };
1721 let sentinel = ShadowSentinel::new(
1722 store,
1723 Box::new(CapturingProbe {
1724 captured: Arc::clone(&captured),
1725 }),
1726 config,
1727 "current-session",
1728 );
1729
1730 let args = serde_json::Value::Object(serde_json::Map::new());
1731 sentinel
1732 .check_tool_call("builtin:shell", &args, 1, "calm")
1733 .await;
1734
1735 let seen = captured.lock().await;
1736 assert!(
1737 seen.iter().any(|e| e.session_id.as_str() == other_session
1738 && e.context_summary.as_deref() == Some("other session ran rm -rf")),
1739 "probe context must include the cross-session tool history event, got: {seen:?}"
1740 );
1741 }
1742
1743 async fn capture_check_tool_call_trajectory(
1748 store: ShadowEventStore,
1749 config: zeph_config::ShadowSentinelConfig,
1750 session_id: &str,
1751 tool_id: &str,
1752 ) -> Vec<SentinelEvent> {
1753 struct CapturingProbe {
1754 captured: Arc<Mutex<Vec<SentinelEvent>>>,
1755 }
1756 impl SafetyProbe for CapturingProbe {
1757 fn evaluate<'a>(
1758 &'a self,
1759 _tool_id: &'a str,
1760 _tool_args: &'a JsonValue,
1761 trajectory: &'a [SentinelEvent],
1762 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>>
1763 {
1764 let captured = Arc::clone(&self.captured);
1765 let trajectory = trajectory.to_vec();
1766 Box::pin(async move {
1767 *captured.lock().await = trajectory;
1768 ProbeVerdict::Allow
1769 })
1770 }
1771 }
1772
1773 let captured: Arc<Mutex<Vec<SentinelEvent>>> = Arc::new(Mutex::new(Vec::new()));
1774 let sentinel = ShadowSentinel::new(
1775 store,
1776 Box::new(CapturingProbe {
1777 captured: Arc::clone(&captured),
1778 }),
1779 config,
1780 session_id,
1781 );
1782 let args = serde_json::Value::Object(serde_json::Map::new());
1783 sentinel.check_tool_call(tool_id, &args, 1, "calm").await;
1784 captured.lock().await.clone()
1785 }
1786
1787 async fn seed_events(
1791 store: &ShadowEventStore,
1792 session_id: &str,
1793 tool_id: &str,
1794 summary_prefix: &str,
1795 base: i64,
1796 count: u32,
1797 ) {
1798 for i in 0..count {
1799 let mut event = make_event(
1800 session_id,
1801 u64::from(i),
1802 tool_id,
1803 &format!("{summary_prefix}-{i}"),
1804 );
1805 event.created_at = base + i64::from(i);
1806 store.record(&event).await.expect("record seeded event");
1807 }
1808 }
1809
1810 #[tokio::test]
1817 async fn check_tool_call_cap_reserves_cross_session_budget_when_session_heavy() {
1818 let store = ShadowEventStore::new(test_pool().await);
1819 let base = unix_now();
1820 seed_events(
1821 &store,
1822 "current-session",
1823 "builtin:shell",
1824 "session",
1825 base,
1826 4,
1827 )
1828 .await;
1829 seed_events(&store, "other-session", "builtin:shell", "cross", base, 3).await;
1830
1831 let config = zeph_config::ShadowSentinelConfig {
1832 enabled: true,
1833 max_context_events: 4,
1834 ..zeph_config::ShadowSentinelConfig::default()
1835 };
1836 let trajectory =
1837 capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1838 .await;
1839
1840 assert_eq!(
1841 trajectory.len(),
1842 4,
1843 "total must be capped at max_context_events"
1844 );
1845 let cross_session_count = trajectory
1846 .iter()
1847 .filter(|e| e.session_id.as_str() == "other-session")
1848 .count();
1849 assert_eq!(
1850 cross_session_count, 2,
1851 "cross-session budget is max_context_events/2 = 2, and must survive even \
1852 though the session's own trajectory alone fills the whole budget; \
1853 got trajectory: {trajectory:?}"
1854 );
1855 }
1856
1857 #[tokio::test]
1860 async fn check_tool_call_cap_cross_session_heavy_case() {
1861 let store = ShadowEventStore::new(test_pool().await);
1862 let base = unix_now();
1863 seed_events(
1864 &store,
1865 "current-session",
1866 "builtin:shell",
1867 "session",
1868 base,
1869 1,
1870 )
1871 .await;
1872 seed_events(&store, "other-session", "builtin:shell", "cross", base, 4).await;
1873
1874 let config = zeph_config::ShadowSentinelConfig {
1875 enabled: true,
1876 max_context_events: 4,
1877 ..zeph_config::ShadowSentinelConfig::default()
1878 };
1879 let trajectory =
1880 capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1881 .await;
1882
1883 let session_count = trajectory
1884 .iter()
1885 .filter(|e| e.session_id.as_str() == "current-session")
1886 .count();
1887 let cross_session_count = trajectory.len() - session_count;
1888 assert_eq!(
1889 session_count, 1,
1890 "session's own (light) trajectory must not be trimmed"
1891 );
1892 assert_eq!(
1893 cross_session_count, 2,
1894 "cross-session budget is max_context_events/2 = 2"
1895 );
1896 }
1897
1898 #[tokio::test]
1901 async fn check_tool_call_cap_boundary_at_exact_limit() {
1902 let store = ShadowEventStore::new(test_pool().await);
1903 let base = unix_now();
1904 seed_events(
1905 &store,
1906 "current-session",
1907 "builtin:shell",
1908 "session",
1909 base,
1910 2,
1911 )
1912 .await;
1913 seed_events(&store, "other-session", "builtin:shell", "cross", base, 2).await;
1914
1915 let config = zeph_config::ShadowSentinelConfig {
1916 enabled: true,
1917 max_context_events: 4,
1918 ..zeph_config::ShadowSentinelConfig::default()
1919 };
1920 let trajectory =
1921 capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1922 .await;
1923
1924 assert_eq!(
1925 trajectory.len(),
1926 4,
1927 "exactly at the limit: nothing should be dropped"
1928 );
1929 }
1930
1931 #[tokio::test]
1935 async fn check_tool_call_cap_boundary_at_limit_plus_one() {
1936 let store = ShadowEventStore::new(test_pool().await);
1937 let base = unix_now();
1938 seed_events(
1939 &store,
1940 "current-session",
1941 "builtin:shell",
1942 "session",
1943 base,
1944 2,
1945 )
1946 .await;
1947 seed_events(&store, "other-session", "builtin:shell", "cross", base, 3).await;
1948
1949 let config = zeph_config::ShadowSentinelConfig {
1950 enabled: true,
1951 max_context_events: 4,
1952 ..zeph_config::ShadowSentinelConfig::default()
1953 };
1954 let trajectory =
1955 capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1956 .await;
1957
1958 assert_eq!(
1959 trajectory.len(),
1960 4,
1961 "limit+1 overall: exactly one event must be dropped"
1962 );
1963 let cross_summaries: Vec<&str> = trajectory
1964 .iter()
1965 .filter(|e| e.session_id.as_str() == "other-session")
1966 .filter_map(|e| e.context_summary.as_deref())
1967 .collect();
1968 assert_eq!(
1969 cross_summaries,
1970 vec!["cross-1", "cross-2"],
1971 "the oldest cross-session event (cross-0) must be the one dropped, \
1972 got: {cross_summaries:?}"
1973 );
1974 }
1975
1976 #[tokio::test]
1981 async fn check_tool_call_excludes_current_session_from_cross_session_merge() {
1982 let store = ShadowEventStore::new(test_pool().await);
1983 let base = unix_now();
1984 seed_events(&store, "current-session", "builtin:shell", "own", base, 2).await;
1985
1986 let config = zeph_config::ShadowSentinelConfig {
1987 enabled: true,
1988 max_context_events: 10,
1989 ..zeph_config::ShadowSentinelConfig::default()
1990 };
1991 let trajectory =
1992 capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1993 .await;
1994
1995 assert_eq!(
1996 trajectory.len(),
1997 2,
1998 "current session's own events must appear exactly once, not duplicated via \
1999 the cross-session merge; got: {trajectory:?}"
2000 );
2001 }
2002
2003 #[tokio::test]
2008 async fn check_tool_call_excludes_probe_result_events_from_cross_session_merge() {
2009 let store = ShadowEventStore::new(test_pool().await);
2010 let base = unix_now();
2011 let mut event = make_event("other-session", 1, "builtin:shell", "probe verdict leaked");
2012 event.event_type = "probe_result".to_owned();
2013 event.created_at = base;
2014 store
2015 .record(&event)
2016 .await
2017 .expect("record probe_result event");
2018
2019 let config = zeph_config::ShadowSentinelConfig {
2020 enabled: true,
2021 max_context_events: 10,
2022 ..zeph_config::ShadowSentinelConfig::default()
2023 };
2024 let trajectory =
2025 capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
2026 .await;
2027
2028 assert!(
2029 trajectory.is_empty(),
2030 "probe_result events from other sessions must never appear in the \
2031 cross-session merge (LLM isolation invariant), got: {trajectory:?}"
2032 );
2033 }
2034
2035 #[tokio::test]
2046 async fn check_tool_call_falls_open_when_both_db_reads_stall() {
2047 use tracing_subscriber::layer::SubscriberExt as _;
2048
2049 let pool = test_pool().await;
2050 let raw_pool = pool.clone();
2051 let store = ShadowEventStore::new(pool);
2052
2053 let base = unix_now();
2056 seed_events(&store, "current-session", "builtin:shell", "own", base, 2).await;
2057 seed_events(&store, "other-session", "builtin:shell", "cross", base, 2).await;
2058
2059 let messages: Arc<std::sync::Mutex<Vec<String>>> =
2060 Arc::new(std::sync::Mutex::new(Vec::new()));
2061 let layer = MessageCaptureLayer {
2062 messages: messages.clone(),
2063 };
2064 let subscriber = tracing_subscriber::registry().with(layer);
2065 let _guard = tracing::subscriber::set_default(subscriber);
2066
2067 let config = zeph_config::ShadowSentinelConfig {
2068 enabled: true,
2069 probe_timeout_ms: 50,
2070 ..zeph_config::ShadowSentinelConfig::default()
2071 };
2072
2073 let tx = zeph_db::begin_write(&raw_pool)
2076 .await
2077 .expect("hold sole in-memory sqlite connection");
2078
2079 let trajectory =
2080 capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
2081 .await;
2082
2083 drop(tx);
2084
2085 assert!(
2086 trajectory.is_empty(),
2087 "trajectory passed to the probe must be empty when both get_trajectory and \
2088 get_tool_history time out, despite real seeded data existing; got: {trajectory:?}"
2089 );
2090
2091 let captured_logs = messages.lock().unwrap();
2092 assert!(
2093 captured_logs
2094 .iter()
2095 .any(|m| m.contains("trajectory load timed out")),
2096 "expected a warn log for the timed-out get_trajectory read, got: {captured_logs:?}"
2097 );
2098 assert!(
2099 captured_logs
2100 .iter()
2101 .any(|m| m.contains("cross-session tool history load timed out")),
2102 "expected a warn log for the timed-out get_tool_history read, got: {captured_logs:?}"
2103 );
2104 }
2105
2106 #[tokio::test]
2109 async fn record_tool_event_persists_event_normal_path() {
2110 let config = zeph_config::ShadowSentinelConfig {
2111 enabled: true,
2112 ..zeph_config::ShadowSentinelConfig::default()
2113 };
2114 let sentinel = make_test_sentinel(config).await;
2115
2116 sentinel
2117 .record_tool_event("builtin:shell", 3, "elevated", "ran `ls -la`")
2118 .await;
2119 sentinel.drain_pending().await;
2120
2121 let events = sentinel
2122 .store
2123 .get_trajectory("test-session", 10)
2124 .await
2125 .expect("get_trajectory");
2126 assert_eq!(events.len(), 1, "expected exactly one persisted event");
2127 assert_eq!(events[0].event_type, "tool_call");
2128 assert_eq!(events[0].tool_id.as_deref(), Some("builtin:shell"));
2129 assert_eq!(events[0].turn_number, 3);
2130 assert_eq!(events[0].risk_level, "elevated");
2131 assert_eq!(events[0].context_summary.as_deref(), Some("ran `ls -la`"));
2132 }
2133
2134 #[tokio::test]
2135 async fn record_tool_event_disabled_does_not_persist() {
2136 let config = zeph_config::ShadowSentinelConfig {
2137 enabled: false,
2138 ..zeph_config::ShadowSentinelConfig::default()
2139 };
2140 let sentinel = make_test_sentinel(config).await;
2141
2142 sentinel
2143 .record_tool_event("builtin:shell", 1, "elevated", "should be skipped")
2144 .await;
2145 sentinel.drain_pending().await;
2146
2147 let events = sentinel
2148 .store
2149 .get_trajectory("test-session", 10)
2150 .await
2151 .expect("get_trajectory");
2152 assert!(
2153 events.is_empty(),
2154 "record_tool_event must be a no-op when the sentinel is disabled"
2155 );
2156 }
2157
2158 struct MessageCaptureLayer {
2163 messages: Arc<std::sync::Mutex<Vec<String>>>,
2164 }
2165
2166 struct MessageVisitor(String);
2167
2168 impl tracing::field::Visit for MessageVisitor {
2169 fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
2170 if field.name() == "message" {
2171 self.0 = format!("{value:?}");
2172 }
2173 }
2174 }
2175
2176 impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for MessageCaptureLayer {
2177 fn on_event(
2178 &self,
2179 event: &tracing::Event<'_>,
2180 _ctx: tracing_subscriber::layer::Context<'_, S>,
2181 ) {
2182 let mut visitor = MessageVisitor(String::new());
2183 event.record(&mut visitor);
2184 self.messages.lock().unwrap().push(visitor.0);
2185 }
2186 }
2187
2188 #[tokio::test]
2191 async fn record_tool_event_persist_failure_logs_warn_with_tool_event_context() {
2192 use tracing_subscriber::layer::SubscriberExt as _;
2193
2194 struct NoopProbe;
2195 impl SafetyProbe for NoopProbe {
2196 fn evaluate<'a>(
2197 &'a self,
2198 _: &'a str,
2199 _: &'a JsonValue,
2200 _: &'a [SentinelEvent],
2201 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>>
2202 {
2203 Box::pin(async { ProbeVerdict::Allow })
2204 }
2205 }
2206
2207 let pool = test_pool().await;
2208 zeph_db::query(zeph_db::sql!("DROP TABLE safety_shadow_events"))
2209 .execute(&pool)
2210 .await
2211 .expect("drop safety_shadow_events table");
2212 let store = ShadowEventStore::new(pool);
2213 let config = zeph_config::ShadowSentinelConfig {
2214 enabled: true,
2215 ..zeph_config::ShadowSentinelConfig::default()
2216 };
2217 let sentinel = ShadowSentinel::new(store, Box::new(NoopProbe), config, "test-session");
2218
2219 let messages: Arc<std::sync::Mutex<Vec<String>>> =
2220 Arc::new(std::sync::Mutex::new(Vec::new()));
2221 let layer = MessageCaptureLayer {
2222 messages: messages.clone(),
2223 };
2224 let subscriber = tracing_subscriber::registry().with(layer);
2225 let _guard = tracing::subscriber::set_default(subscriber);
2226
2227 sentinel
2228 .record_tool_event("builtin:shell", 1, "elevated", "ran a command")
2229 .await;
2230 sentinel.drain_pending().await;
2231
2232 let captured = messages.lock().unwrap();
2233 assert!(
2234 captured
2235 .iter()
2236 .any(|m| m.contains("failed to persist tool event")),
2237 "expected a warn log with 'failed to persist tool event' context, got: {captured:?}"
2238 );
2239 }
2240}