Skip to main content

typesec_agent/interop/
taint.rs

1//! Taint guarded tool output: results enter the prompt as labeled data.
2//!
3//! Tool output flowing back into the model's context is the classic
4//! prompt-injection channel. [`GuardedToolCall::protect_output`] closes the
5//! loop with `typesec-core`'s information-flow machinery: the output of an
6//! *allowed* call is wrapped in a [`SecureValue`] labeled at a caller-chosen
7//! privacy level and tied to the call's resolved resource id. From there the
8//! usual rules apply — the value can be transformed (`map`/`zip`) but
9//! revealing or declassifying it requires a typed capability minted for that
10//! same resource.
11
12use thiserror::Error;
13use typesec_core::SecureValue;
14use typesec_core::resource::GenericResource;
15use typesec_core::secure_value::PrivacyLevel;
16
17use super::call::{GuardedToolCall, ToolCallVerdict};
18
19/// Resource kind attached to tool-output resources.
20pub const TOOL_OUTPUT_KIND: &str = "tool_output";
21
22/// Why tool output could not be protected.
23#[derive(Debug, Error)]
24pub enum TaintError {
25    /// Only allowed calls have output worth labeling; refusing here keeps
26    /// "we ran a denied tool anyway" from being papered over.
27    #[error("cannot protect output of a call that was not allowed: {reason}")]
28    NotAllowed {
29        /// The verdict's rationale.
30        reason: String,
31    },
32    /// The call never resolved a resource (unbound tool or failed
33    /// resource-argument resolution), so there is nothing to tie the label to.
34    #[error("call has no resolved resource to tie the output to")]
35    UnresolvedResource,
36}
37
38impl GuardedToolCall {
39    /// Label the output of an allowed call as protected data tied to the
40    /// call's resolved resource.
41    ///
42    /// The label `L` is chosen by the caller (e.g.
43    /// [`Sensitive`](typesec_core::secure_value::Sensitive)); revealing the
44    /// value downstream requires a `Capability<CanReadSensitive, GenericResource>`
45    /// minted for the *same resource id* this call was checked against.
46    pub fn protect_output<L: PrivacyLevel, T>(
47        &self,
48        output: T,
49    ) -> Result<SecureValue<L, T, GenericResource>, TaintError> {
50        if let ToolCallVerdict::Deny { reason } | ToolCallVerdict::Delegate { reason } =
51            &self.verdict
52        {
53            return Err(TaintError::NotAllowed {
54                reason: reason.clone(),
55            });
56        }
57        let resource_id = self
58            .resource
59            .as_deref()
60            .ok_or(TaintError::UnresolvedResource)?;
61        let resource = GenericResource::new(resource_id, TOOL_OUTPUT_KIND);
62        Ok(SecureValue::protect(output, &resource))
63    }
64}