1use std::collections::HashMap;
2use std::sync::Arc;
3
4use crate::extensions::runtime::process::RegisteredProviderSpec;
5use crate::extensions::runtime::ExtensionHandler;
6
7pub struct RegisteredProvider {
8 pub plugin_id: String,
9 pub provider_id: String,
10 pub runtime_id: String,
11 pub spec: RegisteredProviderSpec,
12 pub handler: Option<Arc<dyn ExtensionHandler>>,
13}
14
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct RegisteredProviderModelSummary {
17 pub runtime_id: String,
18 pub display_name: Option<String>,
19 pub tool_use: bool,
20 pub streaming: bool,
21 pub context_window: Option<u64>,
22}
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct RegisteredProviderSummary {
26 pub runtime_id: String,
27 pub display_name: String,
28 pub models: Vec<RegisteredProviderModelSummary>,
29}
30
31#[derive(Default)]
32pub struct ProviderRegistry {
33 providers: HashMap<String, RegisteredProvider>,
34}
35
36impl ProviderRegistry {
37 pub fn new() -> Self {
38 Self::default()
39 }
40
41 pub fn register(
42 &mut self,
43 plugin_id: &str,
44 spec: RegisteredProviderSpec,
45 ) -> Result<String, String> {
46 self.register_with_handler(plugin_id, spec, None)
47 }
48
49 pub fn register_with_handler(
50 &mut self,
51 plugin_id: &str,
52 spec: RegisteredProviderSpec,
53 handler: Option<Arc<dyn ExtensionHandler>>,
54 ) -> Result<String, String> {
55 let runtime_id = format!("{}:{}", plugin_id, spec.id);
56 if self.providers.contains_key(&runtime_id) {
57 return Err(format!("provider '{}' is already registered", runtime_id));
58 }
59 self.providers.insert(
60 runtime_id.clone(),
61 RegisteredProvider {
62 plugin_id: plugin_id.to_string(),
63 provider_id: spec.id.clone(),
64 runtime_id: runtime_id.clone(),
65 spec,
66 handler,
67 },
68 );
69 Ok(runtime_id)
70 }
71
72 pub fn unregister_plugin(&mut self, plugin_id: &str) {
73 self.providers
74 .retain(|_, provider| provider.plugin_id != plugin_id);
75 }
76
77 pub fn get(&self, runtime_id: &str) -> Option<&RegisteredProvider> {
78 self.providers.get(runtime_id)
79 }
80
81 pub fn list(&self) -> Vec<&RegisteredProvider> {
82 let mut providers: Vec<_> = self.providers.values().collect();
83 providers.sort_by(|a, b| a.runtime_id.cmp(&b.runtime_id));
84 providers
85 }
86
87 pub fn len(&self) -> usize {
88 self.providers.len()
89 }
90
91 pub fn is_empty(&self) -> bool {
92 self.providers.is_empty()
93 }
94
95 pub fn parse_model_id(model: &str) -> Option<(&str, &str, &str)> {
96 let mut parts = model.split(':');
97 let plugin_id = parts.next()?;
98 let provider_id = parts.next()?;
99 let model_id = parts.next()?;
100 if parts.next().is_some()
101 || plugin_id.is_empty()
102 || provider_id.is_empty()
103 || model_id.is_empty()
104 {
105 return None;
106 }
107 Some((plugin_id, provider_id, model_id))
108 }
109
110 pub fn model_runtime_id(plugin_id: &str, provider_id: &str, model_id: &str) -> String {
111 format!("{}:{}:{}", plugin_id, provider_id, model_id)
112 }
113
114 pub fn summaries(&self) -> Vec<RegisteredProviderSummary> {
115 self.list()
116 .into_iter()
117 .map(|provider| RegisteredProviderSummary {
118 runtime_id: provider.runtime_id.clone(),
119 display_name: provider.spec.display_name.clone(),
120 models: provider
121 .spec
122 .models
123 .iter()
124 .map(|model| RegisteredProviderModelSummary {
125 runtime_id: Self::model_runtime_id(
126 &provider.plugin_id,
127 &provider.provider_id,
128 &model.id,
129 ),
130 display_name: model.display_name.clone(),
131 tool_use: model
132 .capabilities
133 .get("tool_use")
134 .and_then(|v| v.as_bool())
135 .unwrap_or(false),
136 streaming: model
137 .capabilities
138 .get("streaming")
139 .and_then(|v| v.as_bool())
140 .unwrap_or(false),
141 context_window: model.context_window,
142 })
143 .collect(),
144 })
145 .collect()
146 }
147}
148
149#[cfg(test)]
150mod tests {
151 use super::*;
152
153 fn spec(id: &str) -> RegisteredProviderSpec {
154 RegisteredProviderSpec {
155 id: id.to_string(),
156 display_name: "Local".to_string(),
157 description: "Local provider".to_string(),
158 models: vec![],
159 config_schema: None,
160 }
161 }
162
163 #[test]
164 fn summaries_include_model_tool_use_capability_and_context_metadata() {
165 let mut spec = spec("local");
166 spec.models = vec![
167 crate::extensions::runtime::process::RegisteredProviderModelSpec {
168 id: "model-a".to_string(),
169 display_name: Some("Model A".to_string()),
170 capabilities: serde_json::json!({"tool_use": true, "streaming": true}),
171 context_window: Some(8192),
172 },
173 ];
174 let mut registry = ProviderRegistry::new();
175 registry.register("plugin", spec).unwrap();
176
177 let summaries = registry.summaries();
178
179 assert_eq!(
180 summaries[0].models,
181 vec![RegisteredProviderModelSummary {
182 runtime_id: "plugin:local:model-a".to_string(),
183 display_name: Some("Model A".to_string()),
184 tool_use: true,
185 streaming: true,
186 context_window: Some(8192),
187 }]
188 );
189 }
190
191 #[test]
192 fn summaries_default_streaming_to_false_when_capability_absent() {
193 let mut spec = spec("local");
194 spec.models = vec![
195 crate::extensions::runtime::process::RegisteredProviderModelSpec {
196 id: "model-b".to_string(),
197 display_name: None,
198 capabilities: serde_json::json!({}),
199 context_window: None,
200 },
201 ];
202 let mut registry = ProviderRegistry::new();
203 registry.register("plugin", spec).unwrap();
204
205 let summaries = registry.summaries();
206 assert!(!summaries[0].models[0].streaming);
207 assert!(!summaries[0].models[0].tool_use);
208 }
209
210 #[test]
211 fn register_namespaces_provider_by_plugin() {
212 let mut registry = ProviderRegistry::new();
213 let id = registry.register("plugin", spec("local")).unwrap();
214
215 assert_eq!(id, "plugin:local");
216 assert!(registry.get("plugin:local").is_some());
217 }
218
219 #[test]
220 fn duplicate_runtime_provider_id_is_rejected() {
221 let mut registry = ProviderRegistry::new();
222 registry.register("plugin", spec("local")).unwrap();
223 let err = registry.register("plugin", spec("local")).unwrap_err();
224
225 assert!(err.contains("already registered"));
226 }
227
228 #[test]
229 fn unregister_plugin_removes_its_providers_only() {
230 let mut registry = ProviderRegistry::new();
231 registry.register("one", spec("local")).unwrap();
232 registry.register("two", spec("local")).unwrap();
233
234 registry.unregister_plugin("one");
235
236 assert!(registry.get("one:local").is_none());
237 assert!(registry.get("two:local").is_some());
238 }
239
240 #[test]
241 fn model_ids_use_three_part_namespace() {
242 assert_eq!(
243 ProviderRegistry::parse_model_id("plugin:local:model-a"),
244 Some(("plugin", "local", "model-a"))
245 );
246 assert_eq!(
247 ProviderRegistry::model_runtime_id("plugin", "local", "model-a"),
248 "plugin:local:model-a"
249 );
250 assert!(ProviderRegistry::parse_model_id("plugin:local").is_none());
251 assert!(ProviderRegistry::parse_model_id("plugin:local:model:extra").is_none());
252 }
253}