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

systemprompt_api/services/gateway/
signature_cache.rs
