Skip to main content

awaken_runtime/
builder.rs

1//! Fluent builder API for constructing `AgentRuntime`.
2
3use std::sync::Arc;
4
5use awaken_contract::StateError;
6use awaken_contract::contract::executor::LlmExecutor;
7use awaken_contract::contract::storage::ThreadRunStore;
8use awaken_contract::contract::tool::Tool;
9use awaken_contract::registry_spec::AgentSpec;
10
11use crate::backend::ExecutionBackendFactory;
12use crate::plugins::Plugin;
13#[cfg(feature = "a2a")]
14use crate::registry::BackendRegistry;
15#[cfg(feature = "a2a")]
16use crate::registry::composite::{CompositeAgentSpecRegistry, RemoteAgentSource};
17#[cfg(feature = "a2a")]
18use crate::registry::memory::MapBackendRegistry;
19use crate::registry::memory::{
20    MapAgentSpecRegistry, MapModelRegistry, MapPluginSource, MapProviderRegistry, MapToolRegistry,
21};
22use crate::registry::snapshot::RegistryHandle;
23use crate::registry::traits::{AgentSpecRegistry, ModelBinding, RegistrySet};
24use crate::runtime::AgentRuntime;
25
26/// Error returned when the builder cannot construct the runtime.
27#[derive(Debug, thiserror::Error)]
28pub enum BuildError {
29    #[error("state error: {0}")]
30    State(#[from] StateError),
31    #[error("agent registry conflict: {0}")]
32    AgentRegistryConflict(String),
33    #[error("tool registry conflict: {0}")]
34    ToolRegistryConflict(String),
35    #[error("model registry conflict: {0}")]
36    ModelRegistryConflict(String),
37    #[error("provider registry conflict: {0}")]
38    ProviderRegistryConflict(String),
39    #[error("plugin registry conflict: {0}")]
40    PluginRegistryConflict(String),
41    #[cfg(feature = "a2a")]
42    #[error("backend registry conflict: {0}")]
43    BackendRegistryConflict(String),
44    #[error("agent validation failed: {0}")]
45    ValidationFailed(String),
46    #[cfg(feature = "a2a")]
47    #[error("discovery failed: {0}")]
48    DiscoveryFailed(#[from] crate::registry::composite::DiscoveryError),
49}
50
51/// Fluent API for constructing an `AgentRuntime`.
52///
53/// Collects agent specs, tools, plugins, models, providers, and optionally
54/// a store, then builds the fully resolved runtime.
55pub struct AgentRuntimeBuilder {
56    agents: MapAgentSpecRegistry,
57    tools: MapToolRegistry,
58    models: MapModelRegistry,
59    providers: MapProviderRegistry,
60    plugins: MapPluginSource,
61    #[cfg(feature = "a2a")]
62    backends: MapBackendRegistry,
63    thread_run_store: Option<Arc<dyn ThreadRunStore>>,
64    profile_store: Option<Arc<dyn awaken_contract::contract::profile_store::ProfileStore>>,
65    errors: Vec<BuildError>,
66    #[cfg(feature = "a2a")]
67    remote_sources: Vec<RemoteAgentSource>,
68}
69
70impl AgentRuntimeBuilder {
71    pub fn new() -> Self {
72        Self {
73            agents: MapAgentSpecRegistry::new(),
74            tools: MapToolRegistry::new(),
75            models: MapModelRegistry::new(),
76            providers: MapProviderRegistry::new(),
77            plugins: MapPluginSource::new(),
78            #[cfg(feature = "a2a")]
79            backends: MapBackendRegistry::with_default_remote_backends(),
80            thread_run_store: None,
81            profile_store: None,
82            errors: Vec::new(),
83            #[cfg(feature = "a2a")]
84            remote_sources: Vec::new(),
85        }
86    }
87
88    /// Register an agent spec.
89    pub fn with_agent_spec(mut self, spec: AgentSpec) -> Self {
90        if let Err(e) = self.agents.register_spec(spec) {
91            self.errors.push(e);
92        }
93        self
94    }
95
96    /// Register multiple agent specs.
97    pub fn with_agent_specs(mut self, specs: impl IntoIterator<Item = AgentSpec>) -> Self {
98        for spec in specs {
99            if let Err(e) = self.agents.register_spec(spec) {
100                self.errors.push(e);
101            }
102        }
103        self
104    }
105
106    /// Register a tool by ID.
107    pub fn with_tool(mut self, id: impl Into<String>, tool: Arc<dyn Tool>) -> Self {
108        if let Err(e) = self.tools.register_tool(id, tool) {
109            self.errors.push(e);
110        }
111        self
112    }
113
114    /// Register a plugin by ID.
115    pub fn with_plugin(mut self, id: impl Into<String>, plugin: Arc<dyn Plugin>) -> Self {
116        if let Err(e) = self.plugins.register_plugin(id, plugin) {
117            self.errors.push(e);
118        }
119        self
120    }
121
122    /// Register a model binding by ID.
123    pub fn with_model_binding(mut self, id: impl Into<String>, binding: ModelBinding) -> Self {
124        if let Err(e) = self.models.register_model(id, binding) {
125            self.errors.push(e);
126        }
127        self
128    }
129
130    /// Register a provider (LLM executor) by ID.
131    pub fn with_provider(mut self, id: impl Into<String>, executor: Arc<dyn LlmExecutor>) -> Self {
132        if let Err(e) = self.providers.register_provider(id, executor) {
133            self.errors.push(e);
134        }
135        self
136    }
137
138    /// Set the thread run store for persistence.
139    pub fn with_thread_run_store(mut self, store: Arc<dyn ThreadRunStore>) -> Self {
140        self.thread_run_store = Some(store);
141        self
142    }
143
144    /// Set the profile store for cross-run key-value persistence.
145    pub fn with_profile_store(
146        mut self,
147        store: Arc<dyn awaken_contract::contract::profile_store::ProfileStore>,
148    ) -> Self {
149        self.profile_store = Some(store);
150        self
151    }
152
153    /// Add a named remote A2A agent source for discovery.
154    ///
155    /// When remote sources are configured, the builder creates a
156    /// [`CompositeAgentSpecRegistry`] that combines local agents with
157    /// agents discovered from remote A2A endpoints. The `name` is used
158    /// for namespaced agent lookup (e.g., `"cloud/translator"`).
159    #[cfg(feature = "a2a")]
160    pub fn with_remote_agents(
161        mut self,
162        name: impl Into<String>,
163        base_url: impl Into<String>,
164        bearer_token: Option<String>,
165    ) -> Self {
166        self.remote_sources.push(RemoteAgentSource {
167            name: name.into(),
168            base_url: base_url.into(),
169            bearer_token,
170        });
171        self
172    }
173
174    /// Register a remote delegate backend factory by its backend kind.
175    #[cfg(feature = "a2a")]
176    pub fn with_agent_backend_factory(mut self, factory: Arc<dyn ExecutionBackendFactory>) -> Self {
177        if let Err(e) = self.backends.register_backend_factory(factory) {
178            self.errors.push(e);
179        }
180        self
181    }
182
183    /// Build the `AgentRuntime` and validate all registered agents can
184    /// resolve successfully.
185    ///
186    /// Performs a dry-run resolve for every registered agent, catching
187    /// configuration errors (missing models, providers, plugins) at build time.
188    /// Use [`build_unchecked()`](Self::build_unchecked) to skip validation.
189    pub fn build(self) -> Result<AgentRuntime, BuildError> {
190        let runtime = self.build_unchecked()?;
191        let resolver = runtime.resolver();
192        #[cfg(feature = "a2a")]
193        let registries = runtime.registry_set();
194        let mut errors = Vec::new();
195        for agent_id in resolver.agent_ids() {
196            #[cfg(feature = "a2a")]
197            {
198                if let Some(spec) = registries
199                    .as_ref()
200                    .and_then(|set| set.agents.get_agent(&agent_id))
201                    && let Some(endpoint) = &spec.endpoint
202                {
203                    let Some(factory) = registries
204                        .as_ref()
205                        .and_then(|set| set.backends.get_backend_factory(&endpoint.backend))
206                    else {
207                        errors.push(format!(
208                            "{agent_id}: unsupported remote backend '{}'",
209                            endpoint.backend
210                        ));
211                        continue;
212                    };
213                    if let Err(error) = factory.validate(endpoint) {
214                        errors.push(format!("{agent_id}: {error}"));
215                    }
216                    continue;
217                }
218            }
219
220            if let Err(e) = resolver.resolve(&agent_id) {
221                errors.push(format!("{agent_id}: {e}"));
222            }
223        }
224        if !errors.is_empty() {
225            return Err(BuildError::ValidationFailed(errors.join("; ")));
226        }
227        Ok(runtime)
228    }
229
230    /// Build the `AgentRuntime` from the accumulated configuration,
231    /// skipping agent validation.
232    ///
233    /// Prefer [`build()`](Self::build) which validates all registered agents
234    /// can resolve successfully at build time.
235    pub fn build_unchecked(mut self) -> Result<AgentRuntime, BuildError> {
236        if !self.errors.is_empty() {
237            return Err(self.errors.remove(0));
238        }
239
240        #[cfg(feature = "a2a")]
241        let (agents, composite_registry): (Arc<dyn AgentSpecRegistry>, _) =
242            if self.remote_sources.is_empty() {
243                (Arc::new(self.agents), None)
244            } else {
245                let mut composite = CompositeAgentSpecRegistry::new(Arc::new(self.agents));
246                for source in self.remote_sources {
247                    composite.add_remote(source);
248                }
249                let arc = Arc::new(composite);
250                (Arc::clone(&arc) as Arc<dyn AgentSpecRegistry>, Some(arc))
251            };
252        #[cfg(not(feature = "a2a"))]
253        let agents: Arc<dyn AgentSpecRegistry> = Arc::new(self.agents);
254
255        let registry_set = RegistrySet {
256            agents,
257            tools: Arc::new(self.tools),
258            models: Arc::new(self.models),
259            providers: Arc::new(self.providers),
260            plugins: Arc::new(self.plugins),
261            #[cfg(feature = "a2a")]
262            backends: Arc::new(self.backends) as Arc<dyn BackendRegistry>,
263        };
264
265        let registry_handle = RegistryHandle::new(registry_set.clone());
266        let resolver: Arc<dyn crate::registry::ExecutionResolver> = Arc::new(
267            crate::registry::resolve::DynamicRegistryResolver::new(registry_handle.clone()),
268        );
269
270        let mut runtime = AgentRuntime::new_with_execution_resolver(resolver)
271            .with_registry_handle(registry_handle);
272
273        #[cfg(feature = "a2a")]
274        if let Some(composite) = composite_registry {
275            runtime = runtime.with_composite_registry(composite);
276        }
277
278        if let Some(store) = self.thread_run_store {
279            runtime = runtime.with_thread_run_store(store);
280        }
281
282        if let Some(store) = self.profile_store {
283            runtime = runtime.with_profile_store(store);
284        }
285
286        Ok(runtime)
287    }
288
289    /// Build and initialize (async). Discovers remote agents after build.
290    #[cfg(feature = "a2a")]
291    pub async fn build_and_discover(self) -> Result<AgentRuntime, BuildError> {
292        let runtime = self.build_unchecked()?;
293        if let Some(composite) = runtime.composite_registry() {
294            composite.discover().await?;
295        }
296        Ok(runtime)
297    }
298}
299
300impl Default for AgentRuntimeBuilder {
301    fn default() -> Self {
302        Self::new()
303    }
304}
305
306#[cfg(test)]
307mod tests {
308    use super::*;
309    use async_trait::async_trait;
310    use awaken_contract::contract::executor::{InferenceExecutionError, InferenceRequest};
311    use awaken_contract::contract::inference::{StopReason, StreamResult, TokenUsage};
312    #[cfg(feature = "a2a")]
313    use awaken_contract::contract::lifecycle::TerminationReason;
314    use awaken_contract::contract::tool::{
315        ToolCallContext, ToolDescriptor, ToolError, ToolOutput, ToolResult,
316    };
317    #[cfg(feature = "a2a")]
318    use awaken_contract::registry_spec::RemoteEndpoint;
319    use serde_json::Value;
320    #[cfg(feature = "a2a")]
321    use std::sync::atomic::{AtomicUsize, Ordering};
322
323    use crate::registry::memory::{
324        MapAgentSpecRegistry, MapModelRegistry, MapPluginSource, MapProviderRegistry,
325        MapToolRegistry,
326    };
327
328    struct MockTool {
329        id: String,
330    }
331
332    #[async_trait]
333    impl Tool for MockTool {
334        fn descriptor(&self) -> ToolDescriptor {
335            ToolDescriptor::new(&self.id, &self.id, "mock tool")
336        }
337
338        async fn execute(
339            &self,
340            _args: Value,
341            _ctx: &ToolCallContext,
342        ) -> Result<ToolOutput, ToolError> {
343            Ok(ToolResult::success(&self.id, Value::Null).into())
344        }
345    }
346
347    struct MockExecutor;
348
349    #[async_trait]
350    impl LlmExecutor for MockExecutor {
351        async fn execute(
352            &self,
353            _request: InferenceRequest,
354        ) -> Result<StreamResult, InferenceExecutionError> {
355            Ok(StreamResult {
356                content: vec![],
357                tool_calls: vec![],
358                usage: Some(TokenUsage::default()),
359                stop_reason: Some(StopReason::EndTurn),
360                has_incomplete_tool_calls: false,
361            })
362        }
363
364        fn name(&self) -> &str {
365            "mock"
366        }
367    }
368
369    #[cfg(feature = "a2a")]
370    struct NoopRemoteBackend;
371
372    #[cfg(feature = "a2a")]
373    #[async_trait]
374    impl crate::backend::ExecutionBackend for NoopRemoteBackend {
375        async fn execute_root(
376            &self,
377            request: crate::backend::BackendRootRunRequest<'_>,
378        ) -> Result<crate::backend::BackendRunResult, crate::backend::ExecutionBackendError>
379        {
380            Ok(crate::backend::BackendRunResult {
381                agent_id: request.agent_id.to_string(),
382                status: crate::backend::BackendRunStatus::Completed,
383                termination: TerminationReason::NaturalEnd,
384                status_reason: None,
385                response: None,
386                output: crate::backend::BackendRunOutput::default(),
387                steps: 0,
388                run_id: None,
389                inbox: None,
390                state: None,
391            })
392        }
393    }
394
395    #[cfg(feature = "a2a")]
396    struct CountingValidationBackendFactory {
397        validate_count: Arc<AtomicUsize>,
398        build_count: Arc<AtomicUsize>,
399    }
400
401    #[cfg(feature = "a2a")]
402    impl crate::backend::ExecutionBackendFactory for CountingValidationBackendFactory {
403        fn backend(&self) -> &str {
404            "counting-remote"
405        }
406
407        fn validate(
408            &self,
409            endpoint: &RemoteEndpoint,
410        ) -> Result<(), crate::backend::ExecutionBackendFactoryError> {
411            self.validate_count.fetch_add(1, Ordering::SeqCst);
412            if endpoint.base_url.trim().is_empty() {
413                return Err(crate::backend::ExecutionBackendFactoryError::InvalidConfig(
414                    "empty base_url".into(),
415                ));
416            }
417            Ok(())
418        }
419
420        fn build(
421            &self,
422            endpoint: &RemoteEndpoint,
423        ) -> Result<
424            Arc<dyn crate::backend::ExecutionBackend>,
425            crate::backend::ExecutionBackendFactoryError,
426        > {
427            self.build_count.fetch_add(1, Ordering::SeqCst);
428            if endpoint.backend != self.backend() {
429                return Err(crate::backend::ExecutionBackendFactoryError::InvalidConfig(
430                    format!("unexpected backend '{}'", endpoint.backend),
431                ));
432            }
433            Ok(Arc::new(NoopRemoteBackend))
434        }
435    }
436
437    fn make_registry_set(agent_id: &str, model_id: &str, upstream_model: &str) -> RegistrySet {
438        let mut agents = MapAgentSpecRegistry::new();
439        agents
440            .register_spec(AgentSpec {
441                id: agent_id.into(),
442                model_id: model_id.into(),
443                system_prompt: format!("system-{agent_id}"),
444                ..Default::default()
445            })
446            .expect("register test agent");
447
448        let mut models = MapModelRegistry::new();
449        models
450            .register_model(
451                model_id,
452                ModelBinding {
453                    provider_id: "mock".into(),
454                    upstream_model: upstream_model.into(),
455                },
456            )
457            .expect("register test model");
458
459        let mut providers = MapProviderRegistry::new();
460        providers
461            .register_provider("mock", Arc::new(MockExecutor))
462            .expect("register test provider");
463
464        RegistrySet {
465            agents: Arc::new(agents),
466            tools: Arc::new(MapToolRegistry::new()),
467            models: Arc::new(models),
468            providers: Arc::new(providers),
469            plugins: Arc::new(MapPluginSource::new()),
470            backends: Arc::new(MapBackendRegistry::new()),
471        }
472    }
473
474    #[test]
475    fn builder_creates_runtime() {
476        let spec = AgentSpec {
477            id: "test-agent".into(),
478            model_id: "test-model".into(),
479            system_prompt: "You are helpful.".into(),
480            ..Default::default()
481        };
482
483        let runtime = AgentRuntimeBuilder::new()
484            .with_agent_spec(spec)
485            .with_tool("echo", Arc::new(MockTool { id: "echo".into() }))
486            .with_model_binding(
487                "test-model",
488                ModelBinding {
489                    provider_id: "mock".into(),
490                    upstream_model: "mock-model".into(),
491                },
492            )
493            .with_provider("mock", Arc::new(MockExecutor))
494            .build();
495
496        assert!(runtime.is_ok());
497    }
498
499    #[test]
500    fn builder_default_creates_empty() {
501        let builder = AgentRuntimeBuilder::default();
502        // Cannot resolve any agent but should build
503        let runtime = builder.build();
504        assert!(runtime.is_ok());
505    }
506
507    #[test]
508    fn builder_with_multiple_agents() {
509        let spec1 = AgentSpec {
510            id: "agent-1".into(),
511            model_id: "m".into(),
512            system_prompt: "sys".into(),
513            ..Default::default()
514        };
515        let spec2 = AgentSpec {
516            id: "agent-2".into(),
517            model_id: "m".into(),
518            system_prompt: "sys".into(),
519            ..Default::default()
520        };
521
522        let runtime = AgentRuntimeBuilder::new()
523            .with_agent_specs(vec![spec1, spec2])
524            .with_model_binding(
525                "m",
526                ModelBinding {
527                    provider_id: "p".into(),
528                    upstream_model: "n".into(),
529                },
530            )
531            .with_provider("p", Arc::new(MockExecutor))
532            .build()
533            .unwrap();
534
535        // Both agents should be resolvable
536        assert!(runtime.resolver().resolve("agent-1").is_ok());
537        assert!(runtime.resolver().resolve("agent-2").is_ok());
538    }
539
540    #[test]
541    fn builder_resolver_returns_correct_config() {
542        let spec = AgentSpec {
543            id: "my-agent".into(),
544            model_id: "test-model".into(),
545            system_prompt: "Be helpful.".into(),
546            max_rounds: 10,
547            ..Default::default()
548        };
549
550        let runtime = AgentRuntimeBuilder::new()
551            .with_agent_spec(spec)
552            .with_tool(
553                "search",
554                Arc::new(MockTool {
555                    id: "search".into(),
556                }),
557            )
558            .with_model_binding(
559                "test-model",
560                ModelBinding {
561                    provider_id: "mock".into(),
562                    upstream_model: "claude-test".into(),
563                },
564            )
565            .with_provider("mock", Arc::new(MockExecutor))
566            .build()
567            .unwrap();
568
569        let resolved = runtime.resolver().resolve("my-agent").unwrap();
570        assert_eq!(resolved.id(), "my-agent");
571        assert_eq!(resolved.upstream_model, "claude-test");
572        assert_eq!(resolved.system_prompt(), "Be helpful.");
573        assert_eq!(resolved.max_rounds(), 10);
574        assert!(resolved.tools.contains_key("search"));
575    }
576
577    #[test]
578    fn builder_missing_agent_errors() {
579        let runtime = AgentRuntimeBuilder::new()
580            .with_model_binding(
581                "m",
582                ModelBinding {
583                    provider_id: "p".into(),
584                    upstream_model: "n".into(),
585                },
586            )
587            .with_provider("p", Arc::new(MockExecutor))
588            .build()
589            .unwrap();
590
591        let err = runtime.resolver().resolve("nonexistent");
592        assert!(err.is_err());
593    }
594
595    // -----------------------------------------------------------------------
596    // Migrated from uncarve: additional builder tests
597    // -----------------------------------------------------------------------
598
599    #[test]
600    fn builder_with_plugin() {
601        use crate::plugins::{Plugin, PluginDescriptor, PluginRegistrar};
602
603        struct TestPlugin;
604        impl Plugin for TestPlugin {
605            fn descriptor(&self) -> PluginDescriptor {
606                PluginDescriptor {
607                    name: "test-builder-plugin",
608                }
609            }
610            fn register(
611                &self,
612                _registrar: &mut PluginRegistrar,
613            ) -> Result<(), awaken_contract::StateError> {
614                Ok(())
615            }
616        }
617
618        let runtime = AgentRuntimeBuilder::new()
619            .with_plugin("test-builder-plugin", Arc::new(TestPlugin))
620            .build()
621            .unwrap();
622        let _ = runtime;
623    }
624
625    #[test]
626    fn builder_chained_tools_all_registered() {
627        let spec = AgentSpec {
628            id: "agent".into(),
629            model_id: "m".into(),
630            system_prompt: "sys".into(),
631            ..Default::default()
632        };
633
634        let runtime = AgentRuntimeBuilder::new()
635            .with_agent_spec(spec)
636            .with_tool("t1", Arc::new(MockTool { id: "t1".into() }))
637            .with_tool("t2", Arc::new(MockTool { id: "t2".into() }))
638            .with_tool("t3", Arc::new(MockTool { id: "t3".into() }))
639            .with_model_binding(
640                "m",
641                ModelBinding {
642                    provider_id: "p".into(),
643                    upstream_model: "n".into(),
644                },
645            )
646            .with_provider("p", Arc::new(MockExecutor))
647            .build()
648            .unwrap();
649
650        let resolved = runtime.resolver().resolve("agent").unwrap();
651        assert!(resolved.tools.contains_key("t1"));
652        assert!(resolved.tools.contains_key("t2"));
653        assert!(resolved.tools.contains_key("t3"));
654    }
655
656    #[test]
657    fn build_catches_missing_model() {
658        let spec = AgentSpec {
659            id: "bad-agent".into(),
660            model_id: "nonexistent-model".into(),
661            system_prompt: "sys".into(),
662            ..Default::default()
663        };
664
665        let result = AgentRuntimeBuilder::new().with_agent_spec(spec).build();
666
667        let err = match result {
668            Err(e) => e.to_string(),
669            Ok(_) => panic!("expected build to fail for missing model"),
670        };
671        assert!(
672            err.contains("bad-agent"),
673            "error should mention the agent ID: {err}"
674        );
675    }
676
677    #[test]
678    fn build_succeeds_with_valid_config() {
679        let spec = AgentSpec {
680            id: "good-agent".into(),
681            model_id: "m".into(),
682            system_prompt: "sys".into(),
683            ..Default::default()
684        };
685
686        let result = AgentRuntimeBuilder::new()
687            .with_agent_spec(spec)
688            .with_model_binding(
689                "m",
690                ModelBinding {
691                    provider_id: "p".into(),
692                    upstream_model: "n".into(),
693                },
694            )
695            .with_provider("p", Arc::new(MockExecutor))
696            .build();
697
698        assert!(result.is_ok());
699    }
700
701    #[test]
702    fn builder_runtime_starts_with_registry_version_one() {
703        let runtime = AgentRuntimeBuilder::new()
704            .with_agent_spec(AgentSpec {
705                id: "versioned-agent".into(),
706                model_id: "m".into(),
707                system_prompt: "sys".into(),
708                ..Default::default()
709            })
710            .with_model_binding(
711                "m",
712                ModelBinding {
713                    provider_id: "mock".into(),
714                    upstream_model: "model-v1".into(),
715                },
716            )
717            .with_provider("mock", Arc::new(MockExecutor))
718            .build()
719            .unwrap();
720
721        assert_eq!(runtime.registry_version(), Some(1));
722        assert!(runtime.registry_handle().is_some());
723    }
724
725    #[test]
726    fn replacing_registry_set_updates_dynamic_resolver() {
727        let runtime = AgentRuntimeBuilder::new()
728            .with_agent_spec(AgentSpec {
729                id: "agent-v1".into(),
730                model_id: "m".into(),
731                system_prompt: "sys".into(),
732                ..Default::default()
733            })
734            .with_model_binding(
735                "m",
736                ModelBinding {
737                    provider_id: "mock".into(),
738                    upstream_model: "model-v1".into(),
739                },
740            )
741            .with_provider("mock", Arc::new(MockExecutor))
742            .build()
743            .unwrap();
744
745        assert!(runtime.resolver().resolve("agent-v1").is_ok());
746        assert!(runtime.resolver().resolve("agent-v2").is_err());
747
748        let version = runtime
749            .replace_registry_set(make_registry_set("agent-v2", "m2", "model-v2"))
750            .expect("builder runtimes should expose a registry handle");
751
752        assert_eq!(version, 2);
753        assert_eq!(runtime.registry_version(), Some(2));
754        assert!(runtime.resolver().resolve("agent-v1").is_err());
755
756        let resolved = runtime.resolver().resolve("agent-v2").unwrap();
757        assert_eq!(resolved.id(), "agent-v2");
758        assert_eq!(resolved.upstream_model, "model-v2");
759    }
760
761    #[test]
762    fn builder_model_binding_provider_name() {
763        let spec = AgentSpec {
764            id: "agent".into(),
765            model_id: "gpt-4".into(),
766            system_prompt: "sys".into(),
767            ..Default::default()
768        };
769
770        let runtime = AgentRuntimeBuilder::new()
771            .with_agent_spec(spec)
772            .with_model_binding(
773                "gpt-4",
774                ModelBinding {
775                    provider_id: "openai".into(),
776                    upstream_model: "gpt-4-turbo".into(),
777                },
778            )
779            .with_provider("openai", Arc::new(MockExecutor))
780            .build()
781            .unwrap();
782
783        let resolved = runtime.resolver().resolve("agent").unwrap();
784        // The model ID should resolve to the upstream model name
785        assert_eq!(resolved.upstream_model, "gpt-4-turbo");
786    }
787
788    #[test]
789    fn builder_with_profile_store() {
790        use awaken_contract::contract::profile_store::{
791            ProfileEntry, ProfileOwner as POwner, ProfileStore,
792        };
793        use awaken_contract::contract::storage::StorageError;
794
795        struct NoOpProfileStore;
796
797        #[async_trait]
798        impl ProfileStore for NoOpProfileStore {
799            async fn get(
800                &self,
801                _owner: &POwner,
802                _key: &str,
803            ) -> Result<Option<ProfileEntry>, StorageError> {
804                Ok(None)
805            }
806            async fn set(
807                &self,
808                _owner: &POwner,
809                _key: &str,
810                _value: Value,
811            ) -> Result<(), StorageError> {
812                Ok(())
813            }
814            async fn delete(&self, _owner: &POwner, _key: &str) -> Result<(), StorageError> {
815                Ok(())
816            }
817            async fn list(&self, _owner: &POwner) -> Result<Vec<ProfileEntry>, StorageError> {
818                Ok(vec![])
819            }
820            async fn clear_owner(&self, _owner: &POwner) -> Result<(), StorageError> {
821                Ok(())
822            }
823        }
824
825        let runtime = AgentRuntimeBuilder::new()
826            .with_profile_store(Arc::new(NoOpProfileStore))
827            .build()
828            .unwrap();
829        assert!(runtime.profile_store.is_some());
830    }
831
832    #[cfg(feature = "a2a")]
833    #[test]
834    fn build_allows_endpoint_agents_when_backend_factory_exists() {
835        let validate_count = Arc::new(AtomicUsize::new(0));
836        let build_count = Arc::new(AtomicUsize::new(0));
837        let runtime = AgentRuntimeBuilder::new()
838            .with_agent_spec(
839                AgentSpec::new("remote-agent")
840                    .with_model_id("remote-model")
841                    .with_system_prompt("remote")
842                    .with_endpoint(RemoteEndpoint {
843                        backend: "counting-remote".into(),
844                        base_url: "https://remote.example.com".into(),
845                        ..Default::default()
846                    }),
847            )
848            .with_agent_backend_factory(Arc::new(CountingValidationBackendFactory {
849                validate_count: validate_count.clone(),
850                build_count: build_count.clone(),
851            }))
852            .build()
853            .expect("endpoint agents should validate through backend factory");
854
855        let spec = runtime
856            .registry_set()
857            .and_then(|set| set.agents.get_agent("remote-agent"))
858            .expect("remote agent should remain registered");
859        assert!(spec.endpoint.is_some());
860        assert_eq!(validate_count.load(Ordering::SeqCst), 1);
861        assert_eq!(build_count.load(Ordering::SeqCst), 0);
862    }
863
864    #[test]
865    fn duplicate_agent_spec_errors_at_build() {
866        let spec = AgentSpec {
867            id: "dup-agent".into(),
868            model_id: "m".into(),
869            system_prompt: "sys".into(),
870            ..Default::default()
871        };
872
873        let result = AgentRuntimeBuilder::new()
874            .with_agent_spec(spec.clone())
875            .with_agent_spec(spec)
876            .with_model_binding(
877                "m",
878                ModelBinding {
879                    provider_id: "p".into(),
880                    upstream_model: "n".into(),
881                },
882            )
883            .with_provider("p", Arc::new(MockExecutor))
884            .build();
885
886        match result {
887            Err(e) => {
888                let err = e.to_string();
889                assert!(
890                    err.contains("dup-agent"),
891                    "error should mention the duplicate agent ID: {err}"
892                );
893            }
894            Ok(_) => panic!("expected build to fail for duplicate agent"),
895        }
896    }
897
898    #[test]
899    fn duplicate_tool_errors_at_build() {
900        let result = AgentRuntimeBuilder::new()
901            .with_tool(
902                "dup-tool",
903                Arc::new(MockTool {
904                    id: "dup-tool".into(),
905                }),
906            )
907            .with_tool(
908                "dup-tool",
909                Arc::new(MockTool {
910                    id: "dup-tool".into(),
911                }),
912            )
913            .build();
914
915        match result {
916            Err(e) => {
917                let err = e.to_string();
918                assert!(
919                    err.contains("dup-tool"),
920                    "error should mention the duplicate tool ID: {err}"
921                );
922            }
923            Ok(_) => panic!("expected build to fail for duplicate tool"),
924        }
925    }
926
927    #[test]
928    fn duplicate_model_errors_at_build() {
929        let result = AgentRuntimeBuilder::new()
930            .with_model_binding(
931                "dup-model",
932                ModelBinding {
933                    provider_id: "p".into(),
934                    upstream_model: "n1".into(),
935                },
936            )
937            .with_model_binding(
938                "dup-model",
939                ModelBinding {
940                    provider_id: "p".into(),
941                    upstream_model: "n2".into(),
942                },
943            )
944            .build();
945
946        match result {
947            Err(e) => {
948                let err = e.to_string();
949                assert!(
950                    err.contains("dup-model"),
951                    "error should mention the duplicate model ID: {err}"
952                );
953            }
954            Ok(_) => panic!("expected build to fail for duplicate model"),
955        }
956    }
957
958    #[test]
959    fn duplicate_provider_errors_at_build() {
960        let result = AgentRuntimeBuilder::new()
961            .with_provider("dup-prov", Arc::new(MockExecutor))
962            .with_provider("dup-prov", Arc::new(MockExecutor))
963            .build();
964
965        match result {
966            Err(e) => {
967                let err = e.to_string();
968                assert!(
969                    err.contains("dup-prov"),
970                    "error should mention the duplicate provider ID: {err}"
971                );
972            }
973            Ok(_) => panic!("expected build to fail for duplicate provider"),
974        }
975    }
976
977    #[cfg(feature = "a2a")]
978    #[test]
979    fn duplicate_backend_factory_errors_at_build() {
980        let result = AgentRuntimeBuilder::new()
981            .with_agent_backend_factory(Arc::new(crate::extensions::a2a::A2aBackendFactory))
982            .build();
983
984        match result {
985            Err(BuildError::BackendRegistryConflict(err)) => {
986                assert!(
987                    err.contains("a2a"),
988                    "error should mention the duplicate backend kind: {err}"
989                );
990            }
991            Err(other) => panic!("expected backend registry conflict, got {other}"),
992            Ok(_) => panic!("expected build to fail for duplicate backend factory"),
993        }
994    }
995}