1use crate::command::{
22 CommandDescriptor, CommandExecutionContext, CommandResult, ExecuteCommandRequest,
23};
24use crate::deployment::DeploymentGrade;
25use crate::events::TokenUsage;
26use crate::mcp_server::{ScopedMcpServers, merge_scoped_mcp_servers};
27use crate::message::Message;
28use crate::message_filter::MessageFilterProvider;
29use crate::runtime_agent::RuntimeAgent;
30use crate::tool_types::{ToolCall, ToolDefinition};
31use crate::tools::{Tool, ToolRegistry};
32use crate::traits::SessionFileSystem;
33use crate::typed_id::SessionId;
34use async_trait::async_trait;
35use serde::{Deserialize, Serialize};
36use std::collections::HashMap;
37use std::sync::Arc;
38
39pub struct IntegrationPlugin {
63 pub experimental_only: bool,
65 pub feature_flag: Option<&'static str>,
68 pub factory: fn() -> Box<dyn Capability>,
70}
71
72inventory::collect!(IntegrationPlugin);
73
74pub use crate::capability_types::{
76 AgentCapabilityConfig, CapabilityId, CapabilityStatus, MountAccess, MountDirectoryBuilder,
77 MountEntry, MountPoint, MountSource,
78};
79
80mod a2a_delegation;
85#[cfg(feature = "ui-capabilities")]
86mod a2ui;
87mod agent_handoff;
88mod agent_instructions;
89pub mod attach_skill;
90mod auto_tool_search;
91mod background_execution;
92mod bashkit_shell;
93mod btw;
94mod budgeting;
95pub mod compaction;
96mod current_time;
97mod data_knowledge;
98mod declarative;
99mod fake_aws;
100mod fake_crm;
101mod fake_financial;
102mod fake_warehouse;
103mod file_system;
104mod human_intent;
105mod infinity_context;
106mod knowledge_base;
107mod loop_detection;
108mod lua;
109mod lua_code_mode;
110pub mod mcp;
111mod memory;
112mod message_metadata;
113mod monitors;
114mod noop;
115mod openai_tool_search;
116#[cfg(feature = "ui-capabilities")]
117mod openui;
118mod platform_management;
119mod prompt_caching;
120mod prompt_canary_guardrail;
121mod research;
122mod sample_data;
123mod self_budget;
124mod session;
125mod session_sandbox;
126mod session_schedule;
127mod session_sql_database;
128mod session_storage;
129mod session_tasks;
130mod skills;
131mod stateless_todo_list;
132mod subagents;
133mod system_commands;
134mod test_math;
135mod test_weather;
136mod tool_output_persistence;
137mod tool_search;
138pub mod user_hooks;
139mod web_fetch;
140mod web_fetch_egress;
141
142pub use a2a_delegation::{
144 A2A_AGENT_DELEGATION_CAPABILITY_ID, A2aAgentDelegationCapability, CancelAgentTool,
145 GetAgentRunsTool, MessageAgentTool, SpawnAgentTool, WaitAgentTool,
146};
147#[cfg(feature = "ui-capabilities")]
148pub use a2ui::{A2UI_CAPABILITY_ID, A2UiCapability};
149pub use agent_handoff::{
150 AGENT_HANDOFF_CAPABILITY_ID, AgentHandoffCapability, GetAgentHandoffsTool,
151 MessageAgentHandoffTool, StartAgentHandoffTool,
152};
153pub use agent_instructions::{
154 AGENT_INSTRUCTIONS_CAPABILITY_ID, AGENTS_MD_PATH, AgentInstructionsCapability,
155 AgentInstructionsConfig, DEFAULT_AGENT_INSTRUCTIONS_FILE, MAX_AGENT_INSTRUCTIONS_FILES,
156 MAX_AGENTS_MD_SIZE, format_agents_md_content, format_instruction_file_content,
157};
158pub use attach_skill::{
159 AttachSkillCapability, SKILL_CAPABILITY_PREFIX, SKILLS_DISCOVERY_PATH, SkillContribution,
160 SkillInstructions, SkillMeta, SkillSource, discover_skills_from_entries, is_skill_capability,
161 parse_skill_capability_id, reconstruct_skill_md, skill_capability_id,
162};
163pub use auto_tool_search::{AUTO_TOOL_SEARCH_CAPABILITY_ID, AutoToolSearchCapability};
164pub use background_execution::{BACKGROUND_EXECUTION_CAPABILITY_ID, BackgroundExecutionCapability};
165pub use btw::{BTW_CAPABILITY_ID, BtwCapability};
166pub use budgeting::BudgetingCapability;
167pub use compaction::{
168 COMPACTION_CAPABILITY_ID, CompactionCapability, CompactionConfig, CompactionStep,
169 CompactionStrategy, CostControlConfig, CostControlMaskingResult, HierarchicalMemoryConfig,
170 MaskingSummaryFormat, MemoryTier, ObservationMaskingConfig, ObservationMaskingResult,
171 SessionCompactionMetrics, SummarizationConfig, aggressive_trim, apply_cost_control_masking,
172 apply_hierarchical_memory, apply_observation_masking, build_model_view_messages,
173 build_summarization_prompt, build_summary_message, classify_memory_tiers, estimate_tokens,
174 estimate_total_tokens, format_messages_for_summarization, should_compact_proactively,
175};
176pub use current_time::{CurrentTimeCapability, GetCurrentTimeTool};
177pub use data_knowledge::DataKnowledgeCapability;
178pub use declarative::{
179 DECLARATIVE_CAPABILITY_PREFIX, DeclarativeCapabilityDefinition, DeclarativeCapabilityFile,
180 DeclarativeCapabilitySkill, declarative_capability_id, declarative_capability_info,
181 hydrate_declarative_capability_config, is_declarative_capability,
182 parse_declarative_capability_id, validate_declarative_capability_definition,
183};
184pub use fake_aws::{
185 AwsCreateEc2InstanceTool, AwsCreateIamUserTool, AwsCreateRdsDatabaseTool,
186 AwsCreateS3BucketTool, AwsGetCloudWatchMetricsTool, AwsListEc2InstancesTool,
187 AwsListIamUsersTool, AwsListRdsDatabasesTool, AwsListS3BucketsTool, AwsListSecurityGroupsTool,
188 AwsStopEc2InstanceTool, FakeAwsCapability,
189};
190pub use fake_crm::{
191 CrmAddInteractionTool, CrmCreateCustomerTool, CrmCreateTicketTool, CrmGetCustomerTool,
192 CrmListCustomersTool, CrmListTicketsTool, CrmSearchCustomersTool, CrmUpdateTicketTool,
193 FakeCrmCapability,
194};
195pub use fake_financial::{
196 FakeFinancialCapability, FinanceCreateBudgetTool, FinanceCreateTransactionTool,
197 FinanceForecastCashFlowTool, FinanceGetBalanceTool, FinanceGetExpenseReportTool,
198 FinanceGetRevenueReportTool, FinanceListBudgetsTool, FinanceListTransactionsTool,
199};
200pub use fake_warehouse::{
201 FakeWarehouseCapability, WarehouseCreateInvoiceTool, WarehouseCreateOrderTool,
202 WarehouseCreateShipmentTool, WarehouseGetInventoryTool, WarehouseInventoryReportTool,
203 WarehouseListOrdersTool, WarehouseListShipmentsTool, WarehouseProcessReturnTool,
204 WarehouseUpdateInventoryTool, WarehouseUpdateShipmentStatusTool,
205};
206pub use file_system::{
207 DeleteFileTool, EditFileTool, FileSystemCapability, GrepFilesTool, ListDirectoryTool,
208 ReadFileTool, StatFileTool, WriteFileTool,
209};
210pub use human_intent::{HUMAN_INTENT_CAPABILITY_ID, HumanIntentCapability};
211pub use infinity_context::{
212 INFINITY_CONTEXT_CAPABILITY_ID, InfinityContextCapability, QueryHistoryTool,
213};
214pub use knowledge_base::{
215 KNOWLEDGE_BASE_CAPABILITY_ID, KnowledgeBaseCapability, KnowledgeBaseConfig,
216 validate_knowledge_base_config,
217};
218pub use loop_detection::LoopDetectionCapability;
219pub use lua::{LUA_CAPABILITY_ID, LuaCapability, LuaTool, LuaVfs, is_code_mode_eligible};
220pub use lua_code_mode::{LUA_CODE_MODE_CAPABILITY_ID, LuaCodeModeCapability};
221pub use mcp::{
222 MCP_CAPABILITY_PREFIX, McpCapability, is_mcp_capability, mcp_capability_id,
223 parse_mcp_capability_id,
224};
225pub use memory::{MEMORY_CAPABILITY_ID, MemoryCapability};
226pub use message_metadata::{
227 MESSAGE_METADATA_CAPABILITY_ID, MessageMetadataCapability, MessageMetadataConfig,
228 MessageMetadataField, render_annotation,
229};
230pub use noop::NoopCapability;
231pub use openai_tool_search::{
232 DEFAULT_TOOL_SEARCH_THRESHOLD, OPENAI_TOOL_SEARCH_CAPABILITY_ID, OpenAiToolSearchCapability,
233 model_supports_native_tool_search,
234};
235#[cfg(feature = "ui-capabilities")]
236pub use openui::{OPENUI_CAPABILITY_ID, OpenUiCapability};
237pub use platform_management::{
238 ManageAgentsTool, ManageHarnessesTool, ManageSessionsTool, PlatformManagementCapability,
239 ReadAgentsTool, ReadCapabilitiesTool, ReadHarnessesTool, ReadSessionsTool,
240 SessionReadMessagesTool, SessionReadResponseTool, SessionSendMessageTool,
241};
242pub use prompt_caching::{PROMPT_CACHING_CAPABILITY_ID, PromptCachingCapability};
243pub use prompt_canary_guardrail::{
244 DEFAULT_REPLACEMENT as PROMPT_CANARY_DEFAULT_REPLACEMENT,
245 PROMPT_CANARY_GUARDRAIL_CAPABILITY_ID, PromptCanaryGuardrailCapability,
246 REASON_CODE_SYSTEM_PROMPT_LEAK,
247};
248pub use research::ResearchCapability;
249pub use sample_data::SampleDataCapability;
250pub use self_budget::{SELF_BUDGET_CAPABILITY_ID, SelfBudgetCapability};
251pub use session::{GetSessionInfoTool, SessionCapability, WriteSessionTitleTool};
252pub use session_sandbox::{
253 SESSION_SANDBOX_CAPABILITY_ID, SandboxExecTool, SandboxManageTool, SandboxReadFileTool,
254 SandboxStatusTool, SandboxWriteFileTool, SessionSandboxCapability,
255};
256pub use session_schedule::{
257 CancelScheduleTool, CreateScheduleTool, ListSchedulesTool, SESSION_SCHEDULE_CAPABILITY_ID,
258 SessionScheduleCapability,
259};
260pub use session_sql_database::{
261 SessionSqlDatabaseCapability, SqlExecuteTool, SqlQueryTool, SqlSchemaTool,
262};
263pub use session_storage::{KvStoreTool, SecretStoreTool, SessionStorageCapability};
264pub use session_tasks::{SESSION_TASKS_CAPABILITY_ID, SessionTasksCapability};
265pub use skills::{SKILLS_CAPABILITY_ID, SkillsCapability};
266pub use stateless_todo_list::{StatelessTodoListCapability, WriteTodosTool};
267pub use subagents::SubagentCapability;
268pub use bashkit_shell::{BashTool, BashkitShellCapability, SessionFileSystemAdapter};
270pub use system_commands::{SYSTEM_COMMANDS_CAPABILITY_ID, SystemCommandsCapability};
271pub use test_math::{AddTool, DivideTool, MultiplyTool, SubtractTool, TestMathCapability};
272pub use test_weather::{GetForecastTool, GetWeatherTool, TestWeatherCapability};
273pub use tool_output_persistence::{PersistOutputHook, ToolOutputPersistenceCapability};
274pub use tool_search::{
275 TOOL_SEARCH_CAPABILITY_ID, TOOL_SEARCH_TOOL_NAME, ToolSearchCapability, ToolSearchTool,
276};
277pub use user_hooks::UserHooksCapability;
278pub use web_fetch::{
279 BotAuthPublicKey, WebFetchCapability, WebFetchTool, derive_bot_auth_public_key,
280};
281
282pub struct SystemPromptContext {
292 pub session_id: SessionId,
294 pub locale: Option<String>,
296 pub file_store: Option<Arc<dyn SessionFileSystem>>,
298 pub model: Option<String>,
304}
305
306impl SystemPromptContext {
307 pub fn without_file_store(session_id: SessionId) -> Self {
309 Self {
310 session_id,
311 locale: None,
312 file_store: None,
313 model: None,
314 }
315 }
316
317 pub fn with_model(mut self, model: impl Into<String>) -> Self {
319 self.model = Some(model.into());
320 self
321 }
322}
323
324#[derive(Debug, Clone)]
376pub struct CapabilityLocalization {
377 pub locale: &'static str,
379 pub name: Option<&'static str>,
381 pub description: Option<&'static str>,
383 pub config_description: Option<&'static str>,
388 pub config_overlay: Option<serde_json::Value>,
394}
395
396impl CapabilityLocalization {
397 pub fn text(locale: &'static str, name: &'static str, description: &'static str) -> Self {
399 Self {
400 locale,
401 name: Some(name),
402 description: Some(description),
403 config_description: None,
404 config_overlay: None,
405 }
406 }
407}
408
409pub fn resolve_localized_field<T>(
413 localizations: &[CapabilityLocalization],
414 locale: Option<&str>,
415 field: impl Fn(&CapabilityLocalization) -> Option<T>,
416) -> Option<T> {
417 let mut candidates: Vec<String> = Vec::new();
418 if let Some(raw) = locale {
419 let normalized = raw.trim().replace('_', "-").to_lowercase();
420 if !normalized.is_empty() {
421 if let Some((language, _)) = normalized.split_once('-') {
422 let language = language.to_string();
423 candidates.push(normalized);
424 candidates.push(language);
425 } else {
426 candidates.push(normalized);
427 }
428 }
429 }
430 candidates.push("en".to_string());
431
432 for candidate in candidates {
433 let hit = localizations
434 .iter()
435 .find(|entry| entry.locale.eq_ignore_ascii_case(&candidate))
436 .and_then(&field);
437 if hit.is_some() {
438 return hit;
439 }
440 }
441 None
442}
443
444#[async_trait]
445pub trait Capability: Send + Sync {
446 fn id(&self) -> &str;
448
449 fn aliases(&self) -> Vec<&'static str> {
458 vec![]
459 }
460
461 fn name(&self) -> &str;
463
464 fn description(&self) -> &str;
466
467 fn localizations(&self) -> Vec<CapabilityLocalization> {
472 vec![]
473 }
474
475 fn localized_name(&self, locale: Option<&str>) -> String {
478 resolve_localized_field(&self.localizations(), locale, |entry| entry.name)
479 .unwrap_or_else(|| self.name())
480 .to_string()
481 }
482
483 fn localized_description(&self, locale: Option<&str>) -> String {
485 resolve_localized_field(&self.localizations(), locale, |entry| entry.description)
486 .unwrap_or_else(|| self.description())
487 .to_string()
488 }
489
490 fn describe_schema(&self, locale: Option<&str>) -> Option<String> {
494 resolve_localized_field(&self.localizations(), locale, |entry| {
495 entry.config_description
496 })
497 .map(str::to_string)
498 }
499
500 fn status(&self) -> CapabilityStatus {
502 CapabilityStatus::Available
503 }
504
505 fn icon(&self) -> Option<&str> {
507 None
508 }
509
510 fn category(&self) -> Option<&str> {
512 None
513 }
514
515 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
526 None
527 }
528
529 fn system_prompt_addition(&self) -> Option<&str> {
549 None
550 }
551
552 async fn system_prompt_contribution(&self, _ctx: &SystemPromptContext) -> Option<String> {
564 self.system_prompt_addition().map(|addition| {
565 format!(
566 "<capability id=\"{}\">\n{}\n</capability>",
567 self.id(),
568 addition
569 )
570 })
571 }
572
573 fn system_prompt_preview(&self) -> Option<String> {
579 self.system_prompt_addition().map(|s| s.to_string())
580 }
581
582 fn tools(&self) -> Vec<Box<dyn Tool>> {
584 vec![]
585 }
586
587 fn tools_with_config(&self, _config: &serde_json::Value) -> Vec<Box<dyn Tool>> {
595 self.tools()
596 }
597
598 async fn system_prompt_contribution_with_config(
605 &self,
606 ctx: &SystemPromptContext,
607 _config: &serde_json::Value,
608 ) -> Option<String> {
609 self.system_prompt_contribution(ctx).await
610 }
611
612 fn tool_definitions(&self) -> Vec<ToolDefinition> {
615 self.tools().iter().map(|t| t.to_definition()).collect()
616 }
617
618 fn mounts(&self) -> Vec<MountPoint> {
626 vec![]
627 }
628
629 fn dependencies(&self) -> Vec<&'static str> {
638 vec![]
639 }
640
641 fn features(&self) -> Vec<&'static str> {
656 vec![]
657 }
658
659 fn config_schema(&self) -> Option<serde_json::Value> {
665 None
666 }
667
668 fn config_ui_schema(&self) -> Option<serde_json::Value> {
673 None
674 }
675
676 fn validate_config(&self, _config: &serde_json::Value) -> Result<(), String> {
682 Ok(())
683 }
684
685 fn mcp_servers(&self) -> ScopedMcpServers {
691 ScopedMcpServers::default()
692 }
693
694 fn mcp_servers_with_config(&self, _config: &serde_json::Value) -> ScopedMcpServers {
696 self.mcp_servers()
697 }
698
699 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
712 None
713 }
714
715 fn model_view_provider(&self) -> Option<Arc<dyn ModelViewProvider>> {
723 None
724 }
725
726 fn pre_tool_use_hooks(&self) -> Vec<Arc<dyn crate::atoms::PreToolUseHook>> {
737 vec![]
738 }
739
740 fn post_tool_exec_hooks(&self) -> Vec<Arc<dyn crate::atoms::PostToolExecHook>> {
748 vec![]
749 }
750
751 fn tool_definition_hooks(&self) -> Vec<Arc<dyn ToolDefinitionHook>> {
760 vec![]
761 }
762
763 fn tool_definition_hooks_with_config(
768 &self,
769 _config: &serde_json::Value,
770 ) -> Vec<Arc<dyn ToolDefinitionHook>> {
771 self.tool_definition_hooks()
772 }
773
774 fn tool_definition_hooks_with_context(
784 &self,
785 _ctx: &SystemPromptContext,
786 config: &serde_json::Value,
787 ) -> Vec<Arc<dyn ToolDefinitionHook>> {
788 self.tool_definition_hooks_with_config(config)
789 }
790
791 fn tool_call_hooks(&self) -> Vec<Arc<dyn ToolCallHook>> {
799 vec![]
800 }
801
802 fn user_hooks(&self) -> Vec<crate::user_hook_types::UserHookSpec> {
818 vec![]
819 }
820
821 fn user_hooks_with_config(
827 &self,
828 _config: &serde_json::Value,
829 ) -> Vec<crate::user_hook_types::UserHookSpec> {
830 self.user_hooks()
831 }
832
833 fn risk_level(&self) -> RiskLevel {
841 RiskLevel::Low
842 }
843
844 fn commands(&self) -> Vec<CommandDescriptor> {
852 vec![]
853 }
854
855 async fn execute_command(
869 &self,
870 request: &ExecuteCommandRequest,
871 _ctx: &CommandExecutionContext,
872 ) -> crate::error::Result<CommandResult> {
873 Err(crate::error::AgentLoopError::config(format!(
874 "capability {} declared command /{} but does not implement execute_command",
875 self.id(),
876 request.name,
877 )))
878 }
879
880 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
888 vec![]
889 }
890
891 fn contribute_skills(&self) -> Vec<SkillContribution> {
901 vec![]
902 }
903
904 fn output_guardrails(&self) -> Vec<Arc<dyn crate::output_guardrail::OutputGuardrail>> {
915 vec![]
916 }
917}
918
919pub trait ToolDefinitionHook: Send + Sync {
920 fn transform(&self, tools: Vec<ToolDefinition>) -> Vec<ToolDefinition>;
921
922 fn applies_with_native_tool_search(&self) -> bool {
927 true
928 }
929}
930
931pub trait ToolCallHook: Send + Sync {
932 fn narration(
933 &self,
934 _tool_def: Option<&ToolDefinition>,
935 _tool_call: &ToolCall,
936 _phase: crate::tool_narration::ToolNarrationPhase,
937 _locale: Option<&str>,
938 ) -> Option<String> {
939 None
940 }
941
942 fn transform_for_execution(&self, tool_call: ToolCall) -> ToolCall {
943 tool_call
944 }
945}
946
947#[derive(
951 Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
952)]
953#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
954#[cfg_attr(feature = "openapi", schema(example = "low"))]
955#[serde(rename_all = "lowercase")]
956pub enum RiskLevel {
957 Low,
959 Medium,
961 High,
963}
964
965#[derive(Debug, Clone, Serialize, Deserialize)]
971#[serde(rename_all = "snake_case")]
972pub enum BlueprintModel {
973 Fixed(String),
975 Default(String),
977 Inherit,
979}
980
981pub struct AgentBlueprint {
987 pub id: &'static str,
989 pub name: &'static str,
991 pub description: &'static str,
993 pub model: BlueprintModel,
995 pub system_prompt: &'static str,
997 pub tools: Vec<Box<dyn Tool>>,
999 pub max_turns: Option<usize>,
1001 pub config_schema: Option<serde_json::Value>,
1003}
1004
1005impl AgentBlueprint {
1006 pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
1008 self.tools.iter().map(|t| t.to_definition()).collect()
1009 }
1010}
1011
1012impl std::fmt::Debug for AgentBlueprint {
1013 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1014 f.debug_struct("AgentBlueprint")
1015 .field("id", &self.id)
1016 .field("name", &self.name)
1017 .field("model", &self.model)
1018 .field("tool_count", &self.tools.len())
1019 .field("max_turns", &self.max_turns)
1020 .finish()
1021 }
1022}
1023
1024#[derive(Clone)]
1051pub struct CapabilityRegistry {
1052 capabilities: HashMap<String, Arc<dyn Capability>>,
1053 aliases: HashMap<String, String>,
1055}
1056
1057impl CapabilityRegistry {
1058 pub fn new() -> Self {
1060 Self {
1061 capabilities: HashMap::new(),
1062 aliases: HashMap::new(),
1063 }
1064 }
1065
1066 pub fn with_builtins() -> Self {
1071 Self::with_builtins_for_grade(DeploymentGrade::from_env())
1072 }
1073
1074 pub fn with_builtins_for_grade(grade: DeploymentGrade) -> Self {
1079 let mut registry = Self::new();
1080
1081 registry.register(AgentInstructionsCapability);
1083 registry.register(HumanIntentCapability);
1084 registry.register(NoopCapability);
1085 registry.register(CurrentTimeCapability);
1086 registry.register(MessageMetadataCapability);
1087 registry.register(ResearchCapability);
1088 registry.register(PlatformManagementCapability);
1089 registry.register(FileSystemCapability);
1090 registry.register(MemoryCapability);
1091 registry.register(SessionStorageCapability);
1092 registry.register(SessionCapability);
1093 registry.register(SessionSqlDatabaseCapability);
1094 registry.register(TestMathCapability);
1095 registry.register(TestWeatherCapability);
1096 registry.register(StatelessTodoListCapability);
1097 registry.register(WebFetchCapability::from_env());
1098 registry.register(BashkitShellCapability);
1099 registry.register(BackgroundExecutionCapability);
1100 registry.register(SessionScheduleCapability);
1101 registry.register(BtwCapability);
1102 registry.register(InfinityContextCapability);
1103 registry.register(budgeting::BudgetingCapability);
1104 registry.register(SelfBudgetCapability);
1105 registry.register(CompactionCapability);
1106
1107 registry.register(OpenAiToolSearchCapability::new());
1109 registry.register(ToolSearchCapability::new());
1111 registry.register(AutoToolSearchCapability::new());
1113 registry.register(PromptCachingCapability::new());
1114
1115 registry.register(SkillsCapability);
1117
1118 registry.register(SubagentCapability);
1120
1121 registry.register(SessionTasksCapability);
1123
1124 if crate::FeatureFlags::from_env(&grade).agent_delegation {
1128 registry.register(AgentHandoffCapability);
1129 registry.register(A2aAgentDelegationCapability);
1130 }
1131
1132 registry.register(SystemCommandsCapability);
1134
1135 registry.register(tool_output_persistence::ToolOutputPersistenceCapability);
1137
1138 registry.register(user_hooks::UserHooksCapability);
1141
1142 registry.register(LoopDetectionCapability);
1144
1145 registry.register(PromptCanaryGuardrailCapability);
1148
1149 #[cfg(feature = "ui-capabilities")]
1151 {
1152 registry.register(OpenUiCapability);
1153 registry.register(A2UiCapability);
1154 }
1155
1156 registry.register(SampleDataCapability);
1158
1159 registry.register(DataKnowledgeCapability);
1161
1162 registry.register(KnowledgeBaseCapability);
1164
1165 registry.register(FakeWarehouseCapability);
1167 registry.register(FakeAwsCapability);
1168 registry.register(FakeCrmCapability);
1169 registry.register(FakeFinancialCapability);
1170
1171 let internal_flags = crate::InternalFeatureFlags::from_env();
1173 if internal_flags.session_sandbox {
1174 registry.register(SessionSandboxCapability);
1175 }
1176
1177 if internal_flags.lua {
1181 registry.register(LuaCapability);
1182 registry.register(LuaCodeModeCapability);
1185 }
1186 for plugin in inventory::iter::<IntegrationPlugin>() {
1187 if (!plugin.experimental_only || grade.experimental_features_enabled())
1188 && plugin
1189 .feature_flag
1190 .is_none_or(|f| internal_flags.is_enabled(f))
1191 {
1192 registry.register_boxed((plugin.factory)());
1193 }
1194 }
1195
1196 registry
1197 }
1198
1199 pub fn register(&mut self, capability: impl Capability + 'static) {
1201 self.register_arc(Arc::new(capability));
1202 }
1203
1204 pub fn register_boxed(&mut self, capability: Box<dyn Capability>) {
1206 self.register_arc(Arc::from(capability));
1207 }
1208
1209 pub fn register_arc(&mut self, capability: Arc<dyn Capability>) {
1211 let canonical = capability.id().to_string();
1212 for alias in capability.aliases() {
1213 self.aliases.insert(alias.to_string(), canonical.clone());
1214 }
1215 self.capabilities.insert(canonical, capability);
1216 }
1217
1218 pub fn get(&self, id: &str) -> Option<&Arc<dyn Capability>> {
1220 self.capabilities
1221 .get(id)
1222 .or_else(|| self.aliases.get(id).and_then(|c| self.capabilities.get(c)))
1223 }
1224
1225 pub fn canonical_id<'a>(&'a self, id: &'a str) -> Option<&'a str> {
1230 if self.capabilities.contains_key(id) {
1231 Some(id)
1232 } else {
1233 self.aliases
1234 .get(id)
1235 .filter(|c| self.capabilities.contains_key(*c))
1236 .map(String::as_str)
1237 }
1238 }
1239
1240 pub fn unregister(&mut self, id: &str) -> Option<Arc<dyn Capability>> {
1242 let canonical = self.canonical_id(id)?.to_string();
1243 let removed = self.capabilities.remove(&canonical);
1244 self.aliases.retain(|_, target| *target != canonical);
1245 removed
1246 }
1247
1248 pub fn has(&self, id: &str) -> bool {
1250 self.get(id).is_some()
1251 }
1252
1253 pub fn list(&self) -> Vec<&Arc<dyn Capability>> {
1255 self.capabilities.values().collect()
1256 }
1257
1258 pub fn len(&self) -> usize {
1260 self.capabilities.len()
1261 }
1262
1263 pub fn is_empty(&self) -> bool {
1265 self.capabilities.is_empty()
1266 }
1267
1268 pub fn builder() -> CapabilityRegistryBuilder {
1270 CapabilityRegistryBuilder::new()
1271 }
1272
1273 pub fn blueprint(&self, id: &str) -> Option<AgentBlueprint> {
1277 for cap in self.capabilities.values() {
1278 for bp in cap.agent_blueprints() {
1279 if bp.id == id {
1280 return Some(bp);
1281 }
1282 }
1283 }
1284 None
1285 }
1286
1287 pub fn blueprint_with_capability(&self, id: &str) -> Option<(String, AgentBlueprint)> {
1291 for (capability_id, cap) in &self.capabilities {
1292 for bp in cap.agent_blueprints() {
1293 if bp.id == id {
1294 return Some((capability_id.clone(), bp));
1295 }
1296 }
1297 }
1298 None
1299 }
1300
1301 pub fn all_blueprints(&self) -> Vec<AgentBlueprint> {
1303 self.capabilities
1304 .values()
1305 .flat_map(|cap| cap.agent_blueprints())
1306 .collect()
1307 }
1308}
1309
1310impl Default for CapabilityRegistry {
1311 fn default() -> Self {
1312 Self::with_builtins()
1313 }
1314}
1315
1316impl std::fmt::Debug for CapabilityRegistry {
1317 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1318 let ids: Vec<_> = self.capabilities.keys().collect();
1319 f.debug_struct("CapabilityRegistry")
1320 .field("capabilities", &ids)
1321 .finish()
1322 }
1323}
1324
1325pub struct CapabilityRegistryBuilder {
1327 registry: CapabilityRegistry,
1328}
1329
1330impl CapabilityRegistryBuilder {
1331 pub fn new() -> Self {
1333 Self {
1334 registry: CapabilityRegistry::new(),
1335 }
1336 }
1337
1338 pub fn with_builtins() -> Self {
1340 Self {
1341 registry: CapabilityRegistry::with_builtins(),
1342 }
1343 }
1344
1345 pub fn capability(mut self, capability: impl Capability + 'static) -> Self {
1347 self.registry.register(capability);
1348 self
1349 }
1350
1351 pub fn build(self) -> CapabilityRegistry {
1353 self.registry
1354 }
1355}
1356
1357impl Default for CapabilityRegistryBuilder {
1358 fn default() -> Self {
1359 Self::new()
1360 }
1361}
1362
1363pub struct ModelViewContext<'a> {
1369 pub session_id: SessionId,
1370 pub prior_usage: Option<&'a TokenUsage>,
1371}
1372
1373pub trait ModelViewProvider: Send + Sync {
1379 fn apply_model_view(
1380 &self,
1381 messages: Vec<Message>,
1382 config: &serde_json::Value,
1383 context: &ModelViewContext<'_>,
1384 ) -> Vec<Message>;
1385
1386 fn priority(&self) -> i32 {
1387 0
1388 }
1389}
1390
1391pub struct CollectedCapabilities {
1396 pub system_prompt_parts: Vec<String>,
1398 pub system_prompt_attributions: Vec<SystemPromptAttribution>,
1400 pub tools: Vec<Box<dyn Tool>>,
1402 pub tool_definitions: Vec<ToolDefinition>,
1404 pub mounts: Vec<MountPoint>,
1406 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1408 pub applied_ids: Vec<String>,
1410 pub tool_search: Option<crate::llm_driver_registry::ToolSearchConfig>,
1412 pub prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig>,
1414 pub tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>>,
1416 pub tool_call_hooks: Vec<Arc<dyn ToolCallHook>>,
1418 pub mcp_servers: ScopedMcpServers,
1420 }
1426
1427#[derive(Debug, Clone, PartialEq, Eq)]
1428pub struct SystemPromptAttribution {
1429 pub capability_id: String,
1430 pub content: String,
1431}
1432
1433impl CollectedCapabilities {
1434 pub fn system_prompt_prefix(&self) -> Option<String> {
1437 if self.system_prompt_parts.is_empty() {
1438 None
1439 } else {
1440 Some(self.system_prompt_parts.join("\n\n"))
1441 }
1442 }
1443
1444 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1448 for (provider, config) in &self.message_filter_providers {
1450 provider.apply_filters(query, config);
1451 }
1452 }
1453
1454 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1457 for (provider, config) in &self.message_filter_providers {
1458 provider.post_load(messages, config);
1459 }
1460 }
1461
1462 pub fn has_message_filters(&self) -> bool {
1464 !self.message_filter_providers.is_empty()
1465 }
1466}
1467
1468pub struct CollectedMessageFilters {
1475 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1477}
1478
1479pub struct CollectedModelViewProviders {
1481 pub model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)>,
1483}
1484
1485impl CollectedMessageFilters {
1491 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1493 for (provider, config) in &self.message_filter_providers {
1494 provider.apply_filters(query, config);
1495 }
1496 }
1497
1498 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1500 for (provider, config) in &self.message_filter_providers {
1501 provider.post_load(messages, config);
1502 }
1503 }
1504}
1505
1506impl CollectedModelViewProviders {
1507 pub fn apply_model_view(
1509 &self,
1510 mut messages: Vec<Message>,
1511 context: &ModelViewContext<'_>,
1512 ) -> Vec<Message> {
1513 for (provider, config) in &self.model_view_providers {
1514 messages = provider.apply_model_view(messages, config, context);
1515 }
1516 messages
1517 }
1518}
1519
1520pub fn collect_message_filters_only(
1526 capability_configs: &[AgentCapabilityConfig],
1527 registry: &CapabilityRegistry,
1528) -> CollectedMessageFilters {
1529 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1530 Vec::new();
1531
1532 for cap_config in capability_configs {
1533 let cap_id = cap_config.capability_ref.as_str();
1534 if let Some(capability) = registry.get(cap_id) {
1535 if capability.status() != CapabilityStatus::Available {
1536 continue;
1537 }
1538 let effective: &dyn Capability = capability
1541 .resolve_for_model(None)
1542 .unwrap_or_else(|| capability.as_ref());
1543 if let Some(provider) = effective.message_filter_provider() {
1544 message_filter_providers.push((provider, cap_config.config.clone()));
1545 }
1546 }
1547 }
1548
1549 message_filter_providers.sort_by_key(|(p, _)| p.priority());
1550
1551 CollectedMessageFilters {
1552 message_filter_providers,
1553 }
1554}
1555
1556pub fn collect_model_view_providers(
1563 capability_configs: &[AgentCapabilityConfig],
1564 registry: &CapabilityRegistry,
1565 model: Option<&str>,
1566) -> CollectedModelViewProviders {
1567 let mut model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)> = Vec::new();
1568
1569 for cap_config in capability_configs {
1570 let cap_id = cap_config.capability_ref.as_str();
1571 if let Some(capability) = registry.get(cap_id) {
1572 if capability.status() != CapabilityStatus::Available {
1573 continue;
1574 }
1575 let effective: &dyn Capability = capability
1576 .resolve_for_model(model)
1577 .unwrap_or_else(|| capability.as_ref());
1578 if let Some(provider) = effective.model_view_provider() {
1579 model_view_providers.push((provider, cap_config.config.clone()));
1580 }
1581 }
1582 }
1583
1584 model_view_providers.sort_by_key(|(p, _)| p.priority());
1585
1586 CollectedModelViewProviders {
1587 model_view_providers,
1588 }
1589}
1590
1591pub fn collect_capability_mcp_servers(
1592 capability_configs: &[AgentCapabilityConfig],
1593 registry: &CapabilityRegistry,
1594) -> ScopedMcpServers {
1595 let mut servers = ScopedMcpServers::default();
1596
1597 for cap_config in capability_configs {
1598 let cap_id = cap_config.capability_ref.as_str();
1599 if is_declarative_capability(cap_id) {
1600 if let Ok(definition) =
1601 serde_json::from_value::<DeclarativeCapabilityDefinition>(cap_config.config.clone())
1602 {
1603 if definition.status != CapabilityStatus::Available {
1604 continue;
1605 }
1606 if let Some(contributed) = definition.mcp_servers {
1607 servers = merge_scoped_mcp_servers(&servers, &contributed);
1608 }
1609 }
1610 continue;
1611 }
1612 if let Some(capability) = registry.get(cap_id) {
1613 if capability.status() != CapabilityStatus::Available {
1614 continue;
1615 }
1616 servers = merge_scoped_mcp_servers(
1617 &servers,
1618 &capability.mcp_servers_with_config(&cap_config.config),
1619 );
1620 }
1621 }
1622
1623 servers
1624}
1625
1626pub const MAX_RESOLVED_CAPABILITIES: usize = 100;
1633
1634#[derive(Debug, Clone, PartialEq, Eq)]
1636pub enum DependencyError {
1637 CircularDependency {
1639 capability_id: String,
1641 chain: Vec<String>,
1643 },
1644 TooManyCapabilities {
1646 count: usize,
1648 max: usize,
1650 },
1651}
1652
1653impl std::fmt::Display for DependencyError {
1654 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1655 match self {
1656 DependencyError::CircularDependency {
1657 capability_id,
1658 chain,
1659 } => {
1660 write!(
1661 f,
1662 "Circular dependency detected: {} depends on itself via chain: {} -> {}",
1663 capability_id,
1664 chain.join(" -> "),
1665 capability_id
1666 )
1667 }
1668 DependencyError::TooManyCapabilities { count, max } => {
1669 write!(
1670 f,
1671 "Too many capabilities after resolution: {} (max: {})",
1672 count, max
1673 )
1674 }
1675 }
1676 }
1677}
1678
1679impl std::error::Error for DependencyError {}
1680
1681#[derive(Debug, Clone)]
1683pub struct ResolvedCapabilities {
1684 pub resolved_ids: Vec<String>,
1687 pub added_as_dependencies: Vec<String>,
1689 pub user_selected: Vec<String>,
1691}
1692
1693pub fn resolve_dependencies(
1713 selected_ids: &[String],
1714 registry: &CapabilityRegistry,
1715) -> Result<ResolvedCapabilities, DependencyError> {
1716 use std::collections::HashSet;
1717
1718 let user_selected: HashSet<String> = selected_ids
1720 .iter()
1721 .map(|id| registry.canonical_id(id).unwrap_or(id).to_string())
1722 .collect();
1723 let mut resolved: Vec<String> = Vec::new();
1724 let mut resolved_set: HashSet<String> = HashSet::new();
1725 let mut added_as_dependencies: Vec<String> = Vec::new();
1726
1727 for cap_id in selected_ids {
1729 resolve_single_capability(
1730 cap_id,
1731 registry,
1732 &mut resolved,
1733 &mut resolved_set,
1734 &mut added_as_dependencies,
1735 &user_selected,
1736 &mut Vec::new(), )?;
1738 }
1739
1740 if resolved.len() > MAX_RESOLVED_CAPABILITIES {
1742 return Err(DependencyError::TooManyCapabilities {
1743 count: resolved.len(),
1744 max: MAX_RESOLVED_CAPABILITIES,
1745 });
1746 }
1747
1748 Ok(ResolvedCapabilities {
1749 resolved_ids: resolved,
1750 added_as_dependencies,
1751 user_selected: selected_ids.to_vec(),
1752 })
1753}
1754
1755pub fn resolve_capability_configs(
1760 selected_configs: &[AgentCapabilityConfig],
1761 registry: &CapabilityRegistry,
1762) -> Result<Vec<AgentCapabilityConfig>, DependencyError> {
1763 let mut selected_ids: Vec<String> = Vec::new();
1764 for config in selected_configs {
1765 if is_declarative_capability(config.capability_id())
1766 && let Ok(definition) =
1767 serde_json::from_value::<DeclarativeCapabilityDefinition>(config.config.clone())
1768 {
1769 selected_ids.extend(definition.dependencies);
1770 }
1771 selected_ids.push(config.capability_id().to_string());
1772 }
1773 let resolved = resolve_dependencies(&selected_ids, registry)?;
1774
1775 let explicit_configs: std::collections::HashMap<String, serde_json::Value> = selected_configs
1778 .iter()
1779 .map(|config| {
1780 let id = config.capability_id();
1781 let id = registry.canonical_id(id).unwrap_or(id);
1782 (id.to_string(), config.config.clone())
1783 })
1784 .collect();
1785
1786 Ok(resolved
1787 .resolved_ids
1788 .into_iter()
1789 .map(|capability_id| {
1790 explicit_configs
1791 .get(&capability_id)
1792 .cloned()
1793 .map(|config| AgentCapabilityConfig::with_config(capability_id.clone(), config))
1794 .unwrap_or_else(|| AgentCapabilityConfig::new(capability_id))
1795 })
1796 .collect())
1797}
1798
1799fn resolve_single_capability(
1801 cap_id: &str,
1802 registry: &CapabilityRegistry,
1803 resolved: &mut Vec<String>,
1804 resolved_set: &mut std::collections::HashSet<String>,
1805 added_as_dependencies: &mut Vec<String>,
1806 user_selected: &std::collections::HashSet<String>,
1807 visiting: &mut Vec<String>,
1808) -> Result<(), DependencyError> {
1809 let cap_id = registry.canonical_id(cap_id).unwrap_or(cap_id);
1813
1814 if resolved_set.contains(cap_id) {
1816 return Ok(());
1817 }
1818
1819 if visiting.contains(&cap_id.to_string()) {
1821 return Err(DependencyError::CircularDependency {
1822 capability_id: cap_id.to_string(),
1823 chain: visiting.clone(),
1824 });
1825 }
1826
1827 let capability = match registry.get(cap_id) {
1829 Some(cap) => cap,
1830 None => {
1831 if is_declarative_capability(cap_id) && !resolved_set.contains(cap_id) {
1832 resolved.push(cap_id.to_string());
1833 resolved_set.insert(cap_id.to_string());
1834 if !user_selected.contains(cap_id) {
1835 added_as_dependencies.push(cap_id.to_string());
1836 }
1837 }
1838 return Ok(());
1839 }
1840 };
1841
1842 visiting.push(cap_id.to_string());
1844
1845 for dep_id in capability.dependencies() {
1847 resolve_single_capability(
1848 dep_id,
1849 registry,
1850 resolved,
1851 resolved_set,
1852 added_as_dependencies,
1853 user_selected,
1854 visiting,
1855 )?;
1856 }
1857
1858 visiting.pop();
1860
1861 if !resolved_set.contains(cap_id) {
1863 resolved.push(cap_id.to_string());
1864 resolved_set.insert(cap_id.to_string());
1865
1866 if !user_selected.contains(cap_id) {
1868 added_as_dependencies.push(cap_id.to_string());
1869 }
1870 }
1871
1872 Ok(())
1873}
1874
1875pub fn compute_features(capability_ids: &[String], registry: &CapabilityRegistry) -> Vec<String> {
1880 use std::collections::HashSet;
1881
1882 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
1883 Ok(resolved) => resolved.resolved_ids,
1884 Err(_) => capability_ids.to_vec(),
1885 };
1886
1887 let mut seen = HashSet::new();
1888 let mut features = Vec::new();
1889 for cap_id in &resolved_ids {
1890 if let Some(cap) = registry.get(cap_id) {
1891 for feature in cap.features() {
1892 if seen.insert(feature) {
1893 features.push(feature.to_string());
1894 }
1895 }
1896 }
1897 }
1898 features
1899}
1900
1901pub fn get_dependencies(cap_id: &str, registry: &CapabilityRegistry) -> Vec<String> {
1904 registry
1905 .get(cap_id)
1906 .map(|cap| cap.dependencies().iter().map(|s| s.to_string()).collect())
1907 .unwrap_or_default()
1908}
1909
1910pub async fn collect_capabilities(
1926 capability_ids: &[String],
1927 registry: &CapabilityRegistry,
1928 ctx: &SystemPromptContext,
1929) -> CollectedCapabilities {
1930 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
1933 Ok(resolved) => resolved.resolved_ids,
1934 Err(e) => {
1935 tracing::warn!("Failed to resolve capability dependencies: {}", e);
1936 capability_ids.to_vec()
1937 }
1938 };
1939
1940 let configs: Vec<AgentCapabilityConfig> = resolved_ids
1942 .iter()
1943 .map(|id| AgentCapabilityConfig {
1944 capability_ref: CapabilityId::new(id),
1945 config: serde_json::Value::Object(serde_json::Map::new()),
1946 })
1947 .collect();
1948
1949 collect_capabilities_with_configs(&configs, registry, ctx).await
1950}
1951
1952pub async fn collect_capabilities_with_configs(
1963 capability_configs: &[AgentCapabilityConfig],
1964 registry: &CapabilityRegistry,
1965 ctx: &SystemPromptContext,
1966) -> CollectedCapabilities {
1967 let mut system_prompt_parts: Vec<String> = Vec::new();
1968 let mut system_prompt_attributions: Vec<SystemPromptAttribution> = Vec::new();
1969 let mut tools: Vec<Box<dyn Tool>> = Vec::new();
1970 let mut tool_definitions: Vec<ToolDefinition> = Vec::new();
1971 let mut mounts: Vec<MountPoint> = Vec::new();
1972 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1973 Vec::new();
1974 let mut applied_ids: Vec<String> = Vec::new();
1975 let mut tool_search: Option<crate::llm_driver_registry::ToolSearchConfig> = None;
1976 let mut prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig> = None;
1977 let mut tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>> = Vec::new();
1978 let mut tool_call_hooks: Vec<Arc<dyn ToolCallHook>> = Vec::new();
1979 let mut mcp_servers = ScopedMcpServers::default();
1980
1981 for cap_config in capability_configs {
1982 let cap_id = cap_config.capability_ref.as_str();
1983 if is_declarative_capability(cap_id) {
1984 match serde_json::from_value::<DeclarativeCapabilityDefinition>(
1985 cap_config.config.clone(),
1986 ) {
1987 Ok(definition) => {
1988 if definition.status != CapabilityStatus::Available {
1989 continue;
1990 }
1991
1992 if let Some(prompt) = definition.system_prompt.as_deref() {
1993 let contribution =
1994 format!("<capability id=\"{}\">\n{}\n</capability>", cap_id, prompt);
1995 system_prompt_attributions.push(SystemPromptAttribution {
1996 capability_id: cap_id.to_string(),
1997 content: contribution.clone(),
1998 });
1999 system_prompt_parts.push(contribution);
2000 }
2001
2002 mounts.extend(definition.mounts(cap_id));
2003 if let Some(ref servers) = definition.mcp_servers {
2004 mcp_servers = merge_scoped_mcp_servers(&mcp_servers, servers);
2005 }
2006 for skill in definition.skill_contributions() {
2007 mounts.push(skill.to_mount(cap_id));
2008 }
2009
2010 applied_ids.push(cap_id.to_string());
2011 }
2012 Err(error) => {
2013 tracing::warn!(
2014 capability_id = %cap_id,
2015 error = %error,
2016 "Skipping invalid declarative capability config"
2017 );
2018 }
2019 }
2020 continue;
2021 }
2022 if let Some(capability) = registry.get(cap_id) {
2023 if capability.status() != CapabilityStatus::Available {
2025 continue;
2026 }
2027
2028 let effective: &dyn Capability =
2040 match capability.resolve_for_model(ctx.model.as_deref()) {
2041 Some(inner) => inner,
2042 None => capability.as_ref(),
2043 };
2044 let effective_id = effective.id();
2045
2046 if let Some(contribution) = effective
2048 .system_prompt_contribution_with_config(ctx, &cap_config.config)
2049 .await
2050 {
2051 system_prompt_attributions.push(SystemPromptAttribution {
2052 capability_id: cap_id.to_string(),
2053 content: contribution.clone(),
2054 });
2055 system_prompt_parts.push(contribution);
2056 }
2057
2058 tools.extend(effective.tools_with_config(&cap_config.config));
2060 tool_definition_hooks
2061 .extend(effective.tool_definition_hooks_with_context(ctx, &cap_config.config));
2062 tool_call_hooks.extend(effective.tool_call_hooks());
2063 let cap_category = effective.category();
2068 for def in effective.tool_definitions() {
2069 let def = match (def.category(), cap_category) {
2070 (None, Some(cat)) => def.with_category(cat),
2071 _ => def,
2072 }
2073 .with_capability_attribution(cap_id, Some(capability.name()));
2074 tool_definitions.push(def);
2075 }
2076
2077 if effective_id == OPENAI_TOOL_SEARCH_CAPABILITY_ID {
2082 let threshold = cap_config
2084 .config
2085 .get("threshold")
2086 .and_then(|v| v.as_u64())
2087 .map(|v| v as usize)
2088 .unwrap_or(DEFAULT_TOOL_SEARCH_THRESHOLD);
2089 tool_search = Some(crate::llm_driver_registry::ToolSearchConfig {
2090 enabled: true,
2091 threshold,
2092 });
2093 }
2094
2095 if cap_id == PROMPT_CACHING_CAPABILITY_ID {
2096 let strategy = cap_config
2097 .config
2098 .get("strategy")
2099 .and_then(|v| v.as_str())
2100 .map(|value| match value {
2101 "auto" => crate::llm_driver_registry::PromptCacheStrategy::Auto,
2102 _ => crate::llm_driver_registry::PromptCacheStrategy::Auto,
2103 })
2104 .unwrap_or(crate::llm_driver_registry::PromptCacheStrategy::Auto);
2105 let gemini_cached_content = cap_config
2106 .config
2107 .get("gemini_cached_content")
2108 .and_then(|v| v.as_str())
2109 .map(str::to_string);
2110 prompt_cache = Some(crate::llm_driver_registry::PromptCacheConfig {
2111 enabled: true,
2112 strategy,
2113 gemini_cached_content,
2114 });
2115 }
2116
2117 mounts.extend(effective.mounts());
2119
2120 mcp_servers = merge_scoped_mcp_servers(
2121 &mcp_servers,
2122 &effective.mcp_servers_with_config(&cap_config.config),
2123 );
2124
2125 for skill in effective.contribute_skills() {
2129 mounts.push(skill.to_mount(cap_id));
2130 }
2131
2132 if let Some(provider) = effective.message_filter_provider() {
2134 message_filter_providers.push((provider, cap_config.config.clone()));
2135 }
2136
2137 applied_ids.push(cap_id.to_string());
2138 }
2139 }
2140
2141 if !applied_ids
2153 .iter()
2154 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID)
2155 && tool_definitions
2156 .iter()
2157 .any(|def| def.hints().supports_background == Some(true))
2158 && let Some(bg_cap) = registry.get(BACKGROUND_EXECUTION_CAPABILITY_ID)
2159 && bg_cap.status() == CapabilityStatus::Available
2160 {
2161 tools.extend(bg_cap.tools());
2162 let cap_category = bg_cap.category();
2163 for def in bg_cap.tool_definitions() {
2164 let def = match (def.category(), cap_category) {
2165 (None, Some(cat)) => def.with_category(cat),
2166 _ => def,
2167 }
2168 .with_capability_attribution(BACKGROUND_EXECUTION_CAPABILITY_ID, Some(bg_cap.name()));
2169 tool_definitions.push(def);
2170 }
2171 applied_ids.push(BACKGROUND_EXECUTION_CAPABILITY_ID.to_string());
2172 }
2173
2174 message_filter_providers.sort_by_key(|(p, _)| p.priority());
2176
2177 CollectedCapabilities {
2178 system_prompt_parts,
2179 system_prompt_attributions,
2180 tools,
2181 tool_definitions,
2182 mounts,
2183 message_filter_providers,
2184 applied_ids,
2185 tool_search,
2186 prompt_cache,
2187 tool_definition_hooks,
2188 tool_call_hooks,
2189 mcp_servers,
2190 }
2191}
2192
2193pub struct AppliedCapabilities {
2199 pub runtime_agent: RuntimeAgent,
2201 pub tool_registry: ToolRegistry,
2203 pub applied_ids: Vec<String>,
2205}
2206
2207pub async fn apply_capabilities(
2244 base_runtime_agent: RuntimeAgent,
2245 capability_ids: &[String],
2246 registry: &CapabilityRegistry,
2247 ctx: &SystemPromptContext,
2248) -> AppliedCapabilities {
2249 let collected = collect_capabilities(capability_ids, registry, ctx).await;
2250
2251 let final_system_prompt = match collected.system_prompt_prefix() {
2253 Some(prefix) => format!(
2254 "{}\n\n<system-prompt>\n{}\n</system-prompt>",
2255 prefix, base_runtime_agent.system_prompt
2256 ),
2257 None => base_runtime_agent.system_prompt,
2258 };
2259
2260 let mut tool_registry = ToolRegistry::new();
2262 for tool in collected.tools {
2263 tool_registry.register_boxed(tool);
2264 }
2265
2266 let mut tools = collected.tool_definitions;
2268 for hook in &collected.tool_definition_hooks {
2269 tools = hook.transform(tools);
2270 }
2271
2272 let runtime_agent = RuntimeAgent {
2273 system_prompt: final_system_prompt,
2274 model: base_runtime_agent.model,
2275 tools,
2276 max_iterations: base_runtime_agent.max_iterations,
2277 temperature: base_runtime_agent.temperature,
2278 max_tokens: base_runtime_agent.max_tokens,
2279 tool_search: collected.tool_search,
2280 prompt_cache: collected.prompt_cache,
2281 network_access: base_runtime_agent.network_access,
2282 };
2283
2284 AppliedCapabilities {
2285 runtime_agent,
2286 tool_registry,
2287 applied_ids: collected.applied_ids,
2288 }
2289}
2290
2291#[cfg(test)]
2296mod tests {
2297 use super::*;
2298 use crate::typed_id::SessionId;
2299 use std::collections::BTreeSet;
2300 use uuid::Uuid;
2301
2302 static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2304
2305 fn lock_env() -> std::sync::MutexGuard<'static, ()> {
2306 ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner())
2307 }
2308
2309 fn test_ctx() -> SystemPromptContext {
2311 SystemPromptContext::without_file_store(SessionId::new())
2312 }
2313
2314 fn expected_core_builtin_ids() -> BTreeSet<&'static str> {
2316 let mut ids = [
2317 "agent_instructions",
2318 "human_intent",
2319 "budgeting",
2320 "self_budget",
2321 "noop",
2322 "current_time",
2323 "research",
2324 "platform_management",
2325 "session_file_system",
2326 "session_storage",
2327 "session",
2328 "session_sql_database",
2329 "test_math",
2330 "test_weather",
2331 "stateless_todo_list",
2332 "web_fetch",
2333 "bashkit_shell",
2334 "background_execution",
2335 "session_schedule",
2336 "btw",
2337 "infinity_context",
2338 "compaction",
2339 "memory",
2340 "message_metadata",
2341 "openai_tool_search",
2342 "tool_search",
2343 "auto_tool_search",
2344 "prompt_caching",
2345 "session_tasks",
2346 "skills",
2347 "subagents",
2348 "system_commands",
2349 "sample_data",
2350 "data_knowledge",
2351 "knowledge_base",
2352 "tool_output_persistence",
2353 "fake_warehouse",
2354 "fake_aws",
2355 "fake_crm",
2356 "fake_financial",
2357 "loop_detection",
2358 "prompt_canary_guardrail",
2359 "user_hooks",
2360 ]
2361 .into_iter()
2362 .collect::<BTreeSet<_>>();
2363 if cfg!(feature = "ui-capabilities") {
2364 ids.insert("openui");
2365 ids.insert("a2ui");
2366 }
2367 ids
2368 }
2369
2370 fn expected_dev_builtin_ids() -> BTreeSet<&'static str> {
2372 let mut ids = expected_core_builtin_ids();
2373 ids.insert("agent_handoff");
2374 ids.insert("a2a_agent_delegation");
2375 ids
2376 }
2377
2378 fn registry_ids(registry: &CapabilityRegistry) -> BTreeSet<&str> {
2379 registry.capabilities.keys().map(String::as_str).collect()
2380 }
2381
2382 #[test]
2392 fn test_capability_registry_with_builtins_dev() {
2393 let _lock = lock_env();
2395 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2396 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
2397 assert_eq!(registry_ids(®istry), expected_dev_builtin_ids());
2398 assert!(registry.has("agent_handoff"));
2399 assert!(registry.has("a2a_agent_delegation"));
2400 }
2401
2402 #[test]
2403 fn test_capability_registry_with_builtins_prod() {
2404 let _lock = lock_env();
2406 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2407 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2408 assert_eq!(registry_ids(®istry), expected_core_builtin_ids());
2409 assert!(!registry.has("docker_container"));
2411 assert!(!registry.has("agent_handoff"));
2412 assert!(!registry.has("a2a_agent_delegation"));
2413 }
2414
2415 #[test]
2416 fn test_agent_delegation_enabled_by_env_in_prod() {
2417 let _lock = lock_env();
2419 unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "true") };
2420 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2421 assert!(registry.has("agent_handoff"));
2422 assert!(registry.has("a2a_agent_delegation"));
2423 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2424 }
2425
2426 #[test]
2427 fn test_agent_delegation_disabled_by_env_in_dev() {
2428 let _lock = lock_env();
2430 unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "false") };
2431 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
2432 assert!(!registry.has("agent_handoff"));
2433 assert!(!registry.has("a2a_agent_delegation"));
2434 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2435 }
2436
2437 #[test]
2438 fn test_capability_registry_get() {
2439 let registry = CapabilityRegistry::with_builtins();
2440
2441 let noop = registry.get("noop").unwrap();
2442 assert_eq!(noop.id(), "noop");
2443 assert_eq!(noop.name(), "No-Op");
2444 assert_eq!(noop.status(), CapabilityStatus::Available);
2445 }
2446
2447 #[test]
2448 fn test_capability_registry_blueprint_with_capability() {
2449 struct BlueprintProviderCapability;
2450
2451 impl Capability for BlueprintProviderCapability {
2452 fn id(&self) -> &str {
2453 "blueprint_provider"
2454 }
2455 fn name(&self) -> &str {
2456 "Blueprint Provider"
2457 }
2458 fn description(&self) -> &str {
2459 "Capability that provides a blueprint for tests"
2460 }
2461 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
2462 vec![AgentBlueprint {
2463 id: "test_blueprint",
2464 name: "Test Blueprint",
2465 description: "Blueprint for capability registry tests",
2466 model: BlueprintModel::Inherit,
2467 system_prompt: "Test prompt",
2468 tools: vec![],
2469 max_turns: None,
2470 config_schema: None,
2471 }]
2472 }
2473 }
2474
2475 let mut registry = CapabilityRegistry::new();
2476 registry.register(BlueprintProviderCapability);
2477
2478 let (capability_id, blueprint) = registry
2479 .blueprint_with_capability("test_blueprint")
2480 .expect("blueprint should resolve with capability id");
2481 assert_eq!(capability_id, "blueprint_provider");
2482 assert_eq!(blueprint.id, "test_blueprint");
2483 }
2484
2485 #[test]
2486 fn test_capability_registry_builder() {
2487 let registry = CapabilityRegistry::builder()
2488 .capability(NoopCapability)
2489 .capability(CurrentTimeCapability)
2490 .build();
2491
2492 assert!(registry.has("noop"));
2493 assert!(registry.has("current_time"));
2494 assert_eq!(registry.len(), 2);
2495 }
2496
2497 #[test]
2498 fn test_capability_status() {
2499 let registry = CapabilityRegistry::with_builtins();
2500
2501 let current_time = registry.get("current_time").unwrap();
2502 assert_eq!(current_time.status(), CapabilityStatus::Available);
2503
2504 let research = registry.get("research").unwrap();
2505 assert_eq!(research.status(), CapabilityStatus::ComingSoon);
2506 }
2507
2508 #[test]
2509 fn test_capability_icons_and_categories() {
2510 let registry = CapabilityRegistry::with_builtins();
2511
2512 let noop = registry.get("noop").unwrap();
2513 assert_eq!(noop.icon(), Some("circle-off"));
2514 assert_eq!(noop.category(), Some("Testing"));
2515
2516 let current_time = registry.get("current_time").unwrap();
2517 assert_eq!(current_time.icon(), Some("clock"));
2518 assert_eq!(current_time.category(), Some("Utilities"));
2519 }
2520
2521 #[test]
2522 fn test_system_prompt_preview_default_delegates_to_addition() {
2523 let registry = CapabilityRegistry::with_builtins();
2524
2525 let test_math = registry.get("test_math").unwrap();
2527 assert_eq!(
2528 test_math.system_prompt_preview().as_deref(),
2529 test_math.system_prompt_addition()
2530 );
2531
2532 let current_time = registry.get("current_time").unwrap();
2534 assert!(current_time.system_prompt_preview().is_none());
2535 assert!(current_time.system_prompt_addition().is_none());
2536 }
2537
2538 #[test]
2539 fn test_system_prompt_preview_dynamic_capability() {
2540 let registry = CapabilityRegistry::with_builtins();
2541 let cap = registry.get("agent_instructions").unwrap();
2542
2543 assert!(cap.system_prompt_addition().is_none());
2545 assert!(cap.system_prompt_preview().is_some());
2546 assert!(cap.system_prompt_preview().unwrap().contains("AGENTS.md"));
2547 }
2548
2549 #[tokio::test]
2554 async fn test_apply_capabilities_empty() {
2555 let registry = CapabilityRegistry::with_builtins();
2556 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2557
2558 let applied =
2559 apply_capabilities(base_runtime_agent.clone(), &[], ®istry, &test_ctx()).await;
2560
2561 assert_eq!(
2562 applied.runtime_agent.system_prompt,
2563 base_runtime_agent.system_prompt
2564 );
2565 assert!(applied.tool_registry.is_empty());
2566 assert!(applied.applied_ids.is_empty());
2567 }
2568
2569 #[tokio::test]
2570 async fn test_apply_capabilities_noop() {
2571 let registry = CapabilityRegistry::with_builtins();
2572 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2573
2574 let applied = apply_capabilities(
2575 base_runtime_agent.clone(),
2576 &["noop".to_string()],
2577 ®istry,
2578 &test_ctx(),
2579 )
2580 .await;
2581
2582 assert_eq!(
2584 applied.runtime_agent.system_prompt,
2585 base_runtime_agent.system_prompt
2586 );
2587 assert!(applied.tool_registry.is_empty());
2588 assert_eq!(applied.applied_ids, vec!["noop"]);
2589 }
2590
2591 #[tokio::test]
2592 async fn test_apply_capabilities_current_time() {
2593 let registry = CapabilityRegistry::with_builtins();
2594 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2595
2596 let applied = apply_capabilities(
2597 base_runtime_agent.clone(),
2598 &["current_time".to_string()],
2599 ®istry,
2600 &test_ctx(),
2601 )
2602 .await;
2603
2604 assert_eq!(
2606 applied.runtime_agent.system_prompt,
2607 base_runtime_agent.system_prompt
2608 );
2609 assert!(applied.tool_registry.has("get_current_time"));
2610 assert_eq!(applied.tool_registry.len(), 1);
2611 assert_eq!(applied.applied_ids, vec!["current_time"]);
2612 }
2613
2614 #[tokio::test]
2615 async fn test_apply_capabilities_skips_coming_soon() {
2616 let registry = CapabilityRegistry::with_builtins();
2617 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2618
2619 let applied = apply_capabilities(
2621 base_runtime_agent.clone(),
2622 &["research".to_string()],
2623 ®istry,
2624 &test_ctx(),
2625 )
2626 .await;
2627
2628 assert_eq!(
2630 applied.runtime_agent.system_prompt,
2631 base_runtime_agent.system_prompt
2632 );
2633 assert!(applied.applied_ids.is_empty()); }
2635
2636 #[tokio::test]
2637 async fn test_apply_capabilities_multiple() {
2638 let registry = CapabilityRegistry::with_builtins();
2639 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2640
2641 let applied = apply_capabilities(
2642 base_runtime_agent.clone(),
2643 &["noop".to_string(), "current_time".to_string()],
2644 ®istry,
2645 &test_ctx(),
2646 )
2647 .await;
2648
2649 assert!(applied.tool_registry.has("get_current_time"));
2650 assert_eq!(applied.applied_ids, vec!["noop", "current_time"]);
2651 }
2652
2653 #[tokio::test]
2654 async fn test_apply_capabilities_preserves_order() {
2655 let registry = CapabilityRegistry::with_builtins();
2656 let base_runtime_agent = RuntimeAgent::new("Base prompt.", "gpt-5.2");
2657
2658 let applied = apply_capabilities(
2660 base_runtime_agent,
2661 &["current_time".to_string(), "noop".to_string()],
2662 ®istry,
2663 &test_ctx(),
2664 )
2665 .await;
2666
2667 assert_eq!(applied.applied_ids, vec!["current_time", "noop"]);
2668 }
2669
2670 #[tokio::test]
2671 async fn test_apply_capabilities_test_math() {
2672 let registry = CapabilityRegistry::with_builtins();
2673 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2674
2675 let applied = apply_capabilities(
2676 base_runtime_agent.clone(),
2677 &["test_math".to_string()],
2678 ®istry,
2679 &test_ctx(),
2680 )
2681 .await;
2682
2683 assert!(
2685 !applied
2686 .runtime_agent
2687 .system_prompt
2688 .contains("<capability id=\"test_math\">")
2689 );
2690 assert!(
2692 applied
2693 .runtime_agent
2694 .system_prompt
2695 .contains("You are a helpful assistant.")
2696 );
2697 assert!(applied.tool_registry.has("add"));
2698 assert!(applied.tool_registry.has("subtract"));
2699 assert!(applied.tool_registry.has("multiply"));
2700 assert!(applied.tool_registry.has("divide"));
2701 assert_eq!(applied.tool_registry.len(), 4);
2702 }
2703
2704 #[tokio::test]
2705 async fn test_apply_capabilities_test_weather() {
2706 let registry = CapabilityRegistry::with_builtins();
2707 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2708
2709 let applied = apply_capabilities(
2710 base_runtime_agent.clone(),
2711 &["test_weather".to_string()],
2712 ®istry,
2713 &test_ctx(),
2714 )
2715 .await;
2716
2717 assert!(
2719 !applied
2720 .runtime_agent
2721 .system_prompt
2722 .contains("<capability id=\"test_weather\">")
2723 );
2724 assert!(applied.tool_registry.has("get_weather"));
2725 assert!(applied.tool_registry.has("get_forecast"));
2726 assert_eq!(applied.tool_registry.len(), 2);
2727 }
2728
2729 #[tokio::test]
2730 async fn test_apply_capabilities_test_math_and_test_weather() {
2731 let registry = CapabilityRegistry::with_builtins();
2732 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2733
2734 let applied = apply_capabilities(
2735 base_runtime_agent.clone(),
2736 &["test_math".to_string(), "test_weather".to_string()],
2737 ®istry,
2738 &test_ctx(),
2739 )
2740 .await;
2741
2742 assert_eq!(applied.tool_registry.len(), 6); assert!(applied.tool_registry.has("add"));
2745 assert!(applied.tool_registry.has("get_weather"));
2746 }
2747
2748 #[tokio::test]
2749 async fn test_apply_capabilities_stateless_todo_list() {
2750 let registry = CapabilityRegistry::with_builtins();
2751 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2752
2753 let applied = apply_capabilities(
2754 base_runtime_agent.clone(),
2755 &["stateless_todo_list".to_string()],
2756 ®istry,
2757 &test_ctx(),
2758 )
2759 .await;
2760
2761 assert!(
2763 applied
2764 .runtime_agent
2765 .system_prompt
2766 .contains("Task Management")
2767 );
2768 assert!(applied.runtime_agent.system_prompt.contains("write_todos"));
2769 assert!(applied.tool_registry.has("write_todos"));
2770 assert_eq!(applied.tool_registry.len(), 1);
2771 }
2772
2773 #[tokio::test]
2774 async fn test_apply_capabilities_web_fetch() {
2775 let registry = CapabilityRegistry::with_builtins();
2776 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2777
2778 let applied = apply_capabilities(
2779 base_runtime_agent.clone(),
2780 &["web_fetch".to_string()],
2781 ®istry,
2782 &test_ctx(),
2783 )
2784 .await;
2785
2786 assert!(
2788 applied
2789 .runtime_agent
2790 .system_prompt
2791 .contains(&base_runtime_agent.system_prompt)
2792 );
2793 assert!(applied.runtime_agent.system_prompt.contains("web_fetch"));
2794 assert!(applied.tool_registry.has("web_fetch"));
2795 assert_eq!(applied.tool_registry.len(), 1);
2796 }
2797
2798 #[tokio::test]
2803 async fn test_xml_tags_wrap_capability_prompts() {
2804 let registry = CapabilityRegistry::with_builtins();
2805 let collected =
2806 collect_capabilities(&["stateless_todo_list".to_string()], ®istry, &test_ctx())
2807 .await;
2808
2809 assert_eq!(collected.system_prompt_parts.len(), 1);
2810 let part = &collected.system_prompt_parts[0];
2811 assert!(part.starts_with("<capability id=\"stateless_todo_list\">"));
2812 assert!(part.ends_with("</capability>"));
2813 assert!(part.contains("Task Management"));
2814 }
2815
2816 #[tokio::test]
2817 async fn test_xml_tags_multiple_capabilities() {
2818 let registry = CapabilityRegistry::with_builtins();
2819 let collected = collect_capabilities(
2820 &[
2821 "stateless_todo_list".to_string(),
2822 "session_schedule".to_string(),
2823 ],
2824 ®istry,
2825 &test_ctx(),
2826 )
2827 .await;
2828
2829 assert_eq!(collected.system_prompt_parts.len(), 2);
2830 assert!(
2831 collected.system_prompt_parts[0].starts_with("<capability id=\"stateless_todo_list\">")
2832 );
2833 assert!(
2834 collected.system_prompt_parts[1].starts_with("<capability id=\"session_schedule\">")
2835 );
2836
2837 let prefix = collected.system_prompt_prefix().unwrap();
2838 assert!(prefix.contains("</capability>\n\n<capability"));
2840 }
2841
2842 #[tokio::test]
2843 async fn test_xml_tags_system_prompt_wrapping() {
2844 let registry = CapabilityRegistry::with_builtins();
2845 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2846
2847 let applied = apply_capabilities(
2848 base,
2849 &["stateless_todo_list".to_string()],
2850 ®istry,
2851 &test_ctx(),
2852 )
2853 .await;
2854
2855 let prompt = &applied.runtime_agent.system_prompt;
2856 assert!(prompt.contains("<capability id=\"stateless_todo_list\">"));
2858 assert!(prompt.contains("</capability>"));
2859 assert!(prompt.contains("<system-prompt>\nYou are helpful.\n</system-prompt>"));
2861 }
2862
2863 #[tokio::test]
2864 async fn test_no_xml_wrapping_without_capabilities() {
2865 let registry = CapabilityRegistry::with_builtins();
2866 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2867
2868 let applied = apply_capabilities(base, &[], ®istry, &test_ctx()).await;
2869
2870 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2872 assert!(
2873 !applied
2874 .runtime_agent
2875 .system_prompt
2876 .contains("<system-prompt>")
2877 );
2878 }
2879
2880 #[tokio::test]
2881 async fn test_no_xml_wrapping_for_noop_capability() {
2882 let registry = CapabilityRegistry::with_builtins();
2883 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2884
2885 let applied = apply_capabilities(base, &["noop".to_string()], ®istry, &test_ctx()).await;
2887
2888 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2889 assert!(
2890 !applied
2891 .runtime_agent
2892 .system_prompt
2893 .contains("<system-prompt>")
2894 );
2895 }
2896
2897 #[tokio::test]
2902 async fn test_collect_capabilities_includes_mounts() {
2903 let registry = CapabilityRegistry::with_builtins();
2904
2905 let collected =
2906 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
2907
2908 assert!(!collected.mounts.is_empty());
2909 assert_eq!(collected.mounts.len(), 1);
2910 assert_eq!(collected.mounts[0].path, "/samples");
2911 assert!(collected.mounts[0].is_readonly());
2912 }
2913
2914 #[tokio::test]
2915 async fn test_collect_capabilities_empty_mounts_by_default() {
2916 let registry = CapabilityRegistry::with_builtins();
2917
2918 let collected =
2920 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
2921
2922 assert!(collected.mounts.is_empty());
2923 }
2924
2925 #[tokio::test]
2926 async fn test_collect_capabilities_combines_mounts() {
2927 let registry = CapabilityRegistry::with_builtins();
2928
2929 let collected = collect_capabilities(
2932 &["sample_data".to_string(), "current_time".to_string()],
2933 ®istry,
2934 &test_ctx(),
2935 )
2936 .await;
2937
2938 assert_eq!(collected.mounts.len(), 1);
2939 assert!(
2941 collected
2942 .applied_ids
2943 .iter()
2944 .any(|id| id == "session_file_system")
2945 );
2946 assert!(collected.applied_ids.iter().any(|id| id == "sample_data"));
2947 assert!(collected.applied_ids.iter().any(|id| id == "current_time"));
2948 }
2949
2950 #[test]
2951 fn test_sample_data_capability() {
2952 let registry = CapabilityRegistry::with_builtins();
2953 let cap = registry.get("sample_data").unwrap();
2954
2955 assert_eq!(cap.id(), "sample_data");
2956 assert_eq!(cap.name(), "Sample Data");
2957 assert_eq!(cap.status(), CapabilityStatus::Available);
2958
2959 assert!(cap.system_prompt_addition().is_some());
2961 assert!(cap.tools().is_empty());
2962
2963 assert!(!cap.mounts().is_empty());
2965 }
2966
2967 #[test]
2972 fn test_resolve_dependencies_empty() {
2973 let registry = CapabilityRegistry::with_builtins();
2974
2975 let resolved = resolve_dependencies(&[], ®istry).unwrap();
2976
2977 assert!(resolved.resolved_ids.is_empty());
2978 assert!(resolved.added_as_dependencies.is_empty());
2979 assert!(resolved.user_selected.is_empty());
2980 }
2981
2982 #[test]
2983 fn test_resolve_dependencies_no_deps() {
2984 let registry = CapabilityRegistry::with_builtins();
2985
2986 let resolved = resolve_dependencies(&["current_time".to_string()], ®istry).unwrap();
2988
2989 assert_eq!(resolved.resolved_ids, vec!["current_time"]);
2990 assert!(resolved.added_as_dependencies.is_empty());
2991 }
2992
2993 #[test]
2994 fn test_resolve_dependencies_with_deps() {
2995 let registry = CapabilityRegistry::with_builtins();
2996
2997 let resolved = resolve_dependencies(&["sample_data".to_string()], ®istry).unwrap();
2999
3000 assert_eq!(resolved.resolved_ids.len(), 2);
3002 let fs_pos = resolved
3003 .resolved_ids
3004 .iter()
3005 .position(|id| id == "session_file_system")
3006 .unwrap();
3007 let sd_pos = resolved
3008 .resolved_ids
3009 .iter()
3010 .position(|id| id == "sample_data")
3011 .unwrap();
3012 assert!(fs_pos < sd_pos, "FileSystem should come before SampleData");
3013
3014 assert_eq!(resolved.added_as_dependencies, vec!["session_file_system"]);
3016 }
3017
3018 #[test]
3019 fn test_resolve_dependencies_already_selected() {
3020 let registry = CapabilityRegistry::with_builtins();
3021
3022 let resolved = resolve_dependencies(
3024 &["session_file_system".to_string(), "sample_data".to_string()],
3025 ®istry,
3026 )
3027 .unwrap();
3028
3029 assert_eq!(resolved.resolved_ids.len(), 2);
3030 assert!(resolved.added_as_dependencies.is_empty());
3032 }
3033
3034 #[test]
3035 fn test_resolve_dependencies_preserves_order() {
3036 let registry = CapabilityRegistry::with_builtins();
3037
3038 let resolved =
3040 resolve_dependencies(&["current_time".to_string(), "noop".to_string()], ®istry)
3041 .unwrap();
3042
3043 assert_eq!(resolved.resolved_ids, vec!["current_time", "noop"]);
3044 }
3045
3046 #[test]
3047 fn test_resolve_dependencies_unknown_capability() {
3048 let registry = CapabilityRegistry::with_builtins();
3049
3050 let resolved =
3052 resolve_dependencies(&["unknown_capability".to_string()], ®istry).unwrap();
3053
3054 assert!(resolved.resolved_ids.is_empty());
3055 }
3056
3057 #[test]
3058 fn test_get_dependencies() {
3059 let registry = CapabilityRegistry::with_builtins();
3060
3061 let deps = get_dependencies("sample_data", ®istry);
3063 assert_eq!(deps, vec!["session_file_system"]);
3064
3065 let deps = get_dependencies("current_time", ®istry);
3067 assert!(deps.is_empty());
3068
3069 let deps = get_dependencies("unknown", ®istry);
3071 assert!(deps.is_empty());
3072 }
3073
3074 #[test]
3075 fn test_sample_data_has_dependency() {
3076 let registry = CapabilityRegistry::with_builtins();
3077 let cap = registry.get("sample_data").unwrap();
3078
3079 let deps = cap.dependencies();
3080 assert_eq!(deps.len(), 1);
3081 assert_eq!(deps[0], "session_file_system");
3082 }
3083
3084 #[test]
3085 fn test_noop_has_no_dependencies() {
3086 let registry = CapabilityRegistry::with_builtins();
3087 let cap = registry.get("noop").unwrap();
3088
3089 assert!(cap.dependencies().is_empty());
3090 }
3091
3092 #[test]
3096 fn test_circular_dependency_error() {
3097 struct CapA;
3099 struct CapB;
3100
3101 impl Capability for CapA {
3102 fn id(&self) -> &str {
3103 "test_cap_a"
3104 }
3105 fn name(&self) -> &str {
3106 "Test A"
3107 }
3108 fn description(&self) -> &str {
3109 "Test capability A"
3110 }
3111 fn dependencies(&self) -> Vec<&'static str> {
3112 vec!["test_cap_b"]
3113 }
3114 }
3115
3116 impl Capability for CapB {
3117 fn id(&self) -> &str {
3118 "test_cap_b"
3119 }
3120 fn name(&self) -> &str {
3121 "Test B"
3122 }
3123 fn description(&self) -> &str {
3124 "Test capability B"
3125 }
3126 fn dependencies(&self) -> Vec<&'static str> {
3127 vec!["test_cap_a"]
3128 }
3129 }
3130
3131 let mut registry = CapabilityRegistry::new();
3132 registry.register(CapA);
3133 registry.register(CapB);
3134
3135 let result = resolve_dependencies(&["test_cap_a".to_string()], ®istry);
3136
3137 assert!(result.is_err());
3138 match result.unwrap_err() {
3139 DependencyError::CircularDependency { capability_id, .. } => {
3140 assert_eq!(capability_id, "test_cap_a");
3141 }
3142 _ => panic!("Expected CircularDependency error"),
3143 }
3144 }
3145
3146 use crate::message_filter::{MessageFilter, MessageFilterProvider, MessageQuery};
3151
3152 struct FilterTestCapability {
3154 priority: i32,
3155 }
3156
3157 impl Capability for FilterTestCapability {
3158 fn id(&self) -> &str {
3159 "filter_test"
3160 }
3161 fn name(&self) -> &str {
3162 "Filter Test"
3163 }
3164 fn description(&self) -> &str {
3165 "Test capability with message filter"
3166 }
3167 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3168 Some(Arc::new(FilterTestProvider {
3169 priority: self.priority,
3170 }))
3171 }
3172 }
3173
3174 struct FilterTestProvider {
3175 priority: i32,
3176 }
3177
3178 impl MessageFilterProvider for FilterTestProvider {
3179 fn apply_filters(&self, query: &mut MessageQuery, config: &serde_json::Value) {
3180 if let Some(search) = config.get("search").and_then(|v| v.as_str()) {
3182 query
3183 .filters
3184 .push(MessageFilter::Search(search.to_string()));
3185 }
3186 }
3187
3188 fn priority(&self) -> i32 {
3189 self.priority
3190 }
3191 }
3192
3193 #[tokio::test]
3194 async fn test_collect_capabilities_with_configs_no_filter_providers() {
3195 let registry = CapabilityRegistry::with_builtins();
3196 let configs = vec![AgentCapabilityConfig {
3197 capability_ref: CapabilityId::new("current_time"),
3198 config: serde_json::json!({}),
3199 }];
3200
3201 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3202
3203 assert!(collected.message_filter_providers.is_empty());
3204 assert!(!collected.has_message_filters());
3205 }
3206
3207 #[tokio::test]
3208 async fn test_collect_capabilities_with_configs_with_filter_provider() {
3209 let mut registry = CapabilityRegistry::new();
3210 registry.register(FilterTestCapability { priority: 0 });
3211
3212 let configs = vec![AgentCapabilityConfig {
3213 capability_ref: CapabilityId::new("filter_test"),
3214 config: serde_json::json!({ "search": "hello" }),
3215 }];
3216
3217 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3218
3219 assert_eq!(collected.message_filter_providers.len(), 1);
3220 assert!(collected.has_message_filters());
3221 }
3222
3223 #[tokio::test]
3224 async fn test_collect_capabilities_with_configs_filter_priority_order() {
3225 struct HighPriorityCapability;
3227 struct LowPriorityCapability;
3228
3229 impl Capability for HighPriorityCapability {
3230 fn id(&self) -> &str {
3231 "high_priority"
3232 }
3233 fn name(&self) -> &str {
3234 "High Priority"
3235 }
3236 fn description(&self) -> &str {
3237 "Test"
3238 }
3239 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3240 Some(Arc::new(FilterTestProvider { priority: 10 }))
3241 }
3242 }
3243
3244 impl Capability for LowPriorityCapability {
3245 fn id(&self) -> &str {
3246 "low_priority"
3247 }
3248 fn name(&self) -> &str {
3249 "Low Priority"
3250 }
3251 fn description(&self) -> &str {
3252 "Test"
3253 }
3254 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3255 Some(Arc::new(FilterTestProvider { priority: -5 }))
3256 }
3257 }
3258
3259 let mut registry = CapabilityRegistry::new();
3260 registry.register(HighPriorityCapability);
3261 registry.register(LowPriorityCapability);
3262
3263 let configs = vec![
3265 AgentCapabilityConfig {
3266 capability_ref: CapabilityId::new("high_priority"),
3267 config: serde_json::json!({}),
3268 },
3269 AgentCapabilityConfig {
3270 capability_ref: CapabilityId::new("low_priority"),
3271 config: serde_json::json!({}),
3272 },
3273 ];
3274
3275 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3276
3277 assert_eq!(collected.message_filter_providers.len(), 2);
3279 assert_eq!(collected.message_filter_providers[0].0.priority(), -5);
3280 assert_eq!(collected.message_filter_providers[1].0.priority(), 10);
3281 }
3282
3283 #[tokio::test]
3284 async fn test_collected_capabilities_apply_message_filters() {
3285 let mut registry = CapabilityRegistry::new();
3286 registry.register(FilterTestCapability { priority: 0 });
3287
3288 let configs = vec![AgentCapabilityConfig {
3289 capability_ref: CapabilityId::new("filter_test"),
3290 config: serde_json::json!({ "search": "test_query" }),
3291 }];
3292
3293 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3294
3295 let session_id: SessionId = Uuid::now_v7().into();
3297 let mut query = MessageQuery::new(session_id);
3298
3299 collected.apply_message_filters(&mut query);
3300
3301 assert_eq!(query.filters.len(), 1);
3303 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3304 }
3305
3306 #[tokio::test]
3307 async fn test_collected_capabilities_apply_multiple_filters_in_priority_order() {
3308 struct SearchCapability {
3309 id: &'static str,
3310 search_term: &'static str,
3311 priority: i32,
3312 }
3313
3314 struct SearchProvider {
3315 search_term: &'static str,
3316 priority: i32,
3317 }
3318
3319 impl MessageFilterProvider for SearchProvider {
3320 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3321 query
3322 .filters
3323 .push(MessageFilter::Search(self.search_term.to_string()));
3324 }
3325
3326 fn priority(&self) -> i32 {
3327 self.priority
3328 }
3329 }
3330
3331 impl Capability for SearchCapability {
3332 fn id(&self) -> &str {
3333 self.id
3334 }
3335 fn name(&self) -> &str {
3336 "Search"
3337 }
3338 fn description(&self) -> &str {
3339 "Test"
3340 }
3341 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3342 Some(Arc::new(SearchProvider {
3343 search_term: self.search_term,
3344 priority: self.priority,
3345 }))
3346 }
3347 }
3348
3349 let mut registry = CapabilityRegistry::new();
3350 registry.register(SearchCapability {
3351 id: "cap_a",
3352 search_term: "alpha",
3353 priority: 5,
3354 });
3355 registry.register(SearchCapability {
3356 id: "cap_b",
3357 search_term: "beta",
3358 priority: 1,
3359 });
3360 registry.register(SearchCapability {
3361 id: "cap_c",
3362 search_term: "gamma",
3363 priority: 10,
3364 });
3365
3366 let configs = vec![
3367 AgentCapabilityConfig {
3368 capability_ref: CapabilityId::new("cap_a"),
3369 config: serde_json::json!({}),
3370 },
3371 AgentCapabilityConfig {
3372 capability_ref: CapabilityId::new("cap_b"),
3373 config: serde_json::json!({}),
3374 },
3375 AgentCapabilityConfig {
3376 capability_ref: CapabilityId::new("cap_c"),
3377 config: serde_json::json!({}),
3378 },
3379 ];
3380
3381 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3382
3383 let session_id: SessionId = Uuid::now_v7().into();
3384 let mut query = MessageQuery::new(session_id);
3385
3386 collected.apply_message_filters(&mut query);
3387
3388 assert_eq!(query.filters.len(), 3);
3390 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3391 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3392 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3393 }
3394
3395 #[test]
3396 fn test_capability_without_message_filter_returns_none() {
3397 let registry = CapabilityRegistry::with_builtins();
3398
3399 let noop = registry.get("noop").unwrap();
3400 assert!(noop.message_filter_provider().is_none());
3401
3402 let current_time = registry.get("current_time").unwrap();
3403 assert!(current_time.message_filter_provider().is_none());
3404 }
3405
3406 #[tokio::test]
3407 async fn test_collect_capabilities_preserves_config_for_filter_provider() {
3408 let mut registry = CapabilityRegistry::new();
3409 registry.register(FilterTestCapability { priority: 0 });
3410
3411 let test_config = serde_json::json!({
3412 "search": "custom_search",
3413 "extra_field": 42
3414 });
3415
3416 let configs = vec![AgentCapabilityConfig {
3417 capability_ref: CapabilityId::new("filter_test"),
3418 config: test_config.clone(),
3419 }];
3420
3421 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3422
3423 assert_eq!(collected.message_filter_providers.len(), 1);
3425 let (_, stored_config) = &collected.message_filter_providers[0];
3426 assert_eq!(*stored_config, test_config);
3427 }
3428
3429 #[test]
3434 fn test_collect_message_filters_only_collects_filters() {
3435 let mut registry = CapabilityRegistry::new();
3436 registry.register(FilterTestCapability { priority: 0 });
3437
3438 let configs = vec![AgentCapabilityConfig {
3439 capability_ref: CapabilityId::new("filter_test"),
3440 config: serde_json::json!({ "search": "test_query" }),
3441 }];
3442
3443 let collected = collect_message_filters_only(&configs, ®istry);
3444
3445 let session_id: SessionId = Uuid::now_v7().into();
3446 let mut query = MessageQuery::new(session_id);
3447 collected.apply_message_filters(&mut query);
3448
3449 assert_eq!(query.filters.len(), 1);
3450 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3451 }
3452
3453 #[test]
3454 fn test_collect_message_filters_only_skips_unknown_capabilities() {
3455 let registry = CapabilityRegistry::new();
3456
3457 let configs = vec![AgentCapabilityConfig {
3458 capability_ref: CapabilityId::new("nonexistent"),
3459 config: serde_json::json!({}),
3460 }];
3461
3462 let collected = collect_message_filters_only(&configs, ®istry);
3463 assert!(collected.message_filter_providers.is_empty());
3464 }
3465
3466 #[test]
3467 fn test_collect_message_filters_only_preserves_priority_order() {
3468 struct PriorityFilterCap {
3469 id: &'static str,
3470 search_term: &'static str,
3471 priority: i32,
3472 }
3473
3474 struct PriorityFilterProvider {
3475 search_term: &'static str,
3476 priority: i32,
3477 }
3478
3479 impl Capability for PriorityFilterCap {
3480 fn id(&self) -> &str {
3481 self.id
3482 }
3483 fn name(&self) -> &str {
3484 self.id
3485 }
3486 fn description(&self) -> &str {
3487 "priority test"
3488 }
3489 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3490 Some(Arc::new(PriorityFilterProvider {
3491 search_term: self.search_term,
3492 priority: self.priority,
3493 }))
3494 }
3495 }
3496
3497 impl MessageFilterProvider for PriorityFilterProvider {
3498 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3499 query
3500 .filters
3501 .push(MessageFilter::Search(self.search_term.to_string()));
3502 }
3503 fn priority(&self) -> i32 {
3504 self.priority
3505 }
3506 }
3507
3508 let mut registry = CapabilityRegistry::new();
3509 registry.register(PriorityFilterCap {
3510 id: "gamma",
3511 search_term: "gamma",
3512 priority: 10,
3513 });
3514 registry.register(PriorityFilterCap {
3515 id: "alpha",
3516 search_term: "alpha",
3517 priority: 5,
3518 });
3519 registry.register(PriorityFilterCap {
3520 id: "beta",
3521 search_term: "beta",
3522 priority: 1,
3523 });
3524
3525 let configs = vec![
3526 AgentCapabilityConfig {
3527 capability_ref: CapabilityId::new("gamma"),
3528 config: serde_json::json!({}),
3529 },
3530 AgentCapabilityConfig {
3531 capability_ref: CapabilityId::new("alpha"),
3532 config: serde_json::json!({}),
3533 },
3534 AgentCapabilityConfig {
3535 capability_ref: CapabilityId::new("beta"),
3536 config: serde_json::json!({}),
3537 },
3538 ];
3539
3540 let collected = collect_message_filters_only(&configs, ®istry);
3541
3542 let session_id: SessionId = Uuid::now_v7().into();
3543 let mut query = MessageQuery::new(session_id);
3544 collected.apply_message_filters(&mut query);
3545
3546 assert_eq!(query.filters.len(), 3);
3548 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3549 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3550 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3551 }
3552
3553 #[test]
3554 fn test_collect_message_filters_only_post_load_invoked() {
3555 use crate::message::Message;
3556
3557 struct PostLoadCap;
3558 struct PostLoadProvider;
3559
3560 impl Capability for PostLoadCap {
3561 fn id(&self) -> &str {
3562 "post_load_test"
3563 }
3564 fn name(&self) -> &str {
3565 "PostLoad Test"
3566 }
3567 fn description(&self) -> &str {
3568 "test"
3569 }
3570 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3571 Some(Arc::new(PostLoadProvider))
3572 }
3573 }
3574
3575 impl MessageFilterProvider for PostLoadProvider {
3576 fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
3577 fn priority(&self) -> i32 {
3578 0
3579 }
3580 fn post_load(&self, messages: &mut Vec<Message>, _config: &serde_json::Value) {
3581 messages.reverse();
3583 }
3584 }
3585
3586 let mut registry = CapabilityRegistry::new();
3587 registry.register(PostLoadCap);
3588
3589 let configs = vec![AgentCapabilityConfig {
3590 capability_ref: CapabilityId::new("post_load_test"),
3591 config: serde_json::json!({}),
3592 }];
3593
3594 let collected = collect_message_filters_only(&configs, ®istry);
3595
3596 let mut messages = vec![Message::user("first"), Message::user("second")];
3597 collected.apply_post_load_filters(&mut messages);
3598
3599 assert_eq!(messages[0].text(), Some("second"));
3601 assert_eq!(messages[1].text(), Some("first"));
3602 }
3603
3604 #[test]
3605 fn test_collect_model_view_providers_respects_compaction_capability_boundary() {
3606 use crate::tool_types::ToolCall;
3607
3608 fn tool_heavy_messages() -> Vec<Message> {
3609 let mut messages = vec![Message::user("inspect files repeatedly")];
3610 for index in 0..9 {
3611 let call_id = format!("call_{index}");
3612 messages.push(Message::assistant_with_tools(
3613 "",
3614 vec![ToolCall {
3615 id: call_id.clone(),
3616 name: "read_file".to_string(),
3617 arguments: serde_json::json!({"path": "/workspace/src/lib.rs"}),
3618 }],
3619 ));
3620 messages.push(Message::tool_result(
3621 call_id,
3622 Some(serde_json::json!({
3623 "path": "/workspace/src/lib.rs",
3624 "content": format!("{}{}", "large file line\n".repeat(1000), index),
3625 "total_lines": 1000,
3626 "lines_shown": {"start": 1, "end": 1000},
3627 "truncated": false
3628 })),
3629 None,
3630 ));
3631 }
3632 messages
3633 }
3634
3635 fn first_tool_result_is_masked(messages: &[Message]) -> bool {
3636 messages[2]
3637 .tool_result_content()
3638 .and_then(|result| result.result.as_ref())
3639 .and_then(|result| result.get("masked"))
3640 .and_then(|masked| masked.as_bool())
3641 .unwrap_or(false)
3642 }
3643
3644 let mut registry = CapabilityRegistry::new();
3645 registry.register(CompactionCapability);
3646 let context = ModelViewContext {
3647 session_id: SessionId::new(),
3648 prior_usage: None,
3649 };
3650
3651 let no_compaction = collect_model_view_providers(&[], ®istry, None);
3652 let unmasked = no_compaction.apply_model_view(tool_heavy_messages(), &context);
3653 assert!(!first_tool_result_is_masked(&unmasked));
3654
3655 let compaction = collect_model_view_providers(
3656 &[AgentCapabilityConfig {
3657 capability_ref: CapabilityId::new(COMPACTION_CAPABILITY_ID),
3658 config: serde_json::json!({}),
3659 }],
3660 ®istry,
3661 None,
3662 );
3663 let masked = compaction.apply_model_view(tool_heavy_messages(), &context);
3664 assert!(first_tool_result_is_masked(&masked));
3665 let last_tool = masked.last().unwrap().tool_result_content().unwrap();
3666 assert!(last_tool.result.as_ref().unwrap().get("content").is_some());
3667 }
3668
3669 struct DelegatingFilterCap {
3672 id: &'static str,
3673 inner: std::sync::Arc<InnerFilterCap>,
3674 }
3675 struct InnerFilterCap;
3676
3677 impl Capability for InnerFilterCap {
3678 fn id(&self) -> &str {
3679 "inner_filter"
3680 }
3681 fn name(&self) -> &str {
3682 "Inner Filter"
3683 }
3684 fn description(&self) -> &str {
3685 "inner"
3686 }
3687 fn message_filter_provider(&self) -> Option<std::sync::Arc<dyn MessageFilterProvider>> {
3688 Some(std::sync::Arc::new(SentinelFilter))
3689 }
3690 }
3691 struct SentinelFilter;
3692 impl MessageFilterProvider for SentinelFilter {
3693 fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
3694 }
3695 impl Capability for DelegatingFilterCap {
3696 fn id(&self) -> &str {
3697 self.id
3698 }
3699 fn name(&self) -> &str {
3700 "Delegating Filter"
3701 }
3702 fn description(&self) -> &str {
3703 "delegating"
3704 }
3705 fn message_filter_provider(&self) -> Option<std::sync::Arc<dyn MessageFilterProvider>> {
3706 None }
3708 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
3709 Some(&*self.inner)
3710 }
3711 }
3712
3713 #[test]
3714 fn test_collect_message_filters_only_honors_resolve_for_model_delegation() {
3715 let inner = std::sync::Arc::new(InnerFilterCap);
3716 let outer = DelegatingFilterCap {
3717 id: "delegating_filter",
3718 inner: inner.clone(),
3719 };
3720
3721 let mut registry = CapabilityRegistry::new();
3722 registry.register(outer);
3723
3724 let configs = vec![AgentCapabilityConfig {
3725 capability_ref: CapabilityId::new("delegating_filter"),
3726 config: serde_json::json!({}),
3727 }];
3728
3729 let collected = collect_message_filters_only(&configs, ®istry);
3732 assert_eq!(
3733 collected.message_filter_providers.len(),
3734 1,
3735 "provider from resolved inner capability must be collected"
3736 );
3737 }
3738
3739 struct DelegatingMvpCap {
3740 id: &'static str,
3741 inner: std::sync::Arc<InnerMvpCap>,
3742 }
3743 struct InnerMvpCap;
3744
3745 impl Capability for InnerMvpCap {
3746 fn id(&self) -> &str {
3747 "inner_mvp"
3748 }
3749 fn name(&self) -> &str {
3750 "Inner MVP"
3751 }
3752 fn description(&self) -> &str {
3753 "inner"
3754 }
3755 fn model_view_provider(
3756 &self,
3757 ) -> Option<std::sync::Arc<dyn crate::capabilities::ModelViewProvider>> {
3758 struct NoopMvp;
3760 impl crate::capabilities::ModelViewProvider for NoopMvp {
3761 fn apply_model_view(
3762 &self,
3763 messages: Vec<Message>,
3764 _config: &serde_json::Value,
3765 _context: &ModelViewContext<'_>,
3766 ) -> Vec<Message> {
3767 messages
3768 }
3769 }
3770 Some(std::sync::Arc::new(NoopMvp))
3771 }
3772 }
3773 impl Capability for DelegatingMvpCap {
3774 fn id(&self) -> &str {
3775 self.id
3776 }
3777 fn name(&self) -> &str {
3778 "Delegating MVP"
3779 }
3780 fn description(&self) -> &str {
3781 "delegating"
3782 }
3783 fn model_view_provider(
3784 &self,
3785 ) -> Option<std::sync::Arc<dyn crate::capabilities::ModelViewProvider>> {
3786 None }
3788 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
3789 Some(&*self.inner)
3790 }
3791 }
3792
3793 #[test]
3794 fn test_collect_model_view_providers_honors_resolve_for_model_delegation() {
3795 let inner = std::sync::Arc::new(InnerMvpCap);
3796 let outer = DelegatingMvpCap {
3797 id: "delegating_mvp",
3798 inner: inner.clone(),
3799 };
3800
3801 let mut registry = CapabilityRegistry::new();
3802 registry.register(outer);
3803
3804 let configs = vec![AgentCapabilityConfig {
3805 capability_ref: CapabilityId::new("delegating_mvp"),
3806 config: serde_json::json!({}),
3807 }];
3808
3809 let collected = collect_model_view_providers(&configs, ®istry, None);
3812 assert_eq!(
3813 collected.model_view_providers.len(),
3814 1,
3815 "provider from resolved inner capability must be collected"
3816 );
3817 }
3818
3819 #[tokio::test]
3829 async fn test_bashkit_shell_capability_produces_bash_tool() {
3830 let registry = CapabilityRegistry::with_builtins();
3831 let collected =
3832 collect_capabilities(&["bashkit_shell".to_string()], ®istry, &test_ctx()).await;
3833
3834 let tool_names: Vec<&str> = collected
3835 .tool_definitions
3836 .iter()
3837 .map(|t| t.name())
3838 .collect();
3839 assert!(
3840 tool_names.contains(&"bash"),
3841 "bashkit_shell capability must produce 'bash' tool, got: {:?}",
3842 tool_names
3843 );
3844 assert!(
3845 !collected.tools.is_empty(),
3846 "bashkit_shell must provide tool implementations"
3847 );
3848 }
3849
3850 #[tokio::test]
3851 async fn test_generic_harness_capability_set_produces_bash_tool() {
3852 let generic_harness_caps = vec![
3855 "session_file_system".to_string(),
3856 "bashkit_shell".to_string(),
3857 "web_fetch".to_string(),
3858 "session_storage".to_string(),
3859 "session".to_string(),
3860 "agent_instructions".to_string(),
3861 "skills".to_string(),
3862 "infinity_context".to_string(),
3863 "auto_tool_search".to_string(),
3864 ];
3865
3866 let registry = CapabilityRegistry::with_builtins();
3867 let collected = collect_capabilities(&generic_harness_caps, ®istry, &test_ctx()).await;
3868
3869 let tool_names: Vec<&str> = collected
3870 .tool_definitions
3871 .iter()
3872 .map(|t| t.name())
3873 .collect();
3874 assert!(
3875 tool_names.contains(&"bash"),
3876 "Generic Harness capabilities must produce 'bash' tool, got: {:?}",
3877 tool_names
3878 );
3879 }
3880
3881 #[tokio::test]
3882 async fn test_collect_capabilities_tool_count_matches_definitions() {
3883 let registry = CapabilityRegistry::with_builtins();
3886 let collected =
3887 collect_capabilities(&["bashkit_shell".to_string()], ®istry, &test_ctx()).await;
3888
3889 assert_eq!(
3890 collected.tools.len(),
3891 collected.tool_definitions.len(),
3892 "tool implementations ({}) must match tool definitions ({})",
3893 collected.tools.len(),
3894 collected.tool_definitions.len(),
3895 );
3896 }
3897
3898 #[tokio::test]
3902 async fn test_collect_capabilities_resolves_dependencies() {
3903 let registry = CapabilityRegistry::with_builtins();
3906 let collected =
3907 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
3908
3909 assert!(
3911 collected
3912 .applied_ids
3913 .iter()
3914 .any(|id| id == "session_file_system"),
3915 "collect_capabilities must apply session_file_system as a dependency; applied_ids: {:?}",
3916 collected.applied_ids
3917 );
3918
3919 let tool_names: Vec<&str> = collected
3920 .tool_definitions
3921 .iter()
3922 .map(|t| t.name())
3923 .collect();
3924
3925 assert!(
3927 tool_names.contains(&"read_file") && tool_names.contains(&"write_file"),
3928 "collect_capabilities must resolve dependencies and include dependency tools, got: {:?}",
3929 tool_names
3930 );
3931
3932 assert_eq!(
3934 collected.tools.len(),
3935 collected.tool_definitions.len(),
3936 "dependency-added tools must have implementations, not just definitions"
3937 );
3938 }
3939
3940 #[test]
3941 fn test_defaults_do_not_include_bash() {
3942 let registry = crate::ToolRegistry::with_defaults();
3945 assert!(
3946 !registry.has("bash"),
3947 "with_defaults() must not include 'bash' — it comes from bashkit_shell capability"
3948 );
3949 }
3950
3951 #[tokio::test]
3958 async fn test_background_execution_auto_activates_with_bashkit_shell() {
3959 let registry = CapabilityRegistry::with_builtins();
3960 let collected =
3961 collect_capabilities(&["bashkit_shell".to_string()], ®istry, &test_ctx()).await;
3962
3963 let tool_names: Vec<&str> = collected
3964 .tool_definitions
3965 .iter()
3966 .map(|t| t.name())
3967 .collect();
3968 assert!(
3969 tool_names.contains(&"spawn_background"),
3970 "spawn_background must be auto-activated when bashkit_shell (a \
3971 background-capable tool) is in the agent's capability set; got: {:?}",
3972 tool_names
3973 );
3974 assert!(
3975 collected
3976 .applied_ids
3977 .iter()
3978 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
3979 "background_execution must be in applied_ids when auto-activated; \
3980 got: {:?}",
3981 collected.applied_ids
3982 );
3983
3984 assert!(
3986 collected
3987 .tools
3988 .iter()
3989 .any(|t| t.name() == "spawn_background"),
3990 "spawn_background tool implementation must be present alongside the \
3991 definition (lockstep contract)"
3992 );
3993 }
3994
3995 #[tokio::test]
3998 async fn test_background_execution_does_not_auto_activate_without_hint() {
3999 let registry = CapabilityRegistry::with_builtins();
4000 let collected =
4002 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
4003
4004 let tool_names: Vec<&str> = collected
4005 .tool_definitions
4006 .iter()
4007 .map(|t| t.name())
4008 .collect();
4009 assert!(
4010 !tool_names.contains(&"spawn_background"),
4011 "spawn_background must NOT be activated without a background-capable \
4012 tool; got: {:?}",
4013 tool_names
4014 );
4015 assert!(
4016 !collected
4017 .applied_ids
4018 .iter()
4019 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
4020 "background_execution must not appear in applied_ids when no \
4021 background-capable tool is present; got: {:?}",
4022 collected.applied_ids
4023 );
4024 }
4025
4026 #[tokio::test]
4030 async fn test_background_execution_explicit_selection_is_idempotent() {
4031 let registry = CapabilityRegistry::with_builtins();
4032 let collected = collect_capabilities(
4033 &[
4034 "bashkit_shell".to_string(),
4035 BACKGROUND_EXECUTION_CAPABILITY_ID.to_string(),
4036 ],
4037 ®istry,
4038 &test_ctx(),
4039 )
4040 .await;
4041
4042 let spawn_background_count = collected
4043 .tool_definitions
4044 .iter()
4045 .filter(|t| t.name() == "spawn_background")
4046 .count();
4047 assert_eq!(
4048 spawn_background_count, 1,
4049 "spawn_background must appear exactly once even when \
4050 background_execution is selected explicitly alongside a \
4051 background-capable tool"
4052 );
4053 let applied_count = collected
4054 .applied_ids
4055 .iter()
4056 .filter(|id| id.as_str() == BACKGROUND_EXECUTION_CAPABILITY_ID)
4057 .count();
4058 assert_eq!(
4059 applied_count, 1,
4060 "background_execution must appear exactly once in applied_ids"
4061 );
4062 }
4063
4064 #[test]
4069 fn test_defaults_do_not_include_spawn_background() {
4070 let registry = crate::ToolRegistry::with_defaults();
4071 assert!(
4072 !registry.has("spawn_background"),
4073 "with_defaults() must not include 'spawn_background' — it comes \
4074 from the background_execution capability (EVE-501)"
4075 );
4076 }
4077
4078 #[test]
4083 fn test_capability_features_default_empty() {
4084 let registry = CapabilityRegistry::with_builtins();
4085
4086 let noop = registry.get("noop").unwrap();
4088 assert!(noop.features().is_empty());
4089
4090 let current_time = registry.get("current_time").unwrap();
4091 assert!(current_time.features().is_empty());
4092 }
4093
4094 #[test]
4095 fn test_file_system_capability_features() {
4096 let registry = CapabilityRegistry::with_builtins();
4097
4098 let fs = registry.get("session_file_system").unwrap();
4099 assert_eq!(fs.features(), vec!["file_system"]);
4100 }
4101
4102 #[test]
4103 fn test_bashkit_shell_capability_features() {
4104 let registry = CapabilityRegistry::with_builtins();
4105
4106 let bash = registry.get("bashkit_shell").unwrap();
4107 assert_eq!(bash.features(), vec!["file_system"]);
4108 }
4109
4110 #[test]
4111 fn test_alias_resolves_to_canonical_capability() {
4112 let registry = CapabilityRegistry::with_builtins();
4113
4114 let via_alias = registry.get("virtual_bash").unwrap();
4116 assert_eq!(via_alias.id(), "bashkit_shell");
4117 assert!(registry.has("virtual_bash"));
4118 assert_eq!(registry.canonical_id("virtual_bash"), Some("bashkit_shell"));
4119 assert_eq!(
4120 registry.canonical_id("bashkit_shell"),
4121 Some("bashkit_shell")
4122 );
4123 assert_eq!(registry.canonical_id("nonexistent"), None);
4124 }
4125
4126 #[test]
4127 fn test_alias_dedupes_with_canonical_in_dependency_resolution() {
4128 let registry = CapabilityRegistry::with_builtins();
4129
4130 let resolved = resolve_dependencies(
4133 &["virtual_bash".to_string(), "bashkit_shell".to_string()],
4134 ®istry,
4135 )
4136 .unwrap();
4137 let bash_ids: Vec<_> = resolved
4138 .resolved_ids
4139 .iter()
4140 .filter(|id| id.as_str() == "bashkit_shell" || id.as_str() == "virtual_bash")
4141 .collect();
4142 assert_eq!(bash_ids, vec!["bashkit_shell"]);
4143 assert!(
4145 !resolved
4146 .added_as_dependencies
4147 .contains(&"bashkit_shell".to_string())
4148 );
4149 }
4150
4151 #[test]
4152 fn test_alias_preserves_explicit_config_in_resolution() {
4153 let registry = CapabilityRegistry::with_builtins();
4154
4155 let configs = vec![AgentCapabilityConfig::with_config(
4156 "virtual_bash".to_string(),
4157 serde_json::json!({"key": "value"}),
4158 )];
4159 let resolved = resolve_capability_configs(&configs, ®istry).unwrap();
4160 let bash = resolved
4161 .iter()
4162 .find(|c| c.capability_id() == "bashkit_shell")
4163 .expect("alias must resolve to canonical bashkit_shell config");
4164 assert_eq!(bash.config, serde_json::json!({"key": "value"}));
4165 }
4166
4167 #[test]
4168 fn test_unregister_by_alias_removes_capability_and_aliases() {
4169 let mut registry = CapabilityRegistry::with_builtins();
4170
4171 assert!(registry.unregister("virtual_bash").is_some());
4172 assert!(!registry.has("bashkit_shell"));
4173 assert!(!registry.has("virtual_bash"));
4174 }
4175
4176 #[test]
4177 fn test_session_storage_capability_features() {
4178 let registry = CapabilityRegistry::with_builtins();
4179
4180 let storage = registry.get("session_storage").unwrap();
4181 let features = storage.features();
4182 assert!(features.contains(&"secrets"));
4183 assert!(features.contains(&"key_value"));
4184 }
4185
4186 #[test]
4187 fn test_session_schedule_capability_features() {
4188 let registry = CapabilityRegistry::with_builtins();
4189
4190 let schedule = registry.get("session_schedule").unwrap();
4191 assert_eq!(schedule.features(), vec!["schedules"]);
4192 }
4193
4194 #[test]
4195 fn test_session_sql_database_capability_features() {
4196 let registry = CapabilityRegistry::with_builtins();
4197
4198 let sql = registry.get("session_sql_database").unwrap();
4199 assert_eq!(sql.features(), vec!["sql_database"]);
4200 }
4201
4202 #[test]
4203 fn test_sample_data_capability_features() {
4204 let registry = CapabilityRegistry::with_builtins();
4205
4206 let sample = registry.get("sample_data").unwrap();
4207 assert_eq!(sample.features(), vec!["file_system"]);
4208 }
4209
4210 #[test]
4211 fn test_compute_features_empty() {
4212 let registry = CapabilityRegistry::with_builtins();
4213
4214 let features = compute_features(&[], ®istry);
4215 assert!(features.is_empty());
4216 }
4217
4218 #[test]
4219 fn test_compute_features_single_capability() {
4220 let registry = CapabilityRegistry::with_builtins();
4221
4222 let features = compute_features(&["session_schedule".to_string()], ®istry);
4223 assert_eq!(features, vec!["schedules"]);
4224 }
4225
4226 #[test]
4227 fn test_compute_features_multiple_capabilities() {
4228 let registry = CapabilityRegistry::with_builtins();
4229
4230 let features = compute_features(
4231 &[
4232 "session_file_system".to_string(),
4233 "session_storage".to_string(),
4234 "session_schedule".to_string(),
4235 ],
4236 ®istry,
4237 );
4238 assert!(features.contains(&"file_system".to_string()));
4239 assert!(features.contains(&"secrets".to_string()));
4240 assert!(features.contains(&"key_value".to_string()));
4241 assert!(features.contains(&"schedules".to_string()));
4242 }
4243
4244 #[test]
4245 fn test_compute_features_deduplicates() {
4246 let registry = CapabilityRegistry::with_builtins();
4247
4248 let features = compute_features(
4250 &[
4251 "session_file_system".to_string(),
4252 "bashkit_shell".to_string(),
4253 ],
4254 ®istry,
4255 );
4256 let file_system_count = features.iter().filter(|f| *f == "file_system").count();
4257 assert_eq!(file_system_count, 1, "file_system should appear only once");
4258 }
4259
4260 #[test]
4261 fn test_compute_features_includes_dependency_features() {
4262 let registry = CapabilityRegistry::with_builtins();
4263
4264 let features = compute_features(&["bashkit_shell".to_string()], ®istry);
4266 assert!(features.contains(&"file_system".to_string()));
4267 }
4268
4269 #[test]
4270 fn test_compute_features_generic_harness_set() {
4271 let registry = CapabilityRegistry::with_builtins();
4272
4273 let features = compute_features(
4275 &[
4276 "session_file_system".to_string(),
4277 "bashkit_shell".to_string(),
4278 "session_storage".to_string(),
4279 "session".to_string(),
4280 "session_schedule".to_string(),
4281 ],
4282 ®istry,
4283 );
4284 assert!(features.contains(&"file_system".to_string()));
4285 assert!(features.contains(&"secrets".to_string()));
4286 assert!(features.contains(&"key_value".to_string()));
4287 assert!(features.contains(&"schedules".to_string()));
4288 }
4289
4290 #[test]
4291 fn test_compute_features_unknown_capability_ignored() {
4292 let registry = CapabilityRegistry::with_builtins();
4293
4294 let features = compute_features(
4295 &["unknown_cap".to_string(), "session_schedule".to_string()],
4296 ®istry,
4297 );
4298 assert_eq!(features, vec!["schedules"]);
4299 }
4300
4301 #[test]
4302 fn test_risk_level_ordering() {
4303 assert!(RiskLevel::Low < RiskLevel::Medium);
4304 assert!(RiskLevel::Medium < RiskLevel::High);
4305 }
4306
4307 #[test]
4308 fn test_risk_level_serde_roundtrip() {
4309 let high = RiskLevel::High;
4310 let json = serde_json::to_string(&high).unwrap();
4311 assert_eq!(json, "\"high\"");
4312 let back: RiskLevel = serde_json::from_str(&json).unwrap();
4313 assert_eq!(back, RiskLevel::High);
4314 }
4315
4316 #[test]
4317 fn test_capability_risk_levels() {
4318 let registry = CapabilityRegistry::with_builtins();
4319
4320 let bash = registry.get("bashkit_shell").unwrap();
4322 assert_eq!(bash.risk_level(), RiskLevel::High);
4323
4324 let fetch = registry.get("web_fetch").unwrap();
4326 assert_eq!(fetch.risk_level(), RiskLevel::High);
4327
4328 let noop = registry.get("noop").unwrap();
4330 assert_eq!(noop.risk_level(), RiskLevel::Low);
4331 }
4332
4333 #[tokio::test]
4338 async fn test_apply_capabilities_openai_tool_search() {
4339 let registry = CapabilityRegistry::with_builtins();
4340 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4341
4342 let applied = apply_capabilities(
4343 base_runtime_agent.clone(),
4344 &["openai_tool_search".to_string()],
4345 ®istry,
4346 &test_ctx(),
4347 )
4348 .await;
4349
4350 assert_eq!(
4352 applied.runtime_agent.system_prompt,
4353 base_runtime_agent.system_prompt
4354 );
4355 assert!(applied.tool_registry.is_empty());
4356 assert_eq!(applied.applied_ids, vec!["openai_tool_search"]);
4357
4358 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
4360 assert!(ts.enabled);
4361 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
4362 }
4363
4364 #[tokio::test]
4365 async fn test_apply_capabilities_openai_tool_search_with_other_capabilities() {
4366 let registry = CapabilityRegistry::with_builtins();
4367 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4368
4369 let applied = apply_capabilities(
4370 base_runtime_agent,
4371 &[
4372 "current_time".to_string(),
4373 "openai_tool_search".to_string(),
4374 "test_math".to_string(),
4375 ],
4376 ®istry,
4377 &test_ctx(),
4378 )
4379 .await;
4380
4381 assert!(applied.tool_registry.has("get_current_time"));
4383 assert!(applied.tool_registry.has("add"));
4384 assert!(applied.tool_registry.has("subtract"));
4385 assert!(applied.tool_registry.has("multiply"));
4386 assert!(applied.tool_registry.has("divide"));
4387
4388 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
4390 assert!(ts.enabled);
4391 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
4392 }
4393
4394 #[tokio::test]
4395 async fn test_collect_capabilities_tool_search_custom_threshold() {
4396 let registry = CapabilityRegistry::with_builtins();
4397
4398 let configs = vec![AgentCapabilityConfig {
4399 capability_ref: CapabilityId::new("openai_tool_search"),
4400 config: serde_json::json!({"threshold": 5}),
4401 }];
4402
4403 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4404
4405 let ts = collected.tool_search.as_ref().unwrap();
4406 assert!(ts.enabled);
4407 assert_eq!(ts.threshold, 5);
4408 }
4409
4410 #[tokio::test]
4411 async fn test_collect_capabilities_auto_tool_search_resolves_to_generic_off_native() {
4412 let registry = CapabilityRegistry::with_builtins();
4413
4414 let configs = vec![
4415 AgentCapabilityConfig {
4416 capability_ref: CapabilityId::new("auto_tool_search"),
4417 config: serde_json::json!({"threshold": 2}),
4418 },
4419 AgentCapabilityConfig {
4420 capability_ref: CapabilityId::new("test_math"),
4421 config: serde_json::json!({}),
4422 },
4423 ];
4424
4425 let ctx = test_ctx().with_model("claude-sonnet-4-5-20250514");
4428 let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
4429
4430 assert!(
4431 collected.tool_search.is_none(),
4432 "auto_tool_search must not set a hosted config on a non-native model"
4433 );
4434 assert!(
4435 collected
4436 .tools
4437 .iter()
4438 .any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
4439 "auto_tool_search must contribute the client-side tool_search tool"
4440 );
4441 assert!(
4442 !collected.tool_definition_hooks.is_empty(),
4443 "auto_tool_search must contribute a client-side deferral hook"
4444 );
4445
4446 let mut transformed = collected.tool_definitions.clone();
4447 for hook in &collected.tool_definition_hooks {
4448 transformed = hook.transform(transformed);
4449 }
4450 let add_tool = transformed
4451 .iter()
4452 .find(|tool| tool.name() == "add")
4453 .expect("test_math contributes add");
4454 assert!(
4455 add_tool.parameters().get("properties").is_none(),
4456 "generic auto_tool_search must honor the configured threshold"
4457 );
4458 }
4459
4460 #[tokio::test]
4461 async fn test_collect_capabilities_auto_tool_search_resolves_to_hosted_on_native() {
4462 let registry = CapabilityRegistry::with_builtins();
4463
4464 let configs = vec![AgentCapabilityConfig {
4465 capability_ref: CapabilityId::new("auto_tool_search"),
4466 config: serde_json::json!({"threshold": 7}),
4467 }];
4468
4469 let ctx = test_ctx().with_model("gpt-5.4");
4472 let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
4473
4474 let ts = collected
4475 .tool_search
4476 .as_ref()
4477 .expect("auto_tool_search must set a hosted config on a native model");
4478 assert!(ts.enabled);
4479 assert_eq!(ts.threshold, 7);
4480 assert!(
4481 !collected
4482 .tools
4483 .iter()
4484 .any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
4485 "hosted mechanism must not contribute the client-side tool_search tool"
4486 );
4487 assert!(
4488 collected.tool_definition_hooks.is_empty(),
4489 "hosted mechanism must not contribute a client-side deferral hook"
4490 );
4491 }
4492
4493 #[tokio::test]
4494 async fn test_collect_capabilities_no_tool_search_without_capability() {
4495 let registry = CapabilityRegistry::with_builtins();
4496
4497 let configs = vec![AgentCapabilityConfig {
4498 capability_ref: CapabilityId::new("current_time"),
4499 config: serde_json::json!({}),
4500 }];
4501
4502 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4503
4504 assert!(collected.tool_search.is_none());
4505 }
4506
4507 #[tokio::test]
4508 async fn test_collect_capabilities_tool_search_category_propagation() {
4509 let registry = CapabilityRegistry::with_builtins();
4510
4511 let configs = vec![
4513 AgentCapabilityConfig {
4514 capability_ref: CapabilityId::new("test_math"),
4515 config: serde_json::json!({}),
4516 },
4517 AgentCapabilityConfig {
4518 capability_ref: CapabilityId::new("openai_tool_search"),
4519 config: serde_json::json!({}),
4520 },
4521 ];
4522
4523 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4524
4525 assert!(collected.tool_search.is_some());
4527
4528 for tool_def in &collected.tool_definitions {
4530 if ["add", "subtract", "multiply", "divide"].contains(&tool_def.name()) {
4532 assert!(
4533 tool_def.category().is_some(),
4534 "Tool {} should have a category from its capability",
4535 tool_def.name()
4536 );
4537 }
4538 }
4539 }
4540
4541 #[tokio::test]
4542 async fn test_apply_capabilities_prompt_caching() {
4543 let registry = CapabilityRegistry::with_builtins();
4544 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4545
4546 let applied = apply_capabilities(
4547 base_runtime_agent.clone(),
4548 &["prompt_caching".to_string()],
4549 ®istry,
4550 &test_ctx(),
4551 )
4552 .await;
4553
4554 assert_eq!(
4555 applied.runtime_agent.system_prompt,
4556 base_runtime_agent.system_prompt
4557 );
4558 assert!(applied.tool_registry.is_empty());
4559 assert_eq!(applied.applied_ids, vec!["prompt_caching"]);
4560
4561 let prompt_cache = applied.runtime_agent.prompt_cache.as_ref().unwrap();
4562 assert!(prompt_cache.enabled);
4563 assert_eq!(
4564 prompt_cache.strategy,
4565 crate::llm_driver_registry::PromptCacheStrategy::Auto
4566 );
4567 assert!(prompt_cache.gemini_cached_content.is_none());
4568 }
4569
4570 #[tokio::test]
4571 async fn test_collect_capabilities_prompt_caching_custom_strategy() {
4572 let registry = CapabilityRegistry::with_builtins();
4573
4574 let configs = vec![AgentCapabilityConfig {
4575 capability_ref: CapabilityId::new("prompt_caching"),
4576 config: serde_json::json!({"strategy": "auto"}),
4577 }];
4578
4579 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4580
4581 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
4582 assert!(prompt_cache.enabled);
4583 assert_eq!(
4584 prompt_cache.strategy,
4585 crate::llm_driver_registry::PromptCacheStrategy::Auto
4586 );
4587 assert!(prompt_cache.gemini_cached_content.is_none());
4588 }
4589
4590 #[tokio::test]
4591 async fn test_collect_capabilities_prompt_caching_gemini_cached_content() {
4592 let registry = CapabilityRegistry::with_builtins();
4593
4594 let configs = vec![AgentCapabilityConfig {
4595 capability_ref: CapabilityId::new("prompt_caching"),
4596 config: serde_json::json!({
4597 "strategy": "auto",
4598 "gemini_cached_content": "cachedContents/demo-cache"
4599 }),
4600 }];
4601
4602 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4603
4604 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
4605 assert_eq!(
4606 prompt_cache.gemini_cached_content.as_deref(),
4607 Some("cachedContents/demo-cache")
4608 );
4609 }
4610
4611 struct SkillContributingCapability;
4616
4617 impl Capability for SkillContributingCapability {
4618 fn id(&self) -> &str {
4619 "contributes_skills"
4620 }
4621 fn name(&self) -> &str {
4622 "Contributes Skills"
4623 }
4624 fn description(&self) -> &str {
4625 "Test capability that contributes skills."
4626 }
4627 fn contribute_skills(&self) -> Vec<SkillContribution> {
4628 vec![
4629 SkillContribution::new("alpha-skill", "Alpha skill desc", "# Alpha\nDo alpha.")
4630 .with_files(vec![(
4631 "scripts/a.sh".to_string(),
4632 "#!/bin/sh\necho a\n".to_string(),
4633 )]),
4634 SkillContribution::new("beta-skill", "Beta skill desc", "# Beta\nDo beta.")
4635 .with_user_invocable(false),
4636 ]
4637 }
4638 }
4639
4640 fn skill_md_from_entries(entries: &HashMap<String, MountEntry>) -> &str {
4641 match &entries.get("SKILL.md").expect("SKILL.md missing").source {
4642 MountSource::InlineFile { content, .. } => content.as_str(),
4643 _ => panic!("Expected InlineFile for SKILL.md"),
4644 }
4645 }
4646
4647 #[tokio::test]
4648 async fn test_contribute_skills_normalized_to_mounts() {
4649 let mut registry = CapabilityRegistry::new();
4650 registry.register(SkillContributingCapability);
4651
4652 let configs = vec![AgentCapabilityConfig {
4653 capability_ref: CapabilityId::new("contributes_skills"),
4654 config: serde_json::json!({}),
4655 }];
4656
4657 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4658
4659 let skill_mounts: Vec<_> = collected
4660 .mounts
4661 .iter()
4662 .filter(|m| m.path.starts_with("/.agents/skills/"))
4663 .collect();
4664 assert_eq!(skill_mounts.len(), 2);
4665
4666 for m in &skill_mounts {
4669 assert!(m.is_readonly());
4670 assert_eq!(m.capability_id, "contributes_skills");
4671 }
4672
4673 let alpha = skill_mounts
4674 .iter()
4675 .find(|m| m.path == "/.agents/skills/alpha-skill")
4676 .expect("alpha-skill mount missing");
4677 match &alpha.source {
4678 MountSource::InlineDirectory { entries } => {
4679 assert!(entries.contains_key("SKILL.md"));
4680 assert!(entries.contains_key("scripts/a.sh"));
4681 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
4682 assert_eq!(parsed.name, "alpha-skill");
4683 assert!(parsed.user_invocable);
4684 }
4685 _ => panic!("Expected InlineDirectory"),
4686 }
4687
4688 let beta = skill_mounts
4689 .iter()
4690 .find(|m| m.path == "/.agents/skills/beta-skill")
4691 .expect("beta-skill mount missing");
4692 match &beta.source {
4693 MountSource::InlineDirectory { entries } => {
4694 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
4695 assert!(!parsed.user_invocable);
4696 }
4697 _ => panic!("Expected InlineDirectory"),
4698 }
4699 }
4700
4701 #[tokio::test]
4702 async fn test_contribute_skills_default_empty() {
4703 let mut registry = CapabilityRegistry::new();
4706 registry.register(FilterTestCapability { priority: 0 });
4707
4708 let configs = vec![AgentCapabilityConfig {
4709 capability_ref: CapabilityId::new("filter_test"),
4710 config: serde_json::json!({}),
4711 }];
4712
4713 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4714 assert!(
4715 collected
4716 .mounts
4717 .iter()
4718 .all(|m| !m.path.starts_with("/.agents/skills/"))
4719 );
4720 }
4721
4722 struct LocalizedCapability;
4723
4724 impl Capability for LocalizedCapability {
4725 fn id(&self) -> &str {
4726 "localized"
4727 }
4728 fn name(&self) -> &str {
4729 "Localized"
4730 }
4731 fn description(&self) -> &str {
4732 "English description"
4733 }
4734 fn localizations(&self) -> Vec<CapabilityLocalization> {
4735 vec![
4736 CapabilityLocalization {
4737 locale: "en",
4738 name: None,
4739 description: None,
4740 config_description: Some("Controls things."),
4741 config_overlay: None,
4742 },
4743 CapabilityLocalization {
4744 locale: "uk",
4745 name: Some("Локалізована"),
4746 description: Some("Український опис"),
4747 config_description: Some("Керує налаштуваннями."),
4748 config_overlay: None,
4749 },
4750 ]
4751 }
4752 }
4753
4754 #[test]
4755 fn localized_name_falls_back_exact_language_then_base() {
4756 let cap = LocalizedCapability;
4757 assert_eq!(cap.localized_name(Some("uk-UA")), "Локалізована");
4759 assert_eq!(cap.localized_name(Some("uk")), "Локалізована");
4760 assert_eq!(cap.localized_name(Some("uk_UA")), "Локалізована");
4762 assert_eq!(cap.localized_name(Some("fr-FR")), "Localized");
4764 assert_eq!(cap.localized_name(None), "Localized");
4765 assert_eq!(cap.localized_description(Some("uk")), "Український опис");
4766 assert_eq!(cap.localized_description(Some("de")), "English description");
4767 }
4768
4769 #[test]
4770 fn describe_schema_resolves_config_description_per_locale() {
4771 let cap = LocalizedCapability;
4772 assert_eq!(
4773 cap.describe_schema(Some("uk-UA")).as_deref(),
4774 Some("Керує налаштуваннями.")
4775 );
4776 assert_eq!(
4778 cap.describe_schema(Some("pl")).as_deref(),
4779 Some("Controls things.")
4780 );
4781 assert_eq!(
4782 cap.describe_schema(None).as_deref(),
4783 Some("Controls things.")
4784 );
4785 assert_eq!(NoopCapability.describe_schema(Some("uk")), None);
4787 }
4788}