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    /// Evaluate a proposed tool call and return a probe verdict.
699    ///
700    /// Returns `ProbeVerdict::Skip` when:
701    /// - The tool is not in a high-risk category.
702    /// - The feature is disabled.
703    /// - The per-turn probe budget (`max_probes_per_turn`) is exhausted.
704    ///
705    /// `ToolRiskCategory::ExfilCapable` calls (#5749) draw from their own independent, higher
706    /// budget (`2 * max_probes_per_turn`) instead of the shared counter, so unrelated earlier
707    /// probes in the same turn can never wave one through — but the budget is still finite, not
708    /// unconditional. `ToolRiskCategory::McpUnclassified` calls (#5750) get a reduced share of
709    /// the shared budget — at least one slot is always reserved for keyword-matched
710    /// (higher-confidence) categories so a burst of low-signal MCP engagement cannot starve them
711    /// out within the same turn, at any `max_probes_per_turn` value.
712    ///
713    /// This method takes `&self` so it can be called from parallel tool dispatch.
714    ///
715    /// # Errors
716    ///
717    /// Does not return errors; probe failures are handled internally (fail-open or
718    /// fail-closed depending on `deny_on_timeout`).
719    #[tracing::instrument(name = "security.shadow.check", skip(self, tool_args), fields(tool_id = %qualified_tool_id))]
720    pub async fn check_tool_call(
721        &self,
722        qualified_tool_id: &str,
723        tool_args: &JsonValue,
724        turn_number: u64,
725        current_risk_level: &str,
726    ) -> ProbeVerdict {
727        if !self.config.enabled {
728            return ProbeVerdict::Skip;
729        }
730
731        let category = self.classify_tool(qualified_tool_id);
732        if category == ToolRiskCategory::Low {
733            return ProbeVerdict::Skip;
734        }
735
736        if self.probe_budget_exhausted(category) {
737            return ProbeVerdict::Skip;
738        }
739
740        // Load recent trajectory for probe context.
741        // Filter out probe_result events — exposing probe verdicts to the LLM would allow
742        // prompt injection attacks that craft tool outputs to manipulate perceived safety.
743        let mut trajectory: Vec<SentinelEvent> = match self
744            .store
745            .get_trajectory(&self.session_id, self.config.max_context_events)
746            .await
747        {
748            Ok(t) => t
749                .into_iter()
750                .filter(|e| e.event_type != "probe_result")
751                .collect(),
752            Err(e) => {
753                tracing::warn!(error = %e, "ShadowSentinel: failed to load trajectory, proceeding without context");
754                vec![]
755            }
756        };
757
758        // Reserve half the total budget for cross-session history so recurring risk patterns
759        // from other sessions always have visibility — even in the busiest sessions, where the
760        // session's own trajectory alone would otherwise fill (and, pre-fix, silently evict
761        // the entire cross-session block from) the whole budget. Enforce the session-side cap
762        // here (trajectory is oldest-first/ASC, so excess is trimmed from the front, keeping
763        // the most recent events).
764        let cross_session_budget = self.config.max_context_events / 2;
765        let session_budget = self.config.max_context_events - cross_session_budget;
766        if trajectory.len() > session_budget {
767            let excess = trajectory.len() - session_budget;
768            trajectory.drain(0..excess);
769        }
770
771        // Load cross-session history for this tool so recurring risk patterns from
772        // other sessions inform the probe, not just the current session (#5449). The
773        // current session is excluded in SQL (not just filtered client-side) so its own
774        // activity can never crowd genuinely cross-session rows out of the LIMIT clip.
775        match self
776            .store
777            .get_tool_history(
778                qualified_tool_id,
779                self.session_id.as_str(),
780                self.config.max_context_events,
781            )
782            .await
783        {
784            Ok(history) => {
785                // get_tool_history is DESC (newest first); reverse to ASC to match
786                // trajectory ordering, then prepend so trajectory stays oldest-first.
787                let mut cross_session: Vec<SentinelEvent> = history
788                    .into_iter()
789                    .filter(|e| e.event_type != "probe_result")
790                    .rev()
791                    .collect();
792                if cross_session.len() > cross_session_budget {
793                    let excess = cross_session.len() - cross_session_budget;
794                    cross_session.drain(0..excess);
795                }
796                trajectory.splice(0..0, cross_session);
797            }
798            Err(e) => {
799                tracing::warn!(error = %e, "ShadowSentinel: failed to load cross-session tool history, proceeding without it");
800            }
801        }
802
803        let verdict = self
804            .probe
805            .evaluate(qualified_tool_id, tool_args, &trajectory)
806            .await;
807
808        // Persist the probe result asynchronously (best-effort — never blocks tool path).
809        let probe_verdict_str = match &verdict {
810            ProbeVerdict::Allow => "allow",
811            ProbeVerdict::Deny { .. } => "deny",
812            ProbeVerdict::Skip => "skip",
813        };
814        let summary = match &verdict {
815            ProbeVerdict::Deny { reason } => {
816                format!("probe denied: {}", &reason[..reason.len().min(120)])
817            }
818            ProbeVerdict::Allow => format!("probe allowed {qualified_tool_id}"),
819            ProbeVerdict::Skip => format!("probe skipped {qualified_tool_id}"),
820        };
821        let event = SentinelEvent {
822            id: 0,
823            session_id: self.session_id.clone(),
824            turn_number,
825            event_type: "probe_result".to_owned(),
826            tool_id: Some(qualified_tool_id.to_owned()),
827            risk_signal: None,
828            risk_level: current_risk_level.to_owned(),
829            probe_verdict: Some(probe_verdict_str.to_owned()),
830            context_summary: Some(summary),
831            created_at: unix_now(),
832        };
833        self.persist_event(event, "probe result").await;
834
835        verdict
836    }
837
838    /// Persist a tool execution event in the shadow stream (fire-and-forget).
839    ///
840    /// Called after a tool finishes execution to maintain the trajectory for future probes.
841    pub async fn record_tool_event(
842        &self,
843        qualified_tool_id: &str,
844        turn_number: u64,
845        risk_level: &str,
846        context_summary: &str,
847    ) {
848        if !self.config.enabled {
849            return;
850        }
851        let event = SentinelEvent {
852            id: 0,
853            session_id: self.session_id.clone(),
854            turn_number,
855            event_type: "tool_call".to_owned(),
856            tool_id: Some(qualified_tool_id.to_owned()),
857            risk_signal: None,
858            risk_level: risk_level.to_owned(),
859            probe_verdict: None,
860            context_summary: Some(context_summary.chars().take(250).collect()),
861            created_at: unix_now(),
862        };
863        self.persist_event(event, "tool event").await;
864    }
865
866    /// Await all queued fire-and-forget persist tasks.
867    ///
868    /// Call once at session shutdown to ensure no DB writes are silently dropped.
869    /// All errors have already been logged inside each task; this method only joins the handles.
870    pub async fn drain_pending(&self) {
871        let mut set = {
872            let mut guard = self.pending_writes.lock().await;
873            std::mem::take(&mut *guard)
874        };
875        while set.join_next().await.is_some() {}
876    }
877
878    /// Clones the store handle and spawns a fire-and-forget persist of `event` via
879    /// [`Self::spawn_persist`], logging `warn_context` on failure. Shared by
880    /// [`Self::check_tool_call`] (probe results) and [`Self::record_tool_event`]
881    /// (tool-call events) — the two call sites differ only in the event they persist
882    /// and the wording of the warn-log context.
883    async fn persist_event(&self, event: SentinelEvent, warn_context: &'static str) {
884        let store = self.store.clone();
885        self.spawn_persist(async move {
886            if let Err(e) = store.record(&event).await {
887                tracing::warn!(error = %e, "ShadowSentinel: failed to persist {warn_context}");
888            }
889        })
890        .await;
891    }
892
893    /// Spawn a background persist task into the bounded `JoinSet`.
894    ///
895    /// Reaps completed handles before spawning to stay within `MAX_PENDING_WRITES`. If the set
896    /// is still at capacity after reaping (all tasks still running), the new task is dropped and
897    /// a debug message is emitted — persistence is best-effort and must never block the tool path.
898    async fn spawn_persist<F>(&self, fut: F)
899    where
900        F: std::future::Future<Output = ()> + Send + 'static,
901    {
902        let mut set = self.pending_writes.lock().await;
903        // Reap only already-finished handles — never block waiting for a running task.
904        // try_join_next() returns immediately if no task has completed yet.
905        while set.try_join_next().is_some() {}
906        if set.len() < MAX_PENDING_WRITES {
907            set.spawn(fut);
908        } else {
909            tracing::debug!(
910                max = MAX_PENDING_WRITES,
911                "ShadowSentinel: pending_writes at capacity, skipping persist"
912            );
913        }
914    }
915
916    /// Reset the per-turn probe counters.
917    ///
918    /// Must be called once per turn BEFORE any tool calls, alongside
919    /// `TrajectorySentinel::advance_turn()`.
920    pub fn advance_turn(&self) {
921        self.probes_this_turn.store(0, Ordering::Release);
922        self.exfil_probes_this_turn.store(0, Ordering::Release);
923    }
924}
925
926// ── Helpers ───────────────────────────────────────────────────────────────────
927
928/// Returns the current Unix timestamp in seconds.
929fn unix_now() -> i64 {
930    std::time::SystemTime::now()
931        .duration_since(std::time::UNIX_EPOCH)
932        .ok()
933        .and_then(|d| i64::try_from(d.as_secs()).ok())
934        .unwrap_or(0)
935}
936
937/// Simple glob matching: `*` matches any sequence of characters except `/`.
938/// `*/` in the pattern matches any single path segment.
939fn glob_matches(pattern: &str, value: &str) -> bool {
940    if pattern == "*" {
941        return true;
942    }
943    // Split on `*` and check each segment is present in order.
944    let parts: Vec<&str> = pattern.split('*').collect();
945    if parts.len() == 1 {
946        return pattern == value;
947    }
948    let mut remaining = value;
949    for (i, part) in parts.iter().enumerate() {
950        if part.is_empty() {
951            continue;
952        }
953        if i == 0 {
954            if !remaining.starts_with(part) {
955                return false;
956            }
957            remaining = &remaining[part.len()..];
958        } else if i == parts.len() - 1 {
959            return remaining.ends_with(part);
960        } else if let Some(pos) = remaining.find(part) {
961            remaining = &remaining[pos + part.len()..];
962        } else {
963            return false;
964        }
965    }
966    true
967}
968
969// ── AgentError extension ──────────────────────────────────────────────────────
970// ShadowEventStore uses AgentError::Db — add that variant if missing.
971// (The actual variant is declared in agent/error.rs; we only reference it here.)
972
973#[cfg(test)]
974mod tests {
975    use super::*;
976
977    #[tokio::test]
978    async fn classify_builtin_shell_is_shell_risk() {
979        let config = zeph_config::ShadowSentinelConfig::default();
980        let sentinel = make_test_sentinel(config).await;
981        assert_eq!(
982            sentinel.classify_tool("builtin:shell"),
983            ToolRiskCategory::Shell
984        );
985        assert_eq!(
986            sentinel.classify_tool("builtin:bash"),
987            ToolRiskCategory::Shell
988        );
989    }
990
991    #[tokio::test]
992    async fn classify_builtin_write_is_file_write_risk() {
993        let config = zeph_config::ShadowSentinelConfig::default();
994        let sentinel = make_test_sentinel(config).await;
995        assert_eq!(
996            sentinel.classify_tool("builtin:write"),
997            ToolRiskCategory::FileWrite
998        );
999        assert_eq!(
1000            sentinel.classify_tool("builtin:edit"),
1001            ToolRiskCategory::FileWrite
1002        );
1003    }
1004
1005    #[tokio::test]
1006    async fn classify_low_risk_returns_low() {
1007        let config = zeph_config::ShadowSentinelConfig::default();
1008        let sentinel = make_test_sentinel(config).await;
1009        assert_eq!(
1010            sentinel.classify_tool("builtin:read"),
1011            ToolRiskCategory::Low
1012        );
1013        assert_eq!(
1014            sentinel.classify_tool("builtin:search"),
1015            ToolRiskCategory::Low
1016        );
1017    }
1018
1019    /// #5750: an MCP-origin tool whose name matches no configured keyword pattern must still
1020    /// be probed (as `McpUnclassified`), not silently fall through to `Low`. Verbs like
1021    /// `frobnicate` can never be fully enumerated in `probe_patterns`.
1022    #[tokio::test]
1023    async fn classify_mcp_tool_with_no_keyword_match_is_mcp_unclassified() {
1024        let config = zeph_config::ShadowSentinelConfig::default();
1025        let sentinel = make_test_sentinel(config).await;
1026        sentinel
1027            .mcp_tool_ids_handle()
1028            .write()
1029            .insert("some-server_frobnicate".to_owned());
1030        assert_eq!(
1031            sentinel.classify_tool("some-server_frobnicate"),
1032            ToolRiskCategory::McpUnclassified
1033        );
1034    }
1035
1036    /// The same non-keyword-matching name for a tool NOT registered as MCP-origin must remain
1037    /// `Low` — engagement is gated on MCP origin, not merely on failing to match Low's fast path.
1038    #[tokio::test]
1039    async fn classify_non_mcp_tool_with_no_keyword_match_stays_low() {
1040        let config = zeph_config::ShadowSentinelConfig::default();
1041        let sentinel = make_test_sentinel(config).await;
1042        assert_eq!(
1043            sentinel.classify_tool("some-server_frobnicate"),
1044            ToolRiskCategory::Low
1045        );
1046    }
1047
1048    #[tokio::test]
1049    async fn classify_bare_shell_names_are_shell_risk() {
1050        let config = zeph_config::ShadowSentinelConfig::default();
1051        let sentinel = make_test_sentinel(config).await;
1052        assert_eq!(sentinel.classify_tool("bash"), ToolRiskCategory::Shell);
1053        assert_eq!(sentinel.classify_tool("shell"), ToolRiskCategory::Shell);
1054        assert_eq!(sentinel.classify_tool("sh"), ToolRiskCategory::Shell);
1055    }
1056
1057    #[tokio::test]
1058    async fn classify_bare_file_write_names_are_file_write_risk() {
1059        let config = zeph_config::ShadowSentinelConfig::default();
1060        let sentinel = make_test_sentinel(config).await;
1061        assert_eq!(sentinel.classify_tool("write"), ToolRiskCategory::FileWrite);
1062        assert_eq!(sentinel.classify_tool("edit"), ToolRiskCategory::FileWrite);
1063        assert_eq!(
1064            sentinel.classify_tool("delete"),
1065            ToolRiskCategory::FileWrite
1066        );
1067    }
1068
1069    /// #5736 regression: MCP-tool escalation must key off the registered `mcp_tool_ids` set
1070    /// (`ToolDef::server_id`-backed), not a `"mcp:"` string prefix — real MCP tool ids are
1071    /// `"{server_id}_{name}"` (`McpTool::sanitized_id`) and never carry that prefix, so the old
1072    /// check silently never escalated any MCP write/edit tool to `ExfilCapable`.
1073    #[tokio::test]
1074    async fn classify_mcp_tool_write_pattern_escalates_to_exfil_capable() {
1075        let config = zeph_config::ShadowSentinelConfig {
1076            probe_patterns: vec!["*edit*".to_owned()],
1077            ..zeph_config::ShadowSentinelConfig::default()
1078        };
1079        let sentinel = make_test_sentinel(config).await;
1080        // Unregistered: a same-shaped id falls to the ordinary FileWrite tier, not ExfilCapable.
1081        assert_eq!(
1082            sentinel.classify_tool("github_edit_file"),
1083            ToolRiskCategory::FileWrite
1084        );
1085        // Register it the same way the real MCP wiring does (via the shared handle) and the
1086        // identical id must now escalate.
1087        sentinel
1088            .mcp_tool_ids_handle()
1089            .write()
1090            .insert("github_edit_file".to_owned());
1091        assert_eq!(
1092            sentinel.classify_tool("github_edit_file"),
1093            ToolRiskCategory::ExfilCapable
1094        );
1095    }
1096
1097    /// #5736 follow-up (CI-1239): `ShadowSentinelConfig::default()` — not a hand-tuned override —
1098    /// must escalate a real MCP write tool. Real MCP tool ids are `"{server_id}_{name}"`
1099    /// (`McpTool::sanitized_id`), e.g. `"fs-test_write_file"`; the shipped default
1100    /// `probe_patterns` (`"mcp:*/file_*"`, `"mcp:*/exec_*"`) assumed a `"mcp:"`-prefixed id
1101    /// shape that no real id ever has, so the outer glob-matching loop in `classify_tool` never
1102    /// even entered the branch containing the `is_mcp_tool()` check — every MCP tool silently
1103    /// fell through to `ToolRiskCategory::Low` (probe skipped entirely), a complete bypass, not
1104    /// just a downgrade to `FileWrite`.
1105    #[tokio::test]
1106    async fn classify_mcp_tool_write_under_default_config_escalates_to_exfil_capable() {
1107        let config = zeph_config::ShadowSentinelConfig::default();
1108        let sentinel = make_test_sentinel(config).await;
1109        sentinel
1110            .mcp_tool_ids_handle()
1111            .write()
1112            .insert("fs-test_write_file".to_owned());
1113        assert_eq!(
1114            sentinel.classify_tool("fs-test_write_file"),
1115            ToolRiskCategory::ExfilCapable
1116        );
1117    }
1118
1119    #[tokio::test]
1120    async fn advance_turn_resets_counter() {
1121        let config = zeph_config::ShadowSentinelConfig::default();
1122        let sentinel = make_test_sentinel(config).await;
1123        sentinel.probes_this_turn.store(3, Ordering::Relaxed);
1124        sentinel.advance_turn();
1125        assert_eq!(sentinel.probes_this_turn.load(Ordering::Relaxed), 0);
1126    }
1127
1128    #[test]
1129    fn glob_matches_star_wildcard() {
1130        assert!(glob_matches("mcp:*/file_*", "mcp:myserver/file_read"));
1131        assert!(glob_matches("mcp:*/file_*", "mcp:other/file_write"));
1132        assert!(!glob_matches("mcp:*/file_*", "builtin:shell"));
1133    }
1134
1135    #[test]
1136    fn glob_matches_exact() {
1137        assert!(glob_matches("builtin:shell", "builtin:shell"));
1138        assert!(!glob_matches("builtin:shell", "builtin:write"));
1139    }
1140
1141    #[test]
1142    fn parse_verdict_allow() {
1143        let v = LlmSafetyProbe::parse_verdict(r#"{"verdict": "allow"}"#);
1144        assert_eq!(v, ProbeVerdict::Allow);
1145    }
1146
1147    #[test]
1148    fn parse_verdict_deny_with_reason() {
1149        let v =
1150            LlmSafetyProbe::parse_verdict(r#"{"verdict": "deny", "reason": "suspicious pattern"}"#);
1151        assert_eq!(
1152            v,
1153            ProbeVerdict::Deny {
1154                reason: "suspicious pattern".to_owned()
1155            }
1156        );
1157    }
1158
1159    #[test]
1160    fn parse_verdict_unparseable_allows() {
1161        let v = LlmSafetyProbe::parse_verdict("I think this is fine");
1162        assert_eq!(v, ProbeVerdict::Allow);
1163    }
1164
1165    #[tokio::test]
1166    async fn check_tool_call_skips_after_budget_exhausted() {
1167        let config = zeph_config::ShadowSentinelConfig {
1168            enabled: true,
1169            max_probes_per_turn: 2,
1170            ..zeph_config::ShadowSentinelConfig::default()
1171        };
1172        let sentinel = make_test_sentinel(config).await;
1173
1174        // First two calls should not be skipped (noop probe returns Allow).
1175        let args = serde_json::Value::Object(serde_json::Map::new());
1176        let v1 = sentinel
1177            .check_tool_call("builtin:shell", &args, 1, "calm")
1178            .await;
1179        let v2 = sentinel
1180            .check_tool_call("builtin:shell", &args, 1, "calm")
1181            .await;
1182        assert_ne!(v1, ProbeVerdict::Skip, "first call within budget");
1183        assert_ne!(v2, ProbeVerdict::Skip, "second call within budget");
1184
1185        // Third call exceeds max_probes_per_turn = 2 → must skip.
1186        let v3 = sentinel
1187            .check_tool_call("builtin:shell", &args, 1, "calm")
1188            .await;
1189        assert_eq!(
1190            v3,
1191            ProbeVerdict::Skip,
1192            "third call must be skipped (budget exhausted)"
1193        );
1194    }
1195
1196    /// #5749: `ExfilCapable` calls must never be skipped due to *shared* per-turn budget
1197    /// exhaustion — they draw from their own independent counter. Exhaust the shared budget with
1198    /// `Shell` calls first, then confirm a subsequent `ExfilCapable` call still probes.
1199    #[tokio::test]
1200    async fn check_tool_call_exfil_capable_bypasses_shared_budget_exhaustion() {
1201        let config = zeph_config::ShadowSentinelConfig {
1202            enabled: true,
1203            max_probes_per_turn: 1,
1204            probe_patterns: vec!["*edit*".to_owned()],
1205            ..zeph_config::ShadowSentinelConfig::default()
1206        };
1207        let sentinel = make_test_sentinel(config).await;
1208        sentinel
1209            .mcp_tool_ids_handle()
1210            .write()
1211            .insert("server_edit_file".to_owned());
1212        assert_eq!(
1213            sentinel.classify_tool("server_edit_file"),
1214            ToolRiskCategory::ExfilCapable
1215        );
1216
1217        let args = serde_json::Value::Object(serde_json::Map::new());
1218
1219        // Exhaust the shared budget (max_probes_per_turn = 1) with a Shell call.
1220        let v1 = sentinel
1221            .check_tool_call("builtin:shell", &args, 1, "calm")
1222            .await;
1223        assert_ne!(v1, ProbeVerdict::Skip, "first Shell call within budget");
1224        let v2 = sentinel
1225            .check_tool_call("builtin:shell", &args, 1, "calm")
1226            .await;
1227        assert_eq!(
1228            v2,
1229            ProbeVerdict::Skip,
1230            "second Shell call must be skipped — budget exhausted"
1231        );
1232
1233        // ExfilCapable call must still probe despite the exhausted shared budget.
1234        let v3 = sentinel
1235            .check_tool_call("server_edit_file", &args, 1, "calm")
1236            .await;
1237        assert_ne!(
1238            v3,
1239            ProbeVerdict::Skip,
1240            "ExfilCapable must not be starved by the shared per-turn budget"
1241        );
1242    }
1243
1244    /// #5749 follow-up (critic SIGNIFICANT-1): `ExfilCapable`'s independent budget must still be
1245    /// finite (`2 * max_probes_per_turn`), not unconditional — otherwise a false-positive
1246    /// keyword match on an MCP-origin tool name generates unbounded LLM probe calls.
1247    #[tokio::test]
1248    async fn check_tool_call_exfil_capable_has_finite_cap() {
1249        let config = zeph_config::ShadowSentinelConfig {
1250            enabled: true,
1251            max_probes_per_turn: 1,
1252            probe_patterns: vec!["*edit*".to_owned()],
1253            ..zeph_config::ShadowSentinelConfig::default()
1254        };
1255        let sentinel = make_test_sentinel(config).await;
1256        sentinel
1257            .mcp_tool_ids_handle()
1258            .write()
1259            .insert("server_edit_file".to_owned());
1260
1261        let args = serde_json::Value::Object(serde_json::Map::new());
1262
1263        // exfil_max = 2 * max_probes_per_turn = 2 — first two calls must probe.
1264        let v1 = sentinel
1265            .check_tool_call("server_edit_file", &args, 1, "calm")
1266            .await;
1267        let v2 = sentinel
1268            .check_tool_call("server_edit_file", &args, 1, "calm")
1269            .await;
1270        assert_ne!(
1271            v1,
1272            ProbeVerdict::Skip,
1273            "first ExfilCapable call within its own budget"
1274        );
1275        assert_ne!(
1276            v2,
1277            ProbeVerdict::Skip,
1278            "second ExfilCapable call within its own budget"
1279        );
1280
1281        // Third call exceeds the independent cap → must skip, not run unbounded.
1282        let v3 = sentinel
1283            .check_tool_call("server_edit_file", &args, 1, "calm")
1284            .await;
1285        assert_eq!(
1286            v3,
1287            ProbeVerdict::Skip,
1288            "ExfilCapable's own budget must still be finite (2 * max_probes_per_turn)"
1289        );
1290    }
1291
1292    /// #5750: `McpUnclassified` calls (engaged by MCP origin alone) must be capped at
1293    /// `max_probes_per_turn - 1`, reserving at least one slot for keyword-matched categories so
1294    /// a burst of low-signal MCP probes cannot starve out a later high-confidence probe.
1295    #[tokio::test]
1296    async fn check_tool_call_mcp_unclassified_reserves_budget_slot() {
1297        let config = zeph_config::ShadowSentinelConfig {
1298            enabled: true,
1299            max_probes_per_turn: 2,
1300            ..zeph_config::ShadowSentinelConfig::default()
1301        };
1302        let sentinel = make_test_sentinel(config).await;
1303        sentinel
1304            .mcp_tool_ids_handle()
1305            .write()
1306            .insert("some-server_frobnicate".to_owned());
1307        assert_eq!(
1308            sentinel.classify_tool("some-server_frobnicate"),
1309            ToolRiskCategory::McpUnclassified
1310        );
1311
1312        let args = serde_json::Value::Object(serde_json::Map::new());
1313
1314        // First McpUnclassified call consumes the one slot it's allowed (max - 1 = 1).
1315        let v1 = sentinel
1316            .check_tool_call("some-server_frobnicate", &args, 1, "calm")
1317            .await;
1318        assert_ne!(
1319            v1,
1320            ProbeVerdict::Skip,
1321            "first McpUnclassified call within reserved share"
1322        );
1323
1324        // Second McpUnclassified call must be skipped — reserved share (1) already used, even
1325        // though the shared counter (1/2) has not reached max_probes_per_turn.
1326        let v2 = sentinel
1327            .check_tool_call("some-server_frobnicate", &args, 1, "calm")
1328            .await;
1329        assert_eq!(
1330            v2,
1331            ProbeVerdict::Skip,
1332            "second McpUnclassified call must be skipped — reserved share exhausted"
1333        );
1334
1335        // A Shell call must still get through using the slot reserved for it.
1336        let v3 = sentinel
1337            .check_tool_call("builtin:shell", &args, 1, "calm")
1338            .await;
1339        assert_ne!(
1340            v3,
1341            ProbeVerdict::Skip,
1342            "Shell call must still probe using the slot reserved for non-McpUnclassified categories"
1343        );
1344    }
1345
1346    /// #5750 follow-up (critic SIGNIFICANT-2): at `max_probes_per_turn == 1` there is only one
1347    /// slot total, so reserving it for `Shell`/`FileWrite` means `McpUnclassified` gets NO share
1348    /// at all (`saturating_sub` floors at 0) rather than the whole budget. This is what makes
1349    /// the "cannot starve a higher-confidence probe" guarantee hold unconditionally, at the cost
1350    /// of never probing `McpUnclassified` when the turn's total budget is 1.
1351    #[tokio::test]
1352    async fn check_tool_call_mcp_unclassified_fully_reserved_out_at_budget_one() {
1353        let config = zeph_config::ShadowSentinelConfig {
1354            enabled: true,
1355            max_probes_per_turn: 1,
1356            ..zeph_config::ShadowSentinelConfig::default()
1357        };
1358        let sentinel = make_test_sentinel(config).await;
1359        sentinel
1360            .mcp_tool_ids_handle()
1361            .write()
1362            .insert("some-server_frobnicate".to_owned());
1363
1364        let args = serde_json::Value::Object(serde_json::Map::new());
1365
1366        // Even the FIRST McpUnclassified call must be skipped — the single slot is fully
1367        // reserved for keyword-matched categories, never handed to McpUnclassified.
1368        let v1 = sentinel
1369            .check_tool_call("some-server_frobnicate", &args, 1, "calm")
1370            .await;
1371        assert_eq!(
1372            v1,
1373            ProbeVerdict::Skip,
1374            "McpUnclassified must get zero share when max_probes_per_turn == 1"
1375        );
1376
1377        // A subsequent Shell call must still probe using the untouched slot — proving the
1378        // McpUnclassified attempt above did not consume or starve it.
1379        let v2 = sentinel
1380            .check_tool_call("builtin:shell", &args, 1, "calm")
1381            .await;
1382        assert_ne!(
1383            v2,
1384            ProbeVerdict::Skip,
1385            "Shell must not be starved by a prior McpUnclassified attempt at max_probes_per_turn == 1"
1386        );
1387    }
1388
1389    /// Boundary: `max_probes_per_turn == 0` must skip every category through the shared budget
1390    /// (`Shell`, `FileWrite`, `McpUnclassified`) AND the independent `ExfilCapable` budget
1391    /// (`2 * 0 == 0`) — an operator disabling the probe budget entirely must not leave any
1392    /// category unbounded.
1393    #[tokio::test]
1394    async fn check_tool_call_all_categories_skip_at_budget_zero() {
1395        let config = zeph_config::ShadowSentinelConfig {
1396            enabled: true,
1397            max_probes_per_turn: 0,
1398            probe_patterns: vec!["*edit*".to_owned()],
1399            ..zeph_config::ShadowSentinelConfig::default()
1400        };
1401        let sentinel = make_test_sentinel(config).await;
1402        sentinel
1403            .mcp_tool_ids_handle()
1404            .write()
1405            .insert("server_edit_file".to_owned());
1406        sentinel
1407            .mcp_tool_ids_handle()
1408            .write()
1409            .insert("some-server_frobnicate".to_owned());
1410        assert_eq!(
1411            sentinel.classify_tool("server_edit_file"),
1412            ToolRiskCategory::ExfilCapable
1413        );
1414        assert_eq!(
1415            sentinel.classify_tool("some-server_frobnicate"),
1416            ToolRiskCategory::McpUnclassified
1417        );
1418
1419        let args = serde_json::Value::Object(serde_json::Map::new());
1420        assert_eq!(
1421            sentinel
1422                .check_tool_call("builtin:shell", &args, 1, "calm")
1423                .await,
1424            ProbeVerdict::Skip,
1425            "Shell must skip when max_probes_per_turn == 0"
1426        );
1427        assert_eq!(
1428            sentinel
1429                .check_tool_call("some-server_frobnicate", &args, 1, "calm")
1430                .await,
1431            ProbeVerdict::Skip,
1432            "McpUnclassified must skip when max_probes_per_turn == 0"
1433        );
1434        assert_eq!(
1435            sentinel
1436                .check_tool_call("server_edit_file", &args, 1, "calm")
1437                .await,
1438            ProbeVerdict::Skip,
1439            "ExfilCapable's independent budget (2 * 0 == 0) must also skip, not run unbounded"
1440        );
1441    }
1442
1443    #[tokio::test]
1444    async fn check_tool_call_returns_skip_when_disabled() {
1445        let config = zeph_config::ShadowSentinelConfig {
1446            enabled: false,
1447            ..zeph_config::ShadowSentinelConfig::default()
1448        };
1449        let sentinel = make_test_sentinel(config).await;
1450        let args = serde_json::Value::Object(serde_json::Map::new());
1451        let verdict = sentinel
1452            .check_tool_call("builtin:shell", &args, 1, "calm")
1453            .await;
1454        assert_eq!(
1455            verdict,
1456            ProbeVerdict::Skip,
1457            "disabled sentinel must always return Skip without calling the probe"
1458        );
1459    }
1460
1461    // ── JoinSet regression tests (#4570) ─────────────────────────────────────
1462
1463    /// `drain_pending` awaits all spawned persist tasks and returns when the set is empty.
1464    #[tokio::test]
1465    async fn drain_pending_awaits_all_tasks() {
1466        use std::sync::atomic::{AtomicU32, Ordering};
1467
1468        let config = zeph_config::ShadowSentinelConfig::default();
1469        let sentinel = make_test_sentinel(config).await;
1470
1471        let counter = Arc::new(AtomicU32::new(0));
1472        for _ in 0..5 {
1473            let c = Arc::clone(&counter);
1474            sentinel
1475                .spawn_persist(async move {
1476                    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1477                    c.fetch_add(1, Ordering::Relaxed);
1478                })
1479                .await;
1480        }
1481
1482        sentinel.drain_pending().await;
1483
1484        assert_eq!(
1485            counter.load(Ordering::Relaxed),
1486            5,
1487            "drain_pending must join all 5 tasks before returning"
1488        );
1489    }
1490
1491    /// When the pending set is at capacity and all running tasks complete before the next
1492    /// `spawn_persist`, the new task IS accepted (the set has room after reaping).
1493    /// Conversely, if we fill the set, drain it, then overfill past capacity while tasks are
1494    /// still running — the implementation drops extras.  We verify the simpler property:
1495    /// `spawn_persist` never panics when called repeatedly beyond `MAX_PENDING_WRITES`.
1496    #[tokio::test]
1497    async fn spawn_persist_beyond_capacity_does_not_panic() {
1498        use std::sync::atomic::{AtomicU32, Ordering};
1499
1500        let config = zeph_config::ShadowSentinelConfig::default();
1501        let sentinel = make_test_sentinel(config).await;
1502        let counter = Arc::new(AtomicU32::new(0));
1503
1504        // Spawn twice the capacity; each task completes instantly.
1505        // spawn_persist will reap completed tasks between spawns, so most will be accepted.
1506        for _ in 0..(MAX_PENDING_WRITES * 2) {
1507            let c = Arc::clone(&counter);
1508            sentinel
1509                .spawn_persist(async move {
1510                    c.fetch_add(1, Ordering::Relaxed);
1511                })
1512                .await;
1513        }
1514
1515        sentinel.drain_pending().await;
1516
1517        // All tasks (or at least MAX_PENDING_WRITES of them) must have run; none panicked.
1518        let ran = counter.load(Ordering::Relaxed);
1519        assert!(
1520            ran >= u32::try_from(MAX_PENDING_WRITES).unwrap(),
1521            "at least MAX_PENDING_WRITES tasks must complete; ran={ran}"
1522        );
1523    }
1524
1525    // Build a minimal ShadowSentinel with a no-op probe for unit tests.
1526    //
1527    // Opens an in-memory SQLite pool. Store methods are never called in these unit
1528    // tests — they test only classification and counter logic.
1529    async fn make_test_sentinel(config: zeph_config::ShadowSentinelConfig) -> ShadowSentinel {
1530        struct NoopProbe;
1531        impl SafetyProbe for NoopProbe {
1532            fn evaluate<'a>(
1533                &'a self,
1534                _: &'a str,
1535                _: &'a JsonValue,
1536                _: &'a [SentinelEvent],
1537            ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>>
1538            {
1539                Box::pin(async { ProbeVerdict::Allow })
1540            }
1541        }
1542        let pool = test_pool().await;
1543        let store = ShadowEventStore::new(pool);
1544        ShadowSentinel::new(store, Box::new(NoopProbe), config, "test-session")
1545    }
1546
1547    // Opens a migrated in-memory SQLite pool (unlike `make_test_sentinel`'s pool, this one
1548    // has the `safety_shadow_events` table from migration 085 and can serve real store queries.
1549    async fn test_pool() -> DbPool {
1550        zeph_db::DbConfig {
1551            url: ":memory:".to_owned(),
1552            ..Default::default()
1553        }
1554        .connect()
1555        .await
1556        .expect("connect + migrate in-memory sqlite pool")
1557    }
1558
1559    fn make_event(
1560        session_id: &str,
1561        turn_number: u64,
1562        tool_id: &str,
1563        summary: &str,
1564    ) -> SentinelEvent {
1565        SentinelEvent {
1566            id: 0,
1567            session_id: SessionId::new(session_id),
1568            turn_number,
1569            event_type: "tool_call".to_owned(),
1570            tool_id: Some(tool_id.to_owned()),
1571            risk_signal: None,
1572            risk_level: "elevated".to_owned(),
1573            probe_verdict: None,
1574            context_summary: Some(summary.to_owned()),
1575            created_at: unix_now(),
1576        }
1577    }
1578
1579    #[tokio::test]
1580    async fn get_tool_history_returns_events_across_sessions() {
1581        let store = ShadowEventStore::new(test_pool().await);
1582
1583        store
1584            .record(&make_event(
1585                "session-a",
1586                1,
1587                "builtin:shell",
1588                "session-a ran a command",
1589            ))
1590            .await
1591            .expect("record session-a event");
1592        store
1593            .record(&make_event(
1594                "session-b",
1595                1,
1596                "builtin:shell",
1597                "session-b ran a command",
1598            ))
1599            .await
1600            .expect("record session-b event");
1601        store
1602            .record(&make_event(
1603                "session-a",
1604                2,
1605                "builtin:write",
1606                "unrelated tool",
1607            ))
1608            .await
1609            .expect("record unrelated-tool event");
1610
1611        let history = store
1612            .get_tool_history("builtin:shell", "unrelated-session", 10)
1613            .await
1614            .expect("get_tool_history");
1615
1616        assert_eq!(
1617            history.len(),
1618            2,
1619            "must return events from both non-excluded sessions for the queried tool_id, \
1620             excluding other tools"
1621        );
1622        assert!(history.iter().any(|e| e.session_id.as_str() == "session-a"));
1623        assert!(history.iter().any(|e| e.session_id.as_str() == "session-b"));
1624
1625        let history_excluding_a = store
1626            .get_tool_history("builtin:shell", "session-a", 10)
1627            .await
1628            .expect("get_tool_history");
1629        assert_eq!(
1630            history_excluding_a.len(),
1631            1,
1632            "exclude_session_id must be applied in SQL, not just usable for client-side \
1633             filtering afterward"
1634        );
1635        assert!(
1636            history_excluding_a
1637                .iter()
1638                .all(|e| e.session_id.as_str() != "session-a")
1639        );
1640    }
1641
1642    /// #5449 regression: `check_tool_call` must fold cross-session `get_tool_history` results
1643    /// into the trajectory passed to the probe, not just the current session's own events.
1644    #[tokio::test]
1645    async fn check_tool_call_incorporates_cross_session_tool_history() {
1646        struct CapturingProbe {
1647            captured: Arc<Mutex<Vec<SentinelEvent>>>,
1648        }
1649        impl SafetyProbe for CapturingProbe {
1650            fn evaluate<'a>(
1651                &'a self,
1652                _tool_id: &'a str,
1653                _tool_args: &'a JsonValue,
1654                trajectory: &'a [SentinelEvent],
1655            ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>>
1656            {
1657                let captured = Arc::clone(&self.captured);
1658                let trajectory = trajectory.to_vec();
1659                Box::pin(async move {
1660                    *captured.lock().await = trajectory;
1661                    ProbeVerdict::Allow
1662                })
1663            }
1664        }
1665
1666        let store = ShadowEventStore::new(test_pool().await);
1667        let other_session = "other-session";
1668        store
1669            .record(&make_event(
1670                other_session,
1671                1,
1672                "builtin:shell",
1673                "other session ran rm -rf",
1674            ))
1675            .await
1676            .expect("record cross-session event");
1677
1678        let captured: Arc<Mutex<Vec<SentinelEvent>>> = Arc::new(Mutex::new(Vec::new()));
1679
1680        let config = zeph_config::ShadowSentinelConfig {
1681            enabled: true,
1682            ..zeph_config::ShadowSentinelConfig::default()
1683        };
1684        let sentinel = ShadowSentinel::new(
1685            store,
1686            Box::new(CapturingProbe {
1687                captured: Arc::clone(&captured),
1688            }),
1689            config,
1690            "current-session",
1691        );
1692
1693        let args = serde_json::Value::Object(serde_json::Map::new());
1694        sentinel
1695            .check_tool_call("builtin:shell", &args, 1, "calm")
1696            .await;
1697
1698        let seen = captured.lock().await;
1699        assert!(
1700            seen.iter().any(|e| e.session_id.as_str() == other_session
1701                && e.context_summary.as_deref() == Some("other session ran rm -rf")),
1702            "probe context must include the cross-session tool history event, got: {seen:?}"
1703        );
1704    }
1705
1706    /// Drives `check_tool_call` for `tool_id` under `session_id` and returns the exact
1707    /// trajectory the probe received, so cap tests can assert WHICH events survive, not
1708    /// just how many — a count-only assertion can pass while the cap silently drops all
1709    /// cross-session data (the bug found in code review of the initial #5449 fix).
1710    async fn capture_check_tool_call_trajectory(
1711        store: ShadowEventStore,
1712        config: zeph_config::ShadowSentinelConfig,
1713        session_id: &str,
1714        tool_id: &str,
1715    ) -> Vec<SentinelEvent> {
1716        struct CapturingProbe {
1717            captured: Arc<Mutex<Vec<SentinelEvent>>>,
1718        }
1719        impl SafetyProbe for CapturingProbe {
1720            fn evaluate<'a>(
1721                &'a self,
1722                _tool_id: &'a str,
1723                _tool_args: &'a JsonValue,
1724                trajectory: &'a [SentinelEvent],
1725            ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>>
1726            {
1727                let captured = Arc::clone(&self.captured);
1728                let trajectory = trajectory.to_vec();
1729                Box::pin(async move {
1730                    *captured.lock().await = trajectory;
1731                    ProbeVerdict::Allow
1732                })
1733            }
1734        }
1735
1736        let captured: Arc<Mutex<Vec<SentinelEvent>>> = Arc::new(Mutex::new(Vec::new()));
1737        let sentinel = ShadowSentinel::new(
1738            store,
1739            Box::new(CapturingProbe {
1740                captured: Arc::clone(&captured),
1741            }),
1742            config,
1743            session_id,
1744        );
1745        let args = serde_json::Value::Object(serde_json::Map::new());
1746        sentinel.check_tool_call(tool_id, &args, 1, "calm").await;
1747        captured.lock().await.clone()
1748    }
1749
1750    /// Seeds `count` events for `session_id`/`tool_id`, with ascending `created_at`
1751    /// timestamps starting at `base`, so cap tests can control which events are "most
1752    /// recent". Summaries are `"{summary_prefix}-{i}"` for index-based assertions.
1753    async fn seed_events(
1754        store: &ShadowEventStore,
1755        session_id: &str,
1756        tool_id: &str,
1757        summary_prefix: &str,
1758        base: i64,
1759        count: u32,
1760    ) {
1761        for i in 0..count {
1762            let mut event = make_event(
1763                session_id,
1764                u64::from(i),
1765                tool_id,
1766                &format!("{summary_prefix}-{i}"),
1767            );
1768            event.created_at = base + i64::from(i);
1769            store.record(&event).await.expect("record seeded event");
1770        }
1771    }
1772
1773    /// Session trajectory and cross-session history are each independently capped at
1774    /// `max_context_events`, so a naive merge can total up to 2x the configured budget.
1775    /// `check_tool_call` must enforce the combined cap AND reserve budget for cross-session
1776    /// data — the original fix trimmed unconditionally from the front, which silently wiped
1777    /// ALL cross-session events whenever the session's own trajectory alone filled the
1778    /// budget (precisely the busiest-session scenario #5449 cares about most).
1779    #[tokio::test]
1780    async fn check_tool_call_cap_reserves_cross_session_budget_when_session_heavy() {
1781        let store = ShadowEventStore::new(test_pool().await);
1782        let base = unix_now();
1783        seed_events(
1784            &store,
1785            "current-session",
1786            "builtin:shell",
1787            "session",
1788            base,
1789            4,
1790        )
1791        .await;
1792        seed_events(&store, "other-session", "builtin:shell", "cross", base, 3).await;
1793
1794        let config = zeph_config::ShadowSentinelConfig {
1795            enabled: true,
1796            max_context_events: 4,
1797            ..zeph_config::ShadowSentinelConfig::default()
1798        };
1799        let trajectory =
1800            capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1801                .await;
1802
1803        assert_eq!(
1804            trajectory.len(),
1805            4,
1806            "total must be capped at max_context_events"
1807        );
1808        let cross_session_count = trajectory
1809            .iter()
1810            .filter(|e| e.session_id.as_str() == "other-session")
1811            .count();
1812        assert_eq!(
1813            cross_session_count, 2,
1814            "cross-session budget is max_context_events/2 = 2, and must survive even \
1815             though the session's own trajectory alone fills the whole budget; \
1816             got trajectory: {trajectory:?}"
1817        );
1818    }
1819
1820    /// Mirror case: cross-session history is the one over budget, session's own trajectory
1821    /// is light. The session-side cap must not trim events that don't need trimming.
1822    #[tokio::test]
1823    async fn check_tool_call_cap_cross_session_heavy_case() {
1824        let store = ShadowEventStore::new(test_pool().await);
1825        let base = unix_now();
1826        seed_events(
1827            &store,
1828            "current-session",
1829            "builtin:shell",
1830            "session",
1831            base,
1832            1,
1833        )
1834        .await;
1835        seed_events(&store, "other-session", "builtin:shell", "cross", base, 4).await;
1836
1837        let config = zeph_config::ShadowSentinelConfig {
1838            enabled: true,
1839            max_context_events: 4,
1840            ..zeph_config::ShadowSentinelConfig::default()
1841        };
1842        let trajectory =
1843            capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1844                .await;
1845
1846        let session_count = trajectory
1847            .iter()
1848            .filter(|e| e.session_id.as_str() == "current-session")
1849            .count();
1850        let cross_session_count = trajectory.len() - session_count;
1851        assert_eq!(
1852            session_count, 1,
1853            "session's own (light) trajectory must not be trimmed"
1854        );
1855        assert_eq!(
1856            cross_session_count, 2,
1857            "cross-session budget is max_context_events/2 = 2"
1858        );
1859    }
1860
1861    /// Boundary: session + cross-session totals exactly `max_context_events` — nothing
1862    /// should be dropped from either side.
1863    #[tokio::test]
1864    async fn check_tool_call_cap_boundary_at_exact_limit() {
1865        let store = ShadowEventStore::new(test_pool().await);
1866        let base = unix_now();
1867        seed_events(
1868            &store,
1869            "current-session",
1870            "builtin:shell",
1871            "session",
1872            base,
1873            2,
1874        )
1875        .await;
1876        seed_events(&store, "other-session", "builtin:shell", "cross", base, 2).await;
1877
1878        let config = zeph_config::ShadowSentinelConfig {
1879            enabled: true,
1880            max_context_events: 4,
1881            ..zeph_config::ShadowSentinelConfig::default()
1882        };
1883        let trajectory =
1884            capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1885                .await;
1886
1887        assert_eq!(
1888            trajectory.len(),
1889            4,
1890            "exactly at the limit: nothing should be dropped"
1891        );
1892    }
1893
1894    /// Boundary: one more cross-session event than the reserved budget — exactly one event
1895    /// must be dropped, and it must be the OLDEST cross-session event (trajectory stays
1896    /// oldest-first/ASC, so the most recent events are kept).
1897    #[tokio::test]
1898    async fn check_tool_call_cap_boundary_at_limit_plus_one() {
1899        let store = ShadowEventStore::new(test_pool().await);
1900        let base = unix_now();
1901        seed_events(
1902            &store,
1903            "current-session",
1904            "builtin:shell",
1905            "session",
1906            base,
1907            2,
1908        )
1909        .await;
1910        seed_events(&store, "other-session", "builtin:shell", "cross", base, 3).await;
1911
1912        let config = zeph_config::ShadowSentinelConfig {
1913            enabled: true,
1914            max_context_events: 4,
1915            ..zeph_config::ShadowSentinelConfig::default()
1916        };
1917        let trajectory =
1918            capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1919                .await;
1920
1921        assert_eq!(
1922            trajectory.len(),
1923            4,
1924            "limit+1 overall: exactly one event must be dropped"
1925        );
1926        let cross_summaries: Vec<&str> = trajectory
1927            .iter()
1928            .filter(|e| e.session_id.as_str() == "other-session")
1929            .filter_map(|e| e.context_summary.as_deref())
1930            .collect();
1931        assert_eq!(
1932            cross_summaries,
1933            vec!["cross-1", "cross-2"],
1934            "the oldest cross-session event (cross-0) must be the one dropped, \
1935             got: {cross_summaries:?}"
1936        );
1937    }
1938
1939    /// The current session's own events must not be double-counted into the cross-session
1940    /// block — `get_tool_history` excludes `exclude_session_id` directly in its SQL
1941    /// (`AND session_id != ?`), and this test confirms that exclusion end-to-end through
1942    /// `check_tool_call`.
1943    #[tokio::test]
1944    async fn check_tool_call_excludes_current_session_from_cross_session_merge() {
1945        let store = ShadowEventStore::new(test_pool().await);
1946        let base = unix_now();
1947        seed_events(&store, "current-session", "builtin:shell", "own", base, 2).await;
1948
1949        let config = zeph_config::ShadowSentinelConfig {
1950            enabled: true,
1951            max_context_events: 10,
1952            ..zeph_config::ShadowSentinelConfig::default()
1953        };
1954        let trajectory =
1955            capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1956                .await;
1957
1958        assert_eq!(
1959            trajectory.len(),
1960            2,
1961            "current session's own events must appear exactly once, not duplicated via \
1962             the cross-session merge; got: {trajectory:?}"
1963        );
1964    }
1965
1966    /// `probe_result` events from OTHER sessions must never leak into the cross-session
1967    /// merge — `get_tool_history`'s SQL does not filter by `event_type`, so this relies
1968    /// entirely on the Rust-side filter (the same LLM-isolation invariant already tested
1969    /// for the same-session trajectory, exercised here on the cross-session path).
1970    #[tokio::test]
1971    async fn check_tool_call_excludes_probe_result_events_from_cross_session_merge() {
1972        let store = ShadowEventStore::new(test_pool().await);
1973        let base = unix_now();
1974        let mut event = make_event("other-session", 1, "builtin:shell", "probe verdict leaked");
1975        event.event_type = "probe_result".to_owned();
1976        event.created_at = base;
1977        store
1978            .record(&event)
1979            .await
1980            .expect("record probe_result event");
1981
1982        let config = zeph_config::ShadowSentinelConfig {
1983            enabled: true,
1984            max_context_events: 10,
1985            ..zeph_config::ShadowSentinelConfig::default()
1986        };
1987        let trajectory =
1988            capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell")
1989                .await;
1990
1991        assert!(
1992            trajectory.is_empty(),
1993            "probe_result events from other sessions must never appear in the \
1994             cross-session merge (LLM isolation invariant), got: {trajectory:?}"
1995        );
1996    }
1997
1998    // ── #5766: record_tool_event had zero test coverage ─────────────────────────
1999
2000    #[tokio::test]
2001    async fn record_tool_event_persists_event_normal_path() {
2002        let config = zeph_config::ShadowSentinelConfig {
2003            enabled: true,
2004            ..zeph_config::ShadowSentinelConfig::default()
2005        };
2006        let sentinel = make_test_sentinel(config).await;
2007
2008        sentinel
2009            .record_tool_event("builtin:shell", 3, "elevated", "ran `ls -la`")
2010            .await;
2011        sentinel.drain_pending().await;
2012
2013        let events = sentinel
2014            .store
2015            .get_trajectory("test-session", 10)
2016            .await
2017            .expect("get_trajectory");
2018        assert_eq!(events.len(), 1, "expected exactly one persisted event");
2019        assert_eq!(events[0].event_type, "tool_call");
2020        assert_eq!(events[0].tool_id.as_deref(), Some("builtin:shell"));
2021        assert_eq!(events[0].turn_number, 3);
2022        assert_eq!(events[0].risk_level, "elevated");
2023        assert_eq!(events[0].context_summary.as_deref(), Some("ran `ls -la`"));
2024    }
2025
2026    #[tokio::test]
2027    async fn record_tool_event_disabled_does_not_persist() {
2028        let config = zeph_config::ShadowSentinelConfig {
2029            enabled: false,
2030            ..zeph_config::ShadowSentinelConfig::default()
2031        };
2032        let sentinel = make_test_sentinel(config).await;
2033
2034        sentinel
2035            .record_tool_event("builtin:shell", 1, "elevated", "should be skipped")
2036            .await;
2037        sentinel.drain_pending().await;
2038
2039        let events = sentinel
2040            .store
2041            .get_trajectory("test-session", 10)
2042            .await
2043            .expect("get_trajectory");
2044        assert!(
2045            events.is_empty(),
2046            "record_tool_event must be a no-op when the sentinel is disabled"
2047        );
2048    }
2049
2050    /// Minimal `tracing_subscriber::Layer` that captures event messages into a shared buffer,
2051    /// used to verify `record_tool_event`'s fire-and-forget persist failure logs the correct
2052    /// warn context ("failed to persist tool event", as opposed to `check_tool_call`'s "failed
2053    /// to persist probe result").
2054    struct MessageCaptureLayer {
2055        messages: Arc<std::sync::Mutex<Vec<String>>>,
2056    }
2057
2058    struct MessageVisitor(String);
2059
2060    impl tracing::field::Visit for MessageVisitor {
2061        fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
2062            if field.name() == "message" {
2063                self.0 = format!("{value:?}");
2064            }
2065        }
2066    }
2067
2068    impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for MessageCaptureLayer {
2069        fn on_event(
2070            &self,
2071            event: &tracing::Event<'_>,
2072            _ctx: tracing_subscriber::layer::Context<'_, S>,
2073        ) {
2074            let mut visitor = MessageVisitor(String::new());
2075            event.record(&mut visitor);
2076            self.messages.lock().unwrap().push(visitor.0);
2077        }
2078    }
2079
2080    // `record()` fails when the backing table is gone. Self-loop-style DB-trigger tampering
2081    // isn't needed here — a real store error is the whole point of this test.
2082    #[tokio::test]
2083    async fn record_tool_event_persist_failure_logs_warn_with_tool_event_context() {
2084        use tracing_subscriber::layer::SubscriberExt as _;
2085
2086        struct NoopProbe;
2087        impl SafetyProbe for NoopProbe {
2088            fn evaluate<'a>(
2089                &'a self,
2090                _: &'a str,
2091                _: &'a JsonValue,
2092                _: &'a [SentinelEvent],
2093            ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ProbeVerdict> + Send + 'a>>
2094            {
2095                Box::pin(async { ProbeVerdict::Allow })
2096            }
2097        }
2098
2099        let pool = test_pool().await;
2100        zeph_db::query(zeph_db::sql!("DROP TABLE safety_shadow_events"))
2101            .execute(&pool)
2102            .await
2103            .expect("drop safety_shadow_events table");
2104        let store = ShadowEventStore::new(pool);
2105        let config = zeph_config::ShadowSentinelConfig {
2106            enabled: true,
2107            ..zeph_config::ShadowSentinelConfig::default()
2108        };
2109        let sentinel = ShadowSentinel::new(store, Box::new(NoopProbe), config, "test-session");
2110
2111        let messages: Arc<std::sync::Mutex<Vec<String>>> =
2112            Arc::new(std::sync::Mutex::new(Vec::new()));
2113        let layer = MessageCaptureLayer {
2114            messages: messages.clone(),
2115        };
2116        let subscriber = tracing_subscriber::registry().with(layer);
2117        let _guard = tracing::subscriber::set_default(subscriber);
2118
2119        sentinel
2120            .record_tool_event("builtin:shell", 1, "elevated", "ran a command")
2121            .await;
2122        sentinel.drain_pending().await;
2123
2124        let captured = messages.lock().unwrap();
2125        assert!(
2126            captured
2127                .iter()
2128                .any(|m| m.contains("failed to persist tool event")),
2129            "expected a warn log with 'failed to persist tool event' context, got: {captured:?}"
2130        );
2131    }
2132}