soma_gateway/gateway/
palette.rs1use 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;