From cdc2da2fedd1cfd8c1a0d3223e244a0025b670ba Mon Sep 17 00:00:00 2001 From: Henry Park Date: Mon, 23 Mar 2026 11:48:15 -0700 Subject: [PATCH] fix: block telegram bot tokens in chat (#1596) --- src/agent/agent_loop.rs | 200 ++++++++++++++++++++++++++++++- src/workspace/seeds/BOOTSTRAP.md | 4 + 2 files changed, 202 insertions(+), 2 deletions(-) diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 5cbd8166..70a3dfef 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -7,9 +7,10 @@ //! - `commands` - System commands and job handlers //! - `thread_ops` - Thread/session operations (user input, undo, approval, persistence) -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use futures::StreamExt; +use regex::Regex; use uuid::Uuid; use crate::agent::context_monitor::ContextMonitor; @@ -62,6 +63,38 @@ pub(crate) fn truncate_for_preview(output: &str, max_chars: usize) -> String { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SensitiveChatCredential { + TelegramBotToken, +} + +impl SensitiveChatCredential { + fn extension_name(self) -> &'static str { + match self { + Self::TelegramBotToken => "telegram", + } + } + + fn redirect_message(self) -> &'static str { + match self { + Self::TelegramBotToken => { + "Telegram bot tokens can't be accepted in normal chat. Use the secure Telegram setup flow instead." + } + } + } +} + +static TELEGRAM_BOT_TOKEN_RE: LazyLock = + LazyLock::new(|| Regex::new(r"^\d{6,}:[A-Za-z0-9_-]{20,}$").expect("TELEGRAM_BOT_TOKEN_RE")); // safety: hardcoded literal + +fn detect_sensitive_chat_credential(content: &str) -> Option { + let trimmed = content.trim(); + if TELEGRAM_BOT_TOKEN_RE.is_match(trimmed) { + return Some(SensitiveChatCredential::TelegramBotToken); + } + None +} + #[cfg(test)] fn resolve_routine_notification_user(metadata: &serde_json::Value) -> Option { resolve_owner_scope_notification_user( @@ -191,6 +224,28 @@ pub struct Agent { } impl Agent { + async fn intercept_sensitive_chat_credential( + &self, + message: &IncomingMessage, + credential: SensitiveChatCredential, + ) -> String { + let instructions = credential.redirect_message().to_string(); + let _ = self + .channels + .send_status( + &message.channel, + crate::channels::StatusUpdate::AuthRequired { + extension_name: credential.extension_name().to_string(), + instructions: Some(instructions.clone()), + auth_url: None, + setup_url: None, + }, + &message.metadata, + ) + .await; + instructions + } + pub(super) fn owner_id(&self) -> &str { if let Some(workspace) = self.deps.workspace.as_ref() { debug_assert_eq!( @@ -977,6 +1032,15 @@ impl Agent { std::any::type_name_of_val(&submission) ); + if let Submission::UserInput { ref content } = submission + && let Some(credential) = detect_sensitive_chat_credential(content) + { + return Ok(Some( + self.intercept_sensitive_chat_credential(message, credential) + .await, + )); + } + // Hook: BeforeInbound — allow hooks to modify or reject user input if let Submission::UserInput { ref content } = submission { let event = crate::hooks::HookEvent::Inbound { @@ -1321,11 +1385,24 @@ impl Agent { #[cfg(test)] mod tests { use super::{ - chat_tool_execution_metadata, resolve_routine_notification_user, + Agent, AgentDeps, SensitiveChatCredential, chat_tool_execution_metadata, + detect_sensitive_chat_credential, resolve_routine_notification_user, should_fallback_routine_notification, truncate_for_preview, }; use crate::channels::IncomingMessage; use crate::error::ChannelError; + use crate::testing::{StubChannel, StubLlm}; + use crate::{ + agent::cost_guard::{CostGuard, CostGuardConfig}, + channels::{ChannelManager, StatusUpdate}, + config::{AgentConfig, SafetyConfig, SkillsConfig}, + context::ContextManager, + hooks::HookRegistry, + safety::SafetyLayer, + tools::ToolRegistry, + }; + use std::sync::Arc; + use std::time::Duration; #[test] fn test_truncate_short_input() { @@ -1483,4 +1560,123 @@ mod tests { assert!(should_fallback_routine_notification(&error)); // safety: test-only assertion } + + #[test] + fn detects_telegram_bot_token_messages() { + let detected = detect_sensitive_chat_credential("123456789:AABBccDDeeFFgg_Test-Token"); + assert_eq!(detected, Some(SensitiveChatCredential::TelegramBotToken)); + } + + #[test] + fn ignores_normal_telegram_setup_messages() { + let detected = detect_sensitive_chat_credential( + "Can you help me connect Telegram without sharing the token here?", + ); + assert_eq!(detected, None); + } + + async fn make_gateway_test_agent( + llm: Arc, + ) -> (Agent, Arc>>) { + let llm_provider: Arc = llm; + let (stub, _sender) = StubChannel::new("gateway"); + let statuses = stub.captured_statuses_handle(); + let channel_manager = ChannelManager::new(); + channel_manager.add(Box::new(stub)).await; + + let deps = AgentDeps { + owner_id: "default".to_string(), + store: None, + llm: llm_provider, + cheap_llm: None, + safety: Arc::new(SafetyLayer::new(&SafetyConfig { + max_output_length: 100_000, + injection_check_enabled: false, + })), + tools: Arc::new(ToolRegistry::new()), + workspace: None, + extension_manager: None, + skill_registry: None, + skill_catalog: None, + skills_config: SkillsConfig::default(), + hooks: Arc::new(HookRegistry::new()), + cost_guard: Arc::new(CostGuard::new(CostGuardConfig::default())), + sse_tx: None, + http_interceptor: None, + transcription: None, + document_extraction: None, + sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig, + builder: None, + }; + + let agent = Agent::new( + AgentConfig { + name: "test-agent".to_string(), + max_parallel_jobs: 1, + job_timeout: Duration::from_secs(60), + stuck_threshold: Duration::from_secs(60), + repair_check_interval: Duration::from_secs(30), + max_repair_attempts: 1, + use_planning: false, + session_idle_timeout: Duration::from_secs(300), + allow_local_tools: false, + max_cost_per_day_cents: None, + max_actions_per_hour: None, + max_tool_iterations: 5, + auto_approve_tools: false, + default_timezone: "UTC".to_string(), + max_tokens_per_job: 0, + }, + deps, + Arc::new(channel_manager), + None, + None, + None, + Some(Arc::new(ContextManager::new(1))), + None, + ); + + (agent, statuses) + } + + #[tokio::test] + async fn telegram_bot_token_messages_are_redirected_before_llm() { + let llm = Arc::new(StubLlm::new("this should never be used")); + let llm_handle = Arc::clone(&llm); + let (agent, statuses) = make_gateway_test_agent(llm).await; + let message = IncomingMessage::new( + "gateway", + "test-user", + "123456789:AABBccDDeeFFgg_Test-Token", + ); + + let response = agent + .handle_message(&message) + .await + .expect("handle_message"); + + assert_eq!( + response.as_deref(), + Some( + "Telegram bot tokens can't be accepted in normal chat. Use the secure Telegram setup flow instead." + ) + ); + assert_eq!(llm_handle.calls(), 0, "LLM should not see raw bot tokens"); + + let statuses = statuses.lock().expect("poisoned"); + assert_eq!(statuses.len(), 1); + assert!(matches!( + &statuses[0], + StatusUpdate::AuthRequired { + extension_name, + instructions, + auth_url: None, + setup_url: None, + } if extension_name == "telegram" + && instructions.as_deref() + == Some( + "Telegram bot tokens can't be accepted in normal chat. Use the secure Telegram setup flow instead." + ) + )); + } } diff --git a/src/workspace/seeds/BOOTSTRAP.md b/src/workspace/seeds/BOOTSTRAP.md index b2b389e8..234fcbac 100644 --- a/src/workspace/seeds/BOOTSTRAP.md +++ b/src/workspace/seeds/BOOTSTRAP.md @@ -35,6 +35,10 @@ If they're interested, set it up right here using the extension tools: 3. Use `tool_auth` to collect credentials (e.g. Telegram bot token from @BotFather) 4. The channel will be hot-activated — no restart needed +Never ask the user to paste tokens, passwords, API keys, or other secrets into +normal chat. If an extension has a secure auth/setup flow, always use that flow +and keep the secret out of the conversation history. + Don't push if they're not interested — note their preference and move on. ## Step 3: Save What You Learned (MANDATORY after 3 user messages)