Skip to main content

chio_kernel/
post_invocation.rs

1//! Post-invocation hook pipeline executed after a tool returns output.
2
3use chio_core::receipt::GuardEvidence;
4use chio_core::{AgentId, ChioScope, ServerId};
5use serde_json::Value;
6
7use crate::runtime::ToolCallRequest;
8
9/// Verdict from a post-invocation hook.
10#[derive(Debug, Clone)]
11pub enum PostInvocationVerdict {
12    Allow,
13    Block(String),
14    Redact(Value),
15    Escalate(String),
16}
17
18/// Context available to post-invocation hooks after a tool has executed.
19#[derive(Clone, Copy, Debug)]
20pub struct PostInvocationContext<'a> {
21    pub tool_name: &'a str,
22    pub request: Option<&'a ToolCallRequest>,
23    pub scope: Option<&'a ChioScope>,
24    pub agent_id: Option<&'a AgentId>,
25    pub server_id: Option<&'a ServerId>,
26    pub matched_grant_index: Option<usize>,
27}
28
29impl<'a> PostInvocationContext<'a> {
30    #[must_use]
31    pub fn synthetic(tool_name: &'a str) -> Self {
32        Self {
33            tool_name,
34            request: None,
35            scope: None,
36            agent_id: None,
37            server_id: None,
38            matched_grant_index: None,
39        }
40    }
41
42    #[must_use]
43    pub fn from_request(request: &'a ToolCallRequest, matched_grant_index: Option<usize>) -> Self {
44        Self {
45            tool_name: request.tool_name.as_str(),
46            request: Some(request),
47            scope: Some(&request.capability.scope),
48            agent_id: Some(&request.agent_id),
49            server_id: Some(&request.server_id),
50            matched_grant_index,
51        }
52    }
53}
54
55/// A hook that inspects tool responses after invocation.
56pub trait PostInvocationHook: Send + Sync {
57    fn name(&self) -> &str;
58
59    fn inspect(&self, ctx: &PostInvocationContext<'_>, response: &Value) -> PostInvocationVerdict;
60
61    fn take_evidence(&self) -> Option<GuardEvidence> {
62        None
63    }
64}
65
66/// Outcome of running the pipeline.
67#[derive(Debug, Clone)]
68pub struct PipelineOutcome {
69    pub verdict: PostInvocationVerdict,
70    pub escalations: Vec<String>,
71    pub evidence: Vec<GuardEvidence>,
72}
73
74/// Pipeline of post-invocation hooks evaluated in registration order.
75pub struct PostInvocationPipeline {
76    hooks: Vec<Box<dyn PostInvocationHook>>,
77}
78
79impl PostInvocationPipeline {
80    #[must_use]
81    pub fn new() -> Self {
82        Self { hooks: Vec::new() }
83    }
84
85    pub fn add(&mut self, hook: Box<dyn PostInvocationHook>) {
86        self.hooks.push(hook);
87    }
88
89    #[must_use]
90    pub fn len(&self) -> usize {
91        self.hooks.len()
92    }
93
94    #[must_use]
95    pub fn is_empty(&self) -> bool {
96        self.hooks.is_empty()
97    }
98
99    #[must_use]
100    pub fn evaluate_with_evidence(&self, tool_name: &str, response: &Value) -> PipelineOutcome {
101        let context = PostInvocationContext::synthetic(tool_name);
102        self.evaluate_with_context_and_evidence(&context, response)
103    }
104
105    #[must_use]
106    pub fn evaluate_with_context_and_evidence(
107        &self,
108        context: &PostInvocationContext<'_>,
109        response: &Value,
110    ) -> PipelineOutcome {
111        let mut current_response = response.clone();
112        let mut escalations = Vec::new();
113        let mut evidence = Vec::new();
114
115        for hook in &self.hooks {
116            let verdict = hook.inspect(context, &current_response);
117            if let Some(ev) = hook.take_evidence() {
118                evidence.push(ev);
119            }
120            match verdict {
121                PostInvocationVerdict::Allow => continue,
122                PostInvocationVerdict::Block(reason) => {
123                    return PipelineOutcome {
124                        verdict: PostInvocationVerdict::Block(reason),
125                        escalations,
126                        evidence,
127                    };
128                }
129                PostInvocationVerdict::Redact(redacted) => {
130                    current_response = redacted;
131                }
132                PostInvocationVerdict::Escalate(message) => {
133                    escalations.push(message);
134                }
135            }
136        }
137
138        let verdict = if current_response != *response {
139            PostInvocationVerdict::Redact(current_response)
140        } else if !escalations.is_empty() {
141            PostInvocationVerdict::Escalate(escalations.join("; "))
142        } else {
143            PostInvocationVerdict::Allow
144        };
145        PipelineOutcome {
146            verdict,
147            escalations,
148            evidence,
149        }
150    }
151
152    #[must_use]
153    pub fn evaluate(
154        &self,
155        tool_name: &str,
156        response: &Value,
157    ) -> (PostInvocationVerdict, Vec<String>) {
158        let outcome = self.evaluate_with_evidence(tool_name, response);
159        (outcome.verdict, outcome.escalations)
160    }
161}
162
163impl Default for PostInvocationPipeline {
164    fn default() -> Self {
165        Self::new()
166    }
167}