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 btw;
91mod budgeting;
92pub mod compaction;
93mod current_time;
94mod data_knowledge;
95mod declarative;
96mod fake_aws;
97mod fake_crm;
98mod fake_financial;
99mod fake_warehouse;
100mod file_system;
101mod human_intent;
102mod infinity_context;
103mod knowledge_base;
104mod loop_detection;
105pub mod mcp;
106mod noop;
107mod openai_tool_search;
108#[cfg(feature = "ui-capabilities")]
109mod openui;
110mod parallel;
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 virtual_bash;
131mod web_fetch;
132mod workspace_volumes;
133
134pub use a2a_delegation::{
136 A2A_AGENT_DELEGATION_CAPABILITY_ID, A2aAgentDelegationCapability, CancelAgentTool,
137 GetAgentRunsTool, MessageAgentTool, SpawnAgentTool, WaitAgentTool,
138};
139#[cfg(feature = "ui-capabilities")]
140pub use a2ui::{A2UI_CAPABILITY_ID, A2UiCapability};
141pub use agent_handoff::{
142 AGENT_HANDOFF_CAPABILITY_ID, AgentHandoffCapability, GetAgentHandoffsTool,
143 MessageAgentHandoffTool, StartAgentHandoffTool,
144};
145pub use agent_instructions::{
146 AGENT_INSTRUCTIONS_CAPABILITY_ID, AGENTS_MD_PATH, AgentInstructionsCapability,
147 AgentInstructionsConfig, DEFAULT_AGENT_INSTRUCTIONS_FILE, MAX_AGENT_INSTRUCTIONS_FILES,
148 MAX_AGENTS_MD_SIZE, format_agents_md_content, format_instruction_file_content,
149};
150pub use attach_skill::{
151 AttachSkillCapability, SKILL_CAPABILITY_PREFIX, SKILLS_DISCOVERY_PATH, SkillContribution,
152 SkillInstructions, SkillMeta, SkillSource, discover_skills_from_entries, is_skill_capability,
153 parse_skill_capability_id, reconstruct_skill_md, skill_capability_id,
154};
155pub use btw::{BTW_CAPABILITY_ID, BtwCapability};
156pub use budgeting::BudgetingCapability;
157pub use compaction::{
158 COMPACTION_CAPABILITY_ID, CompactionCapability, CompactionConfig, CompactionStep,
159 CompactionStrategy, CostControlConfig, CostControlMaskingResult, HierarchicalMemoryConfig,
160 MaskingSummaryFormat, MemoryTier, ObservationMaskingConfig, ObservationMaskingResult,
161 SessionCompactionMetrics, SummarizationConfig, aggressive_trim, apply_cost_control_masking,
162 apply_hierarchical_memory, apply_observation_masking, build_model_view_messages,
163 build_summarization_prompt, build_summary_message, classify_memory_tiers, estimate_tokens,
164 estimate_total_tokens, format_messages_for_summarization, should_compact_proactively,
165};
166pub use current_time::{CurrentTimeCapability, GetCurrentTimeTool};
167pub use data_knowledge::DataKnowledgeCapability;
168pub use declarative::{
169 DECLARATIVE_CAPABILITY_PREFIX, DeclarativeCapabilityDefinition, DeclarativeCapabilityFile,
170 DeclarativeCapabilitySkill, declarative_capability_id, declarative_capability_info,
171 hydrate_declarative_capability_config, is_declarative_capability,
172 parse_declarative_capability_id, validate_declarative_capability_definition,
173};
174pub use fake_aws::{
175 AwsCreateEc2InstanceTool, AwsCreateIamUserTool, AwsCreateRdsDatabaseTool,
176 AwsCreateS3BucketTool, AwsGetCloudWatchMetricsTool, AwsListEc2InstancesTool,
177 AwsListIamUsersTool, AwsListRdsDatabasesTool, AwsListS3BucketsTool, AwsListSecurityGroupsTool,
178 AwsStopEc2InstanceTool, FakeAwsCapability,
179};
180pub use fake_crm::{
181 CrmAddInteractionTool, CrmCreateCustomerTool, CrmCreateTicketTool, CrmGetCustomerTool,
182 CrmListCustomersTool, CrmListTicketsTool, CrmSearchCustomersTool, CrmUpdateTicketTool,
183 FakeCrmCapability,
184};
185pub use fake_financial::{
186 FakeFinancialCapability, FinanceCreateBudgetTool, FinanceCreateTransactionTool,
187 FinanceForecastCashFlowTool, FinanceGetBalanceTool, FinanceGetExpenseReportTool,
188 FinanceGetRevenueReportTool, FinanceListBudgetsTool, FinanceListTransactionsTool,
189};
190pub use fake_warehouse::{
191 FakeWarehouseCapability, WarehouseCreateInvoiceTool, WarehouseCreateOrderTool,
192 WarehouseCreateShipmentTool, WarehouseGetInventoryTool, WarehouseInventoryReportTool,
193 WarehouseListOrdersTool, WarehouseListShipmentsTool, WarehouseProcessReturnTool,
194 WarehouseUpdateInventoryTool, WarehouseUpdateShipmentStatusTool,
195};
196pub use file_system::{
197 DeleteFileTool, EditFileTool, FileSystemCapability, GrepFilesTool, ListDirectoryTool,
198 ReadFileTool, StatFileTool, WriteFileTool,
199};
200pub use human_intent::{HUMAN_INTENT_CAPABILITY_ID, HumanIntentCapability};
201pub use infinity_context::{
202 INFINITY_CONTEXT_CAPABILITY_ID, InfinityContextCapability, QueryHistoryTool,
203};
204pub use knowledge_base::{
205 KNOWLEDGE_BASE_CAPABILITY_ID, KnowledgeBaseCapability, KnowledgeBaseConfig,
206 validate_knowledge_base_config,
207};
208pub use loop_detection::LoopDetectionCapability;
209pub use mcp::{
210 MCP_CAPABILITY_PREFIX, McpCapability, is_mcp_capability, mcp_capability_id,
211 parse_mcp_capability_id,
212};
213pub use noop::NoopCapability;
214pub use openai_tool_search::{
215 DEFAULT_TOOL_SEARCH_THRESHOLD, OPENAI_TOOL_SEARCH_CAPABILITY_ID, OpenAiToolSearchCapability,
216};
217#[cfg(feature = "ui-capabilities")]
218pub use openui::{OPENUI_CAPABILITY_ID, OpenUiCapability};
219pub use parallel::ParallelCapability;
220pub use persistent_memory::{
221 ForgetTool, MEMORY_CAPABILITY_ID, MemoryCapability, MemoryConfig, RecallTool, RememberTool,
222};
223pub use platform_management::{
224 ManageAgentsTool, ManageHarnessesTool, ManageSessionsTool, PlatformManagementCapability,
225 ReadAgentsTool, ReadCapabilitiesTool, ReadHarnessesTool, ReadSessionsTool,
226 SessionReadMessagesTool, SessionReadResponseTool, SessionSendMessageTool,
227};
228pub use prompt_caching::{PROMPT_CACHING_CAPABILITY_ID, PromptCachingCapability};
229pub use prompt_canary_guardrail::{
230 DEFAULT_REPLACEMENT as PROMPT_CANARY_DEFAULT_REPLACEMENT,
231 PROMPT_CANARY_GUARDRAIL_CAPABILITY_ID, PromptCanaryGuardrailCapability,
232 REASON_CODE_SYSTEM_PROMPT_LEAK,
233};
234pub use research::ResearchCapability;
235pub use sample_data::SampleDataCapability;
236pub use self_budget::{SELF_BUDGET_CAPABILITY_ID, SelfBudgetCapability};
237pub use session::{GetSessionInfoTool, SessionCapability, WriteSessionTitleTool};
238pub use session_sandbox::{
239 SESSION_SANDBOX_CAPABILITY_ID, SandboxExecTool, SandboxManageTool, SandboxReadFileTool,
240 SandboxStatusTool, SandboxWriteFileTool, SessionSandboxCapability,
241};
242pub use session_schedule::{
243 CancelScheduleTool, CreateScheduleTool, ListSchedulesTool, SESSION_SCHEDULE_CAPABILITY_ID,
244 SessionScheduleCapability,
245};
246pub use session_sql_database::{
247 SessionSqlDatabaseCapability, SqlExecuteTool, SqlQueryTool, SqlSchemaTool,
248};
249pub use session_storage::{KvStoreTool, SecretStoreTool, SessionStorageCapability};
250pub use skills::{SKILLS_CAPABILITY_ID, SkillsCapability};
251pub use stateless_todo_list::{StatelessTodoListCapability, WriteTodosTool};
252pub use subagents::SubagentCapability;
253pub use system_commands::{SYSTEM_COMMANDS_CAPABILITY_ID, SystemCommandsCapability};
255pub use test_math::{AddTool, DivideTool, MultiplyTool, SubtractTool, TestMathCapability};
256pub use test_weather::{GetForecastTool, GetWeatherTool, TestWeatherCapability};
257pub use tool_output_persistence::{PersistOutputHook, ToolOutputPersistenceCapability};
258pub use virtual_bash::{BashTool, VirtualBashCapability};
259pub use web_fetch::{
260 BotAuthPublicKey, WebFetchCapability, WebFetchTool, derive_bot_auth_public_key,
261};
262pub use workspace_volumes::{WORKSPACE_VOLUMES_CAPABILITY_ID, WorkspaceVolumesCapability};
263
264pub struct SystemPromptContext {
274 pub session_id: SessionId,
276 pub locale: Option<String>,
278 pub file_store: Option<Arc<dyn SessionFileSystem>>,
280}
281
282impl SystemPromptContext {
283 pub fn without_file_store(session_id: SessionId) -> Self {
285 Self {
286 session_id,
287 locale: None,
288 file_store: None,
289 }
290 }
291}
292
293#[async_trait]
340pub trait Capability: Send + Sync {
341 fn id(&self) -> &str;
343
344 fn name(&self) -> &str;
346
347 fn description(&self) -> &str;
349
350 fn status(&self) -> CapabilityStatus {
352 CapabilityStatus::Available
353 }
354
355 fn icon(&self) -> Option<&str> {
357 None
358 }
359
360 fn category(&self) -> Option<&str> {
362 None
363 }
364
365 fn system_prompt_addition(&self) -> Option<&str> {
385 None
386 }
387
388 async fn system_prompt_contribution(&self, _ctx: &SystemPromptContext) -> Option<String> {
400 self.system_prompt_addition().map(|addition| {
401 format!(
402 "<capability id=\"{}\">\n{}\n</capability>",
403 self.id(),
404 addition
405 )
406 })
407 }
408
409 fn system_prompt_preview(&self) -> Option<String> {
415 self.system_prompt_addition().map(|s| s.to_string())
416 }
417
418 fn tools(&self) -> Vec<Box<dyn Tool>> {
420 vec![]
421 }
422
423 fn tools_with_config(&self, _config: &serde_json::Value) -> Vec<Box<dyn Tool>> {
431 self.tools()
432 }
433
434 async fn system_prompt_contribution_with_config(
441 &self,
442 ctx: &SystemPromptContext,
443 _config: &serde_json::Value,
444 ) -> Option<String> {
445 self.system_prompt_contribution(ctx).await
446 }
447
448 fn tool_definitions(&self) -> Vec<ToolDefinition> {
451 self.tools().iter().map(|t| t.to_definition()).collect()
452 }
453
454 fn mounts(&self) -> Vec<MountPoint> {
462 vec![]
463 }
464
465 fn dependencies(&self) -> Vec<&'static str> {
474 vec![]
475 }
476
477 fn features(&self) -> Vec<&'static str> {
492 vec![]
493 }
494
495 fn config_schema(&self) -> Option<serde_json::Value> {
501 None
502 }
503
504 fn config_ui_schema(&self) -> Option<serde_json::Value> {
509 None
510 }
511
512 fn validate_config(&self, _config: &serde_json::Value) -> Result<(), String> {
518 Ok(())
519 }
520
521 fn mcp_servers(&self) -> ScopedMcpServers {
527 ScopedMcpServers::default()
528 }
529
530 fn mcp_servers_with_config(&self, _config: &serde_json::Value) -> ScopedMcpServers {
532 self.mcp_servers()
533 }
534
535 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
548 None
549 }
550
551 fn model_view_provider(&self) -> Option<Arc<dyn ModelViewProvider>> {
559 None
560 }
561
562 fn post_tool_exec_hooks(&self) -> Vec<Arc<dyn crate::atoms::PostToolExecHook>> {
570 vec![]
571 }
572
573 fn tool_definition_hooks(&self) -> Vec<Arc<dyn ToolDefinitionHook>> {
582 vec![]
583 }
584
585 fn tool_call_hooks(&self) -> Vec<Arc<dyn ToolCallHook>> {
593 vec![]
594 }
595
596 fn risk_level(&self) -> RiskLevel {
604 RiskLevel::Low
605 }
606
607 fn commands(&self) -> Vec<CommandDescriptor> {
615 vec![]
616 }
617
618 async fn execute_command(
632 &self,
633 request: &ExecuteCommandRequest,
634 _ctx: &CommandExecutionContext,
635 ) -> crate::error::Result<CommandResult> {
636 Err(crate::error::AgentLoopError::config(format!(
637 "capability {} declared command /{} but does not implement execute_command",
638 self.id(),
639 request.name,
640 )))
641 }
642
643 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
651 vec![]
652 }
653
654 fn contribute_skills(&self) -> Vec<SkillContribution> {
664 vec![]
665 }
666
667 fn output_guardrails(&self) -> Vec<Arc<dyn crate::output_guardrail::OutputGuardrail>> {
678 vec![]
679 }
680}
681
682pub trait ToolDefinitionHook: Send + Sync {
683 fn transform(&self, tools: Vec<ToolDefinition>) -> Vec<ToolDefinition>;
684}
685
686pub trait ToolCallHook: Send + Sync {
687 fn narration(
688 &self,
689 _tool_def: Option<&ToolDefinition>,
690 _tool_call: &ToolCall,
691 _phase: crate::tool_narration::ToolNarrationPhase,
692 _locale: Option<&str>,
693 ) -> Option<String> {
694 None
695 }
696
697 fn transform_for_execution(&self, tool_call: ToolCall) -> ToolCall {
698 tool_call
699 }
700}
701
702#[derive(
706 Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
707)]
708#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
709#[serde(rename_all = "lowercase")]
710pub enum RiskLevel {
711 Low,
713 Medium,
715 High,
717}
718
719#[derive(Debug, Clone, Serialize, Deserialize)]
725#[serde(rename_all = "snake_case")]
726pub enum BlueprintModel {
727 Fixed(String),
729 Default(String),
731 Inherit,
733}
734
735pub struct AgentBlueprint {
741 pub id: &'static str,
743 pub name: &'static str,
745 pub description: &'static str,
747 pub model: BlueprintModel,
749 pub system_prompt: &'static str,
751 pub tools: Vec<Box<dyn Tool>>,
753 pub max_turns: Option<usize>,
755 pub config_schema: Option<serde_json::Value>,
757}
758
759impl AgentBlueprint {
760 pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
762 self.tools.iter().map(|t| t.to_definition()).collect()
763 }
764}
765
766impl std::fmt::Debug for AgentBlueprint {
767 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
768 f.debug_struct("AgentBlueprint")
769 .field("id", &self.id)
770 .field("name", &self.name)
771 .field("model", &self.model)
772 .field("tool_count", &self.tools.len())
773 .field("max_turns", &self.max_turns)
774 .finish()
775 }
776}
777
778#[derive(Clone)]
805pub struct CapabilityRegistry {
806 capabilities: HashMap<String, Arc<dyn Capability>>,
807}
808
809impl CapabilityRegistry {
810 pub fn new() -> Self {
812 Self {
813 capabilities: HashMap::new(),
814 }
815 }
816
817 pub fn with_builtins() -> Self {
822 Self::with_builtins_for_grade(DeploymentGrade::from_env())
823 }
824
825 pub fn with_builtins_for_grade(grade: DeploymentGrade) -> Self {
830 let mut registry = Self::new();
831
832 registry.register(AgentInstructionsCapability);
834 registry.register(HumanIntentCapability);
835 registry.register(NoopCapability);
836 registry.register(CurrentTimeCapability);
837 registry.register(ResearchCapability);
838 registry.register(PlatformManagementCapability);
839 registry.register(FileSystemCapability);
840 registry.register(WorkspaceVolumesCapability);
841 registry.register(SessionStorageCapability);
842 registry.register(SessionCapability);
843 registry.register(SessionSqlDatabaseCapability);
844 registry.register(TestMathCapability);
845 registry.register(TestWeatherCapability);
846 registry.register(StatelessTodoListCapability);
847 registry.register(WebFetchCapability::from_env());
848 registry.register(VirtualBashCapability);
849 registry.register(SessionScheduleCapability);
850 registry.register(BtwCapability);
851 registry.register(InfinityContextCapability);
852 registry.register(budgeting::BudgetingCapability);
853 registry.register(SelfBudgetCapability);
854 registry.register(ParallelCapability);
855 registry.register(CompactionCapability);
856 registry.register(MemoryCapability);
857
858 registry.register(OpenAiToolSearchCapability::new());
860 registry.register(PromptCachingCapability::new());
861
862 registry.register(SkillsCapability);
864
865 registry.register(SubagentCapability);
867 registry.register(AgentHandoffCapability);
868 registry.register(A2aAgentDelegationCapability);
869
870 registry.register(SystemCommandsCapability);
872
873 registry.register(tool_output_persistence::ToolOutputPersistenceCapability);
875
876 registry.register(LoopDetectionCapability);
878
879 registry.register(PromptCanaryGuardrailCapability);
882
883 #[cfg(feature = "ui-capabilities")]
885 {
886 registry.register(OpenUiCapability);
887 registry.register(A2UiCapability);
888 }
889
890 registry.register(SampleDataCapability);
892
893 registry.register(DataKnowledgeCapability);
895
896 registry.register(KnowledgeBaseCapability);
898
899 registry.register(FakeWarehouseCapability);
901 registry.register(FakeAwsCapability);
902 registry.register(FakeCrmCapability);
903 registry.register(FakeFinancialCapability);
904
905 let internal_flags = crate::InternalFeatureFlags::from_env();
907 if internal_flags.session_sandbox {
908 registry.register(SessionSandboxCapability);
909 }
910 for plugin in inventory::iter::<IntegrationPlugin>() {
911 if (!plugin.experimental_only || grade.experimental_features_enabled())
912 && plugin
913 .feature_flag
914 .is_none_or(|f| internal_flags.is_enabled(f))
915 {
916 registry.register_boxed((plugin.factory)());
917 }
918 }
919
920 registry
921 }
922
923 pub fn register(&mut self, capability: impl Capability + 'static) {
925 self.capabilities
926 .insert(capability.id().to_string(), Arc::new(capability));
927 }
928
929 pub fn register_boxed(&mut self, capability: Box<dyn Capability>) {
931 self.capabilities
932 .insert(capability.id().to_string(), Arc::from(capability));
933 }
934
935 pub fn register_arc(&mut self, capability: Arc<dyn Capability>) {
937 self.capabilities
938 .insert(capability.id().to_string(), capability);
939 }
940
941 pub fn get(&self, id: &str) -> Option<&Arc<dyn Capability>> {
943 self.capabilities.get(id)
944 }
945
946 pub fn unregister(&mut self, id: &str) -> Option<Arc<dyn Capability>> {
948 self.capabilities.remove(id)
949 }
950
951 pub fn has(&self, id: &str) -> bool {
953 self.capabilities.contains_key(id)
954 }
955
956 pub fn list(&self) -> Vec<&Arc<dyn Capability>> {
958 self.capabilities.values().collect()
959 }
960
961 pub fn len(&self) -> usize {
963 self.capabilities.len()
964 }
965
966 pub fn is_empty(&self) -> bool {
968 self.capabilities.is_empty()
969 }
970
971 pub fn builder() -> CapabilityRegistryBuilder {
973 CapabilityRegistryBuilder::new()
974 }
975
976 pub fn blueprint(&self, id: &str) -> Option<AgentBlueprint> {
980 for cap in self.capabilities.values() {
981 for bp in cap.agent_blueprints() {
982 if bp.id == id {
983 return Some(bp);
984 }
985 }
986 }
987 None
988 }
989
990 pub fn blueprint_with_capability(&self, id: &str) -> Option<(String, AgentBlueprint)> {
994 for (capability_id, cap) in &self.capabilities {
995 for bp in cap.agent_blueprints() {
996 if bp.id == id {
997 return Some((capability_id.clone(), bp));
998 }
999 }
1000 }
1001 None
1002 }
1003
1004 pub fn all_blueprints(&self) -> Vec<AgentBlueprint> {
1006 self.capabilities
1007 .values()
1008 .flat_map(|cap| cap.agent_blueprints())
1009 .collect()
1010 }
1011}
1012
1013impl Default for CapabilityRegistry {
1014 fn default() -> Self {
1015 Self::with_builtins()
1016 }
1017}
1018
1019impl std::fmt::Debug for CapabilityRegistry {
1020 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1021 let ids: Vec<_> = self.capabilities.keys().collect();
1022 f.debug_struct("CapabilityRegistry")
1023 .field("capabilities", &ids)
1024 .finish()
1025 }
1026}
1027
1028pub struct CapabilityRegistryBuilder {
1030 registry: CapabilityRegistry,
1031}
1032
1033impl CapabilityRegistryBuilder {
1034 pub fn new() -> Self {
1036 Self {
1037 registry: CapabilityRegistry::new(),
1038 }
1039 }
1040
1041 pub fn with_builtins() -> Self {
1043 Self {
1044 registry: CapabilityRegistry::with_builtins(),
1045 }
1046 }
1047
1048 pub fn capability(mut self, capability: impl Capability + 'static) -> Self {
1050 self.registry.register(capability);
1051 self
1052 }
1053
1054 pub fn build(self) -> CapabilityRegistry {
1056 self.registry
1057 }
1058}
1059
1060impl Default for CapabilityRegistryBuilder {
1061 fn default() -> Self {
1062 Self::new()
1063 }
1064}
1065
1066pub struct ModelViewContext<'a> {
1072 pub session_id: SessionId,
1073 pub prior_usage: Option<&'a TokenUsage>,
1074}
1075
1076pub trait ModelViewProvider: Send + Sync {
1082 fn apply_model_view(
1083 &self,
1084 messages: Vec<Message>,
1085 config: &serde_json::Value,
1086 context: &ModelViewContext<'_>,
1087 ) -> Vec<Message>;
1088
1089 fn priority(&self) -> i32 {
1090 0
1091 }
1092}
1093
1094pub struct CollectedCapabilities {
1099 pub system_prompt_parts: Vec<String>,
1101 pub system_prompt_attributions: Vec<SystemPromptAttribution>,
1103 pub tools: Vec<Box<dyn Tool>>,
1105 pub tool_definitions: Vec<ToolDefinition>,
1107 pub mounts: Vec<MountPoint>,
1109 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1111 pub applied_ids: Vec<String>,
1113 pub tool_search: Option<crate::llm_driver_registry::ToolSearchConfig>,
1115 pub prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig>,
1117 pub tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>>,
1119 pub tool_call_hooks: Vec<Arc<dyn ToolCallHook>>,
1121 pub mcp_servers: ScopedMcpServers,
1123 }
1129
1130#[derive(Debug, Clone, PartialEq, Eq)]
1131pub struct SystemPromptAttribution {
1132 pub capability_id: String,
1133 pub content: String,
1134}
1135
1136impl CollectedCapabilities {
1137 pub fn system_prompt_prefix(&self) -> Option<String> {
1140 if self.system_prompt_parts.is_empty() {
1141 None
1142 } else {
1143 Some(self.system_prompt_parts.join("\n\n"))
1144 }
1145 }
1146
1147 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1151 for (provider, config) in &self.message_filter_providers {
1153 provider.apply_filters(query, config);
1154 }
1155 }
1156
1157 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1160 for (provider, config) in &self.message_filter_providers {
1161 provider.post_load(messages, config);
1162 }
1163 }
1164
1165 pub fn has_message_filters(&self) -> bool {
1167 !self.message_filter_providers.is_empty()
1168 }
1169}
1170
1171pub struct CollectedMessageFilters {
1178 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1180}
1181
1182pub struct CollectedModelViewProviders {
1184 pub model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)>,
1186}
1187
1188impl CollectedMessageFilters {
1194 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1196 for (provider, config) in &self.message_filter_providers {
1197 provider.apply_filters(query, config);
1198 }
1199 }
1200
1201 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1203 for (provider, config) in &self.message_filter_providers {
1204 provider.post_load(messages, config);
1205 }
1206 }
1207}
1208
1209impl CollectedModelViewProviders {
1210 pub fn apply_model_view(
1212 &self,
1213 mut messages: Vec<Message>,
1214 context: &ModelViewContext<'_>,
1215 ) -> Vec<Message> {
1216 for (provider, config) in &self.model_view_providers {
1217 messages = provider.apply_model_view(messages, config, context);
1218 }
1219 messages
1220 }
1221}
1222
1223pub fn collect_message_filters_only(
1229 capability_configs: &[AgentCapabilityConfig],
1230 registry: &CapabilityRegistry,
1231) -> CollectedMessageFilters {
1232 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1233 Vec::new();
1234
1235 for cap_config in capability_configs {
1236 let cap_id = cap_config.capability_ref.as_str();
1237 if let Some(capability) = registry.get(cap_id) {
1238 if capability.status() != CapabilityStatus::Available {
1239 continue;
1240 }
1241 if let Some(provider) = capability.message_filter_provider() {
1242 message_filter_providers.push((provider, cap_config.config.clone()));
1243 }
1244 }
1245 }
1246
1247 message_filter_providers.sort_by_key(|(p, _)| p.priority());
1248
1249 CollectedMessageFilters {
1250 message_filter_providers,
1251 }
1252}
1253
1254pub fn collect_model_view_providers(
1256 capability_configs: &[AgentCapabilityConfig],
1257 registry: &CapabilityRegistry,
1258) -> CollectedModelViewProviders {
1259 let mut model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)> = Vec::new();
1260
1261 for cap_config in capability_configs {
1262 let cap_id = cap_config.capability_ref.as_str();
1263 if let Some(capability) = registry.get(cap_id) {
1264 if capability.status() != CapabilityStatus::Available {
1265 continue;
1266 }
1267 if let Some(provider) = capability.model_view_provider() {
1268 model_view_providers.push((provider, cap_config.config.clone()));
1269 }
1270 }
1271 }
1272
1273 model_view_providers.sort_by_key(|(p, _)| p.priority());
1274
1275 CollectedModelViewProviders {
1276 model_view_providers,
1277 }
1278}
1279
1280pub fn collect_capability_mcp_servers(
1281 capability_configs: &[AgentCapabilityConfig],
1282 registry: &CapabilityRegistry,
1283) -> ScopedMcpServers {
1284 let mut servers = ScopedMcpServers::default();
1285
1286 for cap_config in capability_configs {
1287 let cap_id = cap_config.capability_ref.as_str();
1288 if is_declarative_capability(cap_id) {
1289 if let Ok(definition) =
1290 serde_json::from_value::<DeclarativeCapabilityDefinition>(cap_config.config.clone())
1291 {
1292 if definition.status != CapabilityStatus::Available {
1293 continue;
1294 }
1295 if let Some(contributed) = definition.mcp_servers {
1296 servers = merge_scoped_mcp_servers(&servers, &contributed);
1297 }
1298 }
1299 continue;
1300 }
1301 if let Some(capability) = registry.get(cap_id) {
1302 if capability.status() != CapabilityStatus::Available {
1303 continue;
1304 }
1305 servers = merge_scoped_mcp_servers(
1306 &servers,
1307 &capability.mcp_servers_with_config(&cap_config.config),
1308 );
1309 }
1310 }
1311
1312 servers
1313}
1314
1315pub const MAX_RESOLVED_CAPABILITIES: usize = 100;
1322
1323#[derive(Debug, Clone, PartialEq, Eq)]
1325pub enum DependencyError {
1326 CircularDependency {
1328 capability_id: String,
1330 chain: Vec<String>,
1332 },
1333 TooManyCapabilities {
1335 count: usize,
1337 max: usize,
1339 },
1340}
1341
1342impl std::fmt::Display for DependencyError {
1343 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1344 match self {
1345 DependencyError::CircularDependency {
1346 capability_id,
1347 chain,
1348 } => {
1349 write!(
1350 f,
1351 "Circular dependency detected: {} depends on itself via chain: {} -> {}",
1352 capability_id,
1353 chain.join(" -> "),
1354 capability_id
1355 )
1356 }
1357 DependencyError::TooManyCapabilities { count, max } => {
1358 write!(
1359 f,
1360 "Too many capabilities after resolution: {} (max: {})",
1361 count, max
1362 )
1363 }
1364 }
1365 }
1366}
1367
1368impl std::error::Error for DependencyError {}
1369
1370#[derive(Debug, Clone)]
1372pub struct ResolvedCapabilities {
1373 pub resolved_ids: Vec<String>,
1376 pub added_as_dependencies: Vec<String>,
1378 pub user_selected: Vec<String>,
1380}
1381
1382pub fn resolve_dependencies(
1402 selected_ids: &[String],
1403 registry: &CapabilityRegistry,
1404) -> Result<ResolvedCapabilities, DependencyError> {
1405 use std::collections::HashSet;
1406
1407 let user_selected: HashSet<String> = selected_ids.iter().cloned().collect();
1408 let mut resolved: Vec<String> = Vec::new();
1409 let mut resolved_set: HashSet<String> = HashSet::new();
1410 let mut added_as_dependencies: Vec<String> = Vec::new();
1411
1412 for cap_id in selected_ids {
1414 resolve_single_capability(
1415 cap_id,
1416 registry,
1417 &mut resolved,
1418 &mut resolved_set,
1419 &mut added_as_dependencies,
1420 &user_selected,
1421 &mut Vec::new(), )?;
1423 }
1424
1425 if resolved.len() > MAX_RESOLVED_CAPABILITIES {
1427 return Err(DependencyError::TooManyCapabilities {
1428 count: resolved.len(),
1429 max: MAX_RESOLVED_CAPABILITIES,
1430 });
1431 }
1432
1433 Ok(ResolvedCapabilities {
1434 resolved_ids: resolved,
1435 added_as_dependencies,
1436 user_selected: selected_ids.to_vec(),
1437 })
1438}
1439
1440pub fn resolve_capability_configs(
1445 selected_configs: &[AgentCapabilityConfig],
1446 registry: &CapabilityRegistry,
1447) -> Result<Vec<AgentCapabilityConfig>, DependencyError> {
1448 let mut selected_ids: Vec<String> = Vec::new();
1449 for config in selected_configs {
1450 if is_declarative_capability(config.capability_id())
1451 && let Ok(definition) =
1452 serde_json::from_value::<DeclarativeCapabilityDefinition>(config.config.clone())
1453 {
1454 selected_ids.extend(definition.dependencies);
1455 }
1456 selected_ids.push(config.capability_id().to_string());
1457 }
1458 let resolved = resolve_dependencies(&selected_ids, registry)?;
1459
1460 let explicit_configs: std::collections::HashMap<String, serde_json::Value> = selected_configs
1461 .iter()
1462 .map(|config| (config.capability_id().to_string(), config.config.clone()))
1463 .collect();
1464
1465 Ok(resolved
1466 .resolved_ids
1467 .into_iter()
1468 .map(|capability_id| {
1469 explicit_configs
1470 .get(&capability_id)
1471 .cloned()
1472 .map(|config| AgentCapabilityConfig::with_config(capability_id.clone(), config))
1473 .unwrap_or_else(|| AgentCapabilityConfig::new(capability_id))
1474 })
1475 .collect())
1476}
1477
1478fn resolve_single_capability(
1480 cap_id: &str,
1481 registry: &CapabilityRegistry,
1482 resolved: &mut Vec<String>,
1483 resolved_set: &mut std::collections::HashSet<String>,
1484 added_as_dependencies: &mut Vec<String>,
1485 user_selected: &std::collections::HashSet<String>,
1486 visiting: &mut Vec<String>,
1487) -> Result<(), DependencyError> {
1488 if resolved_set.contains(cap_id) {
1490 return Ok(());
1491 }
1492
1493 if visiting.contains(&cap_id.to_string()) {
1495 return Err(DependencyError::CircularDependency {
1496 capability_id: cap_id.to_string(),
1497 chain: visiting.clone(),
1498 });
1499 }
1500
1501 let capability = match registry.get(cap_id) {
1503 Some(cap) => cap,
1504 None => {
1505 if is_declarative_capability(cap_id) && !resolved_set.contains(cap_id) {
1506 resolved.push(cap_id.to_string());
1507 resolved_set.insert(cap_id.to_string());
1508 if !user_selected.contains(cap_id) {
1509 added_as_dependencies.push(cap_id.to_string());
1510 }
1511 }
1512 return Ok(());
1513 }
1514 };
1515
1516 visiting.push(cap_id.to_string());
1518
1519 for dep_id in capability.dependencies() {
1521 resolve_single_capability(
1522 dep_id,
1523 registry,
1524 resolved,
1525 resolved_set,
1526 added_as_dependencies,
1527 user_selected,
1528 visiting,
1529 )?;
1530 }
1531
1532 visiting.pop();
1534
1535 if !resolved_set.contains(cap_id) {
1537 resolved.push(cap_id.to_string());
1538 resolved_set.insert(cap_id.to_string());
1539
1540 if !user_selected.contains(cap_id) {
1542 added_as_dependencies.push(cap_id.to_string());
1543 }
1544 }
1545
1546 Ok(())
1547}
1548
1549pub fn compute_features(capability_ids: &[String], registry: &CapabilityRegistry) -> Vec<String> {
1554 use std::collections::HashSet;
1555
1556 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
1557 Ok(resolved) => resolved.resolved_ids,
1558 Err(_) => capability_ids.to_vec(),
1559 };
1560
1561 let mut seen = HashSet::new();
1562 let mut features = Vec::new();
1563 for cap_id in &resolved_ids {
1564 if let Some(cap) = registry.get(cap_id) {
1565 for feature in cap.features() {
1566 if seen.insert(feature) {
1567 features.push(feature.to_string());
1568 }
1569 }
1570 }
1571 }
1572 features
1573}
1574
1575pub fn get_dependencies(cap_id: &str, registry: &CapabilityRegistry) -> Vec<String> {
1578 registry
1579 .get(cap_id)
1580 .map(|cap| cap.dependencies().iter().map(|s| s.to_string()).collect())
1581 .unwrap_or_default()
1582}
1583
1584pub async fn collect_capabilities(
1600 capability_ids: &[String],
1601 registry: &CapabilityRegistry,
1602 ctx: &SystemPromptContext,
1603) -> CollectedCapabilities {
1604 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
1607 Ok(resolved) => resolved.resolved_ids,
1608 Err(e) => {
1609 tracing::warn!("Failed to resolve capability dependencies: {}", e);
1610 capability_ids.to_vec()
1611 }
1612 };
1613
1614 let configs: Vec<AgentCapabilityConfig> = resolved_ids
1616 .iter()
1617 .map(|id| AgentCapabilityConfig {
1618 capability_ref: CapabilityId::new(id),
1619 config: serde_json::Value::Object(serde_json::Map::new()),
1620 })
1621 .collect();
1622
1623 collect_capabilities_with_configs(&configs, registry, ctx).await
1624}
1625
1626pub async fn collect_capabilities_with_configs(
1637 capability_configs: &[AgentCapabilityConfig],
1638 registry: &CapabilityRegistry,
1639 ctx: &SystemPromptContext,
1640) -> CollectedCapabilities {
1641 let mut system_prompt_parts: Vec<String> = Vec::new();
1642 let mut system_prompt_attributions: Vec<SystemPromptAttribution> = Vec::new();
1643 let mut tools: Vec<Box<dyn Tool>> = Vec::new();
1644 let mut tool_definitions: Vec<ToolDefinition> = Vec::new();
1645 let mut mounts: Vec<MountPoint> = Vec::new();
1646 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1647 Vec::new();
1648 let mut applied_ids: Vec<String> = Vec::new();
1649 let mut tool_search: Option<crate::llm_driver_registry::ToolSearchConfig> = None;
1650 let mut prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig> = None;
1651 let mut tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>> = Vec::new();
1652 let mut tool_call_hooks: Vec<Arc<dyn ToolCallHook>> = Vec::new();
1653 let mut mcp_servers = ScopedMcpServers::default();
1654
1655 for cap_config in capability_configs {
1656 let cap_id = cap_config.capability_ref.as_str();
1657 if is_declarative_capability(cap_id) {
1658 match serde_json::from_value::<DeclarativeCapabilityDefinition>(
1659 cap_config.config.clone(),
1660 ) {
1661 Ok(definition) => {
1662 if definition.status != CapabilityStatus::Available {
1663 continue;
1664 }
1665
1666 if let Some(prompt) = definition.system_prompt.as_deref() {
1667 let contribution =
1668 format!("<capability id=\"{}\">\n{}\n</capability>", cap_id, prompt);
1669 system_prompt_attributions.push(SystemPromptAttribution {
1670 capability_id: cap_id.to_string(),
1671 content: contribution.clone(),
1672 });
1673 system_prompt_parts.push(contribution);
1674 }
1675
1676 mounts.extend(definition.mounts(cap_id));
1677 if let Some(ref servers) = definition.mcp_servers {
1678 mcp_servers = merge_scoped_mcp_servers(&mcp_servers, servers);
1679 }
1680 for skill in definition.skill_contributions() {
1681 mounts.push(skill.to_mount(cap_id));
1682 }
1683
1684 applied_ids.push(cap_id.to_string());
1685 }
1686 Err(error) => {
1687 tracing::warn!(
1688 capability_id = %cap_id,
1689 error = %error,
1690 "Skipping invalid declarative capability config"
1691 );
1692 }
1693 }
1694 continue;
1695 }
1696 if let Some(capability) = registry.get(cap_id) {
1697 if capability.status() != CapabilityStatus::Available {
1699 continue;
1700 }
1701
1702 if let Some(contribution) = capability
1704 .system_prompt_contribution_with_config(ctx, &cap_config.config)
1705 .await
1706 {
1707 system_prompt_attributions.push(SystemPromptAttribution {
1708 capability_id: cap_id.to_string(),
1709 content: contribution.clone(),
1710 });
1711 system_prompt_parts.push(contribution);
1712 }
1713
1714 tools.extend(capability.tools_with_config(&cap_config.config));
1716 tool_definition_hooks.extend(capability.tool_definition_hooks());
1717 tool_call_hooks.extend(capability.tool_call_hooks());
1718 let cap_category = capability.category();
1723 for def in capability.tool_definitions() {
1724 let def = match (def.category(), cap_category) {
1725 (None, Some(cat)) => def.with_category(cat),
1726 _ => def,
1727 }
1728 .with_capability_attribution(cap_id, Some(capability.name()));
1729 tool_definitions.push(def);
1730 }
1731
1732 if cap_id == OPENAI_TOOL_SEARCH_CAPABILITY_ID {
1734 let threshold = cap_config
1736 .config
1737 .get("threshold")
1738 .and_then(|v| v.as_u64())
1739 .map(|v| v as usize)
1740 .unwrap_or(DEFAULT_TOOL_SEARCH_THRESHOLD);
1741 tool_search = Some(crate::llm_driver_registry::ToolSearchConfig {
1742 enabled: true,
1743 threshold,
1744 });
1745 }
1746
1747 if cap_id == PROMPT_CACHING_CAPABILITY_ID {
1748 let strategy = cap_config
1749 .config
1750 .get("strategy")
1751 .and_then(|v| v.as_str())
1752 .map(|value| match value {
1753 "auto" => crate::llm_driver_registry::PromptCacheStrategy::Auto,
1754 _ => crate::llm_driver_registry::PromptCacheStrategy::Auto,
1755 })
1756 .unwrap_or(crate::llm_driver_registry::PromptCacheStrategy::Auto);
1757 let gemini_cached_content = cap_config
1758 .config
1759 .get("gemini_cached_content")
1760 .and_then(|v| v.as_str())
1761 .map(str::to_string);
1762 prompt_cache = Some(crate::llm_driver_registry::PromptCacheConfig {
1763 enabled: true,
1764 strategy,
1765 gemini_cached_content,
1766 });
1767 }
1768
1769 mounts.extend(capability.mounts());
1771
1772 mcp_servers = merge_scoped_mcp_servers(
1773 &mcp_servers,
1774 &capability.mcp_servers_with_config(&cap_config.config),
1775 );
1776
1777 for skill in capability.contribute_skills() {
1781 mounts.push(skill.to_mount(cap_id));
1782 }
1783
1784 if let Some(provider) = capability.message_filter_provider() {
1786 message_filter_providers.push((provider, cap_config.config.clone()));
1787 }
1788
1789 applied_ids.push(cap_id.to_string());
1790 }
1791 }
1792
1793 message_filter_providers.sort_by_key(|(p, _)| p.priority());
1795
1796 CollectedCapabilities {
1797 system_prompt_parts,
1798 system_prompt_attributions,
1799 tools,
1800 tool_definitions,
1801 mounts,
1802 message_filter_providers,
1803 applied_ids,
1804 tool_search,
1805 prompt_cache,
1806 tool_definition_hooks,
1807 tool_call_hooks,
1808 mcp_servers,
1809 }
1810}
1811
1812pub struct AppliedCapabilities {
1818 pub runtime_agent: RuntimeAgent,
1820 pub tool_registry: ToolRegistry,
1822 pub applied_ids: Vec<String>,
1824}
1825
1826pub async fn apply_capabilities(
1863 base_runtime_agent: RuntimeAgent,
1864 capability_ids: &[String],
1865 registry: &CapabilityRegistry,
1866 ctx: &SystemPromptContext,
1867) -> AppliedCapabilities {
1868 let collected = collect_capabilities(capability_ids, registry, ctx).await;
1869
1870 let final_system_prompt = match collected.system_prompt_prefix() {
1872 Some(prefix) => format!(
1873 "{}\n\n<system-prompt>\n{}\n</system-prompt>",
1874 prefix, base_runtime_agent.system_prompt
1875 ),
1876 None => base_runtime_agent.system_prompt,
1877 };
1878
1879 let mut tool_registry = ToolRegistry::new();
1881 for tool in collected.tools {
1882 tool_registry.register_boxed(tool);
1883 }
1884
1885 let mut tools = collected.tool_definitions;
1887 for hook in &collected.tool_definition_hooks {
1888 tools = hook.transform(tools);
1889 }
1890
1891 let runtime_agent = RuntimeAgent {
1892 system_prompt: final_system_prompt,
1893 model: base_runtime_agent.model,
1894 tools,
1895 max_iterations: base_runtime_agent.max_iterations,
1896 temperature: base_runtime_agent.temperature,
1897 max_tokens: base_runtime_agent.max_tokens,
1898 tool_search: collected.tool_search,
1899 prompt_cache: collected.prompt_cache,
1900 network_access: base_runtime_agent.network_access,
1901 };
1902
1903 AppliedCapabilities {
1904 runtime_agent,
1905 tool_registry,
1906 applied_ids: collected.applied_ids,
1907 }
1908}
1909
1910#[cfg(test)]
1915mod tests {
1916 use super::*;
1917 use crate::typed_id::SessionId;
1918 use std::collections::BTreeSet;
1919 use uuid::Uuid;
1920
1921 fn test_ctx() -> SystemPromptContext {
1923 SystemPromptContext::without_file_store(SessionId::new())
1924 }
1925
1926 fn expected_core_builtin_ids() -> BTreeSet<&'static str> {
1927 let mut ids = [
1928 "agent_instructions",
1929 "human_intent",
1930 "budgeting",
1931 "self_budget",
1932 "parallel",
1933 "noop",
1934 "current_time",
1935 "research",
1936 "platform_management",
1937 "session_file_system",
1938 "workspace_volumes",
1939 "session_storage",
1940 "session",
1941 "session_sql_database",
1942 "test_math",
1943 "test_weather",
1944 "stateless_todo_list",
1945 "web_fetch",
1946 "virtual_bash",
1947 "session_schedule",
1948 "btw",
1949 "infinity_context",
1950 "compaction",
1951 "memory",
1952 "openai_tool_search",
1953 "prompt_caching",
1954 "skills",
1955 "subagents",
1956 "agent_handoff",
1957 "a2a_agent_delegation",
1958 "system_commands",
1959 "sample_data",
1960 "data_knowledge",
1961 "knowledge_base",
1962 "tool_output_persistence",
1963 "fake_warehouse",
1964 "fake_aws",
1965 "fake_crm",
1966 "fake_financial",
1967 "loop_detection",
1968 "prompt_canary_guardrail",
1969 ]
1970 .into_iter()
1971 .collect::<BTreeSet<_>>();
1972 if cfg!(feature = "ui-capabilities") {
1973 ids.insert("openui");
1974 ids.insert("a2ui");
1975 }
1976 ids
1977 }
1978
1979 fn registry_ids(registry: &CapabilityRegistry) -> BTreeSet<&str> {
1980 registry.capabilities.keys().map(String::as_str).collect()
1981 }
1982
1983 #[test]
1993 fn test_capability_registry_with_builtins_dev() {
1994 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
1996
1997 assert_eq!(registry_ids(®istry), expected_core_builtin_ids());
1998 }
1999
2000 #[test]
2001 fn test_capability_registry_with_builtins_prod() {
2002 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2004 assert_eq!(registry_ids(®istry), expected_core_builtin_ids());
2005 assert!(!registry.has("docker_container"));
2007 }
2008
2009 #[test]
2010 fn test_capability_registry_get() {
2011 let registry = CapabilityRegistry::with_builtins();
2012
2013 let noop = registry.get("noop").unwrap();
2014 assert_eq!(noop.id(), "noop");
2015 assert_eq!(noop.name(), "No-Op");
2016 assert_eq!(noop.status(), CapabilityStatus::Available);
2017 }
2018
2019 #[test]
2020 fn test_capability_registry_blueprint_with_capability() {
2021 struct BlueprintProviderCapability;
2022
2023 impl Capability for BlueprintProviderCapability {
2024 fn id(&self) -> &str {
2025 "blueprint_provider"
2026 }
2027 fn name(&self) -> &str {
2028 "Blueprint Provider"
2029 }
2030 fn description(&self) -> &str {
2031 "Capability that provides a blueprint for tests"
2032 }
2033 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
2034 vec![AgentBlueprint {
2035 id: "test_blueprint",
2036 name: "Test Blueprint",
2037 description: "Blueprint for capability registry tests",
2038 model: BlueprintModel::Inherit,
2039 system_prompt: "Test prompt",
2040 tools: vec![],
2041 max_turns: None,
2042 config_schema: None,
2043 }]
2044 }
2045 }
2046
2047 let mut registry = CapabilityRegistry::new();
2048 registry.register(BlueprintProviderCapability);
2049
2050 let (capability_id, blueprint) = registry
2051 .blueprint_with_capability("test_blueprint")
2052 .expect("blueprint should resolve with capability id");
2053 assert_eq!(capability_id, "blueprint_provider");
2054 assert_eq!(blueprint.id, "test_blueprint");
2055 }
2056
2057 #[test]
2058 fn test_capability_registry_builder() {
2059 let registry = CapabilityRegistry::builder()
2060 .capability(NoopCapability)
2061 .capability(CurrentTimeCapability)
2062 .build();
2063
2064 assert!(registry.has("noop"));
2065 assert!(registry.has("current_time"));
2066 assert_eq!(registry.len(), 2);
2067 }
2068
2069 #[test]
2070 fn test_capability_status() {
2071 let registry = CapabilityRegistry::with_builtins();
2072
2073 let current_time = registry.get("current_time").unwrap();
2074 assert_eq!(current_time.status(), CapabilityStatus::Available);
2075
2076 let research = registry.get("research").unwrap();
2077 assert_eq!(research.status(), CapabilityStatus::ComingSoon);
2078 }
2079
2080 #[test]
2081 fn test_capability_icons_and_categories() {
2082 let registry = CapabilityRegistry::with_builtins();
2083
2084 let noop = registry.get("noop").unwrap();
2085 assert_eq!(noop.icon(), Some("circle-off"));
2086 assert_eq!(noop.category(), Some("Testing"));
2087
2088 let current_time = registry.get("current_time").unwrap();
2089 assert_eq!(current_time.icon(), Some("clock"));
2090 assert_eq!(current_time.category(), Some("Utilities"));
2091 }
2092
2093 #[test]
2094 fn test_system_prompt_preview_default_delegates_to_addition() {
2095 let registry = CapabilityRegistry::with_builtins();
2096
2097 let test_math = registry.get("test_math").unwrap();
2099 assert_eq!(
2100 test_math.system_prompt_preview().as_deref(),
2101 test_math.system_prompt_addition()
2102 );
2103
2104 let current_time = registry.get("current_time").unwrap();
2106 assert!(current_time.system_prompt_preview().is_none());
2107 assert!(current_time.system_prompt_addition().is_none());
2108 }
2109
2110 #[test]
2111 fn test_system_prompt_preview_dynamic_capability() {
2112 let registry = CapabilityRegistry::with_builtins();
2113 let cap = registry.get("agent_instructions").unwrap();
2114
2115 assert!(cap.system_prompt_addition().is_none());
2117 assert!(cap.system_prompt_preview().is_some());
2118 assert!(cap.system_prompt_preview().unwrap().contains("AGENTS.md"));
2119 }
2120
2121 #[tokio::test]
2126 async fn test_apply_capabilities_empty() {
2127 let registry = CapabilityRegistry::with_builtins();
2128 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2129
2130 let applied =
2131 apply_capabilities(base_runtime_agent.clone(), &[], ®istry, &test_ctx()).await;
2132
2133 assert_eq!(
2134 applied.runtime_agent.system_prompt,
2135 base_runtime_agent.system_prompt
2136 );
2137 assert!(applied.tool_registry.is_empty());
2138 assert!(applied.applied_ids.is_empty());
2139 }
2140
2141 #[tokio::test]
2142 async fn test_apply_capabilities_noop() {
2143 let registry = CapabilityRegistry::with_builtins();
2144 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2145
2146 let applied = apply_capabilities(
2147 base_runtime_agent.clone(),
2148 &["noop".to_string()],
2149 ®istry,
2150 &test_ctx(),
2151 )
2152 .await;
2153
2154 assert_eq!(
2156 applied.runtime_agent.system_prompt,
2157 base_runtime_agent.system_prompt
2158 );
2159 assert!(applied.tool_registry.is_empty());
2160 assert_eq!(applied.applied_ids, vec!["noop"]);
2161 }
2162
2163 #[tokio::test]
2164 async fn test_apply_capabilities_current_time() {
2165 let registry = CapabilityRegistry::with_builtins();
2166 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2167
2168 let applied = apply_capabilities(
2169 base_runtime_agent.clone(),
2170 &["current_time".to_string()],
2171 ®istry,
2172 &test_ctx(),
2173 )
2174 .await;
2175
2176 assert_eq!(
2178 applied.runtime_agent.system_prompt,
2179 base_runtime_agent.system_prompt
2180 );
2181 assert!(applied.tool_registry.has("get_current_time"));
2182 assert_eq!(applied.tool_registry.len(), 1);
2183 assert_eq!(applied.applied_ids, vec!["current_time"]);
2184 }
2185
2186 #[tokio::test]
2187 async fn test_apply_capabilities_skips_coming_soon() {
2188 let registry = CapabilityRegistry::with_builtins();
2189 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2190
2191 let applied = apply_capabilities(
2193 base_runtime_agent.clone(),
2194 &["research".to_string()],
2195 ®istry,
2196 &test_ctx(),
2197 )
2198 .await;
2199
2200 assert_eq!(
2202 applied.runtime_agent.system_prompt,
2203 base_runtime_agent.system_prompt
2204 );
2205 assert!(applied.applied_ids.is_empty()); }
2207
2208 #[tokio::test]
2209 async fn test_apply_capabilities_multiple() {
2210 let registry = CapabilityRegistry::with_builtins();
2211 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2212
2213 let applied = apply_capabilities(
2214 base_runtime_agent.clone(),
2215 &["noop".to_string(), "current_time".to_string()],
2216 ®istry,
2217 &test_ctx(),
2218 )
2219 .await;
2220
2221 assert!(applied.tool_registry.has("get_current_time"));
2222 assert_eq!(applied.applied_ids, vec!["noop", "current_time"]);
2223 }
2224
2225 #[tokio::test]
2226 async fn test_apply_capabilities_preserves_order() {
2227 let registry = CapabilityRegistry::with_builtins();
2228 let base_runtime_agent = RuntimeAgent::new("Base prompt.", "gpt-5.2");
2229
2230 let applied = apply_capabilities(
2232 base_runtime_agent,
2233 &["current_time".to_string(), "noop".to_string()],
2234 ®istry,
2235 &test_ctx(),
2236 )
2237 .await;
2238
2239 assert_eq!(applied.applied_ids, vec!["current_time", "noop"]);
2240 }
2241
2242 #[tokio::test]
2243 async fn test_apply_capabilities_test_math() {
2244 let registry = CapabilityRegistry::with_builtins();
2245 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2246
2247 let applied = apply_capabilities(
2248 base_runtime_agent.clone(),
2249 &["test_math".to_string()],
2250 ®istry,
2251 &test_ctx(),
2252 )
2253 .await;
2254
2255 assert!(
2257 !applied
2258 .runtime_agent
2259 .system_prompt
2260 .contains("<capability id=\"test_math\">")
2261 );
2262 assert!(
2264 applied
2265 .runtime_agent
2266 .system_prompt
2267 .contains("You are a helpful assistant.")
2268 );
2269 assert!(applied.tool_registry.has("add"));
2270 assert!(applied.tool_registry.has("subtract"));
2271 assert!(applied.tool_registry.has("multiply"));
2272 assert!(applied.tool_registry.has("divide"));
2273 assert_eq!(applied.tool_registry.len(), 4);
2274 }
2275
2276 #[tokio::test]
2277 async fn test_apply_capabilities_test_weather() {
2278 let registry = CapabilityRegistry::with_builtins();
2279 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2280
2281 let applied = apply_capabilities(
2282 base_runtime_agent.clone(),
2283 &["test_weather".to_string()],
2284 ®istry,
2285 &test_ctx(),
2286 )
2287 .await;
2288
2289 assert!(
2291 !applied
2292 .runtime_agent
2293 .system_prompt
2294 .contains("<capability id=\"test_weather\">")
2295 );
2296 assert!(applied.tool_registry.has("get_weather"));
2297 assert!(applied.tool_registry.has("get_forecast"));
2298 assert_eq!(applied.tool_registry.len(), 2);
2299 }
2300
2301 #[tokio::test]
2302 async fn test_apply_capabilities_test_math_and_test_weather() {
2303 let registry = CapabilityRegistry::with_builtins();
2304 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2305
2306 let applied = apply_capabilities(
2307 base_runtime_agent.clone(),
2308 &["test_math".to_string(), "test_weather".to_string()],
2309 ®istry,
2310 &test_ctx(),
2311 )
2312 .await;
2313
2314 assert_eq!(applied.tool_registry.len(), 6); assert!(applied.tool_registry.has("add"));
2317 assert!(applied.tool_registry.has("get_weather"));
2318 }
2319
2320 #[tokio::test]
2321 async fn test_apply_capabilities_stateless_todo_list() {
2322 let registry = CapabilityRegistry::with_builtins();
2323 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2324
2325 let applied = apply_capabilities(
2326 base_runtime_agent.clone(),
2327 &["stateless_todo_list".to_string()],
2328 ®istry,
2329 &test_ctx(),
2330 )
2331 .await;
2332
2333 assert!(
2335 applied
2336 .runtime_agent
2337 .system_prompt
2338 .contains("Task Management")
2339 );
2340 assert!(applied.runtime_agent.system_prompt.contains("write_todos"));
2341 assert!(applied.tool_registry.has("write_todos"));
2342 assert_eq!(applied.tool_registry.len(), 1);
2343 }
2344
2345 #[tokio::test]
2346 async fn test_apply_capabilities_web_fetch() {
2347 let registry = CapabilityRegistry::with_builtins();
2348 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2349
2350 let applied = apply_capabilities(
2351 base_runtime_agent.clone(),
2352 &["web_fetch".to_string()],
2353 ®istry,
2354 &test_ctx(),
2355 )
2356 .await;
2357
2358 assert!(
2360 applied
2361 .runtime_agent
2362 .system_prompt
2363 .contains(&base_runtime_agent.system_prompt)
2364 );
2365 assert!(applied.runtime_agent.system_prompt.contains("web_fetch"));
2366 assert!(applied.tool_registry.has("web_fetch"));
2367 assert_eq!(applied.tool_registry.len(), 1);
2368 }
2369
2370 #[tokio::test]
2375 async fn test_xml_tags_wrap_capability_prompts() {
2376 let registry = CapabilityRegistry::with_builtins();
2377 let collected =
2378 collect_capabilities(&["stateless_todo_list".to_string()], ®istry, &test_ctx())
2379 .await;
2380
2381 assert_eq!(collected.system_prompt_parts.len(), 1);
2382 let part = &collected.system_prompt_parts[0];
2383 assert!(part.starts_with("<capability id=\"stateless_todo_list\">"));
2384 assert!(part.ends_with("</capability>"));
2385 assert!(part.contains("Task Management"));
2386 }
2387
2388 #[tokio::test]
2389 async fn test_xml_tags_multiple_capabilities() {
2390 let registry = CapabilityRegistry::with_builtins();
2391 let collected = collect_capabilities(
2392 &[
2393 "stateless_todo_list".to_string(),
2394 "session_schedule".to_string(),
2395 ],
2396 ®istry,
2397 &test_ctx(),
2398 )
2399 .await;
2400
2401 assert_eq!(collected.system_prompt_parts.len(), 2);
2402 assert!(
2403 collected.system_prompt_parts[0].starts_with("<capability id=\"stateless_todo_list\">")
2404 );
2405 assert!(
2406 collected.system_prompt_parts[1].starts_with("<capability id=\"session_schedule\">")
2407 );
2408
2409 let prefix = collected.system_prompt_prefix().unwrap();
2410 assert!(prefix.contains("</capability>\n\n<capability"));
2412 }
2413
2414 #[tokio::test]
2415 async fn test_xml_tags_system_prompt_wrapping() {
2416 let registry = CapabilityRegistry::with_builtins();
2417 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2418
2419 let applied = apply_capabilities(
2420 base,
2421 &["stateless_todo_list".to_string()],
2422 ®istry,
2423 &test_ctx(),
2424 )
2425 .await;
2426
2427 let prompt = &applied.runtime_agent.system_prompt;
2428 assert!(prompt.contains("<capability id=\"stateless_todo_list\">"));
2430 assert!(prompt.contains("</capability>"));
2431 assert!(prompt.contains("<system-prompt>\nYou are helpful.\n</system-prompt>"));
2433 }
2434
2435 #[tokio::test]
2436 async fn test_no_xml_wrapping_without_capabilities() {
2437 let registry = CapabilityRegistry::with_builtins();
2438 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2439
2440 let applied = apply_capabilities(base, &[], ®istry, &test_ctx()).await;
2441
2442 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2444 assert!(
2445 !applied
2446 .runtime_agent
2447 .system_prompt
2448 .contains("<system-prompt>")
2449 );
2450 }
2451
2452 #[tokio::test]
2453 async fn test_no_xml_wrapping_for_noop_capability() {
2454 let registry = CapabilityRegistry::with_builtins();
2455 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2456
2457 let applied = apply_capabilities(base, &["noop".to_string()], ®istry, &test_ctx()).await;
2459
2460 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2461 assert!(
2462 !applied
2463 .runtime_agent
2464 .system_prompt
2465 .contains("<system-prompt>")
2466 );
2467 }
2468
2469 #[tokio::test]
2474 async fn test_collect_capabilities_includes_mounts() {
2475 let registry = CapabilityRegistry::with_builtins();
2476
2477 let collected =
2478 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
2479
2480 assert!(!collected.mounts.is_empty());
2481 assert_eq!(collected.mounts.len(), 1);
2482 assert_eq!(collected.mounts[0].path, "/samples");
2483 assert!(collected.mounts[0].is_readonly());
2484 }
2485
2486 #[tokio::test]
2487 async fn test_collect_capabilities_empty_mounts_by_default() {
2488 let registry = CapabilityRegistry::with_builtins();
2489
2490 let collected =
2492 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
2493
2494 assert!(collected.mounts.is_empty());
2495 }
2496
2497 #[tokio::test]
2498 async fn test_collect_capabilities_combines_mounts() {
2499 let registry = CapabilityRegistry::with_builtins();
2500
2501 let collected = collect_capabilities(
2504 &["sample_data".to_string(), "current_time".to_string()],
2505 ®istry,
2506 &test_ctx(),
2507 )
2508 .await;
2509
2510 assert_eq!(collected.mounts.len(), 1);
2511 assert!(
2513 collected
2514 .applied_ids
2515 .iter()
2516 .any(|id| id == "session_file_system")
2517 );
2518 assert!(collected.applied_ids.iter().any(|id| id == "sample_data"));
2519 assert!(collected.applied_ids.iter().any(|id| id == "current_time"));
2520 }
2521
2522 #[test]
2523 fn test_sample_data_capability() {
2524 let registry = CapabilityRegistry::with_builtins();
2525 let cap = registry.get("sample_data").unwrap();
2526
2527 assert_eq!(cap.id(), "sample_data");
2528 assert_eq!(cap.name(), "Sample Data");
2529 assert_eq!(cap.status(), CapabilityStatus::Available);
2530
2531 assert!(cap.system_prompt_addition().is_some());
2533 assert!(cap.tools().is_empty());
2534
2535 assert!(!cap.mounts().is_empty());
2537 }
2538
2539 #[test]
2544 fn test_resolve_dependencies_empty() {
2545 let registry = CapabilityRegistry::with_builtins();
2546
2547 let resolved = resolve_dependencies(&[], ®istry).unwrap();
2548
2549 assert!(resolved.resolved_ids.is_empty());
2550 assert!(resolved.added_as_dependencies.is_empty());
2551 assert!(resolved.user_selected.is_empty());
2552 }
2553
2554 #[test]
2555 fn test_resolve_dependencies_no_deps() {
2556 let registry = CapabilityRegistry::with_builtins();
2557
2558 let resolved = resolve_dependencies(&["current_time".to_string()], ®istry).unwrap();
2560
2561 assert_eq!(resolved.resolved_ids, vec!["current_time"]);
2562 assert!(resolved.added_as_dependencies.is_empty());
2563 }
2564
2565 #[test]
2566 fn test_resolve_dependencies_with_deps() {
2567 let registry = CapabilityRegistry::with_builtins();
2568
2569 let resolved = resolve_dependencies(&["sample_data".to_string()], ®istry).unwrap();
2571
2572 assert_eq!(resolved.resolved_ids.len(), 2);
2574 let fs_pos = resolved
2575 .resolved_ids
2576 .iter()
2577 .position(|id| id == "session_file_system")
2578 .unwrap();
2579 let sd_pos = resolved
2580 .resolved_ids
2581 .iter()
2582 .position(|id| id == "sample_data")
2583 .unwrap();
2584 assert!(fs_pos < sd_pos, "FileSystem should come before SampleData");
2585
2586 assert_eq!(resolved.added_as_dependencies, vec!["session_file_system"]);
2588 }
2589
2590 #[test]
2591 fn test_resolve_dependencies_already_selected() {
2592 let registry = CapabilityRegistry::with_builtins();
2593
2594 let resolved = resolve_dependencies(
2596 &["session_file_system".to_string(), "sample_data".to_string()],
2597 ®istry,
2598 )
2599 .unwrap();
2600
2601 assert_eq!(resolved.resolved_ids.len(), 2);
2602 assert!(resolved.added_as_dependencies.is_empty());
2604 }
2605
2606 #[test]
2607 fn test_resolve_dependencies_preserves_order() {
2608 let registry = CapabilityRegistry::with_builtins();
2609
2610 let resolved =
2612 resolve_dependencies(&["current_time".to_string(), "noop".to_string()], ®istry)
2613 .unwrap();
2614
2615 assert_eq!(resolved.resolved_ids, vec!["current_time", "noop"]);
2616 }
2617
2618 #[test]
2619 fn test_resolve_dependencies_unknown_capability() {
2620 let registry = CapabilityRegistry::with_builtins();
2621
2622 let resolved =
2624 resolve_dependencies(&["unknown_capability".to_string()], ®istry).unwrap();
2625
2626 assert!(resolved.resolved_ids.is_empty());
2627 }
2628
2629 #[test]
2630 fn test_get_dependencies() {
2631 let registry = CapabilityRegistry::with_builtins();
2632
2633 let deps = get_dependencies("sample_data", ®istry);
2635 assert_eq!(deps, vec!["session_file_system"]);
2636
2637 let deps = get_dependencies("current_time", ®istry);
2639 assert!(deps.is_empty());
2640
2641 let deps = get_dependencies("unknown", ®istry);
2643 assert!(deps.is_empty());
2644 }
2645
2646 #[test]
2647 fn test_sample_data_has_dependency() {
2648 let registry = CapabilityRegistry::with_builtins();
2649 let cap = registry.get("sample_data").unwrap();
2650
2651 let deps = cap.dependencies();
2652 assert_eq!(deps.len(), 1);
2653 assert_eq!(deps[0], "session_file_system");
2654 }
2655
2656 #[test]
2657 fn test_noop_has_no_dependencies() {
2658 let registry = CapabilityRegistry::with_builtins();
2659 let cap = registry.get("noop").unwrap();
2660
2661 assert!(cap.dependencies().is_empty());
2662 }
2663
2664 #[test]
2668 fn test_circular_dependency_error() {
2669 struct CapA;
2671 struct CapB;
2672
2673 impl Capability for CapA {
2674 fn id(&self) -> &str {
2675 "test_cap_a"
2676 }
2677 fn name(&self) -> &str {
2678 "Test A"
2679 }
2680 fn description(&self) -> &str {
2681 "Test capability A"
2682 }
2683 fn dependencies(&self) -> Vec<&'static str> {
2684 vec!["test_cap_b"]
2685 }
2686 }
2687
2688 impl Capability for CapB {
2689 fn id(&self) -> &str {
2690 "test_cap_b"
2691 }
2692 fn name(&self) -> &str {
2693 "Test B"
2694 }
2695 fn description(&self) -> &str {
2696 "Test capability B"
2697 }
2698 fn dependencies(&self) -> Vec<&'static str> {
2699 vec!["test_cap_a"]
2700 }
2701 }
2702
2703 let mut registry = CapabilityRegistry::new();
2704 registry.register(CapA);
2705 registry.register(CapB);
2706
2707 let result = resolve_dependencies(&["test_cap_a".to_string()], ®istry);
2708
2709 assert!(result.is_err());
2710 match result.unwrap_err() {
2711 DependencyError::CircularDependency { capability_id, .. } => {
2712 assert_eq!(capability_id, "test_cap_a");
2713 }
2714 _ => panic!("Expected CircularDependency error"),
2715 }
2716 }
2717
2718 use crate::message_filter::{MessageFilter, MessageFilterProvider, MessageQuery};
2723
2724 struct FilterTestCapability {
2726 priority: i32,
2727 }
2728
2729 impl Capability for FilterTestCapability {
2730 fn id(&self) -> &str {
2731 "filter_test"
2732 }
2733 fn name(&self) -> &str {
2734 "Filter Test"
2735 }
2736 fn description(&self) -> &str {
2737 "Test capability with message filter"
2738 }
2739 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
2740 Some(Arc::new(FilterTestProvider {
2741 priority: self.priority,
2742 }))
2743 }
2744 }
2745
2746 struct FilterTestProvider {
2747 priority: i32,
2748 }
2749
2750 impl MessageFilterProvider for FilterTestProvider {
2751 fn apply_filters(&self, query: &mut MessageQuery, config: &serde_json::Value) {
2752 if let Some(search) = config.get("search").and_then(|v| v.as_str()) {
2754 query
2755 .filters
2756 .push(MessageFilter::Search(search.to_string()));
2757 }
2758 }
2759
2760 fn priority(&self) -> i32 {
2761 self.priority
2762 }
2763 }
2764
2765 #[tokio::test]
2766 async fn test_collect_capabilities_with_configs_no_filter_providers() {
2767 let registry = CapabilityRegistry::with_builtins();
2768 let configs = vec![AgentCapabilityConfig {
2769 capability_ref: CapabilityId::new("current_time"),
2770 config: serde_json::json!({}),
2771 }];
2772
2773 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2774
2775 assert!(collected.message_filter_providers.is_empty());
2776 assert!(!collected.has_message_filters());
2777 }
2778
2779 #[tokio::test]
2780 async fn test_collect_capabilities_with_configs_with_filter_provider() {
2781 let mut registry = CapabilityRegistry::new();
2782 registry.register(FilterTestCapability { priority: 0 });
2783
2784 let configs = vec![AgentCapabilityConfig {
2785 capability_ref: CapabilityId::new("filter_test"),
2786 config: serde_json::json!({ "search": "hello" }),
2787 }];
2788
2789 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2790
2791 assert_eq!(collected.message_filter_providers.len(), 1);
2792 assert!(collected.has_message_filters());
2793 }
2794
2795 #[tokio::test]
2796 async fn test_collect_capabilities_with_configs_filter_priority_order() {
2797 struct HighPriorityCapability;
2799 struct LowPriorityCapability;
2800
2801 impl Capability for HighPriorityCapability {
2802 fn id(&self) -> &str {
2803 "high_priority"
2804 }
2805 fn name(&self) -> &str {
2806 "High Priority"
2807 }
2808 fn description(&self) -> &str {
2809 "Test"
2810 }
2811 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
2812 Some(Arc::new(FilterTestProvider { priority: 10 }))
2813 }
2814 }
2815
2816 impl Capability for LowPriorityCapability {
2817 fn id(&self) -> &str {
2818 "low_priority"
2819 }
2820 fn name(&self) -> &str {
2821 "Low Priority"
2822 }
2823 fn description(&self) -> &str {
2824 "Test"
2825 }
2826 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
2827 Some(Arc::new(FilterTestProvider { priority: -5 }))
2828 }
2829 }
2830
2831 let mut registry = CapabilityRegistry::new();
2832 registry.register(HighPriorityCapability);
2833 registry.register(LowPriorityCapability);
2834
2835 let configs = vec![
2837 AgentCapabilityConfig {
2838 capability_ref: CapabilityId::new("high_priority"),
2839 config: serde_json::json!({}),
2840 },
2841 AgentCapabilityConfig {
2842 capability_ref: CapabilityId::new("low_priority"),
2843 config: serde_json::json!({}),
2844 },
2845 ];
2846
2847 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2848
2849 assert_eq!(collected.message_filter_providers.len(), 2);
2851 assert_eq!(collected.message_filter_providers[0].0.priority(), -5);
2852 assert_eq!(collected.message_filter_providers[1].0.priority(), 10);
2853 }
2854
2855 #[tokio::test]
2856 async fn test_collected_capabilities_apply_message_filters() {
2857 let mut registry = CapabilityRegistry::new();
2858 registry.register(FilterTestCapability { priority: 0 });
2859
2860 let configs = vec![AgentCapabilityConfig {
2861 capability_ref: CapabilityId::new("filter_test"),
2862 config: serde_json::json!({ "search": "test_query" }),
2863 }];
2864
2865 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2866
2867 let session_id: SessionId = Uuid::now_v7().into();
2869 let mut query = MessageQuery::new(session_id);
2870
2871 collected.apply_message_filters(&mut query);
2872
2873 assert_eq!(query.filters.len(), 1);
2875 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
2876 }
2877
2878 #[tokio::test]
2879 async fn test_collected_capabilities_apply_multiple_filters_in_priority_order() {
2880 struct SearchCapability {
2881 id: &'static str,
2882 search_term: &'static str,
2883 priority: i32,
2884 }
2885
2886 struct SearchProvider {
2887 search_term: &'static str,
2888 priority: i32,
2889 }
2890
2891 impl MessageFilterProvider for SearchProvider {
2892 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
2893 query
2894 .filters
2895 .push(MessageFilter::Search(self.search_term.to_string()));
2896 }
2897
2898 fn priority(&self) -> i32 {
2899 self.priority
2900 }
2901 }
2902
2903 impl Capability for SearchCapability {
2904 fn id(&self) -> &str {
2905 self.id
2906 }
2907 fn name(&self) -> &str {
2908 "Search"
2909 }
2910 fn description(&self) -> &str {
2911 "Test"
2912 }
2913 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
2914 Some(Arc::new(SearchProvider {
2915 search_term: self.search_term,
2916 priority: self.priority,
2917 }))
2918 }
2919 }
2920
2921 let mut registry = CapabilityRegistry::new();
2922 registry.register(SearchCapability {
2923 id: "cap_a",
2924 search_term: "alpha",
2925 priority: 5,
2926 });
2927 registry.register(SearchCapability {
2928 id: "cap_b",
2929 search_term: "beta",
2930 priority: 1,
2931 });
2932 registry.register(SearchCapability {
2933 id: "cap_c",
2934 search_term: "gamma",
2935 priority: 10,
2936 });
2937
2938 let configs = vec![
2939 AgentCapabilityConfig {
2940 capability_ref: CapabilityId::new("cap_a"),
2941 config: serde_json::json!({}),
2942 },
2943 AgentCapabilityConfig {
2944 capability_ref: CapabilityId::new("cap_b"),
2945 config: serde_json::json!({}),
2946 },
2947 AgentCapabilityConfig {
2948 capability_ref: CapabilityId::new("cap_c"),
2949 config: serde_json::json!({}),
2950 },
2951 ];
2952
2953 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2954
2955 let session_id: SessionId = Uuid::now_v7().into();
2956 let mut query = MessageQuery::new(session_id);
2957
2958 collected.apply_message_filters(&mut query);
2959
2960 assert_eq!(query.filters.len(), 3);
2962 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
2963 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
2964 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
2965 }
2966
2967 #[test]
2968 fn test_capability_without_message_filter_returns_none() {
2969 let registry = CapabilityRegistry::with_builtins();
2970
2971 let noop = registry.get("noop").unwrap();
2972 assert!(noop.message_filter_provider().is_none());
2973
2974 let current_time = registry.get("current_time").unwrap();
2975 assert!(current_time.message_filter_provider().is_none());
2976 }
2977
2978 #[tokio::test]
2979 async fn test_collect_capabilities_preserves_config_for_filter_provider() {
2980 let mut registry = CapabilityRegistry::new();
2981 registry.register(FilterTestCapability { priority: 0 });
2982
2983 let test_config = serde_json::json!({
2984 "search": "custom_search",
2985 "extra_field": 42
2986 });
2987
2988 let configs = vec![AgentCapabilityConfig {
2989 capability_ref: CapabilityId::new("filter_test"),
2990 config: test_config.clone(),
2991 }];
2992
2993 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
2994
2995 assert_eq!(collected.message_filter_providers.len(), 1);
2997 let (_, stored_config) = &collected.message_filter_providers[0];
2998 assert_eq!(*stored_config, test_config);
2999 }
3000
3001 #[test]
3006 fn test_collect_message_filters_only_collects_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_message_filters_only(&configs, ®istry);
3016
3017 let session_id: SessionId = Uuid::now_v7().into();
3018 let mut query = MessageQuery::new(session_id);
3019 collected.apply_message_filters(&mut query);
3020
3021 assert_eq!(query.filters.len(), 1);
3022 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3023 }
3024
3025 #[test]
3026 fn test_collect_message_filters_only_skips_unknown_capabilities() {
3027 let registry = CapabilityRegistry::new();
3028
3029 let configs = vec![AgentCapabilityConfig {
3030 capability_ref: CapabilityId::new("nonexistent"),
3031 config: serde_json::json!({}),
3032 }];
3033
3034 let collected = collect_message_filters_only(&configs, ®istry);
3035 assert!(collected.message_filter_providers.is_empty());
3036 }
3037
3038 #[test]
3039 fn test_collect_message_filters_only_preserves_priority_order() {
3040 struct PriorityFilterCap {
3041 id: &'static str,
3042 search_term: &'static str,
3043 priority: i32,
3044 }
3045
3046 struct PriorityFilterProvider {
3047 search_term: &'static str,
3048 priority: i32,
3049 }
3050
3051 impl Capability for PriorityFilterCap {
3052 fn id(&self) -> &str {
3053 self.id
3054 }
3055 fn name(&self) -> &str {
3056 self.id
3057 }
3058 fn description(&self) -> &str {
3059 "priority test"
3060 }
3061 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3062 Some(Arc::new(PriorityFilterProvider {
3063 search_term: self.search_term,
3064 priority: self.priority,
3065 }))
3066 }
3067 }
3068
3069 impl MessageFilterProvider for PriorityFilterProvider {
3070 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3071 query
3072 .filters
3073 .push(MessageFilter::Search(self.search_term.to_string()));
3074 }
3075 fn priority(&self) -> i32 {
3076 self.priority
3077 }
3078 }
3079
3080 let mut registry = CapabilityRegistry::new();
3081 registry.register(PriorityFilterCap {
3082 id: "gamma",
3083 search_term: "gamma",
3084 priority: 10,
3085 });
3086 registry.register(PriorityFilterCap {
3087 id: "alpha",
3088 search_term: "alpha",
3089 priority: 5,
3090 });
3091 registry.register(PriorityFilterCap {
3092 id: "beta",
3093 search_term: "beta",
3094 priority: 1,
3095 });
3096
3097 let configs = vec![
3098 AgentCapabilityConfig {
3099 capability_ref: CapabilityId::new("gamma"),
3100 config: serde_json::json!({}),
3101 },
3102 AgentCapabilityConfig {
3103 capability_ref: CapabilityId::new("alpha"),
3104 config: serde_json::json!({}),
3105 },
3106 AgentCapabilityConfig {
3107 capability_ref: CapabilityId::new("beta"),
3108 config: serde_json::json!({}),
3109 },
3110 ];
3111
3112 let collected = collect_message_filters_only(&configs, ®istry);
3113
3114 let session_id: SessionId = Uuid::now_v7().into();
3115 let mut query = MessageQuery::new(session_id);
3116 collected.apply_message_filters(&mut query);
3117
3118 assert_eq!(query.filters.len(), 3);
3120 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3121 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3122 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3123 }
3124
3125 #[test]
3126 fn test_collect_message_filters_only_post_load_invoked() {
3127 use crate::message::Message;
3128
3129 struct PostLoadCap;
3130 struct PostLoadProvider;
3131
3132 impl Capability for PostLoadCap {
3133 fn id(&self) -> &str {
3134 "post_load_test"
3135 }
3136 fn name(&self) -> &str {
3137 "PostLoad Test"
3138 }
3139 fn description(&self) -> &str {
3140 "test"
3141 }
3142 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3143 Some(Arc::new(PostLoadProvider))
3144 }
3145 }
3146
3147 impl MessageFilterProvider for PostLoadProvider {
3148 fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
3149 fn priority(&self) -> i32 {
3150 0
3151 }
3152 fn post_load(&self, messages: &mut Vec<Message>, _config: &serde_json::Value) {
3153 messages.reverse();
3155 }
3156 }
3157
3158 let mut registry = CapabilityRegistry::new();
3159 registry.register(PostLoadCap);
3160
3161 let configs = vec![AgentCapabilityConfig {
3162 capability_ref: CapabilityId::new("post_load_test"),
3163 config: serde_json::json!({}),
3164 }];
3165
3166 let collected = collect_message_filters_only(&configs, ®istry);
3167
3168 let mut messages = vec![Message::user("first"), Message::user("second")];
3169 collected.apply_post_load_filters(&mut messages);
3170
3171 assert_eq!(messages[0].text(), Some("second"));
3173 assert_eq!(messages[1].text(), Some("first"));
3174 }
3175
3176 #[test]
3177 fn test_collect_model_view_providers_respects_compaction_capability_boundary() {
3178 use crate::tool_types::ToolCall;
3179
3180 fn tool_heavy_messages() -> Vec<Message> {
3181 let mut messages = vec![Message::user("inspect files repeatedly")];
3182 for index in 0..9 {
3183 let call_id = format!("call_{index}");
3184 messages.push(Message::assistant_with_tools(
3185 "",
3186 vec![ToolCall {
3187 id: call_id.clone(),
3188 name: "read_file".to_string(),
3189 arguments: serde_json::json!({"path": "/workspace/src/lib.rs"}),
3190 }],
3191 ));
3192 messages.push(Message::tool_result(
3193 call_id,
3194 Some(serde_json::json!({
3195 "path": "/workspace/src/lib.rs",
3196 "content": format!("{}{}", "large file line\n".repeat(1000), index),
3197 "total_lines": 1000,
3198 "lines_shown": {"start": 1, "end": 1000},
3199 "truncated": false
3200 })),
3201 None,
3202 ));
3203 }
3204 messages
3205 }
3206
3207 fn first_tool_result_is_masked(messages: &[Message]) -> bool {
3208 messages[2]
3209 .tool_result_content()
3210 .and_then(|result| result.result.as_ref())
3211 .and_then(|result| result.get("masked"))
3212 .and_then(|masked| masked.as_bool())
3213 .unwrap_or(false)
3214 }
3215
3216 let mut registry = CapabilityRegistry::new();
3217 registry.register(CompactionCapability);
3218 let context = ModelViewContext {
3219 session_id: SessionId::new(),
3220 prior_usage: None,
3221 };
3222
3223 let no_compaction = collect_model_view_providers(&[], ®istry);
3224 let unmasked = no_compaction.apply_model_view(tool_heavy_messages(), &context);
3225 assert!(!first_tool_result_is_masked(&unmasked));
3226
3227 let compaction = collect_model_view_providers(
3228 &[AgentCapabilityConfig {
3229 capability_ref: CapabilityId::new(COMPACTION_CAPABILITY_ID),
3230 config: serde_json::json!({}),
3231 }],
3232 ®istry,
3233 );
3234 let masked = compaction.apply_model_view(tool_heavy_messages(), &context);
3235 assert!(first_tool_result_is_masked(&masked));
3236 let last_tool = masked.last().unwrap().tool_result_content().unwrap();
3237 assert!(last_tool.result.as_ref().unwrap().get("content").is_some());
3238 }
3239
3240 #[tokio::test]
3250 async fn test_virtual_bash_capability_produces_bash_tool() {
3251 let registry = CapabilityRegistry::with_builtins();
3252 let collected =
3253 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3254
3255 let tool_names: Vec<&str> = collected
3256 .tool_definitions
3257 .iter()
3258 .map(|t| t.name())
3259 .collect();
3260 assert!(
3261 tool_names.contains(&"bash"),
3262 "virtual_bash capability must produce 'bash' tool, got: {:?}",
3263 tool_names
3264 );
3265 assert!(
3266 !collected.tools.is_empty(),
3267 "virtual_bash must provide tool implementations"
3268 );
3269 }
3270
3271 #[tokio::test]
3272 async fn test_generic_harness_capability_set_produces_bash_tool() {
3273 let generic_harness_caps = vec![
3276 "session_file_system".to_string(),
3277 "virtual_bash".to_string(),
3278 "web_fetch".to_string(),
3279 "session_storage".to_string(),
3280 "session".to_string(),
3281 "agent_instructions".to_string(),
3282 "skills".to_string(),
3283 "infinity_context".to_string(),
3284 "openai_tool_search".to_string(),
3285 ];
3286
3287 let registry = CapabilityRegistry::with_builtins();
3288 let collected = collect_capabilities(&generic_harness_caps, ®istry, &test_ctx()).await;
3289
3290 let tool_names: Vec<&str> = collected
3291 .tool_definitions
3292 .iter()
3293 .map(|t| t.name())
3294 .collect();
3295 assert!(
3296 tool_names.contains(&"bash"),
3297 "Generic Harness capabilities must produce 'bash' tool, got: {:?}",
3298 tool_names
3299 );
3300 }
3301
3302 #[tokio::test]
3303 async fn test_collect_capabilities_tool_count_matches_definitions() {
3304 let registry = CapabilityRegistry::with_builtins();
3307 let collected =
3308 collect_capabilities(&["virtual_bash".to_string()], ®istry, &test_ctx()).await;
3309
3310 assert_eq!(
3311 collected.tools.len(),
3312 collected.tool_definitions.len(),
3313 "tool implementations ({}) must match tool definitions ({})",
3314 collected.tools.len(),
3315 collected.tool_definitions.len(),
3316 );
3317 }
3318
3319 #[tokio::test]
3323 async fn test_collect_capabilities_resolves_dependencies() {
3324 let registry = CapabilityRegistry::with_builtins();
3327 let collected =
3328 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
3329
3330 assert!(
3332 collected
3333 .applied_ids
3334 .iter()
3335 .any(|id| id == "session_file_system"),
3336 "collect_capabilities must apply session_file_system as a dependency; applied_ids: {:?}",
3337 collected.applied_ids
3338 );
3339
3340 let tool_names: Vec<&str> = collected
3341 .tool_definitions
3342 .iter()
3343 .map(|t| t.name())
3344 .collect();
3345
3346 assert!(
3348 tool_names.contains(&"read_file") && tool_names.contains(&"write_file"),
3349 "collect_capabilities must resolve dependencies and include dependency tools, got: {:?}",
3350 tool_names
3351 );
3352
3353 assert_eq!(
3355 collected.tools.len(),
3356 collected.tool_definitions.len(),
3357 "dependency-added tools must have implementations, not just definitions"
3358 );
3359 }
3360
3361 #[test]
3362 fn test_defaults_do_not_include_bash() {
3363 let registry = crate::ToolRegistry::with_defaults();
3366 assert!(
3367 !registry.has("bash"),
3368 "with_defaults() must not include 'bash' — it comes from virtual_bash capability"
3369 );
3370 }
3371
3372 #[test]
3377 fn test_capability_features_default_empty() {
3378 let registry = CapabilityRegistry::with_builtins();
3379
3380 let noop = registry.get("noop").unwrap();
3382 assert!(noop.features().is_empty());
3383
3384 let current_time = registry.get("current_time").unwrap();
3385 assert!(current_time.features().is_empty());
3386 }
3387
3388 #[test]
3389 fn test_file_system_capability_features() {
3390 let registry = CapabilityRegistry::with_builtins();
3391
3392 let fs = registry.get("session_file_system").unwrap();
3393 assert_eq!(fs.features(), vec!["file_system"]);
3394 }
3395
3396 #[test]
3397 fn test_virtual_bash_capability_features() {
3398 let registry = CapabilityRegistry::with_builtins();
3399
3400 let bash = registry.get("virtual_bash").unwrap();
3401 assert_eq!(bash.features(), vec!["file_system"]);
3402 }
3403
3404 #[test]
3405 fn test_session_storage_capability_features() {
3406 let registry = CapabilityRegistry::with_builtins();
3407
3408 let storage = registry.get("session_storage").unwrap();
3409 let features = storage.features();
3410 assert!(features.contains(&"secrets"));
3411 assert!(features.contains(&"key_value"));
3412 }
3413
3414 #[test]
3415 fn test_session_schedule_capability_features() {
3416 let registry = CapabilityRegistry::with_builtins();
3417
3418 let schedule = registry.get("session_schedule").unwrap();
3419 assert_eq!(schedule.features(), vec!["schedules"]);
3420 }
3421
3422 #[test]
3423 fn test_session_sql_database_capability_features() {
3424 let registry = CapabilityRegistry::with_builtins();
3425
3426 let sql = registry.get("session_sql_database").unwrap();
3427 assert_eq!(sql.features(), vec!["sql_database"]);
3428 }
3429
3430 #[test]
3431 fn test_sample_data_capability_features() {
3432 let registry = CapabilityRegistry::with_builtins();
3433
3434 let sample = registry.get("sample_data").unwrap();
3435 assert_eq!(sample.features(), vec!["file_system"]);
3436 }
3437
3438 #[test]
3439 fn test_compute_features_empty() {
3440 let registry = CapabilityRegistry::with_builtins();
3441
3442 let features = compute_features(&[], ®istry);
3443 assert!(features.is_empty());
3444 }
3445
3446 #[test]
3447 fn test_compute_features_single_capability() {
3448 let registry = CapabilityRegistry::with_builtins();
3449
3450 let features = compute_features(&["session_schedule".to_string()], ®istry);
3451 assert_eq!(features, vec!["schedules"]);
3452 }
3453
3454 #[test]
3455 fn test_compute_features_multiple_capabilities() {
3456 let registry = CapabilityRegistry::with_builtins();
3457
3458 let features = compute_features(
3459 &[
3460 "session_file_system".to_string(),
3461 "session_storage".to_string(),
3462 "session_schedule".to_string(),
3463 ],
3464 ®istry,
3465 );
3466 assert!(features.contains(&"file_system".to_string()));
3467 assert!(features.contains(&"secrets".to_string()));
3468 assert!(features.contains(&"key_value".to_string()));
3469 assert!(features.contains(&"schedules".to_string()));
3470 }
3471
3472 #[test]
3473 fn test_compute_features_deduplicates() {
3474 let registry = CapabilityRegistry::with_builtins();
3475
3476 let features = compute_features(
3478 &[
3479 "session_file_system".to_string(),
3480 "virtual_bash".to_string(),
3481 ],
3482 ®istry,
3483 );
3484 let file_system_count = features.iter().filter(|f| *f == "file_system").count();
3485 assert_eq!(file_system_count, 1, "file_system should appear only once");
3486 }
3487
3488 #[test]
3489 fn test_compute_features_includes_dependency_features() {
3490 let registry = CapabilityRegistry::with_builtins();
3491
3492 let features = compute_features(&["virtual_bash".to_string()], ®istry);
3494 assert!(features.contains(&"file_system".to_string()));
3495 }
3496
3497 #[test]
3498 fn test_compute_features_generic_harness_set() {
3499 let registry = CapabilityRegistry::with_builtins();
3500
3501 let features = compute_features(
3503 &[
3504 "session_file_system".to_string(),
3505 "virtual_bash".to_string(),
3506 "session_storage".to_string(),
3507 "session".to_string(),
3508 "session_schedule".to_string(),
3509 ],
3510 ®istry,
3511 );
3512 assert!(features.contains(&"file_system".to_string()));
3513 assert!(features.contains(&"secrets".to_string()));
3514 assert!(features.contains(&"key_value".to_string()));
3515 assert!(features.contains(&"schedules".to_string()));
3516 }
3517
3518 #[test]
3519 fn test_compute_features_unknown_capability_ignored() {
3520 let registry = CapabilityRegistry::with_builtins();
3521
3522 let features = compute_features(
3523 &["unknown_cap".to_string(), "session_schedule".to_string()],
3524 ®istry,
3525 );
3526 assert_eq!(features, vec!["schedules"]);
3527 }
3528
3529 #[test]
3530 fn test_risk_level_ordering() {
3531 assert!(RiskLevel::Low < RiskLevel::Medium);
3532 assert!(RiskLevel::Medium < RiskLevel::High);
3533 }
3534
3535 #[test]
3536 fn test_risk_level_serde_roundtrip() {
3537 let high = RiskLevel::High;
3538 let json = serde_json::to_string(&high).unwrap();
3539 assert_eq!(json, "\"high\"");
3540 let back: RiskLevel = serde_json::from_str(&json).unwrap();
3541 assert_eq!(back, RiskLevel::High);
3542 }
3543
3544 #[test]
3545 fn test_capability_risk_levels() {
3546 let registry = CapabilityRegistry::with_builtins();
3547
3548 let bash = registry.get("virtual_bash").unwrap();
3550 assert_eq!(bash.risk_level(), RiskLevel::High);
3551
3552 let fetch = registry.get("web_fetch").unwrap();
3554 assert_eq!(fetch.risk_level(), RiskLevel::High);
3555
3556 let noop = registry.get("noop").unwrap();
3558 assert_eq!(noop.risk_level(), RiskLevel::Low);
3559 }
3560
3561 #[tokio::test]
3566 async fn test_apply_capabilities_openai_tool_search() {
3567 let registry = CapabilityRegistry::with_builtins();
3568 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
3569
3570 let applied = apply_capabilities(
3571 base_runtime_agent.clone(),
3572 &["openai_tool_search".to_string()],
3573 ®istry,
3574 &test_ctx(),
3575 )
3576 .await;
3577
3578 assert_eq!(
3580 applied.runtime_agent.system_prompt,
3581 base_runtime_agent.system_prompt
3582 );
3583 assert!(applied.tool_registry.is_empty());
3584 assert_eq!(applied.applied_ids, vec!["openai_tool_search"]);
3585
3586 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
3588 assert!(ts.enabled);
3589 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
3590 }
3591
3592 #[tokio::test]
3593 async fn test_apply_capabilities_openai_tool_search_with_other_capabilities() {
3594 let registry = CapabilityRegistry::with_builtins();
3595 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
3596
3597 let applied = apply_capabilities(
3598 base_runtime_agent,
3599 &[
3600 "current_time".to_string(),
3601 "openai_tool_search".to_string(),
3602 "test_math".to_string(),
3603 ],
3604 ®istry,
3605 &test_ctx(),
3606 )
3607 .await;
3608
3609 assert!(applied.tool_registry.has("get_current_time"));
3611 assert!(applied.tool_registry.has("add"));
3612 assert!(applied.tool_registry.has("subtract"));
3613 assert!(applied.tool_registry.has("multiply"));
3614 assert!(applied.tool_registry.has("divide"));
3615
3616 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
3618 assert!(ts.enabled);
3619 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
3620 }
3621
3622 #[tokio::test]
3623 async fn test_collect_capabilities_tool_search_custom_threshold() {
3624 let registry = CapabilityRegistry::with_builtins();
3625
3626 let configs = vec![AgentCapabilityConfig {
3627 capability_ref: CapabilityId::new("openai_tool_search"),
3628 config: serde_json::json!({"threshold": 5}),
3629 }];
3630
3631 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3632
3633 let ts = collected.tool_search.as_ref().unwrap();
3634 assert!(ts.enabled);
3635 assert_eq!(ts.threshold, 5);
3636 }
3637
3638 #[tokio::test]
3639 async fn test_collect_capabilities_no_tool_search_without_capability() {
3640 let registry = CapabilityRegistry::with_builtins();
3641
3642 let configs = vec![AgentCapabilityConfig {
3643 capability_ref: CapabilityId::new("current_time"),
3644 config: serde_json::json!({}),
3645 }];
3646
3647 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3648
3649 assert!(collected.tool_search.is_none());
3650 }
3651
3652 #[tokio::test]
3653 async fn test_collect_capabilities_tool_search_category_propagation() {
3654 let registry = CapabilityRegistry::with_builtins();
3655
3656 let configs = vec![
3658 AgentCapabilityConfig {
3659 capability_ref: CapabilityId::new("test_math"),
3660 config: serde_json::json!({}),
3661 },
3662 AgentCapabilityConfig {
3663 capability_ref: CapabilityId::new("openai_tool_search"),
3664 config: serde_json::json!({}),
3665 },
3666 ];
3667
3668 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3669
3670 assert!(collected.tool_search.is_some());
3672
3673 for tool_def in &collected.tool_definitions {
3675 if ["add", "subtract", "multiply", "divide"].contains(&tool_def.name()) {
3677 assert!(
3678 tool_def.category().is_some(),
3679 "Tool {} should have a category from its capability",
3680 tool_def.name()
3681 );
3682 }
3683 }
3684 }
3685
3686 #[tokio::test]
3687 async fn test_apply_capabilities_prompt_caching() {
3688 let registry = CapabilityRegistry::with_builtins();
3689 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
3690
3691 let applied = apply_capabilities(
3692 base_runtime_agent.clone(),
3693 &["prompt_caching".to_string()],
3694 ®istry,
3695 &test_ctx(),
3696 )
3697 .await;
3698
3699 assert_eq!(
3700 applied.runtime_agent.system_prompt,
3701 base_runtime_agent.system_prompt
3702 );
3703 assert!(applied.tool_registry.is_empty());
3704 assert_eq!(applied.applied_ids, vec!["prompt_caching"]);
3705
3706 let prompt_cache = applied.runtime_agent.prompt_cache.as_ref().unwrap();
3707 assert!(prompt_cache.enabled);
3708 assert_eq!(
3709 prompt_cache.strategy,
3710 crate::llm_driver_registry::PromptCacheStrategy::Auto
3711 );
3712 assert!(prompt_cache.gemini_cached_content.is_none());
3713 }
3714
3715 #[tokio::test]
3716 async fn test_collect_capabilities_prompt_caching_custom_strategy() {
3717 let registry = CapabilityRegistry::with_builtins();
3718
3719 let configs = vec![AgentCapabilityConfig {
3720 capability_ref: CapabilityId::new("prompt_caching"),
3721 config: serde_json::json!({"strategy": "auto"}),
3722 }];
3723
3724 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3725
3726 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
3727 assert!(prompt_cache.enabled);
3728 assert_eq!(
3729 prompt_cache.strategy,
3730 crate::llm_driver_registry::PromptCacheStrategy::Auto
3731 );
3732 assert!(prompt_cache.gemini_cached_content.is_none());
3733 }
3734
3735 #[tokio::test]
3736 async fn test_collect_capabilities_prompt_caching_gemini_cached_content() {
3737 let registry = CapabilityRegistry::with_builtins();
3738
3739 let configs = vec![AgentCapabilityConfig {
3740 capability_ref: CapabilityId::new("prompt_caching"),
3741 config: serde_json::json!({
3742 "strategy": "auto",
3743 "gemini_cached_content": "cachedContents/demo-cache"
3744 }),
3745 }];
3746
3747 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3748
3749 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
3750 assert_eq!(
3751 prompt_cache.gemini_cached_content.as_deref(),
3752 Some("cachedContents/demo-cache")
3753 );
3754 }
3755
3756 struct SkillContributingCapability;
3761
3762 impl Capability for SkillContributingCapability {
3763 fn id(&self) -> &str {
3764 "contributes_skills"
3765 }
3766 fn name(&self) -> &str {
3767 "Contributes Skills"
3768 }
3769 fn description(&self) -> &str {
3770 "Test capability that contributes skills."
3771 }
3772 fn contribute_skills(&self) -> Vec<SkillContribution> {
3773 vec![
3774 SkillContribution::new("alpha-skill", "Alpha skill desc", "# Alpha\nDo alpha.")
3775 .with_files(vec![(
3776 "scripts/a.sh".to_string(),
3777 "#!/bin/sh\necho a\n".to_string(),
3778 )]),
3779 SkillContribution::new("beta-skill", "Beta skill desc", "# Beta\nDo beta.")
3780 .with_user_invocable(false),
3781 ]
3782 }
3783 }
3784
3785 fn skill_md_from_entries(entries: &HashMap<String, MountEntry>) -> &str {
3786 match &entries.get("SKILL.md").expect("SKILL.md missing").source {
3787 MountSource::InlineFile { content, .. } => content.as_str(),
3788 _ => panic!("Expected InlineFile for SKILL.md"),
3789 }
3790 }
3791
3792 #[tokio::test]
3793 async fn test_contribute_skills_normalized_to_mounts() {
3794 let mut registry = CapabilityRegistry::new();
3795 registry.register(SkillContributingCapability);
3796
3797 let configs = vec![AgentCapabilityConfig {
3798 capability_ref: CapabilityId::new("contributes_skills"),
3799 config: serde_json::json!({}),
3800 }];
3801
3802 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3803
3804 let skill_mounts: Vec<_> = collected
3805 .mounts
3806 .iter()
3807 .filter(|m| m.path.starts_with("/.agents/skills/"))
3808 .collect();
3809 assert_eq!(skill_mounts.len(), 2);
3810
3811 for m in &skill_mounts {
3814 assert!(m.is_readonly());
3815 assert_eq!(m.capability_id, "contributes_skills");
3816 }
3817
3818 let alpha = skill_mounts
3819 .iter()
3820 .find(|m| m.path == "/.agents/skills/alpha-skill")
3821 .expect("alpha-skill mount missing");
3822 match &alpha.source {
3823 MountSource::InlineDirectory { entries } => {
3824 assert!(entries.contains_key("SKILL.md"));
3825 assert!(entries.contains_key("scripts/a.sh"));
3826 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
3827 assert_eq!(parsed.name, "alpha-skill");
3828 assert!(parsed.user_invocable);
3829 }
3830 _ => panic!("Expected InlineDirectory"),
3831 }
3832
3833 let beta = skill_mounts
3834 .iter()
3835 .find(|m| m.path == "/.agents/skills/beta-skill")
3836 .expect("beta-skill mount missing");
3837 match &beta.source {
3838 MountSource::InlineDirectory { entries } => {
3839 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
3840 assert!(!parsed.user_invocable);
3841 }
3842 _ => panic!("Expected InlineDirectory"),
3843 }
3844 }
3845
3846 #[tokio::test]
3847 async fn test_contribute_skills_default_empty() {
3848 let mut registry = CapabilityRegistry::new();
3851 registry.register(FilterTestCapability { priority: 0 });
3852
3853 let configs = vec![AgentCapabilityConfig {
3854 capability_ref: CapabilityId::new("filter_test"),
3855 config: serde_json::json!({}),
3856 }];
3857
3858 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3859 assert!(
3860 collected
3861 .mounts
3862 .iter()
3863 .all(|m| !m.path.starts_with("/.agents/skills/"))
3864 );
3865 }
3866}