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//! The cache is process-local, so every miss is a real failure mode: a cold
16//! process after a restart, or a request routed to a replica that did not serve
17//! the prior turn, leaves the block unsigned. `thought_signature` is omitted
18//! from the outbound wire when absent, and Gemini then rejects the turn — so
19//! misses are counted under `gateway_signature_hydration_total` and warned,
20//! but only when the resolved upstream is [`WireProtocol::Gemini`]; for every
21//! other wire the absent signature is expected and carries no signal.
22//!
23//! Copyright (c) systemprompt.io — Business Source License 1.1.
24//! See <https://systemprompt.io> for licensing details.
25
26use std::collections::HashMap;
27use std::sync::{Mutex, OnceLock, PoisonError};
28use std::time::{Duration, Instant};
29
30use systemprompt_identifiers::GatewayConversationId;
31use systemprompt_models::profile::WireProtocol;
32use systemprompt_models::wire::canonical::{CanonicalContent, CanonicalRequest, CanonicalResponse};
33
34const TTL: Duration = Duration::from_hours(1);
35const MAX_ENTRIES: usize = 10_000;
36const HYDRATION_TOTAL: &str = "gateway_signature_hydration_total";
37const CAPTURE_SKIPPED_TOTAL: &str = "gateway_signature_capture_skipped_total";
38
39struct Entry {
40    signature: String,
41    expires_at: Instant,
42}
43
44type Key = (GatewayConversationId, String);
45
46pub struct ThoughtSignatureCache {
47    entries: Mutex<HashMap<Key, Entry>>,
48    ttl: Duration,
49    max_entries: usize,
50}
51
52impl std::fmt::Debug for ThoughtSignatureCache {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        f.debug_struct("ThoughtSignatureCache")
55            .field("ttl", &self.ttl)
56            .field("max_entries", &self.max_entries)
57            .finish_non_exhaustive()
58    }
59}
60
61impl ThoughtSignatureCache {
62    pub fn global() -> &'static Self {
63        static CACHE: OnceLock<ThoughtSignatureCache> = OnceLock::new();
64        CACHE.get_or_init(|| Self::new(TTL, MAX_ENTRIES))
65    }
66
67    #[must_use]
68    pub fn new(ttl: Duration, max_entries: usize) -> Self {
69        Self {
70            entries: Mutex::new(HashMap::new()),
71            ttl,
72            max_entries,
73        }
74    }
75
76    pub fn store(&self, conversation: &GatewayConversationId, tool_use_id: &str, signature: &str) {
77        let key = (conversation.clone(), tool_use_id.to_owned());
78        let mut entries = self.entries.lock().unwrap_or_else(PoisonError::into_inner);
79        let now = Instant::now();
80        if !entries.contains_key(&key) && entries.len() >= self.max_entries {
81            entries.retain(|_, e| e.expires_at > now);
82            if entries.len() >= self.max_entries
83                && let Some(oldest) = entries
84                    .iter()
85                    .min_by_key(|(_, e)| e.expires_at)
86                    .map(|(k, _)| k.clone())
87            {
88                entries.remove(&oldest);
89            }
90        }
91        entries.insert(
92            key,
93            Entry {
94                signature: signature.to_owned(),
95                expires_at: now + self.ttl,
96            },
97        );
98    }
99
100    pub fn lookup(
101        &self,
102        conversation: &GatewayConversationId,
103        tool_use_id: &str,
104    ) -> Option<String> {
105        let key = (conversation.clone(), tool_use_id.to_owned());
106        let now = Instant::now();
107        let mut entries = self.entries.lock().unwrap_or_else(PoisonError::into_inner);
108        let entry = entries.get_mut(&key)?;
109        if entry.expires_at <= now {
110            entries.remove(&key);
111            return None;
112        }
113        entry.expires_at = now + self.ttl;
114        let signature = entry.signature.clone();
115        drop(entries);
116        Some(signature)
117    }
118
119    pub fn store_from_response(
120        &self,
121        conversation: &GatewayConversationId,
122        response: &CanonicalResponse,
123    ) {
124        for content in &response.content {
125            if let CanonicalContent::ToolUse {
126                id,
127                signature: Some(signature),
128                ..
129            } = content
130            {
131                self.store(conversation, id, signature);
132            }
133        }
134    }
135
136    pub fn hydrate_request(
137        &self,
138        conversation: &GatewayConversationId,
139        request: &mut CanonicalRequest,
140        wire: Option<WireProtocol>,
141    ) {
142        let model = request.model.clone();
143        let signatures_required = wire == Some(WireProtocol::Gemini);
144        for message in &mut request.messages {
145            for content in &mut message.content {
146                let CanonicalContent::ToolUse { id, signature, .. } = content else {
147                    continue;
148                };
149                match signature {
150                    Some(sig) => self.store(conversation, id, sig),
151                    None => match self.lookup(conversation, id) {
152                        Some(cached) => {
153                            *signature = Some(cached);
154                            if signatures_required {
155                                metrics::counter!(HYDRATION_TOTAL, "outcome" => "hit").increment(1);
156                            }
157                        },
158                        None => {
159                            if signatures_required {
160                                metrics::counter!(HYDRATION_TOTAL, "outcome" => "miss")
161                                    .increment(1);
162                                tracing::warn!(
163                                    conversation = %conversation,
164                                    tool_use_id = %id,
165                                    model = %model,
166                                    "no cached thought signature for tool_use; upstream may reject the turn"
167                                );
168                            }
169                        },
170                    },
171                }
172            }
173        }
174    }
175
176    #[cfg(feature = "test-api")]
177    #[expect(
178        clippy::panic,
179        reason = "test-only seam, compiled out unless `test-api` is enabled"
180    )]
181    pub fn poison_lock(&self) {
182        let _guard = self.entries.lock().unwrap_or_else(PoisonError::into_inner);
183        panic!("poisoning the signature cache lock");
184    }
185
186    #[must_use]
187    pub fn signed_tool_use_count(response: &CanonicalResponse) -> usize {
188        response
189            .content
190            .iter()
191            .filter(|c| {
192                matches!(
193                    c,
194                    CanonicalContent::ToolUse {
195                        signature: Some(_),
196                        ..
197                    }
198                )
199            })
200            .count()
201    }
202
203    pub fn note_uncacheable_response(response: &CanonicalResponse, reason: &'static str) {
204        let signed = Self::signed_tool_use_count(response);
205        if signed == 0 {
206            return;
207        }
208        metrics::counter!(CAPTURE_SKIPPED_TOTAL, "reason" => reason).increment(1);
209        tracing::warn!(
210            reason,
211            signed_tool_use_blocks = signed,
212            "thought signatures could not be cached; a later turn in this conversation will miss"
213        );
214    }
215}
systemprompt_api/services/gateway/signature_cache.rs

systemprompt_api/services/gateway/
signature_cache.rs
