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 #[tracing::instrument(name = "security.shadow.check", skip(self, tool_args), fields(tool_id = %qualified_tool_id))]
720 pub async fn check_tool_call(
721 &self,
722 qualified_tool_id: &str,
723 tool_args: &JsonValue,
724 turn_number: u64,
725 current_risk_level: &str,
726 ) -> ProbeVerdict {
727 if !self.config.enabled {
728 return ProbeVerdict::Skip;
729 }
730
731 let category = self.classify_tool(qualified_tool_id);
732 if category == ToolRiskCategory::Low {
733 return ProbeVerdict::Skip;
734 }
735
736 if self.probe_budget_exhausted(category) {
737 return ProbeVerdict::Skip;
738 }
739
740 let mut trajectory: Vec<SentinelEvent> = match self
744 .store
745 .get_trajectory(&self.session_id, self.config.max_context_events)
746 .await
747 {
748 Ok(t) => t
749 .into_iter()
750 .filter(|e| e.event_type != "probe_result")
751 .collect(),
752 Err(e) => {
753 tracing::warn!(error = %e, "ShadowSentinel: failed to load trajectory, proceeding without context");
754 vec![]
755 }
756 };
757
758 let cross_session_budget = self.config.max_context_events / 2;
765 let session_budget = self.config.max_context_events - cross_session_budget;
766 if trajectory.len() > session_budget {
767 let excess = trajectory.len() - session_budget;
768 trajectory.drain(0..excess);
769 }
770
771 match self
776 .store
777 .get_tool_history(
778 qualified_tool_id,
779 self.session_id.as_str(),
780 self.config.max_context_events,
781 )
782 .await
783 {
784 Ok(history) => {
785 let mut cross_session: Vec<SentinelEvent> = history
788 .into_iter()
789 .filter(|e| e.event_type != "probe_result")
790 .rev()
791 .collect();
792 if cross_session.len() > cross_session_budget {
793 let excess = cross_session.len() - cross_session_budget;
794 cross_session.drain(0..excess);
795 }
796 trajectory.splice(0..0, cross_session);
797 }
798 Err(e) => {
799 tracing::warn!(error = %e, "ShadowSentinel: failed to load cross-session tool history, proceeding without it");
800 }
801 }
802
803 let verdict = self
804 .probe
805 .evaluate(qualified_tool_id, tool_args, &trajectory)
806 .await;
807
808 let probe_verdict_str = match &verdict {
810 ProbeVerdict::Allow => "allow",
811 ProbeVerdict::Deny { .. } => "deny",
812 ProbeVerdict::Skip => "skip",
813 };
814 let summary = match &verdict {
815 ProbeVerdict::Deny { reason } => {
816 format!("probe denied: {}", &reason[..reason.len().min(120)])
817 }
818 ProbeVerdict::Allow => format!("probe allowed {qualified_tool_id}"),
819 ProbeVerdict::Skip => format!("probe skipped {qualified_tool_id}"),
820 };
821 let event = SentinelEvent {
822 id: 0,
823 session_id: self.session_id.clone(),
824 turn_number,
825 event_type: "probe_result".to_owned(),
826 tool_id: Some(qualified_tool_id.to_owned()),
827 risk_signal: None,
828 risk_level: current_risk_level.to_owned(),
829 probe_verdict: Some(probe_verdict_str.to_owned()),
830 context_summary: Some(summary),
831 created_at: unix_now(),
832 };
833 self.persist_event(event, "probe result").await;
834
835 verdict
836 }
837
838 pub async fn record_tool_event(
842 &self,
843 qualified_tool_id: &str,
844 turn_number: u64,
845 risk_level: &str,
846 context_summary: &str,
847 ) {
848 if !self.config.enabled {
849 return;
850 }
851 let event = SentinelEvent {
852 id: 0,
853 session_id: self.session_id.clone(),
854 turn_number,
855 event_type: "tool_call".to_owned(),
856 tool_id: Some(qualified_tool_id.to_owned()),
857 risk_signal: None,
858 risk_level: risk_level.to_owned(),
859 probe_verdict: None,
860 context_summary: Some(context_summary.chars().take(250).collect()),
861 created_at: unix_now(),
862 };
863 self.persist_event(event, "tool event").await;
864 }
865
866 pub async fn drain_pending(&self) {
871 let mut set = {
872 let mut guard = self.pending_writes.lock().await;
873 std::mem::take(&mut *guard)
874 };
875 while set.join_next().await.is_some() {}
876 }
877
878 async fn persist_event(&self, event: SentinelEvent, warn_context: &'static str) {
884 let store = self.store.clone();
885 self.spawn_persist(async move {
886 if let Err(e) = store.record(&event).await {
887 tracing::warn!(error = %e, "ShadowSentinel: failed to persist {warn_context}");
888 }
889 })
890 .await;
891 }
892
893 async fn spawn_persist<F>(&self, fut: F)
899 where
900 F: std::future::Future<Output = ()> + Send + 'static,
901 {
902 let mut set = self.pending_writes.lock().await;
903 while set.try_join_next().is_some() {}
906 if set.len() < MAX_PENDING_WRITES {
907 set.spawn(fut);
908 } else {
909 tracing::debug!(
910 max = MAX_PENDING_WRITES,
911 "ShadowSentinel: pending_writes at capacity, skipping persist"
912 );
913 }
914 }
915
916 pub fn advance_turn(&self) {
921 self.probes_this_turn.store(0, Ordering::Release);
922 self.exfil_probes_this_turn.store(0, Ordering::Release);
923 }
924}
925
926fn unix_now() -> i64 {
930 std::time::SystemTime::now()
931 .duration_since(std::time::UNIX_EPOCH)
932 .ok()
933 .and_then(|d| i64::try_from(d.as_secs()).ok())
934 .unwrap_or(0)
935}
936
937fn glob_matches(pattern: &str, value: &str) -> bool {
940 if pattern == "*" {
941 return true;
942 }
943 let parts: Vec<&str> = pattern.split('*').collect();
945 if parts.len() == 1 {
946 return pattern == value;
947 }
948 let mut remaining = value;
949 for (i, part) in parts.iter().enumerate() {
950 if part.is_empty() {
951 continue;
952 }
953 if i == 0 {
954 if !remaining.starts_with(part) {
955 return false;
956 }
957 remaining = &remaining[part.len()..];
958 } else if i == parts.len() - 1 {
959 return remaining.ends_with(part);
960 } else if let Some(pos) = remaining.find(part) {
961 remaining = &remaining[pos + part.len()..];
962 } else {
963 return false;
964 }
965 }
966 true
967}
968
969#[cfg(test)]
974mod tests {
975 use super::*;
976
977 #[tokio::test]
978 async fn classify_builtin_shell_is_shell_risk() {
979 let config = zeph_config::ShadowSentinelConfig::default();
980 let sentinel = make_test_sentinel(config).await;
981 assert_eq!(
982 sentinel.classify_tool("builtin:shell"),
983 ToolRiskCategory::Shell
984 );
985 assert_eq!(
986 sentinel.classify_tool("builtin:bash"),
987 ToolRiskCategory::Shell
988 );
989 }
990
991 #[tokio::test]
992 async fn classify_builtin_write_is_file_write_risk() {
993 let config = zeph_config::ShadowSentinelConfig::default();
994 let sentinel = make_test_sentinel(config).await;
995 assert_eq!(
996 sentinel.classify_tool("builtin:write"),
997 ToolRiskCategory::FileWrite
998 );
999 assert_eq!(
1000 sentinel.classify_tool("builtin:edit"),
1001 ToolRiskCategory::FileWrite
1002 );
1003 }
1004
1005 #[tokio::test]
1006 async fn classify_low_risk_returns_low() {
1007 let config = zeph_config::ShadowSentinelConfig::default();
1008 let sentinel = make_test_sentinel(config).await;
1009 assert_eq!(
1010 sentinel.classify_tool("builtin:read"),
1011 ToolRiskCategory::Low
1012 );
1013 assert_eq!(
1014 sentinel.classify_tool("builtin:search"),
1015 ToolRiskCategory::Low
1016 );
1017 }
1018
1019 #[tokio::test]
1023 async fn classify_mcp_tool_with_no_keyword_match_is_mcp_unclassified() {
1024 let config = zeph_config::ShadowSentinelConfig::default();
1025 let sentinel = make_test_sentinel(config).await;
1026 sentinel
1027 .mcp_tool_ids_handle()
1028 .write()
1029 .insert("some-server_frobnicate".to_owned());
1030 assert_eq!(
1031 sentinel.classify_tool("some-server_frobnicate"),
1032 ToolRiskCategory::McpUnclassified
1033 );
1034 }
1035
1036 #[tokio::test]
1039 async fn classify_non_mcp_tool_with_no_keyword_match_stays_low() {
1040 let config = zeph_config::ShadowSentinelConfig::default();
1041 let sentinel = make_test_sentinel(config).await;
1042 assert_eq!(
1043 sentinel.classify_tool("some-server_frobnicate"),
1044 ToolRiskCategory::Low
1045 );
1046 }
1047
1048 #[tokio::test]
1049 async fn classify_bare_shell_names_are_shell_risk() {
1050 let config = zeph_config::ShadowSentinelConfig::default();
1051 let sentinel = make_test_sentinel(config).await;
1052 assert_eq!(sentinel.classify_tool("bash"), ToolRiskCategory::Shell);
1053 assert_eq!(sentinel.classify_tool("shell"), ToolRiskCategory::Shell);
1054 assert_eq!(sentinel.classify_tool("sh"), ToolRiskCategory::Shell);
1055 }
1056
1057 #[tokio::test]
1058 async fn classify_bare_file_write_names_are_file_write_risk() {
1059 let config = zeph_config::ShadowSentinelConfig::default();
1060 let sentinel = make_test_sentinel(config).await;
1061 assert_eq!(sentinel.classify_tool("write"), ToolRiskCategory::FileWrite);
1062 assert_eq!(sentinel.classify_tool("edit"), ToolRiskCategory::FileWrite);
1063 assert_eq!(
1064 sentinel.classify_tool("delete"),
1065 ToolRiskCategory::FileWrite
1066 );
1067 }
1068
1069 #[tokio::test]
1074 async fn classify_mcp_tool_write_pattern_escalates_to_exfil_capable() {
1075 let config = zeph_config::ShadowSentinelConfig {
1076 probe_patterns: vec!["*edit*".to_owned()],
1077 ..zeph_config::ShadowSentinelConfig::default()
1078 };
1079 let sentinel = make_test_sentinel(config).await;
1080 assert_eq!(
1082 sentinel.classify_tool("github_edit_file"),
1083 ToolRiskCategory::FileWrite
1084 );
1085 sentinel
1088 .mcp_tool_ids_handle()
1089 .write()
1090 .insert("github_edit_file".to_owned());
1091 assert_eq!(
1092 sentinel.classify_tool("github_edit_file"),
1093 ToolRiskCategory::ExfilCapable
1094 );
1095 }
1096
1097 #[tokio::test]
1106 async fn classify_mcp_tool_write_under_default_config_escalates_to_exfil_capable() {
1107 let config = zeph_config::ShadowSentinelConfig::default();
1108 let sentinel = make_test_sentinel(config).await;
1109 sentinel
1110 .mcp_tool_ids_handle()
1111 .write()
1112 .insert("fs-test_write_file".to_owned());
1113 assert_eq!(
1114 sentinel.classify_tool("fs-test_write_file"),
1115 ToolRiskCategory::ExfilCapable
1116 );
1117 }
1118
1119 #[tokio::test]
1120 async fn advance_turn_resets_counter() {
1121 let config = zeph_config::ShadowSentinelConfig::default();
1122 let sentinel = make_test_sentinel(config).await;
1123 sentinel.probes_this_turn.store(3, Ordering::Relaxed);
1124 sentinel.advance_turn();
1125 assert_eq!(sentinel.probes_this_turn.load(Ordering::Relaxed), 0);
1126 }
1127
1128 #[test]
1129 fn glob_matches_star_wildcard() {
1130 assert!(glob_matches("mcp:*/file_*", "mcp:myserver/file_read"));
1131 assert!(glob_matches("mcp:*/file_*", "mcp:other/file_write"));
1132 assert!(!glob_matches("mcp:*/file_*", "builtin:shell"));
1133 }
1134
1135 #[test]
1136 fn glob_matches_exact() {
1137 assert!(glob_matches("builtin:shell", "builtin:shell"));
1138 assert!(!glob_matches("builtin:shell", "builtin:write"));
1139 }
1140
1141 #[test]
1142 fn parse_verdict_allow() {
1143 let v = LlmSafetyProbe::parse_verdict(r#"{"verdict": "allow"}"#);
1144 assert_eq!(v, ProbeVerdict::Allow);
1145 }
1146
1147 #[test]
1148 fn parse_verdict_deny_with_reason() {
1149 let v =
1150 LlmSafetyProbe::parse_verdict(r#"{"verdict": "deny", "reason": "suspicious pattern"}"#);
1151 assert_eq!(
1152 v,
1153 ProbeVerdict::Deny {
1154 reason: "suspicious pattern".to_owned()
1155 }
1156 );
1157 }
1158
1159 #[test]
1160 fn parse_verdict_unparseable_allows() {
1161 let v = LlmSafetyProbe::parse_verdict("I think this is fine");
1162 assert_eq!(v, ProbeVerdict::Allow);
1163 }
1164
1165 #[tokio::test]
1166 async fn check_tool_call_skips_after_budget_exhausted() {
1167 let config = zeph_config::ShadowSentinelConfig {
1168 enabled: true,
1169 max_probes_per_turn: 2,
1170 ..zeph_config::ShadowSentinelConfig::default()
1171 };
1172 let sentinel = make_test_sentinel(config).await;
1173
1174 let args = serde_json::Value::Object(serde_json::Map::new());
1176 let v1 = sentinel
1177 .check_tool_call("builtin:shell", &args, 1, "calm")
1178 .await;
1179 let v2 = sentinel
1180 .check_tool_call("builtin:shell", &args, 1, "calm")
1181 .await;
1182 assert_ne!(v1, ProbeVerdict::Skip, "first call within budget");
1183 assert_ne!(v2, ProbeVerdict::Skip, "second call within budget");
1184
1185 let v3 = sentinel
1187 .check_tool_call("builtin:shell", &args, 1, "calm")
1188 .await;
1189 assert_eq!(
1190 v3,
1191 ProbeVerdict::Skip,
1192 "third call must be skipped (budget exhausted)"
1193 );
1194 }
1195
1196 #[tokio::test]
1200 async fn check_tool_call_exfil_capable_bypasses_shared_budget_exhaustion() {
1201 let config = zeph_config::ShadowSentinelConfig {
1202 enabled: true,
1203 max_probes_per_turn: 1,
1204 probe_patterns: vec!["*edit*".to_owned()],
1205 ..zeph_config::ShadowSentinelConfig::default()
1206 };
1207 let sentinel = make_test_sentinel(config).await;
1208 sentinel
1209 .mcp_tool_ids_handle()
1210 .write()
1211 .insert("server_edit_file".to_owned());
1212 assert_eq!(
1213 sentinel.classify_tool("server_edit_file"),
1214 ToolRiskCategory::ExfilCapable
1215 );
1216
1217 let args = serde_json::Value::Object(serde_json::Map::new());
1218
1219 let v1 = sentinel
1221 .check_tool_call("builtin:shell", &args, 1, "calm")
1222 .await;
1223 assert_ne!(v1, ProbeVerdict::Skip, "first Shell call within budget");
1224 let v2 = sentinel
1225 .check_tool_call("builtin:shell", &args, 1, "calm")
1226 .await;
1227 assert_eq!(
1228 v2,
1229 ProbeVerdict::Skip,
1230 "second Shell call must be skipped — budget exhausted"
1231 );
1232
1233 let v3 = sentinel
1235 .check_tool_call("server_edit_file", &args, 1, "calm")
1236 .await;
1237 assert_ne!(
1238 v3,
1239 ProbeVerdict::Skip,
1240 "ExfilCapable must not be starved by the shared per-turn budget"
1241 );
1242 }
1243
1244 #[tokio::test]
1248 async fn check_tool_call_exfil_capable_has_finite_cap() {
1249 let config = zeph_config::ShadowSentinelConfig {
1250 enabled: true,
1251 max_probes_per_turn: 1,
1252 probe_patterns: vec!["*edit*".to_owned()],
1253 ..zeph_config::ShadowSentinelConfig::default()
1254 };
1255 let sentinel = make_test_sentinel(config).await;
1256 sentinel
1257 .mcp_tool_ids_handle()
1258 .write()
1259 .insert("server_edit_file".to_owned());
1260
1261 let args = serde_json::Value::Object(serde_json::Map::new());
1262
1263 let v1 = sentinel
1265 .check_tool_call("server_edit_file", &args, 1, "calm")
1266 .await;
1267 let v2 = sentinel
1268 .check_tool_call("server_edit_file", &args, 1, "calm")
1269 .await;
1270 assert_ne!(
1271 v1,
1272 ProbeVerdict::Skip,
1273 "first ExfilCapable call within its own budget"
1274 );
1275 assert_ne!(
1276 v2,
1277 ProbeVerdict::Skip,
1278 "second ExfilCapable call within its own budget"
1279 );
1280
1281 let v3 = sentinel
1283 .check_tool_call("server_edit_file", &args, 1, "calm")
1284 .await;
1285 assert_eq!(
1286 v3,
1287 ProbeVerdict::Skip,
1288 "ExfilCapable's own budget must still be finite (2 * max_probes_per_turn)"
1289 );
1290 }
1291
1292 #[tokio::test]
1296 async fn check_tool_call_mcp_unclassified_reserves_budget_slot() {
1297 let config = zeph_config::ShadowSentinelConfig {
1298 enabled: true,
1299 max_probes_per_turn: 2,
1300 ..zeph_config::ShadowSentinelConfig::default()
1301 };
1302 let sentinel = make_test_sentinel(config).await;
1303 sentinel
1304 .mcp_tool_ids_handle()
1305 .write()
1306 .insert("some-server_frobnicate".to_owned());
1307 assert_eq!(
1308 sentinel.classify_tool("some-server_frobnicate"),
1309 ToolRiskCategory::McpUnclassified
1310 );
1311
1312 let args = serde_json::Value::Object(serde_json::Map::new());
1313
1314 let v1 = sentinel
1316 .check_tool_call("some-server_frobnicate", &args, 1, "calm")
1317 .await;
1318 assert_ne!(
1319 v1,
1320 ProbeVerdict::Skip,
1321 "first McpUnclassified call within reserved share"
1322 );
1323
1324 let v2 = sentinel
1327 .check_tool_call("some-server_frobnicate", &args, 1, "calm")
1328 .await;
1329 assert_eq!(
1330 v2,
1331 ProbeVerdict::Skip,
1332 "second McpUnclassified call must be skipped — reserved share exhausted"
1333 );
1334
1335 let v3 = sentinel
1337 .check_tool_call("builtin:shell", &args, 1, "calm")
1338 .await;
1339 assert_ne!(
1340 v3,
1341 ProbeVerdict::Skip,
1342 "Shell call must still probe using the slot reserved for non-McpUnclassified categories"
1343 );
1344 }
1345
1346 #[tokio::test]
1352 async fn check_tool_call_mcp_unclassified_fully_reserved_out_at_budget_one() {
1353 let config = zeph_config::ShadowSentinelConfig {
1354 enabled: true,
1355 max_probes_per_turn: 1,
1356 ..zeph_config::ShadowSentinelConfig::default()
1357 };
1358 let sentinel = make_test_sentinel(config).await;
1359 sentinel
1360 .mcp_tool_ids_handle()
1361 .write()
1362 .insert("some-server_frobnicate".to_owned());
1363
1364 let args = serde_json::Value::Object(serde_json::Map::new());
1365
1366 let v1 = sentinel
1369 .check_tool_call("some-server_frobnicate", &args, 1, "calm")
1370 .await;
1371 assert_eq!(
1372 v1,
1373 ProbeVerdict::Skip,
1374 "McpUnclassified must get zero share when max_probes_per_turn == 1"
1375 );
1376
1377 let v2 = sentinel
1380 .check_tool_call("builtin:shell", &args, 1, "calm")
1381 .await;
1382 assert_ne!(
1383 v2,
1384 ProbeVerdict::Skip,
1385 "Shell must not be starved by a prior McpUnclassified attempt at max_probes_per_turn == 1"
1386 );
1387 }
1388
1389 #[tokio::test]
1394 async fn check_tool_call_all_categories_skip_at_budget_zero() {
1395 let config = zeph_config::ShadowSentinelConfig {
1396 enabled: true,
1397 max_probes_per_turn: 0,
1398 probe_patterns: vec!["*edit*".to_owned()],
1399 ..zeph_config::ShadowSentinelConfig::default()
1400 };
1401 let sentinel = make_test_sentinel(config).await;
1402 sentinel
1403 .mcp_tool_ids_handle()
1404 .write()
1405 .insert("server_edit_file".to_owned());
1406 sentinel
1407 .mcp_tool_ids_handle()
1408 .write()
1409 .insert("some-server_frobnicate".to_owned());
1410 assert_eq!(
1411 sentinel.classify_tool("server_edit_file"),
1412 ToolRiskCategory::ExfilCapable
1413 );
1414 assert_eq!(
1415 sentinel.classify_tool("some-server_frobnicate"),
1416 ToolRiskCategory::McpUnclassified
1417 );
1418
1419 let args = serde_json::Value::Object(serde_json::Map::new());
1420 assert_eq!(
1421 sentinel
1422 .check_tool_call("builtin:shell", &args, 1, "calm")
1423 .await,
1424 ProbeVerdict::Skip,
1425 "Shell must skip when max_probes_per_turn == 0"
1426 );
1427 assert_eq!(
1428 sentinel
1429 .check_tool_call("some-server_frobnicate", &args, 1, "calm")
1430 .await,
1431 ProbeVerdict::Skip,
1432 "McpUnclassified must skip when max_probes_per_turn == 0"
1433 );
1434 assert_eq!(
1435 sentinel
1436 .check_tool_call("server_edit_file", &args, 1, "calm")
1437 .await,
1438 ProbeVerdict::Skip,
1439 "ExfilCapable's independent budget (2 * 0 == 0) must also skip, not run unbounded"
1440 );
1441 }
1442
1443 #[tokio::test]
1444 async fn check_tool_call_returns_skip_when_disabled() {
1445 let config = zeph_config::ShadowSentinelConfig {
1446 enabled: false,
1447 ..zeph_config::ShadowSentinelConfig::default()
1448 };
1449 let sentinel = make_test_sentinel(config).await;
1450 let args = serde_json::Value::Object(serde_json::Map::new());
1451 let verdict = sentinel
1452 .check_tool_call("builtin:shell", &args, 1, "calm")
1453 .await;
1454 assert_eq!(
1455 verdict,
1456 ProbeVerdict::Skip,
1457 "disabled sentinel must always return Skip without calling the probe"
1458 );
1459 }
1460
1461 #[tokio::test]
1465 async fn drain_pending_awaits_all_tasks() {
1466 use std::sync::atomic::{AtomicU32, Ordering};
1467
1468 let config = zeph_config::ShadowSentinelConfig::default();
1469 let sentinel = make_test_sentinel(config).await;
1470
1471 let counter = Arc::new(AtomicU32::new(0));
1472 for _ in 0..5 {
1473 let c = Arc::clone(&counter);
1474 sentinel
1475 .spawn_persist(async move {
1476 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1477 c.fetch_add(1, Ordering::Relaxed);
1478 })
1479 .await;
1480 }
1481
1482 sentinel.drain_pending().await;
1483
1484 assert_eq!(
1485 counter.load(Ordering::Relaxed),
1486 5,
1487 "drain_pending must join all 5 tasks before returning"
1488 );
1489 }
1490
1491 #[tokio::test]
1497 async fn spawn_persist_beyond_capacity_does_not_panic() {
1498 use std::sync::atomic::{AtomicU32, Ordering};
1499
1500 let config = zeph_config::ShadowSentinelConfig::default();
1501 let sentinel = make_test_sentinel(config).await;
1502 let counter = Arc::new(AtomicU32::new(0));
1503
1504 for _ in 0..(MAX_PENDING_WRITES * 2) {
1507 let c = Arc::clone(&counter);
1508 sentinel
1509 .spawn_persist(async move {
1510 c.fetch_add(1, Ordering::Relaxed);
1511 })
1512 .await;
1513 }
1514
1515 sentinel.drain_pending().await;
1516
1517 let ran = counter.load(Ordering::Relaxed);
1519 assert!(
1520 ran >= u32::try_from(MAX_PENDING_WRITES).unwrap(),
1521 "at least MAX_PENDING_WRITES tasks must complete; ran={ran}"
1522 );
1523 }
1524
1525 async fn make_test_sentinel(config: zeph_config::ShadowSentinelConfig) -> ShadowSentinel {
1530 struct NoopProbe;
1531 impl SafetyProbe for NoopProbe {
1532 fn evaluate<'a>(
1533 &'a self,
1534 _: &'a str,
1535 _: &'a JsonValue,
1536 _: &'a [SentinelEvent],
1537 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>>
1538 {
1539 Box::pin(async { ProbeVerdict::Allow })
1540 }
1541 }
1542 let pool = test_pool().await;
1543 let store = ShadowEventStore::new(pool);
1544 ShadowSentinel::new(store, Box::new(NoopProbe), config, "test-session")
1545 }
1546
1547 async fn test_pool() -> DbPool {
1550 zeph_db::DbConfig {
1551 url: ":memory:".to_owned(),
1552 ..Default::default()
1553 }
1554 .connect()
1555 .await
1556 .expect("connect + migrate in-memory sqlite pool")
1557 }
1558
1559 fn make_event(
1560 session_id: &str,
1561 turn_number: u64,
1562 tool_id: &str,
1563 summary: &str,
1564 ) -> SentinelEvent {
1565 SentinelEvent {
1566 id: 0,
1567 session_id: SessionId::new(session_id),
1568 turn_number,
1569 event_type: "tool_call".to_owned(),
1570 tool_id: Some(tool_id.to_owned()),
1571 risk_signal: None,
1572 risk_level: "elevated".to_owned(),
1573 probe_verdict: None,
1574 context_summary: Some(summary.to_owned()),
1575 created_at: unix_now(),
1576 }
1577 }
1578
1579 #[tokio::test]
1580 async fn get_tool_history_returns_events_across_sessions() {
1581 let store = ShadowEventStore::new(test_pool().await);
1582
1583 store
1584 .record(&make_event(
1585 "session-a",
1586 1,
1587 "builtin:shell",
1588 "session-a ran a command",
1589 ))
1590 .await
1591 .expect("record session-a event");
1592 store
1593 .record(&make_event(
1594 "session-b",
1595 1,
1596 "builtin:shell",
1597 "session-b ran a command",
1598 ))
1599 .await
1600 .expect("record session-b event");
1601 store
1602 .record(&make_event(
1603 "session-a",
1604 2,
1605 "builtin:write",
1606 "unrelated tool",
1607 ))
1608 .await
1609 .expect("record unrelated-tool event");
1610
1611 let history = store
1612 .get_tool_history("builtin:shell", "unrelated-session", 10)
1613 .await
1614 .expect("get_tool_history");
1615
1616 assert_eq!(
1617 history.len(),
1618 2,
1619 "must return events from both non-excluded sessions for the queried tool_id, \
1620 excluding other tools"
1621 );
1622 assert!(history.iter().any(|e| e.session_id.as_str() == "session-a"));
1623 assert!(history.iter().any(|e| e.session_id.as_str() == "session-b"));
1624
1625 let history_excluding_a = store
1626 .get_tool_history("builtin:shell", "session-a", 10)
1627 .await
1628 .expect("get_tool_history");
1629 assert_eq!(
1630 history_excluding_a.len(),
1631 1,
1632 "exclude_session_id must be applied in SQL, not just usable for client-side \
1633 filtering afterward"
1634 );
1635 assert!(
1636 history_excluding_a
1637 .iter()
1638 .all(|e| e.session_id.as_str() != "session-a")
1639 );
1640 }
1641
1642 #[tokio::test]
1645 async fn check_tool_call_incorporates_cross_session_tool_history() {
1646 struct CapturingProbe {
1647 captured: Arc<Mutex<Vec<SentinelEvent>>>,
1648 }
1649 impl SafetyProbe for CapturingProbe {
1650 fn evaluate<'a>(
1651 &'a self,
1652 _tool_id: &'a str,
1653 _tool_args: &'a JsonValue,
1654 trajectory: &'a [SentinelEvent],
1655 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>>
1656 {
1657 let captured = Arc::clone(&self.captured);
1658 let trajectory = trajectory.to_vec();
1659 Box::pin(async move {
1660 *captured.lock().await = trajectory;
1661 ProbeVerdict::Allow
1662 })
1663 }
1664 }
1665
1666 let store = ShadowEventStore::new(test_pool().await);
1667 let other_session = "other-session";
1668 store
1669 .record(&make_event(
1670 other_session,
1671 1,
1672 "builtin:shell",
1673 "other session ran rm -rf",
1674 ))
1675 .await
1676 .expect("record cross-session event");
1677
1678 let captured: Arc<Mutex<Vec<SentinelEvent>>> = Arc::new(Mutex::new(Vec::new()));
1679
1680 let config = zeph_config::ShadowSentinelConfig {
1681 enabled: true,
1682 ..zeph_config::ShadowSentinelConfig::default()
1683 };
1684 let sentinel = ShadowSentinel::new(
1685 store,
1686 Box::new(CapturingProbe {
1687 captured: Arc::clone(&captured),
1688 }),
1689 config,
1690 "current-session",
1691 );
1692
1693 let args = serde_json::Value::Object(serde_json::Map::new());
1694 sentinel
1695 .check_tool_call("builtin:shell", &args, 1, "calm")
1696 .await;
1697
1698 let seen = captured.lock().await;
1699 assert!(
1700 seen.iter().any(|e| e.session_id.as_str() == other_session
1701 && e.context_summary.as_deref() == Some("other session ran rm -rf")),
1702 "probe context must include the cross-session tool history event, got: {seen:?}"
1703 );
1704 }
1705
1706 async fn capture_check_tool_call_trajectory(
1711 store: ShadowEventStore,
1712 config: zeph_config::ShadowSentinelConfig,
1713 session_id: &str,
1714 tool_id: &str,
1715 ) -> Vec<SentinelEvent> {
1716 struct CapturingProbe {
1717 captured: Arc<Mutex<Vec<SentinelEvent>>>,
1718 }
1719 impl SafetyProbe for CapturingProbe {
1720 fn evaluate<'a>(
1721 &'a self,
1722 _tool_id: &'a str,
1723 _tool_args: &'a JsonValue,
1724 trajectory: &'a [SentinelEvent],
1725 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>>
1726 {
1727 let captured = Arc::clone(&self.captured);
1728 let trajectory = trajectory.to_vec();
1729 Box::pin(async move {
1730 *captured.lock().await = trajectory;
1731 ProbeVerdict::Allow
1732 })
1733 }
1734 }
1735
1736 let captured: Arc<Mutex<Vec<SentinelEvent>>> = Arc::new(Mutex::new(Vec::new()));
1737 let sentinel = ShadowSentinel::new(
1738 store,
1739 Box::new(CapturingProbe {
1740 captured: Arc::clone(&captured),
1741 }),
1742 config,
1743 session_id,
1744 );
1745 let args = serde_json::Value::Object(serde_json::Map::new());
1746 sentinel.check_tool_call(tool_id, &args, 1, "calm").await;
1747 captured.lock().await.clone()
1748 }
1749
1750 async fn seed_events(
1754 store: &ShadowEventStore,
1755 session_id: &str,
1756 tool_id: &str,
1757 summary_prefix: &str,
1758 base: i64,
1759 count: u32,
1760 ) {
1761 for i in 0..count {
1762 let mut event = make_event(
1763 session_id,
1764 u64::from(i),
1765 tool_id,
1766 &format!("{summary_prefix}-{i}"),
1767 );
1768 event.created_at = base + i64::from(i);
1769 store.record(&event).await.expect("record seeded event");
1770 }
1771 }
1772
1773 #[tokio::test]
1780 async fn check_tool_call_cap_reserves_cross_session_budget_when_session_heavy() {
1781 let store = ShadowEventStore::new(test_pool().await);
1782 let base = unix_now();
1783 seed_events(
1784 &store,
1785 "current-session",
1786 "builtin:shell",
1787 "session",
1788 base,
1789 4,
1790 )
1791 .await;
1792 seed_events(&store, "other-session", "builtin:shell", "cross", base, 3).await;
1793
1794 let config = zeph_config::ShadowSentinelConfig {
1795 enabled: true,
1796 max_context_events: 4,
1797 ..zeph_config::ShadowSentinelConfig::default()
1798 };
1799 let trajectory =
1800 capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1801 .await;
1802
1803 assert_eq!(
1804 trajectory.len(),
1805 4,
1806 "total must be capped at max_context_events"
1807 );
1808 let cross_session_count = trajectory
1809 .iter()
1810 .filter(|e| e.session_id.as_str() == "other-session")
1811 .count();
1812 assert_eq!(
1813 cross_session_count, 2,
1814 "cross-session budget is max_context_events/2 = 2, and must survive even \
1815 though the session's own trajectory alone fills the whole budget; \
1816 got trajectory: {trajectory:?}"
1817 );
1818 }
1819
1820 #[tokio::test]
1823 async fn check_tool_call_cap_cross_session_heavy_case() {
1824 let store = ShadowEventStore::new(test_pool().await);
1825 let base = unix_now();
1826 seed_events(
1827 &store,
1828 "current-session",
1829 "builtin:shell",
1830 "session",
1831 base,
1832 1,
1833 )
1834 .await;
1835 seed_events(&store, "other-session", "builtin:shell", "cross", base, 4).await;
1836
1837 let config = zeph_config::ShadowSentinelConfig {
1838 enabled: true,
1839 max_context_events: 4,
1840 ..zeph_config::ShadowSentinelConfig::default()
1841 };
1842 let trajectory =
1843 capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1844 .await;
1845
1846 let session_count = trajectory
1847 .iter()
1848 .filter(|e| e.session_id.as_str() == "current-session")
1849 .count();
1850 let cross_session_count = trajectory.len() - session_count;
1851 assert_eq!(
1852 session_count, 1,
1853 "session's own (light) trajectory must not be trimmed"
1854 );
1855 assert_eq!(
1856 cross_session_count, 2,
1857 "cross-session budget is max_context_events/2 = 2"
1858 );
1859 }
1860
1861 #[tokio::test]
1864 async fn check_tool_call_cap_boundary_at_exact_limit() {
1865 let store = ShadowEventStore::new(test_pool().await);
1866 let base = unix_now();
1867 seed_events(
1868 &store,
1869 "current-session",
1870 "builtin:shell",
1871 "session",
1872 base,
1873 2,
1874 )
1875 .await;
1876 seed_events(&store, "other-session", "builtin:shell", "cross", base, 2).await;
1877
1878 let config = zeph_config::ShadowSentinelConfig {
1879 enabled: true,
1880 max_context_events: 4,
1881 ..zeph_config::ShadowSentinelConfig::default()
1882 };
1883 let trajectory =
1884 capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1885 .await;
1886
1887 assert_eq!(
1888 trajectory.len(),
1889 4,
1890 "exactly at the limit: nothing should be dropped"
1891 );
1892 }
1893
1894 #[tokio::test]
1898 async fn check_tool_call_cap_boundary_at_limit_plus_one() {
1899 let store = ShadowEventStore::new(test_pool().await);
1900 let base = unix_now();
1901 seed_events(
1902 &store,
1903 "current-session",
1904 "builtin:shell",
1905 "session",
1906 base,
1907 2,
1908 )
1909 .await;
1910 seed_events(&store, "other-session", "builtin:shell", "cross", base, 3).await;
1911
1912 let config = zeph_config::ShadowSentinelConfig {
1913 enabled: true,
1914 max_context_events: 4,
1915 ..zeph_config::ShadowSentinelConfig::default()
1916 };
1917 let trajectory =
1918 capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1919 .await;
1920
1921 assert_eq!(
1922 trajectory.len(),
1923 4,
1924 "limit+1 overall: exactly one event must be dropped"
1925 );
1926 let cross_summaries: Vec<&str> = trajectory
1927 .iter()
1928 .filter(|e| e.session_id.as_str() == "other-session")
1929 .filter_map(|e| e.context_summary.as_deref())
1930 .collect();
1931 assert_eq!(
1932 cross_summaries,
1933 vec!["cross-1", "cross-2"],
1934 "the oldest cross-session event (cross-0) must be the one dropped, \
1935 got: {cross_summaries:?}"
1936 );
1937 }
1938
1939 #[tokio::test]
1944 async fn check_tool_call_excludes_current_session_from_cross_session_merge() {
1945 let store = ShadowEventStore::new(test_pool().await);
1946 let base = unix_now();
1947 seed_events(&store, "current-session", "builtin:shell", "own", base, 2).await;
1948
1949 let config = zeph_config::ShadowSentinelConfig {
1950 enabled: true,
1951 max_context_events: 10,
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 2,
1961 "current session's own events must appear exactly once, not duplicated via \
1962 the cross-session merge; got: {trajectory:?}"
1963 );
1964 }
1965
1966 #[tokio::test]
1971 async fn check_tool_call_excludes_probe_result_events_from_cross_session_merge() {
1972 let store = ShadowEventStore::new(test_pool().await);
1973 let base = unix_now();
1974 let mut event = make_event("other-session", 1, "builtin:shell", "probe verdict leaked");
1975 event.event_type = "probe_result".to_owned();
1976 event.created_at = base;
1977 store
1978 .record(&event)
1979 .await
1980 .expect("record probe_result event");
1981
1982 let config = zeph_config::ShadowSentinelConfig {
1983 enabled: true,
1984 max_context_events: 10,
1985 ..zeph_config::ShadowSentinelConfig::default()
1986 };
1987 let trajectory =
1988 capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1989 .await;
1990
1991 assert!(
1992 trajectory.is_empty(),
1993 "probe_result events from other sessions must never appear in the \
1994 cross-session merge (LLM isolation invariant), got: {trajectory:?}"
1995 );
1996 }
1997
1998 #[tokio::test]
2001 async fn record_tool_event_persists_event_normal_path() {
2002 let config = zeph_config::ShadowSentinelConfig {
2003 enabled: true,
2004 ..zeph_config::ShadowSentinelConfig::default()
2005 };
2006 let sentinel = make_test_sentinel(config).await;
2007
2008 sentinel
2009 .record_tool_event("builtin:shell", 3, "elevated", "ran `ls -la`")
2010 .await;
2011 sentinel.drain_pending().await;
2012
2013 let events = sentinel
2014 .store
2015 .get_trajectory("test-session", 10)
2016 .await
2017 .expect("get_trajectory");
2018 assert_eq!(events.len(), 1, "expected exactly one persisted event");
2019 assert_eq!(events[0].event_type, "tool_call");
2020 assert_eq!(events[0].tool_id.as_deref(), Some("builtin:shell"));
2021 assert_eq!(events[0].turn_number, 3);
2022 assert_eq!(events[0].risk_level, "elevated");
2023 assert_eq!(events[0].context_summary.as_deref(), Some("ran `ls -la`"));
2024 }
2025
2026 #[tokio::test]
2027 async fn record_tool_event_disabled_does_not_persist() {
2028 let config = zeph_config::ShadowSentinelConfig {
2029 enabled: false,
2030 ..zeph_config::ShadowSentinelConfig::default()
2031 };
2032 let sentinel = make_test_sentinel(config).await;
2033
2034 sentinel
2035 .record_tool_event("builtin:shell", 1, "elevated", "should be skipped")
2036 .await;
2037 sentinel.drain_pending().await;
2038
2039 let events = sentinel
2040 .store
2041 .get_trajectory("test-session", 10)
2042 .await
2043 .expect("get_trajectory");
2044 assert!(
2045 events.is_empty(),
2046 "record_tool_event must be a no-op when the sentinel is disabled"
2047 );
2048 }
2049
2050 struct MessageCaptureLayer {
2055 messages: Arc<std::sync::Mutex<Vec<String>>>,
2056 }
2057
2058 struct MessageVisitor(String);
2059
2060 impl tracing::field::Visit for MessageVisitor {
2061 fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
2062 if field.name() == "message" {
2063 self.0 = format!("{value:?}");
2064 }
2065 }
2066 }
2067
2068 impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for MessageCaptureLayer {
2069 fn on_event(
2070 &self,
2071 event: &tracing::Event<'_>,
2072 _ctx: tracing_subscriber::layer::Context<'_, S>,
2073 ) {
2074 let mut visitor = MessageVisitor(String::new());
2075 event.record(&mut visitor);
2076 self.messages.lock().unwrap().push(visitor.0);
2077 }
2078 }
2079
2080 #[tokio::test]
2083 async fn record_tool_event_persist_failure_logs_warn_with_tool_event_context() {
2084 use tracing_subscriber::layer::SubscriberExt as _;
2085
2086 struct NoopProbe;
2087 impl SafetyProbe for NoopProbe {
2088 fn evaluate<'a>(
2089 &'a self,
2090 _: &'a str,
2091 _: &'a JsonValue,
2092 _: &'a [SentinelEvent],
2093 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>>
2094 {
2095 Box::pin(async { ProbeVerdict::Allow })
2096 }
2097 }
2098
2099 let pool = test_pool().await;
2100 zeph_db::query(zeph_db::sql!("DROP TABLE safety_shadow_events"))
2101 .execute(&pool)
2102 .await
2103 .expect("drop safety_shadow_events table");
2104 let store = ShadowEventStore::new(pool);
2105 let config = zeph_config::ShadowSentinelConfig {
2106 enabled: true,
2107 ..zeph_config::ShadowSentinelConfig::default()
2108 };
2109 let sentinel = ShadowSentinel::new(store, Box::new(NoopProbe), config, "test-session");
2110
2111 let messages: Arc<std::sync::Mutex<Vec<String>>> =
2112 Arc::new(std::sync::Mutex::new(Vec::new()));
2113 let layer = MessageCaptureLayer {
2114 messages: messages.clone(),
2115 };
2116 let subscriber = tracing_subscriber::registry().with(layer);
2117 let _guard = tracing::subscriber::set_default(subscriber);
2118
2119 sentinel
2120 .record_tool_event("builtin:shell", 1, "elevated", "ran a command")
2121 .await;
2122 sentinel.drain_pending().await;
2123
2124 let captured = messages.lock().unwrap();
2125 assert!(
2126 captured
2127 .iter()
2128 .any(|m| m.contains("failed to persist tool event")),
2129 "expected a warn log with 'failed to persist tool event' context, got: {captured:?}"
2130 );
2131 }
2132}