1use crate::command::{
22 CommandDescriptor, CommandExecutionContext, CommandResult, ExecuteCommandRequest,
23};
24use crate::deployment::DeploymentGrade;
25use crate::events::TokenUsage;
26use crate::mcp_server::{ScopedMcpServers, merge_scoped_mcp_servers};
27use crate::message::Message;
28use crate::message_filter::MessageFilterProvider;
29use crate::runtime_agent::RuntimeAgent;
30use crate::tool_types::{ToolCall, ToolDefinition};
31use crate::tools::{Tool, ToolRegistry};
32use crate::traits::SessionFileSystem;
33use crate::typed_id::SessionId;
34use async_trait::async_trait;
35use serde::{Deserialize, Serialize};
36use std::collections::HashMap;
37use std::sync::Arc;
38
39pub struct IntegrationPlugin {
63 pub experimental_only: bool,
65 pub feature_flag: Option<&'static str>,
68 pub factory: fn() -> Box<dyn Capability>,
70}
71
72inventory::collect!(IntegrationPlugin);
73
74pub use crate::capability_types::{
76 AgentCapabilityConfig, CapabilityId, CapabilityStatus, MountAccess, MountDirectoryBuilder,
77 MountEntry, MountPoint, MountSource,
78};
79
80mod a2a_delegation;
85#[cfg(feature = "ui-capabilities")]
86mod a2ui;
87mod agent_handoff;
88mod agent_instructions;
89pub mod attach_skill;
90mod auto_tool_search;
91mod background_execution;
92mod btw;
93mod budgeting;
94pub mod compaction;
95mod current_time;
96mod data_knowledge;
97mod declarative;
98mod fake_aws;
99mod fake_crm;
100mod fake_financial;
101mod fake_warehouse;
102mod file_system;
103mod human_intent;
104mod infinity_context;
105mod knowledge_base;
106mod loop_detection;
107mod lua;
108mod lua_code_mode;
109pub mod mcp;
110mod noop;
111mod openai_tool_search;
112#[cfg(feature = "ui-capabilities")]
113mod openui;
114pub mod persistent_memory;
115mod platform_management;
116mod prompt_caching;
117mod prompt_canary_guardrail;
118mod research;
119mod sample_data;
120mod self_budget;
121mod session;
122mod session_sandbox;
123mod session_schedule;
124mod session_sql_database;
125mod session_storage;
126mod skills;
127mod stateless_todo_list;
128mod subagents;
129mod system_commands;
130mod test_math;
131mod test_weather;
132mod tool_output_persistence;
133mod tool_search;
134pub mod user_hooks;
135mod virtual_bash;
136mod web_fetch;
137mod workspace_volumes;
138
139pub use a2a_delegation::{
141 A2A_AGENT_DELEGATION_CAPABILITY_ID, A2aAgentDelegationCapability, CancelAgentTool,
142 GetAgentRunsTool, MessageAgentTool, SpawnAgentTool, WaitAgentTool,
143};
144#[cfg(feature = "ui-capabilities")]
145pub use a2ui::{A2UI_CAPABILITY_ID, A2UiCapability};
146pub use agent_handoff::{
147 AGENT_HANDOFF_CAPABILITY_ID, AgentHandoffCapability, GetAgentHandoffsTool,
148 MessageAgentHandoffTool, StartAgentHandoffTool,
149};
150pub use agent_instructions::{
151 AGENT_INSTRUCTIONS_CAPABILITY_ID, AGENTS_MD_PATH, AgentInstructionsCapability,
152 AgentInstructionsConfig, DEFAULT_AGENT_INSTRUCTIONS_FILE, MAX_AGENT_INSTRUCTIONS_FILES,
153 MAX_AGENTS_MD_SIZE, format_agents_md_content, format_instruction_file_content,
154};
155pub use attach_skill::{
156 AttachSkillCapability, SKILL_CAPABILITY_PREFIX, SKILLS_DISCOVERY_PATH, SkillContribution,
157 SkillInstructions, SkillMeta, SkillSource, discover_skills_from_entries, is_skill_capability,
158 parse_skill_capability_id, reconstruct_skill_md, skill_capability_id,
159};
160pub use auto_tool_search::{AUTO_TOOL_SEARCH_CAPABILITY_ID, AutoToolSearchCapability};
161pub use background_execution::{BACKGROUND_EXECUTION_CAPABILITY_ID, BackgroundExecutionCapability};
162pub use btw::{BTW_CAPABILITY_ID, BtwCapability};
163pub use budgeting::BudgetingCapability;
164pub use compaction::{
165 COMPACTION_CAPABILITY_ID, CompactionCapability, CompactionConfig, CompactionStep,
166 CompactionStrategy, CostControlConfig, CostControlMaskingResult, HierarchicalMemoryConfig,
167 MaskingSummaryFormat, MemoryTier, ObservationMaskingConfig, ObservationMaskingResult,
168 SessionCompactionMetrics, SummarizationConfig, aggressive_trim, apply_cost_control_masking,
169 apply_hierarchical_memory, apply_observation_masking, build_model_view_messages,
170 build_summarization_prompt, build_summary_message, classify_memory_tiers, estimate_tokens,
171 estimate_total_tokens, format_messages_for_summarization, should_compact_proactively,
172};
173pub use current_time::{CurrentTimeCapability, GetCurrentTimeTool};
174pub use data_knowledge::DataKnowledgeCapability;
175pub use declarative::{
176 DECLARATIVE_CAPABILITY_PREFIX, DeclarativeCapabilityDefinition, DeclarativeCapabilityFile,
177 DeclarativeCapabilitySkill, declarative_capability_id, declarative_capability_info,
178 hydrate_declarative_capability_config, is_declarative_capability,
179 parse_declarative_capability_id, validate_declarative_capability_definition,
180};
181pub use fake_aws::{
182 AwsCreateEc2InstanceTool, AwsCreateIamUserTool, AwsCreateRdsDatabaseTool,
183 AwsCreateS3BucketTool, AwsGetCloudWatchMetricsTool, AwsListEc2InstancesTool,
184 AwsListIamUsersTool, AwsListRdsDatabasesTool, AwsListS3BucketsTool, AwsListSecurityGroupsTool,
185 AwsStopEc2InstanceTool, FakeAwsCapability,
186};
187pub use fake_crm::{
188 CrmAddInteractionTool, CrmCreateCustomerTool, CrmCreateTicketTool, CrmGetCustomerTool,
189 CrmListCustomersTool, CrmListTicketsTool, CrmSearchCustomersTool, CrmUpdateTicketTool,
190 FakeCrmCapability,
191};
192pub use fake_financial::{
193 FakeFinancialCapability, FinanceCreateBudgetTool, FinanceCreateTransactionTool,
194 FinanceForecastCashFlowTool, FinanceGetBalanceTool, FinanceGetExpenseReportTool,
195 FinanceGetRevenueReportTool, FinanceListBudgetsTool, FinanceListTransactionsTool,
196};
197pub use fake_warehouse::{
198 FakeWarehouseCapability, WarehouseCreateInvoiceTool, WarehouseCreateOrderTool,
199 WarehouseCreateShipmentTool, WarehouseGetInventoryTool, WarehouseInventoryReportTool,
200 WarehouseListOrdersTool, WarehouseListShipmentsTool, WarehouseProcessReturnTool,
201 WarehouseUpdateInventoryTool, WarehouseUpdateShipmentStatusTool,
202};
203pub use file_system::{
204 DeleteFileTool, EditFileTool, FileSystemCapability, GrepFilesTool, ListDirectoryTool,
205 ReadFileTool, StatFileTool, WriteFileTool,
206};
207pub use human_intent::{HUMAN_INTENT_CAPABILITY_ID, HumanIntentCapability};
208pub use infinity_context::{
209 INFINITY_CONTEXT_CAPABILITY_ID, InfinityContextCapability, QueryHistoryTool,
210};
211pub use knowledge_base::{
212 KNOWLEDGE_BASE_CAPABILITY_ID, KnowledgeBaseCapability, KnowledgeBaseConfig,
213 validate_knowledge_base_config,
214};
215pub use loop_detection::LoopDetectionCapability;
216pub use lua::{LUA_CAPABILITY_ID, LuaCapability, LuaTool, LuaVfs, is_code_mode_eligible};
217pub use lua_code_mode::{LUA_CODE_MODE_CAPABILITY_ID, LuaCodeModeCapability};
218pub use mcp::{
219 MCP_CAPABILITY_PREFIX, McpCapability, is_mcp_capability, mcp_capability_id,
220 parse_mcp_capability_id,
221};
222pub use noop::NoopCapability;
223pub use openai_tool_search::{
224 DEFAULT_TOOL_SEARCH_THRESHOLD, OPENAI_TOOL_SEARCH_CAPABILITY_ID, OpenAiToolSearchCapability,
225 model_supports_native_tool_search,
226};
227#[cfg(feature = "ui-capabilities")]
228pub use openui::{OPENUI_CAPABILITY_ID, OpenUiCapability};
229pub use persistent_memory::{
230 ForgetTool, MEMORY_CAPABILITY_ID, MemoryCapability, MemoryConfig, RecallTool, RememberTool,
231};
232pub use platform_management::{
233 ManageAgentsTool, ManageHarnessesTool, ManageSessionsTool, PlatformManagementCapability,
234 ReadAgentsTool, ReadCapabilitiesTool, ReadHarnessesTool, ReadSessionsTool,
235 SessionReadMessagesTool, SessionReadResponseTool, SessionSendMessageTool,
236};
237pub use prompt_caching::{PROMPT_CACHING_CAPABILITY_ID, PromptCachingCapability};
238pub use prompt_canary_guardrail::{
239 DEFAULT_REPLACEMENT as PROMPT_CANARY_DEFAULT_REPLACEMENT,
240 PROMPT_CANARY_GUARDRAIL_CAPABILITY_ID, PromptCanaryGuardrailCapability,
241 REASON_CODE_SYSTEM_PROMPT_LEAK,
242};
243pub use research::ResearchCapability;
244pub use sample_data::SampleDataCapability;
245pub use self_budget::{SELF_BUDGET_CAPABILITY_ID, SelfBudgetCapability};
246pub use session::{GetSessionInfoTool, SessionCapability, WriteSessionTitleTool};
247pub use session_sandbox::{
248 SESSION_SANDBOX_CAPABILITY_ID, SandboxExecTool, SandboxManageTool, SandboxReadFileTool,
249 SandboxStatusTool, SandboxWriteFileTool, SessionSandboxCapability,
250};
251pub use session_schedule::{
252 CancelScheduleTool, CreateScheduleTool, ListSchedulesTool, SESSION_SCHEDULE_CAPABILITY_ID,
253 SessionScheduleCapability,
254};
255pub use session_sql_database::{
256 SessionSqlDatabaseCapability, SqlExecuteTool, SqlQueryTool, SqlSchemaTool,
257};
258pub use session_storage::{KvStoreTool, SecretStoreTool, SessionStorageCapability};
259pub use skills::{SKILLS_CAPABILITY_ID, SkillsCapability};
260pub use stateless_todo_list::{StatelessTodoListCapability, WriteTodosTool};
261pub use subagents::SubagentCapability;
262pub use system_commands::{SYSTEM_COMMANDS_CAPABILITY_ID, SystemCommandsCapability};
264pub use test_math::{AddTool, DivideTool, MultiplyTool, SubtractTool, TestMathCapability};
265pub use test_weather::{GetForecastTool, GetWeatherTool, TestWeatherCapability};
266pub use tool_output_persistence::{PersistOutputHook, ToolOutputPersistenceCapability};
267pub use tool_search::{
268 TOOL_SEARCH_CAPABILITY_ID, TOOL_SEARCH_TOOL_NAME, ToolSearchCapability, ToolSearchTool,
269};
270pub use user_hooks::UserHooksCapability;
271pub use virtual_bash::{BashTool, SessionFileSystemAdapter, VirtualBashCapability};
272pub use web_fetch::{
273 BotAuthPublicKey, WebFetchCapability, WebFetchTool, derive_bot_auth_public_key,
274};
275pub use workspace_volumes::{WORKSPACE_VOLUMES_CAPABILITY_ID, WorkspaceVolumesCapability};
276
277pub struct SystemPromptContext {
287 pub session_id: SessionId,
289 pub locale: Option<String>,
291 pub file_store: Option<Arc<dyn SessionFileSystem>>,
293 pub model: Option<String>,
299}
300
301impl SystemPromptContext {
302 pub fn without_file_store(session_id: SessionId) -> Self {
304 Self {
305 session_id,
306 locale: None,
307 file_store: None,
308 model: None,
309 }
310 }
311
312 pub fn with_model(mut self, model: impl Into<String>) -> Self {
314 self.model = Some(model.into());
315 self
316 }
317}
318
319#[async_trait]
366pub trait Capability: Send + Sync {
367 fn id(&self) -> &str;
369
370 fn name(&self) -> &str;
372
373 fn description(&self) -> &str;
375
376 fn status(&self) -> CapabilityStatus {
378 CapabilityStatus::Available
379 }
380
381 fn icon(&self) -> Option<&str> {
383 None
384 }
385
386 fn category(&self) -> Option<&str> {
388 None
389 }
390
391 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
402 None
403 }
404
405 fn system_prompt_addition(&self) -> Option<&str> {
425 None
426 }
427
428 async fn system_prompt_contribution(&self, _ctx: &SystemPromptContext) -> Option<String> {
440 self.system_prompt_addition().map(|addition| {
441 format!(
442 "<capability id=\"{}\">\n{}\n</capability>",
443 self.id(),
444 addition
445 )
446 })
447 }
448
449 fn system_prompt_preview(&self) -> Option<String> {
455 self.system_prompt_addition().map(|s| s.to_string())
456 }
457
458 fn tools(&self) -> Vec<Box<dyn Tool>> {
460 vec![]
461 }
462
463 fn tools_with_config(&self, _config: &serde_json::Value) -> Vec<Box<dyn Tool>> {
471 self.tools()
472 }
473
474 async fn system_prompt_contribution_with_config(
481 &self,
482 ctx: &SystemPromptContext,
483 _config: &serde_json::Value,
484 ) -> Option<String> {
485 self.system_prompt_contribution(ctx).await
486 }
487
488 fn tool_definitions(&self) -> Vec<ToolDefinition> {
491 self.tools().iter().map(|t| t.to_definition()).collect()
492 }
493
494 fn mounts(&self) -> Vec<MountPoint> {
502 vec![]
503 }
504
505 fn dependencies(&self) -> Vec<&'static str> {
514 vec![]
515 }
516
517 fn features(&self) -> Vec<&'static str> {
532 vec![]
533 }
534
535 fn config_schema(&self) -> Option<serde_json::Value> {
541 None
542 }
543
544 fn config_ui_schema(&self) -> Option<serde_json::Value> {
549 None
550 }
551
552 fn validate_config(&self, _config: &serde_json::Value) -> Result<(), String> {
558 Ok(())
559 }
560
561 fn mcp_servers(&self) -> ScopedMcpServers {
567 ScopedMcpServers::default()
568 }
569
570 fn mcp_servers_with_config(&self, _config: &serde_json::Value) -> ScopedMcpServers {
572 self.mcp_servers()
573 }
574
575 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
588 None
589 }
590
591 fn model_view_provider(&self) -> Option<Arc<dyn ModelViewProvider>> {
599 None
600 }
601
602 fn pre_tool_use_hooks(&self) -> Vec<Arc<dyn crate::atoms::PreToolUseHook>> {
613 vec![]
614 }
615
616 fn post_tool_exec_hooks(&self) -> Vec<Arc<dyn crate::atoms::PostToolExecHook>> {
624 vec![]
625 }
626
627 fn tool_definition_hooks(&self) -> Vec<Arc<dyn ToolDefinitionHook>> {
636 vec![]
637 }
638
639 fn tool_definition_hooks_with_config(
644 &self,
645 _config: &serde_json::Value,
646 ) -> Vec<Arc<dyn ToolDefinitionHook>> {
647 self.tool_definition_hooks()
648 }
649
650 fn tool_call_hooks(&self) -> Vec<Arc<dyn ToolCallHook>> {
658 vec![]
659 }
660
661 fn user_hooks(&self) -> Vec<crate::user_hook_types::UserHookSpec> {
677 vec![]
678 }
679
680 fn user_hooks_with_config(
686 &self,
687 _config: &serde_json::Value,
688 ) -> Vec<crate::user_hook_types::UserHookSpec> {
689 self.user_hooks()
690 }
691
692 fn risk_level(&self) -> RiskLevel {
700 RiskLevel::Low
701 }
702
703 fn commands(&self) -> Vec<CommandDescriptor> {
711 vec![]
712 }
713
714 async fn execute_command(
728 &self,
729 request: &ExecuteCommandRequest,
730 _ctx: &CommandExecutionContext,
731 ) -> crate::error::Result<CommandResult> {
732 Err(crate::error::AgentLoopError::config(format!(
733 "capability {} declared command /{} but does not implement execute_command",
734 self.id(),
735 request.name,
736 )))
737 }
738
739 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
747 vec![]
748 }
749
750 fn contribute_skills(&self) -> Vec<SkillContribution> {
760 vec![]
761 }
762
763 fn output_guardrails(&self) -> Vec<Arc<dyn crate::output_guardrail::OutputGuardrail>> {
774 vec![]
775 }
776}
777
778pub trait ToolDefinitionHook: Send + Sync {
779 fn transform(&self, tools: Vec<ToolDefinition>) -> Vec<ToolDefinition>;
780
781 fn applies_with_native_tool_search(&self) -> bool {
786 true
787 }
788}
789
790pub trait ToolCallHook: Send + Sync {
791 fn narration(
792 &self,
793 _tool_def: Option<&ToolDefinition>,
794 _tool_call: &ToolCall,
795 _phase: crate::tool_narration::ToolNarrationPhase,
796 _locale: Option<&str>,
797 ) -> Option<String> {
798 None
799 }
800
801 fn transform_for_execution(&self, tool_call: ToolCall) -> ToolCall {
802 tool_call
803 }
804}
805
806#[derive(
810 Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
811)]
812#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
813#[cfg_attr(feature = "openapi", schema(example = "low"))]
814#[serde(rename_all = "lowercase")]
815pub enum RiskLevel {
816 Low,
818 Medium,
820 High,
822}
823
824#[derive(Debug, Clone, Serialize, Deserialize)]
830#[serde(rename_all = "snake_case")]
831pub enum BlueprintModel {
832 Fixed(String),
834 Default(String),
836 Inherit,
838}
839
840pub struct AgentBlueprint {
846 pub id: &'static str,
848 pub name: &'static str,
850 pub description: &'static str,
852 pub model: BlueprintModel,
854 pub system_prompt: &'static str,
856 pub tools: Vec<Box<dyn Tool>>,
858 pub max_turns: Option<usize>,
860 pub config_schema: Option<serde_json::Value>,
862}
863
864impl AgentBlueprint {
865 pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
867 self.tools.iter().map(|t| t.to_definition()).collect()
868 }
869}
870
871impl std::fmt::Debug for AgentBlueprint {
872 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
873 f.debug_struct("AgentBlueprint")
874 .field("id", &self.id)
875 .field("name", &self.name)
876 .field("model", &self.model)
877 .field("tool_count", &self.tools.len())
878 .field("max_turns", &self.max_turns)
879 .finish()
880 }
881}
882
883#[derive(Clone)]
910pub struct CapabilityRegistry {
911 capabilities: HashMap<String, Arc<dyn Capability>>,
912}
913
914impl CapabilityRegistry {
915 pub fn new() -> Self {
917 Self {
918 capabilities: HashMap::new(),
919 }
920 }
921
922 pub fn with_builtins() -> Self {
927 Self::with_builtins_for_grade(DeploymentGrade::from_env())
928 }
929
930 pub fn with_builtins_for_grade(grade: DeploymentGrade) -> Self {
935 let mut registry = Self::new();
936
937 registry.register(AgentInstructionsCapability);
939 registry.register(HumanIntentCapability);
940 registry.register(NoopCapability);
941 registry.register(CurrentTimeCapability);
942 registry.register(ResearchCapability);
943 registry.register(PlatformManagementCapability);
944 registry.register(FileSystemCapability);
945 registry.register(WorkspaceVolumesCapability);
946 registry.register(SessionStorageCapability);
947 registry.register(SessionCapability);
948 registry.register(SessionSqlDatabaseCapability);
949 registry.register(TestMathCapability);
950 registry.register(TestWeatherCapability);
951 registry.register(StatelessTodoListCapability);
952 registry.register(WebFetchCapability::from_env());
953 registry.register(VirtualBashCapability);
954 registry.register(BackgroundExecutionCapability);
955 registry.register(SessionScheduleCapability);
956 registry.register(BtwCapability);
957 registry.register(InfinityContextCapability);
958 registry.register(budgeting::BudgetingCapability);
959 registry.register(SelfBudgetCapability);
960 registry.register(CompactionCapability);
961 registry.register(MemoryCapability);
962
963 registry.register(OpenAiToolSearchCapability::new());
965 registry.register(ToolSearchCapability::new());
967 registry.register(AutoToolSearchCapability::new());
969 registry.register(PromptCachingCapability::new());
970
971 registry.register(SkillsCapability);
973
974 registry.register(SubagentCapability);
976
977 if crate::FeatureFlags::from_env(&grade).agent_delegation {
981 registry.register(AgentHandoffCapability);
982 registry.register(A2aAgentDelegationCapability);
983 }
984
985 registry.register(SystemCommandsCapability);
987
988 registry.register(tool_output_persistence::ToolOutputPersistenceCapability);
990
991 registry.register(user_hooks::UserHooksCapability);
994
995 registry.register(LoopDetectionCapability);
997
998 registry.register(PromptCanaryGuardrailCapability);
1001
1002 #[cfg(feature = "ui-capabilities")]
1004 {
1005 registry.register(OpenUiCapability);
1006 registry.register(A2UiCapability);
1007 }
1008
1009 registry.register(SampleDataCapability);
1011
1012 registry.register(DataKnowledgeCapability);
1014
1015 registry.register(KnowledgeBaseCapability);
1017
1018 registry.register(FakeWarehouseCapability);
1020 registry.register(FakeAwsCapability);
1021 registry.register(FakeCrmCapability);
1022 registry.register(FakeFinancialCapability);
1023
1024 let internal_flags = crate::InternalFeatureFlags::from_env();
1026 if internal_flags.session_sandbox {
1027 registry.register(SessionSandboxCapability);
1028 }
1029
1030 if internal_flags.lua {
1034 registry.register(LuaCapability);
1035 registry.register(LuaCodeModeCapability);
1038 }
1039 for plugin in inventory::iter::<IntegrationPlugin>() {
1040 if (!plugin.experimental_only || grade.experimental_features_enabled())
1041 && plugin
1042 .feature_flag
1043 .is_none_or(|f| internal_flags.is_enabled(f))
1044 {
1045 registry.register_boxed((plugin.factory)());
1046 }
1047 }
1048
1049 registry
1050 }
1051
1052 pub fn register(&mut self, capability: impl Capability + 'static) {
1054 self.capabilities
1055 .insert(capability.id().to_string(), Arc::new(capability));
1056 }
1057
1058 pub fn register_boxed(&mut self, capability: Box<dyn Capability>) {
1060 self.capabilities
1061 .insert(capability.id().to_string(), Arc::from(capability));
1062 }
1063
1064 pub fn register_arc(&mut self, capability: Arc<dyn Capability>) {
1066 self.capabilities
1067 .insert(capability.id().to_string(), capability);
1068 }
1069
1070 pub fn get(&self, id: &str) -> Option<&Arc<dyn Capability>> {
1072 self.capabilities.get(id)
1073 }
1074
1075 pub fn unregister(&mut self, id: &str) -> Option<Arc<dyn Capability>> {
1077 self.capabilities.remove(id)
1078 }
1079
1080 pub fn has(&self, id: &str) -> bool {
1082 self.capabilities.contains_key(id)
1083 }
1084
1085 pub fn list(&self) -> Vec<&Arc<dyn Capability>> {
1087 self.capabilities.values().collect()
1088 }
1089
1090 pub fn len(&self) -> usize {
1092 self.capabilities.len()
1093 }
1094
1095 pub fn is_empty(&self) -> bool {
1097 self.capabilities.is_empty()
1098 }
1099
1100 pub fn builder() -> CapabilityRegistryBuilder {
1102 CapabilityRegistryBuilder::new()
1103 }
1104
1105 pub fn blueprint(&self, id: &str) -> Option<AgentBlueprint> {
1109 for cap in self.capabilities.values() {
1110 for bp in cap.agent_blueprints() {
1111 if bp.id == id {
1112 return Some(bp);
1113 }
1114 }
1115 }
1116 None
1117 }
1118
1119 pub fn blueprint_with_capability(&self, id: &str) -> Option<(String, AgentBlueprint)> {
1123 for (capability_id, cap) in &self.capabilities {
1124 for bp in cap.agent_blueprints() {
1125 if bp.id == id {
1126 return Some((capability_id.clone(), bp));
1127 }
1128 }
1129 }
1130 None
1131 }
1132
1133 pub fn all_blueprints(&self) -> Vec<AgentBlueprint> {
1135 self.capabilities
1136 .values()
1137 .flat_map(|cap| cap.agent_blueprints())
1138 .collect()
1139 }
1140}
1141
1142impl Default for CapabilityRegistry {
1143 fn default() -> Self {
1144 Self::with_builtins()
1145 }
1146}
1147
1148impl std::fmt::Debug for CapabilityRegistry {
1149 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1150 let ids: Vec<_> = self.capabilities.keys().collect();
1151 f.debug_struct("CapabilityRegistry")
1152 .field("capabilities", &ids)
1153 .finish()
1154 }
1155}
1156
1157pub struct CapabilityRegistryBuilder {
1159 registry: CapabilityRegistry,
1160}
1161
1162impl CapabilityRegistryBuilder {
1163 pub fn new() -> Self {
1165 Self {
1166 registry: CapabilityRegistry::new(),
1167 }
1168 }
1169
1170 pub fn with_builtins() -> Self {
1172 Self {
1173 registry: CapabilityRegistry::with_builtins(),
1174 }
1175 }
1176
1177 pub fn capability(mut self, capability: impl Capability + 'static) -> Self {
1179 self.registry.register(capability);
1180 self
1181 }
1182
1183 pub fn build(self) -> CapabilityRegistry {
1185 self.registry
1186 }
1187}
1188
1189impl Default for CapabilityRegistryBuilder {
1190 fn default() -> Self {
1191 Self::new()
1192 }
1193}
1194
1195pub struct ModelViewContext<'a> {
1201 pub session_id: SessionId,
1202 pub prior_usage: Option<&'a TokenUsage>,
1203}
1204
1205pub trait ModelViewProvider: Send + Sync {
1211 fn apply_model_view(
1212 &self,
1213 messages: Vec<Message>,
1214 config: &serde_json::Value,
1215 context: &ModelViewContext<'_>,
1216 ) -> Vec<Message>;
1217
1218 fn priority(&self) -> i32 {
1219 0
1220 }
1221}
1222
1223pub struct CollectedCapabilities {
1228 pub system_prompt_parts: Vec<String>,
1230 pub system_prompt_attributions: Vec<SystemPromptAttribution>,
1232 pub tools: Vec<Box<dyn Tool>>,
1234 pub tool_definitions: Vec<ToolDefinition>,
1236 pub mounts: Vec<MountPoint>,
1238 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1240 pub applied_ids: Vec<String>,
1242 pub tool_search: Option<crate::llm_driver_registry::ToolSearchConfig>,
1244 pub prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig>,
1246 pub tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>>,
1248 pub tool_call_hooks: Vec<Arc<dyn ToolCallHook>>,
1250 pub mcp_servers: ScopedMcpServers,
1252 }
1258
1259#[derive(Debug, Clone, PartialEq, Eq)]
1260pub struct SystemPromptAttribution {
1261 pub capability_id: String,
1262 pub content: String,
1263}
1264
1265impl CollectedCapabilities {
1266 pub fn system_prompt_prefix(&self) -> Option<String> {
1269 if self.system_prompt_parts.is_empty() {
1270 None
1271 } else {
1272 Some(self.system_prompt_parts.join("\n\n"))
1273 }
1274 }
1275
1276 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1280 for (provider, config) in &self.message_filter_providers {
1282 provider.apply_filters(query, config);
1283 }
1284 }
1285
1286 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1289 for (provider, config) in &self.message_filter_providers {
1290 provider.post_load(messages, config);
1291 }
1292 }
1293
1294 pub fn has_message_filters(&self) -> bool {
1296 !self.message_filter_providers.is_empty()
1297 }
1298}
1299
1300pub struct CollectedMessageFilters {
1307 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1309}
1310
1311pub struct CollectedModelViewProviders {
1313 pub model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)>,
1315}
1316
1317impl CollectedMessageFilters {
1323 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1325 for (provider, config) in &self.message_filter_providers {
1326 provider.apply_filters(query, config);
1327 }
1328 }
1329
1330 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1332 for (provider, config) in &self.message_filter_providers {
1333 provider.post_load(messages, config);
1334 }
1335 }
1336}
1337
1338impl CollectedModelViewProviders {
1339 pub fn apply_model_view(
1341 &self,
1342 mut messages: Vec<Message>,
1343 context: &ModelViewContext<'_>,
1344 ) -> Vec<Message> {
1345 for (provider, config) in &self.model_view_providers {
1346 messages = provider.apply_model_view(messages, config, context);
1347 }
1348 messages
1349 }
1350}
1351
1352pub fn collect_message_filters_only(
1358 capability_configs: &[AgentCapabilityConfig],
1359 registry: &CapabilityRegistry,
1360) -> CollectedMessageFilters {
1361 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1362 Vec::new();
1363
1364 for cap_config in capability_configs {
1365 let cap_id = cap_config.capability_ref.as_str();
1366 if let Some(capability) = registry.get(cap_id) {
1367 if capability.status() != CapabilityStatus::Available {
1368 continue;
1369 }
1370 let effective: &dyn Capability = capability
1373 .resolve_for_model(None)
1374 .unwrap_or_else(|| capability.as_ref());
1375 if let Some(provider) = effective.message_filter_provider() {
1376 message_filter_providers.push((provider, cap_config.config.clone()));
1377 }
1378 }
1379 }
1380
1381 message_filter_providers.sort_by_key(|(p, _)| p.priority());
1382
1383 CollectedMessageFilters {
1384 message_filter_providers,
1385 }
1386}
1387
1388pub fn collect_model_view_providers(
1395 capability_configs: &[AgentCapabilityConfig],
1396 registry: &CapabilityRegistry,
1397 model: Option<&str>,
1398) -> CollectedModelViewProviders {
1399 let mut model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)> = Vec::new();
1400
1401 for cap_config in capability_configs {
1402 let cap_id = cap_config.capability_ref.as_str();
1403 if let Some(capability) = registry.get(cap_id) {
1404 if capability.status() != CapabilityStatus::Available {
1405 continue;
1406 }
1407 let effective: &dyn Capability = capability
1408 .resolve_for_model(model)
1409 .unwrap_or_else(|| capability.as_ref());
1410 if let Some(provider) = effective.model_view_provider() {
1411 model_view_providers.push((provider, cap_config.config.clone()));
1412 }
1413 }
1414 }
1415
1416 model_view_providers.sort_by_key(|(p, _)| p.priority());
1417
1418 CollectedModelViewProviders {
1419 model_view_providers,
1420 }
1421}
1422
1423pub fn collect_capability_mcp_servers(
1424 capability_configs: &[AgentCapabilityConfig],
1425 registry: &CapabilityRegistry,
1426) -> ScopedMcpServers {
1427 let mut servers = ScopedMcpServers::default();
1428
1429 for cap_config in capability_configs {
1430 let cap_id = cap_config.capability_ref.as_str();
1431 if is_declarative_capability(cap_id) {
1432 if let Ok(definition) =
1433 serde_json::from_value::<DeclarativeCapabilityDefinition>(cap_config.config.clone())
1434 {
1435 if definition.status != CapabilityStatus::Available {
1436 continue;
1437 }
1438 if let Some(contributed) = definition.mcp_servers {
1439 servers = merge_scoped_mcp_servers(&servers, &contributed);
1440 }
1441 }
1442 continue;
1443 }
1444 if let Some(capability) = registry.get(cap_id) {
1445 if capability.status() != CapabilityStatus::Available {
1446 continue;
1447 }
1448 servers = merge_scoped_mcp_servers(
1449 &servers,
1450 &capability.mcp_servers_with_config(&cap_config.config),
1451 );
1452 }
1453 }
1454
1455 servers
1456}
1457
1458pub const MAX_RESOLVED_CAPABILITIES: usize = 100;
1465
1466#[derive(Debug, Clone, PartialEq, Eq)]
1468pub enum DependencyError {
1469 CircularDependency {
1471 capability_id: String,
1473 chain: Vec<String>,
1475 },
1476 TooManyCapabilities {
1478 count: usize,
1480 max: usize,
1482 },
1483}
1484
1485impl std::fmt::Display for DependencyError {
1486 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1487 match self {
1488 DependencyError::CircularDependency {
1489 capability_id,
1490 chain,
1491 } => {
1492 write!(
1493 f,
1494 "Circular dependency detected: {} depends on itself via chain: {} -> {}",
1495 capability_id,
1496 chain.join(" -> "),
1497 capability_id
1498 )
1499 }
1500 DependencyError::TooManyCapabilities { count, max } => {
1501 write!(
1502 f,
1503 "Too many capabilities after resolution: {} (max: {})",
1504 count, max
1505 )
1506 }
1507 }
1508 }
1509}
1510
1511impl std::error::Error for DependencyError {}
1512
1513#[derive(Debug, Clone)]
1515pub struct ResolvedCapabilities {
1516 pub resolved_ids: Vec<String>,
1519 pub added_as_dependencies: Vec<String>,
1521 pub user_selected: Vec<String>,
1523}
1524
1525pub fn resolve_dependencies(
1545 selected_ids: &[String],
1546 registry: &CapabilityRegistry,
1547) -> Result<ResolvedCapabilities, DependencyError> {
1548 use std::collections::HashSet;
1549
1550 let user_selected: HashSet<String> = selected_ids.iter().cloned().collect();
1551 let mut resolved: Vec<String> = Vec::new();
1552 let mut resolved_set: HashSet<String> = HashSet::new();
1553 let mut added_as_dependencies: Vec<String> = Vec::new();
1554
1555 for cap_id in selected_ids {
1557 resolve_single_capability(
1558 cap_id,
1559 registry,
1560 &mut resolved,
1561 &mut resolved_set,
1562 &mut added_as_dependencies,
1563 &user_selected,
1564 &mut Vec::new(), )?;
1566 }
1567
1568 if resolved.len() > MAX_RESOLVED_CAPABILITIES {
1570 return Err(DependencyError::TooManyCapabilities {
1571 count: resolved.len(),
1572 max: MAX_RESOLVED_CAPABILITIES,
1573 });
1574 }
1575
1576 Ok(ResolvedCapabilities {
1577 resolved_ids: resolved,
1578 added_as_dependencies,
1579 user_selected: selected_ids.to_vec(),
1580 })
1581}
1582
1583pub fn resolve_capability_configs(
1588 selected_configs: &[AgentCapabilityConfig],
1589 registry: &CapabilityRegistry,
1590) -> Result<Vec<AgentCapabilityConfig>, DependencyError> {
1591 let mut selected_ids: Vec<String> = Vec::new();
1592 for config in selected_configs {
1593 if is_declarative_capability(config.capability_id())
1594 && let Ok(definition) =
1595 serde_json::from_value::<DeclarativeCapabilityDefinition>(config.config.clone())
1596 {
1597 selected_ids.extend(definition.dependencies);
1598 }
1599 selected_ids.push(config.capability_id().to_string());
1600 }
1601 let resolved = resolve_dependencies(&selected_ids, registry)?;
1602
1603 let explicit_configs: std::collections::HashMap<String, serde_json::Value> = selected_configs
1604 .iter()
1605 .map(|config| (config.capability_id().to_string(), config.config.clone()))
1606 .collect();
1607
1608 Ok(resolved
1609 .resolved_ids
1610 .into_iter()
1611 .map(|capability_id| {
1612 explicit_configs
1613 .get(&capability_id)
1614 .cloned()
1615 .map(|config| AgentCapabilityConfig::with_config(capability_id.clone(), config))
1616 .unwrap_or_else(|| AgentCapabilityConfig::new(capability_id))
1617 })
1618 .collect())
1619}
1620
1621fn resolve_single_capability(
1623 cap_id: &str,
1624 registry: &CapabilityRegistry,
1625 resolved: &mut Vec<String>,
1626 resolved_set: &mut std::collections::HashSet<String>,
1627 added_as_dependencies: &mut Vec<String>,
1628 user_selected: &std::collections::HashSet<String>,
1629 visiting: &mut Vec<String>,
1630) -> Result<(), DependencyError> {
1631 if resolved_set.contains(cap_id) {
1633 return Ok(());
1634 }
1635
1636 if visiting.contains(&cap_id.to_string()) {
1638 return Err(DependencyError::CircularDependency {
1639 capability_id: cap_id.to_string(),
1640 chain: visiting.clone(),
1641 });
1642 }
1643
1644 let capability = match registry.get(cap_id) {
1646 Some(cap) => cap,
1647 None => {
1648 if is_declarative_capability(cap_id) && !resolved_set.contains(cap_id) {
1649 resolved.push(cap_id.to_string());
1650 resolved_set.insert(cap_id.to_string());
1651 if !user_selected.contains(cap_id) {
1652 added_as_dependencies.push(cap_id.to_string());
1653 }
1654 }
1655 return Ok(());
1656 }
1657 };
1658
1659 visiting.push(cap_id.to_string());
1661
1662 for dep_id in capability.dependencies() {
1664 resolve_single_capability(
1665 dep_id,
1666 registry,
1667 resolved,
1668 resolved_set,
1669 added_as_dependencies,
1670 user_selected,
1671 visiting,
1672 )?;
1673 }
1674
1675 visiting.pop();
1677
1678 if !resolved_set.contains(cap_id) {
1680 resolved.push(cap_id.to_string());
1681 resolved_set.insert(cap_id.to_string());
1682
1683 if !user_selected.contains(cap_id) {
1685 added_as_dependencies.push(cap_id.to_string());
1686 }
1687 }
1688
1689 Ok(())
1690}
1691
1692pub fn compute_features(capability_ids: &[String], registry: &CapabilityRegistry) -> Vec<String> {
1697 use std::collections::HashSet;
1698
1699 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
1700 Ok(resolved) => resolved.resolved_ids,
1701 Err(_) => capability_ids.to_vec(),
1702 };
1703
1704 let mut seen = HashSet::new();
1705 let mut features = Vec::new();
1706 for cap_id in &resolved_ids {
1707 if let Some(cap) = registry.get(cap_id) {
1708 for feature in cap.features() {
1709 if seen.insert(feature) {
1710 features.push(feature.to_string());
1711 }
1712 }
1713 }
1714 }
1715 features
1716}
1717
1718pub fn get_dependencies(cap_id: &str, registry: &CapabilityRegistry) -> Vec<String> {
1721 registry
1722 .get(cap_id)
1723 .map(|cap| cap.dependencies().iter().map(|s| s.to_string()).collect())
1724 .unwrap_or_default()
1725}
1726
1727pub async fn collect_capabilities(
1743 capability_ids: &[String],
1744 registry: &CapabilityRegistry,
1745 ctx: &SystemPromptContext,
1746) -> CollectedCapabilities {
1747 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
1750 Ok(resolved) => resolved.resolved_ids,
1751 Err(e) => {
1752 tracing::warn!("Failed to resolve capability dependencies: {}", e);
1753 capability_ids.to_vec()
1754 }
1755 };
1756
1757 let configs: Vec<AgentCapabilityConfig> = resolved_ids
1759 .iter()
1760 .map(|id| AgentCapabilityConfig {
1761 capability_ref: CapabilityId::new(id),
1762 config: serde_json::Value::Object(serde_json::Map::new()),
1763 })
1764 .collect();
1765
1766 collect_capabilities_with_configs(&configs, registry, ctx).await
1767}
1768
1769pub async fn collect_capabilities_with_configs(
1780 capability_configs: &[AgentCapabilityConfig],
1781 registry: &CapabilityRegistry,
1782 ctx: &SystemPromptContext,
1783) -> CollectedCapabilities {
1784 let mut system_prompt_parts: Vec<String> = Vec::new();
1785 let mut system_prompt_attributions: Vec<SystemPromptAttribution> = Vec::new();
1786 let mut tools: Vec<Box<dyn Tool>> = Vec::new();
1787 let mut tool_definitions: Vec<ToolDefinition> = Vec::new();
1788 let mut mounts: Vec<MountPoint> = Vec::new();
1789 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1790 Vec::new();
1791 let mut applied_ids: Vec<String> = Vec::new();
1792 let mut tool_search: Option<crate::llm_driver_registry::ToolSearchConfig> = None;
1793 let mut prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig> = None;
1794 let mut tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>> = Vec::new();
1795 let mut tool_call_hooks: Vec<Arc<dyn ToolCallHook>> = Vec::new();
1796 let mut mcp_servers = ScopedMcpServers::default();
1797
1798 for cap_config in capability_configs {
1799 let cap_id = cap_config.capability_ref.as_str();
1800 if is_declarative_capability(cap_id) {
1801 match serde_json::from_value::<DeclarativeCapabilityDefinition>(
1802 cap_config.config.clone(),
1803 ) {
1804 Ok(definition) => {
1805 if definition.status != CapabilityStatus::Available {
1806 continue;
1807 }
1808
1809 if let Some(prompt) = definition.system_prompt.as_deref() {
1810 let contribution =
1811 format!("<capability id=\"{}\">\n{}\n</capability>", cap_id, prompt);
1812 system_prompt_attributions.push(SystemPromptAttribution {
1813 capability_id: cap_id.to_string(),
1814 content: contribution.clone(),
1815 });
1816 system_prompt_parts.push(contribution);
1817 }
1818
1819 mounts.extend(definition.mounts(cap_id));
1820 if let Some(ref servers) = definition.mcp_servers {
1821 mcp_servers = merge_scoped_mcp_servers(&mcp_servers, servers);
1822 }
1823 for skill in definition.skill_contributions() {
1824 mounts.push(skill.to_mount(cap_id));
1825 }
1826
1827 applied_ids.push(cap_id.to_string());
1828 }
1829 Err(error) => {
1830 tracing::warn!(
1831 capability_id = %cap_id,
1832 error = %error,
1833 "Skipping invalid declarative capability config"
1834 );
1835 }
1836 }
1837 continue;
1838 }
1839 if let Some(capability) = registry.get(cap_id) {
1840 if capability.status() != CapabilityStatus::Available {
1842 continue;
1843 }
1844
1845 let effective: &dyn Capability =
1857 match capability.resolve_for_model(ctx.model.as_deref()) {
1858 Some(inner) => inner,
1859 None => capability.as_ref(),
1860 };
1861 let effective_id = effective.id();
1862
1863 if let Some(contribution) = effective
1865 .system_prompt_contribution_with_config(ctx, &cap_config.config)
1866 .await
1867 {
1868 system_prompt_attributions.push(SystemPromptAttribution {
1869 capability_id: cap_id.to_string(),
1870 content: contribution.clone(),
1871 });
1872 system_prompt_parts.push(contribution);
1873 }
1874
1875 tools.extend(effective.tools_with_config(&cap_config.config));
1877 tool_definition_hooks
1878 .extend(effective.tool_definition_hooks_with_config(&cap_config.config));
1879 tool_call_hooks.extend(effective.tool_call_hooks());
1880 let cap_category = effective.category();
1885 for def in effective.tool_definitions() {
1886 let def = match (def.category(), cap_category) {
1887 (None, Some(cat)) => def.with_category(cat),
1888 _ => def,
1889 }
1890 .with_capability_attribution(cap_id, Some(capability.name()));
1891 tool_definitions.push(def);
1892 }
1893
1894 if effective_id == OPENAI_TOOL_SEARCH_CAPABILITY_ID {
1899 let threshold = cap_config
1901 .config
1902 .get("threshold")
1903 .and_then(|v| v.as_u64())
1904 .map(|v| v as usize)
1905 .unwrap_or(DEFAULT_TOOL_SEARCH_THRESHOLD);
1906 tool_search = Some(crate::llm_driver_registry::ToolSearchConfig {
1907 enabled: true,
1908 threshold,
1909 });
1910 }
1911
1912 if cap_id == PROMPT_CACHING_CAPABILITY_ID {
1913 let strategy = cap_config
1914 .config
1915 .get("strategy")
1916 .and_then(|v| v.as_str())
1917 .map(|value| match value {
1918 "auto" => crate::llm_driver_registry::PromptCacheStrategy::Auto,
1919 _ => crate::llm_driver_registry::PromptCacheStrategy::Auto,
1920 })
1921 .unwrap_or(crate::llm_driver_registry::PromptCacheStrategy::Auto);
1922 let gemini_cached_content = cap_config
1923 .config
1924 .get("gemini_cached_content")
1925 .and_then(|v| v.as_str())
1926 .map(str::to_string);
1927 prompt_cache = Some(crate::llm_driver_registry::PromptCacheConfig {
1928 enabled: true,
1929 strategy,
1930 gemini_cached_content,
1931 });
1932 }
1933
1934 mounts.extend(effective.mounts());
1936
1937 mcp_servers = merge_scoped_mcp_servers(
1938 &mcp_servers,
1939 &effective.mcp_servers_with_config(&cap_config.config),
1940 );
1941
1942 for skill in effective.contribute_skills() {
1946 mounts.push(skill.to_mount(cap_id));
1947 }
1948
1949 if let Some(provider) = effective.message_filter_provider() {
1951 message_filter_providers.push((provider, cap_config.config.clone()));
1952 }
1953
1954 applied_ids.push(cap_id.to_string());
1955 }
1956 }
1957
1958 if !applied_ids
1970 .iter()
1971 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID)
1972 && tool_definitions
1973 .iter()
1974 .any(|def| def.hints().supports_background == Some(true))
1975 && let Some(bg_cap) = registry.get(BACKGROUND_EXECUTION_CAPABILITY_ID)
1976 && bg_cap.status() == CapabilityStatus::Available
1977 {
1978 tools.extend(bg_cap.tools());
1979 let cap_category = bg_cap.category();
1980 for def in bg_cap.tool_definitions() {
1981 let def = match (def.category(), cap_category) {
1982 (None, Some(cat)) => def.with_category(cat),
1983 _ => def,
1984 }
1985 .with_capability_attribution(BACKGROUND_EXECUTION_CAPABILITY_ID, Some(bg_cap.name()));
1986 tool_definitions.push(def);
1987 }
1988 applied_ids.push(BACKGROUND_EXECUTION_CAPABILITY_ID.to_string());
1989 }
1990
1991 message_filter_providers.sort_by_key(|(p, _)| p.priority());
1993
1994 CollectedCapabilities {
1995 system_prompt_parts,
1996 system_prompt_attributions,
1997 tools,
1998 tool_definitions,
1999 mounts,
2000 message_filter_providers,
2001 applied_ids,
2002 tool_search,
2003 prompt_cache,
2004 tool_definition_hooks,
2005 tool_call_hooks,
2006 mcp_servers,
2007 }
2008}
2009
2010pub struct AppliedCapabilities {
2016 pub runtime_agent: RuntimeAgent,
2018 pub tool_registry: ToolRegistry,
2020 pub applied_ids: Vec<String>,
2022}
2023
2024pub async fn apply_capabilities(
2061 base_runtime_agent: RuntimeAgent,
2062 capability_ids: &[String],
2063 registry: &CapabilityRegistry,
2064 ctx: &SystemPromptContext,
2065) -> AppliedCapabilities {
2066 let collected = collect_capabilities(capability_ids, registry, ctx).await;
2067
2068 let final_system_prompt = match collected.system_prompt_prefix() {
2070 Some(prefix) => format!(
2071 "{}\n\n<system-prompt>\n{}\n</system-prompt>",
2072 prefix, base_runtime_agent.system_prompt
2073 ),
2074 None => base_runtime_agent.system_prompt,
2075 };
2076
2077 let mut tool_registry = ToolRegistry::new();
2079 for tool in collected.tools {
2080 tool_registry.register_boxed(tool);
2081 }
2082
2083 let mut tools = collected.tool_definitions;
2085 for hook in &collected.tool_definition_hooks {
2086 tools = hook.transform(tools);
2087 }
2088
2089 let runtime_agent = RuntimeAgent {
2090 system_prompt: final_system_prompt,
2091 model: base_runtime_agent.model,
2092 tools,
2093 max_iterations: base_runtime_agent.max_iterations,
2094 temperature: base_runtime_agent.temperature,
2095 max_tokens: base_runtime_agent.max_tokens,
2096 tool_search: collected.tool_search,
2097 prompt_cache: collected.prompt_cache,
2098 network_access: base_runtime_agent.network_access,
2099 };
2100
2101 AppliedCapabilities {
2102 runtime_agent,
2103 tool_registry,
2104 applied_ids: collected.applied_ids,
2105 }
2106}
2107
2108#[cfg(test)]
2113mod tests {
2114 use super::*;
2115 use crate::typed_id::SessionId;
2116 use std::collections::BTreeSet;
2117 use uuid::Uuid;
2118
2119 static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2121
2122 fn lock_env() -> std::sync::MutexGuard<'static, ()> {
2123 ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner())
2124 }
2125
2126 fn test_ctx() -> SystemPromptContext {
2128 SystemPromptContext::without_file_store(SessionId::new())
2129 }
2130
2131 fn expected_core_builtin_ids() -> BTreeSet<&'static str> {
2133 let mut ids = [
2134 "agent_instructions",
2135 "human_intent",
2136 "budgeting",
2137 "self_budget",
2138 "noop",
2139 "current_time",
2140 "research",
2141 "platform_management",
2142 "session_file_system",
2143 "workspace_volumes",
2144 "session_storage",
2145 "session",
2146 "session_sql_database",
2147 "test_math",
2148 "test_weather",
2149 "stateless_todo_list",
2150 "web_fetch",
2151 "virtual_bash",
2152 "background_execution",
2153 "session_schedule",
2154 "btw",
2155 "infinity_context",
2156 "compaction",
2157 "memory",
2158 "openai_tool_search",
2159 "tool_search",
2160 "auto_tool_search",
2161 "prompt_caching",
2162 "skills",
2163 "subagents",
2164 "system_commands",
2165 "sample_data",
2166 "data_knowledge",
2167 "knowledge_base",
2168 "tool_output_persistence",
2169 "fake_warehouse",
2170 "fake_aws",
2171 "fake_crm",
2172 "fake_financial",
2173 "loop_detection",
2174 "prompt_canary_guardrail",
2175 "user_hooks",
2176 ]
2177 .into_iter()
2178 .collect::<BTreeSet<_>>();
2179 if cfg!(feature = "ui-capabilities") {
2180 ids.insert("openui");
2181 ids.insert("a2ui");
2182 }
2183 ids
2184 }
2185
2186 fn expected_dev_builtin_ids() -> BTreeSet<&'static str> {
2188 let mut ids = expected_core_builtin_ids();
2189 ids.insert("agent_handoff");
2190 ids.insert("a2a_agent_delegation");
2191 ids
2192 }
2193
2194 fn registry_ids(registry: &CapabilityRegistry) -> BTreeSet<&str> {
2195 registry.capabilities.keys().map(String::as_str).collect()
2196 }
2197
2198 #[test]
2208 fn test_capability_registry_with_builtins_dev() {
2209 let _lock = lock_env();
2211 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2212 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
2213 assert_eq!(registry_ids(®istry), expected_dev_builtin_ids());
2214 assert!(registry.has("agent_handoff"));
2215 assert!(registry.has("a2a_agent_delegation"));
2216 }
2217
2218 #[test]
2219 fn test_capability_registry_with_builtins_prod() {
2220 let _lock = lock_env();
2222 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2223 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2224 assert_eq!(registry_ids(®istry), expected_core_builtin_ids());
2225 assert!(!registry.has("docker_container"));
2227 assert!(!registry.has("agent_handoff"));
2228 assert!(!registry.has("a2a_agent_delegation"));
2229 }
2230
2231 #[test]
2232 fn test_agent_delegation_enabled_by_env_in_prod() {
2233 let _lock = lock_env();
2235 unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "true") };
2236 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2237 assert!(registry.has("agent_handoff"));
2238 assert!(registry.has("a2a_agent_delegation"));
2239 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2240 }
2241
2242 #[test]
2243 fn test_agent_delegation_disabled_by_env_in_dev() {
2244 let _lock = lock_env();
2246 unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "false") };
2247 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
2248 assert!(!registry.has("agent_handoff"));
2249 assert!(!registry.has("a2a_agent_delegation"));
2250 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2251 }
2252
2253 #[test]
2254 fn test_capability_registry_get() {
2255 let registry = CapabilityRegistry::with_builtins();
2256
2257 let noop = registry.get("noop").unwrap();
2258 assert_eq!(noop.id(), "noop");
2259 assert_eq!(noop.name(), "No-Op");
2260 assert_eq!(noop.status(), CapabilityStatus::Available);
2261 }
2262
2263 #[test]
2264 fn test_capability_registry_blueprint_with_capability() {
2265 struct BlueprintProviderCapability;
2266
2267 impl Capability for BlueprintProviderCapability {
2268 fn id(&self) -> &str {
2269 "blueprint_provider"
2270 }
2271 fn name(&self) -> &str {
2272 "Blueprint Provider"
2273 }
2274 fn description(&self) -> &str {
2275 "Capability that provides a blueprint for tests"
2276 }
2277 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
2278 vec![AgentBlueprint {
2279 id: "test_blueprint",
2280 name: "Test Blueprint",
2281 description: "Blueprint for capability registry tests",
2282 model: BlueprintModel::Inherit,
2283 system_prompt: "Test prompt",
2284 tools: vec![],
2285 max_turns: None,
2286 config_schema: None,
2287 }]
2288 }
2289 }
2290
2291 let mut registry = CapabilityRegistry::new();
2292 registry.register(BlueprintProviderCapability);
2293
2294 let (capability_id, blueprint) = registry
2295 .blueprint_with_capability("test_blueprint")
2296 .expect("blueprint should resolve with capability id");
2297 assert_eq!(capability_id, "blueprint_provider");
2298 assert_eq!(blueprint.id, "test_blueprint");
2299 }
2300
2301 #[test]
2302 fn test_capability_registry_builder() {
2303 let registry = CapabilityRegistry::builder()
2304 .capability(NoopCapability)
2305 .capability(CurrentTimeCapability)
2306 .build();
2307
2308 assert!(registry.has("noop"));
2309 assert!(registry.has("current_time"));
2310 assert_eq!(registry.len(), 2);
2311 }
2312
2313 #[test]
2314 fn test_capability_status() {
2315 let registry = CapabilityRegistry::with_builtins();
2316
2317 let current_time = registry.get("current_time").unwrap();
2318 assert_eq!(current_time.status(), CapabilityStatus::Available);
2319
2320 let research = registry.get("research").unwrap();
2321 assert_eq!(research.status(), CapabilityStatus::ComingSoon);
2322 }
2323
2324 #[test]
2325 fn test_capability_icons_and_categories() {
2326 let registry = CapabilityRegistry::with_builtins();
2327
2328 let noop = registry.get("noop").unwrap();
2329 assert_eq!(noop.icon(), Some("circle-off"));
2330 assert_eq!(noop.category(), Some("Testing"));
2331
2332 let current_time = registry.get("current_time").unwrap();
2333 assert_eq!(current_time.icon(), Some("clock"));
2334 assert_eq!(current_time.category(), Some("Utilities"));
2335 }
2336
2337 #[test]
2338 fn test_system_prompt_preview_default_delegates_to_addition() {
2339 let registry = CapabilityRegistry::with_builtins();
2340
2341 let test_math = registry.get("test_math").unwrap();
2343 assert_eq!(
2344 test_math.system_prompt_preview().as_deref(),
2345 test_math.system_prompt_addition()
2346 );
2347
2348 let current_time = registry.get("current_time").unwrap();
2350 assert!(current_time.system_prompt_preview().is_none());
2351 assert!(current_time.system_prompt_addition().is_none());
2352 }
2353
2354 #[test]
2355 fn test_system_prompt_preview_dynamic_capability() {
2356 let registry = CapabilityRegistry::with_builtins();
2357 let cap = registry.get("agent_instructions").unwrap();
2358
2359 assert!(cap.system_prompt_addition().is_none());
2361 assert!(cap.system_prompt_preview().is_some());
2362 assert!(cap.system_prompt_preview().unwrap().contains("AGENTS.md"));
2363 }
2364
2365 #[tokio::test]
2370 async fn test_apply_capabilities_empty() {
2371 let registry = CapabilityRegistry::with_builtins();
2372 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2373
2374 let applied =
2375 apply_capabilities(base_runtime_agent.clone(), &[], ®istry, &test_ctx()).await;
2376
2377 assert_eq!(
2378 applied.runtime_agent.system_prompt,
2379 base_runtime_agent.system_prompt
2380 );
2381 assert!(applied.tool_registry.is_empty());
2382 assert!(applied.applied_ids.is_empty());
2383 }
2384
2385 #[tokio::test]
2386 async fn test_apply_capabilities_noop() {
2387 let registry = CapabilityRegistry::with_builtins();
2388 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2389
2390 let applied = apply_capabilities(
2391 base_runtime_agent.clone(),
2392 &["noop".to_string()],
2393 ®istry,
2394 &test_ctx(),
2395 )
2396 .await;
2397
2398 assert_eq!(
2400 applied.runtime_agent.system_prompt,
2401 base_runtime_agent.system_prompt
2402 );
2403 assert!(applied.tool_registry.is_empty());
2404 assert_eq!(applied.applied_ids, vec!["noop"]);
2405 }
2406
2407 #[tokio::test]
2408 async fn test_apply_capabilities_current_time() {
2409 let registry = CapabilityRegistry::with_builtins();
2410 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2411
2412 let applied = apply_capabilities(
2413 base_runtime_agent.clone(),
2414 &["current_time".to_string()],
2415 ®istry,
2416 &test_ctx(),
2417 )
2418 .await;
2419
2420 assert_eq!(
2422 applied.runtime_agent.system_prompt,
2423 base_runtime_agent.system_prompt
2424 );
2425 assert!(applied.tool_registry.has("get_current_time"));
2426 assert_eq!(applied.tool_registry.len(), 1);
2427 assert_eq!(applied.applied_ids, vec!["current_time"]);
2428 }
2429
2430 #[tokio::test]
2431 async fn test_apply_capabilities_skips_coming_soon() {
2432 let registry = CapabilityRegistry::with_builtins();
2433 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2434
2435 let applied = apply_capabilities(
2437 base_runtime_agent.clone(),
2438 &["research".to_string()],
2439 ®istry,
2440 &test_ctx(),
2441 )
2442 .await;
2443
2444 assert_eq!(
2446 applied.runtime_agent.system_prompt,
2447 base_runtime_agent.system_prompt
2448 );
2449 assert!(applied.applied_ids.is_empty()); }
2451
2452 #[tokio::test]
2453 async fn test_apply_capabilities_multiple() {
2454 let registry = CapabilityRegistry::with_builtins();
2455 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2456
2457 let applied = apply_capabilities(
2458 base_runtime_agent.clone(),
2459 &["noop".to_string(), "current_time".to_string()],
2460 ®istry,
2461 &test_ctx(),
2462 )
2463 .await;
2464
2465 assert!(applied.tool_registry.has("get_current_time"));
2466 assert_eq!(applied.applied_ids, vec!["noop", "current_time"]);
2467 }
2468
2469 #[tokio::test]
2470 async fn test_apply_capabilities_preserves_order() {
2471 let registry = CapabilityRegistry::with_builtins();
2472 let base_runtime_agent = RuntimeAgent::new("Base prompt.", "gpt-5.2");
2473
2474 let applied = apply_capabilities(
2476 base_runtime_agent,
2477 &["current_time".to_string(), "noop".to_string()],
2478 ®istry,
2479 &test_ctx(),
2480 )
2481 .await;
2482
2483 assert_eq!(applied.applied_ids, vec!["current_time", "noop"]);
2484 }
2485
2486 #[tokio::test]
2487 async fn test_apply_capabilities_test_math() {
2488 let registry = CapabilityRegistry::with_builtins();
2489 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2490
2491 let applied = apply_capabilities(
2492 base_runtime_agent.clone(),
2493 &["test_math".to_string()],
2494 ®istry,
2495 &test_ctx(),
2496 )
2497 .await;
2498
2499 assert!(
2501 !applied
2502 .runtime_agent
2503 .system_prompt
2504 .contains("<capability id=\"test_math\">")
2505 );
2506 assert!(
2508 applied
2509 .runtime_agent
2510 .system_prompt
2511 .contains("You are a helpful assistant.")
2512 );
2513 assert!(applied.tool_registry.has("add"));
2514 assert!(applied.tool_registry.has("subtract"));
2515 assert!(applied.tool_registry.has("multiply"));
2516 assert!(applied.tool_registry.has("divide"));
2517 assert_eq!(applied.tool_registry.len(), 4);
2518 }
2519
2520 #[tokio::test]
2521 async fn test_apply_capabilities_test_weather() {
2522 let registry = CapabilityRegistry::with_builtins();
2523 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2524
2525 let applied = apply_capabilities(
2526 base_runtime_agent.clone(),
2527 &["test_weather".to_string()],
2528 ®istry,
2529 &test_ctx(),
2530 )
2531 .await;
2532
2533 assert!(
2535 !applied
2536 .runtime_agent
2537 .system_prompt
2538 .contains("<capability id=\"test_weather\">")
2539 );
2540 assert!(applied.tool_registry.has("get_weather"));
2541 assert!(applied.tool_registry.has("get_forecast"));
2542 assert_eq!(applied.tool_registry.len(), 2);
2543 }
2544
2545 #[tokio::test]
2546 async fn test_apply_capabilities_test_math_and_test_weather() {
2547 let registry = CapabilityRegistry::with_builtins();
2548 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2549
2550 let applied = apply_capabilities(
2551 base_runtime_agent.clone(),
2552 &["test_math".to_string(), "test_weather".to_string()],
2553 ®istry,
2554 &test_ctx(),
2555 )
2556 .await;
2557
2558 assert_eq!(applied.tool_registry.len(), 6); assert!(applied.tool_registry.has("add"));
2561 assert!(applied.tool_registry.has("get_weather"));
2562 }
2563
2564 #[tokio::test]
2565 async fn test_apply_capabilities_stateless_todo_list() {
2566 let registry = CapabilityRegistry::with_builtins();
2567 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2568
2569 let applied = apply_capabilities(
2570 base_runtime_agent.clone(),
2571 &["stateless_todo_list".to_string()],
2572 ®istry,
2573 &test_ctx(),
2574 )
2575 .await;
2576
2577 assert!(
2579 applied
2580 .runtime_agent
2581 .system_prompt
2582 .contains("Task Management")
2583 );
2584 assert!(applied.runtime_agent.system_prompt.contains("write_todos"));
2585 assert!(applied.tool_registry.has("write_todos"));
2586 assert_eq!(applied.tool_registry.len(), 1);
2587 }
2588
2589 #[tokio::test]
2590 async fn test_apply_capabilities_web_fetch() {
2591 let registry = CapabilityRegistry::with_builtins();
2592 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2593
2594 let applied = apply_capabilities(
2595 base_runtime_agent.clone(),
2596 &["web_fetch".to_string()],
2597 ®istry,
2598 &test_ctx(),
2599 )
2600 .await;
2601
2602 assert!(
2604 applied
2605 .runtime_agent
2606 .system_prompt
2607 .contains(&base_runtime_agent.system_prompt)
2608 );
2609 assert!(applied.runtime_agent.system_prompt.contains("web_fetch"));
2610 assert!(applied.tool_registry.has("web_fetch"));
2611 assert_eq!(applied.tool_registry.len(), 1);
2612 }
2613
2614 #[tokio::test]
2619 async fn test_xml_tags_wrap_capability_prompts() {
2620 let registry = CapabilityRegistry::with_builtins();
2621 let collected =
2622 collect_capabilities(&["stateless_todo_list".to_string()], ®istry, &test_ctx())
2623 .await;
2624
2625 assert_eq!(collected.system_prompt_parts.len(), 1);
2626 let part = &collected.system_prompt_parts[0];
2627 assert!(part.starts_with("<capability id=\"stateless_todo_list\">"));
2628 assert!(part.ends_with("</capability>"));
2629 assert!(part.contains("Task Management"));
2630 }
2631
2632 #[tokio::test]
2633 async fn test_xml_tags_multiple_capabilities() {
2634 let registry = CapabilityRegistry::with_builtins();
2635 let collected = collect_capabilities(
2636 &[
2637 "stateless_todo_list".to_string(),
2638 "session_schedule".to_string(),
2639 ],
2640 ®istry,
2641 &test_ctx(),
2642 )
2643 .await;
2644
2645 assert_eq!(collected.system_prompt_parts.len(), 2);
2646 assert!(
2647 collected.system_prompt_parts[0].starts_with("<capability id=\"stateless_todo_list\">")
2648 );
2649 assert!(
2650 collected.system_prompt_parts[1].starts_with("<capability id=\"session_schedule\">")
2651 );
2652
2653 let prefix = collected.system_prompt_prefix().unwrap();
2654 assert!(prefix.contains("</capability>\n\n<capability"));
2656 }
2657
2658 #[tokio::test]
2659 async fn test_xml_tags_system_prompt_wrapping() {
2660 let registry = CapabilityRegistry::with_builtins();
2661 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2662
2663 let applied = apply_capabilities(
2664 base,
2665 &["stateless_todo_list".to_string()],
2666 ®istry,
2667 &test_ctx(),
2668 )
2669 .await;
2670
2671 let prompt = &applied.runtime_agent.system_prompt;
2672 assert!(prompt.contains("<capability id=\"stateless_todo_list\">"));
2674 assert!(prompt.contains("</capability>"));
2675 assert!(prompt.contains("<system-prompt>\nYou are helpful.\n</system-prompt>"));
2677 }
2678
2679 #[tokio::test]
2680 async fn test_no_xml_wrapping_without_capabilities() {
2681 let registry = CapabilityRegistry::with_builtins();
2682 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2683
2684 let applied = apply_capabilities(base, &[], ®istry, &test_ctx()).await;
2685
2686 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2688 assert!(
2689 !applied
2690 .runtime_agent
2691 .system_prompt
2692 .contains("<system-prompt>")
2693 );
2694 }
2695
2696 #[tokio::test]
2697 async fn test_no_xml_wrapping_for_noop_capability() {
2698 let registry = CapabilityRegistry::with_builtins();
2699 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2700
2701 let applied = apply_capabilities(base, &["noop".to_string()], ®istry, &test_ctx()).await;
2703
2704 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2705 assert!(
2706 !applied
2707 .runtime_agent
2708 .system_prompt
2709 .contains("<system-prompt>")
2710 );
2711 }
2712
2713 #[tokio::test]
2718 async fn test_collect_capabilities_includes_mounts() {
2719 let registry = CapabilityRegistry::with_builtins();
2720
2721 let collected =
2722 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
2723
2724 assert!(!collected.mounts.is_empty());
2725 assert_eq!(collected.mounts.len(), 1);
2726 assert_eq!(collected.mounts[0].path, "/samples");
2727 assert!(collected.mounts[0].is_readonly());
2728 }
2729
2730 #[tokio::test]
2731 async fn test_collect_capabilities_empty_mounts_by_default() {
2732 let registry = CapabilityRegistry::with_builtins();
2733
2734 let collected =
2736 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
2737
2738 assert!(collected.mounts.is_empty());
2739 }
2740
2741 #[tokio::test]
2742 async fn test_collect_capabilities_combines_mounts() {
2743 let registry = CapabilityRegistry::with_builtins();
2744
2745 let collected = collect_capabilities(
2748 &["sample_data".to_string(), "current_time".to_string()],
2749 ®istry,
2750 &test_ctx(),
2751 )
2752 .await;
2753
2754 assert_eq!(collected.mounts.len(), 1);
2755 assert!(
2757 collected
2758 .applied_ids
2759 .iter()
2760 .any(|id| id == "session_file_system")
2761 );
2762 assert!(collected.applied_ids.iter().any(|id| id == "sample_data"));
2763 assert!(collected.applied_ids.iter().any(|id| id == "current_time"));
2764 }
2765
2766 #[test]
2767 fn test_sample_data_capability() {
2768 let registry = CapabilityRegistry::with_builtins();
2769 let cap = registry.get("sample_data").unwrap();
2770
2771 assert_eq!(cap.id(), "sample_data");
2772 assert_eq!(cap.name(), "Sample Data");
2773 assert_eq!(cap.status(), CapabilityStatus::Available);
2774
2775 assert!(cap.system_prompt_addition().is_some());
2777 assert!(cap.tools().is_empty());
2778
2779 assert!(!cap.mounts().is_empty());
2781 }
2782
2783 #[test]
2788 fn test_resolve_dependencies_empty() {
2789 let registry = CapabilityRegistry::with_builtins();
2790
2791 let resolved = resolve_dependencies(&[], ®istry).unwrap();
2792
2793 assert!(resolved.resolved_ids.is_empty());
2794 assert!(resolved.added_as_dependencies.is_empty());
2795 assert!(resolved.user_selected.is_empty());
2796 }
2797
2798 #[test]
2799 fn test_resolve_dependencies_no_deps() {
2800 let registry = CapabilityRegistry::with_builtins();
2801
2802 let resolved = resolve_dependencies(&["current_time".to_string()], ®istry).unwrap();
2804
2805 assert_eq!(resolved.resolved_ids, vec!["current_time"]);
2806 assert!(resolved.added_as_dependencies.is_empty());
2807 }
2808
2809 #[test]
2810 fn test_resolve_dependencies_with_deps() {
2811 let registry = CapabilityRegistry::with_builtins();
2812
2813 let resolved = resolve_dependencies(&["sample_data".to_string()], ®istry).unwrap();
2815
2816 assert_eq!(resolved.resolved_ids.len(), 2);
2818 let fs_pos = resolved
2819 .resolved_ids
2820 .iter()
2821 .position(|id| id == "session_file_system")
2822 .unwrap();
2823 let sd_pos = resolved
2824 .resolved_ids
2825 .iter()
2826 .position(|id| id == "sample_data")
2827 .unwrap();
2828 assert!(fs_pos < sd_pos, "FileSystem should come before SampleData");
2829
2830 assert_eq!(resolved.added_as_dependencies, vec!["session_file_system"]);
2832 }
2833
2834 #[test]
2835 fn test_resolve_dependencies_already_selected() {
2836 let registry = CapabilityRegistry::with_builtins();
2837
2838 let resolved = resolve_dependencies(
2840 &["session_file_system".to_string(), "sample_data".to_string()],
2841 ®istry,
2842 )
2843 .unwrap();
2844
2845 assert_eq!(resolved.resolved_ids.len(), 2);
2846 assert!(resolved.added_as_dependencies.is_empty());
2848 }
2849
2850 #[test]
2851 fn test_resolve_dependencies_preserves_order() {
2852 let registry = CapabilityRegistry::with_builtins();
2853
2854 let resolved =
2856 resolve_dependencies(&["current_time".to_string(), "noop".to_string()], ®istry)
2857 .unwrap();
2858
2859 assert_eq!(resolved.resolved_ids, vec!["current_time", "noop"]);
2860 }
2861
2862 #[test]
2863 fn test_resolve_dependencies_unknown_capability() {
2864 let registry = CapabilityRegistry::with_builtins();
2865
2866 let resolved =
2868 resolve_dependencies(&["unknown_capability".to_string()], ®istry).unwrap();
2869
2870 assert!(resolved.resolved_ids.is_empty());
2871 }
2872
2873 #[test]
2874 fn test_get_dependencies() {
2875 let registry = CapabilityRegistry::with_builtins();
2876
2877 let deps = get_dependencies("sample_data", ®istry);
2879 assert_eq!(deps, vec!["session_file_system"]);
2880
2881 let deps = get_dependencies("current_time", ®istry);
2883 assert!(deps.is_empty());
2884
2885 let deps = get_dependencies("unknown", ®istry);
2887 assert!(deps.is_empty());
2888 }
2889
2890 #[test]
2891 fn test_sample_data_has_dependency() {
2892 let registry = CapabilityRegistry::with_builtins();
2893 let cap = registry.get("sample_data").unwrap();
2894
2895 let deps = cap.dependencies();
2896 assert_eq!(deps.len(), 1);
2897 assert_eq!(deps[0], "session_file_system");
2898 }
2899
2900 #[test]
2901 fn test_noop_has_no_dependencies() {
2902 let registry = CapabilityRegistry::with_builtins();
2903 let cap = registry.get("noop").unwrap();
2904
2905 assert!(cap.dependencies().is_empty());
2906 }
2907
2908 #[test]
2912 fn test_circular_dependency_error() {
2913 struct CapA;
2915 struct CapB;
2916
2917 impl Capability for CapA {
2918 fn id(&self) -> &str {
2919 "test_cap_a"
2920 }
2921 fn name(&self) -> &str {
2922 "Test A"
2923 }
2924 fn description(&self) -> &str {
2925 "Test capability A"
2926 }
2927 fn dependencies(&self) -> Vec<&'static str> {
2928 vec!["test_cap_b"]
2929 }
2930 }
2931
2932 impl Capability for CapB {
2933 fn id(&self) -> &str {
2934 "test_cap_b"
2935 }
2936 fn name(&self) -> &str {
2937 "Test B"
2938 }
2939 fn description(&self) -> &str {
2940 "Test capability B"
2941 }
2942 fn dependencies(&self) -> Vec<&'static str> {
2943 vec!["test_cap_a"]
2944 }
2945 }
2946
2947 let mut registry = CapabilityRegistry::new();
2948 registry.register(CapA);
2949 registry.register(CapB);
2950
2951 let result = resolve_dependencies(&["test_cap_a".to_string()], ®istry);
2952
2953 assert!(result.is_err());
2954 match result.unwrap_err() {
2955 DependencyError::CircularDependency { capability_id, .. } => {
2956 assert_eq!(capability_id, "test_cap_a");
2957 }
2958 _ => panic!("Expected CircularDependency error"),
2959 }
2960 }
2961
2962 use crate::message_filter::{MessageFilter, MessageFilterProvider, MessageQuery};
2967
2968 struct FilterTestCapability {
2970 priority: i32,
2971 }
2972
2973 impl Capability for FilterTestCapability {
2974 fn id(&self) -> &str {
2975 "filter_test"
2976 }
2977 fn name(&self) -> &str {
2978 "Filter Test"
2979 }
2980 fn description(&self) -> &str {
2981 "Test capability with message filter"
2982 }
2983 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
2984 Some(Arc::new(FilterTestProvider {
2985 priority: self.priority,
2986 }))
2987 }
2988 }
2989
2990 struct FilterTestProvider {
2991 priority: i32,
2992 }
2993
2994 impl MessageFilterProvider for FilterTestProvider {
2995 fn apply_filters(&self, query: &mut MessageQuery, config: &serde_json::Value) {
2996 if let Some(search) = config.get("search").and_then(|v| v.as_str()) {
2998 query
2999 .filters
3000 .push(MessageFilter::Search(search.to_string()));
3001 }
3002 }
3003
3004 fn priority(&self) -> i32 {
3005 self.priority
3006 }
3007 }
3008
3009 #[tokio::test]
3010 async fn test_collect_capabilities_with_configs_no_filter_providers() {
3011 let registry = CapabilityRegistry::with_builtins();
3012 let configs = vec![AgentCapabilityConfig {
3013 capability_ref: CapabilityId::new("current_time"),
3014 config: serde_json::json!({}),
3015 }];
3016
3017 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3018
3019 assert!(collected.message_filter_providers.is_empty());
3020 assert!(!collected.has_message_filters());
3021 }
3022
3023 #[tokio::test]
3024 async fn test_collect_capabilities_with_configs_with_filter_provider() {
3025 let mut registry = CapabilityRegistry::new();
3026 registry.register(FilterTestCapability { priority: 0 });
3027
3028 let configs = vec![AgentCapabilityConfig {
3029 capability_ref: CapabilityId::new("filter_test"),
3030 config: serde_json::json!({ "search": "hello" }),
3031 }];
3032
3033 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3034
3035 assert_eq!(collected.message_filter_providers.len(), 1);
3036 assert!(collected.has_message_filters());
3037 }
3038
3039 #[tokio::test]
3040 async fn test_collect_capabilities_with_configs_filter_priority_order() {
3041 struct HighPriorityCapability;
3043 struct LowPriorityCapability;
3044
3045 impl Capability for HighPriorityCapability {
3046 fn id(&self) -> &str {
3047 "high_priority"
3048 }
3049 fn name(&self) -> &str {
3050 "High Priority"
3051 }
3052 fn description(&self) -> &str {
3053 "Test"
3054 }
3055 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3056 Some(Arc::new(FilterTestProvider { priority: 10 }))
3057 }
3058 }
3059
3060 impl Capability for LowPriorityCapability {
3061 fn id(&self) -> &str {
3062 "low_priority"
3063 }
3064 fn name(&self) -> &str {
3065 "Low Priority"
3066 }
3067 fn description(&self) -> &str {
3068 "Test"
3069 }
3070 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3071 Some(Arc::new(FilterTestProvider { priority: -5 }))
3072 }
3073 }
3074
3075 let mut registry = CapabilityRegistry::new();
3076 registry.register(HighPriorityCapability);
3077 registry.register(LowPriorityCapability);
3078
3079 let configs = vec![
3081 AgentCapabilityConfig {
3082 capability_ref: CapabilityId::new("high_priority"),
3083 config: serde_json::json!({}),
3084 },
3085 AgentCapabilityConfig {
3086 capability_ref: CapabilityId::new("low_priority"),
3087 config: serde_json::json!({}),
3088 },
3089 ];
3090
3091 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3092
3093 assert_eq!(collected.message_filter_providers.len(), 2);
3095 assert_eq!(collected.message_filter_providers[0].0.priority(), -5);
3096 assert_eq!(collected.message_filter_providers[1].0.priority(), 10);
3097 }
3098
3099 #[tokio::test]
3100 async fn test_collected_capabilities_apply_message_filters() {
3101 let mut registry = CapabilityRegistry::new();
3102 registry.register(FilterTestCapability { priority: 0 });
3103
3104 let configs = vec![AgentCapabilityConfig {
3105 capability_ref: CapabilityId::new("filter_test"),
3106 config: serde_json::json!({ "search": "test_query" }),
3107 }];
3108
3109 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3110
3111 let session_id: SessionId = Uuid::now_v7().into();
3113 let mut query = MessageQuery::new(session_id);
3114
3115 collected.apply_message_filters(&mut query);
3116
3117 assert_eq!(query.filters.len(), 1);
3119 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3120 }
3121
3122 #[tokio::test]
3123 async fn test_collected_capabilities_apply_multiple_filters_in_priority_order() {
3124 struct SearchCapability {
3125 id: &'static str,
3126 search_term: &'static str,
3127 priority: i32,
3128 }
3129
3130 struct SearchProvider {
3131 search_term: &'static str,
3132 priority: i32,
3133 }
3134
3135 impl MessageFilterProvider for SearchProvider {
3136 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3137 query
3138 .filters
3139 .push(MessageFilter::Search(self.search_term.to_string()));
3140 }
3141
3142 fn priority(&self) -> i32 {
3143 self.priority
3144 }
3145 }
3146
3147 impl Capability for SearchCapability {
3148 fn id(&self) -> &str {
3149 self.id
3150 }
3151 fn name(&self) -> &str {
3152 "Search"
3153 }
3154 fn description(&self) -> &str {
3155 "Test"
3156 }
3157 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3158 Some(Arc::new(SearchProvider {
3159 search_term: self.search_term,
3160 priority: self.priority,
3161 }))
3162 }
3163 }
3164
3165 let mut registry = CapabilityRegistry::new();
3166 registry.register(SearchCapability {
3167 id: "cap_a",
3168 search_term: "alpha",
3169 priority: 5,
3170 });
3171 registry.register(SearchCapability {
3172 id: "cap_b",
3173 search_term: "beta",
3174 priority: 1,
3175 });
3176 registry.register(SearchCapability {
3177 id: "cap_c",
3178 search_term: "gamma",
3179 priority: 10,
3180 });
3181
3182 let configs = vec![
3183 AgentCapabilityConfig {
3184 capability_ref: CapabilityId::new("cap_a"),
3185 config: serde_json::json!({}),
3186 },
3187 AgentCapabilityConfig {
3188 capability_ref: CapabilityId::new("cap_b"),
3189 config: serde_json::json!({}),
3190 },
3191 AgentCapabilityConfig {
3192 capability_ref: CapabilityId::new("cap_c"),
3193 config: serde_json::json!({}),
3194 },
3195 ];
3196
3197 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3198
3199 let session_id: SessionId = Uuid::now_v7().into();
3200 let mut query = MessageQuery::new(session_id);
3201
3202 collected.apply_message_filters(&mut query);
3203
3204 assert_eq!(query.filters.len(), 3);
3206 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3207 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3208 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3209 }
3210
3211 #[test]
3212 fn test_capability_without_message_filter_returns_none() {
3213 let registry = CapabilityRegistry::with_builtins();
3214
3215 let noop = registry.get("noop").unwrap();
3216 assert!(noop.message_filter_provider().is_none());
3217
3218 let current_time = registry.get("current_time").unwrap();
3219 assert!(current_time.message_filter_provider().is_none());
3220 }
3221
3222 #[tokio::test]
3223 async fn test_collect_capabilities_preserves_config_for_filter_provider() {
3224 let mut registry = CapabilityRegistry::new();
3225 registry.register(FilterTestCapability { priority: 0 });
3226
3227 let test_config = serde_json::json!({
3228 "search": "custom_search",
3229 "extra_field": 42
3230 });
3231
3232 let configs = vec![AgentCapabilityConfig {
3233 capability_ref: CapabilityId::new("filter_test"),
3234 config: test_config.clone(),
3235 }];
3236
3237 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3238
3239 assert_eq!(collected.message_filter_providers.len(), 1);
3241 let (_, stored_config) = &collected.message_filter_providers[0];
3242 assert_eq!(*stored_config, test_config);
3243 }
3244
3245 #[test]
3250 fn test_collect_message_filters_only_collects_filters() {
3251 let mut registry = CapabilityRegistry::new();
3252 registry.register(FilterTestCapability { priority: 0 });
3253
3254 let configs = vec![AgentCapabilityConfig {
3255 capability_ref: CapabilityId::new("filter_test"),
3256 config: serde_json::json!({ "search": "test_query" }),
3257 }];
3258
3259 let collected = collect_message_filters_only(&configs, ®istry);
3260
3261 let session_id: SessionId = Uuid::now_v7().into();
3262 let mut query = MessageQuery::new(session_id);
3263 collected.apply_message_filters(&mut query);
3264
3265 assert_eq!(query.filters.len(), 1);
3266 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3267 }
3268
3269 #[test]
3270 fn test_collect_message_filters_only_skips_unknown_capabilities() {
3271 let registry = CapabilityRegistry::new();
3272
3273 let configs = vec![AgentCapabilityConfig {
3274 capability_ref: CapabilityId::new("nonexistent"),
3275 config: serde_json::json!({}),
3276 }];
3277
3278 let collected = collect_message_filters_only(&configs, ®istry);
3279 assert!(collected.message_filter_providers.is_empty());
3280 }
3281
3282 #[test]
3283 fn test_collect_message_filters_only_preserves_priority_order() {
3284 struct PriorityFilterCap {
3285 id: &'static str,
3286 search_term: &'static str,
3287 priority: i32,
3288 }
3289
3290 struct PriorityFilterProvider {
3291 search_term: &'static str,
3292 priority: i32,
3293 }
3294
3295 impl Capability for PriorityFilterCap {
3296 fn id(&self) -> &str {
3297 self.id
3298 }
3299 fn name(&self) -> &str {
3300 self.id
3301 }
3302 fn description(&self) -> &str {
3303 "priority test"
3304 }
3305 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3306 Some(Arc::new(PriorityFilterProvider {
3307 search_term: self.search_term,
3308 priority: self.priority,
3309 }))
3310 }
3311 }
3312
3313 impl MessageFilterProvider for PriorityFilterProvider {
3314 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3315 query
3316 .filters
3317 .push(MessageFilter::Search(self.search_term.to_string()));
3318 }
3319 fn priority(&self) -> i32 {
3320 self.priority
3321 }
3322 }
3323
3324 let mut registry = CapabilityRegistry::new();
3325 registry.register(PriorityFilterCap {
3326 id: "gamma",
3327 search_term: "gamma",
3328 priority: 10,
3329 });
3330 registry.register(PriorityFilterCap {
3331 id: "alpha",
3332 search_term: "alpha",
3333 priority: 5,
3334 });
3335 registry.register(PriorityFilterCap {
3336 id: "beta",
3337 search_term: "beta",
3338 priority: 1,
3339 });
3340
3341 let configs = vec![
3342 AgentCapabilityConfig {
3343 capability_ref: CapabilityId::new("gamma"),
3344 config: serde_json::json!({}),
3345 },
3346 AgentCapabilityConfig {
3347 capability_ref: CapabilityId::new("alpha"),
3348 config: serde_json::json!({}),
3349 },
3350 AgentCapabilityConfig {
3351 capability_ref: CapabilityId::new("beta"),
3352 config: serde_json::json!({}),
3353 },
3354 ];
3355
3356 let collected = collect_message_filters_only(&configs, ®istry);
3357
3358 let session_id: SessionId = Uuid::now_v7().into();
3359 let mut query = MessageQuery::new(session_id);
3360 collected.apply_message_filters(&mut query);
3361
3362 assert_eq!(query.filters.len(), 3);
3364 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3365 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3366 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3367 }
3368
3369 #[test]
3370 fn test_collect_message_filters_only_post_load_invoked() {
3371 use crate::message::Message;
3372
3373 struct PostLoadCap;
3374 struct PostLoadProvider;
3375
3376 impl Capability for PostLoadCap {
3377 fn id(&self) -> &str {
3378 "post_load_test"
3379 }
3380 fn name(&self) -> &str {
3381 "PostLoad Test"
3382 }
3383 fn description(&self) -> &str {
3384 "test"
3385 }
3386 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3387 Some(Arc::new(PostLoadProvider))
3388 }
3389 }
3390
3391 impl MessageFilterProvider for PostLoadProvider {
3392 fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
3393 fn priority(&self) -> i32 {
3394 0
3395 }
3396 fn post_load(&self, messages: &mut Vec<Message>, _config: &serde_json::Value) {
3397 messages.reverse();
3399 }
3400 }
3401
3402 let mut registry = CapabilityRegistry::new();
3403 registry.register(PostLoadCap);
3404
3405 let configs = vec![AgentCapabilityConfig {
3406 capability_ref: CapabilityId::new("post_load_test"),
3407 config: serde_json::json!({}),
3408 }];
3409
3410 let collected = collect_message_filters_only(&configs, ®istry);
3411
3412 let mut messages = vec![Message::user("first"), Message::user("second")];
3413 collected.apply_post_load_filters(&mut messages);
3414
3415 assert_eq!(messages[0].text(), Some("second"));
3417 assert_eq!(messages[1].text(), Some("first"));
3418 }
3419
3420 #[test]
3421 fn test_collect_model_view_providers_respects_compaction_capability_boundary() {
3422 use crate::tool_types::ToolCall;
3423
3424 fn tool_heavy_messages() -> Vec<Message> {
3425 let mut messages = vec![Message::user("inspect files repeatedly")];
3426 for index in 0..9 {
3427 let call_id = format!("call_{index}");
3428 messages.push(Message::assistant_with_tools(
3429 "",
3430 vec![ToolCall {
3431 id: call_id.clone(),
3432 name: "read_file".to_string(),
3433 arguments: serde_json::json!({"path": "/workspace/src/lib.rs"}),
3434 }],
3435 ));
3436 messages.push(Message::tool_result(
3437 call_id,
3438 Some(serde_json::json!({
3439 "path": "/workspace/src/lib.rs",
3440 "content": format!("{}{}", "large file line\n".repeat(1000), index),
3441 "total_lines": 1000,
3442 "lines_shown": {"start": 1, "end": 1000},
3443 "truncated": false
3444 })),
3445 None,
3446 ));
3447 }
3448 messages
3449 }
3450
3451 fn first_tool_result_is_masked(messages: &[Message]) -> bool {
3452 messages[2]
3453 .tool_result_content()
3454 .and_then(|result| result.result.as_ref())
3455 .and_then(|result| result.get("masked"))
3456 .and_then(|masked| masked.as_bool())
3457 .unwrap_or(false)
3458 }
3459
3460 let mut registry = CapabilityRegistry::new();
3461 registry.register(CompactionCapability);
3462 let context = ModelViewContext {
3463 session_id: SessionId::new(),
3464 prior_usage: None,
3465 };
3466
3467 let no_compaction = collect_model_view_providers(&[], ®istry, None);
3468 let unmasked = no_compaction.apply_model_view(tool_heavy_messages(), &context);
3469 assert!(!first_tool_result_is_masked(&unmasked));
3470
3471 let compaction = collect_model_view_providers(
3472 &[AgentCapabilityConfig {
3473 capability_ref: CapabilityId::new(COMPACTION_CAPABILITY_ID),
3474 config: serde_json::json!({}),
3475 }],
3476 ®istry,
3477 None,
3478 );
3479 let masked = compaction.apply_model_view(tool_heavy_messages(), &context);
3480 assert!(first_tool_result_is_masked(&masked));
3481 let last_tool = masked.last().unwrap().tool_result_content().unwrap();
3482 assert!(last_tool.result.as_ref().unwrap().get("content").is_some());
3483 }
3484
3485 struct DelegatingFilterCap {
3488 id: &'static str,
3489 inner: std::sync::Arc<InnerFilterCap>,
3490 }
3491 struct InnerFilterCap;
3492
3493 impl Capability for InnerFilterCap {
3494 fn id(&self) -> &str {
3495 "inner_filter"
3496 }
3497 fn name(&self) -> &str {
3498 "Inner Filter"
3499 }
3500 fn description(&self) -> &str {
3501 "inner"
3502 }
3503 fn message_filter_provider(&self) -> Option<std::sync::Arc<dyn MessageFilterProvider>> {
3504 Some(std::sync::Arc::new(SentinelFilter))
3505 }
3506 }
3507 struct SentinelFilter;
3508 impl MessageFilterProvider for SentinelFilter {
3509 fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
3510 }
3511 impl Capability for DelegatingFilterCap {
3512 fn id(&self) -> &str {
3513 self.id
3514 }
3515 fn name(&self) -> &str {
3516 "Delegating Filter"
3517 }
3518 fn description(&self) -> &str {
3519 "delegating"
3520 }
3521 fn message_filter_provider(&self) -> Option<std::sync::Arc<dyn MessageFilterProvider>> {
3522 None }
3524 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
3525 Some(&*self.inner)
3526 }
3527 }
3528
3529 #[test]
3530 fn test_collect_message_filters_only_honors_resolve_for_model_delegation() {
3531 let inner = std::sync::Arc::new(InnerFilterCap);
3532 let outer = DelegatingFilterCap {
3533 id: "delegating_filter",
3534 inner: inner.clone(),
3535 };
3536
3537 let mut registry = CapabilityRegistry::new();
3538 registry.register(outer);
3539
3540 let configs = vec![AgentCapabilityConfig {
3541 capability_ref: CapabilityId::new("delegating_filter"),
3542 config: serde_json::json!({}),
3543 }];
3544
3545 let collected = collect_message_filters_only(&configs, ®istry);
3548 assert_eq!(
3549 collected.message_filter_providers.len(),
3550 1,
3551 "provider from resolved inner capability must be collected"
3552 );
3553 }
3554
3555 struct DelegatingMvpCap {
3556 id: &'static str,
3557 inner: std::sync::Arc<InnerMvpCap>,
3558 }
3559 struct InnerMvpCap;
3560
3561 impl Capability for InnerMvpCap {
3562 fn id(&self) -> &str {
3563 "inner_mvp"
3564 }
3565 fn name(&self) -> &str {
3566 "Inner MVP"
3567 }
3568 fn description(&self) -> &str {
3569 "inner"
3570 }
3571 fn model_view_provider(
3572 &self,
3573 ) -> Option<std::sync::Arc<dyn crate::capabilities::ModelViewProvider>> {
3574 struct NoopMvp;
3576 impl crate::capabilities::ModelViewProvider for NoopMvp {
3577 fn apply_model_view(
3578 &self,
3579 messages: Vec<Message>,
3580 _config: &serde_json::Value,
3581 _context: &ModelViewContext<'_>,
3582 ) -> Vec<Message> {
3583 messages
3584 }
3585 }
3586 Some(std::sync::Arc::new(NoopMvp))
3587 }
3588 }
3589 impl Capability for DelegatingMvpCap {
3590 fn id(&self) -> &str {
3591 self.id
3592 }
3593 fn name(&self) -> &str {
3594 "Delegating MVP"
3595 }
3596 fn description(&self) -> &str {
3597 "delegating"
3598 }
3599 fn model_view_provider(
3600 &self,
3601 ) -> Option<std::sync::Arc<dyn crate::capabilities::ModelViewProvider>> {
3602 None }
3604 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
3605 Some(&*self.inner)
3606 }
3607 }
3608
3609 #[test]
3610 fn test_collect_model_view_providers_honors_resolve_for_model_delegation() {
3611 let inner = std::sync::Arc::new(InnerMvpCap);
3612 let outer = DelegatingMvpCap {
3613 id: "delegating_mvp",
3614 inner: inner.clone(),
3615 };
3616
3617 let mut registry = CapabilityRegistry::new();
3618 registry.register(outer);
3619
3620 let configs = vec![AgentCapabilityConfig {
3621 capability_ref: CapabilityId::new("delegating_mvp"),
3622 config: serde_json::json!({}),
3623 }];
3624
3625 let collected = collect_model_view_providers(&configs, ®istry, None);
3628 assert_eq!(
3629 collected.model_view_providers.len(),
3630 1,
3631 "provider from resolved inner capability must be collected"
3632 );
3633 }
3634
3635 #[tokio::test]
3645 async fn test_virtual_bash_capability_produces_bash_tool() {
3646 let registry = CapabilityRegistry::with_builtins();
3647 let collected =
3648 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3649
3650 let tool_names: Vec<&str> = collected
3651 .tool_definitions
3652 .iter()
3653 .map(|t| t.name())
3654 .collect();
3655 assert!(
3656 tool_names.contains(&"bash"),
3657 "virtual_bash capability must produce 'bash' tool, got: {:?}",
3658 tool_names
3659 );
3660 assert!(
3661 !collected.tools.is_empty(),
3662 "virtual_bash must provide tool implementations"
3663 );
3664 }
3665
3666 #[tokio::test]
3667 async fn test_generic_harness_capability_set_produces_bash_tool() {
3668 let generic_harness_caps = vec![
3671 "session_file_system".to_string(),
3672 "virtual_bash".to_string(),
3673 "web_fetch".to_string(),
3674 "session_storage".to_string(),
3675 "session".to_string(),
3676 "agent_instructions".to_string(),
3677 "skills".to_string(),
3678 "infinity_context".to_string(),
3679 "auto_tool_search".to_string(),
3680 ];
3681
3682 let registry = CapabilityRegistry::with_builtins();
3683 let collected = collect_capabilities(&generic_harness_caps, ®istry, &test_ctx()).await;
3684
3685 let tool_names: Vec<&str> = collected
3686 .tool_definitions
3687 .iter()
3688 .map(|t| t.name())
3689 .collect();
3690 assert!(
3691 tool_names.contains(&"bash"),
3692 "Generic Harness capabilities must produce 'bash' tool, got: {:?}",
3693 tool_names
3694 );
3695 }
3696
3697 #[tokio::test]
3698 async fn test_collect_capabilities_tool_count_matches_definitions() {
3699 let registry = CapabilityRegistry::with_builtins();
3702 let collected =
3703 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3704
3705 assert_eq!(
3706 collected.tools.len(),
3707 collected.tool_definitions.len(),
3708 "tool implementations ({}) must match tool definitions ({})",
3709 collected.tools.len(),
3710 collected.tool_definitions.len(),
3711 );
3712 }
3713
3714 #[tokio::test]
3718 async fn test_collect_capabilities_resolves_dependencies() {
3719 let registry = CapabilityRegistry::with_builtins();
3722 let collected =
3723 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
3724
3725 assert!(
3727 collected
3728 .applied_ids
3729 .iter()
3730 .any(|id| id == "session_file_system"),
3731 "collect_capabilities must apply session_file_system as a dependency; applied_ids: {:?}",
3732 collected.applied_ids
3733 );
3734
3735 let tool_names: Vec<&str> = collected
3736 .tool_definitions
3737 .iter()
3738 .map(|t| t.name())
3739 .collect();
3740
3741 assert!(
3743 tool_names.contains(&"read_file") && tool_names.contains(&"write_file"),
3744 "collect_capabilities must resolve dependencies and include dependency tools, got: {:?}",
3745 tool_names
3746 );
3747
3748 assert_eq!(
3750 collected.tools.len(),
3751 collected.tool_definitions.len(),
3752 "dependency-added tools must have implementations, not just definitions"
3753 );
3754 }
3755
3756 #[test]
3757 fn test_defaults_do_not_include_bash() {
3758 let registry = crate::ToolRegistry::with_defaults();
3761 assert!(
3762 !registry.has("bash"),
3763 "with_defaults() must not include 'bash' — it comes from virtual_bash capability"
3764 );
3765 }
3766
3767 #[tokio::test]
3774 async fn test_background_execution_auto_activates_with_virtual_bash() {
3775 let registry = CapabilityRegistry::with_builtins();
3776 let collected =
3777 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3778
3779 let tool_names: Vec<&str> = collected
3780 .tool_definitions
3781 .iter()
3782 .map(|t| t.name())
3783 .collect();
3784 assert!(
3785 tool_names.contains(&"spawn_background"),
3786 "spawn_background must be auto-activated when virtual_bash (a \
3787 background-capable tool) is in the agent's capability set; got: {:?}",
3788 tool_names
3789 );
3790 assert!(
3791 collected
3792 .applied_ids
3793 .iter()
3794 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
3795 "background_execution must be in applied_ids when auto-activated; \
3796 got: {:?}",
3797 collected.applied_ids
3798 );
3799
3800 assert!(
3802 collected
3803 .tools
3804 .iter()
3805 .any(|t| t.name() == "spawn_background"),
3806 "spawn_background tool implementation must be present alongside the \
3807 definition (lockstep contract)"
3808 );
3809 }
3810
3811 #[tokio::test]
3814 async fn test_background_execution_does_not_auto_activate_without_hint() {
3815 let registry = CapabilityRegistry::with_builtins();
3816 let collected =
3818 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
3819
3820 let tool_names: Vec<&str> = collected
3821 .tool_definitions
3822 .iter()
3823 .map(|t| t.name())
3824 .collect();
3825 assert!(
3826 !tool_names.contains(&"spawn_background"),
3827 "spawn_background must NOT be activated without a background-capable \
3828 tool; got: {:?}",
3829 tool_names
3830 );
3831 assert!(
3832 !collected
3833 .applied_ids
3834 .iter()
3835 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
3836 "background_execution must not appear in applied_ids when no \
3837 background-capable tool is present; got: {:?}",
3838 collected.applied_ids
3839 );
3840 }
3841
3842 #[tokio::test]
3846 async fn test_background_execution_explicit_selection_is_idempotent() {
3847 let registry = CapabilityRegistry::with_builtins();
3848 let collected = collect_capabilities(
3849 &[
3850 "virtual_bash".to_string(),
3851 BACKGROUND_EXECUTION_CAPABILITY_ID.to_string(),
3852 ],
3853 ®istry,
3854 &test_ctx(),
3855 )
3856 .await;
3857
3858 let spawn_background_count = collected
3859 .tool_definitions
3860 .iter()
3861 .filter(|t| t.name() == "spawn_background")
3862 .count();
3863 assert_eq!(
3864 spawn_background_count, 1,
3865 "spawn_background must appear exactly once even when \
3866 background_execution is selected explicitly alongside a \
3867 background-capable tool"
3868 );
3869 let applied_count = collected
3870 .applied_ids
3871 .iter()
3872 .filter(|id| id.as_str() == BACKGROUND_EXECUTION_CAPABILITY_ID)
3873 .count();
3874 assert_eq!(
3875 applied_count, 1,
3876 "background_execution must appear exactly once in applied_ids"
3877 );
3878 }
3879
3880 #[test]
3885 fn test_defaults_do_not_include_spawn_background() {
3886 let registry = crate::ToolRegistry::with_defaults();
3887 assert!(
3888 !registry.has("spawn_background"),
3889 "with_defaults() must not include 'spawn_background' — it comes \
3890 from the background_execution capability (EVE-501)"
3891 );
3892 }
3893
3894 #[test]
3899 fn test_capability_features_default_empty() {
3900 let registry = CapabilityRegistry::with_builtins();
3901
3902 let noop = registry.get("noop").unwrap();
3904 assert!(noop.features().is_empty());
3905
3906 let current_time = registry.get("current_time").unwrap();
3907 assert!(current_time.features().is_empty());
3908 }
3909
3910 #[test]
3911 fn test_file_system_capability_features() {
3912 let registry = CapabilityRegistry::with_builtins();
3913
3914 let fs = registry.get("session_file_system").unwrap();
3915 assert_eq!(fs.features(), vec!["file_system"]);
3916 }
3917
3918 #[test]
3919 fn test_virtual_bash_capability_features() {
3920 let registry = CapabilityRegistry::with_builtins();
3921
3922 let bash = registry.get("virtual_bash").unwrap();
3923 assert_eq!(bash.features(), vec!["file_system"]);
3924 }
3925
3926 #[test]
3927 fn test_session_storage_capability_features() {
3928 let registry = CapabilityRegistry::with_builtins();
3929
3930 let storage = registry.get("session_storage").unwrap();
3931 let features = storage.features();
3932 assert!(features.contains(&"secrets"));
3933 assert!(features.contains(&"key_value"));
3934 }
3935
3936 #[test]
3937 fn test_session_schedule_capability_features() {
3938 let registry = CapabilityRegistry::with_builtins();
3939
3940 let schedule = registry.get("session_schedule").unwrap();
3941 assert_eq!(schedule.features(), vec!["schedules"]);
3942 }
3943
3944 #[test]
3945 fn test_session_sql_database_capability_features() {
3946 let registry = CapabilityRegistry::with_builtins();
3947
3948 let sql = registry.get("session_sql_database").unwrap();
3949 assert_eq!(sql.features(), vec!["sql_database"]);
3950 }
3951
3952 #[test]
3953 fn test_sample_data_capability_features() {
3954 let registry = CapabilityRegistry::with_builtins();
3955
3956 let sample = registry.get("sample_data").unwrap();
3957 assert_eq!(sample.features(), vec!["file_system"]);
3958 }
3959
3960 #[test]
3961 fn test_compute_features_empty() {
3962 let registry = CapabilityRegistry::with_builtins();
3963
3964 let features = compute_features(&[], ®istry);
3965 assert!(features.is_empty());
3966 }
3967
3968 #[test]
3969 fn test_compute_features_single_capability() {
3970 let registry = CapabilityRegistry::with_builtins();
3971
3972 let features = compute_features(&["session_schedule".to_string()], ®istry);
3973 assert_eq!(features, vec!["schedules"]);
3974 }
3975
3976 #[test]
3977 fn test_compute_features_multiple_capabilities() {
3978 let registry = CapabilityRegistry::with_builtins();
3979
3980 let features = compute_features(
3981 &[
3982 "session_file_system".to_string(),
3983 "session_storage".to_string(),
3984 "session_schedule".to_string(),
3985 ],
3986 ®istry,
3987 );
3988 assert!(features.contains(&"file_system".to_string()));
3989 assert!(features.contains(&"secrets".to_string()));
3990 assert!(features.contains(&"key_value".to_string()));
3991 assert!(features.contains(&"schedules".to_string()));
3992 }
3993
3994 #[test]
3995 fn test_compute_features_deduplicates() {
3996 let registry = CapabilityRegistry::with_builtins();
3997
3998 let features = compute_features(
4000 &[
4001 "session_file_system".to_string(),
4002 "virtual_bash".to_string(),
4003 ],
4004 ®istry,
4005 );
4006 let file_system_count = features.iter().filter(|f| *f == "file_system").count();
4007 assert_eq!(file_system_count, 1, "file_system should appear only once");
4008 }
4009
4010 #[test]
4011 fn test_compute_features_includes_dependency_features() {
4012 let registry = CapabilityRegistry::with_builtins();
4013
4014 let features = compute_features(&["virtual_bash".to_string()], ®istry);
4016 assert!(features.contains(&"file_system".to_string()));
4017 }
4018
4019 #[test]
4020 fn test_compute_features_generic_harness_set() {
4021 let registry = CapabilityRegistry::with_builtins();
4022
4023 let features = compute_features(
4025 &[
4026 "session_file_system".to_string(),
4027 "virtual_bash".to_string(),
4028 "session_storage".to_string(),
4029 "session".to_string(),
4030 "session_schedule".to_string(),
4031 ],
4032 ®istry,
4033 );
4034 assert!(features.contains(&"file_system".to_string()));
4035 assert!(features.contains(&"secrets".to_string()));
4036 assert!(features.contains(&"key_value".to_string()));
4037 assert!(features.contains(&"schedules".to_string()));
4038 }
4039
4040 #[test]
4041 fn test_compute_features_unknown_capability_ignored() {
4042 let registry = CapabilityRegistry::with_builtins();
4043
4044 let features = compute_features(
4045 &["unknown_cap".to_string(), "session_schedule".to_string()],
4046 ®istry,
4047 );
4048 assert_eq!(features, vec!["schedules"]);
4049 }
4050
4051 #[test]
4052 fn test_risk_level_ordering() {
4053 assert!(RiskLevel::Low < RiskLevel::Medium);
4054 assert!(RiskLevel::Medium < RiskLevel::High);
4055 }
4056
4057 #[test]
4058 fn test_risk_level_serde_roundtrip() {
4059 let high = RiskLevel::High;
4060 let json = serde_json::to_string(&high).unwrap();
4061 assert_eq!(json, "\"high\"");
4062 let back: RiskLevel = serde_json::from_str(&json).unwrap();
4063 assert_eq!(back, RiskLevel::High);
4064 }
4065
4066 #[test]
4067 fn test_capability_risk_levels() {
4068 let registry = CapabilityRegistry::with_builtins();
4069
4070 let bash = registry.get("virtual_bash").unwrap();
4072 assert_eq!(bash.risk_level(), RiskLevel::High);
4073
4074 let fetch = registry.get("web_fetch").unwrap();
4076 assert_eq!(fetch.risk_level(), RiskLevel::High);
4077
4078 let noop = registry.get("noop").unwrap();
4080 assert_eq!(noop.risk_level(), RiskLevel::Low);
4081 }
4082
4083 #[tokio::test]
4088 async fn test_apply_capabilities_openai_tool_search() {
4089 let registry = CapabilityRegistry::with_builtins();
4090 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4091
4092 let applied = apply_capabilities(
4093 base_runtime_agent.clone(),
4094 &["openai_tool_search".to_string()],
4095 ®istry,
4096 &test_ctx(),
4097 )
4098 .await;
4099
4100 assert_eq!(
4102 applied.runtime_agent.system_prompt,
4103 base_runtime_agent.system_prompt
4104 );
4105 assert!(applied.tool_registry.is_empty());
4106 assert_eq!(applied.applied_ids, vec!["openai_tool_search"]);
4107
4108 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
4110 assert!(ts.enabled);
4111 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
4112 }
4113
4114 #[tokio::test]
4115 async fn test_apply_capabilities_openai_tool_search_with_other_capabilities() {
4116 let registry = CapabilityRegistry::with_builtins();
4117 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4118
4119 let applied = apply_capabilities(
4120 base_runtime_agent,
4121 &[
4122 "current_time".to_string(),
4123 "openai_tool_search".to_string(),
4124 "test_math".to_string(),
4125 ],
4126 ®istry,
4127 &test_ctx(),
4128 )
4129 .await;
4130
4131 assert!(applied.tool_registry.has("get_current_time"));
4133 assert!(applied.tool_registry.has("add"));
4134 assert!(applied.tool_registry.has("subtract"));
4135 assert!(applied.tool_registry.has("multiply"));
4136 assert!(applied.tool_registry.has("divide"));
4137
4138 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
4140 assert!(ts.enabled);
4141 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
4142 }
4143
4144 #[tokio::test]
4145 async fn test_collect_capabilities_tool_search_custom_threshold() {
4146 let registry = CapabilityRegistry::with_builtins();
4147
4148 let configs = vec![AgentCapabilityConfig {
4149 capability_ref: CapabilityId::new("openai_tool_search"),
4150 config: serde_json::json!({"threshold": 5}),
4151 }];
4152
4153 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4154
4155 let ts = collected.tool_search.as_ref().unwrap();
4156 assert!(ts.enabled);
4157 assert_eq!(ts.threshold, 5);
4158 }
4159
4160 #[tokio::test]
4161 async fn test_collect_capabilities_auto_tool_search_resolves_to_generic_off_native() {
4162 let registry = CapabilityRegistry::with_builtins();
4163
4164 let configs = vec![
4165 AgentCapabilityConfig {
4166 capability_ref: CapabilityId::new("auto_tool_search"),
4167 config: serde_json::json!({"threshold": 2}),
4168 },
4169 AgentCapabilityConfig {
4170 capability_ref: CapabilityId::new("test_math"),
4171 config: serde_json::json!({}),
4172 },
4173 ];
4174
4175 let ctx = test_ctx().with_model("claude-sonnet-4-5-20250514");
4178 let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
4179
4180 assert!(
4181 collected.tool_search.is_none(),
4182 "auto_tool_search must not set a hosted config on a non-native model"
4183 );
4184 assert!(
4185 collected
4186 .tools
4187 .iter()
4188 .any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
4189 "auto_tool_search must contribute the client-side tool_search tool"
4190 );
4191 assert!(
4192 !collected.tool_definition_hooks.is_empty(),
4193 "auto_tool_search must contribute a client-side deferral hook"
4194 );
4195
4196 let mut transformed = collected.tool_definitions.clone();
4197 for hook in &collected.tool_definition_hooks {
4198 transformed = hook.transform(transformed);
4199 }
4200 let add_tool = transformed
4201 .iter()
4202 .find(|tool| tool.name() == "add")
4203 .expect("test_math contributes add");
4204 assert!(
4205 add_tool.parameters().get("properties").is_none(),
4206 "generic auto_tool_search must honor the configured threshold"
4207 );
4208 }
4209
4210 #[tokio::test]
4211 async fn test_collect_capabilities_auto_tool_search_resolves_to_hosted_on_native() {
4212 let registry = CapabilityRegistry::with_builtins();
4213
4214 let configs = vec![AgentCapabilityConfig {
4215 capability_ref: CapabilityId::new("auto_tool_search"),
4216 config: serde_json::json!({"threshold": 7}),
4217 }];
4218
4219 let ctx = test_ctx().with_model("gpt-5.4");
4222 let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
4223
4224 let ts = collected
4225 .tool_search
4226 .as_ref()
4227 .expect("auto_tool_search must set a hosted config on a native model");
4228 assert!(ts.enabled);
4229 assert_eq!(ts.threshold, 7);
4230 assert!(
4231 !collected
4232 .tools
4233 .iter()
4234 .any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
4235 "hosted mechanism must not contribute the client-side tool_search tool"
4236 );
4237 assert!(
4238 collected.tool_definition_hooks.is_empty(),
4239 "hosted mechanism must not contribute a client-side deferral hook"
4240 );
4241 }
4242
4243 #[tokio::test]
4244 async fn test_collect_capabilities_no_tool_search_without_capability() {
4245 let registry = CapabilityRegistry::with_builtins();
4246
4247 let configs = vec![AgentCapabilityConfig {
4248 capability_ref: CapabilityId::new("current_time"),
4249 config: serde_json::json!({}),
4250 }];
4251
4252 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4253
4254 assert!(collected.tool_search.is_none());
4255 }
4256
4257 #[tokio::test]
4258 async fn test_collect_capabilities_tool_search_category_propagation() {
4259 let registry = CapabilityRegistry::with_builtins();
4260
4261 let configs = vec![
4263 AgentCapabilityConfig {
4264 capability_ref: CapabilityId::new("test_math"),
4265 config: serde_json::json!({}),
4266 },
4267 AgentCapabilityConfig {
4268 capability_ref: CapabilityId::new("openai_tool_search"),
4269 config: serde_json::json!({}),
4270 },
4271 ];
4272
4273 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4274
4275 assert!(collected.tool_search.is_some());
4277
4278 for tool_def in &collected.tool_definitions {
4280 if ["add", "subtract", "multiply", "divide"].contains(&tool_def.name()) {
4282 assert!(
4283 tool_def.category().is_some(),
4284 "Tool {} should have a category from its capability",
4285 tool_def.name()
4286 );
4287 }
4288 }
4289 }
4290
4291 #[tokio::test]
4292 async fn test_apply_capabilities_prompt_caching() {
4293 let registry = CapabilityRegistry::with_builtins();
4294 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4295
4296 let applied = apply_capabilities(
4297 base_runtime_agent.clone(),
4298 &["prompt_caching".to_string()],
4299 ®istry,
4300 &test_ctx(),
4301 )
4302 .await;
4303
4304 assert_eq!(
4305 applied.runtime_agent.system_prompt,
4306 base_runtime_agent.system_prompt
4307 );
4308 assert!(applied.tool_registry.is_empty());
4309 assert_eq!(applied.applied_ids, vec!["prompt_caching"]);
4310
4311 let prompt_cache = applied.runtime_agent.prompt_cache.as_ref().unwrap();
4312 assert!(prompt_cache.enabled);
4313 assert_eq!(
4314 prompt_cache.strategy,
4315 crate::llm_driver_registry::PromptCacheStrategy::Auto
4316 );
4317 assert!(prompt_cache.gemini_cached_content.is_none());
4318 }
4319
4320 #[tokio::test]
4321 async fn test_collect_capabilities_prompt_caching_custom_strategy() {
4322 let registry = CapabilityRegistry::with_builtins();
4323
4324 let configs = vec![AgentCapabilityConfig {
4325 capability_ref: CapabilityId::new("prompt_caching"),
4326 config: serde_json::json!({"strategy": "auto"}),
4327 }];
4328
4329 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4330
4331 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
4332 assert!(prompt_cache.enabled);
4333 assert_eq!(
4334 prompt_cache.strategy,
4335 crate::llm_driver_registry::PromptCacheStrategy::Auto
4336 );
4337 assert!(prompt_cache.gemini_cached_content.is_none());
4338 }
4339
4340 #[tokio::test]
4341 async fn test_collect_capabilities_prompt_caching_gemini_cached_content() {
4342 let registry = CapabilityRegistry::with_builtins();
4343
4344 let configs = vec![AgentCapabilityConfig {
4345 capability_ref: CapabilityId::new("prompt_caching"),
4346 config: serde_json::json!({
4347 "strategy": "auto",
4348 "gemini_cached_content": "cachedContents/demo-cache"
4349 }),
4350 }];
4351
4352 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4353
4354 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
4355 assert_eq!(
4356 prompt_cache.gemini_cached_content.as_deref(),
4357 Some("cachedContents/demo-cache")
4358 );
4359 }
4360
4361 struct SkillContributingCapability;
4366
4367 impl Capability for SkillContributingCapability {
4368 fn id(&self) -> &str {
4369 "contributes_skills"
4370 }
4371 fn name(&self) -> &str {
4372 "Contributes Skills"
4373 }
4374 fn description(&self) -> &str {
4375 "Test capability that contributes skills."
4376 }
4377 fn contribute_skills(&self) -> Vec<SkillContribution> {
4378 vec![
4379 SkillContribution::new("alpha-skill", "Alpha skill desc", "# Alpha\nDo alpha.")
4380 .with_files(vec![(
4381 "scripts/a.sh".to_string(),
4382 "#!/bin/sh\necho a\n".to_string(),
4383 )]),
4384 SkillContribution::new("beta-skill", "Beta skill desc", "# Beta\nDo beta.")
4385 .with_user_invocable(false),
4386 ]
4387 }
4388 }
4389
4390 fn skill_md_from_entries(entries: &HashMap<String, MountEntry>) -> &str {
4391 match &entries.get("SKILL.md").expect("SKILL.md missing").source {
4392 MountSource::InlineFile { content, .. } => content.as_str(),
4393 _ => panic!("Expected InlineFile for SKILL.md"),
4394 }
4395 }
4396
4397 #[tokio::test]
4398 async fn test_contribute_skills_normalized_to_mounts() {
4399 let mut registry = CapabilityRegistry::new();
4400 registry.register(SkillContributingCapability);
4401
4402 let configs = vec![AgentCapabilityConfig {
4403 capability_ref: CapabilityId::new("contributes_skills"),
4404 config: serde_json::json!({}),
4405 }];
4406
4407 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4408
4409 let skill_mounts: Vec<_> = collected
4410 .mounts
4411 .iter()
4412 .filter(|m| m.path.starts_with("/.agents/skills/"))
4413 .collect();
4414 assert_eq!(skill_mounts.len(), 2);
4415
4416 for m in &skill_mounts {
4419 assert!(m.is_readonly());
4420 assert_eq!(m.capability_id, "contributes_skills");
4421 }
4422
4423 let alpha = skill_mounts
4424 .iter()
4425 .find(|m| m.path == "/.agents/skills/alpha-skill")
4426 .expect("alpha-skill mount missing");
4427 match &alpha.source {
4428 MountSource::InlineDirectory { entries } => {
4429 assert!(entries.contains_key("SKILL.md"));
4430 assert!(entries.contains_key("scripts/a.sh"));
4431 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
4432 assert_eq!(parsed.name, "alpha-skill");
4433 assert!(parsed.user_invocable);
4434 }
4435 _ => panic!("Expected InlineDirectory"),
4436 }
4437
4438 let beta = skill_mounts
4439 .iter()
4440 .find(|m| m.path == "/.agents/skills/beta-skill")
4441 .expect("beta-skill mount missing");
4442 match &beta.source {
4443 MountSource::InlineDirectory { entries } => {
4444 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
4445 assert!(!parsed.user_invocable);
4446 }
4447 _ => panic!("Expected InlineDirectory"),
4448 }
4449 }
4450
4451 #[tokio::test]
4452 async fn test_contribute_skills_default_empty() {
4453 let mut registry = CapabilityRegistry::new();
4456 registry.register(FilterTestCapability { priority: 0 });
4457
4458 let configs = vec![AgentCapabilityConfig {
4459 capability_ref: CapabilityId::new("filter_test"),
4460 config: serde_json::json!({}),
4461 }];
4462
4463 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4464 assert!(
4465 collected
4466 .mounts
4467 .iter()
4468 .all(|m| !m.path.starts_with("/.agents/skills/"))
4469 );
4470 }
4471}