Skip to content

fix(runtime): make turn acceptance retry-safe #940

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Neonforge98 wants to merge 1 commit into from
Closed

fix(runtime): make turn acceptance retry-safe #940

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
163 changes: 124 additions & 39 deletions src-tauri/crates/agent-core/src/core/session/scheduler.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@

use futures::FutureExt;
use std::any::Any;
use std::collections::HashSet;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::sync::{mpsc, Mutex as TokioMutex};
Expand Down Expand Up @@ -163,6 +163,40 @@ pub struct EnqueueResult {
pub duplicate: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ClientMessageClaim {
Claimed,
SameIntentDuplicate,
DifferentIntentDuplicate,
}

fn claim_client_message(
owners: &mut HashMap<String, String>,
client_message_id: &str,
turn_intent_id: &str,
) -> ClientMessageClaim {
match owners.get(client_message_id) {
Some(owner_turn_intent_id) if owner_turn_intent_id == turn_intent_id => {
ClientMessageClaim::SameIntentDuplicate
}
Some(_) => ClientMessageClaim::DifferentIntentDuplicate,
None => {
owners.insert(client_message_id.to_string(), turn_intent_id.to_string());
ClientMessageClaim::Claimed
}
}
}

fn release_client_message(
owners: &mut HashMap<String, String>,
client_message_id: &str,
turn_intent_id: &str,
) {
if owners.get(client_message_id).map(String::as_str) == Some(turn_intent_id) {
owners.remove(client_message_id);
}
}

// ============================================
// DialogScheduler
// ============================================
Expand Down Expand Up @@ -198,7 +232,7 @@ pub struct DialogScheduler {
processing: Arc<std::sync::atomic::AtomicBool>,
/// Whether the job the worker is currently executing is a [`ScheduledKind::Turn`].
processing_turn: Arc<std::sync::atomic::AtomicBool>,
client_message_ids: Arc<TokioMutex<HashSet<String>>>,
client_message_owners: Arc<TokioMutex<HashMap<String, String>>>,
}

impl DialogScheduler {
Expand All @@ -217,7 +251,7 @@ impl DialogScheduler {
generation: Arc::new(AtomicU64::new(0)),
processing: Arc::new(std::sync::atomic::AtomicBool::new(false)),
processing_turn: Arc::new(std::sync::atomic::AtomicBool::new(false)),
client_message_ids: Arc::new(TokioMutex::new(HashSet::new())),
client_message_owners: Arc::new(TokioMutex::new(HashMap::new())),
}
}
/// Ensure the worker is spawned and return a reference to the sender.
Expand All @@ -237,7 +271,7 @@ impl DialogScheduler {
generation: Arc::clone(&self.generation),
processing: Arc::clone(&self.processing),
processing_turn: Arc::clone(&self.processing_turn),
client_message_ids: Arc::clone(&self.client_message_ids),
client_message_owners: Arc::clone(&self.client_message_owners),
};
tokio::spawn(worker.run());

Expand All @@ -262,22 +296,37 @@ impl DialogScheduler {

let message_id = msg.message_id.clone();
if let Some(client_message_id) = msg.client_message_id.as_ref() {
let mut ids = self.client_message_ids.lock().await;
if !ids.insert(client_message_id.clone()) {
// This request minted its own durable intent before enqueue,
// but an equivalent client message is already queued/running.
// The scheduler is the single authority that knows the request
// was coalesced, so it also closes that new intent here.
crate::foundation::session_bridge::update_turn_intent_status(
&self.session_id,
&msg.turn_intent_id,
crate::foundation::session_bridge::TurnIntentBridgeStatus::Coalesced,
);
return Ok(EnqueueResult {
message_id,
queue_position: 0,
duplicate: true,
});
let claim = {
let mut owners = self.client_message_owners.lock().await;
claim_client_message(&mut owners, client_message_id, &msg.turn_intent_id)
};
match claim {
ClientMessageClaim::Claimed => {}
ClientMessageClaim::SameIntentDuplicate => {
// A response-loss retry owns the same durable intent as the
// queued/running message. Keep that row in its original
// state so the retry can reconcile its exact receipt.
return Ok(EnqueueResult {
message_id,
queue_position: 0,
duplicate: true,
});
}
ClientMessageClaim::DifferentIntentDuplicate => {
// A distinct logical intent reused an in-flight idempotency
// key. It will never create a round, so close only the new
// intent as coalesced and preserve the original owner.
crate::foundation::session_bridge::update_turn_intent_status(
&self.session_id,
&msg.turn_intent_id,
crate::foundation::session_bridge::TurnIntentBridgeStatus::Coalesced,
);
return Ok(EnqueueResult {
message_id,
queue_position: 0,
duplicate: true,
});
}
}
}

Expand All @@ -297,10 +346,12 @@ impl DialogScheduler {
Err(mpsc::error::TrySendError::Full(rejected)) => {
self.pending.fetch_sub(1, Ordering::Relaxed);
if let Some(client_message_id) = rejected.client_message_id.as_ref() {
self.client_message_ids
.lock()
.await
.remove(client_message_id);
let mut owners = self.client_message_owners.lock().await;
release_client_message(
&mut owners,
client_message_id,
&rejected.turn_intent_id,
);
}
crate::foundation::session_bridge::update_turn_intent_status(
&self.session_id,
Expand All @@ -315,10 +366,12 @@ impl DialogScheduler {
Err(mpsc::error::TrySendError::Closed(rejected)) => {
self.pending.fetch_sub(1, Ordering::Relaxed);
if let Some(client_message_id) = rejected.client_message_id.as_ref() {
self.client_message_ids
.lock()
.await
.remove(client_message_id);
let mut owners = self.client_message_owners.lock().await;
release_client_message(
&mut owners,
client_message_id,
&rejected.turn_intent_id,
);
}
crate::foundation::session_bridge::update_turn_intent_status(
&self.session_id,
Expand All @@ -339,8 +392,8 @@ impl DialogScheduler {
pub fn invalidate_pending(&self) {
self.generation.fetch_add(1, Ordering::AcqRel);
self.pending.store(0, Ordering::Release);
if let Ok(mut ids) = self.client_message_ids.try_lock() {
ids.clear();
if let Ok(mut owners) = self.client_message_owners.try_lock() {
owners.clear();
}
// Lifecycle: every still-queued / optimistic intent for this
// session walks to `stale`. The worker drops queued-but-stale
Expand Down Expand Up @@ -393,7 +446,7 @@ struct WorkerTask {
generation: Arc<AtomicU64>,
processing: Arc<std::sync::atomic::AtomicBool>,
processing_turn: Arc<std::sync::atomic::AtomicBool>,
client_message_ids: Arc<TokioMutex<HashSet<String>>>,
client_message_owners: Arc<TokioMutex<HashMap<String, String>>>,
}

impl WorkerTask {
Expand All @@ -418,10 +471,8 @@ impl WorkerTask {
crate::foundation::session_bridge::TurnIntentBridgeStatus::Stale,
);
if let Some(client_message_id) = msg.client_message_id.as_ref() {
self.client_message_ids
.lock()
.await
.remove(client_message_id);
let mut owners = self.client_message_owners.lock().await;
release_client_message(&mut owners, client_message_id, &msg.turn_intent_id);
}
self.broadcast_idle_status();
continue;
Expand Down Expand Up @@ -553,10 +604,8 @@ impl WorkerTask {
}

if let Some(client_message_id) = client_message_id.as_ref() {
self.client_message_ids
.lock()
.await
.remove(client_message_id);
let mut owners = self.client_message_owners.lock().await;
release_client_message(&mut owners, client_message_id, &turn_intent_id);
}
self.processing_turn.store(false, Ordering::Relaxed);
self.processing.store(false, Ordering::Relaxed);
Expand All @@ -583,6 +632,42 @@ mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};

#[test]
fn client_message_claim_preserves_the_exact_intent_owner() {
let mut owners = HashMap::new();
assert_eq!(
claim_client_message(&mut owners, "client-1", "intent-owner"),
ClientMessageClaim::Claimed
);
assert_eq!(
claim_client_message(&mut owners, "client-1", "intent-owner"),
ClientMessageClaim::SameIntentDuplicate
);
assert_eq!(
claim_client_message(&mut owners, "client-1", "intent-other"),
ClientMessageClaim::DifferentIntentDuplicate
);
assert_eq!(
owners.get("client-1").map(String::as_str),
Some("intent-owner")
);

// A stale worker must not release a newer owner that claimed the same
// idempotency key after invalidation.
owners.clear();
assert_eq!(
claim_client_message(&mut owners, "client-1", "intent-new"),
ClientMessageClaim::Claimed
);
release_client_message(&mut owners, "client-1", "intent-owner");
assert_eq!(
owners.get("client-1").map(String::as_str),
Some("intent-new")
);
release_client_message(&mut owners, "client-1", "intent-new");
assert!(!owners.contains_key("client-1"));
}

#[tokio::test]
async fn invalidated_pending_message_is_skipped() {
let scheduler = DialogScheduler::new("session-a", 8);
Expand Down
50 changes: 50 additions & 0 deletions src-tauri/crates/agent-core/src/foundation/session_bridge.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -499,6 +499,22 @@ pub type UpsertTurnIntentFn = fn(
status: TurnIntentBridgeStatus,
);

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TurnIntentBridgeClaim {
pub duplicate: bool,
pub status: TurnIntentBridgeStatus,
pub client_message_id: Option<String>,
}

pub type ClaimTurnIntentFn = fn(
session_id: &str,
turn_intent_id: &str,
client_message_id: Option<&str>,
org_run_id: Option<&str>,
source: TurnIntentBridgeSource,
status: TurnIntentBridgeStatus,
) -> Result<TurnIntentBridgeClaim, String>;

pub type UpdateTurnIntentStatusFn =
fn(session_id: &str, turn_intent_id: &str, new_status: TurnIntentBridgeStatus);

Expand All @@ -508,6 +524,7 @@ pub type GetTurnIntentStatusFn =
pub type MarkPendingTurnIntentsStaleFn = fn(session_id: &str);

static UPSERT_TURN_INTENT: OnceLock<UpsertTurnIntentFn> = OnceLock::new();
static CLAIM_TURN_INTENT: OnceLock<ClaimTurnIntentFn> = OnceLock::new();
static UPDATE_TURN_INTENT_STATUS: OnceLock<UpdateTurnIntentStatusFn> = OnceLock::new();
static GET_TURN_INTENT_STATUS: OnceLock<GetTurnIntentStatusFn> = OnceLock::new();
static MARK_PENDING_TURN_INTENTS_STALE: OnceLock<MarkPendingTurnIntentsStaleFn> = OnceLock::new();
Expand All @@ -516,6 +533,10 @@ pub fn register_upsert_turn_intent(implementation: UpsertTurnIntentFn) {
let _ = UPSERT_TURN_INTENT.set(implementation);
}

pub fn register_claim_turn_intent(implementation: ClaimTurnIntentFn) {
let _ = CLAIM_TURN_INTENT.set(implementation);
}

pub fn register_update_turn_intent_status(implementation: UpdateTurnIntentStatusFn) {
let _ = UPDATE_TURN_INTENT_STATUS.set(implementation);
}
Expand Down Expand Up @@ -553,6 +574,35 @@ pub fn upsert_turn_intent(
}
}

/// Atomically reserve an exact logical intent before any turn side effects.
///
/// Unlike the best-effort projection helper above, this is an acceptance
/// boundary: missing registration or persistence failure is returned to the
/// caller so it cannot continue with an unowned execution.
pub fn claim_turn_intent(
session_id: &str,
turn_intent_id: &str,
client_message_id: Option<&str>,
org_run_id: Option<&str>,
source: TurnIntentBridgeSource,
status: TurnIntentBridgeStatus,
) -> Result<TurnIntentBridgeClaim, String> {
if session_id.is_empty() || turn_intent_id.is_empty() {
return Err("turn intent claim requires session_id and turn_intent_id".to_string());
}
let implementation = CLAIM_TURN_INTENT
.get()
.ok_or_else(|| "turn intent claim persistence is not registered".to_string())?;
implementation(
session_id,
turn_intent_id,
client_message_id,
org_run_id,
source,
status,
)
}

/// Patch the status of an existing lifecycle row. Illegal transitions are
/// silently rejected by the implementation — callers do not need to handle
/// the error case.
Expand Down
17 changes: 10 additions & 7 deletions src-tauri/crates/agent-core/src/state/commands/session/message/project_bootstrap.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,8 @@

use crate::foundation::session_bridge::TurnIntentBridgeSource;
use project_management::projects::types::{
EnqueueWorkItemRunRequest, WorkItemRun, WorkItemRunTarget, WorkItemRunTargetSnapshot,
WorkItemRunTrigger, PERSONAL_ORG_ID,
EnqueueWorkItemRunRequest, WorkItemRunTarget, WorkItemRunTargetSnapshot, WorkItemRunTrigger,
PERSONAL_ORG_ID,
};

/// Bootstrap called from the message-accept path. Project mode is an explicit
Expand Down Expand Up @@ -55,7 +55,7 @@ pub(super) async fn enqueue_project_turn_if_needed(
turn_intent_id: &str,
client_message_id: Option<&str>,
source: TurnIntentBridgeSource,
) -> Result<Option<WorkItemRun>, String> {
) -> Result<Option<project_management::work_run_service::EnqueueWorkItemRunReceipt>, String> {
if content.trim().is_empty() || turn_intent_id.starts_with("wir_") {
return Ok(None);
}
Expand Down Expand Up @@ -129,15 +129,18 @@ pub(super) async fn enqueue_project_turn_if_needed(
"content": content,
"displayText": display_text,
"clientMessageId": client_message_id,
"originTurnIntentId": turn_intent_id,
}),
idempotency_key: format!("project-session-turn:{session_id}:{turn_intent_id}"),
max_attempts: 3,
parent_run_id: None,
};
tokio::task::spawn_blocking(move || project_management::work_run_service::enqueue(request))
.await
.map_err(|err| format!("Project WorkItemRun enqueue worker failed: {err}"))?
.map(Some)
tokio::task::spawn_blocking(move || {
project_management::work_run_service::enqueue_with_receipt(request)
})
.await
.map_err(|err| format!("Project WorkItemRun enqueue worker failed: {err}"))?
.map(Some)
}

/// Blocking core, also driven directly by the `Track this` command —
Expand Down
Loading
Loading