1use std::collections::{HashMap, HashSet};
2use std::sync::Arc;
3
4use awaken_runtime_contract::contract::executor::LlmExecutor;
5use awaken_runtime_contract::registry_spec::{AgentSpec, ModelSpec};
6use awaken_runtime_contract::validate_unique_model_ids;
7use serde::Serialize;
8
9#[cfg(feature = "a2a")]
10use super::MapBackendRegistry;
11use super::diagnostics::{RegistryValidationError, validate_registry_set};
12use super::memory::{
13 MapAgentSpecRegistry, MapModelRegistry, MapPluginSource, MapProviderRegistry, MapToolRegistry,
14};
15use super::snapshot::RegistryHandle;
16#[cfg(feature = "a2a")]
17use super::traits::BackendRegistry;
18use super::traits::RegistrySet;
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum ProviderRemovalPolicy {
22 BlockIfReferenced,
23 CascadeUnusedModels,
24}
25
26#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
27pub struct ProviderRemovalPreview {
28 pub provider_id: String,
29 pub model_ids: Vec<String>,
30 pub agent_ids: Vec<String>,
31 pub block_if_referenced_allowed: bool,
32 pub cascade_unused_models_allowed: bool,
33}
34
35impl ProviderRemovalPreview {
36 pub fn new(
37 provider_id: impl Into<String>,
38 mut model_ids: Vec<String>,
39 mut agent_ids: Vec<String>,
40 ) -> Self {
41 model_ids.sort();
42 model_ids.dedup();
43 agent_ids.sort();
44 agent_ids.dedup();
45 Self {
46 provider_id: provider_id.into(),
47 block_if_referenced_allowed: model_ids.is_empty(),
48 cascade_unused_models_allowed: agent_ids.is_empty(),
49 model_ids,
50 agent_ids,
51 }
52 }
53}
54
55#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
56pub struct ProviderRemovalImpact {
57 pub provider_id: String,
58 pub removed_model_ids: Vec<String>,
59 pub affected_agent_ids: Vec<String>,
60}
61
62#[derive(Debug, thiserror::Error)]
63pub enum RegistryUpdateError {
64 #[error("provider already registered: {0}")]
65 ProviderAlreadyExists(String),
66 #[error("provider not found: {0}")]
67 ProviderNotFound(String),
68 #[error(
69 "provider '{provider_id}' is still referenced by models {model_ids:?} and agents {agent_ids:?}"
70 )]
71 ProviderInUse {
72 provider_id: String,
73 model_ids: Vec<String>,
74 agent_ids: Vec<String>,
75 },
76 #[error("registry build failed: {0}")]
77 Build(String),
78 #[error("{0}")]
79 Validation(#[from] RegistryValidationError),
80 #[error("config validation failed: {0}")]
81 ConfigValidation(#[from] awaken_runtime_contract::ConfigValidationError),
82}
83
84pub struct RuntimeRegistryUpdate {
85 pub providers: HashMap<String, Arc<dyn LlmExecutor>>,
86 pub models: Vec<ModelSpec>,
87 pub agents: Vec<AgentSpec>,
88}
89
90impl RegistryHandle {
91 pub fn preview_remove_provider(
92 &self,
93 id: &str,
94 ) -> Result<ProviderRemovalPreview, RegistryUpdateError> {
95 let snapshot = self.snapshot();
96 preview_provider_removal(snapshot.registries(), id)
97 }
98
99 pub fn register_provider(
100 &self,
101 id: impl Into<String>,
102 executor: Arc<dyn LlmExecutor>,
103 ) -> Result<u64, RegistryUpdateError> {
104 let id = id.into();
105 self.update(|registries| {
106 let mut draft = RegistrySetDraft::from_set(registries)?;
107 if draft.providers.contains_key(&id) {
108 return Err(RegistryUpdateError::ProviderAlreadyExists(id));
109 }
110 draft
111 .providers
112 .register_provider(id, executor)
113 .map_err(|error| RegistryUpdateError::Build(error.to_string()))?;
114 draft.into_validated_set()
115 })
116 }
117
118 pub fn replace_provider(
119 &self,
120 id: impl Into<String>,
121 executor: Arc<dyn LlmExecutor>,
122 ) -> Result<u64, RegistryUpdateError> {
123 let id = id.into();
124 self.update(|registries| {
125 let mut draft = RegistrySetDraft::from_set(registries)?;
126 if !draft.providers.contains_key(&id) {
127 return Err(RegistryUpdateError::ProviderNotFound(id));
128 }
129 draft.providers.replace_provider(id, executor);
130 draft.into_validated_set()
131 })
132 }
133
134 pub fn remove_provider(
135 &self,
136 id: &str,
137 policy: ProviderRemovalPolicy,
138 ) -> Result<ProviderRemovalImpact, RegistryUpdateError> {
139 let mut impact = None;
140 self.update(|registries| {
141 let mut draft = RegistrySetDraft::from_set(registries)?;
142 if !draft.providers.contains_key(id) {
143 return Err(RegistryUpdateError::ProviderNotFound(id.to_string()));
144 }
145
146 let preview = preview_provider_removal_from_draft(&draft, id)?;
147
148 match policy {
149 ProviderRemovalPolicy::BlockIfReferenced if !preview.model_ids.is_empty() => {
150 return Err(RegistryUpdateError::ProviderInUse {
151 provider_id: id.to_string(),
152 model_ids: preview.model_ids,
153 agent_ids: preview.agent_ids,
154 });
155 }
156 ProviderRemovalPolicy::CascadeUnusedModels if !preview.agent_ids.is_empty() => {
157 return Err(RegistryUpdateError::ProviderInUse {
158 provider_id: id.to_string(),
159 model_ids: preview.model_ids,
160 agent_ids: preview.agent_ids,
161 });
162 }
163 _ => {}
164 }
165
166 for model_id in &preview.model_ids {
167 draft.models.remove(model_id);
168 }
169 draft.providers.remove_provider(id);
170
171 impact = Some(ProviderRemovalImpact {
172 provider_id: preview.provider_id,
173 removed_model_ids: preview.model_ids,
174 affected_agent_ids: preview.agent_ids,
175 });
176 draft.into_validated_set()
177 })?;
178 impact.ok_or_else(|| RegistryUpdateError::Build("provider removal did not run".into()))
179 }
180}
181
182pub fn preview_provider_removal(
183 registries: &RegistrySet,
184 id: &str,
185) -> Result<ProviderRemovalPreview, RegistryUpdateError> {
186 if registries.providers.get_provider(id).is_none() {
187 return Err(RegistryUpdateError::ProviderNotFound(id.to_string()));
188 }
189 let model_ids = provider_model_ids_from_set(registries, id);
190 let agent_ids = agents_using_models_from_set(registries, &model_ids);
191 Ok(ProviderRemovalPreview::new(id, model_ids, agent_ids))
192}
193
194pub fn rebuild_agent_model_provider_registries(
195 base: &RegistrySet,
196 update: RuntimeRegistryUpdate,
197) -> Result<RegistrySet, RegistryUpdateError> {
198 let mut draft = RegistrySetDraft::from_set(base)?;
199
200 draft.providers = MapProviderRegistry::new();
201 for (id, executor) in update.providers {
202 draft
203 .providers
204 .register_provider(id, executor)
205 .map_err(|error| RegistryUpdateError::Build(error.to_string()))?;
206 }
207
208 validate_unique_model_ids(&update.models)?;
211
212 draft.models = MapModelRegistry::new();
213 for model in update.models {
214 draft
215 .models
216 .register_model(model)
217 .map_err(|error| RegistryUpdateError::Build(error.to_string()))?;
218 }
219
220 draft.agents = MapAgentSpecRegistry::new();
221 for agent in update.agents {
222 draft
223 .agents
224 .register_spec(agent)
225 .map_err(|error| RegistryUpdateError::Build(error.to_string()))?;
226 }
227
228 let registries = draft.into_set();
229 validate_registry_set(®istries)?;
230 Ok(registries)
231}
232
233struct RegistrySetDraft {
234 agents: MapAgentSpecRegistry,
235 tools: MapToolRegistry,
236 models: MapModelRegistry,
237 providers: MapProviderRegistry,
238 plugins: MapPluginSource,
239 #[cfg(feature = "a2a")]
240 backends: MapBackendRegistry,
241}
242
243impl RegistrySetDraft {
244 fn from_set(set: &RegistrySet) -> Result<Self, RegistryUpdateError> {
245 let mut agents = MapAgentSpecRegistry::new();
246 for id in set.agents.agent_ids() {
247 if let Some(agent) = set.agents.get_agent(&id) {
248 agents
249 .register(id, agent, |msg| {
250 crate::builder::BuildError::AgentRegistryConflict(format!("agent {msg}"))
251 })
252 .map_err(|error| RegistryUpdateError::Build(error.to_string()))?;
253 }
254 }
255
256 let mut tools = MapToolRegistry::new();
257 for id in set.tools.tool_ids() {
258 if let Some(tool) = set.tools.get_tool(&id) {
259 tools
260 .register_tool(id, tool)
261 .map_err(|error| RegistryUpdateError::Build(error.to_string()))?;
262 }
263 }
264
265 let mut models = MapModelRegistry::new();
266 for id in set.models.model_ids() {
267 if let Some(model) = set.models.get_model(&id) {
268 models
269 .register_model(model)
270 .map_err(|error| RegistryUpdateError::Build(error.to_string()))?;
271 }
272 }
273
274 let mut providers = MapProviderRegistry::new();
275 for id in set.providers.provider_ids() {
276 if let Some(provider) = set.providers.get_provider(&id) {
277 providers
278 .register_provider(id, provider)
279 .map_err(|error| RegistryUpdateError::Build(error.to_string()))?;
280 }
281 }
282
283 let mut plugins = MapPluginSource::new();
284 for id in set.plugins.plugin_ids() {
285 if let Some(plugin) = set.plugins.get_plugin(&id) {
286 plugins
287 .register_plugin(id, plugin)
288 .map_err(|error| RegistryUpdateError::Build(error.to_string()))?;
289 }
290 }
291
292 #[cfg(feature = "a2a")]
293 let mut backends = MapBackendRegistry::new();
294 #[cfg(feature = "a2a")]
295 for id in set.backends.backend_ids() {
296 if let Some(factory) = set.backends.get_backend_factory(&id) {
297 backends
298 .register_backend_factory(factory)
299 .map_err(|error| RegistryUpdateError::Build(error.to_string()))?;
300 }
301 }
302
303 Ok(Self {
304 agents,
305 tools,
306 models,
307 providers,
308 plugins,
309 #[cfg(feature = "a2a")]
310 backends,
311 })
312 }
313
314 fn into_set(self) -> RegistrySet {
315 RegistrySet {
316 agents: Arc::new(self.agents),
317 tools: Arc::new(self.tools),
318 models: Arc::new(self.models),
319 providers: Arc::new(self.providers),
320 plugins: Arc::new(self.plugins),
321 #[cfg(feature = "a2a")]
322 backends: Arc::new(self.backends) as Arc<dyn BackendRegistry>,
323 }
324 }
325
326 fn into_validated_set(self) -> Result<RegistrySet, RegistryUpdateError> {
327 let registries = self.into_set();
328 validate_registry_set(®istries)?;
329 Ok(registries)
330 }
331}
332
333fn provider_model_ids(models: &MapModelRegistry, provider_id: &str) -> Vec<String> {
334 models
335 .ids()
336 .into_iter()
337 .filter(|model_id| {
338 models
339 .get(model_id)
340 .is_some_and(|model| model.provider_id == provider_id)
341 })
342 .collect()
343}
344
345fn preview_provider_removal_from_draft(
346 draft: &RegistrySetDraft,
347 provider_id: &str,
348) -> Result<ProviderRemovalPreview, RegistryUpdateError> {
349 if !draft.providers.contains_key(provider_id) {
350 return Err(RegistryUpdateError::ProviderNotFound(
351 provider_id.to_string(),
352 ));
353 }
354 let model_ids = provider_model_ids(&draft.models, provider_id);
355 let agent_ids = agents_using_models(&draft.agents, &model_ids);
356 Ok(ProviderRemovalPreview::new(
357 provider_id,
358 model_ids,
359 agent_ids,
360 ))
361}
362
363fn provider_model_ids_from_set(registries: &RegistrySet, provider_id: &str) -> Vec<String> {
364 registries
365 .models
366 .model_ids()
367 .into_iter()
368 .filter(|model_id| {
369 registries
370 .models
371 .get_model(model_id)
372 .is_some_and(|model| model.provider_id == provider_id)
373 })
374 .collect()
375}
376
377fn agents_using_models(agents: &MapAgentSpecRegistry, model_ids: &[String]) -> Vec<String> {
378 let model_ids: HashSet<_> = model_ids.iter().map(String::as_str).collect();
379 agents
380 .ids()
381 .into_iter()
382 .filter(|agent_id| {
383 agents.get(agent_id).is_some_and(|agent| {
384 !agent.uses_remote_backend() && model_ids.contains(agent.model_id.as_str())
385 })
386 })
387 .collect()
388}
389
390fn agents_using_models_from_set(registries: &RegistrySet, model_ids: &[String]) -> Vec<String> {
391 let model_ids: HashSet<_> = model_ids.iter().map(String::as_str).collect();
392 registries
393 .agents
394 .agent_ids()
395 .into_iter()
396 .filter(|agent_id| {
397 registries.agents.get_agent(agent_id).is_some_and(|agent| {
398 !agent.uses_remote_backend() && model_ids.contains(agent.model_id.as_str())
399 })
400 })
401 .collect()
402}
403
404#[cfg(test)]
405mod tests {
406 use super::*;
407 use async_trait::async_trait;
408 use awaken_runtime_contract::contract::executor::{InferenceExecutionError, InferenceRequest};
409 use awaken_runtime_contract::contract::inference::{StopReason, StreamResult, TokenUsage};
410
411 struct StubExecutor;
412
413 #[async_trait]
414 impl LlmExecutor for StubExecutor {
415 async fn execute(
416 &self,
417 _request: InferenceRequest,
418 ) -> Result<StreamResult, InferenceExecutionError> {
419 Ok(StreamResult {
420 content: vec![],
421 tool_calls: vec![],
422 usage: Some(TokenUsage::default()),
423 stop_reason: Some(StopReason::EndTurn),
424 has_incomplete_tool_calls: false,
425 })
426 }
427
428 fn name(&self) -> &str {
429 "stub"
430 }
431 }
432
433 fn executor() -> Arc<dyn LlmExecutor> {
434 Arc::new(StubExecutor)
435 }
436
437 fn registry_set() -> RegistrySet {
438 let mut agents = MapAgentSpecRegistry::new();
439 agents
440 .register_spec(AgentSpec {
441 id: "a".into(),
442 model_id: "m".into(),
443 system_prompt: "s".into(),
444 ..Default::default()
445 })
446 .unwrap();
447
448 let mut models = MapModelRegistry::new();
449 models
450 .register_model(ModelSpec::new("m", "p", "upstream"))
451 .unwrap();
452
453 let mut providers = MapProviderRegistry::new();
454 providers.register_provider("p", executor()).unwrap();
455
456 RegistrySet {
457 agents: Arc::new(agents),
458 tools: Arc::new(MapToolRegistry::new()),
459 models: Arc::new(models),
460 providers: Arc::new(providers),
461 plugins: Arc::new(MapPluginSource::new()),
462 #[cfg(feature = "a2a")]
463 backends: Arc::new(MapBackendRegistry::new()),
464 }
465 }
466
467 #[test]
468 fn remove_provider_blocks_when_model_and_agent_depend_on_it() {
469 let handle = RegistryHandle::new(registry_set());
470 let preview = handle
471 .preview_remove_provider("p")
472 .expect("provider exists");
473 assert_eq!(
474 preview,
475 ProviderRemovalPreview {
476 provider_id: "p".into(),
477 model_ids: vec!["m".into()],
478 agent_ids: vec!["a".into()],
479 block_if_referenced_allowed: false,
480 cascade_unused_models_allowed: false,
481 }
482 );
483
484 let err = handle
485 .remove_provider("p", ProviderRemovalPolicy::CascadeUnusedModels)
486 .expect_err("agent dependency must block removal");
487 assert!(err.to_string().contains("agents [\"a\"]"));
488 }
489
490 #[test]
491 fn remove_provider_cascades_unused_models() {
492 let mut update = RuntimeRegistryUpdate {
493 providers: HashMap::new(),
494 models: vec![ModelSpec::new("m", "p", "upstream")],
495 agents: Vec::new(),
496 };
497 update.providers.insert("p".into(), executor());
498 let base = registry_set();
499 let registries = rebuild_agent_model_provider_registries(&base, update).unwrap();
500 let handle = RegistryHandle::new(registries);
501
502 let impact = handle
503 .remove_provider("p", ProviderRemovalPolicy::CascadeUnusedModels)
504 .expect("unused model can be removed with provider");
505
506 assert_eq!(impact.removed_model_ids, vec!["m"]);
507 let snapshot = handle.snapshot();
508 assert!(snapshot.registries().providers.get_provider("p").is_none());
509 assert!(snapshot.registries().models.get_model("m").is_none());
510 }
511
512 #[test]
513 fn replace_provider_keeps_model_and_agent() {
514 let handle = RegistryHandle::new(registry_set());
515 let version = handle
516 .replace_provider("p", executor())
517 .expect("provider exists");
518 assert_eq!(version, 2);
519 let snapshot = handle.snapshot();
520 assert!(snapshot.registries().providers.get_provider("p").is_some());
521 assert!(snapshot.registries().models.get_model("m").is_some());
522 assert!(snapshot.registries().agents.get_agent("a").is_some());
523 }
524
525 #[test]
526 fn concurrent_provider_registration_preserves_all_updates() {
527 let handle = Arc::new(RegistryHandle::new(registry_set()));
528 let mut threads = Vec::new();
529
530 for index in 0..16 {
531 let handle = Arc::clone(&handle);
532 threads.push(std::thread::spawn(move || {
533 handle
534 .register_provider(format!("p-{index}"), executor())
535 .expect("provider registration must succeed");
536 }));
537 }
538
539 for thread in threads {
540 thread.join().expect("thread must not panic");
541 }
542
543 let snapshot = handle.snapshot();
544 for index in 0..16 {
545 let provider_id = format!("p-{index}");
546 assert!(
547 snapshot
548 .registries()
549 .providers
550 .get_provider(&provider_id)
551 .is_some(),
552 "provider {provider_id} must survive concurrent updates"
553 );
554 }
555 }
556}