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;
111pub mod persistent_memory;
112mod platform_management;
113mod prompt_caching;
114mod prompt_canary_guardrail;
115mod research;
116mod sample_data;
117mod self_budget;
118mod session;
119mod session_sandbox;
120mod session_schedule;
121mod session_sql_database;
122mod session_storage;
123mod skills;
124mod stateless_todo_list;
125mod subagents;
126mod system_commands;
127mod test_math;
128mod test_weather;
129mod tool_output_persistence;
130mod tool_search;
131pub mod user_hooks;
132mod virtual_bash;
133mod web_fetch;
134mod workspace_volumes;
135
136pub use a2a_delegation::{
138 A2A_AGENT_DELEGATION_CAPABILITY_ID, A2aAgentDelegationCapability, CancelAgentTool,
139 GetAgentRunsTool, MessageAgentTool, SpawnAgentTool, WaitAgentTool,
140};
141#[cfg(feature = "ui-capabilities")]
142pub use a2ui::{A2UI_CAPABILITY_ID, A2UiCapability};
143pub use agent_handoff::{
144 AGENT_HANDOFF_CAPABILITY_ID, AgentHandoffCapability, GetAgentHandoffsTool,
145 MessageAgentHandoffTool, StartAgentHandoffTool,
146};
147pub use agent_instructions::{
148 AGENT_INSTRUCTIONS_CAPABILITY_ID, AGENTS_MD_PATH, AgentInstructionsCapability,
149 AgentInstructionsConfig, DEFAULT_AGENT_INSTRUCTIONS_FILE, MAX_AGENT_INSTRUCTIONS_FILES,
150 MAX_AGENTS_MD_SIZE, format_agents_md_content, format_instruction_file_content,
151};
152pub use attach_skill::{
153 AttachSkillCapability, SKILL_CAPABILITY_PREFIX, SKILLS_DISCOVERY_PATH, SkillContribution,
154 SkillInstructions, SkillMeta, SkillSource, discover_skills_from_entries, is_skill_capability,
155 parse_skill_capability_id, reconstruct_skill_md, skill_capability_id,
156};
157pub use background_execution::{BACKGROUND_EXECUTION_CAPABILITY_ID, BackgroundExecutionCapability};
158pub use btw::{BTW_CAPABILITY_ID, BtwCapability};
159pub use budgeting::BudgetingCapability;
160pub use compaction::{
161 COMPACTION_CAPABILITY_ID, CompactionCapability, CompactionConfig, CompactionStep,
162 CompactionStrategy, CostControlConfig, CostControlMaskingResult, HierarchicalMemoryConfig,
163 MaskingSummaryFormat, MemoryTier, ObservationMaskingConfig, ObservationMaskingResult,
164 SessionCompactionMetrics, SummarizationConfig, aggressive_trim, apply_cost_control_masking,
165 apply_hierarchical_memory, apply_observation_masking, build_model_view_messages,
166 build_summarization_prompt, build_summary_message, classify_memory_tiers, estimate_tokens,
167 estimate_total_tokens, format_messages_for_summarization, should_compact_proactively,
168};
169pub use current_time::{CurrentTimeCapability, GetCurrentTimeTool};
170pub use data_knowledge::DataKnowledgeCapability;
171pub use declarative::{
172 DECLARATIVE_CAPABILITY_PREFIX, DeclarativeCapabilityDefinition, DeclarativeCapabilityFile,
173 DeclarativeCapabilitySkill, declarative_capability_id, declarative_capability_info,
174 hydrate_declarative_capability_config, is_declarative_capability,
175 parse_declarative_capability_id, validate_declarative_capability_definition,
176};
177pub use fake_aws::{
178 AwsCreateEc2InstanceTool, AwsCreateIamUserTool, AwsCreateRdsDatabaseTool,
179 AwsCreateS3BucketTool, AwsGetCloudWatchMetricsTool, AwsListEc2InstancesTool,
180 AwsListIamUsersTool, AwsListRdsDatabasesTool, AwsListS3BucketsTool, AwsListSecurityGroupsTool,
181 AwsStopEc2InstanceTool, FakeAwsCapability,
182};
183pub use fake_crm::{
184 CrmAddInteractionTool, CrmCreateCustomerTool, CrmCreateTicketTool, CrmGetCustomerTool,
185 CrmListCustomersTool, CrmListTicketsTool, CrmSearchCustomersTool, CrmUpdateTicketTool,
186 FakeCrmCapability,
187};
188pub use fake_financial::{
189 FakeFinancialCapability, FinanceCreateBudgetTool, FinanceCreateTransactionTool,
190 FinanceForecastCashFlowTool, FinanceGetBalanceTool, FinanceGetExpenseReportTool,
191 FinanceGetRevenueReportTool, FinanceListBudgetsTool, FinanceListTransactionsTool,
192};
193pub use fake_warehouse::{
194 FakeWarehouseCapability, WarehouseCreateInvoiceTool, WarehouseCreateOrderTool,
195 WarehouseCreateShipmentTool, WarehouseGetInventoryTool, WarehouseInventoryReportTool,
196 WarehouseListOrdersTool, WarehouseListShipmentsTool, WarehouseProcessReturnTool,
197 WarehouseUpdateInventoryTool, WarehouseUpdateShipmentStatusTool,
198};
199pub use file_system::{
200 DeleteFileTool, EditFileTool, FileSystemCapability, GrepFilesTool, ListDirectoryTool,
201 ReadFileTool, StatFileTool, WriteFileTool,
202};
203pub use human_intent::{HUMAN_INTENT_CAPABILITY_ID, HumanIntentCapability};
204pub use infinity_context::{
205 INFINITY_CONTEXT_CAPABILITY_ID, InfinityContextCapability, QueryHistoryTool,
206};
207pub use knowledge_base::{
208 KNOWLEDGE_BASE_CAPABILITY_ID, KnowledgeBaseCapability, KnowledgeBaseConfig,
209 validate_knowledge_base_config,
210};
211pub use loop_detection::LoopDetectionCapability;
212pub use mcp::{
213 MCP_CAPABILITY_PREFIX, McpCapability, is_mcp_capability, mcp_capability_id,
214 parse_mcp_capability_id,
215};
216pub use noop::NoopCapability;
217pub use openai_tool_search::{
218 DEFAULT_TOOL_SEARCH_THRESHOLD, OPENAI_TOOL_SEARCH_CAPABILITY_ID, OpenAiToolSearchCapability,
219};
220#[cfg(feature = "ui-capabilities")]
221pub use openui::{OPENUI_CAPABILITY_ID, OpenUiCapability};
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 tool_search::{
261 GenericToolSearchCapability, TOOL_SEARCH_CAPABILITY_ID, TOOL_SEARCH_TOOL_NAME, ToolSearchTool,
262};
263pub use user_hooks::UserHooksCapability;
264pub use virtual_bash::{BashTool, SessionFileSystemAdapter, VirtualBashCapability};
265pub use web_fetch::{
266 BotAuthPublicKey, WebFetchCapability, WebFetchTool, derive_bot_auth_public_key,
267};
268pub use workspace_volumes::{WORKSPACE_VOLUMES_CAPABILITY_ID, WorkspaceVolumesCapability};
269
270pub struct SystemPromptContext {
280 pub session_id: SessionId,
282 pub locale: Option<String>,
284 pub file_store: Option<Arc<dyn SessionFileSystem>>,
286}
287
288impl SystemPromptContext {
289 pub fn without_file_store(session_id: SessionId) -> Self {
291 Self {
292 session_id,
293 locale: None,
294 file_store: None,
295 }
296 }
297}
298
299#[async_trait]
346pub trait Capability: Send + Sync {
347 fn id(&self) -> &str;
349
350 fn name(&self) -> &str;
352
353 fn description(&self) -> &str;
355
356 fn status(&self) -> CapabilityStatus {
358 CapabilityStatus::Available
359 }
360
361 fn icon(&self) -> Option<&str> {
363 None
364 }
365
366 fn category(&self) -> Option<&str> {
368 None
369 }
370
371 fn system_prompt_addition(&self) -> Option<&str> {
391 None
392 }
393
394 async fn system_prompt_contribution(&self, _ctx: &SystemPromptContext) -> Option<String> {
406 self.system_prompt_addition().map(|addition| {
407 format!(
408 "<capability id=\"{}\">\n{}\n</capability>",
409 self.id(),
410 addition
411 )
412 })
413 }
414
415 fn system_prompt_preview(&self) -> Option<String> {
421 self.system_prompt_addition().map(|s| s.to_string())
422 }
423
424 fn tools(&self) -> Vec<Box<dyn Tool>> {
426 vec![]
427 }
428
429 fn tools_with_config(&self, _config: &serde_json::Value) -> Vec<Box<dyn Tool>> {
437 self.tools()
438 }
439
440 async fn system_prompt_contribution_with_config(
447 &self,
448 ctx: &SystemPromptContext,
449 _config: &serde_json::Value,
450 ) -> Option<String> {
451 self.system_prompt_contribution(ctx).await
452 }
453
454 fn tool_definitions(&self) -> Vec<ToolDefinition> {
457 self.tools().iter().map(|t| t.to_definition()).collect()
458 }
459
460 fn mounts(&self) -> Vec<MountPoint> {
468 vec![]
469 }
470
471 fn dependencies(&self) -> Vec<&'static str> {
480 vec![]
481 }
482
483 fn features(&self) -> Vec<&'static str> {
498 vec![]
499 }
500
501 fn config_schema(&self) -> Option<serde_json::Value> {
507 None
508 }
509
510 fn config_ui_schema(&self) -> Option<serde_json::Value> {
515 None
516 }
517
518 fn validate_config(&self, _config: &serde_json::Value) -> Result<(), String> {
524 Ok(())
525 }
526
527 fn mcp_servers(&self) -> ScopedMcpServers {
533 ScopedMcpServers::default()
534 }
535
536 fn mcp_servers_with_config(&self, _config: &serde_json::Value) -> ScopedMcpServers {
538 self.mcp_servers()
539 }
540
541 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
554 None
555 }
556
557 fn model_view_provider(&self) -> Option<Arc<dyn ModelViewProvider>> {
565 None
566 }
567
568 fn pre_tool_use_hooks(&self) -> Vec<Arc<dyn crate::atoms::PreToolUseHook>> {
579 vec![]
580 }
581
582 fn post_tool_exec_hooks(&self) -> Vec<Arc<dyn crate::atoms::PostToolExecHook>> {
590 vec![]
591 }
592
593 fn tool_definition_hooks(&self) -> Vec<Arc<dyn ToolDefinitionHook>> {
602 vec![]
603 }
604
605 fn tool_call_hooks(&self) -> Vec<Arc<dyn ToolCallHook>> {
613 vec![]
614 }
615
616 fn user_hooks(&self) -> Vec<crate::user_hook_types::UserHookSpec> {
632 vec![]
633 }
634
635 fn user_hooks_with_config(
641 &self,
642 _config: &serde_json::Value,
643 ) -> Vec<crate::user_hook_types::UserHookSpec> {
644 self.user_hooks()
645 }
646
647 fn risk_level(&self) -> RiskLevel {
655 RiskLevel::Low
656 }
657
658 fn commands(&self) -> Vec<CommandDescriptor> {
666 vec![]
667 }
668
669 async fn execute_command(
683 &self,
684 request: &ExecuteCommandRequest,
685 _ctx: &CommandExecutionContext,
686 ) -> crate::error::Result<CommandResult> {
687 Err(crate::error::AgentLoopError::config(format!(
688 "capability {} declared command /{} but does not implement execute_command",
689 self.id(),
690 request.name,
691 )))
692 }
693
694 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
702 vec![]
703 }
704
705 fn contribute_skills(&self) -> Vec<SkillContribution> {
715 vec![]
716 }
717
718 fn output_guardrails(&self) -> Vec<Arc<dyn crate::output_guardrail::OutputGuardrail>> {
729 vec![]
730 }
731}
732
733pub trait ToolDefinitionHook: Send + Sync {
734 fn transform(&self, tools: Vec<ToolDefinition>) -> Vec<ToolDefinition>;
735
736 fn applies_with_native_tool_search(&self) -> bool {
741 true
742 }
743}
744
745pub trait ToolCallHook: Send + Sync {
746 fn narration(
747 &self,
748 _tool_def: Option<&ToolDefinition>,
749 _tool_call: &ToolCall,
750 _phase: crate::tool_narration::ToolNarrationPhase,
751 _locale: Option<&str>,
752 ) -> Option<String> {
753 None
754 }
755
756 fn transform_for_execution(&self, tool_call: ToolCall) -> ToolCall {
757 tool_call
758 }
759}
760
761#[derive(
765 Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
766)]
767#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
768#[cfg_attr(feature = "openapi", schema(example = "low"))]
769#[serde(rename_all = "lowercase")]
770pub enum RiskLevel {
771 Low,
773 Medium,
775 High,
777}
778
779#[derive(Debug, Clone, Serialize, Deserialize)]
785#[serde(rename_all = "snake_case")]
786pub enum BlueprintModel {
787 Fixed(String),
789 Default(String),
791 Inherit,
793}
794
795pub struct AgentBlueprint {
801 pub id: &'static str,
803 pub name: &'static str,
805 pub description: &'static str,
807 pub model: BlueprintModel,
809 pub system_prompt: &'static str,
811 pub tools: Vec<Box<dyn Tool>>,
813 pub max_turns: Option<usize>,
815 pub config_schema: Option<serde_json::Value>,
817}
818
819impl AgentBlueprint {
820 pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
822 self.tools.iter().map(|t| t.to_definition()).collect()
823 }
824}
825
826impl std::fmt::Debug for AgentBlueprint {
827 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
828 f.debug_struct("AgentBlueprint")
829 .field("id", &self.id)
830 .field("name", &self.name)
831 .field("model", &self.model)
832 .field("tool_count", &self.tools.len())
833 .field("max_turns", &self.max_turns)
834 .finish()
835 }
836}
837
838#[derive(Clone)]
865pub struct CapabilityRegistry {
866 capabilities: HashMap<String, Arc<dyn Capability>>,
867}
868
869impl CapabilityRegistry {
870 pub fn new() -> Self {
872 Self {
873 capabilities: HashMap::new(),
874 }
875 }
876
877 pub fn with_builtins() -> Self {
882 Self::with_builtins_for_grade(DeploymentGrade::from_env())
883 }
884
885 pub fn with_builtins_for_grade(grade: DeploymentGrade) -> Self {
890 let mut registry = Self::new();
891
892 registry.register(AgentInstructionsCapability);
894 registry.register(HumanIntentCapability);
895 registry.register(NoopCapability);
896 registry.register(CurrentTimeCapability);
897 registry.register(ResearchCapability);
898 registry.register(PlatformManagementCapability);
899 registry.register(FileSystemCapability);
900 registry.register(WorkspaceVolumesCapability);
901 registry.register(SessionStorageCapability);
902 registry.register(SessionCapability);
903 registry.register(SessionSqlDatabaseCapability);
904 registry.register(TestMathCapability);
905 registry.register(TestWeatherCapability);
906 registry.register(StatelessTodoListCapability);
907 registry.register(WebFetchCapability::from_env());
908 registry.register(VirtualBashCapability);
909 registry.register(BackgroundExecutionCapability);
910 registry.register(SessionScheduleCapability);
911 registry.register(BtwCapability);
912 registry.register(InfinityContextCapability);
913 registry.register(budgeting::BudgetingCapability);
914 registry.register(SelfBudgetCapability);
915 registry.register(CompactionCapability);
916 registry.register(MemoryCapability);
917
918 registry.register(OpenAiToolSearchCapability::new());
920 registry.register(GenericToolSearchCapability::new());
922 registry.register(PromptCachingCapability::new());
923
924 registry.register(SkillsCapability);
926
927 registry.register(SubagentCapability);
929
930 if crate::FeatureFlags::from_env(&grade).agent_delegation {
934 registry.register(AgentHandoffCapability);
935 registry.register(A2aAgentDelegationCapability);
936 }
937
938 registry.register(SystemCommandsCapability);
940
941 registry.register(tool_output_persistence::ToolOutputPersistenceCapability);
943
944 registry.register(user_hooks::UserHooksCapability);
947
948 registry.register(LoopDetectionCapability);
950
951 registry.register(PromptCanaryGuardrailCapability);
954
955 #[cfg(feature = "ui-capabilities")]
957 {
958 registry.register(OpenUiCapability);
959 registry.register(A2UiCapability);
960 }
961
962 registry.register(SampleDataCapability);
964
965 registry.register(DataKnowledgeCapability);
967
968 registry.register(KnowledgeBaseCapability);
970
971 registry.register(FakeWarehouseCapability);
973 registry.register(FakeAwsCapability);
974 registry.register(FakeCrmCapability);
975 registry.register(FakeFinancialCapability);
976
977 let internal_flags = crate::InternalFeatureFlags::from_env();
979 if internal_flags.session_sandbox {
980 registry.register(SessionSandboxCapability);
981 }
982 for plugin in inventory::iter::<IntegrationPlugin>() {
983 if (!plugin.experimental_only || grade.experimental_features_enabled())
984 && plugin
985 .feature_flag
986 .is_none_or(|f| internal_flags.is_enabled(f))
987 {
988 registry.register_boxed((plugin.factory)());
989 }
990 }
991
992 registry
993 }
994
995 pub fn register(&mut self, capability: impl Capability + 'static) {
997 self.capabilities
998 .insert(capability.id().to_string(), Arc::new(capability));
999 }
1000
1001 pub fn register_boxed(&mut self, capability: Box<dyn Capability>) {
1003 self.capabilities
1004 .insert(capability.id().to_string(), Arc::from(capability));
1005 }
1006
1007 pub fn register_arc(&mut self, capability: Arc<dyn Capability>) {
1009 self.capabilities
1010 .insert(capability.id().to_string(), capability);
1011 }
1012
1013 pub fn get(&self, id: &str) -> Option<&Arc<dyn Capability>> {
1015 self.capabilities.get(id)
1016 }
1017
1018 pub fn unregister(&mut self, id: &str) -> Option<Arc<dyn Capability>> {
1020 self.capabilities.remove(id)
1021 }
1022
1023 pub fn has(&self, id: &str) -> bool {
1025 self.capabilities.contains_key(id)
1026 }
1027
1028 pub fn list(&self) -> Vec<&Arc<dyn Capability>> {
1030 self.capabilities.values().collect()
1031 }
1032
1033 pub fn len(&self) -> usize {
1035 self.capabilities.len()
1036 }
1037
1038 pub fn is_empty(&self) -> bool {
1040 self.capabilities.is_empty()
1041 }
1042
1043 pub fn builder() -> CapabilityRegistryBuilder {
1045 CapabilityRegistryBuilder::new()
1046 }
1047
1048 pub fn blueprint(&self, id: &str) -> Option<AgentBlueprint> {
1052 for cap in self.capabilities.values() {
1053 for bp in cap.agent_blueprints() {
1054 if bp.id == id {
1055 return Some(bp);
1056 }
1057 }
1058 }
1059 None
1060 }
1061
1062 pub fn blueprint_with_capability(&self, id: &str) -> Option<(String, AgentBlueprint)> {
1066 for (capability_id, cap) in &self.capabilities {
1067 for bp in cap.agent_blueprints() {
1068 if bp.id == id {
1069 return Some((capability_id.clone(), bp));
1070 }
1071 }
1072 }
1073 None
1074 }
1075
1076 pub fn all_blueprints(&self) -> Vec<AgentBlueprint> {
1078 self.capabilities
1079 .values()
1080 .flat_map(|cap| cap.agent_blueprints())
1081 .collect()
1082 }
1083}
1084
1085impl Default for CapabilityRegistry {
1086 fn default() -> Self {
1087 Self::with_builtins()
1088 }
1089}
1090
1091impl std::fmt::Debug for CapabilityRegistry {
1092 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1093 let ids: Vec<_> = self.capabilities.keys().collect();
1094 f.debug_struct("CapabilityRegistry")
1095 .field("capabilities", &ids)
1096 .finish()
1097 }
1098}
1099
1100pub struct CapabilityRegistryBuilder {
1102 registry: CapabilityRegistry,
1103}
1104
1105impl CapabilityRegistryBuilder {
1106 pub fn new() -> Self {
1108 Self {
1109 registry: CapabilityRegistry::new(),
1110 }
1111 }
1112
1113 pub fn with_builtins() -> Self {
1115 Self {
1116 registry: CapabilityRegistry::with_builtins(),
1117 }
1118 }
1119
1120 pub fn capability(mut self, capability: impl Capability + 'static) -> Self {
1122 self.registry.register(capability);
1123 self
1124 }
1125
1126 pub fn build(self) -> CapabilityRegistry {
1128 self.registry
1129 }
1130}
1131
1132impl Default for CapabilityRegistryBuilder {
1133 fn default() -> Self {
1134 Self::new()
1135 }
1136}
1137
1138pub struct ModelViewContext<'a> {
1144 pub session_id: SessionId,
1145 pub prior_usage: Option<&'a TokenUsage>,
1146}
1147
1148pub trait ModelViewProvider: Send + Sync {
1154 fn apply_model_view(
1155 &self,
1156 messages: Vec<Message>,
1157 config: &serde_json::Value,
1158 context: &ModelViewContext<'_>,
1159 ) -> Vec<Message>;
1160
1161 fn priority(&self) -> i32 {
1162 0
1163 }
1164}
1165
1166pub struct CollectedCapabilities {
1171 pub system_prompt_parts: Vec<String>,
1173 pub system_prompt_attributions: Vec<SystemPromptAttribution>,
1175 pub tools: Vec<Box<dyn Tool>>,
1177 pub tool_definitions: Vec<ToolDefinition>,
1179 pub mounts: Vec<MountPoint>,
1181 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1183 pub applied_ids: Vec<String>,
1185 pub tool_search: Option<crate::llm_driver_registry::ToolSearchConfig>,
1187 pub prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig>,
1189 pub tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>>,
1191 pub tool_call_hooks: Vec<Arc<dyn ToolCallHook>>,
1193 pub mcp_servers: ScopedMcpServers,
1195 }
1201
1202#[derive(Debug, Clone, PartialEq, Eq)]
1203pub struct SystemPromptAttribution {
1204 pub capability_id: String,
1205 pub content: String,
1206}
1207
1208impl CollectedCapabilities {
1209 pub fn system_prompt_prefix(&self) -> Option<String> {
1212 if self.system_prompt_parts.is_empty() {
1213 None
1214 } else {
1215 Some(self.system_prompt_parts.join("\n\n"))
1216 }
1217 }
1218
1219 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1223 for (provider, config) in &self.message_filter_providers {
1225 provider.apply_filters(query, config);
1226 }
1227 }
1228
1229 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1232 for (provider, config) in &self.message_filter_providers {
1233 provider.post_load(messages, config);
1234 }
1235 }
1236
1237 pub fn has_message_filters(&self) -> bool {
1239 !self.message_filter_providers.is_empty()
1240 }
1241}
1242
1243pub struct CollectedMessageFilters {
1250 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1252}
1253
1254pub struct CollectedModelViewProviders {
1256 pub model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)>,
1258}
1259
1260impl CollectedMessageFilters {
1266 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1268 for (provider, config) in &self.message_filter_providers {
1269 provider.apply_filters(query, config);
1270 }
1271 }
1272
1273 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1275 for (provider, config) in &self.message_filter_providers {
1276 provider.post_load(messages, config);
1277 }
1278 }
1279}
1280
1281impl CollectedModelViewProviders {
1282 pub fn apply_model_view(
1284 &self,
1285 mut messages: Vec<Message>,
1286 context: &ModelViewContext<'_>,
1287 ) -> Vec<Message> {
1288 for (provider, config) in &self.model_view_providers {
1289 messages = provider.apply_model_view(messages, config, context);
1290 }
1291 messages
1292 }
1293}
1294
1295pub fn collect_message_filters_only(
1301 capability_configs: &[AgentCapabilityConfig],
1302 registry: &CapabilityRegistry,
1303) -> CollectedMessageFilters {
1304 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1305 Vec::new();
1306
1307 for cap_config in capability_configs {
1308 let cap_id = cap_config.capability_ref.as_str();
1309 if let Some(capability) = registry.get(cap_id) {
1310 if capability.status() != CapabilityStatus::Available {
1311 continue;
1312 }
1313 if let Some(provider) = capability.message_filter_provider() {
1314 message_filter_providers.push((provider, cap_config.config.clone()));
1315 }
1316 }
1317 }
1318
1319 message_filter_providers.sort_by_key(|(p, _)| p.priority());
1320
1321 CollectedMessageFilters {
1322 message_filter_providers,
1323 }
1324}
1325
1326pub fn collect_model_view_providers(
1328 capability_configs: &[AgentCapabilityConfig],
1329 registry: &CapabilityRegistry,
1330) -> CollectedModelViewProviders {
1331 let mut model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)> = Vec::new();
1332
1333 for cap_config in capability_configs {
1334 let cap_id = cap_config.capability_ref.as_str();
1335 if let Some(capability) = registry.get(cap_id) {
1336 if capability.status() != CapabilityStatus::Available {
1337 continue;
1338 }
1339 if let Some(provider) = capability.model_view_provider() {
1340 model_view_providers.push((provider, cap_config.config.clone()));
1341 }
1342 }
1343 }
1344
1345 model_view_providers.sort_by_key(|(p, _)| p.priority());
1346
1347 CollectedModelViewProviders {
1348 model_view_providers,
1349 }
1350}
1351
1352pub fn collect_capability_mcp_servers(
1353 capability_configs: &[AgentCapabilityConfig],
1354 registry: &CapabilityRegistry,
1355) -> ScopedMcpServers {
1356 let mut servers = ScopedMcpServers::default();
1357
1358 for cap_config in capability_configs {
1359 let cap_id = cap_config.capability_ref.as_str();
1360 if is_declarative_capability(cap_id) {
1361 if let Ok(definition) =
1362 serde_json::from_value::<DeclarativeCapabilityDefinition>(cap_config.config.clone())
1363 {
1364 if definition.status != CapabilityStatus::Available {
1365 continue;
1366 }
1367 if let Some(contributed) = definition.mcp_servers {
1368 servers = merge_scoped_mcp_servers(&servers, &contributed);
1369 }
1370 }
1371 continue;
1372 }
1373 if let Some(capability) = registry.get(cap_id) {
1374 if capability.status() != CapabilityStatus::Available {
1375 continue;
1376 }
1377 servers = merge_scoped_mcp_servers(
1378 &servers,
1379 &capability.mcp_servers_with_config(&cap_config.config),
1380 );
1381 }
1382 }
1383
1384 servers
1385}
1386
1387pub const MAX_RESOLVED_CAPABILITIES: usize = 100;
1394
1395#[derive(Debug, Clone, PartialEq, Eq)]
1397pub enum DependencyError {
1398 CircularDependency {
1400 capability_id: String,
1402 chain: Vec<String>,
1404 },
1405 TooManyCapabilities {
1407 count: usize,
1409 max: usize,
1411 },
1412}
1413
1414impl std::fmt::Display for DependencyError {
1415 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1416 match self {
1417 DependencyError::CircularDependency {
1418 capability_id,
1419 chain,
1420 } => {
1421 write!(
1422 f,
1423 "Circular dependency detected: {} depends on itself via chain: {} -> {}",
1424 capability_id,
1425 chain.join(" -> "),
1426 capability_id
1427 )
1428 }
1429 DependencyError::TooManyCapabilities { count, max } => {
1430 write!(
1431 f,
1432 "Too many capabilities after resolution: {} (max: {})",
1433 count, max
1434 )
1435 }
1436 }
1437 }
1438}
1439
1440impl std::error::Error for DependencyError {}
1441
1442#[derive(Debug, Clone)]
1444pub struct ResolvedCapabilities {
1445 pub resolved_ids: Vec<String>,
1448 pub added_as_dependencies: Vec<String>,
1450 pub user_selected: Vec<String>,
1452}
1453
1454pub fn resolve_dependencies(
1474 selected_ids: &[String],
1475 registry: &CapabilityRegistry,
1476) -> Result<ResolvedCapabilities, DependencyError> {
1477 use std::collections::HashSet;
1478
1479 let user_selected: HashSet<String> = selected_ids.iter().cloned().collect();
1480 let mut resolved: Vec<String> = Vec::new();
1481 let mut resolved_set: HashSet<String> = HashSet::new();
1482 let mut added_as_dependencies: Vec<String> = Vec::new();
1483
1484 for cap_id in selected_ids {
1486 resolve_single_capability(
1487 cap_id,
1488 registry,
1489 &mut resolved,
1490 &mut resolved_set,
1491 &mut added_as_dependencies,
1492 &user_selected,
1493 &mut Vec::new(), )?;
1495 }
1496
1497 if resolved.len() > MAX_RESOLVED_CAPABILITIES {
1499 return Err(DependencyError::TooManyCapabilities {
1500 count: resolved.len(),
1501 max: MAX_RESOLVED_CAPABILITIES,
1502 });
1503 }
1504
1505 Ok(ResolvedCapabilities {
1506 resolved_ids: resolved,
1507 added_as_dependencies,
1508 user_selected: selected_ids.to_vec(),
1509 })
1510}
1511
1512pub fn resolve_capability_configs(
1517 selected_configs: &[AgentCapabilityConfig],
1518 registry: &CapabilityRegistry,
1519) -> Result<Vec<AgentCapabilityConfig>, DependencyError> {
1520 let mut selected_ids: Vec<String> = Vec::new();
1521 for config in selected_configs {
1522 if is_declarative_capability(config.capability_id())
1523 && let Ok(definition) =
1524 serde_json::from_value::<DeclarativeCapabilityDefinition>(config.config.clone())
1525 {
1526 selected_ids.extend(definition.dependencies);
1527 }
1528 selected_ids.push(config.capability_id().to_string());
1529 }
1530 let resolved = resolve_dependencies(&selected_ids, registry)?;
1531
1532 let explicit_configs: std::collections::HashMap<String, serde_json::Value> = selected_configs
1533 .iter()
1534 .map(|config| (config.capability_id().to_string(), config.config.clone()))
1535 .collect();
1536
1537 Ok(resolved
1538 .resolved_ids
1539 .into_iter()
1540 .map(|capability_id| {
1541 explicit_configs
1542 .get(&capability_id)
1543 .cloned()
1544 .map(|config| AgentCapabilityConfig::with_config(capability_id.clone(), config))
1545 .unwrap_or_else(|| AgentCapabilityConfig::new(capability_id))
1546 })
1547 .collect())
1548}
1549
1550fn resolve_single_capability(
1552 cap_id: &str,
1553 registry: &CapabilityRegistry,
1554 resolved: &mut Vec<String>,
1555 resolved_set: &mut std::collections::HashSet<String>,
1556 added_as_dependencies: &mut Vec<String>,
1557 user_selected: &std::collections::HashSet<String>,
1558 visiting: &mut Vec<String>,
1559) -> Result<(), DependencyError> {
1560 if resolved_set.contains(cap_id) {
1562 return Ok(());
1563 }
1564
1565 if visiting.contains(&cap_id.to_string()) {
1567 return Err(DependencyError::CircularDependency {
1568 capability_id: cap_id.to_string(),
1569 chain: visiting.clone(),
1570 });
1571 }
1572
1573 let capability = match registry.get(cap_id) {
1575 Some(cap) => cap,
1576 None => {
1577 if is_declarative_capability(cap_id) && !resolved_set.contains(cap_id) {
1578 resolved.push(cap_id.to_string());
1579 resolved_set.insert(cap_id.to_string());
1580 if !user_selected.contains(cap_id) {
1581 added_as_dependencies.push(cap_id.to_string());
1582 }
1583 }
1584 return Ok(());
1585 }
1586 };
1587
1588 visiting.push(cap_id.to_string());
1590
1591 for dep_id in capability.dependencies() {
1593 resolve_single_capability(
1594 dep_id,
1595 registry,
1596 resolved,
1597 resolved_set,
1598 added_as_dependencies,
1599 user_selected,
1600 visiting,
1601 )?;
1602 }
1603
1604 visiting.pop();
1606
1607 if !resolved_set.contains(cap_id) {
1609 resolved.push(cap_id.to_string());
1610 resolved_set.insert(cap_id.to_string());
1611
1612 if !user_selected.contains(cap_id) {
1614 added_as_dependencies.push(cap_id.to_string());
1615 }
1616 }
1617
1618 Ok(())
1619}
1620
1621pub fn compute_features(capability_ids: &[String], registry: &CapabilityRegistry) -> Vec<String> {
1626 use std::collections::HashSet;
1627
1628 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
1629 Ok(resolved) => resolved.resolved_ids,
1630 Err(_) => capability_ids.to_vec(),
1631 };
1632
1633 let mut seen = HashSet::new();
1634 let mut features = Vec::new();
1635 for cap_id in &resolved_ids {
1636 if let Some(cap) = registry.get(cap_id) {
1637 for feature in cap.features() {
1638 if seen.insert(feature) {
1639 features.push(feature.to_string());
1640 }
1641 }
1642 }
1643 }
1644 features
1645}
1646
1647pub fn get_dependencies(cap_id: &str, registry: &CapabilityRegistry) -> Vec<String> {
1650 registry
1651 .get(cap_id)
1652 .map(|cap| cap.dependencies().iter().map(|s| s.to_string()).collect())
1653 .unwrap_or_default()
1654}
1655
1656pub async fn collect_capabilities(
1672 capability_ids: &[String],
1673 registry: &CapabilityRegistry,
1674 ctx: &SystemPromptContext,
1675) -> CollectedCapabilities {
1676 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
1679 Ok(resolved) => resolved.resolved_ids,
1680 Err(e) => {
1681 tracing::warn!("Failed to resolve capability dependencies: {}", e);
1682 capability_ids.to_vec()
1683 }
1684 };
1685
1686 let configs: Vec<AgentCapabilityConfig> = resolved_ids
1688 .iter()
1689 .map(|id| AgentCapabilityConfig {
1690 capability_ref: CapabilityId::new(id),
1691 config: serde_json::Value::Object(serde_json::Map::new()),
1692 })
1693 .collect();
1694
1695 collect_capabilities_with_configs(&configs, registry, ctx).await
1696}
1697
1698pub async fn collect_capabilities_with_configs(
1709 capability_configs: &[AgentCapabilityConfig],
1710 registry: &CapabilityRegistry,
1711 ctx: &SystemPromptContext,
1712) -> CollectedCapabilities {
1713 let mut system_prompt_parts: Vec<String> = Vec::new();
1714 let mut system_prompt_attributions: Vec<SystemPromptAttribution> = Vec::new();
1715 let mut tools: Vec<Box<dyn Tool>> = Vec::new();
1716 let mut tool_definitions: Vec<ToolDefinition> = Vec::new();
1717 let mut mounts: Vec<MountPoint> = Vec::new();
1718 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1719 Vec::new();
1720 let mut applied_ids: Vec<String> = Vec::new();
1721 let mut tool_search: Option<crate::llm_driver_registry::ToolSearchConfig> = None;
1722 let mut prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig> = None;
1723 let mut tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>> = Vec::new();
1724 let mut tool_call_hooks: Vec<Arc<dyn ToolCallHook>> = Vec::new();
1725 let mut mcp_servers = ScopedMcpServers::default();
1726
1727 for cap_config in capability_configs {
1728 let cap_id = cap_config.capability_ref.as_str();
1729 if is_declarative_capability(cap_id) {
1730 match serde_json::from_value::<DeclarativeCapabilityDefinition>(
1731 cap_config.config.clone(),
1732 ) {
1733 Ok(definition) => {
1734 if definition.status != CapabilityStatus::Available {
1735 continue;
1736 }
1737
1738 if let Some(prompt) = definition.system_prompt.as_deref() {
1739 let contribution =
1740 format!("<capability id=\"{}\">\n{}\n</capability>", cap_id, prompt);
1741 system_prompt_attributions.push(SystemPromptAttribution {
1742 capability_id: cap_id.to_string(),
1743 content: contribution.clone(),
1744 });
1745 system_prompt_parts.push(contribution);
1746 }
1747
1748 mounts.extend(definition.mounts(cap_id));
1749 if let Some(ref servers) = definition.mcp_servers {
1750 mcp_servers = merge_scoped_mcp_servers(&mcp_servers, servers);
1751 }
1752 for skill in definition.skill_contributions() {
1753 mounts.push(skill.to_mount(cap_id));
1754 }
1755
1756 applied_ids.push(cap_id.to_string());
1757 }
1758 Err(error) => {
1759 tracing::warn!(
1760 capability_id = %cap_id,
1761 error = %error,
1762 "Skipping invalid declarative capability config"
1763 );
1764 }
1765 }
1766 continue;
1767 }
1768 if let Some(capability) = registry.get(cap_id) {
1769 if capability.status() != CapabilityStatus::Available {
1771 continue;
1772 }
1773
1774 if let Some(contribution) = capability
1776 .system_prompt_contribution_with_config(ctx, &cap_config.config)
1777 .await
1778 {
1779 system_prompt_attributions.push(SystemPromptAttribution {
1780 capability_id: cap_id.to_string(),
1781 content: contribution.clone(),
1782 });
1783 system_prompt_parts.push(contribution);
1784 }
1785
1786 tools.extend(capability.tools_with_config(&cap_config.config));
1788 tool_definition_hooks.extend(capability.tool_definition_hooks());
1789 tool_call_hooks.extend(capability.tool_call_hooks());
1790 let cap_category = capability.category();
1795 for def in capability.tool_definitions() {
1796 let def = match (def.category(), cap_category) {
1797 (None, Some(cat)) => def.with_category(cat),
1798 _ => def,
1799 }
1800 .with_capability_attribution(cap_id, Some(capability.name()));
1801 tool_definitions.push(def);
1802 }
1803
1804 if cap_id == OPENAI_TOOL_SEARCH_CAPABILITY_ID {
1806 let threshold = cap_config
1808 .config
1809 .get("threshold")
1810 .and_then(|v| v.as_u64())
1811 .map(|v| v as usize)
1812 .unwrap_or(DEFAULT_TOOL_SEARCH_THRESHOLD);
1813 tool_search = Some(crate::llm_driver_registry::ToolSearchConfig {
1814 enabled: true,
1815 threshold,
1816 });
1817 }
1818
1819 if cap_id == PROMPT_CACHING_CAPABILITY_ID {
1820 let strategy = cap_config
1821 .config
1822 .get("strategy")
1823 .and_then(|v| v.as_str())
1824 .map(|value| match value {
1825 "auto" => crate::llm_driver_registry::PromptCacheStrategy::Auto,
1826 _ => crate::llm_driver_registry::PromptCacheStrategy::Auto,
1827 })
1828 .unwrap_or(crate::llm_driver_registry::PromptCacheStrategy::Auto);
1829 let gemini_cached_content = cap_config
1830 .config
1831 .get("gemini_cached_content")
1832 .and_then(|v| v.as_str())
1833 .map(str::to_string);
1834 prompt_cache = Some(crate::llm_driver_registry::PromptCacheConfig {
1835 enabled: true,
1836 strategy,
1837 gemini_cached_content,
1838 });
1839 }
1840
1841 mounts.extend(capability.mounts());
1843
1844 mcp_servers = merge_scoped_mcp_servers(
1845 &mcp_servers,
1846 &capability.mcp_servers_with_config(&cap_config.config),
1847 );
1848
1849 for skill in capability.contribute_skills() {
1853 mounts.push(skill.to_mount(cap_id));
1854 }
1855
1856 if let Some(provider) = capability.message_filter_provider() {
1858 message_filter_providers.push((provider, cap_config.config.clone()));
1859 }
1860
1861 applied_ids.push(cap_id.to_string());
1862 }
1863 }
1864
1865 if !applied_ids
1877 .iter()
1878 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID)
1879 && tool_definitions
1880 .iter()
1881 .any(|def| def.hints().supports_background == Some(true))
1882 && let Some(bg_cap) = registry.get(BACKGROUND_EXECUTION_CAPABILITY_ID)
1883 && bg_cap.status() == CapabilityStatus::Available
1884 {
1885 tools.extend(bg_cap.tools());
1886 let cap_category = bg_cap.category();
1887 for def in bg_cap.tool_definitions() {
1888 let def = match (def.category(), cap_category) {
1889 (None, Some(cat)) => def.with_category(cat),
1890 _ => def,
1891 }
1892 .with_capability_attribution(BACKGROUND_EXECUTION_CAPABILITY_ID, Some(bg_cap.name()));
1893 tool_definitions.push(def);
1894 }
1895 applied_ids.push(BACKGROUND_EXECUTION_CAPABILITY_ID.to_string());
1896 }
1897
1898 message_filter_providers.sort_by_key(|(p, _)| p.priority());
1900
1901 CollectedCapabilities {
1902 system_prompt_parts,
1903 system_prompt_attributions,
1904 tools,
1905 tool_definitions,
1906 mounts,
1907 message_filter_providers,
1908 applied_ids,
1909 tool_search,
1910 prompt_cache,
1911 tool_definition_hooks,
1912 tool_call_hooks,
1913 mcp_servers,
1914 }
1915}
1916
1917pub struct AppliedCapabilities {
1923 pub runtime_agent: RuntimeAgent,
1925 pub tool_registry: ToolRegistry,
1927 pub applied_ids: Vec<String>,
1929}
1930
1931pub async fn apply_capabilities(
1968 base_runtime_agent: RuntimeAgent,
1969 capability_ids: &[String],
1970 registry: &CapabilityRegistry,
1971 ctx: &SystemPromptContext,
1972) -> AppliedCapabilities {
1973 let collected = collect_capabilities(capability_ids, registry, ctx).await;
1974
1975 let final_system_prompt = match collected.system_prompt_prefix() {
1977 Some(prefix) => format!(
1978 "{}\n\n<system-prompt>\n{}\n</system-prompt>",
1979 prefix, base_runtime_agent.system_prompt
1980 ),
1981 None => base_runtime_agent.system_prompt,
1982 };
1983
1984 let mut tool_registry = ToolRegistry::new();
1986 for tool in collected.tools {
1987 tool_registry.register_boxed(tool);
1988 }
1989
1990 let mut tools = collected.tool_definitions;
1992 for hook in &collected.tool_definition_hooks {
1993 tools = hook.transform(tools);
1994 }
1995
1996 let runtime_agent = RuntimeAgent {
1997 system_prompt: final_system_prompt,
1998 model: base_runtime_agent.model,
1999 tools,
2000 max_iterations: base_runtime_agent.max_iterations,
2001 temperature: base_runtime_agent.temperature,
2002 max_tokens: base_runtime_agent.max_tokens,
2003 tool_search: collected.tool_search,
2004 prompt_cache: collected.prompt_cache,
2005 network_access: base_runtime_agent.network_access,
2006 };
2007
2008 AppliedCapabilities {
2009 runtime_agent,
2010 tool_registry,
2011 applied_ids: collected.applied_ids,
2012 }
2013}
2014
2015#[cfg(test)]
2020mod tests {
2021 use super::*;
2022 use crate::typed_id::SessionId;
2023 use std::collections::BTreeSet;
2024 use uuid::Uuid;
2025
2026 static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2028
2029 fn lock_env() -> std::sync::MutexGuard<'static, ()> {
2030 ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner())
2031 }
2032
2033 fn test_ctx() -> SystemPromptContext {
2035 SystemPromptContext::without_file_store(SessionId::new())
2036 }
2037
2038 fn expected_core_builtin_ids() -> BTreeSet<&'static str> {
2040 let mut ids = [
2041 "agent_instructions",
2042 "human_intent",
2043 "budgeting",
2044 "self_budget",
2045 "noop",
2046 "current_time",
2047 "research",
2048 "platform_management",
2049 "session_file_system",
2050 "workspace_volumes",
2051 "session_storage",
2052 "session",
2053 "session_sql_database",
2054 "test_math",
2055 "test_weather",
2056 "stateless_todo_list",
2057 "web_fetch",
2058 "virtual_bash",
2059 "background_execution",
2060 "session_schedule",
2061 "btw",
2062 "infinity_context",
2063 "compaction",
2064 "memory",
2065 "openai_tool_search",
2066 "tool_search",
2067 "prompt_caching",
2068 "skills",
2069 "subagents",
2070 "system_commands",
2071 "sample_data",
2072 "data_knowledge",
2073 "knowledge_base",
2074 "tool_output_persistence",
2075 "fake_warehouse",
2076 "fake_aws",
2077 "fake_crm",
2078 "fake_financial",
2079 "loop_detection",
2080 "prompt_canary_guardrail",
2081 "user_hooks",
2082 ]
2083 .into_iter()
2084 .collect::<BTreeSet<_>>();
2085 if cfg!(feature = "ui-capabilities") {
2086 ids.insert("openui");
2087 ids.insert("a2ui");
2088 }
2089 ids
2090 }
2091
2092 fn expected_dev_builtin_ids() -> BTreeSet<&'static str> {
2094 let mut ids = expected_core_builtin_ids();
2095 ids.insert("agent_handoff");
2096 ids.insert("a2a_agent_delegation");
2097 ids
2098 }
2099
2100 fn registry_ids(registry: &CapabilityRegistry) -> BTreeSet<&str> {
2101 registry.capabilities.keys().map(String::as_str).collect()
2102 }
2103
2104 #[test]
2114 fn test_capability_registry_with_builtins_dev() {
2115 let _lock = lock_env();
2117 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2118 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
2119 assert_eq!(registry_ids(®istry), expected_dev_builtin_ids());
2120 assert!(registry.has("agent_handoff"));
2121 assert!(registry.has("a2a_agent_delegation"));
2122 }
2123
2124 #[test]
2125 fn test_capability_registry_with_builtins_prod() {
2126 let _lock = lock_env();
2128 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2129 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2130 assert_eq!(registry_ids(®istry), expected_core_builtin_ids());
2131 assert!(!registry.has("docker_container"));
2133 assert!(!registry.has("agent_handoff"));
2134 assert!(!registry.has("a2a_agent_delegation"));
2135 }
2136
2137 #[test]
2138 fn test_agent_delegation_enabled_by_env_in_prod() {
2139 let _lock = lock_env();
2141 unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "true") };
2142 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2143 assert!(registry.has("agent_handoff"));
2144 assert!(registry.has("a2a_agent_delegation"));
2145 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2146 }
2147
2148 #[test]
2149 fn test_agent_delegation_disabled_by_env_in_dev() {
2150 let _lock = lock_env();
2152 unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "false") };
2153 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
2154 assert!(!registry.has("agent_handoff"));
2155 assert!(!registry.has("a2a_agent_delegation"));
2156 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2157 }
2158
2159 #[test]
2160 fn test_capability_registry_get() {
2161 let registry = CapabilityRegistry::with_builtins();
2162
2163 let noop = registry.get("noop").unwrap();
2164 assert_eq!(noop.id(), "noop");
2165 assert_eq!(noop.name(), "No-Op");
2166 assert_eq!(noop.status(), CapabilityStatus::Available);
2167 }
2168
2169 #[test]
2170 fn test_capability_registry_blueprint_with_capability() {
2171 struct BlueprintProviderCapability;
2172
2173 impl Capability for BlueprintProviderCapability {
2174 fn id(&self) -> &str {
2175 "blueprint_provider"
2176 }
2177 fn name(&self) -> &str {
2178 "Blueprint Provider"
2179 }
2180 fn description(&self) -> &str {
2181 "Capability that provides a blueprint for tests"
2182 }
2183 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
2184 vec![AgentBlueprint {
2185 id: "test_blueprint",
2186 name: "Test Blueprint",
2187 description: "Blueprint for capability registry tests",
2188 model: BlueprintModel::Inherit,
2189 system_prompt: "Test prompt",
2190 tools: vec![],
2191 max_turns: None,
2192 config_schema: None,
2193 }]
2194 }
2195 }
2196
2197 let mut registry = CapabilityRegistry::new();
2198 registry.register(BlueprintProviderCapability);
2199
2200 let (capability_id, blueprint) = registry
2201 .blueprint_with_capability("test_blueprint")
2202 .expect("blueprint should resolve with capability id");
2203 assert_eq!(capability_id, "blueprint_provider");
2204 assert_eq!(blueprint.id, "test_blueprint");
2205 }
2206
2207 #[test]
2208 fn test_capability_registry_builder() {
2209 let registry = CapabilityRegistry::builder()
2210 .capability(NoopCapability)
2211 .capability(CurrentTimeCapability)
2212 .build();
2213
2214 assert!(registry.has("noop"));
2215 assert!(registry.has("current_time"));
2216 assert_eq!(registry.len(), 2);
2217 }
2218
2219 #[test]
2220 fn test_capability_status() {
2221 let registry = CapabilityRegistry::with_builtins();
2222
2223 let current_time = registry.get("current_time").unwrap();
2224 assert_eq!(current_time.status(), CapabilityStatus::Available);
2225
2226 let research = registry.get("research").unwrap();
2227 assert_eq!(research.status(), CapabilityStatus::ComingSoon);
2228 }
2229
2230 #[test]
2231 fn test_capability_icons_and_categories() {
2232 let registry = CapabilityRegistry::with_builtins();
2233
2234 let noop = registry.get("noop").unwrap();
2235 assert_eq!(noop.icon(), Some("circle-off"));
2236 assert_eq!(noop.category(), Some("Testing"));
2237
2238 let current_time = registry.get("current_time").unwrap();
2239 assert_eq!(current_time.icon(), Some("clock"));
2240 assert_eq!(current_time.category(), Some("Utilities"));
2241 }
2242
2243 #[test]
2244 fn test_system_prompt_preview_default_delegates_to_addition() {
2245 let registry = CapabilityRegistry::with_builtins();
2246
2247 let test_math = registry.get("test_math").unwrap();
2249 assert_eq!(
2250 test_math.system_prompt_preview().as_deref(),
2251 test_math.system_prompt_addition()
2252 );
2253
2254 let current_time = registry.get("current_time").unwrap();
2256 assert!(current_time.system_prompt_preview().is_none());
2257 assert!(current_time.system_prompt_addition().is_none());
2258 }
2259
2260 #[test]
2261 fn test_system_prompt_preview_dynamic_capability() {
2262 let registry = CapabilityRegistry::with_builtins();
2263 let cap = registry.get("agent_instructions").unwrap();
2264
2265 assert!(cap.system_prompt_addition().is_none());
2267 assert!(cap.system_prompt_preview().is_some());
2268 assert!(cap.system_prompt_preview().unwrap().contains("AGENTS.md"));
2269 }
2270
2271 #[tokio::test]
2276 async fn test_apply_capabilities_empty() {
2277 let registry = CapabilityRegistry::with_builtins();
2278 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2279
2280 let applied =
2281 apply_capabilities(base_runtime_agent.clone(), &[], ®istry, &test_ctx()).await;
2282
2283 assert_eq!(
2284 applied.runtime_agent.system_prompt,
2285 base_runtime_agent.system_prompt
2286 );
2287 assert!(applied.tool_registry.is_empty());
2288 assert!(applied.applied_ids.is_empty());
2289 }
2290
2291 #[tokio::test]
2292 async fn test_apply_capabilities_noop() {
2293 let registry = CapabilityRegistry::with_builtins();
2294 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2295
2296 let applied = apply_capabilities(
2297 base_runtime_agent.clone(),
2298 &["noop".to_string()],
2299 ®istry,
2300 &test_ctx(),
2301 )
2302 .await;
2303
2304 assert_eq!(
2306 applied.runtime_agent.system_prompt,
2307 base_runtime_agent.system_prompt
2308 );
2309 assert!(applied.tool_registry.is_empty());
2310 assert_eq!(applied.applied_ids, vec!["noop"]);
2311 }
2312
2313 #[tokio::test]
2314 async fn test_apply_capabilities_current_time() {
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 &["current_time".to_string()],
2321 ®istry,
2322 &test_ctx(),
2323 )
2324 .await;
2325
2326 assert_eq!(
2328 applied.runtime_agent.system_prompt,
2329 base_runtime_agent.system_prompt
2330 );
2331 assert!(applied.tool_registry.has("get_current_time"));
2332 assert_eq!(applied.tool_registry.len(), 1);
2333 assert_eq!(applied.applied_ids, vec!["current_time"]);
2334 }
2335
2336 #[tokio::test]
2337 async fn test_apply_capabilities_skips_coming_soon() {
2338 let registry = CapabilityRegistry::with_builtins();
2339 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2340
2341 let applied = apply_capabilities(
2343 base_runtime_agent.clone(),
2344 &["research".to_string()],
2345 ®istry,
2346 &test_ctx(),
2347 )
2348 .await;
2349
2350 assert_eq!(
2352 applied.runtime_agent.system_prompt,
2353 base_runtime_agent.system_prompt
2354 );
2355 assert!(applied.applied_ids.is_empty()); }
2357
2358 #[tokio::test]
2359 async fn test_apply_capabilities_multiple() {
2360 let registry = CapabilityRegistry::with_builtins();
2361 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2362
2363 let applied = apply_capabilities(
2364 base_runtime_agent.clone(),
2365 &["noop".to_string(), "current_time".to_string()],
2366 ®istry,
2367 &test_ctx(),
2368 )
2369 .await;
2370
2371 assert!(applied.tool_registry.has("get_current_time"));
2372 assert_eq!(applied.applied_ids, vec!["noop", "current_time"]);
2373 }
2374
2375 #[tokio::test]
2376 async fn test_apply_capabilities_preserves_order() {
2377 let registry = CapabilityRegistry::with_builtins();
2378 let base_runtime_agent = RuntimeAgent::new("Base prompt.", "gpt-5.2");
2379
2380 let applied = apply_capabilities(
2382 base_runtime_agent,
2383 &["current_time".to_string(), "noop".to_string()],
2384 ®istry,
2385 &test_ctx(),
2386 )
2387 .await;
2388
2389 assert_eq!(applied.applied_ids, vec!["current_time", "noop"]);
2390 }
2391
2392 #[tokio::test]
2393 async fn test_apply_capabilities_test_math() {
2394 let registry = CapabilityRegistry::with_builtins();
2395 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2396
2397 let applied = apply_capabilities(
2398 base_runtime_agent.clone(),
2399 &["test_math".to_string()],
2400 ®istry,
2401 &test_ctx(),
2402 )
2403 .await;
2404
2405 assert!(
2407 !applied
2408 .runtime_agent
2409 .system_prompt
2410 .contains("<capability id=\"test_math\">")
2411 );
2412 assert!(
2414 applied
2415 .runtime_agent
2416 .system_prompt
2417 .contains("You are a helpful assistant.")
2418 );
2419 assert!(applied.tool_registry.has("add"));
2420 assert!(applied.tool_registry.has("subtract"));
2421 assert!(applied.tool_registry.has("multiply"));
2422 assert!(applied.tool_registry.has("divide"));
2423 assert_eq!(applied.tool_registry.len(), 4);
2424 }
2425
2426 #[tokio::test]
2427 async fn test_apply_capabilities_test_weather() {
2428 let registry = CapabilityRegistry::with_builtins();
2429 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2430
2431 let applied = apply_capabilities(
2432 base_runtime_agent.clone(),
2433 &["test_weather".to_string()],
2434 ®istry,
2435 &test_ctx(),
2436 )
2437 .await;
2438
2439 assert!(
2441 !applied
2442 .runtime_agent
2443 .system_prompt
2444 .contains("<capability id=\"test_weather\">")
2445 );
2446 assert!(applied.tool_registry.has("get_weather"));
2447 assert!(applied.tool_registry.has("get_forecast"));
2448 assert_eq!(applied.tool_registry.len(), 2);
2449 }
2450
2451 #[tokio::test]
2452 async fn test_apply_capabilities_test_math_and_test_weather() {
2453 let registry = CapabilityRegistry::with_builtins();
2454 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2455
2456 let applied = apply_capabilities(
2457 base_runtime_agent.clone(),
2458 &["test_math".to_string(), "test_weather".to_string()],
2459 ®istry,
2460 &test_ctx(),
2461 )
2462 .await;
2463
2464 assert_eq!(applied.tool_registry.len(), 6); assert!(applied.tool_registry.has("add"));
2467 assert!(applied.tool_registry.has("get_weather"));
2468 }
2469
2470 #[tokio::test]
2471 async fn test_apply_capabilities_stateless_todo_list() {
2472 let registry = CapabilityRegistry::with_builtins();
2473 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2474
2475 let applied = apply_capabilities(
2476 base_runtime_agent.clone(),
2477 &["stateless_todo_list".to_string()],
2478 ®istry,
2479 &test_ctx(),
2480 )
2481 .await;
2482
2483 assert!(
2485 applied
2486 .runtime_agent
2487 .system_prompt
2488 .contains("Task Management")
2489 );
2490 assert!(applied.runtime_agent.system_prompt.contains("write_todos"));
2491 assert!(applied.tool_registry.has("write_todos"));
2492 assert_eq!(applied.tool_registry.len(), 1);
2493 }
2494
2495 #[tokio::test]
2496 async fn test_apply_capabilities_web_fetch() {
2497 let registry = CapabilityRegistry::with_builtins();
2498 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2499
2500 let applied = apply_capabilities(
2501 base_runtime_agent.clone(),
2502 &["web_fetch".to_string()],
2503 ®istry,
2504 &test_ctx(),
2505 )
2506 .await;
2507
2508 assert!(
2510 applied
2511 .runtime_agent
2512 .system_prompt
2513 .contains(&base_runtime_agent.system_prompt)
2514 );
2515 assert!(applied.runtime_agent.system_prompt.contains("web_fetch"));
2516 assert!(applied.tool_registry.has("web_fetch"));
2517 assert_eq!(applied.tool_registry.len(), 1);
2518 }
2519
2520 #[tokio::test]
2525 async fn test_xml_tags_wrap_capability_prompts() {
2526 let registry = CapabilityRegistry::with_builtins();
2527 let collected =
2528 collect_capabilities(&["stateless_todo_list".to_string()], ®istry, &test_ctx())
2529 .await;
2530
2531 assert_eq!(collected.system_prompt_parts.len(), 1);
2532 let part = &collected.system_prompt_parts[0];
2533 assert!(part.starts_with("<capability id=\"stateless_todo_list\">"));
2534 assert!(part.ends_with("</capability>"));
2535 assert!(part.contains("Task Management"));
2536 }
2537
2538 #[tokio::test]
2539 async fn test_xml_tags_multiple_capabilities() {
2540 let registry = CapabilityRegistry::with_builtins();
2541 let collected = collect_capabilities(
2542 &[
2543 "stateless_todo_list".to_string(),
2544 "session_schedule".to_string(),
2545 ],
2546 ®istry,
2547 &test_ctx(),
2548 )
2549 .await;
2550
2551 assert_eq!(collected.system_prompt_parts.len(), 2);
2552 assert!(
2553 collected.system_prompt_parts[0].starts_with("<capability id=\"stateless_todo_list\">")
2554 );
2555 assert!(
2556 collected.system_prompt_parts[1].starts_with("<capability id=\"session_schedule\">")
2557 );
2558
2559 let prefix = collected.system_prompt_prefix().unwrap();
2560 assert!(prefix.contains("</capability>\n\n<capability"));
2562 }
2563
2564 #[tokio::test]
2565 async fn test_xml_tags_system_prompt_wrapping() {
2566 let registry = CapabilityRegistry::with_builtins();
2567 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2568
2569 let applied = apply_capabilities(
2570 base,
2571 &["stateless_todo_list".to_string()],
2572 ®istry,
2573 &test_ctx(),
2574 )
2575 .await;
2576
2577 let prompt = &applied.runtime_agent.system_prompt;
2578 assert!(prompt.contains("<capability id=\"stateless_todo_list\">"));
2580 assert!(prompt.contains("</capability>"));
2581 assert!(prompt.contains("<system-prompt>\nYou are helpful.\n</system-prompt>"));
2583 }
2584
2585 #[tokio::test]
2586 async fn test_no_xml_wrapping_without_capabilities() {
2587 let registry = CapabilityRegistry::with_builtins();
2588 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2589
2590 let applied = apply_capabilities(base, &[], ®istry, &test_ctx()).await;
2591
2592 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2594 assert!(
2595 !applied
2596 .runtime_agent
2597 .system_prompt
2598 .contains("<system-prompt>")
2599 );
2600 }
2601
2602 #[tokio::test]
2603 async fn test_no_xml_wrapping_for_noop_capability() {
2604 let registry = CapabilityRegistry::with_builtins();
2605 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2606
2607 let applied = apply_capabilities(base, &["noop".to_string()], ®istry, &test_ctx()).await;
2609
2610 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2611 assert!(
2612 !applied
2613 .runtime_agent
2614 .system_prompt
2615 .contains("<system-prompt>")
2616 );
2617 }
2618
2619 #[tokio::test]
2624 async fn test_collect_capabilities_includes_mounts() {
2625 let registry = CapabilityRegistry::with_builtins();
2626
2627 let collected =
2628 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
2629
2630 assert!(!collected.mounts.is_empty());
2631 assert_eq!(collected.mounts.len(), 1);
2632 assert_eq!(collected.mounts[0].path, "/samples");
2633 assert!(collected.mounts[0].is_readonly());
2634 }
2635
2636 #[tokio::test]
2637 async fn test_collect_capabilities_empty_mounts_by_default() {
2638 let registry = CapabilityRegistry::with_builtins();
2639
2640 let collected =
2642 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
2643
2644 assert!(collected.mounts.is_empty());
2645 }
2646
2647 #[tokio::test]
2648 async fn test_collect_capabilities_combines_mounts() {
2649 let registry = CapabilityRegistry::with_builtins();
2650
2651 let collected = collect_capabilities(
2654 &["sample_data".to_string(), "current_time".to_string()],
2655 ®istry,
2656 &test_ctx(),
2657 )
2658 .await;
2659
2660 assert_eq!(collected.mounts.len(), 1);
2661 assert!(
2663 collected
2664 .applied_ids
2665 .iter()
2666 .any(|id| id == "session_file_system")
2667 );
2668 assert!(collected.applied_ids.iter().any(|id| id == "sample_data"));
2669 assert!(collected.applied_ids.iter().any(|id| id == "current_time"));
2670 }
2671
2672 #[test]
2673 fn test_sample_data_capability() {
2674 let registry = CapabilityRegistry::with_builtins();
2675 let cap = registry.get("sample_data").unwrap();
2676
2677 assert_eq!(cap.id(), "sample_data");
2678 assert_eq!(cap.name(), "Sample Data");
2679 assert_eq!(cap.status(), CapabilityStatus::Available);
2680
2681 assert!(cap.system_prompt_addition().is_some());
2683 assert!(cap.tools().is_empty());
2684
2685 assert!(!cap.mounts().is_empty());
2687 }
2688
2689 #[test]
2694 fn test_resolve_dependencies_empty() {
2695 let registry = CapabilityRegistry::with_builtins();
2696
2697 let resolved = resolve_dependencies(&[], ®istry).unwrap();
2698
2699 assert!(resolved.resolved_ids.is_empty());
2700 assert!(resolved.added_as_dependencies.is_empty());
2701 assert!(resolved.user_selected.is_empty());
2702 }
2703
2704 #[test]
2705 fn test_resolve_dependencies_no_deps() {
2706 let registry = CapabilityRegistry::with_builtins();
2707
2708 let resolved = resolve_dependencies(&["current_time".to_string()], ®istry).unwrap();
2710
2711 assert_eq!(resolved.resolved_ids, vec!["current_time"]);
2712 assert!(resolved.added_as_dependencies.is_empty());
2713 }
2714
2715 #[test]
2716 fn test_resolve_dependencies_with_deps() {
2717 let registry = CapabilityRegistry::with_builtins();
2718
2719 let resolved = resolve_dependencies(&["sample_data".to_string()], ®istry).unwrap();
2721
2722 assert_eq!(resolved.resolved_ids.len(), 2);
2724 let fs_pos = resolved
2725 .resolved_ids
2726 .iter()
2727 .position(|id| id == "session_file_system")
2728 .unwrap();
2729 let sd_pos = resolved
2730 .resolved_ids
2731 .iter()
2732 .position(|id| id == "sample_data")
2733 .unwrap();
2734 assert!(fs_pos < sd_pos, "FileSystem should come before SampleData");
2735
2736 assert_eq!(resolved.added_as_dependencies, vec!["session_file_system"]);
2738 }
2739
2740 #[test]
2741 fn test_resolve_dependencies_already_selected() {
2742 let registry = CapabilityRegistry::with_builtins();
2743
2744 let resolved = resolve_dependencies(
2746 &["session_file_system".to_string(), "sample_data".to_string()],
2747 ®istry,
2748 )
2749 .unwrap();
2750
2751 assert_eq!(resolved.resolved_ids.len(), 2);
2752 assert!(resolved.added_as_dependencies.is_empty());
2754 }
2755
2756 #[test]
2757 fn test_resolve_dependencies_preserves_order() {
2758 let registry = CapabilityRegistry::with_builtins();
2759
2760 let resolved =
2762 resolve_dependencies(&["current_time".to_string(), "noop".to_string()], ®istry)
2763 .unwrap();
2764
2765 assert_eq!(resolved.resolved_ids, vec!["current_time", "noop"]);
2766 }
2767
2768 #[test]
2769 fn test_resolve_dependencies_unknown_capability() {
2770 let registry = CapabilityRegistry::with_builtins();
2771
2772 let resolved =
2774 resolve_dependencies(&["unknown_capability".to_string()], ®istry).unwrap();
2775
2776 assert!(resolved.resolved_ids.is_empty());
2777 }
2778
2779 #[test]
2780 fn test_get_dependencies() {
2781 let registry = CapabilityRegistry::with_builtins();
2782
2783 let deps = get_dependencies("sample_data", ®istry);
2785 assert_eq!(deps, vec!["session_file_system"]);
2786
2787 let deps = get_dependencies("current_time", ®istry);
2789 assert!(deps.is_empty());
2790
2791 let deps = get_dependencies("unknown", ®istry);
2793 assert!(deps.is_empty());
2794 }
2795
2796 #[test]
2797 fn test_sample_data_has_dependency() {
2798 let registry = CapabilityRegistry::with_builtins();
2799 let cap = registry.get("sample_data").unwrap();
2800
2801 let deps = cap.dependencies();
2802 assert_eq!(deps.len(), 1);
2803 assert_eq!(deps[0], "session_file_system");
2804 }
2805
2806 #[test]
2807 fn test_noop_has_no_dependencies() {
2808 let registry = CapabilityRegistry::with_builtins();
2809 let cap = registry.get("noop").unwrap();
2810
2811 assert!(cap.dependencies().is_empty());
2812 }
2813
2814 #[test]
2818 fn test_circular_dependency_error() {
2819 struct CapA;
2821 struct CapB;
2822
2823 impl Capability for CapA {
2824 fn id(&self) -> &str {
2825 "test_cap_a"
2826 }
2827 fn name(&self) -> &str {
2828 "Test A"
2829 }
2830 fn description(&self) -> &str {
2831 "Test capability A"
2832 }
2833 fn dependencies(&self) -> Vec<&'static str> {
2834 vec!["test_cap_b"]
2835 }
2836 }
2837
2838 impl Capability for CapB {
2839 fn id(&self) -> &str {
2840 "test_cap_b"
2841 }
2842 fn name(&self) -> &str {
2843 "Test B"
2844 }
2845 fn description(&self) -> &str {
2846 "Test capability B"
2847 }
2848 fn dependencies(&self) -> Vec<&'static str> {
2849 vec!["test_cap_a"]
2850 }
2851 }
2852
2853 let mut registry = CapabilityRegistry::new();
2854 registry.register(CapA);
2855 registry.register(CapB);
2856
2857 let result = resolve_dependencies(&["test_cap_a".to_string()], ®istry);
2858
2859 assert!(result.is_err());
2860 match result.unwrap_err() {
2861 DependencyError::CircularDependency { capability_id, .. } => {
2862 assert_eq!(capability_id, "test_cap_a");
2863 }
2864 _ => panic!("Expected CircularDependency error"),
2865 }
2866 }
2867
2868 use crate::message_filter::{MessageFilter, MessageFilterProvider, MessageQuery};
2873
2874 struct FilterTestCapability {
2876 priority: i32,
2877 }
2878
2879 impl Capability for FilterTestCapability {
2880 fn id(&self) -> &str {
2881 "filter_test"
2882 }
2883 fn name(&self) -> &str {
2884 "Filter Test"
2885 }
2886 fn description(&self) -> &str {
2887 "Test capability with message filter"
2888 }
2889 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
2890 Some(Arc::new(FilterTestProvider {
2891 priority: self.priority,
2892 }))
2893 }
2894 }
2895
2896 struct FilterTestProvider {
2897 priority: i32,
2898 }
2899
2900 impl MessageFilterProvider for FilterTestProvider {
2901 fn apply_filters(&self, query: &mut MessageQuery, config: &serde_json::Value) {
2902 if let Some(search) = config.get("search").and_then(|v| v.as_str()) {
2904 query
2905 .filters
2906 .push(MessageFilter::Search(search.to_string()));
2907 }
2908 }
2909
2910 fn priority(&self) -> i32 {
2911 self.priority
2912 }
2913 }
2914
2915 #[tokio::test]
2916 async fn test_collect_capabilities_with_configs_no_filter_providers() {
2917 let registry = CapabilityRegistry::with_builtins();
2918 let configs = vec![AgentCapabilityConfig {
2919 capability_ref: CapabilityId::new("current_time"),
2920 config: serde_json::json!({}),
2921 }];
2922
2923 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2924
2925 assert!(collected.message_filter_providers.is_empty());
2926 assert!(!collected.has_message_filters());
2927 }
2928
2929 #[tokio::test]
2930 async fn test_collect_capabilities_with_configs_with_filter_provider() {
2931 let mut registry = CapabilityRegistry::new();
2932 registry.register(FilterTestCapability { priority: 0 });
2933
2934 let configs = vec![AgentCapabilityConfig {
2935 capability_ref: CapabilityId::new("filter_test"),
2936 config: serde_json::json!({ "search": "hello" }),
2937 }];
2938
2939 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2940
2941 assert_eq!(collected.message_filter_providers.len(), 1);
2942 assert!(collected.has_message_filters());
2943 }
2944
2945 #[tokio::test]
2946 async fn test_collect_capabilities_with_configs_filter_priority_order() {
2947 struct HighPriorityCapability;
2949 struct LowPriorityCapability;
2950
2951 impl Capability for HighPriorityCapability {
2952 fn id(&self) -> &str {
2953 "high_priority"
2954 }
2955 fn name(&self) -> &str {
2956 "High Priority"
2957 }
2958 fn description(&self) -> &str {
2959 "Test"
2960 }
2961 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
2962 Some(Arc::new(FilterTestProvider { priority: 10 }))
2963 }
2964 }
2965
2966 impl Capability for LowPriorityCapability {
2967 fn id(&self) -> &str {
2968 "low_priority"
2969 }
2970 fn name(&self) -> &str {
2971 "Low Priority"
2972 }
2973 fn description(&self) -> &str {
2974 "Test"
2975 }
2976 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
2977 Some(Arc::new(FilterTestProvider { priority: -5 }))
2978 }
2979 }
2980
2981 let mut registry = CapabilityRegistry::new();
2982 registry.register(HighPriorityCapability);
2983 registry.register(LowPriorityCapability);
2984
2985 let configs = vec![
2987 AgentCapabilityConfig {
2988 capability_ref: CapabilityId::new("high_priority"),
2989 config: serde_json::json!({}),
2990 },
2991 AgentCapabilityConfig {
2992 capability_ref: CapabilityId::new("low_priority"),
2993 config: serde_json::json!({}),
2994 },
2995 ];
2996
2997 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2998
2999 assert_eq!(collected.message_filter_providers.len(), 2);
3001 assert_eq!(collected.message_filter_providers[0].0.priority(), -5);
3002 assert_eq!(collected.message_filter_providers[1].0.priority(), 10);
3003 }
3004
3005 #[tokio::test]
3006 async fn test_collected_capabilities_apply_message_filters() {
3007 let mut registry = CapabilityRegistry::new();
3008 registry.register(FilterTestCapability { priority: 0 });
3009
3010 let configs = vec![AgentCapabilityConfig {
3011 capability_ref: CapabilityId::new("filter_test"),
3012 config: serde_json::json!({ "search": "test_query" }),
3013 }];
3014
3015 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3016
3017 let session_id: SessionId = Uuid::now_v7().into();
3019 let mut query = MessageQuery::new(session_id);
3020
3021 collected.apply_message_filters(&mut query);
3022
3023 assert_eq!(query.filters.len(), 1);
3025 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3026 }
3027
3028 #[tokio::test]
3029 async fn test_collected_capabilities_apply_multiple_filters_in_priority_order() {
3030 struct SearchCapability {
3031 id: &'static str,
3032 search_term: &'static str,
3033 priority: i32,
3034 }
3035
3036 struct SearchProvider {
3037 search_term: &'static str,
3038 priority: i32,
3039 }
3040
3041 impl MessageFilterProvider for SearchProvider {
3042 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3043 query
3044 .filters
3045 .push(MessageFilter::Search(self.search_term.to_string()));
3046 }
3047
3048 fn priority(&self) -> i32 {
3049 self.priority
3050 }
3051 }
3052
3053 impl Capability for SearchCapability {
3054 fn id(&self) -> &str {
3055 self.id
3056 }
3057 fn name(&self) -> &str {
3058 "Search"
3059 }
3060 fn description(&self) -> &str {
3061 "Test"
3062 }
3063 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3064 Some(Arc::new(SearchProvider {
3065 search_term: self.search_term,
3066 priority: self.priority,
3067 }))
3068 }
3069 }
3070
3071 let mut registry = CapabilityRegistry::new();
3072 registry.register(SearchCapability {
3073 id: "cap_a",
3074 search_term: "alpha",
3075 priority: 5,
3076 });
3077 registry.register(SearchCapability {
3078 id: "cap_b",
3079 search_term: "beta",
3080 priority: 1,
3081 });
3082 registry.register(SearchCapability {
3083 id: "cap_c",
3084 search_term: "gamma",
3085 priority: 10,
3086 });
3087
3088 let configs = vec![
3089 AgentCapabilityConfig {
3090 capability_ref: CapabilityId::new("cap_a"),
3091 config: serde_json::json!({}),
3092 },
3093 AgentCapabilityConfig {
3094 capability_ref: CapabilityId::new("cap_b"),
3095 config: serde_json::json!({}),
3096 },
3097 AgentCapabilityConfig {
3098 capability_ref: CapabilityId::new("cap_c"),
3099 config: serde_json::json!({}),
3100 },
3101 ];
3102
3103 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3104
3105 let session_id: SessionId = Uuid::now_v7().into();
3106 let mut query = MessageQuery::new(session_id);
3107
3108 collected.apply_message_filters(&mut query);
3109
3110 assert_eq!(query.filters.len(), 3);
3112 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3113 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3114 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3115 }
3116
3117 #[test]
3118 fn test_capability_without_message_filter_returns_none() {
3119 let registry = CapabilityRegistry::with_builtins();
3120
3121 let noop = registry.get("noop").unwrap();
3122 assert!(noop.message_filter_provider().is_none());
3123
3124 let current_time = registry.get("current_time").unwrap();
3125 assert!(current_time.message_filter_provider().is_none());
3126 }
3127
3128 #[tokio::test]
3129 async fn test_collect_capabilities_preserves_config_for_filter_provider() {
3130 let mut registry = CapabilityRegistry::new();
3131 registry.register(FilterTestCapability { priority: 0 });
3132
3133 let test_config = serde_json::json!({
3134 "search": "custom_search",
3135 "extra_field": 42
3136 });
3137
3138 let configs = vec![AgentCapabilityConfig {
3139 capability_ref: CapabilityId::new("filter_test"),
3140 config: test_config.clone(),
3141 }];
3142
3143 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3144
3145 assert_eq!(collected.message_filter_providers.len(), 1);
3147 let (_, stored_config) = &collected.message_filter_providers[0];
3148 assert_eq!(*stored_config, test_config);
3149 }
3150
3151 #[test]
3156 fn test_collect_message_filters_only_collects_filters() {
3157 let mut registry = CapabilityRegistry::new();
3158 registry.register(FilterTestCapability { priority: 0 });
3159
3160 let configs = vec![AgentCapabilityConfig {
3161 capability_ref: CapabilityId::new("filter_test"),
3162 config: serde_json::json!({ "search": "test_query" }),
3163 }];
3164
3165 let collected = collect_message_filters_only(&configs, ®istry);
3166
3167 let session_id: SessionId = Uuid::now_v7().into();
3168 let mut query = MessageQuery::new(session_id);
3169 collected.apply_message_filters(&mut query);
3170
3171 assert_eq!(query.filters.len(), 1);
3172 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3173 }
3174
3175 #[test]
3176 fn test_collect_message_filters_only_skips_unknown_capabilities() {
3177 let registry = CapabilityRegistry::new();
3178
3179 let configs = vec![AgentCapabilityConfig {
3180 capability_ref: CapabilityId::new("nonexistent"),
3181 config: serde_json::json!({}),
3182 }];
3183
3184 let collected = collect_message_filters_only(&configs, ®istry);
3185 assert!(collected.message_filter_providers.is_empty());
3186 }
3187
3188 #[test]
3189 fn test_collect_message_filters_only_preserves_priority_order() {
3190 struct PriorityFilterCap {
3191 id: &'static str,
3192 search_term: &'static str,
3193 priority: i32,
3194 }
3195
3196 struct PriorityFilterProvider {
3197 search_term: &'static str,
3198 priority: i32,
3199 }
3200
3201 impl Capability for PriorityFilterCap {
3202 fn id(&self) -> &str {
3203 self.id
3204 }
3205 fn name(&self) -> &str {
3206 self.id
3207 }
3208 fn description(&self) -> &str {
3209 "priority test"
3210 }
3211 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3212 Some(Arc::new(PriorityFilterProvider {
3213 search_term: self.search_term,
3214 priority: self.priority,
3215 }))
3216 }
3217 }
3218
3219 impl MessageFilterProvider for PriorityFilterProvider {
3220 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3221 query
3222 .filters
3223 .push(MessageFilter::Search(self.search_term.to_string()));
3224 }
3225 fn priority(&self) -> i32 {
3226 self.priority
3227 }
3228 }
3229
3230 let mut registry = CapabilityRegistry::new();
3231 registry.register(PriorityFilterCap {
3232 id: "gamma",
3233 search_term: "gamma",
3234 priority: 10,
3235 });
3236 registry.register(PriorityFilterCap {
3237 id: "alpha",
3238 search_term: "alpha",
3239 priority: 5,
3240 });
3241 registry.register(PriorityFilterCap {
3242 id: "beta",
3243 search_term: "beta",
3244 priority: 1,
3245 });
3246
3247 let configs = vec![
3248 AgentCapabilityConfig {
3249 capability_ref: CapabilityId::new("gamma"),
3250 config: serde_json::json!({}),
3251 },
3252 AgentCapabilityConfig {
3253 capability_ref: CapabilityId::new("alpha"),
3254 config: serde_json::json!({}),
3255 },
3256 AgentCapabilityConfig {
3257 capability_ref: CapabilityId::new("beta"),
3258 config: serde_json::json!({}),
3259 },
3260 ];
3261
3262 let collected = collect_message_filters_only(&configs, ®istry);
3263
3264 let session_id: SessionId = Uuid::now_v7().into();
3265 let mut query = MessageQuery::new(session_id);
3266 collected.apply_message_filters(&mut query);
3267
3268 assert_eq!(query.filters.len(), 3);
3270 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3271 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3272 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3273 }
3274
3275 #[test]
3276 fn test_collect_message_filters_only_post_load_invoked() {
3277 use crate::message::Message;
3278
3279 struct PostLoadCap;
3280 struct PostLoadProvider;
3281
3282 impl Capability for PostLoadCap {
3283 fn id(&self) -> &str {
3284 "post_load_test"
3285 }
3286 fn name(&self) -> &str {
3287 "PostLoad Test"
3288 }
3289 fn description(&self) -> &str {
3290 "test"
3291 }
3292 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3293 Some(Arc::new(PostLoadProvider))
3294 }
3295 }
3296
3297 impl MessageFilterProvider for PostLoadProvider {
3298 fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
3299 fn priority(&self) -> i32 {
3300 0
3301 }
3302 fn post_load(&self, messages: &mut Vec<Message>, _config: &serde_json::Value) {
3303 messages.reverse();
3305 }
3306 }
3307
3308 let mut registry = CapabilityRegistry::new();
3309 registry.register(PostLoadCap);
3310
3311 let configs = vec![AgentCapabilityConfig {
3312 capability_ref: CapabilityId::new("post_load_test"),
3313 config: serde_json::json!({}),
3314 }];
3315
3316 let collected = collect_message_filters_only(&configs, ®istry);
3317
3318 let mut messages = vec![Message::user("first"), Message::user("second")];
3319 collected.apply_post_load_filters(&mut messages);
3320
3321 assert_eq!(messages[0].text(), Some("second"));
3323 assert_eq!(messages[1].text(), Some("first"));
3324 }
3325
3326 #[test]
3327 fn test_collect_model_view_providers_respects_compaction_capability_boundary() {
3328 use crate::tool_types::ToolCall;
3329
3330 fn tool_heavy_messages() -> Vec<Message> {
3331 let mut messages = vec![Message::user("inspect files repeatedly")];
3332 for index in 0..9 {
3333 let call_id = format!("call_{index}");
3334 messages.push(Message::assistant_with_tools(
3335 "",
3336 vec![ToolCall {
3337 id: call_id.clone(),
3338 name: "read_file".to_string(),
3339 arguments: serde_json::json!({"path": "/workspace/src/lib.rs"}),
3340 }],
3341 ));
3342 messages.push(Message::tool_result(
3343 call_id,
3344 Some(serde_json::json!({
3345 "path": "/workspace/src/lib.rs",
3346 "content": format!("{}{}", "large file line\n".repeat(1000), index),
3347 "total_lines": 1000,
3348 "lines_shown": {"start": 1, "end": 1000},
3349 "truncated": false
3350 })),
3351 None,
3352 ));
3353 }
3354 messages
3355 }
3356
3357 fn first_tool_result_is_masked(messages: &[Message]) -> bool {
3358 messages[2]
3359 .tool_result_content()
3360 .and_then(|result| result.result.as_ref())
3361 .and_then(|result| result.get("masked"))
3362 .and_then(|masked| masked.as_bool())
3363 .unwrap_or(false)
3364 }
3365
3366 let mut registry = CapabilityRegistry::new();
3367 registry.register(CompactionCapability);
3368 let context = ModelViewContext {
3369 session_id: SessionId::new(),
3370 prior_usage: None,
3371 };
3372
3373 let no_compaction = collect_model_view_providers(&[], ®istry);
3374 let unmasked = no_compaction.apply_model_view(tool_heavy_messages(), &context);
3375 assert!(!first_tool_result_is_masked(&unmasked));
3376
3377 let compaction = collect_model_view_providers(
3378 &[AgentCapabilityConfig {
3379 capability_ref: CapabilityId::new(COMPACTION_CAPABILITY_ID),
3380 config: serde_json::json!({}),
3381 }],
3382 ®istry,
3383 );
3384 let masked = compaction.apply_model_view(tool_heavy_messages(), &context);
3385 assert!(first_tool_result_is_masked(&masked));
3386 let last_tool = masked.last().unwrap().tool_result_content().unwrap();
3387 assert!(last_tool.result.as_ref().unwrap().get("content").is_some());
3388 }
3389
3390 #[tokio::test]
3400 async fn test_virtual_bash_capability_produces_bash_tool() {
3401 let registry = CapabilityRegistry::with_builtins();
3402 let collected =
3403 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3404
3405 let tool_names: Vec<&str> = collected
3406 .tool_definitions
3407 .iter()
3408 .map(|t| t.name())
3409 .collect();
3410 assert!(
3411 tool_names.contains(&"bash"),
3412 "virtual_bash capability must produce 'bash' tool, got: {:?}",
3413 tool_names
3414 );
3415 assert!(
3416 !collected.tools.is_empty(),
3417 "virtual_bash must provide tool implementations"
3418 );
3419 }
3420
3421 #[tokio::test]
3422 async fn test_generic_harness_capability_set_produces_bash_tool() {
3423 let generic_harness_caps = vec![
3426 "session_file_system".to_string(),
3427 "virtual_bash".to_string(),
3428 "web_fetch".to_string(),
3429 "session_storage".to_string(),
3430 "session".to_string(),
3431 "agent_instructions".to_string(),
3432 "skills".to_string(),
3433 "infinity_context".to_string(),
3434 "openai_tool_search".to_string(),
3435 ];
3436
3437 let registry = CapabilityRegistry::with_builtins();
3438 let collected = collect_capabilities(&generic_harness_caps, ®istry, &test_ctx()).await;
3439
3440 let tool_names: Vec<&str> = collected
3441 .tool_definitions
3442 .iter()
3443 .map(|t| t.name())
3444 .collect();
3445 assert!(
3446 tool_names.contains(&"bash"),
3447 "Generic Harness capabilities must produce 'bash' tool, got: {:?}",
3448 tool_names
3449 );
3450 }
3451
3452 #[tokio::test]
3453 async fn test_collect_capabilities_tool_count_matches_definitions() {
3454 let registry = CapabilityRegistry::with_builtins();
3457 let collected =
3458 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3459
3460 assert_eq!(
3461 collected.tools.len(),
3462 collected.tool_definitions.len(),
3463 "tool implementations ({}) must match tool definitions ({})",
3464 collected.tools.len(),
3465 collected.tool_definitions.len(),
3466 );
3467 }
3468
3469 #[tokio::test]
3473 async fn test_collect_capabilities_resolves_dependencies() {
3474 let registry = CapabilityRegistry::with_builtins();
3477 let collected =
3478 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
3479
3480 assert!(
3482 collected
3483 .applied_ids
3484 .iter()
3485 .any(|id| id == "session_file_system"),
3486 "collect_capabilities must apply session_file_system as a dependency; applied_ids: {:?}",
3487 collected.applied_ids
3488 );
3489
3490 let tool_names: Vec<&str> = collected
3491 .tool_definitions
3492 .iter()
3493 .map(|t| t.name())
3494 .collect();
3495
3496 assert!(
3498 tool_names.contains(&"read_file") && tool_names.contains(&"write_file"),
3499 "collect_capabilities must resolve dependencies and include dependency tools, got: {:?}",
3500 tool_names
3501 );
3502
3503 assert_eq!(
3505 collected.tools.len(),
3506 collected.tool_definitions.len(),
3507 "dependency-added tools must have implementations, not just definitions"
3508 );
3509 }
3510
3511 #[test]
3512 fn test_defaults_do_not_include_bash() {
3513 let registry = crate::ToolRegistry::with_defaults();
3516 assert!(
3517 !registry.has("bash"),
3518 "with_defaults() must not include 'bash' — it comes from virtual_bash capability"
3519 );
3520 }
3521
3522 #[tokio::test]
3529 async fn test_background_execution_auto_activates_with_virtual_bash() {
3530 let registry = CapabilityRegistry::with_builtins();
3531 let collected =
3532 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3533
3534 let tool_names: Vec<&str> = collected
3535 .tool_definitions
3536 .iter()
3537 .map(|t| t.name())
3538 .collect();
3539 assert!(
3540 tool_names.contains(&"spawn_background"),
3541 "spawn_background must be auto-activated when virtual_bash (a \
3542 background-capable tool) is in the agent's capability set; got: {:?}",
3543 tool_names
3544 );
3545 assert!(
3546 collected
3547 .applied_ids
3548 .iter()
3549 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
3550 "background_execution must be in applied_ids when auto-activated; \
3551 got: {:?}",
3552 collected.applied_ids
3553 );
3554
3555 assert!(
3557 collected
3558 .tools
3559 .iter()
3560 .any(|t| t.name() == "spawn_background"),
3561 "spawn_background tool implementation must be present alongside the \
3562 definition (lockstep contract)"
3563 );
3564 }
3565
3566 #[tokio::test]
3569 async fn test_background_execution_does_not_auto_activate_without_hint() {
3570 let registry = CapabilityRegistry::with_builtins();
3571 let collected =
3573 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
3574
3575 let tool_names: Vec<&str> = collected
3576 .tool_definitions
3577 .iter()
3578 .map(|t| t.name())
3579 .collect();
3580 assert!(
3581 !tool_names.contains(&"spawn_background"),
3582 "spawn_background must NOT be activated without a background-capable \
3583 tool; got: {:?}",
3584 tool_names
3585 );
3586 assert!(
3587 !collected
3588 .applied_ids
3589 .iter()
3590 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
3591 "background_execution must not appear in applied_ids when no \
3592 background-capable tool is present; got: {:?}",
3593 collected.applied_ids
3594 );
3595 }
3596
3597 #[tokio::test]
3601 async fn test_background_execution_explicit_selection_is_idempotent() {
3602 let registry = CapabilityRegistry::with_builtins();
3603 let collected = collect_capabilities(
3604 &[
3605 "virtual_bash".to_string(),
3606 BACKGROUND_EXECUTION_CAPABILITY_ID.to_string(),
3607 ],
3608 ®istry,
3609 &test_ctx(),
3610 )
3611 .await;
3612
3613 let spawn_background_count = collected
3614 .tool_definitions
3615 .iter()
3616 .filter(|t| t.name() == "spawn_background")
3617 .count();
3618 assert_eq!(
3619 spawn_background_count, 1,
3620 "spawn_background must appear exactly once even when \
3621 background_execution is selected explicitly alongside a \
3622 background-capable tool"
3623 );
3624 let applied_count = collected
3625 .applied_ids
3626 .iter()
3627 .filter(|id| id.as_str() == BACKGROUND_EXECUTION_CAPABILITY_ID)
3628 .count();
3629 assert_eq!(
3630 applied_count, 1,
3631 "background_execution must appear exactly once in applied_ids"
3632 );
3633 }
3634
3635 #[test]
3640 fn test_defaults_do_not_include_spawn_background() {
3641 let registry = crate::ToolRegistry::with_defaults();
3642 assert!(
3643 !registry.has("spawn_background"),
3644 "with_defaults() must not include 'spawn_background' — it comes \
3645 from the background_execution capability (EVE-501)"
3646 );
3647 }
3648
3649 #[test]
3654 fn test_capability_features_default_empty() {
3655 let registry = CapabilityRegistry::with_builtins();
3656
3657 let noop = registry.get("noop").unwrap();
3659 assert!(noop.features().is_empty());
3660
3661 let current_time = registry.get("current_time").unwrap();
3662 assert!(current_time.features().is_empty());
3663 }
3664
3665 #[test]
3666 fn test_file_system_capability_features() {
3667 let registry = CapabilityRegistry::with_builtins();
3668
3669 let fs = registry.get("session_file_system").unwrap();
3670 assert_eq!(fs.features(), vec!["file_system"]);
3671 }
3672
3673 #[test]
3674 fn test_virtual_bash_capability_features() {
3675 let registry = CapabilityRegistry::with_builtins();
3676
3677 let bash = registry.get("virtual_bash").unwrap();
3678 assert_eq!(bash.features(), vec!["file_system"]);
3679 }
3680
3681 #[test]
3682 fn test_session_storage_capability_features() {
3683 let registry = CapabilityRegistry::with_builtins();
3684
3685 let storage = registry.get("session_storage").unwrap();
3686 let features = storage.features();
3687 assert!(features.contains(&"secrets"));
3688 assert!(features.contains(&"key_value"));
3689 }
3690
3691 #[test]
3692 fn test_session_schedule_capability_features() {
3693 let registry = CapabilityRegistry::with_builtins();
3694
3695 let schedule = registry.get("session_schedule").unwrap();
3696 assert_eq!(schedule.features(), vec!["schedules"]);
3697 }
3698
3699 #[test]
3700 fn test_session_sql_database_capability_features() {
3701 let registry = CapabilityRegistry::with_builtins();
3702
3703 let sql = registry.get("session_sql_database").unwrap();
3704 assert_eq!(sql.features(), vec!["sql_database"]);
3705 }
3706
3707 #[test]
3708 fn test_sample_data_capability_features() {
3709 let registry = CapabilityRegistry::with_builtins();
3710
3711 let sample = registry.get("sample_data").unwrap();
3712 assert_eq!(sample.features(), vec!["file_system"]);
3713 }
3714
3715 #[test]
3716 fn test_compute_features_empty() {
3717 let registry = CapabilityRegistry::with_builtins();
3718
3719 let features = compute_features(&[], ®istry);
3720 assert!(features.is_empty());
3721 }
3722
3723 #[test]
3724 fn test_compute_features_single_capability() {
3725 let registry = CapabilityRegistry::with_builtins();
3726
3727 let features = compute_features(&["session_schedule".to_string()], ®istry);
3728 assert_eq!(features, vec!["schedules"]);
3729 }
3730
3731 #[test]
3732 fn test_compute_features_multiple_capabilities() {
3733 let registry = CapabilityRegistry::with_builtins();
3734
3735 let features = compute_features(
3736 &[
3737 "session_file_system".to_string(),
3738 "session_storage".to_string(),
3739 "session_schedule".to_string(),
3740 ],
3741 ®istry,
3742 );
3743 assert!(features.contains(&"file_system".to_string()));
3744 assert!(features.contains(&"secrets".to_string()));
3745 assert!(features.contains(&"key_value".to_string()));
3746 assert!(features.contains(&"schedules".to_string()));
3747 }
3748
3749 #[test]
3750 fn test_compute_features_deduplicates() {
3751 let registry = CapabilityRegistry::with_builtins();
3752
3753 let features = compute_features(
3755 &[
3756 "session_file_system".to_string(),
3757 "virtual_bash".to_string(),
3758 ],
3759 ®istry,
3760 );
3761 let file_system_count = features.iter().filter(|f| *f == "file_system").count();
3762 assert_eq!(file_system_count, 1, "file_system should appear only once");
3763 }
3764
3765 #[test]
3766 fn test_compute_features_includes_dependency_features() {
3767 let registry = CapabilityRegistry::with_builtins();
3768
3769 let features = compute_features(&["virtual_bash".to_string()], ®istry);
3771 assert!(features.contains(&"file_system".to_string()));
3772 }
3773
3774 #[test]
3775 fn test_compute_features_generic_harness_set() {
3776 let registry = CapabilityRegistry::with_builtins();
3777
3778 let features = compute_features(
3780 &[
3781 "session_file_system".to_string(),
3782 "virtual_bash".to_string(),
3783 "session_storage".to_string(),
3784 "session".to_string(),
3785 "session_schedule".to_string(),
3786 ],
3787 ®istry,
3788 );
3789 assert!(features.contains(&"file_system".to_string()));
3790 assert!(features.contains(&"secrets".to_string()));
3791 assert!(features.contains(&"key_value".to_string()));
3792 assert!(features.contains(&"schedules".to_string()));
3793 }
3794
3795 #[test]
3796 fn test_compute_features_unknown_capability_ignored() {
3797 let registry = CapabilityRegistry::with_builtins();
3798
3799 let features = compute_features(
3800 &["unknown_cap".to_string(), "session_schedule".to_string()],
3801 ®istry,
3802 );
3803 assert_eq!(features, vec!["schedules"]);
3804 }
3805
3806 #[test]
3807 fn test_risk_level_ordering() {
3808 assert!(RiskLevel::Low < RiskLevel::Medium);
3809 assert!(RiskLevel::Medium < RiskLevel::High);
3810 }
3811
3812 #[test]
3813 fn test_risk_level_serde_roundtrip() {
3814 let high = RiskLevel::High;
3815 let json = serde_json::to_string(&high).unwrap();
3816 assert_eq!(json, "\"high\"");
3817 let back: RiskLevel = serde_json::from_str(&json).unwrap();
3818 assert_eq!(back, RiskLevel::High);
3819 }
3820
3821 #[test]
3822 fn test_capability_risk_levels() {
3823 let registry = CapabilityRegistry::with_builtins();
3824
3825 let bash = registry.get("virtual_bash").unwrap();
3827 assert_eq!(bash.risk_level(), RiskLevel::High);
3828
3829 let fetch = registry.get("web_fetch").unwrap();
3831 assert_eq!(fetch.risk_level(), RiskLevel::High);
3832
3833 let noop = registry.get("noop").unwrap();
3835 assert_eq!(noop.risk_level(), RiskLevel::Low);
3836 }
3837
3838 #[tokio::test]
3843 async fn test_apply_capabilities_openai_tool_search() {
3844 let registry = CapabilityRegistry::with_builtins();
3845 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
3846
3847 let applied = apply_capabilities(
3848 base_runtime_agent.clone(),
3849 &["openai_tool_search".to_string()],
3850 ®istry,
3851 &test_ctx(),
3852 )
3853 .await;
3854
3855 assert_eq!(
3857 applied.runtime_agent.system_prompt,
3858 base_runtime_agent.system_prompt
3859 );
3860 assert!(applied.tool_registry.is_empty());
3861 assert_eq!(applied.applied_ids, vec!["openai_tool_search"]);
3862
3863 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
3865 assert!(ts.enabled);
3866 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
3867 }
3868
3869 #[tokio::test]
3870 async fn test_apply_capabilities_openai_tool_search_with_other_capabilities() {
3871 let registry = CapabilityRegistry::with_builtins();
3872 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
3873
3874 let applied = apply_capabilities(
3875 base_runtime_agent,
3876 &[
3877 "current_time".to_string(),
3878 "openai_tool_search".to_string(),
3879 "test_math".to_string(),
3880 ],
3881 ®istry,
3882 &test_ctx(),
3883 )
3884 .await;
3885
3886 assert!(applied.tool_registry.has("get_current_time"));
3888 assert!(applied.tool_registry.has("add"));
3889 assert!(applied.tool_registry.has("subtract"));
3890 assert!(applied.tool_registry.has("multiply"));
3891 assert!(applied.tool_registry.has("divide"));
3892
3893 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
3895 assert!(ts.enabled);
3896 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
3897 }
3898
3899 #[tokio::test]
3900 async fn test_collect_capabilities_tool_search_custom_threshold() {
3901 let registry = CapabilityRegistry::with_builtins();
3902
3903 let configs = vec![AgentCapabilityConfig {
3904 capability_ref: CapabilityId::new("openai_tool_search"),
3905 config: serde_json::json!({"threshold": 5}),
3906 }];
3907
3908 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3909
3910 let ts = collected.tool_search.as_ref().unwrap();
3911 assert!(ts.enabled);
3912 assert_eq!(ts.threshold, 5);
3913 }
3914
3915 #[tokio::test]
3916 async fn test_collect_capabilities_no_tool_search_without_capability() {
3917 let registry = CapabilityRegistry::with_builtins();
3918
3919 let configs = vec![AgentCapabilityConfig {
3920 capability_ref: CapabilityId::new("current_time"),
3921 config: serde_json::json!({}),
3922 }];
3923
3924 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3925
3926 assert!(collected.tool_search.is_none());
3927 }
3928
3929 #[tokio::test]
3930 async fn test_collect_capabilities_tool_search_category_propagation() {
3931 let registry = CapabilityRegistry::with_builtins();
3932
3933 let configs = vec![
3935 AgentCapabilityConfig {
3936 capability_ref: CapabilityId::new("test_math"),
3937 config: serde_json::json!({}),
3938 },
3939 AgentCapabilityConfig {
3940 capability_ref: CapabilityId::new("openai_tool_search"),
3941 config: serde_json::json!({}),
3942 },
3943 ];
3944
3945 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3946
3947 assert!(collected.tool_search.is_some());
3949
3950 for tool_def in &collected.tool_definitions {
3952 if ["add", "subtract", "multiply", "divide"].contains(&tool_def.name()) {
3954 assert!(
3955 tool_def.category().is_some(),
3956 "Tool {} should have a category from its capability",
3957 tool_def.name()
3958 );
3959 }
3960 }
3961 }
3962
3963 #[tokio::test]
3964 async fn test_apply_capabilities_prompt_caching() {
3965 let registry = CapabilityRegistry::with_builtins();
3966 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
3967
3968 let applied = apply_capabilities(
3969 base_runtime_agent.clone(),
3970 &["prompt_caching".to_string()],
3971 ®istry,
3972 &test_ctx(),
3973 )
3974 .await;
3975
3976 assert_eq!(
3977 applied.runtime_agent.system_prompt,
3978 base_runtime_agent.system_prompt
3979 );
3980 assert!(applied.tool_registry.is_empty());
3981 assert_eq!(applied.applied_ids, vec!["prompt_caching"]);
3982
3983 let prompt_cache = applied.runtime_agent.prompt_cache.as_ref().unwrap();
3984 assert!(prompt_cache.enabled);
3985 assert_eq!(
3986 prompt_cache.strategy,
3987 crate::llm_driver_registry::PromptCacheStrategy::Auto
3988 );
3989 assert!(prompt_cache.gemini_cached_content.is_none());
3990 }
3991
3992 #[tokio::test]
3993 async fn test_collect_capabilities_prompt_caching_custom_strategy() {
3994 let registry = CapabilityRegistry::with_builtins();
3995
3996 let configs = vec![AgentCapabilityConfig {
3997 capability_ref: CapabilityId::new("prompt_caching"),
3998 config: serde_json::json!({"strategy": "auto"}),
3999 }];
4000
4001 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4002
4003 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
4004 assert!(prompt_cache.enabled);
4005 assert_eq!(
4006 prompt_cache.strategy,
4007 crate::llm_driver_registry::PromptCacheStrategy::Auto
4008 );
4009 assert!(prompt_cache.gemini_cached_content.is_none());
4010 }
4011
4012 #[tokio::test]
4013 async fn test_collect_capabilities_prompt_caching_gemini_cached_content() {
4014 let registry = CapabilityRegistry::with_builtins();
4015
4016 let configs = vec![AgentCapabilityConfig {
4017 capability_ref: CapabilityId::new("prompt_caching"),
4018 config: serde_json::json!({
4019 "strategy": "auto",
4020 "gemini_cached_content": "cachedContents/demo-cache"
4021 }),
4022 }];
4023
4024 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4025
4026 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
4027 assert_eq!(
4028 prompt_cache.gemini_cached_content.as_deref(),
4029 Some("cachedContents/demo-cache")
4030 );
4031 }
4032
4033 struct SkillContributingCapability;
4038
4039 impl Capability for SkillContributingCapability {
4040 fn id(&self) -> &str {
4041 "contributes_skills"
4042 }
4043 fn name(&self) -> &str {
4044 "Contributes Skills"
4045 }
4046 fn description(&self) -> &str {
4047 "Test capability that contributes skills."
4048 }
4049 fn contribute_skills(&self) -> Vec<SkillContribution> {
4050 vec![
4051 SkillContribution::new("alpha-skill", "Alpha skill desc", "# Alpha\nDo alpha.")
4052 .with_files(vec![(
4053 "scripts/a.sh".to_string(),
4054 "#!/bin/sh\necho a\n".to_string(),
4055 )]),
4056 SkillContribution::new("beta-skill", "Beta skill desc", "# Beta\nDo beta.")
4057 .with_user_invocable(false),
4058 ]
4059 }
4060 }
4061
4062 fn skill_md_from_entries(entries: &HashMap<String, MountEntry>) -> &str {
4063 match &entries.get("SKILL.md").expect("SKILL.md missing").source {
4064 MountSource::InlineFile { content, .. } => content.as_str(),
4065 _ => panic!("Expected InlineFile for SKILL.md"),
4066 }
4067 }
4068
4069 #[tokio::test]
4070 async fn test_contribute_skills_normalized_to_mounts() {
4071 let mut registry = CapabilityRegistry::new();
4072 registry.register(SkillContributingCapability);
4073
4074 let configs = vec![AgentCapabilityConfig {
4075 capability_ref: CapabilityId::new("contributes_skills"),
4076 config: serde_json::json!({}),
4077 }];
4078
4079 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4080
4081 let skill_mounts: Vec<_> = collected
4082 .mounts
4083 .iter()
4084 .filter(|m| m.path.starts_with("/.agents/skills/"))
4085 .collect();
4086 assert_eq!(skill_mounts.len(), 2);
4087
4088 for m in &skill_mounts {
4091 assert!(m.is_readonly());
4092 assert_eq!(m.capability_id, "contributes_skills");
4093 }
4094
4095 let alpha = skill_mounts
4096 .iter()
4097 .find(|m| m.path == "/.agents/skills/alpha-skill")
4098 .expect("alpha-skill mount missing");
4099 match &alpha.source {
4100 MountSource::InlineDirectory { entries } => {
4101 assert!(entries.contains_key("SKILL.md"));
4102 assert!(entries.contains_key("scripts/a.sh"));
4103 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
4104 assert_eq!(parsed.name, "alpha-skill");
4105 assert!(parsed.user_invocable);
4106 }
4107 _ => panic!("Expected InlineDirectory"),
4108 }
4109
4110 let beta = skill_mounts
4111 .iter()
4112 .find(|m| m.path == "/.agents/skills/beta-skill")
4113 .expect("beta-skill mount missing");
4114 match &beta.source {
4115 MountSource::InlineDirectory { entries } => {
4116 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
4117 assert!(!parsed.user_invocable);
4118 }
4119 _ => panic!("Expected InlineDirectory"),
4120 }
4121 }
4122
4123 #[tokio::test]
4124 async fn test_contribute_skills_default_empty() {
4125 let mut registry = CapabilityRegistry::new();
4128 registry.register(FilterTestCapability { priority: 0 });
4129
4130 let configs = vec![AgentCapabilityConfig {
4131 capability_ref: CapabilityId::new("filter_test"),
4132 config: serde_json::json!({}),
4133 }];
4134
4135 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4136 assert!(
4137 collected
4138 .mounts
4139 .iter()
4140 .all(|m| !m.path.starts_with("/.agents/skills/"))
4141 );
4142 }
4143}