From c24705504bec727575b0b112e17b6157ac37ba6e Mon Sep 17 00:00:00 2001 From: Milhous Date: Wed, 8 Jul 2026 15:22:31 +0800 Subject: [PATCH 1/8] feat(runtime): add current turn context to AppState --- crates/puffer-core/lib.rs | 4 ++-- crates/puffer-core/state.rs | 47 +++++++++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 2 deletions(-) diff --git a/crates/puffer-core/lib.rs b/crates/puffer-core/lib.rs index bbc549926..9f6752559 100644 --- a/crates/puffer-core/lib.rs +++ b/crates/puffer-core/lib.rs @@ -117,8 +117,8 @@ pub use runtime::{ }; pub use runtime::{install_observability, observability_handle}; pub use state::{ - AppState, MessageRole, MonitorSourceStampContext, MonitorTaskCreateGateContext, - RenderedAttachment, RenderedMessage, TaskRecord, TaskStatus, + AppState, CurrentTurnContext, MessageRole, MonitorSourceStampContext, + MonitorTaskCreateGateContext, RenderedAttachment, RenderedMessage, TaskRecord, TaskStatus, }; use anyhow::Result; diff --git a/crates/puffer-core/state.rs b/crates/puffer-core/state.rs index ae072d7ff..414f3ef33 100644 --- a/crates/puffer-core/state.rs +++ b/crates/puffer-core/state.rs @@ -381,6 +381,8 @@ pub struct AppState { pub(crate) masked_secrets: Arc>>, /// Trusted exact media discovery entries available to workflow tools. pub(crate) exact_media_discovery_cache: Option, + /// Identifies the live daemon turn this state is executing under, when any. + pub(crate) current_turn_context: Option, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -390,6 +392,13 @@ pub(crate) struct MonitorReplyScope { pub turn_id: String, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CurrentTurnContext { + pub turn_id: String, + pub session_id: String, + pub task_id: Option, +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct MonitorTaskCreateGateContext { pub envelope_id: String, @@ -517,9 +526,22 @@ impl AppState { secret_values: Arc::new(Mutex::new(HashMap::new())), masked_secrets: Arc::new(Mutex::new(HashMap::new())), exact_media_discovery_cache: None, + current_turn_context: None, } } + pub fn set_current_turn_context(&mut self, context: CurrentTurnContext) { + self.current_turn_context = Some(context); + } + + pub fn clear_current_turn_context(&mut self) { + self.current_turn_context = None; + } + + pub fn current_turn_context(&self) -> Option<&CurrentTurnContext> { + self.current_turn_context.as_ref() + } + pub fn set_monitor_reply_scope_for_turn( &mut self, task_id: String, @@ -1422,6 +1444,31 @@ mod tests { } } + #[test] + fn current_turn_context_round_trips() { + let mut state = AppState::new( + PufferConfig::default(), + PathBuf::from("."), + sample_metadata(), + ); + + assert!(state.current_turn_context().is_none()); + + state.set_current_turn_context(CurrentTurnContext { + turn_id: "turn-1".to_string(), + session_id: "session-1".to_string(), + task_id: Some("task-1".to_string()), + }); + + let ctx = state.current_turn_context().expect("turn context"); + assert_eq!(ctx.turn_id, "turn-1"); + assert_eq!(ctx.session_id, "session-1"); + assert_eq!(ctx.task_id.as_deref(), Some("task-1")); + + state.clear_current_turn_context(); + assert!(state.current_turn_context().is_none()); + } + fn stored_image_attachment() -> puffer_session_store::StoredAttachment { puffer_session_store::StoredAttachment { id: "11111111-1111-1111-1111-111111111111".to_string(), From 9bdce0499243f8ae122c66370522c59cf25a6ed0 Mon Sep 17 00:00:00 2001 From: Milhous Date: Wed, 8 Jul 2026 15:46:05 +0800 Subject: [PATCH 2/8] feat(daemon): bound turn interaction waits behind TurnScope --- crates/puffer-cli/src/daemon.rs | 248 +++++++++-------- .../src/daemon_gcal_browser_setup.rs | 26 +- .../src/daemon_gmail_browser_setup.rs | 26 +- .../src/daemon_lark_browser_setup.rs | 25 +- .../src/daemon_slack_browser_setup.rs | 25 +- crates/puffer-cli/src/daemon_turn_scope.rs | 263 ++++++++++++++++++ .../src/daemon_wechat_browser_setup.rs | 21 +- crates/puffer-cli/src/main.rs | 1 + 8 files changed, 465 insertions(+), 170 deletions(-) create mode 100644 crates/puffer-cli/src/daemon_turn_scope.rs diff --git a/crates/puffer-cli/src/daemon.rs b/crates/puffer-cli/src/daemon.rs index c9ba609c6..ceb2d993b 100644 --- a/crates/puffer-cli/src/daemon.rs +++ b/crates/puffer-cli/src/daemon.rs @@ -109,6 +109,7 @@ use crate::daemon_turn_recovery::{ DEFAULT_STALE_TURN_RETRY_AFTER_MS, }; use crate::daemon_turn_routing::persist_explicit_turn_routing; +use crate::daemon_turn_scope::{PendingWait, ResolveInteractionError, TurnFinishReason, TurnScope}; use crate::daemon_ui_state::{ load_file_tabs_state, load_pin_state, load_session_routing_state, set_file_tabs_state, set_pin_state, set_session_routing_state, DesktopFileTab, DesktopFileTabsState, @@ -595,6 +596,17 @@ impl DaemonState { } } +/// Bounded wait for a pending UI interaction before the daemon gives up and +/// resolves it (permission → Deny, question → empty answers). Overridable via +/// `PUFFER_DAEMON_INTERACTION_TIMEOUT_MS`; defaults to 15 minutes. +fn daemon_interaction_timeout() -> std::time::Duration { + std::env::var("PUFFER_DAEMON_INTERACTION_TIMEOUT_MS") + .ok() + .and_then(|raw| raw.parse::().ok()) + .map(std::time::Duration::from_millis) + .unwrap_or_else(|| std::time::Duration::from_secs(15 * 60)) +} + #[derive(Clone)] struct TurnHandle { session_id: Option, @@ -605,9 +617,7 @@ struct TurnHandle { cancel: CancelToken, cancel_reported: Arc, user_prompt_persisted: Arc, - pending: Arc>>>, - pending_questions: - Arc>>>, + scope: Arc, progress: Arc>, } @@ -4641,20 +4651,35 @@ fn handle_resolve_permission(state: &DaemonState, params: &Value) -> Result String { + match error { + ResolveInteractionError::Finished => { + format!("turn finished before `{request_id}` was resolved") + } + ResolveInteractionError::Expired => { + format!("request `{request_id}` expired (interaction timeout)") + } + ResolveInteractionError::Unknown => { + format!("no pending request `{request_id}` on this turn") + } + ResolveInteractionError::WorkerReleased => { + "worker already released the channel".to_string() + } + } +} + fn parse_permission_action(params: &Value) -> Result { let action_str = params .get("action") @@ -4737,21 +4762,20 @@ fn handle_resolve_user_question(state: &DaemonState, params: &Value) -> Result bool { return false; }; handle.cancel.cancel(); - { - let mut pending = handle.pending.lock().unwrap(); - for (_, tx) in pending.drain() { - let _ = tx.send(PermissionPromptAction::Deny); - } - } - { - let mut pending_questions = handle.pending_questions.lock().unwrap(); - for (_, tx) in pending_questions.drain() { - let _ = tx.send(UserQuestionPromptResponse { - answers: serde_json::Map::new(), - annotations: serde_json::Map::new(), - }); - } - } + let _report = handle.scope.finish(TurnFinishReason::CancelledByUser); // Cancellation cleanup is best-effort: never let a failed report block the // cancel (especially on the disconnect path, where no client is waiting). if let (Some(session_uuid), Some(session_id)) = @@ -5613,12 +5623,8 @@ async fn start_turn(state: Arc, params: Value) -> Result { let monitor_reply_scope = resolve_monitor_reply_turn_scope(&state, ¶ms, &message, &session_id, &turn_id)?; - let pending: Arc>>> = - Arc::new(Mutex::new(HashMap::new())); - let pending_questions: Arc< - Mutex>>, - > = Arc::new(Mutex::new(HashMap::new())); let cancel = CancelToken::new(); + let scope = Arc::new(TurnScope::new(cancel.clone(), daemon_interaction_timeout())); let cancel_reported = Arc::new(AtomicBool::new(false)); let user_prompt_persisted = Arc::new(AtomicBool::new(false)); let progress = Arc::new(Mutex::new(TurnProgress::default())); @@ -5642,8 +5648,7 @@ async fn start_turn(state: Arc, params: Value) -> Result { cancel: cancel.clone(), cancel_reported: cancel_reported.clone(), user_prompt_persisted: user_prompt_persisted.clone(), - pending: pending.clone(), - pending_questions: pending_questions.clone(), + scope: scope.clone(), progress: progress.clone(), }, ); @@ -6108,7 +6113,7 @@ async fn start_turn(state: Arc, params: Value) -> Result { let perm_state = setup_state.clone(); let perm_channel = channel_thread.clone(); let perm_turn = turn_id_thread.clone(); - let perm_pending = pending.clone(); + let perm_scope = scope.clone(); let perm_actor = stream_actor.clone(); let perm_cancel = cancel.clone(); let on_permission = move |req: PermissionPromptRequest| -> PermissionPromptAction { @@ -6119,8 +6124,7 @@ async fn start_turn(state: Arc, params: Value) -> Result { return PermissionPromptAction::Deny; } let request_id = next_req_id.fetch_add(1, Ordering::SeqCst).to_string(); - let (tx, rx) = std::sync::mpsc::channel(); - perm_pending.lock().unwrap().insert(request_id.clone(), tx); + let rx = perm_scope.register_permission(request_id.clone()); perm_state.publish_event(ServerEnvelope::Event { event: perm_channel.clone(), @@ -6139,13 +6143,18 @@ async fn start_turn(state: Arc, params: Value) -> Result { ), }); - rx.recv().unwrap_or(PermissionPromptAction::Deny) + match perm_scope.wait_permission(&request_id, rx) { + PendingWait::Resolved(action) => action, + PendingWait::TimedOut | PendingWait::Cancelled | PendingWait::Released => { + PermissionPromptAction::Deny + } + } }; let question_state = setup_state.clone(); let question_channel = channel_thread.clone(); let question_turn = turn_id_thread.clone(); - let question_pending = pending_questions.clone(); + let question_scope = scope.clone(); let question_next_id = setup_state.next_request_id.clone(); let question_actor = stream_actor.clone(); let question_cancel = cancel.clone(); @@ -6161,11 +6170,7 @@ async fn start_turn(state: Arc, params: Value) -> Result { }; } let request_id = question_next_id.fetch_add(1, Ordering::SeqCst).to_string(); - let (tx, rx) = std::sync::mpsc::channel(); - question_pending - .lock() - .unwrap() - .insert(request_id.clone(), tx); + let rx = question_scope.register_user_question(request_id.clone()); question_state.publish_event(ServerEnvelope::Event { event: question_channel.clone(), @@ -6181,10 +6186,18 @@ async fn start_turn(state: Arc, params: Value) -> Result { ), }); - rx.recv().unwrap_or(UserQuestionPromptResponse { - answers: serde_json::Map::new(), - annotations: serde_json::Map::new(), - }) + match question_scope.wait_user_question(&request_id, rx) { + PendingWait::Resolved(response) => response, + PendingWait::TimedOut | PendingWait::Cancelled | PendingWait::Released => { + UserQuestionPromptResponse { + answers: serde_json::Map::new(), + annotations: serde_json::Map::from_iter([( + "_puffer_interaction_status".to_string(), + serde_json::Value::String("timeout_or_cancelled".to_string()), + )]), + } + } + } }; let mut auth_store = inputs.auth_store.clone(); @@ -6342,12 +6355,8 @@ async fn start_slash_command_turn(state: Arc, params: Value) -> Res let session_uuid = Uuid::parse_str(&session_id).context("invalid sessionId")?; let turn_id = Uuid::new_v4().to_string(); let channel = format!("session:{session_id}:event"); - let pending: Arc>>> = - Arc::new(Mutex::new(HashMap::new())); - let pending_questions: Arc< - Mutex>>, - > = Arc::new(Mutex::new(HashMap::new())); let cancel = CancelToken::new(); + let scope = Arc::new(TurnScope::new(cancel.clone(), daemon_interaction_timeout())); let cancel_reported = Arc::new(AtomicBool::new(false)); let user_prompt_persisted = Arc::new(AtomicBool::new(false)); let progress = Arc::new(Mutex::new(TurnProgress::default())); @@ -6371,8 +6380,7 @@ async fn start_slash_command_turn(state: Arc, params: Value) -> Res cancel: cancel.clone(), cancel_reported: cancel_reported.clone(), user_prompt_persisted: user_prompt_persisted.clone(), - pending: pending.clone(), - pending_questions: pending_questions.clone(), + scope: scope.clone(), progress: progress.clone(), }, ); @@ -6443,7 +6451,7 @@ async fn start_slash_command_turn(state: Arc, params: Value) -> Res let question_state = setup_state.clone(); let question_channel = channel_thread.clone(); let question_turn = turn_id_thread.clone(); - let question_pending = pending_questions.clone(); + let question_scope = scope.clone(); let question_next_id = next_req_id.clone(); let question_actor = stream_actor.clone(); let question_cancel = cancel.clone(); @@ -6459,11 +6467,7 @@ async fn start_slash_command_turn(state: Arc, params: Value) -> Res }; } let request_id = question_next_id.fetch_add(1, Ordering::SeqCst).to_string(); - let (tx, rx) = std::sync::mpsc::channel(); - question_pending - .lock() - .unwrap() - .insert(request_id.clone(), tx); + let rx = question_scope.register_user_question(request_id.clone()); question_state.publish_event(ServerEnvelope::Event { event: question_channel.clone(), @@ -6479,10 +6483,18 @@ async fn start_slash_command_turn(state: Arc, params: Value) -> Res ), }); - rx.recv().unwrap_or(UserQuestionPromptResponse { - answers: serde_json::Map::new(), - annotations: serde_json::Map::new(), - }) + match question_scope.wait_user_question(&request_id, rx) { + PendingWait::Resolved(response) => response, + PendingWait::TimedOut | PendingWait::Cancelled | PendingWait::Released => { + UserQuestionPromptResponse { + answers: serde_json::Map::new(), + annotations: serde_json::Map::from_iter([( + "_puffer_interaction_status".to_string(), + serde_json::Value::String("timeout_or_cancelled".to_string()), + )]), + } + } + } }; let mut auth_store = inputs.auth_store.clone(); @@ -6552,7 +6564,7 @@ async fn start_slash_command_turn(state: Arc, params: Value) -> Res cancel_reported.clone(), user_prompt_persisted.clone(), progress.clone(), - pending.clone(), + scope.clone(), ); setup_state.turns.lock().unwrap().remove(&turn_id_thread); }); @@ -6569,12 +6581,8 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R .to_string(); let connect_args = connector_setup_connect_args(&message)?; let channel = format!("connector-setup:{turn_id}:event"); - let pending: Arc>>> = - Arc::new(Mutex::new(HashMap::new())); - let pending_questions: Arc< - Mutex>>, - > = Arc::new(Mutex::new(HashMap::new())); let cancel = CancelToken::new(); + let scope = Arc::new(TurnScope::new(cancel.clone(), daemon_interaction_timeout())); let cancel_reported = Arc::new(AtomicBool::new(false)); let user_prompt_persisted = Arc::new(AtomicBool::new(false)); let progress = Arc::new(Mutex::new(TurnProgress::default())); @@ -6595,8 +6603,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R cancel: cancel.clone(), cancel_reported: cancel_reported.clone(), user_prompt_persisted, - pending, - pending_questions: pending_questions.clone(), + scope: scope.clone(), progress, }, ); @@ -6623,7 +6630,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R turn_id_thread.clone(), connect_args.clone(), next_req_id.clone(), - pending_questions.clone(), + scope.clone(), cancel_thread.clone(), ); match outcome { @@ -6662,7 +6669,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R turn_id_thread.clone(), connect_args.clone(), next_req_id.clone(), - pending_questions.clone(), + scope.clone(), cancel_thread.clone(), ); match outcome { @@ -6701,7 +6708,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R turn_id_thread.clone(), connect_args.clone(), next_req_id.clone(), - pending_questions.clone(), + scope.clone(), cancel_thread.clone(), ); match outcome { @@ -6740,7 +6747,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R turn_id_thread.clone(), connect_args.clone(), next_req_id.clone(), - pending_questions.clone(), + scope.clone(), cancel_thread.clone(), ); match outcome { @@ -6779,7 +6786,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R turn_id_thread.clone(), connect_args.clone(), next_req_id.clone(), - pending_questions.clone(), + scope.clone(), cancel_thread.clone(), ); match outcome { @@ -6833,7 +6840,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R let question_state = setup_state.clone(); let question_channel = channel_thread.clone(); let question_turn = turn_id_thread.clone(); - let question_pending = pending_questions.clone(); + let question_scope = scope.clone(); let question_next_id = next_req_id.clone(); let question_actor = stream_actor.clone(); let question_cancel = cancel.clone(); @@ -6849,11 +6856,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R }; } let request_id = question_next_id.fetch_add(1, Ordering::SeqCst).to_string(); - let (tx, rx) = std::sync::mpsc::channel(); - question_pending - .lock() - .unwrap() - .insert(request_id.clone(), tx); + let rx = question_scope.register_user_question(request_id.clone()); question_state.publish_event(ServerEnvelope::Event { event: question_channel.clone(), @@ -6869,10 +6872,18 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R ), }); - rx.recv().unwrap_or(UserQuestionPromptResponse { - answers: serde_json::Map::new(), - annotations: serde_json::Map::new(), - }) + match question_scope.wait_user_question(&request_id, rx) { + PendingWait::Resolved(response) => response, + PendingWait::TimedOut | PendingWait::Cancelled | PendingWait::Released => { + UserQuestionPromptResponse { + answers: serde_json::Map::new(), + annotations: serde_json::Map::from_iter([( + "_puffer_interaction_status".to_string(), + serde_json::Value::String("timeout_or_cancelled".to_string()), + )]), + } + } + } }; let outcome = with_user_question_prompt_handler(on_user_question, || { @@ -7483,17 +7494,18 @@ mod tests { append_ordered_turn_progress, apply_daemon_yolo_mode, apply_proxy_env_at_startup, apply_turn_model_override, apply_turn_request_options, browser_launch_settings_or_default, browser_permission_payload_json, browser_status_for_turn, cancel_all_active_turns, - connector_setup_connect_args, connector_setup_id, daemon_now_ms, desktop_latency_ms, - file_media_mime_type, generated_video_handler, handle_create_file_media_access, - handle_create_generated_video_access, handle_create_openai_realtime_client_secret, - handle_create_session, handle_generate_media, handle_import_external_credential, - handle_list_lambda_skill_libraries, handle_list_media_capabilities, - handle_list_permissions, handle_list_provider_models, handle_load_session_detail, - handle_local_model_status, handle_login_with_api_key, handle_logout_provider, - handle_read_generated_media_preview, handle_remove_lambda_skill_library, - handle_save_lambda_skill_library, handle_save_permissions, handle_save_proxy_settings, - handle_set_lambda_skill_approval, handle_set_lambda_skill_enabled, handle_update_config, - model_descriptor_dto, parse_single_byte_range, permission_review_payload_json, + connector_setup_connect_args, connector_setup_id, daemon_interaction_timeout, + daemon_now_ms, desktop_latency_ms, file_media_mime_type, generated_video_handler, + handle_create_file_media_access, handle_create_generated_video_access, + handle_create_openai_realtime_client_secret, handle_create_session, handle_generate_media, + handle_import_external_credential, handle_list_lambda_skill_libraries, + handle_list_media_capabilities, handle_list_permissions, handle_list_provider_models, + handle_load_session_detail, handle_local_model_status, handle_login_with_api_key, + handle_logout_provider, handle_read_generated_media_preview, + handle_remove_lambda_skill_library, handle_save_lambda_skill_library, + handle_save_permissions, handle_save_proxy_settings, handle_set_lambda_skill_approval, + handle_set_lambda_skill_enabled, handle_update_config, model_descriptor_dto, + parse_single_byte_range, permission_review_payload_json, realtime_session_config_from_params, report_cancelled_turn, requires_explicit_subscription, resolve_create_session_model_id, resolve_monitor_reply_turn_scope, run_off_runtime, session_used_browser_tool, start_connector_setup_turn, turn_browser_tab_context, @@ -7502,6 +7514,7 @@ mod tests { TurnHandle, TurnProgress, TurnProgressItem, TurnRequestOptions, MONITOR_REPLY_ACTION_PROMPT_SCOPE, }; + use crate::daemon_turn_scope::TurnScope; use axum::{ extract::{Path as AxumPath, State}, http::{header, HeaderMap, HeaderValue, StatusCode}, @@ -7518,7 +7531,7 @@ mod tests { }; use puffer_session_store::{SessionMetadata, SessionStore, TranscriptEvent, TurnBoundaryState}; use serde_json::{json, Value}; - use std::collections::{BTreeMap, HashMap}; + use std::collections::BTreeMap; use std::io::{Read, Write}; use std::net::TcpListener; use std::sync::atomic::{AtomicBool, Ordering}; @@ -8068,11 +8081,10 @@ models: [] channel: "agent".to_string(), message: String::new(), attachments: Vec::new(), - cancel, + cancel: cancel.clone(), cancel_reported: Arc::new(AtomicBool::new(false)), user_prompt_persisted: Arc::new(AtomicBool::new(false)), - pending: Arc::new(Mutex::new(HashMap::new())), - pending_questions: Arc::new(Mutex::new(HashMap::new())), + scope: Arc::new(TurnScope::new(cancel, daemon_interaction_timeout())), progress: Arc::new(Mutex::new(TurnProgress::default())), } } diff --git a/crates/puffer-cli/src/daemon_gcal_browser_setup.rs b/crates/puffer-cli/src/daemon_gcal_browser_setup.rs index f5e1deaa9..759aecd19 100644 --- a/crates/puffer-cli/src/daemon_gcal_browser_setup.rs +++ b/crates/puffer-cli/src/daemon_gcal_browser_setup.rs @@ -5,9 +5,9 @@ use anyhow::{bail, Context, Result}; use puffer_core::{CancelToken, UserQuestionPromptResponse}; use puffer_subscriptions::{ConnectionRecord, ConnectionState}; use serde_json::{json, Map, Value}; -use std::collections::{BTreeSet, HashMap}; +use std::collections::BTreeSet; use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{mpsc, Arc, Mutex}; +use std::sync::Arc; use std::time::{Duration, Instant}; const ACCOUNT_SELECT_QUESTION: &str = @@ -118,7 +118,9 @@ const DISCOVER_ACCOUNTS_SCRIPT: &str = r#" })() "#; -type PendingQuestions = Arc>>>; +use crate::daemon_turn_scope::{PendingWait, TurnScope}; + +type PendingQuestions = Arc; #[derive(Debug, Clone, PartialEq, Eq)] struct SetupTarget { @@ -338,11 +340,9 @@ impl SetupFlow { .next_request_id .fetch_add(1, Ordering::SeqCst) .to_string(); - let (tx, rx) = mpsc::channel(); - self.pending_questions - .lock() - .unwrap() - .insert(request_id.clone(), tx); + let rx = self + .pending_questions + .register_user_question(request_id.clone()); let mut payload = Map::new(); payload.insert("type".to_string(), json!("user-question-request")); @@ -359,8 +359,14 @@ impl SetupFlow { payload: Value::Object(payload), }); - rx.recv() - .map_err(|_| anyhow::anyhow!("connector setup question channel closed")) + match self.pending_questions.wait_user_question(&request_id, rx) { + PendingWait::Resolved(response) => Ok(response), + PendingWait::TimedOut => { + bail!("connector setup question timed out waiting for a response") + } + PendingWait::Cancelled => bail!("connector setup was cancelled"), + PendingWait::Released => bail!("connector setup question channel closed"), + } } } diff --git a/crates/puffer-cli/src/daemon_gmail_browser_setup.rs b/crates/puffer-cli/src/daemon_gmail_browser_setup.rs index 1a096f2ff..fca609287 100644 --- a/crates/puffer-cli/src/daemon_gmail_browser_setup.rs +++ b/crates/puffer-cli/src/daemon_gmail_browser_setup.rs @@ -6,9 +6,9 @@ use anyhow::{bail, Context, Result}; use puffer_core::{CancelToken, UserQuestionPromptResponse}; use puffer_subscriptions::{ConnectionRecord, ConnectionState}; use serde_json::{json, Map, Value}; -use std::collections::{BTreeSet, HashMap}; +use std::collections::BTreeSet; use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{mpsc, Arc, Mutex}; +use std::sync::Arc; use std::time::{Duration, Instant}; const ACCOUNT_SELECT_QUESTION: &str = @@ -119,7 +119,9 @@ const DISCOVER_ACCOUNTS_SCRIPT: &str = r#" })() "#; -type PendingQuestions = Arc>>>; +use crate::daemon_turn_scope::{PendingWait, TurnScope}; + +type PendingQuestions = Arc; #[derive(Debug, Clone, PartialEq, Eq)] struct SetupTarget { @@ -370,11 +372,9 @@ impl SetupFlow { .next_request_id .fetch_add(1, Ordering::SeqCst) .to_string(); - let (tx, rx) = mpsc::channel(); - self.pending_questions - .lock() - .unwrap() - .insert(request_id.clone(), tx); + let rx = self + .pending_questions + .register_user_question(request_id.clone()); let mut payload = Map::new(); payload.insert("type".to_string(), json!("user-question-request")); @@ -391,8 +391,14 @@ impl SetupFlow { payload: Value::Object(payload), }); - rx.recv() - .map_err(|_| anyhow::anyhow!("connector setup question channel closed")) + match self.pending_questions.wait_user_question(&request_id, rx) { + PendingWait::Resolved(response) => Ok(response), + PendingWait::TimedOut => { + bail!("connector setup question timed out waiting for a response") + } + PendingWait::Cancelled => bail!("connector setup was cancelled"), + PendingWait::Released => bail!("connector setup question channel closed"), + } } } diff --git a/crates/puffer-cli/src/daemon_lark_browser_setup.rs b/crates/puffer-cli/src/daemon_lark_browser_setup.rs index 38517f15a..930790e22 100644 --- a/crates/puffer-cli/src/daemon_lark_browser_setup.rs +++ b/crates/puffer-cli/src/daemon_lark_browser_setup.rs @@ -6,9 +6,8 @@ use anyhow::{bail, Context, Result}; use puffer_core::{CancelToken, UserQuestionPromptResponse}; use puffer_subscriptions::{ConnectionRecord, ConnectionState}; use serde_json::{json, Map, Value}; -use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{mpsc, Arc, Mutex}; +use std::sync::Arc; use std::time::{Duration, Instant}; const QR_SIGN_IN_QUESTION: &str = @@ -19,7 +18,9 @@ const BROWSER_HEIGHT: u32 = 820; const LOGIN_POLL_TIMEOUT: Duration = Duration::from_secs(120); const LOGIN_POLL_INTERVAL: Duration = Duration::from_secs(2); -type PendingQuestions = Arc>>>; +use crate::daemon_turn_scope::{PendingWait, TurnScope}; + +type PendingQuestions = Arc; struct SetupFlow { state: Arc, @@ -179,11 +180,9 @@ impl SetupFlow { .next_request_id .fetch_add(1, Ordering::SeqCst) .to_string(); - let (tx, rx) = mpsc::channel(); - self.pending_questions - .lock() - .unwrap() - .insert(request_id.clone(), tx); + let rx = self + .pending_questions + .register_user_question(request_id.clone()); let mut payload = Map::new(); payload.insert("type".to_string(), json!("user-question-request")); @@ -200,8 +199,14 @@ impl SetupFlow { payload: Value::Object(payload), }); - rx.recv() - .map_err(|_| anyhow::anyhow!("connector setup question channel closed")) + match self.pending_questions.wait_user_question(&request_id, rx) { + PendingWait::Resolved(response) => Ok(response), + PendingWait::TimedOut => { + bail!("connector setup question timed out waiting for a response") + } + PendingWait::Cancelled => bail!("connector setup was cancelled"), + PendingWait::Released => bail!("connector setup question channel closed"), + } } } diff --git a/crates/puffer-cli/src/daemon_slack_browser_setup.rs b/crates/puffer-cli/src/daemon_slack_browser_setup.rs index 51907d7d6..2a181f622 100644 --- a/crates/puffer-cli/src/daemon_slack_browser_setup.rs +++ b/crates/puffer-cli/src/daemon_slack_browser_setup.rs @@ -6,9 +6,8 @@ use puffer_core::{CancelToken, UserQuestionPromptResponse}; use puffer_subscriptions::{ConnectionRecord, ConnectionState}; use regex::Regex; use serde_json::{json, Map, Value}; -use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{mpsc, Arc, Mutex}; +use std::sync::Arc; use std::time::{Duration, Instant}; const SIGN_IN_QUESTION: &str = @@ -19,7 +18,9 @@ const BROWSER_HEIGHT: u32 = 820; const LOGIN_POLL_TIMEOUT: Duration = Duration::from_secs(120); const LOGIN_POLL_INTERVAL: Duration = Duration::from_secs(2); -type PendingQuestions = Arc>>>; +use crate::daemon_turn_scope::{PendingWait, TurnScope}; + +type PendingQuestions = Arc; struct SetupFlow { state: Arc, @@ -197,11 +198,9 @@ impl SetupFlow { .next_request_id .fetch_add(1, Ordering::SeqCst) .to_string(); - let (tx, rx) = mpsc::channel(); - self.pending_questions - .lock() - .unwrap() - .insert(request_id.clone(), tx); + let rx = self + .pending_questions + .register_user_question(request_id.clone()); let mut payload = Map::new(); payload.insert("type".to_string(), json!("user-question-request")); @@ -218,8 +217,14 @@ impl SetupFlow { payload: Value::Object(payload), }); - rx.recv() - .map_err(|_| anyhow::anyhow!("connector setup question channel closed")) + match self.pending_questions.wait_user_question(&request_id, rx) { + PendingWait::Resolved(response) => Ok(response), + PendingWait::TimedOut => { + bail!("connector setup question timed out waiting for a response") + } + PendingWait::Cancelled => bail!("connector setup was cancelled"), + PendingWait::Released => bail!("connector setup question channel closed"), + } } } diff --git a/crates/puffer-cli/src/daemon_turn_scope.rs b/crates/puffer-cli/src/daemon_turn_scope.rs new file mode 100644 index 000000000..d93886e71 --- /dev/null +++ b/crates/puffer-cli/src/daemon_turn_scope.rs @@ -0,0 +1,263 @@ +use puffer_core::{CancelToken, PermissionPromptAction, UserQuestionPromptResponse}; +use serde_json::{Map, Value}; +use std::collections::{HashMap, HashSet}; +use std::sync::{mpsc, Mutex}; +use std::time::Duration; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum TurnFinishReason { + Complete, + CancelledByUser, + Error, + ClientDisconnected, +} + +impl TurnFinishReason { + pub(crate) fn as_str(self) -> &'static str { + match self { + Self::Complete => "complete", + Self::CancelledByUser => "cancelled_by_user", + Self::Error => "error", + Self::ClientDisconnected => "client_disconnected", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ResolveInteractionError { + Finished, + Expired, + Unknown, + WorkerReleased, +} + +#[derive(Debug)] +pub(crate) enum PendingWait { + Resolved(T), + TimedOut, + Cancelled, + Released, +} + +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub(crate) struct TurnFinishReport { + pub pending_permissions_resolved: usize, + pub pending_questions_resolved: usize, +} + +pub(crate) struct TurnScope { + cancel: CancelToken, + interaction_timeout: Duration, + pending_permissions: Mutex>>, + pending_questions: Mutex>>, + expired_requests: Mutex>, + finished: Mutex, +} + +impl TurnScope { + pub(crate) fn new(cancel: CancelToken, interaction_timeout: Duration) -> Self { + Self { + cancel, + interaction_timeout, + pending_permissions: Mutex::new(HashMap::new()), + pending_questions: Mutex::new(HashMap::new()), + expired_requests: Mutex::new(HashSet::new()), + finished: Mutex::new(false), + } + } + + pub(crate) fn register_permission( + &self, + request_id: String, + ) -> mpsc::Receiver { + let (tx, rx) = mpsc::channel(); + self.pending_permissions + .lock() + .unwrap() + .insert(request_id, tx); + rx + } + + pub(crate) fn register_user_question( + &self, + request_id: String, + ) -> mpsc::Receiver { + let (tx, rx) = mpsc::channel(); + self.pending_questions + .lock() + .unwrap() + .insert(request_id, tx); + rx + } + + /// Drops a pending user-question responder without resolving it. Used by + /// long-poll setup helpers (WeChat) that re-issue or abandon a prompt from + /// their own `recv_timeout` loop instead of a single bounded wait. + pub(crate) fn deregister_user_question(&self, request_id: &str) { + self.pending_questions.lock().unwrap().remove(request_id); + } + + fn check_resolvable(&self, request_id: &str) -> Result<(), ResolveInteractionError> { + if *self.finished.lock().unwrap() { + return Err(ResolveInteractionError::Finished); + } + if self.expired_requests.lock().unwrap().contains(request_id) { + return Err(ResolveInteractionError::Expired); + } + Ok(()) + } + + pub(crate) fn resolve_permission( + &self, + request_id: &str, + action: PermissionPromptAction, + ) -> Result<(), ResolveInteractionError> { + self.check_resolvable(request_id)?; + let sender = self + .pending_permissions + .lock() + .unwrap() + .remove(request_id) + .ok_or(ResolveInteractionError::Unknown)?; + sender + .send(action) + .map_err(|_| ResolveInteractionError::WorkerReleased) + } + + pub(crate) fn resolve_user_question( + &self, + request_id: &str, + response: UserQuestionPromptResponse, + ) -> Result<(), ResolveInteractionError> { + self.check_resolvable(request_id)?; + let sender = self + .pending_questions + .lock() + .unwrap() + .remove(request_id) + .ok_or(ResolveInteractionError::Unknown)?; + sender + .send(response) + .map_err(|_| ResolveInteractionError::WorkerReleased) + } + + fn wait_generic( + &self, + request_id: &str, + rx: mpsc::Receiver, + pending: &Mutex>>, + ) -> PendingWait { + match rx.recv_timeout(self.interaction_timeout) { + Ok(value) => PendingWait::Resolved(value), + Err(mpsc::RecvTimeoutError::Timeout) => { + pending.lock().unwrap().remove(request_id); + self.expired_requests + .lock() + .unwrap() + .insert(request_id.to_string()); + if self.cancel.is_cancelled() { + PendingWait::Cancelled + } else { + PendingWait::TimedOut + } + } + Err(mpsc::RecvTimeoutError::Disconnected) => PendingWait::Released, + } + } + + pub(crate) fn wait_permission( + &self, + request_id: &str, + rx: mpsc::Receiver, + ) -> PendingWait { + self.wait_generic(request_id, rx, &self.pending_permissions) + } + + pub(crate) fn wait_user_question( + &self, + request_id: &str, + rx: mpsc::Receiver, + ) -> PendingWait { + self.wait_generic(request_id, rx, &self.pending_questions) + } + + pub(crate) fn finish(&self, reason: TurnFinishReason) -> TurnFinishReport { + *self.finished.lock().unwrap() = true; + if !matches!(reason, TurnFinishReason::Complete) { + self.cancel.cancel(); + } + let mut report = TurnFinishReport::default(); + for (_, tx) in self.pending_permissions.lock().unwrap().drain() { + let _ = tx.send(PermissionPromptAction::Deny); + report.pending_permissions_resolved += 1; + } + for (_, tx) in self.pending_questions.lock().unwrap().drain() { + let mut annotations = Map::new(); + annotations.insert( + "_puffer_finish_reason".to_string(), + Value::String(reason.as_str().to_string()), + ); + let _ = tx.send(UserQuestionPromptResponse { + answers: Map::new(), + annotations, + }); + report.pending_questions_resolved += 1; + } + report + } +} + +#[cfg(test)] +mod tests { + use super::*; + use puffer_core::{CancelToken, PermissionPromptAction}; + use std::time::Duration; + + #[test] + fn permission_wait_times_out_and_late_resolve_is_expired() { + let scope = TurnScope::new(CancelToken::new(), Duration::from_millis(5)); + let rx = scope.register_permission("req-1".to_string()); + + let waited = scope.wait_permission("req-1", rx); + assert!(matches!(waited, PendingWait::TimedOut)); + + let err = scope + .resolve_permission("req-1", PermissionPromptAction::AllowOnce) + .unwrap_err(); + assert_eq!(err, ResolveInteractionError::Expired); + } + + #[test] + fn finish_denies_pending_permissions_and_blocks_late_resolves() { + let scope = TurnScope::new(CancelToken::new(), Duration::from_secs(60)); + let rx = scope.register_permission("req-1".to_string()); + + let report = scope.finish(TurnFinishReason::CancelledByUser); + assert_eq!(report.pending_permissions_resolved, 1); + assert_eq!(rx.recv().expect("denied"), PermissionPromptAction::Deny); + + let err = scope + .resolve_permission("req-1", PermissionPromptAction::AllowOnce) + .unwrap_err(); + assert_eq!(err, ResolveInteractionError::Finished); + } + + #[test] + fn resolve_then_wait_delivers_and_duplicate_is_unknown() { + let scope = TurnScope::new(CancelToken::new(), Duration::from_secs(60)); + let rx = scope.register_permission("req-1".to_string()); + scope + .resolve_permission("req-1", PermissionPromptAction::AllowSession) + .unwrap(); + assert!(matches!( + scope.wait_permission("req-1", rx), + PendingWait::Resolved(PermissionPromptAction::AllowSession) + )); + assert_eq!( + scope + .resolve_permission("req-1", PermissionPromptAction::Deny) + .unwrap_err(), + ResolveInteractionError::Unknown + ); + } +} diff --git a/crates/puffer-cli/src/daemon_wechat_browser_setup.rs b/crates/puffer-cli/src/daemon_wechat_browser_setup.rs index 53c33cb25..14be28e90 100644 --- a/crates/puffer-cli/src/daemon_wechat_browser_setup.rs +++ b/crates/puffer-cli/src/daemon_wechat_browser_setup.rs @@ -26,9 +26,8 @@ use anyhow::{bail, Context, Result}; use puffer_core::{CancelToken, UserQuestionPromptResponse}; use puffer_subscriptions::{ConnectionRecord, ConnectionState}; use serde_json::{json, Map, Value}; -use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{mpsc, Arc, Mutex}; +use std::sync::{mpsc, Arc}; use std::time::{Duration, Instant}; /// Connector slug this setup handles (matches the catalog template). @@ -49,7 +48,7 @@ const WINDOW_WAIT_TIMEOUT: Duration = Duration::from_secs(90); /// Total time to wait for the user to scan + log in before giving up. const LOGIN_WAIT_TOTAL: Duration = Duration::from_secs(300); -type PendingQuestions = Arc>>>; +type PendingQuestions = Arc; /// Returns true when connector setup args target the WeChat connector. pub(crate) fn connect_args_are_wechat(connect_args: &str) -> bool { @@ -233,12 +232,12 @@ impl SetupFlow { loop { self.cancel.check()?; if self.rt.block_on(instance.is_logged_in()).unwrap_or(false) { - self.pending_questions.lock().unwrap().remove(&request_id); + self.pending_questions.deregister_user_question(&request_id); self.status("WeChat login detected — finishing up…"); return Ok(()); } if Instant::now() >= deadline { - self.pending_questions.lock().unwrap().remove(&request_id); + self.pending_questions.deregister_user_question(&request_id); bail!("WeChat login was not detected in time; click WeChat ▸ Start to retry."); } // Block up to one poll interval for a manual Continue; on timeout we @@ -246,14 +245,14 @@ impl SetupFlow { match rx.recv_timeout(LOGIN_POLL_INTERVAL) { Ok(_) => { // Manual Continue but not logged in yet — re-show the prompt. - self.pending_questions.lock().unwrap().remove(&request_id); + self.pending_questions.deregister_user_question(&request_id); let (rid, new_rx) = self.publish_scan_question(cfg, authority)?; request_id = rid; rx = new_rx; } Err(mpsc::RecvTimeoutError::Timeout) => {} Err(mpsc::RecvTimeoutError::Disconnected) => { - self.pending_questions.lock().unwrap().remove(&request_id); + self.pending_questions.deregister_user_question(&request_id); bail!("connector setup question channel closed"); } } @@ -311,11 +310,9 @@ impl SetupFlow { .next_request_id .fetch_add(1, Ordering::SeqCst) .to_string(); - let (tx, rx) = mpsc::channel(); - self.pending_questions - .lock() - .unwrap() - .insert(request_id.clone(), tx); + let rx = self + .pending_questions + .register_user_question(request_id.clone()); let mut payload = Map::new(); payload.insert("type".to_string(), json!("user-question-request")); diff --git a/crates/puffer-cli/src/main.rs b/crates/puffer-cli/src/main.rs index 51f2aa57d..e7a2b0d0b 100644 --- a/crates/puffer-cli/src/main.rs +++ b/crates/puffer-cli/src/main.rs @@ -38,6 +38,7 @@ mod daemon_telegram_ranking; mod daemon_title; mod daemon_turn_recovery; mod daemon_turn_routing; +mod daemon_turn_scope; mod daemon_ui_state; #[cfg(unix)] mod daemon_wechat_browser_setup; From 2d23633ce23ce363cd02f5e6c5e9834ec1169cbc Mon Sep 17 00:00:00 2001 From: Milhous Date: Wed, 8 Jul 2026 15:52:03 +0800 Subject: [PATCH 3/8] feat(runtime): owner-aware background task registry with real stop --- .../puffer-core/runtime/background_tasks.rs | 219 ++++++++++++++++-- .../runtime/background_tasks_tests.rs | 97 +++++++- .../runtime/claude_tools/workflow/store.rs | 4 +- 3 files changed, 298 insertions(+), 22 deletions(-) diff --git a/crates/puffer-core/runtime/background_tasks.rs b/crates/puffer-core/runtime/background_tasks.rs index 87c91d506..8f62dec08 100644 --- a/crates/puffer-core/runtime/background_tasks.rs +++ b/crates/puffer-core/runtime/background_tasks.rs @@ -8,6 +8,7 @@ //! - **Auto-backgrounding** (CC): Long-running tasks automatically move to background //! after a configurable timeout budget. +use crate::runtime::CancelToken; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::{Arc, Mutex, OnceLock}; @@ -211,21 +212,79 @@ pub enum BackgroundTaskStatus { Pending, /// Task is currently running. Running, + /// A stop was requested; the worker has not yet acknowledged. + Stopping, /// Task completed successfully. Completed, /// Task failed with an error. Failed, /// Task was cancelled/stopped by the user or system. Stopped, + /// Task was scoped to a turn that finished before the worker acked its stop. + Abandoned, } impl BackgroundTaskStatus { /// Returns true if the task has reached a terminal state. pub fn is_terminal(self) -> bool { - matches!(self, Self::Completed | Self::Failed | Self::Stopped) + matches!( + self, + Self::Completed | Self::Failed | Self::Stopped | Self::Abandoned + ) } } +/// Whether a background task hosts an agent loop or a shell process. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BackgroundTaskKind { + Agent, + Shell, +} + +/// Whether a background task is bound to its owner turn or detached from it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BackgroundTaskDurability { + /// Stopped when the owner turn finishes. + ScopedToTurn, + /// Survives the owner turn; only an explicit `TaskStop` ends it. + DurableAcrossTurns, +} + +impl BackgroundTaskDurability { + /// Maps the tool-facing `durable` boolean to a durability. + pub fn from_durable_flag(durable: bool) -> Self { + if durable { + Self::DurableAcrossTurns + } else { + Self::ScopedToTurn + } + } +} + +/// Full registration record for an owner-aware, stoppable background task. +pub struct BackgroundTaskRegistration { + pub task_id: String, + pub description: String, + pub kind: BackgroundTaskKind, + pub agent_id: Option, + pub output_file: Option, + pub auto_backgrounded: bool, + pub owner_turn_id: Option, + pub owner_session_id: Option, + pub durability: BackgroundTaskDurability, + pub cancel: Option, + pub process_id: Option, +} + +/// Summary of what `stop_scoped_by_turn` did for one turn. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub struct TurnTaskStopReport { + pub stop_requested: usize, + pub abandoned: usize, +} + /// Metadata and state for one background task. #[derive(Debug, Clone, Serialize)] pub struct BackgroundTaskInfo { @@ -241,12 +300,26 @@ pub struct BackgroundTaskInfo { pub output_file: Option, /// True if this task was auto-backgrounded (CC-style timeout). pub auto_backgrounded: bool, + /// Whether this task hosts an agent loop or a shell process. + pub kind: BackgroundTaskKind, + /// Turn that spawned this task, when known. + #[serde(skip_serializing_if = "Option::is_none")] + pub owner_turn_id: Option, + /// Session that spawned this task, when known. + #[serde(skip_serializing_if = "Option::is_none")] + pub owner_session_id: Option, + /// Whether this task is scoped to its owner turn or detached. + pub durability: BackgroundTaskDurability, + /// OS pid of the underlying process, for shell tasks. + #[serde(skip_serializing_if = "Option::is_none")] + pub process_id: Option, } /// Shared mutable state for one tracked background task. struct TrackedTask { info: BackgroundTaskInfo, output: Arc>, + cancel: Option, _start_instant: Instant, } @@ -258,21 +331,17 @@ pub struct BackgroundTaskManager { } impl BackgroundTaskManager { - fn new() -> Self { + pub fn new() -> Self { Self { tasks: Mutex::new(HashMap::new()), } } - /// Registers a new background task. Returns `Err` if the concurrent limit - /// has been reached. - pub fn register( + /// Registers a new owner-aware background task. Returns `Err` if the + /// concurrent limit has been reached. + pub fn register_with_options( &self, - task_id: &str, - description: &str, - agent_id: Option<&str>, - output_file: Option<&str>, - auto_backgrounded: bool, + registration: BackgroundTaskRegistration, ) -> Result>, String> { let mut tasks = self.tasks.lock().unwrap(); @@ -290,31 +359,70 @@ impl BackgroundTaskManager { let output = Arc::new(Mutex::new(HeadTailBuffer::new())); let info = BackgroundTaskInfo { - task_id: task_id.to_string(), - description: description.to_string(), + task_id: registration.task_id.clone(), + description: registration.description, status: BackgroundTaskStatus::Running, created_at: now_ms(), completed_at: None, - agent_id: agent_id.map(ToString::to_string), - output_file: output_file.map(ToString::to_string), - auto_backgrounded, + agent_id: registration.agent_id, + output_file: registration.output_file, + auto_backgrounded: registration.auto_backgrounded, + kind: registration.kind, + owner_turn_id: registration.owner_turn_id, + owner_session_id: registration.owner_session_id, + durability: registration.durability, + process_id: registration.process_id, }; tasks.insert( - task_id.to_string(), + registration.task_id, TrackedTask { info, output: Arc::clone(&output), + cancel: registration.cancel, _start_instant: Instant::now(), }, ); Ok(output) } - /// Marks a task as completed or failed. + /// Registers a new background task. Returns `Err` if the concurrent limit + /// has been reached. Legacy no-owner callers default to a durable agent so + /// nothing new can stop them. + pub fn register( + &self, + task_id: &str, + description: &str, + agent_id: Option<&str>, + output_file: Option<&str>, + auto_backgrounded: bool, + ) -> Result>, String> { + self.register_with_options(BackgroundTaskRegistration { + task_id: task_id.to_string(), + description: description.to_string(), + kind: BackgroundTaskKind::Agent, + agent_id: agent_id.map(str::to_string), + output_file: output_file.map(str::to_string), + auto_backgrounded, + owner_turn_id: None, + owner_session_id: None, + durability: BackgroundTaskDurability::DurableAcrossTurns, + cancel: None, + process_id: None, + }) + } + + /// Marks a task as completed or failed. Terminal statuses are monotonic + /// (a late duplicate never overwrites), and a worker acking a requested + /// stop lands in `Stopped` rather than `Completed`/`Failed`. pub fn complete(&self, task_id: &str, success: bool) { let mut tasks = self.tasks.lock().unwrap(); if let Some(task) = tasks.get_mut(task_id) { - task.info.status = if success { + if task.info.status.is_terminal() { + return; + } + task.info.status = if task.info.status == BackgroundTaskStatus::Stopping { + BackgroundTaskStatus::Stopped + } else if success { BackgroundTaskStatus::Completed } else { BackgroundTaskStatus::Failed @@ -323,6 +431,66 @@ impl BackgroundTaskManager { } } + /// Requests a stop of one task: cancels its token, kills its process if + /// any, and marks it `Stopping`. Returns its info snapshot, or `None` if the + /// task is unknown. A no-op on already-terminal tasks. + pub fn request_stop(&self, task_id: &str) -> Option { + let mut tasks = self.tasks.lock().unwrap(); + let task = tasks.get_mut(task_id)?; + if !task.info.status.is_terminal() { + stop_one(task, "task_stop"); + } + Some(task.info.clone()) + } + + /// Stops every live task scoped to `turn_id`, then waits up to + /// `drain_timeout` for the workers to ack. Any still-live scoped task at the + /// deadline is force-marked `Abandoned`. Durable tasks are never touched. + pub fn stop_scoped_by_turn( + &self, + turn_id: &str, + reason: &str, + drain_timeout: Duration, + ) -> TurnTaskStopReport { + let mut report = TurnTaskStopReport::default(); + let is_scoped_live = |info: &BackgroundTaskInfo| { + info.owner_turn_id.as_deref() == Some(turn_id) + && info.durability == BackgroundTaskDurability::ScopedToTurn + && !info.status.is_terminal() + }; + { + let mut tasks = self.tasks.lock().unwrap(); + for task in tasks.values_mut() { + if is_scoped_live(&task.info) { + stop_one(task, reason); + report.stop_requested += 1; + } + } + } + let start = Instant::now(); + while start.elapsed() < drain_timeout { + let any_live = self + .tasks + .lock() + .unwrap() + .values() + .any(|t| is_scoped_live(&t.info)); + if !any_live { + return report; + } + std::thread::sleep(Duration::from_millis(25)); + } + let mut tasks = self.tasks.lock().unwrap(); + for task in tasks.values_mut() { + if is_scoped_live(&task.info) { + task.info.status = BackgroundTaskStatus::Abandoned; + task.info.completed_at = Some(now_ms()); + report.abandoned += 1; + } + } + report + } + /// Marks a task as stopped (cancelled). pub fn stop(&self, task_id: &str) { let mut tasks = self.tasks.lock().unwrap(); @@ -502,6 +670,21 @@ pub enum AutoBgResult { // Utilities // --------------------------------------------------------------------------- +/// Cancels one task's token, kills its process if any, notes the stop in its +/// output, and marks it `Stopping`. +fn stop_one(task: &mut TrackedTask, reason: &str) { + if let Some(cancel) = &task.cancel { + cancel.cancel(); + } + if let Some(pid) = task.info.process_id { + let _ = crate::runtime::claude_tools::workflow::store::terminate_process(pid); + } + if let Ok(mut output) = task.output.lock() { + output.write_str(&format!("\n[stop requested: {reason}]\n")); + } + task.info.status = BackgroundTaskStatus::Stopping; +} + fn now_ms() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) diff --git a/crates/puffer-core/runtime/background_tasks_tests.rs b/crates/puffer-core/runtime/background_tasks_tests.rs index fbc5b5219..12b1efd3b 100644 --- a/crates/puffer-core/runtime/background_tasks_tests.rs +++ b/crates/puffer-core/runtime/background_tasks_tests.rs @@ -331,10 +331,11 @@ fn task_manager_stop_then_complete_is_idempotent() { assert_eq!(info.status, BackgroundTaskStatus::Stopped); assert_eq!(mgr.active_count(), 0); - // Completing after stop still updates status + // Completing after a terminal stop is a no-op: terminal statuses are + // monotonic, so the task stays Stopped. mgr.complete("s1", true); let info = mgr.get_info("s1").unwrap(); - assert_eq!(info.status, BackgroundTaskStatus::Completed); + assert_eq!(info.status, BackgroundTaskStatus::Stopped); } #[test] @@ -525,3 +526,95 @@ fn end_to_end_background_task_lifecycle() { assert!(mgr.get_info("e2e-agent-1").is_none()); assert_eq!(mgr.active_count(), 0); } + +#[test] +fn complete_on_stopping_yields_stopped_and_terminal_is_monotonic() { + let mgr = BackgroundTaskManager::new(); + let cancel = crate::runtime::CancelToken::new(); + mgr.register_with_options(BackgroundTaskRegistration { + task_id: "agent-1".into(), + description: "background agent".into(), + kind: BackgroundTaskKind::Agent, + agent_id: Some("agent-1".into()), + output_file: None, + auto_backgrounded: false, + owner_turn_id: Some("turn-1".into()), + owner_session_id: Some("session-1".into()), + durability: BackgroundTaskDurability::ScopedToTurn, + cancel: Some(cancel.clone()), + process_id: None, + }) + .unwrap(); + + mgr.request_stop("agent-1").unwrap(); + assert!(cancel.is_cancelled()); + assert_eq!( + mgr.get_info("agent-1").unwrap().status, + BackgroundTaskStatus::Stopping + ); + + mgr.complete("agent-1", true); // worker acks + assert_eq!( + mgr.get_info("agent-1").unwrap().status, + BackgroundTaskStatus::Stopped + ); + + mgr.complete("agent-1", false); // late duplicate must not overwrite + assert_eq!( + mgr.get_info("agent-1").unwrap().status, + BackgroundTaskStatus::Stopped + ); +} + +#[test] +fn stop_scoped_by_turn_spares_durable_and_abandons_unacked() { + let mgr = BackgroundTaskManager::new(); + let scoped_cancel = crate::runtime::CancelToken::new(); + let durable_cancel = crate::runtime::CancelToken::new(); + for (id, durability, cancel) in [ + ( + "scoped", + BackgroundTaskDurability::ScopedToTurn, + &scoped_cancel, + ), + ( + "durable", + BackgroundTaskDurability::DurableAcrossTurns, + &durable_cancel, + ), + ] { + mgr.register_with_options(BackgroundTaskRegistration { + task_id: id.into(), + description: id.into(), + kind: BackgroundTaskKind::Agent, + agent_id: Some(id.into()), + output_file: None, + auto_backgrounded: false, + owner_turn_id: Some("turn-1".into()), + owner_session_id: Some("session-1".into()), + durability, + cancel: Some(cancel.clone()), + process_id: None, + }) + .unwrap(); + } + + let report = mgr.stop_scoped_by_turn( + "turn-1", + "test_finish", + std::time::Duration::from_millis(30), + ); + + assert_eq!(report.stop_requested, 1); + assert_eq!(report.abandoned, 1); // scoped worker never acked within drain + assert!(scoped_cancel.is_cancelled()); + assert!(!durable_cancel.is_cancelled()); + assert_eq!( + mgr.get_info("scoped").unwrap().status, + BackgroundTaskStatus::Abandoned + ); + assert_eq!( + mgr.get_info("durable").unwrap().status, + BackgroundTaskStatus::Running + ); +} diff --git a/crates/puffer-core/runtime/claude_tools/workflow/store.rs b/crates/puffer-core/runtime/claude_tools/workflow/store.rs index 7f2aa1678..be880e1b3 100644 --- a/crates/puffer-core/runtime/claude_tools/workflow/store.rs +++ b/crates/puffer-core/runtime/claude_tools/workflow/store.rs @@ -902,7 +902,7 @@ pub(super) fn process_is_running(pid: u32) -> bool { } /// Waits for the process to exit, returning true when it stops within the timeout. -pub(super) fn wait_for_process_exit(pid: u32, timeout_ms: u64) -> bool { +pub(crate) fn wait_for_process_exit(pid: u32, timeout_ms: u64) -> bool { let deadline = now_ms().saturating_add(timeout_ms); while now_ms() < deadline { if !process_is_running(pid) { @@ -914,7 +914,7 @@ pub(super) fn wait_for_process_exit(pid: u32, timeout_ms: u64) -> bool { } /// Attempts to terminate a background shell process by pid. -pub(super) fn terminate_process(pid: u32) -> Result<()> { +pub(crate) fn terminate_process(pid: u32) -> Result<()> { #[cfg(unix)] { let status = Command::new("kill") From 56139f19ce52fdb14b422cdd8ed659cc59921684 Mon Sep 17 00:00:00 2001 From: Milhous Date: Wed, 8 Jul 2026 16:37:45 +0800 Subject: [PATCH 4/8] feat(runtime): owned, stoppable background agents and shells --- crates/puffer-core/runtime/agents.rs | 49 ++++++++++--- .../puffer-core/runtime/claude_tools/bash.rs | 72 +++++++++++++++++-- .../puffer-core/runtime/claude_tools/mod.rs | 3 + .../claude_tools/workflow/task_tools.rs | 37 ++++++++++ .../runtime/tests/tool_execution/workflow.rs | 37 ++++++++++ 5 files changed, 186 insertions(+), 12 deletions(-) diff --git a/crates/puffer-core/runtime/agents.rs b/crates/puffer-core/runtime/agents.rs index 13dab811f..21d643014 100644 --- a/crates/puffer-core/runtime/agents.rs +++ b/crates/puffer-core/runtime/agents.rs @@ -39,6 +39,8 @@ struct AgentToolInput { #[serde(default)] run_in_background: bool, #[serde(default)] + durable: bool, + #[serde(default)] cwd: Option, #[serde(default)] isolation: Option, @@ -161,6 +163,8 @@ struct PreparedAgentExecution { mode: Option, max_turns: Option, worktree: Option, + durable: bool, + owner_turn_context: Option, } /// Executes the runtime-backed `Agent` tool by running a nested model turn. @@ -429,6 +433,8 @@ fn prepare_agent_execution( .filter(|value| !value.is_empty()), max_turns: input.max_turns.or(agent.value.max_turns), worktree, + durable: input.durable, + owner_turn_context: state.current_turn_context().cloned(), }) } @@ -539,15 +545,32 @@ fn launch_background_agent( ) .with_context(|| format!("failed to initialize {}", output_file.display()))?; - // Register with the centralized task manager for tracking and limit enforcement. + // Register with the centralized task manager for tracking, limit + // enforcement, and owner-scoped stopping. + use super::background_tasks::{ + BackgroundTaskDurability, BackgroundTaskKind, BackgroundTaskRegistration, + }; + let cancel = crate::runtime::CancelToken::new(); let task_output_buf = task_manager() - .register( - &prepared.agent_id, - &prepared.description, - Some(&prepared.agent_id), - Some(&output_file.display().to_string()), - false, // not auto-backgrounded - ) + .register_with_options(BackgroundTaskRegistration { + task_id: prepared.agent_id.clone(), + description: prepared.description.clone(), + kind: BackgroundTaskKind::Agent, + agent_id: Some(prepared.agent_id.clone()), + output_file: Some(output_file.display().to_string()), + auto_backgrounded: false, + owner_turn_id: prepared + .owner_turn_context + .as_ref() + .map(|c| c.turn_id.clone()), + owner_session_id: prepared + .owner_turn_context + .as_ref() + .map(|c| c.session_id.clone()), + durability: BackgroundTaskDurability::from_durable_flag(prepared.durable), + cancel: Some(cancel.clone()), + process_id: None, + }) .map_err(|err| anyhow!(err))?; let response = AgentAsyncOutput { @@ -611,6 +634,11 @@ fn launch_background_agent( let mut failed = false; for outer in 0..max_outer { + if cancel.is_cancelled() { + failed = true; + last_text = "stopped: owner turn finished or TaskStop".to_string(); + break; + } let prompt = if outer == 0 { prepared.prompt.clone() } else { @@ -626,6 +654,11 @@ fn launch_background_agent( &prompt, ) }; + if cancel.is_cancelled() { + failed = true; + last_text = "stopped: owner turn finished or TaskStop".to_string(); + break; + } match result { Ok(turn) => { total_tool_uses += turn.tool_invocations.len(); diff --git a/crates/puffer-core/runtime/claude_tools/bash.rs b/crates/puffer-core/runtime/claude_tools/bash.rs index d1593505f..5c3758690 100644 --- a/crates/puffer-core/runtime/claude_tools/bash.rs +++ b/crates/puffer-core/runtime/claude_tools/bash.rs @@ -73,6 +73,8 @@ pub struct ClaudeBashInput { #[serde(default)] pub run_in_background: bool, #[serde(default)] + pub durable: bool, + #[serde(default)] pub tty: bool, } @@ -131,7 +133,7 @@ pub fn execute_from_value( &std::sync::Arc>, >, ) -> Result { - execute_from_value_with_internal_permissions(cwd, session_id, input, process_store, None) + execute_from_value_with_internal_permissions(cwd, session_id, input, process_store, None, None) } /// Parses JSON input and executes Bash with an internal tool permission callback. @@ -143,10 +145,18 @@ pub(crate) fn execute_from_value_with_internal_permissions( &std::sync::Arc>, >, internal_permissions: Option<&mut InternalPermissionHandler<'_>>, + turn_context: Option, ) -> Result { let typed: ClaudeBashInput = serde_json::from_value(input).context("invalid Bash tool input payload")?; - execute_with_internal_permissions(cwd, session_id, typed, process_store, internal_permissions) + execute_with_internal_permissions( + cwd, + session_id, + typed, + process_store, + internal_permissions, + turn_context, + ) } /// Executes a Claude-style `Bash` tool invocation in the provided working directory. @@ -158,7 +168,7 @@ pub fn execute( &std::sync::Arc>, >, ) -> Result { - execute_with_internal_permissions(cwd, session_id, input, process_store, None) + execute_with_internal_permissions(cwd, session_id, input, process_store, None, None) } fn execute_with_internal_permissions( @@ -169,6 +179,7 @@ fn execute_with_internal_permissions( &std::sync::Arc>, >, internal_permissions: Option<&mut InternalPermissionHandler<'_>>, + turn_context: Option, ) -> Result { if input.tty { if let Some(store) = process_store { @@ -176,7 +187,13 @@ fn execute_with_internal_permissions( } } if input.run_in_background { - return execute_background(cwd, session_id, input, internal_permissions.is_some()); + return execute_background( + cwd, + session_id, + input, + internal_permissions.is_some(), + turn_context, + ); } execute_foreground(cwd, input, internal_permissions) } @@ -268,6 +285,7 @@ fn execute_background( session_id: &Uuid, input: ClaudeBashInput, internal_permission_required: bool, + turn_context: Option, ) -> Result { let output_dir = shell_output_dir(cwd)?; let pending_output_file = @@ -313,6 +331,32 @@ fn execute_background( &output_file, )?; + // Mirror the shell into the owner-aware registry so a finishing turn (or a + // TaskStop that reaches the registry branch) can stop it. `let _`: registry + // capacity refusal must not fail the shell launch — the store registration + // above stays authoritative. + { + use crate::runtime::background_tasks::{ + task_manager, BackgroundTaskDurability, BackgroundTaskKind, BackgroundTaskRegistration, + }; + let _ = task_manager().register_with_options(BackgroundTaskRegistration { + task_id: task_id.clone(), + description: input + .description + .clone() + .unwrap_or_else(|| input.command.clone()), + kind: BackgroundTaskKind::Shell, + agent_id: None, + output_file: Some(output_file.display().to_string()), + auto_backgrounded: false, + owner_turn_id: turn_context.as_ref().map(|c| c.turn_id.clone()), + owner_session_id: turn_context.as_ref().map(|c| c.session_id.clone()), + durability: BackgroundTaskDurability::from_durable_flag(input.durable), + cancel: None, + process_id: Some(pid), + }); + } + // Spawn a reaper thread that calls wait() on the child process. // Without this, the child becomes a zombie after exit because nobody // collects its exit status. The reaper also marks the task as completed @@ -321,8 +365,14 @@ fn execute_background( let reaper_cwd = cwd.to_path_buf(); let reaper_session_id = *session_id; let reaper_task_id = task_id.clone(); + let reaper_registry_task_id = task_id.clone(); thread::spawn(move || { let exit_status = child.wait(); + let exit_ok = exit_status + .as_ref() + .ok() + .map(|status| status.success()) + .unwrap_or(false); let exit_code = exit_status.ok().and_then(|s| s.code()); // Best-effort: mark the stored task as completed. let _ = super::workflow::mark_shell_task_completed( @@ -331,6 +381,10 @@ fn execute_background( &reaper_task_id, exit_code, ); + // Converge the owner-aware registry (guarded by terminal monotonicity — + // a prior stop request stays Stopping→Stopped, not overwritten). + crate::runtime::background_tasks::task_manager() + .complete(&reaper_registry_task_id, exit_ok); }); Ok(ClaudeBashExecution { @@ -639,6 +693,7 @@ mod tests { timeout: None, description: None, run_in_background: false, + durable: false, tty: false, }; assert_eq!(tool_description(&input), "Run shell command"); @@ -651,6 +706,7 @@ mod tests { timeout: None, description: Some("Show greeting".to_string()), run_in_background: false, + durable: false, tty: false, }; assert_eq!(tool_description(&input), "Show greeting"); @@ -668,6 +724,7 @@ mod tests { timeout: Some(5_000), description: None, run_in_background: false, + durable: false, tty: false, }, None, @@ -691,6 +748,7 @@ mod tests { timeout: Some(5_000), description: None, run_in_background: false, + durable: false, tty: false, }, None, @@ -713,6 +771,7 @@ mod tests { timeout: Some(5_000), description: None, run_in_background: false, + durable: false, tty: false, }, None, @@ -748,6 +807,7 @@ mod tests { timeout: Some(20), description: None, run_in_background: false, + durable: false, tty: false, }, None, @@ -771,6 +831,7 @@ mod tests { timeout: Some(5_000), description: None, run_in_background: true, + durable: false, tty: false, }, None, @@ -795,6 +856,7 @@ mod tests { timeout: Some(5_000), description: Some("Sleep briefly".to_string()), run_in_background: true, + durable: false, tty: false, }, None, @@ -859,6 +921,7 @@ mod tests { timeout: None, description: None, run_in_background: false, + durable: false, tty: false, }; let error = summary_line(&input).unwrap_err(); @@ -926,6 +989,7 @@ mod tests { timeout: Some(5_000), description: None, run_in_background: false, + durable: false, tty: false, }, None, diff --git a/crates/puffer-core/runtime/claude_tools/mod.rs b/crates/puffer-core/runtime/claude_tools/mod.rs index ddb828013..4352dfb3a 100644 --- a/crates/puffer-core/runtime/claude_tools/mod.rs +++ b/crates/puffer-core/runtime/claude_tools/mod.rs @@ -120,6 +120,7 @@ pub(crate) fn execute_tool( .unwrap_or_else(puffer_media::ExactMediaDiscoveryCache::empty); let session_id = state.session.id; let process_store = state.process_store.clone(); + let turn_context = state.current_turn_context().cloned(); let mut internal_permission_handler = |request| match request { bash_internal_permissions::InternalToolBrokerRequest::Permission(request) => { bash_internal_permissions::InternalToolBrokerResponse::Permission( @@ -149,6 +150,7 @@ pub(crate) fn execute_tool( input, Some(&process_store), Some(&mut internal_permission_handler), + turn_context, )?; let output = serde_json::to_string_pretty(&execution.output) .context("failed to serialize Bash output")?; @@ -394,6 +396,7 @@ pub(crate) fn execute_parallel_bash_with_media_broker( args, media_ctx.process_store, Some(&mut handler), + None, )?; let output = serde_json::to_string_pretty(&execution.output) .context("failed to serialize Bash output")?; diff --git a/crates/puffer-core/runtime/claude_tools/workflow/task_tools.rs b/crates/puffer-core/runtime/claude_tools/workflow/task_tools.rs index 29aa1c486..87ef04104 100644 --- a/crates/puffer-core/runtime/claude_tools/workflow/task_tools.rs +++ b/crates/puffer-core/runtime/claude_tools/workflow/task_tools.rs @@ -2437,6 +2437,20 @@ pub(super) fn execute_task_stop(state: &mut AppState, _cwd: &Path, input: Value) }))?); } + if let Some(info) = crate::runtime::background_tasks::task_manager().request_stop(&target) { + use crate::runtime::background_tasks::BackgroundTaskKind; + return Ok(serde_json::to_string_pretty(&json!({ + "message": format!("Successfully stopped task: {target}"), + "task_id": target, + "task_type": match info.kind { + BackgroundTaskKind::Agent => "agent", + BackgroundTaskKind::Shell => "shell", + }, + "status": "stopped", + "output_file": info.output_file, + }))?); + } + let mut agents = load_store::(&agents_path(store_cwd))?; if let Some(agent) = agents .agents @@ -2520,6 +2534,29 @@ pub(super) fn execute_task_output( timeout, ); } + if let Some(info) = crate::runtime::background_tasks::task_manager().get_info(&parsed.task_id) { + let output = crate::runtime::background_tasks::task_manager() + .read_output(&parsed.task_id) + .unwrap_or_default(); + let retrieval = if info.status.is_terminal() { + "success" + } else { + "not_ready" + }; + return task_output_response( + retrieval, + json!({ + "task_id": info.task_id, + "status": format!("{:?}", info.status).to_ascii_lowercase(), + "description": info.description, + "output": output, + }), + info.output_file.clone(), + block, + timeout, + ); + } + let agents = load_store::(&agents_path(store_cwd))?; if let Some(agent) = agents .agents diff --git a/crates/puffer-core/runtime/tests/tool_execution/workflow.rs b/crates/puffer-core/runtime/tests/tool_execution/workflow.rs index d55ad663b..5018068a4 100644 --- a/crates/puffer-core/runtime/tests/tool_execution/workflow.rs +++ b/crates/puffer-core/runtime/tests/tool_execution/workflow.rs @@ -1169,3 +1169,40 @@ fn task_stop_rejects_unrecorded_shell_pid() { assert!(error.to_string().contains("unknown task `shell-1`")); } + +#[test] +fn task_stop_stops_live_registry_agent() { + use crate::runtime::background_tasks::{ + task_manager, BackgroundTaskDurability, BackgroundTaskKind, BackgroundTaskRegistration, + }; + use crate::runtime::CancelToken; + + let mut state = temp_state(); + let cwd = state.cwd.clone(); + let cancel = CancelToken::new(); + task_manager() + .register_with_options(BackgroundTaskRegistration { + task_id: "agent-stop-me".into(), + description: "agent".into(), + kind: BackgroundTaskKind::Agent, + agent_id: Some("agent-stop-me".into()), + output_file: None, + auto_backgrounded: false, + owner_turn_id: Some("turn-1".into()), + owner_session_id: Some(state.session.id.to_string()), + durability: BackgroundTaskDurability::DurableAcrossTurns, + cancel: Some(cancel.clone()), + process_id: None, + }) + .unwrap(); + + let output = crate::runtime::claude_tools::workflow::task_stop::execute_task_stop( + &mut state, + &cwd, + json!({ "task_id": "agent-stop-me" }), + ) + .unwrap(); + + assert!(cancel.is_cancelled()); + assert!(output.contains("agent-stop-me")); +} From 66b9be826bec1e3e387906ac19c04c12e5eca8cd Mon Sep 17 00:00:00 2001 From: Milhous Date: Wed, 8 Jul 2026 16:54:31 +0800 Subject: [PATCH 5/8] feat(daemon): single finish_turn path with scoped child harvest Interaction-timeout smoke coverage shipped as the late-resolve-rejection variant (daemon_rejects_resolve_after_turn_finished) plus the Task 2 TurnScope unit tests, per the plan's fallback; the mock SSE stream was not extended to emit AskUserQuestion tool_use. --- crates/puffer-cli/src/daemon.rs | 167 +++++++++++++------ crates/puffer-cli/tests/daemon_turn_smoke.rs | 84 ++++++++++ 2 files changed, 202 insertions(+), 49 deletions(-) diff --git a/crates/puffer-cli/src/daemon.rs b/crates/puffer-cli/src/daemon.rs index ceb2d993b..367dbc98e 100644 --- a/crates/puffer-cli/src/daemon.rs +++ b/crates/puffer-cli/src/daemon.rs @@ -4785,18 +4785,62 @@ fn handle_cancel_turn(state: &DaemonState, params: &Value) -> Result { .or_else(|| params.get("turn_id")) .and_then(|v| v.as_str()) .context("missing turnId")?; - if cancel_turn_by_id(state, turn_id) { + if cancel_turn_by_id(state, turn_id, TurnFinishReason::CancelledByUser) { Ok(json!({"ok": true})) } else { Ok(json!({"ok": false, "error": "turn not found"})) } } -/// Cancels one running turn by id: flips its cancel token, denies any pending -/// permission/question prompts, reports the cancellation to listeners, and -/// removes it from the registry. Returns whether a turn with this id existed. -/// Shared by the `cancel_turn` RPC and the client-disconnect watchdog (#600). -fn cancel_turn_by_id(state: &DaemonState, turn_id: &str) -> bool { +/// Bounded wait for a turn's scoped background children to acknowledge a stop +/// during turn finish. Overridable via `PUFFER_TURN_CHILD_DRAIN_MS`; default 2s. +fn daemon_child_drain_timeout() -> std::time::Duration { + std::env::var("PUFFER_TURN_CHILD_DRAIN_MS") + .ok() + .and_then(|raw| raw.parse::().ok()) + .map(std::time::Duration::from_millis) + .unwrap_or_else(|| std::time::Duration::from_secs(2)) +} + +/// The only path that removes a live turn. Finishes the scope (denying pending +/// interactions), harvests turn-scoped background children, clears the turn's +/// outbound budget, and removes it from the registry. Safe to call on +/// already-removed turns (returns early). +fn finish_turn(state: &DaemonState, turn_id: &str, reason: TurnFinishReason) { + let Some(handle) = state.turns.lock().unwrap().get(turn_id).cloned() else { + return; + }; + if !matches!(reason, TurnFinishReason::Complete) { + handle.cancel.cancel(); + } + let report = handle.scope.finish(reason); + let child_report = puffer_core::background_tasks::task_manager().stop_scoped_by_turn( + turn_id, + reason.as_str(), + daemon_child_drain_timeout(), + ); + // puffer_core::runtime::outbound_budget::budget_registry().clear_turn(turn_id); // enabled in outbound-budget task + if report.pending_permissions_resolved + report.pending_questions_resolved > 0 + || child_report.stop_requested + child_report.abandoned > 0 + { + eprintln!( + "turn {turn_id} finish({}): pending_perms={} pending_questions={} children_stopped={} children_abandoned={}", + reason.as_str(), + report.pending_permissions_resolved, + report.pending_questions_resolved, + child_report.stop_requested, + child_report.abandoned + ); + } + state.turns.lock().unwrap().remove(turn_id); +} + +/// Cancels one running turn by id: routes it through `finish_turn` (which flips +/// the cancel token, denies pending prompts, and harvests scoped children) and +/// reports the cancellation to listeners. Returns whether a turn existed. +/// Shared by the `cancel_turn` RPC (`CancelledByUser`) and the client-disconnect +/// watchdog (`ClientDisconnected`, #600). +fn cancel_turn_by_id(state: &DaemonState, turn_id: &str, reason: TurnFinishReason) -> bool { let handle = { let turns = state.turns.lock().unwrap(); turns.get(turn_id).cloned() @@ -4804,8 +4848,7 @@ fn cancel_turn_by_id(state: &DaemonState, turn_id: &str) -> bool { let Some(handle) = handle else { return false; }; - handle.cancel.cancel(); - let _report = handle.scope.finish(TurnFinishReason::CancelledByUser); + finish_turn(state, turn_id, reason); // Cancellation cleanup is best-effort: never let a failed report block the // cancel (especially on the disconnect path, where no client is waiting). if let (Some(session_uuid), Some(session_id)) = @@ -4826,7 +4869,6 @@ fn cancel_turn_by_id(state: &DaemonState, turn_id: &str) -> bool { } else { report_cancelled_sessionless_turn(state, &handle.channel, turn_id, &handle.cancel_reported); } - state.turns.lock().unwrap().remove(turn_id); true } @@ -4839,7 +4881,7 @@ fn cancel_all_active_turns(state: &DaemonState) -> usize { let turn_ids: Vec = state.turns.lock().unwrap().keys().cloned().collect(); turn_ids .iter() - .filter(|turn_id| cancel_turn_by_id(state, turn_id)) + .filter(|turn_id| cancel_turn_by_id(state, turn_id, TurnFinishReason::ClientDisconnected)) .count() } @@ -5615,7 +5657,7 @@ async fn start_turn(state: Arc, params: Value) -> Result { None, Some("attachment-staging"), ); - state.turns.lock().unwrap().remove(&turn_id); + finish_turn(&state, &turn_id, TurnFinishReason::Error); return Ok(json!({ "turnId": turn_id })); } } @@ -5654,7 +5696,6 @@ async fn start_turn(state: Arc, params: Value) -> Result { ); } - let state_for_thread = state.clone(); let turn_id_thread = turn_id.clone(); let turn_id_resp = turn_id.clone(); let channel_thread = channel.clone(); @@ -5697,7 +5738,7 @@ async fn start_turn(state: Arc, params: Value) -> Result { None, None, ); - setup_state.turns.lock().unwrap().remove(&turn_id_thread); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); return; } }; @@ -5713,7 +5754,7 @@ async fn start_turn(state: Arc, params: Value) -> Result { None, None, ); - setup_state.turns.lock().unwrap().remove(&turn_id_thread); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); return; } }; @@ -5778,7 +5819,7 @@ async fn start_turn(state: Arc, params: Value) -> Result { None, None, ); - setup_state.turns.lock().unwrap().remove(&turn_id_thread); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); return; } } @@ -5792,6 +5833,13 @@ async fn start_turn(state: Arc, params: Value) -> Result { scope.turn_id.clone(), ); } + app_state.set_current_turn_context(puffer_core::CurrentTurnContext { + turn_id: turn_id_thread.clone(), + session_id: session_id_for_thread.clone(), + task_id: monitor_reply_scope_for_thread + .as_ref() + .map(|s| s.task_id.clone()), + }); // Issue #560: reconcile the model's view of the browser with the real // tab registry every turn, instead of letting it trust stale // `connected:true` tool output replayed from the transcript. @@ -5811,7 +5859,7 @@ async fn start_turn(state: Arc, params: Value) -> Result { None, None, ); - setup_state.turns.lock().unwrap().remove(&turn_id_thread); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); return; } match persist_explicit_turn_routing( @@ -5839,7 +5887,7 @@ async fn start_turn(state: Arc, params: Value) -> Result { None, None, ); - setup_state.turns.lock().unwrap().remove(&turn_id_thread); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); return; } } @@ -5859,7 +5907,11 @@ async fn start_turn(state: Arc, params: Value) -> Result { &user_prompt_persisted_thread, &progress_thread, ); - setup_state.turns.lock().unwrap().remove(&turn_id_thread); + finish_turn( + &setup_state, + &turn_id_thread, + TurnFinishReason::CancelledByUser, + ); return; } app_state.set_exact_media_discovery_cache(setup_state.exact_media_discovery_cache(&inputs)); @@ -5896,7 +5948,7 @@ async fn start_turn(state: Arc, params: Value) -> Result { None, None, ); - setup_state.turns.lock().unwrap().remove(&turn_id_thread); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); return; } let _ = inputs @@ -5942,7 +5994,7 @@ async fn start_turn(state: Arc, params: Value) -> Result { None, Some("attachment-hydration"), ); - setup_state.turns.lock().unwrap().remove(&turn_id_thread); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); return; } @@ -6218,14 +6270,15 @@ async fn start_turn(state: Arc, params: Value) -> Result { ) }); + let mut finish_reason = TurnFinishReason::Complete; match outcome { Ok(turn) => { if cancel_reported_thread.load(Ordering::SeqCst) { - state_for_thread - .turns - .lock() - .unwrap() - .remove(&turn_id_thread); + finish_turn( + &setup_state, + &turn_id_thread, + TurnFinishReason::CancelledByUser, + ); return; } if !turn.assistant_text.is_empty() { @@ -6289,11 +6342,11 @@ async fn start_turn(state: Arc, params: Value) -> Result { } Err(err) => { if cancel_reported_thread.load(Ordering::SeqCst) { - state_for_thread - .turns - .lock() - .unwrap() - .remove(&turn_id_thread); + finish_turn( + &setup_state, + &turn_id_thread, + TurnFinishReason::CancelledByUser, + ); return; } eprintln!("turn {turn_id_thread} failed: {err:#}"); @@ -6319,14 +6372,11 @@ async fn start_turn(state: Arc, params: Value) -> Result { Some(raw), Some(category), ); + finish_reason = TurnFinishReason::Error; } } - state_for_thread - .turns - .lock() - .unwrap() - .remove(&turn_id_thread); + finish_turn(&setup_state, &turn_id_thread, finish_reason); }); Ok(json!({"turnId": turn_id_resp})) @@ -6412,7 +6462,7 @@ async fn start_slash_command_turn(state: Arc, params: Value) -> Res None, None, ); - setup_state.turns.lock().unwrap().remove(&turn_id_thread); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); return; } }; @@ -6428,7 +6478,7 @@ async fn start_slash_command_turn(state: Arc, params: Value) -> Res None, None, ); - setup_state.turns.lock().unwrap().remove(&turn_id_thread); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); return; } }; @@ -6442,6 +6492,11 @@ async fn start_slash_command_turn(state: Arc, params: Value) -> Res .unwrap_or(session_uuid) .to_string(); let mut app_state = AppState::from_session_record(cfg_for_turn, record); + app_state.set_current_turn_context(puffer_core::CurrentTurnContext { + turn_id: turn_id_thread.clone(), + session_id: session_id_for_thread.clone(), + task_id: None, + }); app_state.browser_status = browser_status_for_turn( &turn_browser_tab_context(&setup_state, &browser_root_session_id), session_used_browser, @@ -6511,7 +6566,7 @@ async fn start_slash_command_turn(state: Arc, params: Value) -> Res ) }); - match outcome { + let finish_reason = match outcome { Ok(()) => { let assistant_text = app_state .transcript @@ -6541,6 +6596,7 @@ async fn start_slash_command_turn(state: Arc, params: Value) -> Res "sessionId": session_id_for_thread.clone(), }), }); + TurnFinishReason::Complete } Err(err) => { eprintln!("slash command {turn_id_thread} failed: {err:#}"); @@ -6556,8 +6612,9 @@ async fn start_slash_command_turn(state: Arc, params: Value) -> Res None, None, ); + TurnFinishReason::Error } - } + }; let _ = ( cancel.clone(), @@ -6566,7 +6623,7 @@ async fn start_slash_command_turn(state: Arc, params: Value) -> Res progress.clone(), scope.clone(), ); - setup_state.turns.lock().unwrap().remove(&turn_id_thread); + finish_turn(&setup_state, &turn_id_thread, finish_reason); }); Ok(json!({"turnId": turn_id_resp})) @@ -6643,6 +6700,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R "assistantText": assistant_text, }), }); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Complete); } Err(error) => { if !(cancel_thread.is_cancelled() @@ -6656,9 +6714,9 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R None, ); } + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); } } - setup_state.turns.lock().unwrap().remove(&turn_id_thread); return; } @@ -6682,6 +6740,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R "assistantText": assistant_text, }), }); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Complete); } Err(error) => { if !(cancel_thread.is_cancelled() @@ -6695,9 +6754,9 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R None, ); } + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); } } - setup_state.turns.lock().unwrap().remove(&turn_id_thread); return; } @@ -6721,6 +6780,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R "assistantText": assistant_text, }), }); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Complete); } Err(error) => { if !(cancel_thread.is_cancelled() @@ -6734,9 +6794,9 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R None, ); } + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); } } - setup_state.turns.lock().unwrap().remove(&turn_id_thread); return; } @@ -6760,6 +6820,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R "assistantText": assistant_text, }), }); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Complete); } Err(error) => { if !(cancel_thread.is_cancelled() @@ -6773,9 +6834,9 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R None, ); } + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); } } - setup_state.turns.lock().unwrap().remove(&turn_id_thread); return; } @@ -6799,6 +6860,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R "assistantText": assistant_text, }), }); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Complete); } Err(error) => { if !(cancel_thread.is_cancelled() @@ -6812,9 +6874,9 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R None, ); } + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); } } - setup_state.turns.lock().unwrap().remove(&turn_id_thread); return; } @@ -6828,13 +6890,18 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R format!("build_runtime_inputs: {err:#}"), None, ); - setup_state.turns.lock().unwrap().remove(&turn_id_thread); + finish_turn(&setup_state, &turn_id_thread, TurnFinishReason::Error); return; } }; let cfg_for_turn = setup_state.config.lock().unwrap().clone(); let metadata = connector_setup_session_metadata(setup_state.cwd.clone(), Uuid::new_v4()); let mut app_state = AppState::new(cfg_for_turn, setup_state.cwd.clone(), metadata); + app_state.set_current_turn_context(puffer_core::CurrentTurnContext { + turn_id: turn_id_thread.clone(), + session_id: turn_id_thread.clone(), + task_id: None, + }); let stream_actor = app_state.system_actor(); let question_state = setup_state.clone(); @@ -6893,7 +6960,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R Ok::<_, anyhow::Error>(turn) }); - match outcome { + let finish_reason = match outcome { Ok(turn) => { setup_state.publish_event(ServerEnvelope::Event { event: channel_thread.clone(), @@ -6906,6 +6973,7 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R &stream_actor, ), }); + TurnFinishReason::Complete } Err(error) => { if !(cancel_thread.is_cancelled() && cancel_reported_thread.load(Ordering::SeqCst)) @@ -6918,9 +6986,10 @@ async fn start_connector_setup_turn(state: Arc, params: Value) -> R None, ); } + TurnFinishReason::Error } - } - setup_state.turns.lock().unwrap().remove(&turn_id_thread); + }; + finish_turn(&setup_state, &turn_id_thread, finish_reason); }); Ok(json!({"turnId": turn_id_resp, "setupId": turn_id_resp})) diff --git a/crates/puffer-cli/tests/daemon_turn_smoke.rs b/crates/puffer-cli/tests/daemon_turn_smoke.rs index 18180a773..b73789f84 100644 --- a/crates/puffer-cli/tests/daemon_turn_smoke.rs +++ b/crates/puffer-cli/tests/daemon_turn_smoke.rs @@ -976,6 +976,90 @@ fn daemon_rejects_concurrent_turn_for_same_session() { daemon.stop(); } +/// After a turn finishes (removed via `finish_turn`), a late `resolve_user_question` +/// RPC for that turn must be rejected, not silently answered `{"ok": true}`. +/// Also drives the child-drain path with a short `PUFFER_TURN_CHILD_DRAIN_MS`. +#[test] +fn daemon_rejects_resolve_after_turn_finished() { + let mock = MockOpenAiServer::start("Puffer smoke reply"); + let tempdir = tempfile::tempdir().expect("tempdir"); + let workspace = tempdir.path().join("workspace"); + let puffer_home = tempdir.path().join("home"); + let puffer_config = puffer_home.join(".puffer"); + std::fs::create_dir_all(&workspace).expect("workspace"); + std::fs::create_dir_all(&puffer_config).expect("puffer config"); + std::fs::write( + puffer_config.join("auth.json"), + json!({ + "format_version": 1, + "providers": { "openai": { "kind": "api_key", "key": "sk-test" } } + }) + .to_string(), + ) + .expect("auth store"); + let discovery_cache = tempdir.path().join("discovery.json"); + std::fs::write(&discovery_cache, discovery_cache_json()).expect("discovery cache"); + + let mut daemon = DaemonProcess::start_with_env( + &workspace, + &puffer_home, + &discovery_cache, + &[("PUFFER_TURN_CHILD_DRAIN_MS", "25")], + ); + let mut client = DaemonClient::connect(&daemon.handshake); + + client.rpc( + "update_config", + json!({ + "openaiBaseUrl": mock.base_url, + "defaultProvider": "openai", + "defaultModel": "openai/gpt-5", + }), + ); + let session = client.rpc( + "create_session", + json!({ "cwd": workspace.display().to_string() }), + ); + let session_id = session["sessionId"].as_str().expect("session id"); + + let turn = client.rpc( + "run_agent_turn", + json!({ + "sessionId": session_id, + "message": "Say exactly: Puffer smoke reply", + "permissionMode": "read-only", + }), + ); + let turn_id = turn["turnId"].as_str().expect("turn id").to_string(); + let complete = client.wait_for_event(|message| { + message["event"] == format!("session:{session_id}:event") + && message["payload"]["type"] == "turn-complete" + }); + assert_eq!(complete["payload"]["turnId"], turn_id); + + // The turn is finished and removed; a resolve for it must error loudly. + let error = client + .try_rpc( + "resolve_user_question", + json!({ + "turnId": turn_id, + "requestId": "req-never-issued", + "answers": {}, + "annotations": {}, + }), + ) + .expect_err("late resolve on a finished turn must be rejected"); + let serialized = error.to_string(); + assert!( + serialized.contains("no in-flight turn") + || serialized.contains("finished") + || serialized.contains("expired"), + "unexpected resolve error: {serialized}" + ); + + daemon.stop(); +} + struct DaemonProcess { child: Child, handshake: Value, From b15220d2463df1a245d65d066b4f9fd279480c99 Mon Sep 17 00:00:00 2001 From: Milhous Date: Wed, 8 Jul 2026 17:08:09 +0800 Subject: [PATCH 6/8] feat(runtime): per-turn outbound budget and turn-id origin stamping --- crates/puffer-cli/src/daemon.rs | 2 +- crates/puffer-core/lib.rs | 1 + crates/puffer-core/runtime.rs | 1 + .../claude_tools/workflow/connector_tools.rs | 148 +++++++++++++++--- crates/puffer-core/runtime/outbound_budget.rs | 103 ++++++++++++ 5 files changed, 233 insertions(+), 22 deletions(-) create mode 100644 crates/puffer-core/runtime/outbound_budget.rs diff --git a/crates/puffer-cli/src/daemon.rs b/crates/puffer-cli/src/daemon.rs index 367dbc98e..bb28af14f 100644 --- a/crates/puffer-cli/src/daemon.rs +++ b/crates/puffer-cli/src/daemon.rs @@ -4819,7 +4819,7 @@ fn finish_turn(state: &DaemonState, turn_id: &str, reason: TurnFinishReason) { reason.as_str(), daemon_child_drain_timeout(), ); - // puffer_core::runtime::outbound_budget::budget_registry().clear_turn(turn_id); // enabled in outbound-budget task + puffer_core::outbound_budget::budget_registry().clear_turn(turn_id); if report.pending_permissions_resolved + report.pending_questions_resolved > 0 || child_report.stop_requested + child_report.abandoned > 0 { diff --git a/crates/puffer-core/lib.rs b/crates/puffer-core/lib.rs index 9f6752559..729d7c6a0 100644 --- a/crates/puffer-core/lib.rs +++ b/crates/puffer-core/lib.rs @@ -78,6 +78,7 @@ pub use runtime::install_subscription_manager; pub use runtime::invalidate_copilot_bearer; pub use runtime::lambda_gate::LambdaHostConcreteToolBinding; pub use runtime::mcp_discovery; +pub use runtime::outbound_budget; pub use runtime::quota::{QuotaError, QuotaErrorKind, QUOTA_EXIT_CODE}; pub use runtime::resource_watcher; pub use runtime::resource_watcher::ResourceWatcher; diff --git a/crates/puffer-core/runtime.rs b/crates/puffer-core/runtime.rs index d11c0a232..265a47d28 100644 --- a/crates/puffer-core/runtime.rs +++ b/crates/puffer-core/runtime.rs @@ -48,6 +48,7 @@ mod microcompact; mod openai; mod openai_sse; mod openai_ws; +pub mod outbound_budget; pub(crate) mod overflow; mod permission_prompt; mod plan_events; diff --git a/crates/puffer-core/runtime/claude_tools/workflow/connector_tools.rs b/crates/puffer-core/runtime/claude_tools/workflow/connector_tools.rs index 466ab8f5b..5fa70be0e 100644 --- a/crates/puffer-core/runtime/claude_tools/workflow/connector_tools.rs +++ b/crates/puffer-core/runtime/claude_tools/workflow/connector_tools.rs @@ -265,16 +265,33 @@ fn execute_connector_act_with_dispatcher( parsed.action ) })?; + let turn_context = state.current_turn_context().cloned(); + let origin_session_id = turn_context + .as_ref() + .map(|c| c.session_id.clone()) + .unwrap_or_else(|| state.session.id.to_string()); + let origin_turn_id = turn_context + .as_ref() + .map(|c| c.turn_id.clone()) + .or_else(|| { + state + .monitor_reply_scope + .as_ref() + .map(|s| s.turn_id.clone()) + }); + let origin_task_id = turn_context + .as_ref() + .and_then(|c| c.task_id.clone()) + .or_else(|| { + state + .monitor_reply_scope + .as_ref() + .map(|s| s.task_id.clone()) + }); let origin = SendOrigin::LlmInitiated { - session_id: state.session.id.to_string(), - turn_id: state - .monitor_reply_scope - .as_ref() - .map(|scope| scope.turn_id.clone()), - task_id: state - .monitor_reply_scope - .as_ref() - .map(|scope| scope.task_id.clone()), + session_id: origin_session_id, + turn_id: origin_turn_id, + task_id: origin_task_id, }; let decision = puffer_subscriptions::outbound_gate::evaluate( &origin, @@ -300,6 +317,31 @@ fn execute_connector_act_with_dispatcher( parsed.action ); } + // Charge the per-turn outbound budget only for ungated *external* sends + // (gate returned Allowed AND the action has an external side effect — i.e. + // the Telegram `react` whitelist). Internal actions and rule automation are + // never charged (rule automation never reaches this LLM-initiated path). + if let Some(ctx) = &turn_context { + let is_external = puffer_subscriptions::outbound_gate::effective_action_permission( + &parsed.connector_slug, + &template, + &parsed.action, + ) + .map(|permission| permission.external_side_effect) + .unwrap_or(false); + if is_external { + use crate::runtime::outbound_budget::{budget_registry, OutboundBudgetKind}; + budget_registry() + .charge(&ctx.turn_id, OutboundBudgetKind::UngatedExternal) + .map_err(|_| { + anyhow::anyhow!( + "outbound external-action quota exceeded for this turn ({} sends). \ + Slow down; do not keep firing external actions.", + crate::runtime::outbound_budget::UNGATED_EXTERNAL_LIMIT_PER_TURN + ) + })?; + } + } let connection = parsed .connection_slug .clone() @@ -443,13 +485,30 @@ pub fn execute_connector_action_draft( action ) })?; + let turn_context = state.current_turn_context().cloned(); + let origin_session_id = turn_context + .as_ref() + .map(|c| c.session_id.clone()) + .unwrap_or_else(|| state.session.id.to_string()); + let origin_turn_id = turn_context + .as_ref() + .map(|c| c.turn_id.clone()) + .or_else(|| { + state + .monitor_reply_scope + .as_ref() + .map(|s| s.turn_id.clone()) + }); + // A draft's task binding is ONLY the model-supplied `task_id` (already + // scope-validated above). We deliberately do NOT fall back to the turn + // context or monitor scope: omitting `task_id` means a free-form send that + // is not tied to the task (plan-05 spec test-matrix item 4). Auto-stamping + // the scope's task here would relax the human-gated outbound semantics. + let origin_task_id = parsed.task_id.clone(); let origin = SendOrigin::LlmInitiated { - session_id: state.session.id.to_string(), - turn_id: state - .monitor_reply_scope - .as_ref() - .map(|scope| scope.turn_id.clone()), - task_id: parsed.task_id.clone(), + session_id: origin_session_id.clone(), + turn_id: origin_turn_id.clone(), + task_id: origin_task_id.clone(), }; let decision = puffer_subscriptions::outbound_gate::evaluate(&origin, &connector_slug, &template, &action); @@ -469,6 +528,20 @@ pub fn execute_connector_action_draft( ensure_connector_action_draft_connection(&manager, &connector_slug, &connection)?; let message = draft_message_text(&action_input) .context("ConnectorActionDraft requires a message body")?; + // Charge the per-turn draft budget after gate/validation so rejected inputs + // don't consume budget. Existing drafts await human review; don't pile on. + if let Some(ctx) = &turn_context { + use crate::runtime::outbound_budget::{budget_registry, OutboundBudgetKind}; + budget_registry() + .charge(&ctx.turn_id, OutboundBudgetKind::Draft) + .map_err(|_| { + anyhow::anyhow!( + "outbound draft quota exceeded for this turn ({} drafts). \ + Existing drafts are pending human review; do not create more.", + crate::runtime::outbound_budget::DRAFT_LIMIT_PER_TURN + ) + })?; + } let paths = ConfigPaths::discover(cwd); let store = OutboundStore::load(outbound_actions_path(&paths))?; let action_record = store.create_draft(NewOutboundDraft { @@ -480,12 +553,9 @@ pub fn execute_connector_action_draft( recipient_source, message, origin: OutboundOrigin { - session_id: state.session.id.to_string(), - turn_id: state - .monitor_reply_scope - .as_ref() - .map(|scope| scope.turn_id.clone()), - task_id: parsed.task_id.clone(), + session_id: origin_session_id.clone(), + turn_id: origin_turn_id.clone(), + task_id: origin_task_id.clone(), }, ttl_ms: None, })?; @@ -1476,6 +1546,42 @@ mod tests { .starts_with("sha256:")); } + #[test] + fn connector_action_draft_stamps_turn_id_from_current_turn_context() { + let home = tempfile::tempdir().unwrap(); + let _home_override = set_puffer_home_override(home.path()); + let connection_slug = "telegram-user-draft-ctx"; + ensure_connected_test_connection(connection_slug, "telegram-login"); + let (mut state, tmp) = make_state(); + state.set_current_turn_context(crate::CurrentTurnContext { + turn_id: "turn-1".into(), + session_id: state.session.id.to_string(), + task_id: None, + }); + + let raw = execute_connector_action_draft( + &mut state, + tmp.path(), + json!({ + "connector_slug": "telegram-login", + "connection_slug": connection_slug, + "action": "send_message", + "input": { + "chat_id": 123456789, + "message": "stamped by turn context" + } + }), + ) + .expect("draft should be saved"); + let payload: Value = serde_json::from_str(&raw).unwrap(); + let draft_id = payload["draft"]["id"].as_str().expect("draft id"); + + let paths = ConfigPaths::discover(tmp.path()); + let store = OutboundStore::load(outbound_actions_path(&paths)).unwrap(); + let action = store.get(draft_id).unwrap().expect("draft persisted"); + assert_eq!(action.origin.turn_id.as_deref(), Some("turn-1")); + } + #[test] fn connector_action_draft_rejects_deleted_connection_after_disconnect() { let home = tempfile::tempdir().unwrap(); diff --git a/crates/puffer-core/runtime/outbound_budget.rs b/crates/puffer-core/runtime/outbound_budget.rs new file mode 100644 index 000000000..13cf243ec --- /dev/null +++ b/crates/puffer-core/runtime/outbound_budget.rs @@ -0,0 +1,103 @@ +//! Per-turn outbound budget: bounds how many drafts / ungated external sends a +//! single LLM turn may produce, keyed by `turn_id`. Rule automation and +//! human-gated sends are governed by plan-05's outbound gate, not this budget. + +use std::collections::HashMap; +use std::sync::{Mutex, OnceLock}; + +pub const DRAFT_LIMIT_PER_TURN: usize = 5; +pub const UNGATED_EXTERNAL_LIMIT_PER_TURN: usize = 10; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OutboundBudgetKind { + Draft, + UngatedExternal, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OutboundBudgetError { + Exceeded, +} + +#[derive(Debug, Default)] +struct TurnBudget { + drafts: usize, + ungated_external: usize, +} + +pub struct OutboundBudgetRegistry { + budgets: Mutex>, + draft_limit: usize, + ungated_external_limit: usize, +} + +impl OutboundBudgetRegistry { + pub fn with_limits(draft_limit: usize, ungated_external_limit: usize) -> Self { + Self { + budgets: Mutex::new(HashMap::new()), + draft_limit, + ungated_external_limit, + } + } + + pub fn charge( + &self, + turn_id: &str, + kind: OutboundBudgetKind, + ) -> Result<(), OutboundBudgetError> { + let mut budgets = self.budgets.lock().unwrap(); + let budget = budgets.entry(turn_id.to_string()).or_default(); + let (count, limit) = match kind { + OutboundBudgetKind::Draft => (&mut budget.drafts, self.draft_limit), + OutboundBudgetKind::UngatedExternal => { + (&mut budget.ungated_external, self.ungated_external_limit) + } + }; + if *count >= limit { + return Err(OutboundBudgetError::Exceeded); + } + *count += 1; + Ok(()) + } + + pub fn clear_turn(&self, turn_id: &str) { + self.budgets.lock().unwrap().remove(turn_id); + } +} + +static REGISTRY: OnceLock = OnceLock::new(); + +pub fn budget_registry() -> &'static OutboundBudgetRegistry { + REGISTRY.get_or_init(|| { + OutboundBudgetRegistry::with_limits(DRAFT_LIMIT_PER_TURN, UNGATED_EXTERNAL_LIMIT_PER_TURN) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn budget_rejects_after_limit_and_clears_by_turn() { + let registry = OutboundBudgetRegistry::with_limits(2, 1); + registry + .charge("turn-1", OutboundBudgetKind::Draft) + .unwrap(); + registry + .charge("turn-1", OutboundBudgetKind::Draft) + .unwrap(); + assert_eq!( + registry + .charge("turn-1", OutboundBudgetKind::Draft) + .unwrap_err(), + OutboundBudgetError::Exceeded + ); + registry + .charge("turn-2", OutboundBudgetKind::Draft) + .unwrap(); // other turn unaffected + registry.clear_turn("turn-1"); + registry + .charge("turn-1", OutboundBudgetKind::Draft) + .unwrap(); + } +} From f2415a40c1e08b3c4bfbc9b37875733795e32423 Mon Sep 17 00:00:00 2001 From: Milhous Date: Wed, 8 Jul 2026 17:18:13 +0800 Subject: [PATCH 7/8] fix(desktop): remove double-layer resolver no-ops and silent frontend fallback Rust corbina suite green (112 passed incl. REGISTERED_TAURI_COMMANDS meta-test). Playwright E2E not run here (needs a built frontend + browsers); the change is behavior-preserving for the fake-daemon tests, which already exercised the error-propagation path since canInvokeTauri() is false in the browser context. --- apps/puffer-desktop/src-tauri/src/backend.rs | 20 +- apps/puffer-desktop/src-tauri/src/lib.rs | 43 +- apps/puffer-desktop/src-tauri/src/turn.rs | 779 ------------------- apps/puffer-desktop/src/lib/api/desktop.ts | 20 +- 4 files changed, 57 insertions(+), 805 deletions(-) delete mode 100644 apps/puffer-desktop/src-tauri/src/turn.rs diff --git a/apps/puffer-desktop/src-tauri/src/backend.rs b/apps/puffer-desktop/src-tauri/src/backend.rs index 6a373cabd..99413f010 100644 --- a/apps/puffer-desktop/src-tauri/src/backend.rs +++ b/apps/puffer-desktop/src-tauri/src/backend.rs @@ -294,7 +294,9 @@ impl BackendState { | "workflow_runs_list" | "workflow_run_show" => workflow_runtime_unavailable(method), "run_agent_turn" => self.run_agent_turn(events.clone(), params), - "resolve_permission" | "resolve_user_question" => Ok(json!({})), + "resolve_permission" | "resolve_user_question" => anyhow::bail!( + "{method} requires the daemon turn lifecycle; the in-process backend does not host interactive turns" + ), "cancel_turn" => { let turn_id = string_param(¶ms, &["turnId", "turn_id"])?; if let Some(flag) = self.turns.lock().unwrap().get(&turn_id) { @@ -2323,6 +2325,22 @@ mod tests { static TEST_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + #[test] + fn resolver_methods_error_instead_of_noop_success() { + let backend = BackendState::new(); + for method in ["resolve_permission", "resolve_user_question"] { + let result = backend.handle( + EventEmitter::websocket_only(), + method, + serde_json::json!({ + "turnId": "missing-turn", + "requestId": "missing-request" + }), + ); + assert!(result.is_err(), "{method} must not silently succeed"); + } + } + #[test] fn generate_media_result_serializes_artifacts_array() { let result = GenerateMediaResult { diff --git a/apps/puffer-desktop/src-tauri/src/lib.rs b/apps/puffer-desktop/src-tauri/src/lib.rs index 95fffd78c..467ce4a82 100644 --- a/apps/puffer-desktop/src-tauri/src/lib.rs +++ b/apps/puffer-desktop/src-tauri/src/lib.rs @@ -434,21 +434,46 @@ fn run_agent_turn( #[tauri::command] fn resolve_permission( - _turn_id: String, - _request_id: String, - _action: String, + app: AppHandle, + state: State<'_, SharedBackend>, + turn_id: String, + request_id: String, + action: String, ) -> Result<(), String> { - Ok(()) + backend_call( + app, + state, + "resolve_permission", + json!({ + "turnId": turn_id, + "requestId": request_id, + "action": action, + }), + ) + .map(|_| ()) } #[tauri::command] fn resolve_user_question( - _turn_id: String, - _request_id: String, - _answers: Value, - _annotations: Value, + app: AppHandle, + state: State<'_, SharedBackend>, + turn_id: String, + request_id: String, + answers: Value, + annotations: Value, ) -> Result<(), String> { - Ok(()) + backend_call( + app, + state, + "resolve_user_question", + json!({ + "turnId": turn_id, + "requestId": request_id, + "answers": answers, + "annotations": annotations, + }), + ) + .map(|_| ()) } #[tauri::command] diff --git a/apps/puffer-desktop/src-tauri/src/turn.rs b/apps/puffer-desktop/src-tauri/src/turn.rs deleted file mode 100644 index 83faa8ee9..000000000 --- a/apps/puffer-desktop/src-tauri/src/turn.rs +++ /dev/null @@ -1,779 +0,0 @@ -//! Drives a single agent turn from the Tauri host. -//! -//! Mirrors the pattern that `puffer-tui/src/flow.rs::handle_prompt_submit` -//! uses: clone state into a worker thread, stream TurnStreamEvents back via -//! an mpsc channel, block the permission-prompt callback on a one-shot -//! response channel, and persist events after the turn finishes. -//! -//! On the Tauri side we surface three commands: -//! * `run_agent_turn(session_id, message) -> turn_id` — starts a turn and -//! emits `session:{id}:event` events while it runs. -//! * `resolve_permission(turn_id, request_id, action)` — answers a paused -//! permission prompt. -//! * `cancel_turn(turn_id)` — best-effort interrupt (denies any pending -//! permission and drops the emitter so future events are ignored). - -use anyhow::{Context, Result}; -use indexmap::IndexMap; -use puffer_config::{ensure_workspace_dirs, load_config, ConfigPaths}; -use puffer_core::{ - execute_user_turn_streaming_with_permissions_and_cancel, with_user_question_prompt_handler, - AppState, BrowserPermissionPromptActionSet, BrowserPermissionPromptSource, - BrowserPermissionPromptTargetClass, CancelToken, MessageRole, PermissionPromptAction, - PermissionPromptRequest, TurnStreamEvent, UserQuestionPromptRequest, - UserQuestionPromptResponse, -}; -use puffer_provider_registry::{AuthStore, ProviderRegistry}; -use puffer_resources::load_resources; -use puffer_session_store::{MessageActor, SessionStore, TranscriptEvent, TurnBoundaryState}; -use serde::Serialize; -use serde_json::json; -use std::collections::HashMap; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::mpsc; -use std::sync::{Arc, Mutex}; -use std::thread; -use tauri::{AppHandle, Emitter}; -use uuid::Uuid; - -/// Registry of turns currently in flight. Inserted on spawn, removed when the -/// worker finishes. Permission responders are keyed by `(turn_id, request_id)`. -#[derive(Default)] -pub(crate) struct TurnRegistry { - inner: Mutex>, -} - -struct TurnEntry { - cancel: CancelToken, - pending: Arc>>>, - pending_questions: Arc>>>, - // Tracks the next permission-request id so the callback thread can - // hand the UI a stable label to reply to. - next_request_id: Arc, -} - -impl TurnEntry { - #[allow(dead_code)] - fn next_request_id(&self) -> &Arc { - &self.next_request_id - } -} - -impl TurnRegistry { - pub(crate) fn new() -> Self { - Self::default() - } - - fn insert(&self, turn_id: String, entry: TurnEntry) { - self.inner.lock().unwrap().insert(turn_id, entry); - } - - fn remove(&self, turn_id: &str) -> Option { - self.inner.lock().unwrap().remove(turn_id) - } - - fn with(&self, turn_id: &str, f: impl FnOnce(&TurnEntry) -> R) -> Option { - self.inner.lock().unwrap().get(turn_id).map(f) - } -} - -fn append_turn_boundary( - session_store: &SessionStore, - session_uuid: Uuid, - turn_id: &str, - state: TurnBoundaryState, -) -> Result<()> { - session_store.append_event( - session_uuid, - TranscriptEvent::TurnBoundary { - turn_id: turn_id.to_string(), - state, - }, - ) -} - -struct TranscriptTurnGuard<'a> { - session_store: &'a SessionStore, - session_uuid: Uuid, - turn_id: &'a str, - finished: bool, -} - -impl<'a> TranscriptTurnGuard<'a> { - fn start(session_store: &'a SessionStore, session_uuid: Uuid, turn_id: &'a str) -> Result { - append_turn_boundary(session_store, session_uuid, turn_id, TurnBoundaryState::Started)?; - Ok(Self { - session_store, - session_uuid, - turn_id, - finished: false, - }) - } - - fn finish(mut self) -> Result<()> { - append_turn_boundary( - self.session_store, - self.session_uuid, - self.turn_id, - TurnBoundaryState::Finished, - )?; - self.finished = true; - Ok(()) - } -} - -impl Drop for TranscriptTurnGuard<'_> { - fn drop(&mut self) { - if self.finished { - return; - } - let _ = append_turn_boundary( - self.session_store, - self.session_uuid, - self.turn_id, - TurnBoundaryState::Finished, - ); - } -} - -/// Payload emitted on `session:{id}:event`. -#[derive(Serialize, Clone)] -#[serde(tag = "type", rename_all = "kebab-case")] -enum EmittedEvent { - #[serde(rename_all = "camelCase")] - TurnStart { turn_id: String }, - #[serde(rename_all = "camelCase")] - TextDelta { - turn_id: String, - delta: String, - actor: MessageActor, - }, - #[serde(rename_all = "camelCase")] - ThinkingDelta { - turn_id: String, - delta: String, - actor: MessageActor, - }, - #[serde(rename_all = "camelCase")] - ToolCallsRequested { - turn_id: String, - requests: serde_json::Value, - actor: MessageActor, - }, - #[serde(rename_all = "camelCase")] - ToolInvocations { - turn_id: String, - invocations: serde_json::Value, - actor: MessageActor, - }, - #[serde(rename_all = "camelCase")] - PlanUpdated { - turn_id: String, - file_path: String, - content: Option, - actor: MessageActor, - }, - #[serde(rename_all = "camelCase")] - PlanCompleted { - turn_id: String, - file_path: String, - content: Option, - actor: MessageActor, - }, - #[serde(rename_all = "camelCase")] - Usage { - turn_id: String, - report: serde_json::Value, - actor: MessageActor, - }, - #[serde(rename_all = "camelCase")] - ReflectionCheckpoint { - turn_id: String, - summary: String, - actor: MessageActor, - }, - #[serde(rename_all = "camelCase")] - RetryAttempt { - turn_id: String, - attempt: usize, - max_attempts: usize, - error: String, - actor: MessageActor, - }, - #[serde(rename_all = "camelCase")] - PermissionRequest { - turn_id: String, - request_id: String, - tool_id: String, - summary: String, - reason: Option, - browser: Option, - actor: MessageActor, - }, - #[serde(rename_all = "camelCase")] - UserQuestionRequest { - turn_id: String, - request_id: String, - questions: serde_json::Value, - actor: MessageActor, - }, - #[serde(rename_all = "camelCase")] - TurnComplete { - turn_id: String, - assistant_text: String, - actor: MessageActor, - }, - #[serde(rename_all = "camelCase")] - TurnError { turn_id: String, error: String }, -} - -/// `run_agent_turn(session_id, message) -> turn_id`. -/// -/// Loads the session's AppState, resources, providers, and auth store; spawns -/// a worker thread to drive the turn; emits events on `session:{id}:event`. -/// Returns the generated `turn_id` so the caller can correlate incoming events -/// and send permission responses. -pub(crate) fn run_agent_turn( - app: AppHandle, - registry: Arc, - session_id: String, - message: String, -) -> Result { - let turn_id = Uuid::new_v4().to_string(); - let event_channel = format!("session:{session_id}:event"); - - // Load everything the runtime needs up-front so we can fail early instead - // of inside the worker thread. - let (config, resources, providers, auth_store, session_store, state) = - load_context(&session_id).map_err(|err| err.to_string())?; - - let cancel = CancelToken::new(); - let pending: Arc>>> = - Arc::new(Mutex::new(HashMap::new())); - let pending_questions: Arc>>> = - Arc::new(Mutex::new(HashMap::new())); - let next_request_id = Arc::new(AtomicU64::new(0)); - - registry.insert( - turn_id.clone(), - TurnEntry { - cancel: cancel.clone(), - pending: pending.clone(), - pending_questions: pending_questions.clone(), - next_request_id: next_request_id.clone(), - }, - ); - - // Announce start before we spawn so JS can subscribe deterministically. - let _ = app.emit( - &event_channel, - EmittedEvent::TurnStart { - turn_id: turn_id.clone(), - }, - ); - - let spawn_turn_id = turn_id.clone(); - let spawn_session_id = session_id.clone(); - let spawn_actor = state.assistant_actor(); - // Kept on this side of the move so the result handler can tell a - // user-initiated cancel apart from a genuine agent error. - let result_cancel = cancel.clone(); - thread::spawn(move || { - let result = drive_turn( - app.clone(), - event_channel.clone(), - spawn_turn_id.clone(), - spawn_session_id.clone(), - message, - config, - resources, - providers, - auth_store, - session_store, - state, - cancel, - pending, - pending_questions, - next_request_id, - ); - - match result { - Ok(assistant_text) => { - let _ = app.emit( - &event_channel, - EmittedEvent::TurnComplete { - turn_id: spawn_turn_id.clone(), - assistant_text, - actor: spawn_actor.clone(), - }, - ); - } - // A turn the user interrupted returns `Err("cancelled")` from the - // agent loop. Surface that as a clean completion (the UI already - // clears the canceled turn's live state) rather than a scary - // "Agent error" toast. - Err(_) if result_cancel.is_cancelled() => { - let _ = app.emit( - &event_channel, - EmittedEvent::TurnComplete { - turn_id: spawn_turn_id.clone(), - assistant_text: String::new(), - actor: spawn_actor.clone(), - }, - ); - } - Err(err) => { - let _ = app.emit( - &event_channel, - EmittedEvent::TurnError { - turn_id: spawn_turn_id.clone(), - error: err.to_string(), - }, - ); - } - } - - let _ = registry.remove(&spawn_turn_id); - }); - - Ok(turn_id) -} - -/// Pending permission resolution — `action` is one of -/// `"allow_once" | "allow_session" | "allow_all_session" | "deny"`. -pub(crate) fn resolve_permission( - registry: Arc, - turn_id: String, - request_id: String, - action: String, -) -> Result<(), String> { - let decoded = match action.as_str() { - "allow_once" => PermissionPromptAction::AllowOnce, - "allow_session" => PermissionPromptAction::AllowSession, - "allow_all_session" => PermissionPromptAction::AllowAllSession, - "deny" => PermissionPromptAction::Deny, - other => return Err(format!("unknown permission action `{other}`")), - }; - let responder = registry - .with(&turn_id, |entry| { - entry.pending.lock().unwrap().remove(&request_id) - }) - .flatten() - .ok_or_else(|| { - format!("no pending permission request `{request_id}` on turn `{turn_id}`") - })?; - responder - .send(decoded) - .map_err(|_| "agent worker already released the permission channel".to_string()) -} - -/// Resolves a pending `AskUserQuestion` prompt for an in-flight turn. -pub(crate) fn resolve_user_question( - registry: Arc, - turn_id: String, - request_id: String, - answers: serde_json::Map, - annotations: serde_json::Map, -) -> Result<(), String> { - let responder = registry - .with(&turn_id, |entry| { - entry.pending_questions.lock().unwrap().remove(&request_id) - }) - .flatten() - .ok_or_else(|| { - format!("no pending user question request `{request_id}` on turn `{turn_id}`") - })?; - responder - .send(UserQuestionPromptResponse { - answers, - annotations, - }) - .map_err(|_| "agent worker already released the user question channel".to_string()) -} - -/// Best-effort cancel: flips the cancel token and denies any outstanding -/// permission or question prompt. The agent loop observes the token at its -/// next boundary and returns `Err("cancelled")`, so the worker stops instead -/// of running on in the background after the user interrupts. -pub(crate) fn cancel_turn(registry: Arc, turn_id: String) -> Result<(), String> { - registry - .with(&turn_id, |entry| { - entry.cancel.cancel(); - let mut pending = entry.pending.lock().unwrap(); - for (_, tx) in pending.drain() { - let _ = tx.send(PermissionPromptAction::Deny); - } - let mut pending_questions = entry.pending_questions.lock().unwrap(); - for (_, tx) in pending_questions.drain() { - let _ = tx.send(UserQuestionPromptResponse { - answers: serde_json::Map::new(), - annotations: serde_json::Map::new(), - }); - } - }) - .ok_or_else(|| format!("no active turn `{turn_id}`")) -} - -// --------------------------------------------------------------------------- -// Worker-thread driver -// --------------------------------------------------------------------------- - -#[allow(clippy::too_many_arguments)] -fn drive_turn( - app: AppHandle, - event_channel: String, - turn_id: String, - session_id: String, - message: String, - _config: puffer_config::PufferConfig, - resources: puffer_resources::LoadedResources, - providers: ProviderRegistry, - auth_store: AuthStore, - session_store: SessionStore, - mut state: AppState, - cancel: CancelToken, - pending: Arc>>>, - pending_questions: Arc>>>, - next_request_id: Arc, -) -> Result { - // Persist the user message to the session transcript before the turn - // starts so the UI can reload and still see it even if the turn crashes. - state.push_message(MessageRole::User, message.clone()); - let session_uuid = state.session.id; - let turn_guard = TranscriptTurnGuard::start(&session_store, session_uuid, &turn_id)?; - session_store.append_event( - session_uuid, - TranscriptEvent::UserMessage { - text: message.clone(), - attachments: Vec::new(), - actor: Some(state.user_actor()), - }, - )?; - - let mut auth_store = auth_store; - let stream_actor = state.assistant_actor(); - - let on_event_app = app.clone(); - let on_event_channel = event_channel.clone(); - let on_event_turn_id = turn_id.clone(); - let on_event_actor = stream_actor.clone(); - let on_event = move |event: TurnStreamEvent| { - let payload = match event { - TurnStreamEvent::ThinkingDelta(delta) => EmittedEvent::ThinkingDelta { - turn_id: on_event_turn_id.clone(), - delta, - actor: on_event_actor.clone(), - }, - TurnStreamEvent::TextDelta(delta) => EmittedEvent::TextDelta { - turn_id: on_event_turn_id.clone(), - delta, - actor: on_event_actor.clone(), - }, - TurnStreamEvent::ToolCallsRequested(requests) => EmittedEvent::ToolCallsRequested { - turn_id: on_event_turn_id.clone(), - requests: serde_json::Value::Array( - requests - .into_iter() - .map(|r| { - json!({ - "callId": r.call_id, - "toolId": r.tool_id, - "input": r.input, - }) - }) - .collect(), - ), - actor: on_event_actor.clone(), - }, - TurnStreamEvent::ToolInvocations(invocations) => EmittedEvent::ToolInvocations { - turn_id: on_event_turn_id.clone(), - invocations: serde_json::Value::Array( - invocations - .into_iter() - .map(|i| { - json!({ - "callId": i.call_id, - "toolId": i.tool_id, - "input": i.input, - "output": i.output, - "success": i.success, - }) - }) - .collect(), - ), - actor: on_event_actor.clone(), - }, - TurnStreamEvent::PlanUpdated { file_path, content } => EmittedEvent::PlanUpdated { - turn_id: on_event_turn_id.clone(), - file_path, - content, - actor: on_event_actor.clone(), - }, - TurnStreamEvent::PlanCompleted { file_path, content } => EmittedEvent::PlanCompleted { - turn_id: on_event_turn_id.clone(), - file_path, - content, - actor: on_event_actor.clone(), - }, - TurnStreamEvent::Usage(report) => EmittedEvent::Usage { - turn_id: on_event_turn_id.clone(), - report: json!({ - "inputTokens": report.input_tokens, - "outputTokens": report.output_tokens, - "cacheReadTokens": report.cache_read_tokens, - "cacheCreationTokens": report.cache_creation_tokens, - }), - actor: on_event_actor.clone(), - }, - TurnStreamEvent::ReflectionCheckpoint(summary) => EmittedEvent::ReflectionCheckpoint { - turn_id: on_event_turn_id.clone(), - summary, - actor: on_event_actor.clone(), - }, - TurnStreamEvent::ReflectionTrace(_) => return, - TurnStreamEvent::RetryAttempt { - attempt, - max_attempts, - error, - } => EmittedEvent::RetryAttempt { - turn_id: on_event_turn_id.clone(), - attempt, - max_attempts, - error, - actor: on_event_actor.clone(), - }, - }; - let _ = on_event_app.emit(&on_event_channel, payload); - }; - - let on_perm_app = app.clone(); - let on_perm_channel = event_channel.clone(); - let on_perm_turn_id = turn_id.clone(); - let on_perm_pending = pending.clone(); - let on_perm_next_id = next_request_id.clone(); - let on_perm_actor = stream_actor.clone(); - let on_perm_cancel = cancel.clone(); - let on_permission = move |request: PermissionPromptRequest| -> PermissionPromptAction { - // Turn already interrupted: deny without surfacing a fresh prompt. (#671) - if on_perm_cancel.is_cancelled() { - return PermissionPromptAction::Deny; - } - let request_id = on_perm_next_id.fetch_add(1, Ordering::SeqCst).to_string(); - let (tx, rx) = mpsc::channel::(); - on_perm_pending - .lock() - .unwrap() - .insert(request_id.clone(), tx); - - let _ = on_perm_app.emit( - &on_perm_channel, - EmittedEvent::PermissionRequest { - turn_id: on_perm_turn_id.clone(), - request_id, - tool_id: request.tool_id, - summary: request.summary, - reason: request.reason, - browser: request - .browser - .as_ref() - .map(browser_permission_payload_json), - actor: on_perm_actor.clone(), - }, - ); - - // Block until the JS side resolves the prompt (or cancel denies it). - rx.recv().unwrap_or(PermissionPromptAction::Deny) - }; - - let on_question_app = app.clone(); - let on_question_channel = event_channel.clone(); - let on_question_turn_id = turn_id.clone(); - let on_question_pending = pending_questions.clone(); - let on_question_next_id = next_request_id.clone(); - let on_question_actor = stream_actor.clone(); - let on_question_cancel = cancel.clone(); - let on_user_question = move |request: UserQuestionPromptRequest| -> UserQuestionPromptResponse { - // Don't re-surface a question after the user interrupted the turn; - // returning empty lets the loop bail at its next cancel boundary - // instead of popping the prompt back up. (#671) - if on_question_cancel.is_cancelled() { - return UserQuestionPromptResponse { - answers: serde_json::Map::new(), - annotations: serde_json::Map::new(), - }; - } - let request_id = on_question_next_id - .fetch_add(1, Ordering::SeqCst) - .to_string(); - let (tx, rx) = mpsc::channel::(); - on_question_pending - .lock() - .unwrap() - .insert(request_id.clone(), tx); - - let _ = on_question_app.emit( - &on_question_channel, - EmittedEvent::UserQuestionRequest { - turn_id: on_question_turn_id.clone(), - request_id, - questions: request.questions, - actor: on_question_actor.clone(), - }, - ); - - rx.recv().unwrap_or(UserQuestionPromptResponse { - answers: serde_json::Map::new(), - annotations: serde_json::Map::new(), - }) - }; - - let outcome = with_user_question_prompt_handler(on_user_question, || { - execute_user_turn_streaming_with_permissions_and_cancel( - &mut state, - &resources, - &providers, - &mut auth_store, - &message, - None, - &cancel, - on_event, - on_permission, - ) - }) - .with_context(|| format!("run_agent_turn: session={session_id}"))?; - - // Persist tool invocations before the assistant text so reloads read in - // the same order as the work happened. - for invocation in &outcome.tool_invocations { - session_store.append_event( - session_uuid, - TranscriptEvent::ToolInvocation { - call_id: invocation.call_id.clone(), - tool_id: invocation.tool_id.clone(), - input: invocation.input.clone(), - output: invocation.output.clone(), - success: invocation.success, - actor: Some(stream_actor.clone()), - subject: state.tool_subject_actor(&invocation.tool_id, &invocation.output), - }, - )?; - } - if !outcome.assistant_text.is_empty() { - state.push_message(MessageRole::Assistant, outcome.assistant_text.clone()); - session_store.append_event( - session_uuid, - TranscriptEvent::AssistantMessage { - text: outcome.assistant_text.clone(), - actor: Some(stream_actor.clone()), - }, - )?; - } - for trace in &outcome.reflection_traces { - session_store.append_trace_event( - session_uuid, - puffer_session_store::TRACE_RUNTIME, - trace, - )?; - } - turn_guard.finish()?; - - Ok(outcome.assistant_text) -} - -fn browser_permission_payload_json( - payload: &puffer_core::BrowserPermissionPromptPayload, -) -> serde_json::Value { - json!({ - "source": match payload.source { - BrowserPermissionPromptSource::BrowserTool => "browser_tool", - BrowserPermissionPromptSource::BrowserInternalTool => "browser_internal_tool", - }, - "actionSet": match payload.action_set { - BrowserPermissionPromptActionSet::Inspect => "inspect", - BrowserPermissionPromptActionSet::Navigate => "navigate", - BrowserPermissionPromptActionSet::Interact => "interact", - BrowserPermissionPromptActionSet::Evaluate => "evaluate", - }, - "url": payload.url, - "origin": payload.origin, - "host": payload.host, - "targetClass": match payload.target_class { - BrowserPermissionPromptTargetClass::LocalDev => "local_dev", - BrowserPermissionPromptTargetClass::WorkspaceFile => "workspace_file", - BrowserPermissionPromptTargetClass::NonWorkspaceFile => "non_workspace_file", - BrowserPermissionPromptTargetClass::DataUrl => "data_url", - BrowserPermissionPromptTargetClass::OpenWeb => "open_web", - BrowserPermissionPromptTargetClass::Unknown => "unknown", - }, - "tabId": payload.tab_id, - "isCrossSession": payload.is_cross_session, - }) -} - -/// Loads everything the runtime needs, mirroring the setup path in -/// `puffer-cli/src/main.rs`. -fn load_context( - session_id: &str, -) -> Result<( - puffer_config::PufferConfig, - puffer_resources::LoadedResources, - ProviderRegistry, - AuthStore, - SessionStore, - AppState, -)> { - let workspace_root = std::env::current_dir()?; - let paths = ConfigPaths::discover(&workspace_root); - ensure_workspace_dirs(&paths)?; - let config = load_config(&paths)?; - let auth_path = paths.user_config_dir.join("auth.json"); - let auth_store = AuthStore::load(&auth_path)?; - let resources = load_resources(&paths)?; - - let mut providers = ProviderRegistry::new(); - for provider in &resources.providers { - providers.register_with_source( - provider.value.clone().into_descriptor(), - provider.source_info.as_provider_source(), - ); - } - providers.apply_openai_base_url_override(config.openai_base_url.as_deref()); - if !config.openai_headers.is_empty() { - providers.set_openai_headers( - config - .openai_headers - .clone() - .into_iter() - .collect::>(), - ); - } - if !config.openai_query_params.is_empty() { - providers.set_openai_query_params( - config - .openai_query_params - .clone() - .into_iter() - .collect::>(), - ); - } - let _ = providers.discover_and_merge_all(&auth_store); - - let session_store = SessionStore::from_paths(&paths)?; - let session_uuid = Uuid::parse_str(session_id).context("invalid session id")?; - let record = session_store.load_session(session_uuid)?; - let state = AppState::from_session_record(config.clone(), record); - Ok(( - config, - resources, - providers, - auth_store, - session_store, - state, - )) -} - -// Anchor the `json!` import so clippy doesn't trim it while we iterate. -#[allow(dead_code)] -fn _unused_anchor() -> serde_json::Value { - json!({}) -} diff --git a/apps/puffer-desktop/src/lib/api/desktop.ts b/apps/puffer-desktop/src/lib/api/desktop.ts index 7801011a8..139cecc86 100644 --- a/apps/puffer-desktop/src/lib/api/desktop.ts +++ b/apps/puffer-desktop/src/lib/api/desktop.ts @@ -2308,14 +2308,8 @@ export async function resolvePermission( requestId: string, action: PermissionAction ): Promise { - try { - const client = await ensureLocalDaemonClient(); - await client.request("resolve_permission", { turnId, requestId, action }); - return; - } catch (daemonError) { - if (!canInvokeTauri()) throw daemonError; - await invoke("resolve_permission", { turnId, requestId, action }); - } + const client = await ensureLocalDaemonClient(); + await client.request("resolve_permission", { turnId, requestId, action }); } /** Resolves a pending AskUserQuestion prompt for an in-flight turn. */ @@ -2325,14 +2319,8 @@ export async function resolveUserQuestion( answers: UserQuestionAnswers, annotations: UserQuestionAnnotations = {} ): Promise { - try { - const client = await ensureLocalDaemonClient(); - await client.request("resolve_user_question", { turnId, requestId, answers, annotations }); - return; - } catch (daemonError) { - if (!canInvokeTauri()) throw daemonError; - await invoke("resolve_user_question", { turnId, requestId, answers, annotations }); - } + const client = await ensureLocalDaemonClient(); + await client.request("resolve_user_question", { turnId, requestId, answers, annotations }); } /** Best-effort cancel: the current model/tool step completes then the turn From b6815f4ec8ea558c147b7b8bc4299ac1e60c8d22 Mon Sep 17 00:00:00 2001 From: Milhous Date: Wed, 8 Jul 2026 17:36:06 +0800 Subject: [PATCH 8/8] refactor(runtime): drop redundant cancel and dead BackgroundTaskManager::stop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit finish_turn shared the turn's CancelToken with its TurnScope, and TurnScope::finish already flips the token for every non-Complete reason, so the explicit handle.cancel.cancel() in finish_turn was dead duplication — removed, with cancellation ownership documented on scope. BackgroundTaskManager::stop had no production callers: real stops route through request_stop / stop_scoped_by_turn -> stop_one (which cancels the token and kills the process). The bare stop only flipped status, orphaning the worker/process — a footgun the owner-aware system replaced. Removed it and reworked the two unit tests that relied on it to exercise the real convergence paths (request_stop + complete -> Stopped; complete frees a slot). Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/puffer-cli/src/daemon.rs | 5 ++--- crates/puffer-core/runtime/background_tasks.rs | 9 --------- crates/puffer-core/runtime/background_tasks_tests.rs | 7 +++++-- 3 files changed, 7 insertions(+), 14 deletions(-) diff --git a/crates/puffer-cli/src/daemon.rs b/crates/puffer-cli/src/daemon.rs index bb28af14f..33a7d2cc9 100644 --- a/crates/puffer-cli/src/daemon.rs +++ b/crates/puffer-cli/src/daemon.rs @@ -4810,9 +4810,8 @@ fn finish_turn(state: &DaemonState, turn_id: &str, reason: TurnFinishReason) { let Some(handle) = state.turns.lock().unwrap().get(turn_id).cloned() else { return; }; - if !matches!(reason, TurnFinishReason::Complete) { - handle.cancel.cancel(); - } + // `scope` shares the handle's cancel token, and `finish` flips it for every + // non-`Complete` reason — so cancellation is owned there, not duplicated here. let report = handle.scope.finish(reason); let child_report = puffer_core::background_tasks::task_manager().stop_scoped_by_turn( turn_id, diff --git a/crates/puffer-core/runtime/background_tasks.rs b/crates/puffer-core/runtime/background_tasks.rs index 8f62dec08..84d81abf3 100644 --- a/crates/puffer-core/runtime/background_tasks.rs +++ b/crates/puffer-core/runtime/background_tasks.rs @@ -491,15 +491,6 @@ impl BackgroundTaskManager { report } - /// Marks a task as stopped (cancelled). - pub fn stop(&self, task_id: &str) { - let mut tasks = self.tasks.lock().unwrap(); - if let Some(task) = tasks.get_mut(task_id) { - task.info.status = BackgroundTaskStatus::Stopped; - task.info.completed_at = Some(now_ms()); - } - } - /// Returns a snapshot of the task info, if it exists. pub fn get_info(&self, task_id: &str) -> Option { let tasks = self.tasks.lock().unwrap(); diff --git a/crates/puffer-core/runtime/background_tasks_tests.rs b/crates/puffer-core/runtime/background_tasks_tests.rs index 12b1efd3b..bfbb0f433 100644 --- a/crates/puffer-core/runtime/background_tasks_tests.rs +++ b/crates/puffer-core/runtime/background_tasks_tests.rs @@ -326,7 +326,10 @@ fn task_manager_stop_then_complete_is_idempotent() { let mgr = BackgroundTaskManager::new(); let _ = mgr.register("s1", "stop test", None, None, false); - mgr.stop("s1"); + // request_stop marks the task Stopping; the worker ack (`complete`) then + // converges it to the terminal Stopped state. + mgr.request_stop("s1"); + mgr.complete("s1", true); let info = mgr.get_info("s1").unwrap(); assert_eq!(info.status, BackgroundTaskStatus::Stopped); assert_eq!(mgr.active_count(), 0); @@ -375,7 +378,7 @@ fn task_manager_has_capacity_reflects_limit() { } assert!(!mgr.has_capacity()); - mgr.stop("c-0"); + mgr.complete("c-0", true); assert!(mgr.has_capacity()); }