From b026df5e95821111b9db7637f1b0540f2948e3ca Mon Sep 17 00:00:00 2001 From: Neonforge <48338160+Neonforge98@users.noreply.github.com> Date: Tue, 25 Aug 2026 09:17:57 +0800 Subject: [PATCH 1/2] feat(work-runs): hand off remote conversation turns --- .../coordination/conversation_turn_bridge.rs | 630 ++++++++++++++++++ .../agent-core/src/core/coordination/mod.rs | 1 + .../coordination/work_item_run_dispatcher.rs | 412 +++++++++++- .../agent-core/src/state/commands/mod.rs | 1 + .../src/state/commands/work_runs.rs | 136 ++++ .../project-management/src/projects/schema.rs | 2 + .../src/projects/schema_tests.rs | 18 + .../src/work_run_service/dispatch.rs | 197 +++++- .../src/work_run_service/mod.rs | 5 +- .../src/work_run_service/terminal.rs | 48 +- .../src/work_run_service/tests.rs | 179 ++++- src-tauri/src/commands/handler_list.inc | 6 + src/api/tauri/rpc/procedures/index.ts | 1 + src/api/tauri/rpc/procedures/workRuns.ts | 29 + src/api/tauri/rpc/router.ts | 1 + src/api/tauri/rpc/schemas/index.ts | 1 + src/api/tauri/rpc/schemas/workRuns.ts | 33 + 17 files changed, 1676 insertions(+), 24 deletions(-) create mode 100644 src-tauri/crates/agent-core/src/core/coordination/conversation_turn_bridge.rs create mode 100644 src-tauri/crates/agent-core/src/state/commands/work_runs.rs create mode 100644 src/api/tauri/rpc/procedures/workRuns.ts create mode 100644 src/api/tauri/rpc/schemas/workRuns.ts diff --git a/src-tauri/crates/agent-core/src/core/coordination/conversation_turn_bridge.rs b/src-tauri/crates/agent-core/src/core/coordination/conversation_turn_bridge.rs new file mode 100644 index 0000000000..3ed2587c8f --- /dev/null +++ b/src-tauri/crates/agent-core/src/core/coordination/conversation_turn_bridge.rs @@ -0,0 +1,630 @@ +//! Hand-off for Work Item Run turns whose target Session lives on another +//! member's machine. +//! +//! A Discussion comment on a cloud-synced Work Item routes to the Work +//! Item's latest linked Session by id. When that Session is not local, the +//! durable dispatcher cannot resume it; instead the turn is offered to the +//! frontend, which executes it through the conversation-events plane (an +//! invisible local runner whose events are published under the root +//! Session id, design: docs/conversation-events-plane-design-2026-08-21.md). +//! +//! Acceptance is only an in-process claim: it never marks the durable outbox +//! delivered. The accepted frontend first prepares the real local runner, +//! sends the exact turn intent, then explicitly acknowledges that runner. +//! Only a successful durable ack consumes the claim; a crash before ack or a +//! transient ack failure therefore leaves the lease reclaimable/retryable. + +use std::{ + collections::HashMap, + sync::{Mutex, OnceLock}, + time::Duration, +}; + +use serde::{Deserialize, Serialize}; +use tauri::Emitter; +use tokio::sync::oneshot; +use tracing::warn; + +pub const CONVERSATION_TURN_REQUESTED_EVENT: &str = "orgii-work-run-conversation-turn"; +pub const ACCEPT_TIMEOUT: Duration = Duration::from_secs(10); +const CLAIM_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10); +const CLAIM_LEASE_EXTENSION_MS: i64 = 45_000; +// An offer must be accepted inside ACCEPT_TIMEOUT, before the frontend can +// wait for the canonical per-conversation queue. Cover one maximum 15-minute +// turn ahead of this claim, the 2-minute setup bound, and prepare/send/ack +// cushion. Keep it finite so a vanished webview cannot retain a durable +// dispatch forever (process death stops the task immediately). +const CLAIM_HEARTBEAT_MAX_LIFETIME: Duration = Duration::from_secs(20 * 60); + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct ConversationTurnRequest { + pub run_id: String, + pub dispatch_id: String, + pub org_id: String, + pub project_slug: Option, + pub work_item_id: String, + pub work_item_title: Option, + pub assigned_agent_id: Option, + pub root_session_id: String, + pub prepared_runner_session_id: Option, + pub content: String, + pub display_text: Option, + pub discussion_comment_ids: Vec, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct ConversationTurnOffer { + #[serde(flatten)] + request: ConversationTurnRequest, + claim_token: String, +} + +type AcceptSender = oneshot::Sender>; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ConversationTurnClaim { + pub run_id: String, + pub dispatch_id: String, + pub lease_token: String, + pub claim_token: String, + pub root_session_id: String, +} + +struct PendingOffer { + claim: ConversationTurnClaim, + accept_sender: Option, + claimed: bool, + runner_session_id: Option, +} + +static PENDING: OnceLock>> = OnceLock::new(); + +fn pending() -> &'static Mutex> { + PENDING.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn register(claim: ConversationTurnClaim) -> oneshot::Receiver> { + let (tx, rx) = oneshot::channel(); + let mut map = pending() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + map.insert( + claim.run_id.clone(), + PendingOffer { + claim, + accept_sender: Some(tx), + claimed: false, + runner_session_id: None, + }, + ); + rx +} + +fn unregister(run_id: &str, lease_token: &str, claim_token: &str) { + let mut map = pending() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if map.get(run_id).is_some_and(|offer| { + offer.claim.lease_token == lease_token && offer.claim.claim_token == claim_token + }) { + map.remove(run_id); + } +} + +/// Claim or abstain from the current offer. A successful claim only elects +/// one frontend listener; the outbox remains leased until `prepared_claim` +/// supplies the exact prepared local runner Session to the explicit ack. +/// Negative listeners deliberately do not resolve the shared offer: every +/// app window receives it, and an incapable window must not preempt a capable +/// winner. If nobody can claim, the request's bounded timeout drives retry. +pub fn accept(run_id: &str, claim_token: &str, outcome: Result<(), String>) -> Option { + let (sender, accepted_claim) = { + let mut map = pending() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + match outcome { + Ok(()) => { + let offer = map.get_mut(run_id)?; + if offer.claim.claim_token != claim_token || offer.claimed { + return None; + } + offer.claimed = true; + (offer.accept_sender.take(), offer.claim.clone()) + } + Err(_) => return None, + } + }; + let sender = sender?; + if sender.send(Ok(())).is_ok() { + let accepted_claim_token = accepted_claim.claim_token.clone(); + start_claim_heartbeat(accepted_claim); + Some(accepted_claim_token) + } else { + unregister(run_id, &accepted_claim.lease_token, claim_token); + None + } +} + +fn claim_is_current(claim: &ConversationTurnClaim) -> bool { + let map = pending() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + map.get(&claim.run_id).is_some_and(|offer| { + offer.claimed + && offer.claim.lease_token == claim.lease_token + && offer.claim.claim_token == claim.claim_token + }) +} + +fn start_claim_heartbeat(claim: ConversationTurnClaim) { + tauri::async_runtime::spawn(async move { + let deadline = tokio::time::Instant::now() + CLAIM_HEARTBEAT_MAX_LIFETIME; + loop { + tokio::time::sleep(CLAIM_HEARTBEAT_INTERVAL).await; + if !claim_is_current(&claim) { + break; + } + if tokio::time::Instant::now() >= deadline { + // Stop both renewal and in-process authority. A frontend that + // resumes after this bound must lose prepare/ack to the next + // fenced re-offer instead of sending under an expired claim. + unregister(&claim.run_id, &claim.lease_token, &claim.claim_token); + break; + } + let dispatch_id = claim.dispatch_id.clone(); + let lease_token = claim.lease_token.clone(); + let renewed = tokio::task::spawn_blocking(move || { + project_management::work_run_service::renew_dispatch_lease( + &dispatch_id, + &lease_token, + CLAIM_LEASE_EXTENSION_MS, + ) + }) + .await; + match renewed { + Ok(Ok(true)) => {} + Ok(Ok(false)) => { + unregister(&claim.run_id, &claim.lease_token, &claim.claim_token); + break; + } + Ok(Err(error)) => warn!( + run_id = %claim.run_id, + error = %error, + "[conversation-turn-bridge] claim heartbeat failed" + ), + Err(error) => warn!( + run_id = %claim.run_id, + error = %error, + "[conversation-turn-bridge] claim heartbeat task failed" + ), + } + } + }); +} + +/// Record the real local runner selected by the winning frontend without +/// acknowledging the durable outbox. Repeating prepare for the same runner +/// is idempotent; changing runners under one claim is rejected. +pub fn prepare_bound_claim( + run_id: &str, + claim_token: &str, + root_session_id: &str, + runner_session_id: &str, +) -> Result<(), String> { + let mut map = pending() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let Some(offer) = map.get_mut(run_id) else { + return Err(format!( + "conversation turn claim is not pending for run {run_id}" + )); + }; + if !offer.claimed { + return Err(format!( + "conversation turn offer has not been claimed for run {run_id}" + )); + } + if offer.claim.claim_token != claim_token { + return Err(format!( + "conversation turn claim token mismatch for run {run_id}" + )); + } + if offer.claim.root_session_id != root_session_id { + return Err(format!( + "conversation turn root mismatch for run {run_id}: expected {}, got {root_session_id}", + offer.claim.root_session_id + )); + } + if let Some(prepared) = offer.runner_session_id.as_deref() { + if prepared == runner_session_id { + return Ok(()); + } + return Err(format!( + "conversation turn runner mismatch for run {run_id}: prepared {prepared}, got {runner_session_id}" + )); + } + offer.runner_session_id = Some(runner_session_id.to_string()); + Ok(()) +} + +/// Snapshot a prepared claim only when the ack names the exact root and runner +/// selected during prepare. The caller must not consume it until the durable +/// outbox acknowledgement commits. +pub fn prepared_claim( + run_id: &str, + claim_token: &str, + root_session_id: &str, + runner_session_id: &str, +) -> Result { + let map = pending() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let Some(offer) = map.get(run_id) else { + return Err(format!( + "conversation turn claim is not pending for run {run_id}" + )); + }; + if !offer.claimed { + return Err(format!( + "conversation turn offer has not been claimed for run {run_id}" + )); + } + if offer.claim.claim_token != claim_token { + return Err(format!( + "conversation turn claim token mismatch for run {run_id}" + )); + } + if offer.claim.root_session_id != root_session_id { + return Err(format!( + "conversation turn root mismatch for run {run_id}: expected {}, got {root_session_id}", + offer.claim.root_session_id + )); + } + let Some(prepared_runner) = offer.runner_session_id.as_deref() else { + return Err(format!( + "conversation turn runner is not prepared for run {run_id}" + )); + }; + if prepared_runner != runner_session_id { + return Err(format!( + "conversation turn runner mismatch for run {run_id}: prepared {prepared_runner}, got {runner_session_id}" + )); + } + Ok(offer.claim.clone()) +} + +/// Consume a prepared claim after durable acknowledgement, but only if the +/// same lease/root/runner is still current. A late completion from an expired +/// claimant must never remove a newer re-offer. +pub fn consume_prepared_claim( + run_id: &str, + lease_token: &str, + claim_token: &str, + root_session_id: &str, + runner_session_id: &str, +) -> bool { + let mut map = pending() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let matches = map.get(run_id).is_some_and(|offer| { + offer.claimed + && offer.claim.lease_token == lease_token + && offer.claim.claim_token == claim_token + && offer.claim.root_session_id == root_session_id + && offer.runner_session_id.as_deref() == Some(runner_session_id) + }); + if matches { + map.remove(run_id); + } + matches +} + +/// Drop one accepted frontend claim without touching a newer re-offer. The +/// caller uses the returned durable lease identity to shorten the outbox lease; +/// if ack already consumed the claim this is an idempotent no-op. +pub fn release_claim(run_id: &str, claim_token: &str) -> Option { + let mut map = pending() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let matches = map + .get(run_id) + .is_some_and(|offer| offer.claimed && offer.claim.claim_token == claim_token); + if !matches { + return None; + } + map.remove(run_id).map(|offer| offer.claim) +} + +/// Phrased so `work_run_service::classify_failure` files it as a retryable +/// timeout that resumes the same target: the next attempt finds the +/// frontend once it is listening. +pub fn accept_timeout_message(run_id: &str) -> String { + format!( + "conversation turn hand-off timed out: no frontend accepted run {run_id} within {}s", + ACCEPT_TIMEOUT.as_secs() + ) +} + +pub fn rejection_message(reason: &str) -> String { + format!("conversation turn hand-off rejected: {reason}") +} + +/// Offer the turn to the frontend and wait for its acceptance. +pub async fn request_conversation_turn( + app: &tauri::AppHandle, + request: ConversationTurnRequest, + lease_token: String, +) -> Result<(), String> { + let run_id = request.run_id.clone(); + let claim_token = format!("claim_{}", uuid::Uuid::new_v4().simple()); + let claim = ConversationTurnClaim { + run_id: run_id.clone(), + dispatch_id: request.dispatch_id.clone(), + lease_token: lease_token.clone(), + claim_token: claim_token.clone(), + root_session_id: request.root_session_id.clone(), + }; + let receiver = register(claim); + let offer = ConversationTurnOffer { + request, + claim_token: claim_token.clone(), + }; + if let Err(err) = app.emit(CONVERSATION_TURN_REQUESTED_EVENT, offer) { + unregister(&run_id, &lease_token, &claim_token); + return Err(format!("conversation turn hand-off emit failed: {err}")); + } + match tokio::time::timeout(ACCEPT_TIMEOUT, receiver).await { + Ok(Ok(Ok(()))) => Ok(()), + Ok(Ok(Err(reason))) => Err(rejection_message(&reason)), + Ok(Err(_)) => { + unregister(&run_id, &lease_token, &claim_token); + Err("conversation turn hand-off was dropped before acceptance".to_string()) + } + Err(_) => { + unregister(&run_id, &lease_token, &claim_token); + warn!(run_id, "[conversation-turn-bridge] acceptance timed out"); + Err(accept_timeout_message(&run_id)) + } + } +} + +#[cfg(test)] +mod tests { + use project_management::projects::types::{ + WorkItemRunFailureClass, WorkItemRunRetryDisposition, + }; + use project_management::work_run_service::classify_failure; + + use super::*; + + #[test] + fn acceptance_timeout_is_a_retryable_resume() { + let failure = classify_failure(&accept_timeout_message("run-1"), true); + assert_eq!(failure.class, WorkItemRunFailureClass::Timeout); + assert!(failure.retryable); + assert_eq!( + failure.retry_disposition, + WorkItemRunRetryDisposition::ResumeSession + ); + } + + #[test] + fn rejection_is_filed_for_manual_review() { + let failure = classify_failure(&rejection_message("cloud sign-in required"), true); + assert_eq!(failure.class, WorkItemRunFailureClass::Unknown); + assert!(!failure.retryable); + } + + #[test] + fn heartbeat_horizon_covers_a_queued_turn_setup_and_transport_ack() { + const MAX_QUEUED_TURN: Duration = Duration::from_secs(15 * 60); + const MAX_BACKGROUND_SETUP: Duration = Duration::from_secs(2 * 60); + const ACK_CUSHION: Duration = Duration::from_secs(2 * 60); + + assert!( + CLAIM_HEARTBEAT_MAX_LIFETIME >= MAX_QUEUED_TURN + MAX_BACKGROUND_SETUP + ACK_CUSHION + ); + assert!( + Duration::from_millis(CLAIM_LEASE_EXTENSION_MS as u64) > CLAIM_HEARTBEAT_INTERVAL * 2, + "each renewal must survive more than one missed heartbeat" + ); + } + + #[tokio::test] + async fn failed_ack_keeps_prepared_claim_and_successful_ack_consumes_it() { + let receiver = register(ConversationTurnClaim { + run_id: "run-resolve".into(), + dispatch_id: "dispatch-1".into(), + lease_token: "lease-1".into(), + claim_token: "claim-1".into(), + root_session_id: "root-1".into(), + }); + assert_eq!( + accept("run-resolve", "claim-1", Ok(())), + Some("claim-1".to_string()) + ); + assert_eq!(receiver.await.expect("resolved"), Ok(())); + assert_eq!(accept("run-resolve", "claim-1", Ok(())), None); + + prepare_bound_claim("run-resolve", "claim-1", "root-1", "runner-1") + .expect("prepare runner"); + prepare_bound_claim("run-resolve", "claim-1", "root-1", "runner-1") + .expect("same prepare is idempotent"); + assert!(prepared_claim("run-resolve", "claim-1", "root-1", "runner-other").is_err()); + let claim = prepared_claim("run-resolve", "claim-1", "root-1", "runner-1") + .expect("first durable ack attempt snapshots the claim"); + assert_eq!(claim.dispatch_id, "dispatch-1"); + assert_eq!(claim.lease_token, "lease-1"); + + // Simulate a transient durable-store failure: no consume occurs, so + // the exact same prepared ack can be retried. + let retry = prepared_claim("run-resolve", "claim-1", "root-1", "runner-1") + .expect("failed durable ack remains retryable"); + assert_eq!(retry, claim); + assert!(consume_prepared_claim( + "run-resolve", + &claim.lease_token, + "claim-1", + "root-1", + "runner-1" + )); + assert!(!claim_is_current(&claim)); + assert!(prepared_claim("run-resolve", "claim-1", "root-1", "runner-1").is_err()); + assert!(release_claim("run-resolve", "claim-1").is_none()); + } + + #[tokio::test] + async fn a_reoffer_replaces_the_expired_lease_without_stale_cleanup() { + let first = register(ConversationTurnClaim { + run_id: "run-reoffer".into(), + dispatch_id: "dispatch-1".into(), + lease_token: "lease-old".into(), + claim_token: "claim-old".into(), + root_session_id: "root-1".into(), + }); + assert_eq!( + accept("run-reoffer", "claim-old", Ok(())), + Some("claim-old".to_string()) + ); + assert_eq!(first.await.expect("first accepted"), Ok(())); + + let second = register(ConversationTurnClaim { + run_id: "run-reoffer".into(), + dispatch_id: "dispatch-1".into(), + lease_token: "lease-new".into(), + claim_token: "claim-new".into(), + root_session_id: "root-1".into(), + }); + unregister("run-reoffer", "lease-old", "claim-old"); + assert_eq!(accept("run-reoffer", "claim-old", Ok(())), None); + assert_eq!( + accept("run-reoffer", "claim-new", Ok(())), + Some("claim-new".to_string()) + ); + assert_eq!(second.await.expect("second accepted"), Ok(())); + assert!(prepare_bound_claim("run-reoffer", "claim-old", "root-1", "runner-2").is_err()); + prepare_bound_claim("run-reoffer", "claim-new", "root-1", "runner-2").expect("prepare new"); + let claim = + prepared_claim("run-reoffer", "claim-new", "root-1", "runner-2").expect("new claim"); + assert_eq!(claim.lease_token, "lease-new"); + assert!(!consume_prepared_claim( + "run-reoffer", + "lease-old", + "claim-old", + "root-1", + "runner-2" + )); + assert!(prepared_claim("run-reoffer", "claim-new", "root-1", "runner-2").is_ok()); + } + + #[tokio::test] + async fn a_losing_rejection_cannot_clear_the_winning_claim() { + let receiver = register(ConversationTurnClaim { + run_id: "run-reject-race".into(), + dispatch_id: "dispatch-1".into(), + lease_token: "lease-1".into(), + claim_token: "claim-race".into(), + root_session_id: "root-1".into(), + }); + assert_eq!( + accept("run-reject-race", "claim-race", Ok(())), + Some("claim-race".to_string()) + ); + assert_eq!(receiver.await.expect("accepted"), Ok(())); + assert_eq!( + accept( + "run-reject-race", + "claim-race", + Err("signed out in another window".into()) + ), + None + ); + prepare_bound_claim("run-reject-race", "claim-race", "root-1", "runner-1") + .expect("prepare"); + assert!(prepared_claim("run-reject-race", "claim-race", "root-1", "runner-1").is_ok()); + } + + #[tokio::test] + async fn an_incapable_listener_cannot_preempt_a_capable_listener() { + let receiver = register(ConversationTurnClaim { + run_id: "run-multi-window".into(), + dispatch_id: "dispatch-multi-window".into(), + lease_token: "lease-multi-window".into(), + claim_token: "claim-multi-window".into(), + root_session_id: "root-multi-window".into(), + }); + + assert_eq!( + accept( + "run-multi-window", + "claim-multi-window", + Err("cloud sign-in required".into()) + ), + None + ); + assert_eq!( + accept("run-multi-window", "claim-multi-window", Ok(())), + Some("claim-multi-window".to_string()) + ); + assert_eq!(receiver.await.expect("capable listener accepted"), Ok(())); + } + + #[tokio::test] + async fn release_is_claim_token_fenced_and_stops_pre_send_authority() { + let receiver = register(ConversationTurnClaim { + run_id: "run-release".into(), + dispatch_id: "dispatch-release".into(), + lease_token: "lease-release".into(), + claim_token: "claim-release".into(), + root_session_id: "root-release".into(), + }); + assert_eq!( + accept("run-release", "claim-release", Ok(())), + Some("claim-release".to_string()) + ); + assert_eq!(receiver.await.expect("accepted"), Ok(())); + + assert!(release_claim("run-release", "claim-stale").is_none()); + let current = release_claim("run-release", "claim-release").expect("release exact claim"); + assert_eq!(current.lease_token, "lease-release"); + assert!(!claim_is_current(¤t)); + assert!(prepare_bound_claim( + "run-release", + "claim-release", + "root-release", + "runner-release" + ) + .is_err()); + assert!(release_claim("run-release", "claim-release").is_none()); + } + + #[test] + fn request_payload_uses_camel_case_wire_names() { + let value = serde_json::to_value(ConversationTurnOffer { + claim_token: "claim-wire".into(), + request: ConversationTurnRequest { + run_id: "run".into(), + dispatch_id: "dispatch".into(), + org_id: "org".into(), + project_slug: Some("proj".into()), + work_item_id: "WI-1".into(), + work_item_title: None, + assigned_agent_id: Some("agent-a".into()), + root_session_id: "root".into(), + prepared_runner_session_id: Some("runner".into()), + content: "body".into(), + display_text: Some("đź’¬ body".into()), + discussion_comment_ids: vec!["c1".into()], + }, + }) + .expect("serialize"); + assert_eq!(value["runId"], "run"); + assert_eq!(value["assignedAgentId"], "agent-a"); + assert_eq!(value["rootSessionId"], "root"); + assert_eq!(value["preparedRunnerSessionId"], "runner"); + assert_eq!(value["claimToken"], "claim-wire"); + assert_eq!(value["discussionCommentIds"][0], "c1"); + } +} diff --git a/src-tauri/crates/agent-core/src/core/coordination/mod.rs b/src-tauri/crates/agent-core/src/core/coordination/mod.rs index b9b9e3dff8..656266ed3e 100644 --- a/src-tauri/crates/agent-core/src/core/coordination/mod.rs +++ b/src-tauri/crates/agent-core/src/core/coordination/mod.rs @@ -28,6 +28,7 @@ pub mod agent_org_runs; pub mod agent_org_tasks; pub mod agent_org_watchdog; pub mod child_done_wake; +pub mod conversation_turn_bridge; pub mod routine_scheduler; pub mod work_item_recovery; pub mod work_item_run_dispatcher; diff --git a/src-tauri/crates/agent-core/src/core/coordination/work_item_run_dispatcher.rs b/src-tauri/crates/agent-core/src/core/coordination/work_item_run_dispatcher.rs index 4d778e2cb2..80e8821ad6 100644 --- a/src-tauri/crates/agent-core/src/core/coordination/work_item_run_dispatcher.rs +++ b/src-tauri/crates/agent-core/src/core/coordination/work_item_run_dispatcher.rs @@ -12,13 +12,14 @@ use std::{ }; use project_management::projects::types::{ - WorkItemDispatchLease, WorkItemExecutionLockReason, WorkItemRunTarget, WorkItemRunTrigger, - WorkItemRunUsage, + WorkItemDispatchLease, WorkItemExecutionLockReason, WorkItemRunStatus, WorkItemRunTarget, + WorkItemRunTrigger, WorkItemRunUsage, }; use project_management::work_run_service::{self, WorkItemRunTerminalOutcome}; use tauri::Manager; use tracing::{debug, error, info, warn}; +use super::conversation_turn_bridge; use crate::foundation::session_bridge::TurnIntentBridgeStatus; const LEASE_MS: i64 = 30_000; @@ -207,6 +208,33 @@ async fn reconcile_interrupted_session_runs(app: &tauri::AppHandle) { }; for (run, session) in candidates { + if let Some(prepared_runner) = + prepared_remote_runner_for(&run.target_snapshot.target, run.session_id.as_deref()) + { + let intent = + crate::foundation::session_bridge::get_turn_intent_status(prepared_runner, &run.id); + if intent.is_some_and(|status| { + matches!( + status, + TurnIntentBridgeStatus::Completed + | TurnIntentBridgeStatus::Cancelled + | TurnIntentBridgeStatus::Failed + | TurnIntentBridgeStatus::Stale + | TurnIntentBridgeStatus::Coalesced + | TurnIntentBridgeStatus::Rejected + ) + }) { + reconcile_terminal_intent(&run.id, prepared_runner).await; + continue; + } + if run.status == WorkItemRunStatus::Dispatching && intent.is_none() { + // Prepare persisted only a runner hint; transport never + // accepted this Run. Leave the leased outbox reclaimable so + // the next fenced offer can reuse or replace the hint. + continue; + } + } + let Some(session) = session else { warn!( run_id = %run.id, @@ -356,6 +384,33 @@ async fn dispatch_claim( lease: &WorkItemDispatchLease, ) -> Result<(), String> { let run = &lease.run; + + // A previous claimant may have prepared and sent this exact turn before + // crashing ahead of ack. Reclaim by acknowledging the persisted intent; + // never offer/send the same `(runner Session, Run)` twice. + if let Some(prepared_session_id) = + prepared_remote_runner_for(&run.target_snapshot.target, run.session_id.as_deref()) + { + if session_is_local(prepared_session_id).await? + && crate::foundation::session_bridge::get_turn_intent_status( + prepared_session_id, + &run.id, + ) + .is_some() + { + acknowledge_dispatch_delivery( + app, + &run.id, + &lease.dispatch_id, + &lease.lease_token, + None, + prepared_session_id, + ) + .await?; + return Ok(()); + } + } + let session_id = match &run.target_snapshot.target { WorkItemRunTarget::StartWorkItem { account_id, @@ -385,19 +440,74 @@ async fn dispatch_claim( } } WorkItemRunTarget::ResumeSession { session_id } => { - dispatch_session_turn(app, lease, session_id).await?; + if session_is_local(session_id).await? { + dispatch_session_turn(app, lease, session_id).await?; + } else { + request_remote_root_turn(app, lease, session_id).await?; + // Frontend acceptance only claims the offer. The durable + // outbox stays leased until prepare→send→ack names the exact + // local runner selected by the winning window. + return Ok(()); + } session_id.clone() } }; - let dispatch_id = lease.dispatch_id.clone(); - let lease_token = lease.lease_token.clone(); - let ack_session_id = session_id.clone(); - let acknowledged = tokio::task::spawn_blocking(move || { - work_run_service::acknowledge_dispatch_started(&dispatch_id, &lease_token, &ack_session_id) + acknowledge_dispatch_delivery( + app, + &run.id, + &lease.dispatch_id, + &lease.lease_token, + None, + &session_id, + ) + .await +} + +async fn acknowledge_dispatch_delivery( + app: &tauri::AppHandle, + run_id: &str, + dispatch_id: &str, + lease_token: &str, + claim_token: Option<&str>, + session_id: &str, +) -> Result<(), String> { + let dispatch_id = dispatch_id.to_string(); + let lease_token = lease_token.to_string(); + let ack_claim_token = claim_token.map(str::to_string); + let ack_session_id = session_id.to_string(); + let durable_ack = tokio::task::spawn_blocking(move || match ack_claim_token.as_deref() { + Some(claim_token) => work_run_service::acknowledge_claimed_dispatch_started( + &dispatch_id, + &lease_token, + claim_token, + &ack_session_id, + ), + None => work_run_service::acknowledge_dispatch_started( + &dispatch_id, + &lease_token, + &ack_session_id, + ), }) .await - .map_err(|err| format!("dispatch acknowledgement task failed: {err}"))??; + .map_err(|err| format!("dispatch acknowledgement task failed: {err}"))?; + let acknowledged = match durable_ack { + Ok(run) => run, + Err(err) => { + // The exact runtime terminal may win the race and retire the + // still-leased outbox. In that case the desired terminal is + // already durable and a late ack is an idempotent success. + let terminal_won = work_run_service::read(run_id) + .map(|run| { + run.status.is_terminal() && run.session_id.as_deref() == Some(session_id) + }) + .unwrap_or(false); + if terminal_won { + return Ok(()); + } + return Err(err); + } + }; let routine_origin = match &acknowledged.trigger { WorkItemRunTrigger::Routine { @@ -433,7 +543,7 @@ async fn dispatch_claim( if let Some((routine_id, fire_id)) = routine_origin { let fire_id_for_update = fire_id.clone(); let work_item_id = acknowledged.work_item_id.clone(); - let fire_session_id = session_id.clone(); + let fire_session_id = session_id.to_string(); let linked = tokio::task::spawn_blocking(move || { project_management::projects::io::mark_routine_fire_work_item_started( &fire_id_for_update, @@ -471,6 +581,171 @@ async fn dispatch_claim( Ok(()) } +/// Prepare the real local runner for an accepted remote-root hand-off. This +/// does not acknowledge delivery: the frontend prepares before transport +/// send, then calls the separate ack only after the exact intent is accepted. +pub(crate) async fn prepare_remote_conversation_runner( + run_id: &str, + claim_token: &str, + root_session_id: &str, + runner_session_id: &str, +) -> Result<(), String> { + if runner_session_id.trim().is_empty() { + return Err("runner_session_id is required".to_string()); + } + if !session_is_local(runner_session_id).await? { + return Err(format!( + "conversation turn runner {runner_session_id} is not a local Session" + )); + } + conversation_turn_bridge::prepare_bound_claim( + run_id, + claim_token, + root_session_id, + runner_session_id, + )?; + let claim = conversation_turn_bridge::prepared_claim( + run_id, + claim_token, + root_session_id, + runner_session_id, + )?; + let dispatch_id = claim.dispatch_id; + let lease_token = claim.lease_token; + let prepared_claim_token = claim_token.to_string(); + let prepared_session_id = runner_session_id.to_string(); + tokio::task::spawn_blocking(move || { + work_run_service::prepare_dispatch_session( + &dispatch_id, + &lease_token, + &prepared_claim_token, + &prepared_session_id, + ) + }) + .await + .map_err(|err| format!("dispatch prepare task failed: {err}"))??; + Ok(()) +} + +/// Acknowledge only the exact runner prepared under this claim. The in-memory +/// claim is consumed after the durable acknowledgement commits, so response +/// loss can be retried idempotently from the stored claim receipt. +pub(crate) async fn ack_remote_conversation_runner( + app: &tauri::AppHandle, + run_id: &str, + claim_token: &str, + root_session_id: &str, + runner_session_id: &str, +) -> Result<(), String> { + if runner_session_id.trim().is_empty() { + return Err("runner_session_id is required".to_string()); + } + if !session_is_local(runner_session_id).await? { + return Err(format!( + "conversation turn runner {runner_session_id} is not a local Session" + )); + } + let claim = match conversation_turn_bridge::prepared_claim( + run_id, + claim_token, + root_session_id, + runner_session_id, + ) { + Ok(claim) => claim, + Err(claim_error) => { + if durable_remote_ack_is_exact(run_id, claim_token, root_session_id, runner_session_id) + .await? + { + reconcile_terminal_intent(run_id, runner_session_id).await; + return Ok(()); + } + return Err(claim_error); + } + }; + + let acknowledged = acknowledge_dispatch_delivery( + app, + run_id, + &claim.dispatch_id, + &claim.lease_token, + Some(claim_token), + runner_session_id, + ) + .await; + match acknowledged { + Ok(()) => { + conversation_turn_bridge::consume_prepared_claim( + run_id, + &claim.lease_token, + claim_token, + root_session_id, + runner_session_id, + ); + Ok(()) + } + Err(ack_error) => { + if durable_remote_ack_is_exact(run_id, claim_token, root_session_id, runner_session_id) + .await? + { + conversation_turn_bridge::consume_prepared_claim( + run_id, + &claim.lease_token, + claim_token, + root_session_id, + runner_session_id, + ); + reconcile_terminal_intent(run_id, runner_session_id).await; + Ok(()) + } else { + Err(ack_error) + } + } + } +} + +async fn durable_remote_ack_is_exact( + run_id: &str, + claim_token: &str, + root_session_id: &str, + runner_session_id: &str, +) -> Result { + let durable_run_id = run_id.to_string(); + let durable_claim_token = claim_token.to_string(); + let durable_root_session_id = root_session_id.to_string(); + let durable_runner_session_id = runner_session_id.to_string(); + tokio::task::spawn_blocking(move || { + let run = work_run_service::read(&durable_run_id)?; + let delivered = work_run_service::delivered_dispatch_matches_claim( + &durable_run_id, + &durable_claim_token, + )?; + Ok(durable_remote_ack_identity_matches( + delivered, + &run.target_snapshot.target, + run.session_id.as_deref(), + &durable_root_session_id, + &durable_runner_session_id, + )) + }) + .await + .map_err(|err| format!("durable conversation ack lookup task failed: {err}"))? +} + +fn durable_remote_ack_identity_matches( + delivered: bool, + target: &WorkItemRunTarget, + bound_session_id: Option<&str>, + root_session_id: &str, + runner_session_id: &str, +) -> bool { + delivered + && bound_session_id == Some(runner_session_id) + && matches!( + target, + WorkItemRunTarget::ResumeSession { session_id } if session_id == root_session_id + ) +} + async fn dispatch_snapshotted_session_launch( app: &tauri::AppHandle, run: &project_management::projects::types::WorkItemRun, @@ -556,6 +831,88 @@ async fn dispatch_session_turn( .map(|_| ()) } +async fn session_is_local(session_id: &str) -> Result { + let lookup_id = session_id.to_string(); + let persisted = tokio::task::spawn_blocking(move || { + crate::session::persistence::get_session(&lookup_id).map_err(|err| err.to_string()) + }) + .await + .map_err(|err| format!("session lookup task failed: {err}"))??; + if persisted.is_some() { + return Ok(true); + } + Ok(crate::foundation::session_bridge::get_cli_tools_snapshot(session_id)?.is_some()) +} + +fn prepared_remote_runner_for<'a>( + target: &WorkItemRunTarget, + bound_session_id: Option<&'a str>, +) -> Option<&'a str> { + let WorkItemRunTarget::ResumeSession { session_id } = target else { + return None; + }; + bound_session_id.filter(|runner_session_id| *runner_session_id != session_id) +} + +/// The target Session belongs to another member: offer the turn to a local +/// conversation runner instead of treating the remote root as missing. +async fn request_remote_root_turn( + app: &tauri::AppHandle, + lease: &WorkItemDispatchLease, + root_session_id: &str, +) -> Result<(), String> { + let run = &lease.run; + let content = durable_resume_content(&run.input) + .unwrap_or_default() + .to_string(); + if content.trim().is_empty() { + return Err("durable resume dispatch is missing input.content".to_string()); + } + let discussion_comment_ids = run + .input + .get("discussionCommentIds") + .and_then(serde_json::Value::as_array) + .map(|ids| { + ids.iter() + .filter_map(serde_json::Value::as_str) + .map(str::to_string) + .collect() + }) + .unwrap_or_default(); + info!( + run_id = %run.id, + root_session_id, + "[work-run-dispatcher] target Session is not local; offering conversation turn" + ); + conversation_turn_bridge::request_conversation_turn( + app, + conversation_turn_bridge::ConversationTurnRequest { + run_id: run.id.clone(), + dispatch_id: lease.dispatch_id.clone(), + org_id: run.org_id.clone(), + project_slug: run.project_slug.clone(), + work_item_id: run.work_item_id.clone(), + work_item_title: run.target_snapshot.work_item_title.clone(), + assigned_agent_id: run.target_snapshot.agent_definition_id.clone(), + root_session_id: root_session_id.to_string(), + prepared_runner_session_id: prepared_remote_runner_for( + &run.target_snapshot.target, + run.session_id.as_deref(), + ) + .map(str::to_string), + content, + display_text: run + .input + .get("displayText") + .and_then(serde_json::Value::as_str) + .map(str::to_string), + discussion_comment_ids, + }, + lease.lease_token.clone(), + ) + .await +} + fn durable_resume_content(input: &serde_json::Value) -> Option<&str> { ["content", "prompt", "instruction"] .into_iter() @@ -642,8 +999,13 @@ fn lock_reason(trigger: &WorkItemRunTrigger) -> WorkItemExecutionLockReason { #[cfg(test)] mod tests { - use super::{durable_resume_content, lock_reason}; - use project_management::projects::types::{WorkItemExecutionLockReason, WorkItemRunTrigger}; + use super::{ + durable_remote_ack_identity_matches, durable_resume_content, lock_reason, + prepared_remote_runner_for, + }; + use project_management::projects::types::{ + WorkItemExecutionLockReason, WorkItemRunTarget, WorkItemRunTrigger, + }; #[test] fn trigger_maps_to_auditable_lock_reason() { @@ -666,4 +1028,30 @@ mod tests { let input = serde_json::json!({"prompt": "finish the routine"}); assert_eq!(durable_resume_content(&input), Some("finish the routine")); } + + #[test] + fn prepared_remote_runner_never_replaces_the_immutable_root_target() { + let target = WorkItemRunTarget::ResumeSession { + session_id: "shared-root".to_string(), + }; + assert_eq!( + prepared_remote_runner_for(&target, Some("hidden-runner")), + Some("hidden-runner") + ); + assert_eq!(prepared_remote_runner_for(&target, Some("shared-root")), None); + assert!(durable_remote_ack_identity_matches( + true, + &target, + Some("hidden-runner"), + "shared-root", + "hidden-runner", + )); + assert!(!durable_remote_ack_identity_matches( + true, + &target, + Some("runner-other"), + "shared-root", + "hidden-runner", + )); + } } diff --git a/src-tauri/crates/agent-core/src/state/commands/mod.rs b/src-tauri/crates/agent-core/src/state/commands/mod.rs index 78d3b32834..a78c82e8c2 100644 --- a/src-tauri/crates/agent-core/src/state/commands/mod.rs +++ b/src-tauri/crates/agent-core/src/state/commands/mod.rs @@ -5,6 +5,7 @@ pub mod desktop; pub mod routines; pub mod session; pub mod tools; +pub mod work_runs; // `pub use session::*` is required because `tauri::generate_handler!` // in `commands/handler_list.inc` resolves a long list of session-level diff --git a/src-tauri/crates/agent-core/src/state/commands/work_runs.rs b/src-tauri/crates/agent-core/src/state/commands/work_runs.rs new file mode 100644 index 0000000000..3013c9b4af --- /dev/null +++ b/src-tauri/crates/agent-core/src/state/commands/work_runs.rs @@ -0,0 +1,136 @@ +//! Frontend half of the remote-root conversation turn hand-off +//! (`core::coordination::conversation_turn_bridge`). + +use crate::core::coordination::conversation_turn_bridge; + +async fn shorten_conversation_claim_lease( + claim: &conversation_turn_bridge::ConversationTurnClaim, +) -> Result { + let dispatch_id = claim.dispatch_id.clone(); + let lease_token = claim.lease_token.clone(); + tokio::task::spawn_blocking(move || { + // One second is the store's minimum renewal. Shortening rather than + // mutating delivery state preserves normal fenced reclaim semantics. + project_management::work_run_service::renew_dispatch_lease( + &dispatch_id, + &lease_token, + 1_000, + ) + }) + .await + .map_err(|error| format!("conversation turn release task failed: {error}"))? +} + +#[tauri::command] +pub async fn work_run_conversation_turn_accept( + run_id: String, + claim_token: String, + accepted: bool, + reason: Option, +) -> Result, String> { + let outcome = if accepted { + Ok(()) + } else { + Err(reason.unwrap_or_else(|| "declined by frontend".to_string())) + }; + Ok(conversation_turn_bridge::accept( + &run_id, + &claim_token, + outcome, + )) +} + +#[tauri::command] +pub async fn work_run_conversation_turn_release( + run_id: String, + claim_token: String, +) -> Result { + let Some(claim) = conversation_turn_bridge::release_claim(&run_id, &claim_token) else { + return Ok(false); + }; + let shortened = shorten_conversation_claim_lease(&claim).await?; + if shortened { + crate::core::coordination::work_item_run_dispatcher::wake_from_watermark(); + } + Ok(true) +} + +/// File a live frontend failure through the durable dispatch state machine. +/// Unlike a crash-style release, classification decides whether this exact +/// failure is terminal or receives a bounded backoff retry. Taking the claim +/// first fences late prepare/ack calls; an ack that already committed has +/// consumed the claim and makes this an idempotent no-op. +#[tauri::command] +pub async fn work_run_conversation_turn_nack( + app: tauri::AppHandle, + run_id: String, + claim_token: String, + reason: String, +) -> Result { + let Some(claim) = conversation_turn_bridge::release_claim(&run_id, &claim_token) else { + return Ok(false); + }; + let dispatch_id = claim.dispatch_id.clone(); + let lease_token = claim.lease_token.clone(); + let failure_message = if reason.trim().is_empty() { + "conversation turn failed before transport acknowledgement".to_string() + } else { + reason + }; + let recorded = tokio::task::spawn_blocking(move || { + project_management::work_run_service::record_dispatch_failure( + &dispatch_id, + &lease_token, + &failure_message, + ) + }) + .await + .map_err(|error| format!("conversation turn nack task failed: {error}"))?; + let run = match recorded { + Ok(run) => run, + Err(error) => { + // If durable classification failed after the in-memory fence was + // removed, make the lease promptly reclaimable instead of leaving + // it at the heartbeat horizon. + let _ = shorten_conversation_claim_lease(&claim).await; + crate::core::coordination::work_item_run_dispatcher::wake_from_watermark(); + return Err(error); + } + }; + crate::orchestrator_notify::notify_routine_fire_dispatch_terminal(&run, &app).await; + Ok(true) +} + +#[tauri::command] +pub async fn work_run_conversation_turn_prepare_runner( + run_id: String, + claim_token: String, + root_session_id: String, + runner_session_id: String, +) -> Result<(), String> { + crate::core::coordination::work_item_run_dispatcher::prepare_remote_conversation_runner( + &run_id, + &claim_token, + &root_session_id, + &runner_session_id, + ) + .await +} + +#[tauri::command] +pub async fn work_run_conversation_turn_ack_runner( + app: tauri::AppHandle, + run_id: String, + claim_token: String, + root_session_id: String, + runner_session_id: String, +) -> Result<(), String> { + crate::core::coordination::work_item_run_dispatcher::ack_remote_conversation_runner( + &app, + &run_id, + &claim_token, + &root_session_id, + &runner_session_id, + ) + .await +} diff --git a/src-tauri/crates/project-management/src/projects/schema.rs b/src-tauri/crates/project-management/src/projects/schema.rs index 1885475545..6bca871e98 100644 --- a/src-tauri/crates/project-management/src/projects/schema.rs +++ b/src-tauri/crates/project-management/src/projects/schema.rs @@ -211,6 +211,7 @@ pub fn init_pm_service_tables(conn: &Connection) -> SqliteResult<()> { delivery_attempt INTEGER NOT NULL DEFAULT 0, available_at INTEGER NOT NULL, lease_token TEXT, + claim_token TEXT, lease_owner TEXT, lease_expires_at INTEGER, delivered_at INTEGER, @@ -324,6 +325,7 @@ pub fn init_pm_service_tables(conn: &Connection) -> SqliteResult<()> { ON pm_work_item_property_values(scope_key, work_item_id); "#, )?; + ensure_column(conn, "pm_dispatch_outbox", "claim_token", "TEXT")?; Ok(()) } diff --git a/src-tauri/crates/project-management/src/projects/schema_tests.rs b/src-tauri/crates/project-management/src/projects/schema_tests.rs index 362afb2d72..4e8117a3c1 100644 --- a/src-tauri/crates/project-management/src/projects/schema_tests.rs +++ b/src-tauri/crates/project-management/src/projects/schema_tests.rs @@ -73,6 +73,24 @@ fn init_is_idempotent() { init_project_tables(&conn).expect("second init should not fail"); } +#[test] +fn init_adds_claim_token_to_legacy_dispatch_outbox() { + let conn = open_in_memory(); + init_project_tables(&conn).expect("initial schema"); + conn.execute("ALTER TABLE pm_dispatch_outbox DROP COLUMN claim_token", []) + .expect("simulate legacy dispatch outbox"); + + init_project_tables(&conn).expect("upgrade legacy dispatch outbox"); + let columns: Vec = conn + .prepare("PRAGMA table_info(pm_dispatch_outbox)") + .expect("prepare column query") + .query_map([], |row| row.get::<_, String>(1)) + .expect("query columns") + .map(Result::unwrap) + .collect(); + assert!(columns.iter().any(|column| column == "claim_token")); +} + #[test] fn legacy_workitems_schema_is_rebuilt_for_org_level_items() { let conn = open_in_memory(); diff --git a/src-tauri/crates/project-management/src/work_run_service/dispatch.rs b/src-tauri/crates/project-management/src/work_run_service/dispatch.rs index 42933af2d2..e352655817 100644 --- a/src-tauri/crates/project-management/src/work_run_service/dispatch.rs +++ b/src-tauri/crates/project-management/src/work_run_service/dispatch.rs @@ -9,6 +9,8 @@ use super::store::{append_audit, db, iso8601, require_run}; use super::{error, DEFAULT_LEASE_MS}; const MAX_LEASE_MS: i64 = 5 * 60_000; +const MAX_LEASE_RENEWAL_MS: i64 = 60_000; +const PREPARED_LEASE_EXTENSION_MS: i64 = MAX_LEASE_RENEWAL_MS; /// Lease the oldest ready dispatch. Expired leases are reclaimed by the same /// query, so process death cannot strand a Run in `dispatching` forever. @@ -63,6 +65,7 @@ pub fn claim_dispatch_for_run( db(tx.execute( "UPDATE pm_dispatch_outbox SET status = 'leased', delivery_attempt = ?2, lease_token = ?3, + claim_token = NULL, lease_owner = ?4, lease_expires_at = ?5, updated_at = ?1 WHERE id = ?6", params![ @@ -208,6 +211,7 @@ pub fn claim_next_dispatch( db(tx.execute( "UPDATE pm_dispatch_outbox SET status = 'leased', delivery_attempt = ?2, lease_token = ?3, + claim_token = NULL, lease_owner = ?4, lease_expires_at = ?5, updated_at = ?1 WHERE id = ?6", params![ @@ -270,26 +274,213 @@ pub(super) fn leased_run_id( .ok_or_else(|| format!("{}:{}", error::STALE_LEASE, dispatch_id)) } +/// Extend one exact live dispatch lease. The lease token is the durable +/// worker fence, so a reclaimed row rejects a late heartbeat. +pub fn renew_dispatch_lease( + dispatch_id: &str, + lease_token: &str, + requested_extension_ms: i64, +) -> Result { + let extension_ms = requested_extension_ms.clamp(1_000, MAX_LEASE_RENEWAL_MS); + let now = now_ms(); + let connection = conn()?; + let changed = db(connection.execute( + "UPDATE pm_dispatch_outbox + SET lease_expires_at = ?3, updated_at = ?2 + WHERE id = ?1 AND status = 'leased' AND lease_token = ?4", + params![ + dispatch_id, + now, + now.saturating_add(extension_ms), + lease_token + ], + ))?; + Ok(changed == 1) +} + +/// Durably bind the local Session selected for a remote-root hand-off while +/// retaining the exact outbox lease. Prepare is idempotent for the same +/// claimant and Session; a claimant may never switch runners after prepare. +pub fn prepare_dispatch_session( + dispatch_id: &str, + lease_token: &str, + claim_token: &str, + session_id: &str, +) -> Result { + if session_id.trim().is_empty() { + return Err(format!("{}:session_id is required", error::INVALID_REQUEST)); + } + if claim_token.trim().is_empty() { + return Err(format!( + "{}:claim_token is required", + error::INVALID_REQUEST + )); + } + + let mut connection = conn()?; + let tx = db(connection.transaction_with_behavior(TransactionBehavior::Immediate))?; + let run_id = leased_run_id(&tx, dispatch_id, lease_token)?; + let previous_claim_token: Option = db(tx.query_row( + "SELECT claim_token FROM pm_dispatch_outbox + WHERE id = ?1 AND status = 'leased' AND lease_token = ?2", + params![dispatch_id, lease_token], + |row| row.get(0), + ))?; + if previous_claim_token + .as_deref() + .is_some_and(|current| current != claim_token) + { + return Err(format!("{}:{}", error::STALE_LEASE, dispatch_id)); + } + + let prepare_now = now_ms(); + let claimed = db(tx.execute( + "UPDATE pm_dispatch_outbox + SET claim_token = ?3, updated_at = ?4, + lease_expires_at = MAX(lease_expires_at, ?5) + WHERE id = ?1 AND status = 'leased' AND lease_token = ?2 + AND (claim_token IS NULL OR claim_token = ?3)", + params![ + dispatch_id, + lease_token, + claim_token, + prepare_now, + prepare_now.saturating_add(PREPARED_LEASE_EXTENSION_MS) + ], + ))?; + if claimed != 1 { + return Err(format!("{}:{}", error::STALE_LEASE, dispatch_id)); + } + + let existing = require_run(&tx, &run_id)?; + if existing.status != WorkItemRunStatus::Dispatching { + return Err(format!( + "{}:{} cannot prepare from {}", + error::INVALID_TRANSITION, + run_id, + existing.status.as_str() + )); + } + if let Some(bound) = existing.session_id.as_deref() { + if bound != session_id && previous_claim_token.as_deref() == Some(claim_token) { + return Err(format!( + "{}:{} prepared session {}, got {}", + error::INVALID_TRANSITION, + run_id, + bound, + session_id + )); + } + if bound == session_id { + db(tx.commit())?; + return read(&run_id); + } + } + + let now = now_ms(); + let changed = db(tx.execute( + "UPDATE pm_work_item_runs + SET session_id = ?2, updated_at = ?3 + WHERE id = ?1 AND status = 'dispatching'", + params![run_id, session_id, now], + ))?; + if changed != 1 { + return Err(format!( + "{}:{} could not prepare session", + error::INVALID_TRANSITION, + run_id + )); + } + let prepared = require_run(&tx, &run_id)?; + append_audit( + &tx, + &run_id, + "work_run.dispatch_prepared", + prepared.generation as i64, + prepared.project_slug.as_deref(), + &prepared.org_id, + serde_json::json!({ + "dispatchId": dispatch_id, + "sessionId": session_id, + }), + )?; + db(tx.commit())?; + read(&run_id) +} + +/// Query the durable receipt used to make the frontend acknowledgement RPC +/// safe to retry after a successful response was lost. +pub fn delivered_dispatch_matches_claim(run_id: &str, claim_token: &str) -> Result { + let connection = conn()?; + Ok(db(connection + .query_row( + "SELECT 1 FROM pm_dispatch_outbox + WHERE run_id = ?1 AND status = 'delivered' AND claim_token = ?2 + LIMIT 1", + params![run_id, claim_token], + |_| Ok(()), + ) + .optional())? + .is_some()) +} + /// Acknowledge that the runtime accepted the dispatch and materialized a /// Session. This is a Run transition only; Work Item status is untouched. pub fn acknowledge_dispatch_started( dispatch_id: &str, lease_token: &str, session_id: &str, +) -> Result { + acknowledge_dispatch_started_inner(dispatch_id, lease_token, None, session_id) +} + +pub fn acknowledge_claimed_dispatch_started( + dispatch_id: &str, + lease_token: &str, + claim_token: &str, + session_id: &str, +) -> Result { + if claim_token.trim().is_empty() { + return Err(format!( + "{}:claim_token is required", + error::INVALID_REQUEST + )); + } + acknowledge_dispatch_started_inner(dispatch_id, lease_token, Some(claim_token), session_id) +} + +fn acknowledge_dispatch_started_inner( + dispatch_id: &str, + lease_token: &str, + claim_token: Option<&str>, + session_id: &str, ) -> Result { if session_id.trim().is_empty() { return Err(format!("{}:session_id is required", error::INVALID_REQUEST)); } let mut connection = conn()?; let tx = db(connection.transaction_with_behavior(TransactionBehavior::Immediate))?; - let run_id = leased_run_id(&tx, dispatch_id, lease_token)?; + let run_id = match claim_token { + Some(claim_token) => db(tx + .query_row( + "SELECT run_id FROM pm_dispatch_outbox + WHERE id = ?1 AND status = 'leased' AND lease_token = ?2 + AND claim_token = ?3", + params![dispatch_id, lease_token, claim_token], + |row| row.get(0), + ) + .optional())? + .ok_or_else(|| format!("{}:{}", error::STALE_LEASE, dispatch_id))?, + None => leased_run_id(&tx, dispatch_id, lease_token)?, + }; let now = now_ms(); db(tx.execute( "UPDATE pm_dispatch_outbox SET status = 'delivered', delivered_at = ?3, updated_at = ?3, lease_token = NULL, lease_owner = NULL, lease_expires_at = NULL - WHERE id = ?1 AND lease_token = ?2", - params![dispatch_id, lease_token, now], + WHERE id = ?1 AND lease_token = ?2 + AND (?4 IS NULL OR claim_token = ?4)", + params![dispatch_id, lease_token, now, claim_token], ))?; let changed = db(tx.execute( "UPDATE pm_work_item_runs diff --git a/src-tauri/crates/project-management/src/work_run_service/mod.rs b/src-tauri/crates/project-management/src/work_run_service/mod.rs index d111727cf0..00f3859985 100644 --- a/src-tauri/crates/project-management/src/work_run_service/mod.rs +++ b/src-tauri/crates/project-management/src/work_run_service/mod.rs @@ -20,8 +20,9 @@ mod tests; pub use consumer_cursor::{advance_consumer_cursor, initialize_consumer_cursor}; pub use dispatch::{ - acknowledge_dispatch_started, claim_dispatch_for_run, claim_next_dispatch, - has_claimable_dispatch, next_dispatch_due_at_ms, + acknowledge_claimed_dispatch_started, acknowledge_dispatch_started, claim_dispatch_for_run, + claim_next_dispatch, delivered_dispatch_matches_claim, has_claimable_dispatch, + next_dispatch_due_at_ms, prepare_dispatch_session, renew_dispatch_lease, }; pub(crate) use enqueue::enqueue_in_transaction; pub use enqueue::{enqueue, enqueue_for_inline_dispatch, enqueue_with_receipt}; diff --git a/src-tauri/crates/project-management/src/work_run_service/terminal.rs b/src-tauri/crates/project-management/src/work_run_service/terminal.rs index c2965d9a8b..3d5f253c7c 100644 --- a/src-tauri/crates/project-management/src/work_run_service/terminal.rs +++ b/src-tauri/crates/project-management/src/work_run_service/terminal.rs @@ -296,6 +296,7 @@ pub fn record_run_terminal( } } if existing.status.is_terminal() { + close_terminal_dispatch(&tx, run_id, now_ms())?; release_path_lock(&tx, run_id)?; db(tx.commit())?; crate::projects::events::notify_work_item_dispatch_ready(); @@ -338,12 +339,16 @@ pub fn record_run_terminal( let usage_json = serde_json::to_string(&usage) .map_err(|err| format!("work run usage serialization: {err}"))?; let now = now_ms(); - db(tx.execute( + let changed = db(tx.execute( "UPDATE pm_work_item_runs SET status = ?2, failure_json = ?3, usage_json = ?4, session_id = COALESCE(session_id, ?6), completed_at = ?5, updated_at = ?5 - WHERE id = ?1 AND status IN ('running', 'waiting', 'dispatching')", + WHERE id = ?1 + AND ( + status IN ('running', 'waiting', 'dispatching') + OR (status IN ('queued', 'deferred') AND session_id = ?6) + )", params![ run_id, status.as_str(), @@ -353,6 +358,19 @@ pub fn record_run_terminal( expected_session_id ], ))?; + if changed != 1 { + return Err(format!( + "{}:{} cannot record exact terminal from {} (expected session {})", + error::INVALID_TRANSITION, + run_id, + existing.status.as_str(), + expected_session_id.unwrap_or("none") + )); + } + // A runtime terminal proves transport accepted this Run. Retire a still + // leased outbox in the same transaction so it cannot be re-offered if the + // explicit frontend acknowledgement lost its response race. + close_terminal_dispatch(&tx, run_id, now)?; release_path_lock(&tx, run_id)?; let updated = require_run(&tx, run_id)?; append_audit( @@ -381,6 +399,22 @@ pub fn record_run_terminal( Ok(persisted) } +fn close_terminal_dispatch( + tx: &rusqlite::Transaction<'_>, + run_id: &str, + terminal_at: i64, +) -> Result<(), String> { + db(tx.execute( + "UPDATE pm_dispatch_outbox + SET status = 'delivered', delivered_at = COALESCE(delivered_at, ?2), + updated_at = ?2, lease_token = NULL, lease_owner = NULL, + lease_expires_at = NULL + WHERE run_id = ?1 AND status IN ('pending', 'retry_wait', 'leased')", + params![run_id, terminal_at], + ))?; + Ok(()) +} + fn project_succeeded_run_for_review(run: &WorkItemRun) { let projection = match work_service::project_run_success_to_review( run.project_slug.as_deref(), @@ -580,7 +614,15 @@ pub fn retry(run_id: &str, idempotency_key: &str) -> Result run_id ) })?; - target_snapshot.target = WorkItemRunTarget::ResumeSession { session_id }; + let keeps_remote_root = matches!( + &target_snapshot.target, + WorkItemRunTarget::ResumeSession { + session_id: root_session_id + } if root_session_id != &session_id + ); + if !keeps_remote_root { + target_snapshot.target = WorkItemRunTarget::ResumeSession { session_id }; + } } enqueue(EnqueueWorkItemRunRequest { project_slug: previous.project_slug, diff --git a/src-tauri/crates/project-management/src/work_run_service/tests.rs b/src-tauri/crates/project-management/src/work_run_service/tests.rs index 2e010cc693..a2ee4a6db6 100644 --- a/src-tauri/crates/project-management/src/work_run_service/tests.rs +++ b/src-tauri/crates/project-management/src/work_run_service/tests.rs @@ -316,6 +316,129 @@ fn dispatch_claim_is_leased_and_ack_requires_matching_token() { .is_none()); } +#[test] +fn claimed_ack_is_fenced_and_prepare_is_runner_stable() { + let _sandbox = test_env::sandbox(); + seed(); + let run = enqueue(request("manual:claim-fence")).expect("enqueue"); + let lease = claim_next_dispatch("desktop-claim-fence", 30_000) + .expect("claim") + .expect("lease"); + + let prepared = prepare_dispatch_session( + &lease.dispatch_id, + &lease.lease_token, + "claim-winner", + "runner-winner", + ) + .expect("prepare winner"); + assert_eq!(prepared.status, WorkItemRunStatus::Dispatching); + assert_eq!(prepared.session_id.as_deref(), Some("runner-winner")); + prepare_dispatch_session( + &lease.dispatch_id, + &lease.lease_token, + "claim-winner", + "runner-winner", + ) + .expect("same prepare is idempotent"); + assert!(prepare_dispatch_session( + &lease.dispatch_id, + &lease.lease_token, + "claim-winner", + "runner-other", + ) + .expect_err("one claimant cannot switch runners") + .starts_with(error::INVALID_TRANSITION)); + + let stale = acknowledge_claimed_dispatch_started( + &lease.dispatch_id, + &lease.lease_token, + "claim-stale", + "runner-winner", + ) + .expect_err("losing claimant cannot acknowledge"); + assert!(stale.starts_with(error::STALE_LEASE), "{stale}"); + + let started = acknowledge_claimed_dispatch_started( + &lease.dispatch_id, + &lease.lease_token, + "claim-winner", + "runner-winner", + ) + .expect("winner ack"); + assert_eq!(started.status, WorkItemRunStatus::Running); + assert!(delivered_dispatch_matches_claim(&run.id, "claim-winner").expect("durable receipt")); + assert!(!delivered_dispatch_matches_claim(&run.id, "claim-stale").expect("stale receipt")); +} + +#[test] +fn expired_prepared_lease_is_reoffered_with_a_new_fence() { + let _sandbox = test_env::sandbox(); + seed(); + let run = enqueue(request("manual:prepare-crash")).expect("enqueue"); + let first = claim_next_dispatch("desktop-before-crash", 30_000) + .expect("first claim") + .expect("first lease"); + prepare_dispatch_session( + &first.dispatch_id, + &first.lease_token, + "claim-before-crash", + "runner-before-crash", + ) + .expect("prepare before crash"); + + conn() + .expect("connection") + .execute( + "UPDATE pm_dispatch_outbox SET lease_expires_at = ?2 WHERE id = ?1", + rusqlite::params![first.dispatch_id, now_ms().saturating_sub(1)], + ) + .expect("expire first lease"); + let second = claim_next_dispatch("desktop-after-crash", 30_000) + .expect("reclaim") + .expect("reoffered lease"); + assert_eq!(second.run.id, run.id); + assert_ne!(second.lease_token, first.lease_token); + assert_eq!(second.delivery_attempt, 2); + assert_eq!( + second.run.session_id.as_deref(), + Some("runner-before-crash") + ); + prepare_dispatch_session( + &second.dispatch_id, + &second.lease_token, + "claim-after-crash", + "runner-before-crash", + ) + .expect("new claimant reuses prepared hint"); + assert!(acknowledge_claimed_dispatch_started( + &first.dispatch_id, + &first.lease_token, + "claim-before-crash", + "runner-before-crash", + ) + .expect_err("expired claimant stays fenced") + .starts_with(error::STALE_LEASE)); +} + +#[test] +fn lease_heartbeat_renews_only_the_exact_live_lease() { + let _sandbox = test_env::sandbox(); + seed(); + enqueue(request("manual:lease-heartbeat")).expect("enqueue"); + let lease = claim_next_dispatch("desktop-heartbeat", 1_000) + .expect("claim") + .expect("lease"); + assert!( + renew_dispatch_lease(&lease.dispatch_id, &lease.lease_token, 45_000,) + .expect("renew exact lease") + ); + assert!( + !renew_dispatch_lease(&lease.dispatch_id, "lease-stale", 45_000) + .expect("stale heartbeat is a no-op") + ); +} + #[test] fn latest_for_session_returns_attached_execution_episode() { let _sandbox = test_env::sandbox(); @@ -611,6 +734,13 @@ fn turn_can_finish_before_dispatch_ack_without_losing_finality() { let lease = claim_next_dispatch("desktop-1", 30_000) .expect("claim") .expect("dispatch"); + prepare_dispatch_session( + &lease.dispatch_id, + &lease.lease_token, + "claim-fast", + "session-fast", + ) + .expect("prepare runner before send"); let terminal = record_run_terminal( &queued.id, @@ -626,11 +756,15 @@ fn turn_can_finish_before_dispatch_ack_without_losing_finality() { assert_eq!(terminal.status, WorkItemRunStatus::Succeeded); assert_eq!(terminal.session_id.as_deref(), Some("session-fast")); - let acknowledged = + let stale = acknowledge_dispatch_started(&lease.dispatch_id, &lease.lease_token, "session-fast") - .expect("terminal ack is idempotent"); - assert_eq!(acknowledged.status, WorkItemRunStatus::Succeeded); - assert_eq!(acknowledged.usage.total_tokens, 99); + .expect_err("terminal already retired the outbox lease"); + assert!(stale.starts_with(error::STALE_LEASE), "{stale}"); + assert!(claim_next_dispatch("desktop-after-terminal", 30_000) + .expect("claim query") + .is_none()); + let persisted = read(&queued.id).expect("persisted terminal"); + assert_eq!(persisted.usage.total_tokens, 99); } #[test] @@ -688,6 +822,43 @@ fn typed_retry_creates_a_new_run_episode_and_resumes_session() { ); } +#[test] +fn remote_root_retry_never_promotes_the_hidden_runner() { + let _sandbox = test_env::sandbox(); + seed(); + let mut remote_request = request("manual:remote-root-retry"); + remote_request.target_snapshot.target = WorkItemRunTarget::ResumeSession { + session_id: "shared-root".to_string(), + }; + let run = enqueue(remote_request).expect("enqueue remote root"); + let lease = claim_next_dispatch("desktop-remote", 30_000) + .expect("claim") + .expect("lease"); + prepare_dispatch_session( + &lease.dispatch_id, + &lease.lease_token, + "claim-hidden", + "hidden-runner", + ) + .expect("prepare hidden runner"); + let failed = record_run_terminal( + &run.id, + Some("hidden-runner"), + WorkItemRunTerminalOutcome::Failed, + WorkItemRunUsage::default(), + Some("request timed out after process restart"), + ) + .expect("terminal"); + + let retry = retry(&failed.id, "startup-remote-root:test").expect("retry"); + assert_eq!( + retry.target_snapshot.target, + WorkItemRunTarget::ResumeSession { + session_id: "shared-root".to_string(), + } + ); +} + #[test] fn retry_ancestry_preserves_routine_origin() { let _sandbox = test_env::sandbox(); diff --git a/src-tauri/src/commands/handler_list.inc b/src-tauri/src/commands/handler_list.inc index d20b22472b..94ac69b0f1 100644 --- a/src-tauri/src/commands/handler_list.inc +++ b/src-tauri/src/commands/handler_list.inc @@ -1295,5 +1295,11 @@ agent_core::state::commands::session::wingman_stop, agent_core::state::commands::session::wingman_close_windows, agent_core::state::commands::session::wingman_show_desktop_control_test, agent_core::state::commands::session::wingman_list_monitors, +// Durable remote-root Work Item conversation turn hand-off. +agent_core::state::commands::work_runs::work_run_conversation_turn_accept, +agent_core::state::commands::work_runs::work_run_conversation_turn_release, +agent_core::state::commands::work_runs::work_run_conversation_turn_nack, +agent_core::state::commands::work_runs::work_run_conversation_turn_prepare_runner, +agent_core::state::commands::work_runs::work_run_conversation_turn_ack_runner, // Model capability resolution (context window lookup from FAMILY_RULES) agent_core::core::providers::model_capabilities::resolve_model_context_k, diff --git a/src/api/tauri/rpc/procedures/index.ts b/src/api/tauri/rpc/procedures/index.ts index 7cfe6e4a21..098ac379f9 100644 --- a/src/api/tauri/rpc/procedures/index.ts +++ b/src/api/tauri/rpc/procedures/index.ts @@ -18,4 +18,5 @@ export { terminal } from "./terminal"; export { tools } from "./tools"; export { validation } from "./validation"; export { workspaceMemory } from "./workspaceMemory"; +export { workRuns } from "./workRuns"; export { cli } from "./cli"; diff --git a/src/api/tauri/rpc/procedures/workRuns.ts b/src/api/tauri/rpc/procedures/workRuns.ts new file mode 100644 index 0000000000..cf44d793b9 --- /dev/null +++ b/src/api/tauri/rpc/procedures/workRuns.ts @@ -0,0 +1,29 @@ +import { z } from "zod/v4"; + +import { defineProcedure } from "../invoke"; +import * as schemas from "../schemas"; + +export const workRuns = { + conversationTurnAccept: defineProcedure("work_run_conversation_turn_accept") + .input(schemas.workRuns.ConversationTurnAcceptInput) + .output(z.string().nullable()) + .build(), + conversationTurnRelease: defineProcedure("work_run_conversation_turn_release") + .input(schemas.workRuns.ConversationTurnReleaseInput) + .output(z.boolean()) + .build(), + conversationTurnNack: defineProcedure("work_run_conversation_turn_nack") + .input(schemas.workRuns.ConversationTurnNackInput) + .output(z.boolean()) + .build(), + conversationTurnPrepareRunner: defineProcedure( + "work_run_conversation_turn_prepare_runner" + ) + .input(schemas.workRuns.ConversationTurnPrepareRunnerInput) + .build(), + conversationTurnAckRunner: defineProcedure( + "work_run_conversation_turn_ack_runner" + ) + .input(schemas.workRuns.ConversationTurnAckRunnerInput) + .build(), +} as const; diff --git a/src/api/tauri/rpc/router.ts b/src/api/tauri/rpc/router.ts index 374e3b5b84..841bac2d5e 100644 --- a/src/api/tauri/rpc/router.ts +++ b/src/api/tauri/rpc/router.ts @@ -47,6 +47,7 @@ export const procedures = { mcp: p.mcp, flow: p.flow, humanSession: p.humanSession, + workRuns: p.workRuns, cli: p.cli, } as const; diff --git a/src/api/tauri/rpc/schemas/index.ts b/src/api/tauri/rpc/schemas/index.ts index 00846a46b3..9ff29173da 100644 --- a/src/api/tauri/rpc/schemas/index.ts +++ b/src/api/tauri/rpc/schemas/index.ts @@ -27,3 +27,4 @@ export * as humanSession from "./humanSession"; export * as sessionCore from "./sessionCore"; export * as cli from "./cli"; export * as turnIntent from "./turnIntent"; +export * as workRuns from "./workRuns"; diff --git a/src/api/tauri/rpc/schemas/workRuns.ts b/src/api/tauri/rpc/schemas/workRuns.ts new file mode 100644 index 0000000000..8795892ae1 --- /dev/null +++ b/src/api/tauri/rpc/schemas/workRuns.ts @@ -0,0 +1,33 @@ +import { z } from "zod/v4"; + +export const ConversationTurnAcceptInput = z.object({ + runId: z.string(), + claimToken: z.string(), + accepted: z.boolean(), + reason: z.string().optional(), +}); + +export const ConversationTurnReleaseInput = z.object({ + runId: z.string(), + claimToken: z.string(), +}); + +export const ConversationTurnNackInput = z.object({ + runId: z.string(), + claimToken: z.string(), + reason: z.string(), +}); + +export const ConversationTurnPrepareRunnerInput = z.object({ + runId: z.string(), + claimToken: z.string(), + rootSessionId: z.string(), + runnerSessionId: z.string(), +}); + +export const ConversationTurnAckRunnerInput = z.object({ + runId: z.string(), + claimToken: z.string(), + rootSessionId: z.string(), + runnerSessionId: z.string(), +}); From 421587b7da764a1604671cb06dc7910ad3477224 Mon Sep 17 00:00:00 2001 From: Neonforge <48338160+Neonforge98@users.noreply.github.com> Date: Tue, 25 Aug 2026 09:30:58 +0800 Subject: [PATCH 2/2] fix(work-runs): cover cold capability claims --- .../src/core/coordination/conversation_turn_bridge.rs | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src-tauri/crates/agent-core/src/core/coordination/conversation_turn_bridge.rs b/src-tauri/crates/agent-core/src/core/coordination/conversation_turn_bridge.rs index 3ed2587c8f..f8e076b474 100644 --- a/src-tauri/crates/agent-core/src/core/coordination/conversation_turn_bridge.rs +++ b/src-tauri/crates/agent-core/src/core/coordination/conversation_turn_bridge.rs @@ -26,7 +26,10 @@ use tokio::sync::oneshot; use tracing::warn; pub const CONVERSATION_TURN_REQUESTED_EVENT: &str = "orgii-work-run-conversation-turn"; -pub const ACCEPT_TIMEOUT: Duration = Duration::from_secs(10); +// The frontend's cold capability probe is bounded at 15 seconds. Leave room +// for token refresh and local org-alias lookup so a capable cold window can +// still claim before this offer is retried. +pub const ACCEPT_TIMEOUT: Duration = Duration::from_secs(20); const CLAIM_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10); const CLAIM_LEASE_EXTENSION_MS: i64 = 45_000; // An offer must be accepted inside ACCEPT_TIMEOUT, before the frontend can @@ -432,6 +435,12 @@ mod tests { ); } + #[test] + fn acceptance_window_covers_the_cold_capability_probe() { + const FRONTEND_CAPABILITY_TIMEOUT: Duration = Duration::from_secs(15); + assert!(ACCEPT_TIMEOUT > FRONTEND_CAPABILITY_TIMEOUT); + } + #[tokio::test] async fn failed_ack_keeps_prepared_claim_and_successful_ack_consumes_it() { let receiver = register(ConversationTurnClaim {