From a074d8eb3e96c49f89ce8cedfeead1301dbbdf68 Mon Sep 17 00:00:00 2001 From: Abhijay Jain Date: Tue, 28 Jul 2026 01:07:22 +0530 Subject: [PATCH] fix(telegram): prompt for tool approval in gateway sessions (#10613) Signed-off-by: Abhijay Jain --- crates/goose/src/gateway/handler.rs | 194 ++++++++++++++++++++++++++- crates/goose/src/gateway/manager.rs | 8 +- crates/goose/src/gateway/telegram.rs | 11 +- 3 files changed, 208 insertions(+), 5 deletions(-) diff --git a/crates/goose/src/gateway/handler.rs b/crates/goose/src/gateway/handler.rs index d2123a94c..7854dee1b 100644 --- a/crates/goose/src/gateway/handler.rs +++ b/crates/goose/src/gateway/handler.rs @@ -1,16 +1,20 @@ +use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use futures::StreamExt; +use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; -use crate::agents::{AgentEvent, ExtensionConfig, SessionConfig}; +use crate::agents::{Agent, AgentEvent, ExtensionConfig, SessionConfig}; use crate::config::extensions::get_enabled_extensions; use crate::config::paths::Paths; use crate::config::Config; -use crate::conversation::message::{Message, MessageContent}; +use crate::conversation::message::{ActionRequiredData, Message, MessageContent}; use crate::execution::manager::AgentManager; +use crate::permission::permission_confirmation::PrincipalType; +use crate::permission::{Permission, PermissionConfirmation}; use crate::session::SessionType; use crate::session::{EnabledExtensionsState, ExtensionState, Session}; @@ -36,12 +40,21 @@ fn resolve_gateway_max_turns(gateway_override: Option, global_max_turns: Op .unwrap_or(DEFAULT_GATEWAY_MAX_TURNS) } +struct PendingConfirmation { + agent: Arc, + request_id: String, +} + #[derive(Clone)] pub struct GatewayHandler { agent_manager: Arc, pairing_store: Arc, gateway: Arc, config: GatewayConfig, + /// Tracks users who have a tool-confirmation prompt awaiting their reply. + pending_confirmations: Arc>>, + /// Serializes `relay_to_session` per user; confirmation replies bypass this lock. + turn_locks: Arc>>>>, } impl GatewayHandler { @@ -56,6 +69,34 @@ impl GatewayHandler { pairing_store, gateway, config, + pending_confirmations: Arc::new(Mutex::new(HashMap::new())), + turn_locks: Arc::new(Mutex::new(HashMap::new())), + } + } + + pub async fn deny_pending_confirmations(&self) { + let pending: Vec<_> = self.pending_confirmations.lock().await.drain().collect(); + for (_, confirmation) in pending { + confirmation + .agent + .handle_confirmation( + confirmation.request_id, + PermissionConfirmation { + principal_type: PrincipalType::Tool, + permission: Permission::DenyOnce, + }, + ) + .await; + } + } + + async fn prune_turn_lock(&self, user: &PlatformUser) { + let mut locks = self.turn_locks.lock().await; + let in_use = locks + .get(user) + .is_some_and(|lock| Arc::strong_count(lock) > 1); + if !in_use { + locks.remove(user); } } @@ -117,13 +158,81 @@ impl GatewayHandler { } } PairingState::Paired { session_id, .. } => { - self.relay_to_session(&message, &session_id).await?; + if self + .pending_confirmations + .lock() + .await + .contains_key(&message.user) + { + self.handle_pending_confirmation(&message).await?; + } else { + let turn_lock = { + let mut locks = self.turn_locks.lock().await; + Arc::clone( + locks + .entry(message.user.clone()) + .or_insert_with(|| Arc::new(Mutex::new(()))), + ) + }; + let turn_guard = turn_lock.lock().await; + let result = self.relay_to_session(&message, &session_id).await; + drop(turn_guard); + drop(turn_lock); + self.prune_turn_lock(&message.user).await; + result?; + } } } Ok(()) } + async fn handle_pending_confirmation(&self, message: &IncomingMessage) -> anyhow::Result<()> { + let text = message.text.trim().to_lowercase(); + let permission = match text.as_str() { + "approve" | "yes" | "y" => Some(Permission::AllowOnce), + "approve always" => Some(Permission::AlwaysAllow), + "deny" | "no" | "n" => Some(Permission::DenyOnce), + "deny always" => Some(Permission::AlwaysDeny), + _ => None, + }; + + let Some(permission) = permission else { + self.gateway + .send_message( + &message.user, + OutgoingMessage::Text { + body: "Please reply with: approve, approve always, deny, or deny always." + .into(), + }, + ) + .await?; + return Ok(()); + }; + + let Some(pending) = self + .pending_confirmations + .lock() + .await + .remove(&message.user) + else { + return Ok(()); + }; + + pending + .agent + .handle_confirmation( + pending.request_id, + PermissionConfirmation { + principal_type: PrincipalType::Tool, + permission, + }, + ) + .await; + + Ok(()) + } + async fn try_consume_code(&self, text: &str) -> anyhow::Result> { let normalized = text.to_uppercase().replace(['-', ' '], ""); if normalized.len() == 6 @@ -474,6 +583,85 @@ impl GatewayHandler { "gateway stream: tool response" ); } + MessageContent::ActionRequired(action_required) => { + if let ActionRequiredData::ToolConfirmation { + id, + tool_name, + arguments, + prompt, + } = &action_required.data + { + // Flush pending text so the user sees context first. + if !pending_text.is_empty() { + let _ = self + .gateway + .send_message( + &message.user, + OutgoingMessage::Text { + body: std::mem::take(&mut pending_text), + }, + ) + .await; + } + + let args_display = serde_json::to_string_pretty(arguments) + .unwrap_or_default(); + let mut approval_text = format!( + "🔐 Approval required\n\nTool: {tool_name}\nArguments:\n{args_display}" + ); + if let Some(p) = prompt { + approval_text.push_str(&format!("\n\n{p}")); + } + approval_text.push_str( + "\n\nReply with:\n\ + • approve — allow once\n\ + • approve always — always allow\n\ + • deny — deny once\n\ + • deny always — always deny", + ); + + self.pending_confirmations.lock().await.insert( + message.user.clone(), + PendingConfirmation { + agent: agent.clone(), + request_id: id.clone(), + }, + ); + + let send_result = self + .gateway + .send_message( + &message.user, + OutgoingMessage::Text { + body: approval_text, + }, + ) + .await; + + if let Err(e) = send_result { + tracing::error!( + session_id, + error = %e, + "failed to deliver tool approval prompt; denying tool call" + ); + self.pending_confirmations + .lock() + .await + .remove(&message.user); + agent + .handle_confirmation( + id.clone(), + PermissionConfirmation { + principal_type: PrincipalType::Tool, + permission: Permission::DenyOnce, + }, + ) + .await; + } else { + sent_any = true; + } + } + } _ => {} } } diff --git a/crates/goose/src/gateway/manager.rs b/crates/goose/src/gateway/manager.rs index c7050fca9..19e4ad327 100644 --- a/crates/goose/src/gateway/manager.rs +++ b/crates/goose/src/gateway/manager.rs @@ -31,6 +31,7 @@ pub struct GatewayInstance { pub gateway: Arc, pub cancel: CancellationToken, pub handle: tokio::task::JoinHandle<()>, + handler: GatewayHandler, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -129,6 +130,7 @@ impl GatewayManager { .remove(gateway_type) .ok_or_else(|| anyhow::anyhow!("Gateway '{}' is not running", gateway_type))?; + instance.handler.deny_pending_confirmations().await; instance.cancel.cancel(); let _ = instance.handle.await; @@ -140,6 +142,7 @@ impl GatewayManager { pub async fn remove_gateway(&self, gateway_type: &str) -> anyhow::Result<()> { // Stop if running (ignore error if not running). if let Some(instance) = self.gateways.write().await.remove(gateway_type) { + instance.handler.deny_pending_confirmations().await; instance.cancel.cancel(); let _ = instance.handle.await; } @@ -195,13 +198,14 @@ impl GatewayManager { gateway.clone(), config.clone(), ); + let handler_for_task = handler.clone(); let gateway_clone = gateway.clone(); let cancel_clone = cancel.clone(); let gateway_type_for_task = gw_type.clone(); let handle = tokio::spawn(async move { - if let Err(e) = gateway_clone.start(handler, cancel_clone).await { + if let Err(e) = gateway_clone.start(handler_for_task, cancel_clone).await { tracing::error!(gateway = %gateway_type_for_task, error = %e, "gateway stopped with error"); } }); @@ -211,6 +215,7 @@ impl GatewayManager { gateway, cancel, handle, + handler, }; self.gateways.write().await.insert(gw_type, instance); @@ -222,6 +227,7 @@ impl GatewayManager { let instances: Vec<(String, GatewayInstance)> = self.gateways.write().await.drain().collect(); for (gateway_type, instance) in instances { + instance.handler.deny_pending_confirmations().await; instance.cancel.cancel(); let _ = instance.handle.await; tracing::info!(gateway = %gateway_type, "gateway stopped"); diff --git a/crates/goose/src/gateway/telegram.rs b/crates/goose/src/gateway/telegram.rs index 61d30f3ce..2f4cde6f0 100644 --- a/crates/goose/src/gateway/telegram.rs +++ b/crates/goose/src/gateway/telegram.rs @@ -161,14 +161,23 @@ impl TelegramGateway { "Telegram rejected HTML, falling back to plain text" ); for plain_chunk in split_message(text, MAX_MESSAGE_LENGTH) { - self.client + let plain_resp = self + .client .post(self.api_url("sendMessage")) .json(&serde_json::json!({ "chat_id": chat_id, "text": plain_chunk, })) .send() + .await? + .json::>() .await?; + if !plain_resp.ok { + anyhow::bail!( + "Telegram sendMessage failed: {}", + plain_resp.description.unwrap_or_default() + ); + } } return Ok(()); }