Skip to main content

zeph_core/agent/
shadow_sentinel.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! `ShadowSentinel`: persistent safety memory stream + LLM-based pre-execution probe.
5//!
6//! Extends [`TrajectorySentinel`](crate::agent::trajectory) (Phase 1, spec 050) with:
7//!
8//! 1. **Persistent event stream**: `safety_shadow_events` table stores ALL safety-relevant
9//!    events across sessions (not limited to the last 8 turns like the in-memory sentinel).
10//! 2. **[`SafetyProbe`] trait**: before high-risk tool categories (shell, file write, exfil-
11//!    capable MCP tools), an LLM evaluates the full trajectory context and approves/denies.
12//!
13//! `ShadowSentinel` is **defence-in-depth only** — it is NOT the primary security gate.
14//! `PolicyGateExecutor` and `TrajectorySentinel` remain the primary enforcement mechanisms
15//! and continue to run regardless of probe results or timeouts.
16//!
17//! # Fail-open default
18//!
19//! `deny_on_timeout = false` (default) means a probe timeout or LLM error results in
20//! [`ProbeVerdict::Allow`]. This is correct because:
21//!
22//! - `ShadowSentinel` is defence-in-depth; policy gate still runs after it.
23//! - Failing closed on timeout would allow a `DoS`: slow context → every high-risk tool blocked.
24//! - Operators who want fail-closed can set `deny_on_timeout = true` in config.
25//!
26//! # LLM isolation invariant
27//!
28//! The probe prompt MUST NEVER include the `TrajectorySentinel` score or risk level.
29//! Exposing internal risk scores to the LLM would allow prompt injection attacks that
30//! manipulate probe verdicts by crafting tool outputs to lower the perceived risk level.
31
32use 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// ── Risk category ────────────────────────────────────────────────────────────
53
54/// Classifies a tool into a risk tier for probe gating.
55///
56/// `Shell`, `FileWrite`, `ExfilCapable`, and `McpUnclassified` tools trigger a safety probe.
57/// `Low` tools bypass the probe entirely, adding zero latency.
58///
59/// # Why `ExfilCapable` and `FileWrite` are distinct variants
60///
61/// Both are keyword-matched (write/edit/delete) and both trigger a probe, but they draw from
62/// different per-turn budgets in `check_tool_call` (#5749): `ExfilCapable` — a write-capable
63/// tool that also originates from an untrusted MCP server — has its own independent, higher
64/// budget (`2 * max_probes_per_turn`) so it can never be silently waved through
65/// (`ProbeVerdict::Skip`, which behaves like `Allow` under the fail-open default) purely because
66/// unrelated earlier probes in the same turn exhausted the shared counter. That budget is still
67/// finite — not unconditional — because `ExfilCapable` requires only MCP origin plus a keyword
68/// match, so an untrusted MCP server naming its tools accordingly could otherwise trigger
69/// unbounded LLM probe calls. `FileWrite` (builtin-only, no network egress path) has no separate
70/// budget and draws from the ordinary shared counter.
71#[derive(Debug, Clone, Copy, PartialEq, Eq)]
72#[non_exhaustive]
73pub enum ToolRiskCategory {
74    /// Shell execution — arbitrary commands, highest risk.
75    Shell,
76    /// File write or delete operations — persistent side effects.
77    FileWrite,
78    /// Network-capable MCP tools that could exfiltrate data.
79    ExfilCapable,
80    /// MCP-origin tool whose name matched no configured risk keyword pattern.
81    ///
82    /// Engaged unconditionally because MCP origin is itself an untrusted-provenance signal
83    /// (#5750): `probe_patterns` can never fully enumerate every risky verb (`remove`,
84    /// `rename`, `upload`, `spawn`, ...), so gating probe engagement entirely on keyword match
85    /// would silently skip the probe for any MCP tool using a verb outside the list. This tier
86    /// still triggers the probe, but at reduced budget priority relative to the keyword-matched
87    /// tiers above (see `check_tool_call`).
88    McpUnclassified,
89    /// All other tools — probe is skipped.
90    Low,
91}
92
93// ── Probe verdict ─────────────────────────────────────────────────────────────
94
95/// Result of a `SafetyProbe` evaluation.
96#[derive(Debug, Clone, PartialEq, Eq)]
97#[non_exhaustive]
98pub enum ProbeVerdict {
99    /// Tool execution is safe to proceed.
100    Allow,
101    /// Tool execution is denied. The `reason` is LLM-generated and returned to the
102    /// agent loop as the tool result so the model can adapt its strategy.
103    Deny {
104        /// Human-readable explanation from the safety probe.
105        reason: String,
106    },
107    /// Probe was skipped — tool is not in a high-risk category, feature is disabled,
108    /// or the per-turn probe budget was exhausted.
109    Skip,
110}
111
112// ── Sentinel event ───────────────────────────────────────────────────────────
113
114/// A single probe trajectory record in the persistent safety sentinel stream.
115///
116/// Stored in `safety_shadow_events` and retrieved for cross-session probe context.
117#[derive(Debug, Clone)]
118pub struct SentinelEvent {
119    /// Database row id (0 for unsaved records).
120    pub id: i64,
121    /// Agent session identifier.
122    pub session_id: SessionId,
123    /// Turn number within the session.
124    pub turn_number: u64,
125    /// Event category: `"tool_call"`, `"tool_result"`, `"risk_signal"`, `"probe_result"`.
126    pub event_type: String,
127    /// Fully-qualified tool id for tool events, `None` for non-tool events.
128    pub tool_id: Option<String>,
129    /// Serialised risk signal variant (from `TrajectorySentinel`), if applicable.
130    pub risk_signal: Option<String>,
131    /// Risk level at the time of the event: `"calm"`, `"elevated"`, `"high"`, `"critical"`.
132    pub risk_level: String,
133    /// Probe verdict for `probe_result` events: `"allow"`, `"deny"`, `"skip"`.
134    pub probe_verdict: Option<String>,
135    /// Short human-readable summary included in the LLM probe context.
136    pub context_summary: Option<String>,
137    /// Unix timestamp (seconds) when the event was recorded.
138    pub created_at: i64,
139}
140
141// ── SafetyProbe trait ─────────────────────────────────────────────────────────
142
143/// LLM-based pre-execution safety evaluator.
144///
145/// Implementors receive the full trajectory context and the proposed tool call
146/// and return a [`ProbeVerdict`]. The probe runs BEFORE [`zeph_tools::PolicyGateExecutor`].
147///
148/// # Contract
149///
150/// - Probe timeout is mandatory (configured via `probe_timeout_ms`).
151/// - Probe failure (LLM error, timeout when `deny_on_timeout = false`) results in `Allow`.
152/// - Probe results are persisted to `safety_shadow_events` for cross-session learning.
153/// - The probe prompt MUST NOT include the sentinel score or risk level (LLM isolation).
154///
155/// Uses `Pin<Box<dyn Future>>` returns for dyn-compatibility (stored as `Box<dyn SafetyProbe>`).
156pub trait SafetyProbe: Send + Sync {
157    /// Evaluate whether the proposed tool call is safe given the trajectory context.
158    ///
159    /// # Arguments
160    ///
161    /// * `tool_id` — fully-qualified tool identifier (e.g. `"builtin:shell"`).
162    /// * `tool_args` — JSON arguments for the tool call.
163    /// * `trajectory` — recent shadow events for context (last N events from the store).
164    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
172// ── LlmSafetyProbe ───────────────────────────────────────────────────────────
173
174/// LLM-backed implementation of `SafetyProbe`.
175///
176/// Sends a structured prompt to an LLM asking it to evaluate whether the proposed
177/// tool call is safe given the trajectory. The prompt is intentionally minimal and
178/// never includes internal risk scores (LLM isolation invariant).
179pub struct LlmSafetyProbe {
180    provider: Arc<AnyProvider>,
181    timeout_ms: u64,
182    deny_on_timeout: bool,
183}
184
185impl LlmSafetyProbe {
186    /// Create a new `LlmSafetyProbe`.
187    ///
188    /// # Arguments
189    ///
190    /// * `provider` — LLM provider used for probe evaluation.
191    /// * `timeout_ms` — maximum milliseconds to wait for the LLM response.
192    /// * `deny_on_timeout` — when `true`, a timeout results in `Deny`; when `false`, `Allow`.
193    #[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        // Try to extract JSON from the response.
248        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        // Unparseable response → allow (fail-open)
267        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// ── ShadowEventStore ─────────────────────────────────────────────────────────
321
322/// Persistent storage for the safety shadow event stream.
323///
324/// Thin wrapper around [`DbPool`] for the `safety_shadow_events` table.
325/// Methods are `async` and return typed errors.
326#[derive(Clone)]
327pub struct ShadowEventStore {
328    pool: DbPool,
329}
330
331impl ShadowEventStore {
332    /// Create a `ShadowEventStore` backed by the given pool.
333    #[must_use]
334    pub fn new(pool: DbPool) -> Self {
335        Self { pool }
336    }
337
338    /// Persist a shadow event to the database.
339    ///
340    /// The `id` field of the event is ignored; the database assigns a new row id.
341    ///
342    /// # Errors
343    ///
344    /// Returns `AgentError` on database failure.
345    #[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    /// Retrieve the last `limit` events for a session in ascending time order.
370    ///
371    /// Used to build the trajectory context for probe evaluation.
372    ///
373    /// # Errors
374    ///
375    /// Returns `AgentError` on database failure.
376    #[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        // DB returns DESC (newest first); reverse once to get ASC (oldest first) for LLM context.
397        let mut events: Vec<SentinelEvent> = rows.into_iter().map(SentinelEvent::from).collect();
398        events.reverse();
399        Ok(events)
400    }
401
402    /// Retrieve the last `limit` events for a specific tool from sessions OTHER than
403    /// `exclude_session_id`.
404    ///
405    /// Used for cross-session pattern detection. The exclusion is applied in SQL (not just
406    /// filtered client-side afterward) so that a session with heavy recent activity for
407    /// `tool_id` cannot crowd its own rows into the `LIMIT` clip and starve genuinely
408    /// cross-session rows out of the result.
409    ///
410    /// # Errors
411    ///
412    /// Returns `AgentError` on database failure.
413    #[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// Internal sqlx row type for `safety_shadow_events`.
440#[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
471// ── ShadowSentinel ────────────────────────────────────────────────────────────
472
473/// Maximum number of concurrent fire-and-forget persist tasks tracked in `pending_writes`.
474///
475/// When the set is at capacity the oldest completed tasks are reaped before spawning a new one.
476/// If the set is still full after reaping (all tasks are still running), the new spawn is skipped
477/// with a debug log — persistence is best-effort and the sentinel must never block tool dispatch.
478const MAX_PENDING_WRITES: usize = 32;
479
480/// Orchestrates the persistent safety stream and LLM pre-execution probe.
481///
482/// `ShadowSentinel` is wrapped in `Arc` and shared between `ShadowProbeExecutor` instances
483/// when tools run in parallel. All mutable state uses `AtomicU32` to allow `&self` access
484/// from concurrent tool dispatch without a `Mutex`.
485///
486/// # Turn lifecycle
487///
488/// - `advance_turn()` — call once per turn before tool execution; resets the per-turn
489///   probe counter.
490/// - `check_tool_call()` — call before each tool execution to probe high-risk calls.
491/// - `record_tool_event()` — call after tool execution to persist the event.
492/// - `drain_pending()` — call at session shutdown to await all queued persist writes.
493///
494/// # NEVER
495///
496/// Never expose the `ShadowSentinel` state or probe verdicts to LLM-visible context.
497pub struct ShadowSentinel {
498    store: ShadowEventStore,
499    probe: Box<dyn SafetyProbe>,
500    config: zeph_config::ShadowSentinelConfig,
501    /// Counter of `Shell`/`FileWrite`/`McpUnclassified` probe calls made in the current turn.
502    /// Uses `AtomicU32` so all probe-checking methods can take `&self` even under parallel tool
503    /// execution.
504    ///
505    /// `ToolRiskCategory::McpUnclassified` calls are capped at `max_probes_per_turn - 1` so at
506    /// least one slot always stays reserved for `Shell`/`FileWrite` (#5750).
507    /// `ToolRiskCategory::ExfilCapable` never touches this counter — it has its own independent
508    /// budget, see `exfil_probes_this_turn` (#5749).
509    probes_this_turn: AtomicU32,
510    /// Independent per-turn counter for `ToolRiskCategory::ExfilCapable` calls (#5749).
511    ///
512    /// `ExfilCapable` is the highest-confidence risk signal (MCP-origin AND write-capable), so
513    /// it must not compete with — or be starved by — the shared `probes_this_turn` budget. But
514    /// giving it *unconditional* exemption would let a false-positive keyword match on an
515    /// MCP-origin tool (#5750's over-inclusion case) generate unbounded LLM probe calls with no
516    /// cost ceiling at all. This counter gives `ExfilCapable` its own finite cap
517    /// (`2 * max_probes_per_turn`, see `probe_budget_exhausted`) — high priority, but still
518    /// bounded.
519    exfil_probes_this_turn: AtomicU32,
520    session_id: SessionId,
521    /// Bounded set of fire-and-forget DB persist tasks. Prevents unbounded task accumulation
522    /// and ensures panics surface at `drain_pending()` instead of being silently swallowed.
523    pending_writes: Mutex<JoinSet<()>>,
524    /// Sanitized ids (`ToolDef::server_id`-backed) of tools registered by MCP servers.
525    ///
526    /// Mirrors `TrustGateExecutor::mcp_tool_ids` (`zeph_tools::TrustGateExecutor`): empty until
527    /// populated post-construction via [`mcp_tool_ids_handle`](Self::mcp_tool_ids_handle). This
528    /// is the authoritative way to know a `qualified_tool_id` originates from an MCP server —
529    /// real ids are `{server_id}_{name}` (`McpTool::sanitized_id`) and carry no reliable string
530    /// prefix to pattern-match on (#5736).
531    ///
532    /// # Refresh
533    ///
534    /// Populated at startup from the initial `mcp_tools` list (`src/runner.rs`) and refreshed
535    /// on every subsequent tool-list change — `/mcp add`/`/mcp remove` and a live
536    /// `tools/list_changed` notification both route through
537    /// `Agent::refresh_shadow_sentinel_mcp_tool_ids` (`crates/zeph-core/src/agent/mcp.rs`,
538    /// called from `check_tool_refresh` once per turn) — so a server connected mid-session is
539    /// reflected without a process restart.
540    ///
541    /// **Known gap**: `TrustGateExecutor`'s own equivalent set has no such refresh path (it
542    /// lives entirely in the binary crate's tool-executor chain, unreachable from `Agent`) —
543    /// tracked separately (#5747), not fixed by this refresh.
544    mcp_tool_ids: Arc<RwLock<HashSet<String>>>,
545}
546
547impl ShadowSentinel {
548    /// Create a new `ShadowSentinel`.
549    ///
550    /// # Arguments
551    ///
552    /// * `store` — persistent shadow event store.
553    /// * `probe` — safety probe implementation.
554    /// * `config` — subsystem configuration.
555    /// * `session_id` — current agent session identifier.
556    #[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    /// Returns the shared MCP tool-id set so the caller can populate it once MCP servers have
576    /// connected (mirrors `TrustGateExecutor::mcp_tool_ids_handle`).
577    #[must_use]
578    pub fn mcp_tool_ids_handle(&self) -> Arc<RwLock<HashSet<String>>> {
579        Arc::clone(&self.mcp_tool_ids)
580    }
581
582    /// Returns `true` when `tool_id` was registered by an MCP server.
583    fn is_mcp_tool(&self, tool_id: &str) -> bool {
584        self.mcp_tool_ids.read().contains(tool_id)
585    }
586
587    /// Classify a fully-qualified tool id into a risk tier.
588    ///
589    /// Pattern matching is prefix/glob-based against the configured `probe_patterns`.
590    /// For efficiency, we check common built-in names first before falling back to
591    /// glob matching against the configured patterns.
592    #[must_use]
593    pub fn classify_tool(&self, qualified_tool_id: &str) -> ToolRiskCategory {
594        // Fast-path for well-known high-risk builtins.
595        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        // Glob matching against configured patterns.
615        for pattern in &self.config.probe_patterns {
616            if glob_matches(pattern, qualified_tool_id) {
617                // Classify based on the pattern name.
618                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        // #5750: an MCP-origin tool is inherently less trusted, even when its name doesn't
636        // match any configured keyword pattern — engage a lightweight probe unconditionally
637        // rather than gating engagement entirely on keyword match.
638        if self.is_mcp_tool(qualified_tool_id) {
639            return ToolRiskCategory::McpUnclassified;
640        }
641
642        ToolRiskCategory::Low
643    }
644
645    /// Checks and consumes per-turn probe budget for `category`, returning `true` when the
646    /// budget is exhausted (the caller must skip the probe).
647    ///
648    /// `ExfilCapable` (#5749) draws from its own independent, finite budget
649    /// (`2 * max_probes_per_turn`, `exfil_probes_this_turn`) rather than the shared counter: it
650    /// is the highest-confidence risk signal (MCP-origin AND write-capable), so it must not be
651    /// starved by unrelated earlier probes in the same turn — but it must still be bounded, or a
652    /// false-positive keyword match on an MCP-origin tool (#5750's over-inclusion case) could
653    /// generate unbounded LLM probe calls with no cost ceiling.
654    ///
655    /// `McpUnclassified` (#5750) — engaged by MCP origin alone, not a keyword risk signal — is
656    /// capped at `max_probes_per_turn - 1` (saturating), always reserving at least one slot in
657    /// the shared counter for `Shell`/`FileWrite` so a burst of low-signal MCP engagement early
658    /// in the turn can never fully starve out a higher-confidence probe later in the same turn,
659    /// at any `max_probes_per_turn` value (including `0` or `1`).
660    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        // Check per-turn probe budget using relaxed atomics (false sharing is acceptable here).
678        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            // Undo the increment so future fast-path checks are accurate.
687            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    /// Load the trajectory + cross-session tool history used as probe context.
699    ///
700    /// Filters out `probe_result` events — exposing probe verdicts to the LLM would allow
701    /// prompt injection attacks that craft tool outputs to manipulate perceived safety.
702    ///
703    /// Each DB read is independently bounded by `probe_timeout_ms.min(2000)`: a stalled DB
704    /// connection must never block dispatch of every high-risk tool call for the session. A
705    /// timeout or DB error falls back to an empty/partial result (fail-open), matching the
706    /// probe's own fail-open default.
707    ///
708    /// The two reads run sequentially, each with its own independent timeout budget, and the
709    /// LLM probe call in [`check_tool_call`](Self::check_tool_call) has its own separate,
710    /// uncapped `probe_timeout_ms` timeout on top — worst-case `check_tool_call` latency is
711    /// therefore additive across all three: `2 * probe_timeout_ms.min(2000) + probe_timeout_ms`
712    /// (~6s at the 2000ms default), not a single shared ~2s bound.
713    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        // Reserve half the total budget for cross-session history so recurring risk patterns
742        // from other sessions always have visibility — even in the busiest sessions, where the
743        // session's own trajectory alone would otherwise fill (and, pre-fix, silently evict
744        // the entire cross-session block from) the whole budget. Enforce the session-side cap
745        // here (trajectory is oldest-first/ASC, so excess is trimmed from the front, keeping
746        // the most recent events).
747        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        // Load cross-session history for this tool so recurring risk patterns from
755        // other sessions inform the probe, not just the current session (#5449). The
756        // current session is excluded in SQL (not just filtered client-side) so its own
757        // activity can never crowd genuinely cross-session rows out of the LIMIT clip.
758        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                // get_tool_history is DESC (newest first); reverse to ASC to match
770                // trajectory ordering, then prepend so trajectory stays oldest-first.
771                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    /// Evaluate a proposed tool call and return a probe verdict.
797    ///
798    /// Returns `ProbeVerdict::Skip` when:
799    /// - The tool is not in a high-risk category.
800    /// - The feature is disabled.
801    /// - The per-turn probe budget (`max_probes_per_turn`) is exhausted.
802    ///
803    /// `ToolRiskCategory::ExfilCapable` calls (#5749) draw from their own independent, higher
804    /// budget (`2 * max_probes_per_turn`) instead of the shared counter, so unrelated earlier
805    /// probes in the same turn can never wave one through — but the budget is still finite, not
806    /// unconditional. `ToolRiskCategory::McpUnclassified` calls (#5750) get a reduced share of
807    /// the shared budget — at least one slot is always reserved for keyword-matched
808    /// (higher-confidence) categories so a burst of low-signal MCP engagement cannot starve them
809    /// out within the same turn, at any `max_probes_per_turn` value.
810    ///
811    /// This method takes `&self` so it can be called from parallel tool dispatch.
812    ///
813    /// # Errors
814    ///
815    /// Does not return errors; probe failures are handled internally (fail-open or
816    /// fail-closed depending on `deny_on_timeout`).
817    #[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        // Persist the probe result asynchronously (best-effort — never blocks tool path).
846        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    /// Persist a tool execution event in the shadow stream (fire-and-forget).
876    ///
877    /// Called after a tool finishes execution to maintain the trajectory for future probes.
878    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    /// Await all queued fire-and-forget persist tasks.
904    ///
905    /// Call once at session shutdown to ensure no DB writes are silently dropped.
906    /// All errors have already been logged inside each task; this method only joins the handles.
907    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    /// Clones the store handle and spawns a fire-and-forget persist of `event` via
916    /// [`Self::spawn_persist`], logging `warn_context` on failure. Shared by
917    /// [`Self::check_tool_call`] (probe results) and [`Self::record_tool_event`]
918    /// (tool-call events) — the two call sites differ only in the event they persist
919    /// and the wording of the warn-log context.
920    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    /// Spawn a background persist task into the bounded `JoinSet`.
931    ///
932    /// Reaps completed handles before spawning to stay within `MAX_PENDING_WRITES`. If the set
933    /// is still at capacity after reaping (all tasks still running), the new task is dropped and
934    /// a debug message is emitted — persistence is best-effort and must never block the tool path.
935    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        // Reap only already-finished handles — never block waiting for a running task.
941        // try_join_next() returns immediately if no task has completed yet.
942        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    /// Reset the per-turn probe counters.
954    ///
955    /// Must be called once per turn BEFORE any tool calls, alongside
956    /// `TrajectorySentinel::advance_turn()`.
957    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
963// ── Helpers ───────────────────────────────────────────────────────────────────
964
965/// Returns the current Unix timestamp in seconds.
966fn 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
974/// Simple glob matching: `*` matches any sequence of characters except `/`.
975/// `*/` in the pattern matches any single path segment.
976fn glob_matches(pattern: &str, value: &str) -> bool {
977    if pattern == "*" {
978        return true;
979    }
980    // Split on `*` and check each segment is present in order.
981    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// ── AgentError extension ──────────────────────────────────────────────────────
1007// ShadowEventStore uses AgentError::Db — add that variant if missing.
1008// (The actual variant is declared in agent/error.rs; we only reference it here.)
1009
1010#[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    /// #5750: an MCP-origin tool whose name matches no configured keyword pattern must still
1057    /// be probed (as `McpUnclassified`), not silently fall through to `Low`. Verbs like
1058    /// `frobnicate` can never be fully enumerated in `probe_patterns`.
1059    #[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    /// The same non-keyword-matching name for a tool NOT registered as MCP-origin must remain
1074    /// `Low` — engagement is gated on MCP origin, not merely on failing to match Low's fast path.
1075    #[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    /// #5736 regression: MCP-tool escalation must key off the registered `mcp_tool_ids` set
1107    /// (`ToolDef::server_id`-backed), not a `"mcp:"` string prefix — real MCP tool ids are
1108    /// `"{server_id}_{name}"` (`McpTool::sanitized_id`) and never carry that prefix, so the old
1109    /// check silently never escalated any MCP write/edit tool to `ExfilCapable`.
1110    #[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        // Unregistered: a same-shaped id falls to the ordinary FileWrite tier, not ExfilCapable.
1118        assert_eq!(
1119            sentinel.classify_tool("github_edit_file"),
1120            ToolRiskCategory::FileWrite
1121        );
1122        // Register it the same way the real MCP wiring does (via the shared handle) and the
1123        // identical id must now escalate.
1124        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    /// #5736 follow-up (CI-1239): `ShadowSentinelConfig::default()` — not a hand-tuned override —
1135    /// must escalate a real MCP write tool. Real MCP tool ids are `"{server_id}_{name}"`
1136    /// (`McpTool::sanitized_id`), e.g. `"fs-test_write_file"`; the shipped default
1137    /// `probe_patterns` (`"mcp:*/file_*"`, `"mcp:*/exec_*"`) assumed a `"mcp:"`-prefixed id
1138    /// shape that no real id ever has, so the outer glob-matching loop in `classify_tool` never
1139    /// even entered the branch containing the `is_mcp_tool()` check — every MCP tool silently
1140    /// fell through to `ToolRiskCategory::Low` (probe skipped entirely), a complete bypass, not
1141    /// just a downgrade to `FileWrite`.
1142    #[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        // First two calls should not be skipped (noop probe returns Allow).
1212        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        // Third call exceeds max_probes_per_turn = 2 → must skip.
1223        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    /// #5749: `ExfilCapable` calls must never be skipped due to *shared* per-turn budget
1234    /// exhaustion — they draw from their own independent counter. Exhaust the shared budget with
1235    /// `Shell` calls first, then confirm a subsequent `ExfilCapable` call still probes.
1236    #[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        // Exhaust the shared budget (max_probes_per_turn = 1) with a Shell call.
1257        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        // ExfilCapable call must still probe despite the exhausted shared budget.
1271        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    /// #5749 follow-up (critic SIGNIFICANT-1): `ExfilCapable`'s independent budget must still be
1282    /// finite (`2 * max_probes_per_turn`), not unconditional — otherwise a false-positive
1283    /// keyword match on an MCP-origin tool name generates unbounded LLM probe calls.
1284    #[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        // exfil_max = 2 * max_probes_per_turn = 2 — first two calls must probe.
1301        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        // Third call exceeds the independent cap → must skip, not run unbounded.
1319        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    /// #5750: `McpUnclassified` calls (engaged by MCP origin alone) must be capped at
1330    /// `max_probes_per_turn - 1`, reserving at least one slot for keyword-matched categories so
1331    /// a burst of low-signal MCP probes cannot starve out a later high-confidence probe.
1332    #[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        // First McpUnclassified call consumes the one slot it's allowed (max - 1 = 1).
1352        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        // Second McpUnclassified call must be skipped — reserved share (1) already used, even
1362        // though the shared counter (1/2) has not reached max_probes_per_turn.
1363        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        // A Shell call must still get through using the slot reserved for it.
1373        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    /// #5750 follow-up (critic SIGNIFICANT-2): at `max_probes_per_turn == 1` there is only one
1384    /// slot total, so reserving it for `Shell`/`FileWrite` means `McpUnclassified` gets NO share
1385    /// at all (`saturating_sub` floors at 0) rather than the whole budget. This is what makes
1386    /// the "cannot starve a higher-confidence probe" guarantee hold unconditionally, at the cost
1387    /// of never probing `McpUnclassified` when the turn's total budget is 1.
1388    #[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        // Even the FIRST McpUnclassified call must be skipped — the single slot is fully
1404        // reserved for keyword-matched categories, never handed to McpUnclassified.
1405        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        // A subsequent Shell call must still probe using the untouched slot — proving the
1415        // McpUnclassified attempt above did not consume or starve it.
1416        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    /// Boundary: `max_probes_per_turn == 0` must skip every category through the shared budget
1427    /// (`Shell`, `FileWrite`, `McpUnclassified`) AND the independent `ExfilCapable` budget
1428    /// (`2 * 0 == 0`) — an operator disabling the probe budget entirely must not leave any
1429    /// category unbounded.
1430    #[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    // ── JoinSet regression tests (#4570) ─────────────────────────────────────
1499
1500    /// `drain_pending` awaits all spawned persist tasks and returns when the set is empty.
1501    #[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    /// When the pending set is at capacity and all running tasks complete before the next
1529    /// `spawn_persist`, the new task IS accepted (the set has room after reaping).
1530    /// Conversely, if we fill the set, drain it, then overfill past capacity while tasks are
1531    /// still running — the implementation drops extras.  We verify the simpler property:
1532    /// `spawn_persist` never panics when called repeatedly beyond `MAX_PENDING_WRITES`.
1533    #[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        // Spawn twice the capacity; each task completes instantly.
1542        // spawn_persist will reap completed tasks between spawns, so most will be accepted.
1543        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        // All tasks (or at least MAX_PENDING_WRITES of them) must have run; none panicked.
1555        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    // Build a minimal ShadowSentinel with a no-op probe for unit tests.
1563    //
1564    // Opens an in-memory SQLite pool. Store methods are never called in these unit
1565    // tests — they test only classification and counter logic.
1566    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    // Opens a migrated in-memory SQLite pool (unlike `make_test_sentinel`'s pool, this one
1585    // has the `safety_shadow_events` table from migration 085 and can serve real store queries.
1586    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    /// #5449 regression: `check_tool_call` must fold cross-session `get_tool_history` results
1680    /// into the trajectory passed to the probe, not just the current session's own events.
1681    #[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    /// Drives `check_tool_call` for `tool_id` under `session_id` and returns the exact
1744    /// trajectory the probe received, so cap tests can assert WHICH events survive, not
1745    /// just how many — a count-only assertion can pass while the cap silently drops all
1746    /// cross-session data (the bug found in code review of the initial #5449 fix).
1747    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    /// Seeds `count` events for `session_id`/`tool_id`, with ascending `created_at`
1788    /// timestamps starting at `base`, so cap tests can control which events are "most
1789    /// recent". Summaries are `"{summary_prefix}-{i}"` for index-based assertions.
1790    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    /// Session trajectory and cross-session history are each independently capped at
1811    /// `max_context_events`, so a naive merge can total up to 2x the configured budget.
1812    /// `check_tool_call` must enforce the combined cap AND reserve budget for cross-session
1813    /// data — the original fix trimmed unconditionally from the front, which silently wiped
1814    /// ALL cross-session events whenever the session's own trajectory alone filled the
1815    /// budget (precisely the busiest-session scenario #5449 cares about most).
1816    #[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    /// Mirror case: cross-session history is the one over budget, session's own trajectory
1858    /// is light. The session-side cap must not trim events that don't need trimming.
1859    #[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    /// Boundary: session + cross-session totals exactly `max_context_events` — nothing
1899    /// should be dropped from either side.
1900    #[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    /// Boundary: one more cross-session event than the reserved budget — exactly one event
1932    /// must be dropped, and it must be the OLDEST cross-session event (trajectory stays
1933    /// oldest-first/ASC, so the most recent events are kept).
1934    #[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    /// The current session's own events must not be double-counted into the cross-session
1977    /// block — `get_tool_history` excludes `exclude_session_id` directly in its SQL
1978    /// (`AND session_id != ?`), and this test confirms that exclusion end-to-end through
1979    /// `check_tool_call`.
1980    #[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    /// `probe_result` events from OTHER sessions must never leak into the cross-session
2004    /// merge — `get_tool_history`'s SQL does not filter by `event_type`, so this relies
2005    /// entirely on the Rust-side filter (the same LLM-isolation invariant already tested
2006    /// for the same-session trajectory, exercised here on the cross-session path).
2007    #[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    // ── #6269: DB-read timeout fail-open ─────────────────────────────────────
2036
2037    /// #6269 regression: both DB reads inside `load_probe_context` (`get_trajectory` and
2038    /// `get_tool_history`, reached via `check_tool_call`) must fail open when the DB pool
2039    /// stalls, exactly like their existing `Err` (DB-error) branches and the LLM probe's own
2040    /// timeout branch. A real stall is forced — not a synthetic sleep race — by exhausting
2041    /// the in-memory `SQLite` pool's sole connection: `test_pool()` connects with `":memory:"`,
2042    /// which `crates/zeph-db/src/pool.rs` hard-caps at `max_connections(1)`, so holding one
2043    /// `BEGIN IMMEDIATE` transaction open blocks both `fetch_all(&pool)` calls on
2044    /// `pool.acquire()` until `probe_timeout_ms` elapses.
2045    #[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        // Seed real rows so an empty captured trajectory can only be explained by the
2054        // timeout fallback below, not by the store genuinely having nothing to return.
2055        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        // Hold the sole in-memory SQLite connection so both `fetch_all` calls inside
2074        // `load_probe_context` block on `pool.acquire()` for the full 50ms probe timeout.
2075        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    // ── #5766: record_tool_event had zero test coverage ─────────────────────────
2107
2108    #[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    /// Minimal `tracing_subscriber::Layer` that captures event messages into a shared buffer,
2159    /// used to verify `record_tool_event`'s fire-and-forget persist failure logs the correct
2160    /// warn context ("failed to persist tool event", as opposed to `check_tool_call`'s "failed
2161    /// to persist probe result").
2162    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    // `record()` fails when the backing table is gone. Self-loop-style DB-trigger tampering
2189    // isn't needed here — a real store error is the whole point of this test.
2190    #[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}