1use std::collections::hash_map::DefaultHasher;
2use std::collections::{BTreeMap, HashMap, HashSet};
3use std::hash::{Hash, Hasher};
4use std::sync::{Arc, Weak};
5use std::time::Duration;
6
7use async_trait::async_trait;
8use awaken_ext_mcp::{
9 DefaultSamplingHandler, McpServerConnectionConfig, McpServerStatusSnapshot, McpToolRegistry,
10 McpToolRegistryManager, SamplingHandler, SamplingHandlerFactory,
11};
12use awaken_runtime::AgentRuntime;
13use awaken_runtime::engine::GenaiExecutor;
14use awaken_runtime::registry::{AgentSpecRegistry, BackendRegistry, PluginSource, ToolRegistry};
15use awaken_server_contract as server_contract;
16use genai::adapter::AdapterKind;
17use genai::resolver::{AuthData, Endpoint};
18use genai::{Client, ModelIden, ServiceTarget, WebConfig};
19use parking_lot::{Mutex, RwLock};
20use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
21use serde_json::Value;
22use server_contract::contract::config_store::{ConfigChangeNotifier, ConfigStore};
23use server_contract::contract::executor::LlmExecutor;
24use server_contract::contract::storage::StorageError;
25use server_contract::{
26 AgentSpec, ConfigRecord, McpRestartPolicy, McpServerSpec, McpTransportKind, ModelSpec,
27 PeriodicRefresher, ProviderSpec, SkillSpecSink,
28};
29use tokio::runtime::Handle;
30use tokio::sync::oneshot;
31use tokio::task::JoinHandle;
32
33mod a2a_discovery;
34#[cfg(test)]
35mod credential_tests;
36mod discovered_agents;
37mod managed_config;
38mod mcp_inventory;
39mod provider_cache;
40mod provider_capability_discovery;
41mod publish;
42mod registry_compile;
43mod skill_publish;
44#[cfg(test)]
45mod skill_tests;
46mod versioned_publish;
47
48use discovered_agents::{AgentSpecRegistryWithDiscovery, DiscoveredAgentRegistry};
49use managed_config::ManagedConfigSnapshot;
50pub use mcp_inventory::McpServerInventory;
51
52const CONFIG_LOAD_PAGE_SIZE: usize = 1024;
53
54const NS_AGENTS: &str = "agents";
55const NS_MODELS: &str = "models";
56const NS_PROVIDERS: &str = "providers";
57const NS_A2A_SERVERS: &str = "a2a-servers";
58const NS_MCP_SERVERS: &str = "mcp-servers";
59const NS_TOOLS: &str = "tools";
60const NS_SKILLS: &str = "skills";
61
62use provider_cache::{ProviderExecutorCache, ProviderRuntimeCache};
63
64#[derive(Debug, thiserror::Error)]
65pub enum ConfigRuntimeError {
66 #[error("runtime does not expose a configurable registry snapshot")]
67 RuntimeNotConfigurable,
68 #[error(
69 "config store is partially initialized; bootstrap requires all managed namespaces to be empty or all core namespaces populated"
70 )]
71 PartialBootstrap,
72 #[error(
73 "unsupported provider adapter: {0} (valid names mirror genai::adapter::AdapterKind — see https://docs.rs/genai/latest/genai/adapter/enum.AdapterKind.html)"
74 )]
75 UnsupportedProviderAdapter(String),
76 #[error("invalid managed config: {0}")]
77 InvalidConfig(String),
78 #[error("periodic refresh error: {0}")]
79 PeriodicRefresh(String),
80 #[error("config change listener error: {0}")]
81 ChangeListener(String),
82 #[error("versioned registry error: {0}")]
83 VersionedRegistry(String),
84 #[error("storage error: {0}")]
85 Storage(#[from] StorageError),
86}
87
88macro_rules! overlay_registry {
89 ($name:ident, $trait:ident, $get:ident -> $ret:ty, $ids:ident) => {
90 struct $name {
91 base: Arc<dyn $trait>,
92 overlay: Arc<dyn $trait>,
93 }
94
95 impl $name {
96 fn new(base: Arc<dyn $trait>, overlay: Arc<dyn $trait>) -> Self {
97 Self { base, overlay }
98 }
99 }
100
101 impl $trait for $name {
102 fn $get(&self, id: &str) -> $ret {
103 self.base.$get(id).or_else(|| self.overlay.$get(id))
104 }
105
106 fn $ids(&self) -> Vec<String> {
107 let mut ids = self.base.$ids();
108 ids.extend(self.overlay.$ids());
109 ids.sort();
110 ids.dedup();
111 ids
112 }
113 }
114 };
115}
116
117overlay_registry!(OverlayToolRegistry, ToolRegistry, get_tool -> Option<Arc<dyn server_contract::contract::tool::Tool>>, tool_ids);
118
119#[derive(Clone)]
120struct DynamicMcpToolRegistry {
121 registry: McpToolRegistry,
122}
123
124impl DynamicMcpToolRegistry {
125 fn new(registry: McpToolRegistry) -> Self {
126 Self { registry }
127 }
128}
129
130impl ToolRegistry for DynamicMcpToolRegistry {
131 fn get_tool(&self, id: &str) -> Option<Arc<dyn server_contract::contract::tool::Tool>> {
132 self.registry.get(id)
133 }
134
135 fn tool_ids(&self) -> Vec<String> {
136 self.registry.ids()
137 }
138}
139
140pub trait ProviderExecutorFactory: Send + Sync {
141 fn build(&self, spec: &ProviderSpec) -> Result<Arc<dyn LlmExecutor>, ConfigRuntimeError>;
142}
143
144pub struct GenaiProviderExecutorFactory;
152
153pub struct BrokeredGenaiProviderExecutorFactory {
155 broker: Arc<dyn awaken_runtime::credentials::CredentialBroker>,
156}
157
158impl Default for GenaiProviderExecutorFactory {
159 fn default() -> Self {
160 Self
161 }
162}
163
164impl GenaiProviderExecutorFactory {
165 pub fn with_broker(
169 broker: Arc<dyn awaken_runtime::credentials::CredentialBroker>,
170 ) -> BrokeredGenaiProviderExecutorFactory {
171 BrokeredGenaiProviderExecutorFactory { broker }
172 }
173}
174
175impl ProviderExecutorFactory for GenaiProviderExecutorFactory {
176 fn build(&self, spec: &ProviderSpec) -> Result<Arc<dyn LlmExecutor>, ConfigRuntimeError> {
177 build_genai_provider_executor(spec)
178 }
179}
180
181impl ProviderExecutorFactory for BrokeredGenaiProviderExecutorFactory {
182 fn build(&self, spec: &ProviderSpec) -> Result<Arc<dyn LlmExecutor>, ConfigRuntimeError> {
183 build_genai_provider_executor_with_broker(spec, Arc::clone(&self.broker))
184 }
185}
186
187#[async_trait]
188pub trait ManagedMcpRegistry: Send + Sync {
189 fn tool_registry(&self) -> Arc<dyn ToolRegistry>;
190 fn periodic_refresh_running(&self) -> bool;
191 fn start_periodic_refresh(&self, interval: Duration) -> Result<(), ConfigRuntimeError>;
192 async fn stop_periodic_refresh(&self) -> bool;
193 async fn close(&self) -> Result<(), ConfigRuntimeError> {
194 self.stop_periodic_refresh().await;
195 Ok(())
196 }
197 async fn server_status(&self, _server_name: &str) -> Option<McpServerStatusSnapshot> {
202 None
203 }
204 async fn server_prompts(
205 &self,
206 _server_name: &str,
207 ) -> Result<Vec<awaken_ext_mcp::McpPromptEntry>, ConfigRuntimeError> {
208 Ok(Vec::new())
209 }
210 async fn server_resources(
211 &self,
212 _server_name: &str,
213 ) -> Result<Vec<awaken_ext_mcp::McpResourceEntry>, ConfigRuntimeError> {
214 Ok(Vec::new())
215 }
216 async fn reconnect(&self, server_name: &str) -> Result<(), ConfigRuntimeError> {
217 Err(ConfigRuntimeError::InvalidConfig(format!(
218 "MCP registry does not support reconnect for server '{server_name}'"
219 )))
220 }
221}
222
223#[async_trait]
224pub trait McpRegistryFactory: Send + Sync {
225 async fn connect(
226 &self,
227 specs: &[McpServerSpec],
228 ) -> Result<Option<Arc<dyn ManagedMcpRegistry>>, ConfigRuntimeError>;
229}
230
231struct SamplingFactoryCacheState {
256 version: u64,
257 entries: HashMap<(String, String), Arc<dyn SamplingHandler>>,
258}
259
260impl SamplingFactoryCacheState {
261 fn empty() -> Self {
262 Self {
263 version: 0,
264 entries: HashMap::new(),
265 }
266 }
267}
268
269pub(crate) struct RegistryDrivenSamplingHandlerFactory {
270 runtime: Weak<AgentRuntime>,
271 cache: std::sync::Mutex<SamplingFactoryCacheState>,
272}
273
274impl RegistryDrivenSamplingHandlerFactory {
275 pub(crate) fn new(runtime: Weak<AgentRuntime>) -> Self {
276 Self {
277 runtime,
278 cache: std::sync::Mutex::new(SamplingFactoryCacheState::empty()),
279 }
280 }
281}
282
283#[async_trait]
284impl SamplingHandlerFactory for RegistryDrivenSamplingHandlerFactory {
285 async fn for_agent(&self, agent_spec: &AgentSpec) -> Option<Arc<dyn SamplingHandler>> {
286 let runtime = self.runtime.upgrade()?;
287 let snapshot = runtime.registry_snapshot()?;
295 let version = snapshot.version();
296 let key = (agent_spec.id.clone(), agent_spec.model_id.clone());
297
298 {
300 let cache = self.cache.lock().unwrap_or_else(|p| p.into_inner());
301 if cache.version == version
302 && let Some(cached) = cache.entries.get(&key).cloned()
303 {
304 return Some(cached);
305 }
306 }
307
308 let registries = snapshot.registries();
309 let model = registries.models.get_model(&agent_spec.model_id)?;
310 let executor = registries.providers.get_provider(&model.provider_id)?;
311 let handler: Arc<dyn SamplingHandler> =
312 Arc::new(DefaultSamplingHandler::new(executor, model.upstream_model));
313
314 let mut cache = self.cache.lock().unwrap_or_else(|p| p.into_inner());
319 if cache.version != version {
320 cache.entries.clear();
321 cache.version = version;
322 }
323 cache.entries.insert(key, Arc::clone(&handler));
324 Some(handler)
325 }
326}
327
328pub struct DefaultMcpRegistryFactory {
337 sampling_handler_factory: Option<Arc<dyn SamplingHandlerFactory>>,
338}
339
340impl DefaultMcpRegistryFactory {
341 pub fn new() -> Self {
345 Self {
346 sampling_handler_factory: None,
347 }
348 }
349
350 pub fn with_runtime(runtime: Weak<AgentRuntime>) -> Self {
355 Self {
356 sampling_handler_factory: Some(Arc::new(RegistryDrivenSamplingHandlerFactory::new(
357 runtime,
358 ))),
359 }
360 }
361}
362
363impl Default for DefaultMcpRegistryFactory {
364 fn default() -> Self {
365 Self::new()
366 }
367}
368
369#[derive(Clone)]
370struct RealManagedMcpRegistry {
371 manager: McpToolRegistryManager,
372 tool_registry: Arc<dyn ToolRegistry>,
373}
374
375#[async_trait]
376impl ManagedMcpRegistry for RealManagedMcpRegistry {
377 fn tool_registry(&self) -> Arc<dyn ToolRegistry> {
378 Arc::clone(&self.tool_registry)
379 }
380
381 fn periodic_refresh_running(&self) -> bool {
382 self.manager.periodic_refresh_running()
383 }
384
385 fn start_periodic_refresh(&self, interval: Duration) -> Result<(), ConfigRuntimeError> {
386 self.manager
387 .start_periodic_refresh(interval)
388 .map_err(|error| ConfigRuntimeError::InvalidConfig(error.to_string()))
389 }
390
391 async fn stop_periodic_refresh(&self) -> bool {
392 self.manager.stop_periodic_refresh().await
393 }
394
395 async fn close(&self) -> Result<(), ConfigRuntimeError> {
396 self.manager
397 .close_all()
398 .await
399 .map_err(|error| ConfigRuntimeError::InvalidConfig(error.to_string()))
400 }
401
402 async fn server_status(&self, server_name: &str) -> Option<McpServerStatusSnapshot> {
403 self.manager.server_status_snapshot(server_name).await.ok()
404 }
405
406 async fn server_prompts(
407 &self,
408 server_name: &str,
409 ) -> Result<Vec<awaken_ext_mcp::McpPromptEntry>, ConfigRuntimeError> {
410 let prompts = self
411 .manager
412 .list_prompts()
413 .await
414 .map_err(|error| ConfigRuntimeError::InvalidConfig(error.to_string()))?;
415 Ok(prompts
416 .into_iter()
417 .filter(|entry| entry.server_name == server_name)
418 .collect())
419 }
420
421 async fn server_resources(
422 &self,
423 server_name: &str,
424 ) -> Result<Vec<awaken_ext_mcp::McpResourceEntry>, ConfigRuntimeError> {
425 let resources = self
426 .manager
427 .list_resources()
428 .await
429 .map_err(|error| ConfigRuntimeError::InvalidConfig(error.to_string()))?;
430 Ok(resources
431 .into_iter()
432 .filter(|entry| entry.server_name == server_name)
433 .collect())
434 }
435
436 async fn reconnect(&self, server_name: &str) -> Result<(), ConfigRuntimeError> {
437 self.manager
438 .reconnect(server_name)
439 .await
440 .map_err(|e| ConfigRuntimeError::InvalidConfig(e.to_string()))
441 }
442}
443
444#[async_trait]
445impl McpRegistryFactory for DefaultMcpRegistryFactory {
446 async fn connect(
447 &self,
448 specs: &[McpServerSpec],
449 ) -> Result<Option<Arc<dyn ManagedMcpRegistry>>, ConfigRuntimeError> {
450 if specs.is_empty() {
451 return Ok(None);
452 }
453
454 let configs = specs
455 .iter()
456 .map(mcp_spec_to_connection_config)
457 .collect::<Result<Vec<_>, _>>()?;
458 let manager = McpToolRegistryManager::connect_with_sampling_factory(
463 configs,
464 None,
465 self.sampling_handler_factory.clone(),
466 )
467 .await
468 .map_err(|error| {
469 ConfigRuntimeError::InvalidConfig(format!("failed to connect MCP servers: {error}"))
470 })?;
471
472 Ok(Some(Arc::new(RealManagedMcpRegistry {
473 tool_registry: Arc::new(DynamicMcpToolRegistry::new(manager.registry())),
474 manager,
475 }) as Arc<dyn ManagedMcpRegistry>))
476 }
477}
478
479#[derive(Clone)]
480struct ActiveMcpRegistry {
481 specs: Vec<McpServerSpec>,
482 handle: Arc<dyn ManagedMcpRegistry>,
483 tool_registry: Arc<dyn ToolRegistry>,
484}
485
486struct PreparedMcpRegistry {
487 tool_registry: Option<Arc<dyn ToolRegistry>>,
488 next_state: Option<ActiveMcpRegistry>,
489 state_changed: bool,
490}
491
492impl PreparedMcpRegistry {
493 async fn cleanup(self) {
494 if let Some(active) = self.next_state
495 && let Err(error) = active.handle.close().await
496 {
497 tracing::warn!(
498 error = %error,
499 "failed to close prepared MCP registry after publish failure"
500 );
501 }
502 }
503}
504
505struct ChangeListenerRuntime {
506 stop_tx: Option<oneshot::Sender<()>>,
507 join: JoinHandle<()>,
508}
509
510pub struct ConfigRuntimeManager {
511 runtime: Arc<AgentRuntime>,
512 store: Arc<dyn ConfigStore>,
513 tools: Arc<dyn ToolRegistry>,
514 plugins: Arc<dyn PluginSource>,
515 backends: Arc<dyn BackendRegistry>,
516 skill_spec_sink: Option<Arc<dyn SkillSpecSink>>,
517 discovered_agents: Option<Arc<dyn AgentSpecRegistry>>,
521 provider_factory: Arc<dyn ProviderExecutorFactory>,
522 change_notifier: Option<Arc<dyn ConfigChangeNotifier>>,
523 mcp_registry_factory: Arc<dyn McpRegistryFactory>,
524 apply_lock: tokio::sync::Mutex<()>,
525 active_mcp_registry: Mutex<Option<ActiveMcpRegistry>>,
526 last_applied_fingerprint: RwLock<Option<u64>>,
527 provider_cache: Mutex<ProviderRuntimeCache>,
532 periodic_refresh: PeriodicRefresher,
533 change_listener: Mutex<Option<ChangeListenerRuntime>>,
534 mcp_refresh_interval: RwLock<Option<Duration>>,
535 min_apply_interval: Duration,
540 audit_log: Option<Arc<crate::services::audit_log::AuditLogger>>,
543 versioned_registry: Option<versioned_publish::VersionedRegistryPublicationTarget>,
544}
545
546impl ConfigRuntimeManager {
547 pub fn new(
548 runtime: Arc<AgentRuntime>,
549 store: Arc<dyn ConfigStore>,
550 ) -> Result<Self, ConfigRuntimeError> {
551 let registries = runtime
552 .registry_set()
553 .ok_or(ConfigRuntimeError::RuntimeNotConfigurable)?;
554 let discovered_agents = DiscoveredAgentRegistry::from_registry(registries.agents.clone());
555
556 let mcp_registry_factory: Arc<dyn McpRegistryFactory> = Arc::new(
564 DefaultMcpRegistryFactory::with_runtime(Arc::downgrade(&runtime)),
565 );
566
567 Ok(Self {
568 runtime,
569 store,
570 tools: registries.tools,
571 plugins: registries.plugins,
572 backends: registries.backends,
573 skill_spec_sink: None,
574 discovered_agents,
575 provider_factory: Arc::new(GenaiProviderExecutorFactory),
576 change_notifier: None,
577 mcp_registry_factory,
578 apply_lock: tokio::sync::Mutex::new(()),
579 active_mcp_registry: Mutex::new(None),
580 last_applied_fingerprint: RwLock::new(None),
581 provider_cache: Mutex::new(ProviderRuntimeCache::default()),
582 periodic_refresh: PeriodicRefresher::new(),
583 change_listener: Mutex::new(None),
584 mcp_refresh_interval: RwLock::new(None),
585 min_apply_interval: Duration::ZERO,
586 audit_log: None,
587 versioned_registry: None,
588 })
589 }
590
591 #[must_use]
592 pub fn with_provider_factory(
593 mut self,
594 provider_factory: Arc<dyn ProviderExecutorFactory>,
595 ) -> Self {
596 self.provider_factory = provider_factory;
597 self
598 }
599
600 #[must_use]
601 pub fn with_change_notifier(mut self, notifier: Arc<dyn ConfigChangeNotifier>) -> Self {
602 self.change_notifier = Some(notifier);
603 self
604 }
605
606 #[must_use]
607 pub fn with_mcp_registry_factory(mut self, factory: Arc<dyn McpRegistryFactory>) -> Self {
608 self.mcp_registry_factory = factory;
609 self
610 }
611
612 #[must_use]
613 pub fn with_skill_spec_sink(mut self, sink: Arc<dyn SkillSpecSink>) -> Self {
614 self.skill_spec_sink = Some(sink);
615 self
616 }
617
618 #[must_use]
619 pub fn with_mcp_refresh_interval(self, interval: Duration) -> Self {
620 if interval.is_zero() {
621 return self;
622 }
623 *self.mcp_refresh_interval.write() = Some(interval);
624 self
625 }
626
627 #[must_use]
632 pub fn with_min_apply_interval(mut self, interval: Duration) -> Self {
633 self.min_apply_interval = interval;
634 self
635 }
636
637 #[must_use]
640 pub fn with_audit_log(mut self, logger: Arc<crate::services::audit_log::AuditLogger>) -> Self {
641 self.audit_log = Some(logger);
642 self
643 }
644
645 pub async fn apply_seed(
659 &self,
660 seed: &server_contract::BuiltinSeedSet,
661 ) -> Result<crate::services::builtin_seed::SeedReport, ConfigRuntimeError> {
662 let _guard = self.lock_apply().await;
663 let report = crate::services::builtin_seed::apply_builtin_seed(self.store.as_ref(), seed)
664 .await
665 .map_err(map_seed_error)?;
666 if let Some(audit) = &self.audit_log {
667 audit.emit_seed_report(&report).await;
668 }
669 Ok(report)
670 }
671
672 pub async fn bootstrap_if_empty(
677 &self,
678 providers: &[ProviderSpec],
679 models: &[ModelSpec],
680 agents: &[AgentSpec],
681 mcp_servers: &[McpServerSpec],
682 ) -> Result<bool, ConfigRuntimeError> {
683 let has_providers = !self.store.list(NS_PROVIDERS, 0, 1).await?.is_empty();
684 let has_models = !self.store.list(NS_MODELS, 0, 1).await?.is_empty();
685 let has_agents = !self.store.list(NS_AGENTS, 0, 1).await?.is_empty();
686 let has_mcp_servers = !self.store.list(NS_MCP_SERVERS, 0, 1).await?.is_empty();
687
688 if has_providers || has_models || has_agents || has_mcp_servers {
689 if has_providers && has_models && has_agents {
690 return Ok(false);
691 }
692 return Err(ConfigRuntimeError::PartialBootstrap);
693 }
694
695 let specs = providers
696 .iter()
697 .cloned()
698 .map(server_contract::BuiltinSpec::provider)
699 .chain(
700 models
701 .iter()
702 .cloned()
703 .map(server_contract::BuiltinSpec::model),
704 )
705 .chain(
706 agents
707 .iter()
708 .cloned()
709 .map(server_contract::BuiltinSpec::agent),
710 )
711 .chain(
712 mcp_servers
713 .iter()
714 .cloned()
715 .map(server_contract::BuiltinSpec::mcp_server),
716 )
717 .collect();
718
719 let seed = server_contract::BuiltinSeedSet {
720 binary_version: "bootstrap_if_empty".into(),
721 specs,
722 };
723 self.apply_seed(&seed).await?;
724 Ok(true)
725 }
726
727 pub async fn apply(&self) -> Result<u64, ConfigRuntimeError> {
728 let _guard = self.lock_apply().await;
729 self.apply_locked().await
730 }
731
732 pub async fn apply_if_changed(&self) -> Result<Option<u64>, ConfigRuntimeError> {
733 let _guard = self.lock_apply().await;
734 self.apply_if_changed_locked().await
735 }
736
737 pub(crate) async fn lock_apply(&self) -> tokio::sync::MutexGuard<'_, ()> {
738 self.apply_lock.lock().await
739 }
740
741 pub(crate) async fn apply_locked(&self) -> Result<u64, ConfigRuntimeError> {
742 let managed = self.load_managed_config().await?;
743 self.publish(managed).await
744 }
745
746 async fn apply_if_changed_locked(&self) -> Result<Option<u64>, ConfigRuntimeError> {
747 let managed = self.load_managed_config().await?;
748 let current_fingerprint = *self.last_applied_fingerprint.read();
749 if current_fingerprint == Some(managed.fingerprint) {
750 return Ok(None);
751 }
752 self.publish(managed).await.map(Some)
753 }
754
755 pub fn start_periodic_refresh(
756 self: &Arc<Self>,
757 interval: Duration,
758 ) -> Result<(), ConfigRuntimeError> {
759 if interval.is_zero() {
760 return Err(ConfigRuntimeError::PeriodicRefresh(
761 "interval must be non-zero".into(),
762 ));
763 }
764
765 {
766 let mut current_interval = self.mcp_refresh_interval.write();
767 if current_interval.is_none() {
768 *current_interval = Some(interval);
769 }
770 }
771
772 if let Some(active) = self.active_mcp_registry.lock().clone() {
773 self.ensure_mcp_periodic_refresh(&active.handle)?;
774 }
775 self.start_change_listener()?;
776
777 let weak = Arc::downgrade(self);
778 self.periodic_refresh
779 .start(interval, move || {
780 let weak = Weak::clone(&weak);
781 async move {
782 let Some(manager) = weak.upgrade() else {
783 return;
784 };
785 if let Err(error) = manager.apply_if_changed().await {
786 tracing::warn!(error = %error, "config periodic refresh failed");
787 }
788 }
789 })
790 .map_err(ConfigRuntimeError::PeriodicRefresh)
791 }
792
793 pub async fn stop_periodic_refresh(&self) -> bool {
794 let stopped_config = self.periodic_refresh.stop().await;
795 let stopped_listener = self.stop_change_listener().await;
796 let active = self.active_mcp_registry.lock().clone();
797 let stopped_mcp = if let Some(active) = active {
798 active.handle.stop_periodic_refresh().await
799 } else {
800 false
801 };
802 stopped_config || stopped_listener || stopped_mcp
803 }
804
805 pub async fn shutdown(&self) -> Result<(), ConfigRuntimeError> {
806 self.periodic_refresh.stop().await;
807 self.stop_change_listener().await;
808 let active = self.active_mcp_registry.lock().take();
809 if let Some(active) = active {
810 active.handle.close().await?;
811 }
812 Ok(())
813 }
814
815 pub fn periodic_refresh_running(&self) -> bool {
816 self.periodic_refresh.is_running()
817 }
818
819 pub async fn mcp_server_status(&self, server_name: &str) -> Option<McpServerStatusSnapshot> {
828 let handle = {
829 let guard = self.active_mcp_registry.lock();
830 guard.as_ref().map(|active| Arc::clone(&active.handle))
831 };
832 match handle {
833 Some(handle) => handle.server_status(server_name).await,
834 None => None,
835 }
836 }
837
838 pub async fn mcp_server_reconnect(&self, server_name: &str) -> Result<(), ConfigRuntimeError> {
843 let handle = self
844 .active_mcp_registry
845 .lock()
846 .as_ref()
847 .map(|active| Arc::clone(&active.handle));
848 match handle {
849 Some(h) => h.reconnect(server_name).await,
850 None => Err(ConfigRuntimeError::InvalidConfig(
851 "no MCP registry is active".to_string(),
852 )),
853 }
854 }
855
856 pub fn snapshot_tool_specs(&self) -> Vec<server_contract::BuiltinSpec> {
863 let mut out = Vec::new();
864 for id in self.tools.tool_ids() {
865 let Some(tool) = self.tools.get_tool(&id) else {
866 continue;
867 };
868 let descriptor = tool.descriptor();
869 out.push(server_contract::BuiltinSpec::tool(
870 server_contract::ToolSpec {
871 id: descriptor.id,
872 name: descriptor.name,
873 description: descriptor.description,
874 category: descriptor.category,
875 parameters_schema: descriptor.parameters,
876 },
877 ));
878 }
879 out.sort_by(|a, b| a.id().cmp(b.id()));
880 out
881 }
882
883 async fn prepare_mcp_registry(
884 &self,
885 specs: &[McpServerSpec],
886 ) -> Result<PreparedMcpRegistry, ConfigRuntimeError> {
887 let current = self.active_mcp_registry.lock().clone();
888 if let Some(current) = current
889 && current.specs == specs
890 {
891 self.ensure_mcp_periodic_refresh(¤t.handle)?;
892 return Ok(PreparedMcpRegistry {
893 tool_registry: Some(current.tool_registry),
894 next_state: None,
895 state_changed: false,
896 });
897 }
898
899 let mut next_state = self
900 .mcp_registry_factory
901 .connect(specs)
902 .await?
903 .map(|handle| ActiveMcpRegistry {
904 specs: specs.to_vec(),
905 tool_registry: handle.tool_registry(),
906 handle,
907 });
908
909 let refresh_error = next_state
910 .as_ref()
911 .and_then(|active| self.ensure_mcp_periodic_refresh(&active.handle).err());
912 if let Some(error) = refresh_error {
913 if let Some(active) = next_state.take()
914 && let Err(close_error) = active.handle.close().await
915 {
916 tracing::warn!(
917 error = %close_error,
918 "failed to close prepared MCP registry after refresh setup failure"
919 );
920 }
921 return Err(error);
922 }
923
924 Ok(PreparedMcpRegistry {
925 tool_registry: next_state
926 .as_ref()
927 .map(|active| active.tool_registry.clone()),
928 next_state,
929 state_changed: true,
930 })
931 }
932
933 fn ensure_mcp_periodic_refresh(
934 &self,
935 handle: &Arc<dyn ManagedMcpRegistry>,
936 ) -> Result<(), ConfigRuntimeError> {
937 let interval = *self.mcp_refresh_interval.read();
938 let Some(interval) = interval else {
939 return Ok(());
940 };
941 if handle.periodic_refresh_running() {
942 return Ok(());
943 }
944 handle.start_periodic_refresh(interval)
945 }
946
947 fn start_change_listener(self: &Arc<Self>) -> Result<(), ConfigRuntimeError> {
948 let Some(notifier) = self.change_notifier.clone() else {
949 return Ok(());
950 };
951
952 let runtime_handle = Handle::try_current()
953 .map_err(|error| ConfigRuntimeError::ChangeListener(error.to_string()))?;
954
955 let mut guard = self.change_listener.lock();
956 if guard
957 .as_ref()
958 .is_some_and(|runtime| !runtime.join.is_finished())
959 {
960 return Ok(());
961 }
962 if guard
963 .as_ref()
964 .is_some_and(|runtime| runtime.join.is_finished())
965 {
966 *guard = None;
967 }
968
969 let (stop_tx, mut stop_rx) = oneshot::channel();
970 let weak = Arc::downgrade(self);
971 let min_apply_interval = self.min_apply_interval;
972 let join = runtime_handle.spawn(async move {
973 let retry_delay = Duration::from_secs(1);
974 let mut last_applied_at: Option<tokio::time::Instant> = None;
977
978 loop {
979 let mut subscriber = tokio::select! {
980 _ = &mut stop_rx => break,
981 result = notifier.subscribe() => match result {
982 Ok(subscriber) => subscriber,
983 Err(error) => {
984 tracing::warn!(error = %error, "config change listener subscribe failed");
985 tokio::select! {
986 _ = &mut stop_rx => break,
987 _ = tokio::time::sleep(retry_delay) => continue,
988 }
989 }
990 }
991 };
992
993 loop {
994 let event = tokio::select! {
995 _ = &mut stop_rx => return,
996 result = subscriber.next() => result,
997 };
998
999 let event = match event {
1000 Ok(event) => event,
1001 Err(error) => {
1002 tracing::warn!(error = %error, "config change listener receive failed");
1003 break;
1004 }
1005 };
1006
1007 let Some(manager) = weak.upgrade() else {
1008 return;
1009 };
1010
1011 tracing::debug!(
1012 namespace = %event.namespace,
1013 id = %event.id,
1014 kind = ?event.kind,
1015 "config change notification received"
1016 );
1017
1018 if !min_apply_interval.is_zero()
1022 && let Some(last) = last_applied_at
1023 {
1024 let next_allowed = last + min_apply_interval;
1025 let now = tokio::time::Instant::now();
1026 if now < next_allowed {
1027 let wait = next_allowed - now;
1028 tokio::select! {
1029 _ = &mut stop_rx => return,
1030 _ = tokio::time::sleep(wait) => {}
1031 }
1032 loop {
1040 match tokio::time::timeout(
1041 Duration::ZERO,
1042 subscriber.next(),
1043 )
1044 .await
1045 {
1046 Ok(Ok(_event)) => continue,
1047 Ok(Err(error)) => {
1048 tracing::warn!(
1049 error = %error,
1050 "config change listener receive failed while draining debounce window"
1051 );
1052 break;
1053 }
1054 Err(_elapsed) => break,
1055 }
1056 }
1057 }
1058 }
1059
1060 if let Err(error) = manager.apply_if_changed().await {
1061 tracing::warn!(error = %error, "config change apply failed");
1062 }
1063 last_applied_at = Some(tokio::time::Instant::now());
1064 }
1065
1066 tokio::select! {
1067 _ = &mut stop_rx => break,
1068 _ = tokio::time::sleep(retry_delay) => {}
1069 }
1070 }
1071 });
1072
1073 *guard = Some(ChangeListenerRuntime {
1074 stop_tx: Some(stop_tx),
1075 join,
1076 });
1077 Ok(())
1078 }
1079
1080 async fn stop_change_listener(&self) -> bool {
1081 let runtime = {
1082 let mut guard = self.change_listener.lock();
1083 guard.take()
1084 };
1085
1086 let Some(mut runtime) = runtime else {
1087 return false;
1088 };
1089
1090 if let Some(stop_tx) = runtime.stop_tx.take() {
1091 let _ = stop_tx.send(());
1092 }
1093 let _ = runtime.join.await;
1094 true
1095 }
1096
1097 async fn load_namespace_entries(
1098 &self,
1099 namespace: &str,
1100 ) -> Result<Vec<(String, Value)>, ConfigRuntimeError> {
1101 let mut entries = Vec::new();
1102 let mut offset = 0usize;
1103
1104 loop {
1105 let page = self
1106 .store
1107 .list(namespace, offset, CONFIG_LOAD_PAGE_SIZE)
1108 .await?;
1109 if page.is_empty() {
1110 break;
1111 }
1112
1113 offset = offset.saturating_add(page.len());
1114 let reached_end = page.len() < CONFIG_LOAD_PAGE_SIZE;
1115 entries.extend(page);
1116 if reached_end {
1117 break;
1118 }
1119 }
1120
1121 Ok(entries)
1122 }
1123
1124 fn compose_tool_registry(
1125 &self,
1126 dynamic_tools: Option<Arc<dyn ToolRegistry>>,
1127 description_overrides: std::collections::HashMap<String, String>,
1128 ) -> Result<Arc<dyn ToolRegistry>, ConfigRuntimeError> {
1129 let base: Arc<dyn ToolRegistry> = if description_overrides.is_empty() {
1130 Arc::clone(&self.tools)
1131 } else {
1132 Arc::new(
1133 crate::services::tool_overrides::DescriptionOverrideRegistry::new(
1134 Arc::clone(&self.tools),
1135 description_overrides,
1136 ),
1137 ) as Arc<dyn ToolRegistry>
1138 };
1139
1140 let Some(dynamic_tools) = dynamic_tools else {
1141 return Ok(base);
1142 };
1143
1144 let base_ids: HashSet<_> = base.tool_ids().into_iter().collect();
1145 for tool_id in dynamic_tools.tool_ids() {
1146 if base_ids.contains(&tool_id) {
1147 return Err(ConfigRuntimeError::InvalidConfig(format!(
1148 "mcp tool id conflicts with existing tool: {tool_id}"
1149 )));
1150 }
1151 }
1152
1153 Ok(Arc::new(OverlayToolRegistry::new(base, dynamic_tools)) as Arc<dyn ToolRegistry>)
1154 }
1155}
1156
1157pub fn build_genai_provider_executor_with_broker(
1178 spec: &ProviderSpec,
1179 broker: Arc<dyn awaken_runtime::credentials::CredentialBroker>,
1180) -> Result<Arc<dyn LlmExecutor>, ConfigRuntimeError> {
1181 use awaken_runtime::credentials::{
1182 CredentialKind, allow_env_credentials_from_options, build_material,
1183 build_material_allowing_env_fallback,
1184 };
1185
1186 let adapter_kind = parse_adapter_kind(&spec.adapter)?;
1187 let kind = CredentialKind::from_options(&spec.adapter_options)
1188 .map_err(ConfigRuntimeError::InvalidConfig)?;
1189 let allow_env_credentials = allow_env_credentials_from_options(&spec.adapter_options)
1190 .map_err(ConfigRuntimeError::InvalidConfig)?;
1191
1192 let material = if allow_env_credentials {
1200 build_material_allowing_env_fallback(&spec.adapter, kind, spec.api_key.as_ref())
1201 } else {
1202 build_material(&spec.adapter, kind, spec.api_key.as_ref())
1203 }
1204 .map_err(ConfigRuntimeError::InvalidConfig)?;
1205
1206 let mut builder = Client::builder().with_model_mapper_fn(move |model: ModelIden| {
1207 Ok(ModelIden::new(adapter_kind, model.model_name.to_string()))
1208 });
1209
1210 if matches!(kind, CredentialKind::Bearer) {
1211 if let Some(api_key) = spec.api_key.as_ref().filter(|k| !k.is_empty()) {
1216 let key = api_key.expose_secret().to_owned();
1217 builder = builder
1218 .with_auth_resolver_fn(move |_| Ok(Some(AuthData::from_single(key.clone()))));
1219 }
1220 } else if let Some(material) = material {
1223 broker.register(spec.id.clone(), material);
1228
1229 let provider_id: Arc<str> = Arc::from(spec.id.as_str());
1230 let scope: Arc<str> = Arc::from(scopes_from_options(&spec.adapter_options)?);
1231 let broker_for_resolver = Arc::clone(&broker);
1232
1233 type ResolverFuture = std::pin::Pin<
1238 Box<dyn std::future::Future<Output = genai::resolver::Result<Option<AuthData>>> + Send>,
1239 >;
1240 let resolver_fn = move |_iden: ModelIden| -> ResolverFuture {
1241 let broker = Arc::clone(&broker_for_resolver);
1242 let provider_id = Arc::clone(&provider_id);
1243 let scope = Arc::clone(&scope);
1244 Box::pin(async move {
1245 let issued = broker.token_for(&provider_id, &scope).await.map_err(|e| {
1246 genai::resolver::Error::Custom(format!(
1247 "credential broker error for provider '{provider_id}': {e}"
1248 ))
1249 })?;
1250 Ok(Some(AuthData::from_single(issued.bearer().to_owned())))
1251 })
1252 };
1253 builder = builder.with_auth_resolver(
1254 genai::resolver::AuthResolver::from_resolver_async_fn(resolver_fn),
1255 );
1256 }
1257
1258 if let Some(base_url) = spec.base_url.clone().filter(|value| !value.is_empty()) {
1259 let normalized = if base_url.ends_with('/') {
1260 base_url
1261 } else {
1262 format!("{base_url}/")
1263 };
1264 builder = builder.with_service_target_resolver_fn(move |mut target: ServiceTarget| {
1265 target.endpoint = Endpoint::from_owned(normalized.clone());
1266 Ok(target)
1267 });
1268 }
1269
1270 if let Some(headers) = build_default_headers_from_options(&spec.adapter_options)? {
1271 builder = builder.with_web_config(WebConfig::default().with_default_headers(headers));
1272 }
1273
1274 let client = builder.build();
1275 let executor = GenaiExecutor::with_client(client)
1276 .with_timeout(Duration::from_secs(spec.timeout_secs.max(1)));
1277 Ok(Arc::new(executor))
1278}
1279
1280pub fn build_genai_provider_executor(
1286 spec: &ProviderSpec,
1287) -> Result<Arc<dyn LlmExecutor>, ConfigRuntimeError> {
1288 build_genai_provider_executor_with_broker(
1289 spec,
1290 Arc::new(awaken_runtime::credentials::AwakenCredentialBroker::new()),
1291 )
1292}
1293
1294const DEFAULT_OAUTH_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
1297
1298fn scopes_from_options(options: &BTreeMap<String, Value>) -> Result<String, ConfigRuntimeError> {
1302 let Some(value) = options.get("scopes") else {
1303 return Ok(DEFAULT_OAUTH_SCOPE.to_owned());
1304 };
1305 let arr = value.as_array().ok_or_else(|| {
1306 ConfigRuntimeError::InvalidConfig(
1307 "adapter_options.scopes must be an array of strings".into(),
1308 )
1309 })?;
1310 if arr.is_empty() {
1311 return Ok(DEFAULT_OAUTH_SCOPE.to_owned());
1312 }
1313 let mut joined = String::new();
1314 for (i, item) in arr.iter().enumerate() {
1315 let s = item.as_str().ok_or_else(|| {
1316 ConfigRuntimeError::InvalidConfig(
1317 "adapter_options.scopes must be an array of strings".into(),
1318 )
1319 })?;
1320 if i > 0 {
1321 joined.push(' ');
1322 }
1323 joined.push_str(s);
1324 }
1325 Ok(joined)
1326}
1327
1328fn build_default_headers_from_options(
1336 options: &BTreeMap<String, Value>,
1337) -> Result<Option<HeaderMap>, ConfigRuntimeError> {
1338 let Some(headers_value) = options.get("headers") else {
1339 return Ok(None);
1340 };
1341 let entries = headers_value.as_object().ok_or_else(|| {
1342 ConfigRuntimeError::InvalidConfig(
1343 "adapter_options.headers must be an object of string -> string pairs".into(),
1344 )
1345 })?;
1346
1347 let mut map = HeaderMap::with_capacity(entries.len());
1348 for (name, value) in entries {
1349 let value_str = value.as_str().ok_or_else(|| {
1350 ConfigRuntimeError::InvalidConfig(format!(
1351 "adapter_options.headers[{name}] must be a string"
1352 ))
1353 })?;
1354 let header_name = HeaderName::try_from(name).map_err(|err| {
1355 ConfigRuntimeError::InvalidConfig(format!(
1356 "adapter_options.headers[{name}] invalid header name: {err}"
1357 ))
1358 })?;
1359 let header_value = HeaderValue::from_str(value_str).map_err(|err| {
1360 ConfigRuntimeError::InvalidConfig(format!(
1361 "adapter_options.headers[{name}] invalid header value: {err}"
1362 ))
1363 })?;
1364 map.insert(header_name, header_value);
1365 }
1366 Ok(Some(map))
1367}
1368
1369const ADAPTER_CANDIDATES: &[&str] = &[
1386 "anthropic",
1388 "openai",
1389 "openai_resp",
1390 "deepseek",
1391 "gemini",
1392 "ollama",
1393 "ollama_cloud",
1394 "cohere",
1395 "together",
1396 "fireworks",
1397 "groq",
1398 "xai",
1399 "zai",
1400 "bigmodel",
1401 "aliyun",
1402 "mimo",
1403 "nebius",
1404 "vertex",
1405 "github_copilot",
1406 "bedrock",
1408 "azure",
1409 "azure_openai",
1410 "mistral",
1411 "perplexity",
1412 "watsonx",
1413 "huggingface",
1414 "replicate",
1415];
1416
1417pub fn supported_adapters() -> Vec<&'static str> {
1423 ADAPTER_CANDIDATES
1424 .iter()
1425 .copied()
1426 .filter(|name| AdapterKind::from_lower_str(name).is_some())
1427 .collect()
1428}
1429
1430fn parse_adapter_kind(adapter: &str) -> Result<AdapterKind, ConfigRuntimeError> {
1431 let normalized = adapter.trim().to_ascii_lowercase();
1432 if matches!(normalized.as_str(), "openai-resp" | "responses") {
1435 return Ok(AdapterKind::OpenAIResp);
1436 }
1437 AdapterKind::from_lower_str(&normalized)
1438 .ok_or_else(|| ConfigRuntimeError::UnsupportedProviderAdapter(adapter.to_string()))
1439}
1440
1441fn mcp_spec_to_connection_config(
1442 spec: &McpServerSpec,
1443) -> Result<McpServerConnectionConfig, ConfigRuntimeError> {
1444 if spec.id.trim().is_empty() {
1445 return Err(ConfigRuntimeError::InvalidConfig(
1446 "mcp server id cannot be empty".into(),
1447 ));
1448 }
1449
1450 let mut config = match spec.transport {
1451 McpTransportKind::Stdio => {
1452 let command = spec
1453 .command
1454 .as_deref()
1455 .map(str::trim)
1456 .filter(|value| !value.is_empty())
1457 .ok_or_else(|| {
1458 ConfigRuntimeError::InvalidConfig(format!(
1459 "mcp server '{}' requires a non-empty command",
1460 spec.id
1461 ))
1462 })?;
1463 McpServerConnectionConfig::stdio(spec.id.clone(), command, spec.args.clone())
1464 }
1465 McpTransportKind::Http => {
1466 let url = spec
1467 .url
1468 .as_deref()
1469 .map(str::trim)
1470 .filter(|value| !value.is_empty())
1471 .ok_or_else(|| {
1472 ConfigRuntimeError::InvalidConfig(format!(
1473 "mcp server '{}' requires a non-empty url",
1474 spec.id
1475 ))
1476 })?;
1477 McpServerConnectionConfig::http(spec.id.clone(), url)
1478 }
1479 };
1480
1481 config.timeout_secs = spec.timeout_secs.max(1);
1482 config.config = Value::Object(spec.config.clone());
1483 config.env = spec.env.clone().into_iter().collect();
1484 config.restart_policy = restart_policy_to_connection_policy(&spec.restart_policy);
1485 Ok(config)
1486}
1487
1488fn restart_policy_to_connection_policy(policy: &McpRestartPolicy) -> mcp::transport::RestartPolicy {
1489 mcp::transport::RestartPolicy {
1490 enabled: policy.enabled,
1491 max_attempts: policy.max_attempts,
1492 delay_ms: policy.delay_ms,
1493 backoff_multiplier: policy.backoff_multiplier,
1494 max_delay_ms: policy.max_delay_ms,
1495 }
1496}
1497
1498fn deserialize_namespace<T>(entries: &[(String, Value)]) -> Result<Vec<T>, ConfigRuntimeError>
1499where
1500 T: serde::de::DeserializeOwned + server_contract::ConfigRecordMerge,
1501{
1502 let mut out = Vec::with_capacity(entries.len());
1503 for (_, value) in entries {
1504 let raw_record: ConfigRecord<Value> = ConfigRecord::from_value(value.clone())
1505 .map_err(|error| StorageError::Serialization(error.to_string()))
1506 .map_err(ConfigRuntimeError::Storage)?;
1507 if raw_record.meta.hidden {
1508 continue;
1509 }
1510
1511 let record: ConfigRecord<T> = server_contract::validate_config_record(value.clone())
1512 .map_err(|error| StorageError::Serialization(error.to_string()))
1513 .map_err(ConfigRuntimeError::Storage)?;
1514 let effective = crate::services::config_envelope::apply_overrides(
1515 record.spec,
1516 record.meta.user_overrides.as_ref(),
1517 )
1518 .map_err(|error| StorageError::Serialization(error.to_string()))
1519 .map_err(ConfigRuntimeError::Storage)?;
1520 out.push(effective);
1521 }
1522 Ok(out)
1523}
1524
1525fn fingerprint_config(
1526 namespaces: &[(&str, &[(String, Value)])],
1527) -> Result<u64, ConfigRuntimeError> {
1528 let mut hasher = DefaultHasher::new();
1529
1530 for (namespace, entries) in namespaces {
1531 namespace.hash(&mut hasher);
1532 entries.len().hash(&mut hasher);
1533
1534 for (id, value) in *entries {
1535 id.hash(&mut hasher);
1536 let canonical = canonicalize_value(value);
1537 let serialized = serde_json::to_vec(&canonical)
1538 .map_err(|error| StorageError::Serialization(error.to_string()))
1539 .map_err(ConfigRuntimeError::Storage)?;
1540 serialized.hash(&mut hasher);
1541 }
1542 }
1543
1544 Ok(hasher.finish())
1545}
1546
1547fn canonicalize_value(value: &Value) -> Value {
1548 match value {
1549 Value::Array(values) => Value::Array(values.iter().map(canonicalize_value).collect()),
1550 Value::Object(object) => {
1551 let mut keys = object.keys().cloned().collect::<Vec<_>>();
1552 keys.sort();
1553
1554 let mut normalized = serde_json::Map::new();
1555 for key in keys {
1556 if let Some(value) = object.get(&key) {
1557 normalized.insert(key, canonicalize_value(value));
1558 }
1559 }
1560 Value::Object(normalized)
1561 }
1562 _ => value.clone(),
1563 }
1564}
1565
1566fn map_seed_error(error: crate::services::builtin_seed::SeedError) -> ConfigRuntimeError {
1567 use crate::services::builtin_seed::SeedError as E;
1568 use ConfigRuntimeError::{InvalidConfig, Storage};
1569 match error {
1570 E::Storage(e) => Storage(e),
1571 e @ E::Serde(_) => Storage(StorageError::Serialization(e.to_string())),
1572 e => InvalidConfig(e.to_string()),
1573 }
1574}
1575
1576#[cfg(test)]
1577mod tests {
1578 use super::*;
1579 use serde_json::json;
1580 use std::collections::BTreeMap;
1581
1582 fn provider_spec_with_options(adapter_options: BTreeMap<String, Value>) -> ProviderSpec {
1583 ProviderSpec {
1584 id: "test".into(),
1585 adapter: "openai".into(),
1586 api_key: Some("test-secret-key".to_string().into()),
1587 adapter_options,
1588 ..ProviderSpec::default()
1589 }
1590 }
1591
1592 fn test_broker() -> Arc<dyn awaken_runtime::credentials::CredentialBroker> {
1596 Arc::new(awaken_runtime::credentials::AwakenCredentialBroker::new())
1597 }
1598
1599 #[tokio::test]
1606 async fn registry_factory_returns_none_when_runtime_dropped() {
1607 let runtime = Arc::new(
1610 awaken_runtime::AgentRuntimeBuilder::new()
1611 .build()
1612 .expect("minimal runtime builds"),
1613 );
1614 let weak = Arc::downgrade(&runtime);
1615 drop(runtime);
1616
1617 let factory = RegistryDrivenSamplingHandlerFactory::new(weak);
1618 let spec = AgentSpec {
1619 id: "alpha".into(),
1620 model_id: "any-model".into(),
1621 system_prompt: String::new(),
1622 ..AgentSpec::default()
1623 };
1624 assert!(
1625 factory.for_agent(&spec).await.is_none(),
1626 "factory must not produce a handler for a dropped runtime"
1627 );
1628 }
1629
1630 #[tokio::test]
1638 async fn registry_factory_cache_invalidates_on_registry_replace() {
1639 use awaken_runtime::registry::RegistrySet;
1640 use awaken_runtime::registry::memory::{
1641 MapAgentSpecRegistry, MapModelRegistry, MapPluginSource, MapProviderRegistry,
1642 MapToolRegistry,
1643 };
1644 use server_contract::contract::executor::{InferenceExecutionError, InferenceRequest};
1645 use server_contract::contract::inference::{StreamResult, TokenUsage};
1646
1647 struct LabelExecutor {
1652 label: &'static str,
1653 }
1654 #[async_trait::async_trait]
1655 impl LlmExecutor for LabelExecutor {
1656 async fn execute(
1657 &self,
1658 _request: InferenceRequest,
1659 ) -> Result<StreamResult, InferenceExecutionError> {
1660 Ok(StreamResult {
1661 content: vec![server_contract::contract::content::ContentBlock::text(
1662 self.label.to_string(),
1663 )],
1664 tool_calls: vec![],
1665 usage: Some(TokenUsage::default()),
1666 stop_reason: None,
1667 has_incomplete_tool_calls: false,
1668 })
1669 }
1670 fn name(&self) -> &str {
1671 "label"
1672 }
1673 }
1674
1675 fn registry_set_pointing_at(executor_label: &'static str) -> RegistrySet {
1676 let mut agents = MapAgentSpecRegistry::new();
1677 agents
1678 .register_spec(AgentSpec {
1679 id: "alpha".into(),
1680 model_id: "m".into(),
1681 system_prompt: String::new(),
1682 ..Default::default()
1683 })
1684 .unwrap();
1685 let mut models = MapModelRegistry::new();
1686 models
1687 .register_model(ModelSpec::new(
1688 "m",
1689 "p",
1690 format!("upstream-{executor_label}"),
1691 ))
1692 .unwrap();
1693 let mut providers = MapProviderRegistry::new();
1694 providers
1695 .register_provider(
1696 "p",
1697 Arc::new(LabelExecutor {
1698 label: executor_label,
1699 }) as Arc<dyn LlmExecutor>,
1700 )
1701 .unwrap();
1702 RegistrySet {
1703 agents: Arc::new(agents),
1704 tools: Arc::new(MapToolRegistry::new()),
1705 models: Arc::new(models),
1706 providers: Arc::new(providers),
1707 plugins: Arc::new(MapPluginSource::new()),
1708 backends: Arc::new(awaken_runtime::registry::memory::MapBackendRegistry::new()),
1709 }
1710 }
1711
1712 let runtime = Arc::new(
1713 awaken_runtime::AgentRuntimeBuilder::new()
1714 .with_provider("p", Arc::new(LabelExecutor { label: "v1" }))
1715 .with_model(ModelSpec::new("m", "p", "upstream-v1"))
1716 .with_agent_spec(AgentSpec {
1717 id: "alpha".into(),
1718 model_id: "m".into(),
1719 system_prompt: String::new(),
1720 ..Default::default()
1721 })
1722 .build()
1723 .expect("v1 runtime builds"),
1724 );
1725
1726 let v1_version = runtime.registry_version().expect("v1 registered");
1727 let factory = RegistryDrivenSamplingHandlerFactory::new(Arc::downgrade(&runtime));
1728 let spec = AgentSpec {
1729 id: "alpha".into(),
1730 model_id: "m".into(),
1731 system_prompt: String::new(),
1732 ..AgentSpec::default()
1733 };
1734
1735 let handler_v1 = factory
1736 .for_agent(&spec)
1737 .await
1738 .expect("factory yields handler for v1");
1739
1740 let handler_v1_again = factory
1742 .for_agent(&spec)
1743 .await
1744 .expect("factory yields handler for v1 again");
1745 assert!(
1746 Arc::ptr_eq(&handler_v1, &handler_v1_again),
1747 "second lookup at the same registry version must be a cache hit (same Arc)"
1748 );
1749
1750 let new_version = runtime
1755 .replace_registry_set(registry_set_pointing_at("v2"))
1756 .expect("replace_registry_set yields a fresh version");
1757 assert!(
1758 new_version > v1_version,
1759 "replace_registry_set must bump registry_version (was {v1_version}, got {new_version})"
1760 );
1761
1762 let handler_v2 = factory
1763 .for_agent(&spec)
1764 .await
1765 .expect("factory yields handler for v2");
1766 assert!(
1767 !Arc::ptr_eq(&handler_v1, &handler_v2),
1768 "cache must invalidate on registry_version bump — got the same Arc back, \
1769 which means sampling would still route to the OLD executor"
1770 );
1771 }
1772
1773 #[test]
1778 fn default_factory_has_no_sampling_handler_factory() {
1779 let factory = DefaultMcpRegistryFactory::new();
1780 assert!(factory.sampling_handler_factory.is_none());
1781 let factory = DefaultMcpRegistryFactory::default();
1782 assert!(factory.sampling_handler_factory.is_none());
1783 }
1784
1785 #[test]
1786 fn build_genai_with_valid_headers_succeeds() {
1787 let mut options = BTreeMap::new();
1788 options.insert("headers".into(), json!({"OpenAI-Organization": "org-xyz"}));
1789 let spec = provider_spec_with_options(options);
1790 build_genai_provider_executor_with_broker(&spec, test_broker())
1791 .expect("valid headers must build");
1792 }
1793
1794 #[test]
1795 fn build_genai_rejects_non_object_headers() {
1796 let mut options = BTreeMap::new();
1797 options.insert("headers".into(), json!("not-an-object"));
1798 let spec = provider_spec_with_options(options);
1799 let err = match build_genai_provider_executor_with_broker(&spec, test_broker()) {
1800 Ok(_) => panic!("expected build to fail"),
1801 Err(e) => e,
1802 };
1803 assert!(
1804 matches!(err, ConfigRuntimeError::InvalidConfig(ref msg) if msg.contains("headers")),
1805 "expected InvalidConfig mentioning headers, got: {err:?}"
1806 );
1807 }
1808
1809 #[test]
1810 fn build_genai_rejects_non_string_header_value() {
1811 let mut options = BTreeMap::new();
1812 options.insert("headers".into(), json!({"X-Numeric-Value": 42}));
1813 let spec = provider_spec_with_options(options);
1814 let err = match build_genai_provider_executor_with_broker(&spec, test_broker()) {
1815 Ok(_) => panic!("expected build to fail"),
1816 Err(e) => e,
1817 };
1818 assert!(
1819 matches!(err, ConfigRuntimeError::InvalidConfig(ref msg) if msg.contains("X-Numeric-Value")),
1820 "expected InvalidConfig naming the bad header, got: {err:?}"
1821 );
1822 }
1823
1824 #[test]
1825 fn build_genai_ignores_unknown_adapter_options() {
1826 let mut options = BTreeMap::new();
1827 options.insert("future_extension_key".into(), json!({"anything": true}));
1828 let spec = provider_spec_with_options(options);
1829 build_genai_provider_executor_with_broker(&spec, test_broker())
1830 .expect("unknown adapter_options keys must not break the build");
1831 }
1832
1833 #[test]
1834 fn build_default_headers_returns_none_when_absent() {
1835 let result = build_default_headers_from_options(&BTreeMap::new()).unwrap();
1836 assert!(result.is_none());
1837 }
1838
1839 #[test]
1840 fn build_default_headers_parses_string_pairs() {
1841 let mut options = BTreeMap::new();
1842 options.insert(
1843 "headers".into(),
1844 json!({
1845 "OpenAI-Organization": "org-xyz",
1846 "X-Custom": "value",
1847 }),
1848 );
1849 let map = build_default_headers_from_options(&options)
1850 .unwrap()
1851 .expect("headers should be present");
1852 assert_eq!(
1853 map.get("openai-organization").and_then(|v| v.to_str().ok()),
1854 Some("org-xyz")
1855 );
1856 assert_eq!(
1857 map.get("x-custom").and_then(|v| v.to_str().ok()),
1858 Some("value")
1859 );
1860 }
1861
1862 #[test]
1863 fn build_default_headers_rejects_invalid_header_name() {
1864 let mut options = BTreeMap::new();
1865 options.insert("headers".into(), json!({"Invalid Header Name": "value"}));
1866 let err = build_default_headers_from_options(&options).unwrap_err();
1867 assert!(
1868 matches!(err, ConfigRuntimeError::InvalidConfig(ref msg) if msg.contains("Invalid Header Name")),
1869 "expected InvalidConfig naming the bad header, got: {err:?}"
1870 );
1871 }
1872
1873 #[test]
1874 fn supported_adapters_round_trip_through_parser() {
1875 for name in supported_adapters() {
1876 let parsed = parse_adapter_kind(name)
1877 .unwrap_or_else(|err| panic!("supported adapter {name} must parse: {err:?}"));
1878 assert_eq!(
1879 parsed.as_lower_str(),
1880 name,
1881 "as_lower_str round-trip mismatch for {name}"
1882 );
1883 }
1884 }
1885
1886 #[test]
1887 fn scopes_from_options_default_when_absent() {
1888 assert_eq!(
1889 scopes_from_options(&BTreeMap::new()).unwrap(),
1890 DEFAULT_OAUTH_SCOPE
1891 );
1892 }
1893
1894 #[test]
1895 fn scopes_from_options_joins_array_with_spaces() {
1896 let mut options = BTreeMap::new();
1897 options.insert(
1898 "scopes".into(),
1899 json!(["a.googleapis.com/auth/x", "b.googleapis.com/auth/y"]),
1900 );
1901 assert_eq!(
1902 scopes_from_options(&options).unwrap(),
1903 "a.googleapis.com/auth/x b.googleapis.com/auth/y"
1904 );
1905 }
1906
1907 #[test]
1908 fn scopes_from_options_rejects_non_array() {
1909 let mut options = BTreeMap::new();
1910 options.insert("scopes".into(), json!("not-an-array"));
1911 let err = scopes_from_options(&options).unwrap_err();
1912 assert!(matches!(err, ConfigRuntimeError::InvalidConfig(ref m) if m.contains("scopes")));
1913 }
1914
1915 #[test]
1916 fn scopes_from_options_rejects_non_string_entry() {
1917 let mut options = BTreeMap::new();
1918 options.insert("scopes".into(), json!([42]));
1919 let err = scopes_from_options(&options).unwrap_err();
1920 assert!(matches!(err, ConfigRuntimeError::InvalidConfig(ref m) if m.contains("scopes")));
1921 }
1922
1923 #[test]
1924 fn parse_adapter_kind_rejects_unknown() {
1925 let err = parse_adapter_kind("not-a-real-adapter").unwrap_err();
1926 assert!(
1927 matches!(err, ConfigRuntimeError::UnsupportedProviderAdapter(ref s) if s == "not-a-real-adapter"),
1928 "expected UnsupportedProviderAdapter, got: {err:?}"
1929 );
1930 }
1931
1932 #[tokio::test]
1936 async fn apply_seed_writes_builtin_envelope() {
1937 use server_contract::{
1938 BuiltinSeedSet, BuiltinSpec, ConfigRecord, ModelSpec, ProviderSpec, RecordSource,
1939 };
1940
1941 let bin_version = "test-env-ver".to_owned();
1942 let (manager, store) = make_manager_with_store().await;
1943
1944 let seed = BuiltinSeedSet {
1945 binary_version: bin_version.clone(),
1946 specs: vec![
1947 BuiltinSpec::Provider(ProviderSpec {
1948 id: "p1".into(),
1949 adapter: "openai".into(),
1950 ..Default::default()
1951 }),
1952 BuiltinSpec::Model(ModelSpec::new("m1", "p1", "m1-model")),
1953 BuiltinSpec::Agent(Box::new(AgentSpec {
1954 id: "a1".into(),
1955 model_id: "m1".into(),
1956 system_prompt: "seed test".into(),
1957 max_rounds: 1,
1958 ..Default::default()
1959 })),
1960 ],
1961 };
1962
1963 let report = manager.apply_seed(&seed).await.expect("apply_seed");
1964 assert_eq!(report.created.len(), 3, "all three specs must be created");
1965
1966 let raw_p = server_contract::contract::config_store::ConfigStore::get(
1968 store.as_ref(),
1969 "providers",
1970 "p1",
1971 )
1972 .await
1973 .expect("get provider")
1974 .expect("provider present");
1975
1976 let p_obj = raw_p.as_object().expect("must be object");
1977 assert!(p_obj.contains_key("spec"), "provider must have 'spec' key");
1978 assert!(p_obj.contains_key("meta"), "provider must have 'meta' key");
1979 let p_rec: ConfigRecord<serde_json::Value> = ConfigRecord::from_value(raw_p).unwrap();
1980 assert_eq!(
1981 p_rec.meta.source,
1982 RecordSource::Builtin {
1983 binary_version: bin_version.clone()
1984 },
1985 "provider source must be Builtin with correct binary_version"
1986 );
1987
1988 let raw_a = server_contract::contract::config_store::ConfigStore::get(
1990 store.as_ref(),
1991 "agents",
1992 "a1",
1993 )
1994 .await
1995 .expect("get agent")
1996 .expect("agent present");
1997 let a_rec: ConfigRecord<serde_json::Value> = ConfigRecord::from_value(raw_a).unwrap();
1998 assert_eq!(
1999 a_rec.meta.source,
2000 RecordSource::Builtin {
2001 binary_version: bin_version.clone()
2002 },
2003 "agent source must be Builtin"
2004 );
2005
2006 let raw_m = server_contract::contract::config_store::ConfigStore::get(
2008 store.as_ref(),
2009 "models",
2010 "m1",
2011 )
2012 .await
2013 .expect("get model")
2014 .expect("model present");
2015 let m_rec: ConfigRecord<serde_json::Value> = ConfigRecord::from_value(raw_m).unwrap();
2016 assert_eq!(
2017 m_rec.meta.source,
2018 RecordSource::Builtin {
2019 binary_version: bin_version
2020 },
2021 "model source must be Builtin"
2022 );
2023 }
2024
2025 pub(super) async fn make_manager_with_store() -> (
2026 ConfigRuntimeManager,
2027 Arc<dyn server_contract::contract::config_store::ConfigStore>,
2028 ) {
2029 use awaken_stores::InMemoryStore;
2030 use server_contract::contract::executor::{
2031 InferenceExecutionError, InferenceRequest, LlmExecutor,
2032 };
2033 use server_contract::contract::inference::{StopReason, StreamResult, TokenUsage};
2034
2035 struct Stub;
2036 #[async_trait::async_trait]
2037 impl LlmExecutor for Stub {
2038 async fn execute(
2039 &self,
2040 _: InferenceRequest,
2041 ) -> Result<StreamResult, InferenceExecutionError> {
2042 Ok(StreamResult {
2043 content: vec![],
2044 tool_calls: vec![],
2045 usage: Some(TokenUsage::default()),
2046 stop_reason: Some(StopReason::EndTurn),
2047 has_incomplete_tool_calls: false,
2048 })
2049 }
2050 fn name(&self) -> &str {
2051 "stub"
2052 }
2053 }
2054 impl ProviderExecutorFactory for Stub {
2055 fn build(
2056 &self,
2057 _spec: &ProviderSpec,
2058 ) -> Result<Arc<dyn LlmExecutor>, ConfigRuntimeError> {
2059 Ok(Arc::new(Stub))
2060 }
2061 }
2062
2063 let store = Arc::new(InMemoryStore::new())
2064 as Arc<dyn server_contract::contract::config_store::ConfigStore>;
2065 let thread_store = Arc::new(InMemoryStore::new());
2066 let runtime = Arc::new(
2067 awaken_runtime::builder::AgentRuntimeBuilder::new()
2068 .with_provider("boot", Arc::new(Stub))
2069 .with_model(ModelSpec::new("boot", "boot", "boot-model"))
2070 .with_agent_spec(AgentSpec {
2071 id: "boot".into(),
2072 model_id: "boot".into(),
2073 system_prompt: "boot".into(),
2074 max_rounds: 1,
2075 ..Default::default()
2076 })
2077 .with_in_memory_thread_run_store(thread_store.clone())
2078 .build()
2079 .expect("build runtime"),
2080 );
2081 let manager = ConfigRuntimeManager::new(runtime, store.clone())
2082 .expect("manager")
2083 .with_provider_factory(Arc::new(Stub));
2084 (manager, store)
2085 }
2086
2087 #[tokio::test]
2088 async fn publish_closes_prepared_mcp_registry_when_refresh_start_fails() {
2089 use std::sync::atomic::{AtomicUsize, Ordering};
2090
2091 struct RefreshFailingRegistry {
2092 tool_registry: Arc<dyn ToolRegistry>,
2093 close_count: Arc<AtomicUsize>,
2094 }
2095
2096 #[async_trait::async_trait]
2097 impl ManagedMcpRegistry for RefreshFailingRegistry {
2098 fn tool_registry(&self) -> Arc<dyn ToolRegistry> {
2099 Arc::clone(&self.tool_registry)
2100 }
2101
2102 fn periodic_refresh_running(&self) -> bool {
2103 false
2104 }
2105
2106 fn start_periodic_refresh(
2107 &self,
2108 _interval: Duration,
2109 ) -> Result<(), ConfigRuntimeError> {
2110 Err(ConfigRuntimeError::PeriodicRefresh(
2111 "scripted MCP refresh failure".to_string(),
2112 ))
2113 }
2114
2115 async fn stop_periodic_refresh(&self) -> bool {
2116 false
2117 }
2118
2119 async fn close(&self) -> Result<(), ConfigRuntimeError> {
2120 self.close_count.fetch_add(1, Ordering::SeqCst);
2121 Ok(())
2122 }
2123 }
2124
2125 struct RefreshFailingFactory {
2126 close_count: Arc<AtomicUsize>,
2127 }
2128
2129 #[async_trait::async_trait]
2130 impl McpRegistryFactory for RefreshFailingFactory {
2131 async fn connect(
2132 &self,
2133 specs: &[McpServerSpec],
2134 ) -> Result<Option<Arc<dyn ManagedMcpRegistry>>, ConfigRuntimeError> {
2135 assert!(!specs.is_empty(), "test must exercise a real MCP state");
2136 Ok(Some(Arc::new(RefreshFailingRegistry {
2137 tool_registry: Arc::new(
2138 awaken_runtime::registry::memory::MapToolRegistry::new(),
2139 ),
2140 close_count: Arc::clone(&self.close_count),
2141 }) as Arc<dyn ManagedMcpRegistry>))
2142 }
2143 }
2144
2145 let close_count = Arc::new(AtomicUsize::new(0));
2146 let (manager, _) = make_manager_with_store().await;
2147 let manager = manager.with_mcp_registry_factory(Arc::new(RefreshFailingFactory {
2148 close_count: Arc::clone(&close_count),
2149 }));
2150 *manager.mcp_refresh_interval.write() = Some(Duration::from_secs(30));
2151
2152 let err = manager
2153 .publish(ManagedConfigSnapshot {
2154 providers: Vec::new(),
2155 models: Vec::new(),
2156 pools: Vec::new(),
2157 agents: Vec::new(),
2158 a2a_servers: Vec::new(),
2159 mcp_servers: vec![McpServerSpec {
2160 id: "demo".to_string(),
2161 transport: McpTransportKind::Http,
2162 url: Some("http://mcp.example.invalid".to_string()),
2163 ..McpServerSpec::default()
2164 }],
2165 tools: Vec::new(),
2166 skills: Vec::new(),
2167 source_config_revisions: Vec::new(),
2168 fingerprint: 1,
2169 })
2170 .await
2171 .expect_err("refresh setup failure must abort publish");
2172
2173 assert!(
2174 matches!(err, ConfigRuntimeError::PeriodicRefresh(_)),
2175 "unexpected error: {err:?}"
2176 );
2177 assert_eq!(
2178 close_count.load(Ordering::SeqCst),
2179 1,
2180 "prepared MCP registry must be closed when refresh setup fails"
2181 );
2182 }
2183
2184 #[tokio::test]
2185 async fn apply_seed_writes_builtin_records_to_store() {
2186 use server_contract::{
2187 BuiltinSeedSet, BuiltinSpec, ConfigRecord, ModelSpec, ProviderSpec, RecordSource,
2188 };
2189
2190 let (manager, store) = make_manager_with_store().await;
2191
2192 let seed = BuiltinSeedSet {
2193 binary_version: "v1-test".to_owned(),
2194 specs: vec![
2195 BuiltinSpec::Agent(Box::new(AgentSpec {
2196 id: "seed-agent".into(),
2197 model_id: "m".into(),
2198 system_prompt: "hello".into(),
2199 max_rounds: 1,
2200 ..Default::default()
2201 })),
2202 BuiltinSpec::Provider(ProviderSpec {
2203 id: "seed-provider".into(),
2204 adapter: "openai".into(),
2205 ..Default::default()
2206 }),
2207 BuiltinSpec::Model(ModelSpec::new("seed-model", "seed-provider", "gpt-4o")),
2208 ],
2209 };
2210
2211 let report = manager.apply_seed(&seed).await.expect("apply_seed");
2212 assert_eq!(report.created.len(), 3, "expected 3 created");
2213 assert!(report.updated.is_empty());
2214 assert!(report.unchanged.is_empty());
2215
2216 let raw = server_contract::contract::config_store::ConfigStore::get(
2218 store.as_ref(),
2219 "agents",
2220 "seed-agent",
2221 )
2222 .await
2223 .expect("get agent")
2224 .expect("agent must be present");
2225
2226 let rec: ConfigRecord<serde_json::Value> = ConfigRecord::from_value(raw).unwrap();
2227 assert_eq!(
2228 rec.meta.source,
2229 RecordSource::Builtin {
2230 binary_version: "v1-test".to_owned()
2231 },
2232 "source must be Builtin with seed binary_version"
2233 );
2234 }
2235
2236 #[tokio::test]
2237 async fn apply_seed_idempotent() {
2238 use server_contract::{BuiltinSeedSet, BuiltinSpec, ModelSpec, ProviderSpec};
2239
2240 let (manager, _store) = make_manager_with_store().await;
2241
2242 let seed = BuiltinSeedSet {
2243 binary_version: "v1-idem".to_owned(),
2244 specs: vec![
2245 BuiltinSpec::Agent(Box::new(AgentSpec {
2246 id: "idem-agent".into(),
2247 model_id: "m".into(),
2248 system_prompt: "hello".into(),
2249 max_rounds: 1,
2250 ..Default::default()
2251 })),
2252 BuiltinSpec::Provider(ProviderSpec {
2253 id: "idem-provider".into(),
2254 adapter: "openai".into(),
2255 ..Default::default()
2256 }),
2257 BuiltinSpec::Model(ModelSpec::new("idem-model", "idem-provider", "gpt-4o")),
2258 ],
2259 };
2260
2261 manager.apply_seed(&seed).await.expect("first apply_seed");
2262 let report = manager.apply_seed(&seed).await.expect("second apply_seed");
2263
2264 assert_eq!(
2265 report.unchanged.len(),
2266 3,
2267 "second call must report 3 unchanged"
2268 );
2269 assert!(report.created.is_empty());
2270 assert!(report.updated.is_empty());
2271 }
2272
2273 #[tokio::test]
2280 async fn apply_seed_serializes_with_apply_lock() {
2281 use server_contract::{BuiltinSeedSet, BuiltinSpec, ProviderSpec};
2282 use std::sync::Arc;
2283
2284 let (manager, _store) = make_manager_with_store().await;
2285 let manager = Arc::new(manager);
2286
2287 let guard = manager.lock_apply().await;
2289
2290 let manager2 = Arc::clone(&manager);
2291 let seed = BuiltinSeedSet {
2292 binary_version: "lock-test".to_owned(),
2293 specs: vec![BuiltinSpec::Provider(ProviderSpec {
2294 id: "lock-prov".into(),
2295 adapter: "openai".into(),
2296 ..Default::default()
2297 })],
2298 };
2299
2300 let handle = tokio::spawn(async move {
2301 manager2
2302 .apply_seed(&seed)
2303 .await
2304 .expect("apply_seed in task")
2305 });
2306
2307 tokio::task::yield_now().await;
2309 assert!(
2310 !handle.is_finished(),
2311 "apply_seed must block while apply-lock is held"
2312 );
2313
2314 drop(guard);
2316 let report = handle.await.expect("task must not panic");
2317 assert_eq!(
2318 report.created.len(),
2319 1,
2320 "seed record must be created after lock release"
2321 );
2322 }
2323
2324 #[tokio::test]
2327 async fn discovered_agent_overlays_seeded_agent_with_same_id() {
2328 use server_contract::registry_spec::RemoteEndpoint;
2329 use server_contract::{BuiltinSeedSet, BuiltinSpec};
2330
2331 struct Stub;
2332 #[async_trait::async_trait]
2333 impl server_contract::contract::executor::LlmExecutor for Stub {
2334 async fn execute(
2335 &self,
2336 _: server_contract::contract::executor::InferenceRequest,
2337 ) -> Result<
2338 server_contract::contract::inference::StreamResult,
2339 server_contract::contract::executor::InferenceExecutionError,
2340 > {
2341 Ok(server_contract::contract::inference::StreamResult {
2342 content: vec![],
2343 tool_calls: vec![],
2344 usage: Some(server_contract::contract::inference::TokenUsage::default()),
2345 stop_reason: Some(server_contract::contract::inference::StopReason::EndTurn),
2346 has_incomplete_tool_calls: false,
2347 })
2348 }
2349 fn name(&self) -> &str {
2350 "stub"
2351 }
2352 }
2353
2354 let store = Arc::new(awaken_stores::InMemoryStore::new())
2355 as Arc<dyn server_contract::contract::config_store::ConfigStore>;
2356 let thread_store = Arc::new(awaken_stores::InMemoryStore::new());
2357
2358 let shared_discovered = AgentSpec {
2359 id: "shared".into(),
2360 model_id: "boot".into(),
2361 system_prompt: "discovered-prompt".into(),
2362 max_rounds: 1,
2363 endpoint: Some(RemoteEndpoint {
2364 base_url: "http://remote-shared/".into(),
2365 ..Default::default()
2366 }),
2367 ..Default::default()
2368 };
2369 let remote_only = AgentSpec {
2370 id: "remote-only".into(),
2371 model_id: "boot".into(),
2372 system_prompt: "remote-only-prompt".into(),
2373 max_rounds: 1,
2374 endpoint: Some(RemoteEndpoint {
2375 base_url: "http://remote-only/".into(),
2376 ..Default::default()
2377 }),
2378 ..Default::default()
2379 };
2380
2381 let runtime = Arc::new(
2382 awaken_runtime::builder::AgentRuntimeBuilder::new()
2383 .with_provider("boot", Arc::new(Stub))
2384 .with_model(ModelSpec::new("boot", "boot", "boot-model"))
2385 .with_agent_spec(shared_discovered)
2386 .with_agent_spec(remote_only)
2387 .with_in_memory_thread_run_store(thread_store.clone())
2388 .build()
2389 .expect("build runtime"),
2390 );
2391
2392 struct StubFactory;
2393 impl ProviderExecutorFactory for StubFactory {
2394 fn build(
2395 &self,
2396 _spec: &ProviderSpec,
2397 ) -> Result<Arc<dyn server_contract::contract::executor::LlmExecutor>, ConfigRuntimeError>
2398 {
2399 Ok(Arc::new(Stub))
2400 }
2401 }
2402
2403 let manager = ConfigRuntimeManager::new(runtime.clone(), store.clone())
2404 .expect("manager")
2405 .with_provider_factory(Arc::new(StubFactory));
2406
2407 let seed = BuiltinSeedSet {
2410 binary_version: "overlay-test".to_owned(),
2411 specs: vec![
2412 BuiltinSpec::Provider(server_contract::ProviderSpec {
2413 id: "boot-prov".into(),
2414 adapter: "stub".into(),
2415 ..Default::default()
2416 }),
2417 BuiltinSpec::Model(server_contract::ModelSpec::new(
2418 "boot-model",
2419 "boot-prov",
2420 "gpt-4o",
2421 )),
2422 BuiltinSpec::Agent(Box::new(AgentSpec {
2423 id: "shared".into(),
2424 model_id: "boot-model".into(),
2425 system_prompt: "seeded-prompt".into(),
2426 max_rounds: 5,
2427 endpoint: None,
2428 ..Default::default()
2429 })),
2430 ],
2431 };
2432 manager.apply_seed(&seed).await.expect("apply_seed");
2433 manager.apply().await.expect("apply");
2434
2435 let snapshot = runtime.registry_snapshot().expect("registry snapshot");
2437 let registry = &snapshot.registries().agents;
2438
2439 let shared_spec = registry.get_agent("shared").expect("shared must resolve");
2441 let shared_json = serde_json::to_value(&shared_spec).expect("serialize");
2442 assert_eq!(
2443 shared_json["system_prompt"], "seeded-prompt",
2444 "base (seeded) wins: system_prompt must be 'seeded-prompt', got {shared_json}"
2445 );
2446 assert_eq!(
2447 shared_json["max_rounds"], 5,
2448 "base (seeded) wins: max_rounds must be 5"
2449 );
2450
2451 let remote_spec = registry
2453 .get_agent("remote-only")
2454 .expect("remote-only must resolve via overlay");
2455 let remote_json = serde_json::to_value(&remote_spec).expect("serialize");
2456 assert_eq!(
2457 remote_json["system_prompt"], "remote-only-prompt",
2458 "discovery-only agent resolves via overlay"
2459 );
2460 assert_eq!(
2461 remote_json["endpoint"]["base_url"], "http://remote-only/",
2462 "endpoint base_url must be preserved"
2463 );
2464 }
2465
2466 fn builtin_agent_record(
2471 spec: &AgentSpec,
2472 binary_version: &str,
2473 user_overrides: Option<serde_json::Value>,
2474 ) -> serde_json::Value {
2475 use server_contract::{ConfigRecord, RecordMeta};
2476 let mut meta = RecordMeta::new_builtin(binary_version);
2477 meta.user_overrides = user_overrides;
2478 let record = ConfigRecord {
2479 spec: spec.clone(),
2480 meta,
2481 };
2482 record.to_value().expect("envelope serialize must succeed")
2483 }
2484
2485 #[tokio::test]
2486 async fn apply_overrides_merges_at_read() {
2487 use server_contract::{BuiltinSeedSet, BuiltinSpec};
2488
2489 let (manager, store) = make_manager_with_store().await;
2490
2491 let seed = BuiltinSeedSet {
2493 binary_version: "v1".to_owned(),
2494 specs: vec![
2495 BuiltinSpec::Provider(ProviderSpec {
2496 id: "p".into(),
2497 adapter: "openai".into(),
2498 ..Default::default()
2499 }),
2500 BuiltinSpec::Model(ModelSpec::new("m", "p", "gpt-4o")),
2501 BuiltinSpec::Agent(Box::new(AgentSpec {
2502 id: "x".into(),
2503 model_id: "m".into(),
2504 system_prompt: "base-prompt".into(),
2505 max_rounds: 5,
2506 ..Default::default()
2507 })),
2508 ],
2509 };
2510 manager.apply_seed(&seed).await.expect("apply_seed");
2511
2512 let base_spec = AgentSpec {
2514 id: "x".into(),
2515 model_id: "m".into(),
2516 system_prompt: "base-prompt".into(),
2517 max_rounds: 5,
2518 ..Default::default()
2519 };
2520 let envelope =
2521 builtin_agent_record(&base_spec, "v1", Some(json!({"system_prompt": "patched"})));
2522 store
2523 .put("agents", "x", &envelope)
2524 .await
2525 .expect("put must succeed");
2526
2527 manager.apply().await.expect("apply must succeed");
2528
2529 let snapshot = manager
2530 .runtime
2531 .registry_snapshot()
2532 .expect("registry snapshot");
2533 let spec = snapshot
2534 .registries()
2535 .agents
2536 .get_agent("x")
2537 .expect("agent x must resolve");
2538 assert_eq!(
2539 spec.system_prompt, "patched",
2540 "user_overrides must be applied at read time"
2541 );
2542 assert_eq!(spec.max_rounds, 5);
2544 }
2545
2546 #[tokio::test]
2547 async fn failed_candidate_validation_does_not_commit_provider_executor_cache() {
2548 use server_contract::{BuiltinSeedSet, BuiltinSpec, ConfigRecord, RecordMeta};
2549
2550 let (manager, store) = make_manager_with_store().await;
2551
2552 let seed = BuiltinSeedSet {
2553 binary_version: "cache-test".to_owned(),
2554 specs: vec![
2555 BuiltinSpec::Provider(ProviderSpec {
2556 id: "p".into(),
2557 adapter: "openai".into(),
2558 ..Default::default()
2559 }),
2560 BuiltinSpec::Model(ModelSpec::new("m", "p", "gpt-4o")),
2561 BuiltinSpec::Agent(Box::new(AgentSpec {
2562 id: "a".into(),
2563 model_id: "m".into(),
2564 system_prompt: "base".into(),
2565 max_rounds: 1,
2566 ..Default::default()
2567 })),
2568 ],
2569 };
2570 manager.apply_seed(&seed).await.expect("apply seed");
2571 manager.apply().await.expect("initial apply");
2572
2573 let initial_cached_provider = manager
2574 .provider_cache
2575 .lock()
2576 .executor_provider("p")
2577 .expect("provider cache entry");
2578
2579 let changed_provider = ConfigRecord {
2580 spec: ProviderSpec {
2581 id: "p".into(),
2582 adapter: "openai".into(),
2583 base_url: Some("https://provider-cache-candidate.example".into()),
2584 timeout_secs: 17,
2585 ..Default::default()
2586 },
2587 meta: RecordMeta::new_builtin("cache-test"),
2588 };
2589 store
2590 .put(
2591 "providers",
2592 "p",
2593 &changed_provider
2594 .to_value()
2595 .expect("serialize changed provider"),
2596 )
2597 .await
2598 .expect("write changed provider");
2599
2600 let invalid_agent = ConfigRecord {
2601 spec: AgentSpec {
2602 id: "a".into(),
2603 model_id: "missing-model".into(),
2604 system_prompt: "invalid".into(),
2605 max_rounds: 1,
2606 ..Default::default()
2607 },
2608 meta: RecordMeta::new_builtin("cache-test"),
2609 };
2610 store
2611 .put(
2612 "agents",
2613 "a",
2614 &invalid_agent.to_value().expect("serialize invalid agent"),
2615 )
2616 .await
2617 .expect("write invalid agent");
2618
2619 manager
2620 .apply()
2621 .await
2622 .expect_err("invalid candidate must fail validation");
2623
2624 let cached_provider = manager
2625 .provider_cache
2626 .lock()
2627 .executor_provider("p")
2628 .expect("provider cache entry must remain");
2629 assert_eq!(cached_provider, initial_cached_provider);
2630 }
2631
2632 #[tokio::test]
2633 async fn apply_overrides_no_user_overrides_uses_base() {
2634 use server_contract::{BuiltinSeedSet, BuiltinSpec};
2635
2636 let (manager, store) = make_manager_with_store().await;
2637
2638 let seed = BuiltinSeedSet {
2639 binary_version: "v1".to_owned(),
2640 specs: vec![
2641 BuiltinSpec::Provider(ProviderSpec {
2642 id: "p".into(),
2643 adapter: "openai".into(),
2644 ..Default::default()
2645 }),
2646 BuiltinSpec::Model(ModelSpec::new("m", "p", "gpt-4o")),
2647 BuiltinSpec::Agent(Box::new(AgentSpec {
2648 id: "y".into(),
2649 model_id: "m".into(),
2650 system_prompt: "base-prompt".into(),
2651 max_rounds: 3,
2652 ..Default::default()
2653 })),
2654 ],
2655 };
2656 manager.apply_seed(&seed).await.expect("apply_seed");
2657
2658 let raw = store.get("agents", "y").await.unwrap().unwrap();
2660 let rec: server_contract::ConfigRecord<serde_json::Value> =
2661 server_contract::ConfigRecord::from_value(raw).unwrap();
2662 assert!(rec.meta.user_overrides.is_none());
2663
2664 manager.apply().await.expect("apply must succeed");
2665
2666 let snapshot = manager
2667 .runtime
2668 .registry_snapshot()
2669 .expect("registry snapshot");
2670 let spec = snapshot
2671 .registries()
2672 .agents
2673 .get_agent("y")
2674 .expect("agent y must resolve");
2675 assert_eq!(spec.system_prompt, "base-prompt");
2676 assert_eq!(spec.max_rounds, 3);
2677 }
2678
2679 #[tokio::test]
2680 async fn apply_overrides_on_user_record_applies_overrides() {
2681 use server_contract::{BuiltinSeedSet, BuiltinSpec};
2685 use server_contract::{ConfigRecord, RecordMeta};
2686
2687 let (manager, store) = make_manager_with_store().await;
2688
2689 let seed = BuiltinSeedSet {
2691 binary_version: "v1".to_owned(),
2692 specs: vec![
2693 BuiltinSpec::Provider(ProviderSpec {
2694 id: "p".into(),
2695 adapter: "openai".into(),
2696 ..Default::default()
2697 }),
2698 BuiltinSpec::Model(ModelSpec::new("m", "p", "gpt-4o")),
2699 ],
2700 };
2701 manager.apply_seed(&seed).await.expect("apply_seed");
2702
2703 let user_spec = AgentSpec {
2706 id: "z".into(),
2707 model_id: "m".into(),
2708 system_prompt: "user-base".into(),
2709 max_rounds: 2,
2710 ..Default::default()
2711 };
2712 let mut meta = RecordMeta::new_user();
2713 meta.user_overrides = Some(json!({"system_prompt": "user-override"}));
2714 let record = ConfigRecord {
2715 spec: user_spec,
2716 meta,
2717 };
2718 store
2719 .put("agents", "z", &record.to_value().unwrap())
2720 .await
2721 .expect("put must succeed");
2722
2723 manager.apply().await.expect("apply must succeed");
2724
2725 let snapshot = manager
2726 .runtime
2727 .registry_snapshot()
2728 .expect("registry snapshot");
2729 let spec = snapshot
2730 .registries()
2731 .agents
2732 .get_agent("z")
2733 .expect("agent z must resolve");
2734 assert_eq!(
2736 spec.system_prompt, "user-override",
2737 "user_overrides applied at read time regardless of source"
2738 );
2739 }
2740
2741 #[tokio::test]
2742 async fn version_upgrade_preserves_user_overrides() {
2743 use server_contract::{BuiltinSeedSet, BuiltinSpec};
2744
2745 let (manager, store) = make_manager_with_store().await;
2746
2747 let seed_v1 = BuiltinSeedSet {
2749 binary_version: "v1".to_owned(),
2750 specs: vec![
2751 BuiltinSpec::Provider(ProviderSpec {
2752 id: "p".into(),
2753 adapter: "openai".into(),
2754 ..Default::default()
2755 }),
2756 BuiltinSpec::Model(ModelSpec::new("m", "p", "gpt-4o")),
2757 BuiltinSpec::Agent(Box::new(AgentSpec {
2758 id: "a".into(),
2759 model_id: "m".into(),
2760 system_prompt: "v1-prompt".into(),
2761 max_rounds: 5,
2762 ..Default::default()
2763 })),
2764 ],
2765 };
2766 manager.apply_seed(&seed_v1).await.expect("apply_seed v1");
2767
2768 let raw = store.get("agents", "a").await.unwrap().unwrap();
2770 let mut rec: server_contract::ConfigRecord<serde_json::Value> =
2771 server_contract::ConfigRecord::from_value(raw).unwrap();
2772 rec.meta.user_overrides = Some(json!({"system_prompt": "user-prompt"}));
2773 store
2774 .put("agents", "a", &rec.to_value().unwrap())
2775 .await
2776 .expect("put with overrides");
2777
2778 let seed_v2 = BuiltinSeedSet {
2780 binary_version: "v2".to_owned(),
2781 specs: vec![
2782 BuiltinSpec::Provider(ProviderSpec {
2783 id: "p".into(),
2784 adapter: "openai".into(),
2785 ..Default::default()
2786 }),
2787 BuiltinSpec::Model(ModelSpec::new("m", "p", "gpt-4o")),
2788 BuiltinSpec::Agent(Box::new(AgentSpec {
2789 id: "a".into(),
2790 model_id: "m".into(),
2791 system_prompt: "v2-prompt".into(),
2792 max_rounds: 10,
2793 ..Default::default()
2794 })),
2795 ],
2796 };
2797 manager.apply_seed(&seed_v2).await.expect("apply_seed v2");
2798
2799 let raw = store.get("agents", "a").await.unwrap().unwrap();
2801 let stored: server_contract::ConfigRecord<serde_json::Value> =
2802 server_contract::ConfigRecord::from_value(raw).unwrap();
2803 assert_eq!(
2804 stored.meta.source,
2805 server_contract::RecordSource::Builtin {
2806 binary_version: "v2".to_owned()
2807 },
2808 "binary_version must be updated to v2"
2809 );
2810 assert_eq!(
2811 stored.meta.user_overrides,
2812 Some(json!({"system_prompt": "user-prompt"})),
2813 "user_overrides must be preserved across version upgrade"
2814 );
2815 assert_eq!(stored.spec["system_prompt"], "v2-prompt");
2817 assert_eq!(stored.spec["max_rounds"], 10);
2818
2819 manager.apply().await.expect("apply must succeed");
2821 let snapshot = manager
2822 .runtime
2823 .registry_snapshot()
2824 .expect("registry snapshot");
2825 let spec = snapshot
2826 .registries()
2827 .agents
2828 .get_agent("a")
2829 .expect("agent a must resolve");
2830 assert_eq!(
2831 spec.system_prompt, "user-prompt",
2832 "user override for system_prompt must be preserved after version upgrade"
2833 );
2834 assert_eq!(
2835 spec.max_rounds, 10,
2836 "max_rounds must use v2 base (not overridden)"
2837 );
2838 }
2839
2840 async fn bootstrap_with_static_tool(
2847 tool_id: &str,
2848 tool_description: &str,
2849 ) -> (
2850 Arc<ConfigRuntimeManager>,
2851 Arc<awaken_runtime::AgentRuntime>,
2852 Arc<dyn server_contract::contract::config_store::ConfigStore>,
2853 ) {
2854 use awaken_stores::InMemoryStore;
2855 use serde_json::json;
2856 use server_contract::contract::executor::{
2857 InferenceExecutionError, InferenceRequest, LlmExecutor,
2858 };
2859 use server_contract::contract::inference::{StopReason, StreamResult, TokenUsage};
2860 use server_contract::contract::tool::{
2861 Tool, ToolCallContext, ToolDescriptor, ToolError, ToolOutput, ToolResult,
2862 };
2863 use server_contract::{BuiltinSeedSet, BuiltinSpec, ModelSpec, ProviderSpec, ToolSpec};
2864
2865 struct Stub;
2866 #[async_trait::async_trait]
2867 impl LlmExecutor for Stub {
2868 async fn execute(
2869 &self,
2870 _: InferenceRequest,
2871 ) -> Result<StreamResult, InferenceExecutionError> {
2872 Ok(StreamResult {
2873 content: vec![],
2874 tool_calls: vec![],
2875 usage: Some(TokenUsage::default()),
2876 stop_reason: Some(StopReason::EndTurn),
2877 has_incomplete_tool_calls: false,
2878 })
2879 }
2880 fn name(&self) -> &str {
2881 "stub"
2882 }
2883 }
2884
2885 struct StubTool {
2886 id: String,
2887 description: String,
2888 }
2889 #[async_trait::async_trait]
2890 impl Tool for StubTool {
2891 fn descriptor(&self) -> ToolDescriptor {
2892 ToolDescriptor::new(&self.id, &self.id, &self.description)
2893 }
2894 async fn execute(
2895 &self,
2896 _args: serde_json::Value,
2897 _ctx: &ToolCallContext,
2898 ) -> Result<ToolOutput, ToolError> {
2899 Ok(ToolResult::success(&self.id, json!({})).into())
2900 }
2901 }
2902
2903 struct StubFactory;
2904 impl ProviderExecutorFactory for StubFactory {
2905 fn build(
2906 &self,
2907 _spec: &ProviderSpec,
2908 ) -> Result<Arc<dyn server_contract::contract::executor::LlmExecutor>, ConfigRuntimeError>
2909 {
2910 Ok(Arc::new(Stub))
2911 }
2912 }
2913
2914 let store = Arc::new(InMemoryStore::new())
2915 as Arc<dyn server_contract::contract::config_store::ConfigStore>;
2916 let thread_store = Arc::new(InMemoryStore::new());
2917
2918 let runtime = Arc::new(
2919 awaken_runtime::builder::AgentRuntimeBuilder::new()
2920 .with_provider("boot", Arc::new(Stub))
2921 .with_model(ModelSpec::new("boot", "boot", "boot-model"))
2922 .with_tool(
2923 tool_id,
2924 Arc::new(StubTool {
2925 id: tool_id.to_owned(),
2926 description: tool_description.to_owned(),
2927 }),
2928 )
2929 .with_in_memory_thread_run_store(thread_store.clone())
2930 .build()
2931 .expect("build runtime"),
2932 );
2933
2934 let manager = Arc::new(
2935 ConfigRuntimeManager::new(runtime.clone(), store.clone())
2936 .expect("manager")
2937 .with_provider_factory(Arc::new(StubFactory)),
2938 );
2939
2940 let seed = BuiltinSeedSet {
2941 binary_version: "test".to_owned(),
2942 specs: vec![
2943 BuiltinSpec::Provider(ProviderSpec {
2944 id: "test-prov".into(),
2945 adapter: "openai".into(),
2946 ..Default::default()
2947 }),
2948 BuiltinSpec::Model(ModelSpec::new("test-model", "test-prov", "gpt-4o")),
2949 BuiltinSpec::Agent(Box::new(AgentSpec {
2950 id: "agent-using-echo".into(),
2951 model_id: "test-model".into(),
2952 system_prompt: "you are a test".into(),
2953 max_rounds: 1,
2954 allowed_tools: None,
2955 endpoint: None,
2956 ..Default::default()
2957 })),
2958 BuiltinSpec::Tool(ToolSpec {
2959 id: tool_id.to_owned(),
2960 name: tool_id.to_owned(),
2961 description: tool_description.to_owned(),
2962 ..Default::default()
2963 }),
2964 ],
2965 };
2966 manager.apply_seed(&seed).await.expect("apply_seed");
2967
2968 (manager, runtime, store)
2969 }
2970
2971 #[tokio::test]
2972 async fn tool_description_override_applied_to_resolved_agent() {
2973 let (manager, runtime, store) =
2974 bootstrap_with_static_tool("echo", "stock description").await;
2975
2976 manager.apply().await.expect("initial apply");
2977
2978 let envelope = serde_json::json!({
2979 "spec": {
2980 "id": "echo",
2981 "name": "Echo",
2982 "description": "stock description",
2983 "category": null,
2984 "parameters_schema": {}
2985 },
2986 "meta": {
2987 "source": { "kind": "builtin", "binary_version": "test" },
2988 "user_overrides": { "description": "patched description" },
2989 "hidden": false,
2990 "created_at": 1,
2991 "updated_at": 2
2992 }
2993 });
2994 server_contract::contract::config_store::ConfigStore::put(
2995 store.as_ref(),
2996 "tools",
2997 "echo",
2998 &envelope,
2999 )
3000 .await
3001 .expect("write override");
3002
3003 manager.apply().await.expect("apply with override");
3004
3005 let resolver = runtime.resolver_arc();
3006 let resolved = resolver.resolve("agent-using-echo").expect("resolve");
3007 let descs = resolved.tool_descriptors();
3008 let echo = descs
3009 .iter()
3010 .find(|d| d.id == "echo")
3011 .expect("echo descriptor present");
3012 assert_eq!(echo.description, "patched description");
3013 }
3014
3015 #[tokio::test]
3016 async fn snapshot_tool_specs_emits_one_entry_per_registered_tool() {
3017 let (manager, _runtime, _store) =
3018 bootstrap_with_static_tool("echo", "stock description").await;
3019 let specs = manager.snapshot_tool_specs();
3020 assert_eq!(specs.len(), 1);
3021 match &specs[0] {
3022 server_contract::BuiltinSpec::Tool(t) => {
3023 assert_eq!(t.id, "echo");
3024 assert_eq!(t.description, "stock description");
3025 }
3026 other => panic!("expected Tool variant, got {other:?}"),
3027 }
3028 }
3029}