Skip to main content

ferrin_policy/
shadow.rs

1//! Shadow mode: observe policy decisions before enforcing them.
2//!
3//! Derived from the Vercel AI SDK shadow policy (Apache-2.0,
4//! Copyright 2023 Vercel, Inc.), reimplemented with owned Rust tasks.
5
6use std::fmt;
7use std::future::Future;
8use std::panic::AssertUnwindSafe;
9use std::panic::catch_unwind;
10use std::sync::Arc;
11use std::sync::Mutex;
12
13use chrono::DateTime;
14use chrono::Utc;
15use ferrin_core::generate_text::ApprovalContext;
16use ferrin_core::generate_text::ApprovalPolicy;
17use ferrin_core::generate_text::ApprovalStatus;
18use ferrin_core::generate_text::ParsedToolCall;
19use ferrin_spec::BoxFuture;
20use ferrin_spec::JsonValue;
21use ferrin_spec::ToolCallId;
22use ferrin_spec::ToolName;
23use serde::Deserialize;
24use serde::Serialize;
25use tokio::task::JoinSet;
26
27/// Whether a [`Shadow`] policy acts on the decisions it observes.
28#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
29#[non_exhaustive]
30pub enum Enforcement {
31    /// Report decisions and approve execution, overriding tool-defined approval.
32    #[default]
33    Observe,
34    /// Report and return the decisions.
35    Enforce,
36}
37
38/// Identifying information for the tool call a policy evaluated.
39#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
40pub struct PolicyDecisionToolCall {
41    /// Registered name of the tool.
42    pub tool_name: ToolName,
43    /// Identifier of this tool invocation.
44    pub tool_call_id: ToolCallId,
45    /// Parsed tool input submitted for approval.
46    pub input: JsonValue,
47}
48
49/// The observed and effective approval decisions for one tool call.
50#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
51pub struct PolicyDecisionEvent {
52    /// The evaluated tool invocation.
53    pub tool_call: PolicyDecisionToolCall,
54    /// Normalized decision returned by the wrapped policy.
55    pub decision: ApprovalStatus,
56    /// Whether the effective decision enforces the wrapped policy.
57    pub enforced: bool,
58    /// Status the generation loop will act on.
59    pub effective: ApprovalStatus,
60    /// UTC evaluation time, serialized as an ISO 8601 string.
61    pub timestamp: DateTime<Utc>,
62}
63
64/// Asynchronous observer receiving each evaluated and effective decision.
65pub type OnDecisionFn = Arc<dyn Fn(PolicyDecisionEvent) -> BoxFuture<'static, ()> + Send + Sync>;
66
67/// Synchronous low-level observer of the original policy return value.
68pub type OnDecisionSyncFn = Arc<dyn Fn(&ParsedToolCall, Option<&ApprovalStatus>) + Send + Sync>;
69
70enum Observer {
71    Async(OnDecisionFn),
72    Sync(OnDecisionSyncFn),
73}
74
75/// Approval policy created by [`shadow`].
76pub struct Shadow<P> {
77    inner: P,
78    enforcement: Enforcement,
79    on_decision: Option<Observer>,
80    audit_tasks: Mutex<JoinSet<()>>,
81}
82
83/// Evaluates `policy` for every call and reports its decision through
84/// [`Shadow::on_decision`], but only acts on it under
85/// [`Enforcement::Enforce`]. Roll a policy out by observing first and
86/// flipping the enforcement later without changing the wiring.
87pub fn shadow<P: ApprovalPolicy>(policy: P) -> Shadow<P> {
88    Shadow {
89        inner: policy,
90        enforcement: Enforcement::Observe,
91        on_decision: None,
92        audit_tasks: Mutex::new(JoinSet::new()),
93    }
94}
95
96impl<P> Shadow<P> {
97    /// Sets whether decisions are enforced (default: observe only).
98    #[must_use]
99    pub fn enforcement(mut self, enforcement: Enforcement) -> Self {
100        self.enforcement = enforcement;
101        self
102    }
103
104    /// Registers an asynchronous decision observer, replacing any prior observer.
105    ///
106    /// The observer runs independently on the current Tokio runtime; its output
107    /// and task panics are ignored so auditing cannot change approval. Without a
108    /// Tokio runtime the event is skipped. Dropping the policy cancels unfinished
109    /// observers; call [`Self::flush_decisions`] to drain queued events first.
110    ///
111    /// # Examples
112    ///
113    /// ```
114    /// use ferrin_core::generate_text::ApprovalStatus;
115    /// use ferrin_policy::shadow;
116    ///
117    /// let policy = shadow(ApprovalStatus::denied()).on_decision(|event| async move {
118    ///     assert!(!event.enforced);
119    ///     assert_eq!(event.effective, ApprovalStatus::approved());
120    /// });
121    /// ```
122    #[must_use]
123    pub fn on_decision<F, Fut>(mut self, observer: F) -> Self
124    where
125        F: Fn(PolicyDecisionEvent) -> Fut + Send + Sync + 'static,
126        Fut: Future + Send + 'static,
127    {
128        let observer = Arc::new(observer);
129        self.on_decision = Some(Observer::Async(Arc::new(move |event| {
130            let observer = Arc::clone(&observer);
131            Box::pin(async move {
132                let _ = observer(event).await;
133            })
134        })));
135        self
136    }
137
138    /// Registers a synchronous low-level observer, replacing any prior observer.
139    ///
140    /// This Ferrin extension receives the raw optional status and blocks approval
141    /// until it returns. Unwinding observer panics are ignored. Prefer
142    /// [`Self::on_decision`] for normalized events and independent execution.
143    #[must_use]
144    pub fn on_decision_sync(
145        mut self,
146        observer: impl Fn(&ParsedToolCall, Option<&ApprovalStatus>) + Send + Sync + 'static,
147    ) -> Self {
148        self.on_decision = Some(Observer::Sync(Arc::new(observer)));
149        self
150    }
151
152    /// Waits for audit callbacks queued before this call took their task set.
153    ///
154    /// Concurrent approval resolution can queue new events without waiting for
155    /// this flush. Cancelling the flush cancels the callbacks it took ownership of.
156    pub async fn flush_decisions(&self) {
157        let mut pending = {
158            let mut tasks = self
159                .audit_tasks
160                .lock()
161                .unwrap_or_else(std::sync::PoisonError::into_inner);
162            std::mem::take(&mut *tasks)
163        };
164        while pending.join_next().await.is_some() {}
165    }
166
167    fn report(
168        &self,
169        call: &ParsedToolCall,
170        status: Option<&ApprovalStatus>,
171        effective: &ApprovalStatus,
172    ) {
173        match &self.on_decision {
174            Some(Observer::Async(observer)) => {
175                let Ok(runtime) = tokio::runtime::Handle::try_current() else {
176                    return;
177                };
178                let event = PolicyDecisionEvent {
179                    tool_call: PolicyDecisionToolCall {
180                        tool_name: call.tool_name.clone(),
181                        tool_call_id: call.tool_call_id.clone(),
182                        input: call.input.clone(),
183                    },
184                    decision: status.cloned().unwrap_or(ApprovalStatus::NotApplicable),
185                    enforced: self.enforcement == Enforcement::Enforce,
186                    effective: effective.clone(),
187                    timestamp: Utc::now(),
188                };
189                let observer = Arc::clone(observer);
190                let mut tasks = self
191                    .audit_tasks
192                    .lock()
193                    .unwrap_or_else(std::sync::PoisonError::into_inner);
194                while tasks.try_join_next().is_some() {}
195                tasks.spawn_on(async move { observer(event).await }, &runtime);
196            }
197            Some(Observer::Sync(observer)) => {
198                let _ = catch_unwind(AssertUnwindSafe(|| observer(call, status)));
199            }
200            None => {}
201        }
202    }
203}
204
205impl<P> fmt::Debug for Shadow<P> {
206    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
207        f.debug_struct("Shadow")
208            .field("enforcement", &self.enforcement)
209            .field("on_decision", &self.on_decision.is_some())
210            .finish_non_exhaustive()
211    }
212}
213
214impl<P: ApprovalPolicy> ApprovalPolicy for Shadow<P> {
215    fn resolve<'a>(
216        &'a self,
217        call: &'a ParsedToolCall,
218        ctx: ApprovalContext<'a>,
219    ) -> BoxFuture<'a, Option<ApprovalStatus>> {
220        Box::pin(async move {
221            let status = self.inner.resolve(call, ctx).await;
222            tracing::debug!(
223                tool = %call.tool_name,
224                status = crate::diagnostics::status_kind(status.as_ref()),
225                enforcement = ?self.enforcement,
226                "shadow policy decision"
227            );
228            let effective = match self.enforcement {
229                Enforcement::Observe => ApprovalStatus::approved(),
230                Enforcement::Enforce => status.clone().unwrap_or(ApprovalStatus::NotApplicable),
231            };
232            self.report(call, status.as_ref(), &effective);
233            Some(effective)
234        })
235    }
236}