diff --git a/docs/gjc-compat/gap-matrix.md b/docs/gjc-compat/gap-matrix.md index 9eea5ed8..c9e83fbc 100644 --- a/docs/gjc-compat/gap-matrix.md +++ b/docs/gjc-compat/gap-matrix.md @@ -161,52 +161,52 @@ GJC sessions have **three entry points**: ### claw-code Status -**Current** (as of PR #3330): Partial. SDK transport merged (see SDK Wire Protocol). Local session lifecycle exists in `runtime::Session`. Remaining: operation handlers, session router with durable lookup, invocation reconciliation store, terminal scope tracking. +**Current**: The SDK now exposes `sdk::session::SessionRouter` with durable `turn.create`, `turn.prompt`, `turn.steer`, `turn.abort_and_prompt`, and paginated `turn.list` handling. It bridges `runtime::Session`/`ConversationRuntime`, streams ordered lifecycle events, persists invocation outcomes for reconciliation, and stores terminal session scope by correlation key. The WebSocket transport remains a client transport; a remotely hosted broker and automatic recovery of an unknown post-crash outcome remain outside this slice. **Files**: - `rust/crates/runtime/src/session.rs` (1,961 LOC) -- `rust/crates/runtime/src/conversation.rs` (1,878 LOC) +- `rust/crates/runtime/src/conversation.rs` (conversation loop and lifecycle event bridge) - `rust/crates/runtime/src/session_control.rs` (file ops, migration) - `rust/crates/runtime/src/lib.rs` (re-exports `Session`, `SessionStore`, `ConversationRuntime`) +- `rust/crates/sdk/src/session/` (router, event stream, reconciliation, terminal scope) **Lines**: ~3,900 LOC (session + conversation) **Reality**: -- `Session` struct holds in-memory state (messages, usage, context). -- `SessionStore` reads/writes `.claw/sessions/{id}/` directory. -- `ConversationRuntime` drives the local agent loop (sync). -- No WebSocket or remote access. -- No invocation reconciliation or uncertain outcome recovery. -- No terminal scope tracking (per-session-per-correlation). +- `Session` holds persisted conversation state and `SessionStore` resolves session IDs in the workspace namespace. +- `ConversationRuntime` drives the local agent loop and emits lifecycle events to an optional callback. +- `sdk::session::SessionRouter` dispatches the five turn operations and exposes a per-invocation event stream. +- Invocation terminal outcomes and correlation-key terminal scopes are durable; startup recovery marks unfinished invocations uncertain for later reconciliation. +- A remote WebSocket server/broker and inference of an outcome that was never durably recorded are still absent. ### Gap Assessment | Component | GJC | claw-code | Gap | Complexity | |-----------|-----|-----------|-----|------------| -| Remote Session Router | 2,500 LOC | — | 100% | High | -| Invocation Reconciliation | 1,200 LOC | — | 100% | High | -| Terminal Scope Tracking | 600 LOC | — | 100% | Medium | -| Async Event Streaming | 800 LOC | Partial | 60% | Medium | -| Uncertain Outcome Recovery | 1,000 LOC | — | 100% | High | -| **Total** | **6,100 LOC** | **3,900** | **2,200 LOC** | — | +| Remote Session Router | 2,500 LOC | Local router/dispatch | Partial | High | +| Invocation Reconciliation | 1,200 LOC | Durable invocation records | Partial | High | +| Terminal Scope Tracking | 600 LOC | Implemented | — | Medium | +| Async Event Streaming | 800 LOC | Local per-invocation stream | Partial | Medium | +| Uncertain Outcome Recovery | 1,000 LOC | Unfinished records recovered as uncertain | Partial | High | +| **Total** | **6,100 LOC** | **Session lifecycle slice implemented** | **Remote broker remains** | — | ### Proposed Implementation **Owner Crate**: `sdk` (new, or extend `runtime`) -**Module Structure** (add to `sdk/session/`): +**Module Structure** (`sdk/session/`): ``` sdk/session/ -├── router.rs (session routing, durable lookup) +├── mod.rs (session routing and operation dispatch) ├── reconciliation.rs (uncertain → terminal recovery) ├── terminal_scope.rs (per-correlation state retention) -└── events.rs (event publishing, reconciliation) +└── events.rs (ordered lifecycle event types) ``` **Integration Points**: -- Bridge `runtime::Session` into `sdk::Session` wrapper -- Add `runtime::SessionStore` extensions for reconciliation data -- Extend `ConversationRuntime` with event emission +- Bridge `runtime::Session` and `ConversationRuntime` into SDK lifecycle operations +- Persist reconciliation and terminal-scope data alongside the session namespace +- Emit lifecycle callbacks from `ConversationRuntime` **Size Estimate**: 2,500 LOC (Rust, new + modifications) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index f657af9f..6f83ba0a 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2337,8 +2337,10 @@ name = "sdk" version = "0.1.4" dependencies = [ "futures-util", + "runtime", "serde", "serde_json", + "tempfile", "tokio", "tokio-tungstenite", "tracing", diff --git a/rust/crates/runtime/src/conversation.rs b/rust/crates/runtime/src/conversation.rs index 17bcf6f3..0972a698 100644 --- a/rust/crates/runtime/src/conversation.rs +++ b/rust/crates/runtime/src/conversation.rs @@ -1,8 +1,8 @@ -use std::collections::BTreeMap; +use std::collections::{BTreeMap, VecDeque}; use std::fmt::{Display, Formatter}; use std::sync::{ atomic::{AtomicBool, Ordering}, - Arc, + Arc, Mutex, }; use serde_json::{Map, Value}; @@ -47,6 +47,30 @@ pub enum AssistantEvent { MessageStop, } +/// Runtime lifecycle event emitted while processing a conversation turn. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ConversationEvent { + AgentBegin { + session_id: String, + }, + ToolUse { + tool_use_id: String, + tool_name: String, + input: String, + }, + ToolResult { + tool_use_id: String, + tool_name: String, + output: String, + is_error: bool, + }, + AgentEnd { + session_id: String, + interrupted: bool, + output: Option, + }, +} + /// Prompt-cache telemetry captured from the provider response stream. #[derive(Debug, Clone, PartialEq, Eq)] pub struct PromptCacheEvent { @@ -92,6 +116,103 @@ impl TurnInterruptSignal { } } +#[derive(Debug, Default)] +struct TurnControlState { + queued_steers: VecDeque, + closed: bool, + aborted: bool, +} + +/// Shared steering and abort control for a running conversation turn. +/// +/// Clones share one queue and lifecycle state. Steering is accepted until +/// `finish_or_take_steer` closes the signal or `abort` is called. +#[derive(Debug, Clone, Default)] +pub struct TurnControlSignal { + state: Arc>, +} + +impl TurnControlSignal { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Queue a steering prompt while the signal is accepting input. + pub fn steer(&self, prompt: impl Into) -> bool { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.closed || state.aborted { + return false; + } + state.queued_steers.push_back(prompt.into()); + true + } + + /// Remove one queued steering prompt, if present. + #[must_use] + pub fn take_steer(&self) -> Option { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.aborted { + return None; + } + state.queued_steers.pop_front() + } + + /// Drain one queued prompt, or atomically close the signal when empty. + #[must_use] + pub fn finish_or_take_steer(&self) -> Option { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.aborted { + state.queued_steers.clear(); + state.closed = true; + return None; + } + if let Some(prompt) = state.queued_steers.pop_front() { + return Some(prompt); + } + state.closed = true; + None + } + + /// Abort the active turn and reject further steering prompts. + pub fn abort(&self) { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.aborted = true; + state.closed = true; + state.queued_steers.clear(); + } + + /// Close steering without classifying the turn as aborted. + pub fn close(&self) { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.closed = true; + state.queued_steers.clear(); + } + + #[must_use] + pub fn is_aborted(&self) -> bool { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .aborted + } +} + /// Minimal streaming API contract required by [`ConversationRuntime`]. pub trait ApiClient { fn stream(&mut self, request: ApiRequest) -> Result, RuntimeError>; @@ -179,6 +300,8 @@ pub struct ConversationRuntime { auto_compaction_input_tokens_threshold: u32, hook_abort_signal: HookAbortSignal, turn_interrupt_signal: TurnInterruptSignal, + turn_control_signal: Option, + event_callback: Option>, hook_progress_reporter: Option>, session_tracer: Option, } @@ -234,6 +357,8 @@ where auto_compaction_input_tokens_threshold: auto_compaction_threshold_from_env(), hook_abort_signal: HookAbortSignal::default(), turn_interrupt_signal: TurnInterruptSignal::default(), + turn_control_signal: None, + event_callback: None, hook_progress_reporter: None, session_tracer: None, } @@ -273,6 +398,21 @@ where self } + #[must_use] + pub fn with_event_callback( + mut self, + callback: Arc, + ) -> Self { + self.event_callback = Some(callback); + self + } + + #[must_use] + pub fn with_turn_control_signal(mut self, turn_control_signal: TurnControlSignal) -> Self { + self.turn_control_signal = Some(turn_control_signal); + self + } + #[must_use] pub fn with_hook_progress_reporter( mut self, @@ -377,6 +517,58 @@ where } } + fn emit_conversation_event(&self, event: ConversationEvent) { + if let Some(callback) = &self.event_callback { + callback(event); + } + } + + fn turn_control_is_aborted(&self) -> bool { + let is_aborted = self + .turn_control_signal + .as_ref() + .is_some_and(TurnControlSignal::is_aborted); + if is_aborted { + self.turn_interrupt_signal.interrupt(); + } + is_aborted + } + + fn turn_is_interrupted(&self) -> bool { + self.turn_control_is_aborted() || self.turn_interrupt_signal.is_interrupted() + } + + fn append_queued_steer(&mut self) -> Result { + let Some(prompt) = self + .turn_control_signal + .as_ref() + .and_then(TurnControlSignal::take_steer) + else { + return Ok(false); + }; + self.session + .push_user_text(prompt) + .map_err(|error| RuntimeError::new(error.to_string()))?; + Ok(true) + } + + fn emit_tool_result_event(&self, message: &ConversationMessage) { + if let Some(ContentBlock::ToolResult { + tool_use_id, + tool_name, + output, + is_error, + }) = message.blocks.first() + { + self.emit_conversation_event(ConversationEvent::ToolResult { + tool_use_id: tool_use_id.clone(), + tool_name: tool_name.clone(), + output: output.clone(), + is_error: *is_error, + }); + } + } + #[allow(clippy::too_many_lines)] pub fn run_turn( &mut self, @@ -384,246 +576,309 @@ where mut prompter: Option<&mut dyn PermissionPrompter>, ) -> Result { let user_input = user_input.into(); - - // ROADMAP #38: Session-health canary - probe if context was compacted - if self.session.compaction.is_some() { - if let Err(error) = self.run_session_health_probe() { - return Err(RuntimeError::new(format!( - "Session health probe failed after compaction: {error}. \ + let session_id = self.session.session_id.clone(); + self.emit_conversation_event(ConversationEvent::AgentBegin { + 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() { + if let Err(error) = self.run_session_health_probe() { + return Err(RuntimeError::new(format!( + "Session health probe failed after compaction: {error}. \ The session may be in an inconsistent state. \ Consider starting a fresh session with /session new." - ))); - } - } - - self.record_turn_started(&user_input); - self.session - .push_user_text(user_input) - .map_err(|error| RuntimeError::new(error.to_string()))?; - - let mut assistant_messages = Vec::new(); - let mut tool_results = Vec::new(); - let mut prompt_cache_events = Vec::new(); - let mut iterations = 0; - let mut auto_compaction = None; - let mut interrupted = false; - - loop { - if self.turn_interrupt_signal.is_interrupted() { - self.record_turn_interrupted(iterations, "before_request"); - interrupted = true; - break; - } - iterations += 1; - if iterations > self.max_iterations { - let error = RuntimeError::new( - "conversation loop exceeded the maximum number of iterations", - ); - self.record_turn_failed(iterations, &error); - return Err(error); + ))); + } } - let request = ApiRequest { - system_prompt: self.system_prompt.clone(), - messages: self.session.messages.clone(), - }; - let events = match self.api_client.stream(request) { - Ok(events) => events, - Err(error) => { - if self.turn_interrupt_signal.is_interrupted() { - // The client aborted because the user interrupted the - // turn; any partial response is discarded and the stop - // is reported as an interruption rather than a failure. - self.record_turn_interrupted(iterations, "during_request"); - interrupted = true; - break; - } + self.record_turn_started(&user_input); + self.session + .push_user_text(user_input) + .map_err(|error| RuntimeError::new(error.to_string()))?; + + let mut assistant_messages = Vec::new(); + let mut tool_results = Vec::new(); + let mut prompt_cache_events = Vec::new(); + let mut iterations = 0; + let mut auto_compaction = None; + + loop { + if self.turn_is_interrupted() { + self.record_turn_interrupted(iterations, "before_request"); + interrupted = true; + break; + } + iterations += 1; + if iterations > self.max_iterations { + let error = RuntimeError::new( + "conversation loop exceeded the maximum number of iterations", + ); self.record_turn_failed(iterations, &error); return Err(error); } - }; - let (assistant_message, usage, turn_prompt_cache_events) = - match build_assistant_message(events) { - Ok(result) => result, + + let request = ApiRequest { + system_prompt: self.system_prompt.clone(), + messages: self.session.messages.clone(), + }; + let events = match self.api_client.stream(request) { + Ok(events) => events, Err(error) => { + if self.turn_is_interrupted() { + // The client aborted because the user interrupted the + // turn; any partial response is discarded and the stop + // is reported as an interruption rather than a failure. + self.record_turn_interrupted(iterations, "during_request"); + interrupted = true; + break; + } self.record_turn_failed(iterations, &error); return Err(error); } }; - if let Some(usage) = usage { - self.usage_tracker.record(usage); - } - prompt_cache_events.extend(turn_prompt_cache_events); - let pending_tool_uses = assistant_message - .blocks - .iter() - .filter_map(|block| match block { - ContentBlock::ToolUse { id, name, input } => { - Some((id.clone(), name.clone(), input.clone())) + let (assistant_message, usage, turn_prompt_cache_events) = + match build_assistant_message(events) { + Ok(result) => result, + Err(error) => { + self.record_turn_failed(iterations, &error); + return Err(error); + } + }; + final_output = assistant_message_text(&assistant_message); + if let Some(usage) = usage { + self.usage_tracker.record(usage); + } + prompt_cache_events.extend(turn_prompt_cache_events); + let pending_tool_uses = assistant_message + .blocks + .iter() + .filter_map(|block| match block { + ContentBlock::ToolUse { id, name, input } => { + Some((id.clone(), name.clone(), input.clone())) + } + _ => None, + }) + .collect::>(); + self.record_assistant_iteration( + iterations, + &assistant_message, + pending_tool_uses.len(), + ); + + self.session + .push_message(assistant_message.clone()) + .map_err(|error| RuntimeError::new(error.to_string()))?; + assistant_messages.push(assistant_message); + + // Run auto-compaction check before next API call, including on the terminal + // (no-tool) iteration, to prevent unbounded session growth (#3106). + if let Some(compaction) = self.maybe_auto_compact() { + auto_compaction = Some(compaction); + } + + if pending_tool_uses.is_empty() { + if self.turn_control_is_aborted() { + self.record_turn_interrupted(iterations, "after_request"); + interrupted = true; + break; } - _ => None, - }) - .collect::>(); - self.record_assistant_iteration( - iterations, - &assistant_message, - pending_tool_uses.len(), - ); - - self.session - .push_message(assistant_message.clone()) - .map_err(|error| RuntimeError::new(error.to_string()))?; - assistant_messages.push(assistant_message); - - // Run auto-compaction check before next API call, including on the terminal - // (no-tool) iteration, to prevent unbounded session growth (#3106). - if let Some(compaction) = self.maybe_auto_compact() { - auto_compaction = Some(compaction); - } - - if pending_tool_uses.is_empty() { - break; - } - - for (tool_use_id, tool_name, input) in pending_tool_uses { - if interrupted || self.turn_interrupt_signal.is_interrupted() { - // Every pending tool_use must still receive a tool_result - // so the session stays valid for the next request. - if !interrupted { - self.record_turn_interrupted(iterations, "before_tool"); + if self.append_queued_steer()? { + continue; + } + if self.turn_control_is_aborted() { + self.record_turn_interrupted(iterations, "after_request"); interrupted = true; } - let result_message = ConversationMessage::tool_result( - tool_use_id, - tool_name, - "Interrupted by user before this tool could run.", - true, + break; + } + + for (tool_use_id, tool_name, input) in pending_tool_uses { + self.emit_conversation_event(ConversationEvent::ToolUse { + tool_use_id: tool_use_id.clone(), + tool_name: tool_name.clone(), + input: input.clone(), + }); + if interrupted || self.turn_is_interrupted() { + // Every pending tool_use must still receive a tool_result + // so the session stays valid for the next request. + if !interrupted { + self.record_turn_interrupted(iterations, "before_tool"); + interrupted = true; + } + let result_message = ConversationMessage::tool_result( + tool_use_id, + tool_name, + "Interrupted by user before this tool could run.", + true, + ); + self.session + .push_message(result_message.clone()) + .map_err(|error| RuntimeError::new(error.to_string()))?; + self.emit_tool_result_event(&result_message); + tool_results.push(result_message); + continue; + } + let pre_hook_result = self.run_pre_tool_use_hook(&tool_name, &input); + let effective_input = pre_hook_result + .updated_input() + .map_or_else(|| input.clone(), ToOwned::to_owned); + let permission_context = PermissionContext::new( + pre_hook_result.permission_override(), + pre_hook_result.permission_reason().map(ToOwned::to_owned), ); + + let permission_outcome = if pre_hook_result.is_cancelled() { + PermissionOutcome::Deny { + reason: format_hook_message( + &pre_hook_result, + &format!("PreToolUse hook cancelled tool `{tool_name}`"), + ), + } + } else if pre_hook_result.is_failed() { + PermissionOutcome::Deny { + reason: format_hook_message( + &pre_hook_result, + &format!("PreToolUse hook failed for tool `{tool_name}`"), + ), + } + } else if pre_hook_result.is_denied() { + PermissionOutcome::Deny { + reason: format_hook_message( + &pre_hook_result, + &format!("PreToolUse hook denied tool `{tool_name}`"), + ), + } + } else if let Some(prompt) = prompter.as_mut() { + self.permission_policy.authorize_with_context( + &tool_name, + &effective_input, + &permission_context, + Some(*prompt), + ) + } else { + self.permission_policy.authorize_with_context( + &tool_name, + &effective_input, + &permission_context, + None, + ) + }; + + 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); + + 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() + || post_hook_result.is_failed() + || post_hook_result.is_cancelled(), + ); + + ConversationMessage::tool_result( + tool_use_id, + tool_name, + output, + is_error, + ) + } + 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()) .map_err(|error| RuntimeError::new(error.to_string()))?; + self.emit_tool_result_event(&result_message); + self.record_tool_finished(iterations, &result_message); tool_results.push(result_message); + } + + if !interrupted && self.turn_control_is_aborted() { + self.record_turn_interrupted(iterations, "after_tool"); + interrupted = true; + } + if interrupted { + break; + } + if self.append_queued_steer()? { continue; } - let pre_hook_result = self.run_pre_tool_use_hook(&tool_name, &input); - let effective_input = pre_hook_result - .updated_input() - .map_or_else(|| input.clone(), ToOwned::to_owned); - let permission_context = PermissionContext::new( - pre_hook_result.permission_override(), - pre_hook_result.permission_reason().map(ToOwned::to_owned), - ); - - let permission_outcome = if pre_hook_result.is_cancelled() { - PermissionOutcome::Deny { - reason: format_hook_message( - &pre_hook_result, - &format!("PreToolUse hook cancelled tool `{tool_name}`"), - ), - } - } else if pre_hook_result.is_failed() { - PermissionOutcome::Deny { - reason: format_hook_message( - &pre_hook_result, - &format!("PreToolUse hook failed for tool `{tool_name}`"), - ), - } - } else if pre_hook_result.is_denied() { - PermissionOutcome::Deny { - reason: format_hook_message( - &pre_hook_result, - &format!("PreToolUse hook denied tool `{tool_name}`"), - ), - } - } else if let Some(prompt) = prompter.as_mut() { - self.permission_policy.authorize_with_context( - &tool_name, - &effective_input, - &permission_context, - Some(*prompt), - ) - } else { - self.permission_policy.authorize_with_context( - &tool_name, - &effective_input, - &permission_context, - None, - ) - }; - - 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) { - Ok(output) => (output, false), - Err(error) => (error.to_string(), true), - }; - 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() - || post_hook_result.is_failed() - || post_hook_result.is_cancelled(), - ); - - ConversationMessage::tool_result(tool_use_id, tool_name, output, is_error) - } - 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()) - .map_err(|error| RuntimeError::new(error.to_string()))?; - self.record_tool_finished(iterations, &result_message); - tool_results.push(result_message); + if self.turn_control_is_aborted() { + self.record_turn_interrupted(iterations, "after_tool"); + interrupted = true; + break; + } } - if interrupted { - break; + let summary = TurnSummary { + assistant_messages, + tool_results, + prompt_cache_events, + iterations, + usage: self.usage_tracker.cumulative_usage(), + auto_compaction, + interrupted, + }; + if !interrupted { + self.record_turn_completed(&summary); } - } - let summary = TurnSummary { - assistant_messages, - tool_results, - prompt_cache_events, - iterations, - usage: self.usage_tracker.cumulative_usage(), - auto_compaction, + Ok(summary) + })(); + self.emit_conversation_event(ConversationEvent::AgentEnd { + session_id, interrupted, - }; - if !interrupted { - self.record_turn_completed(&summary); - } - - Ok(summary) + output: final_output, + }); + outcome } #[must_use] @@ -885,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 { @@ -956,9 +1223,9 @@ impl ToolExecutor for StaticToolExecutor { mod tests { use super::{ build_assistant_message, parse_auto_compaction_threshold, ApiClient, ApiRequest, - AssistantEvent, AutoCompactionEvent, ConversationRuntime, PromptCacheEvent, RuntimeError, - StaticToolExecutor, ToolExecutor, TurnInterruptSignal, - DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD, + AssistantEvent, AutoCompactionEvent, ConversationEvent, ConversationRuntime, + PromptCacheEvent, RuntimeError, StaticToolExecutor, ToolExecutor, TurnControlSignal, + TurnInterruptSignal, DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD, }; use crate::compact::CompactionConfig; use crate::config::{RuntimeFeatureConfig, RuntimeHookConfig}; @@ -972,7 +1239,7 @@ mod tests { use crate::ToolError; use std::fs; use std::path::PathBuf; - use std::sync::Arc; + use std::sync::{Arc, Mutex}; use std::time::{SystemTime, UNIX_EPOCH}; use telemetry::{MemoryTelemetrySink, SessionTracer, TelemetryEvent}; @@ -1045,6 +1312,300 @@ mod tests { } } + #[test] + fn emits_ordered_lifecycle_events_for_tool_calls() { + struct ToolThenDoneApi { + calls: usize, + } + + impl ApiClient for ToolThenDoneApi { + fn stream( + &mut self, + _request: ApiRequest, + ) -> Result, RuntimeError> { + self.calls += 1; + if self.calls == 1 { + Ok(vec![ + AssistantEvent::ToolUse { + id: "tool-event-1".to_string(), + name: "echo".to_string(), + input: "payload".to_string(), + }, + AssistantEvent::MessageStop, + ]) + } else { + Ok(vec![ + AssistantEvent::TextDelta("done".to_string()), + AssistantEvent::MessageStop, + ]) + } + } + } + + let session = Session::new(); + let session_id = session.session_id.clone(); + let captured_events = Arc::new(Mutex::new(Vec::new())); + let callback_events = Arc::clone(&captured_events); + let mut runtime = ConversationRuntime::new( + session, + ToolThenDoneApi { calls: 0 }, + StaticToolExecutor::new().register("echo", |input| Ok(input.to_string())), + PermissionPolicy::new(PermissionMode::DangerFullAccess), + vec!["system".to_string()], + ) + .with_event_callback(Arc::new(move |event| { + callback_events + .lock() + .expect("event mutex should not be poisoned") + .push(event); + })); + + runtime + .run_turn("start", None) + .expect("tool turn should complete"); + + assert_eq!( + *captured_events + .lock() + .expect("event mutex should not be poisoned"), + vec![ + ConversationEvent::AgentBegin { + session_id: session_id.clone(), + }, + ConversationEvent::ToolUse { + tool_use_id: "tool-event-1".to_string(), + tool_name: "echo".to_string(), + input: "payload".to_string(), + }, + ConversationEvent::ToolResult { + tool_use_id: "tool-event-1".to_string(), + tool_name: "echo".to_string(), + output: "payload".to_string(), + is_error: false, + }, + ConversationEvent::AgentEnd { + session_id, + interrupted: false, + output: Some("done".to_string()), + }, + ] + ); + } + + #[test] + fn consumes_steering_queued_during_a_model_request_and_preserves_late_steers() { + struct SteeringApi { + control: TurnControlSignal, + calls: usize, + } + + impl ApiClient for SteeringApi { + fn stream(&mut self, request: ApiRequest) -> Result, RuntimeError> { + self.calls += 1; + if self.calls == 1 { + // Queue from another thread while this model request is + // still inside `stream`. + let control = self.control.clone(); + assert!(std::thread::spawn(move || control.steer("follow-up")) + .join() + .expect("steering thread should finish")); + return Ok(vec![ + AssistantEvent::ToolUse { + id: "steer-tool".to_string(), + name: "echo".to_string(), + input: "tool input".to_string(), + }, + AssistantEvent::MessageStop, + ]); + } + if self.calls == 2 { + assert!(request + .messages + .iter() + .any(|message| message.role == MessageRole::Tool)); + assert!(request.messages.iter().any(|message| { + message.role == MessageRole::User + && message.blocks.iter().any(|block| { + matches!(block, ContentBlock::Text { text } if text == "follow-up") + }) + })); + return Ok(vec![ + AssistantEvent::TextDelta("second response".to_string()), + AssistantEvent::MessageStop, + ]); + } + unreachable!("the queued steer should only cause one extra request") + } + } + + let control = TurnControlSignal::new(); + let mut runtime = ConversationRuntime::new( + Session::new(), + SteeringApi { + control: control.clone(), + calls: 0, + }, + StaticToolExecutor::new().register("echo", |input| Ok(input.to_string())), + PermissionPolicy::new(PermissionMode::DangerFullAccess), + vec!["system".to_string()], + ) + .with_turn_control_signal(control.clone()); + + let summary = runtime + .run_turn("start", None) + .expect("steering turn should complete"); + + assert_eq!(summary.iterations, 2); + assert!(runtime.session().messages.iter().any(|message| { + message.role == MessageRole::User + && message.blocks.iter().any( + |block| matches!(block, ContentBlock::Text { text } if text == "follow-up"), + ) + })); + assert!(control.steer("late steer")); + assert_eq!( + control.finish_or_take_steer().as_deref(), + Some("late steer") + ); + assert_eq!(control.finish_or_take_steer(), None); + assert!(!control.steer("rejected after finish")); + } + + #[test] + fn aborting_turn_control_uses_existing_interrupted_outcome() { + struct AbortApi { + control: TurnControlSignal, + } + + impl ApiClient for AbortApi { + fn stream( + &mut self, + _request: ApiRequest, + ) -> Result, RuntimeError> { + self.control.abort(); + Err(RuntimeError::new("request aborted")) + } + } + + let control = TurnControlSignal::new(); + let events = Arc::new(Mutex::new(Vec::new())); + let callback_events = Arc::clone(&events); + let mut runtime = ConversationRuntime::new( + Session::new(), + AbortApi { + control: control.clone(), + }, + StaticToolExecutor::new(), + PermissionPolicy::new(PermissionMode::DangerFullAccess), + vec!["system".to_string()], + ) + .with_turn_control_signal(control.clone()) + .with_event_callback(Arc::new(move |event| { + callback_events + .lock() + .expect("event mutex should not be poisoned") + .push(event); + })); + + let summary = runtime + .run_turn("start", None) + .expect("control abort should be reported as an interruption"); + + assert!(summary.interrupted); + assert!(control.is_aborted()); + assert!(!control.steer("too late")); + assert!(matches!( + events + .lock() + .expect("event mutex should not be poisoned") + .last(), + Some(ConversationEvent::AgentEnd { + interrupted: true, + .. + }) + )); + } + + #[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 }; @@ -2127,8 +2688,12 @@ mod tests { // given let interrupt = TurnInterruptSignal::new(); + let session = Session::new(); + let session_id = session.session_id.clone(); + let captured_events = Arc::new(Mutex::new(Vec::new())); + let callback_events = Arc::clone(&captured_events); let mut runtime = ConversationRuntime::new( - Session::new(), + session, ToolUseApi { interrupt: interrupt.clone(), }, @@ -2137,7 +2702,13 @@ mod tests { PermissionPolicy::new(PermissionMode::DangerFullAccess), vec!["system".to_string()], ) - .with_turn_interrupt_signal(interrupt); + .with_turn_interrupt_signal(interrupt) + .with_event_callback(Arc::new(move |event| { + callback_events + .lock() + .expect("event mutex should not be poisoned") + .push(event); + })); // when let summary = runtime @@ -2164,6 +2735,32 @@ mod tests { runtime.session().messages[2].blocks[0], ContentBlock::ToolResult { is_error: true, .. } )); + assert_eq!( + *captured_events + .lock() + .expect("event mutex should not be poisoned"), + vec![ + ConversationEvent::AgentBegin { + session_id: session_id.clone(), + }, + ConversationEvent::ToolUse { + tool_use_id: "tool-1".to_string(), + tool_name: "add".to_string(), + input: "2,2".to_string(), + }, + ConversationEvent::ToolResult { + tool_use_id: "tool-1".to_string(), + tool_name: "add".to_string(), + output: "Interrupted by user before this tool could run.".to_string(), + is_error: true, + }, + ConversationEvent::AgentEnd { + session_id, + interrupted: true, + output: Some("Running the tool.".to_string()), + }, + ] + ); } #[test] diff --git a/rust/crates/runtime/src/lib.rs b/rust/crates/runtime/src/lib.rs index 758e8555..fd2aee4e 100644 --- a/rust/crates/runtime/src/lib.rs +++ b/rust/crates/runtime/src/lib.rs @@ -84,8 +84,8 @@ pub use config_validate::{ }; pub use conversation::{ auto_compaction_threshold_from_env, ApiClient, ApiRequest, AssistantEvent, AutoCompactionEvent, - ConversationRuntime, PromptCacheEvent, RuntimeError, StaticToolExecutor, ToolError, - ToolExecutor, TurnInterruptSignal, TurnSummary, + ConversationEvent, ConversationRuntime, PromptCacheEvent, RuntimeError, StaticToolExecutor, + ToolError, ToolExecutor, TurnControlSignal, TurnInterruptSignal, TurnSummary, }; pub use file_ops::{ edit_file, edit_file_in_workspace, glob_search, glob_search_in_workspace, grep_search, diff --git a/rust/crates/sdk/Cargo.toml b/rust/crates/sdk/Cargo.toml index 9290a864..f542a96c 100644 --- a/rust/crates/sdk/Cargo.toml +++ b/rust/crates/sdk/Cargo.toml @@ -13,3 +13,10 @@ serde = { version = "1", features = ["derive"] } uuid = { version = "1", features = ["v4", "serde"] } tracing = "0.1" url = "2" +runtime = { path = "../runtime" } + +[dev-dependencies] +tempfile = "3" + +[lints] +workspace = true diff --git a/rust/crates/sdk/src/error.rs b/rust/crates/sdk/src/error.rs index b60079f9..9c08601f 100644 --- a/rust/crates/sdk/src/error.rs +++ b/rust/crates/sdk/src/error.rs @@ -15,6 +15,8 @@ pub enum SdkErrorCode { Unauthorized, EndpointCredentialForbidden, Unavailable, + /// Requested session state conflicts with an active invocation + Conflict, ConnectionError, ConnectionClosed, Timeout, @@ -37,6 +39,7 @@ impl SdkErrorCode { Self::Unauthorized => "unauthorized", Self::EndpointCredentialForbidden => "endpoint_credential_forbidden", Self::Unavailable => "unavailable", + Self::Conflict => "conflict", Self::ConnectionError => "connection_error", Self::ConnectionClosed => "connection_closed", Self::Timeout => "timeout", @@ -59,6 +62,7 @@ impl SdkErrorCode { "unauthorized" => Self::Unauthorized, "endpoint_credential_forbidden" => Self::EndpointCredentialForbidden, "unavailable" => Self::Unavailable, + "conflict" => Self::Conflict, "connection_error" => Self::ConnectionError, "connection_closed" => Self::ConnectionClosed, "timeout" => Self::Timeout, @@ -132,6 +136,14 @@ impl SdkError { Self::new(SdkErrorCode::Unavailable, message) } + pub fn not_found(message: impl Into) -> Self { + Self::new(SdkErrorCode::NotFound, message) + } + + pub fn conflict(message: impl Into) -> Self { + Self::new(SdkErrorCode::Conflict, message) + } + pub fn connection_error(message: impl Into) -> Self { Self::new(SdkErrorCode::ConnectionError, message) } @@ -209,6 +221,13 @@ mod tests { ); } + #[test] + fn conflict_error_code_round_trips_on_the_wire() { + let error = SdkError::conflict("session is active"); + assert_eq!(error.code.to_string(), "conflict"); + assert_eq!(SdkErrorCode::from_wire("conflict"), SdkErrorCode::Conflict); + } + #[test] fn unknown_remote_code_is_preserved_with_typed_internal_code() { let error = SdkError::from_remote("future_error", "failed"); diff --git a/rust/crates/sdk/src/lib.rs b/rust/crates/sdk/src/lib.rs index 4703a7cb..672093b7 100644 --- a/rust/crates/sdk/src/lib.rs +++ b/rust/crates/sdk/src/lib.rs @@ -14,9 +14,16 @@ pub mod error; pub mod frame; pub mod registry; +pub mod session; pub mod transport; pub use error::{SdkError, SdkErrorCode, SdkResult, SdkSentRecord}; pub use frame::{FrameType, SdkFrame}; pub use registry::{OperationMetadata, OperationRegistry, Schema}; +pub use session::{ + InvocationOutcome, InvocationReconciliationStore, InvocationRecord, InvocationStatus, + RuntimeSessionRunner, SessionDispatch, SessionEvent, SessionEventStream, SessionEventType, + SessionInvocation, SessionPage, SessionRouter, SessionRunResult, SessionSummary, + SessionTurnRunner, TerminalScopeRecord, TerminalScopeStore, +}; pub use transport::{SdkClient, SdkClientConfig}; diff --git a/rust/crates/sdk/src/registry.rs b/rust/crates/sdk/src/registry.rs index e4e9560e..b52bff0f 100644 --- a/rust/crates/sdk/src/registry.rs +++ b/rust/crates/sdk/src/registry.rs @@ -263,6 +263,89 @@ impl OperationRegistry { )), ); + // SDK session lifecycle operations. + for (name, description, required) in [ + ( + "turn.create", + "Create a durable SDK session", + Vec::<&str>::new(), + ), + ( + "turn.prompt", + "Prompt a session and stream lifecycle events", + Vec::<&str>::new(), + ), + ( + "turn.steer", + "Steer an active session invocation", + vec!["invocationId"], + ), + ( + "turn.abort_and_prompt", + "Abort the active invocation and start a replacement prompt", + Vec::<&str>::new(), + ), + ( + "turn.list", + "List persisted SDK sessions", + Vec::<&str>::new(), + ), + ] { + let mut properties = serde_json::Map::new(); + properties.insert("sessionId".to_string(), json!({ "type": "string" })); + properties.insert("correlationKey".to_string(), json!({ "type": "string" })); + properties.insert("invocationId".to_string(), json!({ "type": "string" })); + properties.insert("prompt".to_string(), json!({ "type": "string" })); + properties.insert( + "override".to_string(), + json!({ "type": ["string", "object", "array"] }), + ); + properties.insert( + "request".to_string(), + json!({ "type": ["string", "object", "array"] }), + ); + properties.insert( + "input".to_string(), + json!({ "type": ["string", "object", "array"] }), + ); + properties.insert("tools".to_string(), json!({ "type": "array" })); + properties.insert("settings".to_string(), json!({ "type": "object" })); + properties.insert("model".to_string(), json!({ "type": "string" })); + properties.insert("cursor".to_string(), json!({ "type": "string" })); + properties.insert( + "limit".to_string(), + json!({ "type": "integer", "minimum": 1 }), + ); + let mut input_schema = json!({ + "type": "object", + "properties": properties, + "required": required, + }); + let prompt_alternatives = json!([ + { "required": ["prompt"] }, + { "required": ["request"] }, + { "required": ["input"] }, + ]); + match name { + "turn.prompt" | "turn.abort_and_prompt" => { + input_schema["anyOf"] = prompt_alternatives; + } + "turn.steer" => { + input_schema["anyOf"] = json!([ + { "required": ["override"] }, + { "required": ["prompt"] }, + { "required": ["request"] }, + { "required": ["input"] }, + ]); + } + _ => {} + } + registry.register( + OperationMetadata::new(name, description) + .with_input_schema(Schema::new(format!("Input for {name}"), input_schema)), + ); + } + registry } @@ -397,6 +480,20 @@ mod tests { assert!(registry.has_operation("events.replay")); } + #[test] + fn standard_registry_has_session_lifecycle_operations() { + let registry = OperationRegistry::with_standard_operations(); + for operation in [ + "turn.create", + "turn.prompt", + "turn.steer", + "turn.abort_and_prompt", + "turn.list", + ] { + assert!(registry.has_operation(operation), "missing {operation}"); + } + } + #[test] fn registry_validation_catches_missing_required_fields() { let registry = OperationRegistry::with_standard_operations(); diff --git a/rust/crates/sdk/src/session/events.rs b/rust/crates/sdk/src/session/events.rs new file mode 100644 index 00000000..56193ee4 --- /dev/null +++ b/rust/crates/sdk/src/session/events.rs @@ -0,0 +1,31 @@ +//! Event stream types for session lifecycle operations. + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +/// Lifecycle events published for one SDK turn invocation. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SessionEventType { + TurnAccepted, + AgentBegin, + ToolUse, + ToolResult, + AgentEnd, + TurnSteered, + TurnEnded, +} + +/// One ordered session lifecycle event. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SessionEvent { + pub event_id: String, + pub event_type: SessionEventType, + pub session_id: String, + pub invocation_id: String, + pub sequence: u64, + pub timestamp_ms: u64, + pub durable: bool, + pub payload: Value, +} diff --git a/rust/crates/sdk/src/session/mod.rs b/rust/crates/sdk/src/session/mod.rs new file mode 100644 index 00000000..c0dc5cbf --- /dev/null +++ b/rust/crates/sdk/src/session/mod.rs @@ -0,0 +1,1309 @@ +//! Durable SDK session routing and turn lifecycle operations. + +pub mod events; +pub mod reconciliation; +pub mod terminal_scope; + +pub use events::{SessionEvent, SessionEventType}; +pub use reconciliation::{ + InvocationOutcome, InvocationReconciliationStore, InvocationRecord, InvocationStatus, +}; +pub use terminal_scope::{TerminalScopeRecord, TerminalScopeStore}; + +use std::collections::{BTreeMap, HashMap}; +use std::fmt::{Debug, Formatter}; +use std::panic::{catch_unwind, AssertUnwindSafe}; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::thread; +use std::time::{SystemTime, UNIX_EPOCH}; + +use runtime::{ + ApiClient, ConversationEvent, ConversationRuntime, Session, SessionStore, ToolExecutor, + TurnControlSignal, +}; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Map, Value}; +use tokio::sync::{mpsc, watch}; +use uuid::Uuid; + +use crate::error::{SdkError, SdkResult}; +use crate::frame::{FrameType, SdkFrame}; + +const DEFAULT_PAGE_SIZE: usize = 20; +const MAX_PAGE_SIZE: usize = 100; +const MAX_SESSION_ID_BYTES: usize = 128; +const MAX_CORRELATION_KEY_BYTES: usize = 512; + +/// Result returned by one session turn implementation. +pub enum SessionRunResult { + Completed(Session), + Aborted(Session), + Failed { session: Session, message: String }, +} + +/// Pluggable synchronous turn driver used by [`SessionRouter`]. +pub trait SessionTurnRunner: Send + Sync + 'static { + fn run( + &self, + session: Session, + prompt: String, + control: TurnControlSignal, + event_callback: Arc, + ) -> SessionRunResult; +} + +/// Bridges a runtime conversation into the SDK session router. +pub struct RuntimeSessionRunner { + build_runtime: Arc ConversationRuntime + Send + Sync>, + marker: std::marker::PhantomData (C, T)>, +} + +impl RuntimeSessionRunner +where + C: ApiClient + 'static, + T: ToolExecutor + 'static, +{ + /// Construct a runner from a factory that supplies the API and tool clients. + pub fn new(build_runtime: F) -> Self + where + F: Fn(Session) -> ConversationRuntime + Send + Sync + 'static, + { + Self { + build_runtime: Arc::new(build_runtime), + marker: std::marker::PhantomData, + } + } +} + +impl SessionTurnRunner for RuntimeSessionRunner +where + C: ApiClient + 'static, + T: ToolExecutor + 'static, +{ + fn run( + &self, + session: Session, + mut prompt: String, + control: TurnControlSignal, + event_callback: Arc, + ) -> SessionRunResult { + let mut runtime = (self.build_runtime)(session) + .with_turn_control_signal(control.clone()) + .with_event_callback(event_callback); + + loop { + match runtime.run_turn(prompt, None) { + Ok(summary) if summary.interrupted || control.is_aborted() => { + return SessionRunResult::Aborted(runtime.into_session()); + } + Ok(_) => match control.finish_or_take_steer() { + Some(steer) => prompt = steer, + None => return SessionRunResult::Completed(runtime.into_session()), + }, + Err(error) => { + control.close(); + return SessionRunResult::Failed { + session: runtime.into_session(), + message: error.to_string(), + }; + } + } + } + } +} + +/// One paginated row returned by `turn.list`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SessionSummary { + pub session_id: String, + pub created_at_ms: u64, + pub updated_at_ms: u64, + pub message_count: usize, +} + +/// Page of persisted sessions. `next_cursor` is an opaque decimal offset. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SessionPage { + pub sessions: Vec, + pub next_cursor: Option, +} + +/// A started invocation and its ordered event stream. +pub struct SessionInvocation { + pub session_id: String, + pub invocation_id: String, + pub events: SessionEventStream, +} + +impl Debug for SessionInvocation { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("SessionInvocation") + .field("session_id", &self.session_id) + .field("invocation_id", &self.invocation_id) + .finish_non_exhaustive() + } +} + +/// Receiver for lifecycle events emitted by an invocation. +pub struct SessionEventStream { + receiver: mpsc::UnboundedReceiver, +} + +impl Debug for SessionEventStream { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("SessionEventStream") + .finish_non_exhaustive() + } +} + +impl SessionEventStream { + /// Receive the next lifecycle event, returning `None` after the stream closes. + pub async fn recv(&mut self) -> Option { + self.receiver.recv().await + } +} + +/// Result of dispatching one SDK broker request. +pub struct SessionDispatch { + pub response: SdkFrame, + pub events: Option, +} + +impl Debug for SessionDispatch { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("SessionDispatch") + .field("response", &self.response) + .field("has_events", &self.events.is_some()) + .finish() + } +} + +#[derive(Clone)] +pub struct SessionRouter { + inner: Arc, +} + +struct SessionRouterInner { + store: SessionStore, + runner: Arc, + reconciliation: InvocationReconciliationStore, + terminal_scopes: TerminalScopeStore, + reconciliation_lock: Mutex<()>, + active: Mutex>, +} + +#[derive(Clone)] +struct ActiveInvocation { + invocation_id: String, + correlation_key: Option, + control: TurnControlSignal, + event_publisher: Arc, + done: watch::Receiver, +} + +struct SessionEventPublisher { + session_id: String, + invocation_id: String, + sequence: std::sync::atomic::AtomicU64, + publication_lock: Mutex<()>, + sender: mpsc::UnboundedSender, +} + +impl SessionEventPublisher { + fn new( + session_id: String, + invocation_id: String, + sender: mpsc::UnboundedSender, + ) -> Self { + Self { + session_id, + invocation_id, + sequence: std::sync::atomic::AtomicU64::new(0), + publication_lock: Mutex::new(()), + sender, + } + } + + fn lock(&self) -> MutexGuard<'_, ()> { + self.publication_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + fn build_locked( + &self, + event_type: SessionEventType, + payload: Value, + durable: bool, + ) -> SessionEvent { + let sequence = self + .sequence + .fetch_add(1, std::sync::atomic::Ordering::SeqCst) + .saturating_add(1); + SessionEvent { + event_id: format!("{}:{sequence}", self.invocation_id), + event_type, + session_id: self.session_id.clone(), + invocation_id: self.invocation_id.clone(), + sequence, + timestamp_ms: now_ms(), + durable, + payload, + } + } + + fn publish(&self, event_type: SessionEventType, payload: Value, durable: bool) -> SessionEvent { + let _guard = self.lock(); + self.publish_locked(event_type, payload, durable) + } + + fn publish_locked( + &self, + event_type: SessionEventType, + payload: Value, + durable: bool, + ) -> SessionEvent { + let event = self.build_locked(event_type, payload, durable); + let _ = self.sender.send(event.clone()); + event + } + + fn publish_steer(&self, control: &TurnControlSignal, prompt: &str) -> bool { + let _guard = self.lock(); + if !control.steer(prompt.to_string()) { + return false; + } + self.publish_locked( + SessionEventType::TurnSteered, + json!({ "prompt": prompt }), + false, + ); + true + } + + fn publish_runtime_event(&self, event: ConversationEvent) { + let (event_type, payload) = match event { + ConversationEvent::AgentBegin { session_id: _ } => { + (SessionEventType::AgentBegin, json!({})) + } + ConversationEvent::ToolUse { + tool_use_id, + tool_name, + input, + } => ( + SessionEventType::ToolUse, + json!({ + "toolUseId": tool_use_id, + "toolName": tool_name, + "input": input, + }), + ), + ConversationEvent::ToolResult { + tool_use_id, + tool_name, + output, + is_error, + } => ( + SessionEventType::ToolResult, + json!({ + "toolUseId": tool_use_id, + "toolName": tool_name, + "output": output, + "isError": is_error, + }), + ), + ConversationEvent::AgentEnd { + session_id: _, + 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); + } +} + +impl SessionRouter { + /// Create a router backed by the supplied durable session store. + /// + /// Interrupted invocations left by an earlier process are moved to the + /// uncertain state so clients can reconcile them against their correlation key. + pub fn new(store: SessionStore, runner: Arc) -> SdkResult { + let sdk_root = store.sessions_dir().join("sdk"); + let reconciliation = InvocationReconciliationStore::new(sdk_root.join("reconciliation")); + 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()), + }), + }) + } + + /// Create and persist an empty session, optionally assigning its model and correlation scope. + pub fn create_session( + &self, + model: Option<&str>, + correlation_key: Option<&str>, + ) -> SdkResult { + validate_optional_correlation_key(correlation_key)?; + let model = model.map(str::trim).filter(|value| !value.is_empty()); + if model.is_some_and(|value| value.len() > 256) { + return Err(SdkError::invalid_input("model name exceeds 256 bytes")); + } + + let mut session = + Session::new().with_workspace_root(self.inner.store.workspace_root().to_path_buf()); + session.model = model.map(str::to_string); + let handle = self.inner.store.create_handle(&session.session_id); + session = session.with_persistence_path(handle.path.clone()); + session + .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) + } + + /// Load a session by its durable ID, rejecting aliases and path-like IDs. + pub fn get_session(&self, session_id: &str) -> SdkResult { + validate_session_id(session_id)?; + let path = self + .inner + .store + .resolve_managed_path(session_id) + .map_err(|error| SdkError::not_found(error.to_string()))?; + let session = Session::load_from_path(&path) + .map_err(|error| SdkError::internal(format!("could not load session: {error}")))?; + if session.session_id != session_id { + return Err(SdkError::not_found(format!( + "session `{session_id}` does not match the requested ID" + ))); + } + Ok(session) + } + + /// 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); + 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 + .inner + .reconciliation + .list_by_correlation_key(correlation_key)? + .iter() + .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)) + } + + /// Start a prompt and return immediately with an ordered event receiver. + pub fn prompt( + &self, + session_id: &str, + prompt: impl Into, + correlation_key: Option<&str>, + ) -> SdkResult { + let prompt = prompt.into(); + validate_session_id(session_id)?; + validate_prompt(&prompt)?; + validate_optional_correlation_key(correlation_key)?; + + let mut active = self + .inner + .active + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if active.contains_key(session_id) { + return Err(SdkError::conflict(format!( + "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(); + let record = InvocationRecord { + invocation_id: invocation_id.clone(), + session_id: session_id.to_string(), + correlation_key: correlation_key.map(str::to_string), + status: InvocationStatus::Accepted, + accepted_at_ms: created_at_ms, + updated_at_ms: created_at_ms, + outcome: None, + uncertainty_reason: None, + terminal_event: None, + terminal_scope_revision: None, + }; + self.inner.reconciliation.insert(record)?; + + let (event_tx, event_rx) = mpsc::unbounded_channel(); + let publisher = Arc::new(SessionEventPublisher::new( + session_id.to_string(), + invocation_id.clone(), + event_tx, + )); + let control = TurnControlSignal::new(); + let (done_tx, done_rx) = watch::channel(false); + active.insert( + session_id.to_string(), + ActiveInvocation { + invocation_id: invocation_id.clone(), + correlation_key: correlation_key.map(str::to_string), + control: control.clone(), + event_publisher: publisher.clone(), + done: done_rx, + }, + ); + drop(active); + + publisher.publish( + SessionEventType::TurnAccepted, + json!({ "correlationKey": correlation_key }), + false, + ); + + let inner = Arc::clone(&self.inner); + let worker_session_id = session_id.to_string(); + let worker_invocation_id = invocation_id.clone(); + let worker_publisher = publisher.clone(); + let worker_control = control.clone(); + let worker_correlation_key = correlation_key.map(str::to_string); + let worker_session = session.clone(); + let worker_done_tx = done_tx.clone(); + let worker = thread::Builder::new() + .name(format!("sdk-session-{}", &invocation_id[..8])) + .spawn(move || { + if let Err(error) = inner.reconciliation.mark_running(&worker_invocation_id) { + finish_uncertain( + &inner, + &worker_invocation_id, + &worker_control, + &worker_publisher, + &error.to_string(), + ); + remove_active(&inner, &worker_session_id, &worker_invocation_id); + let _ = worker_done_tx.send(true); + return; + } + let event_publisher = worker_publisher.clone(); + let callback: Arc = + Arc::new(move |event| event_publisher.publish_runtime_event(event)); + let result = catch_unwind(AssertUnwindSafe(|| { + inner + .runner + .run(worker_session, prompt, worker_control.clone(), callback) + })); + match result { + Ok(run_result) => finish_invocation( + &inner, + &worker_session_id, + &worker_invocation_id, + worker_correlation_key.as_deref(), + &worker_control, + &worker_publisher, + run_result, + ), + Err(panic) => { + let reason = panic_message(panic); + finish_uncertain( + &inner, + &worker_invocation_id, + &worker_control, + &worker_publisher, + &reason, + ); + } + } + remove_active(&inner, &worker_session_id, &worker_invocation_id); + let _ = worker_done_tx.send(true); + }); + + if let Err(error) = worker { + let message = format!("could not start session worker: {error}"); + finish_invocation( + &self.inner, + session_id, + &invocation_id, + correlation_key, + &control, + &publisher, + SessionRunResult::Failed { session, message }, + ); + remove_active(&self.inner, session_id, &invocation_id); + let _ = done_tx.send(true); + } + + Ok(SessionInvocation { + session_id: session_id.to_string(), + invocation_id, + events: SessionEventStream { receiver: event_rx }, + }) + } + + /// Queue a steering prompt for the active invocation of a session. + pub fn steer( + &self, + session_id: &str, + invocation_id: &str, + override_prompt: impl Into, + ) -> SdkResult<()> { + validate_session_id(session_id)?; + validate_invocation_id(invocation_id)?; + let prompt = override_prompt.into(); + validate_prompt(&prompt)?; + let (control, publisher) = { + let active = self + .inner + .active + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let invocation = active + .get(session_id) + .ok_or_else(|| SdkError::not_found("session has no active invocation"))?; + if invocation.invocation_id != invocation_id { + return Err(SdkError::not_found(format!( + "invocation `{invocation_id}` is not active for session `{session_id}`" + ))); + } + ( + invocation.control.clone(), + Arc::clone(&invocation.event_publisher), + ) + }; + if !publisher.publish_steer(&control, &prompt) { + return Err(SdkError::conflict("invocation no longer accepts steering")); + } + Ok(()) + } + + /// Abort the current invocation, wait for its durable terminal outcome, then start a new prompt. + pub async fn abort_and_prompt( + &self, + session_id: &str, + prompt: impl Into, + correlation_key: Option<&str>, + ) -> SdkResult { + validate_session_id(session_id)?; + let prompt = prompt.into(); + validate_prompt(&prompt)?; + validate_optional_correlation_key(correlation_key)?; + + let (mut done, inherited_correlation, previous_invocation_id) = { + let active = self + .inner + .active + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match active.get(session_id) { + Some(invocation) => { + invocation.control.abort(); + ( + Some(invocation.done.clone()), + invocation.correlation_key.clone(), + Some(invocation.invocation_id.clone()), + ) + } + None => (None, None, None), + } + }; + if let Some(receiver) = done.as_mut() { + while !*receiver.borrow() { + receiver.changed().await.map_err(|_| { + SdkError::unavailable("session invocation stopped unexpectedly") + })?; + } + } + + if let Some(invocation_id) = previous_invocation_id { + let record = self + .inner + .reconciliation + .get(&invocation_id)? + .ok_or_else(|| SdkError::not_found("aborted invocation record was not found"))?; + if record.status != InvocationStatus::Terminal { + return Err(SdkError::unavailable(format!( + "invocation `{invocation_id}` has no durable terminal outcome" + ))); + } + } + + let effective_correlation = correlation_key + .map(str::to_string) + .or(inherited_correlation); + self.prompt(session_id, prompt, effective_correlation.as_deref()) + } + + /// List persisted sessions using an opaque offset cursor. + pub fn list(&self, cursor: Option<&str>, limit: Option) -> SdkResult { + let offset = match cursor { + Some(value) => value.parse::().map_err(|_| { + SdkError::invalid_input("cursor must be a non-negative decimal offset") + })?, + None => 0, + }; + let limit = limit.unwrap_or(DEFAULT_PAGE_SIZE); + if !(1..=MAX_PAGE_SIZE).contains(&limit) { + return Err(SdkError::invalid_input(format!( + "limit must be between 1 and {MAX_PAGE_SIZE}" + ))); + } + let sessions = self + .inner + .store + .list_sessions() + .map_err(|error| SdkError::internal(format!("could not list sessions: {error}")))?; + let total = sessions.len(); + let summaries = sessions + .into_iter() + .skip(offset) + .take(limit) + .map(|session| SessionSummary { + session_id: session.id, + created_at_ms: session.created_at_ms, + updated_at_ms: session.updated_at_ms, + message_count: session.message_count, + }) + .collect::>(); + let next_offset = offset.saturating_add(summaries.len()); + let next_cursor = (next_offset < total).then(|| next_offset.to_string()); + Ok(SessionPage { + sessions: summaries, + next_cursor, + }) + } + + /// 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)? + .ok_or_else(|| { + SdkError::not_found(format!("invocation `{invocation_id}` was not found")) + }) + } + + /// Return all invocation outcomes associated with a correlation key in stable ID order. + pub fn reconcile_by_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)?; + self.inner + .reconciliation + .list_by_correlation_key(correlation_key) + } + + /// Dispatch a broker request for `turn.create`, `turn.prompt`, `turn.steer`, + /// `turn.abort_and_prompt`, or `turn.list`. + pub async fn dispatch(&self, request: SdkFrame) -> SdkResult { + if request.frame_type()? != FrameType::BrokerRequest { + return Err(SdkError::invalid_input( + "session router requires a broker_request frame", + )); + } + let request_id = request.id()?.ok_or_else(|| { + SdkError::invalid_input("request frame is missing its correlation ID") + })?; + let operation = request + .operation() + .ok_or_else(|| SdkError::invalid_input("request frame is missing its operation"))?; + let input = request.input().unwrap_or_else(|| json!({})); + let object = input + .as_object() + .ok_or_else(|| SdkError::invalid_input("operation input must be a JSON object"))?; + + match operation.as_str() { + "turn.create" => { + let correlation_key = optional_string(object, "correlationKey")?; + let model = optional_string(object, "model")?; + let initial_prompt = optional_prompt(object, "input")?; + if let Some(prompt) = initial_prompt.as_deref() { + validate_prompt(prompt)?; + } + let session = self.create_session(model.as_deref(), correlation_key.as_deref())?; + if let Some(prompt) = initial_prompt { + let invocation = + self.prompt(&session.session_id, prompt, correlation_key.as_deref())?; + let response = SdkFrame::broker_response_ok( + request_id, + json!({ + "sessionId": session.session_id, + "invocationId": invocation.invocation_id, + }), + ); + Ok(SessionDispatch { + response, + events: Some(invocation.events), + }) + } else { + Ok(SessionDispatch { + response: SdkFrame::broker_response_ok( + request_id, + json!({ "sessionId": session.session_id }), + ), + events: None, + }) + } + } + "turn.prompt" => { + let session_id = self.resolve_session_id(object)?; + let prompt = required_turn_prompt(object)?; + let correlation_key = optional_string(object, "correlationKey")?; + let invocation = self.prompt(&session_id, prompt, correlation_key.as_deref())?; + Ok(SessionDispatch { + response: SdkFrame::broker_response_ok( + request_id, + invocation_result(&invocation), + ), + events: Some(invocation.events), + }) + } + "turn.steer" => { + let session_id = self.resolve_session_id(object)?; + let invocation_id = required_string(object, "invocationId")?; + let prompt = match object.get("override") { + Some(value) => prompt_value(value)?, + None => required_turn_prompt(object)?, + }; + self.steer(&session_id, &invocation_id, prompt)?; + Ok(SessionDispatch { + response: SdkFrame::broker_response_ok(request_id, json!({ "accepted": true })), + events: None, + }) + } + "turn.abort_and_prompt" => { + let session_id = self.resolve_session_id(object)?; + let prompt = required_turn_prompt(object)?; + let correlation_key = optional_string(object, "correlationKey")?; + let invocation = self + .abort_and_prompt(&session_id, prompt, correlation_key.as_deref()) + .await?; + Ok(SessionDispatch { + response: SdkFrame::broker_response_ok( + request_id, + invocation_result(&invocation), + ), + events: Some(invocation.events), + }) + } + "turn.list" => { + let cursor = optional_string(object, "cursor")?; + let limit = optional_usize(object, "limit")?; + let page = self.list(cursor.as_deref(), limit)?; + Ok(SessionDispatch { + response: SdkFrame::broker_response_ok( + request_id, + serde_json::to_value(page).map_err(|error| { + SdkError::internal(format!("could not encode session page: {error}")) + })?, + ), + events: None, + }) + } + _ => Err(SdkError::unknown_operation(format!( + "unknown session operation `{operation}`" + ))), + } + } + + fn resolve_session_id(&self, input: &Map) -> SdkResult { + if let Some(value) = input.get("sessionId") { + let session_id = value + .as_str() + .ok_or_else(|| SdkError::invalid_input("field `sessionId` must be a string"))?; + validate_session_id(session_id)?; + return Ok(session_id.to_string()); + } + if let Some(value) = input.get("correlationKey") { + let correlation_key = value.as_str().ok_or_else(|| { + SdkError::invalid_input("field `correlationKey` must be a string") + })?; + return self + .session_for_correlation_key(correlation_key)? + .ok_or_else(|| { + SdkError::not_found("no session is associated with this correlation key") + }); + } + Err(SdkError::invalid_input( + "operation requires `sessionId` or a known `correlationKey`", + )) + } +} + +fn invocation_result(invocation: &SessionInvocation) -> Value { + json!({ + "sessionId": invocation.session_id, + "invocationId": invocation.invocation_id, + }) +} + +fn finish_invocation( + inner: &SessionRouterInner, + session_id: &str, + invocation_id: &str, + correlation_key: Option<&str>, + control: &TurnControlSignal, + 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(); + + let (outcome, message_count, status, error_message) = match result { + SessionRunResult::Completed(session) => ( + InvocationOutcome::Completed { + message_count: session.messages.len(), + }, + session.messages.len(), + "completed", + None, + ), + SessionRunResult::Aborted(session) => ( + InvocationOutcome::Aborted { + message_count: session.messages.len(), + }, + session.messages.len(), + "aborted", + None, + ), + SessionRunResult::Failed { session, message } => ( + InvocationOutcome::Failed { + message: message.clone(), + message_count: session.messages.len(), + }, + session.messages.len(), + "failed", + Some(message), + ), + }; + let mut payload = json!({ "status": status, "messageCount": message_count }); + 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), + Err(error) => { + let _ = inner + .reconciliation + .mark_uncertain(invocation_id, error.to_string()); + return; + } + }, + None => None, + }; + match inner.reconciliation.record_terminal_with_scope_revision( + invocation_id, + outcome, + Some(stored_event), + terminal_scope_revision, + ) { + Ok(()) => { + let _ = publisher.sender.send(terminal_event); + } + Err(error) => { + let _ = inner + .reconciliation + .mark_uncertain(invocation_id, error.to_string()); + } + } +} + +fn recover_terminal_scopes( + reconciliation: &InvocationReconciliationStore, + terminal_scopes: &TerminalScopeStore, +) -> SdkResult<()> { + let mut latest_by_key = BTreeMap::::new(); + for record in reconciliation.list()? { + if record.status != InvocationStatus::Terminal { + continue; + } + let Some(key) = record.correlation_key.as_ref() else { + continue; + }; + let Some(revision) = record.terminal_scope_revision else { + continue; + }; + match latest_by_key.get(key) { + Some((current_revision, current_session)) if *current_revision == revision => { + if current_session != &record.session_id { + return Err(SdkError::internal(format!( + "terminal scope revision {revision} conflicts for key '{key}'" + ))); + } + } + Some((current_revision, _)) if *current_revision > revision => {} + _ => { + latest_by_key.insert(key.clone(), (revision, record.session_id)); + } + } + } + for (key, (revision, session_id)) in latest_by_key { + validate_correlation_key(&key)?; + terminal_scopes.restore_if_newer(&key, &session_id, revision)?; + } + 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, + control: &TurnControlSignal, + publisher: &SessionEventPublisher, + reason: &str, +) { + let _publication_guard = publisher.lock(); + control.close(); + let _ = inner + .reconciliation + .mark_uncertain(invocation_id, reason.to_string()); +} + +fn remove_active(inner: &SessionRouterInner, session_id: &str, invocation_id: &str) { + let mut active = inner + .active + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if active + .get(session_id) + .is_some_and(|value| value.invocation_id == invocation_id) + { + active.remove(session_id); + } +} + +fn panic_message(panic: Box) -> String { + panic + .downcast_ref::() + .cloned() + .or_else(|| { + panic + .downcast_ref::<&str>() + .map(|value| (*value).to_string()) + }) + .unwrap_or_else(|| "session runner panicked".to_string()) +} + +fn validate_session_id(session_id: &str) -> SdkResult<()> { + let is_valid = !session_id.is_empty() + && session_id.len() <= MAX_SESSION_ID_BYTES + && session_id + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')); + if is_valid { + Ok(()) + } else { + Err(SdkError::invalid_input("sessionId is malformed")) + } +} + +fn validate_invocation_id(invocation_id: &str) -> SdkResult<()> { + let is_uuid = Uuid::parse_str(invocation_id).is_ok(); + if is_uuid { + Ok(()) + } else { + Err(SdkError::invalid_input("invocationId must be a UUID")) + } +} + +fn validate_correlation_key(correlation_key: &str) -> SdkResult<()> { + if correlation_key.len() <= MAX_CORRELATION_KEY_BYTES { + Ok(()) + } else { + Err(SdkError::invalid_input("correlationKey exceeds 512 bytes")) + } +} + +fn validate_optional_correlation_key(correlation_key: Option<&str>) -> SdkResult<()> { + correlation_key.map_or(Ok(()), validate_correlation_key) +} + +fn validate_prompt(prompt: &str) -> SdkResult<()> { + if !prompt.trim().is_empty() && prompt.len() <= 1_000_000 { + Ok(()) + } else { + Err(SdkError::invalid_input( + "prompt must be non-empty and no larger than 1 MiB", + )) + } +} + +fn required_string(object: &Map, key: &str) -> SdkResult { + object + .get(key) + .and_then(Value::as_str) + .map(ToOwned::to_owned) + .ok_or_else(|| SdkError::invalid_input(format!("missing string field `{key}`"))) +} + +fn optional_prompt(object: &Map, key: &str) -> SdkResult> { + match object.get(key) { + None | Some(Value::Null) => Ok(None), + Some(value) => prompt_value(value).map(Some), + } +} + +fn required_turn_prompt(object: &Map) -> SdkResult { + for key in ["prompt", "request", "input"] { + if let Some(value) = object.get(key) { + return prompt_value(value); + } + } + Err(SdkError::invalid_input( + "operation requires a string `prompt` or `request` input", + )) +} + +fn prompt_value(value: &Value) -> SdkResult { + match value { + Value::String(prompt) => Ok(prompt.clone()), + Value::Object(object) => { + if let Some(messages) = object.get("messages").and_then(Value::as_array) { + if let Some(user_message) = messages.iter().rev().find(|message| { + message + .get("role") + .and_then(Value::as_str) + .is_some_and(|role| role == "user") + }) { + if let Some(content) = user_message.get("content") { + return prompt_value(content); + } + } + } + for key in ["prompt", "input", "text", "content"] { + if let Some(value) = object.get(key) { + return prompt_value(value); + } + } + Err(SdkError::invalid_input( + "request object must contain prompt text or a user message", + )) + } + Value::Array(blocks) => { + let mut text = Vec::new(); + for block in blocks { + match block { + Value::String(value) => text.push(value.clone()), + Value::Object(object) => { + if let Some(value) = object.get("text").and_then(Value::as_str) { + text.push(value.to_string()); + } else if let Some(value) = object.get("content") { + text.push(prompt_value(value)?); + } + } + _ => {} + } + } + if text.is_empty() { + Err(SdkError::invalid_input( + "request content has no text blocks", + )) + } else { + Ok(text.join("\n")) + } + } + _ => Err(SdkError::invalid_input( + "prompt/request content must contain text", + )), + } +} + +fn optional_string(object: &Map, key: &str) -> SdkResult> { + match object.get(key) { + None | Some(Value::Null) => Ok(None), + Some(Value::String(value)) => Ok(Some(value.clone())), + Some(_) => Err(SdkError::invalid_input(format!( + "field `{key}` must be a string" + ))), + } +} + +fn optional_usize(object: &Map, key: &str) -> SdkResult> { + match object.get(key) { + None | Some(Value::Null) => Ok(None), + Some(Value::Number(value)) => value + .as_u64() + .and_then(|number| usize::try_from(number).ok()) + .map(Some) + .ok_or_else(|| SdkError::invalid_input(format!("field `{key}` must be positive"))), + Some(_) => Err(SdkError::invalid_input(format!( + "field `{key}` must be an integer" + ))), + } +} + +fn now_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |duration| { + duration.as_millis().min(u128::from(u64::MAX)) as u64 + }) +} diff --git a/rust/crates/sdk/src/session/reconciliation.rs b/rust/crates/sdk/src/session/reconciliation.rs new file mode 100644 index 00000000..92c7496f --- /dev/null +++ b/rust/crates/sdk/src/session/reconciliation.rs @@ -0,0 +1,673 @@ +//! Durable invocation state for recovering SDK sessions across process restarts. + +use std::{ + fs::{self, OpenOptions}, + io::Write, + path::{Path, PathBuf}, + sync::Mutex, + time::{SystemTime, UNIX_EPOCH}, +}; + +#[cfg(unix)] +use std::fs::File; + +use serde::{Deserialize, Serialize}; +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")] +pub enum InvocationStatus { + Accepted, + Running, + Uncertain, + Terminal, +} + +/// Result captured when an invocation reaches a terminal state. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum InvocationOutcome { + Completed { + message_count: usize, + }, + Aborted { + message_count: usize, + }, + Failed { + message: String, + message_count: usize, + }, +} + +/// Persisted lifecycle and reconciliation data for one invocation. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct InvocationRecord { + pub invocation_id: String, + pub session_id: String, + pub correlation_key: Option, + pub status: InvocationStatus, + pub accepted_at_ms: u64, + pub updated_at_ms: u64, + pub outcome: Option, + pub uncertainty_reason: Option, + pub terminal_event: Option, + pub terminal_scope_revision: Option, +} + +/// File-backed store for invocation lifecycle reconciliation. +#[derive(Debug, Clone)] +pub struct InvocationReconciliationStore { + root: PathBuf, +} + +impl InvocationReconciliationStore { + /// Creates a store handle without creating its root directory. + #[must_use] + pub fn new(root: impl Into) -> Self { + Self { root: root.into() } + } + + /// 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!( + "invocation '{}' already exists", + record.invocation_id + ))); + } + atomic_write_json(&path, &record) + } + + /// 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!( + "invocation '{}' cannot start from state {:?}", + invocation_id, record.status + ))); + } + record.status = InvocationStatus::Running; + record.uncertainty_reason = None; + record.updated_at_ms = updated_at_ms(record.updated_at_ms); + self.write_record(&record) + } + + /// Marks a non-terminal invocation as uncertain with a recovery reason. + pub fn mark_uncertain(&self, invocation_id: &str, reason: impl Into) -> SdkResult<()> { + let reason = reason.into(); + self.update_nonterminal(invocation_id, move |record| { + record.status = InvocationStatus::Uncertain; + record.uncertainty_reason = Some(reason); + }) + } + + /// Records a terminal result. Once terminal, a record cannot be changed again. + pub fn record_terminal( + &self, + invocation_id: &str, + outcome: InvocationOutcome, + terminal_event: Option, + ) -> SdkResult<()> { + self.record_terminal_with_scope_revision(invocation_id, outcome, terminal_event, None) + } + + /// Records a terminal result together with its durable correlation-scope revision. + pub fn record_terminal_with_scope_revision( + &self, + invocation_id: &str, + outcome: InvocationOutcome, + 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!( + "invocation '{}' is already terminal", + 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; + record.terminal_event = terminal_event; + record.terminal_scope_revision = terminal_scope_revision; + record.updated_at_ms = updated_at_ms(record.updated_at_ms); + 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)?; + read_record(&path, invocation_id) + } + + /// Loads all invocation records in stable invocation-ID order. + pub fn list(&self) -> SdkResult> { + let entries = match fs::read_dir(&self.root) { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(error) => { + return Err(io_error( + "read invocation store directory", + &self.root, + error, + )) + } + }; + + let mut records = Vec::new(); + for entry in entries { + let entry = entry + .map_err(|error| io_error("read invocation store entry", &self.root, error))?; + if !entry + .file_type() + .map_err(|error| io_error("inspect invocation store entry", &entry.path(), error))? + .is_file() + { + continue; + } + + let path = entry.path(); + if path.extension() != Some(std::ffi::OsStr::new("json")) { + continue; + } + let invocation_id = path + .file_stem() + .and_then(|stem| stem.to_str()) + .ok_or_else(|| SdkError::invalid_input("invocation filename is not valid UTF-8"))?; + validate_invocation_id(invocation_id)?; + if let Some(record) = read_record(&path, invocation_id)? { + records.push(record); + } + } + records.sort_by(|left, right| left.invocation_id.cmp(&right.invocation_id)); + Ok(records) + } + + /// Loads records associated with the exact correlation key. + pub fn list_by_correlation_key(&self, key: &str) -> SdkResult> { + Ok(self + .list()? + .into_iter() + .filter(|record| record.correlation_key.as_deref() == Some(key)) + .collect()) + } + + /// Durably marks accepted and running records uncertain after a restart. + /// + /// 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!( + record.status, + InvocationStatus::Accepted | InvocationStatus::Running + ) { + record.status = InvocationStatus::Uncertain; + record.uncertainty_reason = + Some("recovered unfinished invocation after restart".into()); + record.updated_at_ms = updated_at_ms(record.updated_at_ms); + self.write_record(&record)?; + recovered.push(record); + } + } + Ok(recovered) + } + + fn record_path(&self, invocation_id: &str) -> SdkResult { + validate_invocation_id(invocation_id)?; + Ok(self.root.join(format!("{invocation_id}.json"))) + } + + fn required_record(&self, invocation_id: &str) -> SdkResult { + self.get(invocation_id)?.ok_or_else(|| { + SdkError::invalid_input(format!("invocation '{invocation_id}' does not exist")) + }) + } + + fn update_nonterminal( + &self, + 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!( + "invocation '{}' is already terminal", + invocation_id + ))); + } + update(&mut record); + record.updated_at_ms = updated_at_ms(record.updated_at_ms); + self.write_record(&record) + } + + fn write_record(&self, record: &InvocationRecord) -> SdkResult<()> { + let path = self.record_path(&record.invocation_id)?; + atomic_write_json(&path, record) + } +} + +/// Atomically serializes a value beside its target and renames it into place. +pub(super) fn atomic_write_json(path: &Path, value: &T) -> SdkResult<()> { + let bytes = serde_json::to_vec(value) + .map_err(|error| SdkError::internal(format!("serialize {}: {error}", path.display())))?; + let parent = path + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + fs::create_dir_all(parent) + .map_err(|error| io_error("create store directory", parent, error))?; + let filename = path + .file_name() + .ok_or_else(|| SdkError::invalid_input("store path has no filename"))?; + let temporary_path = parent.join(format!( + ".{}.{}.tmp", + filename.to_string_lossy(), + Uuid::new_v4() + )); + + let write_result = (|| -> std::io::Result<()> { + let mut temporary = OpenOptions::new() + .write(true) + .create_new(true) + .open(&temporary_path)?; + temporary.write_all(&bytes)?; + temporary.sync_all()?; + drop(temporary); + fs::rename(&temporary_path, path)?; + #[cfg(unix)] + File::open(parent)?.sync_all()?; + Ok(()) + })(); + + match write_result { + Ok(()) => Ok(()), + Err(error) => { + let _ = fs::remove_file(&temporary_path); + Err(io_error("atomically write JSON", path, error)) + } + } +} + +fn read_record(path: &Path, expected_id: &str) -> SdkResult> { + let bytes = match fs::read(path) { + Ok(bytes) => bytes, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(io_error("read invocation record", path, error)), + }; + let record: InvocationRecord = serde_json::from_slice(&bytes).map_err(|error| { + SdkError::internal(format!( + "parse invocation record {}: {error}", + path.display() + )) + })?; + validate_invocation_id(&record.invocation_id)?; + if record.invocation_id != expected_id { + return Err(SdkError::internal(format!( + "invocation record ID '{}' does not match filename ID '{expected_id}'", + record.invocation_id + ))); + } + Ok(Some(record)) +} + +fn validate_invocation_id(invocation_id: &str) -> SdkResult<()> { + if invocation_id.is_empty() + || invocation_id.len() > 200 + || !invocation_id + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'_') + { + return Err(SdkError::invalid_input(format!( + "unsafe invocation ID '{invocation_id}'" + ))); + } + Ok(()) +} + +fn updated_at_ms(previous: u64) -> u64 { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .min(u128::from(u64::MAX)) as u64; + now.max(previous.saturating_add(1)) +} + +fn io_error(action: &str, path: &Path, error: std::io::Error) -> SdkError { + SdkError::internal(format!("{action} '{}': {error}", path.display())) +} + +#[cfg(test)] +mod tests { + use super::*; + + struct TestDirectory(PathBuf); + + impl TestDirectory { + fn new() -> Self { + let path = std::env::temp_dir().join(format!("sdk-session-{}", Uuid::new_v4())); + fs::create_dir(&path).expect("create test directory"); + Self(path) + } + } + + impl Drop for TestDirectory { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } + } + + fn accepted_record(invocation_id: &str) -> InvocationRecord { + InvocationRecord { + invocation_id: invocation_id.into(), + session_id: "session-a".into(), + correlation_key: Some("request-a".into()), + status: InvocationStatus::Accepted, + accepted_at_ms: 10, + updated_at_ms: 10, + outcome: None, + uncertainty_reason: None, + terminal_event: None, + terminal_scope_revision: None, + } + } + + #[test] + fn constructor_has_no_filesystem_side_effects() { + let root = std::env::temp_dir().join(format!("sdk-session-{}", Uuid::new_v4())); + let _store = InvocationReconciliationStore::new(&root); + assert!(!root.exists()); + } + + #[test] + fn invocation_records_survive_store_restart() { + let root = TestDirectory::new(); + let record = accepted_record("invocation-a"); + InvocationReconciliationStore::new(&root.0) + .insert(record.clone()) + .expect("insert record"); + + let loaded = InvocationReconciliationStore::new(&root.0) + .get("invocation-a") + .expect("load record") + .expect("record exists"); + assert_eq!(loaded, record); + } + + #[test] + fn running_transition_only_accepts_an_accepted_invocation() { + let root = TestDirectory::new(); + let store = InvocationReconciliationStore::new(&root.0); + store + .insert(accepted_record("invocation-a")) + .expect("insert"); + store.mark_running("invocation-a").expect("mark running"); + assert_eq!( + store.get("invocation-a").unwrap().unwrap().status, + InvocationStatus::Running + ); + assert!(store.mark_running("invocation-a").is_err()); + store + .mark_uncertain("invocation-a", "recovered") + .expect("mark uncertain"); + assert!(store.mark_running("invocation-a").is_err()); + assert_eq!( + store.get("invocation-a").unwrap().unwrap().status, + InvocationStatus::Uncertain + ); + } + + #[test] + fn uncertain_invocation_can_become_terminal_only_once() { + let root = TestDirectory::new(); + let store = InvocationReconciliationStore::new(&root.0); + store + .insert(accepted_record("invocation-a")) + .expect("insert"); + store + .mark_uncertain("invocation-a", "runtime disconnected") + .expect("mark uncertain"); + store + .record_terminal( + "invocation-a", + InvocationOutcome::Completed { message_count: 2 }, + Some(serde_json::json!({"type": "completed"})), + ) + .expect("record terminal"); + + let terminal = store.get("invocation-a").expect("load").expect("exists"); + assert_eq!(terminal.status, InvocationStatus::Terminal); + assert_eq!(terminal.uncertainty_reason, None); + assert_eq!( + terminal.outcome, + Some(InvocationOutcome::Completed { message_count: 2 }) + ); + assert!(store + .record_terminal( + "invocation-a", + InvocationOutcome::Aborted { message_count: 0 }, + None, + ) + .is_err()); + assert_eq!(store.get("invocation-a").expect("reload"), Some(terminal)); + } + + #[test] + fn recovery_marks_only_accepted_and_running_records_uncertain() { + let root = TestDirectory::new(); + let store = InvocationReconciliationStore::new(&root.0); + store + .insert(accepted_record("accepted")) + .expect("insert accepted"); + let mut running = accepted_record("running"); + running.status = InvocationStatus::Running; + store.insert(running).expect("insert running"); + let mut terminal = accepted_record("terminal"); + terminal.status = InvocationStatus::Terminal; + terminal.outcome = Some(InvocationOutcome::Aborted { message_count: 0 }); + store.insert(terminal.clone()).expect("insert terminal"); + + let recovered = store.reconcile_unfinished().expect("reconcile"); + assert_eq!(recovered.len(), 2); + assert!(recovered + .iter() + .all(|record| record.status == InvocationStatus::Uncertain)); + assert_eq!( + store.get("terminal").expect("load terminal"), + Some(terminal) + ); + assert!(store + .reconcile_unfinished() + .expect("repeat recovery") + .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(); + let store = InvocationReconciliationStore::new(&root.0); + assert!(store.get("../escape").is_err()); + assert!(store.insert(accepted_record("bad/id")).is_err()); + + fs::write(root.0.join("malformed.json"), "{").expect("write malformed record"); + assert!(store.get("malformed").is_err()); + assert!(store.list().is_err()); + } +} diff --git a/rust/crates/sdk/src/session/terminal_scope.rs b/rust/crates/sdk/src/session/terminal_scope.rs new file mode 100644 index 00000000..d15acaf8 --- /dev/null +++ b/rust/crates/sdk/src/session/terminal_scope.rs @@ -0,0 +1,310 @@ +//! Durable mapping from SDK correlation keys to their active terminal session. + +use std::{ + fs, + path::PathBuf, + sync::Mutex, + time::{SystemTime, UNIX_EPOCH}, +}; + +use serde::{Deserialize, Serialize}; + +use crate::error::{SdkError, SdkResult}; + +use super::reconciliation::atomic_write_json; + +const MAX_CORRELATION_KEY_BYTES: usize = 512; +static TERMINAL_SCOPE_WRITE_LOCK: Mutex<()> = Mutex::new(()); + +/// Persisted terminal-session scope for one correlation key. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct TerminalScopeRecord { + pub correlation_key: String, + pub session_id: String, + pub revision: u64, + pub updated_at_ms: u64, +} + +/// File-backed mapping of correlation keys to terminal sessions. +#[derive(Debug, Clone)] +pub struct TerminalScopeStore { + root: PathBuf, +} + +impl TerminalScopeStore { + /// Creates a store handle without creating its root directory. + #[must_use] + pub fn new(root: impl Into) -> Self { + Self { root: root.into() } + } + + /// Sets the session associated with a correlation key. + pub fn set( + &self, + correlation_key: impl Into, + session_id: impl Into, + ) -> SdkResult { + let correlation_key = correlation_key.into(); + validate_correlation_key(&correlation_key)?; + let _guard = TERMINAL_SCOPE_WRITE_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let previous = self.read_record(&correlation_key)?; + let revision = previous + .as_ref() + .map_or(Some(1), |record| record.revision.checked_add(1)) + .ok_or_else(|| SdkError::internal("terminal scope revision is exhausted"))?; + let previous_updated_at = previous.map_or(0, |record| record.updated_at_ms); + let record = TerminalScopeRecord { + correlation_key: correlation_key.clone(), + session_id: session_id.into(), + revision, + updated_at_ms: current_timestamp_ms().max(previous_updated_at.saturating_add(1)), + }; + atomic_write_json(&self.record_path(&correlation_key), &record)?; + Ok(record) + } + + /// Rebuilds a mapping from durable invocation history without moving it backward. + pub fn restore_if_newer( + &self, + correlation_key: &str, + session_id: &str, + revision: u64, + ) -> SdkResult<()> { + validate_correlation_key(correlation_key)?; + let _guard = TERMINAL_SCOPE_WRITE_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let previous = self.read_record(correlation_key)?; + match previous.as_ref() { + Some(record) if record.revision > revision => return Ok(()), + Some(record) if record.revision == revision && record.session_id == session_id => { + return Ok(()) + } + Some(record) if record.revision == revision => { + return Err(SdkError::internal(format!( + "terminal scope revision {revision} conflicts for key '{correlation_key}'" + ))); + } + _ => {} + } + let record = TerminalScopeRecord { + 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) + .saturating_add(1), + ), + }; + atomic_write_json(&self.record_path(correlation_key), &record) + } + + /// Loads the session associated with a correlation key, if present. + pub fn get(&self, correlation_key: &str) -> SdkResult> { + validate_correlation_key(correlation_key)?; + self.read_record(correlation_key) + } + + fn read_record(&self, correlation_key: &str) -> SdkResult> { + let path = self.record_path(correlation_key); + let bytes = match fs::read(&path) { + Ok(bytes) => bytes, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => { + return Err(SdkError::internal(format!( + "read terminal scope '{}': {error}", + path.display() + ))); + } + }; + let record: TerminalScopeRecord = serde_json::from_slice(&bytes).map_err(|error| { + SdkError::internal(format!( + "parse terminal scope '{}': {error}", + path.display() + )) + })?; + if record.correlation_key != correlation_key { + return Err(SdkError::internal(format!( + "terminal scope key does not match requested key for '{}'", + path.display() + ))); + } + Ok(Some(record)) + } + + fn record_path(&self, correlation_key: &str) -> PathBuf { + const MAX_COMPONENT_BYTES: usize = 180; + let encoded = encode_key_as_hex(correlation_key); + let mut path = self.root.clone(); + let mut remaining = encoded.as_str(); + while remaining.len() > MAX_COMPONENT_BYTES { + let (component, tail) = remaining.split_at(MAX_COMPONENT_BYTES); + path.push(component); + remaining = tail; + } + path.push(format!("{remaining}.json")); + path + } +} + +fn validate_correlation_key(correlation_key: &str) -> SdkResult<()> { + if correlation_key.len() <= MAX_CORRELATION_KEY_BYTES { + Ok(()) + } else { + Err(SdkError::invalid_input("correlation key exceeds 512 bytes")) + } +} + +fn encode_key_as_hex(key: &str) -> String { + const HEX_DIGITS: &[u8; 16] = b"0123456789abcdef"; + let mut encoded = String::with_capacity(key.len() * 2); + for byte in key.as_bytes() { + encoded.push(char::from(HEX_DIGITS[usize::from(*byte >> 4)])); + encoded.push(char::from(HEX_DIGITS[usize::from(*byte & 0x0f)])); + } + encoded +} + +fn current_timestamp_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .min(u128::from(u64::MAX)) as u64 +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use uuid::Uuid; + + struct TestDirectory(PathBuf); + + impl TestDirectory { + fn new() -> Self { + let path = std::env::temp_dir().join(format!("sdk-terminal-scope-{}", Uuid::new_v4())); + fs::create_dir(&path).expect("create test directory"); + Self(path) + } + } + + impl Drop for TestDirectory { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } + } + + #[test] + fn terminal_scope_survives_store_restart() { + let root = TestDirectory::new(); + TerminalScopeStore::new(&root.0) + .set("request-1", "session-1") + .expect("set terminal scope"); + + let loaded = TerminalScopeStore::new(&root.0) + .get("request-1") + .expect("load terminal scope") + .expect("scope exists"); + assert_eq!(loaded.correlation_key, "request-1"); + assert_eq!(loaded.session_id, "session-1"); + } + + #[test] + fn scope_revisions_increase_and_recovery_never_moves_backward() { + let root = TestDirectory::new(); + let store = TerminalScopeStore::new(&root.0); + 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) + .expect("stale scope restoration should be ignored"); + assert_eq!( + store.get("request-1").unwrap().unwrap().session_id, + "session-2" + ); + assert!(store + .restore_if_newer("request-1", "other-session", second.revision) + .is_err()); + } + + #[test] + fn encoded_keys_are_isolated_and_cannot_traverse_paths() { + let root = TestDirectory::new(); + let store = TerminalScopeStore::new(&root.0); + store.set("same", "session-a").expect("set first scope"); + store + .set("../same", "session-b") + .expect("set path-like key"); + store + .set("../sämé", "session-c") + .expect("set Unicode path-like key"); + + assert_eq!( + store + .get("same") + .expect("get first") + .expect("exists") + .session_id, + "session-a" + ); + assert_eq!( + store + .get("../same") + .expect("get path-like key") + .expect("exists") + .session_id, + "session-b" + ); + assert_eq!( + store + .get("../sämé") + .expect("get Unicode path-like key") + .expect("exists") + .session_id, + "session-c" + ); + assert_eq!(fs::read_dir(&root.0).expect("read root").count(), 3); + } + + #[test] + fn malformed_terminal_scope_json_is_rejected() { + let root = TestDirectory::new(); + let store = TerminalScopeStore::new(&root.0); + let path = root + .0 + .join(format!("{}.json", encode_key_as_hex("request"))); + fs::write(path, "{").expect("write malformed scope"); + assert!(store.get("request").is_err()); + } + + #[test] + fn long_hex_keys_use_safe_path_components() { + let root = TestDirectory::new(); + let store = TerminalScopeStore::new(&root.0); + let key = "k".repeat(512); + store + .set(&key, "session-long") + .expect("long key should persist"); + assert_eq!( + store.get(&key).expect("get long key").unwrap().session_id, + "session-long" + ); + } + + #[test] + fn oversized_correlation_keys_are_rejected() { + let root = TestDirectory::new(); + let store = TerminalScopeStore::new(&root.0); + let key = "k".repeat(513); + assert!(store.set(&key, "session-long").is_err()); + assert!(store.get(&key).is_err()); + assert!(!root.0.exists() || fs::read_dir(&root.0).unwrap().next().is_none()); + } +} diff --git a/rust/crates/sdk/tests/session_lifecycle.rs b/rust/crates/sdk/tests/session_lifecycle.rs new file mode 100644 index 00000000..f8889be1 --- /dev/null +++ b/rust/crates/sdk/tests/session_lifecycle.rs @@ -0,0 +1,1920 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{mpsc, Arc, Barrier, Mutex}; +use std::thread; +use std::time::Duration; + +use runtime::{ConversationEvent, Session, SessionStore, TurnControlSignal}; +use sdk::error::SdkErrorCode as ErrorCode; +use sdk::session::reconciliation::{InvocationOutcome, InvocationRecord, InvocationStatus}; +use sdk::{ + RuntimeSessionRunner, SdkFrame as Frame, SessionEvent, SessionEventType, SessionRouter, + SessionRunResult, SessionTurnRunner, +}; +use serde_json::{json, Value}; +use tempfile::TempDir; +use uuid::Uuid; + +struct EchoRunner; + +impl SessionTurnRunner for EchoRunner { + fn run( + &self, + mut session: Session, + prompt: String, + _control: TurnControlSignal, + callback: Arc, + ) -> SessionRunResult { + let session_id = session.session_id.clone(); + callback(ConversationEvent::AgentBegin { + session_id: session_id.clone(), + }); + if let Err(error) = session.push_user_text(prompt) { + return SessionRunResult::Failed { + session, + message: error.to_string(), + }; + } + callback(ConversationEvent::ToolUse { + tool_use_id: "tool-1".to_string(), + tool_name: "inspect".to_string(), + input: "{}".to_string(), + }); + callback(ConversationEvent::ToolResult { + tool_use_id: "tool-1".to_string(), + tool_name: "inspect".to_string(), + output: "done".to_string(), + is_error: false, + }); + callback(ConversationEvent::AgentEnd { + session_id, + interrupted: false, + output: None, + }); + SessionRunResult::Completed(session) + } +} + +struct FailingRunner; + +impl SessionTurnRunner for FailingRunner { + fn run( + &self, + session: Session, + _prompt: String, + _control: TurnControlSignal, + _callback: Arc, + ) -> SessionRunResult { + SessionRunResult::Failed { + session, + message: "upstream failed".to_string(), + } + } +} + +struct PanicRunner; + +impl SessionTurnRunner for PanicRunner { + fn run( + &self, + _session: Session, + _prompt: String, + _control: TurnControlSignal, + _callback: Arc, + ) -> SessionRunResult { + panic!("injected runner panic"); + } +} + +struct GateRunner { + started: Mutex>>, + release: Mutex>, +} + +impl SessionTurnRunner for GateRunner { + fn run( + &self, + mut session: Session, + prompt: String, + control: TurnControlSignal, + callback: Arc, + ) -> SessionRunResult { + let session_id = session.session_id.clone(); + callback(ConversationEvent::AgentBegin { + session_id: session_id.clone(), + }); + if let Some(started) = self + .started + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() + { + let _ = started.send(()); + } + let _ = self + .release + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .recv(); + if let Err(error) = session.push_user_text(prompt) { + return SessionRunResult::Failed { + session, + message: error.to_string(), + }; + } + while let Some(steer) = control.take_steer() { + if let Err(error) = session.push_user_text(steer) { + return SessionRunResult::Failed { + session, + message: error.to_string(), + }; + } + } + let _ = control.finish_or_take_steer(); + callback(ConversationEvent::AgentEnd { + session_id, + interrupted: false, + output: None, + }); + SessionRunResult::Completed(session) + } +} + +struct AbortRunner { + calls: AtomicUsize, +} + +impl SessionTurnRunner for AbortRunner { + fn run( + &self, + mut session: Session, + prompt: String, + control: TurnControlSignal, + callback: Arc, + ) -> SessionRunResult { + let call = self.calls.fetch_add(1, Ordering::SeqCst); + let session_id = session.session_id.clone(); + callback(ConversationEvent::AgentBegin { + session_id: session_id.clone(), + }); + if call == 0 { + while !control.is_aborted() { + thread::sleep(Duration::from_millis(1)); + } + callback(ConversationEvent::AgentEnd { + session_id, + interrupted: true, + output: None, + }); + return SessionRunResult::Aborted(session); + } + if let Err(error) = session.push_user_text(prompt) { + return SessionRunResult::Failed { + session, + message: error.to_string(), + }; + } + callback(ConversationEvent::AgentEnd { + session_id, + interrupted: false, + output: None, + }); + SessionRunResult::Completed(session) + } +} + +fn setup() -> (TempDir, SessionStore, SessionRouter) { + setup_with(Arc::new(EchoRunner)) +} + +fn setup_with(runner: Arc) -> (TempDir, SessionStore, SessionRouter) { + let temp = TempDir::new().expect("temporary directory"); + let workspace = temp.path().join("workspace"); + std::fs::create_dir_all(&workspace).expect("workspace directory"); + let store = + SessionStore::from_data_dir(temp.path().join("data"), &workspace).expect("session store"); + let router = SessionRouter::new(store.clone(), runner).expect("session router"); + (temp, store, router) +} + +fn broker_request(operation: &str, input: Value) -> Frame { + Frame::broker_request(Uuid::new_v4(), operation.to_string(), input) +} + +async fn collect_events(mut invocation: sdk::SessionInvocation) -> Vec { + let mut events = Vec::new(); + while let Some(event) = invocation.events.recv().await { + events.push(event); + } + events +} + +async fn collect_stream(mut stream: sdk::SessionEventStream) -> Vec { + let mut events = Vec::new(); + while let Some(event) = stream.recv().await { + events.push(event); + } + events +} + +fn assert_invalid_input(error: sdk::SdkError) { + assert_eq!(error.code, ErrorCode::InvalidInput); +} + +#[test] +fn create_persists_session_for_durable_lookup() { + let (_temp, _store, router) = setup(); + let session = router + .create_session(None, None) + .expect("session should persist"); + let restored = router + .get_session(&session.session_id) + .expect("stored session should resolve by ID"); + assert_eq!(restored.session_id, session.session_id); +} + +#[test] +fn create_persists_model_metadata() { + let (_temp, _store, router) = setup(); + let session = router + .create_session(Some("claude-test"), None) + .expect("session should persist"); + let restored = router + .get_session(&session.session_id) + .expect("stored session should reload"); + assert_eq!(restored.model.as_deref(), Some("claude-test")); +} + +#[test] +fn create_trims_model_name() { + let (_temp, _store, router) = setup(); + let session = router + .create_session(Some(" model-x "), None) + .expect("session should persist"); + assert_eq!(session.model.as_deref(), Some("model-x")); +} + +#[test] +fn create_maps_correlation_key_to_session() { + let (_temp, _store, router) = setup(); + let session = router + .create_session(None, Some("discord:conversation:7")) + .expect("session should persist"); + assert_eq!( + router + .session_for_correlation_key("discord:conversation:7") + .expect("scope read"), + Some(session.session_id) + ); +} + +#[test] +fn create_without_correlation_key_does_not_create_scope() { + let (_temp, _store, router) = setup(); + router + .create_session(None, None) + .expect("session should persist"); + assert_eq!( + router + .session_for_correlation_key("missing") + .expect("scope read"), + None + ); +} + +#[test] +fn create_without_filesystem_side_effects_until_first_session() { + let temp = TempDir::new().expect("temporary directory"); + let workspace = temp.path().join("workspace"); + std::fs::create_dir_all(&workspace).expect("workspace directory"); + let store = + SessionStore::from_data_dir(temp.path().join("data"), &workspace).expect("session store"); + let sdk_root = store.sessions_dir().join("sdk"); + let router = SessionRouter::new(store, Arc::new(EchoRunner)).expect("router"); + assert!(!sdk_root.exists()); + drop(router); + assert!(!sdk_root.exists()); +} + +#[test] +fn session_lookup_returns_not_found_for_unknown_id() { + let (_temp, _store, router) = setup(); + let error = router.get_session("session-0000000000000-99").unwrap_err(); + assert_eq!(error.code, ErrorCode::NotFound); +} + +#[tokio::test] +async fn prompt_publishes_lifecycle_events_in_order() { + let (_temp, _store, router) = setup(); + let session = router + .create_session(None, None) + .expect("session should persist"); + let invocation = router + .prompt(&session.session_id, "hello", None) + .expect("prompt should start"); + let events = collect_events(invocation).await; + let kinds = events + .iter() + .map(|event| event.event_type) + .collect::>(); + assert_eq!( + kinds, + vec![ + SessionEventType::TurnAccepted, + SessionEventType::AgentBegin, + SessionEventType::ToolUse, + SessionEventType::ToolResult, + SessionEventType::AgentEnd, + SessionEventType::TurnEnded, + ] + ); +} + +#[tokio::test] +async fn prompt_assigns_monotonic_event_sequences() { + let (_temp, _store, router) = setup(); + let session = router + .create_session(None, None) + .expect("session should persist"); + let events = collect_events( + router + .prompt(&session.session_id, "hello", None) + .expect("prompt should start"), + ) + .await; + assert_eq!( + events + .iter() + .map(|event| event.sequence) + .collect::>(), + vec![1, 2, 3, 4, 5, 6] + ); + assert!(events + .windows(2) + .all(|pair| pair[0].timestamp_ms <= pair[1].timestamp_ms)); +} + +#[tokio::test] +async fn prompt_persists_durable_terminal_event() { + let (_temp, _store, router) = setup(); + let session = router + .create_session(None, None) + .expect("session should persist"); + let invocation = router + .prompt(&session.session_id, "hello", Some("corr-terminal")) + .expect("prompt should start"); + let invocation_id = invocation.invocation_id.clone(); + let events = collect_events(invocation).await; + let terminal = events.last().expect("terminal event"); + assert_eq!(terminal.event_type, SessionEventType::TurnEnded); + assert!(terminal.durable); + let record = router + .reconcile(&invocation_id) + .expect("terminal record should reconcile"); + assert_eq!(record.status, InvocationStatus::Terminal); + assert_eq!( + record.outcome, + Some(InvocationOutcome::Completed { message_count: 1 }) + ); + assert_eq!( + record.terminal_event, + Some(serde_json::to_value(terminal).unwrap()) + ); +} + +#[tokio::test] +async fn prompt_updates_terminal_scope_when_turn_finishes() { + let (_temp, _store, router) = setup(); + let session = router + .create_session(None, None) + .expect("session should persist"); + let invocation = router + .prompt(&session.session_id, "hello", Some("corr-finish")) + .expect("prompt should start"); + collect_events(invocation).await; + assert_eq!( + router + .session_for_correlation_key("corr-finish") + .expect("scope read"), + Some(session.session_id) + ); +} + +#[tokio::test] +async fn prompt_rejects_second_active_invocation_for_session() { + 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 session = router + .create_session(None, None) + .expect("session should persist"); + let first = router + .prompt(&session.session_id, "first", None) + .expect("first prompt should start"); + started_rx + .recv_timeout(Duration::from_secs(2)) + .expect("runner should start"); + let error = router + .prompt(&session.session_id, "second", None) + .expect_err("parallel invocation should be rejected"); + assert_eq!(error.code, ErrorCode::Conflict); + release_tx.send(()).expect("release runner"); + collect_events(first).await; +} + +#[tokio::test] +async fn steer_is_accepted_for_the_current_active_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 session = router + .create_session(None, None) + .expect("session should persist"); + let invocation = router + .prompt(&session.session_id, "first", None) + .expect("prompt should start"); + started_rx + .recv_timeout(Duration::from_secs(2)) + .expect("runner should start"); + router + .steer(&session.session_id, &invocation.invocation_id, "redirect") + .expect("active invocation should accept steering"); + release_tx.send(()).expect("release runner"); + let events = collect_events(invocation).await; + assert!(events + .iter() + .any(|event| event.event_type == SessionEventType::TurnSteered)); + let session = router + .get_session(&session.session_id) + .expect("session should remain loadable"); + assert_eq!(session.messages.len(), 2); +} + +#[test] +fn steer_rejects_unknown_invocation_id() { + let (_temp, _store, router) = setup(); + let session = router + .create_session(None, None) + .expect("session should persist"); + let error = router + .steer(&session.session_id, &Uuid::new_v4().to_string(), "redirect") + .unwrap_err(); + assert_eq!(error.code, ErrorCode::NotFound); +} + +#[test] +fn steer_rejects_malformed_invocation_id() { + let (_temp, _store, router) = setup(); + let session = router + .create_session(None, None) + .expect("session should persist"); + let error = router + .steer(&session.session_id, "not-a-uuid", "redirect") + .unwrap_err(); + assert_eq!(error.code, ErrorCode::InvalidInput); +} + +#[tokio::test] +async fn abort_and_prompt_persists_old_terminal_outcome_and_starts_new_turn() { + let runner = Arc::new(AbortRunner { + calls: AtomicUsize::new(0), + }); + let (_temp, _store, router) = setup_with(runner); + let session = router + .create_session(None, None) + .expect("session should persist"); + let previous = router + .prompt(&session.session_id, "long-running", None) + .expect("initial invocation"); + let previous_id = previous.invocation_id.clone(); + let next = router + .abort_and_prompt(&session.session_id, "replacement", None) + .await + .expect("abort should finish before replacement starts"); + assert_ne!(next.invocation_id, previous_id); + let old_events = collect_events(previous).await; + assert_eq!( + old_events.last().unwrap().event_type, + SessionEventType::TurnEnded + ); + assert_eq!( + router.reconcile(&previous_id).unwrap().outcome, + Some(InvocationOutcome::Aborted { message_count: 0 }) + ); + let new_events = collect_events(next).await; + assert_eq!( + new_events.last().unwrap().event_type, + SessionEventType::TurnEnded + ); + assert_eq!( + router + .reconcile(&new_events[0].invocation_id) + .unwrap() + .status, + InvocationStatus::Terminal + ); +} + +#[tokio::test] +async fn abort_and_prompt_without_active_turn_starts_prompt() { + let (_temp, _store, router) = setup(); + let session = router + .create_session(None, None) + .expect("session should persist"); + let invocation = router + .abort_and_prompt(&session.session_id, "new prompt", None) + .await + .expect("new prompt should start"); + let events = collect_events(invocation).await; + assert_eq!(events[0].event_type, SessionEventType::TurnAccepted); + assert_eq!( + events.last().unwrap().event_type, + SessionEventType::TurnEnded + ); +} + +#[tokio::test] +async fn runner_failure_is_a_terminal_failed_outcome() { + let (_temp, _store, router) = setup_with(Arc::new(FailingRunner)); + let session = router + .create_session(None, None) + .expect("session should persist"); + let invocation = router + .prompt(&session.session_id, "hello", None) + .expect("prompt should start"); + let invocation_id = invocation.invocation_id.clone(); + let events = collect_events(invocation).await; + assert_eq!( + events.last().unwrap().event_type, + SessionEventType::TurnEnded + ); + assert_eq!( + router.reconcile(&invocation_id).unwrap().outcome, + Some(InvocationOutcome::Failed { + message: "upstream failed".to_string(), + message_count: 0, + }) + ); +} + +#[tokio::test] +async fn runner_panic_marks_invocation_uncertain_for_recovery() { + let (_temp, _store, router) = setup_with(Arc::new(PanicRunner)); + let session = router + .create_session(None, None) + .expect("session should persist"); + let invocation = router + .prompt(&session.session_id, "hello", None) + .expect("prompt should start"); + let invocation_id = invocation.invocation_id.clone(); + let events = collect_events(invocation).await; + assert_eq!(events.len(), 1); + let record = router + .reconcile(&invocation_id) + .expect("record should remain"); + assert_eq!(record.status, InvocationStatus::Uncertain); + assert!(record + .uncertainty_reason + .as_deref() + .is_some_and(|reason| reason.contains("injected runner panic"))); +} + +#[test] +fn list_returns_empty_page_without_sessions() { + let (_temp, _store, router) = setup(); + let page = router.list(None, None).expect("empty page"); + assert!(page.sessions.is_empty()); + assert_eq!(page.next_cursor, None); +} + +#[test] +fn list_paginates_without_duplicate_sessions() { + let (_temp, _store, router) = setup(); + let first = router.create_session(None, None).expect("first session"); + let second = router.create_session(None, None).expect("second session"); + let third = router.create_session(None, None).expect("third session"); + let first_page = router.list(None, Some(2)).expect("first page"); + assert_eq!(first_page.sessions.len(), 2); + let cursor = first_page.next_cursor.expect("next cursor"); + let second_page = router.list(Some(&cursor), Some(2)).expect("second page"); + assert_eq!(second_page.sessions.len(), 1); + assert_eq!(second_page.next_cursor, None); + let mut ids = first_page + .sessions + .iter() + .chain(second_page.sessions.iter()) + .map(|session| session.session_id.clone()) + .collect::>(); + ids.sort(); + let mut expected = vec![first.session_id, second.session_id, third.session_id]; + expected.sort(); + assert_eq!(ids, expected); +} + +#[test] +fn list_rejects_zero_limit() { + let (_temp, _store, router) = setup(); + assert_invalid_input(router.list(None, Some(0)).unwrap_err()); +} + +#[test] +fn list_rejects_limit_above_maximum() { + let (_temp, _store, router) = setup(); + assert_invalid_input(router.list(None, Some(101)).unwrap_err()); +} + +#[test] +fn list_rejects_malformed_cursor() { + let (_temp, _store, router) = setup(); + assert_invalid_input(router.list(Some("offset"), None).unwrap_err()); +} + +#[tokio::test] +async fn dispatch_turn_create_returns_durable_session_id() { + let (_temp, _store, router) = setup(); + let dispatch = router + .dispatch(broker_request( + "turn.create", + json!({ "model": "model-y", "correlationKey": "thread-1" }), + )) + .await + .expect("create dispatch"); + assert!(dispatch.events.is_none()); + let session_id = dispatch.response.result().unwrap()["sessionId"] + .as_str() + .unwrap() + .to_string(); + assert_eq!( + router.get_session(&session_id).unwrap().model.as_deref(), + Some("model-y") + ); + assert_eq!( + router.session_for_correlation_key("thread-1").unwrap(), + Some(session_id) + ); +} + +#[tokio::test] +async fn dispatch_turn_create_with_input_starts_stream() { + let (_temp, _store, router) = setup(); + let dispatch = router + .dispatch(broker_request("turn.create", json!({ "input": "start" }))) + .await + .expect("create dispatch"); + assert!(dispatch.response.result().unwrap()["invocationId"].is_string()); + let mut stream = dispatch.events.expect("event stream"); + let first = stream.recv().await.expect("accepted event"); + assert_eq!(first.event_type, SessionEventType::TurnAccepted); + while stream.recv().await.is_some() {} +} + +#[tokio::test] +async fn dispatch_turn_create_accepts_object_prompt_input() { + let (_temp, _store, router) = setup(); + let dispatch = router + .dispatch(broker_request( + "turn.create", + json!({ "input": { "prompt": "start from object" } }), + )) + .await + .expect("create dispatch"); + let events = collect_stream(dispatch.events.expect("event stream")).await; + assert_eq!( + events.last().unwrap().event_type, + SessionEventType::TurnEnded + ); +} + +#[tokio::test] +async fn dispatch_turn_prompt_returns_events() { + let (_temp, _store, router) = setup(); + let session = router + .create_session(None, None) + .expect("session should persist"); + let dispatch = router + .dispatch(broker_request( + "turn.prompt", + json!({ "sessionId": session.session_id, "prompt": "hello" }), + )) + .await + .expect("prompt dispatch"); + assert!(dispatch.response.result().unwrap()["invocationId"].is_string()); + let events = collect_stream(dispatch.events.expect("event stream")).await; + assert_eq!( + events.last().unwrap().event_type, + SessionEventType::TurnEnded + ); +} + +#[tokio::test] +async fn dispatch_turn_prompt_accepts_request_message_shape() { + let (_temp, _store, router) = setup(); + let session = router.create_session(None, None).unwrap(); + let dispatch = router + .dispatch(broker_request( + "turn.prompt", + json!({ + "sessionId": session.session_id, + "request": { + "messages": [ + { "role": "assistant", "content": "prior context" }, + { "role": "user", "content": "prompt from request" } + ] + }, + }), + )) + .await + .expect("request-shaped prompt should dispatch"); + let events = collect_stream(dispatch.events.expect("event stream")).await; + assert_eq!( + events.last().unwrap().event_type, + SessionEventType::TurnEnded + ); + let restored = router.get_session(&session.session_id).unwrap(); + assert!(restored.messages.iter().any(|message| { + message.role == runtime::MessageRole::User + && message.blocks.iter().any(|block| { + matches!(block, runtime::ContentBlock::Text { text } if text == "prompt from request") + }) + })); +} + +#[tokio::test] +async fn dispatch_turn_prompt_resolves_terminal_scope() { + let (_temp, _store, router) = setup(); + let session = router + .create_session(None, Some("scope-prompt")) + .expect("session should persist"); + let dispatch = router + .dispatch(broker_request( + "turn.prompt", + json!({ "correlationKey": "scope-prompt", "prompt": "hello" }), + )) + .await + .expect("scoped prompt"); + assert_eq!( + dispatch.response.result().unwrap()["sessionId"], + session.session_id + ); + collect_stream(dispatch.events.expect("event stream")).await; +} + +#[tokio::test] +async fn dispatch_turn_steer_accepts_override() { + 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 session = router + .create_session(None, None) + .expect("session should persist"); + let invocation = router + .prompt(&session.session_id, "initial", None) + .expect("initial turn"); + started_rx.recv_timeout(Duration::from_secs(2)).unwrap(); + let response = router + .dispatch(broker_request( + "turn.steer", + json!({ + "sessionId": session.session_id, + "invocationId": invocation.invocation_id, + "override": { "prompt": "change direction" }, + }), + )) + .await + .expect("steer dispatch"); + assert_eq!(response.response.result().unwrap()["accepted"], true); + release_tx.send(()).unwrap(); + collect_events(invocation).await; +} + +#[tokio::test] +async fn dispatch_turn_abort_and_prompt_returns_replacement_invocation() { + let runner = Arc::new(AbortRunner { + calls: AtomicUsize::new(0), + }); + let (_temp, _store, router) = setup_with(runner); + let session = router + .create_session(None, None) + .expect("session should persist"); + let previous = router + .prompt(&session.session_id, "first", None) + .expect("initial invocation"); + let response = router + .dispatch(broker_request( + "turn.abort_and_prompt", + json!({ "sessionId": session.session_id, "prompt": "replacement" }), + )) + .await + .expect("abort and prompt dispatch"); + assert_ne!( + response.response.result().unwrap()["invocationId"], + previous.invocation_id + ); + let old_events = collect_events(previous).await; + assert_eq!( + old_events.last().unwrap().event_type, + SessionEventType::TurnEnded + ); + let new_events = collect_stream(response.events.expect("replacement event stream")).await; + assert_eq!( + new_events.last().unwrap().event_type, + SessionEventType::TurnEnded + ); +} + +#[tokio::test] +async fn dispatch_turn_list_serializes_pagination() { + let (_temp, _store, router) = setup(); + router.create_session(None, None).unwrap(); + let response = router + .dispatch(broker_request("turn.list", json!({ "limit": 1 }))) + .await + .expect("list dispatch"); + let result = response.response.result().unwrap(); + assert_eq!(result["sessions"].as_array().unwrap().len(), 1); + assert!(result["sessions"][0]["sessionId"].is_string()); +} + +#[tokio::test] +async fn dispatch_rejects_unknown_session_operation() { + let (_temp, _store, router) = setup(); + let error = router + .dispatch(broker_request("turn.unknown", json!({}))) + .await + .unwrap_err(); + assert_eq!(error.code, ErrorCode::UnknownOperation); +} + +#[tokio::test] +async fn dispatch_rejects_non_broker_frame() { + let (_temp, _store, router) = setup(); + let error = router.dispatch(Frame::broker_hello()).await.unwrap_err(); + assert_eq!(error.code, ErrorCode::InvalidInput); +} + +#[tokio::test] +async fn dispatch_rejects_missing_request_id() { + let (_temp, _store, router) = setup(); + let mut data = serde_json::Map::new(); + data.insert( + "type".to_string(), + Value::String("broker_request".to_string()), + ); + data.insert( + "operation".to_string(), + Value::String("turn.list".to_string()), + ); + data.insert("input".to_string(), json!({})); + let error = router.dispatch(Frame::new(data)).await.unwrap_err(); + assert_eq!(error.code, ErrorCode::InvalidInput); +} + +#[tokio::test] +async fn dispatch_rejects_non_object_operation_input() { + let (_temp, _store, router) = setup(); + let error = router + .dispatch(broker_request("turn.list", json!("invalid"))) + .await + .unwrap_err(); + assert_eq!(error.code, ErrorCode::InvalidInput); +} + +#[test] +fn reconciliation_returns_not_found_for_unknown_invocation() { + let (_temp, _store, router) = setup(); + let error = router.reconcile(&Uuid::new_v4().to_string()).unwrap_err(); + assert_eq!(error.code, ErrorCode::NotFound); +} + +#[test] +fn startup_recovers_running_invocations_as_uncertain() { + let (_temp, store, router) = setup(); + let session = router + .create_session(None, Some("recover-me")) + .expect("session should persist"); + let invocation_id = Uuid::new_v4().to_string(); + let reconciliation = sdk::InvocationReconciliationStore::new( + store.sessions_dir().join("sdk").join("reconciliation"), + ); + reconciliation + .insert(InvocationRecord { + invocation_id: invocation_id.clone(), + session_id: session.session_id, + correlation_key: Some("recover-me".to_string()), + status: InvocationStatus::Running, + accepted_at_ms: 1, + updated_at_ms: 1, + outcome: None, + uncertainty_reason: None, + terminal_event: None, + terminal_scope_revision: None, + }) + .expect("running record"); + let restarted = SessionRouter::new(store, Arc::new(EchoRunner)).expect("new router"); + let record = restarted.reconcile(&invocation_id).unwrap(); + assert_eq!(record.status, InvocationStatus::Uncertain); + let error = restarted + .session_for_correlation_key("recover-me") + .unwrap_err(); + assert_eq!(error.code, ErrorCode::Unavailable); +} + +#[test] +fn terminal_scope_survives_router_restart() { + let (temp, store, router) = setup(); + let session = router + .create_session(None, Some("restart-scope")) + .expect("session should persist"); + drop(router); + let restarted = SessionRouter::new(store, Arc::new(EchoRunner)).expect("new router"); + assert_eq!( + restarted + .session_for_correlation_key("restart-scope") + .unwrap(), + Some(session.session_id) + ); + assert!(temp.path().exists()); +} + +#[tokio::test] +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 { + started: Mutex::new(Some(started_tx)), + release: Mutex::new(release_rx), + }); + let (_temp, store, router) = setup_with(runner); + let stale_session = router.create_session(None, Some("repair-key")).unwrap(); + let current_session = router.create_session(None, None).unwrap(); + let invocation = router + .prompt(¤t_session.session_id, "complete", Some("repair-key")) + .unwrap(); + let invocation_id = invocation.invocation_id.clone(); + started_rx.recv_timeout(Duration::from_secs(2)).unwrap(); + + let scope_root = store.sessions_dir().join("sdk").join("terminal-scopes"); + std::fs::remove_dir_all(&scope_root).unwrap(); + std::fs::write(&scope_root, "block the scope store").unwrap(); + release_tx.send(()).unwrap(); + + let events = collect_events(invocation).await; + assert_eq!( + 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!( + 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); +} + +#[test] +fn terminal_scope_keys_do_not_collide() { + let (_temp, _store, router) = setup(); + let first = router + .create_session(None, Some("key-one")) + .expect("first session"); + let second = router + .create_session(None, Some("key-two")) + .expect("second session"); + assert_ne!(first.session_id, second.session_id); + assert_eq!( + router.session_for_correlation_key("key-one").unwrap(), + Some(first.session_id) + ); + assert_eq!( + router.session_for_correlation_key("key-two").unwrap(), + Some(second.session_id) + ); +} + +#[tokio::test] +async fn correlation_lookup_rejects_multiple_active_sessions() { + let (started_tx, started_rx) = mpsc::sync_channel(1); + let (release_tx, release_rx) = mpsc::sync_channel(2); + let runner = Arc::new(GateRunner { + started: Mutex::new(Some(started_tx)), + release: Mutex::new(release_rx), + }); + let (_temp, _store, router) = setup_with(runner); + let first = router.create_session(None, Some("active-key")).unwrap(); + let second = router.create_session(None, Some("active-key")).unwrap(); + let first_invocation = router + .prompt(&first.session_id, "first", Some("active-key")) + .unwrap(); + let second_invocation = router + .prompt(&second.session_id, "second", Some("active-key")) + .unwrap(); + started_rx.recv_timeout(Duration::from_secs(2)).unwrap(); + + assert_eq!( + router + .session_for_correlation_key("active-key") + .unwrap_err() + .code, + ErrorCode::Conflict + ); + release_tx.send(()).unwrap(); + release_tx.send(()).unwrap(); + collect_events(first_invocation).await; + 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); + 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 cloned = router.clone(); + let session = router.create_session(None, None).unwrap(); + let invocation = router.prompt(&session.session_id, "hello", None).unwrap(); + started_rx.recv_timeout(Duration::from_secs(2)).unwrap(); + assert_eq!( + cloned + .prompt(&session.session_id, "parallel", None) + .unwrap_err() + .code, + ErrorCode::Conflict + ); + release_tx.send(()).unwrap(); + collect_events(invocation).await; +} + +#[test] +fn reconciliation_terminal_outcome_retains_correlation_key() { + let (_temp, _store, router) = setup(); + let session = router.create_session(None, None).unwrap(); + let invocation = router + .prompt(&session.session_id, "hello", Some("corr-record")) + .unwrap(); + let invocation_id = invocation.invocation_id.clone(); + drop(invocation); + // The runner is intentionally asynchronous; its durable record is awaited by polling below. + for _ in 0..200 { + if router + .reconcile(&invocation_id) + .is_ok_and(|record| record.status == InvocationStatus::Terminal) + { + break; + } + thread::sleep(Duration::from_millis(2)); + } + let record = router.reconcile(&invocation_id).unwrap(); + assert_eq!(record.correlation_key.as_deref(), Some("corr-record")); +} + +#[tokio::test] +async fn reconciliation_can_resume_all_invocations_for_a_correlation_key() { + let (_temp, _store, router) = setup(); + let session = router.create_session(None, None).unwrap(); + let invocation = router + .prompt(&session.session_id, "hello", Some("corr-history")) + .unwrap(); + let invocation_id = invocation.invocation_id.clone(); + collect_events(invocation).await; + let records = router + .reconcile_by_correlation_key("corr-history") + .expect("correlation reconciliation"); + assert_eq!(records.len(), 1); + assert_eq!(records[0].invocation_id, invocation_id); + assert_eq!(records[0].status, InvocationStatus::Terminal); +} + +#[test] +fn create_rejects_model_longer_than_limit() { + let (_temp, _store, router) = setup(); + let error = router + .create_session(Some(&"m".repeat(257)), None) + .unwrap_err(); + assert_eq!(error.code, ErrorCode::InvalidInput); +} + +macro_rules! malformed_session_id_tests { + ($($name:ident => $session_id:expr),+ $(,)?) => { + $( + #[test] + fn $name() { + let (_temp, _store, router) = setup(); + assert_invalid_input(router.get_session($session_id).unwrap_err()); + } + )+ + }; +} + +malformed_session_id_tests! { + malformed_session_id_empty => "", + malformed_session_id_dot => ".", + malformed_session_id_dot_dot => "..", + malformed_session_id_forward_slash => "folder/session", + malformed_session_id_parent_traversal => "../outside", + malformed_session_id_absolute_path => "/tmp/session", + malformed_session_id_windows_path => "C:\\temp\\session", + malformed_session_id_backslash => "folder\\session", + malformed_session_id_space => "session id", + malformed_session_id_unicode => "sëssion", + malformed_session_id_nul => "session\0id", + malformed_session_id_too_long => "s".repeat(129).as_str(), +} + +macro_rules! malformed_invocation_id_tests { + ($($name:ident => $invocation_id:expr),+ $(,)?) => { + $( + #[test] + fn $name() { + let (_temp, _store, router) = setup(); + assert_invalid_input(router.reconcile($invocation_id).unwrap_err()); + } + )+ + }; +} + +malformed_invocation_id_tests! { + malformed_invocation_id_empty => "", + malformed_invocation_id_path => "../invocation", + malformed_invocation_id_short => "123", + malformed_invocation_id_invalid_character => "00000000-0000-0000-0000-00000000000g", + malformed_invocation_id_braces => "{00000000-0000-0000-0000-000000000000}", + malformed_invocation_id_whitespace => " 00000000-0000-0000-0000-000000000000", + malformed_invocation_id_extra_suffix => "00000000-0000-0000-0000-000000000000.json", + malformed_invocation_id_too_many_sections => "00000000-0000-0000-0000-000000000000-0000", +} + +macro_rules! invalid_prompt_tests { + ($($name:ident => $prompt:expr),+ $(,)?) => { + $( + #[test] + fn $name() { + let (_temp, _store, router) = setup(); + let session = router.create_session(None, None).unwrap(); + assert_invalid_input( + router.prompt(&session.session_id, $prompt, None).unwrap_err(), + ); + } + )+ + }; +} + +invalid_prompt_tests! { + prompt_rejects_empty_text => "", + prompt_rejects_spaces_only => " ", + prompt_rejects_tabs_only => "\t\t", + prompt_rejects_newlines_only => "\n", + prompt_rejects_oversized_input => "x".repeat(1_000_001).as_str(), +} + +macro_rules! invalid_correlation_key_tests { + ($($name:ident => $key:expr),+ $(,)?) => { + $( + #[test] + fn $name() { + let (_temp, _store, router) = setup(); + assert_invalid_input( + router.session_for_correlation_key($key).unwrap_err(), + ); + } + )+ + }; +} + +invalid_correlation_key_tests! { + correlation_key_too_long_513_bytes => "k".repeat(513).as_str(), + correlation_key_too_long_1024_bytes => "k".repeat(1024).as_str(), +} + +#[tokio::test] +async fn dispatch_rejects_missing_prompt_field() { + let (_temp, _store, router) = setup(); + let session = router.create_session(None, None).unwrap(); + let error = router + .dispatch(broker_request( + "turn.prompt", + json!({ "sessionId": session.session_id }), + )) + .await + .unwrap_err(); + assert_eq!(error.code, ErrorCode::InvalidInput); +} + +#[tokio::test] +async fn dispatch_rejects_non_string_model() { + let (_temp, _store, router) = setup(); + let error = router + .dispatch(broker_request("turn.create", json!({ "model": 42 }))) + .await + .unwrap_err(); + assert_eq!(error.code, ErrorCode::InvalidInput); +} + +#[tokio::test] +async fn dispatch_rejects_non_string_correlation_key() { + let (_temp, _store, router) = setup(); + let error = router + .dispatch(broker_request( + "turn.create", + json!({ "correlationKey": [] }), + )) + .await + .unwrap_err(); + assert_eq!(error.code, ErrorCode::InvalidInput); +} + +#[tokio::test] +async fn dispatch_rejects_malformed_list_limit() { + let (_temp, _store, router) = setup(); + let error = router + .dispatch(broker_request("turn.list", json!({ "limit": "many" }))) + .await + .unwrap_err(); + assert_eq!(error.code, ErrorCode::InvalidInput); +} + +#[test] +fn invocation_terminal_outcome_serializes_as_durable_status() { + let outcome = InvocationOutcome::Completed { message_count: 3 }; + let json = serde_json::to_value(outcome).expect("outcome should serialize"); + assert_eq!(json["completed"]["message_count"], 3); +} + +#[tokio::test] +async fn prompt_completion_does_not_leave_session_busy() { + let (_temp, _store, router) = setup(); + let session = router.create_session(None, None).unwrap(); + let first = router.prompt(&session.session_id, "first", None).unwrap(); + let first_id = first.invocation_id.clone(); + collect_events(first).await; + let second = router.prompt(&session.session_id, "second", None).unwrap(); + assert_ne!(first_id, second.invocation_id); + collect_events(second).await; +} + +#[tokio::test] +async fn event_payloads_carry_tool_identity_and_output() { + let (_temp, _store, router) = setup(); + let session = router.create_session(None, None).unwrap(); + let events = collect_events(router.prompt(&session.session_id, "hello", None).unwrap()).await; + let tool_use = events + .iter() + .find(|event| event.event_type == SessionEventType::ToolUse) + .unwrap(); + assert_eq!(tool_use.payload["toolName"], "inspect"); + let tool_result = events + .iter() + .find(|event| event.event_type == SessionEventType::ToolResult) + .unwrap(); + assert_eq!(tool_result.payload["output"], "done"); +} + +#[tokio::test] +async fn accepted_event_is_transient_while_terminal_event_is_durable() { + let (_temp, _store, router) = setup(); + let session = router.create_session(None, None).unwrap(); + let events = collect_events(router.prompt(&session.session_id, "hello", None).unwrap()).await; + assert!(!events.first().unwrap().durable); + assert!(events.last().unwrap().durable); +} + +#[test] +fn store_root_is_scoped_to_the_workspace_session_namespace() { + let (_temp, store, _router) = setup(); + assert!(store.sessions_dir().to_string_lossy().contains("sessions")); +} + +#[tokio::test] +async fn turn_list_accepts_cursor_at_end_of_collection() { + let (_temp, _store, router) = setup(); + router.create_session(None, None).unwrap(); + let page = router.list(Some("1"), Some(1)).unwrap(); + assert!(page.sessions.is_empty()); + assert_eq!(page.next_cursor, None); +} + +#[test] +fn turn_list_accepts_maximum_page_size() { + let (_temp, _store, router) = setup(); + assert!(router.list(None, Some(100)).is_ok()); +} + +#[test] +fn turn_list_rejects_overflowing_cursor() { + let (_temp, _store, router) = setup(); + assert_invalid_input( + router + .list(Some("184467440737095516160"), None) + .unwrap_err(), + ); +} + +#[test] +fn terminal_scope_lookup_rejects_oversized_key() { + let (_temp, _store, router) = setup(); + assert_invalid_input( + router + .session_for_correlation_key(&"x".repeat(513)) + .unwrap_err(), + ); +} + +#[test] +fn new_router_recovers_only_unfinished_invocations() { + let (_temp, store, router) = setup(); + let session = router.create_session(None, None).unwrap(); + let store_path = store.sessions_dir().join("sdk").join("reconciliation"); + let reconciliation = sdk::InvocationReconciliationStore::new(store_path); + let terminal_id = Uuid::new_v4().to_string(); + reconciliation + .insert(InvocationRecord { + invocation_id: terminal_id.clone(), + session_id: session.session_id.clone(), + correlation_key: Some("restored-terminal-scope".to_string()), + status: InvocationStatus::Terminal, + accepted_at_ms: 1, + updated_at_ms: 2, + outcome: Some(InvocationOutcome::Completed { message_count: 0 }), + uncertainty_reason: None, + terminal_event: Some(json!({ "type": "turn_ended" })), + terminal_scope_revision: Some(1), + }) + .unwrap(); + let restarted = SessionRouter::new(store, Arc::new(EchoRunner)).unwrap(); + assert_eq!( + restarted.reconcile(&terminal_id).unwrap().status, + InvocationStatus::Terminal + ); + assert_eq!( + restarted + .session_for_correlation_key("restored-terminal-scope") + .unwrap(), + Some(session.session_id) + ); +} + +#[test] +fn terminal_scope_recovery_uses_revision_when_timestamps_tie() { + let (_temp, store, router) = setup(); + let first = router.create_session(None, Some("revision-tie")).unwrap(); + let latest = router.create_session(None, None).unwrap(); + let reconciliation = sdk::InvocationReconciliationStore::new( + store.sessions_dir().join("sdk").join("reconciliation"), + ); + for (invocation_id, session_id, revision) in [ + ("invocation-a", first.session_id.as_str(), 1), + ("invocation-z", latest.session_id.as_str(), 2), + ] { + reconciliation + .insert(InvocationRecord { + invocation_id: invocation_id.to_string(), + session_id: session_id.to_string(), + correlation_key: Some("revision-tie".to_string()), + status: InvocationStatus::Terminal, + accepted_at_ms: 10, + updated_at_ms: 20, + outcome: Some(InvocationOutcome::Completed { message_count: 1 }), + uncertainty_reason: None, + terminal_event: Some(json!({ "type": "turn_ended" })), + terminal_scope_revision: Some(revision), + }) + .unwrap(); + } + + let restarted = SessionRouter::new(store, Arc::new(EchoRunner)).unwrap(); + assert_eq!( + restarted + .session_for_correlation_key("revision-tie") + .unwrap(), + Some(latest.session_id) + ); +} + +#[test] +fn runtime_session_runner_type_is_available_for_runtime_integration() { + struct EmptyApi; + impl runtime::ApiClient for EmptyApi { + fn stream( + &mut self, + _request: runtime::ApiRequest, + ) -> Result, runtime::RuntimeError> { + Ok(vec![runtime::AssistantEvent::TextDelta("done".to_string())]) + } + } + let _runner: RuntimeSessionRunner = + RuntimeSessionRunner::new(|session| { + runtime::ConversationRuntime::new( + session, + EmptyApi, + runtime::StaticToolExecutor::default(), + runtime::PermissionPolicy::new(runtime::PermissionMode::ReadOnly), + Vec::new(), + ) + }); +} + +#[tokio::test] +async fn runtime_runner_streams_real_tool_lifecycle_through_the_router() { + struct ToolApi { + calls: usize, + } + + impl runtime::ApiClient for ToolApi { + fn stream( + &mut self, + _request: runtime::ApiRequest, + ) -> Result, runtime::RuntimeError> { + self.calls += 1; + if self.calls == 1 { + Ok(vec![ + runtime::AssistantEvent::ToolUse { + id: "call-1".to_string(), + name: "echo".to_string(), + input: "payload".to_string(), + }, + runtime::AssistantEvent::MessageStop, + ]) + } else { + Ok(vec![ + runtime::AssistantEvent::TextDelta("complete".to_string()), + runtime::AssistantEvent::MessageStop, + ]) + } + } + } + + let runner = RuntimeSessionRunner::new(|session| { + runtime::ConversationRuntime::new( + session, + ToolApi { calls: 0 }, + runtime::StaticToolExecutor::new().register("echo", |input| Ok(input.to_string())), + runtime::PermissionPolicy::new(runtime::PermissionMode::DangerFullAccess), + Vec::new(), + ) + }); + let (_temp, store, router) = setup_with(Arc::new(runner)); + let session = router.create_session(None, None).unwrap(); + let invocation = router + .prompt(&session.session_id, "run tool", None) + .unwrap(); + let invocation_id = invocation.invocation_id.clone(); + let events = collect_events(invocation).await; + assert_eq!( + events + .iter() + .map(|event| event.event_type) + .collect::>(), + vec![ + SessionEventType::TurnAccepted, + SessionEventType::AgentBegin, + SessionEventType::ToolUse, + SessionEventType::ToolResult, + SessionEventType::AgentEnd, + SessionEventType::TurnEnded, + ] + ); + 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 }) + ); + assert_eq!( + Session::load_from_path(store.create_handle(&session.session_id).path) + .unwrap() + .messages + .len(), + 4 + ); +} + +#[tokio::test] +async fn router_steer_reaches_an_active_runtime_before_its_next_model_request() { + struct BlockingApi { + calls: usize, + entered: Arc>>>, + release: Arc>>, + } + + impl runtime::ApiClient for BlockingApi { + fn stream( + &mut self, + request: runtime::ApiRequest, + ) -> Result, runtime::RuntimeError> { + self.calls += 1; + if self.calls == 1 { + if let Some(entered) = self + .entered + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() + { + let _ = entered.send(()); + } + self.release + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .recv() + .map_err(|error| runtime::RuntimeError::new(error.to_string()))?; + return Ok(vec![ + runtime::AssistantEvent::TextDelta("first response".to_string()), + runtime::AssistantEvent::MessageStop, + ]); + } + + let steer_is_in_history = request.messages.iter().any(|message| { + message.role == runtime::MessageRole::User + && message.blocks.iter().any(|block| { + matches!(block, runtime::ContentBlock::Text { text } if text == "midflight steer") + }) + }); + if !steer_is_in_history { + return Err(runtime::RuntimeError::new( + "steering prompt was missing from the next API request", + )); + } + Ok(vec![ + runtime::AssistantEvent::TextDelta("steered response".to_string()), + runtime::AssistantEvent::MessageStop, + ]) + } + } + + let (entered_tx, entered_rx) = mpsc::sync_channel(1); + let (release_tx, release_rx) = mpsc::sync_channel(1); + let entered = Arc::new(Mutex::new(Some(entered_tx))); + let release = Arc::new(Mutex::new(release_rx)); + let runner = RuntimeSessionRunner::new({ + let entered = Arc::clone(&entered); + let release = Arc::clone(&release); + move |session| { + runtime::ConversationRuntime::new( + session, + BlockingApi { + calls: 0, + entered: Arc::clone(&entered), + release: Arc::clone(&release), + }, + runtime::StaticToolExecutor::default(), + runtime::PermissionPolicy::new(runtime::PermissionMode::ReadOnly), + Vec::new(), + ) + } + }); + let (_temp, _store, router) = setup_with(Arc::new(runner)); + let session = router.create_session(None, None).unwrap(); + let invocation = router + .prompt(&session.session_id, "initial prompt", None) + .unwrap(); + let invocation_id = invocation.invocation_id.clone(); + entered_rx + .recv_timeout(Duration::from_secs(2)) + .expect("first model request should block"); + router + .steer(&session.session_id, &invocation_id, "midflight steer") + .expect("active model request should accept steering"); + release_tx.send(()).expect("release first model request"); + + let events = collect_events(invocation).await; + assert_eq!( + events.last().unwrap().event_type, + SessionEventType::TurnEnded + ); + assert!(events + .iter() + .any(|event| event.event_type == SessionEventType::TurnSteered)); + assert_eq!( + router.reconcile(&invocation_id).unwrap().status, + InvocationStatus::Terminal + ); + assert!(router + .get_session(&session.session_id) + .unwrap() + .messages + .iter() + .any(|message| { + message.role == runtime::MessageRole::User + && message.blocks.iter().any(|block| { + matches!(block, runtime::ContentBlock::Text { text } if text == "midflight steer") + }) + })); +} + +macro_rules! more_invalid_page_cursors { + ($($name:ident => $cursor:expr),+ $(,)?) => { + $( + #[test] + fn $name() { + let (_temp, _store, router) = setup(); + assert_invalid_input(router.list(Some($cursor), None).unwrap_err()); + } + )+ + }; +} + +more_invalid_page_cursors! { + turn_list_rejects_alphabetic_cursor => "next", + turn_list_rejects_negative_cursor => "-1", + turn_list_rejects_fractional_cursor => "1.5", + turn_list_rejects_blank_cursor => "", +} + +#[tokio::test] +async fn invocation_events_use_the_same_id_as_reconciliation_records() { + let (_temp, _store, router) = setup(); + let session = router.create_session(None, None).unwrap(); + let events = collect_events(router.prompt(&session.session_id, "hello", None).unwrap()).await; + let invocation_id = events[0].invocation_id.clone(); + assert!(events + .iter() + .all(|event| event.invocation_id == invocation_id)); + assert_eq!( + router.reconcile(&invocation_id).unwrap().invocation_id, + invocation_id + ); +} + +#[test] +fn create_rejects_oversized_correlation_key_before_writing_session() { + let (_temp, store, router) = setup(); + let error = router + .create_session(None, Some(&"c".repeat(513))) + .unwrap_err(); + assert_eq!(error.code, ErrorCode::InvalidInput); + assert!(store.list_sessions().unwrap().is_empty()); +} + +#[tokio::test] +async fn turn_prompt_terminal_outcome_is_available_after_event_receiver_drops() { + let (_temp, _store, router) = setup(); + let session = router.create_session(None, None).unwrap(); + let invocation = router.prompt(&session.session_id, "hello", None).unwrap(); + let invocation_id = invocation.invocation_id.clone(); + drop(invocation); + 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; + } + }) + .await + .expect("invocation should eventually become terminal"); + assert_eq!(terminal.status, InvocationStatus::Terminal); +} + +#[tokio::test] +async fn turn_prompt_event_stream_closes_after_terminal_event() { + let (_temp, _store, router) = setup(); + let session = router.create_session(None, None).unwrap(); + let events = collect_events(router.prompt(&session.session_id, "hello", None).unwrap()).await; + assert_eq!( + events.last().unwrap().event_type, + SessionEventType::TurnEnded + ); +} + +#[tokio::test] +async fn abort_marks_terminal_event_durable() { + let runner = Arc::new(AbortRunner { + calls: AtomicUsize::new(0), + }); + let (_temp, _store, router) = setup_with(runner); + let session = router.create_session(None, None).unwrap(); + let previous = router.prompt(&session.session_id, "long", None).unwrap(); + let previous_id = previous.invocation_id.clone(); + let replacement = router + .abort_and_prompt(&session.session_id, "next", None) + .await + .unwrap(); + let old_events = collect_events(previous).await; + let terminal = old_events.last().unwrap(); + assert_eq!(terminal.event_type, SessionEventType::TurnEnded); + assert!(terminal.durable); + assert_eq!( + router.reconcile(&previous_id).unwrap().outcome, + Some(InvocationOutcome::Aborted { message_count: 0 }) + ); + drop(replacement); +} + +#[tokio::test] +async fn prompt_session_messages_are_durable_before_terminal_event() { + let (_temp, _store, router) = setup(); + let session = router.create_session(None, None).unwrap(); + let invocation = router + .prompt(&session.session_id, "persist me", None) + .unwrap(); + collect_events(invocation).await; + let loaded = router.get_session(&session.session_id).unwrap(); + assert_eq!(loaded.messages.len(), 1); +} + +#[test] +fn session_id_lookup_rejects_uuid_that_is_not_a_persisted_session() { + let (_temp, _store, router) = setup(); + assert_eq!( + router + .get_session(&Uuid::new_v4().to_string()) + .unwrap_err() + .code, + ErrorCode::NotFound + ); +} + +#[tokio::test] +async fn turn_create_empty_prompt_is_rejected_without_orphaning_session() { + let (_temp, store, router) = setup(); + let error = router + .dispatch(broker_request("turn.create", json!({ "input": " " }))) + .await + .unwrap_err(); + assert_eq!(error.code, ErrorCode::InvalidInput); + assert!(store.list_sessions().unwrap().is_empty()); +} + +#[tokio::test] +async fn turn_steer_event_is_sequenced_after_acceptance() { + 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 session = router.create_session(None, None).unwrap(); + let invocation = router.prompt(&session.session_id, "initial", None).unwrap(); + started_rx.recv_timeout(Duration::from_secs(2)).unwrap(); + router + .steer(&session.session_id, &invocation.invocation_id, "override") + .unwrap(); + release_tx.send(()).unwrap(); + let events = collect_events(invocation).await; + assert!(events[0].sequence < events[1].sequence); + let steer_index = events + .iter() + .position(|event| event.event_type == SessionEventType::TurnSteered) + .unwrap(); + assert!(steer_index > 0); +} + +#[tokio::test] +async fn concurrent_steer_and_completion_preserve_sequence_and_terminal_order() { + for _ in 0..16 { + 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 session = router.create_session(None, None).unwrap(); + let invocation = router.prompt(&session.session_id, "initial", None).unwrap(); + let invocation_id = invocation.invocation_id.clone(); + started_rx.recv_timeout(Duration::from_secs(2)).unwrap(); + + let barrier = Arc::new(Barrier::new(3)); + let steer_router = router.clone(); + let steer_session_id = session.session_id.clone(); + let steer_invocation_id = invocation_id.clone(); + let steer_barrier = Arc::clone(&barrier); + let steer_thread = thread::spawn(move || { + steer_barrier.wait(); + steer_router + .steer(&steer_session_id, &steer_invocation_id, "racing steer") + .is_ok() + }); + let finish_barrier = Arc::clone(&barrier); + let finish_thread = thread::spawn(move || { + finish_barrier.wait(); + release_tx.send(()).unwrap(); + }); + barrier.wait(); + + let steer_accepted = steer_thread.join().unwrap(); + finish_thread.join().unwrap(); + let events = collect_events(invocation).await; + assert!(events + .windows(2) + .all(|pair| pair[0].sequence < pair[1].sequence)); + assert_eq!( + events.last().unwrap().event_type, + SessionEventType::TurnEnded + ); + let steer_index = events + .iter() + .position(|event| event.event_type == SessionEventType::TurnSteered); + assert_eq!(steer_accepted, steer_index.is_some()); + if let Some(index) = steer_index { + assert!(index < events.len() - 1); + } + } +} + +#[test] +fn session_page_types_round_trip_as_json() { + let page = sdk::SessionPage { + sessions: vec![sdk::SessionSummary { + session_id: "session-1".to_string(), + created_at_ms: 1, + updated_at_ms: 2, + message_count: 3, + }], + next_cursor: Some("1".to_string()), + }; + let encoded = serde_json::to_value(&page).unwrap(); + let decoded: sdk::SessionPage = serde_json::from_value(encoded).unwrap(); + assert_eq!(decoded, page); +} + +#[test] +fn session_event_type_uses_protocol_snake_case() { + assert_eq!( + serde_json::to_string(&SessionEventType::TurnAccepted).unwrap(), + "\"turn_accepted\"" + ); +}