1use crate::capability_types::is_plugin_capability;
22use crate::command::{
23 CommandDescriptor, CommandExecutionContext, CommandResult, ExecuteCommandRequest,
24};
25use crate::deployment::DeploymentGrade;
26use crate::events::TokenUsage;
27use crate::mcp_server::{ScopedMcpServers, merge_scoped_mcp_servers};
28use crate::message::Message;
29use crate::message_filter::MessageFilterProvider;
30use crate::runtime_agent::RuntimeAgent;
31use crate::tool_types::{ToolCall, ToolDefinition};
32use crate::tools::{Tool, ToolRegistry};
33use crate::traits::SessionFileSystem;
34use crate::typed_id::SessionId;
35use async_trait::async_trait;
36use serde::{Deserialize, Serialize};
37use std::collections::HashMap;
38use std::sync::Arc;
39
40pub struct IntegrationPlugin {
64 pub experimental_only: bool,
66 pub feature_flag: Option<&'static str>,
69 pub factory: fn() -> Box<dyn Capability>,
71}
72
73inventory::collect!(IntegrationPlugin);
74
75pub use crate::capability_types::{
77 AgentCapabilityConfig, CapabilityId, CapabilityStatus, MountAccess, MountDirectoryBuilder,
78 MountEntry, MountPoint, MountSource,
79};
80
81mod a2a_delegation;
86#[cfg(feature = "ui-capabilities")]
87mod a2ui;
88mod agent_handoff;
89mod agent_instructions;
90pub mod attach_skill;
91mod auto_tool_search;
92mod background_execution;
93mod bashkit_shell;
94mod btw;
95mod budgeting;
96mod claude_tool_search;
97pub mod compaction;
98mod current_time;
99mod data_knowledge;
100mod declarative;
101mod error_disclosure;
102mod fake_aws;
103mod fake_crm;
104mod fake_financial;
105mod fake_warehouse;
106mod file_system;
107mod guardrails;
108mod human_intent;
109mod infinity_context;
110mod knowledge_base;
111mod loop_detection;
112mod lua;
113mod lua_code_mode;
114pub mod mcp;
115mod memory;
116mod message_metadata;
117mod model_scout;
118mod monitors;
119mod noop;
120mod openai_tool_search;
121mod openrouter_workspace;
122#[cfg(feature = "ui-capabilities")]
123mod openui;
124mod platform_management;
125mod prompt_caching;
126mod prompt_canary_guardrail;
127mod research;
128mod sample_data;
129mod self_budget;
130mod session;
131mod session_sandbox;
132mod session_schedule;
133mod session_sql_database;
134mod session_storage;
135mod session_tasks;
136mod skills;
137mod skills_scoped;
138mod stateless_todo_list;
139mod subagents;
140mod system_commands;
141mod test_math;
142mod test_weather;
143mod tool_output_distillation;
144mod tool_output_persistence;
145mod tool_search;
146pub mod user_hooks;
147mod web_fetch;
148
149pub use a2a_delegation::{
151 A2A_AGENT_DELEGATION_CAPABILITY_ID, A2aAgentDelegationCapability, SpawnAgentTool,
152};
153#[cfg(feature = "ui-capabilities")]
154pub use a2ui::{A2UI_CAPABILITY_ID, A2UiCapability};
155pub use agent_handoff::{
156 AGENT_HANDOFF_CAPABILITY_ID, AgentHandoffCapability, GetAgentHandoffsTool,
157 MessageAgentHandoffTool, StartAgentHandoffTool,
158};
159pub use agent_instructions::{
160 AGENT_INSTRUCTIONS_CAPABILITY_ID, AGENTS_MD_PATH, AgentInstructionsCapability,
161 AgentInstructionsConfig, DEFAULT_AGENT_INSTRUCTIONS_FILE, MAX_AGENT_INSTRUCTIONS_FILES,
162 MAX_AGENTS_MD_SIZE, format_agents_md_content, format_instruction_file_content,
163};
164pub use attach_skill::{
165 AttachSkillCapability, SKILL_CAPABILITY_PREFIX, SKILLS_DISCOVERY_PATH, SkillContribution,
166 SkillInstructions, SkillMeta, SkillSource, discover_skills_from_entries, is_skill_capability,
167 parse_skill_capability_id, reconstruct_skill_md, skill_capability_id,
168};
169pub use auto_tool_search::{AUTO_TOOL_SEARCH_CAPABILITY_ID, AutoToolSearchCapability};
170pub use background_execution::{BACKGROUND_EXECUTION_CAPABILITY_ID, BackgroundExecutionCapability};
171pub use btw::{BTW_CAPABILITY_ID, BtwCapability};
172pub use budgeting::{BUDGETING_CAPABILITY_ID, BudgetingCapability};
173pub use claude_tool_search::{CLAUDE_TOOL_SEARCH_CAPABILITY_ID, ClaudeToolSearchCapability};
174pub use compaction::{
175 COMPACTION_CAPABILITY_ID, CompactionCapability, CompactionConfig, CompactionStep,
176 CompactionStrategy, CostControlConfig, CostControlMaskingResult, HierarchicalMemoryConfig,
177 MaskingSummaryFormat, MemoryTier, ObservationMaskingConfig, ObservationMaskingResult,
178 SessionCompactionMetrics, SummarizationConfig, aggressive_trim, apply_cost_control_masking,
179 apply_hierarchical_memory, apply_observation_masking, build_model_view_messages,
180 build_summarization_prompt, build_summary_message, classify_memory_tiers, estimate_tokens,
181 estimate_total_tokens, format_messages_for_summarization, should_compact_proactively,
182};
183pub use current_time::{CURRENT_TIME_CAPABILITY_ID, CurrentTimeCapability, GetCurrentTimeTool};
184pub use data_knowledge::{DATA_KNOWLEDGE_CAPABILITY_ID, DataKnowledgeCapability};
185pub use declarative::{
186 DECLARATIVE_CAPABILITY_PREFIX, DeclarativeCapabilityDefinition, DeclarativeCapabilityFile,
187 DeclarativeCapabilitySkill, DeclarativeCapabilitySkillFile, declarative_capability_id,
188 declarative_capability_info, hydrate_declarative_capability_config,
189 hydrate_plugin_capability_config, is_declarative_capability, parse_declarative_capability_id,
190 plugin_capability_info, validate_declarative_capability_definition,
191};
192pub use error_disclosure::{
193 ERROR_DISCLOSURE_CAPABILITY_ID, ErrorDisclosureCapability, resolve_error_disclosure,
194};
195pub use fake_aws::{
196 AwsCreateEc2InstanceTool, AwsCreateIamUserTool, AwsCreateRdsDatabaseTool,
197 AwsCreateS3BucketTool, AwsGetCloudWatchMetricsTool, AwsListEc2InstancesTool,
198 AwsListIamUsersTool, AwsListRdsDatabasesTool, AwsListS3BucketsTool, AwsListSecurityGroupsTool,
199 AwsStopEc2InstanceTool, FAKE_AWS_CAPABILITY_ID, FakeAwsCapability,
200};
201pub use fake_crm::{
202 CrmAddInteractionTool, CrmCreateCustomerTool, CrmCreateTicketTool, CrmGetCustomerTool,
203 CrmListCustomersTool, CrmListTicketsTool, CrmSearchCustomersTool, CrmUpdateTicketTool,
204 FAKE_CRM_CAPABILITY_ID, FakeCrmCapability,
205};
206pub use fake_financial::{
207 FAKE_FINANCIAL_CAPABILITY_ID, FakeFinancialCapability, FinanceCreateBudgetTool,
208 FinanceCreateTransactionTool, FinanceForecastCashFlowTool, FinanceGetBalanceTool,
209 FinanceGetExpenseReportTool, FinanceGetRevenueReportTool, FinanceListBudgetsTool,
210 FinanceListTransactionsTool,
211};
212pub use fake_warehouse::{
213 FAKE_WAREHOUSE_CAPABILITY_ID, FakeWarehouseCapability, WarehouseCreateInvoiceTool,
214 WarehouseCreateOrderTool, WarehouseCreateShipmentTool, WarehouseGetInventoryTool,
215 WarehouseInventoryReportTool, WarehouseListOrdersTool, WarehouseListShipmentsTool,
216 WarehouseProcessReturnTool, WarehouseUpdateInventoryTool, WarehouseUpdateShipmentStatusTool,
217};
218pub use file_system::{
219 DeleteFileTool, EditFileTool, FileSystemCapability, GrepFilesTool, ListDirectoryTool,
220 ReadFileTool, SESSION_FILE_SYSTEM_CAPABILITY_ID, StatFileTool, WriteFileTool,
221};
222pub use guardrails::{GUARDRAILS_CAPABILITY_ID, GuardrailsCapability};
223pub use human_intent::{HUMAN_INTENT_CAPABILITY_ID, HumanIntentCapability};
224pub use infinity_context::{
225 INFINITY_CONTEXT_CAPABILITY_ID, InfinityContextCapability, QueryHistoryTool,
226};
227pub use knowledge_base::{
228 KNOWLEDGE_BASE_CAPABILITY_ID, KnowledgeBaseCapability, KnowledgeBaseConfig,
229 validate_knowledge_base_config,
230};
231pub use loop_detection::{LOOP_DETECTION_CAPABILITY_ID, LoopDetectionCapability};
232pub use lua::{LUA_CAPABILITY_ID, LuaCapability, LuaTool, LuaVfs, is_code_mode_eligible};
233pub use lua_code_mode::{LUA_CODE_MODE_CAPABILITY_ID, LuaCodeModeCapability};
234pub use mcp::{
235 MCP_CAPABILITY_PREFIX, McpCapability, is_mcp_capability, mcp_capability_id,
236 parse_mcp_capability_id,
237};
238pub use memory::{MEMORY_CAPABILITY_ID, MemoryCapability};
239pub use message_metadata::{
240 MESSAGE_METADATA_CAPABILITY_ID, MessageMetadataCapability, MessageMetadataConfig,
241 MessageMetadataField, render_annotation,
242};
243pub use model_scout::{
244 MODEL_SCOUT_CAPABILITY_ID, ModelRanking, ModelScoutCapability, ProbeResult, ProbeTask,
245 RouterUpdateProposal, compute_score, rank_results,
246};
247pub use noop::{NOOP_CAPABILITY_ID, NoopCapability};
248pub use openai_tool_search::{
249 DEFAULT_TOOL_SEARCH_THRESHOLD, OPENAI_TOOL_SEARCH_CAPABILITY_ID, OpenAiToolSearchCapability,
250 model_supports_native_tool_search,
251};
252pub use openrouter_workspace::{
253 OPENROUTER_WORKSPACE_CAPABILITY_ID, OpenRouterKeyInfo, OpenRouterRateLimit,
254 OpenRouterWorkspaceCapability, PolicyCompatibilityReport, WorkspacePolicyDrift,
255 detect_policy_drift,
256};
257#[cfg(feature = "ui-capabilities")]
258pub use openui::{OPENUI_CAPABILITY_ID, OpenUiCapability};
259pub use platform_management::{
260 ManageAgentsTool, ManageHarnessesTool, ManageSessionsTool, PLATFORM_MANAGEMENT_CAPABILITY_ID,
261 PlatformManagementCapability, ReadAgentsTool, ReadCapabilitiesTool, ReadHarnessesTool,
262 ReadSessionsTool, SessionReadMessagesTool, SessionReadResponseTool, SessionSendMessageTool,
263};
264pub use prompt_caching::{PROMPT_CACHING_CAPABILITY_ID, PromptCachingCapability};
265pub use prompt_canary_guardrail::{
266 DEFAULT_REPLACEMENT as PROMPT_CANARY_DEFAULT_REPLACEMENT,
267 PROMPT_CANARY_GUARDRAIL_CAPABILITY_ID, PromptCanaryGuardrailCapability,
268 REASON_CODE_SYSTEM_PROMPT_LEAK,
269};
270pub use research::{RESEARCH_CAPABILITY_ID, ResearchCapability};
271pub use sample_data::{SAMPLE_DATA_CAPABILITY_ID, SampleDataCapability};
272pub use self_budget::{SELF_BUDGET_CAPABILITY_ID, SelfBudgetCapability};
273pub use session::{
274 GetSessionInfoTool, SESSION_CAPABILITY_ID, SessionCapability, WriteSessionTitleTool,
275};
276pub use session_sandbox::{
277 SESSION_SANDBOX_CAPABILITY_ID, SandboxExecTool, SandboxManageTool, SandboxReadFileTool,
278 SandboxStatusTool, SandboxWriteFileTool, SessionSandboxCapability,
279};
280pub use session_schedule::{
281 CancelScheduleTool, CreateScheduleTool, ListSchedulesTool, SESSION_SCHEDULE_CAPABILITY_ID,
282 SessionScheduleCapability,
283};
284pub use session_sql_database::{
285 SESSION_SQL_DATABASE_CAPABILITY_ID, SessionSqlDatabaseCapability, SqlExecuteTool, SqlQueryTool,
286 SqlSchemaTool,
287};
288pub use session_storage::{
289 KvStoreTool, SESSION_STORAGE_CAPABILITY_ID, SecretStoreTool, SessionStorageCapability,
290};
291pub use session_tasks::{SESSION_TASKS_CAPABILITY_ID, SessionTasksCapability};
292pub use skills::{SKILLS_CAPABILITY_ID, SkillsCapability};
293pub use skills_scoped::{
294 ScopedSkillsCapability, SkillDirResolver, SkillScope, SkillsConfig, VfsSkillDirResolver,
295};
296pub use stateless_todo_list::{
297 STATELESS_TODO_LIST_CAPABILITY_ID, StatelessTodoListCapability, WriteTodosTool,
298};
299pub use subagents::{SUBAGENTS_CAPABILITY_ID, SubagentCapability};
300pub use bashkit_shell::{
302 BASHKIT_SHELL_CAPABILITY_ID, BashTool, BashkitShellCapability, SessionFileSystemAdapter,
303};
304pub use system_commands::{SYSTEM_COMMANDS_CAPABILITY_ID, SystemCommandsCapability};
305pub use test_math::{
306 AddTool, DivideTool, MultiplyTool, SubtractTool, TEST_MATH_CAPABILITY_ID, TestMathCapability,
307};
308pub use test_weather::{
309 GetForecastTool, GetWeatherTool, TEST_WEATHER_CAPABILITY_ID, TestWeatherCapability,
310};
311pub use tool_output_distillation::{
312 DistillOutputHook, TOOL_OUTPUT_DISTILLATION_CAPABILITY_ID, ToolOutputDistillationCapability,
313};
314pub use tool_output_persistence::{
315 PersistOutputHook, TOOL_OUTPUT_PERSISTENCE_CAPABILITY_ID, ToolOutputPersistenceCapability,
316};
317pub use tool_search::{
318 TOOL_SEARCH_CAPABILITY_ID, TOOL_SEARCH_TOOL_NAME, ToolSearchCapability, ToolSearchTool,
319};
320pub use user_hooks::{USER_HOOKS_CAPABILITY_ID, UserHooksCapability};
321pub use web_fetch::{
322 BotAuthPublicKey, WEB_FETCH_CAPABILITY_ID, WebFetchCapability, WebFetchTool,
323 derive_bot_auth_public_key,
324};
325
326pub struct SystemPromptContext {
336 pub session_id: SessionId,
338 pub locale: Option<String>,
340 pub file_store: Option<Arc<dyn SessionFileSystem>>,
342 pub model: Option<String>,
348}
349
350impl SystemPromptContext {
351 pub fn without_file_store(session_id: SessionId) -> Self {
353 Self {
354 session_id,
355 locale: None,
356 file_store: None,
357 model: None,
358 }
359 }
360
361 pub fn with_model(mut self, model: impl Into<String>) -> Self {
363 self.model = Some(model.into());
364 self
365 }
366}
367
368#[derive(Debug, Clone)]
420pub struct CapabilityLocalization {
421 pub locale: &'static str,
423 pub name: Option<&'static str>,
425 pub description: Option<&'static str>,
427 pub config_description: Option<&'static str>,
432 pub config_overlay: Option<serde_json::Value>,
438}
439
440impl CapabilityLocalization {
441 pub fn text(locale: &'static str, name: &'static str, description: &'static str) -> Self {
443 Self {
444 locale,
445 name: Some(name),
446 description: Some(description),
447 config_description: None,
448 config_overlay: None,
449 }
450 }
451}
452
453pub fn resolve_localized_field<T>(
457 localizations: &[CapabilityLocalization],
458 locale: Option<&str>,
459 field: impl Fn(&CapabilityLocalization) -> Option<T>,
460) -> Option<T> {
461 let mut candidates: Vec<String> = Vec::new();
462 if let Some(raw) = locale {
463 let normalized = raw.trim().replace('_', "-").to_lowercase();
464 if !normalized.is_empty() {
465 if let Some((language, _)) = normalized.split_once('-') {
466 let language = language.to_string();
467 candidates.push(normalized);
468 candidates.push(language);
469 } else {
470 candidates.push(normalized);
471 }
472 }
473 }
474 candidates.push("en".to_string());
475
476 for candidate in candidates {
477 let hit = localizations
478 .iter()
479 .find(|entry| entry.locale.eq_ignore_ascii_case(&candidate))
480 .and_then(&field);
481 if hit.is_some() {
482 return hit;
483 }
484 }
485 None
486}
487
488#[async_trait]
489pub trait Capability: Send + Sync {
490 fn id(&self) -> &str;
492
493 fn aliases(&self) -> Vec<&'static str> {
502 vec![]
503 }
504
505 fn name(&self) -> &str;
507
508 fn description(&self) -> &str;
510
511 fn localizations(&self) -> Vec<CapabilityLocalization> {
516 vec![]
517 }
518
519 fn localized_name(&self, locale: Option<&str>) -> String {
522 resolve_localized_field(&self.localizations(), locale, |entry| entry.name)
523 .unwrap_or_else(|| self.name())
524 .to_string()
525 }
526
527 fn localized_description(&self, locale: Option<&str>) -> String {
529 resolve_localized_field(&self.localizations(), locale, |entry| entry.description)
530 .unwrap_or_else(|| self.description())
531 .to_string()
532 }
533
534 fn describe_schema(&self, locale: Option<&str>) -> Option<String> {
538 resolve_localized_field(&self.localizations(), locale, |entry| {
539 entry.config_description
540 })
541 .map(str::to_string)
542 }
543
544 fn status(&self) -> CapabilityStatus {
546 CapabilityStatus::Available
547 }
548
549 fn icon(&self) -> Option<&str> {
551 None
552 }
553
554 fn category(&self) -> Option<&str> {
556 None
557 }
558
559 fn is_guardrail(&self) -> bool {
564 false
565 }
566
567 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
578 None
579 }
580
581 fn system_prompt_addition(&self) -> Option<&str> {
601 None
602 }
603
604 async fn system_prompt_contribution(&self, _ctx: &SystemPromptContext) -> Option<String> {
616 self.system_prompt_addition().map(|addition| {
617 format!(
618 "<capability id=\"{}\">\n{}\n</capability>",
619 self.id(),
620 addition
621 )
622 })
623 }
624
625 fn system_prompt_preview(&self) -> Option<String> {
631 self.system_prompt_addition().map(|s| s.to_string())
632 }
633
634 fn tools(&self) -> Vec<Box<dyn Tool>> {
636 vec![]
637 }
638
639 fn tools_with_config(&self, _config: &serde_json::Value) -> Vec<Box<dyn Tool>> {
647 self.tools()
648 }
649
650 async fn system_prompt_contribution_with_config(
657 &self,
658 ctx: &SystemPromptContext,
659 _config: &serde_json::Value,
660 ) -> Option<String> {
661 self.system_prompt_contribution(ctx).await
662 }
663
664 fn tool_definitions(&self) -> Vec<ToolDefinition> {
667 self.tools().iter().map(|t| t.to_definition()).collect()
668 }
669
670 fn mounts(&self) -> Vec<MountPoint> {
678 vec![]
679 }
680
681 fn dependencies(&self) -> Vec<&'static str> {
690 vec![]
691 }
692
693 fn features(&self) -> Vec<&'static str> {
708 vec![]
709 }
710
711 fn config_schema(&self) -> Option<serde_json::Value> {
717 None
718 }
719
720 fn config_ui_schema(&self) -> Option<serde_json::Value> {
725 None
726 }
727
728 fn validate_config(&self, _config: &serde_json::Value) -> Result<(), String> {
734 Ok(())
735 }
736
737 fn mcp_servers(&self) -> ScopedMcpServers {
743 ScopedMcpServers::default()
744 }
745
746 fn mcp_servers_with_config(&self, _config: &serde_json::Value) -> ScopedMcpServers {
748 self.mcp_servers()
749 }
750
751 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
764 None
765 }
766
767 fn model_view_provider(&self) -> Option<Arc<dyn ModelViewProvider>> {
775 None
776 }
777
778 fn pre_tool_use_hooks(&self) -> Vec<Arc<dyn crate::atoms::PreToolUseHook>> {
789 vec![]
790 }
791
792 fn pre_tool_use_hooks_with_config(
797 &self,
798 _config: &serde_json::Value,
799 ) -> Vec<Arc<dyn crate::atoms::PreToolUseHook>> {
800 self.pre_tool_use_hooks()
801 }
802
803 fn post_tool_exec_hooks(&self) -> Vec<Arc<dyn crate::atoms::PostToolExecHook>> {
811 vec![]
812 }
813
814 fn post_tool_exec_hooks_with_config(
819 &self,
820 _config: &serde_json::Value,
821 ) -> Vec<Arc<dyn crate::atoms::PostToolExecHook>> {
822 self.post_tool_exec_hooks()
823 }
824
825 fn tool_definition_hooks(&self) -> Vec<Arc<dyn ToolDefinitionHook>> {
834 vec![]
835 }
836
837 fn tool_definition_hooks_with_config(
842 &self,
843 _config: &serde_json::Value,
844 ) -> Vec<Arc<dyn ToolDefinitionHook>> {
845 self.tool_definition_hooks()
846 }
847
848 fn tool_definition_hooks_with_context(
858 &self,
859 _ctx: &SystemPromptContext,
860 config: &serde_json::Value,
861 ) -> Vec<Arc<dyn ToolDefinitionHook>> {
862 self.tool_definition_hooks_with_config(config)
863 }
864
865 fn tool_call_hooks(&self) -> Vec<Arc<dyn ToolCallHook>> {
873 vec![]
874 }
875
876 fn user_hooks(&self) -> Vec<crate::user_hook_types::UserHookSpec> {
892 vec![]
893 }
894
895 fn user_hooks_with_config(
901 &self,
902 _config: &serde_json::Value,
903 ) -> Vec<crate::user_hook_types::UserHookSpec> {
904 self.user_hooks()
905 }
906
907 fn risk_level(&self) -> RiskLevel {
915 RiskLevel::Low
916 }
917
918 fn commands(&self) -> Vec<CommandDescriptor> {
926 vec![]
927 }
928
929 async fn execute_command(
943 &self,
944 request: &ExecuteCommandRequest,
945 _ctx: &CommandExecutionContext,
946 ) -> crate::error::Result<CommandResult> {
947 Err(crate::error::AgentLoopError::config(format!(
948 "capability {} declared command /{} but does not implement execute_command",
949 self.id(),
950 request.name,
951 )))
952 }
953
954 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
962 vec![]
963 }
964
965 fn contribute_skills(&self) -> Vec<SkillContribution> {
975 vec![]
976 }
977
978 fn output_guardrails(&self) -> Vec<Arc<dyn crate::output_guardrail::OutputGuardrail>> {
989 vec![]
990 }
991}
992
993pub trait ToolDefinitionHook: Send + Sync {
994 fn transform(&self, tools: Vec<ToolDefinition>) -> Vec<ToolDefinition>;
995
996 fn applies_with_native_tool_search(&self) -> bool {
1001 true
1002 }
1003}
1004
1005pub trait ToolCallHook: Send + Sync {
1006 fn narration(
1007 &self,
1008 _tool_def: Option<&ToolDefinition>,
1009 _tool_call: &ToolCall,
1010 _phase: crate::tool_narration::ToolNarrationPhase,
1011 _locale: Option<&str>,
1012 ) -> Option<String> {
1013 None
1014 }
1015
1016 fn transform_for_execution(&self, tool_call: ToolCall) -> ToolCall {
1017 tool_call
1018 }
1019}
1020
1021#[derive(
1025 Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
1026)]
1027#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
1028#[cfg_attr(feature = "openapi", schema(example = "low"))]
1029#[serde(rename_all = "lowercase")]
1030pub enum RiskLevel {
1031 Low,
1033 Medium,
1035 High,
1037}
1038
1039#[derive(Debug, Clone, Serialize, Deserialize)]
1045#[serde(rename_all = "snake_case")]
1046pub enum BlueprintModel {
1047 Fixed(String),
1049 Default(String),
1051 Inherit,
1053}
1054
1055pub struct AgentBlueprint {
1061 pub id: &'static str,
1063 pub name: &'static str,
1065 pub description: &'static str,
1067 pub model: BlueprintModel,
1069 pub system_prompt: &'static str,
1071 pub tools: Vec<Box<dyn Tool>>,
1073 pub max_turns: Option<usize>,
1075 pub config_schema: Option<serde_json::Value>,
1077}
1078
1079impl AgentBlueprint {
1080 pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
1082 self.tools.iter().map(|t| t.to_definition()).collect()
1083 }
1084}
1085
1086impl std::fmt::Debug for AgentBlueprint {
1087 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1088 f.debug_struct("AgentBlueprint")
1089 .field("id", &self.id)
1090 .field("name", &self.name)
1091 .field("model", &self.model)
1092 .field("tool_count", &self.tools.len())
1093 .field("max_turns", &self.max_turns)
1094 .finish()
1095 }
1096}
1097
1098#[derive(Clone)]
1125pub struct CapabilityRegistry {
1126 capabilities: HashMap<String, Arc<dyn Capability>>,
1127 aliases: HashMap<String, String>,
1129}
1130
1131impl CapabilityRegistry {
1132 pub fn new() -> Self {
1134 Self {
1135 capabilities: HashMap::new(),
1136 aliases: HashMap::new(),
1137 }
1138 }
1139
1140 pub fn with_builtins() -> Self {
1145 Self::with_builtins_for_grade(DeploymentGrade::from_env())
1146 }
1147
1148 pub fn with_builtins_for_grade(grade: DeploymentGrade) -> Self {
1153 let mut registry = Self::new();
1154
1155 registry.register(AgentInstructionsCapability);
1157 registry.register(HumanIntentCapability);
1158 registry.register(NoopCapability);
1159 registry.register(CurrentTimeCapability);
1160 registry.register(MessageMetadataCapability);
1161 registry.register(ResearchCapability);
1162 registry.register(ModelScoutCapability);
1163 registry.register(OpenRouterWorkspaceCapability);
1164 registry.register(PlatformManagementCapability);
1165 registry.register(FileSystemCapability);
1166 registry.register(MemoryCapability);
1167 registry.register(SessionStorageCapability);
1168 registry.register(SessionCapability);
1169 registry.register(SessionSqlDatabaseCapability);
1170 registry.register(TestMathCapability);
1171 registry.register(TestWeatherCapability);
1172 registry.register(StatelessTodoListCapability);
1173 registry.register(WebFetchCapability::from_env());
1174 registry.register(BashkitShellCapability);
1175 registry.register(BackgroundExecutionCapability);
1176 registry.register(SessionScheduleCapability);
1177 registry.register(BtwCapability);
1178 registry.register(InfinityContextCapability);
1179 registry.register(budgeting::BudgetingCapability);
1180 registry.register(SelfBudgetCapability);
1181 registry.register(CompactionCapability);
1182 registry.register(ErrorDisclosureCapability);
1183
1184 registry.register(OpenAiToolSearchCapability::new());
1186 registry.register(ClaudeToolSearchCapability::new());
1188 registry.register(ToolSearchCapability::new());
1190 registry.register(AutoToolSearchCapability::new());
1192 registry.register(PromptCachingCapability::new());
1193
1194 registry.register(SkillsCapability);
1196
1197 registry.register(SubagentCapability);
1199
1200 registry.register(SessionTasksCapability);
1202
1203 if crate::FeatureFlags::from_env(&grade).agent_delegation {
1207 registry.register(AgentHandoffCapability);
1208 registry.register(A2aAgentDelegationCapability);
1209 }
1210
1211 registry.register(SystemCommandsCapability);
1213
1214 registry.register(tool_output_persistence::ToolOutputPersistenceCapability);
1216 registry.register(tool_output_distillation::ToolOutputDistillationCapability);
1217
1218 registry.register(user_hooks::UserHooksCapability);
1221
1222 registry.register(LoopDetectionCapability);
1224
1225 registry.register(PromptCanaryGuardrailCapability);
1228
1229 registry.register(GuardrailsCapability);
1232
1233 #[cfg(feature = "ui-capabilities")]
1235 {
1236 registry.register(OpenUiCapability);
1237 registry.register(A2UiCapability);
1238 }
1239
1240 registry.register(SampleDataCapability);
1242
1243 registry.register(DataKnowledgeCapability);
1245
1246 registry.register(KnowledgeBaseCapability);
1248
1249 registry.register(FakeWarehouseCapability);
1251 registry.register(FakeAwsCapability);
1252 registry.register(FakeCrmCapability);
1253 registry.register(FakeFinancialCapability);
1254
1255 let internal_flags = crate::InternalFeatureFlags::from_env();
1257 if internal_flags.session_sandbox {
1258 registry.register(SessionSandboxCapability);
1259 }
1260
1261 if internal_flags.lua {
1265 registry.register(LuaCapability);
1266 registry.register(LuaCodeModeCapability);
1269 }
1270 for plugin in inventory::iter::<IntegrationPlugin>() {
1271 if (!plugin.experimental_only || grade.experimental_features_enabled())
1272 && plugin
1273 .feature_flag
1274 .is_none_or(|f| internal_flags.is_enabled(f))
1275 {
1276 registry.register_boxed((plugin.factory)());
1277 }
1278 }
1279
1280 registry
1281 }
1282
1283 pub fn register(&mut self, capability: impl Capability + 'static) {
1285 self.register_arc(Arc::new(capability));
1286 }
1287
1288 pub fn register_boxed(&mut self, capability: Box<dyn Capability>) {
1290 self.register_arc(Arc::from(capability));
1291 }
1292
1293 pub fn register_arc(&mut self, capability: Arc<dyn Capability>) {
1295 let canonical = capability.id().to_string();
1296 for alias in capability.aliases() {
1297 self.aliases.insert(alias.to_string(), canonical.clone());
1298 }
1299 self.capabilities.insert(canonical, capability);
1300 }
1301
1302 pub fn get(&self, id: &str) -> Option<&Arc<dyn Capability>> {
1304 self.capabilities
1305 .get(id)
1306 .or_else(|| self.aliases.get(id).and_then(|c| self.capabilities.get(c)))
1307 }
1308
1309 pub fn canonical_id<'a>(&'a self, id: &'a str) -> Option<&'a str> {
1314 if self.capabilities.contains_key(id) {
1315 Some(id)
1316 } else {
1317 self.aliases
1318 .get(id)
1319 .filter(|c| self.capabilities.contains_key(*c))
1320 .map(String::as_str)
1321 }
1322 }
1323
1324 pub fn unregister(&mut self, id: &str) -> Option<Arc<dyn Capability>> {
1326 let canonical = self.canonical_id(id)?.to_string();
1327 let removed = self.capabilities.remove(&canonical);
1328 self.aliases.retain(|_, target| *target != canonical);
1329 removed
1330 }
1331
1332 pub fn has(&self, id: &str) -> bool {
1334 self.get(id).is_some()
1335 }
1336
1337 pub fn list(&self) -> Vec<&Arc<dyn Capability>> {
1339 self.capabilities.values().collect()
1340 }
1341
1342 pub fn len(&self) -> usize {
1344 self.capabilities.len()
1345 }
1346
1347 pub fn is_empty(&self) -> bool {
1349 self.capabilities.is_empty()
1350 }
1351
1352 pub fn builder() -> CapabilityRegistryBuilder {
1354 CapabilityRegistryBuilder::new()
1355 }
1356
1357 pub fn blueprint(&self, id: &str) -> Option<AgentBlueprint> {
1361 for cap in self.capabilities.values() {
1362 for bp in cap.agent_blueprints() {
1363 if bp.id == id {
1364 return Some(bp);
1365 }
1366 }
1367 }
1368 None
1369 }
1370
1371 pub fn blueprint_with_capability(&self, id: &str) -> Option<(String, AgentBlueprint)> {
1375 for (capability_id, cap) in &self.capabilities {
1376 for bp in cap.agent_blueprints() {
1377 if bp.id == id {
1378 return Some((capability_id.clone(), bp));
1379 }
1380 }
1381 }
1382 None
1383 }
1384
1385 pub fn all_blueprints(&self) -> Vec<AgentBlueprint> {
1387 self.capabilities
1388 .values()
1389 .flat_map(|cap| cap.agent_blueprints())
1390 .collect()
1391 }
1392}
1393
1394impl Default for CapabilityRegistry {
1395 fn default() -> Self {
1396 Self::with_builtins()
1397 }
1398}
1399
1400impl std::fmt::Debug for CapabilityRegistry {
1401 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1402 let ids: Vec<_> = self.capabilities.keys().collect();
1403 f.debug_struct("CapabilityRegistry")
1404 .field("capabilities", &ids)
1405 .finish()
1406 }
1407}
1408
1409pub struct CapabilityRegistryBuilder {
1411 registry: CapabilityRegistry,
1412}
1413
1414impl CapabilityRegistryBuilder {
1415 pub fn new() -> Self {
1417 Self {
1418 registry: CapabilityRegistry::new(),
1419 }
1420 }
1421
1422 pub fn with_builtins() -> Self {
1424 Self {
1425 registry: CapabilityRegistry::with_builtins(),
1426 }
1427 }
1428
1429 pub fn capability(mut self, capability: impl Capability + 'static) -> Self {
1431 self.registry.register(capability);
1432 self
1433 }
1434
1435 pub fn build(self) -> CapabilityRegistry {
1437 self.registry
1438 }
1439}
1440
1441impl Default for CapabilityRegistryBuilder {
1442 fn default() -> Self {
1443 Self::new()
1444 }
1445}
1446
1447pub struct ModelViewContext<'a> {
1453 pub session_id: SessionId,
1454 pub prior_usage: Option<&'a TokenUsage>,
1455}
1456
1457pub trait ModelViewProvider: Send + Sync {
1463 fn apply_model_view(
1464 &self,
1465 messages: Vec<Message>,
1466 config: &serde_json::Value,
1467 context: &ModelViewContext<'_>,
1468 ) -> Vec<Message>;
1469
1470 fn priority(&self) -> i32 {
1471 0
1472 }
1473}
1474
1475pub struct CollectedCapabilities {
1480 pub system_prompt_parts: Vec<String>,
1482 pub system_prompt_attributions: Vec<SystemPromptAttribution>,
1484 pub tools: Vec<Box<dyn Tool>>,
1486 pub tool_definitions: Vec<ToolDefinition>,
1488 pub mounts: Vec<MountPoint>,
1490 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1492 pub applied_ids: Vec<String>,
1494 pub tool_search: Option<crate::llm_driver_registry::ToolSearchConfig>,
1496 pub prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig>,
1498 pub tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>>,
1500 pub tool_call_hooks: Vec<Arc<dyn ToolCallHook>>,
1502 pub mcp_servers: ScopedMcpServers,
1504 }
1510
1511#[derive(Debug, Clone, PartialEq, Eq)]
1512pub struct SystemPromptAttribution {
1513 pub capability_id: String,
1514 pub content: String,
1515}
1516
1517impl CollectedCapabilities {
1518 pub fn system_prompt_prefix(&self) -> Option<String> {
1521 if self.system_prompt_parts.is_empty() {
1522 None
1523 } else {
1524 Some(self.system_prompt_parts.join("\n\n"))
1525 }
1526 }
1527
1528 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1532 for (provider, config) in &self.message_filter_providers {
1534 provider.apply_filters(query, config);
1535 }
1536 }
1537
1538 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1541 for (provider, config) in &self.message_filter_providers {
1542 provider.post_load(messages, config);
1543 }
1544 }
1545
1546 pub fn has_message_filters(&self) -> bool {
1548 !self.message_filter_providers.is_empty()
1549 }
1550}
1551
1552pub struct CollectedMessageFilters {
1559 pub message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)>,
1561}
1562
1563pub struct CollectedModelViewProviders {
1565 pub model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)>,
1567}
1568
1569impl CollectedMessageFilters {
1575 pub fn apply_message_filters(&self, query: &mut crate::message_filter::MessageQuery) {
1577 for (provider, config) in &self.message_filter_providers {
1578 provider.apply_filters(query, config);
1579 }
1580 }
1581
1582 pub fn apply_post_load_filters(&self, messages: &mut Vec<crate::message::Message>) {
1584 for (provider, config) in &self.message_filter_providers {
1585 provider.post_load(messages, config);
1586 }
1587 }
1588}
1589
1590impl CollectedModelViewProviders {
1591 pub fn apply_model_view(
1593 &self,
1594 mut messages: Vec<Message>,
1595 context: &ModelViewContext<'_>,
1596 ) -> Vec<Message> {
1597 for (provider, config) in &self.model_view_providers {
1598 messages = provider.apply_model_view(messages, config, context);
1599 }
1600 messages
1601 }
1602}
1603
1604pub fn collect_message_filters_only(
1610 capability_configs: &[AgentCapabilityConfig],
1611 registry: &CapabilityRegistry,
1612) -> CollectedMessageFilters {
1613 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1614 Vec::new();
1615
1616 for cap_config in capability_configs {
1617 let cap_id = cap_config.capability_ref.as_str();
1618 if let Some(capability) = registry.get(cap_id) {
1619 if capability.status() != CapabilityStatus::Available {
1620 continue;
1621 }
1622 let effective: &dyn Capability = capability
1625 .resolve_for_model(None)
1626 .unwrap_or_else(|| capability.as_ref());
1627 if let Some(provider) = effective.message_filter_provider() {
1628 message_filter_providers.push((provider, cap_config.config.clone()));
1629 }
1630 }
1631 }
1632
1633 message_filter_providers.sort_by_key(|(p, _)| p.priority());
1634
1635 CollectedMessageFilters {
1636 message_filter_providers,
1637 }
1638}
1639
1640pub fn collect_model_view_providers(
1647 capability_configs: &[AgentCapabilityConfig],
1648 registry: &CapabilityRegistry,
1649 model: Option<&str>,
1650) -> CollectedModelViewProviders {
1651 let mut model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)> = Vec::new();
1652
1653 for cap_config in capability_configs {
1654 let cap_id = cap_config.capability_ref.as_str();
1655 if let Some(capability) = registry.get(cap_id) {
1656 if capability.status() != CapabilityStatus::Available {
1657 continue;
1658 }
1659 let effective: &dyn Capability = capability
1660 .resolve_for_model(model)
1661 .unwrap_or_else(|| capability.as_ref());
1662 if let Some(provider) = effective.model_view_provider() {
1663 model_view_providers.push((provider, cap_config.config.clone()));
1664 }
1665 }
1666 }
1667
1668 model_view_providers.sort_by_key(|(p, _)| p.priority());
1669
1670 CollectedModelViewProviders {
1671 model_view_providers,
1672 }
1673}
1674
1675pub fn collect_capability_mcp_servers(
1676 capability_configs: &[AgentCapabilityConfig],
1677 registry: &CapabilityRegistry,
1678) -> ScopedMcpServers {
1679 let mut servers = ScopedMcpServers::default();
1680
1681 for cap_config in capability_configs {
1682 let cap_id = cap_config.capability_ref.as_str();
1683 if is_declarative_capability(cap_id) || is_plugin_capability(cap_id) {
1686 if let Ok(definition) =
1687 serde_json::from_value::<DeclarativeCapabilityDefinition>(cap_config.config.clone())
1688 {
1689 if definition.status != CapabilityStatus::Available {
1690 continue;
1691 }
1692 if let Some(contributed) = definition.mcp_servers {
1693 servers = merge_scoped_mcp_servers(&servers, &contributed);
1694 }
1695 }
1696 continue;
1697 }
1698 if let Some(capability) = registry.get(cap_id) {
1699 if capability.status() != CapabilityStatus::Available {
1700 continue;
1701 }
1702 servers = merge_scoped_mcp_servers(
1703 &servers,
1704 &capability.mcp_servers_with_config(&cap_config.config),
1705 );
1706 }
1707 }
1708
1709 servers
1710}
1711
1712pub const MAX_RESOLVED_CAPABILITIES: usize = 100;
1719
1720#[derive(Debug, Clone, PartialEq, Eq)]
1722pub enum DependencyError {
1723 CircularDependency {
1725 capability_id: String,
1727 chain: Vec<String>,
1729 },
1730 TooManyCapabilities {
1732 count: usize,
1734 max: usize,
1736 },
1737}
1738
1739impl std::fmt::Display for DependencyError {
1740 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1741 match self {
1742 DependencyError::CircularDependency {
1743 capability_id,
1744 chain,
1745 } => {
1746 write!(
1747 f,
1748 "Circular dependency detected: {} depends on itself via chain: {} -> {}",
1749 capability_id,
1750 chain.join(" -> "),
1751 capability_id
1752 )
1753 }
1754 DependencyError::TooManyCapabilities { count, max } => {
1755 write!(
1756 f,
1757 "Too many capabilities after resolution: {} (max: {})",
1758 count, max
1759 )
1760 }
1761 }
1762 }
1763}
1764
1765impl std::error::Error for DependencyError {}
1766
1767#[derive(Debug, Clone)]
1769pub struct ResolvedCapabilities {
1770 pub resolved_ids: Vec<String>,
1773 pub added_as_dependencies: Vec<String>,
1775 pub user_selected: Vec<String>,
1777}
1778
1779pub fn resolve_dependencies(
1799 selected_ids: &[String],
1800 registry: &CapabilityRegistry,
1801) -> Result<ResolvedCapabilities, DependencyError> {
1802 use std::collections::HashSet;
1803
1804 let user_selected: HashSet<String> = selected_ids
1806 .iter()
1807 .map(|id| registry.canonical_id(id).unwrap_or(id).to_string())
1808 .collect();
1809 let mut resolved: Vec<String> = Vec::new();
1810 let mut resolved_set: HashSet<String> = HashSet::new();
1811 let mut added_as_dependencies: Vec<String> = Vec::new();
1812
1813 for cap_id in selected_ids {
1815 resolve_single_capability(
1816 cap_id,
1817 registry,
1818 &mut resolved,
1819 &mut resolved_set,
1820 &mut added_as_dependencies,
1821 &user_selected,
1822 &mut Vec::new(), )?;
1824 }
1825
1826 if resolved.len() > MAX_RESOLVED_CAPABILITIES {
1828 return Err(DependencyError::TooManyCapabilities {
1829 count: resolved.len(),
1830 max: MAX_RESOLVED_CAPABILITIES,
1831 });
1832 }
1833
1834 Ok(ResolvedCapabilities {
1835 resolved_ids: resolved,
1836 added_as_dependencies,
1837 user_selected: selected_ids.to_vec(),
1838 })
1839}
1840
1841pub fn resolve_capability_configs(
1846 selected_configs: &[AgentCapabilityConfig],
1847 registry: &CapabilityRegistry,
1848) -> Result<Vec<AgentCapabilityConfig>, DependencyError> {
1849 let mut selected_ids: Vec<String> = Vec::new();
1850 for config in selected_configs {
1851 if (is_declarative_capability(config.capability_id())
1854 || is_plugin_capability(config.capability_id()))
1855 && let Ok(definition) =
1856 serde_json::from_value::<DeclarativeCapabilityDefinition>(config.config.clone())
1857 {
1858 selected_ids.extend(definition.dependencies);
1859 }
1860 selected_ids.push(config.capability_id().to_string());
1861 }
1862 let resolved = resolve_dependencies(&selected_ids, registry)?;
1863
1864 let explicit_configs: std::collections::HashMap<String, serde_json::Value> = selected_configs
1867 .iter()
1868 .map(|config| {
1869 let id = config.capability_id();
1870 let id = registry.canonical_id(id).unwrap_or(id);
1871 (id.to_string(), config.config.clone())
1872 })
1873 .collect();
1874
1875 Ok(resolved
1876 .resolved_ids
1877 .into_iter()
1878 .map(|capability_id| {
1879 explicit_configs
1880 .get(&capability_id)
1881 .cloned()
1882 .map(|config| AgentCapabilityConfig::with_config(capability_id.clone(), config))
1883 .unwrap_or_else(|| AgentCapabilityConfig::new(capability_id))
1884 })
1885 .collect())
1886}
1887
1888fn resolve_single_capability(
1890 cap_id: &str,
1891 registry: &CapabilityRegistry,
1892 resolved: &mut Vec<String>,
1893 resolved_set: &mut std::collections::HashSet<String>,
1894 added_as_dependencies: &mut Vec<String>,
1895 user_selected: &std::collections::HashSet<String>,
1896 visiting: &mut Vec<String>,
1897) -> Result<(), DependencyError> {
1898 let cap_id = registry.canonical_id(cap_id).unwrap_or(cap_id);
1902
1903 if resolved_set.contains(cap_id) {
1905 return Ok(());
1906 }
1907
1908 if visiting.contains(&cap_id.to_string()) {
1910 return Err(DependencyError::CircularDependency {
1911 capability_id: cap_id.to_string(),
1912 chain: visiting.clone(),
1913 });
1914 }
1915
1916 let capability = match registry.get(cap_id) {
1918 Some(cap) => cap,
1919 None => {
1920 if (is_declarative_capability(cap_id) || is_plugin_capability(cap_id))
1924 && !resolved_set.contains(cap_id)
1925 {
1926 resolved.push(cap_id.to_string());
1927 resolved_set.insert(cap_id.to_string());
1928 if !user_selected.contains(cap_id) {
1929 added_as_dependencies.push(cap_id.to_string());
1930 }
1931 }
1932 return Ok(());
1933 }
1934 };
1935
1936 visiting.push(cap_id.to_string());
1938
1939 for dep_id in capability.dependencies() {
1941 resolve_single_capability(
1942 dep_id,
1943 registry,
1944 resolved,
1945 resolved_set,
1946 added_as_dependencies,
1947 user_selected,
1948 visiting,
1949 )?;
1950 }
1951
1952 visiting.pop();
1954
1955 if !resolved_set.contains(cap_id) {
1957 resolved.push(cap_id.to_string());
1958 resolved_set.insert(cap_id.to_string());
1959
1960 if !user_selected.contains(cap_id) {
1962 added_as_dependencies.push(cap_id.to_string());
1963 }
1964 }
1965
1966 Ok(())
1967}
1968
1969pub fn compute_features(capability_ids: &[String], registry: &CapabilityRegistry) -> Vec<String> {
1974 use std::collections::HashSet;
1975
1976 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
1977 Ok(resolved) => resolved.resolved_ids,
1978 Err(_) => capability_ids.to_vec(),
1979 };
1980
1981 let mut seen = HashSet::new();
1982 let mut features = Vec::new();
1983 for cap_id in &resolved_ids {
1984 if let Some(cap) = registry.get(cap_id) {
1985 for feature in cap.features() {
1986 if seen.insert(feature) {
1987 features.push(feature.to_string());
1988 }
1989 }
1990 }
1991 }
1992 features
1993}
1994
1995pub fn get_dependencies(cap_id: &str, registry: &CapabilityRegistry) -> Vec<String> {
1998 registry
1999 .get(cap_id)
2000 .map(|cap| cap.dependencies().iter().map(|s| s.to_string()).collect())
2001 .unwrap_or_default()
2002}
2003
2004pub async fn collect_capabilities(
2020 capability_ids: &[String],
2021 registry: &CapabilityRegistry,
2022 ctx: &SystemPromptContext,
2023) -> CollectedCapabilities {
2024 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
2027 Ok(resolved) => resolved.resolved_ids,
2028 Err(e) => {
2029 tracing::warn!("Failed to resolve capability dependencies: {}", e);
2030 capability_ids.to_vec()
2031 }
2032 };
2033
2034 let configs: Vec<AgentCapabilityConfig> = resolved_ids
2036 .iter()
2037 .map(|id| AgentCapabilityConfig {
2038 capability_ref: CapabilityId::new(id),
2039 config: serde_json::Value::Object(serde_json::Map::new()),
2040 })
2041 .collect();
2042
2043 collect_capabilities_with_configs(&configs, registry, ctx).await
2044}
2045
2046pub async fn collect_capabilities_with_configs(
2057 capability_configs: &[AgentCapabilityConfig],
2058 registry: &CapabilityRegistry,
2059 ctx: &SystemPromptContext,
2060) -> CollectedCapabilities {
2061 let mut system_prompt_parts: Vec<String> = Vec::new();
2062 let mut system_prompt_attributions: Vec<SystemPromptAttribution> = Vec::new();
2063 let mut tools: Vec<Box<dyn Tool>> = Vec::new();
2064 let mut tool_definitions: Vec<ToolDefinition> = Vec::new();
2065 let mut mounts: Vec<MountPoint> = Vec::new();
2066 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
2067 Vec::new();
2068 let mut applied_ids: Vec<String> = Vec::new();
2069 let mut tool_search: Option<crate::llm_driver_registry::ToolSearchConfig> = None;
2070 let mut prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig> = None;
2071 let mut tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>> = Vec::new();
2072 let mut tool_call_hooks: Vec<Arc<dyn ToolCallHook>> = Vec::new();
2073 let mut mcp_servers = ScopedMcpServers::default();
2074
2075 for cap_config in capability_configs {
2076 let cap_id = cap_config.capability_ref.as_str();
2077 if is_declarative_capability(cap_id) || is_plugin_capability(cap_id) {
2082 match serde_json::from_value::<DeclarativeCapabilityDefinition>(
2083 cap_config.config.clone(),
2084 ) {
2085 Ok(definition) => {
2086 if definition.status != CapabilityStatus::Available {
2087 continue;
2088 }
2089
2090 if let Some(prompt) = definition.system_prompt.as_deref() {
2091 let contribution =
2092 format!("<capability id=\"{}\">\n{}\n</capability>", cap_id, prompt);
2093 system_prompt_attributions.push(SystemPromptAttribution {
2094 capability_id: cap_id.to_string(),
2095 content: contribution.clone(),
2096 });
2097 system_prompt_parts.push(contribution);
2098 }
2099
2100 mounts.extend(definition.mounts(cap_id));
2101 if let Some(ref servers) = definition.mcp_servers {
2102 mcp_servers = merge_scoped_mcp_servers(&mcp_servers, servers);
2103 }
2104 for skill in definition.skill_contributions() {
2105 mounts.push(skill.to_mount(cap_id));
2106 }
2107
2108 applied_ids.push(cap_id.to_string());
2109 }
2110 Err(error) => {
2111 tracing::warn!(
2112 capability_id = %cap_id,
2113 error = %error,
2114 "Skipping invalid declarative/plugin capability config"
2115 );
2116 }
2117 }
2118 continue;
2119 }
2120 if let Some(capability) = registry.get(cap_id) {
2121 if capability.status() != CapabilityStatus::Available {
2123 continue;
2124 }
2125
2126 let effective: &dyn Capability =
2138 match capability.resolve_for_model(ctx.model.as_deref()) {
2139 Some(inner) => inner,
2140 None => capability.as_ref(),
2141 };
2142 let effective_id = effective.id();
2143
2144 if let Some(contribution) = effective
2146 .system_prompt_contribution_with_config(ctx, &cap_config.config)
2147 .await
2148 {
2149 system_prompt_attributions.push(SystemPromptAttribution {
2150 capability_id: cap_id.to_string(),
2151 content: contribution.clone(),
2152 });
2153 system_prompt_parts.push(contribution);
2154 }
2155
2156 tools.extend(effective.tools_with_config(&cap_config.config));
2158 tool_definition_hooks
2159 .extend(effective.tool_definition_hooks_with_context(ctx, &cap_config.config));
2160 tool_call_hooks.extend(effective.tool_call_hooks());
2161 let cap_category = effective.category();
2166 for def in effective.tool_definitions() {
2167 let def = match (def.category(), cap_category) {
2168 (None, Some(cat)) => def.with_category(cat),
2169 _ => def,
2170 }
2171 .with_capability_attribution(cap_id, Some(capability.name()));
2172 tool_definitions.push(def);
2173 }
2174
2175 if effective_id == OPENAI_TOOL_SEARCH_CAPABILITY_ID
2183 || effective_id == CLAUDE_TOOL_SEARCH_CAPABILITY_ID
2184 {
2185 let threshold = cap_config
2187 .config
2188 .get("threshold")
2189 .and_then(|v| v.as_u64())
2190 .map(|v| v as usize)
2191 .unwrap_or(DEFAULT_TOOL_SEARCH_THRESHOLD);
2192 tool_search = Some(crate::llm_driver_registry::ToolSearchConfig {
2193 enabled: true,
2194 threshold,
2195 });
2196 }
2197
2198 if cap_id == PROMPT_CACHING_CAPABILITY_ID {
2199 let strategy = cap_config
2200 .config
2201 .get("strategy")
2202 .and_then(|v| v.as_str())
2203 .map(|value| match value {
2204 "auto" => crate::llm_driver_registry::PromptCacheStrategy::Auto,
2205 _ => crate::llm_driver_registry::PromptCacheStrategy::Auto,
2206 })
2207 .unwrap_or(crate::llm_driver_registry::PromptCacheStrategy::Auto);
2208 let gemini_cached_content = cap_config
2209 .config
2210 .get("gemini_cached_content")
2211 .and_then(|v| v.as_str())
2212 .map(str::to_string);
2213 prompt_cache = Some(crate::llm_driver_registry::PromptCacheConfig {
2214 enabled: true,
2215 strategy,
2216 gemini_cached_content,
2217 });
2218 }
2219
2220 mounts.extend(effective.mounts());
2222
2223 mcp_servers = merge_scoped_mcp_servers(
2224 &mcp_servers,
2225 &effective.mcp_servers_with_config(&cap_config.config),
2226 );
2227
2228 for skill in effective.contribute_skills() {
2232 mounts.push(skill.to_mount(cap_id));
2233 }
2234
2235 if let Some(provider) = effective.message_filter_provider() {
2237 message_filter_providers.push((provider, cap_config.config.clone()));
2238 }
2239
2240 applied_ids.push(cap_id.to_string());
2241 }
2242 }
2243
2244 if !applied_ids
2256 .iter()
2257 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID)
2258 && tool_definitions
2259 .iter()
2260 .any(|def| def.hints().supports_background == Some(true))
2261 && let Some(bg_cap) = registry.get(BACKGROUND_EXECUTION_CAPABILITY_ID)
2262 && bg_cap.status() == CapabilityStatus::Available
2263 {
2264 tools.extend(bg_cap.tools());
2265 let cap_category = bg_cap.category();
2266 for def in bg_cap.tool_definitions() {
2267 let def = match (def.category(), cap_category) {
2268 (None, Some(cat)) => def.with_category(cat),
2269 _ => def,
2270 }
2271 .with_capability_attribution(BACKGROUND_EXECUTION_CAPABILITY_ID, Some(bg_cap.name()));
2272 tool_definitions.push(def);
2273 }
2274 applied_ids.push(BACKGROUND_EXECUTION_CAPABILITY_ID.to_string());
2275 }
2276
2277 message_filter_providers.sort_by_key(|(p, _)| p.priority());
2279
2280 CollectedCapabilities {
2281 system_prompt_parts,
2282 system_prompt_attributions,
2283 tools,
2284 tool_definitions,
2285 mounts,
2286 message_filter_providers,
2287 applied_ids,
2288 tool_search,
2289 prompt_cache,
2290 tool_definition_hooks,
2291 tool_call_hooks,
2292 mcp_servers,
2293 }
2294}
2295
2296pub struct AppliedCapabilities {
2302 pub runtime_agent: RuntimeAgent,
2304 pub tool_registry: ToolRegistry,
2306 pub applied_ids: Vec<String>,
2308}
2309
2310pub async fn apply_capabilities(
2347 base_runtime_agent: RuntimeAgent,
2348 capability_ids: &[String],
2349 registry: &CapabilityRegistry,
2350 ctx: &SystemPromptContext,
2351) -> AppliedCapabilities {
2352 let collected = collect_capabilities(capability_ids, registry, ctx).await;
2353
2354 let final_system_prompt = match collected.system_prompt_prefix() {
2356 Some(prefix) => format!(
2357 "{}\n\n<system-prompt>\n{}\n</system-prompt>",
2358 prefix, base_runtime_agent.system_prompt
2359 ),
2360 None => base_runtime_agent.system_prompt,
2361 };
2362
2363 let mut tool_registry = ToolRegistry::new();
2365 for tool in collected.tools {
2366 tool_registry.register_boxed(tool);
2367 }
2368
2369 let mut tools = collected.tool_definitions;
2371 for hook in &collected.tool_definition_hooks {
2372 tools = hook.transform(tools);
2373 }
2374
2375 let runtime_agent = RuntimeAgent {
2376 system_prompt: final_system_prompt,
2377 model: base_runtime_agent.model,
2378 tools,
2379 max_iterations: base_runtime_agent.max_iterations,
2380 temperature: base_runtime_agent.temperature,
2381 max_tokens: base_runtime_agent.max_tokens,
2382 tool_search: collected.tool_search,
2383 prompt_cache: collected.prompt_cache,
2384 network_access: base_runtime_agent.network_access,
2385 };
2386
2387 AppliedCapabilities {
2388 runtime_agent,
2389 tool_registry,
2390 applied_ids: collected.applied_ids,
2391 }
2392}
2393
2394#[cfg(test)]
2399mod tests {
2400 use super::*;
2401 use crate::typed_id::SessionId;
2402 use std::collections::BTreeSet;
2403 use uuid::Uuid;
2404
2405 static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2407
2408 fn lock_env() -> std::sync::MutexGuard<'static, ()> {
2409 ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner())
2410 }
2411
2412 fn test_ctx() -> SystemPromptContext {
2414 SystemPromptContext::without_file_store(SessionId::new())
2415 }
2416
2417 fn expected_core_builtin_ids() -> BTreeSet<&'static str> {
2419 let mut ids = [
2420 "agent_instructions",
2421 "human_intent",
2422 "budgeting",
2423 "self_budget",
2424 "noop",
2425 "current_time",
2426 "research",
2427 "platform_management",
2428 "session_file_system",
2429 "session_storage",
2430 "session",
2431 "session_sql_database",
2432 "test_math",
2433 "test_weather",
2434 "stateless_todo_list",
2435 "web_fetch",
2436 "bashkit_shell",
2437 "background_execution",
2438 "session_schedule",
2439 "btw",
2440 "infinity_context",
2441 "compaction",
2442 "memory",
2443 "message_metadata",
2444 "openai_tool_search",
2445 "claude_tool_search",
2446 "tool_search",
2447 "auto_tool_search",
2448 "prompt_caching",
2449 "session_tasks",
2450 "skills",
2451 "subagents",
2452 "system_commands",
2453 "sample_data",
2454 "data_knowledge",
2455 "knowledge_base",
2456 "tool_output_persistence",
2457 "tool_output_distillation",
2458 "fake_warehouse",
2459 "fake_aws",
2460 "fake_crm",
2461 "fake_financial",
2462 "loop_detection",
2463 "error_disclosure",
2464 "prompt_canary_guardrail",
2465 "guardrails",
2466 "user_hooks",
2467 "model_scout",
2468 "openrouter_workspace",
2469 ]
2470 .into_iter()
2471 .collect::<BTreeSet<_>>();
2472 if cfg!(feature = "ui-capabilities") {
2473 ids.insert("openui");
2474 ids.insert("a2ui");
2475 }
2476 ids
2477 }
2478
2479 fn expected_dev_builtin_ids() -> BTreeSet<&'static str> {
2481 let mut ids = expected_core_builtin_ids();
2482 ids.insert("agent_handoff");
2483 ids.insert("a2a_agent_delegation");
2484 ids
2485 }
2486
2487 fn registry_ids(registry: &CapabilityRegistry) -> BTreeSet<&str> {
2488 registry.capabilities.keys().map(String::as_str).collect()
2489 }
2490
2491 #[test]
2501 fn test_capability_registry_with_builtins_dev() {
2502 let _lock = lock_env();
2504 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2505 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
2506 assert_eq!(registry_ids(®istry), expected_dev_builtin_ids());
2507 assert!(registry.has("agent_handoff"));
2508 assert!(registry.has("a2a_agent_delegation"));
2509 }
2510
2511 #[test]
2512 fn test_capability_registry_with_builtins_prod() {
2513 let _lock = lock_env();
2515 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2516 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2517 assert_eq!(registry_ids(®istry), expected_core_builtin_ids());
2518 assert!(!registry.has("docker_container"));
2520 assert!(!registry.has("agent_handoff"));
2521 assert!(!registry.has("a2a_agent_delegation"));
2522 }
2523
2524 #[test]
2525 fn test_agent_delegation_enabled_by_env_in_prod() {
2526 let _lock = lock_env();
2528 unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "true") };
2529 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2530 assert!(registry.has("agent_handoff"));
2531 assert!(registry.has("a2a_agent_delegation"));
2532 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2533 }
2534
2535 #[test]
2536 fn test_agent_delegation_disabled_by_env_in_dev() {
2537 let _lock = lock_env();
2539 unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "false") };
2540 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
2541 assert!(!registry.has("agent_handoff"));
2542 assert!(!registry.has("a2a_agent_delegation"));
2543 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2544 }
2545
2546 #[test]
2547 fn test_capability_registry_get() {
2548 let registry = CapabilityRegistry::with_builtins();
2549
2550 let noop = registry.get("noop").unwrap();
2551 assert_eq!(noop.id(), "noop");
2552 assert_eq!(noop.name(), "No-Op");
2553 assert_eq!(noop.status(), CapabilityStatus::Available);
2554 }
2555
2556 #[test]
2557 fn test_capability_registry_blueprint_with_capability() {
2558 struct BlueprintProviderCapability;
2559
2560 impl Capability for BlueprintProviderCapability {
2561 fn id(&self) -> &str {
2562 "blueprint_provider"
2563 }
2564 fn name(&self) -> &str {
2565 "Blueprint Provider"
2566 }
2567 fn description(&self) -> &str {
2568 "Capability that provides a blueprint for tests"
2569 }
2570 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
2571 vec![AgentBlueprint {
2572 id: "test_blueprint",
2573 name: "Test Blueprint",
2574 description: "Blueprint for capability registry tests",
2575 model: BlueprintModel::Inherit,
2576 system_prompt: "Test prompt",
2577 tools: vec![],
2578 max_turns: None,
2579 config_schema: None,
2580 }]
2581 }
2582 }
2583
2584 let mut registry = CapabilityRegistry::new();
2585 registry.register(BlueprintProviderCapability);
2586
2587 let (capability_id, blueprint) = registry
2588 .blueprint_with_capability("test_blueprint")
2589 .expect("blueprint should resolve with capability id");
2590 assert_eq!(capability_id, "blueprint_provider");
2591 assert_eq!(blueprint.id, "test_blueprint");
2592 }
2593
2594 #[test]
2595 fn test_capability_registry_builder() {
2596 let registry = CapabilityRegistry::builder()
2597 .capability(NoopCapability)
2598 .capability(CurrentTimeCapability)
2599 .build();
2600
2601 assert!(registry.has("noop"));
2602 assert!(registry.has("current_time"));
2603 assert_eq!(registry.len(), 2);
2604 }
2605
2606 #[test]
2607 fn test_capability_status() {
2608 let registry = CapabilityRegistry::with_builtins();
2609
2610 let current_time = registry.get("current_time").unwrap();
2611 assert_eq!(current_time.status(), CapabilityStatus::Available);
2612
2613 let research = registry.get("research").unwrap();
2614 assert_eq!(research.status(), CapabilityStatus::ComingSoon);
2615 }
2616
2617 #[test]
2618 fn test_capability_icons_and_categories() {
2619 let registry = CapabilityRegistry::with_builtins();
2620
2621 let noop = registry.get("noop").unwrap();
2622 assert_eq!(noop.icon(), Some("circle-off"));
2623 assert_eq!(noop.category(), Some("Testing"));
2624
2625 let current_time = registry.get("current_time").unwrap();
2626 assert_eq!(current_time.icon(), Some("clock"));
2627 assert_eq!(current_time.category(), Some("Core"));
2628 }
2629
2630 #[test]
2631 fn test_system_prompt_preview_default_delegates_to_addition() {
2632 let registry = CapabilityRegistry::with_builtins();
2633
2634 let test_math = registry.get("test_math").unwrap();
2636 assert_eq!(
2637 test_math.system_prompt_preview().as_deref(),
2638 test_math.system_prompt_addition()
2639 );
2640
2641 let current_time = registry.get("current_time").unwrap();
2643 assert!(current_time.system_prompt_preview().is_none());
2644 assert!(current_time.system_prompt_addition().is_none());
2645 }
2646
2647 #[test]
2648 fn test_system_prompt_preview_dynamic_capability() {
2649 let registry = CapabilityRegistry::with_builtins();
2650 let cap = registry.get("agent_instructions").unwrap();
2651
2652 assert!(cap.system_prompt_addition().is_none());
2654 assert!(cap.system_prompt_preview().is_some());
2655 assert!(cap.system_prompt_preview().unwrap().contains("AGENTS.md"));
2656 }
2657
2658 #[tokio::test]
2663 async fn test_apply_capabilities_empty() {
2664 let registry = CapabilityRegistry::with_builtins();
2665 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2666
2667 let applied =
2668 apply_capabilities(base_runtime_agent.clone(), &[], ®istry, &test_ctx()).await;
2669
2670 assert_eq!(
2671 applied.runtime_agent.system_prompt,
2672 base_runtime_agent.system_prompt
2673 );
2674 assert!(applied.tool_registry.is_empty());
2675 assert!(applied.applied_ids.is_empty());
2676 }
2677
2678 #[tokio::test]
2679 async fn test_apply_capabilities_noop() {
2680 let registry = CapabilityRegistry::with_builtins();
2681 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2682
2683 let applied = apply_capabilities(
2684 base_runtime_agent.clone(),
2685 &["noop".to_string()],
2686 ®istry,
2687 &test_ctx(),
2688 )
2689 .await;
2690
2691 assert_eq!(
2693 applied.runtime_agent.system_prompt,
2694 base_runtime_agent.system_prompt
2695 );
2696 assert!(applied.tool_registry.is_empty());
2697 assert_eq!(applied.applied_ids, vec!["noop"]);
2698 }
2699
2700 #[tokio::test]
2701 async fn test_apply_capabilities_current_time() {
2702 let registry = CapabilityRegistry::with_builtins();
2703 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2704
2705 let applied = apply_capabilities(
2706 base_runtime_agent.clone(),
2707 &["current_time".to_string()],
2708 ®istry,
2709 &test_ctx(),
2710 )
2711 .await;
2712
2713 assert_eq!(
2715 applied.runtime_agent.system_prompt,
2716 base_runtime_agent.system_prompt
2717 );
2718 assert!(applied.tool_registry.has("get_current_time"));
2719 assert_eq!(applied.tool_registry.len(), 1);
2720 assert_eq!(applied.applied_ids, vec!["current_time"]);
2721 }
2722
2723 #[tokio::test]
2724 async fn test_apply_capabilities_skips_coming_soon() {
2725 let registry = CapabilityRegistry::with_builtins();
2726 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2727
2728 let applied = apply_capabilities(
2730 base_runtime_agent.clone(),
2731 &["research".to_string()],
2732 ®istry,
2733 &test_ctx(),
2734 )
2735 .await;
2736
2737 assert_eq!(
2739 applied.runtime_agent.system_prompt,
2740 base_runtime_agent.system_prompt
2741 );
2742 assert!(applied.applied_ids.is_empty()); }
2744
2745 #[tokio::test]
2746 async fn test_apply_capabilities_multiple() {
2747 let registry = CapabilityRegistry::with_builtins();
2748 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2749
2750 let applied = apply_capabilities(
2751 base_runtime_agent.clone(),
2752 &["noop".to_string(), "current_time".to_string()],
2753 ®istry,
2754 &test_ctx(),
2755 )
2756 .await;
2757
2758 assert!(applied.tool_registry.has("get_current_time"));
2759 assert_eq!(applied.applied_ids, vec!["noop", "current_time"]);
2760 }
2761
2762 #[tokio::test]
2763 async fn test_apply_capabilities_preserves_order() {
2764 let registry = CapabilityRegistry::with_builtins();
2765 let base_runtime_agent = RuntimeAgent::new("Base prompt.", "gpt-5.2");
2766
2767 let applied = apply_capabilities(
2769 base_runtime_agent,
2770 &["current_time".to_string(), "noop".to_string()],
2771 ®istry,
2772 &test_ctx(),
2773 )
2774 .await;
2775
2776 assert_eq!(applied.applied_ids, vec!["current_time", "noop"]);
2777 }
2778
2779 #[tokio::test]
2780 async fn test_apply_capabilities_test_math() {
2781 let registry = CapabilityRegistry::with_builtins();
2782 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2783
2784 let applied = apply_capabilities(
2785 base_runtime_agent.clone(),
2786 &["test_math".to_string()],
2787 ®istry,
2788 &test_ctx(),
2789 )
2790 .await;
2791
2792 assert!(
2794 !applied
2795 .runtime_agent
2796 .system_prompt
2797 .contains("<capability id=\"test_math\">")
2798 );
2799 assert!(
2801 applied
2802 .runtime_agent
2803 .system_prompt
2804 .contains("You are a helpful assistant.")
2805 );
2806 assert!(applied.tool_registry.has("add"));
2807 assert!(applied.tool_registry.has("subtract"));
2808 assert!(applied.tool_registry.has("multiply"));
2809 assert!(applied.tool_registry.has("divide"));
2810 assert_eq!(applied.tool_registry.len(), 4);
2811 }
2812
2813 #[tokio::test]
2814 async fn test_apply_capabilities_test_weather() {
2815 let registry = CapabilityRegistry::with_builtins();
2816 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2817
2818 let applied = apply_capabilities(
2819 base_runtime_agent.clone(),
2820 &["test_weather".to_string()],
2821 ®istry,
2822 &test_ctx(),
2823 )
2824 .await;
2825
2826 assert!(
2828 !applied
2829 .runtime_agent
2830 .system_prompt
2831 .contains("<capability id=\"test_weather\">")
2832 );
2833 assert!(applied.tool_registry.has("get_weather"));
2834 assert!(applied.tool_registry.has("get_forecast"));
2835 assert_eq!(applied.tool_registry.len(), 2);
2836 }
2837
2838 #[tokio::test]
2839 async fn test_apply_capabilities_test_math_and_test_weather() {
2840 let registry = CapabilityRegistry::with_builtins();
2841 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2842
2843 let applied = apply_capabilities(
2844 base_runtime_agent.clone(),
2845 &["test_math".to_string(), "test_weather".to_string()],
2846 ®istry,
2847 &test_ctx(),
2848 )
2849 .await;
2850
2851 assert_eq!(applied.tool_registry.len(), 6); assert!(applied.tool_registry.has("add"));
2854 assert!(applied.tool_registry.has("get_weather"));
2855 }
2856
2857 #[tokio::test]
2858 async fn test_apply_capabilities_stateless_todo_list() {
2859 let registry = CapabilityRegistry::with_builtins();
2860 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2861
2862 let applied = apply_capabilities(
2863 base_runtime_agent.clone(),
2864 &["stateless_todo_list".to_string()],
2865 ®istry,
2866 &test_ctx(),
2867 )
2868 .await;
2869
2870 assert!(
2872 applied
2873 .runtime_agent
2874 .system_prompt
2875 .contains("Task Management")
2876 );
2877 assert!(applied.runtime_agent.system_prompt.contains("write_todos"));
2878 assert!(applied.tool_registry.has("write_todos"));
2879 assert_eq!(applied.tool_registry.len(), 1);
2880 }
2881
2882 #[tokio::test]
2883 async fn test_apply_capabilities_web_fetch() {
2884 let registry = CapabilityRegistry::with_builtins();
2885 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2886
2887 let applied = apply_capabilities(
2888 base_runtime_agent.clone(),
2889 &["web_fetch".to_string()],
2890 ®istry,
2891 &test_ctx(),
2892 )
2893 .await;
2894
2895 assert!(
2897 applied
2898 .runtime_agent
2899 .system_prompt
2900 .contains(&base_runtime_agent.system_prompt)
2901 );
2902 assert!(applied.runtime_agent.system_prompt.contains("web_fetch"));
2903 assert!(applied.tool_registry.has("web_fetch"));
2904 assert_eq!(applied.tool_registry.len(), 1);
2905 }
2906
2907 #[tokio::test]
2912 async fn test_xml_tags_wrap_capability_prompts() {
2913 let registry = CapabilityRegistry::with_builtins();
2914 let collected =
2915 collect_capabilities(&["stateless_todo_list".to_string()], ®istry, &test_ctx())
2916 .await;
2917
2918 assert_eq!(collected.system_prompt_parts.len(), 1);
2919 let part = &collected.system_prompt_parts[0];
2920 assert!(part.starts_with("<capability id=\"stateless_todo_list\">"));
2921 assert!(part.ends_with("</capability>"));
2922 assert!(part.contains("Task Management"));
2923 }
2924
2925 #[tokio::test]
2926 async fn test_xml_tags_multiple_capabilities() {
2927 let registry = CapabilityRegistry::with_builtins();
2928 let collected = collect_capabilities(
2929 &[
2930 "stateless_todo_list".to_string(),
2931 "session_schedule".to_string(),
2932 ],
2933 ®istry,
2934 &test_ctx(),
2935 )
2936 .await;
2937
2938 assert_eq!(collected.system_prompt_parts.len(), 2);
2939 assert!(
2940 collected.system_prompt_parts[0].starts_with("<capability id=\"stateless_todo_list\">")
2941 );
2942 assert!(
2943 collected.system_prompt_parts[1].starts_with("<capability id=\"session_schedule\">")
2944 );
2945
2946 let prefix = collected.system_prompt_prefix().unwrap();
2947 assert!(prefix.contains("</capability>\n\n<capability"));
2949 }
2950
2951 #[tokio::test]
2952 async fn test_xml_tags_system_prompt_wrapping() {
2953 let registry = CapabilityRegistry::with_builtins();
2954 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2955
2956 let applied = apply_capabilities(
2957 base,
2958 &["stateless_todo_list".to_string()],
2959 ®istry,
2960 &test_ctx(),
2961 )
2962 .await;
2963
2964 let prompt = &applied.runtime_agent.system_prompt;
2965 assert!(prompt.contains("<capability id=\"stateless_todo_list\">"));
2967 assert!(prompt.contains("</capability>"));
2968 assert!(prompt.contains("<system-prompt>\nYou are helpful.\n</system-prompt>"));
2970 }
2971
2972 #[tokio::test]
2973 async fn test_no_xml_wrapping_without_capabilities() {
2974 let registry = CapabilityRegistry::with_builtins();
2975 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2976
2977 let applied = apply_capabilities(base, &[], ®istry, &test_ctx()).await;
2978
2979 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2981 assert!(
2982 !applied
2983 .runtime_agent
2984 .system_prompt
2985 .contains("<system-prompt>")
2986 );
2987 }
2988
2989 #[tokio::test]
2990 async fn test_no_xml_wrapping_for_noop_capability() {
2991 let registry = CapabilityRegistry::with_builtins();
2992 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
2993
2994 let applied = apply_capabilities(base, &["noop".to_string()], ®istry, &test_ctx()).await;
2996
2997 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
2998 assert!(
2999 !applied
3000 .runtime_agent
3001 .system_prompt
3002 .contains("<system-prompt>")
3003 );
3004 }
3005
3006 #[tokio::test]
3011 async fn test_collect_capabilities_includes_mounts() {
3012 let registry = CapabilityRegistry::with_builtins();
3013
3014 let collected =
3015 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
3016
3017 assert!(!collected.mounts.is_empty());
3018 assert_eq!(collected.mounts.len(), 1);
3019 assert_eq!(collected.mounts[0].path, "/samples");
3020 assert!(collected.mounts[0].is_readonly());
3021 }
3022
3023 #[tokio::test]
3024 async fn test_collect_capabilities_empty_mounts_by_default() {
3025 let registry = CapabilityRegistry::with_builtins();
3026
3027 let collected =
3029 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
3030
3031 assert!(collected.mounts.is_empty());
3032 }
3033
3034 #[tokio::test]
3035 async fn test_collect_capabilities_combines_mounts() {
3036 let registry = CapabilityRegistry::with_builtins();
3037
3038 let collected = collect_capabilities(
3041 &["sample_data".to_string(), "current_time".to_string()],
3042 ®istry,
3043 &test_ctx(),
3044 )
3045 .await;
3046
3047 assert_eq!(collected.mounts.len(), 1);
3048 assert!(
3050 collected
3051 .applied_ids
3052 .iter()
3053 .any(|id| id == "session_file_system")
3054 );
3055 assert!(collected.applied_ids.iter().any(|id| id == "sample_data"));
3056 assert!(collected.applied_ids.iter().any(|id| id == "current_time"));
3057 }
3058
3059 #[test]
3060 fn test_sample_data_capability() {
3061 let registry = CapabilityRegistry::with_builtins();
3062 let cap = registry.get("sample_data").unwrap();
3063
3064 assert_eq!(cap.id(), "sample_data");
3065 assert_eq!(cap.name(), "Sample Data");
3066 assert_eq!(cap.status(), CapabilityStatus::Available);
3067
3068 assert!(cap.system_prompt_addition().is_some());
3070 assert!(cap.tools().is_empty());
3071
3072 assert!(!cap.mounts().is_empty());
3074 }
3075
3076 #[test]
3081 fn test_resolve_dependencies_empty() {
3082 let registry = CapabilityRegistry::with_builtins();
3083
3084 let resolved = resolve_dependencies(&[], ®istry).unwrap();
3085
3086 assert!(resolved.resolved_ids.is_empty());
3087 assert!(resolved.added_as_dependencies.is_empty());
3088 assert!(resolved.user_selected.is_empty());
3089 }
3090
3091 #[test]
3092 fn test_resolve_dependencies_no_deps() {
3093 let registry = CapabilityRegistry::with_builtins();
3094
3095 let resolved = resolve_dependencies(&["current_time".to_string()], ®istry).unwrap();
3097
3098 assert_eq!(resolved.resolved_ids, vec!["current_time"]);
3099 assert!(resolved.added_as_dependencies.is_empty());
3100 }
3101
3102 #[test]
3103 fn test_resolve_dependencies_with_deps() {
3104 let registry = CapabilityRegistry::with_builtins();
3105
3106 let resolved = resolve_dependencies(&["sample_data".to_string()], ®istry).unwrap();
3108
3109 assert_eq!(resolved.resolved_ids.len(), 2);
3111 let fs_pos = resolved
3112 .resolved_ids
3113 .iter()
3114 .position(|id| id == "session_file_system")
3115 .unwrap();
3116 let sd_pos = resolved
3117 .resolved_ids
3118 .iter()
3119 .position(|id| id == "sample_data")
3120 .unwrap();
3121 assert!(fs_pos < sd_pos, "FileSystem should come before SampleData");
3122
3123 assert_eq!(resolved.added_as_dependencies, vec!["session_file_system"]);
3125 }
3126
3127 #[test]
3128 fn test_resolve_dependencies_already_selected() {
3129 let registry = CapabilityRegistry::with_builtins();
3130
3131 let resolved = resolve_dependencies(
3133 &["session_file_system".to_string(), "sample_data".to_string()],
3134 ®istry,
3135 )
3136 .unwrap();
3137
3138 assert_eq!(resolved.resolved_ids.len(), 2);
3139 assert!(resolved.added_as_dependencies.is_empty());
3141 }
3142
3143 #[test]
3144 fn test_resolve_dependencies_preserves_order() {
3145 let registry = CapabilityRegistry::with_builtins();
3146
3147 let resolved =
3149 resolve_dependencies(&["current_time".to_string(), "noop".to_string()], ®istry)
3150 .unwrap();
3151
3152 assert_eq!(resolved.resolved_ids, vec!["current_time", "noop"]);
3153 }
3154
3155 #[test]
3156 fn test_resolve_dependencies_unknown_capability() {
3157 let registry = CapabilityRegistry::with_builtins();
3158
3159 let resolved =
3161 resolve_dependencies(&["unknown_capability".to_string()], ®istry).unwrap();
3162
3163 assert!(resolved.resolved_ids.is_empty());
3164 }
3165
3166 #[test]
3167 fn test_get_dependencies() {
3168 let registry = CapabilityRegistry::with_builtins();
3169
3170 let deps = get_dependencies("sample_data", ®istry);
3172 assert_eq!(deps, vec!["session_file_system"]);
3173
3174 let deps = get_dependencies("current_time", ®istry);
3176 assert!(deps.is_empty());
3177
3178 let deps = get_dependencies("unknown", ®istry);
3180 assert!(deps.is_empty());
3181 }
3182
3183 #[test]
3184 fn test_sample_data_has_dependency() {
3185 let registry = CapabilityRegistry::with_builtins();
3186 let cap = registry.get("sample_data").unwrap();
3187
3188 let deps = cap.dependencies();
3189 assert_eq!(deps.len(), 1);
3190 assert_eq!(deps[0], "session_file_system");
3191 }
3192
3193 #[test]
3194 fn test_noop_has_no_dependencies() {
3195 let registry = CapabilityRegistry::with_builtins();
3196 let cap = registry.get("noop").unwrap();
3197
3198 assert!(cap.dependencies().is_empty());
3199 }
3200
3201 #[test]
3205 fn test_circular_dependency_error() {
3206 struct CapA;
3208 struct CapB;
3209
3210 impl Capability for CapA {
3211 fn id(&self) -> &str {
3212 "test_cap_a"
3213 }
3214 fn name(&self) -> &str {
3215 "Test A"
3216 }
3217 fn description(&self) -> &str {
3218 "Test capability A"
3219 }
3220 fn dependencies(&self) -> Vec<&'static str> {
3221 vec!["test_cap_b"]
3222 }
3223 }
3224
3225 impl Capability for CapB {
3226 fn id(&self) -> &str {
3227 "test_cap_b"
3228 }
3229 fn name(&self) -> &str {
3230 "Test B"
3231 }
3232 fn description(&self) -> &str {
3233 "Test capability B"
3234 }
3235 fn dependencies(&self) -> Vec<&'static str> {
3236 vec!["test_cap_a"]
3237 }
3238 }
3239
3240 let mut registry = CapabilityRegistry::new();
3241 registry.register(CapA);
3242 registry.register(CapB);
3243
3244 let result = resolve_dependencies(&["test_cap_a".to_string()], ®istry);
3245
3246 assert!(result.is_err());
3247 match result.unwrap_err() {
3248 DependencyError::CircularDependency { capability_id, .. } => {
3249 assert_eq!(capability_id, "test_cap_a");
3250 }
3251 _ => panic!("Expected CircularDependency error"),
3252 }
3253 }
3254
3255 use crate::message_filter::{MessageFilter, MessageFilterProvider, MessageQuery};
3260
3261 struct FilterTestCapability {
3263 priority: i32,
3264 }
3265
3266 impl Capability for FilterTestCapability {
3267 fn id(&self) -> &str {
3268 "filter_test"
3269 }
3270 fn name(&self) -> &str {
3271 "Filter Test"
3272 }
3273 fn description(&self) -> &str {
3274 "Test capability with message filter"
3275 }
3276 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3277 Some(Arc::new(FilterTestProvider {
3278 priority: self.priority,
3279 }))
3280 }
3281 }
3282
3283 struct FilterTestProvider {
3284 priority: i32,
3285 }
3286
3287 impl MessageFilterProvider for FilterTestProvider {
3288 fn apply_filters(&self, query: &mut MessageQuery, config: &serde_json::Value) {
3289 if let Some(search) = config.get("search").and_then(|v| v.as_str()) {
3291 query
3292 .filters
3293 .push(MessageFilter::Search(search.to_string()));
3294 }
3295 }
3296
3297 fn priority(&self) -> i32 {
3298 self.priority
3299 }
3300 }
3301
3302 #[tokio::test]
3303 async fn test_collect_capabilities_with_configs_no_filter_providers() {
3304 let registry = CapabilityRegistry::with_builtins();
3305 let configs = vec![AgentCapabilityConfig {
3306 capability_ref: CapabilityId::new("current_time"),
3307 config: serde_json::json!({}),
3308 }];
3309
3310 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3311
3312 assert!(collected.message_filter_providers.is_empty());
3313 assert!(!collected.has_message_filters());
3314 }
3315
3316 #[tokio::test]
3317 async fn test_collect_capabilities_with_configs_with_filter_provider() {
3318 let mut registry = CapabilityRegistry::new();
3319 registry.register(FilterTestCapability { priority: 0 });
3320
3321 let configs = vec![AgentCapabilityConfig {
3322 capability_ref: CapabilityId::new("filter_test"),
3323 config: serde_json::json!({ "search": "hello" }),
3324 }];
3325
3326 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3327
3328 assert_eq!(collected.message_filter_providers.len(), 1);
3329 assert!(collected.has_message_filters());
3330 }
3331
3332 #[tokio::test]
3333 async fn test_collect_capabilities_with_configs_filter_priority_order() {
3334 struct HighPriorityCapability;
3336 struct LowPriorityCapability;
3337
3338 impl Capability for HighPriorityCapability {
3339 fn id(&self) -> &str {
3340 "high_priority"
3341 }
3342 fn name(&self) -> &str {
3343 "High Priority"
3344 }
3345 fn description(&self) -> &str {
3346 "Test"
3347 }
3348 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3349 Some(Arc::new(FilterTestProvider { priority: 10 }))
3350 }
3351 }
3352
3353 impl Capability for LowPriorityCapability {
3354 fn id(&self) -> &str {
3355 "low_priority"
3356 }
3357 fn name(&self) -> &str {
3358 "Low Priority"
3359 }
3360 fn description(&self) -> &str {
3361 "Test"
3362 }
3363 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3364 Some(Arc::new(FilterTestProvider { priority: -5 }))
3365 }
3366 }
3367
3368 let mut registry = CapabilityRegistry::new();
3369 registry.register(HighPriorityCapability);
3370 registry.register(LowPriorityCapability);
3371
3372 let configs = vec![
3374 AgentCapabilityConfig {
3375 capability_ref: CapabilityId::new("high_priority"),
3376 config: serde_json::json!({}),
3377 },
3378 AgentCapabilityConfig {
3379 capability_ref: CapabilityId::new("low_priority"),
3380 config: serde_json::json!({}),
3381 },
3382 ];
3383
3384 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3385
3386 assert_eq!(collected.message_filter_providers.len(), 2);
3388 assert_eq!(collected.message_filter_providers[0].0.priority(), -5);
3389 assert_eq!(collected.message_filter_providers[1].0.priority(), 10);
3390 }
3391
3392 #[tokio::test]
3393 async fn test_collected_capabilities_apply_message_filters() {
3394 let mut registry = CapabilityRegistry::new();
3395 registry.register(FilterTestCapability { priority: 0 });
3396
3397 let configs = vec![AgentCapabilityConfig {
3398 capability_ref: CapabilityId::new("filter_test"),
3399 config: serde_json::json!({ "search": "test_query" }),
3400 }];
3401
3402 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3403
3404 let session_id: SessionId = Uuid::now_v7().into();
3406 let mut query = MessageQuery::new(session_id);
3407
3408 collected.apply_message_filters(&mut query);
3409
3410 assert_eq!(query.filters.len(), 1);
3412 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3413 }
3414
3415 #[tokio::test]
3416 async fn test_collected_capabilities_apply_multiple_filters_in_priority_order() {
3417 struct SearchCapability {
3418 id: &'static str,
3419 search_term: &'static str,
3420 priority: i32,
3421 }
3422
3423 struct SearchProvider {
3424 search_term: &'static str,
3425 priority: i32,
3426 }
3427
3428 impl MessageFilterProvider for SearchProvider {
3429 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3430 query
3431 .filters
3432 .push(MessageFilter::Search(self.search_term.to_string()));
3433 }
3434
3435 fn priority(&self) -> i32 {
3436 self.priority
3437 }
3438 }
3439
3440 impl Capability for SearchCapability {
3441 fn id(&self) -> &str {
3442 self.id
3443 }
3444 fn name(&self) -> &str {
3445 "Search"
3446 }
3447 fn description(&self) -> &str {
3448 "Test"
3449 }
3450 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3451 Some(Arc::new(SearchProvider {
3452 search_term: self.search_term,
3453 priority: self.priority,
3454 }))
3455 }
3456 }
3457
3458 let mut registry = CapabilityRegistry::new();
3459 registry.register(SearchCapability {
3460 id: "cap_a",
3461 search_term: "alpha",
3462 priority: 5,
3463 });
3464 registry.register(SearchCapability {
3465 id: "cap_b",
3466 search_term: "beta",
3467 priority: 1,
3468 });
3469 registry.register(SearchCapability {
3470 id: "cap_c",
3471 search_term: "gamma",
3472 priority: 10,
3473 });
3474
3475 let configs = vec![
3476 AgentCapabilityConfig {
3477 capability_ref: CapabilityId::new("cap_a"),
3478 config: serde_json::json!({}),
3479 },
3480 AgentCapabilityConfig {
3481 capability_ref: CapabilityId::new("cap_b"),
3482 config: serde_json::json!({}),
3483 },
3484 AgentCapabilityConfig {
3485 capability_ref: CapabilityId::new("cap_c"),
3486 config: serde_json::json!({}),
3487 },
3488 ];
3489
3490 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3491
3492 let session_id: SessionId = Uuid::now_v7().into();
3493 let mut query = MessageQuery::new(session_id);
3494
3495 collected.apply_message_filters(&mut query);
3496
3497 assert_eq!(query.filters.len(), 3);
3499 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3500 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3501 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3502 }
3503
3504 #[test]
3505 fn test_capability_without_message_filter_returns_none() {
3506 let registry = CapabilityRegistry::with_builtins();
3507
3508 let noop = registry.get("noop").unwrap();
3509 assert!(noop.message_filter_provider().is_none());
3510
3511 let current_time = registry.get("current_time").unwrap();
3512 assert!(current_time.message_filter_provider().is_none());
3513 }
3514
3515 #[tokio::test]
3516 async fn test_collect_capabilities_preserves_config_for_filter_provider() {
3517 let mut registry = CapabilityRegistry::new();
3518 registry.register(FilterTestCapability { priority: 0 });
3519
3520 let test_config = serde_json::json!({
3521 "search": "custom_search",
3522 "extra_field": 42
3523 });
3524
3525 let configs = vec![AgentCapabilityConfig {
3526 capability_ref: CapabilityId::new("filter_test"),
3527 config: test_config.clone(),
3528 }];
3529
3530 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3531
3532 assert_eq!(collected.message_filter_providers.len(), 1);
3534 let (_, stored_config) = &collected.message_filter_providers[0];
3535 assert_eq!(*stored_config, test_config);
3536 }
3537
3538 #[test]
3543 fn test_collect_message_filters_only_collects_filters() {
3544 let mut registry = CapabilityRegistry::new();
3545 registry.register(FilterTestCapability { priority: 0 });
3546
3547 let configs = vec![AgentCapabilityConfig {
3548 capability_ref: CapabilityId::new("filter_test"),
3549 config: serde_json::json!({ "search": "test_query" }),
3550 }];
3551
3552 let collected = collect_message_filters_only(&configs, ®istry);
3553
3554 let session_id: SessionId = Uuid::now_v7().into();
3555 let mut query = MessageQuery::new(session_id);
3556 collected.apply_message_filters(&mut query);
3557
3558 assert_eq!(query.filters.len(), 1);
3559 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3560 }
3561
3562 #[test]
3563 fn test_collect_message_filters_only_skips_unknown_capabilities() {
3564 let registry = CapabilityRegistry::new();
3565
3566 let configs = vec![AgentCapabilityConfig {
3567 capability_ref: CapabilityId::new("nonexistent"),
3568 config: serde_json::json!({}),
3569 }];
3570
3571 let collected = collect_message_filters_only(&configs, ®istry);
3572 assert!(collected.message_filter_providers.is_empty());
3573 }
3574
3575 #[test]
3576 fn test_collect_message_filters_only_preserves_priority_order() {
3577 struct PriorityFilterCap {
3578 id: &'static str,
3579 search_term: &'static str,
3580 priority: i32,
3581 }
3582
3583 struct PriorityFilterProvider {
3584 search_term: &'static str,
3585 priority: i32,
3586 }
3587
3588 impl Capability for PriorityFilterCap {
3589 fn id(&self) -> &str {
3590 self.id
3591 }
3592 fn name(&self) -> &str {
3593 self.id
3594 }
3595 fn description(&self) -> &str {
3596 "priority test"
3597 }
3598 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3599 Some(Arc::new(PriorityFilterProvider {
3600 search_term: self.search_term,
3601 priority: self.priority,
3602 }))
3603 }
3604 }
3605
3606 impl MessageFilterProvider for PriorityFilterProvider {
3607 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3608 query
3609 .filters
3610 .push(MessageFilter::Search(self.search_term.to_string()));
3611 }
3612 fn priority(&self) -> i32 {
3613 self.priority
3614 }
3615 }
3616
3617 let mut registry = CapabilityRegistry::new();
3618 registry.register(PriorityFilterCap {
3619 id: "gamma",
3620 search_term: "gamma",
3621 priority: 10,
3622 });
3623 registry.register(PriorityFilterCap {
3624 id: "alpha",
3625 search_term: "alpha",
3626 priority: 5,
3627 });
3628 registry.register(PriorityFilterCap {
3629 id: "beta",
3630 search_term: "beta",
3631 priority: 1,
3632 });
3633
3634 let configs = vec![
3635 AgentCapabilityConfig {
3636 capability_ref: CapabilityId::new("gamma"),
3637 config: serde_json::json!({}),
3638 },
3639 AgentCapabilityConfig {
3640 capability_ref: CapabilityId::new("alpha"),
3641 config: serde_json::json!({}),
3642 },
3643 AgentCapabilityConfig {
3644 capability_ref: CapabilityId::new("beta"),
3645 config: serde_json::json!({}),
3646 },
3647 ];
3648
3649 let collected = collect_message_filters_only(&configs, ®istry);
3650
3651 let session_id: SessionId = Uuid::now_v7().into();
3652 let mut query = MessageQuery::new(session_id);
3653 collected.apply_message_filters(&mut query);
3654
3655 assert_eq!(query.filters.len(), 3);
3657 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3658 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3659 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3660 }
3661
3662 #[test]
3663 fn test_collect_message_filters_only_post_load_invoked() {
3664 use crate::message::Message;
3665
3666 struct PostLoadCap;
3667 struct PostLoadProvider;
3668
3669 impl Capability for PostLoadCap {
3670 fn id(&self) -> &str {
3671 "post_load_test"
3672 }
3673 fn name(&self) -> &str {
3674 "PostLoad Test"
3675 }
3676 fn description(&self) -> &str {
3677 "test"
3678 }
3679 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3680 Some(Arc::new(PostLoadProvider))
3681 }
3682 }
3683
3684 impl MessageFilterProvider for PostLoadProvider {
3685 fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
3686 fn priority(&self) -> i32 {
3687 0
3688 }
3689 fn post_load(&self, messages: &mut Vec<Message>, _config: &serde_json::Value) {
3690 messages.reverse();
3692 }
3693 }
3694
3695 let mut registry = CapabilityRegistry::new();
3696 registry.register(PostLoadCap);
3697
3698 let configs = vec![AgentCapabilityConfig {
3699 capability_ref: CapabilityId::new("post_load_test"),
3700 config: serde_json::json!({}),
3701 }];
3702
3703 let collected = collect_message_filters_only(&configs, ®istry);
3704
3705 let mut messages = vec![Message::user("first"), Message::user("second")];
3706 collected.apply_post_load_filters(&mut messages);
3707
3708 assert_eq!(messages[0].text(), Some("second"));
3710 assert_eq!(messages[1].text(), Some("first"));
3711 }
3712
3713 #[test]
3714 fn test_collect_model_view_providers_respects_compaction_capability_boundary() {
3715 use crate::tool_types::ToolCall;
3716
3717 fn tool_heavy_messages() -> Vec<Message> {
3718 let mut messages = vec![Message::user("inspect files repeatedly")];
3719 for index in 0..9 {
3720 let call_id = format!("call_{index}");
3721 messages.push(Message::assistant_with_tools(
3722 "",
3723 vec![ToolCall {
3724 id: call_id.clone(),
3725 name: "read_file".to_string(),
3726 arguments: serde_json::json!({"path": "/workspace/src/lib.rs"}),
3727 }],
3728 ));
3729 messages.push(Message::tool_result(
3730 call_id,
3731 Some(serde_json::json!({
3732 "path": "/workspace/src/lib.rs",
3733 "content": format!("{}{}", "large file line\n".repeat(1000), index),
3734 "total_lines": 1000,
3735 "lines_shown": {"start": 1, "end": 1000},
3736 "truncated": false
3737 })),
3738 None,
3739 ));
3740 }
3741 messages
3742 }
3743
3744 fn first_tool_result_is_masked(messages: &[Message]) -> bool {
3745 messages[2]
3746 .tool_result_content()
3747 .and_then(|result| result.result.as_ref())
3748 .and_then(|result| result.get("masked"))
3749 .and_then(|masked| masked.as_bool())
3750 .unwrap_or(false)
3751 }
3752
3753 let mut registry = CapabilityRegistry::new();
3754 registry.register(CompactionCapability);
3755 let context = ModelViewContext {
3756 session_id: SessionId::new(),
3757 prior_usage: None,
3758 };
3759
3760 let no_compaction = collect_model_view_providers(&[], ®istry, None);
3761 let unmasked = no_compaction.apply_model_view(tool_heavy_messages(), &context);
3762 assert!(!first_tool_result_is_masked(&unmasked));
3763
3764 let compaction = collect_model_view_providers(
3765 &[AgentCapabilityConfig {
3766 capability_ref: CapabilityId::new(COMPACTION_CAPABILITY_ID),
3767 config: serde_json::json!({}),
3768 }],
3769 ®istry,
3770 None,
3771 );
3772 let masked = compaction.apply_model_view(tool_heavy_messages(), &context);
3773 assert!(first_tool_result_is_masked(&masked));
3774 let last_tool = masked.last().unwrap().tool_result_content().unwrap();
3775 assert!(last_tool.result.as_ref().unwrap().get("content").is_some());
3776 }
3777
3778 struct DelegatingFilterCap {
3781 id: &'static str,
3782 inner: std::sync::Arc<InnerFilterCap>,
3783 }
3784 struct InnerFilterCap;
3785
3786 impl Capability for InnerFilterCap {
3787 fn id(&self) -> &str {
3788 "inner_filter"
3789 }
3790 fn name(&self) -> &str {
3791 "Inner Filter"
3792 }
3793 fn description(&self) -> &str {
3794 "inner"
3795 }
3796 fn message_filter_provider(&self) -> Option<std::sync::Arc<dyn MessageFilterProvider>> {
3797 Some(std::sync::Arc::new(SentinelFilter))
3798 }
3799 }
3800 struct SentinelFilter;
3801 impl MessageFilterProvider for SentinelFilter {
3802 fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
3803 }
3804 impl Capability for DelegatingFilterCap {
3805 fn id(&self) -> &str {
3806 self.id
3807 }
3808 fn name(&self) -> &str {
3809 "Delegating Filter"
3810 }
3811 fn description(&self) -> &str {
3812 "delegating"
3813 }
3814 fn message_filter_provider(&self) -> Option<std::sync::Arc<dyn MessageFilterProvider>> {
3815 None }
3817 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
3818 Some(&*self.inner)
3819 }
3820 }
3821
3822 #[test]
3823 fn test_collect_message_filters_only_honors_resolve_for_model_delegation() {
3824 let inner = std::sync::Arc::new(InnerFilterCap);
3825 let outer = DelegatingFilterCap {
3826 id: "delegating_filter",
3827 inner: inner.clone(),
3828 };
3829
3830 let mut registry = CapabilityRegistry::new();
3831 registry.register(outer);
3832
3833 let configs = vec![AgentCapabilityConfig {
3834 capability_ref: CapabilityId::new("delegating_filter"),
3835 config: serde_json::json!({}),
3836 }];
3837
3838 let collected = collect_message_filters_only(&configs, ®istry);
3841 assert_eq!(
3842 collected.message_filter_providers.len(),
3843 1,
3844 "provider from resolved inner capability must be collected"
3845 );
3846 }
3847
3848 struct DelegatingMvpCap {
3849 id: &'static str,
3850 inner: std::sync::Arc<InnerMvpCap>,
3851 }
3852 struct InnerMvpCap;
3853
3854 impl Capability for InnerMvpCap {
3855 fn id(&self) -> &str {
3856 "inner_mvp"
3857 }
3858 fn name(&self) -> &str {
3859 "Inner MVP"
3860 }
3861 fn description(&self) -> &str {
3862 "inner"
3863 }
3864 fn model_view_provider(
3865 &self,
3866 ) -> Option<std::sync::Arc<dyn crate::capabilities::ModelViewProvider>> {
3867 struct NoopMvp;
3869 impl crate::capabilities::ModelViewProvider for NoopMvp {
3870 fn apply_model_view(
3871 &self,
3872 messages: Vec<Message>,
3873 _config: &serde_json::Value,
3874 _context: &ModelViewContext<'_>,
3875 ) -> Vec<Message> {
3876 messages
3877 }
3878 }
3879 Some(std::sync::Arc::new(NoopMvp))
3880 }
3881 }
3882 impl Capability for DelegatingMvpCap {
3883 fn id(&self) -> &str {
3884 self.id
3885 }
3886 fn name(&self) -> &str {
3887 "Delegating MVP"
3888 }
3889 fn description(&self) -> &str {
3890 "delegating"
3891 }
3892 fn model_view_provider(
3893 &self,
3894 ) -> Option<std::sync::Arc<dyn crate::capabilities::ModelViewProvider>> {
3895 None }
3897 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
3898 Some(&*self.inner)
3899 }
3900 }
3901
3902 #[test]
3903 fn test_collect_model_view_providers_honors_resolve_for_model_delegation() {
3904 let inner = std::sync::Arc::new(InnerMvpCap);
3905 let outer = DelegatingMvpCap {
3906 id: "delegating_mvp",
3907 inner: inner.clone(),
3908 };
3909
3910 let mut registry = CapabilityRegistry::new();
3911 registry.register(outer);
3912
3913 let configs = vec![AgentCapabilityConfig {
3914 capability_ref: CapabilityId::new("delegating_mvp"),
3915 config: serde_json::json!({}),
3916 }];
3917
3918 let collected = collect_model_view_providers(&configs, ®istry, None);
3921 assert_eq!(
3922 collected.model_view_providers.len(),
3923 1,
3924 "provider from resolved inner capability must be collected"
3925 );
3926 }
3927
3928 #[tokio::test]
3938 async fn test_bashkit_shell_capability_produces_bash_tool() {
3939 let registry = CapabilityRegistry::with_builtins();
3940 let collected =
3941 collect_capabilities(&["bashkit_shell".to_string()], ®istry, &test_ctx()).await;
3942
3943 let tool_names: Vec<&str> = collected
3944 .tool_definitions
3945 .iter()
3946 .map(|t| t.name())
3947 .collect();
3948 assert!(
3949 tool_names.contains(&"bash"),
3950 "bashkit_shell capability must produce 'bash' tool, got: {:?}",
3951 tool_names
3952 );
3953 assert!(
3954 !collected.tools.is_empty(),
3955 "bashkit_shell must provide tool implementations"
3956 );
3957 }
3958
3959 #[tokio::test]
3960 async fn test_generic_harness_capability_set_produces_bash_tool() {
3961 let generic_harness_caps = vec![
3964 "session_file_system".to_string(),
3965 "bashkit_shell".to_string(),
3966 "web_fetch".to_string(),
3967 "session_storage".to_string(),
3968 "session".to_string(),
3969 "agent_instructions".to_string(),
3970 "skills".to_string(),
3971 "infinity_context".to_string(),
3972 "auto_tool_search".to_string(),
3973 ];
3974
3975 let registry = CapabilityRegistry::with_builtins();
3976 let collected = collect_capabilities(&generic_harness_caps, ®istry, &test_ctx()).await;
3977
3978 let tool_names: Vec<&str> = collected
3979 .tool_definitions
3980 .iter()
3981 .map(|t| t.name())
3982 .collect();
3983 assert!(
3984 tool_names.contains(&"bash"),
3985 "Generic Harness capabilities must produce 'bash' tool, got: {:?}",
3986 tool_names
3987 );
3988 }
3989
3990 #[tokio::test]
3991 async fn test_collect_capabilities_tool_count_matches_definitions() {
3992 let registry = CapabilityRegistry::with_builtins();
3995 let collected =
3996 collect_capabilities(&["bashkit_shell".to_string()], ®istry, &test_ctx()).await;
3997
3998 assert_eq!(
3999 collected.tools.len(),
4000 collected.tool_definitions.len(),
4001 "tool implementations ({}) must match tool definitions ({})",
4002 collected.tools.len(),
4003 collected.tool_definitions.len(),
4004 );
4005 }
4006
4007 #[tokio::test]
4011 async fn test_collect_capabilities_resolves_dependencies() {
4012 let registry = CapabilityRegistry::with_builtins();
4015 let collected =
4016 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
4017
4018 assert!(
4020 collected
4021 .applied_ids
4022 .iter()
4023 .any(|id| id == "session_file_system"),
4024 "collect_capabilities must apply session_file_system as a dependency; applied_ids: {:?}",
4025 collected.applied_ids
4026 );
4027
4028 let tool_names: Vec<&str> = collected
4029 .tool_definitions
4030 .iter()
4031 .map(|t| t.name())
4032 .collect();
4033
4034 assert!(
4036 tool_names.contains(&"read_file") && tool_names.contains(&"write_file"),
4037 "collect_capabilities must resolve dependencies and include dependency tools, got: {:?}",
4038 tool_names
4039 );
4040
4041 assert_eq!(
4043 collected.tools.len(),
4044 collected.tool_definitions.len(),
4045 "dependency-added tools must have implementations, not just definitions"
4046 );
4047 }
4048
4049 #[test]
4050 fn test_defaults_do_not_include_bash() {
4051 let registry = crate::ToolRegistry::with_defaults();
4054 assert!(
4055 !registry.has("bash"),
4056 "with_defaults() must not include 'bash' — it comes from bashkit_shell capability"
4057 );
4058 }
4059
4060 #[tokio::test]
4067 async fn test_background_execution_auto_activates_with_bashkit_shell() {
4068 let registry = CapabilityRegistry::with_builtins();
4069 let collected =
4070 collect_capabilities(&["bashkit_shell".to_string()], ®istry, &test_ctx()).await;
4071
4072 let tool_names: Vec<&str> = collected
4073 .tool_definitions
4074 .iter()
4075 .map(|t| t.name())
4076 .collect();
4077 assert!(
4078 tool_names.contains(&"spawn_background"),
4079 "spawn_background must be auto-activated when bashkit_shell (a \
4080 background-capable tool) is in the agent's capability set; got: {:?}",
4081 tool_names
4082 );
4083 assert!(
4084 collected
4085 .applied_ids
4086 .iter()
4087 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
4088 "background_execution must be in applied_ids when auto-activated; \
4089 got: {:?}",
4090 collected.applied_ids
4091 );
4092
4093 assert!(
4095 collected
4096 .tools
4097 .iter()
4098 .any(|t| t.name() == "spawn_background"),
4099 "spawn_background tool implementation must be present alongside the \
4100 definition (lockstep contract)"
4101 );
4102 }
4103
4104 #[tokio::test]
4107 async fn test_background_execution_does_not_auto_activate_without_hint() {
4108 let registry = CapabilityRegistry::with_builtins();
4109 let collected =
4111 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
4112
4113 let tool_names: Vec<&str> = collected
4114 .tool_definitions
4115 .iter()
4116 .map(|t| t.name())
4117 .collect();
4118 assert!(
4119 !tool_names.contains(&"spawn_background"),
4120 "spawn_background must NOT be activated without a background-capable \
4121 tool; got: {:?}",
4122 tool_names
4123 );
4124 assert!(
4125 !collected
4126 .applied_ids
4127 .iter()
4128 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
4129 "background_execution must not appear in applied_ids when no \
4130 background-capable tool is present; got: {:?}",
4131 collected.applied_ids
4132 );
4133 }
4134
4135 #[tokio::test]
4139 async fn test_background_execution_explicit_selection_is_idempotent() {
4140 let registry = CapabilityRegistry::with_builtins();
4141 let collected = collect_capabilities(
4142 &[
4143 "bashkit_shell".to_string(),
4144 BACKGROUND_EXECUTION_CAPABILITY_ID.to_string(),
4145 ],
4146 ®istry,
4147 &test_ctx(),
4148 )
4149 .await;
4150
4151 let spawn_background_count = collected
4152 .tool_definitions
4153 .iter()
4154 .filter(|t| t.name() == "spawn_background")
4155 .count();
4156 assert_eq!(
4157 spawn_background_count, 1,
4158 "spawn_background must appear exactly once even when \
4159 background_execution is selected explicitly alongside a \
4160 background-capable tool"
4161 );
4162 let applied_count = collected
4163 .applied_ids
4164 .iter()
4165 .filter(|id| id.as_str() == BACKGROUND_EXECUTION_CAPABILITY_ID)
4166 .count();
4167 assert_eq!(
4168 applied_count, 1,
4169 "background_execution must appear exactly once in applied_ids"
4170 );
4171 }
4172
4173 #[test]
4178 fn test_defaults_do_not_include_spawn_background() {
4179 let registry = crate::ToolRegistry::with_defaults();
4180 assert!(
4181 !registry.has("spawn_background"),
4182 "with_defaults() must not include 'spawn_background' — it comes \
4183 from the background_execution capability (EVE-501)"
4184 );
4185 }
4186
4187 #[test]
4192 fn test_capability_features_default_empty() {
4193 let registry = CapabilityRegistry::with_builtins();
4194
4195 let noop = registry.get("noop").unwrap();
4197 assert!(noop.features().is_empty());
4198
4199 let current_time = registry.get("current_time").unwrap();
4200 assert!(current_time.features().is_empty());
4201 }
4202
4203 #[test]
4204 fn test_file_system_capability_features() {
4205 let registry = CapabilityRegistry::with_builtins();
4206
4207 let fs = registry.get("session_file_system").unwrap();
4208 assert_eq!(fs.features(), vec!["file_system"]);
4209 }
4210
4211 #[test]
4212 fn test_bashkit_shell_capability_features() {
4213 let registry = CapabilityRegistry::with_builtins();
4214
4215 let bash = registry.get("bashkit_shell").unwrap();
4216 assert_eq!(bash.features(), vec!["file_system"]);
4217 }
4218
4219 #[test]
4220 fn test_alias_resolves_to_canonical_capability() {
4221 let registry = CapabilityRegistry::with_builtins();
4222
4223 let via_alias = registry.get("virtual_bash").unwrap();
4225 assert_eq!(via_alias.id(), "bashkit_shell");
4226 assert!(registry.has("virtual_bash"));
4227 assert_eq!(registry.canonical_id("virtual_bash"), Some("bashkit_shell"));
4228 assert_eq!(
4229 registry.canonical_id("bashkit_shell"),
4230 Some("bashkit_shell")
4231 );
4232 assert_eq!(registry.canonical_id("nonexistent"), None);
4233 }
4234
4235 #[test]
4236 fn test_alias_dedupes_with_canonical_in_dependency_resolution() {
4237 let registry = CapabilityRegistry::with_builtins();
4238
4239 let resolved = resolve_dependencies(
4242 &["virtual_bash".to_string(), "bashkit_shell".to_string()],
4243 ®istry,
4244 )
4245 .unwrap();
4246 let bash_ids: Vec<_> = resolved
4247 .resolved_ids
4248 .iter()
4249 .filter(|id| id.as_str() == "bashkit_shell" || id.as_str() == "virtual_bash")
4250 .collect();
4251 assert_eq!(bash_ids, vec!["bashkit_shell"]);
4252 assert!(
4254 !resolved
4255 .added_as_dependencies
4256 .contains(&"bashkit_shell".to_string())
4257 );
4258 }
4259
4260 #[test]
4261 fn test_alias_preserves_explicit_config_in_resolution() {
4262 let registry = CapabilityRegistry::with_builtins();
4263
4264 let configs = vec![AgentCapabilityConfig::with_config(
4265 "virtual_bash".to_string(),
4266 serde_json::json!({"key": "value"}),
4267 )];
4268 let resolved = resolve_capability_configs(&configs, ®istry).unwrap();
4269 let bash = resolved
4270 .iter()
4271 .find(|c| c.capability_id() == "bashkit_shell")
4272 .expect("alias must resolve to canonical bashkit_shell config");
4273 assert_eq!(bash.config, serde_json::json!({"key": "value"}));
4274 }
4275
4276 #[test]
4277 fn test_unregister_by_alias_removes_capability_and_aliases() {
4278 let mut registry = CapabilityRegistry::with_builtins();
4279
4280 assert!(registry.unregister("virtual_bash").is_some());
4281 assert!(!registry.has("bashkit_shell"));
4282 assert!(!registry.has("virtual_bash"));
4283 }
4284
4285 #[test]
4286 fn test_session_storage_capability_features() {
4287 let registry = CapabilityRegistry::with_builtins();
4288
4289 let storage = registry.get("session_storage").unwrap();
4290 let features = storage.features();
4291 assert!(features.contains(&"secrets"));
4292 assert!(features.contains(&"key_value"));
4293 }
4294
4295 #[test]
4296 fn test_session_schedule_capability_features() {
4297 let registry = CapabilityRegistry::with_builtins();
4298
4299 let schedule = registry.get("session_schedule").unwrap();
4300 assert_eq!(schedule.features(), vec!["schedules"]);
4301 }
4302
4303 #[test]
4304 fn test_session_sql_database_capability_features() {
4305 let registry = CapabilityRegistry::with_builtins();
4306
4307 let sql = registry.get("session_sql_database").unwrap();
4308 assert_eq!(sql.features(), vec!["sql_database"]);
4309 }
4310
4311 #[test]
4312 fn test_sample_data_capability_features() {
4313 let registry = CapabilityRegistry::with_builtins();
4314
4315 let sample = registry.get("sample_data").unwrap();
4316 assert_eq!(sample.features(), vec!["file_system"]);
4317 }
4318
4319 #[test]
4320 fn test_compute_features_empty() {
4321 let registry = CapabilityRegistry::with_builtins();
4322
4323 let features = compute_features(&[], ®istry);
4324 assert!(features.is_empty());
4325 }
4326
4327 #[test]
4328 fn test_compute_features_single_capability() {
4329 let registry = CapabilityRegistry::with_builtins();
4330
4331 let features = compute_features(&["session_schedule".to_string()], ®istry);
4332 assert_eq!(features, vec!["schedules"]);
4333 }
4334
4335 #[test]
4336 fn test_compute_features_multiple_capabilities() {
4337 let registry = CapabilityRegistry::with_builtins();
4338
4339 let features = compute_features(
4340 &[
4341 "session_file_system".to_string(),
4342 "session_storage".to_string(),
4343 "session_schedule".to_string(),
4344 ],
4345 ®istry,
4346 );
4347 assert!(features.contains(&"file_system".to_string()));
4348 assert!(features.contains(&"secrets".to_string()));
4349 assert!(features.contains(&"key_value".to_string()));
4350 assert!(features.contains(&"schedules".to_string()));
4351 }
4352
4353 #[test]
4354 fn test_compute_features_deduplicates() {
4355 let registry = CapabilityRegistry::with_builtins();
4356
4357 let features = compute_features(
4359 &[
4360 "session_file_system".to_string(),
4361 "bashkit_shell".to_string(),
4362 ],
4363 ®istry,
4364 );
4365 let file_system_count = features.iter().filter(|f| *f == "file_system").count();
4366 assert_eq!(file_system_count, 1, "file_system should appear only once");
4367 }
4368
4369 #[test]
4370 fn test_compute_features_includes_dependency_features() {
4371 let registry = CapabilityRegistry::with_builtins();
4372
4373 let features = compute_features(&["bashkit_shell".to_string()], ®istry);
4375 assert!(features.contains(&"file_system".to_string()));
4376 }
4377
4378 #[test]
4379 fn test_compute_features_generic_harness_set() {
4380 let registry = CapabilityRegistry::with_builtins();
4381
4382 let features = compute_features(
4384 &[
4385 "session_file_system".to_string(),
4386 "bashkit_shell".to_string(),
4387 "session_storage".to_string(),
4388 "session".to_string(),
4389 "session_schedule".to_string(),
4390 ],
4391 ®istry,
4392 );
4393 assert!(features.contains(&"file_system".to_string()));
4394 assert!(features.contains(&"secrets".to_string()));
4395 assert!(features.contains(&"key_value".to_string()));
4396 assert!(features.contains(&"schedules".to_string()));
4397 }
4398
4399 #[test]
4400 fn test_compute_features_unknown_capability_ignored() {
4401 let registry = CapabilityRegistry::with_builtins();
4402
4403 let features = compute_features(
4404 &["unknown_cap".to_string(), "session_schedule".to_string()],
4405 ®istry,
4406 );
4407 assert_eq!(features, vec!["schedules"]);
4408 }
4409
4410 #[test]
4411 fn test_risk_level_ordering() {
4412 assert!(RiskLevel::Low < RiskLevel::Medium);
4413 assert!(RiskLevel::Medium < RiskLevel::High);
4414 }
4415
4416 #[test]
4417 fn test_risk_level_serde_roundtrip() {
4418 let high = RiskLevel::High;
4419 let json = serde_json::to_string(&high).unwrap();
4420 assert_eq!(json, "\"high\"");
4421 let back: RiskLevel = serde_json::from_str(&json).unwrap();
4422 assert_eq!(back, RiskLevel::High);
4423 }
4424
4425 #[test]
4426 fn test_capability_risk_levels() {
4427 let registry = CapabilityRegistry::with_builtins();
4428
4429 let bash = registry.get("bashkit_shell").unwrap();
4431 assert_eq!(bash.risk_level(), RiskLevel::High);
4432
4433 let fetch = registry.get("web_fetch").unwrap();
4435 assert_eq!(fetch.risk_level(), RiskLevel::High);
4436
4437 let noop = registry.get("noop").unwrap();
4439 assert_eq!(noop.risk_level(), RiskLevel::Low);
4440 }
4441
4442 #[tokio::test]
4447 async fn test_apply_capabilities_openai_tool_search() {
4448 let registry = CapabilityRegistry::with_builtins();
4449 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4450
4451 let applied = apply_capabilities(
4452 base_runtime_agent.clone(),
4453 &["openai_tool_search".to_string()],
4454 ®istry,
4455 &test_ctx(),
4456 )
4457 .await;
4458
4459 assert_eq!(
4461 applied.runtime_agent.system_prompt,
4462 base_runtime_agent.system_prompt
4463 );
4464 assert!(applied.tool_registry.is_empty());
4465 assert_eq!(applied.applied_ids, vec!["openai_tool_search"]);
4466
4467 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
4469 assert!(ts.enabled);
4470 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
4471 }
4472
4473 #[tokio::test]
4474 async fn test_apply_capabilities_openai_tool_search_with_other_capabilities() {
4475 let registry = CapabilityRegistry::with_builtins();
4476 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4477
4478 let applied = apply_capabilities(
4479 base_runtime_agent,
4480 &[
4481 "current_time".to_string(),
4482 "openai_tool_search".to_string(),
4483 "test_math".to_string(),
4484 ],
4485 ®istry,
4486 &test_ctx(),
4487 )
4488 .await;
4489
4490 assert!(applied.tool_registry.has("get_current_time"));
4492 assert!(applied.tool_registry.has("add"));
4493 assert!(applied.tool_registry.has("subtract"));
4494 assert!(applied.tool_registry.has("multiply"));
4495 assert!(applied.tool_registry.has("divide"));
4496
4497 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
4499 assert!(ts.enabled);
4500 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
4501 }
4502
4503 #[tokio::test]
4504 async fn test_collect_capabilities_tool_search_custom_threshold() {
4505 let registry = CapabilityRegistry::with_builtins();
4506
4507 let configs = vec![AgentCapabilityConfig {
4508 capability_ref: CapabilityId::new("openai_tool_search"),
4509 config: serde_json::json!({"threshold": 5}),
4510 }];
4511
4512 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4513
4514 let ts = collected.tool_search.as_ref().unwrap();
4515 assert!(ts.enabled);
4516 assert_eq!(ts.threshold, 5);
4517 }
4518
4519 #[tokio::test]
4520 async fn test_collect_capabilities_auto_tool_search_resolves_to_generic_off_native() {
4521 let registry = CapabilityRegistry::with_builtins();
4522
4523 let configs = vec![
4524 AgentCapabilityConfig {
4525 capability_ref: CapabilityId::new("auto_tool_search"),
4526 config: serde_json::json!({"threshold": 2}),
4527 },
4528 AgentCapabilityConfig {
4529 capability_ref: CapabilityId::new("test_math"),
4530 config: serde_json::json!({}),
4531 },
4532 ];
4533
4534 let ctx = test_ctx().with_model("claude-3-5-haiku");
4538 let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
4539
4540 assert!(
4541 collected.tool_search.is_none(),
4542 "auto_tool_search must not set a hosted config on a non-native model"
4543 );
4544 assert!(
4545 collected
4546 .tools
4547 .iter()
4548 .any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
4549 "auto_tool_search must contribute the client-side tool_search tool"
4550 );
4551 assert!(
4552 !collected.tool_definition_hooks.is_empty(),
4553 "auto_tool_search must contribute a client-side deferral hook"
4554 );
4555
4556 let mut transformed = collected.tool_definitions.clone();
4557 for hook in &collected.tool_definition_hooks {
4558 transformed = hook.transform(transformed);
4559 }
4560 let add_tool = transformed
4561 .iter()
4562 .find(|tool| tool.name() == "add")
4563 .expect("test_math contributes add");
4564 assert!(
4565 add_tool.parameters().get("properties").is_none(),
4566 "generic auto_tool_search must honor the configured threshold"
4567 );
4568 }
4569
4570 #[tokio::test]
4571 async fn test_collect_capabilities_auto_tool_search_resolves_to_hosted_on_native() {
4572 let registry = CapabilityRegistry::with_builtins();
4573
4574 let configs = vec![AgentCapabilityConfig {
4575 capability_ref: CapabilityId::new("auto_tool_search"),
4576 config: serde_json::json!({"threshold": 7}),
4577 }];
4578
4579 let ctx = test_ctx().with_model("gpt-5.4");
4582 let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
4583
4584 let ts = collected
4585 .tool_search
4586 .as_ref()
4587 .expect("auto_tool_search must set a hosted config on a native model");
4588 assert!(ts.enabled);
4589 assert_eq!(ts.threshold, 7);
4590 assert!(
4591 !collected
4592 .tools
4593 .iter()
4594 .any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
4595 "hosted mechanism must not contribute the client-side tool_search tool"
4596 );
4597 assert!(
4598 collected.tool_definition_hooks.is_empty(),
4599 "hosted mechanism must not contribute a client-side deferral hook"
4600 );
4601 }
4602
4603 #[tokio::test]
4604 async fn test_collect_capabilities_auto_tool_search_resolves_to_hosted_on_anthropic() {
4605 let registry = CapabilityRegistry::with_builtins();
4606
4607 let configs = vec![AgentCapabilityConfig {
4608 capability_ref: CapabilityId::new("auto_tool_search"),
4609 config: serde_json::json!({"threshold": 9}),
4610 }];
4611
4612 let ctx = test_ctx().with_model("claude-opus-4-8");
4615 let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
4616
4617 let ts = collected
4618 .tool_search
4619 .as_ref()
4620 .expect("auto_tool_search must set a hosted config on a native Claude model");
4621 assert!(ts.enabled);
4622 assert_eq!(ts.threshold, 9);
4623 assert!(
4624 !collected
4625 .tools
4626 .iter()
4627 .any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
4628 "hosted mechanism must not contribute the client-side tool_search tool"
4629 );
4630 assert!(
4631 collected.tool_definition_hooks.is_empty(),
4632 "hosted mechanism must not contribute a client-side deferral hook"
4633 );
4634 }
4635
4636 #[tokio::test]
4637 async fn test_collect_capabilities_no_tool_search_without_capability() {
4638 let registry = CapabilityRegistry::with_builtins();
4639
4640 let configs = vec![AgentCapabilityConfig {
4641 capability_ref: CapabilityId::new("current_time"),
4642 config: serde_json::json!({}),
4643 }];
4644
4645 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4646
4647 assert!(collected.tool_search.is_none());
4648 }
4649
4650 #[tokio::test]
4651 async fn test_collect_capabilities_tool_search_category_propagation() {
4652 let registry = CapabilityRegistry::with_builtins();
4653
4654 let configs = vec![
4656 AgentCapabilityConfig {
4657 capability_ref: CapabilityId::new("test_math"),
4658 config: serde_json::json!({}),
4659 },
4660 AgentCapabilityConfig {
4661 capability_ref: CapabilityId::new("openai_tool_search"),
4662 config: serde_json::json!({}),
4663 },
4664 ];
4665
4666 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4667
4668 assert!(collected.tool_search.is_some());
4670
4671 for tool_def in &collected.tool_definitions {
4673 if ["add", "subtract", "multiply", "divide"].contains(&tool_def.name()) {
4675 assert!(
4676 tool_def.category().is_some(),
4677 "Tool {} should have a category from its capability",
4678 tool_def.name()
4679 );
4680 }
4681 }
4682 }
4683
4684 #[tokio::test]
4685 async fn test_apply_capabilities_prompt_caching() {
4686 let registry = CapabilityRegistry::with_builtins();
4687 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4688
4689 let applied = apply_capabilities(
4690 base_runtime_agent.clone(),
4691 &["prompt_caching".to_string()],
4692 ®istry,
4693 &test_ctx(),
4694 )
4695 .await;
4696
4697 assert_eq!(
4698 applied.runtime_agent.system_prompt,
4699 base_runtime_agent.system_prompt
4700 );
4701 assert!(applied.tool_registry.is_empty());
4702 assert_eq!(applied.applied_ids, vec!["prompt_caching"]);
4703
4704 let prompt_cache = applied.runtime_agent.prompt_cache.as_ref().unwrap();
4705 assert!(prompt_cache.enabled);
4706 assert_eq!(
4707 prompt_cache.strategy,
4708 crate::llm_driver_registry::PromptCacheStrategy::Auto
4709 );
4710 assert!(prompt_cache.gemini_cached_content.is_none());
4711 }
4712
4713 #[tokio::test]
4714 async fn test_collect_capabilities_prompt_caching_custom_strategy() {
4715 let registry = CapabilityRegistry::with_builtins();
4716
4717 let configs = vec![AgentCapabilityConfig {
4718 capability_ref: CapabilityId::new("prompt_caching"),
4719 config: serde_json::json!({"strategy": "auto"}),
4720 }];
4721
4722 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4723
4724 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
4725 assert!(prompt_cache.enabled);
4726 assert_eq!(
4727 prompt_cache.strategy,
4728 crate::llm_driver_registry::PromptCacheStrategy::Auto
4729 );
4730 assert!(prompt_cache.gemini_cached_content.is_none());
4731 }
4732
4733 #[tokio::test]
4734 async fn test_collect_capabilities_prompt_caching_gemini_cached_content() {
4735 let registry = CapabilityRegistry::with_builtins();
4736
4737 let configs = vec![AgentCapabilityConfig {
4738 capability_ref: CapabilityId::new("prompt_caching"),
4739 config: serde_json::json!({
4740 "strategy": "auto",
4741 "gemini_cached_content": "cachedContents/demo-cache"
4742 }),
4743 }];
4744
4745 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4746
4747 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
4748 assert_eq!(
4749 prompt_cache.gemini_cached_content.as_deref(),
4750 Some("cachedContents/demo-cache")
4751 );
4752 }
4753
4754 struct SkillContributingCapability;
4759
4760 impl Capability for SkillContributingCapability {
4761 fn id(&self) -> &str {
4762 "contributes_skills"
4763 }
4764 fn name(&self) -> &str {
4765 "Contributes Skills"
4766 }
4767 fn description(&self) -> &str {
4768 "Test capability that contributes skills."
4769 }
4770 fn contribute_skills(&self) -> Vec<SkillContribution> {
4771 vec![
4772 SkillContribution::new("alpha-skill", "Alpha skill desc", "# Alpha\nDo alpha.")
4773 .with_files(vec![(
4774 "scripts/a.sh".to_string(),
4775 "#!/bin/sh\necho a\n".to_string(),
4776 )]),
4777 SkillContribution::new("beta-skill", "Beta skill desc", "# Beta\nDo beta.")
4778 .with_user_invocable(false),
4779 ]
4780 }
4781 }
4782
4783 fn skill_md_from_entries(entries: &HashMap<String, MountEntry>) -> &str {
4784 match &entries.get("SKILL.md").expect("SKILL.md missing").source {
4785 MountSource::InlineFile { content, .. } => content.as_str(),
4786 _ => panic!("Expected InlineFile for SKILL.md"),
4787 }
4788 }
4789
4790 #[tokio::test]
4791 async fn test_contribute_skills_normalized_to_mounts() {
4792 let mut registry = CapabilityRegistry::new();
4793 registry.register(SkillContributingCapability);
4794
4795 let configs = vec![AgentCapabilityConfig {
4796 capability_ref: CapabilityId::new("contributes_skills"),
4797 config: serde_json::json!({}),
4798 }];
4799
4800 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4801
4802 let skill_mounts: Vec<_> = collected
4803 .mounts
4804 .iter()
4805 .filter(|m| m.path.starts_with("/.agents/skills/"))
4806 .collect();
4807 assert_eq!(skill_mounts.len(), 2);
4808
4809 for m in &skill_mounts {
4812 assert!(m.is_readonly());
4813 assert_eq!(m.capability_id, "contributes_skills");
4814 }
4815
4816 let alpha = skill_mounts
4817 .iter()
4818 .find(|m| m.path == "/.agents/skills/alpha-skill")
4819 .expect("alpha-skill mount missing");
4820 match &alpha.source {
4821 MountSource::InlineDirectory { entries } => {
4822 assert!(entries.contains_key("SKILL.md"));
4823 assert!(entries.contains_key("scripts/a.sh"));
4824 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
4825 assert_eq!(parsed.name, "alpha-skill");
4826 assert!(parsed.user_invocable);
4827 }
4828 _ => panic!("Expected InlineDirectory"),
4829 }
4830
4831 let beta = skill_mounts
4832 .iter()
4833 .find(|m| m.path == "/.agents/skills/beta-skill")
4834 .expect("beta-skill mount missing");
4835 match &beta.source {
4836 MountSource::InlineDirectory { entries } => {
4837 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
4838 assert!(!parsed.user_invocable);
4839 }
4840 _ => panic!("Expected InlineDirectory"),
4841 }
4842 }
4843
4844 #[tokio::test]
4845 async fn test_contribute_skills_default_empty() {
4846 let mut registry = CapabilityRegistry::new();
4849 registry.register(FilterTestCapability { priority: 0 });
4850
4851 let configs = vec![AgentCapabilityConfig {
4852 capability_ref: CapabilityId::new("filter_test"),
4853 config: serde_json::json!({}),
4854 }];
4855
4856 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4857 assert!(
4858 collected
4859 .mounts
4860 .iter()
4861 .all(|m| !m.path.starts_with("/.agents/skills/"))
4862 );
4863 }
4864
4865 struct LocalizedCapability;
4866
4867 impl Capability for LocalizedCapability {
4868 fn id(&self) -> &str {
4869 "localized"
4870 }
4871 fn name(&self) -> &str {
4872 "Localized"
4873 }
4874 fn description(&self) -> &str {
4875 "English description"
4876 }
4877 fn localizations(&self) -> Vec<CapabilityLocalization> {
4878 vec![
4879 CapabilityLocalization {
4880 locale: "en",
4881 name: None,
4882 description: None,
4883 config_description: Some("Controls things."),
4884 config_overlay: None,
4885 },
4886 CapabilityLocalization {
4887 locale: "uk",
4888 name: Some("Локалізована"),
4889 description: Some("Український опис"),
4890 config_description: Some("Керує налаштуваннями."),
4891 config_overlay: None,
4892 },
4893 ]
4894 }
4895 }
4896
4897 #[test]
4898 fn localized_name_falls_back_exact_language_then_base() {
4899 let cap = LocalizedCapability;
4900 assert_eq!(cap.localized_name(Some("uk-UA")), "Локалізована");
4902 assert_eq!(cap.localized_name(Some("uk")), "Локалізована");
4903 assert_eq!(cap.localized_name(Some("uk_UA")), "Локалізована");
4905 assert_eq!(cap.localized_name(Some("fr-FR")), "Localized");
4907 assert_eq!(cap.localized_name(None), "Localized");
4908 assert_eq!(cap.localized_description(Some("uk")), "Український опис");
4909 assert_eq!(cap.localized_description(Some("de")), "English description");
4910 }
4911
4912 #[test]
4913 fn describe_schema_resolves_config_description_per_locale() {
4914 let cap = LocalizedCapability;
4915 assert_eq!(
4916 cap.describe_schema(Some("uk-UA")).as_deref(),
4917 Some("Керує налаштуваннями.")
4918 );
4919 assert_eq!(
4921 cap.describe_schema(Some("pl")).as_deref(),
4922 Some("Controls things.")
4923 );
4924 assert_eq!(
4925 cap.describe_schema(None).as_deref(),
4926 Some("Controls things.")
4927 );
4928 assert_eq!(NoopCapability.describe_schema(Some("uk")), None);
4930 }
4931}