From b06b6b32d42e440a514ce1e0e0b0ba863d6edc91 Mon Sep 17 00:00:00 2001 From: clawdbot Date: Tue, 13 Oct 2026 05:27:58 +0000 Subject: [PATCH] fix(sdk): preserve lifecycle outcomes and output --- rust/crates/runtime/src/conversation.rs | 196 ++++++++++++++---- rust/crates/sdk/src/session/mod.rs | 178 +++++++++++++--- rust/crates/sdk/src/session/reconciliation.rs | 172 ++++++++++++++- rust/crates/sdk/src/session/terminal_scope.rs | 10 +- rust/crates/sdk/tests/session_lifecycle.rs | 103 +++++++-- 5 files changed, 570 insertions(+), 89 deletions(-) diff --git a/rust/crates/runtime/src/conversation.rs b/rust/crates/runtime/src/conversation.rs index b5fc2a66..0972a698 100644 --- a/rust/crates/runtime/src/conversation.rs +++ b/rust/crates/runtime/src/conversation.rs @@ -67,6 +67,7 @@ pub enum ConversationEvent { AgentEnd { session_id: String, interrupted: bool, + output: Option, }, } @@ -580,6 +581,7 @@ where session_id: session_id.clone(), }); let mut interrupted = false; + let mut final_output = None; let outcome = (|| { // ROADMAP #38: Session-health canary - probe if context was compacted if self.session.compaction.is_some() { @@ -645,6 +647,7 @@ where return Err(error); } }; + final_output = assistant_message_text(&assistant_message); if let Some(usage) = usage { self.usage_tracker.record(usage); } @@ -764,57 +767,71 @@ where ) }; - let result_message = match permission_outcome { - PermissionOutcome::Allow => { - self.record_tool_started(iterations, &tool_name); - let (mut output, mut is_error) = - match self.tool_executor.execute(&tool_name, &effective_input) { + let result_message = if self.turn_is_interrupted() { + self.record_turn_interrupted(iterations, "before_tool_dispatch"); + interrupted = true; + ConversationMessage::tool_result( + tool_use_id, + tool_name, + "Interrupted by user before this tool could run.", + true, + ) + } else { + match permission_outcome { + PermissionOutcome::Allow => { + self.record_tool_started(iterations, &tool_name); + let (mut output, mut is_error) = match self + .tool_executor + .execute(&tool_name, &effective_input) + { Ok(output) => (output, false), Err(error) => (error.to_string(), true), }; - output = merge_hook_feedback(pre_hook_result.messages(), output, false); + output = + merge_hook_feedback(pre_hook_result.messages(), output, false); - let post_hook_result = if is_error { - self.run_post_tool_use_failure_hook( - &tool_name, - &effective_input, - &output, - ) - } else { - self.run_post_tool_use_hook( - &tool_name, - &effective_input, - &output, - false, - ) - }; - if post_hook_result.is_denied() - || post_hook_result.is_failed() - || post_hook_result.is_cancelled() - { - is_error = true; - } - output = merge_hook_feedback( - post_hook_result.messages(), - output, - post_hook_result.is_denied() + let post_hook_result = if is_error { + self.run_post_tool_use_failure_hook( + &tool_name, + &effective_input, + &output, + ) + } else { + self.run_post_tool_use_hook( + &tool_name, + &effective_input, + &output, + false, + ) + }; + if post_hook_result.is_denied() || post_hook_result.is_failed() - || post_hook_result.is_cancelled(), - ); + || post_hook_result.is_cancelled() + { + is_error = true; + } + output = merge_hook_feedback( + post_hook_result.messages(), + output, + post_hook_result.is_denied() + || post_hook_result.is_failed() + || post_hook_result.is_cancelled(), + ); - ConversationMessage::tool_result( + ConversationMessage::tool_result( + tool_use_id, + tool_name, + output, + is_error, + ) + } + PermissionOutcome::Deny { reason } => ConversationMessage::tool_result( tool_use_id, tool_name, - output, - is_error, - ) + merge_hook_feedback(pre_hook_result.messages(), reason, true), + true, + ), } - PermissionOutcome::Deny { reason } => ConversationMessage::tool_result( - tool_use_id, - tool_name, - merge_hook_feedback(pre_hook_result.messages(), reason, true), - true, - ), }; self.session .push_message(result_message.clone()) @@ -859,6 +876,7 @@ where self.emit_conversation_event(ConversationEvent::AgentEnd { session_id, interrupted, + output: final_output, }); outcome } @@ -1122,6 +1140,18 @@ fn build_assistant_message( )) } +fn assistant_message_text(message: &ConversationMessage) -> Option { + let output = message + .blocks + .iter() + .filter_map(|block| match block { + ContentBlock::Text { text } => Some(text.as_str()), + _ => None, + }) + .collect::(); + (!output.is_empty()).then_some(output) +} + fn flush_text_block(text: &mut String, blocks: &mut Vec) { if !text.is_empty() { blocks.push(ContentBlock::Text { @@ -1356,6 +1386,7 @@ mod tests { ConversationEvent::AgentEnd { session_id, interrupted: false, + output: Some("done".to_string()), }, ] ); @@ -1495,6 +1526,86 @@ mod tests { )); } + #[test] + fn abort_during_permission_prompt_prevents_tool_dispatch() { + struct ToolApi; + + impl ApiClient for ToolApi { + fn stream( + &mut self, + _request: ApiRequest, + ) -> Result, RuntimeError> { + Ok(vec![ + AssistantEvent::ToolUse { + id: "pending-tool".to_string(), + name: "dangerous".to_string(), + input: "{}".to_string(), + }, + AssistantEvent::MessageStop, + ]) + } + } + + struct BlockingAllowPrompter { + started: std::sync::mpsc::SyncSender<()>, + release: std::sync::mpsc::Receiver<()>, + } + + impl PermissionPrompter for BlockingAllowPrompter { + fn decide(&mut self, _request: &PermissionRequest) -> PermissionPromptDecision { + self.started.send(()).expect("notify permission prompt"); + self.release.recv().expect("wait for approval"); + PermissionPromptDecision::Allow + } + } + + let control = TurnControlSignal::new(); + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let tool_calls = Arc::clone(&calls); + let executor = StaticToolExecutor::new().register("dangerous", move |_| { + tool_calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok("executed".to_string()) + }); + let (started_tx, started_rx) = std::sync::mpsc::sync_channel(1); + let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); + let abort_control = control.clone(); + let abort_thread = std::thread::spawn(move || { + started_rx + .recv_timeout(std::time::Duration::from_secs(2)) + .expect("permission prompt should start"); + abort_control.abort(); + release_tx.send(()).expect("release permission prompt"); + }); + let mut prompter = BlockingAllowPrompter { + started: started_tx, + release: release_rx, + }; + let mut runtime = ConversationRuntime::new( + Session::new(), + ToolApi, + executor, + PermissionPolicy::new(PermissionMode::Prompt), + vec!["system".to_string()], + ) + .with_turn_control_signal(control); + let summary = runtime + .run_turn("start", Some(&mut prompter)) + .expect("aborted approval should be a normal interruption"); + abort_thread.join().expect("abort control thread"); + let messages = runtime.session().messages.clone(); + + assert!(summary.interrupted); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 0); + assert!(matches!( + &messages[2].blocks[0], + ContentBlock::ToolResult { + is_error: true, + output, + .. + } if output.contains("Interrupted by user") + )); + } + #[test] fn runs_user_to_tool_to_result_loop_end_to_end_and_tracks_usage() { let api_client = ScriptedApiClient { call_count: 0 }; @@ -2646,6 +2757,7 @@ mod tests { ConversationEvent::AgentEnd { session_id, interrupted: true, + output: Some("Running the tool.".to_string()), }, ] ); diff --git a/rust/crates/sdk/src/session/mod.rs b/rust/crates/sdk/src/session/mod.rs index e8923ed5..c0dc5cbf 100644 --- a/rust/crates/sdk/src/session/mod.rs +++ b/rust/crates/sdk/src/session/mod.rs @@ -193,6 +193,7 @@ struct SessionRouterInner { runner: Arc, reconciliation: InvocationReconciliationStore, terminal_scopes: TerminalScopeStore, + reconciliation_lock: Mutex<()>, active: Mutex>, } @@ -319,10 +320,14 @@ impl SessionEventPublisher { ConversationEvent::AgentEnd { session_id: _, interrupted, - } => ( - SessionEventType::AgentEnd, - json!({ "interrupted": interrupted }), - ), + output, + } => { + let mut payload = json!({ "interrupted": interrupted }); + if let Some(output) = output { + payload["output"] = Value::String(output); + } + (SessionEventType::AgentEnd, payload) + } }; self.publish(event_type, payload, false); } @@ -339,12 +344,14 @@ impl SessionRouter { let terminal_scopes = TerminalScopeStore::new(sdk_root.join("terminal-scopes")); reconciliation.reconcile_unfinished()?; recover_terminal_scopes(&reconciliation, &terminal_scopes)?; + recover_pending_terminal_outcomes(&reconciliation, &terminal_scopes)?; Ok(Self { inner: Arc::new(SessionRouterInner { store, runner, reconciliation, terminal_scopes, + reconciliation_lock: Mutex::new(()), active: Mutex::new(HashMap::new()), }), }) @@ -371,6 +378,11 @@ impl SessionRouter { .save_to_path(&handle.path) .map_err(|error| SdkError::internal(format!("could not persist session: {error}")))?; if let Some(key) = correlation_key { + let _reconciliation_guard = self + .inner + .reconciliation_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); self.inner.terminal_scopes.set(key, &session.session_id)?; } Ok(session) @@ -397,23 +409,32 @@ impl SessionRouter { /// Resolve the most recently recorded session for a correlation key. pub fn session_for_correlation_key(&self, correlation_key: &str) -> SdkResult> { validate_correlation_key(correlation_key)?; - { + let _reconciliation_guard = self + .inner + .reconciliation_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + recover_pending_terminal_outcomes(&self.inner.reconciliation, &self.inner.terminal_scopes)?; + let active_matches = { let active = self .inner .active .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - let mut matches = active.iter().filter(|(_, invocation)| { - invocation.correlation_key.as_deref() == Some(correlation_key) - }); - if let Some((session_id, _)) = matches.next() { - if matches.next().is_some() { - return Err(SdkError::conflict(format!( - "correlation key '{correlation_key}' has multiple active sessions" - ))); - } - return Ok(Some(session_id.clone())); - } + active + .iter() + .filter(|(_, invocation)| { + invocation.correlation_key.as_deref() == Some(correlation_key) + }) + .map(|(session_id, invocation)| { + (session_id.clone(), invocation.invocation_id.clone()) + }) + .collect::>() + }; + if active_matches.len() > 1 { + return Err(SdkError::conflict(format!( + "correlation key '{correlation_key}' has multiple active sessions" + ))); } let has_unfinished_invocation = self @@ -421,12 +442,20 @@ impl SessionRouter { .reconciliation .list_by_correlation_key(correlation_key)? .iter() - .any(|record| record.status != InvocationStatus::Terminal); + .any(|record| { + record.status != InvocationStatus::Terminal + && active_matches + .first() + .is_none_or(|(_, active_id)| record.invocation_id != *active_id) + }); if has_unfinished_invocation { return Err(SdkError::unavailable(format!( "correlation key '{correlation_key}' has an invocation awaiting reconciliation" ))); } + if let Some((session_id, _)) = active_matches.first() { + return Ok(Some(session_id.clone())); + } let scope = self.inner.terminal_scopes.get(correlation_key)?; Ok(scope.map(|scope| scope.session_id)) } @@ -442,7 +471,6 @@ impl SessionRouter { validate_session_id(session_id)?; validate_prompt(&prompt)?; validate_optional_correlation_key(correlation_key)?; - let session = self.get_session(session_id)?; let mut active = self .inner @@ -454,6 +482,9 @@ impl SessionRouter { "session `{session_id}` already has an active invocation" ))); } + // Load after reserving the session slot. Loading before this lock lets a + // waiting prompt retain a stale snapshot and overwrite the prior turn. + let session = self.get_session(session_id)?; let invocation_id = Uuid::new_v4().to_string(); let created_at_ms = now_ms(); @@ -710,6 +741,12 @@ impl SessionRouter { /// Return a durable invocation record for recovery after a disconnected client. pub fn reconcile(&self, invocation_id: &str) -> SdkResult { validate_invocation_id(invocation_id)?; + let _reconciliation_guard = self + .inner + .reconciliation_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + recover_pending_terminal_outcomes(&self.inner.reconciliation, &self.inner.terminal_scopes)?; self.inner .reconciliation .get(invocation_id)? @@ -724,6 +761,12 @@ impl SessionRouter { correlation_key: &str, ) -> SdkResult> { validate_correlation_key(correlation_key)?; + let _reconciliation_guard = self + .inner + .reconciliation_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + recover_pending_terminal_outcomes(&self.inner.reconciliation, &self.inner.terminal_scopes)?; self.inner .reconciliation .list_by_correlation_key(correlation_key) @@ -882,6 +925,10 @@ fn finish_invocation( publisher: &SessionEventPublisher, result: SessionRunResult, ) { + let _reconciliation_guard = inner + .reconciliation_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); let _publication_guard = publisher.lock(); control.close(); @@ -916,6 +963,26 @@ fn finish_invocation( if let Some(error) = error_message { payload["error"] = Value::String(error); } + let terminal_event = publisher.build_locked(SessionEventType::TurnEnded, payload, true); + let stored_event = match serde_json::to_value(&terminal_event) { + Ok(event) => event, + Err(error) => { + let _ = inner + .reconciliation + .mark_uncertain(invocation_id, error.to_string()); + return; + } + }; + if let Err(error) = inner.reconciliation.record_pending_terminal( + invocation_id, + outcome.clone(), + stored_event.clone(), + ) { + let _ = inner + .reconciliation + .mark_uncertain(invocation_id, error.to_string()); + return; + } let terminal_scope_revision = match correlation_key { Some(key) => match inner.terminal_scopes.set(key, session_id) { Ok(scope) => Some(scope.revision), @@ -928,16 +995,6 @@ fn finish_invocation( }, None => None, }; - let terminal_event = publisher.build_locked(SessionEventType::TurnEnded, payload, true); - let stored_event = match serde_json::to_value(&terminal_event) { - Ok(event) => event, - Err(error) => { - let _ = inner - .reconciliation - .mark_uncertain(invocation_id, error.to_string()); - return; - } - }; match inner.reconciliation.record_terminal_with_scope_revision( invocation_id, outcome, @@ -991,6 +1048,71 @@ fn recover_terminal_scopes( Ok(()) } +fn recover_pending_terminal_outcomes( + reconciliation: &InvocationReconciliationStore, + terminal_scopes: &TerminalScopeStore, +) -> SdkResult<()> { + let records = reconciliation.list()?; + let mut latest_terminal_update = BTreeMap::::new(); + let mut original_scope_update = BTreeMap::>::new(); + let mut pending = Vec::new(); + for record in records { + if record.status == InvocationStatus::Terminal { + if let Some(key) = record.correlation_key.as_ref() { + latest_terminal_update + .entry(key.clone()) + .and_modify(|updated_at| *updated_at = (*updated_at).max(record.updated_at_ms)) + .or_insert(record.updated_at_ms); + } + } else if record.outcome.is_some() && record.terminal_event.is_some() { + pending.push(record); + } + } + pending.sort_by(|left, right| { + left.updated_at_ms + .cmp(&right.updated_at_ms) + .then_with(|| left.invocation_id.cmp(&right.invocation_id)) + }); + + for record in pending { + let scope_revision = match record.correlation_key.as_deref() { + Some(key) => { + let scope_updated_at = if let Some(updated_at) = original_scope_update.get(key) { + *updated_at + } else { + let updated_at = terminal_scopes.get(key)?.map(|scope| scope.updated_at_ms); + original_scope_update.insert(key.to_string(), updated_at); + updated_at + }; + let newer_terminal = latest_terminal_update + .get(key) + .is_some_and(|updated_at| *updated_at >= record.updated_at_ms); + let newer_scope = + scope_updated_at.is_some_and(|updated_at| updated_at >= record.updated_at_ms); + if newer_terminal || newer_scope { + None + } else { + Some(terminal_scopes.set(key, &record.session_id)?.revision) + } + } + None => None, + }; + reconciliation.record_terminal_with_scope_revision( + &record.invocation_id, + record + .outcome + .ok_or_else(|| SdkError::internal("pending terminal outcome is missing"))?, + Some( + record + .terminal_event + .ok_or_else(|| SdkError::internal("pending terminal event is missing"))?, + ), + scope_revision, + )?; + } + Ok(()) +} + fn finish_uncertain( inner: &SessionRouterInner, invocation_id: &str, diff --git a/rust/crates/sdk/src/session/reconciliation.rs b/rust/crates/sdk/src/session/reconciliation.rs index da1a74c6..92c7496f 100644 --- a/rust/crates/sdk/src/session/reconciliation.rs +++ b/rust/crates/sdk/src/session/reconciliation.rs @@ -4,6 +4,7 @@ use std::{ fs::{self, OpenOptions}, io::Write, path::{Path, PathBuf}, + sync::Mutex, time::{SystemTime, UNIX_EPOCH}, }; @@ -15,6 +16,8 @@ use uuid::Uuid; use crate::error::{SdkError, SdkResult}; +static RECONCILIATION_WRITE_LOCK: Mutex<()> = Mutex::new(()); + /// Lifecycle state persisted for an invocation. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] @@ -71,6 +74,9 @@ impl InvocationReconciliationStore { /// Inserts a new invocation record. Existing IDs are not overwritten. pub fn insert(&self, record: InvocationRecord) -> SdkResult<()> { + let _guard = RECONCILIATION_WRITE_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); let path = self.record_path(&record.invocation_id)?; if self.get(&record.invocation_id)?.is_some() { return Err(SdkError::invalid_input(format!( @@ -83,6 +89,9 @@ impl InvocationReconciliationStore { /// Marks a non-terminal invocation as running. pub fn mark_running(&self, invocation_id: &str) -> SdkResult<()> { + let _guard = RECONCILIATION_WRITE_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); let mut record = self.required_record(invocation_id)?; if record.status != InvocationStatus::Accepted { return Err(SdkError::conflict(format!( @@ -123,6 +132,9 @@ impl InvocationReconciliationStore { terminal_event: Option, terminal_scope_revision: Option, ) -> SdkResult<()> { + let _guard = RECONCILIATION_WRITE_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); let mut record = self.required_record(invocation_id)?; if record.status == InvocationStatus::Terminal { return Err(SdkError::invalid_input(format!( @@ -130,6 +142,20 @@ impl InvocationReconciliationStore { invocation_id ))); } + if record + .outcome + .as_ref() + .is_some_and(|pending| pending != &outcome) + || record + .terminal_event + .as_ref() + .is_some_and(|pending| Some(pending) != terminal_event.as_ref()) + { + return Err(SdkError::conflict(format!( + "invocation '{}' has a different pending terminal result", + invocation_id + ))); + } record.status = InvocationStatus::Terminal; record.outcome = Some(outcome); record.uncertainty_reason = None; @@ -139,6 +165,39 @@ impl InvocationReconciliationStore { self.write_record(&record) } + /// Persists a known turn result before its correlation scope is updated. + /// + /// The record remains non-terminal until the scope update succeeds, allowing + /// startup or a later reconciliation read to safely finish the transition. + pub fn record_pending_terminal( + &self, + invocation_id: &str, + outcome: InvocationOutcome, + terminal_event: serde_json::Value, + ) -> SdkResult<()> { + let _guard = RECONCILIATION_WRITE_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut record = self.required_record(invocation_id)?; + if record.status == InvocationStatus::Terminal { + return Err(SdkError::invalid_input(format!( + "invocation '{}' is already terminal", + invocation_id + ))); + } + if record.outcome.is_some() || record.terminal_event.is_some() { + return Err(SdkError::conflict(format!( + "invocation '{}' already has a pending terminal result", + invocation_id + ))); + } + record.outcome = Some(outcome); + record.terminal_event = Some(terminal_event); + record.terminal_scope_revision = None; + record.updated_at_ms = updated_at_ms(record.updated_at_ms); + self.write_record(&record) + } + /// Loads one invocation record, if present. pub fn get(&self, invocation_id: &str) -> SdkResult> { let path = self.record_path(invocation_id)?; @@ -202,6 +261,9 @@ impl InvocationReconciliationStore { /// Returns the records changed by this recovery pass; already uncertain and /// terminal records are left untouched and are not included in the result. pub fn reconcile_unfinished(&self) -> SdkResult> { + let _guard = RECONCILIATION_WRITE_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); let mut recovered = Vec::new(); for mut record in self.list()? { if matches!( @@ -235,6 +297,9 @@ impl InvocationReconciliationStore { invocation_id: &str, update: impl FnOnce(&mut InvocationRecord), ) -> SdkResult<()> { + let _guard = RECONCILIATION_WRITE_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); let mut record = self.required_record(invocation_id)?; if record.status == InvocationStatus::Terminal { return Err(SdkError::invalid_input(format!( @@ -337,7 +402,7 @@ fn updated_at_ms(previous: u64) -> u64 { .unwrap_or_default() .as_millis() .min(u128::from(u64::MAX)) as u64; - now.max(previous) + now.max(previous.saturating_add(1)) } fn io_error(action: &str, path: &Path, error: std::io::Error) -> SdkError { @@ -489,6 +554,111 @@ mod tests { .is_empty()); } + #[test] + fn pending_terminal_result_survives_uncertainty_and_can_be_finalized() { + let root = TestDirectory::new(); + let store = InvocationReconciliationStore::new(&root.0); + store + .insert(accepted_record("pending")) + .expect("insert accepted record"); + store.mark_running("pending").expect("mark running"); + store + .record_pending_terminal( + "pending", + InvocationOutcome::Completed { message_count: 3 }, + serde_json::json!({"eventType": "turn_ended"}), + ) + .expect("record pending result"); + store + .mark_uncertain("pending", "terminal scope store was unavailable") + .expect("mark uncertain"); + + let pending = store.get("pending").unwrap().unwrap(); + assert_eq!(pending.status, InvocationStatus::Uncertain); + assert_eq!( + pending.outcome, + Some(InvocationOutcome::Completed { message_count: 3 }) + ); + store + .record_terminal_with_scope_revision( + "pending", + pending.outcome.unwrap(), + pending.terminal_event, + Some(1), + ) + .expect("finalize after recovery"); + assert_eq!( + store.get("pending").unwrap().unwrap().status, + InvocationStatus::Terminal + ); + } + + #[test] + fn concurrent_terminal_and_uncertain_transitions_preserve_terminal_result() { + let root = TestDirectory::new(); + let store = InvocationReconciliationStore::new(&root.0); + store + .insert(accepted_record("concurrent")) + .expect("insert accepted record"); + store.mark_running("concurrent").expect("mark running"); + + let barrier = std::sync::Arc::new(std::sync::Barrier::new(3)); + let terminal_store = store.clone(); + let terminal_barrier = std::sync::Arc::clone(&barrier); + let terminal = std::thread::spawn(move || { + terminal_barrier.wait(); + terminal_store.record_terminal( + "concurrent", + InvocationOutcome::Completed { message_count: 1 }, + None, + ) + }); + let uncertain_store = store.clone(); + let uncertain_barrier = std::sync::Arc::clone(&barrier); + let uncertain = std::thread::spawn(move || { + uncertain_barrier.wait(); + uncertain_store.mark_uncertain("concurrent", "raced with terminal write") + }); + barrier.wait(); + let _ = terminal.join().expect("terminal writer thread"); + let _ = uncertain.join().expect("uncertain writer thread"); + + let record = store.get("concurrent").unwrap().unwrap(); + assert_eq!(record.status, InvocationStatus::Terminal); + assert_eq!( + record.outcome, + Some(InvocationOutcome::Completed { message_count: 1 }) + ); + } + + #[test] + fn concurrent_duplicate_inserts_have_only_one_winner() { + let root = TestDirectory::new(); + let store = InvocationReconciliationStore::new(&root.0); + let barrier = std::sync::Arc::new(std::sync::Barrier::new(3)); + let first_store = store.clone(); + let first_barrier = std::sync::Arc::clone(&barrier); + let first = std::thread::spawn(move || { + first_barrier.wait(); + first_store.insert(accepted_record("duplicate")) + }); + let second_store = store.clone(); + let second_barrier = std::sync::Arc::clone(&barrier); + let second = std::thread::spawn(move || { + second_barrier.wait(); + second_store.insert(accepted_record("duplicate")) + }); + barrier.wait(); + let first = first.join().expect("first insert thread"); + let second = second.join().expect("second insert thread"); + + assert_ne!(first.is_ok(), second.is_ok()); + assert_eq!( + store.get("duplicate").unwrap().unwrap().invocation_id, + "duplicate" + ); + } + #[test] fn malformed_ids_and_json_are_rejected() { let root = TestDirectory::new(); diff --git a/rust/crates/sdk/src/session/terminal_scope.rs b/rust/crates/sdk/src/session/terminal_scope.rs index 8b2ca87a..d15acaf8 100644 --- a/rust/crates/sdk/src/session/terminal_scope.rs +++ b/rust/crates/sdk/src/session/terminal_scope.rs @@ -59,7 +59,7 @@ impl TerminalScopeStore { correlation_key: correlation_key.clone(), session_id: session_id.into(), revision, - updated_at_ms: current_timestamp_ms().max(previous_updated_at), + updated_at_ms: current_timestamp_ms().max(previous_updated_at.saturating_add(1)), }; atomic_write_json(&self.record_path(&correlation_key), &record)?; Ok(record) @@ -93,8 +93,11 @@ impl TerminalScopeStore { correlation_key: correlation_key.to_string(), session_id: session_id.to_string(), revision, - updated_at_ms: current_timestamp_ms() - .max(previous.map_or(0, |record| record.updated_at_ms)), + updated_at_ms: current_timestamp_ms().max( + previous + .map_or(0, |record| record.updated_at_ms) + .saturating_add(1), + ), }; atomic_write_json(&self.record_path(correlation_key), &record) } @@ -217,6 +220,7 @@ mod tests { let first = store.set("request-1", "session-1").unwrap(); let second = store.set("request-1", "session-2").unwrap(); assert_eq!(second.revision, first.revision + 1); + assert!(second.updated_at_ms > first.updated_at_ms); store .restore_if_newer("request-1", "session-1", first.revision) diff --git a/rust/crates/sdk/tests/session_lifecycle.rs b/rust/crates/sdk/tests/session_lifecycle.rs index 7f368b1f..f8889be1 100644 --- a/rust/crates/sdk/tests/session_lifecycle.rs +++ b/rust/crates/sdk/tests/session_lifecycle.rs @@ -48,6 +48,7 @@ impl SessionTurnRunner for EchoRunner { callback(ConversationEvent::AgentEnd { session_id, interrupted: false, + output: None, }); SessionRunResult::Completed(session) } @@ -132,6 +133,7 @@ impl SessionTurnRunner for GateRunner { callback(ConversationEvent::AgentEnd { session_id, interrupted: false, + output: None, }); SessionRunResult::Completed(session) } @@ -161,6 +163,7 @@ impl SessionTurnRunner for AbortRunner { callback(ConversationEvent::AgentEnd { session_id, interrupted: true, + output: None, }); return SessionRunResult::Aborted(session); } @@ -173,6 +176,7 @@ impl SessionTurnRunner for AbortRunner { callback(ConversationEvent::AgentEnd { session_id, interrupted: false, + output: None, }); SessionRunResult::Completed(session) } @@ -944,7 +948,7 @@ fn terminal_scope_survives_router_restart() { } #[tokio::test] -async fn terminal_scope_write_failure_keeps_outcome_uncertain_and_blocks_stale_routing() { +async fn terminal_scope_write_failure_preserves_outcome_for_reconciliation() { let (started_tx, started_rx) = mpsc::sync_channel(1); let (release_tx, release_rx) = mpsc::sync_channel(1); let runner = Arc::new(GateRunner { @@ -970,11 +974,34 @@ async fn terminal_scope_write_failure_keeps_outcome_uncertain_and_blocks_stale_r events.last().unwrap().event_type, SessionEventType::AgentEnd ); + let reconciliation = sdk::InvocationReconciliationStore::new( + store.sessions_dir().join("sdk").join("reconciliation"), + ); + let pending = reconciliation + .get(&invocation_id) + .unwrap() + .expect("pending result should be durable"); + assert_eq!(pending.status, InvocationStatus::Uncertain); assert_eq!( - router.reconcile(&invocation_id).unwrap().status, - InvocationStatus::Uncertain + pending.outcome, + Some(InvocationOutcome::Completed { message_count: 1 }) + ); + assert_eq!( + pending.terminal_event.as_ref().unwrap()["eventType"], + "turn_ended" ); assert!(router.session_for_correlation_key("repair-key").is_err()); + + std::fs::remove_file(&scope_root).unwrap(); + std::fs::create_dir_all(&scope_root).unwrap(); + drop(router); + let restarted = SessionRouter::new(store, Arc::new(EchoRunner)).unwrap(); + let recovered = restarted.reconcile(&invocation_id).unwrap(); + assert_eq!(recovered.status, InvocationStatus::Terminal); + assert_eq!( + restarted.session_for_correlation_key("repair-key").unwrap(), + Some(current_session.session_id.clone()) + ); assert_ne!(stale_session.session_id, current_session.session_id); } @@ -1030,6 +1057,51 @@ async fn correlation_lookup_rejects_multiple_active_sessions() { collect_events(second_invocation).await; } +#[tokio::test] +async fn correlation_lookup_does_not_bypass_an_older_uncertain_invocation() { + let (started_tx, started_rx) = mpsc::sync_channel(1); + let (release_tx, release_rx) = mpsc::sync_channel(1); + let runner = Arc::new(GateRunner { + started: Mutex::new(Some(started_tx)), + release: Mutex::new(release_rx), + }); + let (_temp, store, router) = setup_with(runner); + let older_session = router.create_session(None, None).unwrap(); + let current_session = router.create_session(None, None).unwrap(); + let reconciliation = sdk::InvocationReconciliationStore::new( + store.sessions_dir().join("sdk").join("reconciliation"), + ); + reconciliation + .insert(InvocationRecord { + invocation_id: Uuid::new_v4().to_string(), + session_id: older_session.session_id, + correlation_key: Some("shared-key".to_string()), + status: InvocationStatus::Uncertain, + accepted_at_ms: 1, + updated_at_ms: 2, + outcome: None, + uncertainty_reason: Some("recovered unfinished invocation".to_string()), + terminal_event: None, + terminal_scope_revision: None, + }) + .unwrap(); + + let invocation = router + .prompt(¤t_session.session_id, "current", Some("shared-key")) + .unwrap(); + started_rx.recv_timeout(Duration::from_secs(2)).unwrap(); + assert_eq!( + router + .session_for_correlation_key("shared-key") + .unwrap_err() + .code, + ErrorCode::Unavailable + ); + + release_tx.send(()).unwrap(); + collect_events(invocation).await; +} + #[tokio::test] async fn session_router_clones_share_active_invocations() { let (started_tx, started_rx) = mpsc::sync_channel(1); @@ -1483,6 +1555,7 @@ async fn runtime_runner_streams_real_tool_lifecycle_through_the_router() { ] ); assert_eq!(events[3].payload["output"], "payload"); + assert_eq!(events[4].payload["output"], "complete"); assert_eq!( router.reconcile(&invocation_id).unwrap().outcome, Some(InvocationOutcome::Completed { message_count: 4 }) @@ -1659,19 +1732,19 @@ async fn turn_prompt_terminal_outcome_is_available_after_event_receiver_drops() let invocation = router.prompt(&session.session_id, "hello", None).unwrap(); let invocation_id = invocation.invocation_id.clone(); drop(invocation); - for _ in 0..200 { - if router - .reconcile(&invocation_id) - .is_ok_and(|record| record.status == InvocationStatus::Terminal) - { - break; + let terminal = tokio::time::timeout(Duration::from_secs(5), async { + loop { + if let Ok(record) = router.reconcile(&invocation_id) { + if record.status == InvocationStatus::Terminal { + return record; + } + } + tokio::time::sleep(Duration::from_millis(10)).await; } - thread::sleep(Duration::from_millis(2)); - } - assert_eq!( - router.reconcile(&invocation_id).unwrap().status, - InvocationStatus::Terminal - ); + }) + .await + .expect("invocation should eventually become terminal"); + assert_eq!(terminal.status, InvocationStatus::Terminal); } #[tokio::test]