soma_provider_core/registry/
index.rs

1use std::{collections::BTreeMap, sync::Arc};
2
3use jsonschema::Validator;
4
5use crate::{
6    ProviderCatalog, ProviderId, ProviderValidationError, ToolSpec, validate_provider_manifest,
7};
8
9#[derive(Clone)]
10pub struct RegisteredTool {
11    provider_id: ProviderId,
12    tool: ToolSpec,
13    pub(super) input_validator: Arc<Validator>,
14    pub(super) output_validator: Option<Arc<Validator>>,
15}
16
17impl RegisteredTool {
18    pub fn provider_id(&self) -> &ProviderId {
19        &self.provider_id
20    }
21
22    pub fn spec(&self) -> &ToolSpec {
23        &self.tool
24    }
25}
26
27#[derive(Clone, Default)]
28pub struct ProviderIndexes {
29    tools: BTreeMap<String, RegisteredTool>,
30    rest: BTreeMap<(String, String), String>,
31    cli: BTreeMap<String, String>,
32    primitives: BTreeMap<String, &'static str>,
33}
34
35impl ProviderIndexes {
36    pub(super) fn build(catalogs: &[ProviderCatalog]) -> Result<Self, ProviderValidationError> {
37        let mut indexes = Self::default();
38        for catalog in catalogs {
39            validate_provider_manifest(catalog)?;
40            let provider_id = ProviderId::new(&catalog.provider.name).map_err(|error| {
41                ProviderValidationError::new("invalid_provider_name", error.to_string())
42            })?;
43            for tool in &catalog.tools {
44                indexes.insert_tool(provider_id.clone(), tool.clone())?;
45            }
46            for prompt in &catalog.prompts {
47                indexes.insert_primitive(&prompt.name, "prompt")?;
48            }
49            for resource in &catalog.resources {
50                indexes.insert_primitive(&resource.name, "resource")?;
51            }
52            for task in &catalog.tasks {
53                indexes.insert_primitive(&task.name, "task")?;
54            }
55            for elicitation in &catalog.elicitation {
56                indexes.insert_primitive(&elicitation.name, "elicitation")?;
57            }
58        }
59        Ok(indexes)
60    }
61
62    fn insert_tool(
63        &mut self,
64        provider_id: ProviderId,
65        tool: ToolSpec,
66    ) -> Result<(), ProviderValidationError> {
67        let input_validator = Arc::new(jsonschema::validator_for(&tool.input_schema).map_err(
68            |error| {
69                ProviderValidationError::new(
70                    "input_schema_invalid",
71                    format!("tool `{}` has invalid input_schema: {error}", tool.name),
72                )
73            },
74        )?);
75        let output_validator = tool
76            .output_schema
77            .as_ref()
78            .map(jsonschema::validator_for)
79            .transpose()
80            .map_err(|error| {
81                ProviderValidationError::new(
82                    "output_schema_invalid",
83                    format!("tool `{}` has invalid output_schema: {error}", tool.name),
84                )
85            })?
86            .map(Arc::new);
87        let name = tool.name.clone();
88        let entry = RegisteredTool {
89            provider_id,
90            tool: tool.clone(),
91            input_validator,
92            output_validator,
93        };
94        if self.tools.insert(name.clone(), entry).is_some() {
95            return Err(ProviderValidationError::new(
96                "duplicate_tool_name",
97                format!("duplicate action `{name}`"),
98            ));
99        }
100
101        if let Some(rest) = &tool.rest
102            && rest.enabled
103        {
104            let method = rest.method.clone().unwrap_or_else(|| "POST".to_owned());
105            let path = rest.path.clone().unwrap_or_else(|| format!("/v1/{name}"));
106            if self
107                .rest
108                .insert((method.clone(), path.clone()), name.clone())
109                .is_some()
110            {
111                return Err(ProviderValidationError::new(
112                    "duplicate_rest_route",
113                    format!("duplicate REST route {method} {path}"),
114                ));
115            }
116        }
117        if let Some(cli) = &tool.cli
118            && cli.enabled
119        {
120            let command = cli.command.clone().unwrap_or_else(|| name.clone());
121            self.insert_cli(command, &name)?;
122            for alias in &cli.aliases {
123                self.insert_cli(alias.clone(), &name)?;
124            }
125        }
126        Ok(())
127    }
128
129    fn insert_cli(&mut self, command: String, action: &str) -> Result<(), ProviderValidationError> {
130        if self
131            .cli
132            .insert(command.clone(), action.to_owned())
133            .is_some()
134        {
135            return Err(ProviderValidationError::new(
136                "duplicate_cli_command",
137                format!("duplicate CLI command `{command}`"),
138            ));
139        }
140        Ok(())
141    }
142
143    fn insert_primitive(
144        &mut self,
145        name: &str,
146        kind: &'static str,
147    ) -> Result<(), ProviderValidationError> {
148        if self.primitives.insert(name.to_owned(), kind).is_some() {
149            return Err(ProviderValidationError::new(
150                "duplicate_mcp_primitive",
151                format!("duplicate MCP primitive `{name}`"),
152            ));
153        }
154        Ok(())
155    }
156
157    pub fn tool(&self, action: &str) -> Option<&RegisteredTool> {
158        self.tools.get(action)
159    }
160
161    pub fn action_names(&self) -> impl Iterator<Item = &str> {
162        self.tools.keys().map(String::as_str)
163    }
164
165    pub fn route_action(&self, method: &str, path: &str) -> Option<&str> {
166        self.rest
167            .get(&(method.to_owned(), path.to_owned()))
168            .map(String::as_str)
169    }
170
171    pub fn cli_action(&self, command: &str) -> Option<&str> {
172        self.cli.get(command).map(String::as_str)
173    }
174
175    pub fn primitive_kind(&self, name: &str) -> Option<&str> {
176        self.primitives.get(name).copied()
177    }
178
179    pub fn rest_routes(&self) -> impl Iterator<Item = (&str, &str, &str)> {
180        self.rest
181            .iter()
182            .map(|((method, path), action)| (method.as_str(), path.as_str(), action.as_str()))
183    }
184
185    pub fn compiled_validator_count(&self) -> usize {
186        self.tools
187            .values()
188            .map(|tool| 1 + usize::from(tool.output_validator.is_some()))
189            .sum()
190    }
191}
soma_provider_core/registry/index.rs

soma_provider_core/registry/
index.rs
