chio_kernel/
post_invocation.rs

1//! Post-invocation hook pipeline executed after a tool returns output.
2
3use chio_core::receipt::metadata::GuardEvidence;
4use chio_core::{capability::scope::ChioScope, AgentId, ServerId};
5use serde::Serialize;
6use serde_json::Value;
7use std::sync::{Mutex, MutexGuard};
8
9use crate::runtime::ToolCallRequest;
10
11/// Verdict from a post-invocation hook.
12#[derive(Debug, Clone)]
13pub enum PostInvocationVerdict {
14    Allow,
15    Block(String),
16    Redact(Value),
17    Escalate(String),
18}
19
20/// Context available to post-invocation hooks after a tool has executed.
21#[derive(Clone, Copy, Debug)]
22pub struct PostInvocationContext<'a> {
23    pub tool_name: &'a str,
24    pub request: Option<&'a ToolCallRequest>,
25    pub scope: Option<&'a ChioScope>,
26    pub agent_id: Option<&'a AgentId>,
27    pub server_id: Option<&'a ServerId>,
28    pub matched_grant_index: Option<usize>,
29}
30
31impl<'a> PostInvocationContext<'a> {
32    #[must_use]
33    pub fn synthetic(tool_name: &'a str) -> Self {
34        Self {
35            tool_name,
36            request: None,
37            scope: None,
38            agent_id: None,
39            server_id: None,
40            matched_grant_index: None,
41        }
42    }
43
44    #[must_use]
45    pub fn from_request(request: &'a ToolCallRequest, matched_grant_index: Option<usize>) -> Self {
46        Self {
47            tool_name: request.tool_name.as_str(),
48            request: Some(request),
49            scope: Some(&request.capability.scope),
50            agent_id: Some(&request.agent_id),
51            server_id: Some(&request.server_id),
52            matched_grant_index,
53        }
54    }
55}
56
57/// A hook that inspects tool responses after invocation.
58pub trait PostInvocationHook: Send + Sync {
59    fn name(&self) -> &str;
60
61    fn inspect(&self, ctx: &PostInvocationContext<'_>, response: &Value) -> PostInvocationVerdict;
62
63    /// Stable identity for deterministic, non-blocking durable evaluation.
64    ///
65    /// Returning an identity asserts that `inspect` is a pure function of the
66    /// request context and response, and that it returns only `Allow`,
67    /// `Redact`, or `Escalate`. Hooks without this contract remain usable on
68    /// non-durable paths and are rejected before a durable dispatch starts.
69    fn durable_identity(&self) -> Result<Option<PostInvocationHookIdentity>, String> {
70        Ok(None)
71    }
72
73    fn take_evidence(&self) -> Option<GuardEvidence> {
74        None
75    }
76}
77
78/// Canonical implementation and configuration identity for a durable hook.
79#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
80#[serde(deny_unknown_fields)]
81pub struct PostInvocationHookIdentity {
82    component_id: String,
83    component_version: String,
84    implementation_digest: String,
85}
86
87impl PostInvocationHookIdentity {
88    /// Bind an implementation tag and its complete effective configuration.
89    pub fn from_canonical_config<T: Serialize>(
90        component_id: impl Into<String>,
91        component_version: impl Into<String>,
92        implementation_tag: &str,
93        configuration: &T,
94    ) -> Result<Self, String> {
95        #[derive(Serialize)]
96        struct IdentityPreimage<'a, T> {
97            schema: &'static str,
98            component_id: &'a str,
99            component_version: &'a str,
100            implementation_tag: &'a str,
101            configuration: &'a T,
102        }
103
104        let component_id = component_id.into();
105        let component_version = component_version.into();
106        for (field, value) in [
107            ("component_id", component_id.as_str()),
108            ("component_version", component_version.as_str()),
109            ("implementation_tag", implementation_tag),
110        ] {
111            if value.trim().is_empty() || value.len() > 128 {
112                return Err(format!(
113                    "post-invocation {field} must contain 1 to 128 characters"
114                ));
115            }
116        }
117        let preimage = IdentityPreimage {
118            schema: "chio.post-invocation-hook-identity.v1",
119            component_id: &component_id,
120            component_version: &component_version,
121            implementation_tag,
122            configuration,
123        };
124        let canonical = crate::canonical_json_bytes(&preimage)
125            .map_err(|error| format!("post-invocation identity is not canonical JSON: {error}"))?;
126        Ok(Self {
127            component_id,
128            component_version,
129            implementation_digest: crate::sha256_hex(&canonical),
130        })
131    }
132
133    #[must_use]
134    pub fn component_id(&self) -> &str {
135        &self.component_id
136    }
137
138    #[must_use]
139    pub fn component_version(&self) -> &str {
140        &self.component_version
141    }
142
143    #[must_use]
144    pub fn implementation_digest(&self) -> &str {
145        &self.implementation_digest
146    }
147}
148
149#[derive(Debug, Clone, Serialize)]
150#[serde(tag = "verdict", rename_all = "snake_case")]
151enum DurablePostInvocationVerdict {
152    Allow,
153    Redact,
154    Escalate { message: String },
155}
156
157#[derive(Debug, Clone, Serialize)]
158#[serde(deny_unknown_fields)]
159pub(crate) struct DurablePipelineStepResult {
160    schema: &'static str,
161    identity: PostInvocationHookIdentity,
162    verdict: DurablePostInvocationVerdict,
163    response: Value,
164    evidence: Option<GuardEvidence>,
165}
166
167pub(crate) struct DurablePipelineOutcome {
168    pub(crate) outcome: PipelineOutcome,
169    pub(crate) step_results: Vec<DurablePipelineStepResult>,
170}
171
172/// Outcome of running the pipeline.
173#[derive(Debug, Clone)]
174pub struct PipelineOutcome {
175    pub verdict: PostInvocationVerdict,
176    pub escalations: Vec<String>,
177    pub evidence: Vec<GuardEvidence>,
178}
179
180/// Pipeline of post-invocation hooks evaluated in registration order.
181pub struct PostInvocationPipeline {
182    hooks: Vec<Box<dyn PostInvocationHook>>,
183    evaluation_lock: Mutex<()>,
184}
185
186impl PostInvocationPipeline {
187    #[must_use]
188    pub fn new() -> Self {
189        Self {
190            hooks: Vec::new(),
191            evaluation_lock: Mutex::new(()),
192        }
193    }
194
195    pub fn add(&mut self, hook: Box<dyn PostInvocationHook>) {
196        self.hooks.push(hook);
197    }
198
199    pub fn append(&mut self, mut other: Self) {
200        self.hooks.append(&mut other.hooks);
201    }
202
203    #[must_use]
204    pub fn len(&self) -> usize {
205        self.hooks.len()
206    }
207
208    #[must_use]
209    pub fn is_empty(&self) -> bool {
210        self.hooks.is_empty()
211    }
212
213    pub(crate) fn durable_identities(&self) -> Result<Vec<PostInvocationHookIdentity>, String> {
214        self.hooks
215            .iter()
216            .map(|hook| {
217                hook.durable_identity()?.ok_or_else(|| {
218                    format!(
219                        "post-invocation hook {:?} has no durable implementation identity",
220                        hook.name()
221                    )
222                })
223            })
224            .collect()
225    }
226
227    pub(crate) fn evaluate_durable_with_context_and_evidence(
228        &self,
229        context: &PostInvocationContext<'_>,
230        response: &Value,
231        expected_identities: &[PostInvocationHookIdentity],
232    ) -> Result<DurablePipelineOutcome, String> {
233        let _evaluation = self.lock_evaluation();
234        if self.hooks.len() != expected_identities.len() {
235            return Err("post-invocation pipeline changed after durable admission".to_owned());
236        }
237
238        let mut current_response = response.clone();
239        let mut escalations = Vec::new();
240        let mut evidence = Vec::new();
241        let mut step_results = Vec::with_capacity(self.hooks.len());
242
243        for (hook, expected_identity) in self.hooks.iter().zip(expected_identities) {
244            let identity = hook.durable_identity()?.ok_or_else(|| {
245                format!(
246                    "post-invocation hook {:?} lost its durable implementation identity",
247                    hook.name()
248                )
249            })?;
250            if &identity != expected_identity {
251                return Err(format!(
252                    "post-invocation hook {:?} changed after durable admission",
253                    hook.name()
254                ));
255            }
256            let verdict = hook.inspect(context, &current_response);
257            let hook_evidence = hook.take_evidence();
258            if let Some(item) = hook_evidence.clone() {
259                evidence.push(item);
260            }
261            let durable_verdict = match verdict {
262                PostInvocationVerdict::Allow => DurablePostInvocationVerdict::Allow,
263                PostInvocationVerdict::Redact(redacted) => {
264                    current_response = redacted;
265                    DurablePostInvocationVerdict::Redact
266                }
267                PostInvocationVerdict::Escalate(message) => {
268                    escalations.push(message.clone());
269                    DurablePostInvocationVerdict::Escalate { message }
270                }
271                PostInvocationVerdict::Block(_) => {
272                    return Err(format!(
273                        "durable post-invocation hook {:?} violated its non-blocking contract",
274                        hook.name()
275                    ));
276                }
277            };
278            step_results.push(DurablePipelineStepResult {
279                schema: "chio.durable-post-invocation-step-result.v1",
280                identity,
281                verdict: durable_verdict,
282                response: current_response.clone(),
283                evidence: hook_evidence,
284            });
285        }
286
287        let verdict = if current_response != *response {
288            PostInvocationVerdict::Redact(current_response)
289        } else if !escalations.is_empty() {
290            PostInvocationVerdict::Escalate(escalations.join("; "))
291        } else {
292            PostInvocationVerdict::Allow
293        };
294        Ok(DurablePipelineOutcome {
295            outcome: PipelineOutcome {
296                verdict,
297                escalations,
298                evidence,
299            },
300            step_results,
301        })
302    }
303
304    #[must_use]
305    pub fn evaluate_with_evidence(&self, tool_name: &str, response: &Value) -> PipelineOutcome {
306        let context = PostInvocationContext::synthetic(tool_name);
307        self.evaluate_with_context_and_evidence(&context, response)
308    }
309
310    #[must_use]
311    pub fn evaluate_with_context_and_evidence(
312        &self,
313        context: &PostInvocationContext<'_>,
314        response: &Value,
315    ) -> PipelineOutcome {
316        let _evaluation = self.lock_evaluation();
317        let mut current_response = response.clone();
318        let mut escalations = Vec::new();
319        let mut evidence = Vec::new();
320
321        for hook in &self.hooks {
322            let verdict = hook.inspect(context, &current_response);
323            if let Some(ev) = hook.take_evidence() {
324                evidence.push(ev);
325            }
326            match verdict {
327                PostInvocationVerdict::Allow => continue,
328                PostInvocationVerdict::Block(reason) => {
329                    return PipelineOutcome {
330                        verdict: PostInvocationVerdict::Block(reason),
331                        escalations,
332                        evidence,
333                    };
334                }
335                PostInvocationVerdict::Redact(redacted) => {
336                    current_response = redacted;
337                }
338                PostInvocationVerdict::Escalate(message) => {
339                    escalations.push(message);
340                }
341            }
342        }
343
344        let verdict = if current_response != *response {
345            PostInvocationVerdict::Redact(current_response)
346        } else if !escalations.is_empty() {
347            PostInvocationVerdict::Escalate(escalations.join("; "))
348        } else {
349            PostInvocationVerdict::Allow
350        };
351        PipelineOutcome {
352            verdict,
353            escalations,
354            evidence,
355        }
356    }
357
358    fn lock_evaluation(&self) -> MutexGuard<'_, ()> {
359        match self.evaluation_lock.lock() {
360            Ok(guard) => guard,
361            Err(poisoned) => poisoned.into_inner(),
362        }
363    }
364
365    #[must_use]
366    pub fn evaluate(
367        &self,
368        tool_name: &str,
369        response: &Value,
370    ) -> (PostInvocationVerdict, Vec<String>) {
371        let outcome = self.evaluate_with_evidence(tool_name, response);
372        (outcome.verdict, outcome.escalations)
373    }
374}
375
376impl Default for PostInvocationPipeline {
377    fn default() -> Self {
378        Self::new()
379    }
380}
381
382#[cfg(test)]
383mod tests {
384    use super::*;
385    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
386    use std::sync::{Arc, Barrier};
387    use std::thread;
388    use std::time::Duration;
389
390    struct OverlapDetectingHook {
391        active: Arc<AtomicUsize>,
392        overlapped: Arc<AtomicBool>,
393    }
394
395    impl PostInvocationHook for OverlapDetectingHook {
396        fn name(&self) -> &str {
397            "overlap-detecting-hook"
398        }
399
400        fn inspect(
401            &self,
402            _ctx: &PostInvocationContext<'_>,
403            _response: &Value,
404        ) -> PostInvocationVerdict {
405            if self.active.fetch_add(1, Ordering::SeqCst) != 0 {
406                self.overlapped.store(true, Ordering::SeqCst);
407            }
408            thread::sleep(Duration::from_millis(25));
409            self.active.fetch_sub(1, Ordering::SeqCst);
410            PostInvocationVerdict::Allow
411        }
412    }
413
414    #[test]
415    fn pipeline_serializes_hook_evaluation_and_evidence_side_channels() {
416        let active = Arc::new(AtomicUsize::new(0));
417        let overlapped = Arc::new(AtomicBool::new(false));
418        let mut pipeline = PostInvocationPipeline::new();
419        pipeline.add(Box::new(OverlapDetectingHook {
420            active,
421            overlapped: overlapped.clone(),
422        }));
423        let pipeline = Arc::new(pipeline);
424        let start = Arc::new(Barrier::new(3));
425        let workers = (0..2)
426            .map(|index| {
427                let pipeline = pipeline.clone();
428                let start = start.clone();
429                thread::spawn(move || {
430                    start.wait();
431                    let _ = pipeline
432                        .evaluate_with_evidence("tool", &serde_json::json!({"index": index}));
433                })
434            })
435            .collect::<Vec<_>>();
436        start.wait();
437        for worker in workers {
438            worker.join().expect("post-invocation worker");
439        }
440
441        assert!(!overlapped.load(Ordering::SeqCst));
442    }
443}
chio_kernel/post_invocation.rs

chio_kernel/
post_invocation.rs
