use super::*;
use crate::contracts::runtime::tool_call::Tool;
use crate::contracts::runtime::AgentBehavior;
use crate::contracts::storage::ThreadStore;
#[cfg(feature = "skills")]
use crate::runtime::wiring::SkillsSystemWiring;
use crate::runtime::StopPolicy;
use crate::runtime::{AgentOs, RuntimeServices};
use genai::Client;
use std::collections::HashMap;
use std::sync::Arc;
#[cfg(feature = "skills")]
use std::time::Duration;
#[cfg(all(feature = "skills", feature = "mcp"))]
use tirea_extension_mcp::McpToolRegistryManager;
#[cfg(all(feature = "skills", feature = "mcp"))]
use tirea_extension_skills::McpPromptSkillRegistryManager;
#[cfg(feature = "skills")]
use tirea_extension_skills::{
CompositeSkillRegistry, InMemorySkillRegistry, Skill, SkillRegistry, SkillRegistryManagerError,
};
pub struct AgentOsBuilder {
pub(crate) client: Option<Client>,
pub(crate) bundles: Vec<Arc<dyn RegistryBundle>>,
pub(crate) agents: HashMap<String, AgentDefinition>,
pub(crate) agent_registries: Vec<Arc<dyn AgentRegistry>>,
pub(crate) resolved_agents: HashMap<String, ResolvedAgent>,
pub(crate) agent_catalogs: Vec<Arc<dyn AgentCatalog>>,
pub(crate) base_tools: HashMap<String, Arc<dyn Tool>>,
pub(crate) base_tool_registries: Vec<Arc<dyn ToolRegistry>>,
pub(crate) behaviors: HashMap<String, Arc<dyn AgentBehavior>>,
pub(crate) behavior_registries: Vec<Arc<dyn BehaviorRegistry>>,
pub(crate) stop_policies: HashMap<String, Arc<dyn StopPolicy>>,
pub(crate) stop_policy_registries: Vec<Arc<dyn StopPolicyRegistry>>,
pub(crate) providers: HashMap<String, Client>,
pub(crate) provider_registries: Vec<Arc<dyn ProviderRegistry>>,
pub(crate) models: HashMap<String, ModelDefinition>,
pub(crate) model_registries: Vec<Arc<dyn ModelRegistry>>,
#[cfg(feature = "skills")]
pub(crate) skills: Vec<Arc<dyn Skill>>,
#[cfg(feature = "skills")]
pub(crate) skill_registries: Vec<Arc<dyn SkillRegistry>>,
#[cfg(feature = "skills")]
pub(crate) skills_refresh_interval: Option<Duration>,
#[cfg(feature = "skills")]
pub(crate) skills_config: SkillsConfig,
pub(crate) system_wirings: Vec<Arc<dyn SystemWiring>>,
pub(crate) agent_tools: AgentToolsConfig,
pub(crate) agent_state_store: Option<Arc<dyn ThreadStore>>,
}
fn merge_registry<R: ?Sized, M>(
memory: M,
mut external: Vec<Arc<R>>,
is_memory_empty: impl Fn(&M) -> bool,
into_memory_registry: impl FnOnce(M) -> Arc<R>,
compose: impl FnOnce(Vec<Arc<R>>) -> Result<Arc<R>, AgentOsBuildError>,
) -> Result<Arc<R>, AgentOsBuildError> {
if external.is_empty() {
return Ok(into_memory_registry(memory));
}
let mut registries: Vec<Arc<R>> = Vec::new();
if !is_memory_empty(&memory) {
registries.push(into_memory_registry(memory));
}
registries.append(&mut external);
if registries.len() == 1 {
Ok(registries.pop().expect("single registry must exist"))
} else {
compose(registries)
}
}
impl std::fmt::Debug for AgentOs {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut s = f.debug_struct("AgentOs");
s.field("default_client", &"[genai::Client]")
.field("agents", &self.agents.len())
.field("agent_catalog", &self.agent_catalog.len())
.field("base_tools", &self.base_tools.len())
.field("behaviors", &self.behaviors.len())
.field("stop_policies", &self.stop_policies.len())
.field("providers", &self.providers.len())
.field("models", &self.models.len());
#[cfg(feature = "skills")]
s.field(
"skills",
&self.skills_registry.as_ref().map(|registry| registry.len()),
);
s.field("system_wirings", &self.system_wirings.len())
.field("active_runs", &"[internal]")
.field("agent_tools", &self.agent_tools)
.field("agent_state_store", &self.agent_state_store.is_some())
.finish()
}
}
impl std::fmt::Debug for AgentOsBuilder {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut s = f.debug_struct("AgentOsBuilder");
s.field("client", &self.client.is_some())
.field("bundles", &self.bundles.len())
.field("agents", &self.agents.len())
.field("resolved_agents", &self.resolved_agents.len())
.field("agent_catalogs", &self.agent_catalogs.len())
.field("base_tools", &self.base_tools.len())
.field("behaviors", &self.behaviors.len())
.field("stop_policies", &self.stop_policies.len())
.field("stop_policy_registries", &self.stop_policy_registries.len())
.field("providers", &self.providers.len())
.field("models", &self.models.len());
#[cfg(feature = "skills")]
{
s.field("skills", &self.skills.len())
.field("skill_registries", &self.skill_registries.len())
.field("skills_refresh_interval", &self.skills_refresh_interval)
.field("skills_config", &self.skills_config);
}
s.field("system_wirings", &self.system_wirings.len())
.field("agent_tools", &self.agent_tools)
.field("agent_state_store", &self.agent_state_store.is_some())
.finish()
}
}
impl AgentOsBuilder {
fn insert_local_agent_definition(&mut self, agent_id: String, mut definition: AgentDefinition) {
definition.id = agent_id.clone();
self.agents.insert(agent_id, definition);
}
fn insert_resolved_agent_definition(&mut self, agent_id: String, mut agent: ResolvedAgent) {
agent.descriptor.id = agent_id.clone();
if agent.descriptor.name.trim().is_empty() {
agent.descriptor.name = agent_id.clone();
}
self.resolved_agents.insert(agent_id, agent);
}
fn insert_agent_spec(&mut self, spec: AgentDefinitionSpec) {
match spec {
AgentDefinitionSpec::Local(definition) => {
let definition = *definition;
let agent_id = definition.id.clone();
self.insert_local_agent_definition(agent_id, definition);
}
AgentDefinitionSpec::Remote(definition) => {
let agent_id = definition.id().to_string();
self.insert_resolved_agent_definition(agent_id, definition.into_resolved_agent());
}
}
}
pub fn new() -> Self {
Self {
client: None,
bundles: Vec::new(),
agents: HashMap::new(),
agent_registries: Vec::new(),
resolved_agents: HashMap::new(),
agent_catalogs: Vec::new(),
base_tools: HashMap::new(),
base_tool_registries: Vec::new(),
behaviors: HashMap::new(),
behavior_registries: Vec::new(),
stop_policies: HashMap::new(),
stop_policy_registries: Vec::new(),
providers: HashMap::new(),
provider_registries: Vec::new(),
models: HashMap::new(),
model_registries: Vec::new(),
#[cfg(feature = "skills")]
skills: Vec::new(),
#[cfg(feature = "skills")]
skill_registries: Vec::new(),
#[cfg(feature = "skills")]
skills_refresh_interval: None,
#[cfg(feature = "skills")]
skills_config: SkillsConfig::default(),
system_wirings: Vec::new(),
agent_tools: AgentToolsConfig::default(),
agent_state_store: None,
}
}
pub fn with_client(mut self, client: Client) -> Self {
self.client = Some(client);
self
}
pub fn with_bundle(mut self, bundle: Arc<dyn RegistryBundle>) -> Self {
self.bundles.push(bundle);
self
}
pub fn with_agent_spec(mut self, spec: AgentDefinitionSpec) -> Self {
self.insert_agent_spec(spec);
self
}
pub fn with_agent_specs(
mut self,
specs: impl IntoIterator<Item = AgentDefinitionSpec>,
) -> Self {
for spec in specs {
self = self.with_agent_spec(spec);
}
self
}
pub fn with_agent_registry(mut self, registry: Arc<dyn AgentRegistry>) -> Self {
self.agent_registries.push(registry);
self
}
pub fn with_agent_catalog(mut self, catalog: Arc<dyn AgentCatalog>) -> Self {
self.agent_catalogs.push(catalog);
self
}
pub fn with_tools(mut self, tools: HashMap<String, Arc<dyn Tool>>) -> Self {
self.base_tools = tools;
self
}
pub fn with_tool_registry(mut self, registry: Arc<dyn ToolRegistry>) -> Self {
self.base_tool_registries.push(registry);
self
}
pub fn with_registered_behavior(
mut self,
behavior_id: impl Into<String>,
behavior: Arc<dyn AgentBehavior>,
) -> Self {
self.behaviors.insert(behavior_id.into(), behavior);
self
}
pub fn with_behavior_registry(mut self, registry: Arc<dyn BehaviorRegistry>) -> Self {
self.behavior_registries.push(registry);
self
}
pub fn with_stop_policy(mut self, id: impl Into<String>, policy: Arc<dyn StopPolicy>) -> Self {
self.stop_policies.insert(id.into(), policy);
self
}
pub fn with_stop_policy_registry(mut self, registry: Arc<dyn StopPolicyRegistry>) -> Self {
self.stop_policy_registries.push(registry);
self
}
pub fn with_provider(mut self, provider_id: impl Into<String>, client: Client) -> Self {
self.providers.insert(provider_id.into(), client);
self
}
pub fn with_provider_registry(mut self, registry: Arc<dyn ProviderRegistry>) -> Self {
self.provider_registries.push(registry);
self
}
pub fn with_model(mut self, model_id: impl Into<String>, def: ModelDefinition) -> Self {
self.models.insert(model_id.into(), def);
self
}
pub fn with_models(mut self, defs: HashMap<String, ModelDefinition>) -> Self {
self.models = defs;
self
}
pub fn with_model_registry(mut self, registry: Arc<dyn ModelRegistry>) -> Self {
self.model_registries.push(registry);
self
}
#[cfg(feature = "skills")]
pub fn with_skills(mut self, skills: Vec<Arc<dyn Skill>>) -> Self {
self.skills = skills;
self
}
#[cfg(feature = "skills")]
pub fn with_skill_registry(mut self, registry: Arc<dyn SkillRegistry>) -> Self {
self.skill_registries.push(registry);
self
}
#[cfg(all(feature = "skills", feature = "mcp"))]
pub async fn with_mcp_prompt_skills(
mut self,
manager: Arc<McpToolRegistryManager>,
) -> Result<Self, SkillRegistryManagerError> {
let registry: Arc<dyn SkillRegistry> =
Arc::new(McpPromptSkillRegistryManager::discover(manager).await?);
self.skill_registries.push(registry);
Ok(self)
}
#[cfg(feature = "skills")]
pub fn with_skill_registry_refresh_interval(mut self, interval: Duration) -> Self {
self.skills_refresh_interval = Some(interval);
self
}
#[cfg(feature = "skills")]
pub fn with_skills_config(mut self, cfg: SkillsConfig) -> Self {
self.skills_config = cfg;
self
}
pub fn with_system_wiring(mut self, wiring: Arc<dyn SystemWiring>) -> Self {
self.system_wirings.push(wiring);
self
}
pub fn with_agent_tools_config(mut self, cfg: AgentToolsConfig) -> Self {
self.agent_tools = cfg;
self
}
pub fn with_agent_state_store(mut self, agent_state_store: Arc<dyn ThreadStore>) -> Self {
self.agent_state_store = Some(agent_state_store);
self
}
pub fn build(self) -> Result<AgentOs, AgentOsBuildError> {
let AgentOsBuilder {
client,
bundles,
agents: mut agents_defs,
mut agent_registries,
resolved_agents: mut resolved_agent_defs,
mut agent_catalogs,
base_tools: mut base_tools_defs,
mut base_tool_registries,
behaviors: mut behavior_defs,
mut behavior_registries,
stop_policies: stop_policy_defs,
stop_policy_registries,
providers: mut provider_defs,
mut provider_registries,
models: mut model_defs,
mut model_registries,
#[cfg(feature = "skills")]
skills,
#[cfg(feature = "skills")]
mut skill_registries,
#[cfg(feature = "skills")]
skills_refresh_interval,
#[cfg(feature = "skills")]
skills_config,
system_wirings,
agent_tools,
agent_state_store,
} = self;
BundleComposer::apply(
&bundles,
BundleRegistryAccumulator {
agent_definitions: &mut agents_defs,
agent_registries: &mut agent_registries,
tool_definitions: &mut base_tools_defs,
tool_registries: &mut base_tool_registries,
behavior_definitions: &mut behavior_defs,
behavior_registries: &mut behavior_registries,
provider_definitions: &mut provider_defs,
provider_registries: &mut provider_registries,
model_definitions: &mut model_defs,
model_registries: &mut model_registries,
},
)?;
#[allow(unused_mut)]
let mut system_wirings = system_wirings;
#[cfg(feature = "skills")]
let skills_registry = {
if skills_config.enabled && skills.is_empty() && skill_registries.is_empty() {
return Err(AgentOsBuildError::SkillsNotConfigured);
}
let mut in_memory_skills = InMemorySkillRegistry::new();
in_memory_skills.extend_upsert(skills);
let registry = if in_memory_skills.is_empty() && skill_registries.is_empty() {
None
} else {
Some(merge_registry(
in_memory_skills,
std::mem::take(&mut skill_registries),
|reg: &InMemorySkillRegistry| reg.is_empty(),
|reg| Arc::new(reg),
|regs| Ok(Arc::new(CompositeSkillRegistry::try_new(regs)?)),
)?)
};
if let (Some(r), Some(interval)) = (®istry, skills_refresh_interval) {
match r.start_periodic_refresh(interval) {
Ok(()) | Err(SkillRegistryManagerError::PeriodicRefreshAlreadyRunning) => {}
Err(err) => return Err(err.into()),
}
}
if skills_config.enabled {
if let Some(ref reg) = registry {
system_wirings.push(Arc::new(SkillsSystemWiring::new(
reg.clone(),
skills_config.clone(),
)));
}
}
registry
};
let mut base_tools = InMemoryToolRegistry::new();
base_tools.extend_named(base_tools_defs)?;
let base_tools: Arc<dyn ToolRegistry> = merge_registry(
base_tools,
base_tool_registries,
|reg: &InMemoryToolRegistry| reg.is_empty(),
|reg| Arc::new(reg),
|regs| Ok(Arc::new(CompositeToolRegistry::try_new(regs)?)),
)?;
let mut behaviors = InMemoryBehaviorRegistry::new();
behaviors.extend_named(behavior_defs)?;
let behaviors: Arc<dyn BehaviorRegistry> = merge_registry(
behaviors,
behavior_registries,
|reg: &InMemoryBehaviorRegistry| reg.is_empty(),
|reg| Arc::new(reg),
|regs| Ok(Arc::new(CompositeBehaviorRegistry::try_new(regs)?)),
)?;
let mut stop_policies_mem = InMemoryStopPolicyRegistry::new();
stop_policies_mem.extend_named(stop_policy_defs)?;
let stop_policies: Arc<dyn StopPolicyRegistry> = merge_registry(
stop_policies_mem,
stop_policy_registries,
|reg: &InMemoryStopPolicyRegistry| reg.is_empty(),
|reg| Arc::new(reg),
|regs| Ok(Arc::new(CompositeStopPolicyRegistry::try_new(regs)?)),
)?;
{
let reserved = AgentOs::reserved_behavior_ids(&system_wirings);
for (agent_id, def) in &agents_defs {
let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
for id in &def.behavior_ids {
let id = id.trim();
if id.is_empty() {
return Err(AgentOsBuildError::AgentEmptyBehaviorRef {
agent_id: agent_id.clone(),
});
}
if reserved.contains(&id) {
return Err(AgentOsBuildError::AgentReservedBehaviorId {
agent_id: agent_id.clone(),
behavior_id: id.to_string(),
});
}
if !seen.insert(id.to_string()) {
return Err(AgentOsBuildError::AgentDuplicateBehaviorRef {
agent_id: agent_id.clone(),
behavior_id: id.to_string(),
});
}
if behaviors.get(id).is_none() {
return Err(AgentOsBuildError::AgentBehaviorNotFound {
agent_id: agent_id.clone(),
behavior_id: id.to_string(),
});
}
}
let mut sc_seen: std::collections::HashSet<String> =
std::collections::HashSet::new();
for sc_id in &def.stop_condition_ids {
let sc_id = sc_id.trim();
if sc_id.is_empty() {
return Err(AgentOsBuildError::AgentEmptyStopConditionRef {
agent_id: agent_id.clone(),
});
}
if !sc_seen.insert(sc_id.to_string()) {
return Err(AgentOsBuildError::AgentDuplicateStopConditionRef {
agent_id: agent_id.clone(),
stop_condition_id: sc_id.to_string(),
});
}
if stop_policies.get(sc_id).is_none() {
return Err(AgentOsBuildError::AgentStopConditionNotFound {
agent_id: agent_id.clone(),
stop_condition_id: sc_id.to_string(),
});
}
}
}
}
let mut providers = InMemoryProviderRegistry::new();
providers.extend(provider_defs)?;
let providers: Arc<dyn ProviderRegistry> = merge_registry(
providers,
provider_registries,
|reg: &InMemoryProviderRegistry| reg.is_empty(),
|reg| Arc::new(reg),
|regs| Ok(Arc::new(CompositeProviderRegistry::try_new(regs)?)),
)?;
let mut models = InMemoryModelRegistry::new();
models.extend(model_defs.clone())?;
let models: Arc<dyn ModelRegistry> = merge_registry(
models,
model_registries,
|reg: &InMemoryModelRegistry| reg.is_empty(),
|reg| Arc::new(reg),
|regs| Ok(Arc::new(CompositeModelRegistry::try_new(regs)?)),
)?;
if !models.is_empty() && providers.is_empty() {
return Err(AgentOsBuildError::ProvidersNotConfigured);
}
for (model_id, def) in models.snapshot() {
if providers.get(&def.provider).is_none() {
return Err(AgentOsBuildError::ProviderNotFound {
provider_id: def.provider,
model_id,
});
}
}
let mut agents = InMemoryAgentRegistry::new();
agents.extend_upsert(agents_defs);
let agents: Arc<dyn AgentRegistry> = merge_registry(
agents,
agent_registries,
|reg: &InMemoryAgentRegistry| reg.is_empty(),
|reg| Arc::new(reg),
|regs| Ok(Arc::new(CompositeAgentRegistry::try_new(regs)?)),
)?;
let mut static_agents = InMemoryAgentCatalog::new();
static_agents.extend_upsert(std::mem::take(&mut resolved_agent_defs));
agent_catalogs.insert(
0,
Arc::new(HostedAgentCatalog::new(agents.clone())) as Arc<dyn AgentCatalog>,
);
let agent_catalog: Arc<dyn AgentCatalog> = merge_registry(
static_agents,
agent_catalogs,
|catalog: &InMemoryAgentCatalog| catalog.is_empty(),
|catalog| Arc::new(catalog),
|catalogs| Ok(Arc::new(CompositeAgentCatalog::try_new(catalogs)?)),
)?;
let registries = RegistrySet::new(
agents,
base_tools,
behaviors,
providers,
models,
stop_policies,
#[cfg(feature = "skills")]
skills_registry,
);
let services = RuntimeServices {
default_client: client.unwrap_or_default(),
system_wirings,
agent_tools,
agent_state_store,
agent_catalog,
};
Ok(AgentOs::from_registry_set(registries, services))
}
}
impl Default for AgentOsBuilder {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn with_agent_spec_routes_local_and_remote_into_runtime_surfaces() {
let mut local_definition = AgentDefinition::new("mock");
local_definition.id = "local-worker".to_string();
let os = AgentOsBuilder::new()
.with_agent_spec(AgentDefinitionSpec::local(
local_definition
.with_name("Local Worker")
.with_description("Hosted locally"),
))
.with_agent_spec(AgentDefinitionSpec::a2a(
AgentDescriptor::new("remote-worker")
.with_name("Remote Worker")
.with_description("Delegated over A2A"),
A2aAgentBinding::new("https://example.test/v1/a2a", "remote-worker"),
))
.build()
.expect("builder should accept unified agent specs");
assert_eq!(
os.agent("local-worker")
.expect("local agent should be registered")
.display_name(),
"Local Worker"
);
let remote = os
.agent_catalog()
.get("remote-worker")
.expect("remote agent should be discoverable");
assert_eq!(remote.descriptor.name, "Remote Worker");
assert!(matches!(remote.binding, AgentBinding::A2a(_)));
}
#[test]
fn unified_builder_entrypoints_normalize_ids_and_names() {
let os = AgentOsBuilder::new()
.with_agent_spec(AgentDefinitionSpec::local_with_id(
"local-worker",
AgentDefinition::new("mock")
.with_name("Local Worker")
.with_description("Hosted locally"),
))
.with_agent_spec(AgentDefinitionSpec::a2a(
AgentDescriptor::new("remote-worker")
.with_name(" ")
.with_description("Delegated over A2A"),
A2aAgentBinding::new("https://example.test/v1/a2a", "remote-worker"),
))
.build()
.expect("unified builder entrypoints should build");
assert_eq!(
os.agent("local-worker")
.expect("local agent should be registered")
.display_name(),
"Local Worker"
);
let remote = os
.agent_catalog()
.get("remote-worker")
.expect("remote agent should be discoverable");
assert_eq!(remote.descriptor.id, "remote-worker");
assert_eq!(remote.descriptor.name, "remote-worker");
assert_eq!(remote.descriptor.description, "Delegated over A2A");
}
#[test]
fn build_rejects_duplicate_local_and_remote_agent_ids() {
let err = AgentOsBuilder::new()
.with_agent_spec(AgentDefinitionSpec::local_with_id(
"worker",
AgentDefinition::new("mock"),
))
.with_agent_spec(AgentDefinitionSpec::a2a_with_id(
"worker",
A2aAgentBinding::new("https://example.test/v1/a2a", "worker"),
))
.build()
.expect_err("duplicate local/remote ids should be rejected");
assert!(matches!(
err,
AgentOsBuildError::AgentCatalog(AgentCatalogError::AgentIdConflict(id)) if id == "worker"
));
}
}