rig_core/providers/gemini/
cached_content.rs

1//! Gemini explicit context caching through the `cachedContents` resource.
2//! Upload content once and reuse its handle across requests. Storage is billed
3//! until the cache expires or is deleted.
4//!
5//! Requests using a cache must not supply their own system instruction, tools,
6//! or tool configuration. Cached function declarations require a caller-managed
7//! tool loop; provider-hosted tools execute on Gemini.
8//!
9//! ```no_run
10//! use rig_core::providers::gemini::cached_content::{CacheExpiry, NewCachedContent};
11//! use rig_core::providers::gemini::completion::GEMINI_2_5_FLASH;
12//! use std::time::Duration;
13//!
14//! let corpus = NewCachedContent::new(GEMINI_2_5_FLASH)
15//!     .content("A reusable document corpus")
16//!     .expiry(CacheExpiry::ttl(Duration::from_secs(600)));
17//! ```
18
19use std::time::Duration;
20
21use serde::{Deserialize, Serialize};
22
23use serde_json::{Value, json};
24
25use crate::error::EncodeError;
26use crate::error::ProviderError;
27use crate::json_utils::Lenient;
28use crate::operation::Whole;
29use crate::providers::internal::{
30    wire::{classify_or, classify_untyped_line},
31    with_query_pairs,
32};
33use crate::wire::{
34    Body, Call, Decoder, Descriptor, Encoded, Flow, Framing, Free, Mode, Operation, Out, Wire,
35    WireEvent, WireFrame,
36};
37
38/// The `cachedContents` collection path.
39const CACHED_CONTENTS_PATH: &str = "/v1beta/cachedContents";
40
41/// Gemini caps a page of `cachedContents` at 1000.
42const MAX_PAGE_SIZE: usize = 1000;
43
44/// Converts a 403 or 404 reply for the existing handle `name` to
45/// [`ProviderError::CacheExpired`], keeping the reply. Other failures are
46/// unchanged. Call only for existing handles, never for cache creation.
47pub(crate) fn on_handle(error: ProviderError, name: &str) -> ProviderError {
48    match error {
49        ProviderError::ProviderResponse(response)
50            if matches!(
51                response.status,
52                Some(http::StatusCode::FORBIDDEN | http::StatusCode::NOT_FOUND)
53            ) =>
54        {
55            ProviderError::CacheExpired {
56                name: name.to_owned(),
57                response,
58            }
59        }
60        other => other,
61    }
62}
63
64/// A relative lifetime or absolute expiry time for cached content.
65#[derive(Clone, Debug, PartialEq, Eq)]
66pub enum CacheExpiry {
67    /// Expire this long after creation. Serialized as Gemini's duration string
68    /// (`"600s"`).
69    Ttl(Duration),
70    /// Expire at an absolute RFC 3339 timestamp.
71    ExpireTime(String),
72}
73
74impl CacheExpiry {
75    pub fn ttl(ttl: Duration) -> Self {
76        Self::Ttl(ttl)
77    }
78
79    pub fn expire_time(timestamp: impl Into<String>) -> Self {
80        Self::ExpireTime(timestamp.into())
81    }
82
83    /// Gemini's duration encoding: fractional seconds with an `s` suffix.
84    fn ttl_string(ttl: Duration) -> String {
85        format!("{}.{:09}s", ttl.as_secs(), ttl.subsec_nanos())
86    }
87}
88
89/// Content and configuration for creating a cache.
90/// Supply content or a system instruction before creation. The expiry builder
91/// keeps relative and absolute expiry mutually exclusive.
92#[derive(Debug, Default, Serialize)]
93#[serde(rename_all = "camelCase")]
94pub struct NewCachedContent {
95    /// Fully qualified model name (`models/gemini-2.5-flash`). A request that
96    /// uses the cache must name the same model.
97    model: String,
98    #[serde(skip_serializing_if = "Vec::is_empty")]
99    contents: Vec<Value>,
100    #[serde(skip_serializing_if = "Option::is_none")]
101    system_instruction: Option<Value>,
102    #[serde(skip_serializing_if = "Option::is_none")]
103    tools: Option<Vec<Value>>,
104    #[serde(skip_serializing_if = "Option::is_none")]
105    tool_config: Option<Value>,
106    #[serde(skip_serializing_if = "Option::is_none")]
107    display_name: Option<String>,
108    #[serde(skip_serializing_if = "Option::is_none")]
109    ttl: Option<String>,
110    #[serde(skip_serializing_if = "Option::is_none")]
111    expire_time: Option<String>,
112}
113
114impl NewCachedContent {
115    /// Start a cached content for `model`.
116    ///
117    /// Accepts either the bare id (`gemini-2.5-flash`) or the qualified name
118    /// (`models/gemini-2.5-flash`) and normalizes to the latter, which is what
119    /// the API returns and what a `generateContent` request must match.
120    pub fn new(model: impl AsRef<str>) -> Self {
121        Self {
122            model: qualify_model(model.as_ref()),
123            ..Default::default()
124        }
125    }
126
127    /// Append a user-role text content block.
128    pub fn content(mut self, text: impl Into<String>) -> Self {
129        let text = text.into();
130        self.contents
131            .push(json!({ "parts": [{ "text": text, "thought": false }], "role": "user" }));
132        self
133    }
134
135    /// Append a content block in Gemini's REST JSON (multimodal payloads).
136    pub fn content_block(mut self, content: Value) -> Self {
137        self.contents.push(content);
138        self
139    }
140
141    pub fn system_instruction(mut self, text: impl Into<String>) -> Self {
142        let text = text.into();
143        self.system_instruction =
144            Some(json!({ "parts": [{ "text": text, "thought": false }], "role": "model" }));
145        self
146    }
147
148    /// Set the tools inherited by requests using this cache, in Gemini's
149    /// REST JSON.
150    /// Requests must not supply their own tools. Cached function declarations
151    /// require a caller-managed tool loop; provider-hosted tools do not.
152    pub fn tools(mut self, tools: Vec<Value>) -> Self {
153        self.tools = Some(tools);
154        self
155    }
156
157    /// Set the tool configuration inherited by requests using this cache,
158    /// in Gemini's REST JSON.
159    /// Requests must not supply their own tool configuration. May be set
160    /// without a tool set.
161    pub fn tool_config(mut self, tool_config: Value) -> Self {
162        self.tool_config = Some(tool_config);
163        self
164    }
165
166    pub fn display_name(mut self, name: impl Into<String>) -> Self {
167        self.display_name = Some(name.into());
168        self
169    }
170
171    /// Set the expiry. Setting it twice replaces the previous value rather than
172    /// sending both, which the API rejects.
173    pub fn expiry(mut self, expiry: CacheExpiry) -> Self {
174        match expiry {
175            CacheExpiry::Ttl(ttl) => {
176                self.ttl = Some(CacheExpiry::ttl_string(ttl));
177                self.expire_time = None;
178            }
179            CacheExpiry::ExpireTime(at) => {
180                self.expire_time = Some(at);
181                self.ttl = None;
182            }
183        }
184        self
185    }
186
187    fn validate(&self) -> Result<(), EncodeError> {
188        if self.contents.is_empty() && self.system_instruction.is_none() {
189            return Err(EncodeError::request(
190                "a cached content needs contents or a system instruction; an empty cache would \
191                 bill for storage and cache nothing",
192            ));
193        }
194        Ok(())
195    }
196}
197
198/// Storage accounting Gemini reports for a cached content.
199#[derive(Clone, Debug, Default, Deserialize, Serialize)]
200#[serde(rename_all = "camelCase")]
201pub struct CachedContentUsage {
202    /// Tokens held by this cache. This is what storage is billed on, and it is
203    /// also the ceiling on what a request against the handle can read back.
204    #[serde(default)]
205    pub total_token_count: u64,
206}
207
208/// A cached content resource as Gemini reports it.
209#[derive(Clone, Debug, Deserialize, Serialize)]
210#[serde(rename_all = "camelCase")]
211pub struct CachedContent {
212    /// Server-assigned handle, `cachedContents/<id>`. This is what
213    /// [`super::completion::GenerateContent::with_cached_content`] takes.
214    pub name: String,
215    /// Qualified model this cache is bound to.
216    #[serde(default)]
217    pub model: String,
218    #[serde(default, skip_serializing_if = "Option::is_none")]
219    pub display_name: Option<String>,
220    #[serde(default, skip_serializing_if = "Option::is_none")]
221    pub create_time: Option<String>,
222    #[serde(default, skip_serializing_if = "Option::is_none")]
223    pub update_time: Option<String>,
224    /// When this cache lapses. After it does, using the handle fails with
225    /// [`ProviderError::CacheExpired`].
226    #[serde(default, skip_serializing_if = "Option::is_none")]
227    pub expire_time: Option<String>,
228    #[serde(default, skip_serializing_if = "Option::is_none")]
229    pub usage_metadata: Option<CachedContentUsage>,
230}
231
232/// One `cachedContents` verb: what [`ContextCache`] sends.
233#[derive(Debug)]
234pub enum CachedContentRequest {
235    /// `POST /v1beta/cachedContents`; answers with the resource.
236    Create(NewCachedContent),
237    /// `GET /v1beta/cachedContents/<id>`; answers with the resource.
238    Get(String),
239    /// `GET /v1beta/cachedContents?pageSize=…`, after `page_token` when
240    /// continuing; answers with one page.
241    List {
242        /// The previous page's `nextPageToken`; `None` for the first page.
243        page_token: Option<String>,
244    },
245    /// `PATCH /v1beta/cachedContents/<id>?updateMask=…`; answers with the
246    /// resource.
247    UpdateExpiry { name: String, expiry: CacheExpiry },
248    /// `DELETE /v1beta/cachedContents/<id>`; answers with `{}`.
249    Delete(String),
250}
251
252/// One page of a `cachedContents` listing.
253#[derive(Clone, Debug, Deserialize)]
254#[serde(rename_all = "camelCase")]
255pub struct CachedContentPage {
256    /// Entries in arrival order, concatenated across pages when folded.
257    /// Required during deserialization to distinguish pages from resources.
258    pub cached_contents: Vec<CachedContent>,
259    /// The cursor naming the next page, when the listing has one. The
260    /// decoder takes it before the page reaches the fold, so a folded
261    /// reply's is always `None`.
262    #[serde(default)]
263    pub next_page_token: Option<String>,
264}
265
266/// A resource, listing page, or empty acknowledgement from `cachedContents`.
267/// Malformed bodies fail decoding rather than representing absent resources.
268#[derive(Clone, Debug, Default)]
269pub enum CachedContentReply {
270    /// `create`, `get` and `update_expiry`: the resource.
271    Resource(CachedContent),
272    /// `list`: one page of the collection.
273    Page(CachedContentPage),
274    /// An empty successful reply for deletion or an empty collection.
275    /// Also the default for a fold that receives no replies.
276    #[default]
277    Acknowledged,
278}
279
280impl CachedContentReply {
281    /// Extract the resource returned by creation, lookup, or expiry update.
282    /// Return a response error for a page or acknowledgement.
283    pub fn resource(self) -> Result<CachedContent, ProviderError> {
284        match self {
285            Self::Resource(resource) => Ok(resource),
286            other => Err(other.mismatch("one cached content")),
287        }
288    }
289
290    /// The cursor of the page after this one. An empty cursor counts as
291    /// absent: re-sending an empty `pageToken` returns the same page forever.
292    pub fn next_page_token(&self) -> Option<String> {
293        match self {
294            Self::Page(page) => page
295                .next_page_token
296                .clone()
297                .filter(|token| !token.is_empty()),
298            _ => None,
299        }
300    }
301
302    /// The entries of one listing page. An empty collection is answered
303    /// with the empty object, which is [`Self::Acknowledged`].
304    pub fn entries(self) -> Result<Vec<CachedContent>, ProviderError> {
305        match self {
306            Self::Page(page) => Ok(page.cached_contents),
307            Self::Acknowledged => Ok(Vec::new()),
308            other => Err(other.mismatch("a listing page")),
309        }
310    }
311
312    /// Build a response error naming the actual and expected reply shapes.
313    fn mismatch(&self, wanted: &str) -> ProviderError {
314        let carried = match self {
315            Self::Resource(_) => "one cached content",
316            Self::Page(_) => "a listing page",
317            Self::Acknowledged => "nothing to read, only a success status",
318        };
319        ProviderError::Response(format!("the reply carried {carried}, not {wanted}"))
320    }
321}
322
323/// Creates, reads, lists, updates expiry, or deletes explicit context caches.
324/// Requests use [`CachedContentRequest`]; a listing reads one page per call.
325#[derive(Debug, Clone, Copy, PartialEq, Eq)]
326pub struct ContextCache;
327
328impl Operation for ContextCache {
329    type Request = CachedContentRequest;
330    type Event = std::convert::Infallible;
331    type End = CachedContentReply;
332    type Response = CachedContentReply;
333    type Fold = Whole<Self>;
334    type Emit = Free;
335
336    fn fold(_request: &Self::Request, _call: &mut Call<'_>) -> Self::Fold {
337        Whole::new()
338    }
339}
340
341/// Gemini's `cachedContents` resource: the wire for
342/// [`ContextCache`].
343///
344/// Built by [`Gemini::cached_contents`](super::Gemini::cached_contents); the
345/// calls are the inherent methods of a [`Model`](crate::Model) over it.
346#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
347pub struct CachedContents {
348    /// The provider this wire speaks to.
349    pub provider: super::GeminiConfig,
350    /// Requested entries per listing page. Defaults to Gemini's cap of 1,000.
351    pub page_size: usize,
352}
353
354impl CachedContents {
355    /// The wire over `provider`, listing a full page at a time.
356    pub fn new(provider: super::GeminiConfig) -> Self {
357        Self {
358            provider,
359            page_size: MAX_PAGE_SIZE,
360        }
361    }
362
363    /// List `page_size` entries per request.
364    pub fn with_page_size(mut self, page_size: usize) -> Self {
365        self.page_size = page_size;
366        self
367    }
368
369    /// Build a listing request, optionally continuing after `page_token`.
370    /// Percent-encoding preserves the cursor and prevents query injection.
371    fn list_request(&self, page_token: Option<&str>) -> Result<http::Request<Body>, http::Error> {
372        let page_size = self.page_size.to_string();
373        let mut pairs = vec![("pageSize", page_size.as_str())];
374        if let Some(token) = page_token {
375            pairs.push(("pageToken", token));
376        }
377        let path = with_query_pairs(CACHED_CONTENTS_PATH, &pairs);
378        http::Request::get(self.provider.uri(&path)).body(Body::empty())
379    }
380}
381
382/// Set `name`, a `cachedContents/<id>` handle, as the prefix `body` reads.
383///
384/// # Errors
385///
386/// When `name` is not a handle, `body` names another one, or it sets a
387/// system instruction, tools or a tool choice, which the cache owns.
388pub fn with_cached_content(
389    body: &mut serde_json::Map<String, Value>,
390    name: &str,
391) -> Result<(), EncodeError> {
392    let fail = |message: String| Err(EncodeError::request(message));
393    if !name.starts_with("cachedContents/") {
394        return fail(format!(
395            "gemini cached content handle should look like `cachedContents/<id>`, got `{name}`"
396        ));
397    }
398    if let Some(existing) = body
399        .get("cachedContent")
400        .and_then(Value::as_str)
401        .filter(|existing| *existing != name)
402    {
403        return fail(format!(
404            "a Gemini request set cached content twice, to `{existing}` and `{name}`: set it \
405             one way or the other"
406        ));
407    }
408    let conflicts: Vec<&str> = [
409        (
410            &super::completion::SYSTEM_INSTRUCTION[..],
411            "a system instruction (preamble)",
412        ),
413        (&["tools"][..], "tools"),
414        (&super::completion::TOOL_CONFIG[..], "a tool choice"),
415    ]
416    .into_iter()
417    .filter_map(|(spellings, what)| super::completion::present(body, spellings).map(|_| what))
418    .collect();
419    if conflicts.is_empty() {
420        body.insert("cachedContent".to_owned(), Value::String(name.to_owned()));
421        return Ok(());
422    }
423    // Cached function declarations need caller-side dispatch; hosted tools
424    // run on Gemini's side.
425    let tools = body
426        .get("tools")
427        .and_then(Value::as_array)
428        .map_or(&[][..], Vec::as_slice);
429    let declares_functions = tools.iter().any(|tool| {
430        !tool.arr("functionDeclarations").is_empty()
431            || !tool.arr("function_declarations").is_empty()
432    });
433    let caveat = match declares_functions {
434        true => {
435            " Function declarations in a cache are declarations only: an `Agent` dispatches \
436             only tools it advertised, so a cached function tool runs only when you drive \
437             `GenerateContent` yourself. Hosted tools such as `codeExecution` are fine to cache."
438        }
439        false => "",
440    };
441    fail(format!(
442        "a Gemini request using cached content `{name}` also set {}. The cached content owns \
443         the system instruction, tools and tool choice of every request that uses it: move \
444         them into the cache, or drop the cache handle.{caveat}",
445        conflicts.join(" and ")
446    ))
447}
448
449impl super::GeminiConfig {
450    /// Build a wire for Gemini's explicit context cache (`cachedContents`).
451    pub(crate) fn cached_contents(&self) -> CachedContents {
452        CachedContents::new(self.clone())
453    }
454}
455
456impl Wire for CachedContents {
457    type Op = ContextCache;
458    type Payload = crate::wire::Encoded;
459    type Frame = crate::wire::WireFrame;
460    type Decoder<'id> = CachedContentsDecoder;
461    type Reassembler = crate::wire::document::Unreassembled;
462
463    fn describe(&self) -> Descriptor<'_> {
464        Descriptor::new(super::PROVIDER_NAME)
465    }
466
467    /// A resource call never streams, so both modes send the one request.
468    /// A handle is validated by `resource_path` before anything is built,
469    /// and an empty cache is refused before it bills.
470    fn encode(&self, request: CachedContentRequest, _mode: Mode) -> Result<Encoded, EncodeError> {
471        let request = match request {
472            CachedContentRequest::Create(new) => {
473                new.validate()?;
474                http::Request::post(self.provider.uri(CACHED_CONTENTS_PATH))
475                    .body(Body::Bytes(serde_json::to_vec(&new)?))?
476            }
477            CachedContentRequest::Get(name) => {
478                http::Request::get(self.provider.uri(&resource_path(&name)?)).body(Body::empty())?
479            }
480            CachedContentRequest::List { page_token } => {
481                self.list_request(page_token.as_deref())?
482            }
483            CachedContentRequest::UpdateExpiry { name, expiry } => {
484                let (patch, mask) = expiry_patch(expiry)?;
485                // Handle validation prevents query injection and retargeting the patch.
486                let path = format!("{}?updateMask={mask}", resource_path(&name)?);
487                http::Request::patch(self.provider.uri(&path)).body(Body::Bytes(patch))?
488            }
489            CachedContentRequest::Delete(name) => {
490                http::Request::delete(self.provider.uri(&resource_path(&name)?))
491                    .body(Body::empty())?
492            }
493        };
494        Ok(Encoded::new(request, Framing::Whole))
495    }
496
497    fn decoder<'id>(&self) -> Self::Decoder<'id> {
498        CachedContentsDecoder
499    }
500}
501
502/// Decodes one `cachedContents` reply.
503pub struct CachedContentsDecoder;
504
505impl<'id> Decoder<'id, ContextCache> for CachedContentsDecoder {
506    type Event = CachedContentReply;
507
508    /// Classify a page, empty acknowledgement, or resource.
509    /// Malformed bodies remain decoding failures.
510    fn classify(&self, frame: WireFrame) -> WireEvent<Self::Event> {
511        let body = frame.as_str();
512        classify_or(&body, as_page, |data| {
513            classify_or(data, as_acknowledgement, as_resource)
514        })
515    }
516
517    fn decode(
518        &mut self,
519        reply: Self::Event,
520        out: Out<'id, ContextCache>,
521    ) -> Result<Flow, ProviderError> {
522        Ok(out.end(reply))
523    }
524
525    /// A reply with no body at all is an acknowledgement: the status
526    /// already answered.
527    fn eof(&mut self, out: Out<'id, ContextCache>) -> Result<Flow, ProviderError> {
528        Ok(out.end(CachedContentReply::Acknowledged))
529    }
530}
531
532/// One listing page, or a body that is not one.
533fn as_page(data: &str) -> WireEvent<CachedContentReply> {
534    classify_untyped_line::<CachedContentPage>(data.as_bytes()).map(CachedContentReply::Page)
535}
536
537/// Classify an empty object as a deletion acknowledgement; reject any fields.
538fn as_acknowledgement(data: &str) -> WireEvent<CachedContentReply> {
539    #[derive(Deserialize)]
540    #[serde(deny_unknown_fields)]
541    struct Acknowledgement {}
542
543    classify_untyped_line::<Acknowledgement>(data.as_bytes())
544        .map(|_| CachedContentReply::Acknowledged)
545}
546
547/// One cached content, or a body that is not one.
548fn as_resource(data: &str) -> WireEvent<CachedContentReply> {
549    classify_untyped_line::<CachedContent>(data.as_bytes()).map(CachedContentReply::Resource)
550}
551
552/// Serialize the expiry patch with an update mask naming its only field.
553fn expiry_patch(expiry: CacheExpiry) -> Result<(Vec<u8>, &'static str), EncodeError> {
554    let (field, value) = match expiry {
555        CacheExpiry::Ttl(ttl) => ("ttl", CacheExpiry::ttl_string(ttl)),
556        CacheExpiry::ExpireTime(at) => ("expireTime", at),
557    };
558    let patch = serde_json::Map::from_iter([(field.to_owned(), serde_json::Value::String(value))]);
559    Ok((serde_json::to_vec(&patch)?, field))
560}
561
562/// `models/x` from `x`, idempotently.
563fn qualify_model(model: &str) -> String {
564    if model.starts_with("models/") {
565        model.to_owned()
566    } else {
567        format!("models/{model}")
568    }
569}
570
571/// Build `/v1beta/cachedContents/<id>` from a bare id or prefixed handle.
572/// Reject empty ids and characters other than ASCII letters, digits, `-`, and
573/// `_` to prevent path traversal, query injection, or resource retargeting.
574fn resource_path(name: &str) -> Result<String, EncodeError> {
575    let id = name.strip_prefix("cachedContents/").unwrap_or(name);
576    let is_id_char = |ch: char| ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_');
577    if id.is_empty() || !id.chars().all(is_id_char) {
578        return Err(EncodeError::request(format!(
579            "`{name}` is not a cached content handle; expected `cachedContents/<id>` or a bare \
580                 `<id>` of letters, digits, `-` and `_`. The id is spliced into the request path, \
581                 where a `?`, `#` or `/` silently retargets the call at a different resource — \
582                 and this is the path that deletes"
583        )));
584    }
585    Ok(format!("{CACHED_CONTENTS_PATH}/{id}"))
586}
587
588#[cfg(test)]
589mod tests;
590
591#[cfg(test)]
592mod exhaustive_validation_tests;
593
594#[cfg(test)]
595mod status_triage_tests;
rig_core/providers/gemini/cached_content.rs

rig_core/providers/gemini/
cached_content.rs
