llm_tool/types/
surrogate.rs

1//! Host-pinned surrogate token vault (`<surrogate:kind:id>`) for prompt-invisible
2//! secret/capability binding and multi-turn continuity verification.
3//!
4//! Inspired by Meta AI's Muse multi-turn continuity & capability-binding architecture:
5//! sensitive credentials, session handles, or capability tokens are replaced with
6//! opaque `<surrogate:kind:id>` handles before tool outputs enter the LLM context
7//! window, and resolved back to raw values on the host boundary before tool execution.
8//! Any forged or unknown `<surrogate:...>` token in model-supplied arguments fails
9//! closed with a [`ToolError`].
10
11use alloc::{
12    format,
13    string::{String, ToString},
14    sync::Arc,
15    vec::Vec,
16};
17
18use serde::{Deserialize, Serialize};
19
20use super::ToolError;
21use crate::compat::{HashMap, RwLock, read_lock, write_lock};
22
23/// Prefix of every serialized surrogate token.
24pub const SURROGATE_TOKEN_PREFIX: &str = "<surrogate:";
25
26/// Suffix of every serialized surrogate token.
27pub const SURROGATE_TOKEN_SUFFIX: &str = ">";
28
29/// Strongly-typed opaque surrogate handle minted by a [`SurrogateVault`].
30#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
31pub struct SurrogateHandle {
32    kind: String,
33    id: u64,
34    token: String,
35}
36
37impl SurrogateHandle {
38    /// Construct a handle for `(kind, id)`.
39    #[must_use]
40    fn new(kind: String, id: u64) -> Self {
41        let token = format!("{SURROGATE_TOKEN_PREFIX}{kind}:{id}{SURROGATE_TOKEN_SUFFIX}");
42        Self { kind, id, token }
43    }
44
45    /// Return the namespace/kind of this surrogate handle (e.g. `"api_key"`, `"file_cap"`).
46    #[must_use]
47    pub fn kind(&self) -> &str {
48        &self.kind
49    }
50
51    /// Return the numeric ID assigned within the vault.
52    #[must_use]
53    pub const fn id(&self) -> u64 {
54        self.id
55    }
56
57    /// Return the wire token string (`"<surrogate:kind:id>"`) shown to the model.
58    #[must_use]
59    pub fn token(&self) -> &str {
60        &self.token
61    }
62}
63
64impl core::fmt::Display for SurrogateHandle {
65    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
66        f.write_str(&self.token)
67    }
68}
69
70#[derive(Default)]
71struct SurrogateVaultInner {
72    next_id: u64,
73    /// Maps `"<surrogate:kind:id>"` -> raw value.
74    token_to_raw: HashMap<String, String>,
75    /// Maps `(kind, raw_value)` -> `SurrogateHandle` for deterministic deduplication.
76    raw_to_handle: HashMap<(String, String), SurrogateHandle>,
77    /// Pre-sorted `(raw_value, token)` pairs (longest `raw_value` first) for `redact_str`.
78    sorted_pairs: Vec<(String, String)>,
79}
80
81/// Thread-safe, host-pinned vault mapping opaque `<surrogate:kind:id>` tokens to raw values.
82#[derive(Clone, Default)]
83pub struct SurrogateVault(Arc<RwLock<SurrogateVaultInner>>);
84
85impl core::fmt::Debug for SurrogateVault {
86    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
87        let count = match read_lock(&self.0) {
88            Ok(g) => g.token_to_raw.len(),
89            Err(e) => {
90                tracing::warn!(error = %e, "SurrogateVault::fmt: lock poisoned");
91                0
92            }
93        };
94        f.debug_struct("SurrogateVault")
95            .field("entry_count", &count)
96            .finish()
97    }
98}
99
100fn is_valid_surrogate_kind(kind: &str) -> bool {
101    !kind.is_empty()
102        && kind
103            .bytes()
104            .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-')
105}
106
107impl SurrogateVault {
108    /// Create a new, empty [`SurrogateVault`].
109    #[must_use]
110    pub fn new() -> Self {
111        Self::default()
112    }
113
114    /// Mint or retrieve an existing [`SurrogateHandle`] for `(kind, raw_value)`.
115    ///
116    /// # Errors
117    ///
118    /// Returns [`ToolError`] if `kind` is empty or contains characters outside `[a-zA-Z0-9_-]`,
119    /// `raw_value` is empty, or the vault lock is poisoned.
120    pub fn mint(
121        &self,
122        kind: &str,
123        raw_value: impl Into<String>,
124    ) -> Result<SurrogateHandle, ToolError> {
125        if !is_valid_surrogate_kind(kind) {
126            return Err(ToolError::new(format!(
127                "SurrogateVault::mint: invalid kind '{kind}' (must match [a-zA-Z0-9_-]+)"
128            )));
129        }
130        let raw = raw_value.into();
131        if raw.is_empty() {
132            return Err(ToolError::new(
133                "SurrogateVault::mint: raw_value must not be empty",
134            ));
135        }
136        let mut guard = write_lock(&self.0).map_err(|e| {
137            let msg = format!("SurrogateVault::mint: lock poisoned: {e}");
138            tracing::warn!("{msg}");
139            ToolError::new(msg)
140        })?;
141        let key = (kind.to_string(), raw.clone());
142        if let Some(existing) = guard.raw_to_handle.get(&key) {
143            return Ok(existing.clone());
144        }
145        guard.next_id += 1;
146        let id = guard.next_id;
147        let handle = SurrogateHandle::new(kind.to_string(), id);
148        let insert_idx = guard
149            .sorted_pairs
150            .binary_search_by(|(probe_raw, _)| {
151                raw.len()
152                    .cmp(&probe_raw.len())
153                    .then_with(|| probe_raw.as_str().cmp(raw.as_str()))
154            })
155            .unwrap_or_else(|idx| idx);
156        guard
157            .sorted_pairs
158            .insert(insert_idx, (raw.clone(), handle.token.clone()));
159        guard.token_to_raw.insert(handle.token.clone(), raw);
160        guard.raw_to_handle.insert(key, handle.clone());
161        drop(guard);
162        Ok(handle)
163    }
164
165    /// Resolve an exact `<surrogate:kind:id>` token to its bound raw value.
166    ///
167    /// Fails closed with [`ToolError`] if the token was not minted in this vault.
168    ///
169    /// # Errors
170    ///
171    /// Returns [`ToolError`] if `token` is unknown or the vault lock is poisoned.
172    pub fn resolve_token(&self, token: &str) -> Result<String, ToolError> {
173        let guard = read_lock(&self.0).map_err(|e| {
174            let msg = format!("SurrogateVault::resolve_token: lock poisoned: {e}");
175            tracing::warn!("{msg}");
176            ToolError::new(msg)
177        })?;
178        guard.token_to_raw.get(token).cloned().ok_or_else(|| {
179            ToolError::new(format!(
180                "unknown or forged surrogate token '{token}' rejected by continuity check"
181            ))
182        })
183    }
184
185    /// Resolve all `<surrogate:kind:id>` tokens embedded in `input`.
186    ///
187    /// Returns `Ok((resolved_string, replacement_count))` or fails closed with
188    /// [`ToolError`] if any `<surrogate:...>` token in `input` is malformed or unknown.
189    ///
190    /// # Errors
191    ///
192    /// Returns [`ToolError`] if an unknown/malformed surrogate token is encountered
193    /// or the vault lock is poisoned.
194    pub fn resolve_in_str_counted(&self, input: &str) -> Result<(String, usize), ToolError> {
195        if !input.contains(SURROGATE_TOKEN_PREFIX) {
196            return Ok((input.to_string(), 0));
197        }
198        let guard = read_lock(&self.0).map_err(|e| {
199            let msg = format!("SurrogateVault::resolve_in_str: lock poisoned: {e}");
200            tracing::warn!("{msg}");
201            ToolError::new(msg)
202        })?;
203
204        let mut out = String::with_capacity(input.len());
205        let mut rest = input;
206        let mut count = 0usize;
207
208        while let Some(start) = rest.find(SURROGATE_TOKEN_PREFIX) {
209            out.push_str(&rest[..start]);
210            let after_prefix = &rest[start..];
211            let Some(end_rel) = after_prefix.find(SURROGATE_TOKEN_SUFFIX) else {
212                return Err(ToolError::new(
213                    "unclosed '<surrogate:' token in tool input rejected by continuity check",
214                ));
215            };
216            let candidate = &after_prefix[..=end_rel];
217            let Some(raw) = guard.token_to_raw.get(candidate) else {
218                return Err(ToolError::new(format!(
219                    "unknown or forged surrogate token '{candidate}' rejected by continuity check"
220                )));
221            };
222            out.push_str(raw);
223            count += 1;
224            rest = &after_prefix[end_rel + 1..];
225        }
226        drop(guard);
227        out.push_str(rest);
228        Ok((out, count))
229    }
230
231    /// Resolve all `<surrogate:kind:id>` tokens embedded in `input`.
232    ///
233    /// # Errors
234    ///
235    /// Returns [`ToolError`] if an unknown/malformed surrogate token is encountered
236    /// or the vault lock is poisoned.
237    pub fn resolve_in_str(&self, input: &str) -> Result<String, ToolError> {
238        self.resolve_in_str_counted(input).map(|(s, _)| s)
239    }
240
241    /// Recursively resolve all `<surrogate:kind:id>` tokens across string values inside `value`.
242    ///
243    /// Returns the total number of resolved tokens. Fails closed with [`ToolError`] if any
244    /// surrogate token is unknown or malformed.
245    ///
246    /// # Errors
247    ///
248    /// Returns [`ToolError`] if an unknown/malformed surrogate token is encountered
249    /// or the vault lock is poisoned.
250    pub fn resolve_in_json(&self, value: &mut serde_json::Value) -> Result<usize, ToolError> {
251        match value {
252            serde_json::Value::String(s) => {
253                let (resolved, count) = self.resolve_in_str_counted(s)?;
254                if count > 0 {
255                    *s = resolved;
256                }
257                Ok(count)
258            }
259            serde_json::Value::Array(items) => {
260                let mut total = 0usize;
261                for item in items {
262                    total += self.resolve_in_json(item)?;
263                }
264                Ok(total)
265            }
266            serde_json::Value::Object(map) => {
267                let mut total = 0usize;
268                for val in map.values_mut() {
269                    total += self.resolve_in_json(val)?;
270                }
271                Ok(total)
272            }
273            serde_json::Value::Null | serde_json::Value::Bool(_) | serde_json::Value::Number(_) => {
274                Ok(0)
275            }
276        }
277    }
278
279    /// Replace any registered raw values occurring in `output` with their `<surrogate:kind:id>`
280    /// tokens (longest raw values first so overlapping prefixes cannot leak partial secrets).
281    ///
282    /// # Errors
283    ///
284    /// Returns [`ToolError`] if the vault lock is poisoned.
285    pub fn redact_str(&self, output: &str) -> Result<String, ToolError> {
286        if output.is_empty() {
287            return Ok(String::new());
288        }
289        let guard = read_lock(&self.0).map_err(|e| {
290            let msg = format!("SurrogateVault::redact_str: lock poisoned: {e}");
291            tracing::warn!("{msg}");
292            ToolError::new(msg)
293        })?;
294        if guard.sorted_pairs.is_empty() {
295            return Ok(output.to_string());
296        }
297        let matching_pairs: Vec<(String, String)> = guard
298            .sorted_pairs
299            .iter()
300            .filter(|(raw, _)| output.contains(raw.as_str()))
301            .cloned()
302            .collect();
303        drop(guard);
304        if matching_pairs.is_empty() {
305            return Ok(output.to_string());
306        }
307
308        let mut redacted = output.to_string();
309        for (raw, token) in &matching_pairs {
310            if redacted.contains(raw.as_str()) {
311                redacted = redacted.replace(raw.as_str(), token.as_str());
312            }
313        }
314        Ok(redacted)
315    }
316}
317
318#[cfg(test)]
319mod tests {
320    use super::*;
321    use crate::ToolContext;
322
323    #[test]
324    fn surrogate_vault_mint_deduplicate_resolve_and_redact() {
325        let ctx = ToolContext::new();
326        let h1 = ctx.mint_surrogate("api_key", "sk-live-secret-123").unwrap();
327        let h2 = ctx.mint_surrogate("api_key", "sk-live-secret-123").unwrap();
328        assert_eq!(h1, h2);
329        assert_eq!(h1.token(), "<surrogate:api_key:1>");
330        assert_eq!(h1.kind(), "api_key");
331        assert_eq!(h1.id(), 1);
332
333        let redacted = ctx
334            .redact_surrogates("Authorization: Bearer sk-live-secret-123")
335            .unwrap();
336        assert_eq!(redacted, "Authorization: Bearer <surrogate:api_key:1>");
337
338        let resolved = ctx.resolve_surrogates(&redacted).unwrap();
339        assert_eq!(resolved, "Authorization: Bearer sk-live-secret-123");
340
341        let mut json_val = serde_json::json!({
342            "headers": ["Bearer <surrogate:api_key:1>"],
343            "count": 1
344        });
345        let count = ctx
346            .surrogate_vault()
347            .unwrap()
348            .resolve_in_json(&mut json_val)
349            .unwrap();
350        assert_eq!(count, 1);
351        assert_eq!(json_val["headers"][0], "Bearer sk-live-secret-123");
352    }
353
354    #[test]
355    fn surrogate_vault_rejects_forged_or_unclosed_tokens() {
356        let ctx = ToolContext::new();
357        ctx.mint_surrogate("cap", "real-cap").unwrap();
358
359        let forged = ctx
360            .resolve_surrogates("use <surrogate:cap:999>")
361            .unwrap_err();
362        assert!(
363            forged
364                .to_string()
365                .contains("unknown or forged surrogate token")
366        );
367
368        let unclosed = ctx.resolve_surrogates("use <surrogate:cap:1").unwrap_err();
369        assert!(unclosed.to_string().contains("unclosed '<surrogate:'"));
370    }
371
372    #[test]
373    fn shared_surrogate_vault_across_contexts_and_validation_rules() {
374        let vault = SurrogateVault::new();
375        assert!(vault.mint("", "secret").is_err());
376        assert!(vault.mint("bad:kind", "secret").is_err());
377        assert!(vault.mint("api_key", "").is_err());
378
379        // Mint overlapping prefix secrets to verify longest-match-first redaction.
380        let short = vault.mint("api_key", "sk-live").unwrap();
381        let long = vault.mint("api_key", "sk-live-extended-999").unwrap();
382
383        let parent_ctx = ToolContext::new().with_surrogate_vault(vault.clone());
384        let child_ctx = ToolContext::new().with_surrogate_vault(vault);
385
386        let redacted = parent_ctx
387            .redact_surrogates("keys: sk-live-extended-999 and sk-live")
388            .unwrap();
389        assert_eq!(
390            redacted,
391            format!("keys: {} and {}", long.token(), short.token())
392        );
393
394        let resolved = child_ctx.resolve_surrogates(&redacted).unwrap();
395        assert_eq!(resolved, "keys: sk-live-extended-999 and sk-live");
396    }
397}
llm_tool/types/surrogate.rs

llm_tool/types/
surrogate.rs
