use crate::command::{
CommandDescriptor, CommandExecutionContext, CommandResult, ExecuteCommandRequest,
};
use crate::deployment::DeploymentGrade;
use crate::events::TokenUsage;
use crate::mcp_server::{ScopedMcpServers, merge_scoped_mcp_servers};
use crate::message::Message;
use crate::message_filter::MessageFilterProvider;
use crate::runtime_agent::RuntimeAgent;
use crate::tool_types::{ToolCall, ToolDefinition};
use crate::tools::{Tool, ToolRegistry};
use crate::traits::SessionFileSystem;
use crate::typed_id::SessionId;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
pub struct IntegrationPlugin {
pub experimental_only: bool,
pub feature_flag: Option<&'static str>,
pub factory: fn() -> Box<dyn Capability>,
}
inventory::collect!(IntegrationPlugin);
pub use crate::capability_types::{
AgentCapabilityConfig, CapabilityId, CapabilityStatus, MountAccess, MountDirectoryBuilder,
MountEntry, MountPoint, MountSource,
};
mod a2a_delegation;
#[cfg(feature = "ui-capabilities")]
mod a2ui;
mod agent_handoff;
mod agent_instructions;
pub mod attach_skill;
mod auto_tool_search;
mod background_execution;
mod btw;
mod budgeting;
pub mod compaction;
mod current_time;
mod data_knowledge;
mod declarative;
mod fake_aws;
mod fake_crm;
mod fake_financial;
mod fake_warehouse;
mod file_system;
mod human_intent;
mod infinity_context;
mod knowledge_base;
mod loop_detection;
mod lua;
mod lua_code_mode;
pub mod mcp;
mod noop;
mod openai_tool_search;
#[cfg(feature = "ui-capabilities")]
mod openui;
pub mod persistent_memory;
mod platform_management;
mod prompt_caching;
mod prompt_canary_guardrail;
mod research;
mod sample_data;
mod self_budget;
mod session;
mod session_sandbox;
mod session_schedule;
mod session_sql_database;
mod session_storage;
mod skills;
mod stateless_todo_list;
mod subagents;
mod system_commands;
mod test_math;
mod test_weather;
mod tool_output_persistence;
mod tool_search;
pub mod user_hooks;
mod virtual_bash;
mod web_fetch;
mod workspace_volumes;
pub use a2a_delegation::{
A2A_AGENT_DELEGATION_CAPABILITY_ID, A2aAgentDelegationCapability, CancelAgentTool,
GetAgentRunsTool, MessageAgentTool, SpawnAgentTool, WaitAgentTool,
};
#[cfg(feature = "ui-capabilities")]
pub use a2ui::{A2UI_CAPABILITY_ID, A2UiCapability};
pub use agent_handoff::{
AGENT_HANDOFF_CAPABILITY_ID, AgentHandoffCapability, GetAgentHandoffsTool,
MessageAgentHandoffTool, StartAgentHandoffTool,
};
pub use agent_instructions::{
AGENT_INSTRUCTIONS_CAPABILITY_ID, AGENTS_MD_PATH, AgentInstructionsCapability,
AgentInstructionsConfig, DEFAULT_AGENT_INSTRUCTIONS_FILE, MAX_AGENT_INSTRUCTIONS_FILES,
MAX_AGENTS_MD_SIZE, format_agents_md_content, format_instruction_file_content,
};
pub use attach_skill::{
AttachSkillCapability, SKILL_CAPABILITY_PREFIX, SKILLS_DISCOVERY_PATH, SkillContribution,
SkillInstructions, SkillMeta, SkillSource, discover_skills_from_entries, is_skill_capability,
parse_skill_capability_id, reconstruct_skill_md, skill_capability_id,
};
pub use auto_tool_search::{AUTO_TOOL_SEARCH_CAPABILITY_ID, AutoToolSearchCapability};
pub use background_execution::{BACKGROUND_EXECUTION_CAPABILITY_ID, BackgroundExecutionCapability};
pub use btw::{BTW_CAPABILITY_ID, BtwCapability};
pub use budgeting::BudgetingCapability;
pub use compaction::{
COMPACTION_CAPABILITY_ID, CompactionCapability, CompactionConfig, CompactionStep,
CompactionStrategy, CostControlConfig, CostControlMaskingResult, HierarchicalMemoryConfig,
MaskingSummaryFormat, MemoryTier, ObservationMaskingConfig, ObservationMaskingResult,
SessionCompactionMetrics, SummarizationConfig, aggressive_trim, apply_cost_control_masking,
apply_hierarchical_memory, apply_observation_masking, build_model_view_messages,
build_summarization_prompt, build_summary_message, classify_memory_tiers, estimate_tokens,
estimate_total_tokens, format_messages_for_summarization, should_compact_proactively,
};
pub use current_time::{CurrentTimeCapability, GetCurrentTimeTool};
pub use data_knowledge::DataKnowledgeCapability;
pub use declarative::{
DECLARATIVE_CAPABILITY_PREFIX, DeclarativeCapabilityDefinition, DeclarativeCapabilityFile,
DeclarativeCapabilitySkill, declarative_capability_id, declarative_capability_info,
hydrate_declarative_capability_config, is_declarative_capability,
parse_declarative_capability_id, validate_declarative_capability_definition,
};
pub use fake_aws::{
AwsCreateEc2InstanceTool, AwsCreateIamUserTool, AwsCreateRdsDatabaseTool,
AwsCreateS3BucketTool, AwsGetCloudWatchMetricsTool, AwsListEc2InstancesTool,
AwsListIamUsersTool, AwsListRdsDatabasesTool, AwsListS3BucketsTool, AwsListSecurityGroupsTool,
AwsStopEc2InstanceTool, FakeAwsCapability,
};
pub use fake_crm::{
CrmAddInteractionTool, CrmCreateCustomerTool, CrmCreateTicketTool, CrmGetCustomerTool,
CrmListCustomersTool, CrmListTicketsTool, CrmSearchCustomersTool, CrmUpdateTicketTool,
FakeCrmCapability,
};
pub use fake_financial::{
FakeFinancialCapability, FinanceCreateBudgetTool, FinanceCreateTransactionTool,
FinanceForecastCashFlowTool, FinanceGetBalanceTool, FinanceGetExpenseReportTool,
FinanceGetRevenueReportTool, FinanceListBudgetsTool, FinanceListTransactionsTool,
};
pub use fake_warehouse::{
FakeWarehouseCapability, WarehouseCreateInvoiceTool, WarehouseCreateOrderTool,
WarehouseCreateShipmentTool, WarehouseGetInventoryTool, WarehouseInventoryReportTool,
WarehouseListOrdersTool, WarehouseListShipmentsTool, WarehouseProcessReturnTool,
WarehouseUpdateInventoryTool, WarehouseUpdateShipmentStatusTool,
};
pub use file_system::{
DeleteFileTool, EditFileTool, FileSystemCapability, GrepFilesTool, ListDirectoryTool,
ReadFileTool, StatFileTool, WriteFileTool,
};
pub use human_intent::{HUMAN_INTENT_CAPABILITY_ID, HumanIntentCapability};
pub use infinity_context::{
INFINITY_CONTEXT_CAPABILITY_ID, InfinityContextCapability, QueryHistoryTool,
};
pub use knowledge_base::{
KNOWLEDGE_BASE_CAPABILITY_ID, KnowledgeBaseCapability, KnowledgeBaseConfig,
validate_knowledge_base_config,
};
pub use loop_detection::LoopDetectionCapability;
pub use lua::{LUA_CAPABILITY_ID, LuaCapability, LuaTool, LuaVfs, is_code_mode_eligible};
pub use lua_code_mode::{LUA_CODE_MODE_CAPABILITY_ID, LuaCodeModeCapability};
pub use mcp::{
MCP_CAPABILITY_PREFIX, McpCapability, is_mcp_capability, mcp_capability_id,
parse_mcp_capability_id,
};
pub use noop::NoopCapability;
pub use openai_tool_search::{
DEFAULT_TOOL_SEARCH_THRESHOLD, OPENAI_TOOL_SEARCH_CAPABILITY_ID, OpenAiToolSearchCapability,
model_supports_native_tool_search,
};
#[cfg(feature = "ui-capabilities")]
pub use openui::{OPENUI_CAPABILITY_ID, OpenUiCapability};
pub use persistent_memory::{
ForgetTool, MEMORY_CAPABILITY_ID, MemoryCapability, MemoryConfig, RecallTool, RememberTool,
};
pub use platform_management::{
ManageAgentsTool, ManageHarnessesTool, ManageSessionsTool, PlatformManagementCapability,
ReadAgentsTool, ReadCapabilitiesTool, ReadHarnessesTool, ReadSessionsTool,
SessionReadMessagesTool, SessionReadResponseTool, SessionSendMessageTool,
};
pub use prompt_caching::{PROMPT_CACHING_CAPABILITY_ID, PromptCachingCapability};
pub use prompt_canary_guardrail::{
DEFAULT_REPLACEMENT as PROMPT_CANARY_DEFAULT_REPLACEMENT,
PROMPT_CANARY_GUARDRAIL_CAPABILITY_ID, PromptCanaryGuardrailCapability,
REASON_CODE_SYSTEM_PROMPT_LEAK,
};
pub use research::ResearchCapability;
pub use sample_data::SampleDataCapability;
pub use self_budget::{SELF_BUDGET_CAPABILITY_ID, SelfBudgetCapability};
pub use session::{GetSessionInfoTool, SessionCapability, WriteSessionTitleTool};
pub use session_sandbox::{
SESSION_SANDBOX_CAPABILITY_ID, SandboxExecTool, SandboxManageTool, SandboxReadFileTool,
SandboxStatusTool, SandboxWriteFileTool, SessionSandboxCapability,
};
pub use session_schedule::{
CancelScheduleTool, CreateScheduleTool, ListSchedulesTool, SESSION_SCHEDULE_CAPABILITY_ID,
SessionScheduleCapability,
};
pub use session_sql_database::{
SessionSqlDatabaseCapability, SqlExecuteTool, SqlQueryTool, SqlSchemaTool,
};
pub use session_storage::{KvStoreTool, SecretStoreTool, SessionStorageCapability};
pub use skills::{SKILLS_CAPABILITY_ID, SkillsCapability};
pub use stateless_todo_list::{StatelessTodoListCapability, WriteTodosTool};
pub use subagents::SubagentCapability;
pub use system_commands::{SYSTEM_COMMANDS_CAPABILITY_ID, SystemCommandsCapability};
pub use test_math::{AddTool, DivideTool, MultiplyTool, SubtractTool, TestMathCapability};
pub use test_weather::{GetForecastTool, GetWeatherTool, TestWeatherCapability};
pub use tool_output_persistence::{PersistOutputHook, ToolOutputPersistenceCapability};
pub use tool_search::{
TOOL_SEARCH_CAPABILITY_ID, TOOL_SEARCH_TOOL_NAME, ToolSearchCapability, ToolSearchTool,
};
pub use user_hooks::UserHooksCapability;
pub use virtual_bash::{BashTool, SessionFileSystemAdapter, VirtualBashCapability};
pub use web_fetch::{
BotAuthPublicKey, WebFetchCapability, WebFetchTool, derive_bot_auth_public_key,
};
pub use workspace_volumes::{WORKSPACE_VOLUMES_CAPABILITY_ID, WorkspaceVolumesCapability};
pub struct SystemPromptContext {
pub session_id: SessionId,
pub locale: Option<String>,
pub file_store: Option<Arc<dyn SessionFileSystem>>,
pub model: Option<String>,
}
impl SystemPromptContext {
pub fn without_file_store(session_id: SessionId) -> Self {
Self {
session_id,
locale: None,
file_store: None,
model: None,
}
}
pub fn with_model(mut self, model: impl Into<String>) -> Self {
self.model = Some(model.into());
self
}
}
#[async_trait]
pub trait Capability: Send + Sync {
fn id(&self) -> &str;
fn name(&self) -> &str;
fn description(&self) -> &str;
fn status(&self) -> CapabilityStatus {
CapabilityStatus::Available
}
fn icon(&self) -> Option<&str> {
None
}
fn category(&self) -> Option<&str> {
None
}
fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
None
}
fn system_prompt_addition(&self) -> Option<&str> {
None
}
async fn system_prompt_contribution(&self, _ctx: &SystemPromptContext) -> Option<String> {
self.system_prompt_addition().map(|addition| {
format!(
"<capability id=\"{}\">\n{}\n</capability>",
self.id(),
addition
)
})
}
fn system_prompt_preview(&self) -> Option<String> {
self.system_prompt_addition().map(|s| s.to_string())
}
fn tools(&self) -> Vec<Box<dyn Tool>> {
vec![]
}
fn tools_with_config(&self, _config: &serde_json::Value) -> Vec<Box<dyn Tool>> {
self.tools()
}
async fn system_prompt_contribution_with_config(
&self,
ctx: &SystemPromptContext,
_config: &serde_json::Value,
) -> Option<String> {
self.system_prompt_contribution(ctx).await
}
fn tool_definitions(&self) -> Vec<ToolDefinition> {
self.tools().iter().map(|t| t.to_definition()).collect()
}
fn mounts(&self) -> Vec<MountPoint> {
vec![]
}
fn dependencies(&self) -> Vec<&'static str> {
vec![]
}
fn features(&self) -> Vec<&'static str> {
vec![]
}
fn config_schema(&self) -> Option<serde_json::Value> {
None
}
fn config_ui_schema(&self) -> Option<serde_json::Value> {
None
}
fn validate_config(&self, _config: &serde_json::Value) -> Result<(), String> {
Ok(())
}
fn mcp_servers(&self) -> ScopedMcpServers {
ScopedMcpServers::default()
}
fn mcp_servers_with_config(&self, _config: &serde_json::Value) -> ScopedMcpServers {
self.mcp_servers()
}
fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
None
}
fn model_view_provider(&self) -> Option<Arc<dyn ModelViewProvider>> {
None
}
fn pre_tool_use_hooks(&self) -> Vec<Arc<dyn crate::atoms::PreToolUseHook>> {
vec![]
}
fn post_tool_exec_hooks(&self) -> Vec<Arc<dyn crate::atoms::PostToolExecHook>> {
vec![]
}
fn tool_definition_hooks(&self) -> Vec<Arc<dyn ToolDefinitionHook>> {
vec![]
}
fn tool_definition_hooks_with_config(
&self,
_config: &serde_json::Value,
) -> Vec<Arc<dyn ToolDefinitionHook>> {
self.tool_definition_hooks()
}
fn tool_call_hooks(&self) -> Vec<Arc<dyn ToolCallHook>> {
vec![]
}
fn user_hooks(&self) -> Vec<crate::user_hook_types::UserHookSpec> {
vec![]
}
fn user_hooks_with_config(
&self,
_config: &serde_json::Value,
) -> Vec<crate::user_hook_types::UserHookSpec> {
self.user_hooks()
}
fn risk_level(&self) -> RiskLevel {
RiskLevel::Low
}
fn commands(&self) -> Vec<CommandDescriptor> {
vec![]
}
async fn execute_command(
&self,
request: &ExecuteCommandRequest,
_ctx: &CommandExecutionContext,
) -> crate::error::Result<CommandResult> {
Err(crate::error::AgentLoopError::config(format!(
"capability {} declared command /{} but does not implement execute_command",
self.id(),
request.name,
)))
}
fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
vec![]
}
fn contribute_skills(&self) -> Vec<SkillContribution> {
vec![]
}
fn output_guardrails(&self) -> Vec<Arc<dyn crate::output_guardrail::OutputGuardrail>> {
vec![]
}
}
pub trait ToolDefinitionHook: Send + Sync {
fn transform(&self, tools: Vec<ToolDefinition>) -> Vec<ToolDefinition>;
fn applies_with_native_tool_search(&self) -> bool {
true
}
}
pub trait ToolCallHook: Send + Sync {
fn narration(
&self,
_tool_def: Option<&ToolDefinition>,
_tool_call: &ToolCall,
_phase: crate::tool_narration::ToolNarrationPhase,
_locale: Option<&str>,
) -> Option<String> {
None
}
fn transform_for_execution(&self, tool_call: ToolCall) -> ToolCall {
tool_call
}
}
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
)]
#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
#[cfg_attr(feature = "openapi", schema(example = "low"))]
#[serde(rename_all = "lowercase")]
pub enum RiskLevel {
Low,
Medium,
High,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BlueprintModel {
Fixed(String),
Default(String),
Inherit,
}
pub struct AgentBlueprint {
pub id: &'static str,
pub name: &'static str,
pub description: &'static str,
pub model: BlueprintModel,
pub system_prompt: &'static str,
pub tools: Vec<Box<dyn Tool>>,
pub max_turns: Option<usize>,
pub config_schema: Option<serde_json::Value>,
}
impl AgentBlueprint {
pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
self.tools.iter().map(|t| t.to_definition()).collect()
}
}
impl std::fmt::Debug for AgentBlueprint {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AgentBlueprint")
.field("id", &self.id)
.field("name", &self.name)
.field("model", &self.model)
.field("tool_count", &self.tools.len())
.field("max_turns", &self.max_turns)
.finish()
}
}
#[derive(Clone)]
pub struct CapabilityRegistry {
capabilities: HashMap<String, Arc<dyn Capability>>,
}
impl CapabilityRegistry {
pub fn new() -> Self {
Self {
capabilities: HashMap::new(),
}
}
pub fn with_builtins() -> Self {
Self::with_builtins_for_grade(DeploymentGrade::from_env())
}
pub fn with_builtins_for_grade(grade: DeploymentGrade) -> Self {
let mut registry = Self::new();
registry.register(AgentInstructionsCapability);
registry.register(HumanIntentCapability);
registry.register(NoopCapability);
registry.register(CurrentTimeCapability);
registry.register(ResearchCapability);
registry.register(PlatformManagementCapability);
registry.register(FileSystemCapability);
registry.register(WorkspaceVolumesCapability);
registry.register(SessionStorageCapability);
registry.register(SessionCapability);
registry.register(SessionSqlDatabaseCapability);
registry.register(TestMathCapability);
registry.register(TestWeatherCapability);
registry.register(StatelessTodoListCapability);
registry.register(WebFetchCapability::from_env());
registry.register(VirtualBashCapability);
registry.register(BackgroundExecutionCapability);
registry.register(SessionScheduleCapability);
registry.register(BtwCapability);
registry.register(InfinityContextCapability);
registry.register(budgeting::BudgetingCapability);
registry.register(SelfBudgetCapability);
registry.register(CompactionCapability);
registry.register(MemoryCapability);
registry.register(OpenAiToolSearchCapability::new());
registry.register(ToolSearchCapability::new());
registry.register(AutoToolSearchCapability::new());
registry.register(PromptCachingCapability::new());
registry.register(SkillsCapability);
registry.register(SubagentCapability);
if crate::FeatureFlags::from_env(&grade).agent_delegation {
registry.register(AgentHandoffCapability);
registry.register(A2aAgentDelegationCapability);
}
registry.register(SystemCommandsCapability);
registry.register(tool_output_persistence::ToolOutputPersistenceCapability);
registry.register(user_hooks::UserHooksCapability);
registry.register(LoopDetectionCapability);
registry.register(PromptCanaryGuardrailCapability);
#[cfg(feature = "ui-capabilities")]
{
registry.register(OpenUiCapability);
registry.register(A2UiCapability);
}
registry.register(SampleDataCapability);
registry.register(DataKnowledgeCapability);
registry.register(KnowledgeBaseCapability);
registry.register(FakeWarehouseCapability);
registry.register(FakeAwsCapability);
registry.register(FakeCrmCapability);
registry.register(FakeFinancialCapability);
let internal_flags = crate::InternalFeatureFlags::from_env();
if internal_flags.session_sandbox {
registry.register(SessionSandboxCapability);
}
if internal_flags.lua {
registry.register(LuaCapability);
registry.register(LuaCodeModeCapability);
}
for plugin in inventory::iter::<IntegrationPlugin>() {
if (!plugin.experimental_only || grade.experimental_features_enabled())
&& plugin
.feature_flag
.is_none_or(|f| internal_flags.is_enabled(f))
{
registry.register_boxed((plugin.factory)());
}
}
registry
}
pub fn register(&mut self, capability: impl Capability + 'static) {
self.capabilities
.insert(capability.id().to_string(), Arc::new(capability));
}
pub fn register_boxed(&mut self, capability: Box<dyn Capability>) {
self.capabilities
.insert(capability.id().to_string(), Arc::from(capability));
}
pub fn register_arc(&mut self, capability: Arc<dyn Capability>) {
self.capabilities
.insert(capability.id().to_string(), capability);
}
pub fn get(&self, id: &str) -> Option<&Arc<dyn Capability>> {
self.capabilities.get(id)
}
pub fn unregister(&mut self, id: &str) -> Option<Arc<dyn Capability>> {
self.capabilities.remove(id)
}
pub fn has(&self, id: &str) -> bool {
self.capabilities.contains_key(id)
}
pub fn list(&self) -> Vec<&Arc<dyn Capability>> {
self.capabilities.values().collect()
}
pub fn len(&self) -> usize {
self.capabilities.len()
}
pub fn is_empty(&self) -> bool {
self.capabilities.is_empty()
}
pub fn builder() -> CapabilityRegistryBuilder {
CapabilityRegistryBuilder::new()
}
pub fn blueprint(&self, id: &str) -> Option<AgentBlueprint> {
for cap in self.capabilities.values() {
for bp in cap.agent_blueprints() {
if bp.id == id {
return Some(bp);
}
}
}
None
}
pub fn blueprint_with_capability(&self, id: &str) -> Option<(String, AgentBlueprint)> {
for (capability_id, cap) in &self.capabilities {
for bp in cap.agent_blueprints() {
if bp.id == id {
return Some((capability_id.clone(), bp));
}
}
}
None
}
pub fn all_blueprints(&self) -> Vec<AgentBlueprint> {
self.capabilities
.values()
.flat_map(|cap| cap.agent_blueprints())
.collect()
}
}
impl Default for CapabilityRegistry {
fn default() -> Self {
Self::with_builtins()
}
}
impl std::fmt::Debug for CapabilityRegistry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let ids: Vec<_> = self.capabilities.keys().collect();
f.debug_struct("CapabilityRegistry")
.field("capabilities", &ids)
.finish()
}
}
pub struct CapabilityRegistryBuilder {
registry: CapabilityRegistry,
}
impl CapabilityRegistryBuilder {
pub fn new() -> Self {
Self {
registry: CapabilityRegistry::new(),
}
}
pub fn with_builtins() -> Self {
Self {
registry: CapabilityRegistry::with_builtins(),
}
}
pub fn capability(mut self, capability: impl Capability + 'static) -> Self {
self.registry.register(capability);
self
}
pub fn build(self) -> CapabilityRegistry {
self.registry
}
}
impl Default for CapabilityRegistryBuilder {
fn default() -> Self {
Self::new()
}
}
pub struct ModelViewContext<'a> {
pub session_id: SessionId,
pub prior_usage: Option<&'a TokenUsage>,
}
pub trait ModelViewProvider: Send + Sync {
fn apply_model_view(
&self,
messages: Vec<Message>,
config: &serde_json::Value,
context: &ModelViewContext<'_>,
) -> Vec<Message>;
fn priority(&self) -> i32 {
0
}
}
pub struct CollectedCapabilities {
pub system_prompt_parts: Vec<String>,
pub system_prompt_attributions: Vec<SystemPromptAttribution>,
pub tools: Vec<Box<dyn Tool>>,
pub tool_definitions: Vec<ToolDefinition>,
pub mounts: Vec<MountPoint>,
pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
pub applied_ids: Vec<String>,
pub tool_search: Option<crate::llm_driver_registry::ToolSearchConfig>,
pub prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig>,
pub tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>>,
pub tool_call_hooks: Vec<Arc<dyn ToolCallHook>>,
pub mcp_servers: ScopedMcpServers,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SystemPromptAttribution {
pub capability_id: String,
pub content: String,
}
impl CollectedCapabilities {
pub fn system_prompt_prefix(&self) -> Option<String> {
if self.system_prompt_parts.is_empty() {
None
} else {
Some(self.system_prompt_parts.join("\n\n"))
}
}
pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
for (provider, config) in &self.message_filter_providers {
provider.apply_filters(query, config);
}
}
pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
for (provider, config) in &self.message_filter_providers {
provider.post_load(messages, config);
}
}
pub fn has_message_filters(&self) -> bool {
!self.message_filter_providers.is_empty()
}
}
pub struct CollectedMessageFilters {
pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
}
pub struct CollectedModelViewProviders {
pub model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)>,
}
impl CollectedMessageFilters {
pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
for (provider, config) in &self.message_filter_providers {
provider.apply_filters(query, config);
}
}
pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
for (provider, config) in &self.message_filter_providers {
provider.post_load(messages, config);
}
}
}
impl CollectedModelViewProviders {
pub fn apply_model_view(
&self,
mut messages: Vec<Message>,
context: &ModelViewContext<'_>,
) -> Vec<Message> {
for (provider, config) in &self.model_view_providers {
messages = provider.apply_model_view(messages, config, context);
}
messages
}
}
pub fn collect_message_filters_only(
capability_configs: &[AgentCapabilityConfig],
registry: &CapabilityRegistry,
) -> CollectedMessageFilters {
let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
Vec::new();
for cap_config in capability_configs {
let cap_id = cap_config.capability_ref.as_str();
if let Some(capability) = registry.get(cap_id) {
if capability.status() != CapabilityStatus::Available {
continue;
}
let effective: &dyn Capability = capability
.resolve_for_model(None)
.unwrap_or_else(|| capability.as_ref());
if let Some(provider) = effective.message_filter_provider() {
message_filter_providers.push((provider, cap_config.config.clone()));
}
}
}
message_filter_providers.sort_by_key(|(p, _)| p.priority());
CollectedMessageFilters {
message_filter_providers,
}
}
pub fn collect_model_view_providers(
capability_configs: &[AgentCapabilityConfig],
registry: &CapabilityRegistry,
model: Option<&str>,
) -> CollectedModelViewProviders {
let mut model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)> = Vec::new();
for cap_config in capability_configs {
let cap_id = cap_config.capability_ref.as_str();
if let Some(capability) = registry.get(cap_id) {
if capability.status() != CapabilityStatus::Available {
continue;
}
let effective: &dyn Capability = capability
.resolve_for_model(model)
.unwrap_or_else(|| capability.as_ref());
if let Some(provider) = effective.model_view_provider() {
model_view_providers.push((provider, cap_config.config.clone()));
}
}
}
model_view_providers.sort_by_key(|(p, _)| p.priority());
CollectedModelViewProviders {
model_view_providers,
}
}
pub fn collect_capability_mcp_servers(
capability_configs: &[AgentCapabilityConfig],
registry: &CapabilityRegistry,
) -> ScopedMcpServers {
let mut servers = ScopedMcpServers::default();
for cap_config in capability_configs {
let cap_id = cap_config.capability_ref.as_str();
if is_declarative_capability(cap_id) {
if let Ok(definition) =
serde_json::from_value::<DeclarativeCapabilityDefinition>(cap_config.config.clone())
{
if definition.status != CapabilityStatus::Available {
continue;
}
if let Some(contributed) = definition.mcp_servers {
servers = merge_scoped_mcp_servers(&servers, &contributed);
}
}
continue;
}
if let Some(capability) = registry.get(cap_id) {
if capability.status() != CapabilityStatus::Available {
continue;
}
servers = merge_scoped_mcp_servers(
&servers,
&capability.mcp_servers_with_config(&cap_config.config),
);
}
}
servers
}
pub const MAX_RESOLVED_CAPABILITIES: usize = 100;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DependencyError {
CircularDependency {
capability_id: String,
chain: Vec<String>,
},
TooManyCapabilities {
count: usize,
max: usize,
},
}
impl std::fmt::Display for DependencyError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
DependencyError::CircularDependency {
capability_id,
chain,
} => {
write!(
f,
"Circular dependency detected: {} depends on itself via chain: {} -> {}",
capability_id,
chain.join(" -> "),
capability_id
)
}
DependencyError::TooManyCapabilities { count, max } => {
write!(
f,
"Too many capabilities after resolution: {} (max: {})",
count, max
)
}
}
}
}
impl std::error::Error for DependencyError {}
#[derive(Debug, Clone)]
pub struct ResolvedCapabilities {
pub resolved_ids: Vec<String>,
pub added_as_dependencies: Vec<String>,
pub user_selected: Vec<String>,
}
pub fn resolve_dependencies(
selected_ids: &[String],
registry: &CapabilityRegistry,
) -> Result<ResolvedCapabilities, DependencyError> {
use std::collections::HashSet;
let user_selected: HashSet<String> = selected_ids.iter().cloned().collect();
let mut resolved: Vec<String> = Vec::new();
let mut resolved_set: HashSet<String> = HashSet::new();
let mut added_as_dependencies: Vec<String> = Vec::new();
for cap_id in selected_ids {
resolve_single_capability(
cap_id,
registry,
&mut resolved,
&mut resolved_set,
&mut added_as_dependencies,
&user_selected,
&mut Vec::new(), )?;
}
if resolved.len() > MAX_RESOLVED_CAPABILITIES {
return Err(DependencyError::TooManyCapabilities {
count: resolved.len(),
max: MAX_RESOLVED_CAPABILITIES,
});
}
Ok(ResolvedCapabilities {
resolved_ids: resolved,
added_as_dependencies,
user_selected: selected_ids.to_vec(),
})
}
pub fn resolve_capability_configs(
selected_configs: &[AgentCapabilityConfig],
registry: &CapabilityRegistry,
) -> Result<Vec<AgentCapabilityConfig>, DependencyError> {
let mut selected_ids: Vec<String> = Vec::new();
for config in selected_configs {
if is_declarative_capability(config.capability_id())
&& let Ok(definition) =
serde_json::from_value::<DeclarativeCapabilityDefinition>(config.config.clone())
{
selected_ids.extend(definition.dependencies);
}
selected_ids.push(config.capability_id().to_string());
}
let resolved = resolve_dependencies(&selected_ids, registry)?;
let explicit_configs: std::collections::HashMap<String, serde_json::Value> = selected_configs
.iter()
.map(|config| (config.capability_id().to_string(), config.config.clone()))
.collect();
Ok(resolved
.resolved_ids
.into_iter()
.map(|capability_id| {
explicit_configs
.get(&capability_id)
.cloned()
.map(|config| AgentCapabilityConfig::with_config(capability_id.clone(), config))
.unwrap_or_else(|| AgentCapabilityConfig::new(capability_id))
})
.collect())
}
fn resolve_single_capability(
cap_id: &str,
registry: &CapabilityRegistry,
resolved: &mut Vec<String>,
resolved_set: &mut std::collections::HashSet<String>,
added_as_dependencies: &mut Vec<String>,
user_selected: &std::collections::HashSet<String>,
visiting: &mut Vec<String>,
) -> Result<(), DependencyError> {
if resolved_set.contains(cap_id) {
return Ok(());
}
if visiting.contains(&cap_id.to_string()) {
return Err(DependencyError::CircularDependency {
capability_id: cap_id.to_string(),
chain: visiting.clone(),
});
}
let capability = match registry.get(cap_id) {
Some(cap) => cap,
None => {
if is_declarative_capability(cap_id) && !resolved_set.contains(cap_id) {
resolved.push(cap_id.to_string());
resolved_set.insert(cap_id.to_string());
if !user_selected.contains(cap_id) {
added_as_dependencies.push(cap_id.to_string());
}
}
return Ok(());
}
};
visiting.push(cap_id.to_string());
for dep_id in capability.dependencies() {
resolve_single_capability(
dep_id,
registry,
resolved,
resolved_set,
added_as_dependencies,
user_selected,
visiting,
)?;
}
visiting.pop();
if !resolved_set.contains(cap_id) {
resolved.push(cap_id.to_string());
resolved_set.insert(cap_id.to_string());
if !user_selected.contains(cap_id) {
added_as_dependencies.push(cap_id.to_string());
}
}
Ok(())
}
pub fn compute_features(capability_ids: &[String], registry: &CapabilityRegistry) -> Vec<String> {
use std::collections::HashSet;
let resolved_ids = match resolve_dependencies(capability_ids, registry) {
Ok(resolved) => resolved.resolved_ids,
Err(_) => capability_ids.to_vec(),
};
let mut seen = HashSet::new();
let mut features = Vec::new();
for cap_id in &resolved_ids {
if let Some(cap) = registry.get(cap_id) {
for feature in cap.features() {
if seen.insert(feature) {
features.push(feature.to_string());
}
}
}
}
features
}
pub fn get_dependencies(cap_id: &str, registry: &CapabilityRegistry) -> Vec<String> {
registry
.get(cap_id)
.map(|cap| cap.dependencies().iter().map(|s| s.to_string()).collect())
.unwrap_or_default()
}
pub async fn collect_capabilities(
capability_ids: &[String],
registry: &CapabilityRegistry,
ctx: &SystemPromptContext,
) -> CollectedCapabilities {
let resolved_ids = match resolve_dependencies(capability_ids, registry) {
Ok(resolved) => resolved.resolved_ids,
Err(e) => {
tracing::warn!("Failed to resolve capability dependencies: {}", e);
capability_ids.to_vec()
}
};
let configs: Vec<AgentCapabilityConfig> = resolved_ids
.iter()
.map(|id| AgentCapabilityConfig {
capability_ref: CapabilityId::new(id),
config: serde_json::Value::Object(serde_json::Map::new()),
})
.collect();
collect_capabilities_with_configs(&configs, registry, ctx).await
}
pub async fn collect_capabilities_with_configs(
capability_configs: &[AgentCapabilityConfig],
registry: &CapabilityRegistry,
ctx: &SystemPromptContext,
) -> CollectedCapabilities {
let mut system_prompt_parts: Vec<String> = Vec::new();
let mut system_prompt_attributions: Vec<SystemPromptAttribution> = Vec::new();
let mut tools: Vec<Box<dyn Tool>> = Vec::new();
let mut tool_definitions: Vec<ToolDefinition> = Vec::new();
let mut mounts: Vec<MountPoint> = Vec::new();
let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
Vec::new();
let mut applied_ids: Vec<String> = Vec::new();
let mut tool_search: Option<crate::llm_driver_registry::ToolSearchConfig> = None;
let mut prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig> = None;
let mut tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>> = Vec::new();
let mut tool_call_hooks: Vec<Arc<dyn ToolCallHook>> = Vec::new();
let mut mcp_servers = ScopedMcpServers::default();
for cap_config in capability_configs {
let cap_id = cap_config.capability_ref.as_str();
if is_declarative_capability(cap_id) {
match serde_json::from_value::<DeclarativeCapabilityDefinition>(
cap_config.config.clone(),
) {
Ok(definition) => {
if definition.status != CapabilityStatus::Available {
continue;
}
if let Some(prompt) = definition.system_prompt.as_deref() {
let contribution =
format!("<capability id=\"{}\">\n{}\n</capability>", cap_id, prompt);
system_prompt_attributions.push(SystemPromptAttribution {
capability_id: cap_id.to_string(),
content: contribution.clone(),
});
system_prompt_parts.push(contribution);
}
mounts.extend(definition.mounts(cap_id));
if let Some(ref servers) = definition.mcp_servers {
mcp_servers = merge_scoped_mcp_servers(&mcp_servers, servers);
}
for skill in definition.skill_contributions() {
mounts.push(skill.to_mount(cap_id));
}
applied_ids.push(cap_id.to_string());
}
Err(error) => {
tracing::warn!(
capability_id = %cap_id,
error = %error,
"Skipping invalid declarative capability config"
);
}
}
continue;
}
if let Some(capability) = registry.get(cap_id) {
if capability.status() != CapabilityStatus::Available {
continue;
}
let effective: &dyn Capability =
match capability.resolve_for_model(ctx.model.as_deref()) {
Some(inner) => inner,
None => capability.as_ref(),
};
let effective_id = effective.id();
if let Some(contribution) = effective
.system_prompt_contribution_with_config(ctx, &cap_config.config)
.await
{
system_prompt_attributions.push(SystemPromptAttribution {
capability_id: cap_id.to_string(),
content: contribution.clone(),
});
system_prompt_parts.push(contribution);
}
tools.extend(effective.tools_with_config(&cap_config.config));
tool_definition_hooks
.extend(effective.tool_definition_hooks_with_config(&cap_config.config));
tool_call_hooks.extend(effective.tool_call_hooks());
let cap_category = effective.category();
for def in effective.tool_definitions() {
let def = match (def.category(), cap_category) {
(None, Some(cat)) => def.with_category(cat),
_ => def,
}
.with_capability_attribution(cap_id, Some(capability.name()));
tool_definitions.push(def);
}
if effective_id == OPENAI_TOOL_SEARCH_CAPABILITY_ID {
let threshold = cap_config
.config
.get("threshold")
.and_then(|v| v.as_u64())
.map(|v| v as usize)
.unwrap_or(DEFAULT_TOOL_SEARCH_THRESHOLD);
tool_search = Some(crate::llm_driver_registry::ToolSearchConfig {
enabled: true,
threshold,
});
}
if cap_id == PROMPT_CACHING_CAPABILITY_ID {
let strategy = cap_config
.config
.get("strategy")
.and_then(|v| v.as_str())
.map(|value| match value {
"auto" => crate::llm_driver_registry::PromptCacheStrategy::Auto,
_ => crate::llm_driver_registry::PromptCacheStrategy::Auto,
})
.unwrap_or(crate::llm_driver_registry::PromptCacheStrategy::Auto);
let gemini_cached_content = cap_config
.config
.get("gemini_cached_content")
.and_then(|v| v.as_str())
.map(str::to_string);
prompt_cache = Some(crate::llm_driver_registry::PromptCacheConfig {
enabled: true,
strategy,
gemini_cached_content,
});
}
mounts.extend(effective.mounts());
mcp_servers = merge_scoped_mcp_servers(
&mcp_servers,
&effective.mcp_servers_with_config(&cap_config.config),
);
for skill in effective.contribute_skills() {
mounts.push(skill.to_mount(cap_id));
}
if let Some(provider) = effective.message_filter_provider() {
message_filter_providers.push((provider, cap_config.config.clone()));
}
applied_ids.push(cap_id.to_string());
}
}
if !applied_ids
.iter()
.any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID)
&& tool_definitions
.iter()
.any(|def| def.hints().supports_background == Some(true))
&& let Some(bg_cap) = registry.get(BACKGROUND_EXECUTION_CAPABILITY_ID)
&& bg_cap.status() == CapabilityStatus::Available
{
tools.extend(bg_cap.tools());
let cap_category = bg_cap.category();
for def in bg_cap.tool_definitions() {
let def = match (def.category(), cap_category) {
(None, Some(cat)) => def.with_category(cat),
_ => def,
}
.with_capability_attribution(BACKGROUND_EXECUTION_CAPABILITY_ID, Some(bg_cap.name()));
tool_definitions.push(def);
}
applied_ids.push(BACKGROUND_EXECUTION_CAPABILITY_ID.to_string());
}
message_filter_providers.sort_by_key(|(p, _)| p.priority());
CollectedCapabilities {
system_prompt_parts,
system_prompt_attributions,
tools,
tool_definitions,
mounts,
message_filter_providers,
applied_ids,
tool_search,
prompt_cache,
tool_definition_hooks,
tool_call_hooks,
mcp_servers,
}
}
pub struct AppliedCapabilities {
pub runtime_agent: RuntimeAgent,
pub tool_registry: ToolRegistry,
pub applied_ids: Vec<String>,
}
pub async fn apply_capabilities(
base_runtime_agent: RuntimeAgent,
capability_ids: &[String],
registry: &CapabilityRegistry,
ctx: &SystemPromptContext,
) -> AppliedCapabilities {
let collected = collect_capabilities(capability_ids, registry, ctx).await;
let final_system_prompt = match collected.system_prompt_prefix() {
Some(prefix) => format!(
"{}\n\n<system-prompt>\n{}\n</system-prompt>",
prefix, base_runtime_agent.system_prompt
),
None => base_runtime_agent.system_prompt,
};
let mut tool_registry = ToolRegistry::new();
for tool in collected.tools {
tool_registry.register_boxed(tool);
}
let mut tools = collected.tool_definitions;
for hook in &collected.tool_definition_hooks {
tools = hook.transform(tools);
}
let runtime_agent = RuntimeAgent {
system_prompt: final_system_prompt,
model: base_runtime_agent.model,
tools,
max_iterations: base_runtime_agent.max_iterations,
temperature: base_runtime_agent.temperature,
max_tokens: base_runtime_agent.max_tokens,
tool_search: collected.tool_search,
prompt_cache: collected.prompt_cache,
network_access: base_runtime_agent.network_access,
};
AppliedCapabilities {
runtime_agent,
tool_registry,
applied_ids: collected.applied_ids,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::typed_id::SessionId;
use std::collections::BTreeSet;
use uuid::Uuid;
static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
fn lock_env() -> std::sync::MutexGuard<'static, ()> {
ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner())
}
fn test_ctx() -> SystemPromptContext {
SystemPromptContext::without_file_store(SessionId::new())
}
fn expected_core_builtin_ids() -> BTreeSet<&'static str> {
let mut ids = [
"agent_instructions",
"human_intent",
"budgeting",
"self_budget",
"noop",
"current_time",
"research",
"platform_management",
"session_file_system",
"workspace_volumes",
"session_storage",
"session",
"session_sql_database",
"test_math",
"test_weather",
"stateless_todo_list",
"web_fetch",
"virtual_bash",
"background_execution",
"session_schedule",
"btw",
"infinity_context",
"compaction",
"memory",
"openai_tool_search",
"tool_search",
"auto_tool_search",
"prompt_caching",
"skills",
"subagents",
"system_commands",
"sample_data",
"data_knowledge",
"knowledge_base",
"tool_output_persistence",
"fake_warehouse",
"fake_aws",
"fake_crm",
"fake_financial",
"loop_detection",
"prompt_canary_guardrail",
"user_hooks",
]
.into_iter()
.collect::<BTreeSet<_>>();
if cfg!(feature = "ui-capabilities") {
ids.insert("openui");
ids.insert("a2ui");
}
ids
}
fn expected_dev_builtin_ids() -> BTreeSet<&'static str> {
let mut ids = expected_core_builtin_ids();
ids.insert("agent_handoff");
ids.insert("a2a_agent_delegation");
ids
}
fn registry_ids(registry: &CapabilityRegistry) -> BTreeSet<&str> {
registry.capabilities.keys().map(String::as_str).collect()
}
#[test]
fn test_capability_registry_with_builtins_dev() {
let _lock = lock_env();
unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
assert_eq!(registry_ids(®istry), expected_dev_builtin_ids());
assert!(registry.has("agent_handoff"));
assert!(registry.has("a2a_agent_delegation"));
}
#[test]
fn test_capability_registry_with_builtins_prod() {
let _lock = lock_env();
unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
assert_eq!(registry_ids(®istry), expected_core_builtin_ids());
assert!(!registry.has("docker_container"));
assert!(!registry.has("agent_handoff"));
assert!(!registry.has("a2a_agent_delegation"));
}
#[test]
fn test_agent_delegation_enabled_by_env_in_prod() {
let _lock = lock_env();
unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "true") };
let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
assert!(registry.has("agent_handoff"));
assert!(registry.has("a2a_agent_delegation"));
unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
}
#[test]
fn test_agent_delegation_disabled_by_env_in_dev() {
let _lock = lock_env();
unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "false") };
let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
assert!(!registry.has("agent_handoff"));
assert!(!registry.has("a2a_agent_delegation"));
unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
}
#[test]
fn test_capability_registry_get() {
let registry = CapabilityRegistry::with_builtins();
let noop = registry.get("noop").unwrap();
assert_eq!(noop.id(), "noop");
assert_eq!(noop.name(), "No-Op");
assert_eq!(noop.status(), CapabilityStatus::Available);
}
#[test]
fn test_capability_registry_blueprint_with_capability() {
struct BlueprintProviderCapability;
impl Capability for BlueprintProviderCapability {
fn id(&self) -> &str {
"blueprint_provider"
}
fn name(&self) -> &str {
"Blueprint Provider"
}
fn description(&self) -> &str {
"Capability that provides a blueprint for tests"
}
fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
vec![AgentBlueprint {
id: "test_blueprint",
name: "Test Blueprint",
description: "Blueprint for capability registry tests",
model: BlueprintModel::Inherit,
system_prompt: "Test prompt",
tools: vec![],
max_turns: None,
config_schema: None,
}]
}
}
let mut registry = CapabilityRegistry::new();
registry.register(BlueprintProviderCapability);
let (capability_id, blueprint) = registry
.blueprint_with_capability("test_blueprint")
.expect("blueprint should resolve with capability id");
assert_eq!(capability_id, "blueprint_provider");
assert_eq!(blueprint.id, "test_blueprint");
}
#[test]
fn test_capability_registry_builder() {
let registry = CapabilityRegistry::builder()
.capability(NoopCapability)
.capability(CurrentTimeCapability)
.build();
assert!(registry.has("noop"));
assert!(registry.has("current_time"));
assert_eq!(registry.len(), 2);
}
#[test]
fn test_capability_status() {
let registry = CapabilityRegistry::with_builtins();
let current_time = registry.get("current_time").unwrap();
assert_eq!(current_time.status(), CapabilityStatus::Available);
let research = registry.get("research").unwrap();
assert_eq!(research.status(), CapabilityStatus::ComingSoon);
}
#[test]
fn test_capability_icons_and_categories() {
let registry = CapabilityRegistry::with_builtins();
let noop = registry.get("noop").unwrap();
assert_eq!(noop.icon(), Some("circle-off"));
assert_eq!(noop.category(), Some("Testing"));
let current_time = registry.get("current_time").unwrap();
assert_eq!(current_time.icon(), Some("clock"));
assert_eq!(current_time.category(), Some("Utilities"));
}
#[test]
fn test_system_prompt_preview_default_delegates_to_addition() {
let registry = CapabilityRegistry::with_builtins();
let test_math = registry.get("test_math").unwrap();
assert_eq!(
test_math.system_prompt_preview().as_deref(),
test_math.system_prompt_addition()
);
let current_time = registry.get("current_time").unwrap();
assert!(current_time.system_prompt_preview().is_none());
assert!(current_time.system_prompt_addition().is_none());
}
#[test]
fn test_system_prompt_preview_dynamic_capability() {
let registry = CapabilityRegistry::with_builtins();
let cap = registry.get("agent_instructions").unwrap();
assert!(cap.system_prompt_addition().is_none());
assert!(cap.system_prompt_preview().is_some());
assert!(cap.system_prompt_preview().unwrap().contains("AGENTS.md"));
}
#[tokio::test]
async fn test_apply_capabilities_empty() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
let applied =
apply_capabilities(base_runtime_agent.clone(), &[], ®istry, &test_ctx()).await;
assert_eq!(
applied.runtime_agent.system_prompt,
base_runtime_agent.system_prompt
);
assert!(applied.tool_registry.is_empty());
assert!(applied.applied_ids.is_empty());
}
#[tokio::test]
async fn test_apply_capabilities_noop() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
let applied = apply_capabilities(
base_runtime_agent.clone(),
&["noop".to_string()],
®istry,
&test_ctx(),
)
.await;
assert_eq!(
applied.runtime_agent.system_prompt,
base_runtime_agent.system_prompt
);
assert!(applied.tool_registry.is_empty());
assert_eq!(applied.applied_ids, vec!["noop"]);
}
#[tokio::test]
async fn test_apply_capabilities_current_time() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
let applied = apply_capabilities(
base_runtime_agent.clone(),
&["current_time".to_string()],
®istry,
&test_ctx(),
)
.await;
assert_eq!(
applied.runtime_agent.system_prompt,
base_runtime_agent.system_prompt
);
assert!(applied.tool_registry.has("get_current_time"));
assert_eq!(applied.tool_registry.len(), 1);
assert_eq!(applied.applied_ids, vec!["current_time"]);
}
#[tokio::test]
async fn test_apply_capabilities_skips_coming_soon() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
let applied = apply_capabilities(
base_runtime_agent.clone(),
&["research".to_string()],
®istry,
&test_ctx(),
)
.await;
assert_eq!(
applied.runtime_agent.system_prompt,
base_runtime_agent.system_prompt
);
assert!(applied.applied_ids.is_empty()); }
#[tokio::test]
async fn test_apply_capabilities_multiple() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
let applied = apply_capabilities(
base_runtime_agent.clone(),
&["noop".to_string(), "current_time".to_string()],
®istry,
&test_ctx(),
)
.await;
assert!(applied.tool_registry.has("get_current_time"));
assert_eq!(applied.applied_ids, vec!["noop", "current_time"]);
}
#[tokio::test]
async fn test_apply_capabilities_preserves_order() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("Base prompt.", "gpt-5.2");
let applied = apply_capabilities(
base_runtime_agent,
&["current_time".to_string(), "noop".to_string()],
®istry,
&test_ctx(),
)
.await;
assert_eq!(applied.applied_ids, vec!["current_time", "noop"]);
}
#[tokio::test]
async fn test_apply_capabilities_test_math() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
let applied = apply_capabilities(
base_runtime_agent.clone(),
&["test_math".to_string()],
®istry,
&test_ctx(),
)
.await;
assert!(
!applied
.runtime_agent
.system_prompt
.contains("<capability id=\"test_math\">")
);
assert!(
applied
.runtime_agent
.system_prompt
.contains("You are a helpful assistant.")
);
assert!(applied.tool_registry.has("add"));
assert!(applied.tool_registry.has("subtract"));
assert!(applied.tool_registry.has("multiply"));
assert!(applied.tool_registry.has("divide"));
assert_eq!(applied.tool_registry.len(), 4);
}
#[tokio::test]
async fn test_apply_capabilities_test_weather() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
let applied = apply_capabilities(
base_runtime_agent.clone(),
&["test_weather".to_string()],
®istry,
&test_ctx(),
)
.await;
assert!(
!applied
.runtime_agent
.system_prompt
.contains("<capability id=\"test_weather\">")
);
assert!(applied.tool_registry.has("get_weather"));
assert!(applied.tool_registry.has("get_forecast"));
assert_eq!(applied.tool_registry.len(), 2);
}
#[tokio::test]
async fn test_apply_capabilities_test_math_and_test_weather() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
let applied = apply_capabilities(
base_runtime_agent.clone(),
&["test_math".to_string(), "test_weather".to_string()],
®istry,
&test_ctx(),
)
.await;
assert_eq!(applied.tool_registry.len(), 6); assert!(applied.tool_registry.has("add"));
assert!(applied.tool_registry.has("get_weather"));
}
#[tokio::test]
async fn test_apply_capabilities_stateless_todo_list() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
let applied = apply_capabilities(
base_runtime_agent.clone(),
&["stateless_todo_list".to_string()],
®istry,
&test_ctx(),
)
.await;
assert!(
applied
.runtime_agent
.system_prompt
.contains("Task Management")
);
assert!(applied.runtime_agent.system_prompt.contains("write_todos"));
assert!(applied.tool_registry.has("write_todos"));
assert_eq!(applied.tool_registry.len(), 1);
}
#[tokio::test]
async fn test_apply_capabilities_web_fetch() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
let applied = apply_capabilities(
base_runtime_agent.clone(),
&["web_fetch".to_string()],
®istry,
&test_ctx(),
)
.await;
assert!(
applied
.runtime_agent
.system_prompt
.contains(&base_runtime_agent.system_prompt)
);
assert!(applied.runtime_agent.system_prompt.contains("web_fetch"));
assert!(applied.tool_registry.has("web_fetch"));
assert_eq!(applied.tool_registry.len(), 1);
}
#[tokio::test]
async fn test_xml_tags_wrap_capability_prompts() {
let registry = CapabilityRegistry::with_builtins();
let collected =
collect_capabilities(&["stateless_todo_list".to_string()], ®istry, &test_ctx())
.await;
assert_eq!(collected.system_prompt_parts.len(), 1);
let part = &collected.system_prompt_parts[0];
assert!(part.starts_with("<capability id=\"stateless_todo_list\">"));
assert!(part.ends_with("</capability>"));
assert!(part.contains("Task Management"));
}
#[tokio::test]
async fn test_xml_tags_multiple_capabilities() {
let registry = CapabilityRegistry::with_builtins();
let collected = collect_capabilities(
&[
"stateless_todo_list".to_string(),
"session_schedule".to_string(),
],
®istry,
&test_ctx(),
)
.await;
assert_eq!(collected.system_prompt_parts.len(), 2);
assert!(
collected.system_prompt_parts[0].starts_with("<capability id=\"stateless_todo_list\">")
);
assert!(
collected.system_prompt_parts[1].starts_with("<capability id=\"session_schedule\">")
);
let prefix = collected.system_prompt_prefix().unwrap();
assert!(prefix.contains("</capability>\n\n<capability"));
}
#[tokio::test]
async fn test_xml_tags_system_prompt_wrapping() {
let registry = CapabilityRegistry::with_builtins();
let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
let applied = apply_capabilities(
base,
&["stateless_todo_list".to_string()],
®istry,
&test_ctx(),
)
.await;
let prompt = &applied.runtime_agent.system_prompt;
assert!(prompt.contains("<capability id=\"stateless_todo_list\">"));
assert!(prompt.contains("</capability>"));
assert!(prompt.contains("<system-prompt>\nYou are helpful.\n</system-prompt>"));
}
#[tokio::test]
async fn test_no_xml_wrapping_without_capabilities() {
let registry = CapabilityRegistry::with_builtins();
let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
let applied = apply_capabilities(base, &[], ®istry, &test_ctx()).await;
assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
assert!(
!applied
.runtime_agent
.system_prompt
.contains("<system-prompt>")
);
}
#[tokio::test]
async fn test_no_xml_wrapping_for_noop_capability() {
let registry = CapabilityRegistry::with_builtins();
let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
let applied = apply_capabilities(base, &["noop".to_string()], ®istry, &test_ctx()).await;
assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
assert!(
!applied
.runtime_agent
.system_prompt
.contains("<system-prompt>")
);
}
#[tokio::test]
async fn test_collect_capabilities_includes_mounts() {
let registry = CapabilityRegistry::with_builtins();
let collected =
collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
assert!(!collected.mounts.is_empty());
assert_eq!(collected.mounts.len(), 1);
assert_eq!(collected.mounts[0].path, "/samples");
assert!(collected.mounts[0].is_readonly());
}
#[tokio::test]
async fn test_collect_capabilities_empty_mounts_by_default() {
let registry = CapabilityRegistry::with_builtins();
let collected =
collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
assert!(collected.mounts.is_empty());
}
#[tokio::test]
async fn test_collect_capabilities_combines_mounts() {
let registry = CapabilityRegistry::with_builtins();
let collected = collect_capabilities(
&["sample_data".to_string(), "current_time".to_string()],
®istry,
&test_ctx(),
)
.await;
assert_eq!(collected.mounts.len(), 1);
assert!(
collected
.applied_ids
.iter()
.any(|id| id == "session_file_system")
);
assert!(collected.applied_ids.iter().any(|id| id == "sample_data"));
assert!(collected.applied_ids.iter().any(|id| id == "current_time"));
}
#[test]
fn test_sample_data_capability() {
let registry = CapabilityRegistry::with_builtins();
let cap = registry.get("sample_data").unwrap();
assert_eq!(cap.id(), "sample_data");
assert_eq!(cap.name(), "Sample Data");
assert_eq!(cap.status(), CapabilityStatus::Available);
assert!(cap.system_prompt_addition().is_some());
assert!(cap.tools().is_empty());
assert!(!cap.mounts().is_empty());
}
#[test]
fn test_resolve_dependencies_empty() {
let registry = CapabilityRegistry::with_builtins();
let resolved = resolve_dependencies(&[], ®istry).unwrap();
assert!(resolved.resolved_ids.is_empty());
assert!(resolved.added_as_dependencies.is_empty());
assert!(resolved.user_selected.is_empty());
}
#[test]
fn test_resolve_dependencies_no_deps() {
let registry = CapabilityRegistry::with_builtins();
let resolved = resolve_dependencies(&["current_time".to_string()], ®istry).unwrap();
assert_eq!(resolved.resolved_ids, vec!["current_time"]);
assert!(resolved.added_as_dependencies.is_empty());
}
#[test]
fn test_resolve_dependencies_with_deps() {
let registry = CapabilityRegistry::with_builtins();
let resolved = resolve_dependencies(&["sample_data".to_string()], ®istry).unwrap();
assert_eq!(resolved.resolved_ids.len(), 2);
let fs_pos = resolved
.resolved_ids
.iter()
.position(|id| id == "session_file_system")
.unwrap();
let sd_pos = resolved
.resolved_ids
.iter()
.position(|id| id == "sample_data")
.unwrap();
assert!(fs_pos < sd_pos, "FileSystem should come before SampleData");
assert_eq!(resolved.added_as_dependencies, vec!["session_file_system"]);
}
#[test]
fn test_resolve_dependencies_already_selected() {
let registry = CapabilityRegistry::with_builtins();
let resolved = resolve_dependencies(
&["session_file_system".to_string(), "sample_data".to_string()],
®istry,
)
.unwrap();
assert_eq!(resolved.resolved_ids.len(), 2);
assert!(resolved.added_as_dependencies.is_empty());
}
#[test]
fn test_resolve_dependencies_preserves_order() {
let registry = CapabilityRegistry::with_builtins();
let resolved =
resolve_dependencies(&["current_time".to_string(), "noop".to_string()], ®istry)
.unwrap();
assert_eq!(resolved.resolved_ids, vec!["current_time", "noop"]);
}
#[test]
fn test_resolve_dependencies_unknown_capability() {
let registry = CapabilityRegistry::with_builtins();
let resolved =
resolve_dependencies(&["unknown_capability".to_string()], ®istry).unwrap();
assert!(resolved.resolved_ids.is_empty());
}
#[test]
fn test_get_dependencies() {
let registry = CapabilityRegistry::with_builtins();
let deps = get_dependencies("sample_data", ®istry);
assert_eq!(deps, vec!["session_file_system"]);
let deps = get_dependencies("current_time", ®istry);
assert!(deps.is_empty());
let deps = get_dependencies("unknown", ®istry);
assert!(deps.is_empty());
}
#[test]
fn test_sample_data_has_dependency() {
let registry = CapabilityRegistry::with_builtins();
let cap = registry.get("sample_data").unwrap();
let deps = cap.dependencies();
assert_eq!(deps.len(), 1);
assert_eq!(deps[0], "session_file_system");
}
#[test]
fn test_noop_has_no_dependencies() {
let registry = CapabilityRegistry::with_builtins();
let cap = registry.get("noop").unwrap();
assert!(cap.dependencies().is_empty());
}
#[test]
fn test_circular_dependency_error() {
struct CapA;
struct CapB;
impl Capability for CapA {
fn id(&self) -> &str {
"test_cap_a"
}
fn name(&self) -> &str {
"Test A"
}
fn description(&self) -> &str {
"Test capability A"
}
fn dependencies(&self) -> Vec<&'static str> {
vec!["test_cap_b"]
}
}
impl Capability for CapB {
fn id(&self) -> &str {
"test_cap_b"
}
fn name(&self) -> &str {
"Test B"
}
fn description(&self) -> &str {
"Test capability B"
}
fn dependencies(&self) -> Vec<&'static str> {
vec!["test_cap_a"]
}
}
let mut registry = CapabilityRegistry::new();
registry.register(CapA);
registry.register(CapB);
let result = resolve_dependencies(&["test_cap_a".to_string()], ®istry);
assert!(result.is_err());
match result.unwrap_err() {
DependencyError::CircularDependency { capability_id, .. } => {
assert_eq!(capability_id, "test_cap_a");
}
_ => panic!("Expected CircularDependency error"),
}
}
use crate::message_filter::{MessageFilter, MessageFilterProvider, MessageQuery};
struct FilterTestCapability {
priority: i32,
}
impl Capability for FilterTestCapability {
fn id(&self) -> &str {
"filter_test"
}
fn name(&self) -> &str {
"Filter Test"
}
fn description(&self) -> &str {
"Test capability with message filter"
}
fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
Some(Arc::new(FilterTestProvider {
priority: self.priority,
}))
}
}
struct FilterTestProvider {
priority: i32,
}
impl MessageFilterProvider for FilterTestProvider {
fn apply_filters(&self, query: &mut MessageQuery, config: &serde_json::Value) {
if let Some(search) = config.get("search").and_then(|v| v.as_str()) {
query
.filters
.push(MessageFilter::Search(search.to_string()));
}
}
fn priority(&self) -> i32 {
self.priority
}
}
#[tokio::test]
async fn test_collect_capabilities_with_configs_no_filter_providers() {
let registry = CapabilityRegistry::with_builtins();
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("current_time"),
config: serde_json::json!({}),
}];
let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
assert!(collected.message_filter_providers.is_empty());
assert!(!collected.has_message_filters());
}
#[tokio::test]
async fn test_collect_capabilities_with_configs_with_filter_provider() {
let mut registry = CapabilityRegistry::new();
registry.register(FilterTestCapability { priority: 0 });
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("filter_test"),
config: serde_json::json!({ "search": "hello" }),
}];
let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
assert_eq!(collected.message_filter_providers.len(), 1);
assert!(collected.has_message_filters());
}
#[tokio::test]
async fn test_collect_capabilities_with_configs_filter_priority_order() {
struct HighPriorityCapability;
struct LowPriorityCapability;
impl Capability for HighPriorityCapability {
fn id(&self) -> &str {
"high_priority"
}
fn name(&self) -> &str {
"High Priority"
}
fn description(&self) -> &str {
"Test"
}
fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
Some(Arc::new(FilterTestProvider { priority: 10 }))
}
}
impl Capability for LowPriorityCapability {
fn id(&self) -> &str {
"low_priority"
}
fn name(&self) -> &str {
"Low Priority"
}
fn description(&self) -> &str {
"Test"
}
fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
Some(Arc::new(FilterTestProvider { priority: -5 }))
}
}
let mut registry = CapabilityRegistry::new();
registry.register(HighPriorityCapability);
registry.register(LowPriorityCapability);
let configs = vec![
AgentCapabilityConfig {
capability_ref: CapabilityId::new("high_priority"),
config: serde_json::json!({}),
},
AgentCapabilityConfig {
capability_ref: CapabilityId::new("low_priority"),
config: serde_json::json!({}),
},
];
let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
assert_eq!(collected.message_filter_providers.len(), 2);
assert_eq!(collected.message_filter_providers[0].0.priority(), -5);
assert_eq!(collected.message_filter_providers[1].0.priority(), 10);
}
#[tokio::test]
async fn test_collected_capabilities_apply_message_filters() {
let mut registry = CapabilityRegistry::new();
registry.register(FilterTestCapability { priority: 0 });
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("filter_test"),
config: serde_json::json!({ "search": "test_query" }),
}];
let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
let session_id: SessionId = Uuid::now_v7().into();
let mut query = MessageQuery::new(session_id);
collected.apply_message_filters(&mut query);
assert_eq!(query.filters.len(), 1);
assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
}
#[tokio::test]
async fn test_collected_capabilities_apply_multiple_filters_in_priority_order() {
struct SearchCapability {
id: &'static str,
search_term: &'static str,
priority: i32,
}
struct SearchProvider {
search_term: &'static str,
priority: i32,
}
impl MessageFilterProvider for SearchProvider {
fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
query
.filters
.push(MessageFilter::Search(self.search_term.to_string()));
}
fn priority(&self) -> i32 {
self.priority
}
}
impl Capability for SearchCapability {
fn id(&self) -> &str {
self.id
}
fn name(&self) -> &str {
"Search"
}
fn description(&self) -> &str {
"Test"
}
fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
Some(Arc::new(SearchProvider {
search_term: self.search_term,
priority: self.priority,
}))
}
}
let mut registry = CapabilityRegistry::new();
registry.register(SearchCapability {
id: "cap_a",
search_term: "alpha",
priority: 5,
});
registry.register(SearchCapability {
id: "cap_b",
search_term: "beta",
priority: 1,
});
registry.register(SearchCapability {
id: "cap_c",
search_term: "gamma",
priority: 10,
});
let configs = vec![
AgentCapabilityConfig {
capability_ref: CapabilityId::new("cap_a"),
config: serde_json::json!({}),
},
AgentCapabilityConfig {
capability_ref: CapabilityId::new("cap_b"),
config: serde_json::json!({}),
},
AgentCapabilityConfig {
capability_ref: CapabilityId::new("cap_c"),
config: serde_json::json!({}),
},
];
let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
let session_id: SessionId = Uuid::now_v7().into();
let mut query = MessageQuery::new(session_id);
collected.apply_message_filters(&mut query);
assert_eq!(query.filters.len(), 3);
assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
}
#[test]
fn test_capability_without_message_filter_returns_none() {
let registry = CapabilityRegistry::with_builtins();
let noop = registry.get("noop").unwrap();
assert!(noop.message_filter_provider().is_none());
let current_time = registry.get("current_time").unwrap();
assert!(current_time.message_filter_provider().is_none());
}
#[tokio::test]
async fn test_collect_capabilities_preserves_config_for_filter_provider() {
let mut registry = CapabilityRegistry::new();
registry.register(FilterTestCapability { priority: 0 });
let test_config = serde_json::json!({
"search": "custom_search",
"extra_field": 42
});
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("filter_test"),
config: test_config.clone(),
}];
let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
assert_eq!(collected.message_filter_providers.len(), 1);
let (_, stored_config) = &collected.message_filter_providers[0];
assert_eq!(*stored_config, test_config);
}
#[test]
fn test_collect_message_filters_only_collects_filters() {
let mut registry = CapabilityRegistry::new();
registry.register(FilterTestCapability { priority: 0 });
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("filter_test"),
config: serde_json::json!({ "search": "test_query" }),
}];
let collected = collect_message_filters_only(&configs, ®istry);
let session_id: SessionId = Uuid::now_v7().into();
let mut query = MessageQuery::new(session_id);
collected.apply_message_filters(&mut query);
assert_eq!(query.filters.len(), 1);
assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
}
#[test]
fn test_collect_message_filters_only_skips_unknown_capabilities() {
let registry = CapabilityRegistry::new();
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("nonexistent"),
config: serde_json::json!({}),
}];
let collected = collect_message_filters_only(&configs, ®istry);
assert!(collected.message_filter_providers.is_empty());
}
#[test]
fn test_collect_message_filters_only_preserves_priority_order() {
struct PriorityFilterCap {
id: &'static str,
search_term: &'static str,
priority: i32,
}
struct PriorityFilterProvider {
search_term: &'static str,
priority: i32,
}
impl Capability for PriorityFilterCap {
fn id(&self) -> &str {
self.id
}
fn name(&self) -> &str {
self.id
}
fn description(&self) -> &str {
"priority test"
}
fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
Some(Arc::new(PriorityFilterProvider {
search_term: self.search_term,
priority: self.priority,
}))
}
}
impl MessageFilterProvider for PriorityFilterProvider {
fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
query
.filters
.push(MessageFilter::Search(self.search_term.to_string()));
}
fn priority(&self) -> i32 {
self.priority
}
}
let mut registry = CapabilityRegistry::new();
registry.register(PriorityFilterCap {
id: "gamma",
search_term: "gamma",
priority: 10,
});
registry.register(PriorityFilterCap {
id: "alpha",
search_term: "alpha",
priority: 5,
});
registry.register(PriorityFilterCap {
id: "beta",
search_term: "beta",
priority: 1,
});
let configs = vec![
AgentCapabilityConfig {
capability_ref: CapabilityId::new("gamma"),
config: serde_json::json!({}),
},
AgentCapabilityConfig {
capability_ref: CapabilityId::new("alpha"),
config: serde_json::json!({}),
},
AgentCapabilityConfig {
capability_ref: CapabilityId::new("beta"),
config: serde_json::json!({}),
},
];
let collected = collect_message_filters_only(&configs, ®istry);
let session_id: SessionId = Uuid::now_v7().into();
let mut query = MessageQuery::new(session_id);
collected.apply_message_filters(&mut query);
assert_eq!(query.filters.len(), 3);
assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
}
#[test]
fn test_collect_message_filters_only_post_load_invoked() {
use crate::message::Message;
struct PostLoadCap;
struct PostLoadProvider;
impl Capability for PostLoadCap {
fn id(&self) -> &str {
"post_load_test"
}
fn name(&self) -> &str {
"PostLoad Test"
}
fn description(&self) -> &str {
"test"
}
fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
Some(Arc::new(PostLoadProvider))
}
}
impl MessageFilterProvider for PostLoadProvider {
fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
fn priority(&self) -> i32 {
0
}
fn post_load(&self, messages: &mut Vec<Message>, _config: &serde_json::Value) {
messages.reverse();
}
}
let mut registry = CapabilityRegistry::new();
registry.register(PostLoadCap);
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("post_load_test"),
config: serde_json::json!({}),
}];
let collected = collect_message_filters_only(&configs, ®istry);
let mut messages = vec![Message::user("first"), Message::user("second")];
collected.apply_post_load_filters(&mut messages);
assert_eq!(messages[0].text(), Some("second"));
assert_eq!(messages[1].text(), Some("first"));
}
#[test]
fn test_collect_model_view_providers_respects_compaction_capability_boundary() {
use crate::tool_types::ToolCall;
fn tool_heavy_messages() -> Vec<Message> {
let mut messages = vec![Message::user("inspect files repeatedly")];
for index in 0..9 {
let call_id = format!("call_{index}");
messages.push(Message::assistant_with_tools(
"",
vec![ToolCall {
id: call_id.clone(),
name: "read_file".to_string(),
arguments: serde_json::json!({"path": "/workspace/src/lib.rs"}),
}],
));
messages.push(Message::tool_result(
call_id,
Some(serde_json::json!({
"path": "/workspace/src/lib.rs",
"content": format!("{}{}", "large file line\n".repeat(1000), index),
"total_lines": 1000,
"lines_shown": {"start": 1, "end": 1000},
"truncated": false
})),
None,
));
}
messages
}
fn first_tool_result_is_masked(messages: &[Message]) -> bool {
messages[2]
.tool_result_content()
.and_then(|result| result.result.as_ref())
.and_then(|result| result.get("masked"))
.and_then(|masked| masked.as_bool())
.unwrap_or(false)
}
let mut registry = CapabilityRegistry::new();
registry.register(CompactionCapability);
let context = ModelViewContext {
session_id: SessionId::new(),
prior_usage: None,
};
let no_compaction = collect_model_view_providers(&[], ®istry, None);
let unmasked = no_compaction.apply_model_view(tool_heavy_messages(), &context);
assert!(!first_tool_result_is_masked(&unmasked));
let compaction = collect_model_view_providers(
&[AgentCapabilityConfig {
capability_ref: CapabilityId::new(COMPACTION_CAPABILITY_ID),
config: serde_json::json!({}),
}],
®istry,
None,
);
let masked = compaction.apply_model_view(tool_heavy_messages(), &context);
assert!(first_tool_result_is_masked(&masked));
let last_tool = masked.last().unwrap().tool_result_content().unwrap();
assert!(last_tool.result.as_ref().unwrap().get("content").is_some());
}
struct DelegatingFilterCap {
id: &'static str,
inner: std::sync::Arc<InnerFilterCap>,
}
struct InnerFilterCap;
impl Capability for InnerFilterCap {
fn id(&self) -> &str {
"inner_filter"
}
fn name(&self) -> &str {
"Inner Filter"
}
fn description(&self) -> &str {
"inner"
}
fn message_filter_provider(&self) -> Option<std::sync::Arc<dyn MessageFilterProvider>> {
Some(std::sync::Arc::new(SentinelFilter))
}
}
struct SentinelFilter;
impl MessageFilterProvider for SentinelFilter {
fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
}
impl Capability for DelegatingFilterCap {
fn id(&self) -> &str {
self.id
}
fn name(&self) -> &str {
"Delegating Filter"
}
fn description(&self) -> &str {
"delegating"
}
fn message_filter_provider(&self) -> Option<std::sync::Arc<dyn MessageFilterProvider>> {
None }
fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
Some(&*self.inner)
}
}
#[test]
fn test_collect_message_filters_only_honors_resolve_for_model_delegation() {
let inner = std::sync::Arc::new(InnerFilterCap);
let outer = DelegatingFilterCap {
id: "delegating_filter",
inner: inner.clone(),
};
let mut registry = CapabilityRegistry::new();
registry.register(outer);
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("delegating_filter"),
config: serde_json::json!({}),
}];
let collected = collect_message_filters_only(&configs, ®istry);
assert_eq!(
collected.message_filter_providers.len(),
1,
"provider from resolved inner capability must be collected"
);
}
struct DelegatingMvpCap {
id: &'static str,
inner: std::sync::Arc<InnerMvpCap>,
}
struct InnerMvpCap;
impl Capability for InnerMvpCap {
fn id(&self) -> &str {
"inner_mvp"
}
fn name(&self) -> &str {
"Inner MVP"
}
fn description(&self) -> &str {
"inner"
}
fn model_view_provider(
&self,
) -> Option<std::sync::Arc<dyn crate::capabilities::ModelViewProvider>> {
struct NoopMvp;
impl crate::capabilities::ModelViewProvider for NoopMvp {
fn apply_model_view(
&self,
messages: Vec<Message>,
_config: &serde_json::Value,
_context: &ModelViewContext<'_>,
) -> Vec<Message> {
messages
}
}
Some(std::sync::Arc::new(NoopMvp))
}
}
impl Capability for DelegatingMvpCap {
fn id(&self) -> &str {
self.id
}
fn name(&self) -> &str {
"Delegating MVP"
}
fn description(&self) -> &str {
"delegating"
}
fn model_view_provider(
&self,
) -> Option<std::sync::Arc<dyn crate::capabilities::ModelViewProvider>> {
None }
fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
Some(&*self.inner)
}
}
#[test]
fn test_collect_model_view_providers_honors_resolve_for_model_delegation() {
let inner = std::sync::Arc::new(InnerMvpCap);
let outer = DelegatingMvpCap {
id: "delegating_mvp",
inner: inner.clone(),
};
let mut registry = CapabilityRegistry::new();
registry.register(outer);
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("delegating_mvp"),
config: serde_json::json!({}),
}];
let collected = collect_model_view_providers(&configs, ®istry, None);
assert_eq!(
collected.model_view_providers.len(),
1,
"provider from resolved inner capability must be collected"
);
}
#[tokio::test]
async fn test_virtual_bash_capability_produces_bash_tool() {
let registry = CapabilityRegistry::with_builtins();
let collected =
collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
let tool_names: Vec<&str> = collected
.tool_definitions
.iter()
.map(|t| t.name())
.collect();
assert!(
tool_names.contains(&"bash"),
"virtual_bash capability must produce 'bash' tool, got: {:?}",
tool_names
);
assert!(
!collected.tools.is_empty(),
"virtual_bash must provide tool implementations"
);
}
#[tokio::test]
async fn test_generic_harness_capability_set_produces_bash_tool() {
let generic_harness_caps = vec![
"session_file_system".to_string(),
"virtual_bash".to_string(),
"web_fetch".to_string(),
"session_storage".to_string(),
"session".to_string(),
"agent_instructions".to_string(),
"skills".to_string(),
"infinity_context".to_string(),
"auto_tool_search".to_string(),
];
let registry = CapabilityRegistry::with_builtins();
let collected = collect_capabilities(&generic_harness_caps, ®istry, &test_ctx()).await;
let tool_names: Vec<&str> = collected
.tool_definitions
.iter()
.map(|t| t.name())
.collect();
assert!(
tool_names.contains(&"bash"),
"Generic Harness capabilities must produce 'bash' tool, got: {:?}",
tool_names
);
}
#[tokio::test]
async fn test_collect_capabilities_tool_count_matches_definitions() {
let registry = CapabilityRegistry::with_builtins();
let collected =
collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
assert_eq!(
collected.tools.len(),
collected.tool_definitions.len(),
"tool implementations ({}) must match tool definitions ({})",
collected.tools.len(),
collected.tool_definitions.len(),
);
}
#[tokio::test]
async fn test_collect_capabilities_resolves_dependencies() {
let registry = CapabilityRegistry::with_builtins();
let collected =
collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
assert!(
collected
.applied_ids
.iter()
.any(|id| id == "session_file_system"),
"collect_capabilities must apply session_file_system as a dependency; applied_ids: {:?}",
collected.applied_ids
);
let tool_names: Vec<&str> = collected
.tool_definitions
.iter()
.map(|t| t.name())
.collect();
assert!(
tool_names.contains(&"read_file") && tool_names.contains(&"write_file"),
"collect_capabilities must resolve dependencies and include dependency tools, got: {:?}",
tool_names
);
assert_eq!(
collected.tools.len(),
collected.tool_definitions.len(),
"dependency-added tools must have implementations, not just definitions"
);
}
#[test]
fn test_defaults_do_not_include_bash() {
let registry = crate::ToolRegistry::with_defaults();
assert!(
!registry.has("bash"),
"with_defaults() must not include 'bash' — it comes from virtual_bash capability"
);
}
#[tokio::test]
async fn test_background_execution_auto_activates_with_virtual_bash() {
let registry = CapabilityRegistry::with_builtins();
let collected =
collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
let tool_names: Vec<&str> = collected
.tool_definitions
.iter()
.map(|t| t.name())
.collect();
assert!(
tool_names.contains(&"spawn_background"),
"spawn_background must be auto-activated when virtual_bash (a \
background-capable tool) is in the agent's capability set; got: {:?}",
tool_names
);
assert!(
collected
.applied_ids
.iter()
.any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
"background_execution must be in applied_ids when auto-activated; \
got: {:?}",
collected.applied_ids
);
assert!(
collected
.tools
.iter()
.any(|t| t.name() == "spawn_background"),
"spawn_background tool implementation must be present alongside the \
definition (lockstep contract)"
);
}
#[tokio::test]
async fn test_background_execution_does_not_auto_activate_without_hint() {
let registry = CapabilityRegistry::with_builtins();
let collected =
collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
let tool_names: Vec<&str> = collected
.tool_definitions
.iter()
.map(|t| t.name())
.collect();
assert!(
!tool_names.contains(&"spawn_background"),
"spawn_background must NOT be activated without a background-capable \
tool; got: {:?}",
tool_names
);
assert!(
!collected
.applied_ids
.iter()
.any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
"background_execution must not appear in applied_ids when no \
background-capable tool is present; got: {:?}",
collected.applied_ids
);
}
#[tokio::test]
async fn test_background_execution_explicit_selection_is_idempotent() {
let registry = CapabilityRegistry::with_builtins();
let collected = collect_capabilities(
&[
"virtual_bash".to_string(),
BACKGROUND_EXECUTION_CAPABILITY_ID.to_string(),
],
®istry,
&test_ctx(),
)
.await;
let spawn_background_count = collected
.tool_definitions
.iter()
.filter(|t| t.name() == "spawn_background")
.count();
assert_eq!(
spawn_background_count, 1,
"spawn_background must appear exactly once even when \
background_execution is selected explicitly alongside a \
background-capable tool"
);
let applied_count = collected
.applied_ids
.iter()
.filter(|id| id.as_str() == BACKGROUND_EXECUTION_CAPABILITY_ID)
.count();
assert_eq!(
applied_count, 1,
"background_execution must appear exactly once in applied_ids"
);
}
#[test]
fn test_defaults_do_not_include_spawn_background() {
let registry = crate::ToolRegistry::with_defaults();
assert!(
!registry.has("spawn_background"),
"with_defaults() must not include 'spawn_background' — it comes \
from the background_execution capability (EVE-501)"
);
}
#[test]
fn test_capability_features_default_empty() {
let registry = CapabilityRegistry::with_builtins();
let noop = registry.get("noop").unwrap();
assert!(noop.features().is_empty());
let current_time = registry.get("current_time").unwrap();
assert!(current_time.features().is_empty());
}
#[test]
fn test_file_system_capability_features() {
let registry = CapabilityRegistry::with_builtins();
let fs = registry.get("session_file_system").unwrap();
assert_eq!(fs.features(), vec!["file_system"]);
}
#[test]
fn test_virtual_bash_capability_features() {
let registry = CapabilityRegistry::with_builtins();
let bash = registry.get("virtual_bash").unwrap();
assert_eq!(bash.features(), vec!["file_system"]);
}
#[test]
fn test_session_storage_capability_features() {
let registry = CapabilityRegistry::with_builtins();
let storage = registry.get("session_storage").unwrap();
let features = storage.features();
assert!(features.contains(&"secrets"));
assert!(features.contains(&"key_value"));
}
#[test]
fn test_session_schedule_capability_features() {
let registry = CapabilityRegistry::with_builtins();
let schedule = registry.get("session_schedule").unwrap();
assert_eq!(schedule.features(), vec!["schedules"]);
}
#[test]
fn test_session_sql_database_capability_features() {
let registry = CapabilityRegistry::with_builtins();
let sql = registry.get("session_sql_database").unwrap();
assert_eq!(sql.features(), vec!["sql_database"]);
}
#[test]
fn test_sample_data_capability_features() {
let registry = CapabilityRegistry::with_builtins();
let sample = registry.get("sample_data").unwrap();
assert_eq!(sample.features(), vec!["file_system"]);
}
#[test]
fn test_compute_features_empty() {
let registry = CapabilityRegistry::with_builtins();
let features = compute_features(&[], ®istry);
assert!(features.is_empty());
}
#[test]
fn test_compute_features_single_capability() {
let registry = CapabilityRegistry::with_builtins();
let features = compute_features(&["session_schedule".to_string()], ®istry);
assert_eq!(features, vec!["schedules"]);
}
#[test]
fn test_compute_features_multiple_capabilities() {
let registry = CapabilityRegistry::with_builtins();
let features = compute_features(
&[
"session_file_system".to_string(),
"session_storage".to_string(),
"session_schedule".to_string(),
],
®istry,
);
assert!(features.contains(&"file_system".to_string()));
assert!(features.contains(&"secrets".to_string()));
assert!(features.contains(&"key_value".to_string()));
assert!(features.contains(&"schedules".to_string()));
}
#[test]
fn test_compute_features_deduplicates() {
let registry = CapabilityRegistry::with_builtins();
let features = compute_features(
&[
"session_file_system".to_string(),
"virtual_bash".to_string(),
],
®istry,
);
let file_system_count = features.iter().filter(|f| *f == "file_system").count();
assert_eq!(file_system_count, 1, "file_system should appear only once");
}
#[test]
fn test_compute_features_includes_dependency_features() {
let registry = CapabilityRegistry::with_builtins();
let features = compute_features(&["virtual_bash".to_string()], ®istry);
assert!(features.contains(&"file_system".to_string()));
}
#[test]
fn test_compute_features_generic_harness_set() {
let registry = CapabilityRegistry::with_builtins();
let features = compute_features(
&[
"session_file_system".to_string(),
"virtual_bash".to_string(),
"session_storage".to_string(),
"session".to_string(),
"session_schedule".to_string(),
],
®istry,
);
assert!(features.contains(&"file_system".to_string()));
assert!(features.contains(&"secrets".to_string()));
assert!(features.contains(&"key_value".to_string()));
assert!(features.contains(&"schedules".to_string()));
}
#[test]
fn test_compute_features_unknown_capability_ignored() {
let registry = CapabilityRegistry::with_builtins();
let features = compute_features(
&["unknown_cap".to_string(), "session_schedule".to_string()],
®istry,
);
assert_eq!(features, vec!["schedules"]);
}
#[test]
fn test_risk_level_ordering() {
assert!(RiskLevel::Low < RiskLevel::Medium);
assert!(RiskLevel::Medium < RiskLevel::High);
}
#[test]
fn test_risk_level_serde_roundtrip() {
let high = RiskLevel::High;
let json = serde_json::to_string(&high).unwrap();
assert_eq!(json, "\"high\"");
let back: RiskLevel = serde_json::from_str(&json).unwrap();
assert_eq!(back, RiskLevel::High);
}
#[test]
fn test_capability_risk_levels() {
let registry = CapabilityRegistry::with_builtins();
let bash = registry.get("virtual_bash").unwrap();
assert_eq!(bash.risk_level(), RiskLevel::High);
let fetch = registry.get("web_fetch").unwrap();
assert_eq!(fetch.risk_level(), RiskLevel::High);
let noop = registry.get("noop").unwrap();
assert_eq!(noop.risk_level(), RiskLevel::Low);
}
#[tokio::test]
async fn test_apply_capabilities_openai_tool_search() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
let applied = apply_capabilities(
base_runtime_agent.clone(),
&["openai_tool_search".to_string()],
®istry,
&test_ctx(),
)
.await;
assert_eq!(
applied.runtime_agent.system_prompt,
base_runtime_agent.system_prompt
);
assert!(applied.tool_registry.is_empty());
assert_eq!(applied.applied_ids, vec!["openai_tool_search"]);
let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
assert!(ts.enabled);
assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
}
#[tokio::test]
async fn test_apply_capabilities_openai_tool_search_with_other_capabilities() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
let applied = apply_capabilities(
base_runtime_agent,
&[
"current_time".to_string(),
"openai_tool_search".to_string(),
"test_math".to_string(),
],
®istry,
&test_ctx(),
)
.await;
assert!(applied.tool_registry.has("get_current_time"));
assert!(applied.tool_registry.has("add"));
assert!(applied.tool_registry.has("subtract"));
assert!(applied.tool_registry.has("multiply"));
assert!(applied.tool_registry.has("divide"));
let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
assert!(ts.enabled);
assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
}
#[tokio::test]
async fn test_collect_capabilities_tool_search_custom_threshold() {
let registry = CapabilityRegistry::with_builtins();
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("openai_tool_search"),
config: serde_json::json!({"threshold": 5}),
}];
let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
let ts = collected.tool_search.as_ref().unwrap();
assert!(ts.enabled);
assert_eq!(ts.threshold, 5);
}
#[tokio::test]
async fn test_collect_capabilities_auto_tool_search_resolves_to_generic_off_native() {
let registry = CapabilityRegistry::with_builtins();
let configs = vec![
AgentCapabilityConfig {
capability_ref: CapabilityId::new("auto_tool_search"),
config: serde_json::json!({"threshold": 2}),
},
AgentCapabilityConfig {
capability_ref: CapabilityId::new("test_math"),
config: serde_json::json!({}),
},
];
let ctx = test_ctx().with_model("claude-sonnet-4-5-20250514");
let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
assert!(
collected.tool_search.is_none(),
"auto_tool_search must not set a hosted config on a non-native model"
);
assert!(
collected
.tools
.iter()
.any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
"auto_tool_search must contribute the client-side tool_search tool"
);
assert!(
!collected.tool_definition_hooks.is_empty(),
"auto_tool_search must contribute a client-side deferral hook"
);
let mut transformed = collected.tool_definitions.clone();
for hook in &collected.tool_definition_hooks {
transformed = hook.transform(transformed);
}
let add_tool = transformed
.iter()
.find(|tool| tool.name() == "add")
.expect("test_math contributes add");
assert!(
add_tool.parameters().get("properties").is_none(),
"generic auto_tool_search must honor the configured threshold"
);
}
#[tokio::test]
async fn test_collect_capabilities_auto_tool_search_resolves_to_hosted_on_native() {
let registry = CapabilityRegistry::with_builtins();
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("auto_tool_search"),
config: serde_json::json!({"threshold": 7}),
}];
let ctx = test_ctx().with_model("gpt-5.4");
let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
let ts = collected
.tool_search
.as_ref()
.expect("auto_tool_search must set a hosted config on a native model");
assert!(ts.enabled);
assert_eq!(ts.threshold, 7);
assert!(
!collected
.tools
.iter()
.any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
"hosted mechanism must not contribute the client-side tool_search tool"
);
assert!(
collected.tool_definition_hooks.is_empty(),
"hosted mechanism must not contribute a client-side deferral hook"
);
}
#[tokio::test]
async fn test_collect_capabilities_no_tool_search_without_capability() {
let registry = CapabilityRegistry::with_builtins();
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("current_time"),
config: serde_json::json!({}),
}];
let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
assert!(collected.tool_search.is_none());
}
#[tokio::test]
async fn test_collect_capabilities_tool_search_category_propagation() {
let registry = CapabilityRegistry::with_builtins();
let configs = vec![
AgentCapabilityConfig {
capability_ref: CapabilityId::new("test_math"),
config: serde_json::json!({}),
},
AgentCapabilityConfig {
capability_ref: CapabilityId::new("openai_tool_search"),
config: serde_json::json!({}),
},
];
let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
assert!(collected.tool_search.is_some());
for tool_def in &collected.tool_definitions {
if ["add", "subtract", "multiply", "divide"].contains(&tool_def.name()) {
assert!(
tool_def.category().is_some(),
"Tool {} should have a category from its capability",
tool_def.name()
);
}
}
}
#[tokio::test]
async fn test_apply_capabilities_prompt_caching() {
let registry = CapabilityRegistry::with_builtins();
let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
let applied = apply_capabilities(
base_runtime_agent.clone(),
&["prompt_caching".to_string()],
®istry,
&test_ctx(),
)
.await;
assert_eq!(
applied.runtime_agent.system_prompt,
base_runtime_agent.system_prompt
);
assert!(applied.tool_registry.is_empty());
assert_eq!(applied.applied_ids, vec!["prompt_caching"]);
let prompt_cache = applied.runtime_agent.prompt_cache.as_ref().unwrap();
assert!(prompt_cache.enabled);
assert_eq!(
prompt_cache.strategy,
crate::llm_driver_registry::PromptCacheStrategy::Auto
);
assert!(prompt_cache.gemini_cached_content.is_none());
}
#[tokio::test]
async fn test_collect_capabilities_prompt_caching_custom_strategy() {
let registry = CapabilityRegistry::with_builtins();
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("prompt_caching"),
config: serde_json::json!({"strategy": "auto"}),
}];
let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
let prompt_cache = collected.prompt_cache.as_ref().unwrap();
assert!(prompt_cache.enabled);
assert_eq!(
prompt_cache.strategy,
crate::llm_driver_registry::PromptCacheStrategy::Auto
);
assert!(prompt_cache.gemini_cached_content.is_none());
}
#[tokio::test]
async fn test_collect_capabilities_prompt_caching_gemini_cached_content() {
let registry = CapabilityRegistry::with_builtins();
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("prompt_caching"),
config: serde_json::json!({
"strategy": "auto",
"gemini_cached_content": "cachedContents/demo-cache"
}),
}];
let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
let prompt_cache = collected.prompt_cache.as_ref().unwrap();
assert_eq!(
prompt_cache.gemini_cached_content.as_deref(),
Some("cachedContents/demo-cache")
);
}
struct SkillContributingCapability;
impl Capability for SkillContributingCapability {
fn id(&self) -> &str {
"contributes_skills"
}
fn name(&self) -> &str {
"Contributes Skills"
}
fn description(&self) -> &str {
"Test capability that contributes skills."
}
fn contribute_skills(&self) -> Vec<SkillContribution> {
vec![
SkillContribution::new("alpha-skill", "Alpha skill desc", "# Alpha\nDo alpha.")
.with_files(vec![(
"scripts/a.sh".to_string(),
"#!/bin/sh\necho a\n".to_string(),
)]),
SkillContribution::new("beta-skill", "Beta skill desc", "# Beta\nDo beta.")
.with_user_invocable(false),
]
}
}
fn skill_md_from_entries(entries: &HashMap<String, MountEntry>) -> &str {
match &entries.get("SKILL.md").expect("SKILL.md missing").source {
MountSource::InlineFile { content, .. } => content.as_str(),
_ => panic!("Expected InlineFile for SKILL.md"),
}
}
#[tokio::test]
async fn test_contribute_skills_normalized_to_mounts() {
let mut registry = CapabilityRegistry::new();
registry.register(SkillContributingCapability);
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("contributes_skills"),
config: serde_json::json!({}),
}];
let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
let skill_mounts: Vec<_> = collected
.mounts
.iter()
.filter(|m| m.path.starts_with("/.agents/skills/"))
.collect();
assert_eq!(skill_mounts.len(), 2);
for m in &skill_mounts {
assert!(m.is_readonly());
assert_eq!(m.capability_id, "contributes_skills");
}
let alpha = skill_mounts
.iter()
.find(|m| m.path == "/.agents/skills/alpha-skill")
.expect("alpha-skill mount missing");
match &alpha.source {
MountSource::InlineDirectory { entries } => {
assert!(entries.contains_key("SKILL.md"));
assert!(entries.contains_key("scripts/a.sh"));
let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
assert_eq!(parsed.name, "alpha-skill");
assert!(parsed.user_invocable);
}
_ => panic!("Expected InlineDirectory"),
}
let beta = skill_mounts
.iter()
.find(|m| m.path == "/.agents/skills/beta-skill")
.expect("beta-skill mount missing");
match &beta.source {
MountSource::InlineDirectory { entries } => {
let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
assert!(!parsed.user_invocable);
}
_ => panic!("Expected InlineDirectory"),
}
}
#[tokio::test]
async fn test_contribute_skills_default_empty() {
let mut registry = CapabilityRegistry::new();
registry.register(FilterTestCapability { priority: 0 });
let configs = vec![AgentCapabilityConfig {
capability_ref: CapabilityId::new("filter_test"),
config: serde_json::json!({}),
}];
let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
assert!(
collected
.mounts
.iter()
.all(|m| !m.path.starts_with("/.agents/skills/"))
);
}
}