systemprompt_api/services/gateway/
signature_cache.rs

1//! Server-side cache of Gemini `thoughtSignature` values keyed by
2//! conversation and `tool_use` id.
3//!
4//! Gemini attaches an opaque signature to each function call that must be
5//! echoed back verbatim on the next turn. The gateway forwards it to
6//! Anthropic-protocol clients as a non-standard `signature` field on
7//! `tool_use` blocks, but strict clients drop unknown fields when replaying
8//! history. This cache captures signatures as responses pass through and
9//! re-injects them on inbound requests whose `tool_use` blocks arrive without
10//! one, so any faithful Anthropic client works against Gemini upstreams.
11//! Keys are scoped by [`GatewayConversationId`] because `tool_use` ids on
12//! inbound requests are client-supplied: without the scope, a caller could
13//! read another conversation's cached signatures by guessing ids.
14//!
15//! Copyright (c) systemprompt.io — Business Source License 1.1.
16//! See <https://systemprompt.io> for licensing details.
17
18use std::collections::HashMap;
19use std::sync::{Mutex, OnceLock};
20use std::time::{Duration, Instant};
21
22use systemprompt_identifiers::GatewayConversationId;
23use systemprompt_models::wire::canonical::{CanonicalContent, CanonicalRequest, CanonicalResponse};
24
25const TTL: Duration = Duration::from_hours(1);
26const MAX_ENTRIES: usize = 10_000;
27
28struct Entry {
29    signature: String,
30    expires_at: Instant,
31}
32
33type Key = (GatewayConversationId, String);
34
35pub struct ThoughtSignatureCache {
36    entries: Mutex<HashMap<Key, Entry>>,
37    ttl: Duration,
38    max_entries: usize,
39}
40
41impl std::fmt::Debug for ThoughtSignatureCache {
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        f.debug_struct("ThoughtSignatureCache")
44            .field("ttl", &self.ttl)
45            .field("max_entries", &self.max_entries)
46            .finish_non_exhaustive()
47    }
48}
49
50impl ThoughtSignatureCache {
51    pub fn global() -> &'static Self {
52        static CACHE: OnceLock<ThoughtSignatureCache> = OnceLock::new();
53        CACHE.get_or_init(|| Self::new(TTL, MAX_ENTRIES))
54    }
55
56    #[must_use]
57    pub fn new(ttl: Duration, max_entries: usize) -> Self {
58        Self {
59            entries: Mutex::new(HashMap::new()),
60            ttl,
61            max_entries,
62        }
63    }
64
65    pub fn store(&self, conversation: &GatewayConversationId, tool_use_id: &str, signature: &str) {
66        let key = (conversation.clone(), tool_use_id.to_owned());
67        let Ok(mut entries) = self.entries.lock() else {
68            return;
69        };
70        let now = Instant::now();
71        if !entries.contains_key(&key) && entries.len() >= self.max_entries {
72            entries.retain(|_, e| e.expires_at > now);
73            if entries.len() >= self.max_entries
74                && let Some(oldest) = entries
75                    .iter()
76                    .min_by_key(|(_, e)| e.expires_at)
77                    .map(|(k, _)| k.clone())
78            {
79                entries.remove(&oldest);
80            }
81        }
82        entries.insert(
83            key,
84            Entry {
85                signature: signature.to_owned(),
86                expires_at: now + self.ttl,
87            },
88        );
89    }
90
91    pub fn lookup(
92        &self,
93        conversation: &GatewayConversationId,
94        tool_use_id: &str,
95    ) -> Option<String> {
96        let key = (conversation.clone(), tool_use_id.to_owned());
97        let now = Instant::now();
98        let mut entries = self.entries.lock().ok()?;
99        let entry = entries.get_mut(&key)?;
100        if entry.expires_at <= now {
101            entries.remove(&key);
102            return None;
103        }
104        entry.expires_at = now + self.ttl;
105        let signature = entry.signature.clone();
106        drop(entries);
107        Some(signature)
108    }
109
110    pub fn store_from_response(
111        &self,
112        conversation: &GatewayConversationId,
113        response: &CanonicalResponse,
114    ) {
115        for content in &response.content {
116            if let CanonicalContent::ToolUse {
117                id,
118                signature: Some(signature),
119                ..
120            } = content
121            {
122                self.store(conversation, id, signature);
123            }
124        }
125    }
126
127    pub fn hydrate_request(
128        &self,
129        conversation: &GatewayConversationId,
130        request: &mut CanonicalRequest,
131    ) {
132        for message in &mut request.messages {
133            for content in &mut message.content {
134                let CanonicalContent::ToolUse { id, signature, .. } = content else {
135                    continue;
136                };
137                match signature {
138                    Some(sig) => self.store(conversation, id, sig),
139                    None => {
140                        if let Some(cached) = self.lookup(conversation, id) {
141                            tracing::debug!(
142                                tool_use_id = %id,
143                                "re-injected cached thought signature into tool_use block"
144                            );
145                            *signature = Some(cached);
146                        }
147                    },
148                }
149            }
150        }
151    }
152}
systemprompt_api/services/gateway/signature_cache.rs

systemprompt_api/services/gateway/
signature_cache.rs
