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;
107pub mod mcp;
108mod noop;
109mod openai_tool_search;
110#[cfg(feature = "ui-capabilities")]
111mod openui;
112pub mod persistent_memory;
113mod platform_management;
114mod prompt_caching;
115mod prompt_canary_guardrail;
116mod research;
117mod sample_data;
118mod self_budget;
119mod session;
120mod session_sandbox;
121mod session_schedule;
122mod session_sql_database;
123mod session_storage;
124mod skills;
125mod stateless_todo_list;
126mod subagents;
127mod system_commands;
128mod test_math;
129mod test_weather;
130mod tool_output_persistence;
131mod tool_search;
132pub mod user_hooks;
133mod virtual_bash;
134mod web_fetch;
135mod workspace_volumes;
136
137pub use a2a_delegation::{
139 A2A_AGENT_DELEGATION_CAPABILITY_ID, A2aAgentDelegationCapability, CancelAgentTool,
140 GetAgentRunsTool, MessageAgentTool, SpawnAgentTool, WaitAgentTool,
141};
142#[cfg(feature = "ui-capabilities")]
143pub use a2ui::{A2UI_CAPABILITY_ID, A2UiCapability};
144pub use agent_handoff::{
145 AGENT_HANDOFF_CAPABILITY_ID, AgentHandoffCapability, GetAgentHandoffsTool,
146 MessageAgentHandoffTool, StartAgentHandoffTool,
147};
148pub use agent_instructions::{
149 AGENT_INSTRUCTIONS_CAPABILITY_ID, AGENTS_MD_PATH, AgentInstructionsCapability,
150 AgentInstructionsConfig, DEFAULT_AGENT_INSTRUCTIONS_FILE, MAX_AGENT_INSTRUCTIONS_FILES,
151 MAX_AGENTS_MD_SIZE, format_agents_md_content, format_instruction_file_content,
152};
153pub use attach_skill::{
154 AttachSkillCapability, SKILL_CAPABILITY_PREFIX, SKILLS_DISCOVERY_PATH, SkillContribution,
155 SkillInstructions, SkillMeta, SkillSource, discover_skills_from_entries, is_skill_capability,
156 parse_skill_capability_id, reconstruct_skill_md, skill_capability_id,
157};
158pub use auto_tool_search::{AUTO_TOOL_SEARCH_CAPABILITY_ID, AutoToolSearchCapability};
159pub use background_execution::{BACKGROUND_EXECUTION_CAPABILITY_ID, BackgroundExecutionCapability};
160pub use btw::{BTW_CAPABILITY_ID, BtwCapability};
161pub use budgeting::BudgetingCapability;
162pub use compaction::{
163 COMPACTION_CAPABILITY_ID, CompactionCapability, CompactionConfig, CompactionStep,
164 CompactionStrategy, CostControlConfig, CostControlMaskingResult, HierarchicalMemoryConfig,
165 MaskingSummaryFormat, MemoryTier, ObservationMaskingConfig, ObservationMaskingResult,
166 SessionCompactionMetrics, SummarizationConfig, aggressive_trim, apply_cost_control_masking,
167 apply_hierarchical_memory, apply_observation_masking, build_model_view_messages,
168 build_summarization_prompt, build_summary_message, classify_memory_tiers, estimate_tokens,
169 estimate_total_tokens, format_messages_for_summarization, should_compact_proactively,
170};
171pub use current_time::{CurrentTimeCapability, GetCurrentTimeTool};
172pub use data_knowledge::DataKnowledgeCapability;
173pub use declarative::{
174 DECLARATIVE_CAPABILITY_PREFIX, DeclarativeCapabilityDefinition, DeclarativeCapabilityFile,
175 DeclarativeCapabilitySkill, declarative_capability_id, declarative_capability_info,
176 hydrate_declarative_capability_config, is_declarative_capability,
177 parse_declarative_capability_id, validate_declarative_capability_definition,
178};
179pub use fake_aws::{
180 AwsCreateEc2InstanceTool, AwsCreateIamUserTool, AwsCreateRdsDatabaseTool,
181 AwsCreateS3BucketTool, AwsGetCloudWatchMetricsTool, AwsListEc2InstancesTool,
182 AwsListIamUsersTool, AwsListRdsDatabasesTool, AwsListS3BucketsTool, AwsListSecurityGroupsTool,
183 AwsStopEc2InstanceTool, FakeAwsCapability,
184};
185pub use fake_crm::{
186 CrmAddInteractionTool, CrmCreateCustomerTool, CrmCreateTicketTool, CrmGetCustomerTool,
187 CrmListCustomersTool, CrmListTicketsTool, CrmSearchCustomersTool, CrmUpdateTicketTool,
188 FakeCrmCapability,
189};
190pub use fake_financial::{
191 FakeFinancialCapability, FinanceCreateBudgetTool, FinanceCreateTransactionTool,
192 FinanceForecastCashFlowTool, FinanceGetBalanceTool, FinanceGetExpenseReportTool,
193 FinanceGetRevenueReportTool, FinanceListBudgetsTool, FinanceListTransactionsTool,
194};
195pub use fake_warehouse::{
196 FakeWarehouseCapability, WarehouseCreateInvoiceTool, WarehouseCreateOrderTool,
197 WarehouseCreateShipmentTool, WarehouseGetInventoryTool, WarehouseInventoryReportTool,
198 WarehouseListOrdersTool, WarehouseListShipmentsTool, WarehouseProcessReturnTool,
199 WarehouseUpdateInventoryTool, WarehouseUpdateShipmentStatusTool,
200};
201pub use file_system::{
202 DeleteFileTool, EditFileTool, FileSystemCapability, GrepFilesTool, ListDirectoryTool,
203 ReadFileTool, StatFileTool, WriteFileTool,
204};
205pub use human_intent::{HUMAN_INTENT_CAPABILITY_ID, HumanIntentCapability};
206pub use infinity_context::{
207 INFINITY_CONTEXT_CAPABILITY_ID, InfinityContextCapability, QueryHistoryTool,
208};
209pub use knowledge_base::{
210 KNOWLEDGE_BASE_CAPABILITY_ID, KnowledgeBaseCapability, KnowledgeBaseConfig,
211 validate_knowledge_base_config,
212};
213pub use loop_detection::LoopDetectionCapability;
214pub use mcp::{
215 MCP_CAPABILITY_PREFIX, McpCapability, is_mcp_capability, mcp_capability_id,
216 parse_mcp_capability_id,
217};
218pub use noop::NoopCapability;
219pub use openai_tool_search::{
220 DEFAULT_TOOL_SEARCH_THRESHOLD, OPENAI_TOOL_SEARCH_CAPABILITY_ID, OpenAiToolSearchCapability,
221 model_supports_native_tool_search,
222};
223#[cfg(feature = "ui-capabilities")]
224pub use openui::{OPENUI_CAPABILITY_ID, OpenUiCapability};
225pub use persistent_memory::{
226 ForgetTool, MEMORY_CAPABILITY_ID, MemoryCapability, MemoryConfig, RecallTool, RememberTool,
227};
228pub use platform_management::{
229 ManageAgentsTool, ManageHarnessesTool, ManageSessionsTool, PlatformManagementCapability,
230 ReadAgentsTool, ReadCapabilitiesTool, ReadHarnessesTool, ReadSessionsTool,
231 SessionReadMessagesTool, SessionReadResponseTool, SessionSendMessageTool,
232};
233pub use prompt_caching::{PROMPT_CACHING_CAPABILITY_ID, PromptCachingCapability};
234pub use prompt_canary_guardrail::{
235 DEFAULT_REPLACEMENT as PROMPT_CANARY_DEFAULT_REPLACEMENT,
236 PROMPT_CANARY_GUARDRAIL_CAPABILITY_ID, PromptCanaryGuardrailCapability,
237 REASON_CODE_SYSTEM_PROMPT_LEAK,
238};
239pub use research::ResearchCapability;
240pub use sample_data::SampleDataCapability;
241pub use self_budget::{SELF_BUDGET_CAPABILITY_ID, SelfBudgetCapability};
242pub use session::{GetSessionInfoTool, SessionCapability, WriteSessionTitleTool};
243pub use session_sandbox::{
244 SESSION_SANDBOX_CAPABILITY_ID, SandboxExecTool, SandboxManageTool, SandboxReadFileTool,
245 SandboxStatusTool, SandboxWriteFileTool, SessionSandboxCapability,
246};
247pub use session_schedule::{
248 CancelScheduleTool, CreateScheduleTool, ListSchedulesTool, SESSION_SCHEDULE_CAPABILITY_ID,
249 SessionScheduleCapability,
250};
251pub use session_sql_database::{
252 SessionSqlDatabaseCapability, SqlExecuteTool, SqlQueryTool, SqlSchemaTool,
253};
254pub use session_storage::{KvStoreTool, SecretStoreTool, SessionStorageCapability};
255pub use skills::{SKILLS_CAPABILITY_ID, SkillsCapability};
256pub use stateless_todo_list::{StatelessTodoListCapability, WriteTodosTool};
257pub use subagents::SubagentCapability;
258pub use system_commands::{SYSTEM_COMMANDS_CAPABILITY_ID, SystemCommandsCapability};
260pub use test_math::{AddTool, DivideTool, MultiplyTool, SubtractTool, TestMathCapability};
261pub use test_weather::{GetForecastTool, GetWeatherTool, TestWeatherCapability};
262pub use tool_output_persistence::{PersistOutputHook, ToolOutputPersistenceCapability};
263pub use tool_search::{
264 TOOL_SEARCH_CAPABILITY_ID, TOOL_SEARCH_TOOL_NAME, ToolSearchCapability, ToolSearchTool,
265};
266pub use user_hooks::UserHooksCapability;
267pub use virtual_bash::{BashTool, SessionFileSystemAdapter, VirtualBashCapability};
268pub use web_fetch::{
269 BotAuthPublicKey, WebFetchCapability, WebFetchTool, derive_bot_auth_public_key,
270};
271pub use workspace_volumes::{WORKSPACE_VOLUMES_CAPABILITY_ID, WorkspaceVolumesCapability};
272
273pub struct SystemPromptContext {
283 pub session_id: SessionId,
285 pub locale: Option<String>,
287 pub file_store: Option<Arc<dyn SessionFileSystem>>,
289 pub model: Option<String>,
295}
296
297impl SystemPromptContext {
298 pub fn without_file_store(session_id: SessionId) -> Self {
300 Self {
301 session_id,
302 locale: None,
303 file_store: None,
304 model: None,
305 }
306 }
307
308 pub fn with_model(mut self, model: impl Into<String>) -> Self {
310 self.model = Some(model.into());
311 self
312 }
313}
314
315#[async_trait]
362pub trait Capability: Send + Sync {
363 fn id(&self) -> &str;
365
366 fn name(&self) -> &str;
368
369 fn description(&self) -> &str;
371
372 fn status(&self) -> CapabilityStatus {
374 CapabilityStatus::Available
375 }
376
377 fn icon(&self) -> Option<&str> {
379 None
380 }
381
382 fn category(&self) -> Option<&str> {
384 None
385 }
386
387 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
398 None
399 }
400
401 fn system_prompt_addition(&self) -> Option<&str> {
421 None
422 }
423
424 async fn system_prompt_contribution(&self, _ctx: &SystemPromptContext) -> Option<String> {
436 self.system_prompt_addition().map(|addition| {
437 format!(
438 "<capability id=\"{}\">\n{}\n</capability>",
439 self.id(),
440 addition
441 )
442 })
443 }
444
445 fn system_prompt_preview(&self) -> Option<String> {
451 self.system_prompt_addition().map(|s| s.to_string())
452 }
453
454 fn tools(&self) -> Vec<Box<dyn Tool>> {
456 vec![]
457 }
458
459 fn tools_with_config(&self, _config: &serde_json::Value) -> Vec<Box<dyn Tool>> {
467 self.tools()
468 }
469
470 async fn system_prompt_contribution_with_config(
477 &self,
478 ctx: &SystemPromptContext,
479 _config: &serde_json::Value,
480 ) -> Option<String> {
481 self.system_prompt_contribution(ctx).await
482 }
483
484 fn tool_definitions(&self) -> Vec<ToolDefinition> {
487 self.tools().iter().map(|t| t.to_definition()).collect()
488 }
489
490 fn mounts(&self) -> Vec<MountPoint> {
498 vec![]
499 }
500
501 fn dependencies(&self) -> Vec<&'static str> {
510 vec![]
511 }
512
513 fn features(&self) -> Vec<&'static str> {
528 vec![]
529 }
530
531 fn config_schema(&self) -> Option<serde_json::Value> {
537 None
538 }
539
540 fn config_ui_schema(&self) -> Option<serde_json::Value> {
545 None
546 }
547
548 fn validate_config(&self, _config: &serde_json::Value) -> Result<(), String> {
554 Ok(())
555 }
556
557 fn mcp_servers(&self) -> ScopedMcpServers {
563 ScopedMcpServers::default()
564 }
565
566 fn mcp_servers_with_config(&self, _config: &serde_json::Value) -> ScopedMcpServers {
568 self.mcp_servers()
569 }
570
571 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
584 None
585 }
586
587 fn model_view_provider(&self) -> Option<Arc<dyn ModelViewProvider>> {
595 None
596 }
597
598 fn pre_tool_use_hooks(&self) -> Vec<Arc<dyn crate::atoms::PreToolUseHook>> {
609 vec![]
610 }
611
612 fn post_tool_exec_hooks(&self) -> Vec<Arc<dyn crate::atoms::PostToolExecHook>> {
620 vec![]
621 }
622
623 fn tool_definition_hooks(&self) -> Vec<Arc<dyn ToolDefinitionHook>> {
632 vec![]
633 }
634
635 fn tool_definition_hooks_with_config(
640 &self,
641 _config: &serde_json::Value,
642 ) -> Vec<Arc<dyn ToolDefinitionHook>> {
643 self.tool_definition_hooks()
644 }
645
646 fn tool_call_hooks(&self) -> Vec<Arc<dyn ToolCallHook>> {
654 vec![]
655 }
656
657 fn user_hooks(&self) -> Vec<crate::user_hook_types::UserHookSpec> {
673 vec![]
674 }
675
676 fn user_hooks_with_config(
682 &self,
683 _config: &serde_json::Value,
684 ) -> Vec<crate::user_hook_types::UserHookSpec> {
685 self.user_hooks()
686 }
687
688 fn risk_level(&self) -> RiskLevel {
696 RiskLevel::Low
697 }
698
699 fn commands(&self) -> Vec<CommandDescriptor> {
707 vec![]
708 }
709
710 async fn execute_command(
724 &self,
725 request: &ExecuteCommandRequest,
726 _ctx: &CommandExecutionContext,
727 ) -> crate::error::Result<CommandResult> {
728 Err(crate::error::AgentLoopError::config(format!(
729 "capability {} declared command /{} but does not implement execute_command",
730 self.id(),
731 request.name,
732 )))
733 }
734
735 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
743 vec![]
744 }
745
746 fn contribute_skills(&self) -> Vec<SkillContribution> {
756 vec![]
757 }
758
759 fn output_guardrails(&self) -> Vec<Arc<dyn crate::output_guardrail::OutputGuardrail>> {
770 vec![]
771 }
772}
773
774pub trait ToolDefinitionHook: Send + Sync {
775 fn transform(&self, tools: Vec<ToolDefinition>) -> Vec<ToolDefinition>;
776
777 fn applies_with_native_tool_search(&self) -> bool {
782 true
783 }
784}
785
786pub trait ToolCallHook: Send + Sync {
787 fn narration(
788 &self,
789 _tool_def: Option<&ToolDefinition>,
790 _tool_call: &ToolCall,
791 _phase: crate::tool_narration::ToolNarrationPhase,
792 _locale: Option<&str>,
793 ) -> Option<String> {
794 None
795 }
796
797 fn transform_for_execution(&self, tool_call: ToolCall) -> ToolCall {
798 tool_call
799 }
800}
801
802#[derive(
806 Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
807)]
808#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
809#[cfg_attr(feature = "openapi", schema(example = "low"))]
810#[serde(rename_all = "lowercase")]
811pub enum RiskLevel {
812 Low,
814 Medium,
816 High,
818}
819
820#[derive(Debug, Clone, Serialize, Deserialize)]
826#[serde(rename_all = "snake_case")]
827pub enum BlueprintModel {
828 Fixed(String),
830 Default(String),
832 Inherit,
834}
835
836pub struct AgentBlueprint {
842 pub id: &'static str,
844 pub name: &'static str,
846 pub description: &'static str,
848 pub model: BlueprintModel,
850 pub system_prompt: &'static str,
852 pub tools: Vec<Box<dyn Tool>>,
854 pub max_turns: Option<usize>,
856 pub config_schema: Option<serde_json::Value>,
858}
859
860impl AgentBlueprint {
861 pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
863 self.tools.iter().map(|t| t.to_definition()).collect()
864 }
865}
866
867impl std::fmt::Debug for AgentBlueprint {
868 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
869 f.debug_struct("AgentBlueprint")
870 .field("id", &self.id)
871 .field("name", &self.name)
872 .field("model", &self.model)
873 .field("tool_count", &self.tools.len())
874 .field("max_turns", &self.max_turns)
875 .finish()
876 }
877}
878
879#[derive(Clone)]
906pub struct CapabilityRegistry {
907 capabilities: HashMap<String, Arc<dyn Capability>>,
908}
909
910impl CapabilityRegistry {
911 pub fn new() -> Self {
913 Self {
914 capabilities: HashMap::new(),
915 }
916 }
917
918 pub fn with_builtins() -> Self {
923 Self::with_builtins_for_grade(DeploymentGrade::from_env())
924 }
925
926 pub fn with_builtins_for_grade(grade: DeploymentGrade) -> Self {
931 let mut registry = Self::new();
932
933 registry.register(AgentInstructionsCapability);
935 registry.register(HumanIntentCapability);
936 registry.register(NoopCapability);
937 registry.register(CurrentTimeCapability);
938 registry.register(ResearchCapability);
939 registry.register(PlatformManagementCapability);
940 registry.register(FileSystemCapability);
941 registry.register(WorkspaceVolumesCapability);
942 registry.register(SessionStorageCapability);
943 registry.register(SessionCapability);
944 registry.register(SessionSqlDatabaseCapability);
945 registry.register(TestMathCapability);
946 registry.register(TestWeatherCapability);
947 registry.register(StatelessTodoListCapability);
948 registry.register(WebFetchCapability::from_env());
949 registry.register(VirtualBashCapability);
950 registry.register(BackgroundExecutionCapability);
951 registry.register(SessionScheduleCapability);
952 registry.register(BtwCapability);
953 registry.register(InfinityContextCapability);
954 registry.register(budgeting::BudgetingCapability);
955 registry.register(SelfBudgetCapability);
956 registry.register(CompactionCapability);
957 registry.register(MemoryCapability);
958
959 registry.register(OpenAiToolSearchCapability::new());
961 registry.register(ToolSearchCapability::new());
963 registry.register(AutoToolSearchCapability::new());
965 registry.register(PromptCachingCapability::new());
966
967 registry.register(SkillsCapability);
969
970 registry.register(SubagentCapability);
972
973 if crate::FeatureFlags::from_env(&grade).agent_delegation {
977 registry.register(AgentHandoffCapability);
978 registry.register(A2aAgentDelegationCapability);
979 }
980
981 registry.register(SystemCommandsCapability);
983
984 registry.register(tool_output_persistence::ToolOutputPersistenceCapability);
986
987 registry.register(user_hooks::UserHooksCapability);
990
991 registry.register(LoopDetectionCapability);
993
994 registry.register(PromptCanaryGuardrailCapability);
997
998 #[cfg(feature = "ui-capabilities")]
1000 {
1001 registry.register(OpenUiCapability);
1002 registry.register(A2UiCapability);
1003 }
1004
1005 registry.register(SampleDataCapability);
1007
1008 registry.register(DataKnowledgeCapability);
1010
1011 registry.register(KnowledgeBaseCapability);
1013
1014 registry.register(FakeWarehouseCapability);
1016 registry.register(FakeAwsCapability);
1017 registry.register(FakeCrmCapability);
1018 registry.register(FakeFinancialCapability);
1019
1020 let internal_flags = crate::InternalFeatureFlags::from_env();
1022 if internal_flags.session_sandbox {
1023 registry.register(SessionSandboxCapability);
1024 }
1025 for plugin in inventory::iter::<IntegrationPlugin>() {
1026 if (!plugin.experimental_only || grade.experimental_features_enabled())
1027 && plugin
1028 .feature_flag
1029 .is_none_or(|f| internal_flags.is_enabled(f))
1030 {
1031 registry.register_boxed((plugin.factory)());
1032 }
1033 }
1034
1035 registry
1036 }
1037
1038 pub fn register(&mut self, capability: impl Capability + 'static) {
1040 self.capabilities
1041 .insert(capability.id().to_string(), Arc::new(capability));
1042 }
1043
1044 pub fn register_boxed(&mut self, capability: Box<dyn Capability>) {
1046 self.capabilities
1047 .insert(capability.id().to_string(), Arc::from(capability));
1048 }
1049
1050 pub fn register_arc(&mut self, capability: Arc<dyn Capability>) {
1052 self.capabilities
1053 .insert(capability.id().to_string(), capability);
1054 }
1055
1056 pub fn get(&self, id: &str) -> Option<&Arc<dyn Capability>> {
1058 self.capabilities.get(id)
1059 }
1060
1061 pub fn unregister(&mut self, id: &str) -> Option<Arc<dyn Capability>> {
1063 self.capabilities.remove(id)
1064 }
1065
1066 pub fn has(&self, id: &str) -> bool {
1068 self.capabilities.contains_key(id)
1069 }
1070
1071 pub fn list(&self) -> Vec<&Arc<dyn Capability>> {
1073 self.capabilities.values().collect()
1074 }
1075
1076 pub fn len(&self) -> usize {
1078 self.capabilities.len()
1079 }
1080
1081 pub fn is_empty(&self) -> bool {
1083 self.capabilities.is_empty()
1084 }
1085
1086 pub fn builder() -> CapabilityRegistryBuilder {
1088 CapabilityRegistryBuilder::new()
1089 }
1090
1091 pub fn blueprint(&self, id: &str) -> Option<AgentBlueprint> {
1095 for cap in self.capabilities.values() {
1096 for bp in cap.agent_blueprints() {
1097 if bp.id == id {
1098 return Some(bp);
1099 }
1100 }
1101 }
1102 None
1103 }
1104
1105 pub fn blueprint_with_capability(&self, id: &str) -> Option<(String, AgentBlueprint)> {
1109 for (capability_id, cap) in &self.capabilities {
1110 for bp in cap.agent_blueprints() {
1111 if bp.id == id {
1112 return Some((capability_id.clone(), bp));
1113 }
1114 }
1115 }
1116 None
1117 }
1118
1119 pub fn all_blueprints(&self) -> Vec<AgentBlueprint> {
1121 self.capabilities
1122 .values()
1123 .flat_map(|cap| cap.agent_blueprints())
1124 .collect()
1125 }
1126}
1127
1128impl Default for CapabilityRegistry {
1129 fn default() -> Self {
1130 Self::with_builtins()
1131 }
1132}
1133
1134impl std::fmt::Debug for CapabilityRegistry {
1135 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1136 let ids: Vec<_> = self.capabilities.keys().collect();
1137 f.debug_struct("CapabilityRegistry")
1138 .field("capabilities", &ids)
1139 .finish()
1140 }
1141}
1142
1143pub struct CapabilityRegistryBuilder {
1145 registry: CapabilityRegistry,
1146}
1147
1148impl CapabilityRegistryBuilder {
1149 pub fn new() -> Self {
1151 Self {
1152 registry: CapabilityRegistry::new(),
1153 }
1154 }
1155
1156 pub fn with_builtins() -> Self {
1158 Self {
1159 registry: CapabilityRegistry::with_builtins(),
1160 }
1161 }
1162
1163 pub fn capability(mut self, capability: impl Capability + 'static) -> Self {
1165 self.registry.register(capability);
1166 self
1167 }
1168
1169 pub fn build(self) -> CapabilityRegistry {
1171 self.registry
1172 }
1173}
1174
1175impl Default for CapabilityRegistryBuilder {
1176 fn default() -> Self {
1177 Self::new()
1178 }
1179}
1180
1181pub struct ModelViewContext<'a> {
1187 pub session_id: SessionId,
1188 pub prior_usage: Option<&'a TokenUsage>,
1189}
1190
1191pub trait ModelViewProvider: Send + Sync {
1197 fn apply_model_view(
1198 &self,
1199 messages: Vec<Message>,
1200 config: &serde_json::Value,
1201 context: &ModelViewContext<'_>,
1202 ) -> Vec<Message>;
1203
1204 fn priority(&self) -> i32 {
1205 0
1206 }
1207}
1208
1209pub struct CollectedCapabilities {
1214 pub system_prompt_parts: Vec<String>,
1216 pub system_prompt_attributions: Vec<SystemPromptAttribution>,
1218 pub tools: Vec<Box<dyn Tool>>,
1220 pub tool_definitions: Vec<ToolDefinition>,
1222 pub mounts: Vec<MountPoint>,
1224 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1226 pub applied_ids: Vec<String>,
1228 pub tool_search: Option<crate::llm_driver_registry::ToolSearchConfig>,
1230 pub prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig>,
1232 pub tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>>,
1234 pub tool_call_hooks: Vec<Arc<dyn ToolCallHook>>,
1236 pub mcp_servers: ScopedMcpServers,
1238 }
1244
1245#[derive(Debug, Clone, PartialEq, Eq)]
1246pub struct SystemPromptAttribution {
1247 pub capability_id: String,
1248 pub content: String,
1249}
1250
1251impl CollectedCapabilities {
1252 pub fn system_prompt_prefix(&self) -> Option<String> {
1255 if self.system_prompt_parts.is_empty() {
1256 None
1257 } else {
1258 Some(self.system_prompt_parts.join("\n\n"))
1259 }
1260 }
1261
1262 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1266 for (provider, config) in &self.message_filter_providers {
1268 provider.apply_filters(query, config);
1269 }
1270 }
1271
1272 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1275 for (provider, config) in &self.message_filter_providers {
1276 provider.post_load(messages, config);
1277 }
1278 }
1279
1280 pub fn has_message_filters(&self) -> bool {
1282 !self.message_filter_providers.is_empty()
1283 }
1284}
1285
1286pub struct CollectedMessageFilters {
1293 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1295}
1296
1297pub struct CollectedModelViewProviders {
1299 pub model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)>,
1301}
1302
1303impl CollectedMessageFilters {
1309 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1311 for (provider, config) in &self.message_filter_providers {
1312 provider.apply_filters(query, config);
1313 }
1314 }
1315
1316 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1318 for (provider, config) in &self.message_filter_providers {
1319 provider.post_load(messages, config);
1320 }
1321 }
1322}
1323
1324impl CollectedModelViewProviders {
1325 pub fn apply_model_view(
1327 &self,
1328 mut messages: Vec<Message>,
1329 context: &ModelViewContext<'_>,
1330 ) -> Vec<Message> {
1331 for (provider, config) in &self.model_view_providers {
1332 messages = provider.apply_model_view(messages, config, context);
1333 }
1334 messages
1335 }
1336}
1337
1338pub fn collect_message_filters_only(
1344 capability_configs: &[AgentCapabilityConfig],
1345 registry: &CapabilityRegistry,
1346) -> CollectedMessageFilters {
1347 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1348 Vec::new();
1349
1350 for cap_config in capability_configs {
1351 let cap_id = cap_config.capability_ref.as_str();
1352 if let Some(capability) = registry.get(cap_id) {
1353 if capability.status() != CapabilityStatus::Available {
1354 continue;
1355 }
1356 let effective: &dyn Capability = capability
1359 .resolve_for_model(None)
1360 .unwrap_or_else(|| capability.as_ref());
1361 if let Some(provider) = effective.message_filter_provider() {
1362 message_filter_providers.push((provider, cap_config.config.clone()));
1363 }
1364 }
1365 }
1366
1367 message_filter_providers.sort_by_key(|(p, _)| p.priority());
1368
1369 CollectedMessageFilters {
1370 message_filter_providers,
1371 }
1372}
1373
1374pub fn collect_model_view_providers(
1381 capability_configs: &[AgentCapabilityConfig],
1382 registry: &CapabilityRegistry,
1383 model: Option<&str>,
1384) -> CollectedModelViewProviders {
1385 let mut model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)> = Vec::new();
1386
1387 for cap_config in capability_configs {
1388 let cap_id = cap_config.capability_ref.as_str();
1389 if let Some(capability) = registry.get(cap_id) {
1390 if capability.status() != CapabilityStatus::Available {
1391 continue;
1392 }
1393 let effective: &dyn Capability = capability
1394 .resolve_for_model(model)
1395 .unwrap_or_else(|| capability.as_ref());
1396 if let Some(provider) = effective.model_view_provider() {
1397 model_view_providers.push((provider, cap_config.config.clone()));
1398 }
1399 }
1400 }
1401
1402 model_view_providers.sort_by_key(|(p, _)| p.priority());
1403
1404 CollectedModelViewProviders {
1405 model_view_providers,
1406 }
1407}
1408
1409pub fn collect_capability_mcp_servers(
1410 capability_configs: &[AgentCapabilityConfig],
1411 registry: &CapabilityRegistry,
1412) -> ScopedMcpServers {
1413 let mut servers = ScopedMcpServers::default();
1414
1415 for cap_config in capability_configs {
1416 let cap_id = cap_config.capability_ref.as_str();
1417 if is_declarative_capability(cap_id) {
1418 if let Ok(definition) =
1419 serde_json::from_value::<DeclarativeCapabilityDefinition>(cap_config.config.clone())
1420 {
1421 if definition.status != CapabilityStatus::Available {
1422 continue;
1423 }
1424 if let Some(contributed) = definition.mcp_servers {
1425 servers = merge_scoped_mcp_servers(&servers, &contributed);
1426 }
1427 }
1428 continue;
1429 }
1430 if let Some(capability) = registry.get(cap_id) {
1431 if capability.status() != CapabilityStatus::Available {
1432 continue;
1433 }
1434 servers = merge_scoped_mcp_servers(
1435 &servers,
1436 &capability.mcp_servers_with_config(&cap_config.config),
1437 );
1438 }
1439 }
1440
1441 servers
1442}
1443
1444pub const MAX_RESOLVED_CAPABILITIES: usize = 100;
1451
1452#[derive(Debug, Clone, PartialEq, Eq)]
1454pub enum DependencyError {
1455 CircularDependency {
1457 capability_id: String,
1459 chain: Vec<String>,
1461 },
1462 TooManyCapabilities {
1464 count: usize,
1466 max: usize,
1468 },
1469}
1470
1471impl std::fmt::Display for DependencyError {
1472 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1473 match self {
1474 DependencyError::CircularDependency {
1475 capability_id,
1476 chain,
1477 } => {
1478 write!(
1479 f,
1480 "Circular dependency detected: {} depends on itself via chain: {} -> {}",
1481 capability_id,
1482 chain.join(" -> "),
1483 capability_id
1484 )
1485 }
1486 DependencyError::TooManyCapabilities { count, max } => {
1487 write!(
1488 f,
1489 "Too many capabilities after resolution: {} (max: {})",
1490 count, max
1491 )
1492 }
1493 }
1494 }
1495}
1496
1497impl std::error::Error for DependencyError {}
1498
1499#[derive(Debug, Clone)]
1501pub struct ResolvedCapabilities {
1502 pub resolved_ids: Vec<String>,
1505 pub added_as_dependencies: Vec<String>,
1507 pub user_selected: Vec<String>,
1509}
1510
1511pub fn resolve_dependencies(
1531 selected_ids: &[String],
1532 registry: &CapabilityRegistry,
1533) -> Result<ResolvedCapabilities, DependencyError> {
1534 use std::collections::HashSet;
1535
1536 let user_selected: HashSet<String> = selected_ids.iter().cloned().collect();
1537 let mut resolved: Vec<String> = Vec::new();
1538 let mut resolved_set: HashSet<String> = HashSet::new();
1539 let mut added_as_dependencies: Vec<String> = Vec::new();
1540
1541 for cap_id in selected_ids {
1543 resolve_single_capability(
1544 cap_id,
1545 registry,
1546 &mut resolved,
1547 &mut resolved_set,
1548 &mut added_as_dependencies,
1549 &user_selected,
1550 &mut Vec::new(), )?;
1552 }
1553
1554 if resolved.len() > MAX_RESOLVED_CAPABILITIES {
1556 return Err(DependencyError::TooManyCapabilities {
1557 count: resolved.len(),
1558 max: MAX_RESOLVED_CAPABILITIES,
1559 });
1560 }
1561
1562 Ok(ResolvedCapabilities {
1563 resolved_ids: resolved,
1564 added_as_dependencies,
1565 user_selected: selected_ids.to_vec(),
1566 })
1567}
1568
1569pub fn resolve_capability_configs(
1574 selected_configs: &[AgentCapabilityConfig],
1575 registry: &CapabilityRegistry,
1576) -> Result<Vec<AgentCapabilityConfig>, DependencyError> {
1577 let mut selected_ids: Vec<String> = Vec::new();
1578 for config in selected_configs {
1579 if is_declarative_capability(config.capability_id())
1580 && let Ok(definition) =
1581 serde_json::from_value::<DeclarativeCapabilityDefinition>(config.config.clone())
1582 {
1583 selected_ids.extend(definition.dependencies);
1584 }
1585 selected_ids.push(config.capability_id().to_string());
1586 }
1587 let resolved = resolve_dependencies(&selected_ids, registry)?;
1588
1589 let explicit_configs: std::collections::HashMap<String, serde_json::Value> = selected_configs
1590 .iter()
1591 .map(|config| (config.capability_id().to_string(), config.config.clone()))
1592 .collect();
1593
1594 Ok(resolved
1595 .resolved_ids
1596 .into_iter()
1597 .map(|capability_id| {
1598 explicit_configs
1599 .get(&capability_id)
1600 .cloned()
1601 .map(|config| AgentCapabilityConfig::with_config(capability_id.clone(), config))
1602 .unwrap_or_else(|| AgentCapabilityConfig::new(capability_id))
1603 })
1604 .collect())
1605}
1606
1607fn resolve_single_capability(
1609 cap_id: &str,
1610 registry: &CapabilityRegistry,
1611 resolved: &mut Vec<String>,
1612 resolved_set: &mut std::collections::HashSet<String>,
1613 added_as_dependencies: &mut Vec<String>,
1614 user_selected: &std::collections::HashSet<String>,
1615 visiting: &mut Vec<String>,
1616) -> Result<(), DependencyError> {
1617 if resolved_set.contains(cap_id) {
1619 return Ok(());
1620 }
1621
1622 if visiting.contains(&cap_id.to_string()) {
1624 return Err(DependencyError::CircularDependency {
1625 capability_id: cap_id.to_string(),
1626 chain: visiting.clone(),
1627 });
1628 }
1629
1630 let capability = match registry.get(cap_id) {
1632 Some(cap) => cap,
1633 None => {
1634 if is_declarative_capability(cap_id) && !resolved_set.contains(cap_id) {
1635 resolved.push(cap_id.to_string());
1636 resolved_set.insert(cap_id.to_string());
1637 if !user_selected.contains(cap_id) {
1638 added_as_dependencies.push(cap_id.to_string());
1639 }
1640 }
1641 return Ok(());
1642 }
1643 };
1644
1645 visiting.push(cap_id.to_string());
1647
1648 for dep_id in capability.dependencies() {
1650 resolve_single_capability(
1651 dep_id,
1652 registry,
1653 resolved,
1654 resolved_set,
1655 added_as_dependencies,
1656 user_selected,
1657 visiting,
1658 )?;
1659 }
1660
1661 visiting.pop();
1663
1664 if !resolved_set.contains(cap_id) {
1666 resolved.push(cap_id.to_string());
1667 resolved_set.insert(cap_id.to_string());
1668
1669 if !user_selected.contains(cap_id) {
1671 added_as_dependencies.push(cap_id.to_string());
1672 }
1673 }
1674
1675 Ok(())
1676}
1677
1678pub fn compute_features(capability_ids: &[String], registry: &CapabilityRegistry) -> Vec<String> {
1683 use std::collections::HashSet;
1684
1685 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
1686 Ok(resolved) => resolved.resolved_ids,
1687 Err(_) => capability_ids.to_vec(),
1688 };
1689
1690 let mut seen = HashSet::new();
1691 let mut features = Vec::new();
1692 for cap_id in &resolved_ids {
1693 if let Some(cap) = registry.get(cap_id) {
1694 for feature in cap.features() {
1695 if seen.insert(feature) {
1696 features.push(feature.to_string());
1697 }
1698 }
1699 }
1700 }
1701 features
1702}
1703
1704pub fn get_dependencies(cap_id: &str, registry: &CapabilityRegistry) -> Vec<String> {
1707 registry
1708 .get(cap_id)
1709 .map(|cap| cap.dependencies().iter().map(|s| s.to_string()).collect())
1710 .unwrap_or_default()
1711}
1712
1713pub async fn collect_capabilities(
1729 capability_ids: &[String],
1730 registry: &CapabilityRegistry,
1731 ctx: &SystemPromptContext,
1732) -> CollectedCapabilities {
1733 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
1736 Ok(resolved) => resolved.resolved_ids,
1737 Err(e) => {
1738 tracing::warn!("Failed to resolve capability dependencies: {}", e);
1739 capability_ids.to_vec()
1740 }
1741 };
1742
1743 let configs: Vec<AgentCapabilityConfig> = resolved_ids
1745 .iter()
1746 .map(|id| AgentCapabilityConfig {
1747 capability_ref: CapabilityId::new(id),
1748 config: serde_json::Value::Object(serde_json::Map::new()),
1749 })
1750 .collect();
1751
1752 collect_capabilities_with_configs(&configs, registry, ctx).await
1753}
1754
1755pub async fn collect_capabilities_with_configs(
1766 capability_configs: &[AgentCapabilityConfig],
1767 registry: &CapabilityRegistry,
1768 ctx: &SystemPromptContext,
1769) -> CollectedCapabilities {
1770 let mut system_prompt_parts: Vec<String> = Vec::new();
1771 let mut system_prompt_attributions: Vec<SystemPromptAttribution> = Vec::new();
1772 let mut tools: Vec<Box<dyn Tool>> = Vec::new();
1773 let mut tool_definitions: Vec<ToolDefinition> = Vec::new();
1774 let mut mounts: Vec<MountPoint> = Vec::new();
1775 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1776 Vec::new();
1777 let mut applied_ids: Vec<String> = Vec::new();
1778 let mut tool_search: Option<crate::llm_driver_registry::ToolSearchConfig> = None;
1779 let mut prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig> = None;
1780 let mut tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>> = Vec::new();
1781 let mut tool_call_hooks: Vec<Arc<dyn ToolCallHook>> = Vec::new();
1782 let mut mcp_servers = ScopedMcpServers::default();
1783
1784 for cap_config in capability_configs {
1785 let cap_id = cap_config.capability_ref.as_str();
1786 if is_declarative_capability(cap_id) {
1787 match serde_json::from_value::<DeclarativeCapabilityDefinition>(
1788 cap_config.config.clone(),
1789 ) {
1790 Ok(definition) => {
1791 if definition.status != CapabilityStatus::Available {
1792 continue;
1793 }
1794
1795 if let Some(prompt) = definition.system_prompt.as_deref() {
1796 let contribution =
1797 format!("<capability id=\"{}\">\n{}\n</capability>", cap_id, prompt);
1798 system_prompt_attributions.push(SystemPromptAttribution {
1799 capability_id: cap_id.to_string(),
1800 content: contribution.clone(),
1801 });
1802 system_prompt_parts.push(contribution);
1803 }
1804
1805 mounts.extend(definition.mounts(cap_id));
1806 if let Some(ref servers) = definition.mcp_servers {
1807 mcp_servers = merge_scoped_mcp_servers(&mcp_servers, servers);
1808 }
1809 for skill in definition.skill_contributions() {
1810 mounts.push(skill.to_mount(cap_id));
1811 }
1812
1813 applied_ids.push(cap_id.to_string());
1814 }
1815 Err(error) => {
1816 tracing::warn!(
1817 capability_id = %cap_id,
1818 error = %error,
1819 "Skipping invalid declarative capability config"
1820 );
1821 }
1822 }
1823 continue;
1824 }
1825 if let Some(capability) = registry.get(cap_id) {
1826 if capability.status() != CapabilityStatus::Available {
1828 continue;
1829 }
1830
1831 let effective: &dyn Capability =
1843 match capability.resolve_for_model(ctx.model.as_deref()) {
1844 Some(inner) => inner,
1845 None => capability.as_ref(),
1846 };
1847 let effective_id = effective.id();
1848
1849 if let Some(contribution) = effective
1851 .system_prompt_contribution_with_config(ctx, &cap_config.config)
1852 .await
1853 {
1854 system_prompt_attributions.push(SystemPromptAttribution {
1855 capability_id: cap_id.to_string(),
1856 content: contribution.clone(),
1857 });
1858 system_prompt_parts.push(contribution);
1859 }
1860
1861 tools.extend(effective.tools_with_config(&cap_config.config));
1863 tool_definition_hooks
1864 .extend(effective.tool_definition_hooks_with_config(&cap_config.config));
1865 tool_call_hooks.extend(effective.tool_call_hooks());
1866 let cap_category = effective.category();
1871 for def in effective.tool_definitions() {
1872 let def = match (def.category(), cap_category) {
1873 (None, Some(cat)) => def.with_category(cat),
1874 _ => def,
1875 }
1876 .with_capability_attribution(cap_id, Some(capability.name()));
1877 tool_definitions.push(def);
1878 }
1879
1880 if effective_id == OPENAI_TOOL_SEARCH_CAPABILITY_ID {
1885 let threshold = cap_config
1887 .config
1888 .get("threshold")
1889 .and_then(|v| v.as_u64())
1890 .map(|v| v as usize)
1891 .unwrap_or(DEFAULT_TOOL_SEARCH_THRESHOLD);
1892 tool_search = Some(crate::llm_driver_registry::ToolSearchConfig {
1893 enabled: true,
1894 threshold,
1895 });
1896 }
1897
1898 if cap_id == PROMPT_CACHING_CAPABILITY_ID {
1899 let strategy = cap_config
1900 .config
1901 .get("strategy")
1902 .and_then(|v| v.as_str())
1903 .map(|value| match value {
1904 "auto" => crate::llm_driver_registry::PromptCacheStrategy::Auto,
1905 _ => crate::llm_driver_registry::PromptCacheStrategy::Auto,
1906 })
1907 .unwrap_or(crate::llm_driver_registry::PromptCacheStrategy::Auto);
1908 let gemini_cached_content = cap_config
1909 .config
1910 .get("gemini_cached_content")
1911 .and_then(|v| v.as_str())
1912 .map(str::to_string);
1913 prompt_cache = Some(crate::llm_driver_registry::PromptCacheConfig {
1914 enabled: true,
1915 strategy,
1916 gemini_cached_content,
1917 });
1918 }
1919
1920 mounts.extend(effective.mounts());
1922
1923 mcp_servers = merge_scoped_mcp_servers(
1924 &mcp_servers,
1925 &effective.mcp_servers_with_config(&cap_config.config),
1926 );
1927
1928 for skill in effective.contribute_skills() {
1932 mounts.push(skill.to_mount(cap_id));
1933 }
1934
1935 if let Some(provider) = effective.message_filter_provider() {
1937 message_filter_providers.push((provider, cap_config.config.clone()));
1938 }
1939
1940 applied_ids.push(cap_id.to_string());
1941 }
1942 }
1943
1944 if !applied_ids
1956 .iter()
1957 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID)
1958 && tool_definitions
1959 .iter()
1960 .any(|def| def.hints().supports_background == Some(true))
1961 && let Some(bg_cap) = registry.get(BACKGROUND_EXECUTION_CAPABILITY_ID)
1962 && bg_cap.status() == CapabilityStatus::Available
1963 {
1964 tools.extend(bg_cap.tools());
1965 let cap_category = bg_cap.category();
1966 for def in bg_cap.tool_definitions() {
1967 let def = match (def.category(), cap_category) {
1968 (None, Some(cat)) => def.with_category(cat),
1969 _ => def,
1970 }
1971 .with_capability_attribution(BACKGROUND_EXECUTION_CAPABILITY_ID, Some(bg_cap.name()));
1972 tool_definitions.push(def);
1973 }
1974 applied_ids.push(BACKGROUND_EXECUTION_CAPABILITY_ID.to_string());
1975 }
1976
1977 message_filter_providers.sort_by_key(|(p, _)| p.priority());
1979
1980 CollectedCapabilities {
1981 system_prompt_parts,
1982 system_prompt_attributions,
1983 tools,
1984 tool_definitions,
1985 mounts,
1986 message_filter_providers,
1987 applied_ids,
1988 tool_search,
1989 prompt_cache,
1990 tool_definition_hooks,
1991 tool_call_hooks,
1992 mcp_servers,
1993 }
1994}
1995
1996pub struct AppliedCapabilities {
2002 pub runtime_agent: RuntimeAgent,
2004 pub tool_registry: ToolRegistry,
2006 pub applied_ids: Vec<String>,
2008}
2009
2010pub async fn apply_capabilities(
2047 base_runtime_agent: RuntimeAgent,
2048 capability_ids: &[String],
2049 registry: &CapabilityRegistry,
2050 ctx: &SystemPromptContext,
2051) -> AppliedCapabilities {
2052 let collected = collect_capabilities(capability_ids, registry, ctx).await;
2053
2054 let final_system_prompt = match collected.system_prompt_prefix() {
2056 Some(prefix) => format!(
2057 "{}\n\n<system-prompt>\n{}\n</system-prompt>",
2058 prefix, base_runtime_agent.system_prompt
2059 ),
2060 None => base_runtime_agent.system_prompt,
2061 };
2062
2063 let mut tool_registry = ToolRegistry::new();
2065 for tool in collected.tools {
2066 tool_registry.register_boxed(tool);
2067 }
2068
2069 let mut tools = collected.tool_definitions;
2071 for hook in &collected.tool_definition_hooks {
2072 tools = hook.transform(tools);
2073 }
2074
2075 let runtime_agent = RuntimeAgent {
2076 system_prompt: final_system_prompt,
2077 model: base_runtime_agent.model,
2078 tools,
2079 max_iterations: base_runtime_agent.max_iterations,
2080 temperature: base_runtime_agent.temperature,
2081 max_tokens: base_runtime_agent.max_tokens,
2082 tool_search: collected.tool_search,
2083 prompt_cache: collected.prompt_cache,
2084 network_access: base_runtime_agent.network_access,
2085 };
2086
2087 AppliedCapabilities {
2088 runtime_agent,
2089 tool_registry,
2090 applied_ids: collected.applied_ids,
2091 }
2092}
2093
2094#[cfg(test)]
2099mod tests {
2100 use super::*;
2101 use crate::typed_id::SessionId;
2102 use std::collections::BTreeSet;
2103 use uuid::Uuid;
2104
2105 static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2107
2108 fn lock_env() -> std::sync::MutexGuard<'static, ()> {
2109 ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner())
2110 }
2111
2112 fn test_ctx() -> SystemPromptContext {
2114 SystemPromptContext::without_file_store(SessionId::new())
2115 }
2116
2117 fn expected_core_builtin_ids() -> BTreeSet<&'static str> {
2119 let mut ids = [
2120 "agent_instructions",
2121 "human_intent",
2122 "budgeting",
2123 "self_budget",
2124 "noop",
2125 "current_time",
2126 "research",
2127 "platform_management",
2128 "session_file_system",
2129 "workspace_volumes",
2130 "session_storage",
2131 "session",
2132 "session_sql_database",
2133 "test_math",
2134 "test_weather",
2135 "stateless_todo_list",
2136 "web_fetch",
2137 "virtual_bash",
2138 "background_execution",
2139 "session_schedule",
2140 "btw",
2141 "infinity_context",
2142 "compaction",
2143 "memory",
2144 "openai_tool_search",
2145 "tool_search",
2146 "auto_tool_search",
2147 "prompt_caching",
2148 "skills",
2149 "subagents",
2150 "system_commands",
2151 "sample_data",
2152 "data_knowledge",
2153 "knowledge_base",
2154 "tool_output_persistence",
2155 "fake_warehouse",
2156 "fake_aws",
2157 "fake_crm",
2158 "fake_financial",
2159 "loop_detection",
2160 "prompt_canary_guardrail",
2161 "user_hooks",
2162 ]
2163 .into_iter()
2164 .collect::<BTreeSet<_>>();
2165 if cfg!(feature = "ui-capabilities") {
2166 ids.insert("openui");
2167 ids.insert("a2ui");
2168 }
2169 ids
2170 }
2171
2172 fn expected_dev_builtin_ids() -> BTreeSet<&'static str> {
2174 let mut ids = expected_core_builtin_ids();
2175 ids.insert("agent_handoff");
2176 ids.insert("a2a_agent_delegation");
2177 ids
2178 }
2179
2180 fn registry_ids(registry: &CapabilityRegistry) -> BTreeSet<&str> {
2181 registry.capabilities.keys().map(String::as_str).collect()
2182 }
2183
2184 #[test]
2194 fn test_capability_registry_with_builtins_dev() {
2195 let _lock = lock_env();
2197 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2198 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
2199 assert_eq!(registry_ids(®istry), expected_dev_builtin_ids());
2200 assert!(registry.has("agent_handoff"));
2201 assert!(registry.has("a2a_agent_delegation"));
2202 }
2203
2204 #[test]
2205 fn test_capability_registry_with_builtins_prod() {
2206 let _lock = lock_env();
2208 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2209 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2210 assert_eq!(registry_ids(®istry), expected_core_builtin_ids());
2211 assert!(!registry.has("docker_container"));
2213 assert!(!registry.has("agent_handoff"));
2214 assert!(!registry.has("a2a_agent_delegation"));
2215 }
2216
2217 #[test]
2218 fn test_agent_delegation_enabled_by_env_in_prod() {
2219 let _lock = lock_env();
2221 unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "true") };
2222 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2223 assert!(registry.has("agent_handoff"));
2224 assert!(registry.has("a2a_agent_delegation"));
2225 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2226 }
2227
2228 #[test]
2229 fn test_agent_delegation_disabled_by_env_in_dev() {
2230 let _lock = lock_env();
2232 unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "false") };
2233 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
2234 assert!(!registry.has("agent_handoff"));
2235 assert!(!registry.has("a2a_agent_delegation"));
2236 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2237 }
2238
2239 #[test]
2240 fn test_capability_registry_get() {
2241 let registry = CapabilityRegistry::with_builtins();
2242
2243 let noop = registry.get("noop").unwrap();
2244 assert_eq!(noop.id(), "noop");
2245 assert_eq!(noop.name(), "No-Op");
2246 assert_eq!(noop.status(), CapabilityStatus::Available);
2247 }
2248
2249 #[test]
2250 fn test_capability_registry_blueprint_with_capability() {
2251 struct BlueprintProviderCapability;
2252
2253 impl Capability for BlueprintProviderCapability {
2254 fn id(&self) -> &str {
2255 "blueprint_provider"
2256 }
2257 fn name(&self) -> &str {
2258 "Blueprint Provider"
2259 }
2260 fn description(&self) -> &str {
2261 "Capability that provides a blueprint for tests"
2262 }
2263 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
2264 vec![AgentBlueprint {
2265 id: "test_blueprint",
2266 name: "Test Blueprint",
2267 description: "Blueprint for capability registry tests",
2268 model: BlueprintModel::Inherit,
2269 system_prompt: "Test prompt",
2270 tools: vec![],
2271 max_turns: None,
2272 config_schema: None,
2273 }]
2274 }
2275 }
2276
2277 let mut registry = CapabilityRegistry::new();
2278 registry.register(BlueprintProviderCapability);
2279
2280 let (capability_id, blueprint) = registry
2281 .blueprint_with_capability("test_blueprint")
2282 .expect("blueprint should resolve with capability id");
2283 assert_eq!(capability_id, "blueprint_provider");
2284 assert_eq!(blueprint.id, "test_blueprint");
2285 }
2286
2287 #[test]
2288 fn test_capability_registry_builder() {
2289 let registry = CapabilityRegistry::builder()
2290 .capability(NoopCapability)
2291 .capability(CurrentTimeCapability)
2292 .build();
2293
2294 assert!(registry.has("noop"));
2295 assert!(registry.has("current_time"));
2296 assert_eq!(registry.len(), 2);
2297 }
2298
2299 #[test]
2300 fn test_capability_status() {
2301 let registry = CapabilityRegistry::with_builtins();
2302
2303 let current_time = registry.get("current_time").unwrap();
2304 assert_eq!(current_time.status(), CapabilityStatus::Available);
2305
2306 let research = registry.get("research").unwrap();
2307 assert_eq!(research.status(), CapabilityStatus::ComingSoon);
2308 }
2309
2310 #[test]
2311 fn test_capability_icons_and_categories() {
2312 let registry = CapabilityRegistry::with_builtins();
2313
2314 let noop = registry.get("noop").unwrap();
2315 assert_eq!(noop.icon(), Some("circle-off"));
2316 assert_eq!(noop.category(), Some("Testing"));
2317
2318 let current_time = registry.get("current_time").unwrap();
2319 assert_eq!(current_time.icon(), Some("clock"));
2320 assert_eq!(current_time.category(), Some("Utilities"));
2321 }
2322
2323 #[test]
2324 fn test_system_prompt_preview_default_delegates_to_addition() {
2325 let registry = CapabilityRegistry::with_builtins();
2326
2327 let test_math = registry.get("test_math").unwrap();
2329 assert_eq!(
2330 test_math.system_prompt_preview().as_deref(),
2331 test_math.system_prompt_addition()
2332 );
2333
2334 let current_time = registry.get("current_time").unwrap();
2336 assert!(current_time.system_prompt_preview().is_none());
2337 assert!(current_time.system_prompt_addition().is_none());
2338 }
2339
2340 #[test]
2341 fn test_system_prompt_preview_dynamic_capability() {
2342 let registry = CapabilityRegistry::with_builtins();
2343 let cap = registry.get("agent_instructions").unwrap();
2344
2345 assert!(cap.system_prompt_addition().is_none());
2347 assert!(cap.system_prompt_preview().is_some());
2348 assert!(cap.system_prompt_preview().unwrap().contains("AGENTS.md"));
2349 }
2350
2351 #[tokio::test]
2356 async fn test_apply_capabilities_empty() {
2357 let registry = CapabilityRegistry::with_builtins();
2358 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2359
2360 let applied =
2361 apply_capabilities(base_runtime_agent.clone(), &[], ®istry, &test_ctx()).await;
2362
2363 assert_eq!(
2364 applied.runtime_agent.system_prompt,
2365 base_runtime_agent.system_prompt
2366 );
2367 assert!(applied.tool_registry.is_empty());
2368 assert!(applied.applied_ids.is_empty());
2369 }
2370
2371 #[tokio::test]
2372 async fn test_apply_capabilities_noop() {
2373 let registry = CapabilityRegistry::with_builtins();
2374 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2375
2376 let applied = apply_capabilities(
2377 base_runtime_agent.clone(),
2378 &["noop".to_string()],
2379 ®istry,
2380 &test_ctx(),
2381 )
2382 .await;
2383
2384 assert_eq!(
2386 applied.runtime_agent.system_prompt,
2387 base_runtime_agent.system_prompt
2388 );
2389 assert!(applied.tool_registry.is_empty());
2390 assert_eq!(applied.applied_ids, vec!["noop"]);
2391 }
2392
2393 #[tokio::test]
2394 async fn test_apply_capabilities_current_time() {
2395 let registry = CapabilityRegistry::with_builtins();
2396 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2397
2398 let applied = apply_capabilities(
2399 base_runtime_agent.clone(),
2400 &["current_time".to_string()],
2401 ®istry,
2402 &test_ctx(),
2403 )
2404 .await;
2405
2406 assert_eq!(
2408 applied.runtime_agent.system_prompt,
2409 base_runtime_agent.system_prompt
2410 );
2411 assert!(applied.tool_registry.has("get_current_time"));
2412 assert_eq!(applied.tool_registry.len(), 1);
2413 assert_eq!(applied.applied_ids, vec!["current_time"]);
2414 }
2415
2416 #[tokio::test]
2417 async fn test_apply_capabilities_skips_coming_soon() {
2418 let registry = CapabilityRegistry::with_builtins();
2419 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2420
2421 let applied = apply_capabilities(
2423 base_runtime_agent.clone(),
2424 &["research".to_string()],
2425 ®istry,
2426 &test_ctx(),
2427 )
2428 .await;
2429
2430 assert_eq!(
2432 applied.runtime_agent.system_prompt,
2433 base_runtime_agent.system_prompt
2434 );
2435 assert!(applied.applied_ids.is_empty()); }
2437
2438 #[tokio::test]
2439 async fn test_apply_capabilities_multiple() {
2440 let registry = CapabilityRegistry::with_builtins();
2441 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2442
2443 let applied = apply_capabilities(
2444 base_runtime_agent.clone(),
2445 &["noop".to_string(), "current_time".to_string()],
2446 ®istry,
2447 &test_ctx(),
2448 )
2449 .await;
2450
2451 assert!(applied.tool_registry.has("get_current_time"));
2452 assert_eq!(applied.applied_ids, vec!["noop", "current_time"]);
2453 }
2454
2455 #[tokio::test]
2456 async fn test_apply_capabilities_preserves_order() {
2457 let registry = CapabilityRegistry::with_builtins();
2458 let base_runtime_agent = RuntimeAgent::new("Base prompt.", "gpt-5.2");
2459
2460 let applied = apply_capabilities(
2462 base_runtime_agent,
2463 &["current_time".to_string(), "noop".to_string()],
2464 ®istry,
2465 &test_ctx(),
2466 )
2467 .await;
2468
2469 assert_eq!(applied.applied_ids, vec!["current_time", "noop"]);
2470 }
2471
2472 #[tokio::test]
2473 async fn test_apply_capabilities_test_math() {
2474 let registry = CapabilityRegistry::with_builtins();
2475 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2476
2477 let applied = apply_capabilities(
2478 base_runtime_agent.clone(),
2479 &["test_math".to_string()],
2480 ®istry,
2481 &test_ctx(),
2482 )
2483 .await;
2484
2485 assert!(
2487 !applied
2488 .runtime_agent
2489 .system_prompt
2490 .contains("<capability id=\"test_math\">")
2491 );
2492 assert!(
2494 applied
2495 .runtime_agent
2496 .system_prompt
2497 .contains("You are a helpful assistant.")
2498 );
2499 assert!(applied.tool_registry.has("add"));
2500 assert!(applied.tool_registry.has("subtract"));
2501 assert!(applied.tool_registry.has("multiply"));
2502 assert!(applied.tool_registry.has("divide"));
2503 assert_eq!(applied.tool_registry.len(), 4);
2504 }
2505
2506 #[tokio::test]
2507 async fn test_apply_capabilities_test_weather() {
2508 let registry = CapabilityRegistry::with_builtins();
2509 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2510
2511 let applied = apply_capabilities(
2512 base_runtime_agent.clone(),
2513 &["test_weather".to_string()],
2514 ®istry,
2515 &test_ctx(),
2516 )
2517 .await;
2518
2519 assert!(
2521 !applied
2522 .runtime_agent
2523 .system_prompt
2524 .contains("<capability id=\"test_weather\">")
2525 );
2526 assert!(applied.tool_registry.has("get_weather"));
2527 assert!(applied.tool_registry.has("get_forecast"));
2528 assert_eq!(applied.tool_registry.len(), 2);
2529 }
2530
2531 #[tokio::test]
2532 async fn test_apply_capabilities_test_math_and_test_weather() {
2533 let registry = CapabilityRegistry::with_builtins();
2534 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2535
2536 let applied = apply_capabilities(
2537 base_runtime_agent.clone(),
2538 &["test_math".to_string(), "test_weather".to_string()],
2539 ®istry,
2540 &test_ctx(),
2541 )
2542 .await;
2543
2544 assert_eq!(applied.tool_registry.len(), 6); assert!(applied.tool_registry.has("add"));
2547 assert!(applied.tool_registry.has("get_weather"));
2548 }
2549
2550 #[tokio::test]
2551 async fn test_apply_capabilities_stateless_todo_list() {
2552 let registry = CapabilityRegistry::with_builtins();
2553 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2554
2555 let applied = apply_capabilities(
2556 base_runtime_agent.clone(),
2557 &["stateless_todo_list".to_string()],
2558 ®istry,
2559 &test_ctx(),
2560 )
2561 .await;
2562
2563 assert!(
2565 applied
2566 .runtime_agent
2567 .system_prompt
2568 .contains("Task Management")
2569 );
2570 assert!(applied.runtime_agent.system_prompt.contains("write_todos"));
2571 assert!(applied.tool_registry.has("write_todos"));
2572 assert_eq!(applied.tool_registry.len(), 1);
2573 }
2574
2575 #[tokio::test]
2576 async fn test_apply_capabilities_web_fetch() {
2577 let registry = CapabilityRegistry::with_builtins();
2578 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2579
2580 let applied = apply_capabilities(
2581 base_runtime_agent.clone(),
2582 &["web_fetch".to_string()],
2583 ®istry,
2584 &test_ctx(),
2585 )
2586 .await;
2587
2588 assert!(
2590 applied
2591 .runtime_agent
2592 .system_prompt
2593 .contains(&base_runtime_agent.system_prompt)
2594 );
2595 assert!(applied.runtime_agent.system_prompt.contains("web_fetch"));
2596 assert!(applied.tool_registry.has("web_fetch"));
2597 assert_eq!(applied.tool_registry.len(), 1);
2598 }
2599
2600 #[tokio::test]
2605 async fn test_xml_tags_wrap_capability_prompts() {
2606 let registry = CapabilityRegistry::with_builtins();
2607 let collected =
2608 collect_capabilities(&["stateless_todo_list".to_string()], ®istry, &test_ctx())
2609 .await;
2610
2611 assert_eq!(collected.system_prompt_parts.len(), 1);
2612 let part = &collected.system_prompt_parts[0];
2613 assert!(part.starts_with("<capability id=\"stateless_todo_list\">"));
2614 assert!(part.ends_with("</capability>"));
2615 assert!(part.contains("Task Management"));
2616 }
2617
2618 #[tokio::test]
2619 async fn test_xml_tags_multiple_capabilities() {
2620 let registry = CapabilityRegistry::with_builtins();
2621 let collected = collect_capabilities(
2622 &[
2623 "stateless_todo_list".to_string(),
2624 "session_schedule".to_string(),
2625 ],
2626 ®istry,
2627 &test_ctx(),
2628 )
2629 .await;
2630
2631 assert_eq!(collected.system_prompt_parts.len(), 2);
2632 assert!(
2633 collected.system_prompt_parts[0].starts_with("<capability id=\"stateless_todo_list\">")
2634 );
2635 assert!(
2636 collected.system_prompt_parts[1].starts_with("<capability id=\"session_schedule\">")
2637 );
2638
2639 let prefix = collected.system_prompt_prefix().unwrap();
2640 assert!(prefix.contains("</capability>\n\n<capability"));
2642 }
2643
2644 #[tokio::test]
2645 async fn test_xml_tags_system_prompt_wrapping() {
2646 let registry = CapabilityRegistry::with_builtins();
2647 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2648
2649 let applied = apply_capabilities(
2650 base,
2651 &["stateless_todo_list".to_string()],
2652 ®istry,
2653 &test_ctx(),
2654 )
2655 .await;
2656
2657 let prompt = &applied.runtime_agent.system_prompt;
2658 assert!(prompt.contains("<capability id=\"stateless_todo_list\">"));
2660 assert!(prompt.contains("</capability>"));
2661 assert!(prompt.contains("<system-prompt>\nYou are helpful.\n</system-prompt>"));
2663 }
2664
2665 #[tokio::test]
2666 async fn test_no_xml_wrapping_without_capabilities() {
2667 let registry = CapabilityRegistry::with_builtins();
2668 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2669
2670 let applied = apply_capabilities(base, &[], ®istry, &test_ctx()).await;
2671
2672 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2674 assert!(
2675 !applied
2676 .runtime_agent
2677 .system_prompt
2678 .contains("<system-prompt>")
2679 );
2680 }
2681
2682 #[tokio::test]
2683 async fn test_no_xml_wrapping_for_noop_capability() {
2684 let registry = CapabilityRegistry::with_builtins();
2685 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2686
2687 let applied = apply_capabilities(base, &["noop".to_string()], ®istry, &test_ctx()).await;
2689
2690 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2691 assert!(
2692 !applied
2693 .runtime_agent
2694 .system_prompt
2695 .contains("<system-prompt>")
2696 );
2697 }
2698
2699 #[tokio::test]
2704 async fn test_collect_capabilities_includes_mounts() {
2705 let registry = CapabilityRegistry::with_builtins();
2706
2707 let collected =
2708 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
2709
2710 assert!(!collected.mounts.is_empty());
2711 assert_eq!(collected.mounts.len(), 1);
2712 assert_eq!(collected.mounts[0].path, "/samples");
2713 assert!(collected.mounts[0].is_readonly());
2714 }
2715
2716 #[tokio::test]
2717 async fn test_collect_capabilities_empty_mounts_by_default() {
2718 let registry = CapabilityRegistry::with_builtins();
2719
2720 let collected =
2722 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
2723
2724 assert!(collected.mounts.is_empty());
2725 }
2726
2727 #[tokio::test]
2728 async fn test_collect_capabilities_combines_mounts() {
2729 let registry = CapabilityRegistry::with_builtins();
2730
2731 let collected = collect_capabilities(
2734 &["sample_data".to_string(), "current_time".to_string()],
2735 ®istry,
2736 &test_ctx(),
2737 )
2738 .await;
2739
2740 assert_eq!(collected.mounts.len(), 1);
2741 assert!(
2743 collected
2744 .applied_ids
2745 .iter()
2746 .any(|id| id == "session_file_system")
2747 );
2748 assert!(collected.applied_ids.iter().any(|id| id == "sample_data"));
2749 assert!(collected.applied_ids.iter().any(|id| id == "current_time"));
2750 }
2751
2752 #[test]
2753 fn test_sample_data_capability() {
2754 let registry = CapabilityRegistry::with_builtins();
2755 let cap = registry.get("sample_data").unwrap();
2756
2757 assert_eq!(cap.id(), "sample_data");
2758 assert_eq!(cap.name(), "Sample Data");
2759 assert_eq!(cap.status(), CapabilityStatus::Available);
2760
2761 assert!(cap.system_prompt_addition().is_some());
2763 assert!(cap.tools().is_empty());
2764
2765 assert!(!cap.mounts().is_empty());
2767 }
2768
2769 #[test]
2774 fn test_resolve_dependencies_empty() {
2775 let registry = CapabilityRegistry::with_builtins();
2776
2777 let resolved = resolve_dependencies(&[], ®istry).unwrap();
2778
2779 assert!(resolved.resolved_ids.is_empty());
2780 assert!(resolved.added_as_dependencies.is_empty());
2781 assert!(resolved.user_selected.is_empty());
2782 }
2783
2784 #[test]
2785 fn test_resolve_dependencies_no_deps() {
2786 let registry = CapabilityRegistry::with_builtins();
2787
2788 let resolved = resolve_dependencies(&["current_time".to_string()], ®istry).unwrap();
2790
2791 assert_eq!(resolved.resolved_ids, vec!["current_time"]);
2792 assert!(resolved.added_as_dependencies.is_empty());
2793 }
2794
2795 #[test]
2796 fn test_resolve_dependencies_with_deps() {
2797 let registry = CapabilityRegistry::with_builtins();
2798
2799 let resolved = resolve_dependencies(&["sample_data".to_string()], ®istry).unwrap();
2801
2802 assert_eq!(resolved.resolved_ids.len(), 2);
2804 let fs_pos = resolved
2805 .resolved_ids
2806 .iter()
2807 .position(|id| id == "session_file_system")
2808 .unwrap();
2809 let sd_pos = resolved
2810 .resolved_ids
2811 .iter()
2812 .position(|id| id == "sample_data")
2813 .unwrap();
2814 assert!(fs_pos < sd_pos, "FileSystem should come before SampleData");
2815
2816 assert_eq!(resolved.added_as_dependencies, vec!["session_file_system"]);
2818 }
2819
2820 #[test]
2821 fn test_resolve_dependencies_already_selected() {
2822 let registry = CapabilityRegistry::with_builtins();
2823
2824 let resolved = resolve_dependencies(
2826 &["session_file_system".to_string(), "sample_data".to_string()],
2827 ®istry,
2828 )
2829 .unwrap();
2830
2831 assert_eq!(resolved.resolved_ids.len(), 2);
2832 assert!(resolved.added_as_dependencies.is_empty());
2834 }
2835
2836 #[test]
2837 fn test_resolve_dependencies_preserves_order() {
2838 let registry = CapabilityRegistry::with_builtins();
2839
2840 let resolved =
2842 resolve_dependencies(&["current_time".to_string(), "noop".to_string()], ®istry)
2843 .unwrap();
2844
2845 assert_eq!(resolved.resolved_ids, vec!["current_time", "noop"]);
2846 }
2847
2848 #[test]
2849 fn test_resolve_dependencies_unknown_capability() {
2850 let registry = CapabilityRegistry::with_builtins();
2851
2852 let resolved =
2854 resolve_dependencies(&["unknown_capability".to_string()], ®istry).unwrap();
2855
2856 assert!(resolved.resolved_ids.is_empty());
2857 }
2858
2859 #[test]
2860 fn test_get_dependencies() {
2861 let registry = CapabilityRegistry::with_builtins();
2862
2863 let deps = get_dependencies("sample_data", ®istry);
2865 assert_eq!(deps, vec!["session_file_system"]);
2866
2867 let deps = get_dependencies("current_time", ®istry);
2869 assert!(deps.is_empty());
2870
2871 let deps = get_dependencies("unknown", ®istry);
2873 assert!(deps.is_empty());
2874 }
2875
2876 #[test]
2877 fn test_sample_data_has_dependency() {
2878 let registry = CapabilityRegistry::with_builtins();
2879 let cap = registry.get("sample_data").unwrap();
2880
2881 let deps = cap.dependencies();
2882 assert_eq!(deps.len(), 1);
2883 assert_eq!(deps[0], "session_file_system");
2884 }
2885
2886 #[test]
2887 fn test_noop_has_no_dependencies() {
2888 let registry = CapabilityRegistry::with_builtins();
2889 let cap = registry.get("noop").unwrap();
2890
2891 assert!(cap.dependencies().is_empty());
2892 }
2893
2894 #[test]
2898 fn test_circular_dependency_error() {
2899 struct CapA;
2901 struct CapB;
2902
2903 impl Capability for CapA {
2904 fn id(&self) -> &str {
2905 "test_cap_a"
2906 }
2907 fn name(&self) -> &str {
2908 "Test A"
2909 }
2910 fn description(&self) -> &str {
2911 "Test capability A"
2912 }
2913 fn dependencies(&self) -> Vec<&'static str> {
2914 vec!["test_cap_b"]
2915 }
2916 }
2917
2918 impl Capability for CapB {
2919 fn id(&self) -> &str {
2920 "test_cap_b"
2921 }
2922 fn name(&self) -> &str {
2923 "Test B"
2924 }
2925 fn description(&self) -> &str {
2926 "Test capability B"
2927 }
2928 fn dependencies(&self) -> Vec<&'static str> {
2929 vec!["test_cap_a"]
2930 }
2931 }
2932
2933 let mut registry = CapabilityRegistry::new();
2934 registry.register(CapA);
2935 registry.register(CapB);
2936
2937 let result = resolve_dependencies(&["test_cap_a".to_string()], ®istry);
2938
2939 assert!(result.is_err());
2940 match result.unwrap_err() {
2941 DependencyError::CircularDependency { capability_id, .. } => {
2942 assert_eq!(capability_id, "test_cap_a");
2943 }
2944 _ => panic!("Expected CircularDependency error"),
2945 }
2946 }
2947
2948 use crate::message_filter::{MessageFilter, MessageFilterProvider, MessageQuery};
2953
2954 struct FilterTestCapability {
2956 priority: i32,
2957 }
2958
2959 impl Capability for FilterTestCapability {
2960 fn id(&self) -> &str {
2961 "filter_test"
2962 }
2963 fn name(&self) -> &str {
2964 "Filter Test"
2965 }
2966 fn description(&self) -> &str {
2967 "Test capability with message filter"
2968 }
2969 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
2970 Some(Arc::new(FilterTestProvider {
2971 priority: self.priority,
2972 }))
2973 }
2974 }
2975
2976 struct FilterTestProvider {
2977 priority: i32,
2978 }
2979
2980 impl MessageFilterProvider for FilterTestProvider {
2981 fn apply_filters(&self, query: &mut MessageQuery, config: &serde_json::Value) {
2982 if let Some(search) = config.get("search").and_then(|v| v.as_str()) {
2984 query
2985 .filters
2986 .push(MessageFilter::Search(search.to_string()));
2987 }
2988 }
2989
2990 fn priority(&self) -> i32 {
2991 self.priority
2992 }
2993 }
2994
2995 #[tokio::test]
2996 async fn test_collect_capabilities_with_configs_no_filter_providers() {
2997 let registry = CapabilityRegistry::with_builtins();
2998 let configs = vec![AgentCapabilityConfig {
2999 capability_ref: CapabilityId::new("current_time"),
3000 config: serde_json::json!({}),
3001 }];
3002
3003 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3004
3005 assert!(collected.message_filter_providers.is_empty());
3006 assert!(!collected.has_message_filters());
3007 }
3008
3009 #[tokio::test]
3010 async fn test_collect_capabilities_with_configs_with_filter_provider() {
3011 let mut registry = CapabilityRegistry::new();
3012 registry.register(FilterTestCapability { priority: 0 });
3013
3014 let configs = vec![AgentCapabilityConfig {
3015 capability_ref: CapabilityId::new("filter_test"),
3016 config: serde_json::json!({ "search": "hello" }),
3017 }];
3018
3019 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3020
3021 assert_eq!(collected.message_filter_providers.len(), 1);
3022 assert!(collected.has_message_filters());
3023 }
3024
3025 #[tokio::test]
3026 async fn test_collect_capabilities_with_configs_filter_priority_order() {
3027 struct HighPriorityCapability;
3029 struct LowPriorityCapability;
3030
3031 impl Capability for HighPriorityCapability {
3032 fn id(&self) -> &str {
3033 "high_priority"
3034 }
3035 fn name(&self) -> &str {
3036 "High Priority"
3037 }
3038 fn description(&self) -> &str {
3039 "Test"
3040 }
3041 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3042 Some(Arc::new(FilterTestProvider { priority: 10 }))
3043 }
3044 }
3045
3046 impl Capability for LowPriorityCapability {
3047 fn id(&self) -> &str {
3048 "low_priority"
3049 }
3050 fn name(&self) -> &str {
3051 "Low Priority"
3052 }
3053 fn description(&self) -> &str {
3054 "Test"
3055 }
3056 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3057 Some(Arc::new(FilterTestProvider { priority: -5 }))
3058 }
3059 }
3060
3061 let mut registry = CapabilityRegistry::new();
3062 registry.register(HighPriorityCapability);
3063 registry.register(LowPriorityCapability);
3064
3065 let configs = vec![
3067 AgentCapabilityConfig {
3068 capability_ref: CapabilityId::new("high_priority"),
3069 config: serde_json::json!({}),
3070 },
3071 AgentCapabilityConfig {
3072 capability_ref: CapabilityId::new("low_priority"),
3073 config: serde_json::json!({}),
3074 },
3075 ];
3076
3077 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3078
3079 assert_eq!(collected.message_filter_providers.len(), 2);
3081 assert_eq!(collected.message_filter_providers[0].0.priority(), -5);
3082 assert_eq!(collected.message_filter_providers[1].0.priority(), 10);
3083 }
3084
3085 #[tokio::test]
3086 async fn test_collected_capabilities_apply_message_filters() {
3087 let mut registry = CapabilityRegistry::new();
3088 registry.register(FilterTestCapability { priority: 0 });
3089
3090 let configs = vec![AgentCapabilityConfig {
3091 capability_ref: CapabilityId::new("filter_test"),
3092 config: serde_json::json!({ "search": "test_query" }),
3093 }];
3094
3095 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3096
3097 let session_id: SessionId = Uuid::now_v7().into();
3099 let mut query = MessageQuery::new(session_id);
3100
3101 collected.apply_message_filters(&mut query);
3102
3103 assert_eq!(query.filters.len(), 1);
3105 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3106 }
3107
3108 #[tokio::test]
3109 async fn test_collected_capabilities_apply_multiple_filters_in_priority_order() {
3110 struct SearchCapability {
3111 id: &'static str,
3112 search_term: &'static str,
3113 priority: i32,
3114 }
3115
3116 struct SearchProvider {
3117 search_term: &'static str,
3118 priority: i32,
3119 }
3120
3121 impl MessageFilterProvider for SearchProvider {
3122 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3123 query
3124 .filters
3125 .push(MessageFilter::Search(self.search_term.to_string()));
3126 }
3127
3128 fn priority(&self) -> i32 {
3129 self.priority
3130 }
3131 }
3132
3133 impl Capability for SearchCapability {
3134 fn id(&self) -> &str {
3135 self.id
3136 }
3137 fn name(&self) -> &str {
3138 "Search"
3139 }
3140 fn description(&self) -> &str {
3141 "Test"
3142 }
3143 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3144 Some(Arc::new(SearchProvider {
3145 search_term: self.search_term,
3146 priority: self.priority,
3147 }))
3148 }
3149 }
3150
3151 let mut registry = CapabilityRegistry::new();
3152 registry.register(SearchCapability {
3153 id: "cap_a",
3154 search_term: "alpha",
3155 priority: 5,
3156 });
3157 registry.register(SearchCapability {
3158 id: "cap_b",
3159 search_term: "beta",
3160 priority: 1,
3161 });
3162 registry.register(SearchCapability {
3163 id: "cap_c",
3164 search_term: "gamma",
3165 priority: 10,
3166 });
3167
3168 let configs = vec![
3169 AgentCapabilityConfig {
3170 capability_ref: CapabilityId::new("cap_a"),
3171 config: serde_json::json!({}),
3172 },
3173 AgentCapabilityConfig {
3174 capability_ref: CapabilityId::new("cap_b"),
3175 config: serde_json::json!({}),
3176 },
3177 AgentCapabilityConfig {
3178 capability_ref: CapabilityId::new("cap_c"),
3179 config: serde_json::json!({}),
3180 },
3181 ];
3182
3183 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3184
3185 let session_id: SessionId = Uuid::now_v7().into();
3186 let mut query = MessageQuery::new(session_id);
3187
3188 collected.apply_message_filters(&mut query);
3189
3190 assert_eq!(query.filters.len(), 3);
3192 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3193 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3194 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3195 }
3196
3197 #[test]
3198 fn test_capability_without_message_filter_returns_none() {
3199 let registry = CapabilityRegistry::with_builtins();
3200
3201 let noop = registry.get("noop").unwrap();
3202 assert!(noop.message_filter_provider().is_none());
3203
3204 let current_time = registry.get("current_time").unwrap();
3205 assert!(current_time.message_filter_provider().is_none());
3206 }
3207
3208 #[tokio::test]
3209 async fn test_collect_capabilities_preserves_config_for_filter_provider() {
3210 let mut registry = CapabilityRegistry::new();
3211 registry.register(FilterTestCapability { priority: 0 });
3212
3213 let test_config = serde_json::json!({
3214 "search": "custom_search",
3215 "extra_field": 42
3216 });
3217
3218 let configs = vec![AgentCapabilityConfig {
3219 capability_ref: CapabilityId::new("filter_test"),
3220 config: test_config.clone(),
3221 }];
3222
3223 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3224
3225 assert_eq!(collected.message_filter_providers.len(), 1);
3227 let (_, stored_config) = &collected.message_filter_providers[0];
3228 assert_eq!(*stored_config, test_config);
3229 }
3230
3231 #[test]
3236 fn test_collect_message_filters_only_collects_filters() {
3237 let mut registry = CapabilityRegistry::new();
3238 registry.register(FilterTestCapability { priority: 0 });
3239
3240 let configs = vec![AgentCapabilityConfig {
3241 capability_ref: CapabilityId::new("filter_test"),
3242 config: serde_json::json!({ "search": "test_query" }),
3243 }];
3244
3245 let collected = collect_message_filters_only(&configs, ®istry);
3246
3247 let session_id: SessionId = Uuid::now_v7().into();
3248 let mut query = MessageQuery::new(session_id);
3249 collected.apply_message_filters(&mut query);
3250
3251 assert_eq!(query.filters.len(), 1);
3252 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3253 }
3254
3255 #[test]
3256 fn test_collect_message_filters_only_skips_unknown_capabilities() {
3257 let registry = CapabilityRegistry::new();
3258
3259 let configs = vec![AgentCapabilityConfig {
3260 capability_ref: CapabilityId::new("nonexistent"),
3261 config: serde_json::json!({}),
3262 }];
3263
3264 let collected = collect_message_filters_only(&configs, ®istry);
3265 assert!(collected.message_filter_providers.is_empty());
3266 }
3267
3268 #[test]
3269 fn test_collect_message_filters_only_preserves_priority_order() {
3270 struct PriorityFilterCap {
3271 id: &'static str,
3272 search_term: &'static str,
3273 priority: i32,
3274 }
3275
3276 struct PriorityFilterProvider {
3277 search_term: &'static str,
3278 priority: i32,
3279 }
3280
3281 impl Capability for PriorityFilterCap {
3282 fn id(&self) -> &str {
3283 self.id
3284 }
3285 fn name(&self) -> &str {
3286 self.id
3287 }
3288 fn description(&self) -> &str {
3289 "priority test"
3290 }
3291 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3292 Some(Arc::new(PriorityFilterProvider {
3293 search_term: self.search_term,
3294 priority: self.priority,
3295 }))
3296 }
3297 }
3298
3299 impl MessageFilterProvider for PriorityFilterProvider {
3300 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3301 query
3302 .filters
3303 .push(MessageFilter::Search(self.search_term.to_string()));
3304 }
3305 fn priority(&self) -> i32 {
3306 self.priority
3307 }
3308 }
3309
3310 let mut registry = CapabilityRegistry::new();
3311 registry.register(PriorityFilterCap {
3312 id: "gamma",
3313 search_term: "gamma",
3314 priority: 10,
3315 });
3316 registry.register(PriorityFilterCap {
3317 id: "alpha",
3318 search_term: "alpha",
3319 priority: 5,
3320 });
3321 registry.register(PriorityFilterCap {
3322 id: "beta",
3323 search_term: "beta",
3324 priority: 1,
3325 });
3326
3327 let configs = vec![
3328 AgentCapabilityConfig {
3329 capability_ref: CapabilityId::new("gamma"),
3330 config: serde_json::json!({}),
3331 },
3332 AgentCapabilityConfig {
3333 capability_ref: CapabilityId::new("alpha"),
3334 config: serde_json::json!({}),
3335 },
3336 AgentCapabilityConfig {
3337 capability_ref: CapabilityId::new("beta"),
3338 config: serde_json::json!({}),
3339 },
3340 ];
3341
3342 let collected = collect_message_filters_only(&configs, ®istry);
3343
3344 let session_id: SessionId = Uuid::now_v7().into();
3345 let mut query = MessageQuery::new(session_id);
3346 collected.apply_message_filters(&mut query);
3347
3348 assert_eq!(query.filters.len(), 3);
3350 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3351 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3352 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3353 }
3354
3355 #[test]
3356 fn test_collect_message_filters_only_post_load_invoked() {
3357 use crate::message::Message;
3358
3359 struct PostLoadCap;
3360 struct PostLoadProvider;
3361
3362 impl Capability for PostLoadCap {
3363 fn id(&self) -> &str {
3364 "post_load_test"
3365 }
3366 fn name(&self) -> &str {
3367 "PostLoad Test"
3368 }
3369 fn description(&self) -> &str {
3370 "test"
3371 }
3372 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3373 Some(Arc::new(PostLoadProvider))
3374 }
3375 }
3376
3377 impl MessageFilterProvider for PostLoadProvider {
3378 fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
3379 fn priority(&self) -> i32 {
3380 0
3381 }
3382 fn post_load(&self, messages: &mut Vec<Message>, _config: &serde_json::Value) {
3383 messages.reverse();
3385 }
3386 }
3387
3388 let mut registry = CapabilityRegistry::new();
3389 registry.register(PostLoadCap);
3390
3391 let configs = vec![AgentCapabilityConfig {
3392 capability_ref: CapabilityId::new("post_load_test"),
3393 config: serde_json::json!({}),
3394 }];
3395
3396 let collected = collect_message_filters_only(&configs, ®istry);
3397
3398 let mut messages = vec![Message::user("first"), Message::user("second")];
3399 collected.apply_post_load_filters(&mut messages);
3400
3401 assert_eq!(messages[0].text(), Some("second"));
3403 assert_eq!(messages[1].text(), Some("first"));
3404 }
3405
3406 #[test]
3407 fn test_collect_model_view_providers_respects_compaction_capability_boundary() {
3408 use crate::tool_types::ToolCall;
3409
3410 fn tool_heavy_messages() -> Vec<Message> {
3411 let mut messages = vec![Message::user("inspect files repeatedly")];
3412 for index in 0..9 {
3413 let call_id = format!("call_{index}");
3414 messages.push(Message::assistant_with_tools(
3415 "",
3416 vec![ToolCall {
3417 id: call_id.clone(),
3418 name: "read_file".to_string(),
3419 arguments: serde_json::json!({"path": "/workspace/src/lib.rs"}),
3420 }],
3421 ));
3422 messages.push(Message::tool_result(
3423 call_id,
3424 Some(serde_json::json!({
3425 "path": "/workspace/src/lib.rs",
3426 "content": format!("{}{}", "large file line\n".repeat(1000), index),
3427 "total_lines": 1000,
3428 "lines_shown": {"start": 1, "end": 1000},
3429 "truncated": false
3430 })),
3431 None,
3432 ));
3433 }
3434 messages
3435 }
3436
3437 fn first_tool_result_is_masked(messages: &[Message]) -> bool {
3438 messages[2]
3439 .tool_result_content()
3440 .and_then(|result| result.result.as_ref())
3441 .and_then(|result| result.get("masked"))
3442 .and_then(|masked| masked.as_bool())
3443 .unwrap_or(false)
3444 }
3445
3446 let mut registry = CapabilityRegistry::new();
3447 registry.register(CompactionCapability);
3448 let context = ModelViewContext {
3449 session_id: SessionId::new(),
3450 prior_usage: None,
3451 };
3452
3453 let no_compaction = collect_model_view_providers(&[], ®istry, None);
3454 let unmasked = no_compaction.apply_model_view(tool_heavy_messages(), &context);
3455 assert!(!first_tool_result_is_masked(&unmasked));
3456
3457 let compaction = collect_model_view_providers(
3458 &[AgentCapabilityConfig {
3459 capability_ref: CapabilityId::new(COMPACTION_CAPABILITY_ID),
3460 config: serde_json::json!({}),
3461 }],
3462 ®istry,
3463 None,
3464 );
3465 let masked = compaction.apply_model_view(tool_heavy_messages(), &context);
3466 assert!(first_tool_result_is_masked(&masked));
3467 let last_tool = masked.last().unwrap().tool_result_content().unwrap();
3468 assert!(last_tool.result.as_ref().unwrap().get("content").is_some());
3469 }
3470
3471 struct DelegatingFilterCap {
3474 id: &'static str,
3475 inner: std::sync::Arc<InnerFilterCap>,
3476 }
3477 struct InnerFilterCap;
3478
3479 impl Capability for InnerFilterCap {
3480 fn id(&self) -> &str {
3481 "inner_filter"
3482 }
3483 fn name(&self) -> &str {
3484 "Inner Filter"
3485 }
3486 fn description(&self) -> &str {
3487 "inner"
3488 }
3489 fn message_filter_provider(&self) -> Option<std::sync::Arc<dyn MessageFilterProvider>> {
3490 Some(std::sync::Arc::new(SentinelFilter))
3491 }
3492 }
3493 struct SentinelFilter;
3494 impl MessageFilterProvider for SentinelFilter {
3495 fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
3496 }
3497 impl Capability for DelegatingFilterCap {
3498 fn id(&self) -> &str {
3499 self.id
3500 }
3501 fn name(&self) -> &str {
3502 "Delegating Filter"
3503 }
3504 fn description(&self) -> &str {
3505 "delegating"
3506 }
3507 fn message_filter_provider(&self) -> Option<std::sync::Arc<dyn MessageFilterProvider>> {
3508 None }
3510 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
3511 Some(&*self.inner)
3512 }
3513 }
3514
3515 #[test]
3516 fn test_collect_message_filters_only_honors_resolve_for_model_delegation() {
3517 let inner = std::sync::Arc::new(InnerFilterCap);
3518 let outer = DelegatingFilterCap {
3519 id: "delegating_filter",
3520 inner: inner.clone(),
3521 };
3522
3523 let mut registry = CapabilityRegistry::new();
3524 registry.register(outer);
3525
3526 let configs = vec![AgentCapabilityConfig {
3527 capability_ref: CapabilityId::new("delegating_filter"),
3528 config: serde_json::json!({}),
3529 }];
3530
3531 let collected = collect_message_filters_only(&configs, ®istry);
3534 assert_eq!(
3535 collected.message_filter_providers.len(),
3536 1,
3537 "provider from resolved inner capability must be collected"
3538 );
3539 }
3540
3541 struct DelegatingMvpCap {
3542 id: &'static str,
3543 inner: std::sync::Arc<InnerMvpCap>,
3544 }
3545 struct InnerMvpCap;
3546
3547 impl Capability for InnerMvpCap {
3548 fn id(&self) -> &str {
3549 "inner_mvp"
3550 }
3551 fn name(&self) -> &str {
3552 "Inner MVP"
3553 }
3554 fn description(&self) -> &str {
3555 "inner"
3556 }
3557 fn model_view_provider(
3558 &self,
3559 ) -> Option<std::sync::Arc<dyn crate::capabilities::ModelViewProvider>> {
3560 struct NoopMvp;
3562 impl crate::capabilities::ModelViewProvider for NoopMvp {
3563 fn apply_model_view(
3564 &self,
3565 messages: Vec<Message>,
3566 _config: &serde_json::Value,
3567 _context: &ModelViewContext<'_>,
3568 ) -> Vec<Message> {
3569 messages
3570 }
3571 }
3572 Some(std::sync::Arc::new(NoopMvp))
3573 }
3574 }
3575 impl Capability for DelegatingMvpCap {
3576 fn id(&self) -> &str {
3577 self.id
3578 }
3579 fn name(&self) -> &str {
3580 "Delegating MVP"
3581 }
3582 fn description(&self) -> &str {
3583 "delegating"
3584 }
3585 fn model_view_provider(
3586 &self,
3587 ) -> Option<std::sync::Arc<dyn crate::capabilities::ModelViewProvider>> {
3588 None }
3590 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
3591 Some(&*self.inner)
3592 }
3593 }
3594
3595 #[test]
3596 fn test_collect_model_view_providers_honors_resolve_for_model_delegation() {
3597 let inner = std::sync::Arc::new(InnerMvpCap);
3598 let outer = DelegatingMvpCap {
3599 id: "delegating_mvp",
3600 inner: inner.clone(),
3601 };
3602
3603 let mut registry = CapabilityRegistry::new();
3604 registry.register(outer);
3605
3606 let configs = vec![AgentCapabilityConfig {
3607 capability_ref: CapabilityId::new("delegating_mvp"),
3608 config: serde_json::json!({}),
3609 }];
3610
3611 let collected = collect_model_view_providers(&configs, ®istry, None);
3614 assert_eq!(
3615 collected.model_view_providers.len(),
3616 1,
3617 "provider from resolved inner capability must be collected"
3618 );
3619 }
3620
3621 #[tokio::test]
3631 async fn test_virtual_bash_capability_produces_bash_tool() {
3632 let registry = CapabilityRegistry::with_builtins();
3633 let collected =
3634 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3635
3636 let tool_names: Vec<&str> = collected
3637 .tool_definitions
3638 .iter()
3639 .map(|t| t.name())
3640 .collect();
3641 assert!(
3642 tool_names.contains(&"bash"),
3643 "virtual_bash capability must produce 'bash' tool, got: {:?}",
3644 tool_names
3645 );
3646 assert!(
3647 !collected.tools.is_empty(),
3648 "virtual_bash must provide tool implementations"
3649 );
3650 }
3651
3652 #[tokio::test]
3653 async fn test_generic_harness_capability_set_produces_bash_tool() {
3654 let generic_harness_caps = vec![
3657 "session_file_system".to_string(),
3658 "virtual_bash".to_string(),
3659 "web_fetch".to_string(),
3660 "session_storage".to_string(),
3661 "session".to_string(),
3662 "agent_instructions".to_string(),
3663 "skills".to_string(),
3664 "infinity_context".to_string(),
3665 "auto_tool_search".to_string(),
3666 ];
3667
3668 let registry = CapabilityRegistry::with_builtins();
3669 let collected = collect_capabilities(&generic_harness_caps, ®istry, &test_ctx()).await;
3670
3671 let tool_names: Vec<&str> = collected
3672 .tool_definitions
3673 .iter()
3674 .map(|t| t.name())
3675 .collect();
3676 assert!(
3677 tool_names.contains(&"bash"),
3678 "Generic Harness capabilities must produce 'bash' tool, got: {:?}",
3679 tool_names
3680 );
3681 }
3682
3683 #[tokio::test]
3684 async fn test_collect_capabilities_tool_count_matches_definitions() {
3685 let registry = CapabilityRegistry::with_builtins();
3688 let collected =
3689 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3690
3691 assert_eq!(
3692 collected.tools.len(),
3693 collected.tool_definitions.len(),
3694 "tool implementations ({}) must match tool definitions ({})",
3695 collected.tools.len(),
3696 collected.tool_definitions.len(),
3697 );
3698 }
3699
3700 #[tokio::test]
3704 async fn test_collect_capabilities_resolves_dependencies() {
3705 let registry = CapabilityRegistry::with_builtins();
3708 let collected =
3709 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
3710
3711 assert!(
3713 collected
3714 .applied_ids
3715 .iter()
3716 .any(|id| id == "session_file_system"),
3717 "collect_capabilities must apply session_file_system as a dependency; applied_ids: {:?}",
3718 collected.applied_ids
3719 );
3720
3721 let tool_names: Vec<&str> = collected
3722 .tool_definitions
3723 .iter()
3724 .map(|t| t.name())
3725 .collect();
3726
3727 assert!(
3729 tool_names.contains(&"read_file") && tool_names.contains(&"write_file"),
3730 "collect_capabilities must resolve dependencies and include dependency tools, got: {:?}",
3731 tool_names
3732 );
3733
3734 assert_eq!(
3736 collected.tools.len(),
3737 collected.tool_definitions.len(),
3738 "dependency-added tools must have implementations, not just definitions"
3739 );
3740 }
3741
3742 #[test]
3743 fn test_defaults_do_not_include_bash() {
3744 let registry = crate::ToolRegistry::with_defaults();
3747 assert!(
3748 !registry.has("bash"),
3749 "with_defaults() must not include 'bash' — it comes from virtual_bash capability"
3750 );
3751 }
3752
3753 #[tokio::test]
3760 async fn test_background_execution_auto_activates_with_virtual_bash() {
3761 let registry = CapabilityRegistry::with_builtins();
3762 let collected =
3763 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3764
3765 let tool_names: Vec<&str> = collected
3766 .tool_definitions
3767 .iter()
3768 .map(|t| t.name())
3769 .collect();
3770 assert!(
3771 tool_names.contains(&"spawn_background"),
3772 "spawn_background must be auto-activated when virtual_bash (a \
3773 background-capable tool) is in the agent's capability set; got: {:?}",
3774 tool_names
3775 );
3776 assert!(
3777 collected
3778 .applied_ids
3779 .iter()
3780 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
3781 "background_execution must be in applied_ids when auto-activated; \
3782 got: {:?}",
3783 collected.applied_ids
3784 );
3785
3786 assert!(
3788 collected
3789 .tools
3790 .iter()
3791 .any(|t| t.name() == "spawn_background"),
3792 "spawn_background tool implementation must be present alongside the \
3793 definition (lockstep contract)"
3794 );
3795 }
3796
3797 #[tokio::test]
3800 async fn test_background_execution_does_not_auto_activate_without_hint() {
3801 let registry = CapabilityRegistry::with_builtins();
3802 let collected =
3804 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
3805
3806 let tool_names: Vec<&str> = collected
3807 .tool_definitions
3808 .iter()
3809 .map(|t| t.name())
3810 .collect();
3811 assert!(
3812 !tool_names.contains(&"spawn_background"),
3813 "spawn_background must NOT be activated without a background-capable \
3814 tool; got: {:?}",
3815 tool_names
3816 );
3817 assert!(
3818 !collected
3819 .applied_ids
3820 .iter()
3821 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
3822 "background_execution must not appear in applied_ids when no \
3823 background-capable tool is present; got: {:?}",
3824 collected.applied_ids
3825 );
3826 }
3827
3828 #[tokio::test]
3832 async fn test_background_execution_explicit_selection_is_idempotent() {
3833 let registry = CapabilityRegistry::with_builtins();
3834 let collected = collect_capabilities(
3835 &[
3836 "virtual_bash".to_string(),
3837 BACKGROUND_EXECUTION_CAPABILITY_ID.to_string(),
3838 ],
3839 ®istry,
3840 &test_ctx(),
3841 )
3842 .await;
3843
3844 let spawn_background_count = collected
3845 .tool_definitions
3846 .iter()
3847 .filter(|t| t.name() == "spawn_background")
3848 .count();
3849 assert_eq!(
3850 spawn_background_count, 1,
3851 "spawn_background must appear exactly once even when \
3852 background_execution is selected explicitly alongside a \
3853 background-capable tool"
3854 );
3855 let applied_count = collected
3856 .applied_ids
3857 .iter()
3858 .filter(|id| id.as_str() == BACKGROUND_EXECUTION_CAPABILITY_ID)
3859 .count();
3860 assert_eq!(
3861 applied_count, 1,
3862 "background_execution must appear exactly once in applied_ids"
3863 );
3864 }
3865
3866 #[test]
3871 fn test_defaults_do_not_include_spawn_background() {
3872 let registry = crate::ToolRegistry::with_defaults();
3873 assert!(
3874 !registry.has("spawn_background"),
3875 "with_defaults() must not include 'spawn_background' — it comes \
3876 from the background_execution capability (EVE-501)"
3877 );
3878 }
3879
3880 #[test]
3885 fn test_capability_features_default_empty() {
3886 let registry = CapabilityRegistry::with_builtins();
3887
3888 let noop = registry.get("noop").unwrap();
3890 assert!(noop.features().is_empty());
3891
3892 let current_time = registry.get("current_time").unwrap();
3893 assert!(current_time.features().is_empty());
3894 }
3895
3896 #[test]
3897 fn test_file_system_capability_features() {
3898 let registry = CapabilityRegistry::with_builtins();
3899
3900 let fs = registry.get("session_file_system").unwrap();
3901 assert_eq!(fs.features(), vec!["file_system"]);
3902 }
3903
3904 #[test]
3905 fn test_virtual_bash_capability_features() {
3906 let registry = CapabilityRegistry::with_builtins();
3907
3908 let bash = registry.get("virtual_bash").unwrap();
3909 assert_eq!(bash.features(), vec!["file_system"]);
3910 }
3911
3912 #[test]
3913 fn test_session_storage_capability_features() {
3914 let registry = CapabilityRegistry::with_builtins();
3915
3916 let storage = registry.get("session_storage").unwrap();
3917 let features = storage.features();
3918 assert!(features.contains(&"secrets"));
3919 assert!(features.contains(&"key_value"));
3920 }
3921
3922 #[test]
3923 fn test_session_schedule_capability_features() {
3924 let registry = CapabilityRegistry::with_builtins();
3925
3926 let schedule = registry.get("session_schedule").unwrap();
3927 assert_eq!(schedule.features(), vec!["schedules"]);
3928 }
3929
3930 #[test]
3931 fn test_session_sql_database_capability_features() {
3932 let registry = CapabilityRegistry::with_builtins();
3933
3934 let sql = registry.get("session_sql_database").unwrap();
3935 assert_eq!(sql.features(), vec!["sql_database"]);
3936 }
3937
3938 #[test]
3939 fn test_sample_data_capability_features() {
3940 let registry = CapabilityRegistry::with_builtins();
3941
3942 let sample = registry.get("sample_data").unwrap();
3943 assert_eq!(sample.features(), vec!["file_system"]);
3944 }
3945
3946 #[test]
3947 fn test_compute_features_empty() {
3948 let registry = CapabilityRegistry::with_builtins();
3949
3950 let features = compute_features(&[], ®istry);
3951 assert!(features.is_empty());
3952 }
3953
3954 #[test]
3955 fn test_compute_features_single_capability() {
3956 let registry = CapabilityRegistry::with_builtins();
3957
3958 let features = compute_features(&["session_schedule".to_string()], ®istry);
3959 assert_eq!(features, vec!["schedules"]);
3960 }
3961
3962 #[test]
3963 fn test_compute_features_multiple_capabilities() {
3964 let registry = CapabilityRegistry::with_builtins();
3965
3966 let features = compute_features(
3967 &[
3968 "session_file_system".to_string(),
3969 "session_storage".to_string(),
3970 "session_schedule".to_string(),
3971 ],
3972 ®istry,
3973 );
3974 assert!(features.contains(&"file_system".to_string()));
3975 assert!(features.contains(&"secrets".to_string()));
3976 assert!(features.contains(&"key_value".to_string()));
3977 assert!(features.contains(&"schedules".to_string()));
3978 }
3979
3980 #[test]
3981 fn test_compute_features_deduplicates() {
3982 let registry = CapabilityRegistry::with_builtins();
3983
3984 let features = compute_features(
3986 &[
3987 "session_file_system".to_string(),
3988 "virtual_bash".to_string(),
3989 ],
3990 ®istry,
3991 );
3992 let file_system_count = features.iter().filter(|f| *f == "file_system").count();
3993 assert_eq!(file_system_count, 1, "file_system should appear only once");
3994 }
3995
3996 #[test]
3997 fn test_compute_features_includes_dependency_features() {
3998 let registry = CapabilityRegistry::with_builtins();
3999
4000 let features = compute_features(&["virtual_bash".to_string()], ®istry);
4002 assert!(features.contains(&"file_system".to_string()));
4003 }
4004
4005 #[test]
4006 fn test_compute_features_generic_harness_set() {
4007 let registry = CapabilityRegistry::with_builtins();
4008
4009 let features = compute_features(
4011 &[
4012 "session_file_system".to_string(),
4013 "virtual_bash".to_string(),
4014 "session_storage".to_string(),
4015 "session".to_string(),
4016 "session_schedule".to_string(),
4017 ],
4018 ®istry,
4019 );
4020 assert!(features.contains(&"file_system".to_string()));
4021 assert!(features.contains(&"secrets".to_string()));
4022 assert!(features.contains(&"key_value".to_string()));
4023 assert!(features.contains(&"schedules".to_string()));
4024 }
4025
4026 #[test]
4027 fn test_compute_features_unknown_capability_ignored() {
4028 let registry = CapabilityRegistry::with_builtins();
4029
4030 let features = compute_features(
4031 &["unknown_cap".to_string(), "session_schedule".to_string()],
4032 ®istry,
4033 );
4034 assert_eq!(features, vec!["schedules"]);
4035 }
4036
4037 #[test]
4038 fn test_risk_level_ordering() {
4039 assert!(RiskLevel::Low < RiskLevel::Medium);
4040 assert!(RiskLevel::Medium < RiskLevel::High);
4041 }
4042
4043 #[test]
4044 fn test_risk_level_serde_roundtrip() {
4045 let high = RiskLevel::High;
4046 let json = serde_json::to_string(&high).unwrap();
4047 assert_eq!(json, "\"high\"");
4048 let back: RiskLevel = serde_json::from_str(&json).unwrap();
4049 assert_eq!(back, RiskLevel::High);
4050 }
4051
4052 #[test]
4053 fn test_capability_risk_levels() {
4054 let registry = CapabilityRegistry::with_builtins();
4055
4056 let bash = registry.get("virtual_bash").unwrap();
4058 assert_eq!(bash.risk_level(), RiskLevel::High);
4059
4060 let fetch = registry.get("web_fetch").unwrap();
4062 assert_eq!(fetch.risk_level(), RiskLevel::High);
4063
4064 let noop = registry.get("noop").unwrap();
4066 assert_eq!(noop.risk_level(), RiskLevel::Low);
4067 }
4068
4069 #[tokio::test]
4074 async fn test_apply_capabilities_openai_tool_search() {
4075 let registry = CapabilityRegistry::with_builtins();
4076 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4077
4078 let applied = apply_capabilities(
4079 base_runtime_agent.clone(),
4080 &["openai_tool_search".to_string()],
4081 ®istry,
4082 &test_ctx(),
4083 )
4084 .await;
4085
4086 assert_eq!(
4088 applied.runtime_agent.system_prompt,
4089 base_runtime_agent.system_prompt
4090 );
4091 assert!(applied.tool_registry.is_empty());
4092 assert_eq!(applied.applied_ids, vec!["openai_tool_search"]);
4093
4094 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
4096 assert!(ts.enabled);
4097 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
4098 }
4099
4100 #[tokio::test]
4101 async fn test_apply_capabilities_openai_tool_search_with_other_capabilities() {
4102 let registry = CapabilityRegistry::with_builtins();
4103 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4104
4105 let applied = apply_capabilities(
4106 base_runtime_agent,
4107 &[
4108 "current_time".to_string(),
4109 "openai_tool_search".to_string(),
4110 "test_math".to_string(),
4111 ],
4112 ®istry,
4113 &test_ctx(),
4114 )
4115 .await;
4116
4117 assert!(applied.tool_registry.has("get_current_time"));
4119 assert!(applied.tool_registry.has("add"));
4120 assert!(applied.tool_registry.has("subtract"));
4121 assert!(applied.tool_registry.has("multiply"));
4122 assert!(applied.tool_registry.has("divide"));
4123
4124 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
4126 assert!(ts.enabled);
4127 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
4128 }
4129
4130 #[tokio::test]
4131 async fn test_collect_capabilities_tool_search_custom_threshold() {
4132 let registry = CapabilityRegistry::with_builtins();
4133
4134 let configs = vec![AgentCapabilityConfig {
4135 capability_ref: CapabilityId::new("openai_tool_search"),
4136 config: serde_json::json!({"threshold": 5}),
4137 }];
4138
4139 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4140
4141 let ts = collected.tool_search.as_ref().unwrap();
4142 assert!(ts.enabled);
4143 assert_eq!(ts.threshold, 5);
4144 }
4145
4146 #[tokio::test]
4147 async fn test_collect_capabilities_auto_tool_search_resolves_to_generic_off_native() {
4148 let registry = CapabilityRegistry::with_builtins();
4149
4150 let configs = vec![
4151 AgentCapabilityConfig {
4152 capability_ref: CapabilityId::new("auto_tool_search"),
4153 config: serde_json::json!({"threshold": 2}),
4154 },
4155 AgentCapabilityConfig {
4156 capability_ref: CapabilityId::new("test_math"),
4157 config: serde_json::json!({}),
4158 },
4159 ];
4160
4161 let ctx = test_ctx().with_model("claude-sonnet-4-5-20250514");
4164 let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
4165
4166 assert!(
4167 collected.tool_search.is_none(),
4168 "auto_tool_search must not set a hosted config on a non-native model"
4169 );
4170 assert!(
4171 collected
4172 .tools
4173 .iter()
4174 .any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
4175 "auto_tool_search must contribute the client-side tool_search tool"
4176 );
4177 assert!(
4178 !collected.tool_definition_hooks.is_empty(),
4179 "auto_tool_search must contribute a client-side deferral hook"
4180 );
4181
4182 let mut transformed = collected.tool_definitions.clone();
4183 for hook in &collected.tool_definition_hooks {
4184 transformed = hook.transform(transformed);
4185 }
4186 let add_tool = transformed
4187 .iter()
4188 .find(|tool| tool.name() == "add")
4189 .expect("test_math contributes add");
4190 assert!(
4191 add_tool.parameters().get("properties").is_none(),
4192 "generic auto_tool_search must honor the configured threshold"
4193 );
4194 }
4195
4196 #[tokio::test]
4197 async fn test_collect_capabilities_auto_tool_search_resolves_to_hosted_on_native() {
4198 let registry = CapabilityRegistry::with_builtins();
4199
4200 let configs = vec![AgentCapabilityConfig {
4201 capability_ref: CapabilityId::new("auto_tool_search"),
4202 config: serde_json::json!({"threshold": 7}),
4203 }];
4204
4205 let ctx = test_ctx().with_model("gpt-5.4");
4208 let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
4209
4210 let ts = collected
4211 .tool_search
4212 .as_ref()
4213 .expect("auto_tool_search must set a hosted config on a native model");
4214 assert!(ts.enabled);
4215 assert_eq!(ts.threshold, 7);
4216 assert!(
4217 !collected
4218 .tools
4219 .iter()
4220 .any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
4221 "hosted mechanism must not contribute the client-side tool_search tool"
4222 );
4223 assert!(
4224 collected.tool_definition_hooks.is_empty(),
4225 "hosted mechanism must not contribute a client-side deferral hook"
4226 );
4227 }
4228
4229 #[tokio::test]
4230 async fn test_collect_capabilities_no_tool_search_without_capability() {
4231 let registry = CapabilityRegistry::with_builtins();
4232
4233 let configs = vec![AgentCapabilityConfig {
4234 capability_ref: CapabilityId::new("current_time"),
4235 config: serde_json::json!({}),
4236 }];
4237
4238 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4239
4240 assert!(collected.tool_search.is_none());
4241 }
4242
4243 #[tokio::test]
4244 async fn test_collect_capabilities_tool_search_category_propagation() {
4245 let registry = CapabilityRegistry::with_builtins();
4246
4247 let configs = vec![
4249 AgentCapabilityConfig {
4250 capability_ref: CapabilityId::new("test_math"),
4251 config: serde_json::json!({}),
4252 },
4253 AgentCapabilityConfig {
4254 capability_ref: CapabilityId::new("openai_tool_search"),
4255 config: serde_json::json!({}),
4256 },
4257 ];
4258
4259 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4260
4261 assert!(collected.tool_search.is_some());
4263
4264 for tool_def in &collected.tool_definitions {
4266 if ["add", "subtract", "multiply", "divide"].contains(&tool_def.name()) {
4268 assert!(
4269 tool_def.category().is_some(),
4270 "Tool {} should have a category from its capability",
4271 tool_def.name()
4272 );
4273 }
4274 }
4275 }
4276
4277 #[tokio::test]
4278 async fn test_apply_capabilities_prompt_caching() {
4279 let registry = CapabilityRegistry::with_builtins();
4280 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4281
4282 let applied = apply_capabilities(
4283 base_runtime_agent.clone(),
4284 &["prompt_caching".to_string()],
4285 ®istry,
4286 &test_ctx(),
4287 )
4288 .await;
4289
4290 assert_eq!(
4291 applied.runtime_agent.system_prompt,
4292 base_runtime_agent.system_prompt
4293 );
4294 assert!(applied.tool_registry.is_empty());
4295 assert_eq!(applied.applied_ids, vec!["prompt_caching"]);
4296
4297 let prompt_cache = applied.runtime_agent.prompt_cache.as_ref().unwrap();
4298 assert!(prompt_cache.enabled);
4299 assert_eq!(
4300 prompt_cache.strategy,
4301 crate::llm_driver_registry::PromptCacheStrategy::Auto
4302 );
4303 assert!(prompt_cache.gemini_cached_content.is_none());
4304 }
4305
4306 #[tokio::test]
4307 async fn test_collect_capabilities_prompt_caching_custom_strategy() {
4308 let registry = CapabilityRegistry::with_builtins();
4309
4310 let configs = vec![AgentCapabilityConfig {
4311 capability_ref: CapabilityId::new("prompt_caching"),
4312 config: serde_json::json!({"strategy": "auto"}),
4313 }];
4314
4315 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4316
4317 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
4318 assert!(prompt_cache.enabled);
4319 assert_eq!(
4320 prompt_cache.strategy,
4321 crate::llm_driver_registry::PromptCacheStrategy::Auto
4322 );
4323 assert!(prompt_cache.gemini_cached_content.is_none());
4324 }
4325
4326 #[tokio::test]
4327 async fn test_collect_capabilities_prompt_caching_gemini_cached_content() {
4328 let registry = CapabilityRegistry::with_builtins();
4329
4330 let configs = vec![AgentCapabilityConfig {
4331 capability_ref: CapabilityId::new("prompt_caching"),
4332 config: serde_json::json!({
4333 "strategy": "auto",
4334 "gemini_cached_content": "cachedContents/demo-cache"
4335 }),
4336 }];
4337
4338 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4339
4340 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
4341 assert_eq!(
4342 prompt_cache.gemini_cached_content.as_deref(),
4343 Some("cachedContents/demo-cache")
4344 );
4345 }
4346
4347 struct SkillContributingCapability;
4352
4353 impl Capability for SkillContributingCapability {
4354 fn id(&self) -> &str {
4355 "contributes_skills"
4356 }
4357 fn name(&self) -> &str {
4358 "Contributes Skills"
4359 }
4360 fn description(&self) -> &str {
4361 "Test capability that contributes skills."
4362 }
4363 fn contribute_skills(&self) -> Vec<SkillContribution> {
4364 vec![
4365 SkillContribution::new("alpha-skill", "Alpha skill desc", "# Alpha\nDo alpha.")
4366 .with_files(vec![(
4367 "scripts/a.sh".to_string(),
4368 "#!/bin/sh\necho a\n".to_string(),
4369 )]),
4370 SkillContribution::new("beta-skill", "Beta skill desc", "# Beta\nDo beta.")
4371 .with_user_invocable(false),
4372 ]
4373 }
4374 }
4375
4376 fn skill_md_from_entries(entries: &HashMap<String, MountEntry>) -> &str {
4377 match &entries.get("SKILL.md").expect("SKILL.md missing").source {
4378 MountSource::InlineFile { content, .. } => content.as_str(),
4379 _ => panic!("Expected InlineFile for SKILL.md"),
4380 }
4381 }
4382
4383 #[tokio::test]
4384 async fn test_contribute_skills_normalized_to_mounts() {
4385 let mut registry = CapabilityRegistry::new();
4386 registry.register(SkillContributingCapability);
4387
4388 let configs = vec![AgentCapabilityConfig {
4389 capability_ref: CapabilityId::new("contributes_skills"),
4390 config: serde_json::json!({}),
4391 }];
4392
4393 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4394
4395 let skill_mounts: Vec<_> = collected
4396 .mounts
4397 .iter()
4398 .filter(|m| m.path.starts_with("/.agents/skills/"))
4399 .collect();
4400 assert_eq!(skill_mounts.len(), 2);
4401
4402 for m in &skill_mounts {
4405 assert!(m.is_readonly());
4406 assert_eq!(m.capability_id, "contributes_skills");
4407 }
4408
4409 let alpha = skill_mounts
4410 .iter()
4411 .find(|m| m.path == "/.agents/skills/alpha-skill")
4412 .expect("alpha-skill mount missing");
4413 match &alpha.source {
4414 MountSource::InlineDirectory { entries } => {
4415 assert!(entries.contains_key("SKILL.md"));
4416 assert!(entries.contains_key("scripts/a.sh"));
4417 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
4418 assert_eq!(parsed.name, "alpha-skill");
4419 assert!(parsed.user_invocable);
4420 }
4421 _ => panic!("Expected InlineDirectory"),
4422 }
4423
4424 let beta = skill_mounts
4425 .iter()
4426 .find(|m| m.path == "/.agents/skills/beta-skill")
4427 .expect("beta-skill mount missing");
4428 match &beta.source {
4429 MountSource::InlineDirectory { entries } => {
4430 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
4431 assert!(!parsed.user_invocable);
4432 }
4433 _ => panic!("Expected InlineDirectory"),
4434 }
4435 }
4436
4437 #[tokio::test]
4438 async fn test_contribute_skills_default_empty() {
4439 let mut registry = CapabilityRegistry::new();
4442 registry.register(FilterTestCapability { priority: 0 });
4443
4444 let configs = vec![AgentCapabilityConfig {
4445 capability_ref: CapabilityId::new("filter_test"),
4446 config: serde_json::json!({}),
4447 }];
4448
4449 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4450 assert!(
4451 collected
4452 .mounts
4453 .iter()
4454 .all(|m| !m.path.starts_with("/.agents/skills/"))
4455 );
4456 }
4457}