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 background_execution;
91mod btw;
92mod budgeting;
93pub mod compaction;
94mod current_time;
95mod data_knowledge;
96mod declarative;
97mod fake_aws;
98mod fake_crm;
99mod fake_financial;
100mod fake_warehouse;
101mod file_system;
102mod human_intent;
103mod infinity_context;
104mod knowledge_base;
105mod loop_detection;
106pub mod mcp;
107mod noop;
108mod openai_tool_search;
109#[cfg(feature = "ui-capabilities")]
110mod openui;
111mod parallel;
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 virtual_bash;
132mod web_fetch;
133mod workspace_volumes;
134
135pub use a2a_delegation::{
137 A2A_AGENT_DELEGATION_CAPABILITY_ID, A2aAgentDelegationCapability, CancelAgentTool,
138 GetAgentRunsTool, MessageAgentTool, SpawnAgentTool, WaitAgentTool,
139};
140#[cfg(feature = "ui-capabilities")]
141pub use a2ui::{A2UI_CAPABILITY_ID, A2UiCapability};
142pub use agent_handoff::{
143 AGENT_HANDOFF_CAPABILITY_ID, AgentHandoffCapability, GetAgentHandoffsTool,
144 MessageAgentHandoffTool, StartAgentHandoffTool,
145};
146pub use agent_instructions::{
147 AGENT_INSTRUCTIONS_CAPABILITY_ID, AGENTS_MD_PATH, AgentInstructionsCapability,
148 AgentInstructionsConfig, DEFAULT_AGENT_INSTRUCTIONS_FILE, MAX_AGENT_INSTRUCTIONS_FILES,
149 MAX_AGENTS_MD_SIZE, format_agents_md_content, format_instruction_file_content,
150};
151pub use attach_skill::{
152 AttachSkillCapability, SKILL_CAPABILITY_PREFIX, SKILLS_DISCOVERY_PATH, SkillContribution,
153 SkillInstructions, SkillMeta, SkillSource, discover_skills_from_entries, is_skill_capability,
154 parse_skill_capability_id, reconstruct_skill_md, skill_capability_id,
155};
156pub use background_execution::{BACKGROUND_EXECUTION_CAPABILITY_ID, BackgroundExecutionCapability};
157pub use btw::{BTW_CAPABILITY_ID, BtwCapability};
158pub use budgeting::BudgetingCapability;
159pub use compaction::{
160 COMPACTION_CAPABILITY_ID, CompactionCapability, CompactionConfig, CompactionStep,
161 CompactionStrategy, CostControlConfig, CostControlMaskingResult, HierarchicalMemoryConfig,
162 MaskingSummaryFormat, MemoryTier, ObservationMaskingConfig, ObservationMaskingResult,
163 SessionCompactionMetrics, SummarizationConfig, aggressive_trim, apply_cost_control_masking,
164 apply_hierarchical_memory, apply_observation_masking, build_model_view_messages,
165 build_summarization_prompt, build_summary_message, classify_memory_tiers, estimate_tokens,
166 estimate_total_tokens, format_messages_for_summarization, should_compact_proactively,
167};
168pub use current_time::{CurrentTimeCapability, GetCurrentTimeTool};
169pub use data_knowledge::DataKnowledgeCapability;
170pub use declarative::{
171 DECLARATIVE_CAPABILITY_PREFIX, DeclarativeCapabilityDefinition, DeclarativeCapabilityFile,
172 DeclarativeCapabilitySkill, declarative_capability_id, declarative_capability_info,
173 hydrate_declarative_capability_config, is_declarative_capability,
174 parse_declarative_capability_id, validate_declarative_capability_definition,
175};
176pub use fake_aws::{
177 AwsCreateEc2InstanceTool, AwsCreateIamUserTool, AwsCreateRdsDatabaseTool,
178 AwsCreateS3BucketTool, AwsGetCloudWatchMetricsTool, AwsListEc2InstancesTool,
179 AwsListIamUsersTool, AwsListRdsDatabasesTool, AwsListS3BucketsTool, AwsListSecurityGroupsTool,
180 AwsStopEc2InstanceTool, FakeAwsCapability,
181};
182pub use fake_crm::{
183 CrmAddInteractionTool, CrmCreateCustomerTool, CrmCreateTicketTool, CrmGetCustomerTool,
184 CrmListCustomersTool, CrmListTicketsTool, CrmSearchCustomersTool, CrmUpdateTicketTool,
185 FakeCrmCapability,
186};
187pub use fake_financial::{
188 FakeFinancialCapability, FinanceCreateBudgetTool, FinanceCreateTransactionTool,
189 FinanceForecastCashFlowTool, FinanceGetBalanceTool, FinanceGetExpenseReportTool,
190 FinanceGetRevenueReportTool, FinanceListBudgetsTool, FinanceListTransactionsTool,
191};
192pub use fake_warehouse::{
193 FakeWarehouseCapability, WarehouseCreateInvoiceTool, WarehouseCreateOrderTool,
194 WarehouseCreateShipmentTool, WarehouseGetInventoryTool, WarehouseInventoryReportTool,
195 WarehouseListOrdersTool, WarehouseListShipmentsTool, WarehouseProcessReturnTool,
196 WarehouseUpdateInventoryTool, WarehouseUpdateShipmentStatusTool,
197};
198pub use file_system::{
199 DeleteFileTool, EditFileTool, FileSystemCapability, GrepFilesTool, ListDirectoryTool,
200 ReadFileTool, StatFileTool, WriteFileTool,
201};
202pub use human_intent::{HUMAN_INTENT_CAPABILITY_ID, HumanIntentCapability};
203pub use infinity_context::{
204 INFINITY_CONTEXT_CAPABILITY_ID, InfinityContextCapability, QueryHistoryTool,
205};
206pub use knowledge_base::{
207 KNOWLEDGE_BASE_CAPABILITY_ID, KnowledgeBaseCapability, KnowledgeBaseConfig,
208 validate_knowledge_base_config,
209};
210pub use loop_detection::LoopDetectionCapability;
211pub use mcp::{
212 MCP_CAPABILITY_PREFIX, McpCapability, is_mcp_capability, mcp_capability_id,
213 parse_mcp_capability_id,
214};
215pub use noop::NoopCapability;
216pub use openai_tool_search::{
217 DEFAULT_TOOL_SEARCH_THRESHOLD, OPENAI_TOOL_SEARCH_CAPABILITY_ID, OpenAiToolSearchCapability,
218};
219#[cfg(feature = "ui-capabilities")]
220pub use openui::{OPENUI_CAPABILITY_ID, OpenUiCapability};
221pub use parallel::ParallelCapability;
222pub use persistent_memory::{
223 ForgetTool, MEMORY_CAPABILITY_ID, MemoryCapability, MemoryConfig, RecallTool, RememberTool,
224};
225pub use platform_management::{
226 ManageAgentsTool, ManageHarnessesTool, ManageSessionsTool, PlatformManagementCapability,
227 ReadAgentsTool, ReadCapabilitiesTool, ReadHarnessesTool, ReadSessionsTool,
228 SessionReadMessagesTool, SessionReadResponseTool, SessionSendMessageTool,
229};
230pub use prompt_caching::{PROMPT_CACHING_CAPABILITY_ID, PromptCachingCapability};
231pub use prompt_canary_guardrail::{
232 DEFAULT_REPLACEMENT as PROMPT_CANARY_DEFAULT_REPLACEMENT,
233 PROMPT_CANARY_GUARDRAIL_CAPABILITY_ID, PromptCanaryGuardrailCapability,
234 REASON_CODE_SYSTEM_PROMPT_LEAK,
235};
236pub use research::ResearchCapability;
237pub use sample_data::SampleDataCapability;
238pub use self_budget::{SELF_BUDGET_CAPABILITY_ID, SelfBudgetCapability};
239pub use session::{GetSessionInfoTool, SessionCapability, WriteSessionTitleTool};
240pub use session_sandbox::{
241 SESSION_SANDBOX_CAPABILITY_ID, SandboxExecTool, SandboxManageTool, SandboxReadFileTool,
242 SandboxStatusTool, SandboxWriteFileTool, SessionSandboxCapability,
243};
244pub use session_schedule::{
245 CancelScheduleTool, CreateScheduleTool, ListSchedulesTool, SESSION_SCHEDULE_CAPABILITY_ID,
246 SessionScheduleCapability,
247};
248pub use session_sql_database::{
249 SessionSqlDatabaseCapability, SqlExecuteTool, SqlQueryTool, SqlSchemaTool,
250};
251pub use session_storage::{KvStoreTool, SecretStoreTool, SessionStorageCapability};
252pub use skills::{SKILLS_CAPABILITY_ID, SkillsCapability};
253pub use stateless_todo_list::{StatelessTodoListCapability, WriteTodosTool};
254pub use subagents::SubagentCapability;
255pub use system_commands::{SYSTEM_COMMANDS_CAPABILITY_ID, SystemCommandsCapability};
257pub use test_math::{AddTool, DivideTool, MultiplyTool, SubtractTool, TestMathCapability};
258pub use test_weather::{GetForecastTool, GetWeatherTool, TestWeatherCapability};
259pub use tool_output_persistence::{PersistOutputHook, ToolOutputPersistenceCapability};
260pub use virtual_bash::{BashTool, VirtualBashCapability};
261pub use web_fetch::{
262 BotAuthPublicKey, WebFetchCapability, WebFetchTool, derive_bot_auth_public_key,
263};
264pub use workspace_volumes::{WORKSPACE_VOLUMES_CAPABILITY_ID, WorkspaceVolumesCapability};
265
266pub struct SystemPromptContext {
276 pub session_id: SessionId,
278 pub locale: Option<String>,
280 pub file_store: Option<Arc<dyn SessionFileSystem>>,
282}
283
284impl SystemPromptContext {
285 pub fn without_file_store(session_id: SessionId) -> Self {
287 Self {
288 session_id,
289 locale: None,
290 file_store: None,
291 }
292 }
293}
294
295#[async_trait]
342pub trait Capability: Send + Sync {
343 fn id(&self) -> &str;
345
346 fn name(&self) -> &str;
348
349 fn description(&self) -> &str;
351
352 fn status(&self) -> CapabilityStatus {
354 CapabilityStatus::Available
355 }
356
357 fn icon(&self) -> Option<&str> {
359 None
360 }
361
362 fn category(&self) -> Option<&str> {
364 None
365 }
366
367 fn system_prompt_addition(&self) -> Option<&str> {
387 None
388 }
389
390 async fn system_prompt_contribution(&self, _ctx: &SystemPromptContext) -> Option<String> {
402 self.system_prompt_addition().map(|addition| {
403 format!(
404 "<capability id=\"{}\">\n{}\n</capability>",
405 self.id(),
406 addition
407 )
408 })
409 }
410
411 fn system_prompt_preview(&self) -> Option<String> {
417 self.system_prompt_addition().map(|s| s.to_string())
418 }
419
420 fn tools(&self) -> Vec<Box<dyn Tool>> {
422 vec![]
423 }
424
425 fn tools_with_config(&self, _config: &serde_json::Value) -> Vec<Box<dyn Tool>> {
433 self.tools()
434 }
435
436 async fn system_prompt_contribution_with_config(
443 &self,
444 ctx: &SystemPromptContext,
445 _config: &serde_json::Value,
446 ) -> Option<String> {
447 self.system_prompt_contribution(ctx).await
448 }
449
450 fn tool_definitions(&self) -> Vec<ToolDefinition> {
453 self.tools().iter().map(|t| t.to_definition()).collect()
454 }
455
456 fn mounts(&self) -> Vec<MountPoint> {
464 vec![]
465 }
466
467 fn dependencies(&self) -> Vec<&'static str> {
476 vec![]
477 }
478
479 fn features(&self) -> Vec<&'static str> {
494 vec![]
495 }
496
497 fn config_schema(&self) -> Option<serde_json::Value> {
503 None
504 }
505
506 fn config_ui_schema(&self) -> Option<serde_json::Value> {
511 None
512 }
513
514 fn validate_config(&self, _config: &serde_json::Value) -> Result<(), String> {
520 Ok(())
521 }
522
523 fn mcp_servers(&self) -> ScopedMcpServers {
529 ScopedMcpServers::default()
530 }
531
532 fn mcp_servers_with_config(&self, _config: &serde_json::Value) -> ScopedMcpServers {
534 self.mcp_servers()
535 }
536
537 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
550 None
551 }
552
553 fn model_view_provider(&self) -> Option<Arc<dyn ModelViewProvider>> {
561 None
562 }
563
564 fn post_tool_exec_hooks(&self) -> Vec<Arc<dyn crate::atoms::PostToolExecHook>> {
572 vec![]
573 }
574
575 fn tool_definition_hooks(&self) -> Vec<Arc<dyn ToolDefinitionHook>> {
584 vec![]
585 }
586
587 fn tool_call_hooks(&self) -> Vec<Arc<dyn ToolCallHook>> {
595 vec![]
596 }
597
598 fn risk_level(&self) -> RiskLevel {
606 RiskLevel::Low
607 }
608
609 fn commands(&self) -> Vec<CommandDescriptor> {
617 vec![]
618 }
619
620 async fn execute_command(
634 &self,
635 request: &ExecuteCommandRequest,
636 _ctx: &CommandExecutionContext,
637 ) -> crate::error::Result<CommandResult> {
638 Err(crate::error::AgentLoopError::config(format!(
639 "capability {} declared command /{} but does not implement execute_command",
640 self.id(),
641 request.name,
642 )))
643 }
644
645 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
653 vec![]
654 }
655
656 fn contribute_skills(&self) -> Vec<SkillContribution> {
666 vec![]
667 }
668
669 fn output_guardrails(&self) -> Vec<Arc<dyn crate::output_guardrail::OutputGuardrail>> {
680 vec![]
681 }
682}
683
684pub trait ToolDefinitionHook: Send + Sync {
685 fn transform(&self, tools: Vec<ToolDefinition>) -> Vec<ToolDefinition>;
686}
687
688pub trait ToolCallHook: Send + Sync {
689 fn narration(
690 &self,
691 _tool_def: Option<&ToolDefinition>,
692 _tool_call: &ToolCall,
693 _phase: crate::tool_narration::ToolNarrationPhase,
694 _locale: Option<&str>,
695 ) -> Option<String> {
696 None
697 }
698
699 fn transform_for_execution(&self, tool_call: ToolCall) -> ToolCall {
700 tool_call
701 }
702}
703
704#[derive(
708 Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
709)]
710#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
711#[serde(rename_all = "lowercase")]
712pub enum RiskLevel {
713 Low,
715 Medium,
717 High,
719}
720
721#[derive(Debug, Clone, Serialize, Deserialize)]
727#[serde(rename_all = "snake_case")]
728pub enum BlueprintModel {
729 Fixed(String),
731 Default(String),
733 Inherit,
735}
736
737pub struct AgentBlueprint {
743 pub id: &'static str,
745 pub name: &'static str,
747 pub description: &'static str,
749 pub model: BlueprintModel,
751 pub system_prompt: &'static str,
753 pub tools: Vec<Box<dyn Tool>>,
755 pub max_turns: Option<usize>,
757 pub config_schema: Option<serde_json::Value>,
759}
760
761impl AgentBlueprint {
762 pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
764 self.tools.iter().map(|t| t.to_definition()).collect()
765 }
766}
767
768impl std::fmt::Debug for AgentBlueprint {
769 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
770 f.debug_struct("AgentBlueprint")
771 .field("id", &self.id)
772 .field("name", &self.name)
773 .field("model", &self.model)
774 .field("tool_count", &self.tools.len())
775 .field("max_turns", &self.max_turns)
776 .finish()
777 }
778}
779
780#[derive(Clone)]
807pub struct CapabilityRegistry {
808 capabilities: HashMap<String, Arc<dyn Capability>>,
809}
810
811impl CapabilityRegistry {
812 pub fn new() -> Self {
814 Self {
815 capabilities: HashMap::new(),
816 }
817 }
818
819 pub fn with_builtins() -> Self {
824 Self::with_builtins_for_grade(DeploymentGrade::from_env())
825 }
826
827 pub fn with_builtins_for_grade(grade: DeploymentGrade) -> Self {
832 let mut registry = Self::new();
833
834 registry.register(AgentInstructionsCapability);
836 registry.register(HumanIntentCapability);
837 registry.register(NoopCapability);
838 registry.register(CurrentTimeCapability);
839 registry.register(ResearchCapability);
840 registry.register(PlatformManagementCapability);
841 registry.register(FileSystemCapability);
842 registry.register(WorkspaceVolumesCapability);
843 registry.register(SessionStorageCapability);
844 registry.register(SessionCapability);
845 registry.register(SessionSqlDatabaseCapability);
846 registry.register(TestMathCapability);
847 registry.register(TestWeatherCapability);
848 registry.register(StatelessTodoListCapability);
849 registry.register(WebFetchCapability::from_env());
850 registry.register(VirtualBashCapability);
851 registry.register(BackgroundExecutionCapability);
852 registry.register(SessionScheduleCapability);
853 registry.register(BtwCapability);
854 registry.register(InfinityContextCapability);
855 registry.register(budgeting::BudgetingCapability);
856 registry.register(SelfBudgetCapability);
857 registry.register(ParallelCapability);
858 registry.register(CompactionCapability);
859 registry.register(MemoryCapability);
860
861 registry.register(OpenAiToolSearchCapability::new());
863 registry.register(PromptCachingCapability::new());
864
865 registry.register(SkillsCapability);
867
868 registry.register(SubagentCapability);
870 registry.register(AgentHandoffCapability);
871 registry.register(A2aAgentDelegationCapability);
872
873 registry.register(SystemCommandsCapability);
875
876 registry.register(tool_output_persistence::ToolOutputPersistenceCapability);
878
879 registry.register(LoopDetectionCapability);
881
882 registry.register(PromptCanaryGuardrailCapability);
885
886 #[cfg(feature = "ui-capabilities")]
888 {
889 registry.register(OpenUiCapability);
890 registry.register(A2UiCapability);
891 }
892
893 registry.register(SampleDataCapability);
895
896 registry.register(DataKnowledgeCapability);
898
899 registry.register(KnowledgeBaseCapability);
901
902 registry.register(FakeWarehouseCapability);
904 registry.register(FakeAwsCapability);
905 registry.register(FakeCrmCapability);
906 registry.register(FakeFinancialCapability);
907
908 let internal_flags = crate::InternalFeatureFlags::from_env();
910 if internal_flags.session_sandbox {
911 registry.register(SessionSandboxCapability);
912 }
913 for plugin in inventory::iter::<IntegrationPlugin>() {
914 if (!plugin.experimental_only || grade.experimental_features_enabled())
915 && plugin
916 .feature_flag
917 .is_none_or(|f| internal_flags.is_enabled(f))
918 {
919 registry.register_boxed((plugin.factory)());
920 }
921 }
922
923 registry
924 }
925
926 pub fn register(&mut self, capability: impl Capability + 'static) {
928 self.capabilities
929 .insert(capability.id().to_string(), Arc::new(capability));
930 }
931
932 pub fn register_boxed(&mut self, capability: Box<dyn Capability>) {
934 self.capabilities
935 .insert(capability.id().to_string(), Arc::from(capability));
936 }
937
938 pub fn register_arc(&mut self, capability: Arc<dyn Capability>) {
940 self.capabilities
941 .insert(capability.id().to_string(), capability);
942 }
943
944 pub fn get(&self, id: &str) -> Option<&Arc<dyn Capability>> {
946 self.capabilities.get(id)
947 }
948
949 pub fn unregister(&mut self, id: &str) -> Option<Arc<dyn Capability>> {
951 self.capabilities.remove(id)
952 }
953
954 pub fn has(&self, id: &str) -> bool {
956 self.capabilities.contains_key(id)
957 }
958
959 pub fn list(&self) -> Vec<&Arc<dyn Capability>> {
961 self.capabilities.values().collect()
962 }
963
964 pub fn len(&self) -> usize {
966 self.capabilities.len()
967 }
968
969 pub fn is_empty(&self) -> bool {
971 self.capabilities.is_empty()
972 }
973
974 pub fn builder() -> CapabilityRegistryBuilder {
976 CapabilityRegistryBuilder::new()
977 }
978
979 pub fn blueprint(&self, id: &str) -> Option<AgentBlueprint> {
983 for cap in self.capabilities.values() {
984 for bp in cap.agent_blueprints() {
985 if bp.id == id {
986 return Some(bp);
987 }
988 }
989 }
990 None
991 }
992
993 pub fn blueprint_with_capability(&self, id: &str) -> Option<(String, AgentBlueprint)> {
997 for (capability_id, cap) in &self.capabilities {
998 for bp in cap.agent_blueprints() {
999 if bp.id == id {
1000 return Some((capability_id.clone(), bp));
1001 }
1002 }
1003 }
1004 None
1005 }
1006
1007 pub fn all_blueprints(&self) -> Vec<AgentBlueprint> {
1009 self.capabilities
1010 .values()
1011 .flat_map(|cap| cap.agent_blueprints())
1012 .collect()
1013 }
1014}
1015
1016impl Default for CapabilityRegistry {
1017 fn default() -> Self {
1018 Self::with_builtins()
1019 }
1020}
1021
1022impl std::fmt::Debug for CapabilityRegistry {
1023 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1024 let ids: Vec<_> = self.capabilities.keys().collect();
1025 f.debug_struct("CapabilityRegistry")
1026 .field("capabilities", &ids)
1027 .finish()
1028 }
1029}
1030
1031pub struct CapabilityRegistryBuilder {
1033 registry: CapabilityRegistry,
1034}
1035
1036impl CapabilityRegistryBuilder {
1037 pub fn new() -> Self {
1039 Self {
1040 registry: CapabilityRegistry::new(),
1041 }
1042 }
1043
1044 pub fn with_builtins() -> Self {
1046 Self {
1047 registry: CapabilityRegistry::with_builtins(),
1048 }
1049 }
1050
1051 pub fn capability(mut self, capability: impl Capability + 'static) -> Self {
1053 self.registry.register(capability);
1054 self
1055 }
1056
1057 pub fn build(self) -> CapabilityRegistry {
1059 self.registry
1060 }
1061}
1062
1063impl Default for CapabilityRegistryBuilder {
1064 fn default() -> Self {
1065 Self::new()
1066 }
1067}
1068
1069pub struct ModelViewContext<'a> {
1075 pub session_id: SessionId,
1076 pub prior_usage: Option<&'a TokenUsage>,
1077}
1078
1079pub trait ModelViewProvider: Send + Sync {
1085 fn apply_model_view(
1086 &self,
1087 messages: Vec<Message>,
1088 config: &serde_json::Value,
1089 context: &ModelViewContext<'_>,
1090 ) -> Vec<Message>;
1091
1092 fn priority(&self) -> i32 {
1093 0
1094 }
1095}
1096
1097pub struct CollectedCapabilities {
1102 pub system_prompt_parts: Vec<String>,
1104 pub system_prompt_attributions: Vec<SystemPromptAttribution>,
1106 pub tools: Vec<Box<dyn Tool>>,
1108 pub tool_definitions: Vec<ToolDefinition>,
1110 pub mounts: Vec<MountPoint>,
1112 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1114 pub applied_ids: Vec<String>,
1116 pub tool_search: Option<crate::llm_driver_registry::ToolSearchConfig>,
1118 pub prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig>,
1120 pub tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>>,
1122 pub tool_call_hooks: Vec<Arc<dyn ToolCallHook>>,
1124 pub mcp_servers: ScopedMcpServers,
1126 }
1132
1133#[derive(Debug, Clone, PartialEq, Eq)]
1134pub struct SystemPromptAttribution {
1135 pub capability_id: String,
1136 pub content: String,
1137}
1138
1139impl CollectedCapabilities {
1140 pub fn system_prompt_prefix(&self) -> Option<String> {
1143 if self.system_prompt_parts.is_empty() {
1144 None
1145 } else {
1146 Some(self.system_prompt_parts.join("\n\n"))
1147 }
1148 }
1149
1150 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1154 for (provider, config) in &self.message_filter_providers {
1156 provider.apply_filters(query, config);
1157 }
1158 }
1159
1160 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1163 for (provider, config) in &self.message_filter_providers {
1164 provider.post_load(messages, config);
1165 }
1166 }
1167
1168 pub fn has_message_filters(&self) -> bool {
1170 !self.message_filter_providers.is_empty()
1171 }
1172}
1173
1174pub struct CollectedMessageFilters {
1181 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1183}
1184
1185pub struct CollectedModelViewProviders {
1187 pub model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)>,
1189}
1190
1191impl CollectedMessageFilters {
1197 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1199 for (provider, config) in &self.message_filter_providers {
1200 provider.apply_filters(query, config);
1201 }
1202 }
1203
1204 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1206 for (provider, config) in &self.message_filter_providers {
1207 provider.post_load(messages, config);
1208 }
1209 }
1210}
1211
1212impl CollectedModelViewProviders {
1213 pub fn apply_model_view(
1215 &self,
1216 mut messages: Vec<Message>,
1217 context: &ModelViewContext<'_>,
1218 ) -> Vec<Message> {
1219 for (provider, config) in &self.model_view_providers {
1220 messages = provider.apply_model_view(messages, config, context);
1221 }
1222 messages
1223 }
1224}
1225
1226pub fn collect_message_filters_only(
1232 capability_configs: &[AgentCapabilityConfig],
1233 registry: &CapabilityRegistry,
1234) -> CollectedMessageFilters {
1235 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1236 Vec::new();
1237
1238 for cap_config in capability_configs {
1239 let cap_id = cap_config.capability_ref.as_str();
1240 if let Some(capability) = registry.get(cap_id) {
1241 if capability.status() != CapabilityStatus::Available {
1242 continue;
1243 }
1244 if let Some(provider) = capability.message_filter_provider() {
1245 message_filter_providers.push((provider, cap_config.config.clone()));
1246 }
1247 }
1248 }
1249
1250 message_filter_providers.sort_by_key(|(p, _)| p.priority());
1251
1252 CollectedMessageFilters {
1253 message_filter_providers,
1254 }
1255}
1256
1257pub fn collect_model_view_providers(
1259 capability_configs: &[AgentCapabilityConfig],
1260 registry: &CapabilityRegistry,
1261) -> CollectedModelViewProviders {
1262 let mut model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)> = Vec::new();
1263
1264 for cap_config in capability_configs {
1265 let cap_id = cap_config.capability_ref.as_str();
1266 if let Some(capability) = registry.get(cap_id) {
1267 if capability.status() != CapabilityStatus::Available {
1268 continue;
1269 }
1270 if let Some(provider) = capability.model_view_provider() {
1271 model_view_providers.push((provider, cap_config.config.clone()));
1272 }
1273 }
1274 }
1275
1276 model_view_providers.sort_by_key(|(p, _)| p.priority());
1277
1278 CollectedModelViewProviders {
1279 model_view_providers,
1280 }
1281}
1282
1283pub fn collect_capability_mcp_servers(
1284 capability_configs: &[AgentCapabilityConfig],
1285 registry: &CapabilityRegistry,
1286) -> ScopedMcpServers {
1287 let mut servers = ScopedMcpServers::default();
1288
1289 for cap_config in capability_configs {
1290 let cap_id = cap_config.capability_ref.as_str();
1291 if is_declarative_capability(cap_id) {
1292 if let Ok(definition) =
1293 serde_json::from_value::<DeclarativeCapabilityDefinition>(cap_config.config.clone())
1294 {
1295 if definition.status != CapabilityStatus::Available {
1296 continue;
1297 }
1298 if let Some(contributed) = definition.mcp_servers {
1299 servers = merge_scoped_mcp_servers(&servers, &contributed);
1300 }
1301 }
1302 continue;
1303 }
1304 if let Some(capability) = registry.get(cap_id) {
1305 if capability.status() != CapabilityStatus::Available {
1306 continue;
1307 }
1308 servers = merge_scoped_mcp_servers(
1309 &servers,
1310 &capability.mcp_servers_with_config(&cap_config.config),
1311 );
1312 }
1313 }
1314
1315 servers
1316}
1317
1318pub const MAX_RESOLVED_CAPABILITIES: usize = 100;
1325
1326#[derive(Debug, Clone, PartialEq, Eq)]
1328pub enum DependencyError {
1329 CircularDependency {
1331 capability_id: String,
1333 chain: Vec<String>,
1335 },
1336 TooManyCapabilities {
1338 count: usize,
1340 max: usize,
1342 },
1343}
1344
1345impl std::fmt::Display for DependencyError {
1346 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1347 match self {
1348 DependencyError::CircularDependency {
1349 capability_id,
1350 chain,
1351 } => {
1352 write!(
1353 f,
1354 "Circular dependency detected: {} depends on itself via chain: {} -> {}",
1355 capability_id,
1356 chain.join(" -> "),
1357 capability_id
1358 )
1359 }
1360 DependencyError::TooManyCapabilities { count, max } => {
1361 write!(
1362 f,
1363 "Too many capabilities after resolution: {} (max: {})",
1364 count, max
1365 )
1366 }
1367 }
1368 }
1369}
1370
1371impl std::error::Error for DependencyError {}
1372
1373#[derive(Debug, Clone)]
1375pub struct ResolvedCapabilities {
1376 pub resolved_ids: Vec<String>,
1379 pub added_as_dependencies: Vec<String>,
1381 pub user_selected: Vec<String>,
1383}
1384
1385pub fn resolve_dependencies(
1405 selected_ids: &[String],
1406 registry: &CapabilityRegistry,
1407) -> Result<ResolvedCapabilities, DependencyError> {
1408 use std::collections::HashSet;
1409
1410 let user_selected: HashSet<String> = selected_ids.iter().cloned().collect();
1411 let mut resolved: Vec<String> = Vec::new();
1412 let mut resolved_set: HashSet<String> = HashSet::new();
1413 let mut added_as_dependencies: Vec<String> = Vec::new();
1414
1415 for cap_id in selected_ids {
1417 resolve_single_capability(
1418 cap_id,
1419 registry,
1420 &mut resolved,
1421 &mut resolved_set,
1422 &mut added_as_dependencies,
1423 &user_selected,
1424 &mut Vec::new(), )?;
1426 }
1427
1428 if resolved.len() > MAX_RESOLVED_CAPABILITIES {
1430 return Err(DependencyError::TooManyCapabilities {
1431 count: resolved.len(),
1432 max: MAX_RESOLVED_CAPABILITIES,
1433 });
1434 }
1435
1436 Ok(ResolvedCapabilities {
1437 resolved_ids: resolved,
1438 added_as_dependencies,
1439 user_selected: selected_ids.to_vec(),
1440 })
1441}
1442
1443pub fn resolve_capability_configs(
1448 selected_configs: &[AgentCapabilityConfig],
1449 registry: &CapabilityRegistry,
1450) -> Result<Vec<AgentCapabilityConfig>, DependencyError> {
1451 let mut selected_ids: Vec<String> = Vec::new();
1452 for config in selected_configs {
1453 if is_declarative_capability(config.capability_id())
1454 && let Ok(definition) =
1455 serde_json::from_value::<DeclarativeCapabilityDefinition>(config.config.clone())
1456 {
1457 selected_ids.extend(definition.dependencies);
1458 }
1459 selected_ids.push(config.capability_id().to_string());
1460 }
1461 let resolved = resolve_dependencies(&selected_ids, registry)?;
1462
1463 let explicit_configs: std::collections::HashMap<String, serde_json::Value> = selected_configs
1464 .iter()
1465 .map(|config| (config.capability_id().to_string(), config.config.clone()))
1466 .collect();
1467
1468 Ok(resolved
1469 .resolved_ids
1470 .into_iter()
1471 .map(|capability_id| {
1472 explicit_configs
1473 .get(&capability_id)
1474 .cloned()
1475 .map(|config| AgentCapabilityConfig::with_config(capability_id.clone(), config))
1476 .unwrap_or_else(|| AgentCapabilityConfig::new(capability_id))
1477 })
1478 .collect())
1479}
1480
1481fn resolve_single_capability(
1483 cap_id: &str,
1484 registry: &CapabilityRegistry,
1485 resolved: &mut Vec<String>,
1486 resolved_set: &mut std::collections::HashSet<String>,
1487 added_as_dependencies: &mut Vec<String>,
1488 user_selected: &std::collections::HashSet<String>,
1489 visiting: &mut Vec<String>,
1490) -> Result<(), DependencyError> {
1491 if resolved_set.contains(cap_id) {
1493 return Ok(());
1494 }
1495
1496 if visiting.contains(&cap_id.to_string()) {
1498 return Err(DependencyError::CircularDependency {
1499 capability_id: cap_id.to_string(),
1500 chain: visiting.clone(),
1501 });
1502 }
1503
1504 let capability = match registry.get(cap_id) {
1506 Some(cap) => cap,
1507 None => {
1508 if is_declarative_capability(cap_id) && !resolved_set.contains(cap_id) {
1509 resolved.push(cap_id.to_string());
1510 resolved_set.insert(cap_id.to_string());
1511 if !user_selected.contains(cap_id) {
1512 added_as_dependencies.push(cap_id.to_string());
1513 }
1514 }
1515 return Ok(());
1516 }
1517 };
1518
1519 visiting.push(cap_id.to_string());
1521
1522 for dep_id in capability.dependencies() {
1524 resolve_single_capability(
1525 dep_id,
1526 registry,
1527 resolved,
1528 resolved_set,
1529 added_as_dependencies,
1530 user_selected,
1531 visiting,
1532 )?;
1533 }
1534
1535 visiting.pop();
1537
1538 if !resolved_set.contains(cap_id) {
1540 resolved.push(cap_id.to_string());
1541 resolved_set.insert(cap_id.to_string());
1542
1543 if !user_selected.contains(cap_id) {
1545 added_as_dependencies.push(cap_id.to_string());
1546 }
1547 }
1548
1549 Ok(())
1550}
1551
1552pub fn compute_features(capability_ids: &[String], registry: &CapabilityRegistry) -> Vec<String> {
1557 use std::collections::HashSet;
1558
1559 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
1560 Ok(resolved) => resolved.resolved_ids,
1561 Err(_) => capability_ids.to_vec(),
1562 };
1563
1564 let mut seen = HashSet::new();
1565 let mut features = Vec::new();
1566 for cap_id in &resolved_ids {
1567 if let Some(cap) = registry.get(cap_id) {
1568 for feature in cap.features() {
1569 if seen.insert(feature) {
1570 features.push(feature.to_string());
1571 }
1572 }
1573 }
1574 }
1575 features
1576}
1577
1578pub fn get_dependencies(cap_id: &str, registry: &CapabilityRegistry) -> Vec<String> {
1581 registry
1582 .get(cap_id)
1583 .map(|cap| cap.dependencies().iter().map(|s| s.to_string()).collect())
1584 .unwrap_or_default()
1585}
1586
1587pub async fn collect_capabilities(
1603 capability_ids: &[String],
1604 registry: &CapabilityRegistry,
1605 ctx: &SystemPromptContext,
1606) -> CollectedCapabilities {
1607 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
1610 Ok(resolved) => resolved.resolved_ids,
1611 Err(e) => {
1612 tracing::warn!("Failed to resolve capability dependencies: {}", e);
1613 capability_ids.to_vec()
1614 }
1615 };
1616
1617 let configs: Vec<AgentCapabilityConfig> = resolved_ids
1619 .iter()
1620 .map(|id| AgentCapabilityConfig {
1621 capability_ref: CapabilityId::new(id),
1622 config: serde_json::Value::Object(serde_json::Map::new()),
1623 })
1624 .collect();
1625
1626 collect_capabilities_with_configs(&configs, registry, ctx).await
1627}
1628
1629pub async fn collect_capabilities_with_configs(
1640 capability_configs: &[AgentCapabilityConfig],
1641 registry: &CapabilityRegistry,
1642 ctx: &SystemPromptContext,
1643) -> CollectedCapabilities {
1644 let mut system_prompt_parts: Vec<String> = Vec::new();
1645 let mut system_prompt_attributions: Vec<SystemPromptAttribution> = Vec::new();
1646 let mut tools: Vec<Box<dyn Tool>> = Vec::new();
1647 let mut tool_definitions: Vec<ToolDefinition> = Vec::new();
1648 let mut mounts: Vec<MountPoint> = Vec::new();
1649 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1650 Vec::new();
1651 let mut applied_ids: Vec<String> = Vec::new();
1652 let mut tool_search: Option<crate::llm_driver_registry::ToolSearchConfig> = None;
1653 let mut prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig> = None;
1654 let mut tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>> = Vec::new();
1655 let mut tool_call_hooks: Vec<Arc<dyn ToolCallHook>> = Vec::new();
1656 let mut mcp_servers = ScopedMcpServers::default();
1657
1658 for cap_config in capability_configs {
1659 let cap_id = cap_config.capability_ref.as_str();
1660 if is_declarative_capability(cap_id) {
1661 match serde_json::from_value::<DeclarativeCapabilityDefinition>(
1662 cap_config.config.clone(),
1663 ) {
1664 Ok(definition) => {
1665 if definition.status != CapabilityStatus::Available {
1666 continue;
1667 }
1668
1669 if let Some(prompt) = definition.system_prompt.as_deref() {
1670 let contribution =
1671 format!("<capability id=\"{}\">\n{}\n</capability>", cap_id, prompt);
1672 system_prompt_attributions.push(SystemPromptAttribution {
1673 capability_id: cap_id.to_string(),
1674 content: contribution.clone(),
1675 });
1676 system_prompt_parts.push(contribution);
1677 }
1678
1679 mounts.extend(definition.mounts(cap_id));
1680 if let Some(ref servers) = definition.mcp_servers {
1681 mcp_servers = merge_scoped_mcp_servers(&mcp_servers, servers);
1682 }
1683 for skill in definition.skill_contributions() {
1684 mounts.push(skill.to_mount(cap_id));
1685 }
1686
1687 applied_ids.push(cap_id.to_string());
1688 }
1689 Err(error) => {
1690 tracing::warn!(
1691 capability_id = %cap_id,
1692 error = %error,
1693 "Skipping invalid declarative capability config"
1694 );
1695 }
1696 }
1697 continue;
1698 }
1699 if let Some(capability) = registry.get(cap_id) {
1700 if capability.status() != CapabilityStatus::Available {
1702 continue;
1703 }
1704
1705 if let Some(contribution) = capability
1707 .system_prompt_contribution_with_config(ctx, &cap_config.config)
1708 .await
1709 {
1710 system_prompt_attributions.push(SystemPromptAttribution {
1711 capability_id: cap_id.to_string(),
1712 content: contribution.clone(),
1713 });
1714 system_prompt_parts.push(contribution);
1715 }
1716
1717 tools.extend(capability.tools_with_config(&cap_config.config));
1719 tool_definition_hooks.extend(capability.tool_definition_hooks());
1720 tool_call_hooks.extend(capability.tool_call_hooks());
1721 let cap_category = capability.category();
1726 for def in capability.tool_definitions() {
1727 let def = match (def.category(), cap_category) {
1728 (None, Some(cat)) => def.with_category(cat),
1729 _ => def,
1730 }
1731 .with_capability_attribution(cap_id, Some(capability.name()));
1732 tool_definitions.push(def);
1733 }
1734
1735 if cap_id == OPENAI_TOOL_SEARCH_CAPABILITY_ID {
1737 let threshold = cap_config
1739 .config
1740 .get("threshold")
1741 .and_then(|v| v.as_u64())
1742 .map(|v| v as usize)
1743 .unwrap_or(DEFAULT_TOOL_SEARCH_THRESHOLD);
1744 tool_search = Some(crate::llm_driver_registry::ToolSearchConfig {
1745 enabled: true,
1746 threshold,
1747 });
1748 }
1749
1750 if cap_id == PROMPT_CACHING_CAPABILITY_ID {
1751 let strategy = cap_config
1752 .config
1753 .get("strategy")
1754 .and_then(|v| v.as_str())
1755 .map(|value| match value {
1756 "auto" => crate::llm_driver_registry::PromptCacheStrategy::Auto,
1757 _ => crate::llm_driver_registry::PromptCacheStrategy::Auto,
1758 })
1759 .unwrap_or(crate::llm_driver_registry::PromptCacheStrategy::Auto);
1760 let gemini_cached_content = cap_config
1761 .config
1762 .get("gemini_cached_content")
1763 .and_then(|v| v.as_str())
1764 .map(str::to_string);
1765 prompt_cache = Some(crate::llm_driver_registry::PromptCacheConfig {
1766 enabled: true,
1767 strategy,
1768 gemini_cached_content,
1769 });
1770 }
1771
1772 mounts.extend(capability.mounts());
1774
1775 mcp_servers = merge_scoped_mcp_servers(
1776 &mcp_servers,
1777 &capability.mcp_servers_with_config(&cap_config.config),
1778 );
1779
1780 for skill in capability.contribute_skills() {
1784 mounts.push(skill.to_mount(cap_id));
1785 }
1786
1787 if let Some(provider) = capability.message_filter_provider() {
1789 message_filter_providers.push((provider, cap_config.config.clone()));
1790 }
1791
1792 applied_ids.push(cap_id.to_string());
1793 }
1794 }
1795
1796 if !applied_ids
1808 .iter()
1809 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID)
1810 && tool_definitions
1811 .iter()
1812 .any(|def| def.hints().supports_background == Some(true))
1813 && let Some(bg_cap) = registry.get(BACKGROUND_EXECUTION_CAPABILITY_ID)
1814 && bg_cap.status() == CapabilityStatus::Available
1815 {
1816 tools.extend(bg_cap.tools());
1817 let cap_category = bg_cap.category();
1818 for def in bg_cap.tool_definitions() {
1819 let def = match (def.category(), cap_category) {
1820 (None, Some(cat)) => def.with_category(cat),
1821 _ => def,
1822 }
1823 .with_capability_attribution(BACKGROUND_EXECUTION_CAPABILITY_ID, Some(bg_cap.name()));
1824 tool_definitions.push(def);
1825 }
1826 applied_ids.push(BACKGROUND_EXECUTION_CAPABILITY_ID.to_string());
1827 }
1828
1829 message_filter_providers.sort_by_key(|(p, _)| p.priority());
1831
1832 CollectedCapabilities {
1833 system_prompt_parts,
1834 system_prompt_attributions,
1835 tools,
1836 tool_definitions,
1837 mounts,
1838 message_filter_providers,
1839 applied_ids,
1840 tool_search,
1841 prompt_cache,
1842 tool_definition_hooks,
1843 tool_call_hooks,
1844 mcp_servers,
1845 }
1846}
1847
1848pub struct AppliedCapabilities {
1854 pub runtime_agent: RuntimeAgent,
1856 pub tool_registry: ToolRegistry,
1858 pub applied_ids: Vec<String>,
1860}
1861
1862pub async fn apply_capabilities(
1899 base_runtime_agent: RuntimeAgent,
1900 capability_ids: &[String],
1901 registry: &CapabilityRegistry,
1902 ctx: &SystemPromptContext,
1903) -> AppliedCapabilities {
1904 let collected = collect_capabilities(capability_ids, registry, ctx).await;
1905
1906 let final_system_prompt = match collected.system_prompt_prefix() {
1908 Some(prefix) => format!(
1909 "{}\n\n<system-prompt>\n{}\n</system-prompt>",
1910 prefix, base_runtime_agent.system_prompt
1911 ),
1912 None => base_runtime_agent.system_prompt,
1913 };
1914
1915 let mut tool_registry = ToolRegistry::new();
1917 for tool in collected.tools {
1918 tool_registry.register_boxed(tool);
1919 }
1920
1921 let mut tools = collected.tool_definitions;
1923 for hook in &collected.tool_definition_hooks {
1924 tools = hook.transform(tools);
1925 }
1926
1927 let runtime_agent = RuntimeAgent {
1928 system_prompt: final_system_prompt,
1929 model: base_runtime_agent.model,
1930 tools,
1931 max_iterations: base_runtime_agent.max_iterations,
1932 temperature: base_runtime_agent.temperature,
1933 max_tokens: base_runtime_agent.max_tokens,
1934 tool_search: collected.tool_search,
1935 prompt_cache: collected.prompt_cache,
1936 network_access: base_runtime_agent.network_access,
1937 };
1938
1939 AppliedCapabilities {
1940 runtime_agent,
1941 tool_registry,
1942 applied_ids: collected.applied_ids,
1943 }
1944}
1945
1946#[cfg(test)]
1951mod tests {
1952 use super::*;
1953 use crate::typed_id::SessionId;
1954 use std::collections::BTreeSet;
1955 use uuid::Uuid;
1956
1957 fn test_ctx() -> SystemPromptContext {
1959 SystemPromptContext::without_file_store(SessionId::new())
1960 }
1961
1962 fn expected_core_builtin_ids() -> BTreeSet<&'static str> {
1963 let mut ids = [
1964 "agent_instructions",
1965 "human_intent",
1966 "budgeting",
1967 "self_budget",
1968 "parallel",
1969 "noop",
1970 "current_time",
1971 "research",
1972 "platform_management",
1973 "session_file_system",
1974 "workspace_volumes",
1975 "session_storage",
1976 "session",
1977 "session_sql_database",
1978 "test_math",
1979 "test_weather",
1980 "stateless_todo_list",
1981 "web_fetch",
1982 "virtual_bash",
1983 "background_execution",
1984 "session_schedule",
1985 "btw",
1986 "infinity_context",
1987 "compaction",
1988 "memory",
1989 "openai_tool_search",
1990 "prompt_caching",
1991 "skills",
1992 "subagents",
1993 "agent_handoff",
1994 "a2a_agent_delegation",
1995 "system_commands",
1996 "sample_data",
1997 "data_knowledge",
1998 "knowledge_base",
1999 "tool_output_persistence",
2000 "fake_warehouse",
2001 "fake_aws",
2002 "fake_crm",
2003 "fake_financial",
2004 "loop_detection",
2005 "prompt_canary_guardrail",
2006 ]
2007 .into_iter()
2008 .collect::<BTreeSet<_>>();
2009 if cfg!(feature = "ui-capabilities") {
2010 ids.insert("openui");
2011 ids.insert("a2ui");
2012 }
2013 ids
2014 }
2015
2016 fn registry_ids(registry: &CapabilityRegistry) -> BTreeSet<&str> {
2017 registry.capabilities.keys().map(String::as_str).collect()
2018 }
2019
2020 #[test]
2030 fn test_capability_registry_with_builtins_dev() {
2031 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
2033
2034 assert_eq!(registry_ids(®istry), expected_core_builtin_ids());
2035 }
2036
2037 #[test]
2038 fn test_capability_registry_with_builtins_prod() {
2039 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2041 assert_eq!(registry_ids(®istry), expected_core_builtin_ids());
2042 assert!(!registry.has("docker_container"));
2044 }
2045
2046 #[test]
2047 fn test_capability_registry_get() {
2048 let registry = CapabilityRegistry::with_builtins();
2049
2050 let noop = registry.get("noop").unwrap();
2051 assert_eq!(noop.id(), "noop");
2052 assert_eq!(noop.name(), "No-Op");
2053 assert_eq!(noop.status(), CapabilityStatus::Available);
2054 }
2055
2056 #[test]
2057 fn test_capability_registry_blueprint_with_capability() {
2058 struct BlueprintProviderCapability;
2059
2060 impl Capability for BlueprintProviderCapability {
2061 fn id(&self) -> &str {
2062 "blueprint_provider"
2063 }
2064 fn name(&self) -> &str {
2065 "Blueprint Provider"
2066 }
2067 fn description(&self) -> &str {
2068 "Capability that provides a blueprint for tests"
2069 }
2070 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
2071 vec![AgentBlueprint {
2072 id: "test_blueprint",
2073 name: "Test Blueprint",
2074 description: "Blueprint for capability registry tests",
2075 model: BlueprintModel::Inherit,
2076 system_prompt: "Test prompt",
2077 tools: vec![],
2078 max_turns: None,
2079 config_schema: None,
2080 }]
2081 }
2082 }
2083
2084 let mut registry = CapabilityRegistry::new();
2085 registry.register(BlueprintProviderCapability);
2086
2087 let (capability_id, blueprint) = registry
2088 .blueprint_with_capability("test_blueprint")
2089 .expect("blueprint should resolve with capability id");
2090 assert_eq!(capability_id, "blueprint_provider");
2091 assert_eq!(blueprint.id, "test_blueprint");
2092 }
2093
2094 #[test]
2095 fn test_capability_registry_builder() {
2096 let registry = CapabilityRegistry::builder()
2097 .capability(NoopCapability)
2098 .capability(CurrentTimeCapability)
2099 .build();
2100
2101 assert!(registry.has("noop"));
2102 assert!(registry.has("current_time"));
2103 assert_eq!(registry.len(), 2);
2104 }
2105
2106 #[test]
2107 fn test_capability_status() {
2108 let registry = CapabilityRegistry::with_builtins();
2109
2110 let current_time = registry.get("current_time").unwrap();
2111 assert_eq!(current_time.status(), CapabilityStatus::Available);
2112
2113 let research = registry.get("research").unwrap();
2114 assert_eq!(research.status(), CapabilityStatus::ComingSoon);
2115 }
2116
2117 #[test]
2118 fn test_capability_icons_and_categories() {
2119 let registry = CapabilityRegistry::with_builtins();
2120
2121 let noop = registry.get("noop").unwrap();
2122 assert_eq!(noop.icon(), Some("circle-off"));
2123 assert_eq!(noop.category(), Some("Testing"));
2124
2125 let current_time = registry.get("current_time").unwrap();
2126 assert_eq!(current_time.icon(), Some("clock"));
2127 assert_eq!(current_time.category(), Some("Utilities"));
2128 }
2129
2130 #[test]
2131 fn test_system_prompt_preview_default_delegates_to_addition() {
2132 let registry = CapabilityRegistry::with_builtins();
2133
2134 let test_math = registry.get("test_math").unwrap();
2136 assert_eq!(
2137 test_math.system_prompt_preview().as_deref(),
2138 test_math.system_prompt_addition()
2139 );
2140
2141 let current_time = registry.get("current_time").unwrap();
2143 assert!(current_time.system_prompt_preview().is_none());
2144 assert!(current_time.system_prompt_addition().is_none());
2145 }
2146
2147 #[test]
2148 fn test_system_prompt_preview_dynamic_capability() {
2149 let registry = CapabilityRegistry::with_builtins();
2150 let cap = registry.get("agent_instructions").unwrap();
2151
2152 assert!(cap.system_prompt_addition().is_none());
2154 assert!(cap.system_prompt_preview().is_some());
2155 assert!(cap.system_prompt_preview().unwrap().contains("AGENTS.md"));
2156 }
2157
2158 #[tokio::test]
2163 async fn test_apply_capabilities_empty() {
2164 let registry = CapabilityRegistry::with_builtins();
2165 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2166
2167 let applied =
2168 apply_capabilities(base_runtime_agent.clone(), &[], ®istry, &test_ctx()).await;
2169
2170 assert_eq!(
2171 applied.runtime_agent.system_prompt,
2172 base_runtime_agent.system_prompt
2173 );
2174 assert!(applied.tool_registry.is_empty());
2175 assert!(applied.applied_ids.is_empty());
2176 }
2177
2178 #[tokio::test]
2179 async fn test_apply_capabilities_noop() {
2180 let registry = CapabilityRegistry::with_builtins();
2181 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2182
2183 let applied = apply_capabilities(
2184 base_runtime_agent.clone(),
2185 &["noop".to_string()],
2186 ®istry,
2187 &test_ctx(),
2188 )
2189 .await;
2190
2191 assert_eq!(
2193 applied.runtime_agent.system_prompt,
2194 base_runtime_agent.system_prompt
2195 );
2196 assert!(applied.tool_registry.is_empty());
2197 assert_eq!(applied.applied_ids, vec!["noop"]);
2198 }
2199
2200 #[tokio::test]
2201 async fn test_apply_capabilities_current_time() {
2202 let registry = CapabilityRegistry::with_builtins();
2203 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2204
2205 let applied = apply_capabilities(
2206 base_runtime_agent.clone(),
2207 &["current_time".to_string()],
2208 ®istry,
2209 &test_ctx(),
2210 )
2211 .await;
2212
2213 assert_eq!(
2215 applied.runtime_agent.system_prompt,
2216 base_runtime_agent.system_prompt
2217 );
2218 assert!(applied.tool_registry.has("get_current_time"));
2219 assert_eq!(applied.tool_registry.len(), 1);
2220 assert_eq!(applied.applied_ids, vec!["current_time"]);
2221 }
2222
2223 #[tokio::test]
2224 async fn test_apply_capabilities_skips_coming_soon() {
2225 let registry = CapabilityRegistry::with_builtins();
2226 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2227
2228 let applied = apply_capabilities(
2230 base_runtime_agent.clone(),
2231 &["research".to_string()],
2232 ®istry,
2233 &test_ctx(),
2234 )
2235 .await;
2236
2237 assert_eq!(
2239 applied.runtime_agent.system_prompt,
2240 base_runtime_agent.system_prompt
2241 );
2242 assert!(applied.applied_ids.is_empty()); }
2244
2245 #[tokio::test]
2246 async fn test_apply_capabilities_multiple() {
2247 let registry = CapabilityRegistry::with_builtins();
2248 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2249
2250 let applied = apply_capabilities(
2251 base_runtime_agent.clone(),
2252 &["noop".to_string(), "current_time".to_string()],
2253 ®istry,
2254 &test_ctx(),
2255 )
2256 .await;
2257
2258 assert!(applied.tool_registry.has("get_current_time"));
2259 assert_eq!(applied.applied_ids, vec!["noop", "current_time"]);
2260 }
2261
2262 #[tokio::test]
2263 async fn test_apply_capabilities_preserves_order() {
2264 let registry = CapabilityRegistry::with_builtins();
2265 let base_runtime_agent = RuntimeAgent::new("Base prompt.", "gpt-5.2");
2266
2267 let applied = apply_capabilities(
2269 base_runtime_agent,
2270 &["current_time".to_string(), "noop".to_string()],
2271 ®istry,
2272 &test_ctx(),
2273 )
2274 .await;
2275
2276 assert_eq!(applied.applied_ids, vec!["current_time", "noop"]);
2277 }
2278
2279 #[tokio::test]
2280 async fn test_apply_capabilities_test_math() {
2281 let registry = CapabilityRegistry::with_builtins();
2282 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2283
2284 let applied = apply_capabilities(
2285 base_runtime_agent.clone(),
2286 &["test_math".to_string()],
2287 ®istry,
2288 &test_ctx(),
2289 )
2290 .await;
2291
2292 assert!(
2294 !applied
2295 .runtime_agent
2296 .system_prompt
2297 .contains("<capability id=\"test_math\">")
2298 );
2299 assert!(
2301 applied
2302 .runtime_agent
2303 .system_prompt
2304 .contains("You are a helpful assistant.")
2305 );
2306 assert!(applied.tool_registry.has("add"));
2307 assert!(applied.tool_registry.has("subtract"));
2308 assert!(applied.tool_registry.has("multiply"));
2309 assert!(applied.tool_registry.has("divide"));
2310 assert_eq!(applied.tool_registry.len(), 4);
2311 }
2312
2313 #[tokio::test]
2314 async fn test_apply_capabilities_test_weather() {
2315 let registry = CapabilityRegistry::with_builtins();
2316 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2317
2318 let applied = apply_capabilities(
2319 base_runtime_agent.clone(),
2320 &["test_weather".to_string()],
2321 ®istry,
2322 &test_ctx(),
2323 )
2324 .await;
2325
2326 assert!(
2328 !applied
2329 .runtime_agent
2330 .system_prompt
2331 .contains("<capability id=\"test_weather\">")
2332 );
2333 assert!(applied.tool_registry.has("get_weather"));
2334 assert!(applied.tool_registry.has("get_forecast"));
2335 assert_eq!(applied.tool_registry.len(), 2);
2336 }
2337
2338 #[tokio::test]
2339 async fn test_apply_capabilities_test_math_and_test_weather() {
2340 let registry = CapabilityRegistry::with_builtins();
2341 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2342
2343 let applied = apply_capabilities(
2344 base_runtime_agent.clone(),
2345 &["test_math".to_string(), "test_weather".to_string()],
2346 ®istry,
2347 &test_ctx(),
2348 )
2349 .await;
2350
2351 assert_eq!(applied.tool_registry.len(), 6); assert!(applied.tool_registry.has("add"));
2354 assert!(applied.tool_registry.has("get_weather"));
2355 }
2356
2357 #[tokio::test]
2358 async fn test_apply_capabilities_stateless_todo_list() {
2359 let registry = CapabilityRegistry::with_builtins();
2360 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2361
2362 let applied = apply_capabilities(
2363 base_runtime_agent.clone(),
2364 &["stateless_todo_list".to_string()],
2365 ®istry,
2366 &test_ctx(),
2367 )
2368 .await;
2369
2370 assert!(
2372 applied
2373 .runtime_agent
2374 .system_prompt
2375 .contains("Task Management")
2376 );
2377 assert!(applied.runtime_agent.system_prompt.contains("write_todos"));
2378 assert!(applied.tool_registry.has("write_todos"));
2379 assert_eq!(applied.tool_registry.len(), 1);
2380 }
2381
2382 #[tokio::test]
2383 async fn test_apply_capabilities_web_fetch() {
2384 let registry = CapabilityRegistry::with_builtins();
2385 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2386
2387 let applied = apply_capabilities(
2388 base_runtime_agent.clone(),
2389 &["web_fetch".to_string()],
2390 ®istry,
2391 &test_ctx(),
2392 )
2393 .await;
2394
2395 assert!(
2397 applied
2398 .runtime_agent
2399 .system_prompt
2400 .contains(&base_runtime_agent.system_prompt)
2401 );
2402 assert!(applied.runtime_agent.system_prompt.contains("web_fetch"));
2403 assert!(applied.tool_registry.has("web_fetch"));
2404 assert_eq!(applied.tool_registry.len(), 1);
2405 }
2406
2407 #[tokio::test]
2412 async fn test_xml_tags_wrap_capability_prompts() {
2413 let registry = CapabilityRegistry::with_builtins();
2414 let collected =
2415 collect_capabilities(&["stateless_todo_list".to_string()], ®istry, &test_ctx())
2416 .await;
2417
2418 assert_eq!(collected.system_prompt_parts.len(), 1);
2419 let part = &collected.system_prompt_parts[0];
2420 assert!(part.starts_with("<capability id=\"stateless_todo_list\">"));
2421 assert!(part.ends_with("</capability>"));
2422 assert!(part.contains("Task Management"));
2423 }
2424
2425 #[tokio::test]
2426 async fn test_xml_tags_multiple_capabilities() {
2427 let registry = CapabilityRegistry::with_builtins();
2428 let collected = collect_capabilities(
2429 &[
2430 "stateless_todo_list".to_string(),
2431 "session_schedule".to_string(),
2432 ],
2433 ®istry,
2434 &test_ctx(),
2435 )
2436 .await;
2437
2438 assert_eq!(collected.system_prompt_parts.len(), 2);
2439 assert!(
2440 collected.system_prompt_parts[0].starts_with("<capability id=\"stateless_todo_list\">")
2441 );
2442 assert!(
2443 collected.system_prompt_parts[1].starts_with("<capability id=\"session_schedule\">")
2444 );
2445
2446 let prefix = collected.system_prompt_prefix().unwrap();
2447 assert!(prefix.contains("</capability>\n\n<capability"));
2449 }
2450
2451 #[tokio::test]
2452 async fn test_xml_tags_system_prompt_wrapping() {
2453 let registry = CapabilityRegistry::with_builtins();
2454 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2455
2456 let applied = apply_capabilities(
2457 base,
2458 &["stateless_todo_list".to_string()],
2459 ®istry,
2460 &test_ctx(),
2461 )
2462 .await;
2463
2464 let prompt = &applied.runtime_agent.system_prompt;
2465 assert!(prompt.contains("<capability id=\"stateless_todo_list\">"));
2467 assert!(prompt.contains("</capability>"));
2468 assert!(prompt.contains("<system-prompt>\nYou are helpful.\n</system-prompt>"));
2470 }
2471
2472 #[tokio::test]
2473 async fn test_no_xml_wrapping_without_capabilities() {
2474 let registry = CapabilityRegistry::with_builtins();
2475 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2476
2477 let applied = apply_capabilities(base, &[], ®istry, &test_ctx()).await;
2478
2479 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2481 assert!(
2482 !applied
2483 .runtime_agent
2484 .system_prompt
2485 .contains("<system-prompt>")
2486 );
2487 }
2488
2489 #[tokio::test]
2490 async fn test_no_xml_wrapping_for_noop_capability() {
2491 let registry = CapabilityRegistry::with_builtins();
2492 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2493
2494 let applied = apply_capabilities(base, &["noop".to_string()], ®istry, &test_ctx()).await;
2496
2497 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2498 assert!(
2499 !applied
2500 .runtime_agent
2501 .system_prompt
2502 .contains("<system-prompt>")
2503 );
2504 }
2505
2506 #[tokio::test]
2511 async fn test_collect_capabilities_includes_mounts() {
2512 let registry = CapabilityRegistry::with_builtins();
2513
2514 let collected =
2515 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
2516
2517 assert!(!collected.mounts.is_empty());
2518 assert_eq!(collected.mounts.len(), 1);
2519 assert_eq!(collected.mounts[0].path, "/samples");
2520 assert!(collected.mounts[0].is_readonly());
2521 }
2522
2523 #[tokio::test]
2524 async fn test_collect_capabilities_empty_mounts_by_default() {
2525 let registry = CapabilityRegistry::with_builtins();
2526
2527 let collected =
2529 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
2530
2531 assert!(collected.mounts.is_empty());
2532 }
2533
2534 #[tokio::test]
2535 async fn test_collect_capabilities_combines_mounts() {
2536 let registry = CapabilityRegistry::with_builtins();
2537
2538 let collected = collect_capabilities(
2541 &["sample_data".to_string(), "current_time".to_string()],
2542 ®istry,
2543 &test_ctx(),
2544 )
2545 .await;
2546
2547 assert_eq!(collected.mounts.len(), 1);
2548 assert!(
2550 collected
2551 .applied_ids
2552 .iter()
2553 .any(|id| id == "session_file_system")
2554 );
2555 assert!(collected.applied_ids.iter().any(|id| id == "sample_data"));
2556 assert!(collected.applied_ids.iter().any(|id| id == "current_time"));
2557 }
2558
2559 #[test]
2560 fn test_sample_data_capability() {
2561 let registry = CapabilityRegistry::with_builtins();
2562 let cap = registry.get("sample_data").unwrap();
2563
2564 assert_eq!(cap.id(), "sample_data");
2565 assert_eq!(cap.name(), "Sample Data");
2566 assert_eq!(cap.status(), CapabilityStatus::Available);
2567
2568 assert!(cap.system_prompt_addition().is_some());
2570 assert!(cap.tools().is_empty());
2571
2572 assert!(!cap.mounts().is_empty());
2574 }
2575
2576 #[test]
2581 fn test_resolve_dependencies_empty() {
2582 let registry = CapabilityRegistry::with_builtins();
2583
2584 let resolved = resolve_dependencies(&[], ®istry).unwrap();
2585
2586 assert!(resolved.resolved_ids.is_empty());
2587 assert!(resolved.added_as_dependencies.is_empty());
2588 assert!(resolved.user_selected.is_empty());
2589 }
2590
2591 #[test]
2592 fn test_resolve_dependencies_no_deps() {
2593 let registry = CapabilityRegistry::with_builtins();
2594
2595 let resolved = resolve_dependencies(&["current_time".to_string()], ®istry).unwrap();
2597
2598 assert_eq!(resolved.resolved_ids, vec!["current_time"]);
2599 assert!(resolved.added_as_dependencies.is_empty());
2600 }
2601
2602 #[test]
2603 fn test_resolve_dependencies_with_deps() {
2604 let registry = CapabilityRegistry::with_builtins();
2605
2606 let resolved = resolve_dependencies(&["sample_data".to_string()], ®istry).unwrap();
2608
2609 assert_eq!(resolved.resolved_ids.len(), 2);
2611 let fs_pos = resolved
2612 .resolved_ids
2613 .iter()
2614 .position(|id| id == "session_file_system")
2615 .unwrap();
2616 let sd_pos = resolved
2617 .resolved_ids
2618 .iter()
2619 .position(|id| id == "sample_data")
2620 .unwrap();
2621 assert!(fs_pos < sd_pos, "FileSystem should come before SampleData");
2622
2623 assert_eq!(resolved.added_as_dependencies, vec!["session_file_system"]);
2625 }
2626
2627 #[test]
2628 fn test_resolve_dependencies_already_selected() {
2629 let registry = CapabilityRegistry::with_builtins();
2630
2631 let resolved = resolve_dependencies(
2633 &["session_file_system".to_string(), "sample_data".to_string()],
2634 ®istry,
2635 )
2636 .unwrap();
2637
2638 assert_eq!(resolved.resolved_ids.len(), 2);
2639 assert!(resolved.added_as_dependencies.is_empty());
2641 }
2642
2643 #[test]
2644 fn test_resolve_dependencies_preserves_order() {
2645 let registry = CapabilityRegistry::with_builtins();
2646
2647 let resolved =
2649 resolve_dependencies(&["current_time".to_string(), "noop".to_string()], ®istry)
2650 .unwrap();
2651
2652 assert_eq!(resolved.resolved_ids, vec!["current_time", "noop"]);
2653 }
2654
2655 #[test]
2656 fn test_resolve_dependencies_unknown_capability() {
2657 let registry = CapabilityRegistry::with_builtins();
2658
2659 let resolved =
2661 resolve_dependencies(&["unknown_capability".to_string()], ®istry).unwrap();
2662
2663 assert!(resolved.resolved_ids.is_empty());
2664 }
2665
2666 #[test]
2667 fn test_get_dependencies() {
2668 let registry = CapabilityRegistry::with_builtins();
2669
2670 let deps = get_dependencies("sample_data", ®istry);
2672 assert_eq!(deps, vec!["session_file_system"]);
2673
2674 let deps = get_dependencies("current_time", ®istry);
2676 assert!(deps.is_empty());
2677
2678 let deps = get_dependencies("unknown", ®istry);
2680 assert!(deps.is_empty());
2681 }
2682
2683 #[test]
2684 fn test_sample_data_has_dependency() {
2685 let registry = CapabilityRegistry::with_builtins();
2686 let cap = registry.get("sample_data").unwrap();
2687
2688 let deps = cap.dependencies();
2689 assert_eq!(deps.len(), 1);
2690 assert_eq!(deps[0], "session_file_system");
2691 }
2692
2693 #[test]
2694 fn test_noop_has_no_dependencies() {
2695 let registry = CapabilityRegistry::with_builtins();
2696 let cap = registry.get("noop").unwrap();
2697
2698 assert!(cap.dependencies().is_empty());
2699 }
2700
2701 #[test]
2705 fn test_circular_dependency_error() {
2706 struct CapA;
2708 struct CapB;
2709
2710 impl Capability for CapA {
2711 fn id(&self) -> &str {
2712 "test_cap_a"
2713 }
2714 fn name(&self) -> &str {
2715 "Test A"
2716 }
2717 fn description(&self) -> &str {
2718 "Test capability A"
2719 }
2720 fn dependencies(&self) -> Vec<&'static str> {
2721 vec!["test_cap_b"]
2722 }
2723 }
2724
2725 impl Capability for CapB {
2726 fn id(&self) -> &str {
2727 "test_cap_b"
2728 }
2729 fn name(&self) -> &str {
2730 "Test B"
2731 }
2732 fn description(&self) -> &str {
2733 "Test capability B"
2734 }
2735 fn dependencies(&self) -> Vec<&'static str> {
2736 vec!["test_cap_a"]
2737 }
2738 }
2739
2740 let mut registry = CapabilityRegistry::new();
2741 registry.register(CapA);
2742 registry.register(CapB);
2743
2744 let result = resolve_dependencies(&["test_cap_a".to_string()], ®istry);
2745
2746 assert!(result.is_err());
2747 match result.unwrap_err() {
2748 DependencyError::CircularDependency { capability_id, .. } => {
2749 assert_eq!(capability_id, "test_cap_a");
2750 }
2751 _ => panic!("Expected CircularDependency error"),
2752 }
2753 }
2754
2755 use crate::message_filter::{MessageFilter, MessageFilterProvider, MessageQuery};
2760
2761 struct FilterTestCapability {
2763 priority: i32,
2764 }
2765
2766 impl Capability for FilterTestCapability {
2767 fn id(&self) -> &str {
2768 "filter_test"
2769 }
2770 fn name(&self) -> &str {
2771 "Filter Test"
2772 }
2773 fn description(&self) -> &str {
2774 "Test capability with message filter"
2775 }
2776 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
2777 Some(Arc::new(FilterTestProvider {
2778 priority: self.priority,
2779 }))
2780 }
2781 }
2782
2783 struct FilterTestProvider {
2784 priority: i32,
2785 }
2786
2787 impl MessageFilterProvider for FilterTestProvider {
2788 fn apply_filters(&self, query: &mut MessageQuery, config: &serde_json::Value) {
2789 if let Some(search) = config.get("search").and_then(|v| v.as_str()) {
2791 query
2792 .filters
2793 .push(MessageFilter::Search(search.to_string()));
2794 }
2795 }
2796
2797 fn priority(&self) -> i32 {
2798 self.priority
2799 }
2800 }
2801
2802 #[tokio::test]
2803 async fn test_collect_capabilities_with_configs_no_filter_providers() {
2804 let registry = CapabilityRegistry::with_builtins();
2805 let configs = vec![AgentCapabilityConfig {
2806 capability_ref: CapabilityId::new("current_time"),
2807 config: serde_json::json!({}),
2808 }];
2809
2810 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2811
2812 assert!(collected.message_filter_providers.is_empty());
2813 assert!(!collected.has_message_filters());
2814 }
2815
2816 #[tokio::test]
2817 async fn test_collect_capabilities_with_configs_with_filter_provider() {
2818 let mut registry = CapabilityRegistry::new();
2819 registry.register(FilterTestCapability { priority: 0 });
2820
2821 let configs = vec![AgentCapabilityConfig {
2822 capability_ref: CapabilityId::new("filter_test"),
2823 config: serde_json::json!({ "search": "hello" }),
2824 }];
2825
2826 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2827
2828 assert_eq!(collected.message_filter_providers.len(), 1);
2829 assert!(collected.has_message_filters());
2830 }
2831
2832 #[tokio::test]
2833 async fn test_collect_capabilities_with_configs_filter_priority_order() {
2834 struct HighPriorityCapability;
2836 struct LowPriorityCapability;
2837
2838 impl Capability for HighPriorityCapability {
2839 fn id(&self) -> &str {
2840 "high_priority"
2841 }
2842 fn name(&self) -> &str {
2843 "High Priority"
2844 }
2845 fn description(&self) -> &str {
2846 "Test"
2847 }
2848 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
2849 Some(Arc::new(FilterTestProvider { priority: 10 }))
2850 }
2851 }
2852
2853 impl Capability for LowPriorityCapability {
2854 fn id(&self) -> &str {
2855 "low_priority"
2856 }
2857 fn name(&self) -> &str {
2858 "Low Priority"
2859 }
2860 fn description(&self) -> &str {
2861 "Test"
2862 }
2863 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
2864 Some(Arc::new(FilterTestProvider { priority: -5 }))
2865 }
2866 }
2867
2868 let mut registry = CapabilityRegistry::new();
2869 registry.register(HighPriorityCapability);
2870 registry.register(LowPriorityCapability);
2871
2872 let configs = vec![
2874 AgentCapabilityConfig {
2875 capability_ref: CapabilityId::new("high_priority"),
2876 config: serde_json::json!({}),
2877 },
2878 AgentCapabilityConfig {
2879 capability_ref: CapabilityId::new("low_priority"),
2880 config: serde_json::json!({}),
2881 },
2882 ];
2883
2884 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2885
2886 assert_eq!(collected.message_filter_providers.len(), 2);
2888 assert_eq!(collected.message_filter_providers[0].0.priority(), -5);
2889 assert_eq!(collected.message_filter_providers[1].0.priority(), 10);
2890 }
2891
2892 #[tokio::test]
2893 async fn test_collected_capabilities_apply_message_filters() {
2894 let mut registry = CapabilityRegistry::new();
2895 registry.register(FilterTestCapability { priority: 0 });
2896
2897 let configs = vec![AgentCapabilityConfig {
2898 capability_ref: CapabilityId::new("filter_test"),
2899 config: serde_json::json!({ "search": "test_query" }),
2900 }];
2901
2902 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2903
2904 let session_id: SessionId = Uuid::now_v7().into();
2906 let mut query = MessageQuery::new(session_id);
2907
2908 collected.apply_message_filters(&mut query);
2909
2910 assert_eq!(query.filters.len(), 1);
2912 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
2913 }
2914
2915 #[tokio::test]
2916 async fn test_collected_capabilities_apply_multiple_filters_in_priority_order() {
2917 struct SearchCapability {
2918 id: &'static str,
2919 search_term: &'static str,
2920 priority: i32,
2921 }
2922
2923 struct SearchProvider {
2924 search_term: &'static str,
2925 priority: i32,
2926 }
2927
2928 impl MessageFilterProvider for SearchProvider {
2929 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
2930 query
2931 .filters
2932 .push(MessageFilter::Search(self.search_term.to_string()));
2933 }
2934
2935 fn priority(&self) -> i32 {
2936 self.priority
2937 }
2938 }
2939
2940 impl Capability for SearchCapability {
2941 fn id(&self) -> &str {
2942 self.id
2943 }
2944 fn name(&self) -> &str {
2945 "Search"
2946 }
2947 fn description(&self) -> &str {
2948 "Test"
2949 }
2950 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
2951 Some(Arc::new(SearchProvider {
2952 search_term: self.search_term,
2953 priority: self.priority,
2954 }))
2955 }
2956 }
2957
2958 let mut registry = CapabilityRegistry::new();
2959 registry.register(SearchCapability {
2960 id: "cap_a",
2961 search_term: "alpha",
2962 priority: 5,
2963 });
2964 registry.register(SearchCapability {
2965 id: "cap_b",
2966 search_term: "beta",
2967 priority: 1,
2968 });
2969 registry.register(SearchCapability {
2970 id: "cap_c",
2971 search_term: "gamma",
2972 priority: 10,
2973 });
2974
2975 let configs = vec![
2976 AgentCapabilityConfig {
2977 capability_ref: CapabilityId::new("cap_a"),
2978 config: serde_json::json!({}),
2979 },
2980 AgentCapabilityConfig {
2981 capability_ref: CapabilityId::new("cap_b"),
2982 config: serde_json::json!({}),
2983 },
2984 AgentCapabilityConfig {
2985 capability_ref: CapabilityId::new("cap_c"),
2986 config: serde_json::json!({}),
2987 },
2988 ];
2989
2990 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2991
2992 let session_id: SessionId = Uuid::now_v7().into();
2993 let mut query = MessageQuery::new(session_id);
2994
2995 collected.apply_message_filters(&mut query);
2996
2997 assert_eq!(query.filters.len(), 3);
2999 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3000 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3001 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3002 }
3003
3004 #[test]
3005 fn test_capability_without_message_filter_returns_none() {
3006 let registry = CapabilityRegistry::with_builtins();
3007
3008 let noop = registry.get("noop").unwrap();
3009 assert!(noop.message_filter_provider().is_none());
3010
3011 let current_time = registry.get("current_time").unwrap();
3012 assert!(current_time.message_filter_provider().is_none());
3013 }
3014
3015 #[tokio::test]
3016 async fn test_collect_capabilities_preserves_config_for_filter_provider() {
3017 let mut registry = CapabilityRegistry::new();
3018 registry.register(FilterTestCapability { priority: 0 });
3019
3020 let test_config = serde_json::json!({
3021 "search": "custom_search",
3022 "extra_field": 42
3023 });
3024
3025 let configs = vec![AgentCapabilityConfig {
3026 capability_ref: CapabilityId::new("filter_test"),
3027 config: test_config.clone(),
3028 }];
3029
3030 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3031
3032 assert_eq!(collected.message_filter_providers.len(), 1);
3034 let (_, stored_config) = &collected.message_filter_providers[0];
3035 assert_eq!(*stored_config, test_config);
3036 }
3037
3038 #[test]
3043 fn test_collect_message_filters_only_collects_filters() {
3044 let mut registry = CapabilityRegistry::new();
3045 registry.register(FilterTestCapability { priority: 0 });
3046
3047 let configs = vec![AgentCapabilityConfig {
3048 capability_ref: CapabilityId::new("filter_test"),
3049 config: serde_json::json!({ "search": "test_query" }),
3050 }];
3051
3052 let collected = collect_message_filters_only(&configs, ®istry);
3053
3054 let session_id: SessionId = Uuid::now_v7().into();
3055 let mut query = MessageQuery::new(session_id);
3056 collected.apply_message_filters(&mut query);
3057
3058 assert_eq!(query.filters.len(), 1);
3059 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3060 }
3061
3062 #[test]
3063 fn test_collect_message_filters_only_skips_unknown_capabilities() {
3064 let registry = CapabilityRegistry::new();
3065
3066 let configs = vec![AgentCapabilityConfig {
3067 capability_ref: CapabilityId::new("nonexistent"),
3068 config: serde_json::json!({}),
3069 }];
3070
3071 let collected = collect_message_filters_only(&configs, ®istry);
3072 assert!(collected.message_filter_providers.is_empty());
3073 }
3074
3075 #[test]
3076 fn test_collect_message_filters_only_preserves_priority_order() {
3077 struct PriorityFilterCap {
3078 id: &'static str,
3079 search_term: &'static str,
3080 priority: i32,
3081 }
3082
3083 struct PriorityFilterProvider {
3084 search_term: &'static str,
3085 priority: i32,
3086 }
3087
3088 impl Capability for PriorityFilterCap {
3089 fn id(&self) -> &str {
3090 self.id
3091 }
3092 fn name(&self) -> &str {
3093 self.id
3094 }
3095 fn description(&self) -> &str {
3096 "priority test"
3097 }
3098 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3099 Some(Arc::new(PriorityFilterProvider {
3100 search_term: self.search_term,
3101 priority: self.priority,
3102 }))
3103 }
3104 }
3105
3106 impl MessageFilterProvider for PriorityFilterProvider {
3107 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3108 query
3109 .filters
3110 .push(MessageFilter::Search(self.search_term.to_string()));
3111 }
3112 fn priority(&self) -> i32 {
3113 self.priority
3114 }
3115 }
3116
3117 let mut registry = CapabilityRegistry::new();
3118 registry.register(PriorityFilterCap {
3119 id: "gamma",
3120 search_term: "gamma",
3121 priority: 10,
3122 });
3123 registry.register(PriorityFilterCap {
3124 id: "alpha",
3125 search_term: "alpha",
3126 priority: 5,
3127 });
3128 registry.register(PriorityFilterCap {
3129 id: "beta",
3130 search_term: "beta",
3131 priority: 1,
3132 });
3133
3134 let configs = vec![
3135 AgentCapabilityConfig {
3136 capability_ref: CapabilityId::new("gamma"),
3137 config: serde_json::json!({}),
3138 },
3139 AgentCapabilityConfig {
3140 capability_ref: CapabilityId::new("alpha"),
3141 config: serde_json::json!({}),
3142 },
3143 AgentCapabilityConfig {
3144 capability_ref: CapabilityId::new("beta"),
3145 config: serde_json::json!({}),
3146 },
3147 ];
3148
3149 let collected = collect_message_filters_only(&configs, ®istry);
3150
3151 let session_id: SessionId = Uuid::now_v7().into();
3152 let mut query = MessageQuery::new(session_id);
3153 collected.apply_message_filters(&mut query);
3154
3155 assert_eq!(query.filters.len(), 3);
3157 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3158 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3159 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3160 }
3161
3162 #[test]
3163 fn test_collect_message_filters_only_post_load_invoked() {
3164 use crate::message::Message;
3165
3166 struct PostLoadCap;
3167 struct PostLoadProvider;
3168
3169 impl Capability for PostLoadCap {
3170 fn id(&self) -> &str {
3171 "post_load_test"
3172 }
3173 fn name(&self) -> &str {
3174 "PostLoad Test"
3175 }
3176 fn description(&self) -> &str {
3177 "test"
3178 }
3179 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3180 Some(Arc::new(PostLoadProvider))
3181 }
3182 }
3183
3184 impl MessageFilterProvider for PostLoadProvider {
3185 fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
3186 fn priority(&self) -> i32 {
3187 0
3188 }
3189 fn post_load(&self, messages: &mut Vec<Message>, _config: &serde_json::Value) {
3190 messages.reverse();
3192 }
3193 }
3194
3195 let mut registry = CapabilityRegistry::new();
3196 registry.register(PostLoadCap);
3197
3198 let configs = vec![AgentCapabilityConfig {
3199 capability_ref: CapabilityId::new("post_load_test"),
3200 config: serde_json::json!({}),
3201 }];
3202
3203 let collected = collect_message_filters_only(&configs, ®istry);
3204
3205 let mut messages = vec![Message::user("first"), Message::user("second")];
3206 collected.apply_post_load_filters(&mut messages);
3207
3208 assert_eq!(messages[0].text(), Some("second"));
3210 assert_eq!(messages[1].text(), Some("first"));
3211 }
3212
3213 #[test]
3214 fn test_collect_model_view_providers_respects_compaction_capability_boundary() {
3215 use crate::tool_types::ToolCall;
3216
3217 fn tool_heavy_messages() -> Vec<Message> {
3218 let mut messages = vec![Message::user("inspect files repeatedly")];
3219 for index in 0..9 {
3220 let call_id = format!("call_{index}");
3221 messages.push(Message::assistant_with_tools(
3222 "",
3223 vec![ToolCall {
3224 id: call_id.clone(),
3225 name: "read_file".to_string(),
3226 arguments: serde_json::json!({"path": "/workspace/src/lib.rs"}),
3227 }],
3228 ));
3229 messages.push(Message::tool_result(
3230 call_id,
3231 Some(serde_json::json!({
3232 "path": "/workspace/src/lib.rs",
3233 "content": format!("{}{}", "large file line\n".repeat(1000), index),
3234 "total_lines": 1000,
3235 "lines_shown": {"start": 1, "end": 1000},
3236 "truncated": false
3237 })),
3238 None,
3239 ));
3240 }
3241 messages
3242 }
3243
3244 fn first_tool_result_is_masked(messages: &[Message]) -> bool {
3245 messages[2]
3246 .tool_result_content()
3247 .and_then(|result| result.result.as_ref())
3248 .and_then(|result| result.get("masked"))
3249 .and_then(|masked| masked.as_bool())
3250 .unwrap_or(false)
3251 }
3252
3253 let mut registry = CapabilityRegistry::new();
3254 registry.register(CompactionCapability);
3255 let context = ModelViewContext {
3256 session_id: SessionId::new(),
3257 prior_usage: None,
3258 };
3259
3260 let no_compaction = collect_model_view_providers(&[], ®istry);
3261 let unmasked = no_compaction.apply_model_view(tool_heavy_messages(), &context);
3262 assert!(!first_tool_result_is_masked(&unmasked));
3263
3264 let compaction = collect_model_view_providers(
3265 &[AgentCapabilityConfig {
3266 capability_ref: CapabilityId::new(COMPACTION_CAPABILITY_ID),
3267 config: serde_json::json!({}),
3268 }],
3269 ®istry,
3270 );
3271 let masked = compaction.apply_model_view(tool_heavy_messages(), &context);
3272 assert!(first_tool_result_is_masked(&masked));
3273 let last_tool = masked.last().unwrap().tool_result_content().unwrap();
3274 assert!(last_tool.result.as_ref().unwrap().get("content").is_some());
3275 }
3276
3277 #[tokio::test]
3287 async fn test_virtual_bash_capability_produces_bash_tool() {
3288 let registry = CapabilityRegistry::with_builtins();
3289 let collected =
3290 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3291
3292 let tool_names: Vec<&str> = collected
3293 .tool_definitions
3294 .iter()
3295 .map(|t| t.name())
3296 .collect();
3297 assert!(
3298 tool_names.contains(&"bash"),
3299 "virtual_bash capability must produce 'bash' tool, got: {:?}",
3300 tool_names
3301 );
3302 assert!(
3303 !collected.tools.is_empty(),
3304 "virtual_bash must provide tool implementations"
3305 );
3306 }
3307
3308 #[tokio::test]
3309 async fn test_generic_harness_capability_set_produces_bash_tool() {
3310 let generic_harness_caps = vec![
3313 "session_file_system".to_string(),
3314 "virtual_bash".to_string(),
3315 "web_fetch".to_string(),
3316 "session_storage".to_string(),
3317 "session".to_string(),
3318 "agent_instructions".to_string(),
3319 "skills".to_string(),
3320 "infinity_context".to_string(),
3321 "openai_tool_search".to_string(),
3322 ];
3323
3324 let registry = CapabilityRegistry::with_builtins();
3325 let collected = collect_capabilities(&generic_harness_caps, ®istry, &test_ctx()).await;
3326
3327 let tool_names: Vec<&str> = collected
3328 .tool_definitions
3329 .iter()
3330 .map(|t| t.name())
3331 .collect();
3332 assert!(
3333 tool_names.contains(&"bash"),
3334 "Generic Harness capabilities must produce 'bash' tool, got: {:?}",
3335 tool_names
3336 );
3337 }
3338
3339 #[tokio::test]
3340 async fn test_collect_capabilities_tool_count_matches_definitions() {
3341 let registry = CapabilityRegistry::with_builtins();
3344 let collected =
3345 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3346
3347 assert_eq!(
3348 collected.tools.len(),
3349 collected.tool_definitions.len(),
3350 "tool implementations ({}) must match tool definitions ({})",
3351 collected.tools.len(),
3352 collected.tool_definitions.len(),
3353 );
3354 }
3355
3356 #[tokio::test]
3360 async fn test_collect_capabilities_resolves_dependencies() {
3361 let registry = CapabilityRegistry::with_builtins();
3364 let collected =
3365 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
3366
3367 assert!(
3369 collected
3370 .applied_ids
3371 .iter()
3372 .any(|id| id == "session_file_system"),
3373 "collect_capabilities must apply session_file_system as a dependency; applied_ids: {:?}",
3374 collected.applied_ids
3375 );
3376
3377 let tool_names: Vec<&str> = collected
3378 .tool_definitions
3379 .iter()
3380 .map(|t| t.name())
3381 .collect();
3382
3383 assert!(
3385 tool_names.contains(&"read_file") && tool_names.contains(&"write_file"),
3386 "collect_capabilities must resolve dependencies and include dependency tools, got: {:?}",
3387 tool_names
3388 );
3389
3390 assert_eq!(
3392 collected.tools.len(),
3393 collected.tool_definitions.len(),
3394 "dependency-added tools must have implementations, not just definitions"
3395 );
3396 }
3397
3398 #[test]
3399 fn test_defaults_do_not_include_bash() {
3400 let registry = crate::ToolRegistry::with_defaults();
3403 assert!(
3404 !registry.has("bash"),
3405 "with_defaults() must not include 'bash' — it comes from virtual_bash capability"
3406 );
3407 }
3408
3409 #[tokio::test]
3416 async fn test_background_execution_auto_activates_with_virtual_bash() {
3417 let registry = CapabilityRegistry::with_builtins();
3418 let collected =
3419 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3420
3421 let tool_names: Vec<&str> = collected
3422 .tool_definitions
3423 .iter()
3424 .map(|t| t.name())
3425 .collect();
3426 assert!(
3427 tool_names.contains(&"spawn_background"),
3428 "spawn_background must be auto-activated when virtual_bash (a \
3429 background-capable tool) is in the agent's capability set; got: {:?}",
3430 tool_names
3431 );
3432 assert!(
3433 collected
3434 .applied_ids
3435 .iter()
3436 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
3437 "background_execution must be in applied_ids when auto-activated; \
3438 got: {:?}",
3439 collected.applied_ids
3440 );
3441
3442 assert!(
3444 collected
3445 .tools
3446 .iter()
3447 .any(|t| t.name() == "spawn_background"),
3448 "spawn_background tool implementation must be present alongside the \
3449 definition (lockstep contract)"
3450 );
3451 }
3452
3453 #[tokio::test]
3456 async fn test_background_execution_does_not_auto_activate_without_hint() {
3457 let registry = CapabilityRegistry::with_builtins();
3458 let collected =
3460 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
3461
3462 let tool_names: Vec<&str> = collected
3463 .tool_definitions
3464 .iter()
3465 .map(|t| t.name())
3466 .collect();
3467 assert!(
3468 !tool_names.contains(&"spawn_background"),
3469 "spawn_background must NOT be activated without a background-capable \
3470 tool; got: {:?}",
3471 tool_names
3472 );
3473 assert!(
3474 !collected
3475 .applied_ids
3476 .iter()
3477 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
3478 "background_execution must not appear in applied_ids when no \
3479 background-capable tool is present; got: {:?}",
3480 collected.applied_ids
3481 );
3482 }
3483
3484 #[tokio::test]
3488 async fn test_background_execution_explicit_selection_is_idempotent() {
3489 let registry = CapabilityRegistry::with_builtins();
3490 let collected = collect_capabilities(
3491 &[
3492 "virtual_bash".to_string(),
3493 BACKGROUND_EXECUTION_CAPABILITY_ID.to_string(),
3494 ],
3495 ®istry,
3496 &test_ctx(),
3497 )
3498 .await;
3499
3500 let spawn_background_count = collected
3501 .tool_definitions
3502 .iter()
3503 .filter(|t| t.name() == "spawn_background")
3504 .count();
3505 assert_eq!(
3506 spawn_background_count, 1,
3507 "spawn_background must appear exactly once even when \
3508 background_execution is selected explicitly alongside a \
3509 background-capable tool"
3510 );
3511 let applied_count = collected
3512 .applied_ids
3513 .iter()
3514 .filter(|id| id.as_str() == BACKGROUND_EXECUTION_CAPABILITY_ID)
3515 .count();
3516 assert_eq!(
3517 applied_count, 1,
3518 "background_execution must appear exactly once in applied_ids"
3519 );
3520 }
3521
3522 #[test]
3527 fn test_defaults_do_not_include_spawn_background() {
3528 let registry = crate::ToolRegistry::with_defaults();
3529 assert!(
3530 !registry.has("spawn_background"),
3531 "with_defaults() must not include 'spawn_background' — it comes \
3532 from the background_execution capability (EVE-501)"
3533 );
3534 }
3535
3536 #[test]
3541 fn test_capability_features_default_empty() {
3542 let registry = CapabilityRegistry::with_builtins();
3543
3544 let noop = registry.get("noop").unwrap();
3546 assert!(noop.features().is_empty());
3547
3548 let current_time = registry.get("current_time").unwrap();
3549 assert!(current_time.features().is_empty());
3550 }
3551
3552 #[test]
3553 fn test_file_system_capability_features() {
3554 let registry = CapabilityRegistry::with_builtins();
3555
3556 let fs = registry.get("session_file_system").unwrap();
3557 assert_eq!(fs.features(), vec!["file_system"]);
3558 }
3559
3560 #[test]
3561 fn test_virtual_bash_capability_features() {
3562 let registry = CapabilityRegistry::with_builtins();
3563
3564 let bash = registry.get("virtual_bash").unwrap();
3565 assert_eq!(bash.features(), vec!["file_system"]);
3566 }
3567
3568 #[test]
3569 fn test_session_storage_capability_features() {
3570 let registry = CapabilityRegistry::with_builtins();
3571
3572 let storage = registry.get("session_storage").unwrap();
3573 let features = storage.features();
3574 assert!(features.contains(&"secrets"));
3575 assert!(features.contains(&"key_value"));
3576 }
3577
3578 #[test]
3579 fn test_session_schedule_capability_features() {
3580 let registry = CapabilityRegistry::with_builtins();
3581
3582 let schedule = registry.get("session_schedule").unwrap();
3583 assert_eq!(schedule.features(), vec!["schedules"]);
3584 }
3585
3586 #[test]
3587 fn test_session_sql_database_capability_features() {
3588 let registry = CapabilityRegistry::with_builtins();
3589
3590 let sql = registry.get("session_sql_database").unwrap();
3591 assert_eq!(sql.features(), vec!["sql_database"]);
3592 }
3593
3594 #[test]
3595 fn test_sample_data_capability_features() {
3596 let registry = CapabilityRegistry::with_builtins();
3597
3598 let sample = registry.get("sample_data").unwrap();
3599 assert_eq!(sample.features(), vec!["file_system"]);
3600 }
3601
3602 #[test]
3603 fn test_compute_features_empty() {
3604 let registry = CapabilityRegistry::with_builtins();
3605
3606 let features = compute_features(&[], ®istry);
3607 assert!(features.is_empty());
3608 }
3609
3610 #[test]
3611 fn test_compute_features_single_capability() {
3612 let registry = CapabilityRegistry::with_builtins();
3613
3614 let features = compute_features(&["session_schedule".to_string()], ®istry);
3615 assert_eq!(features, vec!["schedules"]);
3616 }
3617
3618 #[test]
3619 fn test_compute_features_multiple_capabilities() {
3620 let registry = CapabilityRegistry::with_builtins();
3621
3622 let features = compute_features(
3623 &[
3624 "session_file_system".to_string(),
3625 "session_storage".to_string(),
3626 "session_schedule".to_string(),
3627 ],
3628 ®istry,
3629 );
3630 assert!(features.contains(&"file_system".to_string()));
3631 assert!(features.contains(&"secrets".to_string()));
3632 assert!(features.contains(&"key_value".to_string()));
3633 assert!(features.contains(&"schedules".to_string()));
3634 }
3635
3636 #[test]
3637 fn test_compute_features_deduplicates() {
3638 let registry = CapabilityRegistry::with_builtins();
3639
3640 let features = compute_features(
3642 &[
3643 "session_file_system".to_string(),
3644 "virtual_bash".to_string(),
3645 ],
3646 ®istry,
3647 );
3648 let file_system_count = features.iter().filter(|f| *f == "file_system").count();
3649 assert_eq!(file_system_count, 1, "file_system should appear only once");
3650 }
3651
3652 #[test]
3653 fn test_compute_features_includes_dependency_features() {
3654 let registry = CapabilityRegistry::with_builtins();
3655
3656 let features = compute_features(&["virtual_bash".to_string()], ®istry);
3658 assert!(features.contains(&"file_system".to_string()));
3659 }
3660
3661 #[test]
3662 fn test_compute_features_generic_harness_set() {
3663 let registry = CapabilityRegistry::with_builtins();
3664
3665 let features = compute_features(
3667 &[
3668 "session_file_system".to_string(),
3669 "virtual_bash".to_string(),
3670 "session_storage".to_string(),
3671 "session".to_string(),
3672 "session_schedule".to_string(),
3673 ],
3674 ®istry,
3675 );
3676 assert!(features.contains(&"file_system".to_string()));
3677 assert!(features.contains(&"secrets".to_string()));
3678 assert!(features.contains(&"key_value".to_string()));
3679 assert!(features.contains(&"schedules".to_string()));
3680 }
3681
3682 #[test]
3683 fn test_compute_features_unknown_capability_ignored() {
3684 let registry = CapabilityRegistry::with_builtins();
3685
3686 let features = compute_features(
3687 &["unknown_cap".to_string(), "session_schedule".to_string()],
3688 ®istry,
3689 );
3690 assert_eq!(features, vec!["schedules"]);
3691 }
3692
3693 #[test]
3694 fn test_risk_level_ordering() {
3695 assert!(RiskLevel::Low < RiskLevel::Medium);
3696 assert!(RiskLevel::Medium < RiskLevel::High);
3697 }
3698
3699 #[test]
3700 fn test_risk_level_serde_roundtrip() {
3701 let high = RiskLevel::High;
3702 let json = serde_json::to_string(&high).unwrap();
3703 assert_eq!(json, "\"high\"");
3704 let back: RiskLevel = serde_json::from_str(&json).unwrap();
3705 assert_eq!(back, RiskLevel::High);
3706 }
3707
3708 #[test]
3709 fn test_capability_risk_levels() {
3710 let registry = CapabilityRegistry::with_builtins();
3711
3712 let bash = registry.get("virtual_bash").unwrap();
3714 assert_eq!(bash.risk_level(), RiskLevel::High);
3715
3716 let fetch = registry.get("web_fetch").unwrap();
3718 assert_eq!(fetch.risk_level(), RiskLevel::High);
3719
3720 let noop = registry.get("noop").unwrap();
3722 assert_eq!(noop.risk_level(), RiskLevel::Low);
3723 }
3724
3725 #[tokio::test]
3730 async fn test_apply_capabilities_openai_tool_search() {
3731 let registry = CapabilityRegistry::with_builtins();
3732 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
3733
3734 let applied = apply_capabilities(
3735 base_runtime_agent.clone(),
3736 &["openai_tool_search".to_string()],
3737 ®istry,
3738 &test_ctx(),
3739 )
3740 .await;
3741
3742 assert_eq!(
3744 applied.runtime_agent.system_prompt,
3745 base_runtime_agent.system_prompt
3746 );
3747 assert!(applied.tool_registry.is_empty());
3748 assert_eq!(applied.applied_ids, vec!["openai_tool_search"]);
3749
3750 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
3752 assert!(ts.enabled);
3753 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
3754 }
3755
3756 #[tokio::test]
3757 async fn test_apply_capabilities_openai_tool_search_with_other_capabilities() {
3758 let registry = CapabilityRegistry::with_builtins();
3759 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
3760
3761 let applied = apply_capabilities(
3762 base_runtime_agent,
3763 &[
3764 "current_time".to_string(),
3765 "openai_tool_search".to_string(),
3766 "test_math".to_string(),
3767 ],
3768 ®istry,
3769 &test_ctx(),
3770 )
3771 .await;
3772
3773 assert!(applied.tool_registry.has("get_current_time"));
3775 assert!(applied.tool_registry.has("add"));
3776 assert!(applied.tool_registry.has("subtract"));
3777 assert!(applied.tool_registry.has("multiply"));
3778 assert!(applied.tool_registry.has("divide"));
3779
3780 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
3782 assert!(ts.enabled);
3783 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
3784 }
3785
3786 #[tokio::test]
3787 async fn test_collect_capabilities_tool_search_custom_threshold() {
3788 let registry = CapabilityRegistry::with_builtins();
3789
3790 let configs = vec![AgentCapabilityConfig {
3791 capability_ref: CapabilityId::new("openai_tool_search"),
3792 config: serde_json::json!({"threshold": 5}),
3793 }];
3794
3795 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3796
3797 let ts = collected.tool_search.as_ref().unwrap();
3798 assert!(ts.enabled);
3799 assert_eq!(ts.threshold, 5);
3800 }
3801
3802 #[tokio::test]
3803 async fn test_collect_capabilities_no_tool_search_without_capability() {
3804 let registry = CapabilityRegistry::with_builtins();
3805
3806 let configs = vec![AgentCapabilityConfig {
3807 capability_ref: CapabilityId::new("current_time"),
3808 config: serde_json::json!({}),
3809 }];
3810
3811 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3812
3813 assert!(collected.tool_search.is_none());
3814 }
3815
3816 #[tokio::test]
3817 async fn test_collect_capabilities_tool_search_category_propagation() {
3818 let registry = CapabilityRegistry::with_builtins();
3819
3820 let configs = vec![
3822 AgentCapabilityConfig {
3823 capability_ref: CapabilityId::new("test_math"),
3824 config: serde_json::json!({}),
3825 },
3826 AgentCapabilityConfig {
3827 capability_ref: CapabilityId::new("openai_tool_search"),
3828 config: serde_json::json!({}),
3829 },
3830 ];
3831
3832 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3833
3834 assert!(collected.tool_search.is_some());
3836
3837 for tool_def in &collected.tool_definitions {
3839 if ["add", "subtract", "multiply", "divide"].contains(&tool_def.name()) {
3841 assert!(
3842 tool_def.category().is_some(),
3843 "Tool {} should have a category from its capability",
3844 tool_def.name()
3845 );
3846 }
3847 }
3848 }
3849
3850 #[tokio::test]
3851 async fn test_apply_capabilities_prompt_caching() {
3852 let registry = CapabilityRegistry::with_builtins();
3853 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
3854
3855 let applied = apply_capabilities(
3856 base_runtime_agent.clone(),
3857 &["prompt_caching".to_string()],
3858 ®istry,
3859 &test_ctx(),
3860 )
3861 .await;
3862
3863 assert_eq!(
3864 applied.runtime_agent.system_prompt,
3865 base_runtime_agent.system_prompt
3866 );
3867 assert!(applied.tool_registry.is_empty());
3868 assert_eq!(applied.applied_ids, vec!["prompt_caching"]);
3869
3870 let prompt_cache = applied.runtime_agent.prompt_cache.as_ref().unwrap();
3871 assert!(prompt_cache.enabled);
3872 assert_eq!(
3873 prompt_cache.strategy,
3874 crate::llm_driver_registry::PromptCacheStrategy::Auto
3875 );
3876 assert!(prompt_cache.gemini_cached_content.is_none());
3877 }
3878
3879 #[tokio::test]
3880 async fn test_collect_capabilities_prompt_caching_custom_strategy() {
3881 let registry = CapabilityRegistry::with_builtins();
3882
3883 let configs = vec![AgentCapabilityConfig {
3884 capability_ref: CapabilityId::new("prompt_caching"),
3885 config: serde_json::json!({"strategy": "auto"}),
3886 }];
3887
3888 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3889
3890 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
3891 assert!(prompt_cache.enabled);
3892 assert_eq!(
3893 prompt_cache.strategy,
3894 crate::llm_driver_registry::PromptCacheStrategy::Auto
3895 );
3896 assert!(prompt_cache.gemini_cached_content.is_none());
3897 }
3898
3899 #[tokio::test]
3900 async fn test_collect_capabilities_prompt_caching_gemini_cached_content() {
3901 let registry = CapabilityRegistry::with_builtins();
3902
3903 let configs = vec![AgentCapabilityConfig {
3904 capability_ref: CapabilityId::new("prompt_caching"),
3905 config: serde_json::json!({
3906 "strategy": "auto",
3907 "gemini_cached_content": "cachedContents/demo-cache"
3908 }),
3909 }];
3910
3911 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3912
3913 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
3914 assert_eq!(
3915 prompt_cache.gemini_cached_content.as_deref(),
3916 Some("cachedContents/demo-cache")
3917 );
3918 }
3919
3920 struct SkillContributingCapability;
3925
3926 impl Capability for SkillContributingCapability {
3927 fn id(&self) -> &str {
3928 "contributes_skills"
3929 }
3930 fn name(&self) -> &str {
3931 "Contributes Skills"
3932 }
3933 fn description(&self) -> &str {
3934 "Test capability that contributes skills."
3935 }
3936 fn contribute_skills(&self) -> Vec<SkillContribution> {
3937 vec![
3938 SkillContribution::new("alpha-skill", "Alpha skill desc", "# Alpha\nDo alpha.")
3939 .with_files(vec![(
3940 "scripts/a.sh".to_string(),
3941 "#!/bin/sh\necho a\n".to_string(),
3942 )]),
3943 SkillContribution::new("beta-skill", "Beta skill desc", "# Beta\nDo beta.")
3944 .with_user_invocable(false),
3945 ]
3946 }
3947 }
3948
3949 fn skill_md_from_entries(entries: &HashMap<String, MountEntry>) -> &str {
3950 match &entries.get("SKILL.md").expect("SKILL.md missing").source {
3951 MountSource::InlineFile { content, .. } => content.as_str(),
3952 _ => panic!("Expected InlineFile for SKILL.md"),
3953 }
3954 }
3955
3956 #[tokio::test]
3957 async fn test_contribute_skills_normalized_to_mounts() {
3958 let mut registry = CapabilityRegistry::new();
3959 registry.register(SkillContributingCapability);
3960
3961 let configs = vec![AgentCapabilityConfig {
3962 capability_ref: CapabilityId::new("contributes_skills"),
3963 config: serde_json::json!({}),
3964 }];
3965
3966 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3967
3968 let skill_mounts: Vec<_> = collected
3969 .mounts
3970 .iter()
3971 .filter(|m| m.path.starts_with("/.agents/skills/"))
3972 .collect();
3973 assert_eq!(skill_mounts.len(), 2);
3974
3975 for m in &skill_mounts {
3978 assert!(m.is_readonly());
3979 assert_eq!(m.capability_id, "contributes_skills");
3980 }
3981
3982 let alpha = skill_mounts
3983 .iter()
3984 .find(|m| m.path == "/.agents/skills/alpha-skill")
3985 .expect("alpha-skill mount missing");
3986 match &alpha.source {
3987 MountSource::InlineDirectory { entries } => {
3988 assert!(entries.contains_key("SKILL.md"));
3989 assert!(entries.contains_key("scripts/a.sh"));
3990 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
3991 assert_eq!(parsed.name, "alpha-skill");
3992 assert!(parsed.user_invocable);
3993 }
3994 _ => panic!("Expected InlineDirectory"),
3995 }
3996
3997 let beta = skill_mounts
3998 .iter()
3999 .find(|m| m.path == "/.agents/skills/beta-skill")
4000 .expect("beta-skill mount missing");
4001 match &beta.source {
4002 MountSource::InlineDirectory { entries } => {
4003 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
4004 assert!(!parsed.user_invocable);
4005 }
4006 _ => panic!("Expected InlineDirectory"),
4007 }
4008 }
4009
4010 #[tokio::test]
4011 async fn test_contribute_skills_default_empty() {
4012 let mut registry = CapabilityRegistry::new();
4015 registry.register(FilterTestCapability { priority: 0 });
4016
4017 let configs = vec![AgentCapabilityConfig {
4018 capability_ref: CapabilityId::new("filter_test"),
4019 config: serde_json::json!({}),
4020 }];
4021
4022 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4023 assert!(
4024 collected
4025 .mounts
4026 .iter()
4027 .all(|m| !m.path.starts_with("/.agents/skills/"))
4028 );
4029 }
4030}