Skip to main content

soma_gateway/gateway/
palette.rs

1use std::collections::BTreeMap;
2use std::time::{Duration, Instant};
3
4use serde_json::Value;
5use thiserror::Error;
6
7use crate::gateway::code_mode::catalog::CodeModeCatalog;
8
9pub const PALETTE_SCHEMA_CAP_BYTES: usize = 64 * 1024;
10
11#[derive(Debug, Clone, PartialEq, Eq)]
12pub struct PaletteCatalog {
13    pub tools: Vec<String>,
14}
15
16#[derive(Debug, Error, PartialEq, Eq)]
17pub enum PaletteError {
18    #[error("schema exceeds palette cap")]
19    SchemaTooLarge,
20    #[error("schema not found")]
21    NotFound,
22}
23
24#[derive(Debug, Clone)]
25pub struct PaletteCache {
26    freshness: Duration,
27    cached_at: Option<Instant>,
28    catalog: PaletteCatalog,
29    reprobe_count: usize,
30    schemas: BTreeMap<String, Value>,
31}
32
33impl PaletteCache {
34    #[must_use]
35    pub fn new(freshness: Duration) -> Self {
36        Self {
37            freshness,
38            cached_at: None,
39            catalog: PaletteCatalog { tools: Vec::new() },
40            reprobe_count: 0,
41            schemas: BTreeMap::new(),
42        }
43    }
44
45    pub fn catalog_from_code_mode(
46        &mut self,
47        catalog: &CodeModeCatalog,
48        now: Instant,
49    ) -> PaletteCatalog {
50        if self
51            .cached_at
52            .is_some_and(|cached_at| now.duration_since(cached_at) < self.freshness)
53        {
54            return self.catalog.clone();
55        }
56        self.reprobe_count += 1;
57        self.catalog = PaletteCatalog {
58            tools: catalog.names_only(),
59        };
60        self.cached_at = Some(now);
61        self.catalog.clone()
62    }
63
64    pub fn set_schema(&mut self, id: impl Into<String>, schema: Value) {
65        self.schemas.insert(id.into(), schema);
66    }
67
68    pub fn schema(&self, id: &str) -> Result<Value, PaletteError> {
69        let schema = self.schemas.get(id).ok_or(PaletteError::NotFound)?;
70        let bytes = serde_json::to_vec(schema).map_or(usize::MAX, |bytes| bytes.len());
71        if bytes > PALETTE_SCHEMA_CAP_BYTES {
72            return Err(PaletteError::SchemaTooLarge);
73        }
74        Ok(schema.clone())
75    }
76
77    #[must_use]
78    pub fn reprobe_count(&self) -> usize {
79        self.reprobe_count
80    }
81}
82
83#[cfg(test)]
84#[path = "palette_tests.rs"]
85mod tests;