diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index bccffba0c..a705e815b 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -26,9 +26,7 @@ use goose::agents::platform_extensions::developer::shell::{ }; use goose::agents::AgentEvent; use goose::agents::SUBAGENT_TOOL_REQUEST_TYPE; -use goose::permission::permission_confirmation::PrincipalType; use goose::permission::Permission; -use goose::permission::PermissionConfirmation; use goose::providers::base::Provider; use goose::providers::base::ProviderUsage; use goose::utils::safe_truncate; @@ -1553,7 +1551,6 @@ impl CliSession { .messages .last() .ok_or_else(|| anyhow::anyhow!("No user message"))?; - let cancel_token_interrupt = cancel_token.clone(); let handle = tokio::spawn(async move { if ctrl_c().await.is_ok() { @@ -1590,7 +1587,7 @@ impl CliSession { first_token_at = Some(Instant::now()); } if let Some((id, security_prompt)) = find_tool_confirmation(&message) { - let permission = if interactive { + let selected_permission = if interactive { prompt_tool_confirmation(&security_prompt)? } else { // Non-interactive/headless mode: refuse to run in @@ -1613,30 +1610,39 @@ impl CliSession { Permission::AllowOnce }; - if permission == Permission::Cancel { + let cancelled_by_user = selected_permission == Permission::Cancel; + if cancelled_by_user { output::render_text("Tool call cancelled. Returning to chat...", Some(Color::Yellow), true); - self.agent.handle_confirmation(id.clone(), PermissionConfirmation { - principal_type: PrincipalType::Tool, - permission: Permission::DenyOnce, - }).await; + } + self.agent + .submit_tool_confirmation( + &self.session_id, + &id, + selected_permission, + ) + .await?; + if cancelled_by_user { let mut response_message = Message::user(); response_message.content.push(MessageContent::tool_response( id, Err(ErrorData { code: ErrorCode::INVALID_REQUEST, - message: std::borrow::Cow::from("Tool call cancelled by user"), + message: std::borrow::Cow::from( + "Tool call cancelled by user", + ), data: None, }), )); + self.agent + .config + .session_manager + .add_message(&self.session_id, &response_message) + .await?; self.messages.push(response_message); cancel_token_clone.cancel(); drop(stream); break; } - self.agent.handle_confirmation(id, PermissionConfirmation { - principal_type: PrincipalType::Tool, - permission, - }).await; } else if let Some((elicitation_id, elicitation_message, schema)) = find_elicitation_request(&message) { if !interactive { // Non-interactive/headless mode: cannot collect user input diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index 7e0338e3a..c21a4b5b8 100644 --- a/crates/goose/src/acp/server.rs +++ b/crates/goose/src/acp/server.rs @@ -12,6 +12,9 @@ use crate::acp::{PermissionDecision, ACP_CURRENT_MODEL}; use crate::agents::extension::{Envs, PLATFORM_EXTENSIONS}; use crate::agents::mcp_client::{GooseMcpHostInfo, McpClientTrait}; use crate::agents::platform_extensions::developer::DeveloperClient; +use crate::agents::state_machine::{ + has_unapplied_tool_confirmation_response, pending_tool_confirmations, +}; use crate::agents::{ Agent, AgentConfig, ExtensionConfig, ExtensionLoadResult, GoosePlatform, SessionConfig, }; @@ -22,7 +25,7 @@ use crate::config::permission::PermissionManager; use crate::config::{Config, GooseMode}; use crate::conversation::message::{ ActionRequiredData, Message, MessageContent, SystemNotificationContent, SystemNotificationType, - ToolRequest, ToolResponse, + ToolConfirmationRequest, ToolRequest, ToolResponse, }; use crate::conversation::Conversation; use crate::execution::manager::{AgentManager, AgentManagerGetResult, RuntimeContext}; @@ -64,7 +67,7 @@ use agent_client_protocol::{ use anyhow::Result; use fs_err as fs; use futures::future::{BoxFuture, FutureExt}; -use futures::stream::{self, StreamExt}; +use futures::stream::{self, BoxStream, StreamExt}; use goose_providers::errors::ProviderError; use rmcp::model::{ Annotations as RmcpAnnotations, ImageContent as RmcpImageContent, Role, @@ -243,6 +246,11 @@ pub struct ActivePromptRun { agent: Arc, } +struct AgentStreamOutcome { + was_cancelled: bool, + output_token_limit_reached: bool, +} + /// Per-session active-run registry, shared by every `GooseAcpAgent` created /// from one `AcpServer`. Roaming spawns a fresh agent per connection, so two /// paired clients loading the same session get distinct agents; sharing this @@ -251,11 +259,11 @@ pub struct ActivePromptRun { /// one session. pub type ActiveRunRegistry = Arc>>; -/// Releases a registry entry if the owning `on_prompt` future is dropped +/// Releases a registry entry if the task consuming an agent stream is dropped /// without reaching its explicit `clear_active_run` — e.g. a roaming -/// connection is revoked or lost mid-prompt and the transport drops the -/// request future. Without this, the shared registry retains the run forever -/// and every later connection gets "session already has active run". +/// connection is revoked or lost mid-turn. Without this, the shared registry +/// retains the run forever and every later connection gets "session already +/// has active run". /// /// The explicit clear still runs on normal paths; this drop is then a no-op /// because the entry (matched by run id) is already gone. @@ -741,6 +749,20 @@ fn prompt_stop_reason(was_cancelled: bool, output_token_limit_reached: bool) -> } } +#[derive(Clone)] +struct SessionAgentTarget { + agent: Arc, + session_id: String, + cancel_token: Option, +} + +struct PendingToolPermission { + request_id: String, + tool_name: String, + arguments: serde_json::Map, + prompt: Option, +} + fn update_output_token_limit_reached(output_token_limit_reached: &mut bool, message: &Message) { if message.role == Role::Assistant { *output_token_limit_reached = message.metadata.output_token_limit_reached; @@ -1347,7 +1369,7 @@ impl GooseAcpAgent { content_item: &MessageContent, message: &Message, session_id: &SessionId, - agent: &Arc, + target: &SessionAgentTarget, tool_requests: &HashMap, cx: &ConnectionTo, ) -> Result<(), agent_client_protocol::Error> { @@ -1366,7 +1388,7 @@ impl GooseAcpAgent { cx.send_notification(SessionNotification::new(session_id.clone(), update))?; } MessageContent::ToolRequest(tool_request) => { - self.handle_tool_request(tool_request, message, session_id, agent, cx) + self.handle_tool_request(tool_request, message, session_id, &target.agent, cx) .await?; } MessageContent::ToolResponse(tool_response) => { @@ -1396,12 +1418,14 @@ impl GooseAcpAgent { } => { self.handle_tool_permission_request( cx, - agent, session_id, - id.clone(), - tool_name.clone(), - arguments.clone(), - prompt.clone(), + PendingToolPermission { + request_id: id.clone(), + tool_name: tool_name.clone(), + arguments: arguments.clone(), + prompt: prompt.clone(), + }, + target.clone(), )?; } ActionRequiredData::Elicitation { @@ -1543,23 +1567,22 @@ impl GooseAcpAgent { Ok(()) } - #[allow(clippy::too_many_arguments)] fn handle_tool_permission_request( &self, cx: &ConnectionTo, - agent: &Arc, session_id: &SessionId, - request_id: String, - tool_name: String, - arguments: serde_json::Map, - prompt: Option, + request: PendingToolPermission, + target: SessionAgentTarget, ) -> Result<(), agent_client_protocol::Error> { let cx = cx.clone(); - let agent = agent.clone(); let session_id = session_id.clone(); - let tool_call_update = - build_permission_tool_call_update(&request_id, &tool_name, arguments, prompt); + let tool_call_update = build_permission_tool_call_update( + &request.request_id, + &request.tool_name, + request.arguments, + request.prompt, + ); fn option(kind: PermissionOptionKind) -> PermissionOption { let id = serde_json::to_value(kind) @@ -1578,33 +1601,34 @@ impl GooseAcpAgent { let permission_request = RequestPermissionRequest::new(session_id, tool_call_update, options); + let request_id = request.request_id; cx.send_request(permission_request) .on_receiving_result(move |result| async move { - match result { - Ok(response) => { - agent - .handle_confirmation( - request_id, - outcome_to_confirmation(&response.outcome), - ) - .await; - Ok(()) - } + let permission = match result { + Ok(response) => outcome_to_confirmation(&response.outcome).permission, Err(e) => { error!(error = ?e, "permission request failed"); - agent - .handle_confirmation( - request_id, - PermissionConfirmation { - principal_type: PrincipalType::Tool, - permission: Permission::Cancel, - }, - ) - .await; - Ok(()) + Permission::Cancel + } + }; + + if let Err(error) = target + .agent + .submit_tool_confirmation(&target.session_id, &request_id, permission) + .await + { + error!( + session_id = %target.session_id, + request_id = %request_id, + %error, + "failed to submit tool confirmation" + ); + if let Some(cancel_token) = target.cancel_token { + cancel_token.cancel(); } } + Ok(()) })?; Ok(()) @@ -2060,9 +2084,166 @@ impl GooseAcpAgent { ) } - async fn on_load_session( + async fn send_session_usage_updates( &self, cx: &ConnectionTo, + acp_session_id: &SessionId, + session_id: &str, + agent: &Arc, + ) -> Result { + let session = self + .session_manager + .get_session(session_id, false) + .await + .internal_err_ctx("Failed to load session")?; + let totals = self + .session_manager + .get_session_usage_totals(session_id) + .await + .unwrap_or_default(); + let provider = agent + .provider() + .await + .internal_err_ctx("Failed to resolve session provider")?; + let model = session.model_config.as_ref().ok_or_else(|| { + agent_client_protocol::Error::internal_error().data("Session has no model") + })?; + let context_limit = + crate::context_limit::get_context_limit(provider.as_ref(), &model.model_name) + .await + .internal_err_ctx("Failed to resolve context limit")?; + let updates = build_usage_updates(&session, &totals, context_limit); + if self.supports_goose_custom_notifications() { + cx.send_notification(updates.custom)?; + } + cx.send_notification(SessionNotification::new( + acp_session_id.clone(), + SessionUpdate::UsageUpdate(updates.standard), + ))?; + + Ok(session) + } + + async fn forward_agent_stream( + &self, + cx: &ConnectionTo, + acp_session_id: &SessionId, + session_id: &str, + agent: &Arc, + cancel_token: &CancellationToken, + mut stream: BoxStream<'_, Result>, + ) -> Result { + let mut was_cancelled = false; + let mut output_token_limit_reached = false; + let mut tool_requests = HashMap::new(); + let mut chain_tracker = ToolChainTracker::default(); + let target = SessionAgentTarget { + agent: agent.clone(), + session_id: session_id.to_string(), + cancel_token: Some(cancel_token.clone()), + }; + + while let Some(event) = stream.next().await { + if cancel_token.is_cancelled() { + was_cancelled = true; + break; + } + + match event { + Ok(crate::agents::AgentEvent::Message(mut message)) => { + update_output_token_limit_reached(&mut output_token_limit_reached, &message); + + let sessions = self.sessions.lock().await; + if !sessions.contains_key(session_id) { + return Err(agent_client_protocol::Error::invalid_params() + .data(format!("Session not found: {session_id}"))); + } + drop(sessions); + + populate_output_token_limit_content(&mut message); + for content_item in &message.content { + if let Some(error) = prompt_error_from_message_content(content_item) { + return Err(error); + } + + if let MessageContent::ToolRequest(tool_request) = content_item { + tool_requests.insert(tool_request.id.clone(), tool_request.clone()); + } + + self.handle_message_content( + content_item, + &message, + acp_session_id, + &target, + &tool_requests, + cx, + ) + .await?; + + let ready_chain = match content_item { + MessageContent::ToolRequest(tool_request) => { + chain_tracker.record_request(tool_request.clone()); + None + } + MessageContent::ToolResponse(tool_response) => { + chain_tracker.record_response(&tool_response.id) + } + content if breaks_consecutive_tool_calls(content) => { + chain_tracker.close_current_chain() + } + _ => None, + }; + + if let Some(chain) = ready_chain { + self.spawn_ready_chain_summary(chain, agent, acp_session_id, cx); + } + } + } + Ok(crate::agents::AgentEvent::McpNotification((request_id, notification))) => { + if let Some(update) = + tool_notifications::tool_notification_update(request_id, notification) + { + let tool_call_notifier = ToolCallNotifier::new(cx, acp_session_id); + tool_call_notifier.send_update(update)?; + } + } + Ok(crate::agents::AgentEvent::MessageUsage { message_id, usage }) => { + if self.supports_goose_custom_notifications() { + cx.send_notification(GooseSessionNotification { + session_id: session_id.to_string(), + update: GooseSessionUpdate::MessageUsage(message_usage_update( + message_id, &usage, + )), + })?; + } + } + Ok(_) => {} + Err(error) => { + return Err(agent_client_protocol::Error::internal_error() + .data(format!("Error in agent response stream: {error}"))); + } + } + } + + if cancel_token.is_cancelled() { + was_cancelled = true; + } + + if !was_cancelled { + if let Some(chain) = chain_tracker.close_current_chain() { + self.spawn_ready_chain_summary(chain, agent, acp_session_id, cx); + } + } + + Ok(AgentStreamOutcome { + was_cancelled, + output_token_limit_reached, + }) + } + + async fn on_load_session( + self: &Arc, + cx: &ConnectionTo, args: LoadSessionRequest, ) -> Result { self.handle_load_session(cx, args).await @@ -2122,7 +2303,6 @@ impl GooseAcpAgent { } let user_message = Self::convert_acp_prompt_to_message(&args.prompt); - let session_config = SessionConfig { id: session_id.clone(), schedule_id: None, @@ -2130,7 +2310,7 @@ impl GooseAcpAgent { retry_config: None, }; - let mut stream = match agent + let stream = match agent .reply(user_message, session_config, Some(cancel_token.clone())) .await { @@ -2142,152 +2322,26 @@ impl GooseAcpAgent { .data(format!("Error getting agent reply: {error}"))); } }; - - let mut was_cancelled = false; - let mut output_token_limit_reached = false; - let mut tool_requests = HashMap::new(); - let mut chain_tracker = ToolChainTracker::default(); - let mut stream_error = None; - - while let Some(event) = stream.next().await { - if cancel_token.is_cancelled() { - was_cancelled = true; - break; - } - - match event { - Ok(crate::agents::AgentEvent::Message(mut message)) => { - update_output_token_limit_reached(&mut output_token_limit_reached, &message); - - let sessions = self.sessions.lock().await; - if !sessions.contains_key(&session_id) { - stream_error = Some( - agent_client_protocol::Error::invalid_params() - .data(format!("Session not found: {}", session_id)), - ); - break; - } - - populate_output_token_limit_content(&mut message); - for content_item in &message.content { - if let Some(error) = prompt_error_from_message_content(content_item) { - stream_error = Some(error); - break; - } - - if let MessageContent::ToolRequest(tool_request) = content_item { - tool_requests.insert(tool_request.id.clone(), tool_request.clone()); - } - - if let Err(error) = self - .handle_message_content( - content_item, - &message, - &args.session_id, - &agent, - &tool_requests, - cx, - ) - .await - { - stream_error = Some(error); - break; - } - - let ready_chain = match content_item { - MessageContent::ToolRequest(tool_request) => { - chain_tracker.record_request(tool_request.clone()); - None - } - MessageContent::ToolResponse(tool_response) => { - chain_tracker.record_response(&tool_response.id) - } - content if breaks_consecutive_tool_calls(content) => { - chain_tracker.close_current_chain() - } - _ => None, - }; - - if let Some(chain) = ready_chain { - self.spawn_ready_chain_summary(chain, &agent, &args.session_id, cx); - } - } - - if stream_error.is_some() { - break; - } - } - Ok(crate::agents::AgentEvent::McpNotification((request_id, notification))) => { - if let Some(update) = - tool_notifications::tool_notification_update(request_id, notification) - { - let tool_call_notifier = ToolCallNotifier::new(cx, &args.session_id); - tool_call_notifier.send_update(update)?; - } - } - Ok(crate::agents::AgentEvent::MessageUsage { message_id, usage }) => { - if self.supports_goose_custom_notifications() { - cx.send_notification(GooseSessionNotification { - session_id: session_id.clone(), - update: GooseSessionUpdate::MessageUsage(message_usage_update( - message_id, &usage, - )), - })?; - } - } - Ok(_) => {} - Err(e) => { - stream_error = Some( - agent_client_protocol::Error::internal_error() - .data(format!("Error in agent response stream: {}", e)), - ); - break; - } - } - } - - if !was_cancelled && stream_error.is_none() { - if let Some(chain) = chain_tracker.close_current_chain() { - self.spawn_ready_chain_summary(chain, &agent, &args.session_id, cx); - } - } + let stream_result = self + .forward_agent_stream( + cx, + &args.session_id, + &session_id, + &agent, + &cancel_token, + stream, + ) + .await; self.clear_active_run(&session_id, &run_id).await; Self::send_active_run_update(cx, &args.session_id, None)?; - if let Some(error) = stream_error { - return Err(error); - } + let outcome = stream_result?; let session = self - .session_manager - .get_session(&session_id, false) - .await - .internal_err_ctx("Failed to load session")?; - let totals = self - .session_manager - .get_session_usage_totals(&session_id) - .await - .unwrap_or_default(); - let provider = agent - .provider() - .await - .internal_err_ctx("Failed to resolve session provider")?; - let model = session.model_config.as_ref().ok_or_else(|| { - agent_client_protocol::Error::internal_error().data("Session has no model") - })?; - let context_limit = - crate::context_limit::get_context_limit(provider.as_ref(), &model.model_name) - .await - .internal_err_ctx("Failed to resolve context limit")?; - let updates = build_usage_updates(&session, &totals, context_limit); - if self.supports_goose_custom_notifications() { - cx.send_notification(updates.custom)?; - } - cx.send_notification(SessionNotification::new( - args.session_id.clone(), - SessionUpdate::UsageUpdate(updates.standard), - ))?; + .send_session_usage_updates(cx, &args.session_id, &session_id, &agent) + .await?; - let stop_reason = prompt_stop_reason(was_cancelled, output_token_limit_reached); + let stop_reason = + prompt_stop_reason(outcome.was_cancelled, outcome.output_token_limit_reached); let mut response = PromptResponse::new(stop_reason); if let Some(usage) = build_prompt_usage(&session) { diff --git a/crates/goose/src/acp/server/load_session.rs b/crates/goose/src/acp/server/load_session.rs index 376fa55a8..25b211fbb 100644 --- a/crates/goose/src/acp/server/load_session.rs +++ b/crates/goose/src/acp/server/load_session.rs @@ -35,17 +35,6 @@ fn messages_for_acp_replay(conversation: &Conversation) -> Vec { .collect() } -fn active_turn_messages(conversation: &Conversation) -> &[Message] { - let messages = conversation.messages(); - messages - .iter() - .rposition(|message| { - message.role == Role::User && message.is_user_visible() && !message.is_tool_response() - }) - .map(|start| &messages[start..]) - .unwrap_or(messages) -} - fn send_replay_content_chunk( cx: &ConnectionTo, session_id: &SessionId, @@ -233,67 +222,164 @@ impl GooseAcpAgent { &self, cx: &ConnectionTo, agent: &Arc, - session: &Session, + session_id: &str, + requests: &[ToolConfirmationRequest], + cancel_token: Option, ) -> Result<(), agent_client_protocol::Error> { - let session_id = SessionId::new(session.id.clone()); - let messages = session - .conversation - .as_ref() - .map(active_turn_messages) - .unwrap_or(&[]); - - let mut answered = HashSet::new(); - let mut responses = HashSet::new(); - let mut requests = Vec::new(); - - for message in messages { - for content in &message.content { - match content { - MessageContent::ToolResponse(response) => { - answered.insert(response.id.clone()); - } - MessageContent::ActionRequired(action) => match &action.data { - ActionRequiredData::ToolConfirmation { - id, - tool_name, - arguments, - prompt, - } => requests.push(( - id.clone(), - tool_name.clone(), - arguments.clone(), - prompt.clone(), - )), - ActionRequiredData::ToolConfirmationResponse { id, .. } => { - responses.insert(id.clone()); - } - _ => {} - }, - _ => {} - } - } - } - - for (id, tool_name, arguments, prompt) in requests { - if answered.contains(&id) || responses.contains(&id) { - continue; - } + let acp_session_id = SessionId::new(session_id.to_string()); + for request in requests { self.handle_tool_permission_request( cx, - agent, - &session_id, - id, - tool_name, - arguments, - prompt, + &acp_session_id, + PendingToolPermission { + request_id: request.id.clone(), + tool_name: request.tool_name.clone(), + arguments: request.arguments.clone(), + prompt: request.prompt.clone(), + }, + SessionAgentTarget { + agent: agent.clone(), + session_id: session_id.to_string(), + cancel_token: cancel_token.clone(), + }, )?; } Ok(()) } + async fn start_resumed_state_machine_turn( + self: &Arc, + cx: &ConnectionTo, + agent: &Arc, + session_id: &str, + requests: &[ToolConfirmationRequest], + ) -> Result<(), agent_client_protocol::Error> { + let run_id = format!("resume_{}", Uuid::new_v4()); + let cancel_token = CancellationToken::new(); + self.start_active_run( + session_id, + run_id.clone(), + cancel_token.clone(), + agent.clone(), + ) + .await?; + + let acp_session_id = SessionId::new(session_id.to_string()); + if let Err(error) = Self::send_active_run_update(cx, &acp_session_id, Some(&run_id)) { + self.clear_active_run(session_id, &run_id).await; + return Err(error); + } + + let session_config = SessionConfig { + id: session_id.to_string(), + schedule_id: None, + max_turns: None, + retry_config: None, + }; + let stream = match agent + .resume_state_machine_turn(session_config, cancel_token.clone()) + .await + { + Ok(Some(stream)) => stream, + Ok(None) => { + self.clear_active_run(session_id, &run_id).await; + Self::send_active_run_update(cx, &acp_session_id, None)?; + return Ok(()); + } + Err(error) => { + self.clear_active_run(session_id, &run_id).await; + let _ = Self::send_active_run_update(cx, &acp_session_id, None); + return Err(agent_client_protocol::Error::internal_error().data(format!( + "Failed to resume pending tool confirmation: {error}" + ))); + } + }; + + let server = Arc::clone(self); + let task_cx = cx.clone(); + let task_agent = agent.clone(); + let task_session_id = session_id.to_string(); + let task_run_id = run_id.clone(); + let task_cancel_token = cancel_token.clone(); + let task_acp_session_id = acp_session_id.clone(); + if let Err(error) = cx.spawn(async move { + let _run_guard = ActiveRunDropGuard { + registry: server.active_prompt_runs.clone(), + session_id: task_session_id.clone(), + run_id: task_run_id.clone(), + cancel_token: task_cancel_token.clone(), + }; + let result = server + .forward_agent_stream( + &task_cx, + &task_acp_session_id, + &task_session_id, + &task_agent, + &task_cancel_token, + stream, + ) + .await; + if result.is_ok() { + if let Err(error) = server + .send_session_usage_updates( + &task_cx, + &task_acp_session_id, + &task_session_id, + &task_agent, + ) + .await + { + warn!( + session_id = task_session_id, + ?error, + "Failed to update usage after resumed ACP turn" + ); + } + } + server + .clear_active_run(&task_session_id, &task_run_id) + .await; + if let Err(error) = Self::send_active_run_update(&task_cx, &task_acp_session_id, None) { + warn!( + session_id = task_session_id, + ?error, + "Failed to clear resumed ACP run status" + ); + } + if let Err(error) = result { + warn!( + session_id = task_session_id, + ?error, + "Resumed ACP state-machine turn failed" + ); + } + Ok(()) + }) { + cancel_token.cancel(); + self.clear_active_run(session_id, &run_id).await; + let _ = Self::send_active_run_update(cx, &SessionId::new(session_id), None); + return Err(error); + } + + if let Err(error) = self.resend_pending_tool_permissions( + cx, + agent, + session_id, + requests, + Some(cancel_token.clone()), + ) { + cancel_token.cancel(); + self.clear_active_run(session_id, &run_id).await; + let _ = Self::send_active_run_update(cx, &acp_session_id, None); + return Err(error); + } + + Ok(()) + } + pub(super) async fn handle_load_session( - &self, + self: &Arc, cx: &ConnectionTo, args: LoadSessionRequest, ) -> Result { @@ -333,8 +419,6 @@ impl GooseAcpAgent { .await .internal_err_ctx("Failed to get provider while loading ACP session")?; resume_saved_provider_session(&provider, session.conversation.as_ref()).await; - self.resend_pending_tool_permissions(cx, &agent, &session)?; - session = self .session_manager .get_session(&session_id_str, false) @@ -367,6 +451,35 @@ impl GooseAcpAgent { } response = response.meta(meta); + let pending_confirmations = session + .conversation + .as_ref() + .map(pending_tool_confirmations) + .unwrap_or_default(); + let should_resume_state_machine = crate::agents::state_machine::enabled() + && (!pending_confirmations.is_empty() + || session + .conversation + .as_ref() + .is_some_and(has_unapplied_tool_confirmation_response)); + if should_resume_state_machine { + self.start_resumed_state_machine_turn( + cx, + &agent, + &session_id_str, + &pending_confirmations, + ) + .await?; + } else { + self.resend_pending_tool_permissions( + cx, + &agent, + &session_id_str, + &pending_confirmations, + None, + )?; + } + self.closed_session_ids.lock().await.remove(&session_id_str); Ok(response) } @@ -643,21 +756,13 @@ mod tests { approval("current"), ]); - let active = active_turn_messages(&conversation); - let approval_ids = active - .iter() - .flat_map(|message| &message.content) - .filter_map(|content| match content { - MessageContent::ActionRequired(action) => match &action.data { - ActionRequiredData::ToolConfirmation { id, .. } => Some(id.as_str()), - _ => None, - }, - _ => None, - }) + let approval_ids = pending_tool_confirmations(&conversation) + .into_iter() + .map(|request| request.id) .collect::>(); assert_eq!(approval_ids, ["current"]); let no_kickoff = Conversation::new_unvalidated([approval("orphan")]); - assert_eq!(active_turn_messages(&no_kickoff).len(), 1); + assert_eq!(pending_tool_confirmations(&no_kickoff).len(), 1); } } diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 80aee5079..d8abdf0fe 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -12,6 +12,9 @@ use super::container::Container; use super::final_output_tool::FinalOutputTool; use super::gen_ai_telemetry; use super::mcp_client::GooseMcpHostInfo; +use super::tool_confirmation_coordinator::{ + ActiveTurnGuard, ConfirmationAnswer, ToolConfirmationCoordinator, +}; use super::tool_confirmation_router::ToolConfirmationRouter; use super::tool_execution::{ tool_stream, ToolCallResult, ToolStream, ToolStreamItem, CHAT_MODE_TOOL_SKIPPED_RESPONSE, @@ -30,12 +33,14 @@ use crate::agents::platform_extensions::MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE; use crate::agents::prompt_manager::PromptManager; use crate::agents::retry::{RetryManager, RetryResult}; use crate::agents::state_machine::{ - run_goose, BangShellOperation, CompactionOperation, DoctorOperation, Emitter, - EntryHookOperation, ExitOnErrorOperation, GooseEffect, GooseInferenceProvider, - GooseInferenceRequestPreparer, InferenceRunner, MaxTurnsOperation, Operation, ProjectOperation, - RecipeOperation, RetryOperation, SkillOperation, SlashCommandOperation, StateMachine, - StatusOperation, SteerOperation, SteerQueue, Step, StopHookOperation, ToolApprovalOperation, - ToolExecutionOperation, ToolPairCompactionOperation, UnknownToolOperation, MAX_TURNS_MESSAGE, + has_unapplied_tool_confirmation_response, pending_tool_confirmations, + persist_tool_confirmation_decision, run_goose, BangShellOperation, CompactionOperation, + DoctorOperation, Emitter, EntryHookOperation, ExitOnErrorOperation, GooseEffect, + GooseInferenceProvider, GooseInferenceRequestPreparer, InferenceRunner, MaxTurnsOperation, + Operation, ProjectOperation, RecipeOperation, RetryOperation, SkillOperation, + SlashCommandOperation, StateMachine, StatusOperation, SteerOperation, SteerQueue, Step, + StopHookOperation, ToolApprovalOperation, ToolExecutionOperation, ToolPairCompactionOperation, + UnknownToolOperation, MAX_TURNS_MESSAGE, }; use crate::agents::types::{ SessionConfig, SharedProvider, DEFAULT_ON_FAILURE_TIMEOUT_SECONDS, @@ -57,7 +62,7 @@ use crate::conversation::{ }; use crate::permission::permission_inspector::PermissionInspector; use crate::permission::permission_judge::PermissionCheckResult; -use crate::permission::PermissionConfirmation; +use crate::permission::{Permission, PermissionConfirmation}; use crate::providers::base::{PermissionRouting, Provider}; use crate::recipe::{Author, Recipe, Response, Settings}; use crate::scheduler_trait::SchedulerTrait; @@ -287,7 +292,8 @@ pub struct Agent { pub extension_manager: Arc, pub(super) final_output_tool: Arc>>, pub(super) prompt_manager: Mutex, - pub tool_confirmation_router: ToolConfirmationRouter, + pub(super) tool_confirmation_router: ToolConfirmationRouter, + tool_confirmation_coordinator: ToolConfirmationCoordinator, pub(super) retry_manager: RetryManager, pub(super) tool_inspection_manager: ToolInspectionManager, @@ -449,6 +455,7 @@ impl Agent { final_output_tool: Arc::new(Mutex::new(None)), prompt_manager: Mutex::new(PromptManager::new()), tool_confirmation_router: ToolConfirmationRouter::new(), + tool_confirmation_coordinator: ToolConfirmationCoordinator::new(), retry_manager: RetryManager::new(), tool_inspection_manager: Self::create_tool_inspection_manager( permission_manager, @@ -1502,29 +1509,102 @@ impl Agent { self.extension_manager.get_extension_configs().await } - /// Handle a confirmation response for a tool request - pub async fn handle_confirmation( + pub async fn submit_tool_confirmation( &self, - request_id: String, - confirmation: PermissionConfirmation, - ) { + session_id: &str, + request_id: &str, + permission: Permission, + ) -> Result<()> { + self.config + .session_manager + .get_session(session_id, false) + .await?; + + let state = self.tool_confirmation_coordinator.session(session_id); + let _confirmation_submission_guard = state.confirmation_submission_lock.lock().await; + let state_machine_permission = if permission == Permission::Cancel { + Permission::DenyOnce + } else { + permission.clone() + }; + + if let Some(answer) = state.answer(request_id) { + return match answer { + ConfirmationAnswer::LiveHandled => { + Err(anyhow!("tool confirmation request was already answered")) + } + ConfirmationAnswer::StateMachine(previous) + if previous == state_machine_permission => + { + Ok(()) + } + ConfirmationAnswer::StateMachine(_) => Err(anyhow!( + "tool confirmation request already has a different decision" + )), + }; + } + + let confirmation = PermissionConfirmation { + principal_type: crate::permission::permission_confirmation::PrincipalType::Tool, + permission: permission.clone(), + }; + if self + .try_route_tool_confirmation_to_provider(request_id, &confirmation) + .await + { + if state.contains_request(request_id) { + state.record_answer(request_id, ConfirmationAnswer::LiveHandled)?; + } + return Ok(()); + } + + if self + .tool_confirmation_router + .deliver(session_id, request_id, confirmation) + .await + { + if state.contains_request(request_id) { + state.record_answer(request_id, ConfirmationAnswer::LiveHandled)?; + } + return Ok(()); + } + + if state.contains_request(request_id) { + persist_tool_confirmation_decision( + self.config.session_manager.as_ref(), + session_id, + request_id, + &state_machine_permission, + ) + .await?; + state.record_answer( + request_id, + ConfirmationAnswer::StateMachine(state_machine_permission), + )?; + return Ok(()); + } + + Err(anyhow!( + "unknown or stale tool confirmation request {request_id} for session {session_id}" + )) + } + + async fn try_route_tool_confirmation_to_provider( + &self, + request_id: &str, + confirmation: &PermissionConfirmation, + ) -> bool { let provider = self.provider.lock().await.clone(); if let Some(provider) = provider.as_ref() { if provider.permission_routing() == PermissionRouting::ActionRequired && provider - .handle_permission_confirmation(&request_id, &confirmation) + .handle_permission_confirmation(request_id, confirmation) .await { - return; + return true; } } - if !self - .tool_confirmation_router - .deliver(request_id, confirmation) - .await - { - error!("Failed to deliver confirmation"); - } + false } pub async fn supports_action_required_permissions(&self) -> bool { @@ -1668,10 +1748,12 @@ impl Agent { cancel_token: Option, ) -> Result>> { let session_manager = self.config.session_manager.clone(); - let cancel = cancel_token.unwrap_or_default(); let session_id = session_config.id.clone(); + let turn_guard = self + .tool_confirmation_coordinator + .session(&session_id) + .try_start_turn()?; - let entry_session = session_manager.get_session(&session_id, false).await?; if let Some(schedule_id) = session_config.schedule_id.clone() { session_manager .update(&session_id) @@ -1683,14 +1765,13 @@ impl Agent { .add_message(&session_config.id, &user_message) .await?; - let provider = self - .provider - .lock() - .await - .clone() - .ok_or_else(|| anyhow!("Provider not set"))?; - if !self.config.disable_session_naming { + let provider = self + .provider + .lock() + .await + .clone() + .ok_or_else(|| anyhow!("Provider not set"))?; let manager = session_manager.clone(); let tx = self.config.session_name_update_tx.clone(); let id = session_id.clone(); @@ -1710,6 +1791,158 @@ impl Agent { }); } + let cancel = cancel_token.unwrap_or_default(); + let initial_stream = self + .stream_state_machine_session(session_config.clone(), cancel.clone()) + .await?; + Ok( + self.stream_state_machine_turn( + session_config, + cancel, + turn_guard, + Some(initial_stream), + ), + ) + } + + pub(crate) async fn resume_state_machine_turn( + self: &Arc, + session_config: SessionConfig, + cancel: CancellationToken, + ) -> Result>>> { + if !super::state_machine::enabled() { + return Ok(None); + } + + let session = self + .config + .session_manager + .get_session(&session_config.id, true) + .await?; + let conversation = session + .conversation + .as_ref() + .ok_or_else(|| anyhow!("Session {} has no conversation", session_config.id))?; + let pending_confirmations = pending_tool_confirmations(conversation); + let resume_from_persisted_response = pending_confirmations.is_empty() + && has_unapplied_tool_confirmation_response(conversation); + if pending_confirmations.is_empty() && !resume_from_persisted_response { + return Ok(None); + } + + let turn_guard = self + .tool_confirmation_coordinator + .session(&session_config.id) + .try_start_turn()?; + for request in pending_confirmations { + turn_guard.state().register_request(request.id); + } + + let agent = Arc::clone(self); + Ok(Some(Box::pin(async_stream::try_stream! { + let initial_stream = if resume_from_persisted_response { + Some( + agent + .stream_state_machine_session( + session_config.clone(), + cancel.clone(), + ) + .await?, + ) + } else { + None + }; + let mut stream = agent.stream_state_machine_turn( + session_config, + cancel, + turn_guard, + initial_stream, + ); + while let Some(event) = stream.next().await { + yield event?; + } + }))) + } + + fn tool_confirmation_request_ids(event: &AgentEvent) -> Vec { + let AgentEvent::Message(message) = event else { + return Vec::new(); + }; + + message + .content + .iter() + .filter_map(|content| { + let MessageContent::ActionRequired(action) = content else { + return None; + }; + let ActionRequiredData::ToolConfirmation { id, .. } = &action.data else { + return None; + }; + Some(id.clone()) + }) + .collect() + } + + fn stream_state_machine_turn<'a>( + &'a self, + session_config: SessionConfig, + cancel: CancellationToken, + turn_guard: ActiveTurnGuard, + initial_stream: Option>>, + ) -> BoxStream<'a, Result> { + Box::pin(async_stream::try_stream! { + let mut stream = initial_stream; + loop { + if let Some(active_stream) = stream.as_mut() { + let mut has_confirmations = false; + while let Some(event) = active_stream.next().await { + let event = event?; + for request_id in Self::tool_confirmation_request_ids(&event) { + turn_guard.state().register_request(request_id); + has_confirmations = true; + } + yield event; + } + + if !has_confirmations { + return; + } + } + + let has_state_machine_answer = turn_guard + .state() + .wait_for_all_confirmation_answers(&cancel) + .await?; + if !has_state_machine_answer { + return; + } + turn_guard.state().clear_confirmations(); + stream = Some( + self.stream_state_machine_session( + session_config.clone(), + cancel.clone(), + ) + .await?, + ); + } + }) + } + + async fn stream_state_machine_session( + &self, + session_config: SessionConfig, + cancel: CancellationToken, + ) -> Result>> { + let session_manager = self.config.session_manager.clone(); + let session_id = session_config.id.clone(); + let entry_session = session_manager.get_session(&session_id, false).await?; + let provider = self + .provider + .lock() + .await + .clone() + .ok_or_else(|| anyhow!("Provider not set"))?; let model_config = match entry_session.model_config { Some(model_config) => model_config, None => { @@ -3992,7 +4225,6 @@ fn recipe_conversation_history(messages: &Conversation) -> Vec { mod tests { use super::*; use crate::agents::gen_ai_telemetry::{self, test_support::SpanFieldCapture}; - use crate::permission::permission_confirmation::PrincipalType; use crate::plugins::discovery::{DiscoveredPlugin, PluginScope}; use crate::providers::base::{stream_from_single_message, MessageStream, PermissionRouting}; use crate::recipe::Response; @@ -4334,39 +4566,37 @@ mod tests { } #[tokio::test] - async fn test_handle_confirmation_routes_to_provider() { - let agent = Agent::new(); + async fn test_submit_tool_confirmation_routes_to_provider() { + let (agent, session, _data_dir) = tracing_test_agent_and_session().await; let provider = Arc::new(ActionRequiredProvider::new()); *agent.provider.lock().await = Some(provider.clone() as Arc); // Known request_id → provider handles it, confirmation_router NOT called agent - .handle_confirmation( - "known".to_string(), - PermissionConfirmation { - principal_type: PrincipalType::Tool, - permission: crate::permission::Permission::AllowOnce, - }, + .submit_tool_confirmation( + &session.id, + "known", + crate::permission::Permission::AllowOnce, ) - .await; + .await + .unwrap(); assert_eq!(provider.handled.lock().await.len(), 1); // Unknown request_id → provider returns false, falls through to confirmation_router // Register first so deliver() has somewhere to send let rx = agent .tool_confirmation_router - .register("unknown".to_string()) + .register(session.id.clone(), "unknown".to_string()) .await; agent - .handle_confirmation( - "unknown".to_string(), - PermissionConfirmation { - principal_type: PrincipalType::Tool, - permission: crate::permission::Permission::DenyOnce, - }, + .submit_tool_confirmation( + &session.id, + "unknown", + crate::permission::Permission::DenyOnce, ) - .await; + .await + .unwrap(); assert_eq!(provider.handled.lock().await.len(), 2); // Verify the fallthrough went to confirmation_router let conf = rx.await.unwrap(); @@ -4374,23 +4604,18 @@ mod tests { } #[tokio::test] - async fn test_handle_confirmation_noop_provider() { - let agent = Agent::new(); + async fn test_submit_tool_confirmation_routes_to_legacy() { + let (agent, session, _data_dir) = tracing_test_agent_and_session().await; // No provider set → Noop routing, goes straight to confirmation_router // Register first so deliver() has somewhere to send let rx = agent .tool_confirmation_router - .register("any".to_string()) + .register(session.id.clone(), "any".to_string()) .await; agent - .handle_confirmation( - "any".to_string(), - PermissionConfirmation { - principal_type: PrincipalType::Tool, - permission: crate::permission::Permission::AllowOnce, - }, - ) - .await; + .submit_tool_confirmation(&session.id, "any", crate::permission::Permission::AllowOnce) + .await + .unwrap(); let conf = rx.await.unwrap(); assert_eq!(conf.permission, crate::permission::Permission::AllowOnce); diff --git a/crates/goose/src/agents/mod.rs b/crates/goose/src/agents/mod.rs index 10269dd2c..498199212 100644 --- a/crates/goose/src/agents/mod.rs +++ b/crates/goose/src/agents/mod.rs @@ -19,6 +19,7 @@ pub mod state_machine; pub mod subagent_execution_tool; pub(crate) mod subagent_handler; pub(crate) mod subagent_task_config; +mod tool_confirmation_coordinator; mod tool_confirmation_router; pub mod tool_execution; mod tool_schema_normalize; diff --git a/crates/goose/src/agents/state_machine/mod.rs b/crates/goose/src/agents/state_machine/mod.rs index 802f4cef8..63ac285bd 100644 --- a/crates/goose/src/agents/state_machine/mod.rs +++ b/crates/goose/src/agents/state_machine/mod.rs @@ -28,6 +28,7 @@ mod ops_toolcalling; mod ops_unknown_tool; mod session; pub(crate) use session::run as run_goose; +mod tool_confirmation; mod usage; #[cfg(test)] @@ -42,9 +43,14 @@ pub use goose_agent::operation::{ not_applicable, trailing_error, yielded, yielded_with, ConversationEffect, Emitter, Inference, InferenceInput, MachineEffect, Operation, OperationResult, SlashCommand, StepResult, }; +pub(crate) use tool_confirmation::{ + has_unapplied_tool_confirmation_response, pending_tool_confirmations, + persist_tool_confirmation_decision, +}; pub(super) use inference_preparation::GooseInferenceRequestPreparer; -pub(super) use ops_bang_shell::{bang_shell_command, BangShellOperation}; +pub(crate) use ops_bang_shell::bang_shell_command; +pub(super) use ops_bang_shell::BangShellOperation; pub(super) use ops_compaction::CompactionOperation; pub(super) use ops_doctor::DoctorOperation; pub(super) use ops_entry_hook::EntryHookOperation; diff --git a/crates/goose/src/agents/state_machine/ops_tool_approval.rs b/crates/goose/src/agents/state_machine/ops_tool_approval.rs index aec3ba9fe..113bfecee 100644 --- a/crates/goose/src/agents/state_machine/ops_tool_approval.rs +++ b/crates/goose/src/agents/state_machine/ops_tool_approval.rs @@ -50,7 +50,7 @@ impl Operation for ToolApprovalOperation<'_> { &self, session: &Session, conversation: &Conversation, - emit: &Emitter, + _emit: &Emitter, ) -> Result> { let goose_mode = *self.goose_mode.lock().await; if goose_mode == GooseMode::Chat { @@ -130,7 +130,6 @@ impl Operation for ToolApprovalOperation<'_> { security_message, ) .user_only(); - let action_required = emit.message(action_required).await; effects.push(action_required.into()); if let Some(finding_id) = @@ -162,7 +161,7 @@ struct PendingResponse { executable: Option, } -struct ApprovalState { +pub(super) struct ApprovalState { answered: HashSet, approval_requests: HashSet, approval_responses: HashMap, @@ -170,7 +169,7 @@ struct ApprovalState { } impl ApprovalState { - fn from_messages(messages: &[Message]) -> Self { + pub(super) fn from_messages(messages: &[Message]) -> Self { let mut answered = HashSet::new(); let mut approval_requests = HashSet::new(); let mut approval_responses = HashMap::new(); @@ -249,6 +248,30 @@ impl ApprovalState { }) .collect() } + + pub(super) fn has_confirmation_request(&self, request_id: &str) -> bool { + self.approval_requests.contains(request_id) + && self + .tool_requests + .iter() + .any(|request| request.id == request_id) + } + + pub(super) fn has_unapplied_confirmation_response(&self) -> bool { + self.tool_requests.iter().any(|request| { + !self.answered.contains(&request.id) + && self.approval_requests.contains(&request.id) + && self.approval_responses.contains_key(&request.id) + }) + } + + pub(super) fn confirmation_response(&self, request_id: &str) -> Option<&Permission> { + self.approval_responses.get(request_id) + } + + pub(super) fn has_tool_response(&self, request_id: &str) -> bool { + self.answered.contains(request_id) + } } pub fn request_executable(request: &ToolRequest) -> Option { diff --git a/crates/goose/src/agents/state_machine/session.rs b/crates/goose/src/agents/state_machine/session.rs index 779c94f81..8cbddefad 100644 --- a/crates/goose/src/agents/state_machine/session.rs +++ b/crates/goose/src/agents/state_machine/session.rs @@ -4,11 +4,22 @@ use async_trait::async_trait; use crate::agents::state_machine::effects::GooseEffect; use crate::agents::state_machine::usage; use crate::agents::AgentEvent; +use crate::conversation::message::{ActionRequiredData, Message, MessageContent}; use crate::conversation::Conversation; use crate::session::{Session, SessionManager}; use goose_agent::machine::{EffectHandler, EffectUsage, MachineSession, SessionLoader}; use goose_agent::operation::{ConversationEffect, Emitter, MachineEffect}; +fn contains_tool_confirmation_request(message: &Message) -> bool { + message.content.iter().any(|content| { + matches!( + content, + MessageContent::ActionRequired(action) + if matches!(&action.data, ActionRequiredData::ToolConfirmation { .. }) + ) + }) +} + impl MachineSession for Session { fn id(&self) -> &str { &self.id @@ -105,6 +116,10 @@ impl EffectHandler for SessionManager { for effect in effects { match effect { GooseEffect::Conversation(ConversationEffect::AppendMessage(message)) => { + if contains_tool_confirmation_request(message) { + // Responses can arrive immediately, so publish only after the persistence pass. + emit.emit(AgentEvent::Message(message.clone())).await; + } if let Some(usage) = message .metadata .usage diff --git a/crates/goose/src/agents/state_machine/tests/agent_reply.rs b/crates/goose/src/agents/state_machine/tests/agent_reply.rs index cf79f4d76..8a0d5be53 100644 --- a/crates/goose/src/agents/state_machine/tests/agent_reply.rs +++ b/crates/goose/src/agents/state_machine/tests/agent_reply.rs @@ -13,12 +13,16 @@ use anyhow::Result; use futures::StreamExt; use tokio_util::sync::CancellationToken; +use super::calculator_extension::{value, CalculatorExtension, ADD}; use super::dummy_api::{DummyApi, ProviderFeatures}; use crate::acp::server::GooseAcpAgent; +use crate::agents::extension::ExtensionConfig; +use crate::agents::mcp_client::McpClientTrait; use crate::agents::{Agent, AgentConfig, AgentEvent, GoosePlatform, SessionConfig}; use crate::config::permission::PermissionManager; use crate::config::GooseMode; -use crate::conversation::message::{Message, MessageContent}; +use crate::conversation::message::{ActionRequiredData, Message, MessageContent}; +use crate::permission::Permission; use crate::providers::base::Provider; use crate::session::{SessionManager, SessionType}; use goose_providers::model::ModelConfig; @@ -48,7 +52,7 @@ async fn agent_with_dummy_api() -> Result<(Agent, Arc, String, tempfil .await?; let agent = Agent::with_config(AgentConfig::new( session_manager, - PermissionManager::instance(), + Arc::new(PermissionManager::new(temp_dir.path().join("permissions"))), None, GooseMode::Auto, true, @@ -66,6 +70,220 @@ async fn agent_with_dummy_api() -> Result<(Agent, Arc, String, tempfil Ok((agent, api, session.id, temp_dir)) } +async fn agent_with_calculator() -> Result<( + Agent, + Arc, + String, + Arc, + tempfile::TempDir, +)> { + let (agent, api, session_id, temp_dir) = agent_with_dummy_api().await?; + agent + .update_goose_mode(GooseMode::Approve, &session_id) + .await?; + let calculator = Arc::new(CalculatorExtension::new( + agent.config.session_manager.action_required(), + )); + agent + .extension_manager + .add_client( + "calculator".to_string(), + ExtensionConfig::Platform { + name: "calculator".to_string(), + description: "Stateful test calculator".to_string(), + display_name: None, + bundled: None, + available_tools: vec![], + }, + calculator.clone(), + calculator.get_info().cloned(), + ) + .await; + Ok((agent, api, session_id, calculator, temp_dir)) +} + +fn confirmation_ids(messages: &[Message]) -> Vec { + messages + .iter() + .flat_map(|message| &message.content) + .filter_map(|content| match content { + MessageContent::ActionRequired(action) => match &action.data { + ActionRequiredData::ToolConfirmation { id, .. } => Some(id.clone()), + _ => None, + }, + _ => None, + }) + .collect() +} + +async fn stream_messages( + mut stream: futures::stream::BoxStream<'_, Result>, +) -> Result> { + let mut messages = Vec::new(); + while let Some(event) = stream.next().await { + if let AgentEvent::Message(message) = event? { + messages.push(message); + } + } + Ok(messages) +} + +#[tokio::test] +async fn state_machine_confirmation_through_agent_resumes_tool_call() -> Result<()> { + let _guard = env_lock::lock_env([("GOOSE_STATE_MACHINE", Some("1"))]); + let (agent, api, session_id, calculator, _temp_dir) = agent_with_calculator().await?; + let agent = Arc::new(agent); + + api.on("add one").call(ADD, value(1)); + api.on("result: 1").reply("the result is one"); + + let session_config = SessionConfig { + id: session_id, + schedule_id: None, + max_turns: Some(2), + retry_config: None, + }; + let mut stream = agent + .reply( + Message::user().with_text("add one"), + session_config.clone(), + Some(CancellationToken::new()), + ) + .await?; + let mut messages = Vec::new(); + let confirmation_id = loop { + let event = stream + .next() + .await + .expect("state machine should request confirmation")?; + if let AgentEvent::Message(message) = event { + let confirmation_id = confirmation_ids(std::slice::from_ref(&message)).pop(); + messages.push(message); + if let Some(confirmation_id) = confirmation_id { + break confirmation_id; + } + } + }; + assert_eq!(calculator.total(), 0); + { + let session = agent + .config + .session_manager + .get_session(&session_config.id, true) + .await?; + assert!(confirmation_ids( + session + .conversation + .as_ref() + .expect("session conversation") + .messages() + ) + .contains(&confirmation_id)); + } + + agent + .submit_tool_confirmation(&session_config.id, &confirmation_id, Permission::AllowOnce) + .await?; + { + let session = agent + .config + .session_manager + .get_session(&session_config.id, true) + .await?; + assert!(session + .conversation + .as_ref() + .expect("session conversation") + .messages() + .iter() + .any(|message| { + message.content.iter().any(|content| { + matches!( + content, + MessageContent::ActionRequired(action) + if matches!( + &action.data, + ActionRequiredData::ToolConfirmationResponse { id, permission } + if id == &confirmation_id && permission == &Permission::AllowOnce + ) + ) + }) + })); + } + agent + .submit_tool_confirmation(&session_config.id, &confirmation_id, Permission::AllowOnce) + .await?; + assert!(agent + .submit_tool_confirmation(&session_config.id, &confirmation_id, Permission::DenyOnce) + .await + .is_err()); + drop(stream); + let stream = agent + .resume_state_machine_turn(session_config.clone(), CancellationToken::new()) + .await? + .expect("persisted confirmation response should resume the state-machine turn"); + messages.extend(stream_messages(stream).await?); + assert!(messages.iter().any(|message| message + .get_tool_response_ids() + .contains(&confirmation_id.as_str()))); + assert_eq!(calculator.total(), 1); + assert_eq!(api.call_count(), 2); + + assert!(agent + .submit_tool_confirmation(&session_config.id, &confirmation_id, Permission::AllowOnce) + .await + .is_err()); + assert_eq!(calculator.total(), 1); + + assert!(agent + .submit_tool_confirmation(&session_config.id, "stale-request", Permission::AllowOnce) + .await + .is_err()); + + let session = agent + .config + .session_manager + .get_session(&session_config.id, true) + .await?; + let messages = session + .conversation + .as_ref() + .expect("session conversation") + .messages(); + let confirmation_responses = messages + .iter() + .filter(|message| { + message.content.iter().any(|content| { + matches!( + content, + MessageContent::ActionRequired(action) + if matches!( + &action.data, + ActionRequiredData::ToolConfirmationResponse { id, .. } + if id == &confirmation_id + ) + ) + }) + }) + .collect::>(); + assert_eq!(confirmation_responses.len(), 1); + assert!(!confirmation_responses[0].is_user_visible()); + assert!(!confirmation_responses[0].is_agent_visible()); + assert_eq!( + messages + .iter() + .filter(|message| { + message.role == rmcp::model::Role::User + && message.is_user_visible() + && !message.is_tool_response() + }) + .count(), + 1 + ); + + Ok(()) +} + #[tokio::test] async fn reply_streams_the_turn_and_ends() -> Result<()> { let (agent, api, session_id, _temp_dir) = agent_with_dummy_api().await?; diff --git a/crates/goose/src/agents/state_machine/tool_confirmation.rs b/crates/goose/src/agents/state_machine/tool_confirmation.rs new file mode 100644 index 000000000..8cdc4b377 --- /dev/null +++ b/crates/goose/src/agents/state_machine/tool_confirmation.rs @@ -0,0 +1,117 @@ +use std::collections::HashSet; + +use anyhow::{anyhow, Result}; + +use crate::agents::state_machine::messages_since_kickoff; +use crate::agents::state_machine::ops_tool_approval::ApprovalState; +use crate::conversation::message::{ + ActionRequiredData, Message, MessageContent, ToolConfirmationRequest, +}; +use crate::conversation::Conversation; +use crate::permission::Permission; +use crate::session::SessionManager; + +fn active_turn_messages(conversation: &Conversation) -> &[Message] { + let messages = conversation.messages(); + messages + .iter() + .rposition(|message| { + message.role == rmcp::model::Role::User + && message.is_user_visible() + && !message.is_tool_response() + }) + .map(|start| &messages[start..]) + .unwrap_or(messages) +} + +pub(crate) fn pending_tool_confirmations( + conversation: &Conversation, +) -> Vec { + let mut answered = HashSet::new(); + let mut confirmation_responses = HashSet::new(); + let mut confirmation_requests = Vec::new(); + + for message in active_turn_messages(conversation) { + for content in &message.content { + match content { + MessageContent::ToolResponse(response) => { + answered.insert(response.id.clone()); + } + MessageContent::ActionRequired(action) => match &action.data { + ActionRequiredData::ToolConfirmation { + id, + tool_name, + arguments, + prompt, + } => confirmation_requests.push(ToolConfirmationRequest { + id: id.clone(), + tool_name: tool_name.clone(), + arguments: arguments.clone(), + prompt: prompt.clone(), + }), + ActionRequiredData::ToolConfirmationResponse { id, .. } => { + confirmation_responses.insert(id.clone()); + } + _ => {} + }, + _ => {} + } + } + } + + confirmation_requests + .into_iter() + .filter(|request| { + !answered.contains(&request.id) && !confirmation_responses.contains(&request.id) + }) + .collect() +} + +pub(crate) fn has_unapplied_tool_confirmation_response(conversation: &Conversation) -> bool { + ApprovalState::from_messages(active_turn_messages(conversation)) + .has_unapplied_confirmation_response() +} + +pub(crate) async fn persist_tool_confirmation_decision( + session_manager: &SessionManager, + session_id: &str, + request_id: &str, + permission: &Permission, +) -> Result<()> { + let session = session_manager.get_session(session_id, true).await?; + let conversation = session + .conversation + .as_ref() + .ok_or_else(|| anyhow!("Session {session_id} has no conversation"))?; + let approval_state = ApprovalState::from_messages(messages_since_kickoff(conversation)?); + + if !approval_state.has_confirmation_request(request_id) { + return Err(anyhow!( + "Tool request {request_id} is not awaiting confirmation in session {session_id}" + )); + } + + if let Some(persisted_permission) = approval_state.confirmation_response(request_id) { + return if persisted_permission == permission { + Ok(()) + } else { + Err(anyhow!( + "Tool request {request_id} already has a different confirmation decision" + )) + }; + } + + if approval_state.has_tool_response(request_id) { + return Err(anyhow!("Tool request {request_id} already has a response")); + } + + let response_message = Message::user().with_visibility(false, false).with_content( + MessageContent::action_required_tool_confirmation_response( + request_id.to_string(), + permission.clone(), + ), + ); + session_manager + .add_message(session_id, &response_message) + .await +} diff --git a/crates/goose/src/agents/tool_confirmation_coordinator.rs b/crates/goose/src/agents/tool_confirmation_coordinator.rs new file mode 100644 index 000000000..10ebb9829 --- /dev/null +++ b/crates/goose/src/agents/tool_confirmation_coordinator.rs @@ -0,0 +1,245 @@ +use std::collections::HashMap; +use std::sync::{Arc, Mutex as StdMutex}; + +use anyhow::{anyhow, Result}; +use tokio::sync::{Mutex, Notify, OwnedMutexGuard}; +use tokio_util::sync::CancellationToken; + +use crate::permission::Permission; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) enum ConfirmationAnswer { + LiveHandled, + StateMachine(Permission), +} + +pub(super) struct SessionToolConfirmationState { + // Held for the lifetime of one Agent::reply stream, including confirmation waits and resumes. + turn_lock: Arc>, + // Serializes submit_tool_confirmation so concurrent answers cannot both be accepted. + pub(super) confirmation_submission_lock: Mutex<()>, + // Tracks requests from the current confirmation pause; None means still unanswered. + confirmations: StdMutex>>, + // Wakes wait_for_all_confirmation_answers; confirmations remains the source of truth. + confirmation_answered: Notify, +} + +impl SessionToolConfirmationState { + fn new() -> Self { + Self { + turn_lock: Arc::new(Mutex::new(())), + confirmation_submission_lock: Mutex::new(()), + confirmations: StdMutex::new(HashMap::new()), + confirmation_answered: Notify::new(), + } + } + + pub(super) fn try_start_turn(self: &Arc) -> Result { + let turn_lock_guard = self + .turn_lock + .clone() + .try_lock_owned() + .map_err(|_| anyhow!("session already has an active turn"))?; + Ok(ActiveTurnGuard { + state: self.clone(), + _turn_lock_guard: turn_lock_guard, + }) + } + + pub(super) fn register_request(&self, request_id: String) { + self.confirmations + .lock() + .expect("tool confirmation state unavailable") + .entry(request_id) + .or_insert(None); + } + + pub(super) fn answer(&self, request_id: &str) -> Option { + self.confirmations + .lock() + .expect("tool confirmation state unavailable") + .get(request_id) + .and_then(Clone::clone) + } + + pub(super) fn contains_request(&self, request_id: &str) -> bool { + self.confirmations + .lock() + .expect("tool confirmation state unavailable") + .contains_key(request_id) + } + + pub(super) fn record_answer(&self, request_id: &str, answer: ConfirmationAnswer) -> Result<()> { + let mut confirmations = self + .confirmations + .lock() + .expect("tool confirmation state unavailable"); + match confirmations.get_mut(request_id) { + None => return Err(anyhow!("tool confirmation request is no longer active")), + Some(Some(_)) => return Err(anyhow!("tool confirmation request was already answered")), + Some(slot @ None) => *slot = Some(answer), + } + drop(confirmations); + self.confirmation_answered.notify_waiters(); + Ok(()) + } + + pub(super) async fn wait_for_all_confirmation_answers( + &self, + cancel: &CancellationToken, + ) -> Result { + loop { + let answer_received = self.confirmation_answered.notified(); + tokio::pin!(answer_received); + answer_received.as_mut().enable(); + + let completed = { + let confirmations = self + .confirmations + .lock() + .expect("tool confirmation state unavailable"); + (!confirmations.is_empty() && confirmations.values().all(Option::is_some)).then( + || { + confirmations.values().any(|answer| { + matches!(answer, Some(ConfirmationAnswer::StateMachine(_))) + }) + }, + ) + }; + if let Some(has_state_machine_answer) = completed { + return Ok(has_state_machine_answer); + } + + tokio::select! { + _ = answer_received => {} + _ = cancel.cancelled() => return Err(anyhow!("state-machine turn cancelled")), + } + } + } + + pub(super) fn clear_confirmations(&self) { + self.confirmations + .lock() + .expect("tool confirmation state unavailable") + .clear(); + } +} + +pub(super) struct ActiveTurnGuard { + state: Arc, + _turn_lock_guard: OwnedMutexGuard<()>, +} + +impl ActiveTurnGuard { + pub(super) fn state(&self) -> &Arc { + &self.state + } +} + +impl Drop for ActiveTurnGuard { + fn drop(&mut self) { + self.state.clear_confirmations(); + } +} + +pub(super) struct ToolConfirmationCoordinator { + sessions: StdMutex>>, +} + +impl ToolConfirmationCoordinator { + pub(super) fn new() -> Self { + Self { + sessions: StdMutex::new(HashMap::new()), + } + } + + pub(super) fn session(&self, session_id: &str) -> Arc { + self.sessions + .lock() + .expect("tool confirmation coordinator unavailable") + .entry(session_id.to_string()) + .or_insert_with(|| Arc::new(SessionToolConfirmationState::new())) + .clone() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_a_second_active_turn_and_releases_on_drop() { + let coordinator = ToolConfirmationCoordinator::new(); + let session = coordinator.session("session"); + let guard = session.try_start_turn().unwrap(); + + assert!(session.try_start_turn().is_err()); + + drop(guard); + assert!(session.try_start_turn().is_ok()); + } + + #[tokio::test] + async fn waits_for_every_confirmation_in_the_batch() { + let coordinator = ToolConfirmationCoordinator::new(); + let session = coordinator.session("session"); + let _guard = session.try_start_turn().unwrap(); + session.register_request("request-1".to_string()); + session.register_request("request-2".to_string()); + session + .record_answer( + "request-1", + ConfirmationAnswer::StateMachine(Permission::AllowOnce), + ) + .unwrap(); + session + .record_answer("request-2", ConfirmationAnswer::LiveHandled) + .unwrap(); + + let has_state_machine_answer = session + .wait_for_all_confirmation_answers(&CancellationToken::new()) + .await + .unwrap(); + + assert!(has_state_machine_answer); + } + + #[test] + fn active_turn_drop_clears_pending_requests() { + let coordinator = ToolConfirmationCoordinator::new(); + let session = coordinator.session("session"); + let guard = session.try_start_turn().unwrap(); + session.register_request("request".to_string()); + assert!(session.contains_request("request")); + + drop(guard); + + assert!(!session.contains_request("request")); + assert!(session.answer("request").is_none()); + } + + #[test] + fn first_answer_is_immutable() { + let coordinator = ToolConfirmationCoordinator::new(); + let session = coordinator.session("session"); + let _guard = session.try_start_turn().unwrap(); + session.register_request("request".to_string()); + session + .record_answer( + "request", + ConfirmationAnswer::StateMachine(Permission::AllowOnce), + ) + .unwrap(); + + assert!(session + .record_answer( + "request", + ConfirmationAnswer::StateMachine(Permission::DenyOnce), + ) + .is_err()); + assert_eq!( + session.answer("request"), + Some(ConfirmationAnswer::StateMachine(Permission::AllowOnce)) + ); + } +} diff --git a/crates/goose/src/agents/tool_confirmation_router.rs b/crates/goose/src/agents/tool_confirmation_router.rs index 2148e5058..31283424b 100644 --- a/crates/goose/src/agents/tool_confirmation_router.rs +++ b/crates/goose/src/agents/tool_confirmation_router.rs @@ -5,27 +5,37 @@ use tracing::warn; use crate::permission::PermissionConfirmation; -pub struct ToolConfirmationRouter { - pending: Mutex>>, +pub(super) struct ToolConfirmationRouter { + pending: Mutex>>, } impl ToolConfirmationRouter { - pub fn new() -> Self { + pub(super) fn new() -> Self { Self { pending: Mutex::new(HashMap::new()), } } - pub async fn register(&self, request_id: String) -> oneshot::Receiver { + pub(super) async fn register( + &self, + session_id: String, + request_id: String, + ) -> oneshot::Receiver { let (tx, rx) = oneshot::channel(); let mut pending = self.pending.lock().await; pending.retain(|_, sender| !sender.is_closed()); - pending.insert(request_id, tx); + pending.insert((session_id, request_id), tx); rx } - pub async fn deliver(&self, request_id: String, confirmation: PermissionConfirmation) -> bool { - if let Some(tx) = self.pending.lock().await.remove(&request_id) { + pub(super) async fn deliver( + &self, + session_id: &str, + request_id: &str, + confirmation: PermissionConfirmation, + ) -> bool { + let key = (session_id.to_string(), request_id.to_string()); + if let Some(tx) = self.pending.lock().await.remove(&key) { if tx.send(confirmation).is_err() { warn!( request_id = %request_id, @@ -36,10 +46,6 @@ impl ToolConfirmationRouter { true } } else { - warn!( - request_id = %request_id, - "No task waiting for confirmation" - ); false } } @@ -61,10 +67,12 @@ mod tests { #[tokio::test] async fn test_register_then_deliver() { let router = ToolConfirmationRouter::new(); - let rx = router.register("req_1".to_string()).await; + let rx = router + .register("session_1".to_string(), "req_1".to_string()) + .await; assert!( router - .deliver("req_1".to_string(), test_confirmation()) + .deliver("session_1", "req_1", test_confirmation()) .await ); let confirmation = rx.await.unwrap(); @@ -76,19 +84,36 @@ mod tests { let router = ToolConfirmationRouter::new(); assert!( !router - .deliver("unknown".to_string(), test_confirmation()) + .deliver("session_1", "unknown", test_confirmation()) .await ); } + #[tokio::test] + async fn test_request_cannot_be_delivered_from_another_session() { + let router = ToolConfirmationRouter::new(); + let _rx = router + .register("session_1".to_string(), "req_1".to_string()) + .await; + + assert!( + !router + .deliver("session_2", "req_1", test_confirmation()) + .await + ); + assert_eq!(router.pending.lock().await.len(), 1); + } + #[tokio::test] async fn test_cancelled_receiver() { let router = ToolConfirmationRouter::new(); - let rx = router.register("req_1".to_string()).await; + let rx = router + .register("session_1".to_string(), "req_1".to_string()) + .await; drop(rx); // simulate task cancellation assert!( !router - .deliver("req_1".to_string(), test_confirmation()) + .deliver("session_1", "req_1", test_confirmation()) .await ); } @@ -96,14 +121,22 @@ mod tests { #[tokio::test] async fn test_stale_entries_pruned_on_register() { let router = ToolConfirmationRouter::new(); - let rx = router.register("req_1".to_string()).await; + let rx = router + .register("session_1".to_string(), "req_1".to_string()) + .await; drop(rx); // simulate task cancellation — entry is now stale assert_eq!(router.pending.lock().await.len(), 1); - let _rx2 = router.register("req_2".to_string()).await; + let _rx2 = router + .register("session_1".to_string(), "req_2".to_string()) + .await; assert_eq!(router.pending.lock().await.len(), 1); // only req_2 remains - assert!(router.pending.lock().await.contains_key("req_2")); + assert!(router + .pending + .lock() + .await + .contains_key(&("session_1".to_string(), "req_2".to_string()))); } #[tokio::test] @@ -113,14 +146,19 @@ mod tests { let router = Arc::new(ToolConfirmationRouter::new()); // Register two requests - let rx1 = router.register("req_1".to_string()).await; - let rx2 = router.register("req_2".to_string()).await; + let rx1 = router + .register("session_1".to_string(), "req_1".to_string()) + .await; + let rx2 = router + .register("session_1".to_string(), "req_2".to_string()) + .await; // Deliver in reverse order assert!( router .deliver( - "req_2".to_string(), + "session_1", + "req_2", PermissionConfirmation { principal_type: PrincipalType::Tool, permission: Permission::DenyOnce, @@ -131,7 +169,7 @@ mod tests { assert_eq!(router.pending.lock().await.len(), 1); assert!( router - .deliver("req_1".to_string(), test_confirmation()) + .deliver("session_1", "req_1", test_confirmation()) .await ); assert_eq!(router.pending.lock().await.len(), 0); diff --git a/crates/goose/src/agents/tool_execution.rs b/crates/goose/src/agents/tool_execution.rs index 9639d3d76..482ce7da3 100644 --- a/crates/goose/src/agents/tool_execution.rs +++ b/crates/goose/src/agents/tool_execution.rs @@ -168,7 +168,10 @@ impl Agent { } }); - let confirmation_rx = self.tool_confirmation_router.register(request.id.clone()).await; + let confirmation_rx = self + .tool_confirmation_router + .register(session.id.clone(), request.id.clone()) + .await; let action_required_msg = Message::assistant() .with_action_required( diff --git a/crates/goose/src/gateway/handler.rs b/crates/goose/src/gateway/handler.rs index 0d88a600f..fb9f0351c 100644 --- a/crates/goose/src/gateway/handler.rs +++ b/crates/goose/src/gateway/handler.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::{HashMap, VecDeque}; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -13,8 +13,7 @@ use crate::config::paths::Paths; use crate::config::Config; 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::permission::Permission; use crate::session::SessionType; use crate::session::{EnabledExtensionsState, ExtensionState, Session}; @@ -42,9 +41,13 @@ fn resolve_gateway_max_turns(gateway_override: Option, global_max_turns: Op struct PendingConfirmation { agent: Arc, - request_id: String, + session_id: String, + request_ids: VecDeque, + cancel_token: CancellationToken, } +type PerUserLocks = Arc>>>>; + #[derive(Clone)] pub struct GatewayHandler { agent_manager: Arc, @@ -54,7 +57,9 @@ pub struct GatewayHandler { /// 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>>>>, + turn_locks: PerUserLocks, + /// Serializes confirmation replies without waiting for the active turn to finish. + confirmation_reply_locks: PerUserLocks, } impl GatewayHandler { @@ -71,27 +76,42 @@ impl GatewayHandler { config, pending_confirmations: Arc::new(Mutex::new(HashMap::new())), turn_locks: Arc::new(Mutex::new(HashMap::new())), + confirmation_reply_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; + for (_, pending) in pending { + let PendingConfirmation { + agent, + session_id, + request_ids, + cancel_token, + } = pending; + for request_id in request_ids { + if let Err(error) = agent + .submit_tool_confirmation(&session_id, &request_id, Permission::DenyOnce) + .await + { + tracing::error!(%error, %request_id, "failed to deny pending gateway confirmation"); + cancel_token.cancel(); + } + } } } - async fn prune_turn_lock(&self, user: &PlatformUser) { - let mut locks = self.turn_locks.lock().await; + async fn per_user_lock(locks: &PerUserLocks, user: &PlatformUser) -> Arc> { + let mut locks = locks.lock().await; + Arc::clone( + locks + .entry(user.clone()) + .or_insert_with(|| Arc::new(Mutex::new(()))), + ) + } + + async fn prune_per_user_lock(locks: &PerUserLocks, user: &PlatformUser) { + let mut locks = locks.lock().await; let in_use = locks .get(user) .is_some_and(|lock| Arc::strong_count(lock) > 1); @@ -166,19 +186,12 @@ impl GatewayHandler { { 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_lock = Self::per_user_lock(&self.turn_locks, &message.user).await; 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; + Self::prune_per_user_lock(&self.turn_locks, &message.user).await; result?; } } @@ -210,25 +223,55 @@ impl GatewayHandler { return Ok(()); }; - let Some(pending) = self - .pending_confirmations - .lock() - .await - .remove(&message.user) - else { + let reply_lock = Self::per_user_lock(&self.confirmation_reply_locks, &message.user).await; + let reply_guard = reply_lock.lock().await; + let result = self.submit_pending_confirmation(message, permission).await; + drop(reply_guard); + drop(reply_lock); + Self::prune_per_user_lock(&self.confirmation_reply_locks, &message.user).await; + result + } + + async fn submit_pending_confirmation( + &self, + message: &IncomingMessage, + permission: Permission, + ) -> anyhow::Result<()> { + let pending_confirmations = self.pending_confirmations.lock().await; + let Some(pending) = pending_confirmations.get(&message.user) else { return Ok(()); }; + let Some(request_id) = pending.request_ids.front().cloned() else { + return Ok(()); + }; + let agent = pending.agent.clone(); + let session_id = pending.session_id.clone(); + let cancel_token = pending.cancel_token.clone(); + drop(pending_confirmations); + if let Err(error) = agent + .submit_tool_confirmation(&session_id, &request_id, permission) + .await + { + cancel_token.cancel(); + self.pending_confirmations + .lock() + .await + .remove(&message.user); + return Err(error); + } - pending - .agent - .handle_confirmation( - pending.request_id, - PermissionConfirmation { - principal_type: PrincipalType::Tool, - permission, - }, - ) - .await; + let mut pending_confirmations = self.pending_confirmations.lock().await; + let remove_pending = pending_confirmations + .get_mut(&message.user) + .is_some_and(|pending| { + if pending.request_ids.front() == Some(&request_id) { + pending.request_ids.pop_front(); + } + pending.request_ids.is_empty() + }); + if remove_pending { + pending_confirmations.remove(&message.user); + } Ok(()) } @@ -490,7 +533,7 @@ impl GatewayHandler { }; let mut stream = match agent - .reply(user_message, session_config, Some(cancel)) + .reply(user_message, session_config, Some(cancel.clone())) .await { Ok(s) => s, @@ -506,7 +549,6 @@ impl GatewayHandler { return Ok(()); } }; - // Telegram stops showing "typing…" after ~5 seconds. Re-send the // indicator every 4 s so the user always sees activity while the // agent is working (tool calls, LLM round-trips, etc.). @@ -635,13 +677,19 @@ impl GatewayHandler { • deny always — always deny", ); - self.pending_confirmations.lock().await.insert( - message.user.clone(), - PendingConfirmation { + self.pending_confirmations + .lock() + .await + .entry(message.user.clone()) + .and_modify(|pending| { + pending.request_ids.push_back(id.clone()); + }) + .or_insert_with(|| PendingConfirmation { agent: agent.clone(), - request_id: id.clone(), - }, - ); + session_id: session_id.to_string(), + request_ids: VecDeque::from([id.clone()]), + cancel_token: cancel.clone(), + }); let send_result = self .gateway @@ -659,19 +707,41 @@ impl GatewayHandler { 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, - }, + let mut pending = + self.pending_confirmations.lock().await; + let remove_pending = pending + .get_mut(&message.user) + .is_some_and(|pending| { + pending + .request_ids + .retain(|request_id| request_id != id); + pending.request_ids.is_empty() + }); + if remove_pending { + pending.remove(&message.user); + } + drop(pending); + if let Err(error) = agent + .submit_tool_confirmation( + session_id, + id, + Permission::DenyOnce, ) - .await; + .await + { + tracing::error!( + session_id, + request_id = %id, + %error, + "failed to deny undeliverable tool approval" + ); + cancel.cancel(); + self.pending_confirmations + .lock() + .await + .remove(&message.user); + } + sent_any = true; } else { sent_any = true; } diff --git a/crates/goose/tests/agent.rs b/crates/goose/tests/agent.rs index 13f541539..722345bd5 100644 --- a/crates/goose/tests/agent.rs +++ b/crates/goose/tests/agent.rs @@ -604,8 +604,9 @@ mod tests { .update_provider(provider, ModelConfig::new("mock-model"), &session.id) .await?; + let session_id = session.id; let session_config = SessionConfig { - id: session.id, + id: session_id.clone(), schedule_id: None, max_turns: Some(1), retry_config: None, @@ -622,13 +623,13 @@ mod tests { response.content.first() { if let goose::conversation::message::ActionRequiredData::ToolConfirmation { id, .. } = &action.data { - agent.handle_confirmation( - id.clone(), - goose::permission::PermissionConfirmation { - principal_type: goose::permission::permission_confirmation::PrincipalType::Tool, - permission: goose::permission::Permission::AllowOnce, - } - ).await; + agent + .submit_tool_confirmation( + &session_id, + id, + goose::permission::Permission::AllowOnce, + ) + .await?; } } responses.push(response); diff --git a/crates/goose/tests/providers.rs b/crates/goose/tests/providers.rs index 4ed94c0b7..3f7794240 100644 --- a/crates/goose/tests/providers.rs +++ b/crates/goose/tests/providers.rs @@ -5,8 +5,7 @@ use goose::acp::ACP_CURRENT_MODEL; use goose::agents::{Agent, AgentConfig, AgentEvent, GoosePlatform, PromptManager, SessionConfig}; use goose::config::{ExtensionConfig, GooseMode, PermissionManager}; use goose::conversation::message::{ActionRequiredData, Message, MessageContent}; -use goose::permission::permission_confirmation::PrincipalType; -use goose::permission::{Permission, PermissionConfirmation}; +use goose::permission::Permission; use goose::providers::anthropic::ANTHROPIC_DEFAULT_MODEL; use goose::providers::azure::AZURE_DEFAULT_MODEL; use goose::providers::base::Provider; @@ -600,14 +599,8 @@ impl ProviderFixture { if let ActionRequiredData::ToolConfirmation { ref id, .. } = ar.data { saw_action_required = true; self.agent - .handle_confirmation( - id.clone(), - PermissionConfirmation { - principal_type: PrincipalType::Tool, - permission: permission.clone(), - }, - ) - .await; + .submit_tool_confirmation(&self.session_id, id, permission.clone()) + .await?; } } }