From 060ce8de2522cf9b3290aab67dceb888b63311a2 Mon Sep 17 00:00:00 2001 From: Henry Park Date: Mon, 23 Mar 2026 11:47:58 -0700 Subject: [PATCH] fix(agent): surface active extension state to the LLM --- src/agent/dispatcher.rs | 12 +++ src/extensions/manager.rs | 216 ++++++++++++++++++++++++++++++++++++++ src/llm/reasoning.rs | 104 +++++++++++++++--- 3 files changed, 320 insertions(+), 12 deletions(-) diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 03548219..726ee2a2 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -131,6 +131,18 @@ impl Agent { } } + if let Some(extension_manager) = self.deps.extension_manager.as_ref() { + match extension_manager.llm_extension_state_summary().await { + Ok(Some(summary)) => { + reasoning = reasoning.with_extension_state_summary(summary); + } + Ok(None) => {} + Err(e) => { + tracing::debug!("Could not load extension state summary: {}", e); + } + } + } + if let Some(prompt) = system_prompt { reasoning = reasoning.with_system_prompt(prompt); } diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index df5de72d..79756e66 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -1498,6 +1498,52 @@ impl ExtensionManager { Ok(extensions) } + /// Build a compact, deterministic extension snapshot for LLM prompt context. + pub async fn llm_extension_state_summary(&self) -> Result, ExtensionError> { + let mut extensions = self.list(None, false).await?; + extensions.sort_by(|a, b| { + llm_extension_sort_key(a.kind) + .cmp(&llm_extension_sort_key(b.kind)) + .then_with(|| a.name.cmp(&b.name)) + }); + + let mut channels = Vec::new(); + let mut tools = Vec::new(); + let mut servers = Vec::new(); + + for extension in extensions { + let owner_bound = matches!(extension.kind, ExtensionKind::WasmChannel) + && self.has_wasm_channel_owner_binding(&extension.name).await; + if !(extension.active || extension.authenticated || owner_bound) { + continue; + } + + let item = llm_extension_summary_item(&extension, owner_bound); + match extension.kind { + ExtensionKind::WasmChannel | ExtensionKind::ChannelRelay => channels.push(item), + ExtensionKind::WasmTool => tools.push(item), + ExtensionKind::McpServer => servers.push(item), + } + } + + let mut lines = Vec::new(); + if !channels.is_empty() { + lines.push(format!("- Channels: {}", channels.join("; "))); + } + if !tools.is_empty() { + lines.push(format!("- Tools: {}", tools.join("; "))); + } + if !servers.is_empty() { + lines.push(format!("- MCP servers: {}", servers.join("; "))); + } + + if lines.is_empty() { + Ok(None) + } else { + Ok(Some(lines.join("\n"))) + } + } + /// Remove an installed extension. pub async fn remove(&self, name: &str) -> Result { Self::validate_extension_name(name)?; @@ -5613,6 +5659,40 @@ fn combine_install_errors( } } +fn llm_extension_sort_key(kind: ExtensionKind) -> u8 { + match kind { + ExtensionKind::WasmChannel | ExtensionKind::ChannelRelay => 0, + ExtensionKind::WasmTool => 1, + ExtensionKind::McpServer => 2, + } +} + +fn llm_extension_summary_item(extension: &InstalledExtension, owner_bound: bool) -> String { + let mut states = Vec::new(); + if extension.authenticated { + states.push("authenticated".to_string()); + } + if extension.active { + states.push("active".to_string()); + } else if extension.authenticated { + states.push("inactive".to_string()); + } + if owner_bound { + states.push("owner-bound".to_string()); + } + if !extension.tools.is_empty() { + let mut tool_names = extension.tools.clone(); + tool_names.sort(); + states.push(format!("tools: {}", tool_names.join(", "))); + } + + if states.is_empty() { + extension.name.clone() + } else { + format!("{} ({})", extension.name, states.join(", ")) + } +} + #[cfg(test)] mod tests { use std::fmt::Debug; @@ -6571,6 +6651,142 @@ mod tests { ) } + #[cfg(feature = "libsql")] + #[tokio::test] + async fn test_llm_extension_state_summary_reports_active_owner_bound_telegram() + -> Result<(), String> { + let dir = tempfile::tempdir().map_err(|err| format!("temp dir: {err}"))?; + let channels_dir = dir.path().join("channels"); + std::fs::create_dir_all(&channels_dir).map_err(|err| format!("channels dir: {err}"))?; + std::fs::write(channels_dir.join("telegram.wasm"), b"mock") + .map_err(|err| format!("write wasm: {err}"))?; + std::fs::write( + channels_dir.join("telegram.capabilities.json"), + serde_json::to_vec(&serde_json::json!({ + "type": "channel", + "name": "telegram", + "setup": { + "required_secrets": [ + { + "name": "telegram_bot_token", + "prompt": "Enter your Telegram Bot API token (from @BotFather)", + "optional": false + } + ] + }, + "capabilities": { + "channel": { + "allowed_paths": ["/webhook/telegram"] + } + }, + "config": { + "owner_id": null + } + })) + .map_err(|err| format!("serialize capabilities: {err}"))?, + ) + .map_err(|err| format!("write capabilities: {err}"))?; + + let (db, _db_tmp) = crate::testing::test_db().await; + let manager = { + use crate::secrets::{InMemorySecretsStore, SecretsCrypto}; + use crate::testing::credentials::TEST_CRYPTO_KEY; + use crate::tools::ToolRegistry; + use crate::tools::mcp::process::McpProcessManager; + use crate::tools::mcp::session::McpSessionManager; + + let master_key = secrecy::SecretString::from(TEST_CRYPTO_KEY.to_string()); + let crypto = Arc::new( + SecretsCrypto::new(master_key) + .map_err(|err| format!("failed to construct test crypto: {err}"))?, + ); + + ExtensionManager::new( + Arc::new(McpSessionManager::new()), + Arc::new(McpProcessManager::new()), + Arc::new(InMemorySecretsStore::new(crypto)), + Arc::new(ToolRegistry::new()), + None, + None, + dir.path().join("tools"), + channels_dir.clone(), + None, + "test".to_string(), + Some(db), + Vec::new(), + ) + }; + + let channel_manager = Arc::new(ChannelManager::new()); + let runtime = Arc::new( + WasmChannelRuntime::new(WasmChannelRuntimeConfig::for_testing()) + .map_err(|err| format!("runtime: {err}"))?, + ); + let pairing_store = Arc::new(PairingStore::with_base_dir( + dir.path().join("pairing-state"), + )); + let router = Arc::new(WasmChannelRouter::new()); + manager + .set_channel_runtime( + Arc::clone(&channel_manager), + Arc::clone(&runtime), + Arc::clone(&pairing_store), + Arc::clone(&router), + std::collections::HashMap::new(), + ) + .await; + manager + .set_test_wasm_channel_loader(Arc::new({ + let runtime = Arc::clone(&runtime); + let pairing_store = Arc::clone(&pairing_store); + move |name| { + Ok(make_test_loaded_channel( + Arc::clone(&runtime), + name, + Arc::clone(&pairing_store), + )) + } + })) + .await; + manager + .set_test_telegram_binding_resolver(Arc::new(|_token, existing_owner_id| { + if existing_owner_id.is_some() { + return Err(ExtensionError::Other( + "owner binding should be derived during setup".to_string(), + )); + } + Ok(TelegramBindingResult::Bound(TelegramBindingData { + owner_id: 424242, + bot_username: Some("test_hot_bot".to_string()), + binding_state: TelegramOwnerBindingState::VerifiedNow, + })) + })) + .await; + + manager + .configure( + "telegram", + &std::collections::HashMap::from([( + "telegram_bot_token".to_string(), + "123456789:ABCdefGhI".to_string(), + )]), + &std::collections::HashMap::new(), + ) + .await + .map_err(|err| format!("configure succeeds: {err}"))?; + + let summary = manager + .llm_extension_state_summary() + .await + .map_err(|err| format!("summary: {err}"))? + .ok_or_else(|| "expected extension summary".to_string())?; + + require( + summary.contains("- Channels: telegram (authenticated, active, owner-bound)"), + format!("unexpected summary: {summary}"), + ) + } + #[tokio::test] async fn test_telegram_hot_activation_returns_verification_challenge_before_binding() -> Result<(), String> { diff --git a/src/llm/reasoning.rs b/src/llm/reasoning.rs index cbec297b..9d1aee0d 100644 --- a/src/llm/reasoning.rs +++ b/src/llm/reasoning.rs @@ -353,6 +353,8 @@ pub struct Reasoning { workspace_system_prompt: Option, /// Optional skill context block to inject into system prompt. skill_context: Option, + /// Optional snapshot of connected/active extensions for the current user. + extension_state_summary: Option, /// Channel name (e.g. "discord", "telegram") for formatting hints. channel: Option, /// Model name for runtime context. @@ -371,6 +373,7 @@ impl Reasoning { llm, workspace_system_prompt: None, skill_context: None, + extension_state_summary: None, channel: None, model_name: None, is_group_chat: false, @@ -400,6 +403,14 @@ impl Reasoning { self } + /// Set extension runtime context to inject into the system prompt. + pub fn with_extension_state_summary(mut self, summary: String) -> Self { + if !summary.is_empty() { + self.extension_state_summary = Some(summary); + } + self + } + /// Set the channel name for channel-specific formatting hints. pub fn with_channel(mut self, channel: impl Into) -> Self { let ch = channel.into(); @@ -932,21 +943,54 @@ Example: } fn build_extensions_section_for_tools(&self, tools: &[ToolDefinition]) -> String { - // Only include when the extension management tools are available - let has_ext_tools = tools.iter().any(|t| t.name == "tool_search"); - if !has_ext_tools { + let has_search = tools.iter().any(|t| t.name == "tool_search"); + let has_list = tools.iter().any(|t| t.name == "tool_list"); + let has_info = tools.iter().any(|t| t.name == "extension_info"); + + if self.extension_state_summary.is_none() && !has_search && !has_list && !has_info { return String::new(); } - "\n\n## Extensions\n\ - You can search, install, and activate extensions to add new capabilities:\n\ - - **Channels** (Telegram, Slack, Discord) — messaging integrations. \ - When users ask about connecting a messaging platform, search for it as a channel.\n\ - - **Tools** — sandboxed functions that extend your abilities.\n\ - - **MCP servers** — external API integrations via the Model Context Protocol.\n\n\ - Use `tool_search` to find extensions by name. Refer to them by their kind \ - (channel, tool, or server) — not as \"MCP server\" generically." - .to_string() + let mut blocks = Vec::new(); + if let Some(ref summary) = self.extension_state_summary { + blocks.push(format!( + "Current extension state for this user:\n{}", + summary + )); + } + + if has_search || has_list || has_info { + let mut guidance = String::from( + "You can search, install, and activate extensions to add new capabilities:\n\ + - **Channels** (Telegram, Slack, Discord) — messaging integrations. \ + When users ask about connecting a messaging platform, search for it as a channel.\n\ + - **Tools** — sandboxed functions that extend your abilities.\n\ + - **MCP servers** — external API integrations via the Model Context Protocol.", + ); + + if has_list { + guidance.push_str( + "\n\nBefore telling the user to connect, activate, or re-enable an extension, \ + inspect the current state with `tool_list`.", + ); + } + if has_info { + guidance.push_str( + "\nUse `extension_info` when you need deeper compatibility or runtime details \ + for an installed extension.", + ); + } + if has_search { + guidance.push_str( + "\nUse `tool_search` to find extensions by name. Refer to them by their kind \ + (channel, tool, or server) — not as \"MCP server\" generically.", + ); + } + + blocks.push(guidance); + } + + format!("\n\n## Extensions\n{}", blocks.join("\n\n")) } fn build_channel_section(&self) -> String { @@ -2285,6 +2329,42 @@ That's my plan."#; ); } + #[test] + fn test_system_prompt_includes_extension_runtime_summary() { + let reasoning = make_test_reasoning().with_extension_state_summary( + "- Channels: telegram (authenticated, active, owner-bound)".to_string(), + ); + + let prompt = reasoning.build_system_prompt_with_tools(&[]); + assert!( + prompt.contains("## Extensions"), + "Prompt should contain an Extensions section when runtime state is present" + ); + assert!( + prompt.contains("telegram (authenticated, active, owner-bound)"), + "Prompt should include the injected extension runtime summary" + ); + } + + #[test] + fn test_system_prompt_extension_guidance_prefers_inspection_before_reconnect() { + let reasoning = make_test_reasoning(); + let prompt = reasoning.build_system_prompt_with_tools(&make_tools(&[ + "tool_search", + "tool_list", + "extension_info", + ])); + + assert!( + prompt.contains("inspect the current state with `tool_list`"), + "Prompt should direct the model to inspect installed/active state first" + ); + assert!( + prompt.contains("Use `extension_info` when you need deeper compatibility"), + "Prompt should mention extension_info for deeper extension details" + ); + } + // ---- plan/evaluate bypass clean_response (Bug #564-2) ---- #[test]