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
1604fn compaction_is_enabled(
1610 capability_configs: &[AgentCapabilityConfig],
1611 registry: &CapabilityRegistry,
1612) -> bool {
1613 capability_configs.iter().any(|cap_config| {
1614 cap_config.capability_ref.as_str() == COMPACTION_CAPABILITY_ID
1615 && registry
1616 .get(cap_config.capability_ref.as_str())
1617 .is_some_and(|cap| cap.status() == CapabilityStatus::Available)
1618 })
1619}
1620
1621fn message_filter_config_for(
1630 cap_id: &str,
1631 base: &serde_json::Value,
1632 compaction_on: bool,
1633) -> serde_json::Value {
1634 if cap_id != INFINITY_CONTEXT_CAPABILITY_ID || !compaction_on {
1635 return base.clone();
1636 }
1637 let mut config = base.clone();
1638 match config.as_object_mut() {
1639 Some(map) => {
1640 map.insert(
1641 "compaction_active".to_string(),
1642 serde_json::Value::Bool(true),
1643 );
1644 }
1645 None => {
1646 config = serde_json::json!({ "compaction_active": true });
1647 }
1648 }
1649 config
1650}
1651
1652pub fn collect_message_filters_only(
1658 capability_configs: &[AgentCapabilityConfig],
1659 registry: &CapabilityRegistry,
1660) -> CollectedMessageFilters {
1661 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
1662 Vec::new();
1663 let compaction_on = compaction_is_enabled(capability_configs, registry);
1664
1665 for cap_config in capability_configs {
1666 let cap_id = cap_config.capability_ref.as_str();
1667 if let Some(capability) = registry.get(cap_id) {
1668 if capability.status() != CapabilityStatus::Available {
1669 continue;
1670 }
1671 let effective: &dyn Capability = capability
1674 .resolve_for_model(None)
1675 .unwrap_or_else(|| capability.as_ref());
1676 if let Some(provider) = effective.message_filter_provider() {
1677 let config = message_filter_config_for(cap_id, &cap_config.config, compaction_on);
1678 message_filter_providers.push((provider, config));
1679 }
1680 }
1681 }
1682
1683 message_filter_providers.sort_by_key(|(p, _)| p.priority());
1684
1685 CollectedMessageFilters {
1686 message_filter_providers,
1687 }
1688}
1689
1690pub fn collect_model_view_providers(
1697 capability_configs: &[AgentCapabilityConfig],
1698 registry: &CapabilityRegistry,
1699 model: Option<&str>,
1700) -> CollectedModelViewProviders {
1701 let mut model_view_providers: Vec<(Arc<dyn ModelViewProvider>, serde_json::Value)> = Vec::new();
1702
1703 for cap_config in capability_configs {
1704 let cap_id = cap_config.capability_ref.as_str();
1705 if let Some(capability) = registry.get(cap_id) {
1706 if capability.status() != CapabilityStatus::Available {
1707 continue;
1708 }
1709 let effective: &dyn Capability = capability
1710 .resolve_for_model(model)
1711 .unwrap_or_else(|| capability.as_ref());
1712 if let Some(provider) = effective.model_view_provider() {
1713 model_view_providers.push((provider, cap_config.config.clone()));
1714 }
1715 }
1716 }
1717
1718 model_view_providers.sort_by_key(|(p, _)| p.priority());
1719
1720 CollectedModelViewProviders {
1721 model_view_providers,
1722 }
1723}
1724
1725pub fn collect_capability_mcp_servers(
1726 capability_configs: &[AgentCapabilityConfig],
1727 registry: &CapabilityRegistry,
1728) -> ScopedMcpServers {
1729 let mut servers = ScopedMcpServers::default();
1730
1731 for cap_config in capability_configs {
1732 let cap_id = cap_config.capability_ref.as_str();
1733 if is_declarative_capability(cap_id) || is_plugin_capability(cap_id) {
1736 if let Ok(definition) =
1737 serde_json::from_value::<DeclarativeCapabilityDefinition>(cap_config.config.clone())
1738 {
1739 if definition.status != CapabilityStatus::Available {
1740 continue;
1741 }
1742 if let Some(contributed) = definition.mcp_servers {
1743 servers = merge_scoped_mcp_servers(&servers, &contributed);
1744 }
1745 }
1746 continue;
1747 }
1748 if let Some(capability) = registry.get(cap_id) {
1749 if capability.status() != CapabilityStatus::Available {
1750 continue;
1751 }
1752 servers = merge_scoped_mcp_servers(
1753 &servers,
1754 &capability.mcp_servers_with_config(&cap_config.config),
1755 );
1756 }
1757 }
1758
1759 servers
1760}
1761
1762pub const MAX_RESOLVED_CAPABILITIES: usize = 100;
1769
1770#[derive(Debug, Clone, PartialEq, Eq)]
1772pub enum DependencyError {
1773 CircularDependency {
1775 capability_id: String,
1777 chain: Vec<String>,
1779 },
1780 TooManyCapabilities {
1782 count: usize,
1784 max: usize,
1786 },
1787}
1788
1789impl std::fmt::Display for DependencyError {
1790 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1791 match self {
1792 DependencyError::CircularDependency {
1793 capability_id,
1794 chain,
1795 } => {
1796 write!(
1797 f,
1798 "Circular dependency detected: {} depends on itself via chain: {} -> {}",
1799 capability_id,
1800 chain.join(" -> "),
1801 capability_id
1802 )
1803 }
1804 DependencyError::TooManyCapabilities { count, max } => {
1805 write!(
1806 f,
1807 "Too many capabilities after resolution: {} (max: {})",
1808 count, max
1809 )
1810 }
1811 }
1812 }
1813}
1814
1815impl std::error::Error for DependencyError {}
1816
1817#[derive(Debug, Clone)]
1819pub struct ResolvedCapabilities {
1820 pub resolved_ids: Vec<String>,
1823 pub added_as_dependencies: Vec<String>,
1825 pub user_selected: Vec<String>,
1827}
1828
1829pub fn resolve_dependencies(
1849 selected_ids: &[String],
1850 registry: &CapabilityRegistry,
1851) -> Result<ResolvedCapabilities, DependencyError> {
1852 use std::collections::HashSet;
1853
1854 let user_selected: HashSet<String> = selected_ids
1856 .iter()
1857 .map(|id| registry.canonical_id(id).unwrap_or(id).to_string())
1858 .collect();
1859 let mut resolved: Vec<String> = Vec::new();
1860 let mut resolved_set: HashSet<String> = HashSet::new();
1861 let mut added_as_dependencies: Vec<String> = Vec::new();
1862
1863 for cap_id in selected_ids {
1865 resolve_single_capability(
1866 cap_id,
1867 registry,
1868 &mut resolved,
1869 &mut resolved_set,
1870 &mut added_as_dependencies,
1871 &user_selected,
1872 &mut Vec::new(), )?;
1874 }
1875
1876 if resolved.len() > MAX_RESOLVED_CAPABILITIES {
1878 return Err(DependencyError::TooManyCapabilities {
1879 count: resolved.len(),
1880 max: MAX_RESOLVED_CAPABILITIES,
1881 });
1882 }
1883
1884 Ok(ResolvedCapabilities {
1885 resolved_ids: resolved,
1886 added_as_dependencies,
1887 user_selected: selected_ids.to_vec(),
1888 })
1889}
1890
1891pub fn resolve_capability_configs(
1896 selected_configs: &[AgentCapabilityConfig],
1897 registry: &CapabilityRegistry,
1898) -> Result<Vec<AgentCapabilityConfig>, DependencyError> {
1899 let mut selected_ids: Vec<String> = Vec::new();
1900 for config in selected_configs {
1901 if (is_declarative_capability(config.capability_id())
1904 || is_plugin_capability(config.capability_id()))
1905 && let Ok(definition) =
1906 serde_json::from_value::<DeclarativeCapabilityDefinition>(config.config.clone())
1907 {
1908 selected_ids.extend(definition.dependencies);
1909 }
1910 selected_ids.push(config.capability_id().to_string());
1911 }
1912 let resolved = resolve_dependencies(&selected_ids, registry)?;
1913
1914 let explicit_configs: std::collections::HashMap<String, serde_json::Value> = selected_configs
1917 .iter()
1918 .map(|config| {
1919 let id = config.capability_id();
1920 let id = registry.canonical_id(id).unwrap_or(id);
1921 (id.to_string(), config.config.clone())
1922 })
1923 .collect();
1924
1925 Ok(resolved
1926 .resolved_ids
1927 .into_iter()
1928 .map(|capability_id| {
1929 explicit_configs
1930 .get(&capability_id)
1931 .cloned()
1932 .map(|config| AgentCapabilityConfig::with_config(capability_id.clone(), config))
1933 .unwrap_or_else(|| AgentCapabilityConfig::new(capability_id))
1934 })
1935 .collect())
1936}
1937
1938fn resolve_single_capability(
1940 cap_id: &str,
1941 registry: &CapabilityRegistry,
1942 resolved: &mut Vec<String>,
1943 resolved_set: &mut std::collections::HashSet<String>,
1944 added_as_dependencies: &mut Vec<String>,
1945 user_selected: &std::collections::HashSet<String>,
1946 visiting: &mut Vec<String>,
1947) -> Result<(), DependencyError> {
1948 let cap_id = registry.canonical_id(cap_id).unwrap_or(cap_id);
1952
1953 if resolved_set.contains(cap_id) {
1955 return Ok(());
1956 }
1957
1958 if visiting.contains(&cap_id.to_string()) {
1960 return Err(DependencyError::CircularDependency {
1961 capability_id: cap_id.to_string(),
1962 chain: visiting.clone(),
1963 });
1964 }
1965
1966 let capability = match registry.get(cap_id) {
1968 Some(cap) => cap,
1969 None => {
1970 if (is_declarative_capability(cap_id) || is_plugin_capability(cap_id))
1974 && !resolved_set.contains(cap_id)
1975 {
1976 resolved.push(cap_id.to_string());
1977 resolved_set.insert(cap_id.to_string());
1978 if !user_selected.contains(cap_id) {
1979 added_as_dependencies.push(cap_id.to_string());
1980 }
1981 }
1982 return Ok(());
1983 }
1984 };
1985
1986 visiting.push(cap_id.to_string());
1988
1989 for dep_id in capability.dependencies() {
1991 resolve_single_capability(
1992 dep_id,
1993 registry,
1994 resolved,
1995 resolved_set,
1996 added_as_dependencies,
1997 user_selected,
1998 visiting,
1999 )?;
2000 }
2001
2002 visiting.pop();
2004
2005 if !resolved_set.contains(cap_id) {
2007 resolved.push(cap_id.to_string());
2008 resolved_set.insert(cap_id.to_string());
2009
2010 if !user_selected.contains(cap_id) {
2012 added_as_dependencies.push(cap_id.to_string());
2013 }
2014 }
2015
2016 Ok(())
2017}
2018
2019pub fn compute_features(capability_ids: &[String], registry: &CapabilityRegistry) -> Vec<String> {
2024 use std::collections::HashSet;
2025
2026 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
2027 Ok(resolved) => resolved.resolved_ids,
2028 Err(_) => capability_ids.to_vec(),
2029 };
2030
2031 let mut seen = HashSet::new();
2032 let mut features = Vec::new();
2033 for cap_id in &resolved_ids {
2034 if let Some(cap) = registry.get(cap_id) {
2035 for feature in cap.features() {
2036 if seen.insert(feature) {
2037 features.push(feature.to_string());
2038 }
2039 }
2040 }
2041 }
2042 features
2043}
2044
2045pub fn get_dependencies(cap_id: &str, registry: &CapabilityRegistry) -> Vec<String> {
2048 registry
2049 .get(cap_id)
2050 .map(|cap| cap.dependencies().iter().map(|s| s.to_string()).collect())
2051 .unwrap_or_default()
2052}
2053
2054pub async fn collect_capabilities(
2070 capability_ids: &[String],
2071 registry: &CapabilityRegistry,
2072 ctx: &SystemPromptContext,
2073) -> CollectedCapabilities {
2074 let resolved_ids = match resolve_dependencies(capability_ids, registry) {
2077 Ok(resolved) => resolved.resolved_ids,
2078 Err(e) => {
2079 tracing::warn!("Failed to resolve capability dependencies: {}", e);
2080 capability_ids.to_vec()
2081 }
2082 };
2083
2084 let configs: Vec<AgentCapabilityConfig> = resolved_ids
2086 .iter()
2087 .map(|id| AgentCapabilityConfig {
2088 capability_ref: CapabilityId::new(id),
2089 config: serde_json::Value::Object(serde_json::Map::new()),
2090 })
2091 .collect();
2092
2093 collect_capabilities_with_configs(&configs, registry, ctx).await
2094}
2095
2096pub async fn collect_capabilities_with_configs(
2107 capability_configs: &[AgentCapabilityConfig],
2108 registry: &CapabilityRegistry,
2109 ctx: &SystemPromptContext,
2110) -> CollectedCapabilities {
2111 let mut system_prompt_parts: Vec<String> = Vec::new();
2112 let mut system_prompt_attributions: Vec<SystemPromptAttribution> = Vec::new();
2113 let mut tools: Vec<Box<dyn Tool>> = Vec::new();
2114 let mut tool_definitions: Vec<ToolDefinition> = Vec::new();
2115 let mut mounts: Vec<MountPoint> = Vec::new();
2116 let mut message_filter_providers: Vec<(Arc<dyn MessageFilterProvider>, serde_json::Value)> =
2117 Vec::new();
2118 let mut applied_ids: Vec<String> = Vec::new();
2119 let mut tool_search: Option<crate::llm_driver_registry::ToolSearchConfig> = None;
2120 let mut prompt_cache: Option<crate::llm_driver_registry::PromptCacheConfig> = None;
2121 let mut tool_definition_hooks: Vec<Arc<dyn ToolDefinitionHook>> = Vec::new();
2122 let mut tool_call_hooks: Vec<Arc<dyn ToolCallHook>> = Vec::new();
2123 let mut mcp_servers = ScopedMcpServers::default();
2124 let compaction_on = compaction_is_enabled(capability_configs, registry);
2125
2126 for cap_config in capability_configs {
2127 let cap_id = cap_config.capability_ref.as_str();
2128 if is_declarative_capability(cap_id) || is_plugin_capability(cap_id) {
2133 match serde_json::from_value::<DeclarativeCapabilityDefinition>(
2134 cap_config.config.clone(),
2135 ) {
2136 Ok(definition) => {
2137 if definition.status != CapabilityStatus::Available {
2138 continue;
2139 }
2140
2141 if let Some(prompt) = definition.system_prompt.as_deref() {
2142 let contribution =
2143 format!("<capability id=\"{}\">\n{}\n</capability>", cap_id, prompt);
2144 system_prompt_attributions.push(SystemPromptAttribution {
2145 capability_id: cap_id.to_string(),
2146 content: contribution.clone(),
2147 });
2148 system_prompt_parts.push(contribution);
2149 }
2150
2151 mounts.extend(definition.mounts(cap_id));
2152 if let Some(ref servers) = definition.mcp_servers {
2153 mcp_servers = merge_scoped_mcp_servers(&mcp_servers, servers);
2154 }
2155 for skill in definition.skill_contributions() {
2156 mounts.push(skill.to_mount(cap_id));
2157 }
2158
2159 applied_ids.push(cap_id.to_string());
2160 }
2161 Err(error) => {
2162 tracing::warn!(
2163 capability_id = %cap_id,
2164 error = %error,
2165 "Skipping invalid declarative/plugin capability config"
2166 );
2167 }
2168 }
2169 continue;
2170 }
2171 if let Some(capability) = registry.get(cap_id) {
2172 if capability.status() != CapabilityStatus::Available {
2174 continue;
2175 }
2176
2177 let effective: &dyn Capability =
2189 match capability.resolve_for_model(ctx.model.as_deref()) {
2190 Some(inner) => inner,
2191 None => capability.as_ref(),
2192 };
2193 let effective_id = effective.id();
2194
2195 if let Some(contribution) = effective
2197 .system_prompt_contribution_with_config(ctx, &cap_config.config)
2198 .await
2199 {
2200 system_prompt_attributions.push(SystemPromptAttribution {
2201 capability_id: cap_id.to_string(),
2202 content: contribution.clone(),
2203 });
2204 system_prompt_parts.push(contribution);
2205 }
2206
2207 tools.extend(effective.tools_with_config(&cap_config.config));
2209 tool_definition_hooks
2210 .extend(effective.tool_definition_hooks_with_context(ctx, &cap_config.config));
2211 tool_call_hooks.extend(effective.tool_call_hooks());
2212 let cap_category = effective.category();
2217 for def in effective.tool_definitions() {
2218 let def = match (def.category(), cap_category) {
2219 (None, Some(cat)) => def.with_category(cat),
2220 _ => def,
2221 }
2222 .with_capability_attribution(cap_id, Some(capability.name()));
2223 tool_definitions.push(def);
2224 }
2225
2226 if effective_id == OPENAI_TOOL_SEARCH_CAPABILITY_ID
2234 || effective_id == CLAUDE_TOOL_SEARCH_CAPABILITY_ID
2235 {
2236 let threshold = cap_config
2238 .config
2239 .get("threshold")
2240 .and_then(|v| v.as_u64())
2241 .map(|v| v as usize)
2242 .unwrap_or(DEFAULT_TOOL_SEARCH_THRESHOLD);
2243 tool_search = Some(crate::llm_driver_registry::ToolSearchConfig {
2244 enabled: true,
2245 threshold,
2246 });
2247 }
2248
2249 if cap_id == PROMPT_CACHING_CAPABILITY_ID {
2250 let strategy = cap_config
2251 .config
2252 .get("strategy")
2253 .and_then(|v| v.as_str())
2254 .map(|value| match value {
2255 "auto" => crate::llm_driver_registry::PromptCacheStrategy::Auto,
2256 _ => crate::llm_driver_registry::PromptCacheStrategy::Auto,
2257 })
2258 .unwrap_or(crate::llm_driver_registry::PromptCacheStrategy::Auto);
2259 let gemini_cached_content = cap_config
2260 .config
2261 .get("gemini_cached_content")
2262 .and_then(|v| v.as_str())
2263 .map(str::to_string);
2264 prompt_cache = Some(crate::llm_driver_registry::PromptCacheConfig {
2265 enabled: true,
2266 strategy,
2267 gemini_cached_content,
2268 });
2269 }
2270
2271 mounts.extend(effective.mounts());
2273
2274 mcp_servers = merge_scoped_mcp_servers(
2275 &mcp_servers,
2276 &effective.mcp_servers_with_config(&cap_config.config),
2277 );
2278
2279 for skill in effective.contribute_skills() {
2283 mounts.push(skill.to_mount(cap_id));
2284 }
2285
2286 if let Some(provider) = effective.message_filter_provider() {
2288 let config = message_filter_config_for(cap_id, &cap_config.config, compaction_on);
2289 message_filter_providers.push((provider, config));
2290 }
2291
2292 applied_ids.push(cap_id.to_string());
2293 }
2294 }
2295
2296 if !applied_ids
2308 .iter()
2309 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID)
2310 && tool_definitions
2311 .iter()
2312 .any(|def| def.hints().supports_background == Some(true))
2313 && let Some(bg_cap) = registry.get(BACKGROUND_EXECUTION_CAPABILITY_ID)
2314 && bg_cap.status() == CapabilityStatus::Available
2315 {
2316 tools.extend(bg_cap.tools());
2317 let cap_category = bg_cap.category();
2318 for def in bg_cap.tool_definitions() {
2319 let def = match (def.category(), cap_category) {
2320 (None, Some(cat)) => def.with_category(cat),
2321 _ => def,
2322 }
2323 .with_capability_attribution(BACKGROUND_EXECUTION_CAPABILITY_ID, Some(bg_cap.name()));
2324 tool_definitions.push(def);
2325 }
2326 applied_ids.push(BACKGROUND_EXECUTION_CAPABILITY_ID.to_string());
2327 }
2328
2329 message_filter_providers.sort_by_key(|(p, _)| p.priority());
2331
2332 CollectedCapabilities {
2333 system_prompt_parts,
2334 system_prompt_attributions,
2335 tools,
2336 tool_definitions,
2337 mounts,
2338 message_filter_providers,
2339 applied_ids,
2340 tool_search,
2341 prompt_cache,
2342 tool_definition_hooks,
2343 tool_call_hooks,
2344 mcp_servers,
2345 }
2346}
2347
2348pub struct AppliedCapabilities {
2354 pub runtime_agent: RuntimeAgent,
2356 pub tool_registry: ToolRegistry,
2358 pub applied_ids: Vec<String>,
2360}
2361
2362pub async fn apply_capabilities(
2399 base_runtime_agent: RuntimeAgent,
2400 capability_ids: &[String],
2401 registry: &CapabilityRegistry,
2402 ctx: &SystemPromptContext,
2403) -> AppliedCapabilities {
2404 let collected = collect_capabilities(capability_ids, registry, ctx).await;
2405
2406 let final_system_prompt = match collected.system_prompt_prefix() {
2408 Some(prefix) => format!(
2409 "{}\n\n<system-prompt>\n{}\n</system-prompt>",
2410 prefix, base_runtime_agent.system_prompt
2411 ),
2412 None => base_runtime_agent.system_prompt,
2413 };
2414
2415 let mut tool_registry = ToolRegistry::new();
2417 for tool in collected.tools {
2418 tool_registry.register_boxed(tool);
2419 }
2420
2421 let mut tools = collected.tool_definitions;
2423 for hook in &collected.tool_definition_hooks {
2424 tools = hook.transform(tools);
2425 }
2426
2427 let runtime_agent = RuntimeAgent {
2428 system_prompt: final_system_prompt,
2429 model: base_runtime_agent.model,
2430 tools,
2431 max_iterations: base_runtime_agent.max_iterations,
2432 temperature: base_runtime_agent.temperature,
2433 max_tokens: base_runtime_agent.max_tokens,
2434 tool_search: collected.tool_search,
2435 prompt_cache: collected.prompt_cache,
2436 network_access: base_runtime_agent.network_access,
2437 };
2438
2439 AppliedCapabilities {
2440 runtime_agent,
2441 tool_registry,
2442 applied_ids: collected.applied_ids,
2443 }
2444}
2445
2446#[cfg(test)]
2451mod tests {
2452 use super::*;
2453 use crate::typed_id::SessionId;
2454 use std::collections::BTreeSet;
2455 use uuid::Uuid;
2456
2457 static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2459
2460 fn lock_env() -> std::sync::MutexGuard<'static, ()> {
2461 ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner())
2462 }
2463
2464 fn test_ctx() -> SystemPromptContext {
2466 SystemPromptContext::without_file_store(SessionId::new())
2467 }
2468
2469 fn expected_core_builtin_ids() -> BTreeSet<&'static str> {
2471 let mut ids = [
2472 "agent_instructions",
2473 "human_intent",
2474 "budgeting",
2475 "self_budget",
2476 "noop",
2477 "current_time",
2478 "research",
2479 "platform_management",
2480 "session_file_system",
2481 "session_storage",
2482 "session",
2483 "session_sql_database",
2484 "test_math",
2485 "test_weather",
2486 "stateless_todo_list",
2487 "web_fetch",
2488 "bashkit_shell",
2489 "background_execution",
2490 "session_schedule",
2491 "btw",
2492 "infinity_context",
2493 "compaction",
2494 "memory",
2495 "message_metadata",
2496 "openai_tool_search",
2497 "claude_tool_search",
2498 "tool_search",
2499 "auto_tool_search",
2500 "prompt_caching",
2501 "session_tasks",
2502 "skills",
2503 "subagents",
2504 "system_commands",
2505 "sample_data",
2506 "data_knowledge",
2507 "knowledge_base",
2508 "tool_output_persistence",
2509 "tool_output_distillation",
2510 "fake_warehouse",
2511 "fake_aws",
2512 "fake_crm",
2513 "fake_financial",
2514 "loop_detection",
2515 "error_disclosure",
2516 "prompt_canary_guardrail",
2517 "guardrails",
2518 "user_hooks",
2519 "model_scout",
2520 "openrouter_workspace",
2521 ]
2522 .into_iter()
2523 .collect::<BTreeSet<_>>();
2524 if cfg!(feature = "ui-capabilities") {
2525 ids.insert("openui");
2526 ids.insert("a2ui");
2527 }
2528 ids
2529 }
2530
2531 fn expected_dev_builtin_ids() -> BTreeSet<&'static str> {
2533 let mut ids = expected_core_builtin_ids();
2534 ids.insert("agent_handoff");
2535 ids.insert("a2a_agent_delegation");
2536 ids
2537 }
2538
2539 fn registry_ids(registry: &CapabilityRegistry) -> BTreeSet<&str> {
2540 registry.capabilities.keys().map(String::as_str).collect()
2541 }
2542
2543 #[test]
2553 fn test_capability_registry_with_builtins_dev() {
2554 let _lock = lock_env();
2556 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2557 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
2558 assert_eq!(registry_ids(®istry), expected_dev_builtin_ids());
2559 assert!(registry.has("agent_handoff"));
2560 assert!(registry.has("a2a_agent_delegation"));
2561 }
2562
2563 #[test]
2564 fn test_capability_registry_with_builtins_prod() {
2565 let _lock = lock_env();
2567 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2568 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2569 assert_eq!(registry_ids(®istry), expected_core_builtin_ids());
2570 assert!(!registry.has("docker_container"));
2572 assert!(!registry.has("agent_handoff"));
2573 assert!(!registry.has("a2a_agent_delegation"));
2574 }
2575
2576 #[test]
2577 fn test_agent_delegation_enabled_by_env_in_prod() {
2578 let _lock = lock_env();
2580 unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "true") };
2581 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Prod);
2582 assert!(registry.has("agent_handoff"));
2583 assert!(registry.has("a2a_agent_delegation"));
2584 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2585 }
2586
2587 #[test]
2588 fn test_agent_delegation_disabled_by_env_in_dev() {
2589 let _lock = lock_env();
2591 unsafe { std::env::set_var("FEATURE_AGENT_DELEGATION", "false") };
2592 let registry = CapabilityRegistry::with_builtins_for_grade(DeploymentGrade::Dev);
2593 assert!(!registry.has("agent_handoff"));
2594 assert!(!registry.has("a2a_agent_delegation"));
2595 unsafe { std::env::remove_var("FEATURE_AGENT_DELEGATION") };
2596 }
2597
2598 #[test]
2599 fn test_capability_registry_get() {
2600 let registry = CapabilityRegistry::with_builtins();
2601
2602 let noop = registry.get("noop").unwrap();
2603 assert_eq!(noop.id(), "noop");
2604 assert_eq!(noop.name(), "No-Op");
2605 assert_eq!(noop.status(), CapabilityStatus::Available);
2606 }
2607
2608 #[test]
2609 fn test_capability_registry_blueprint_with_capability() {
2610 struct BlueprintProviderCapability;
2611
2612 impl Capability for BlueprintProviderCapability {
2613 fn id(&self) -> &str {
2614 "blueprint_provider"
2615 }
2616 fn name(&self) -> &str {
2617 "Blueprint Provider"
2618 }
2619 fn description(&self) -> &str {
2620 "Capability that provides a blueprint for tests"
2621 }
2622 fn agent_blueprints(&self) -> Vec<AgentBlueprint> {
2623 vec![AgentBlueprint {
2624 id: "test_blueprint",
2625 name: "Test Blueprint",
2626 description: "Blueprint for capability registry tests",
2627 model: BlueprintModel::Inherit,
2628 system_prompt: "Test prompt",
2629 tools: vec![],
2630 max_turns: None,
2631 config_schema: None,
2632 }]
2633 }
2634 }
2635
2636 let mut registry = CapabilityRegistry::new();
2637 registry.register(BlueprintProviderCapability);
2638
2639 let (capability_id, blueprint) = registry
2640 .blueprint_with_capability("test_blueprint")
2641 .expect("blueprint should resolve with capability id");
2642 assert_eq!(capability_id, "blueprint_provider");
2643 assert_eq!(blueprint.id, "test_blueprint");
2644 }
2645
2646 #[test]
2647 fn test_capability_registry_builder() {
2648 let registry = CapabilityRegistry::builder()
2649 .capability(NoopCapability)
2650 .capability(CurrentTimeCapability)
2651 .build();
2652
2653 assert!(registry.has("noop"));
2654 assert!(registry.has("current_time"));
2655 assert_eq!(registry.len(), 2);
2656 }
2657
2658 #[test]
2659 fn test_capability_status() {
2660 let registry = CapabilityRegistry::with_builtins();
2661
2662 let current_time = registry.get("current_time").unwrap();
2663 assert_eq!(current_time.status(), CapabilityStatus::Available);
2664
2665 let research = registry.get("research").unwrap();
2666 assert_eq!(research.status(), CapabilityStatus::ComingSoon);
2667 }
2668
2669 #[test]
2670 fn test_capability_icons_and_categories() {
2671 let registry = CapabilityRegistry::with_builtins();
2672
2673 let noop = registry.get("noop").unwrap();
2674 assert_eq!(noop.icon(), Some("circle-off"));
2675 assert_eq!(noop.category(), Some("Testing"));
2676
2677 let current_time = registry.get("current_time").unwrap();
2678 assert_eq!(current_time.icon(), Some("clock"));
2679 assert_eq!(current_time.category(), Some("Core"));
2680 }
2681
2682 #[test]
2683 fn test_system_prompt_preview_default_delegates_to_addition() {
2684 let registry = CapabilityRegistry::with_builtins();
2685
2686 let test_math = registry.get("test_math").unwrap();
2688 assert_eq!(
2689 test_math.system_prompt_preview().as_deref(),
2690 test_math.system_prompt_addition()
2691 );
2692
2693 let current_time = registry.get("current_time").unwrap();
2695 assert!(current_time.system_prompt_preview().is_none());
2696 assert!(current_time.system_prompt_addition().is_none());
2697 }
2698
2699 #[test]
2700 fn test_system_prompt_preview_dynamic_capability() {
2701 let registry = CapabilityRegistry::with_builtins();
2702 let cap = registry.get("agent_instructions").unwrap();
2703
2704 assert!(cap.system_prompt_addition().is_none());
2706 assert!(cap.system_prompt_preview().is_some());
2707 assert!(cap.system_prompt_preview().unwrap().contains("AGENTS.md"));
2708 }
2709
2710 #[tokio::test]
2715 async fn test_apply_capabilities_empty() {
2716 let registry = CapabilityRegistry::with_builtins();
2717 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2718
2719 let applied =
2720 apply_capabilities(base_runtime_agent.clone(), &[], ®istry, &test_ctx()).await;
2721
2722 assert_eq!(
2723 applied.runtime_agent.system_prompt,
2724 base_runtime_agent.system_prompt
2725 );
2726 assert!(applied.tool_registry.is_empty());
2727 assert!(applied.applied_ids.is_empty());
2728 }
2729
2730 #[tokio::test]
2731 async fn test_apply_capabilities_noop() {
2732 let registry = CapabilityRegistry::with_builtins();
2733 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2734
2735 let applied = apply_capabilities(
2736 base_runtime_agent.clone(),
2737 &["noop".to_string()],
2738 ®istry,
2739 &test_ctx(),
2740 )
2741 .await;
2742
2743 assert_eq!(
2745 applied.runtime_agent.system_prompt,
2746 base_runtime_agent.system_prompt
2747 );
2748 assert!(applied.tool_registry.is_empty());
2749 assert_eq!(applied.applied_ids, vec!["noop"]);
2750 }
2751
2752 #[tokio::test]
2753 async fn test_apply_capabilities_current_time() {
2754 let registry = CapabilityRegistry::with_builtins();
2755 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2756
2757 let applied = apply_capabilities(
2758 base_runtime_agent.clone(),
2759 &["current_time".to_string()],
2760 ®istry,
2761 &test_ctx(),
2762 )
2763 .await;
2764
2765 assert_eq!(
2767 applied.runtime_agent.system_prompt,
2768 base_runtime_agent.system_prompt
2769 );
2770 assert!(applied.tool_registry.has("get_current_time"));
2771 assert_eq!(applied.tool_registry.len(), 1);
2772 assert_eq!(applied.applied_ids, vec!["current_time"]);
2773 }
2774
2775 #[tokio::test]
2776 async fn test_apply_capabilities_skips_coming_soon() {
2777 let registry = CapabilityRegistry::with_builtins();
2778 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2779
2780 let applied = apply_capabilities(
2782 base_runtime_agent.clone(),
2783 &["research".to_string()],
2784 ®istry,
2785 &test_ctx(),
2786 )
2787 .await;
2788
2789 assert_eq!(
2791 applied.runtime_agent.system_prompt,
2792 base_runtime_agent.system_prompt
2793 );
2794 assert!(applied.applied_ids.is_empty()); }
2796
2797 #[tokio::test]
2798 async fn test_apply_capabilities_multiple() {
2799 let registry = CapabilityRegistry::with_builtins();
2800 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2801
2802 let applied = apply_capabilities(
2803 base_runtime_agent.clone(),
2804 &["noop".to_string(), "current_time".to_string()],
2805 ®istry,
2806 &test_ctx(),
2807 )
2808 .await;
2809
2810 assert!(applied.tool_registry.has("get_current_time"));
2811 assert_eq!(applied.applied_ids, vec!["noop", "current_time"]);
2812 }
2813
2814 #[tokio::test]
2815 async fn test_apply_capabilities_preserves_order() {
2816 let registry = CapabilityRegistry::with_builtins();
2817 let base_runtime_agent = RuntimeAgent::new("Base prompt.", "gpt-5.2");
2818
2819 let applied = apply_capabilities(
2821 base_runtime_agent,
2822 &["current_time".to_string(), "noop".to_string()],
2823 ®istry,
2824 &test_ctx(),
2825 )
2826 .await;
2827
2828 assert_eq!(applied.applied_ids, vec!["current_time", "noop"]);
2829 }
2830
2831 #[tokio::test]
2832 async fn test_apply_capabilities_test_math() {
2833 let registry = CapabilityRegistry::with_builtins();
2834 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2835
2836 let applied = apply_capabilities(
2837 base_runtime_agent.clone(),
2838 &["test_math".to_string()],
2839 ®istry,
2840 &test_ctx(),
2841 )
2842 .await;
2843
2844 assert!(
2846 !applied
2847 .runtime_agent
2848 .system_prompt
2849 .contains("<capability id=\"test_math\">")
2850 );
2851 assert!(
2853 applied
2854 .runtime_agent
2855 .system_prompt
2856 .contains("You are a helpful assistant.")
2857 );
2858 assert!(applied.tool_registry.has("add"));
2859 assert!(applied.tool_registry.has("subtract"));
2860 assert!(applied.tool_registry.has("multiply"));
2861 assert!(applied.tool_registry.has("divide"));
2862 assert_eq!(applied.tool_registry.len(), 4);
2863 }
2864
2865 #[tokio::test]
2866 async fn test_apply_capabilities_test_weather() {
2867 let registry = CapabilityRegistry::with_builtins();
2868 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2869
2870 let applied = apply_capabilities(
2871 base_runtime_agent.clone(),
2872 &["test_weather".to_string()],
2873 ®istry,
2874 &test_ctx(),
2875 )
2876 .await;
2877
2878 assert!(
2880 !applied
2881 .runtime_agent
2882 .system_prompt
2883 .contains("<capability id=\"test_weather\">")
2884 );
2885 assert!(applied.tool_registry.has("get_weather"));
2886 assert!(applied.tool_registry.has("get_forecast"));
2887 assert_eq!(applied.tool_registry.len(), 2);
2888 }
2889
2890 #[tokio::test]
2891 async fn test_apply_capabilities_test_math_and_test_weather() {
2892 let registry = CapabilityRegistry::with_builtins();
2893 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2894
2895 let applied = apply_capabilities(
2896 base_runtime_agent.clone(),
2897 &["test_math".to_string(), "test_weather".to_string()],
2898 ®istry,
2899 &test_ctx(),
2900 )
2901 .await;
2902
2903 assert_eq!(applied.tool_registry.len(), 6); assert!(applied.tool_registry.has("add"));
2906 assert!(applied.tool_registry.has("get_weather"));
2907 }
2908
2909 #[tokio::test]
2910 async fn test_apply_capabilities_stateless_todo_list() {
2911 let registry = CapabilityRegistry::with_builtins();
2912 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2913
2914 let applied = apply_capabilities(
2915 base_runtime_agent.clone(),
2916 &["stateless_todo_list".to_string()],
2917 ®istry,
2918 &test_ctx(),
2919 )
2920 .await;
2921
2922 assert!(
2924 applied
2925 .runtime_agent
2926 .system_prompt
2927 .contains("Task Management")
2928 );
2929 assert!(applied.runtime_agent.system_prompt.contains("write_todos"));
2930 assert!(applied.tool_registry.has("write_todos"));
2931 assert_eq!(applied.tool_registry.len(), 1);
2932 }
2933
2934 #[tokio::test]
2935 async fn test_apply_capabilities_web_fetch() {
2936 let registry = CapabilityRegistry::with_builtins();
2937 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.2");
2938
2939 let applied = apply_capabilities(
2940 base_runtime_agent.clone(),
2941 &["web_fetch".to_string()],
2942 ®istry,
2943 &test_ctx(),
2944 )
2945 .await;
2946
2947 assert!(
2949 applied
2950 .runtime_agent
2951 .system_prompt
2952 .contains(&base_runtime_agent.system_prompt)
2953 );
2954 assert!(applied.runtime_agent.system_prompt.contains("web_fetch"));
2955 assert!(applied.tool_registry.has("web_fetch"));
2956 assert_eq!(applied.tool_registry.len(), 1);
2957 }
2958
2959 #[tokio::test]
2964 async fn test_xml_tags_wrap_capability_prompts() {
2965 let registry = CapabilityRegistry::with_builtins();
2966 let collected =
2967 collect_capabilities(&["stateless_todo_list".to_string()], ®istry, &test_ctx())
2968 .await;
2969
2970 assert_eq!(collected.system_prompt_parts.len(), 1);
2971 let part = &collected.system_prompt_parts[0];
2972 assert!(part.starts_with("<capability id=\"stateless_todo_list\">"));
2973 assert!(part.ends_with("</capability>"));
2974 assert!(part.contains("Task Management"));
2975 }
2976
2977 #[tokio::test]
2978 async fn test_xml_tags_multiple_capabilities() {
2979 let registry = CapabilityRegistry::with_builtins();
2980 let collected = collect_capabilities(
2981 &[
2982 "stateless_todo_list".to_string(),
2983 "session_schedule".to_string(),
2984 ],
2985 ®istry,
2986 &test_ctx(),
2987 )
2988 .await;
2989
2990 assert_eq!(collected.system_prompt_parts.len(), 2);
2991 assert!(
2992 collected.system_prompt_parts[0].starts_with("<capability id=\"stateless_todo_list\">")
2993 );
2994 assert!(
2995 collected.system_prompt_parts[1].starts_with("<capability id=\"session_schedule\">")
2996 );
2997
2998 let prefix = collected.system_prompt_prefix().unwrap();
2999 assert!(prefix.contains("</capability>\n\n<capability"));
3001 }
3002
3003 #[tokio::test]
3004 async fn test_xml_tags_system_prompt_wrapping() {
3005 let registry = CapabilityRegistry::with_builtins();
3006 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
3007
3008 let applied = apply_capabilities(
3009 base,
3010 &["stateless_todo_list".to_string()],
3011 ®istry,
3012 &test_ctx(),
3013 )
3014 .await;
3015
3016 let prompt = &applied.runtime_agent.system_prompt;
3017 assert!(prompt.contains("<capability id=\"stateless_todo_list\">"));
3019 assert!(prompt.contains("</capability>"));
3020 assert!(prompt.contains("<system-prompt>\nYou are helpful.\n</system-prompt>"));
3022 }
3023
3024 #[tokio::test]
3025 async fn test_no_xml_wrapping_without_capabilities() {
3026 let registry = CapabilityRegistry::with_builtins();
3027 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
3028
3029 let applied = apply_capabilities(base, &[], ®istry, &test_ctx()).await;
3030
3031 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
3033 assert!(
3034 !applied
3035 .runtime_agent
3036 .system_prompt
3037 .contains("<system-prompt>")
3038 );
3039 }
3040
3041 #[tokio::test]
3042 async fn test_no_xml_wrapping_for_noop_capability() {
3043 let registry = CapabilityRegistry::with_builtins();
3044 let base = RuntimeAgent::new("You are helpful.", "gpt-5.2");
3045
3046 let applied = apply_capabilities(base, &["noop".to_string()], ®istry, &test_ctx()).await;
3048
3049 assert_eq!(applied.runtime_agent.system_prompt, "You are helpful.");
3050 assert!(
3051 !applied
3052 .runtime_agent
3053 .system_prompt
3054 .contains("<system-prompt>")
3055 );
3056 }
3057
3058 #[tokio::test]
3063 async fn test_collect_capabilities_includes_mounts() {
3064 let registry = CapabilityRegistry::with_builtins();
3065
3066 let collected =
3067 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
3068
3069 assert!(!collected.mounts.is_empty());
3070 assert_eq!(collected.mounts.len(), 1);
3071 assert_eq!(collected.mounts[0].path, "/samples");
3072 assert!(collected.mounts[0].is_readonly());
3073 }
3074
3075 #[tokio::test]
3076 async fn test_collect_capabilities_empty_mounts_by_default() {
3077 let registry = CapabilityRegistry::with_builtins();
3078
3079 let collected =
3081 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
3082
3083 assert!(collected.mounts.is_empty());
3084 }
3085
3086 #[tokio::test]
3087 async fn test_collect_capabilities_combines_mounts() {
3088 let registry = CapabilityRegistry::with_builtins();
3089
3090 let collected = collect_capabilities(
3093 &["sample_data".to_string(), "current_time".to_string()],
3094 ®istry,
3095 &test_ctx(),
3096 )
3097 .await;
3098
3099 assert_eq!(collected.mounts.len(), 1);
3100 assert!(
3102 collected
3103 .applied_ids
3104 .iter()
3105 .any(|id| id == "session_file_system")
3106 );
3107 assert!(collected.applied_ids.iter().any(|id| id == "sample_data"));
3108 assert!(collected.applied_ids.iter().any(|id| id == "current_time"));
3109 }
3110
3111 #[test]
3112 fn test_sample_data_capability() {
3113 let registry = CapabilityRegistry::with_builtins();
3114 let cap = registry.get("sample_data").unwrap();
3115
3116 assert_eq!(cap.id(), "sample_data");
3117 assert_eq!(cap.name(), "Sample Data");
3118 assert_eq!(cap.status(), CapabilityStatus::Available);
3119
3120 assert!(cap.system_prompt_addition().is_some());
3122 assert!(cap.tools().is_empty());
3123
3124 assert!(!cap.mounts().is_empty());
3126 }
3127
3128 #[test]
3133 fn test_resolve_dependencies_empty() {
3134 let registry = CapabilityRegistry::with_builtins();
3135
3136 let resolved = resolve_dependencies(&[], ®istry).unwrap();
3137
3138 assert!(resolved.resolved_ids.is_empty());
3139 assert!(resolved.added_as_dependencies.is_empty());
3140 assert!(resolved.user_selected.is_empty());
3141 }
3142
3143 #[test]
3144 fn test_resolve_dependencies_no_deps() {
3145 let registry = CapabilityRegistry::with_builtins();
3146
3147 let resolved = resolve_dependencies(&["current_time".to_string()], ®istry).unwrap();
3149
3150 assert_eq!(resolved.resolved_ids, vec!["current_time"]);
3151 assert!(resolved.added_as_dependencies.is_empty());
3152 }
3153
3154 #[test]
3155 fn test_resolve_dependencies_with_deps() {
3156 let registry = CapabilityRegistry::with_builtins();
3157
3158 let resolved = resolve_dependencies(&["sample_data".to_string()], ®istry).unwrap();
3160
3161 assert_eq!(resolved.resolved_ids.len(), 2);
3163 let fs_pos = resolved
3164 .resolved_ids
3165 .iter()
3166 .position(|id| id == "session_file_system")
3167 .unwrap();
3168 let sd_pos = resolved
3169 .resolved_ids
3170 .iter()
3171 .position(|id| id == "sample_data")
3172 .unwrap();
3173 assert!(fs_pos < sd_pos, "FileSystem should come before SampleData");
3174
3175 assert_eq!(resolved.added_as_dependencies, vec!["session_file_system"]);
3177 }
3178
3179 #[test]
3180 fn test_resolve_dependencies_already_selected() {
3181 let registry = CapabilityRegistry::with_builtins();
3182
3183 let resolved = resolve_dependencies(
3185 &["session_file_system".to_string(), "sample_data".to_string()],
3186 ®istry,
3187 )
3188 .unwrap();
3189
3190 assert_eq!(resolved.resolved_ids.len(), 2);
3191 assert!(resolved.added_as_dependencies.is_empty());
3193 }
3194
3195 #[test]
3196 fn test_resolve_dependencies_preserves_order() {
3197 let registry = CapabilityRegistry::with_builtins();
3198
3199 let resolved =
3201 resolve_dependencies(&["current_time".to_string(), "noop".to_string()], ®istry)
3202 .unwrap();
3203
3204 assert_eq!(resolved.resolved_ids, vec!["current_time", "noop"]);
3205 }
3206
3207 #[test]
3208 fn test_resolve_dependencies_unknown_capability() {
3209 let registry = CapabilityRegistry::with_builtins();
3210
3211 let resolved =
3213 resolve_dependencies(&["unknown_capability".to_string()], ®istry).unwrap();
3214
3215 assert!(resolved.resolved_ids.is_empty());
3216 }
3217
3218 #[test]
3219 fn test_get_dependencies() {
3220 let registry = CapabilityRegistry::with_builtins();
3221
3222 let deps = get_dependencies("sample_data", ®istry);
3224 assert_eq!(deps, vec!["session_file_system"]);
3225
3226 let deps = get_dependencies("current_time", ®istry);
3228 assert!(deps.is_empty());
3229
3230 let deps = get_dependencies("unknown", ®istry);
3232 assert!(deps.is_empty());
3233 }
3234
3235 #[test]
3236 fn test_sample_data_has_dependency() {
3237 let registry = CapabilityRegistry::with_builtins();
3238 let cap = registry.get("sample_data").unwrap();
3239
3240 let deps = cap.dependencies();
3241 assert_eq!(deps.len(), 1);
3242 assert_eq!(deps[0], "session_file_system");
3243 }
3244
3245 #[test]
3246 fn test_noop_has_no_dependencies() {
3247 let registry = CapabilityRegistry::with_builtins();
3248 let cap = registry.get("noop").unwrap();
3249
3250 assert!(cap.dependencies().is_empty());
3251 }
3252
3253 #[test]
3257 fn test_circular_dependency_error() {
3258 struct CapA;
3260 struct CapB;
3261
3262 impl Capability for CapA {
3263 fn id(&self) -> &str {
3264 "test_cap_a"
3265 }
3266 fn name(&self) -> &str {
3267 "Test A"
3268 }
3269 fn description(&self) -> &str {
3270 "Test capability A"
3271 }
3272 fn dependencies(&self) -> Vec<&'static str> {
3273 vec!["test_cap_b"]
3274 }
3275 }
3276
3277 impl Capability for CapB {
3278 fn id(&self) -> &str {
3279 "test_cap_b"
3280 }
3281 fn name(&self) -> &str {
3282 "Test B"
3283 }
3284 fn description(&self) -> &str {
3285 "Test capability B"
3286 }
3287 fn dependencies(&self) -> Vec<&'static str> {
3288 vec!["test_cap_a"]
3289 }
3290 }
3291
3292 let mut registry = CapabilityRegistry::new();
3293 registry.register(CapA);
3294 registry.register(CapB);
3295
3296 let result = resolve_dependencies(&["test_cap_a".to_string()], ®istry);
3297
3298 assert!(result.is_err());
3299 match result.unwrap_err() {
3300 DependencyError::CircularDependency { capability_id, .. } => {
3301 assert_eq!(capability_id, "test_cap_a");
3302 }
3303 _ => panic!("Expected CircularDependency error"),
3304 }
3305 }
3306
3307 use crate::message_filter::{MessageFilter, MessageFilterProvider, MessageQuery};
3312
3313 struct FilterTestCapability {
3315 priority: i32,
3316 }
3317
3318 impl Capability for FilterTestCapability {
3319 fn id(&self) -> &str {
3320 "filter_test"
3321 }
3322 fn name(&self) -> &str {
3323 "Filter Test"
3324 }
3325 fn description(&self) -> &str {
3326 "Test capability with message filter"
3327 }
3328 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3329 Some(Arc::new(FilterTestProvider {
3330 priority: self.priority,
3331 }))
3332 }
3333 }
3334
3335 struct FilterTestProvider {
3336 priority: i32,
3337 }
3338
3339 impl MessageFilterProvider for FilterTestProvider {
3340 fn apply_filters(&self, query: &mut MessageQuery, config: &serde_json::Value) {
3341 if let Some(search) = config.get("search").and_then(|v| v.as_str()) {
3343 query
3344 .filters
3345 .push(MessageFilter::Search(search.to_string()));
3346 }
3347 }
3348
3349 fn priority(&self) -> i32 {
3350 self.priority
3351 }
3352 }
3353
3354 #[tokio::test]
3355 async fn test_collect_capabilities_with_configs_no_filter_providers() {
3356 let registry = CapabilityRegistry::with_builtins();
3357 let configs = vec![AgentCapabilityConfig {
3358 capability_ref: CapabilityId::new("current_time"),
3359 config: serde_json::json!({}),
3360 }];
3361
3362 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3363
3364 assert!(collected.message_filter_providers.is_empty());
3365 assert!(!collected.has_message_filters());
3366 }
3367
3368 #[tokio::test]
3369 async fn test_collect_capabilities_with_configs_with_filter_provider() {
3370 let mut registry = CapabilityRegistry::new();
3371 registry.register(FilterTestCapability { priority: 0 });
3372
3373 let configs = vec![AgentCapabilityConfig {
3374 capability_ref: CapabilityId::new("filter_test"),
3375 config: serde_json::json!({ "search": "hello" }),
3376 }];
3377
3378 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3379
3380 assert_eq!(collected.message_filter_providers.len(), 1);
3381 assert!(collected.has_message_filters());
3382 }
3383
3384 #[tokio::test]
3385 async fn test_collect_capabilities_with_configs_filter_priority_order() {
3386 struct HighPriorityCapability;
3388 struct LowPriorityCapability;
3389
3390 impl Capability for HighPriorityCapability {
3391 fn id(&self) -> &str {
3392 "high_priority"
3393 }
3394 fn name(&self) -> &str {
3395 "High Priority"
3396 }
3397 fn description(&self) -> &str {
3398 "Test"
3399 }
3400 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3401 Some(Arc::new(FilterTestProvider { priority: 10 }))
3402 }
3403 }
3404
3405 impl Capability for LowPriorityCapability {
3406 fn id(&self) -> &str {
3407 "low_priority"
3408 }
3409 fn name(&self) -> &str {
3410 "Low Priority"
3411 }
3412 fn description(&self) -> &str {
3413 "Test"
3414 }
3415 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3416 Some(Arc::new(FilterTestProvider { priority: -5 }))
3417 }
3418 }
3419
3420 let mut registry = CapabilityRegistry::new();
3421 registry.register(HighPriorityCapability);
3422 registry.register(LowPriorityCapability);
3423
3424 let configs = vec![
3426 AgentCapabilityConfig {
3427 capability_ref: CapabilityId::new("high_priority"),
3428 config: serde_json::json!({}),
3429 },
3430 AgentCapabilityConfig {
3431 capability_ref: CapabilityId::new("low_priority"),
3432 config: serde_json::json!({}),
3433 },
3434 ];
3435
3436 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3437
3438 assert_eq!(collected.message_filter_providers.len(), 2);
3440 assert_eq!(collected.message_filter_providers[0].0.priority(), -5);
3441 assert_eq!(collected.message_filter_providers[1].0.priority(), 10);
3442 }
3443
3444 #[tokio::test]
3445 async fn test_collected_capabilities_apply_message_filters() {
3446 let mut registry = CapabilityRegistry::new();
3447 registry.register(FilterTestCapability { priority: 0 });
3448
3449 let configs = vec![AgentCapabilityConfig {
3450 capability_ref: CapabilityId::new("filter_test"),
3451 config: serde_json::json!({ "search": "test_query" }),
3452 }];
3453
3454 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3455
3456 let session_id: SessionId = Uuid::now_v7().into();
3458 let mut query = MessageQuery::new(session_id);
3459
3460 collected.apply_message_filters(&mut query);
3461
3462 assert_eq!(query.filters.len(), 1);
3464 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3465 }
3466
3467 #[tokio::test]
3468 async fn test_collected_capabilities_apply_multiple_filters_in_priority_order() {
3469 struct SearchCapability {
3470 id: &'static str,
3471 search_term: &'static str,
3472 priority: i32,
3473 }
3474
3475 struct SearchProvider {
3476 search_term: &'static str,
3477 priority: i32,
3478 }
3479
3480 impl MessageFilterProvider for SearchProvider {
3481 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3482 query
3483 .filters
3484 .push(MessageFilter::Search(self.search_term.to_string()));
3485 }
3486
3487 fn priority(&self) -> i32 {
3488 self.priority
3489 }
3490 }
3491
3492 impl Capability for SearchCapability {
3493 fn id(&self) -> &str {
3494 self.id
3495 }
3496 fn name(&self) -> &str {
3497 "Search"
3498 }
3499 fn description(&self) -> &str {
3500 "Test"
3501 }
3502 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3503 Some(Arc::new(SearchProvider {
3504 search_term: self.search_term,
3505 priority: self.priority,
3506 }))
3507 }
3508 }
3509
3510 let mut registry = CapabilityRegistry::new();
3511 registry.register(SearchCapability {
3512 id: "cap_a",
3513 search_term: "alpha",
3514 priority: 5,
3515 });
3516 registry.register(SearchCapability {
3517 id: "cap_b",
3518 search_term: "beta",
3519 priority: 1,
3520 });
3521 registry.register(SearchCapability {
3522 id: "cap_c",
3523 search_term: "gamma",
3524 priority: 10,
3525 });
3526
3527 let configs = vec![
3528 AgentCapabilityConfig {
3529 capability_ref: CapabilityId::new("cap_a"),
3530 config: serde_json::json!({}),
3531 },
3532 AgentCapabilityConfig {
3533 capability_ref: CapabilityId::new("cap_b"),
3534 config: serde_json::json!({}),
3535 },
3536 AgentCapabilityConfig {
3537 capability_ref: CapabilityId::new("cap_c"),
3538 config: serde_json::json!({}),
3539 },
3540 ];
3541
3542 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3543
3544 let session_id: SessionId = Uuid::now_v7().into();
3545 let mut query = MessageQuery::new(session_id);
3546
3547 collected.apply_message_filters(&mut query);
3548
3549 assert_eq!(query.filters.len(), 3);
3551 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3552 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3553 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3554 }
3555
3556 #[test]
3557 fn test_capability_without_message_filter_returns_none() {
3558 let registry = CapabilityRegistry::with_builtins();
3559
3560 let noop = registry.get("noop").unwrap();
3561 assert!(noop.message_filter_provider().is_none());
3562
3563 let current_time = registry.get("current_time").unwrap();
3564 assert!(current_time.message_filter_provider().is_none());
3565 }
3566
3567 #[tokio::test]
3568 async fn test_collect_capabilities_preserves_config_for_filter_provider() {
3569 let mut registry = CapabilityRegistry::new();
3570 registry.register(FilterTestCapability { priority: 0 });
3571
3572 let test_config = serde_json::json!({
3573 "search": "custom_search",
3574 "extra_field": 42
3575 });
3576
3577 let configs = vec![AgentCapabilityConfig {
3578 capability_ref: CapabilityId::new("filter_test"),
3579 config: test_config.clone(),
3580 }];
3581
3582 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
3583
3584 assert_eq!(collected.message_filter_providers.len(), 1);
3586 let (_, stored_config) = &collected.message_filter_providers[0];
3587 assert_eq!(*stored_config, test_config);
3588 }
3589
3590 #[test]
3595 fn test_collect_message_filters_only_collects_filters() {
3596 let mut registry = CapabilityRegistry::new();
3597 registry.register(FilterTestCapability { priority: 0 });
3598
3599 let configs = vec![AgentCapabilityConfig {
3600 capability_ref: CapabilityId::new("filter_test"),
3601 config: serde_json::json!({ "search": "test_query" }),
3602 }];
3603
3604 let collected = collect_message_filters_only(&configs, ®istry);
3605
3606 let session_id: SessionId = Uuid::now_v7().into();
3607 let mut query = MessageQuery::new(session_id);
3608 collected.apply_message_filters(&mut query);
3609
3610 assert_eq!(query.filters.len(), 1);
3611 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "test_query"));
3612 }
3613
3614 #[test]
3615 fn test_message_filter_config_injects_compaction_active_for_infinity_context() {
3616 let base = serde_json::json!({ "context_budget_tokens": 1000 });
3617
3618 let with = message_filter_config_for(INFINITY_CONTEXT_CAPABILITY_ID, &base, true);
3620 assert_eq!(with["compaction_active"], serde_json::json!(true));
3621 assert_eq!(with["context_budget_tokens"], serde_json::json!(1000));
3622
3623 let without = message_filter_config_for(INFINITY_CONTEXT_CAPABILITY_ID, &base, false);
3624 assert!(without.get("compaction_active").is_none());
3625
3626 let other = message_filter_config_for("other", &base, true);
3628 assert!(other.get("compaction_active").is_none());
3629
3630 let null_base = message_filter_config_for(
3632 INFINITY_CONTEXT_CAPABILITY_ID,
3633 &serde_json::Value::Null,
3634 true,
3635 );
3636 assert_eq!(null_base["compaction_active"], serde_json::json!(true));
3637 }
3638
3639 #[test]
3640 fn test_infinity_context_defers_to_compaction_end_to_end() {
3641 use crate::message::Message;
3642
3643 let mut registry = CapabilityRegistry::new();
3644 registry.register(InfinityContextCapability);
3645 registry.register(CompactionCapability);
3646
3647 let tight = serde_json::json!({
3648 "context_budget_tokens": 1,
3649 "min_recent_messages": 1
3650 });
3651
3652 let solo = vec![AgentCapabilityConfig {
3654 capability_ref: CapabilityId::new(INFINITY_CONTEXT_CAPABILITY_ID),
3655 config: tight.clone(),
3656 }];
3657 let mut messages = vec![
3658 Message::user("task"),
3659 Message::assistant("old ".repeat(400)),
3660 Message::user("recent"),
3661 ];
3662 collect_message_filters_only(&solo, ®istry).apply_post_load_filters(&mut messages);
3663 assert!(
3664 messages
3665 .iter()
3666 .any(|m| m.text().is_some_and(|t| t.contains("NOT visible"))),
3667 "infinity context alone should trim and notice"
3668 );
3669
3670 let both = vec![
3672 AgentCapabilityConfig {
3673 capability_ref: CapabilityId::new(INFINITY_CONTEXT_CAPABILITY_ID),
3674 config: tight,
3675 },
3676 AgentCapabilityConfig {
3677 capability_ref: CapabilityId::new(COMPACTION_CAPABILITY_ID),
3678 config: serde_json::json!({}),
3679 },
3680 ];
3681 let mut messages = vec![
3682 Message::user("task"),
3683 Message::assistant("old ".repeat(400)),
3684 Message::user("recent"),
3685 ];
3686 collect_message_filters_only(&both, ®istry).apply_post_load_filters(&mut messages);
3687 assert_eq!(messages.len(), 3, "compaction owns reduction; no eviction");
3688 assert!(
3689 messages
3690 .iter()
3691 .all(|m| !m.text().is_some_and(|t| t.contains("NOT visible"))),
3692 "no hidden-history notice when compaction is the active reducer"
3693 );
3694 }
3695
3696 #[test]
3697 fn test_compaction_is_enabled_detects_compaction() {
3698 let mut registry = CapabilityRegistry::new();
3699 registry.register(CompactionCapability);
3700
3701 let with_compaction = vec![AgentCapabilityConfig {
3702 capability_ref: CapabilityId::new(COMPACTION_CAPABILITY_ID),
3703 config: serde_json::json!({}),
3704 }];
3705 assert!(compaction_is_enabled(&with_compaction, ®istry));
3706
3707 let without = vec![AgentCapabilityConfig {
3708 capability_ref: CapabilityId::new("current_time"),
3709 config: serde_json::json!({}),
3710 }];
3711 assert!(!compaction_is_enabled(&without, ®istry));
3712 }
3713
3714 #[test]
3715 fn test_collect_message_filters_only_skips_unknown_capabilities() {
3716 let registry = CapabilityRegistry::new();
3717
3718 let configs = vec![AgentCapabilityConfig {
3719 capability_ref: CapabilityId::new("nonexistent"),
3720 config: serde_json::json!({}),
3721 }];
3722
3723 let collected = collect_message_filters_only(&configs, ®istry);
3724 assert!(collected.message_filter_providers.is_empty());
3725 }
3726
3727 #[test]
3728 fn test_collect_message_filters_only_preserves_priority_order() {
3729 struct PriorityFilterCap {
3730 id: &'static str,
3731 search_term: &'static str,
3732 priority: i32,
3733 }
3734
3735 struct PriorityFilterProvider {
3736 search_term: &'static str,
3737 priority: i32,
3738 }
3739
3740 impl Capability for PriorityFilterCap {
3741 fn id(&self) -> &str {
3742 self.id
3743 }
3744 fn name(&self) -> &str {
3745 self.id
3746 }
3747 fn description(&self) -> &str {
3748 "priority test"
3749 }
3750 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3751 Some(Arc::new(PriorityFilterProvider {
3752 search_term: self.search_term,
3753 priority: self.priority,
3754 }))
3755 }
3756 }
3757
3758 impl MessageFilterProvider for PriorityFilterProvider {
3759 fn apply_filters(&self, query: &mut MessageQuery, _config: &serde_json::Value) {
3760 query
3761 .filters
3762 .push(MessageFilter::Search(self.search_term.to_string()));
3763 }
3764 fn priority(&self) -> i32 {
3765 self.priority
3766 }
3767 }
3768
3769 let mut registry = CapabilityRegistry::new();
3770 registry.register(PriorityFilterCap {
3771 id: "gamma",
3772 search_term: "gamma",
3773 priority: 10,
3774 });
3775 registry.register(PriorityFilterCap {
3776 id: "alpha",
3777 search_term: "alpha",
3778 priority: 5,
3779 });
3780 registry.register(PriorityFilterCap {
3781 id: "beta",
3782 search_term: "beta",
3783 priority: 1,
3784 });
3785
3786 let configs = vec![
3787 AgentCapabilityConfig {
3788 capability_ref: CapabilityId::new("gamma"),
3789 config: serde_json::json!({}),
3790 },
3791 AgentCapabilityConfig {
3792 capability_ref: CapabilityId::new("alpha"),
3793 config: serde_json::json!({}),
3794 },
3795 AgentCapabilityConfig {
3796 capability_ref: CapabilityId::new("beta"),
3797 config: serde_json::json!({}),
3798 },
3799 ];
3800
3801 let collected = collect_message_filters_only(&configs, ®istry);
3802
3803 let session_id: SessionId = Uuid::now_v7().into();
3804 let mut query = MessageQuery::new(session_id);
3805 collected.apply_message_filters(&mut query);
3806
3807 assert_eq!(query.filters.len(), 3);
3809 assert!(matches!(&query.filters[0], MessageFilter::Search(s) if s == "beta"));
3810 assert!(matches!(&query.filters[1], MessageFilter::Search(s) if s == "alpha"));
3811 assert!(matches!(&query.filters[2], MessageFilter::Search(s) if s == "gamma"));
3812 }
3813
3814 #[test]
3815 fn test_collect_message_filters_only_post_load_invoked() {
3816 use crate::message::Message;
3817
3818 struct PostLoadCap;
3819 struct PostLoadProvider;
3820
3821 impl Capability for PostLoadCap {
3822 fn id(&self) -> &str {
3823 "post_load_test"
3824 }
3825 fn name(&self) -> &str {
3826 "PostLoad Test"
3827 }
3828 fn description(&self) -> &str {
3829 "test"
3830 }
3831 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
3832 Some(Arc::new(PostLoadProvider))
3833 }
3834 }
3835
3836 impl MessageFilterProvider for PostLoadProvider {
3837 fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
3838 fn priority(&self) -> i32 {
3839 0
3840 }
3841 fn post_load(&self, messages: &mut Vec<Message>, _config: &serde_json::Value) {
3842 messages.reverse();
3844 }
3845 }
3846
3847 let mut registry = CapabilityRegistry::new();
3848 registry.register(PostLoadCap);
3849
3850 let configs = vec![AgentCapabilityConfig {
3851 capability_ref: CapabilityId::new("post_load_test"),
3852 config: serde_json::json!({}),
3853 }];
3854
3855 let collected = collect_message_filters_only(&configs, ®istry);
3856
3857 let mut messages = vec![Message::user("first"), Message::user("second")];
3858 collected.apply_post_load_filters(&mut messages);
3859
3860 assert_eq!(messages[0].text(), Some("second"));
3862 assert_eq!(messages[1].text(), Some("first"));
3863 }
3864
3865 #[test]
3866 fn test_collect_model_view_providers_respects_compaction_capability_boundary() {
3867 use crate::tool_types::ToolCall;
3868
3869 fn tool_heavy_messages() -> Vec<Message> {
3870 let mut messages = vec![Message::user("inspect files repeatedly")];
3871 for index in 0..9 {
3872 let call_id = format!("call_{index}");
3873 messages.push(Message::assistant_with_tools(
3874 "",
3875 vec![ToolCall {
3876 id: call_id.clone(),
3877 name: "read_file".to_string(),
3878 arguments: serde_json::json!({"path": "/workspace/src/lib.rs"}),
3879 }],
3880 ));
3881 messages.push(Message::tool_result(
3882 call_id,
3883 Some(serde_json::json!({
3884 "path": "/workspace/src/lib.rs",
3885 "content": format!("{}{}", "large file line\n".repeat(1000), index),
3886 "total_lines": 1000,
3887 "lines_shown": {"start": 1, "end": 1000},
3888 "truncated": false
3889 })),
3890 None,
3891 ));
3892 }
3893 messages
3894 }
3895
3896 fn first_tool_result_is_masked(messages: &[Message]) -> bool {
3897 messages[2]
3898 .tool_result_content()
3899 .and_then(|result| result.result.as_ref())
3900 .and_then(|result| result.get("masked"))
3901 .and_then(|masked| masked.as_bool())
3902 .unwrap_or(false)
3903 }
3904
3905 let mut registry = CapabilityRegistry::new();
3906 registry.register(CompactionCapability);
3907 let context = ModelViewContext {
3908 session_id: SessionId::new(),
3909 prior_usage: None,
3910 };
3911
3912 let no_compaction = collect_model_view_providers(&[], ®istry, None);
3913 let unmasked = no_compaction.apply_model_view(tool_heavy_messages(), &context);
3914 assert!(!first_tool_result_is_masked(&unmasked));
3915
3916 let compaction = collect_model_view_providers(
3917 &[AgentCapabilityConfig {
3918 capability_ref: CapabilityId::new(COMPACTION_CAPABILITY_ID),
3919 config: serde_json::json!({}),
3920 }],
3921 ®istry,
3922 None,
3923 );
3924 let masked = compaction.apply_model_view(tool_heavy_messages(), &context);
3925 assert!(first_tool_result_is_masked(&masked));
3926 let last_tool = masked.last().unwrap().tool_result_content().unwrap();
3927 assert!(last_tool.result.as_ref().unwrap().get("content").is_some());
3928 }
3929
3930 struct DelegatingFilterCap {
3933 id: &'static str,
3934 inner: std::sync::Arc<InnerFilterCap>,
3935 }
3936 struct InnerFilterCap;
3937
3938 impl Capability for InnerFilterCap {
3939 fn id(&self) -> &str {
3940 "inner_filter"
3941 }
3942 fn name(&self) -> &str {
3943 "Inner Filter"
3944 }
3945 fn description(&self) -> &str {
3946 "inner"
3947 }
3948 fn message_filter_provider(&self) -> Option<std::sync::Arc<dyn MessageFilterProvider>> {
3949 Some(std::sync::Arc::new(SentinelFilter))
3950 }
3951 }
3952 struct SentinelFilter;
3953 impl MessageFilterProvider for SentinelFilter {
3954 fn apply_filters(&self, _query: &mut MessageQuery, _config: &serde_json::Value) {}
3955 }
3956 impl Capability for DelegatingFilterCap {
3957 fn id(&self) -> &str {
3958 self.id
3959 }
3960 fn name(&self) -> &str {
3961 "Delegating Filter"
3962 }
3963 fn description(&self) -> &str {
3964 "delegating"
3965 }
3966 fn message_filter_provider(&self) -> Option<std::sync::Arc<dyn MessageFilterProvider>> {
3967 None }
3969 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
3970 Some(&*self.inner)
3971 }
3972 }
3973
3974 #[test]
3975 fn test_collect_message_filters_only_honors_resolve_for_model_delegation() {
3976 let inner = std::sync::Arc::new(InnerFilterCap);
3977 let outer = DelegatingFilterCap {
3978 id: "delegating_filter",
3979 inner: inner.clone(),
3980 };
3981
3982 let mut registry = CapabilityRegistry::new();
3983 registry.register(outer);
3984
3985 let configs = vec![AgentCapabilityConfig {
3986 capability_ref: CapabilityId::new("delegating_filter"),
3987 config: serde_json::json!({}),
3988 }];
3989
3990 let collected = collect_message_filters_only(&configs, ®istry);
3993 assert_eq!(
3994 collected.message_filter_providers.len(),
3995 1,
3996 "provider from resolved inner capability must be collected"
3997 );
3998 }
3999
4000 struct DelegatingMvpCap {
4001 id: &'static str,
4002 inner: std::sync::Arc<InnerMvpCap>,
4003 }
4004 struct InnerMvpCap;
4005
4006 impl Capability for InnerMvpCap {
4007 fn id(&self) -> &str {
4008 "inner_mvp"
4009 }
4010 fn name(&self) -> &str {
4011 "Inner MVP"
4012 }
4013 fn description(&self) -> &str {
4014 "inner"
4015 }
4016 fn model_view_provider(
4017 &self,
4018 ) -> Option<std::sync::Arc<dyn crate::capabilities::ModelViewProvider>> {
4019 struct NoopMvp;
4021 impl crate::capabilities::ModelViewProvider for NoopMvp {
4022 fn apply_model_view(
4023 &self,
4024 messages: Vec<Message>,
4025 _config: &serde_json::Value,
4026 _context: &ModelViewContext<'_>,
4027 ) -> Vec<Message> {
4028 messages
4029 }
4030 }
4031 Some(std::sync::Arc::new(NoopMvp))
4032 }
4033 }
4034 impl Capability for DelegatingMvpCap {
4035 fn id(&self) -> &str {
4036 self.id
4037 }
4038 fn name(&self) -> &str {
4039 "Delegating MVP"
4040 }
4041 fn description(&self) -> &str {
4042 "delegating"
4043 }
4044 fn model_view_provider(
4045 &self,
4046 ) -> Option<std::sync::Arc<dyn crate::capabilities::ModelViewProvider>> {
4047 None }
4049 fn resolve_for_model(&self, _model: Option<&str>) -> Option<&dyn Capability> {
4050 Some(&*self.inner)
4051 }
4052 }
4053
4054 #[test]
4055 fn test_collect_model_view_providers_honors_resolve_for_model_delegation() {
4056 let inner = std::sync::Arc::new(InnerMvpCap);
4057 let outer = DelegatingMvpCap {
4058 id: "delegating_mvp",
4059 inner: inner.clone(),
4060 };
4061
4062 let mut registry = CapabilityRegistry::new();
4063 registry.register(outer);
4064
4065 let configs = vec![AgentCapabilityConfig {
4066 capability_ref: CapabilityId::new("delegating_mvp"),
4067 config: serde_json::json!({}),
4068 }];
4069
4070 let collected = collect_model_view_providers(&configs, ®istry, None);
4073 assert_eq!(
4074 collected.model_view_providers.len(),
4075 1,
4076 "provider from resolved inner capability must be collected"
4077 );
4078 }
4079
4080 #[tokio::test]
4090 async fn test_bashkit_shell_capability_produces_bash_tool() {
4091 let registry = CapabilityRegistry::with_builtins();
4092 let collected =
4093 collect_capabilities(&["bashkit_shell".to_string()], ®istry, &test_ctx()).await;
4094
4095 let tool_names: Vec<&str> = collected
4096 .tool_definitions
4097 .iter()
4098 .map(|t| t.name())
4099 .collect();
4100 assert!(
4101 tool_names.contains(&"bash"),
4102 "bashkit_shell capability must produce 'bash' tool, got: {:?}",
4103 tool_names
4104 );
4105 assert!(
4106 !collected.tools.is_empty(),
4107 "bashkit_shell must provide tool implementations"
4108 );
4109 }
4110
4111 #[tokio::test]
4112 async fn test_generic_harness_capability_set_produces_bash_tool() {
4113 let generic_harness_caps = vec![
4116 "session_file_system".to_string(),
4117 "bashkit_shell".to_string(),
4118 "web_fetch".to_string(),
4119 "session_storage".to_string(),
4120 "session".to_string(),
4121 "agent_instructions".to_string(),
4122 "skills".to_string(),
4123 "infinity_context".to_string(),
4124 "auto_tool_search".to_string(),
4125 ];
4126
4127 let registry = CapabilityRegistry::with_builtins();
4128 let collected = collect_capabilities(&generic_harness_caps, ®istry, &test_ctx()).await;
4129
4130 let tool_names: Vec<&str> = collected
4131 .tool_definitions
4132 .iter()
4133 .map(|t| t.name())
4134 .collect();
4135 assert!(
4136 tool_names.contains(&"bash"),
4137 "Generic Harness capabilities must produce 'bash' tool, got: {:?}",
4138 tool_names
4139 );
4140 }
4141
4142 #[tokio::test]
4143 async fn test_collect_capabilities_tool_count_matches_definitions() {
4144 let registry = CapabilityRegistry::with_builtins();
4147 let collected =
4148 collect_capabilities(&["bashkit_shell".to_string()], ®istry, &test_ctx()).await;
4149
4150 assert_eq!(
4151 collected.tools.len(),
4152 collected.tool_definitions.len(),
4153 "tool implementations ({}) must match tool definitions ({})",
4154 collected.tools.len(),
4155 collected.tool_definitions.len(),
4156 );
4157 }
4158
4159 #[tokio::test]
4163 async fn test_collect_capabilities_resolves_dependencies() {
4164 let registry = CapabilityRegistry::with_builtins();
4167 let collected =
4168 collect_capabilities(&["sample_data".to_string()], ®istry, &test_ctx()).await;
4169
4170 assert!(
4172 collected
4173 .applied_ids
4174 .iter()
4175 .any(|id| id == "session_file_system"),
4176 "collect_capabilities must apply session_file_system as a dependency; applied_ids: {:?}",
4177 collected.applied_ids
4178 );
4179
4180 let tool_names: Vec<&str> = collected
4181 .tool_definitions
4182 .iter()
4183 .map(|t| t.name())
4184 .collect();
4185
4186 assert!(
4188 tool_names.contains(&"read_file") && tool_names.contains(&"write_file"),
4189 "collect_capabilities must resolve dependencies and include dependency tools, got: {:?}",
4190 tool_names
4191 );
4192
4193 assert_eq!(
4195 collected.tools.len(),
4196 collected.tool_definitions.len(),
4197 "dependency-added tools must have implementations, not just definitions"
4198 );
4199 }
4200
4201 #[test]
4202 fn test_defaults_do_not_include_bash() {
4203 let registry = crate::ToolRegistry::with_defaults();
4206 assert!(
4207 !registry.has("bash"),
4208 "with_defaults() must not include 'bash' — it comes from bashkit_shell capability"
4209 );
4210 }
4211
4212 #[tokio::test]
4219 async fn test_background_execution_auto_activates_with_bashkit_shell() {
4220 let registry = CapabilityRegistry::with_builtins();
4221 let collected =
4222 collect_capabilities(&["bashkit_shell".to_string()], ®istry, &test_ctx()).await;
4223
4224 let tool_names: Vec<&str> = collected
4225 .tool_definitions
4226 .iter()
4227 .map(|t| t.name())
4228 .collect();
4229 assert!(
4230 tool_names.contains(&"spawn_background"),
4231 "spawn_background must be auto-activated when bashkit_shell (a \
4232 background-capable tool) is in the agent's capability set; got: {:?}",
4233 tool_names
4234 );
4235 assert!(
4236 collected
4237 .applied_ids
4238 .iter()
4239 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
4240 "background_execution must be in applied_ids when auto-activated; \
4241 got: {:?}",
4242 collected.applied_ids
4243 );
4244
4245 assert!(
4247 collected
4248 .tools
4249 .iter()
4250 .any(|t| t.name() == "spawn_background"),
4251 "spawn_background tool implementation must be present alongside the \
4252 definition (lockstep contract)"
4253 );
4254 }
4255
4256 #[tokio::test]
4259 async fn test_background_execution_does_not_auto_activate_without_hint() {
4260 let registry = CapabilityRegistry::with_builtins();
4261 let collected =
4263 collect_capabilities(&["current_time".to_string()], ®istry, &test_ctx()).await;
4264
4265 let tool_names: Vec<&str> = collected
4266 .tool_definitions
4267 .iter()
4268 .map(|t| t.name())
4269 .collect();
4270 assert!(
4271 !tool_names.contains(&"spawn_background"),
4272 "spawn_background must NOT be activated without a background-capable \
4273 tool; got: {:?}",
4274 tool_names
4275 );
4276 assert!(
4277 !collected
4278 .applied_ids
4279 .iter()
4280 .any(|id| id == BACKGROUND_EXECUTION_CAPABILITY_ID),
4281 "background_execution must not appear in applied_ids when no \
4282 background-capable tool is present; got: {:?}",
4283 collected.applied_ids
4284 );
4285 }
4286
4287 #[tokio::test]
4291 async fn test_background_execution_explicit_selection_is_idempotent() {
4292 let registry = CapabilityRegistry::with_builtins();
4293 let collected = collect_capabilities(
4294 &[
4295 "bashkit_shell".to_string(),
4296 BACKGROUND_EXECUTION_CAPABILITY_ID.to_string(),
4297 ],
4298 ®istry,
4299 &test_ctx(),
4300 )
4301 .await;
4302
4303 let spawn_background_count = collected
4304 .tool_definitions
4305 .iter()
4306 .filter(|t| t.name() == "spawn_background")
4307 .count();
4308 assert_eq!(
4309 spawn_background_count, 1,
4310 "spawn_background must appear exactly once even when \
4311 background_execution is selected explicitly alongside a \
4312 background-capable tool"
4313 );
4314 let applied_count = collected
4315 .applied_ids
4316 .iter()
4317 .filter(|id| id.as_str() == BACKGROUND_EXECUTION_CAPABILITY_ID)
4318 .count();
4319 assert_eq!(
4320 applied_count, 1,
4321 "background_execution must appear exactly once in applied_ids"
4322 );
4323 }
4324
4325 #[test]
4330 fn test_defaults_do_not_include_spawn_background() {
4331 let registry = crate::ToolRegistry::with_defaults();
4332 assert!(
4333 !registry.has("spawn_background"),
4334 "with_defaults() must not include 'spawn_background' — it comes \
4335 from the background_execution capability (EVE-501)"
4336 );
4337 }
4338
4339 #[test]
4344 fn test_capability_features_default_empty() {
4345 let registry = CapabilityRegistry::with_builtins();
4346
4347 let noop = registry.get("noop").unwrap();
4349 assert!(noop.features().is_empty());
4350
4351 let current_time = registry.get("current_time").unwrap();
4352 assert!(current_time.features().is_empty());
4353 }
4354
4355 #[test]
4356 fn test_file_system_capability_features() {
4357 let registry = CapabilityRegistry::with_builtins();
4358
4359 let fs = registry.get("session_file_system").unwrap();
4360 assert_eq!(fs.features(), vec!["file_system"]);
4361 }
4362
4363 #[test]
4364 fn test_bashkit_shell_capability_features() {
4365 let registry = CapabilityRegistry::with_builtins();
4366
4367 let bash = registry.get("bashkit_shell").unwrap();
4368 assert_eq!(bash.features(), vec!["file_system"]);
4369 }
4370
4371 #[test]
4372 fn test_alias_resolves_to_canonical_capability() {
4373 let registry = CapabilityRegistry::with_builtins();
4374
4375 let via_alias = registry.get("virtual_bash").unwrap();
4377 assert_eq!(via_alias.id(), "bashkit_shell");
4378 assert!(registry.has("virtual_bash"));
4379 assert_eq!(registry.canonical_id("virtual_bash"), Some("bashkit_shell"));
4380 assert_eq!(
4381 registry.canonical_id("bashkit_shell"),
4382 Some("bashkit_shell")
4383 );
4384 assert_eq!(registry.canonical_id("nonexistent"), None);
4385 }
4386
4387 #[test]
4388 fn test_alias_dedupes_with_canonical_in_dependency_resolution() {
4389 let registry = CapabilityRegistry::with_builtins();
4390
4391 let resolved = resolve_dependencies(
4394 &["virtual_bash".to_string(), "bashkit_shell".to_string()],
4395 ®istry,
4396 )
4397 .unwrap();
4398 let bash_ids: Vec<_> = resolved
4399 .resolved_ids
4400 .iter()
4401 .filter(|id| id.as_str() == "bashkit_shell" || id.as_str() == "virtual_bash")
4402 .collect();
4403 assert_eq!(bash_ids, vec!["bashkit_shell"]);
4404 assert!(
4406 !resolved
4407 .added_as_dependencies
4408 .contains(&"bashkit_shell".to_string())
4409 );
4410 }
4411
4412 #[test]
4413 fn test_alias_preserves_explicit_config_in_resolution() {
4414 let registry = CapabilityRegistry::with_builtins();
4415
4416 let configs = vec![AgentCapabilityConfig::with_config(
4417 "virtual_bash".to_string(),
4418 serde_json::json!({"key": "value"}),
4419 )];
4420 let resolved = resolve_capability_configs(&configs, ®istry).unwrap();
4421 let bash = resolved
4422 .iter()
4423 .find(|c| c.capability_id() == "bashkit_shell")
4424 .expect("alias must resolve to canonical bashkit_shell config");
4425 assert_eq!(bash.config, serde_json::json!({"key": "value"}));
4426 }
4427
4428 #[test]
4429 fn test_unregister_by_alias_removes_capability_and_aliases() {
4430 let mut registry = CapabilityRegistry::with_builtins();
4431
4432 assert!(registry.unregister("virtual_bash").is_some());
4433 assert!(!registry.has("bashkit_shell"));
4434 assert!(!registry.has("virtual_bash"));
4435 }
4436
4437 #[test]
4438 fn test_session_storage_capability_features() {
4439 let registry = CapabilityRegistry::with_builtins();
4440
4441 let storage = registry.get("session_storage").unwrap();
4442 let features = storage.features();
4443 assert!(features.contains(&"secrets"));
4444 assert!(features.contains(&"key_value"));
4445 }
4446
4447 #[test]
4448 fn test_session_schedule_capability_features() {
4449 let registry = CapabilityRegistry::with_builtins();
4450
4451 let schedule = registry.get("session_schedule").unwrap();
4452 assert_eq!(schedule.features(), vec!["schedules"]);
4453 }
4454
4455 #[test]
4456 fn test_session_sql_database_capability_features() {
4457 let registry = CapabilityRegistry::with_builtins();
4458
4459 let sql = registry.get("session_sql_database").unwrap();
4460 assert_eq!(sql.features(), vec!["sql_database"]);
4461 }
4462
4463 #[test]
4464 fn test_sample_data_capability_features() {
4465 let registry = CapabilityRegistry::with_builtins();
4466
4467 let sample = registry.get("sample_data").unwrap();
4468 assert_eq!(sample.features(), vec!["file_system"]);
4469 }
4470
4471 #[test]
4472 fn test_compute_features_empty() {
4473 let registry = CapabilityRegistry::with_builtins();
4474
4475 let features = compute_features(&[], ®istry);
4476 assert!(features.is_empty());
4477 }
4478
4479 #[test]
4480 fn test_compute_features_single_capability() {
4481 let registry = CapabilityRegistry::with_builtins();
4482
4483 let features = compute_features(&["session_schedule".to_string()], ®istry);
4484 assert_eq!(features, vec!["schedules"]);
4485 }
4486
4487 #[test]
4488 fn test_compute_features_multiple_capabilities() {
4489 let registry = CapabilityRegistry::with_builtins();
4490
4491 let features = compute_features(
4492 &[
4493 "session_file_system".to_string(),
4494 "session_storage".to_string(),
4495 "session_schedule".to_string(),
4496 ],
4497 ®istry,
4498 );
4499 assert!(features.contains(&"file_system".to_string()));
4500 assert!(features.contains(&"secrets".to_string()));
4501 assert!(features.contains(&"key_value".to_string()));
4502 assert!(features.contains(&"schedules".to_string()));
4503 }
4504
4505 #[test]
4506 fn test_compute_features_deduplicates() {
4507 let registry = CapabilityRegistry::with_builtins();
4508
4509 let features = compute_features(
4511 &[
4512 "session_file_system".to_string(),
4513 "bashkit_shell".to_string(),
4514 ],
4515 ®istry,
4516 );
4517 let file_system_count = features.iter().filter(|f| *f == "file_system").count();
4518 assert_eq!(file_system_count, 1, "file_system should appear only once");
4519 }
4520
4521 #[test]
4522 fn test_compute_features_includes_dependency_features() {
4523 let registry = CapabilityRegistry::with_builtins();
4524
4525 let features = compute_features(&["bashkit_shell".to_string()], ®istry);
4527 assert!(features.contains(&"file_system".to_string()));
4528 }
4529
4530 #[test]
4531 fn test_compute_features_generic_harness_set() {
4532 let registry = CapabilityRegistry::with_builtins();
4533
4534 let features = compute_features(
4536 &[
4537 "session_file_system".to_string(),
4538 "bashkit_shell".to_string(),
4539 "session_storage".to_string(),
4540 "session".to_string(),
4541 "session_schedule".to_string(),
4542 ],
4543 ®istry,
4544 );
4545 assert!(features.contains(&"file_system".to_string()));
4546 assert!(features.contains(&"secrets".to_string()));
4547 assert!(features.contains(&"key_value".to_string()));
4548 assert!(features.contains(&"schedules".to_string()));
4549 }
4550
4551 #[test]
4552 fn test_compute_features_unknown_capability_ignored() {
4553 let registry = CapabilityRegistry::with_builtins();
4554
4555 let features = compute_features(
4556 &["unknown_cap".to_string(), "session_schedule".to_string()],
4557 ®istry,
4558 );
4559 assert_eq!(features, vec!["schedules"]);
4560 }
4561
4562 #[test]
4563 fn test_risk_level_ordering() {
4564 assert!(RiskLevel::Low < RiskLevel::Medium);
4565 assert!(RiskLevel::Medium < RiskLevel::High);
4566 }
4567
4568 #[test]
4569 fn test_risk_level_serde_roundtrip() {
4570 let high = RiskLevel::High;
4571 let json = serde_json::to_string(&high).unwrap();
4572 assert_eq!(json, "\"high\"");
4573 let back: RiskLevel = serde_json::from_str(&json).unwrap();
4574 assert_eq!(back, RiskLevel::High);
4575 }
4576
4577 #[test]
4578 fn test_capability_risk_levels() {
4579 let registry = CapabilityRegistry::with_builtins();
4580
4581 let bash = registry.get("bashkit_shell").unwrap();
4583 assert_eq!(bash.risk_level(), RiskLevel::High);
4584
4585 let fetch = registry.get("web_fetch").unwrap();
4587 assert_eq!(fetch.risk_level(), RiskLevel::High);
4588
4589 let noop = registry.get("noop").unwrap();
4591 assert_eq!(noop.risk_level(), RiskLevel::Low);
4592 }
4593
4594 #[tokio::test]
4599 async fn test_apply_capabilities_openai_tool_search() {
4600 let registry = CapabilityRegistry::with_builtins();
4601 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4602
4603 let applied = apply_capabilities(
4604 base_runtime_agent.clone(),
4605 &["openai_tool_search".to_string()],
4606 ®istry,
4607 &test_ctx(),
4608 )
4609 .await;
4610
4611 assert_eq!(
4613 applied.runtime_agent.system_prompt,
4614 base_runtime_agent.system_prompt
4615 );
4616 assert!(applied.tool_registry.is_empty());
4617 assert_eq!(applied.applied_ids, vec!["openai_tool_search"]);
4618
4619 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
4621 assert!(ts.enabled);
4622 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
4623 }
4624
4625 #[tokio::test]
4626 async fn test_apply_capabilities_openai_tool_search_with_other_capabilities() {
4627 let registry = CapabilityRegistry::with_builtins();
4628 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4629
4630 let applied = apply_capabilities(
4631 base_runtime_agent,
4632 &[
4633 "current_time".to_string(),
4634 "openai_tool_search".to_string(),
4635 "test_math".to_string(),
4636 ],
4637 ®istry,
4638 &test_ctx(),
4639 )
4640 .await;
4641
4642 assert!(applied.tool_registry.has("get_current_time"));
4644 assert!(applied.tool_registry.has("add"));
4645 assert!(applied.tool_registry.has("subtract"));
4646 assert!(applied.tool_registry.has("multiply"));
4647 assert!(applied.tool_registry.has("divide"));
4648
4649 let ts = applied.runtime_agent.tool_search.as_ref().unwrap();
4651 assert!(ts.enabled);
4652 assert_eq!(ts.threshold, DEFAULT_TOOL_SEARCH_THRESHOLD);
4653 }
4654
4655 #[tokio::test]
4656 async fn test_collect_capabilities_tool_search_custom_threshold() {
4657 let registry = CapabilityRegistry::with_builtins();
4658
4659 let configs = vec![AgentCapabilityConfig {
4660 capability_ref: CapabilityId::new("openai_tool_search"),
4661 config: serde_json::json!({"threshold": 5}),
4662 }];
4663
4664 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4665
4666 let ts = collected.tool_search.as_ref().unwrap();
4667 assert!(ts.enabled);
4668 assert_eq!(ts.threshold, 5);
4669 }
4670
4671 #[tokio::test]
4672 async fn test_collect_capabilities_auto_tool_search_resolves_to_generic_off_native() {
4673 let registry = CapabilityRegistry::with_builtins();
4674
4675 let configs = vec![
4676 AgentCapabilityConfig {
4677 capability_ref: CapabilityId::new("auto_tool_search"),
4678 config: serde_json::json!({"threshold": 2}),
4679 },
4680 AgentCapabilityConfig {
4681 capability_ref: CapabilityId::new("test_math"),
4682 config: serde_json::json!({}),
4683 },
4684 ];
4685
4686 let ctx = test_ctx().with_model("claude-3-5-haiku");
4690 let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
4691
4692 assert!(
4693 collected.tool_search.is_none(),
4694 "auto_tool_search must not set a hosted config on a non-native model"
4695 );
4696 assert!(
4697 collected
4698 .tools
4699 .iter()
4700 .any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
4701 "auto_tool_search must contribute the client-side tool_search tool"
4702 );
4703 assert!(
4704 !collected.tool_definition_hooks.is_empty(),
4705 "auto_tool_search must contribute a client-side deferral hook"
4706 );
4707
4708 let mut transformed = collected.tool_definitions.clone();
4709 for hook in &collected.tool_definition_hooks {
4710 transformed = hook.transform(transformed);
4711 }
4712 let add_tool = transformed
4713 .iter()
4714 .find(|tool| tool.name() == "add")
4715 .expect("test_math contributes add");
4716 assert!(
4717 add_tool.parameters().get("properties").is_none(),
4718 "generic auto_tool_search must honor the configured threshold"
4719 );
4720 }
4721
4722 #[tokio::test]
4723 async fn test_collect_capabilities_auto_tool_search_resolves_to_hosted_on_native() {
4724 let registry = CapabilityRegistry::with_builtins();
4725
4726 let configs = vec![AgentCapabilityConfig {
4727 capability_ref: CapabilityId::new("auto_tool_search"),
4728 config: serde_json::json!({"threshold": 7}),
4729 }];
4730
4731 let ctx = test_ctx().with_model("gpt-5.4");
4734 let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
4735
4736 let ts = collected
4737 .tool_search
4738 .as_ref()
4739 .expect("auto_tool_search must set a hosted config on a native model");
4740 assert!(ts.enabled);
4741 assert_eq!(ts.threshold, 7);
4742 assert!(
4743 !collected
4744 .tools
4745 .iter()
4746 .any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
4747 "hosted mechanism must not contribute the client-side tool_search tool"
4748 );
4749 assert!(
4750 collected.tool_definition_hooks.is_empty(),
4751 "hosted mechanism must not contribute a client-side deferral hook"
4752 );
4753 }
4754
4755 #[tokio::test]
4756 async fn test_collect_capabilities_auto_tool_search_resolves_to_hosted_on_anthropic() {
4757 let registry = CapabilityRegistry::with_builtins();
4758
4759 let configs = vec![AgentCapabilityConfig {
4760 capability_ref: CapabilityId::new("auto_tool_search"),
4761 config: serde_json::json!({"threshold": 9}),
4762 }];
4763
4764 let ctx = test_ctx().with_model("claude-opus-4-8");
4767 let collected = collect_capabilities_with_configs(&configs, ®istry, &ctx).await;
4768
4769 let ts = collected
4770 .tool_search
4771 .as_ref()
4772 .expect("auto_tool_search must set a hosted config on a native Claude model");
4773 assert!(ts.enabled);
4774 assert_eq!(ts.threshold, 9);
4775 assert!(
4776 !collected
4777 .tools
4778 .iter()
4779 .any(|t| t.name() == TOOL_SEARCH_TOOL_NAME),
4780 "hosted mechanism must not contribute the client-side tool_search tool"
4781 );
4782 assert!(
4783 collected.tool_definition_hooks.is_empty(),
4784 "hosted mechanism must not contribute a client-side deferral hook"
4785 );
4786 }
4787
4788 #[tokio::test]
4789 async fn test_collect_capabilities_no_tool_search_without_capability() {
4790 let registry = CapabilityRegistry::with_builtins();
4791
4792 let configs = vec![AgentCapabilityConfig {
4793 capability_ref: CapabilityId::new("current_time"),
4794 config: serde_json::json!({}),
4795 }];
4796
4797 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4798
4799 assert!(collected.tool_search.is_none());
4800 }
4801
4802 #[tokio::test]
4803 async fn test_collect_capabilities_tool_search_category_propagation() {
4804 let registry = CapabilityRegistry::with_builtins();
4805
4806 let configs = vec![
4808 AgentCapabilityConfig {
4809 capability_ref: CapabilityId::new("test_math"),
4810 config: serde_json::json!({}),
4811 },
4812 AgentCapabilityConfig {
4813 capability_ref: CapabilityId::new("openai_tool_search"),
4814 config: serde_json::json!({}),
4815 },
4816 ];
4817
4818 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4819
4820 assert!(collected.tool_search.is_some());
4822
4823 for tool_def in &collected.tool_definitions {
4825 if ["add", "subtract", "multiply", "divide"].contains(&tool_def.name()) {
4827 assert!(
4828 tool_def.category().is_some(),
4829 "Tool {} should have a category from its capability",
4830 tool_def.name()
4831 );
4832 }
4833 }
4834 }
4835
4836 #[tokio::test]
4837 async fn test_apply_capabilities_prompt_caching() {
4838 let registry = CapabilityRegistry::with_builtins();
4839 let base_runtime_agent = RuntimeAgent::new("You are a helpful assistant.", "gpt-5.4");
4840
4841 let applied = apply_capabilities(
4842 base_runtime_agent.clone(),
4843 &["prompt_caching".to_string()],
4844 ®istry,
4845 &test_ctx(),
4846 )
4847 .await;
4848
4849 assert_eq!(
4850 applied.runtime_agent.system_prompt,
4851 base_runtime_agent.system_prompt
4852 );
4853 assert!(applied.tool_registry.is_empty());
4854 assert_eq!(applied.applied_ids, vec!["prompt_caching"]);
4855
4856 let prompt_cache = applied.runtime_agent.prompt_cache.as_ref().unwrap();
4857 assert!(prompt_cache.enabled);
4858 assert_eq!(
4859 prompt_cache.strategy,
4860 crate::llm_driver_registry::PromptCacheStrategy::Auto
4861 );
4862 assert!(prompt_cache.gemini_cached_content.is_none());
4863 }
4864
4865 #[tokio::test]
4866 async fn test_collect_capabilities_prompt_caching_custom_strategy() {
4867 let registry = CapabilityRegistry::with_builtins();
4868
4869 let configs = vec![AgentCapabilityConfig {
4870 capability_ref: CapabilityId::new("prompt_caching"),
4871 config: serde_json::json!({"strategy": "auto"}),
4872 }];
4873
4874 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4875
4876 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
4877 assert!(prompt_cache.enabled);
4878 assert_eq!(
4879 prompt_cache.strategy,
4880 crate::llm_driver_registry::PromptCacheStrategy::Auto
4881 );
4882 assert!(prompt_cache.gemini_cached_content.is_none());
4883 }
4884
4885 #[tokio::test]
4886 async fn test_collect_capabilities_prompt_caching_gemini_cached_content() {
4887 let registry = CapabilityRegistry::with_builtins();
4888
4889 let configs = vec![AgentCapabilityConfig {
4890 capability_ref: CapabilityId::new("prompt_caching"),
4891 config: serde_json::json!({
4892 "strategy": "auto",
4893 "gemini_cached_content": "cachedContents/demo-cache"
4894 }),
4895 }];
4896
4897 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4898
4899 let prompt_cache = collected.prompt_cache.as_ref().unwrap();
4900 assert_eq!(
4901 prompt_cache.gemini_cached_content.as_deref(),
4902 Some("cachedContents/demo-cache")
4903 );
4904 }
4905
4906 struct SkillContributingCapability;
4911
4912 impl Capability for SkillContributingCapability {
4913 fn id(&self) -> &str {
4914 "contributes_skills"
4915 }
4916 fn name(&self) -> &str {
4917 "Contributes Skills"
4918 }
4919 fn description(&self) -> &str {
4920 "Test capability that contributes skills."
4921 }
4922 fn contribute_skills(&self) -> Vec<SkillContribution> {
4923 vec![
4924 SkillContribution::new("alpha-skill", "Alpha skill desc", "# Alpha\nDo alpha.")
4925 .with_files(vec![(
4926 "scripts/a.sh".to_string(),
4927 "#!/bin/sh\necho a\n".to_string(),
4928 )]),
4929 SkillContribution::new("beta-skill", "Beta skill desc", "# Beta\nDo beta.")
4930 .with_user_invocable(false),
4931 ]
4932 }
4933 }
4934
4935 fn skill_md_from_entries(entries: &HashMap<String, MountEntry>) -> &str {
4936 match &entries.get("SKILL.md").expect("SKILL.md missing").source {
4937 MountSource::InlineFile { content, .. } => content.as_str(),
4938 _ => panic!("Expected InlineFile for SKILL.md"),
4939 }
4940 }
4941
4942 #[tokio::test]
4943 async fn test_contribute_skills_normalized_to_mounts() {
4944 let mut registry = CapabilityRegistry::new();
4945 registry.register(SkillContributingCapability);
4946
4947 let configs = vec![AgentCapabilityConfig {
4948 capability_ref: CapabilityId::new("contributes_skills"),
4949 config: serde_json::json!({}),
4950 }];
4951
4952 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
4953
4954 let skill_mounts: Vec<_> = collected
4955 .mounts
4956 .iter()
4957 .filter(|m| m.path.starts_with("/.agents/skills/"))
4958 .collect();
4959 assert_eq!(skill_mounts.len(), 2);
4960
4961 for m in &skill_mounts {
4964 assert!(m.is_readonly());
4965 assert_eq!(m.capability_id, "contributes_skills");
4966 }
4967
4968 let alpha = skill_mounts
4969 .iter()
4970 .find(|m| m.path == "/.agents/skills/alpha-skill")
4971 .expect("alpha-skill mount missing");
4972 match &alpha.source {
4973 MountSource::InlineDirectory { entries } => {
4974 assert!(entries.contains_key("SKILL.md"));
4975 assert!(entries.contains_key("scripts/a.sh"));
4976 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
4977 assert_eq!(parsed.name, "alpha-skill");
4978 assert!(parsed.user_invocable);
4979 }
4980 _ => panic!("Expected InlineDirectory"),
4981 }
4982
4983 let beta = skill_mounts
4984 .iter()
4985 .find(|m| m.path == "/.agents/skills/beta-skill")
4986 .expect("beta-skill mount missing");
4987 match &beta.source {
4988 MountSource::InlineDirectory { entries } => {
4989 let parsed = crate::skill::parse_skill_md(skill_md_from_entries(entries)).unwrap();
4990 assert!(!parsed.user_invocable);
4991 }
4992 _ => panic!("Expected InlineDirectory"),
4993 }
4994 }
4995
4996 #[tokio::test]
4997 async fn test_contribute_skills_default_empty() {
4998 let mut registry = CapabilityRegistry::new();
5001 registry.register(FilterTestCapability { priority: 0 });
5002
5003 let configs = vec![AgentCapabilityConfig {
5004 capability_ref: CapabilityId::new("filter_test"),
5005 config: serde_json::json!({}),
5006 }];
5007
5008 let collected = collect_capabilities_with_configs(&configs, ®istry, &test_ctx()).await;
5009 assert!(
5010 collected
5011 .mounts
5012 .iter()
5013 .all(|m| !m.path.starts_with("/.agents/skills/"))
5014 );
5015 }
5016
5017 struct LocalizedCapability;
5018
5019 impl Capability for LocalizedCapability {
5020 fn id(&self) -> &str {
5021 "localized"
5022 }
5023 fn name(&self) -> &str {
5024 "Localized"
5025 }
5026 fn description(&self) -> &str {
5027 "English description"
5028 }
5029 fn localizations(&self) -> Vec<CapabilityLocalization> {
5030 vec![
5031 CapabilityLocalization {
5032 locale: "en",
5033 name: None,
5034 description: None,
5035 config_description: Some("Controls things."),
5036 config_overlay: None,
5037 },
5038 CapabilityLocalization {
5039 locale: "uk",
5040 name: Some("Локалізована"),
5041 description: Some("Український опис"),
5042 config_description: Some("Керує налаштуваннями."),
5043 config_overlay: None,
5044 },
5045 ]
5046 }
5047 }
5048
5049 #[test]
5050 fn localized_name_falls_back_exact_language_then_base() {
5051 let cap = LocalizedCapability;
5052 assert_eq!(cap.localized_name(Some("uk-UA")), "Локалізована");
5054 assert_eq!(cap.localized_name(Some("uk")), "Локалізована");
5055 assert_eq!(cap.localized_name(Some("uk_UA")), "Локалізована");
5057 assert_eq!(cap.localized_name(Some("fr-FR")), "Localized");
5059 assert_eq!(cap.localized_name(None), "Localized");
5060 assert_eq!(cap.localized_description(Some("uk")), "Український опис");
5061 assert_eq!(cap.localized_description(Some("de")), "English description");
5062 }
5063
5064 #[test]
5065 fn describe_schema_resolves_config_description_per_locale() {
5066 let cap = LocalizedCapability;
5067 assert_eq!(
5068 cap.describe_schema(Some("uk-UA")).as_deref(),
5069 Some("Керує налаштуваннями.")
5070 );
5071 assert_eq!(
5073 cap.describe_schema(Some("pl")).as_deref(),
5074 Some("Controls things.")
5075 );
5076 assert_eq!(
5077 cap.describe_schema(None).as_deref(),
5078 Some("Controls things.")
5079 );
5080 assert_eq!(NoopCapability.describe_schema(Some("uk")), None);
5082 }
5083}