Skip to main content

soma_gateway/gateway/code_mode/
catalog.rs

1use std::collections::BTreeMap;
2
3use serde_json::Value;
4use thiserror::Error;
5
6use crate::upstream::UpstreamSnapshot;
7
8pub const CODEMODE_SCHEMA_CAP_BYTES: usize = 512 * 1024;
9
10#[derive(Debug, Clone, PartialEq)]
11pub struct CodeModeCatalogEntry {
12    pub id: String,
13    pub namespace: String,
14    pub tool: String,
15    pub description: Option<String>,
16    pub input_schema: Option<Value>,
17    pub ui_link: Option<String>,
18}
19
20#[derive(Debug, Clone, Default, PartialEq)]
21pub struct CodeModeCatalog {
22    entries: BTreeMap<String, CodeModeCatalogEntry>,
23}
24
25#[derive(Debug, Error, PartialEq, Eq)]
26pub enum CatalogError {
27    #[error("schema exceeds Code Mode cap")]
28    SchemaTooLarge,
29    #[error("schema not found")]
30    NotFound,
31}
32
33impl CodeModeCatalog {
34    #[must_use]
35    pub fn from_snapshots(snapshots: &[UpstreamSnapshot]) -> Self {
36        let mut entries = BTreeMap::new();
37        for snapshot in snapshots {
38            for tool in &snapshot.tools {
39                let descriptor = soma_codemode::ToolDescriptor::tool(
40                    &snapshot.name,
41                    &tool.name,
42                    tool.description.as_deref().unwrap_or_default(),
43                    tool.input_schema.clone(),
44                    tool.output_schema.clone(),
45                );
46                let id = descriptor.id;
47                entries.insert(
48                    id.clone(),
49                    CodeModeCatalogEntry {
50                        id,
51                        namespace: snapshot.name.clone(),
52                        tool: tool.name.clone(),
53                        description: tool.description.clone(),
54                        input_schema: descriptor.schema,
55                        ui_link: None,
56                    },
57                );
58            }
59        }
60        Self { entries }
61    }
62
63    #[must_use]
64    pub fn names_only(&self) -> Vec<String> {
65        self.entries.keys().cloned().collect()
66    }
67
68    pub fn schema_for(&self, id: &str) -> Result<Option<Value>, CatalogError> {
69        let entry = self.entries.get(id).ok_or(CatalogError::NotFound)?;
70        let Some(schema) = &entry.input_schema else {
71            return Ok(None);
72        };
73        let bytes = serde_json::to_vec(schema).map_or(usize::MAX, |bytes| bytes.len());
74        if bytes > CODEMODE_SCHEMA_CAP_BYTES {
75            return Err(CatalogError::SchemaTooLarge);
76        }
77        Ok(Some(schema.clone()))
78    }
79}
80
81#[cfg(test)]
82#[path = "catalog_tests.rs"]
83mod tests;