diff --git a/src-rust/crates/cli/src/main.rs b/src-rust/crates/cli/src/main.rs index 4b7d798..e8f6cc7 100644 --- a/src-rust/crates/cli/src/main.rs +++ b/src-rust/crates/cli/src/main.rs @@ -2009,6 +2009,72 @@ fn permission_request_from_core( } } +#[derive(Debug, Clone)] +struct GoalTurnBaseline { + goal_id: String, + tracker_total: u64, + started_at: std::time::Instant, +} + +#[derive(Debug, PartialEq, Eq)] +enum GoalTurnAccountingDecision { + Skip, + Record { total_tokens_used: u64 }, +} + +fn capture_goal_turn_baseline( + session_id: &str, + tracker_total: u64, +) -> Result, claurst_core::GoalError> { + let store = claurst_core::GoalStore::open(&claurst_core::GoalStore::default_path())?; + let goal = store.try_get_goal(session_id)?; + Ok(goal + .filter(|goal| goal.status == claurst_core::GoalStatus::Active) + .map(|goal| GoalTurnBaseline { + goal_id: goal.id, + tracker_total, + started_at: std::time::Instant::now(), + })) +} + +fn decide_goal_turn_accounting( + baseline: Option<&GoalTurnBaseline>, + current_goal: Option<&claurst_core::Goal>, + tracker_total: u64, +) -> Result { + match (baseline, current_goal) { + (None, None) => Ok(GoalTurnAccountingDecision::Skip), + (None, Some(goal)) if goal.status != claurst_core::GoalStatus::Active => { + Ok(GoalTurnAccountingDecision::Skip) + } + (None, Some(_)) => Err("missing active-goal turn baseline".to_string()), + (Some(_), None) => { + Err("goal changed during the turn: original goal is missing".to_string()) + } + (Some(baseline), Some(goal)) if baseline.goal_id != goal.id => Err(format!( + "goal changed during the turn: expected {}, found {}", + baseline.goal_id, goal.id + )), + (Some(_), Some(goal)) if goal.status != claurst_core::GoalStatus::Active => { + Ok(GoalTurnAccountingDecision::Skip) + } + (Some(baseline), Some(goal)) => { + let turn_tokens = tracker_total + .checked_sub(baseline.tracker_total) + .ok_or_else(|| { + format!( + "goal token tracker moved backwards: start {}, end {}", + baseline.tracker_total, tracker_total + ) + })?; + let total_tokens_used = goal.tokens_used.checked_add(turn_tokens).ok_or_else(|| { + "goal token accounting overflow while adding completed turn".to_string() + })?; + Ok(GoalTurnAccountingDecision::Record { total_tokens_used }) + } + } +} + // Allowed: this is the main interactive-mode bootstrap. Its inputs are all // distinct, separately-constructed services (config, settings, HTTP client, // tool registry, tool execution context, query loop knobs, etc.) and @@ -2377,8 +2443,10 @@ async fn run_interactive( // Active effort level (None = use model default / High). // Tracks the user's /effort selection; flows into qcfg each turn. let mut current_effort: Option = None; - // Timestamp of when the most recent query turn was dispatched (for goal elapsed tracking). - let mut goal_turn_start: std::time::Instant = std::time::Instant::now(); + // Start time and session-wide tracker baseline for the most recent real + // query turn. Goal accounting converts the tracker delta into a durable, + // goal-absolute token total only after a successful EndTurn. + let mut goal_turn_baseline: Option = None; // Background update check: spawned once at startup; result delivered via channel. let (update_tx, mut update_rx) = tokio::sync::mpsc::channel::>(1); @@ -3286,7 +3354,23 @@ async fn run_interactive( let tracker = cost_tracker.clone(); let tx = event_tx.clone(); let client_clone = client.clone(); - goal_turn_start = std::time::Instant::now(); + goal_turn_baseline = if claurst_core::goals_enabled() { + match capture_goal_turn_baseline( + &session.id, + cost_tracker.total_tokens(), + ) { + Ok(baseline) => baseline, + Err(error) => { + app.status_message = Some(format!( + "Goal error: could not capture turn baseline: {}", + error + )); + None + } + } + } else { + None + }; let handle = tokio::spawn(async move { let mut msgs = msgs_arc_clone.lock().await.clone(); @@ -3650,6 +3734,7 @@ async fn run_interactive( claurst_api::effective_model_for_config(&cmd_ctx.config, &model_registry); let client_clone = client.clone(); app.is_streaming = true; + goal_turn_baseline = None; let handle = tokio::spawn(async move { if ct.is_cancelled() { @@ -3813,6 +3898,23 @@ async fn run_interactive( let tracker = cost_tracker.clone(); let tx = event_tx.clone(); let client_clone = client.clone(); + goal_turn_baseline = if claurst_core::goals_enabled() { + match capture_goal_turn_baseline( + &session.id, + cost_tracker.total_tokens(), + ) { + Ok(baseline) => baseline, + Err(error) => { + app.status_message = Some(format!( + "Goal error: could not capture turn baseline: {}", + error + )); + None + } + } + } else { + None + }; let handle = tokio::spawn(async move { let mut msgs = msgs_arc_clone.lock().await.clone(); let outcome = claurst_query::run_query_loop( @@ -3925,6 +4027,20 @@ async fn run_interactive( let tracker = cost_tracker.clone(); let tx = event_tx.clone(); let client_clone = client.clone(); + goal_turn_baseline = if claurst_core::goals_enabled() { + match capture_goal_turn_baseline(&session.id, cost_tracker.total_tokens()) { + Ok(baseline) => baseline, + Err(error) => { + app.status_message = Some(format!( + "Goal error: could not capture turn baseline: {}", + error + )); + None + } + } + } else { + None + }; let handle = tokio::spawn(async move { let mut msgs = msgs_arc_clone.lock().await.clone(); let outcome = claurst_query::run_query_loop( @@ -4297,7 +4413,20 @@ async fn run_interactive( let tracker = cost_tracker.clone(); let tx = event_tx.clone(); let client_clone = client.clone(); - goal_turn_start = std::time::Instant::now(); + goal_turn_baseline = if claurst_core::goals_enabled() { + match capture_goal_turn_baseline(&session.id, cost_tracker.total_tokens()) { + Ok(baseline) => baseline, + Err(error) => { + app.status_message = Some(format!( + "Goal error: could not capture turn baseline: {}", + error + )); + None + } + } + } else { + None + }; let handle = tokio::spawn(async move { let mut msgs = msgs_arc_clone.lock().await.clone(); @@ -4376,8 +4505,9 @@ async fn run_interactive( } _ => {} }; - let auto_compact_succeeded = - matches!(query_outcome, Ok(QueryOutcome::EndTurn { .. })); + let completed_end_turn = matches!(&query_outcome, Ok(QueryOutcome::EndTurn { .. })); + let was_auto_compact = app.auto_compact_running; + let auto_compact_succeeded = completed_end_turn; // Sync the updated conversation back to our local vector messages = msgs_arc.lock().await.clone(); session.messages = messages.clone(); @@ -4451,14 +4581,46 @@ async fn run_interactive( // After every completed turn check if there is an active goal. // If so, inject a continuation user message and dispatch another turn // without waiting for user input. - if !app.auto_compact_running && claurst_core::goals_enabled() { - let elapsed_secs = goal_turn_start.elapsed().as_secs(); - let total_tokens = cost_tracker.total_tokens(); - match claurst_query::check_and_continue_goal( - &session.id, - total_tokens, - elapsed_secs, - ) { + if completed_end_turn && !was_auto_compact && claurst_core::goals_enabled() { + let baseline = goal_turn_baseline.take(); + let current_goal = + claurst_core::GoalStore::open(&claurst_core::GoalStore::default_path()) + .and_then(|store| store.try_get_goal(&session.id)); + let continuation = match current_goal { + Ok(goal) => match decide_goal_turn_accounting( + baseline.as_ref(), + goal.as_ref(), + cost_tracker.total_tokens(), + ) { + Ok(GoalTurnAccountingDecision::Skip) => { + claurst_query::GoalContinuation::NoGoal + } + Ok(GoalTurnAccountingDecision::Record { total_tokens_used }) => { + match baseline.as_ref() { + Some(baseline) => { + claurst_query::check_and_continue_goal_for_goal( + &session.id, + &baseline.goal_id, + total_tokens_used, + baseline.started_at.elapsed().as_secs(), + ) + } + None => claurst_query::GoalContinuation::Stop { + reason: claurst_query::StopReason::Error( + "missing active-goal turn baseline".to_string(), + ), + }, + } + } + Err(error) => claurst_query::GoalContinuation::Stop { + reason: claurst_query::StopReason::Error(error), + }, + }, + Err(error) => claurst_query::GoalContinuation::Stop { + reason: claurst_query::StopReason::Error(error.to_string()), + }, + }; + match continuation { claurst_query::GoalContinuation::Continue { message } => { // Show a subtle status notice. app.status_message = Some( @@ -4546,7 +4708,23 @@ async fn run_interactive( let tracker = cost_tracker.clone(); let tx = event_tx.clone(); let client_clone = client.clone(); - goal_turn_start = std::time::Instant::now(); + goal_turn_baseline = if claurst_core::goals_enabled() { + match capture_goal_turn_baseline( + &session.id, + cost_tracker.total_tokens(), + ) { + Ok(baseline) => baseline, + Err(error) => { + app.status_message = Some(format!( + "Goal error: could not capture turn baseline: {}", + error + )); + None + } + } + } else { + None + }; let handle = tokio::spawn(async move { let mut msgs = msgs_arc_clone.lock().await.clone(); @@ -4577,6 +4755,8 @@ async fn run_interactive( app.active_goal_badge = None; } } + } else { + goal_turn_baseline = None; } } } @@ -5813,4 +5993,95 @@ mod tests { // Piped stdout/stderr (e.g. CI, script) assert!(!should_show_engine_notice(true, false, false)); } + + fn accounting_goal( + id: &str, + status: claurst_core::GoalStatus, + tokens_used: u64, + ) -> claurst_core::Goal { + claurst_core::Goal { + id: id.to_string(), + session_id: "session".to_string(), + objective: "finish".to_string(), + status, + token_budget: None, + tokens_used, + time_used_secs: 0, + turns_used: 0, + created_at_ms: 1, + updated_at_ms: 1, + } + } + + fn accounting_baseline(id: &str, tracker_total: u64) -> GoalTurnBaseline { + GoalTurnBaseline { + goal_id: id.to_string(), + tracker_total, + started_at: std::time::Instant::now(), + } + } + + #[test] + fn goal_accounting_excludes_pre_goal_tracker_usage() { + let baseline = accounting_baseline("goal-1", 1_000); + let goal = accounting_goal("goal-1", claurst_core::GoalStatus::Active, 30); + + assert_eq!( + decide_goal_turn_accounting(Some(&baseline), Some(&goal), 1_025).unwrap(), + GoalTurnAccountingDecision::Record { + total_tokens_used: 55 + } + ); + } + + #[test] + fn goal_accounting_skips_goal_that_became_terminal_during_turn() { + let baseline = accounting_baseline("goal-1", 100); + let goal = accounting_goal("goal-1", claurst_core::GoalStatus::Complete, 9); + + assert_eq!( + decide_goal_turn_accounting(Some(&baseline), Some(&goal), 110).unwrap(), + GoalTurnAccountingDecision::Skip + ); + } + + #[test] + fn goal_accounting_skips_terminal_goal_without_active_baseline() { + let goal = accounting_goal("goal-1", claurst_core::GoalStatus::Paused, 9); + + assert_eq!( + decide_goal_turn_accounting(None, Some(&goal), 110).unwrap(), + GoalTurnAccountingDecision::Skip + ); + } + + #[test] + fn goal_accounting_rejects_active_goal_without_baseline() { + let goal = accounting_goal("goal-1", claurst_core::GoalStatus::Active, 0); + + let error = decide_goal_turn_accounting(None, Some(&goal), 110).unwrap_err(); + + assert!(error.contains("missing active-goal turn baseline")); + } + + #[test] + fn goal_accounting_rejects_replacement_goal_identity() { + let baseline = accounting_baseline("old-goal", 100); + let replacement = accounting_goal("new-goal", claurst_core::GoalStatus::Active, 0); + + let error = + decide_goal_turn_accounting(Some(&baseline), Some(&replacement), 110).unwrap_err(); + + assert!(error.contains("goal changed during the turn")); + } + + #[test] + fn goal_accounting_reports_tracker_underflow() { + let baseline = accounting_baseline("goal-1", 100); + let goal = accounting_goal("goal-1", claurst_core::GoalStatus::Active, 0); + + let error = decide_goal_turn_accounting(Some(&baseline), Some(&goal), 99).unwrap_err(); + + assert!(error.contains("tracker moved backwards")); + } } diff --git a/src-rust/crates/core/src/goal.rs b/src-rust/crates/core/src/goal.rs index d8412c7..289fab6 100644 --- a/src-rust/crates/core/src/goal.rs +++ b/src-rust/crates/core/src/goal.rs @@ -8,6 +8,8 @@ use std::path::PathBuf; use std::time::{SystemTime, UNIX_EPOCH}; +use rusqlite::{OptionalExtension, Row}; + /// Maximum number of characters allowed in an objective (matches Codex MAX_THREAD_GOAL_OBJECTIVE_CHARS). pub const MAX_OBJECTIVE_CHARS: usize = 4000; @@ -108,18 +110,67 @@ impl Goal { // Error type // --------------------------------------------------------------------------- -#[derive(Debug)] +#[derive(Debug, PartialEq, Eq)] pub enum GoalError { - ObjectiveTooLong { len: usize, max: usize }, + ObjectiveEmpty, + ObjectiveTooLong { + len: usize, + max: usize, + }, + TokenBudgetTooLarge { + budget: u64, + max: u64, + }, + NotFound { + session_id: String, + }, + NotActive { + session_id: String, + }, + Replaced { + session_id: String, + expected_goal_id: String, + actual_goal_id: String, + }, + ValueTooLarge { + field: &'static str, + value: u64, + max: u64, + }, + InvalidStoredValue { + field: &'static str, + value: String, + }, Db(String), } impl std::fmt::Display for GoalError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { + GoalError::ObjectiveEmpty => write!(f, "Goal objective must not be empty"), GoalError::ObjectiveTooLong { len, max } => { write!(f, "Objective too long: {} chars (max {})", len, max) } + GoalError::TokenBudgetTooLarge { budget, max } => { + write!(f, "Token budget {} exceeds SQLite maximum {}", budget, max) + } + GoalError::NotFound { session_id } => write!(f, "Goal not found: {}", session_id), + GoalError::NotActive { session_id } => write!(f, "Goal is not active: {}", session_id), + GoalError::Replaced { + session_id, + expected_goal_id, + actual_goal_id, + } => write!( + f, + "Goal was replaced for session {}: expected {}, found {}", + session_id, expected_goal_id, actual_goal_id + ), + GoalError::ValueTooLarge { field, value, max } => { + write!(f, "Goal {} {} exceeds maximum {}", field, value, max) + } + GoalError::InvalidStoredValue { field, value } => { + write!(f, "Invalid stored goal {}: {}", field, value) + } GoalError::Db(msg) => write!(f, "Goal DB error: {}", msg), } } @@ -135,6 +186,36 @@ pub struct GoalStore { conn: rusqlite::Connection, } +struct StoredGoal { + id: String, + session_id: String, + objective: String, + status: String, + token_budget: Option, + tokens_used: i64, + time_used_secs: i64, + turns_used: i64, + created_at_ms: i64, + updated_at_ms: i64, +} + +impl StoredGoal { + fn from_row(row: &Row<'_>) -> rusqlite::Result { + Ok(Self { + id: row.get(0)?, + session_id: row.get(1)?, + objective: row.get(2)?, + status: row.get(3)?, + token_budget: row.get(4)?, + tokens_used: row.get(5)?, + time_used_secs: row.get(6)?, + turns_used: row.get(7)?, + created_at_ms: row.get(8)?, + updated_at_ms: row.get(9)?, + }) + } +} + impl GoalStore { /// Open (or create) the goal database. pub fn open(db_path: &std::path::Path) -> Result { @@ -177,6 +258,131 @@ impl GoalStore { .as_millis() as u64 } + fn sqlite_i64(value: u64, field: &'static str) -> Result { + i64::try_from(value).map_err(|_| GoalError::ValueTooLarge { + field, + value, + max: i64::MAX as u64, + }) + } + + fn stored_u64(value: i64, field: &'static str) -> Result { + u64::try_from(value).map_err(|_| GoalError::InvalidStoredValue { + field, + value: value.to_string(), + }) + } + + fn stored_u32(value: i64, field: &'static str) -> Result { + u32::try_from(value).map_err(|_| GoalError::InvalidStoredValue { + field, + value: value.to_string(), + }) + } + + fn decode_goal_row(row: StoredGoal) -> Result { + let status = + GoalStatus::parse(&row.status).ok_or_else(|| GoalError::InvalidStoredValue { + field: "status", + value: row.status.clone(), + })?; + let token_budget = row + .token_budget + .map(|budget| Self::stored_u64(budget, "token_budget")) + .transpose()?; + Ok(Goal { + id: row.id, + session_id: row.session_id, + objective: row.objective, + status, + token_budget, + tokens_used: Self::stored_u64(row.tokens_used, "tokens_used")?, + time_used_secs: Self::stored_u64(row.time_used_secs, "time_used_secs")?, + turns_used: Self::stored_u32(row.turns_used, "turns_used")?, + created_at_ms: Self::stored_u64(row.created_at_ms, "created_at_ms")?, + updated_at_ms: Self::stored_u64(row.updated_at_ms, "updated_at_ms")?, + }) + } + + fn query_goals

(&self, query: &str, params: P) -> Result, GoalError> + where + P: rusqlite::Params, + { + let mut statement = self + .conn + .prepare(query) + .map_err(|err| GoalError::Db(err.to_string()))?; + let mut rows = statement + .query(params) + .map_err(|err| GoalError::Db(err.to_string()))?; + let mut goals = Vec::new(); + while let Some(row) = rows.next().map_err(|err| GoalError::Db(err.to_string()))? { + goals.push(Self::decode_goal_row( + StoredGoal::from_row(row).map_err(|err| GoalError::Db(err.to_string()))?, + )?); + } + Ok(goals) + } + + fn transition_error(&self, session_id: &str) -> Result { + if self.try_get_goal(session_id)?.is_some() { + Ok(GoalError::NotActive { + session_id: session_id.to_string(), + }) + } else { + Ok(GoalError::NotFound { + session_id: session_id.to_string(), + }) + } + } + + fn immediate_transaction(&self) -> Result, GoalError> { + rusqlite::Transaction::new_unchecked(&self.conn, rusqlite::TransactionBehavior::Immediate) + .map_err(|err| GoalError::Db(err.to_string())) + } + + fn transaction_goal( + transaction: &rusqlite::Transaction<'_>, + session_id: &str, + ) -> Result, GoalError> { + let stored = transaction + .query_row( + "SELECT id, session_id, objective, status, token_budget, + tokens_used, time_used_secs, turns_used, + created_at_ms, updated_at_ms + FROM goals WHERE session_id = ?1", + [session_id], + StoredGoal::from_row, + ) + .optional() + .map_err(|err| GoalError::Db(err.to_string()))?; + stored.map(Self::decode_goal_row).transpose() + } + + fn checked_sqlite_sum( + current: u64, + increment: u64, + field: &'static str, + ) -> Result { + let total = current + .checked_add(increment) + .ok_or(GoalError::ValueTooLarge { + field, + value: u64::MAX, + max: i64::MAX as u64, + })?; + Self::sqlite_i64(total, field) + } + + fn next_turn(turns_used: u32) -> Result { + let next = turns_used.checked_add(1).ok_or(GoalError::ValueTooLarge { + field: "turns_used", + value: u64::from(u32::MAX) + 1, + max: u64::from(u32::MAX), + })?; + Ok(i64::from(next)) + } + /// Create or replace the active goal for a session. pub fn set_goal( &self, @@ -184,6 +390,9 @@ impl GoalStore { objective: &str, token_budget: Option, ) -> Result { + if objective.trim().is_empty() { + return Err(GoalError::ObjectiveEmpty); + } if objective.chars().count() > MAX_OBJECTIVE_CHARS { return Err(GoalError::ObjectiveTooLong { len: objective.chars().count(), @@ -192,22 +401,37 @@ impl GoalStore { } let now = Self::now_ms(); + let sqlite_now = Self::sqlite_i64(now, "timestamp")?; + let sqlite_budget = token_budget + .map(|budget| { + i64::try_from(budget).map_err(|_| GoalError::TokenBudgetTooLarge { + budget, + max: i64::MAX as u64, + }) + }) + .transpose()?; let id = uuid_v4(); - // Remove any pre-existing goal for this session first. - self.conn + let transaction = self + .conn + .unchecked_transaction() + .map_err(|err| GoalError::Db(err.to_string()))?; + transaction .execute("DELETE FROM goals WHERE session_id = ?1", [session_id]) - .map_err(|e| GoalError::Db(e.to_string()))?; + .map_err(|err| GoalError::Db(err.to_string()))?; - self.conn + transaction .execute( "INSERT INTO goals (id, session_id, objective, status, token_budget, tokens_used, time_used_secs, turns_used, created_at_ms, updated_at_ms) VALUES (?1, ?2, ?3, 'active', ?4, 0, 0, 0, ?5, ?5)", - rusqlite::params![id, session_id, objective, token_budget, now], + rusqlite::params![&id, session_id, objective, sqlite_budget, sqlite_now], ) - .map_err(|e| GoalError::Db(e.to_string()))?; + .map_err(|err| GoalError::Db(err.to_string()))?; + transaction + .commit() + .map_err(|err| GoalError::Db(err.to_string()))?; Ok(Goal { id, @@ -223,32 +447,26 @@ impl GoalStore { }) } - /// Get the current goal for a session (any status). - pub fn get_goal(&self, session_id: &str) -> Option { - self.conn + /// Get the current goal for a session while preserving read/conversion errors. + pub fn try_get_goal(&self, session_id: &str) -> Result, GoalError> { + let stored = self + .conn .query_row( "SELECT id, session_id, objective, status, token_budget, tokens_used, time_used_secs, turns_used, created_at_ms, updated_at_ms FROM goals WHERE session_id = ?1", [session_id], - |row| { - let status_str: String = row.get(3)?; - Ok(Goal { - id: row.get(0)?, - session_id: row.get(1)?, - objective: row.get(2)?, - status: GoalStatus::parse(&status_str).unwrap_or(GoalStatus::Paused), - token_budget: row.get(4)?, - tokens_used: row.get::<_, i64>(5)? as u64, - time_used_secs: row.get::<_, i64>(6)? as u64, - turns_used: row.get::<_, i64>(7)? as u32, - created_at_ms: row.get::<_, i64>(8)? as u64, - updated_at_ms: row.get::<_, i64>(9)? as u64, - }) - }, + StoredGoal::from_row, ) - .ok() + .optional() + .map_err(|err| GoalError::Db(err.to_string()))?; + stored.map(Self::decode_goal_row).transpose() + } + + /// Get the current goal for a session (any status). + pub fn get_goal(&self, session_id: &str) -> Option { + self.try_get_goal(session_id).ok().flatten() } /// Get the active goal for a session (status = 'active' only). @@ -260,12 +478,22 @@ impl GoalStore { /// Update the status of the goal for a session. pub fn set_status(&self, session_id: &str, status: GoalStatus) -> Result<(), GoalError> { let now = Self::now_ms(); - self.conn + let updated = self + .conn .execute( "UPDATE goals SET status = ?1, updated_at_ms = ?2 WHERE session_id = ?3", - rusqlite::params![status.as_str(), now, session_id], + rusqlite::params![ + status.as_str(), + Self::sqlite_i64(now, "timestamp")?, + session_id + ], ) - .map_err(|e| GoalError::Db(e.to_string()))?; + .map_err(|err| GoalError::Db(err.to_string()))?; + if updated == 0 { + return Err(GoalError::NotFound { + session_id: session_id.to_string(), + }); + } Ok(()) } @@ -279,33 +507,383 @@ impl GoalStore { /// Record one completed turn: increment turns_used, add elapsed seconds. pub fn record_turn(&self, session_id: &str, elapsed_secs: u64) -> Result<(), GoalError> { - let now = Self::now_ms(); - self.conn + let transaction = self.immediate_transaction()?; + let goal = Self::transaction_goal(&transaction, session_id)?.ok_or_else(|| { + GoalError::NotFound { + session_id: session_id.to_string(), + } + })?; + let next_time_used_secs = + Self::checked_sqlite_sum(goal.time_used_secs, elapsed_secs, "time_used_secs")?; + let next_turns_used = Self::next_turn(goal.turns_used)?; + let updated = transaction .execute( "UPDATE goals - SET turns_used = turns_used + 1, - time_used_secs = time_used_secs + ?1, - updated_at_ms = ?2 - WHERE session_id = ?3", - rusqlite::params![elapsed_secs, now, session_id], + SET turns_used = ?1, + time_used_secs = ?2, + updated_at_ms = ?3 + WHERE session_id = ?4", + rusqlite::params![ + next_turns_used, + next_time_used_secs, + Self::sqlite_i64(Self::now_ms(), "timestamp")?, + session_id + ], ) - .map_err(|e| GoalError::Db(e.to_string()))?; + .map_err(|err| GoalError::Db(err.to_string()))?; + if updated == 0 { + return Err(GoalError::NotFound { + session_id: session_id.to_string(), + }); + } + transaction + .commit() + .map_err(|err| GoalError::Db(err.to_string()))?; Ok(()) } /// Add token usage (used to enforce soft budget). pub fn add_tokens(&self, session_id: &str, tokens: u64) -> Result<(), GoalError> { - let now = Self::now_ms(); - self.conn + let transaction = self.immediate_transaction()?; + let goal = Self::transaction_goal(&transaction, session_id)?.ok_or_else(|| { + GoalError::NotFound { + session_id: session_id.to_string(), + } + })?; + let next_tokens_used = Self::checked_sqlite_sum(goal.tokens_used, tokens, "tokens_used")?; + let updated = transaction .execute( "UPDATE goals - SET tokens_used = tokens_used + ?1, updated_at_ms = ?2 + SET tokens_used = ?1, updated_at_ms = ?2 WHERE session_id = ?3", - rusqlite::params![tokens, now, session_id], + rusqlite::params![ + next_tokens_used, + Self::sqlite_i64(Self::now_ms(), "timestamp")?, + session_id + ], ) - .map_err(|e| GoalError::Db(e.to_string()))?; + .map_err(|err| GoalError::Db(err.to_string()))?; + if updated == 0 { + return Err(GoalError::NotFound { + session_id: session_id.to_string(), + }); + } + transaction + .commit() + .map_err(|err| GoalError::Db(err.to_string()))?; Ok(()) } + + /// Record an absolute token total and elapsed duration for one completed turn. + pub fn record_completed_turn( + &self, + session_id: &str, + total_tokens_used: u64, + elapsed_secs: u64, + ) -> Result<(), GoalError> { + self.record_completed_turn_inner(session_id, None, total_tokens_used, elapsed_secs) + } + + /// Record one completed turn only when the session still owns `expected_goal_id`. + pub fn record_completed_turn_for_goal( + &self, + session_id: &str, + expected_goal_id: &str, + total_tokens_used: u64, + elapsed_secs: u64, + ) -> Result<(), GoalError> { + self.record_completed_turn_inner( + session_id, + Some(expected_goal_id), + total_tokens_used, + elapsed_secs, + ) + } + + fn record_completed_turn_inner( + &self, + session_id: &str, + expected_goal_id: Option<&str>, + total_tokens_used: u64, + elapsed_secs: u64, + ) -> Result<(), GoalError> { + let total_tokens_used = Self::sqlite_i64(total_tokens_used, "total_tokens_used")?; + let transaction = self.immediate_transaction()?; + let goal = Self::transaction_goal(&transaction, session_id)?.ok_or_else(|| { + GoalError::NotFound { + session_id: session_id.to_string(), + } + })?; + if let Some(expected_goal_id) = expected_goal_id { + if goal.id != expected_goal_id { + return Err(GoalError::Replaced { + session_id: session_id.to_string(), + expected_goal_id: expected_goal_id.to_string(), + actual_goal_id: goal.id, + }); + } + } + if goal.status != GoalStatus::Active { + return Err(GoalError::NotActive { + session_id: session_id.to_string(), + }); + } + let next_time_used_secs = + Self::checked_sqlite_sum(goal.time_used_secs, elapsed_secs, "time_used_secs")?; + let next_turns_used = Self::next_turn(goal.turns_used)?; + let updated = transaction + .execute( + "UPDATE goals + SET tokens_used = MAX(tokens_used, ?1), + time_used_secs = ?2, + turns_used = ?3, + updated_at_ms = ?4 + WHERE session_id = ?5 + AND status = 'active' + AND (?6 IS NULL OR id = ?6)", + rusqlite::params![ + total_tokens_used, + next_time_used_secs, + next_turns_used, + Self::sqlite_i64(Self::now_ms(), "timestamp")?, + session_id, + expected_goal_id, + ], + ) + .map_err(|err| GoalError::Db(err.to_string()))?; + if updated == 0 { + return Err(GoalError::NotActive { + session_id: session_id.to_string(), + }); + } + transaction + .commit() + .map_err(|err| GoalError::Db(err.to_string()))?; + Ok(()) + } + + /// Mark an active goal complete without changing paused or terminal goals. + pub fn complete_active_goal(&self, session_id: &str) -> Result<(), GoalError> { + let updated = self + .conn + .execute( + "UPDATE goals SET status = 'complete', updated_at_ms = ?1 + WHERE session_id = ?2 AND status = 'active'", + rusqlite::params![Self::sqlite_i64(Self::now_ms(), "timestamp")?, session_id], + ) + .map_err(|err| GoalError::Db(err.to_string()))?; + if updated == 0 { + return Err(self.transition_error(session_id)?); + } + Ok(()) + } + + /// Pause an active goal. Repeating a pause for a paused goal is idempotent. + pub fn pause_active_goal(&self, session_id: &str) -> Result<(), GoalError> { + let updated = self + .conn + .execute( + "UPDATE goals SET status = 'paused', updated_at_ms = ?1 + WHERE session_id = ?2 AND status = 'active'", + rusqlite::params![Self::sqlite_i64(Self::now_ms(), "timestamp")?, session_id], + ) + .map_err(|err| GoalError::Db(err.to_string()))?; + if updated > 0 { + return Ok(()); + } + match self.try_get_goal(session_id)? { + Some(goal) if goal.status == GoalStatus::Paused => Ok(()), + Some(_) => Err(GoalError::NotActive { + session_id: session_id.to_string(), + }), + None => Err(GoalError::NotFound { + session_id: session_id.to_string(), + }), + } + } + + /// Pause the goal only if `expected_goal_id` still owns the session row. + pub fn pause_active_goal_for_goal( + &self, + session_id: &str, + expected_goal_id: &str, + ) -> Result<(), GoalError> { + let transaction = self.immediate_transaction()?; + let goal = Self::transaction_goal(&transaction, session_id)?.ok_or_else(|| { + GoalError::NotFound { + session_id: session_id.to_string(), + } + })?; + if goal.id != expected_goal_id { + return Err(GoalError::Replaced { + session_id: session_id.to_string(), + expected_goal_id: expected_goal_id.to_string(), + actual_goal_id: goal.id, + }); + } + if goal.status == GoalStatus::Paused { + return Ok(()); + } + if goal.status != GoalStatus::Active { + return Err(GoalError::NotActive { + session_id: session_id.to_string(), + }); + } + let updated = transaction + .execute( + "UPDATE goals SET status = 'paused', updated_at_ms = ?1 + WHERE session_id = ?2 AND id = ?3 AND status = 'active'", + rusqlite::params![ + Self::sqlite_i64(Self::now_ms(), "timestamp")?, + session_id, + expected_goal_id, + ], + ) + .map_err(|err| GoalError::Db(err.to_string()))?; + if updated == 0 { + return Err(GoalError::NotActive { + session_id: session_id.to_string(), + }); + } + transaction + .commit() + .map_err(|err| GoalError::Db(err.to_string()))?; + Ok(()) + } + + /// Resume a paused goal if it has not reached the automatic turn cap. + pub fn resume_paused_goal(&self, session_id: &str) -> Result<(), GoalError> { + let updated = self + .conn + .execute( + "UPDATE goals SET status = 'active', updated_at_ms = ?1 + WHERE session_id = ?2 AND status = 'paused' AND turns_used < ?3", + rusqlite::params![ + Self::sqlite_i64(Self::now_ms(), "timestamp")?, + session_id, + i64::from(MAX_GOAL_TURNS), + ], + ) + .map_err(|err| GoalError::Db(err.to_string()))?; + if updated == 0 { + return Err(self.transition_error(session_id)?); + } + Ok(()) + } + + /// Stop an active goal because its soft budget has been reached. + pub fn budget_limit_active_goal(&self, session_id: &str) -> Result<(), GoalError> { + let updated = self + .conn + .execute( + "UPDATE goals SET status = 'budget_limited', updated_at_ms = ?1 + WHERE session_id = ?2 AND status = 'active'", + rusqlite::params![Self::sqlite_i64(Self::now_ms(), "timestamp")?, session_id], + ) + .map_err(|err| GoalError::Db(err.to_string()))?; + if updated == 0 { + return Err(self.transition_error(session_id)?); + } + Ok(()) + } + + /// Budget-limit the goal only if `expected_goal_id` still owns the session row. + pub fn budget_limit_active_goal_for_goal( + &self, + session_id: &str, + expected_goal_id: &str, + ) -> Result<(), GoalError> { + let transaction = self.immediate_transaction()?; + let goal = Self::transaction_goal(&transaction, session_id)?.ok_or_else(|| { + GoalError::NotFound { + session_id: session_id.to_string(), + } + })?; + if goal.id != expected_goal_id { + return Err(GoalError::Replaced { + session_id: session_id.to_string(), + expected_goal_id: expected_goal_id.to_string(), + actual_goal_id: goal.id, + }); + } + if goal.status != GoalStatus::Active { + return Err(GoalError::NotActive { + session_id: session_id.to_string(), + }); + } + let updated = transaction + .execute( + "UPDATE goals SET status = 'budget_limited', updated_at_ms = ?1 + WHERE session_id = ?2 AND id = ?3 AND status = 'active'", + rusqlite::params![ + Self::sqlite_i64(Self::now_ms(), "timestamp")?, + session_id, + expected_goal_id, + ], + ) + .map_err(|err| GoalError::Db(err.to_string()))?; + if updated == 0 { + return Err(GoalError::NotActive { + session_id: session_id.to_string(), + }); + } + transaction + .commit() + .map_err(|err| GoalError::Db(err.to_string()))?; + Ok(()) + } + + /// Return every persisted goal while propagating row conversion errors. + pub fn list_goals(&self) -> Result, GoalError> { + self.query_goals( + "SELECT id, session_id, objective, status, token_budget, + tokens_used, time_used_secs, turns_used, + created_at_ms, updated_at_ms + FROM goals ORDER BY created_at_ms ASC, id ASC", + [], + ) + } + + /// Reconcile active goals after launch by pausing exactly the rows that were active. + pub fn pause_active_goals(&mut self) -> Result, GoalError> { + let now = Self::now_ms(); + let sqlite_now = Self::sqlite_i64(now, "timestamp")?; + let transaction = self.immediate_transaction()?; + let mut active_goals = { + let mut statement = transaction + .prepare( + "SELECT id, session_id, objective, status, token_budget, + tokens_used, time_used_secs, turns_used, + created_at_ms, updated_at_ms + FROM goals WHERE status = 'active' ORDER BY created_at_ms ASC, id ASC", + ) + .map_err(|err| GoalError::Db(err.to_string()))?; + let mut rows = statement + .query([]) + .map_err(|err| GoalError::Db(err.to_string()))?; + let mut goals = Vec::new(); + while let Some(row) = rows.next().map_err(|err| GoalError::Db(err.to_string()))? { + goals.push(Self::decode_goal_row( + StoredGoal::from_row(row).map_err(|err| GoalError::Db(err.to_string()))?, + )?); + } + goals + }; + if !active_goals.is_empty() { + transaction.execute( + "UPDATE goals SET status = 'paused', updated_at_ms = ?1 WHERE status = 'active'", + [sqlite_now], + ).map_err(|err| GoalError::Db(err.to_string()))?; + } + transaction + .commit() + .map_err(|err| GoalError::Db(err.to_string()))?; + for goal in &mut active_goals { + goal.status = GoalStatus::Paused; + goal.updated_at_ms = now; + } + Ok(active_goals) + } } // --------------------------------------------------------------------------- @@ -416,6 +994,25 @@ pub fn goal_continuation_message(goal: &Goal) -> String { mod tests { use super::*; use std::path::Path; + use std::sync::{mpsc, Condvar, Mutex, OnceLock}; + use std::thread; + use std::time::Duration; + + static BUSY_HANDLER_STATE: OnceLock<(Mutex, Condvar)> = OnceLock::new(); + + fn busy_handler_state() -> &'static (Mutex, Condvar) { + BUSY_HANDLER_STATE.get_or_init(|| (Mutex::new(false), Condvar::new())) + } + + fn observe_busy_handler(_: i32) -> bool { + let (lock, condvar) = busy_handler_state(); + let mut observed = lock.lock().unwrap(); + *observed = true; + condvar.notify_all(); + drop(observed); + thread::sleep(Duration::from_millis(10)); + true + } fn open_tmp() -> GoalStore { GoalStore::open(Path::new(":memory:")).unwrap() @@ -441,6 +1038,548 @@ mod tests { assert!(matches!(result, Err(GoalError::ObjectiveTooLong { .. }))); } + #[test] + fn empty_objectives_are_rejected_before_length_validation() { + let store = open_tmp(); + assert!(matches!( + store.set_goal("sess1", " \n\t ", None), + Err(GoalError::ObjectiveEmpty) + )); + } + + #[test] + fn missing_goal_mutations_are_errors() { + let store = open_tmp(); + assert!(matches!( + store.set_status("missing", GoalStatus::Paused), + Err(GoalError::NotFound { session_id }) if session_id == "missing" + )); + assert!(matches!( + store.record_completed_turn("missing", 10, 2), + Err(GoalError::NotFound { session_id }) if session_id == "missing" + )); + } + + #[test] + fn completed_turn_records_monotonic_absolute_progress_atomically() { + let store = open_tmp(); + store.set_goal("sess1", "ship the feature", None).unwrap(); + + store.record_completed_turn("sess1", 700, 11).unwrap(); + store.record_completed_turn("sess1", 650, 7).unwrap(); + + let goal = store.try_get_goal("sess1").unwrap().unwrap(); + assert_eq!(goal.tokens_used, 700); + assert_eq!(goal.time_used_secs, 18); + assert_eq!(goal.turns_used, 2); + } + + #[test] + fn completed_turn_rejects_elapsed_overflow_without_mutating_goal() { + let store = open_tmp(); + store + .set_goal("sess1", "preserve valid counters", None) + .unwrap(); + store + .conn + .execute( + "UPDATE goals + SET tokens_used = 44, time_used_secs = ?1, turns_used = 4 + WHERE session_id = ?2", + rusqlite::params![i64::MAX, "sess1"], + ) + .unwrap(); + + assert!(matches!( + store.record_completed_turn("sess1", 99, 1), + Err(GoalError::ValueTooLarge { + field: "time_used_secs", + .. + }) + )); + + let goal = store.try_get_goal("sess1").unwrap().unwrap(); + assert_eq!(goal.tokens_used, 44); + assert_eq!(goal.time_used_secs, i64::MAX as u64); + assert_eq!(goal.turns_used, 4); + } + + #[test] + fn completed_turn_rejects_turn_overflow_without_mutating_goal() { + let store = open_tmp(); + store + .set_goal("sess1", "preserve valid counters", None) + .unwrap(); + store + .conn + .execute( + "UPDATE goals + SET tokens_used = 44, time_used_secs = 50, turns_used = ?1 + WHERE session_id = ?2", + rusqlite::params![i64::from(u32::MAX), "sess1"], + ) + .unwrap(); + + assert!(matches!( + store.record_completed_turn("sess1", 99, 1), + Err(GoalError::ValueTooLarge { + field: "turns_used", + .. + }) + )); + + let goal = store.try_get_goal("sess1").unwrap().unwrap(); + assert_eq!(goal.tokens_used, 44); + assert_eq!(goal.time_used_secs, 50); + assert_eq!(goal.turns_used, u32::MAX); + } + + #[test] + fn completed_turn_serializes_concurrent_writers_without_losing_accounting() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("goals.sqlite"); + let first = GoalStore::open(&path).unwrap(); + first.set_goal("sess1", "serialize writers", None).unwrap(); + let second = GoalStore::open(&path).unwrap(); + let transaction = rusqlite::Transaction::new_unchecked( + &first.conn, + rusqlite::TransactionBehavior::Immediate, + ) + .unwrap(); + transaction + .execute( + "UPDATE goals + SET tokens_used = 300, time_used_secs = 5, turns_used = 1 + WHERE session_id = ?1", + ["sess1"], + ) + .unwrap(); + + let (busy_lock, busy_condvar) = busy_handler_state(); + *busy_lock.lock().unwrap() = false; + second + .conn + .busy_handler(Some(observe_busy_handler)) + .unwrap(); + let (finished_tx, finished_rx) = mpsc::channel(); + let writer = thread::spawn(move || { + finished_tx + .send(second.record_completed_turn("sess1", 500, 7)) + .unwrap(); + }); + + let (busy_observed, _) = busy_condvar + .wait_timeout_while( + busy_lock.lock().unwrap(), + Duration::from_secs(1), + |observed| !*observed, + ) + .unwrap(); + assert!( + *busy_observed, + "second writer never contended on SQLite's busy handler" + ); + drop(busy_observed); + transaction.commit().unwrap(); + assert!(finished_rx + .recv_timeout(Duration::from_secs(1)) + .unwrap() + .is_ok()); + writer.join().unwrap(); + + let goal = first.try_get_goal("sess1").unwrap().unwrap(); + assert_eq!(goal.tokens_used, 500); + assert_eq!(goal.time_used_secs, 12); + assert_eq!(goal.turns_used, 2); + } + + #[test] + fn record_turn_rejects_elapsed_overflow_without_mutating_goal() { + let store = open_tmp(); + store.set_goal("sess1", "preserve counters", None).unwrap(); + store + .conn + .execute( + "UPDATE goals SET time_used_secs = ?1, turns_used = 4 WHERE session_id = ?2", + rusqlite::params![i64::MAX, "sess1"], + ) + .unwrap(); + + assert!(matches!( + store.record_turn("sess1", 1), + Err(GoalError::ValueTooLarge { + field: "time_used_secs", + .. + }) + )); + let goal = store.try_get_goal("sess1").unwrap().unwrap(); + assert_eq!(goal.time_used_secs, i64::MAX as u64); + assert_eq!(goal.turns_used, 4); + } + + #[test] + fn record_turn_rejects_turn_overflow_without_mutating_goal() { + let store = open_tmp(); + store.set_goal("sess1", "preserve counters", None).unwrap(); + store + .conn + .execute( + "UPDATE goals SET time_used_secs = 50, turns_used = ?1 WHERE session_id = ?2", + rusqlite::params![i64::from(u32::MAX), "sess1"], + ) + .unwrap(); + + assert!(matches!( + store.record_turn("sess1", 1), + Err(GoalError::ValueTooLarge { + field: "turns_used", + .. + }) + )); + let goal = store.try_get_goal("sess1").unwrap().unwrap(); + assert_eq!(goal.time_used_secs, 50); + assert_eq!(goal.turns_used, u32::MAX); + } + + #[test] + fn add_tokens_rejects_accumulated_overflow_without_mutating_goal() { + let store = open_tmp(); + store.set_goal("sess1", "preserve counters", None).unwrap(); + store + .conn + .execute( + "UPDATE goals SET tokens_used = ?1 WHERE session_id = ?2", + rusqlite::params![i64::MAX, "sess1"], + ) + .unwrap(); + + assert!(matches!( + store.add_tokens("sess1", 1), + Err(GoalError::ValueTooLarge { + field: "tokens_used", + .. + }) + )); + let goal = store.try_get_goal("sess1").unwrap().unwrap(); + assert_eq!(goal.tokens_used, i64::MAX as u64); + } + + #[test] + fn invalid_replacement_budget_preserves_the_existing_goal() { + let store = open_tmp(); + store + .set_goal("sess1", "keep this goal", Some(100)) + .unwrap(); + + assert!(matches!( + store.set_goal("sess1", "replacement", Some(u64::MAX)), + Err(GoalError::TokenBudgetTooLarge { + budget: u64::MAX, + .. + }) + )); + + let goal = store.try_get_goal("sess1").unwrap().unwrap(); + assert_eq!(goal.objective, "keep this goal"); + assert_eq!(goal.token_budget, Some(100)); + } + + #[test] + fn complete_active_goal_rejects_paused_or_missing_goal() { + let store = open_tmp(); + store.set_goal("paused", "wait here", None).unwrap(); + store.set_status("paused", GoalStatus::Paused).unwrap(); + + assert!(matches!( + store.complete_active_goal("paused"), + Err(GoalError::NotActive { session_id }) if session_id == "paused" + )); + assert!(matches!( + store.complete_active_goal("missing"), + Err(GoalError::NotFound { session_id }) if session_id == "missing" + )); + } + + #[test] + fn pause_and_resume_are_status_guarded() { + let store = open_tmp(); + store.set_goal("active", "work", None).unwrap(); + store.pause_active_goal("active").unwrap(); + assert_eq!( + store.try_get_goal("active").unwrap().unwrap().status, + GoalStatus::Paused + ); + store.pause_active_goal("active").unwrap(); + store.resume_paused_goal("active").unwrap(); + assert_eq!( + store.try_get_goal("active").unwrap().unwrap().status, + GoalStatus::Active + ); + + store.set_goal("complete", "finished", None).unwrap(); + store.complete_active_goal("complete").unwrap(); + assert!(matches!( + store.pause_active_goal("complete"), + Err(GoalError::NotActive { .. }) + )); + assert!(matches!( + store.resume_paused_goal("complete"), + Err(GoalError::NotActive { .. }) + )); + assert_eq!( + store.try_get_goal("complete").unwrap().unwrap().status, + GoalStatus::Complete + ); + + store.set_goal("limited", "budget exhausted", None).unwrap(); + store.budget_limit_active_goal("limited").unwrap(); + assert!(matches!( + store.pause_active_goal("limited"), + Err(GoalError::NotActive { .. }) + )); + assert_eq!( + store.try_get_goal("limited").unwrap().unwrap().status, + GoalStatus::BudgetLimited + ); + + store + .set_goal("turn-cap", "stop automatically", None) + .unwrap(); + store + .conn + .execute( + "UPDATE goals SET status = 'paused', turns_used = ?1 WHERE session_id = ?2", + rusqlite::params![i64::from(MAX_GOAL_TURNS), "turn-cap"], + ) + .unwrap(); + assert!(matches!( + store.resume_paused_goal("turn-cap"), + Err(GoalError::NotActive { .. }) + )); + } + + #[test] + fn guarded_transitions_reject_missing_and_non_active_goals() { + let store = open_tmp(); + for transition in [ + store.pause_active_goal("missing"), + store.resume_paused_goal("missing"), + store.budget_limit_active_goal("missing"), + ] { + assert!(matches!( + transition, + Err(GoalError::NotFound { session_id }) if session_id == "missing" + )); + } + + store.set_goal("paused", "wait", None).unwrap(); + store.pause_active_goal("paused").unwrap(); + assert!(matches!( + store.budget_limit_active_goal("paused"), + Err(GoalError::NotActive { session_id }) if session_id == "paused" + )); + } + + #[test] + fn expected_goal_pause_rejects_replacement_without_mutating_it() { + let store = open_tmp(); + let original = store.set_goal("session", "original", None).unwrap(); + let replacement = store.set_goal("session", "replacement", None).unwrap(); + + let error = store + .pause_active_goal_for_goal("session", &original.id) + .unwrap_err(); + + assert!(matches!( + error, + GoalError::Replaced { + expected_goal_id, + actual_goal_id, + .. + } if expected_goal_id == original.id && actual_goal_id == replacement.id + )); + let current = store.try_get_goal("session").unwrap().unwrap(); + assert_eq!(current.id, replacement.id); + assert_eq!(current.status, GoalStatus::Active); + assert_eq!(current.tokens_used, 0); + assert_eq!(current.time_used_secs, 0); + assert_eq!(current.turns_used, 0); + } + + #[test] + fn expected_goal_budget_limit_rejects_replacement_without_mutating_it() { + let store = open_tmp(); + let original = store.set_goal("session", "original", Some(1)).unwrap(); + let replacement = store.set_goal("session", "replacement", Some(1)).unwrap(); + + let error = store + .budget_limit_active_goal_for_goal("session", &original.id) + .unwrap_err(); + + assert!(matches!( + error, + GoalError::Replaced { + expected_goal_id, + actual_goal_id, + .. + } if expected_goal_id == original.id && actual_goal_id == replacement.id + )); + let current = store.try_get_goal("session").unwrap().unwrap(); + assert_eq!(current.id, replacement.id); + assert_eq!(current.status, GoalStatus::Active); + assert_eq!(current.tokens_used, 0); + assert_eq!(current.time_used_secs, 0); + assert_eq!(current.turns_used, 0); + } + + #[test] + fn expected_goal_turn_recording_rejects_terminal_goal_without_mutating_it() { + let store = open_tmp(); + let goal = store.set_goal("session", "finish", None).unwrap(); + store.complete_active_goal("session").unwrap(); + + assert!(matches!( + store.record_completed_turn_for_goal("session", &goal.id, 100, 5), + Err(GoalError::NotActive { session_id }) if session_id == "session" + )); + + let current = store.try_get_goal("session").unwrap().unwrap(); + assert_eq!(current.status, GoalStatus::Complete); + assert_eq!(current.tokens_used, 0); + assert_eq!(current.time_used_secs, 0); + assert_eq!(current.turns_used, 0); + } + + #[test] + fn pause_active_goals_returns_only_reconciled_rows() { + let mut store = open_tmp(); + store.set_goal("active", "resume later", None).unwrap(); + store.set_goal("complete", "already done", None).unwrap(); + store.complete_active_goal("complete").unwrap(); + + let paused = store.pause_active_goals().unwrap(); + assert_eq!(paused.len(), 1); + assert_eq!(paused[0].session_id, "active"); + assert_eq!(paused[0].status, GoalStatus::Paused); + let persisted = store.try_get_goal("active").unwrap().unwrap(); + assert_eq!(persisted.status, GoalStatus::Paused); + assert_eq!(paused[0].updated_at_ms, persisted.updated_at_ms); + assert_eq!( + store.try_get_goal("complete").unwrap().unwrap().status, + GoalStatus::Complete + ); + assert!(store.pause_active_goals().unwrap().is_empty()); + } + + #[test] + fn fallible_reads_reject_invalid_persisted_status_and_numeric_values() { + let store = open_tmp(); + store + .set_goal("sess1", "validate stored values", None) + .unwrap(); + + store + .conn + .execute( + "UPDATE goals SET tokens_used = -1 WHERE session_id = ?1", + ["sess1"], + ) + .unwrap(); + assert!(matches!( + store.try_get_goal("sess1"), + Err(GoalError::InvalidStoredValue { + field: "tokens_used", + .. + }) + )); + + store + .conn + .execute( + "UPDATE goals SET tokens_used = 0, turns_used = ?1 WHERE session_id = ?2", + rusqlite::params![i64::from(u32::MAX) + 1, "sess1"], + ) + .unwrap(); + assert!(matches!( + store.try_get_goal("sess1"), + Err(GoalError::InvalidStoredValue { + field: "turns_used", + .. + }) + )); + + store + .conn + .execute( + "UPDATE goals SET turns_used = 0, status = 'unknown' WHERE session_id = ?1", + ["sess1"], + ) + .unwrap(); + assert!(matches!( + store.try_get_goal("sess1"), + Err(GoalError::InvalidStoredValue { + field: "status", + .. + }) + )); + assert!(store.get_goal("sess1").is_none()); + } + + #[test] + fn oversized_completed_turn_progress_returns_an_explicit_conversion_error() { + let store = open_tmp(); + store + .set_goal("sess1", "avoid opaque sqlite errors", None) + .unwrap(); + + assert!(matches!( + store.record_completed_turn("sess1", i64::MAX as u64 + 1, 1), + Err(GoalError::ValueTooLarge { + field: "total_tokens_used", + .. + }) + )); + + let goal = store.try_get_goal("sess1").unwrap().unwrap(); + assert_eq!(goal.tokens_used, 0); + assert_eq!(goal.time_used_secs, 0); + assert_eq!(goal.turns_used, 0); + } + + #[test] + fn list_goals_is_fallible_and_returns_all_persisted_goals() { + let store = open_tmp(); + store.set_goal("alpha", "first", None).unwrap(); + store.set_goal("beta", "second", None).unwrap(); + + let goals = store.list_goals().unwrap(); + assert_eq!(goals.len(), 2); + assert!(goals.iter().any(|goal| goal.session_id == "alpha")); + assert!(goals.iter().any(|goal| goal.session_id == "beta")); + } + + #[test] + fn list_goals_rejects_corrupt_rows_instead_of_dropping_them() { + let store = open_tmp(); + store.set_goal("valid", "valid goal", None).unwrap(); + store.set_goal("corrupt", "corrupt goal", None).unwrap(); + store + .conn + .execute( + "UPDATE goals SET status = 'invalid' WHERE session_id = ?1", + ["corrupt"], + ) + .unwrap(); + + assert!(matches!( + store.list_goals(), + Err(GoalError::InvalidStoredValue { + field: "status", + .. + }) + )); + } + #[test] fn test_status_transitions() { let store = open_tmp(); diff --git a/src-rust/crates/query/src/goal_loop.rs b/src-rust/crates/query/src/goal_loop.rs index f3bb614..27b9c6a 100644 --- a/src-rust/crates/query/src/goal_loop.rs +++ b/src-rust/crates/query/src/goal_loop.rs @@ -2,17 +2,22 @@ // // `check_and_continue_goal` is called by the CLI REPL after each query loop // turn completes. When an active goal exists it: -// 1. Checks runaway / budget guards -// 2. Records the turn in the GoalStore +// 1. Records the turn in the GoalStore +// 2. Checks runaway / budget guards // 3. Returns `GoalContinuation::Continue { message }` with the continuation // user message to inject, signalling the caller to dispatch another turn. // // The caller (cli/src/main.rs) is responsible for the actual dispatch so that // TUI event handling and cancellation tokens stay in the right place. -use claurst_core::{goal_continuation_message, GoalStatus, GoalStore, MAX_GOAL_TURNS}; +use std::path::Path; + +use claurst_core::{ + goal_continuation_message, Goal, GoalError, GoalStatus, GoalStore, MAX_GOAL_TURNS, +}; /// Result returned to the caller after a completed query loop turn. +#[derive(Debug)] pub enum GoalContinuation { /// Inject this user message and run another turn. Continue { message: String }, @@ -37,11 +42,11 @@ impl StopReason { StopReason::GoalComplete => Some("Goal marked complete by the model.".to_string()), StopReason::Paused => None, // user-initiated, no extra message needed StopReason::BudgetLimited => Some( - "Soft token budget reached — goal paused. Use /goal resume to continue." + "Soft token budget reached — goal paused. Start a new goal with a new budget." .to_string(), ), StopReason::RunawayGuard { turns_used } => Some(format!( - "Goal paused after {} turns (runaway guard). Use /goal resume to continue.", + "Goal paused after {} turns (runaway guard). Start a new goal to continue.", turns_used )), StopReason::Error(msg) => Some(format!("Goal error: {}", msg)), @@ -52,84 +57,359 @@ impl StopReason { /// Inspect the current goal for `session_id` after a completed turn and decide /// whether to continue. /// -/// `total_tokens_used` is the session-wide cumulative token count from the -/// cost tracker (used to enforce soft budgets). +/// `total_tokens_used` is the goal-wide cumulative token count (used to +/// enforce soft budgets). /// `turn_elapsed_secs` is how long this turn took (for time accounting). pub fn check_and_continue_goal( session_id: &str, total_tokens_used: u64, turn_elapsed_secs: u64, ) -> GoalContinuation { - let store = match GoalStore::open_default() { - Some(s) => s, - None => return GoalContinuation::NoGoal, - }; + match GoalStore::open(&GoalStore::default_path()).and_then(|store| { + check_and_continue_goal_in_store(&store, session_id, total_tokens_used, turn_elapsed_secs) + }) { + Ok(decision) => decision, + Err(GoalError::NotFound { .. }) => GoalContinuation::NoGoal, + Err(error) => GoalContinuation::Stop { + reason: StopReason::Error(error.to_string()), + }, + } +} + +/// Expected-ID-aware continuation check used by turn accounting baselines. +pub fn check_and_continue_goal_for_goal( + session_id: &str, + expected_goal_id: &str, + total_tokens_used: u64, + turn_elapsed_secs: u64, +) -> GoalContinuation { + match GoalStore::open(&GoalStore::default_path()).and_then(|store| { + check_and_continue_goal_in_store_for_goal( + &store, + session_id, + expected_goal_id, + total_tokens_used, + turn_elapsed_secs, + ) + }) { + Ok(decision) => decision, + Err(GoalError::NotFound { .. }) => GoalContinuation::NoGoal, + Err(error) => GoalContinuation::Stop { + reason: StopReason::Error(error.to_string()), + }, + } +} - let goal = match store.get_goal(session_id) { - Some(g) => g, - None => return GoalContinuation::NoGoal, +fn stop_for_terminal_status(goal: &Goal) -> Option { + let reason = match goal.status { + GoalStatus::Complete => StopReason::GoalComplete, + GoalStatus::Paused => StopReason::Paused, + GoalStatus::BudgetLimited => StopReason::BudgetLimited, + GoalStatus::Active => return None, }; + Some(GoalContinuation::Stop { reason }) +} - // If model (or user) already marked complete/paused, stop. - match goal.status { - GoalStatus::Complete => { - return GoalContinuation::Stop { - reason: StopReason::GoalComplete, - }; +fn reload_after_racing_transition( + store: &GoalStore, + session_id: &str, + transition_error: GoalError, +) -> Result { + let goal = store + .try_get_goal(session_id)? + .ok_or_else(|| GoalError::NotFound { + session_id: session_id.to_string(), + })?; + stop_for_terminal_status(&goal).ok_or(transition_error) +} + +fn check_and_continue_goal_in_store( + store: &GoalStore, + session_id: &str, + total_tokens_used: u64, + turn_elapsed_secs: u64, +) -> Result { + match store.record_completed_turn(session_id, total_tokens_used, turn_elapsed_secs) { + Ok(()) => check_goal_guards(store, session_id, None), + Err(error @ GoalError::NotActive { .. }) => { + reload_after_racing_transition(store, session_id, error) } - GoalStatus::Paused => { - return GoalContinuation::Stop { - reason: StopReason::Paused, - }; + Err(error) => Err(error), + } +} + +fn check_and_continue_goal_in_store_for_goal( + store: &GoalStore, + session_id: &str, + expected_goal_id: &str, + total_tokens_used: u64, + turn_elapsed_secs: u64, +) -> Result { + match store.record_completed_turn_for_goal( + session_id, + expected_goal_id, + total_tokens_used, + turn_elapsed_secs, + ) { + Ok(()) => check_goal_guards(store, session_id, Some(expected_goal_id)), + Err(error @ GoalError::NotActive { .. }) => { + reload_after_racing_transition(store, session_id, error) } - GoalStatus::BudgetLimited => { - return GoalContinuation::Stop { - reason: StopReason::BudgetLimited, - }; + Err(error) => Err(error), + } +} + +fn check_goal_guards( + store: &GoalStore, + session_id: &str, + expected_goal_id: Option<&str>, +) -> Result { + let goal = store + .try_get_goal(session_id)? + .ok_or_else(|| GoalError::NotFound { + session_id: session_id.to_string(), + })?; + if let Some(expected_goal_id) = expected_goal_id { + if goal.id != expected_goal_id { + return Err(GoalError::Replaced { + session_id: session_id.to_string(), + expected_goal_id: expected_goal_id.to_string(), + actual_goal_id: goal.id, + }); } - GoalStatus::Active => {} } - // Runaway guard: check before incrementing so first fire is at MAX_GOAL_TURNS. + if let Some(decision) = stop_for_terminal_status(&goal) { + return Ok(decision); + } + if goal.turns_used >= MAX_GOAL_TURNS { - let _ = store.set_status(session_id, GoalStatus::Paused); - return GoalContinuation::Stop { + let transition = match expected_goal_id { + Some(expected_goal_id) => { + store.pause_active_goal_for_goal(session_id, expected_goal_id) + } + None => store.pause_active_goal(session_id), + }; + if let Err(error) = transition { + return match error { + error @ GoalError::NotActive { .. } => { + reload_after_racing_transition(store, session_id, error) + } + other => Err(other), + }; + } + return Ok(GoalContinuation::Stop { reason: StopReason::RunawayGuard { turns_used: goal.turns_used, }, - }; + }); } - // Soft token budget check. - if goal.is_over_budget(total_tokens_used) { - let _ = store.set_status(session_id, GoalStatus::BudgetLimited); - return GoalContinuation::Stop { - reason: StopReason::BudgetLimited, + if goal.is_over_budget(goal.tokens_used) { + let transition = match expected_goal_id { + Some(expected_goal_id) => { + store.budget_limit_active_goal_for_goal(session_id, expected_goal_id) + } + None => store.budget_limit_active_goal(session_id), }; + if let Err(error) = transition { + return match error { + error @ GoalError::NotActive { .. } => { + reload_after_racing_transition(store, session_id, error) + } + other => Err(other), + }; + } + return Ok(GoalContinuation::Stop { + reason: StopReason::BudgetLimited, + }); } - // Record this turn. - if let Err(e) = store.record_turn(session_id, turn_elapsed_secs) { - return GoalContinuation::Stop { - reason: StopReason::Error(e.to_string()), - }; - } + Ok(GoalContinuation::Continue { + message: goal_continuation_message(&goal), + }) +} - // Reload after the update so turns_used is current. - let goal = match store.get_goal(session_id) { - Some(g) => g, - None => return GoalContinuation::NoGoal, - }; +pub fn check_and_continue_goal_at_path( + goal_db_path: &Path, + session_id: &str, + total_tokens_used: u64, + turn_elapsed_secs: u64, +) -> Result { + let store = GoalStore::open(goal_db_path)?; + check_and_continue_goal_in_store(&store, session_id, total_tokens_used, turn_elapsed_secs) +} - // Build the continuation message. - let message = goal_continuation_message(&goal); - GoalContinuation::Continue { message } +pub fn check_and_continue_goal_at_path_for_goal( + goal_db_path: &Path, + session_id: &str, + expected_goal_id: &str, + total_tokens_used: u64, + turn_elapsed_secs: u64, +) -> Result { + let store = GoalStore::open(goal_db_path)?; + check_and_continue_goal_in_store_for_goal( + &store, + session_id, + expected_goal_id, + total_tokens_used, + turn_elapsed_secs, + ) } /// Called by GoalCompleteTool to mark the goal complete. pub fn mark_goal_complete(session_id: &str) -> Result<(), String> { let store = GoalStore::open_default().ok_or_else(|| "Could not open goal store".to_string())?; store - .set_status(session_id, GoalStatus::Complete) + .complete_active_goal(session_id) .map_err(|e| e.to_string()) } + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use super::*; + use claurst_core::{GoalError, GoalStore}; + + fn goal_path(dir: &tempfile::TempDir) -> PathBuf { + dir.path().join("goals.sqlite") + } + + #[test] + fn explicit_path_records_progress_before_budget_stop() { + let dir = tempfile::tempdir().unwrap(); + let path = goal_path(&dir); + GoalStore::open(&path) + .unwrap() + .set_goal("session", "finish", Some(100)) + .unwrap(); + + let decision = check_and_continue_goal_at_path(&path, "session", 100, 9).unwrap(); + assert!(matches!( + decision, + GoalContinuation::Stop { + reason: StopReason::BudgetLimited + } + )); + let goal = GoalStore::open(&path) + .unwrap() + .try_get_goal("session") + .unwrap() + .unwrap(); + assert_eq!(goal.tokens_used, 100); + assert_eq!(goal.time_used_secs, 9); + assert_eq!(goal.turns_used, 1); + } + + #[test] + fn explicit_path_stops_at_the_recorded_runaway_boundary() { + let dir = tempfile::tempdir().unwrap(); + let path = goal_path(&dir); + let store = GoalStore::open(&path).unwrap(); + store.set_goal("session", "finish", None).unwrap(); + for turn in 1..MAX_GOAL_TURNS { + store + .record_completed_turn("session", u64::from(turn), 1) + .unwrap(); + } + + let decision = check_and_continue_goal_at_path(&path, "session", 999, 1).unwrap(); + assert!(matches!( + decision, + GoalContinuation::Stop { + reason: StopReason::RunawayGuard { + turns_used: MAX_GOAL_TURNS + } + } + )); + } + + #[test] + fn explicit_path_reports_missing_goal() { + let dir = tempfile::tempdir().unwrap(); + let error = check_and_continue_goal_at_path(&goal_path(&dir), "missing", 0, 0).unwrap_err(); + assert!(matches!(error, GoalError::NotFound { .. })); + } + + #[test] + fn explicit_path_terminal_goals_stop_without_recording_progress() { + let dir = tempfile::tempdir().unwrap(); + let path = goal_path(&dir); + let store = GoalStore::open(&path).unwrap(); + + for (session_id, status, expected_reason) in [ + ("complete", GoalStatus::Complete, StopReason::GoalComplete), + ("paused", GoalStatus::Paused, StopReason::Paused), + ( + "budget-limited", + GoalStatus::BudgetLimited, + StopReason::BudgetLimited, + ), + ] { + store.set_goal(session_id, "finish", None).unwrap(); + match status { + GoalStatus::Complete => store.complete_active_goal(session_id).unwrap(), + GoalStatus::Paused => store.pause_active_goal(session_id).unwrap(), + GoalStatus::BudgetLimited => store.budget_limit_active_goal(session_id).unwrap(), + GoalStatus::Active => unreachable!("only terminal states are under test"), + } + + let decision = check_and_continue_goal_at_path(&path, session_id, 999, 8).unwrap(); + assert!(matches!( + (decision, expected_reason), + ( + GoalContinuation::Stop { + reason: StopReason::GoalComplete + }, + StopReason::GoalComplete + ) | ( + GoalContinuation::Stop { + reason: StopReason::Paused + }, + StopReason::Paused + ) | ( + GoalContinuation::Stop { + reason: StopReason::BudgetLimited + }, + StopReason::BudgetLimited + ) + )); + + let goal = store.try_get_goal(session_id).unwrap().unwrap(); + assert_eq!(goal.tokens_used, 0); + assert_eq!(goal.time_used_secs, 0); + assert_eq!(goal.turns_used, 0); + } + } + + #[test] + fn expected_goal_id_rejects_replacement_without_recording_progress() { + let dir = tempfile::tempdir().unwrap(); + let path = goal_path(&dir); + let store = GoalStore::open(&path).unwrap(); + let original = store + .set_goal("session", "original objective", None) + .unwrap(); + + let replacement = store + .set_goal("session", "replacement objective", None) + .unwrap(); + assert_ne!(original.id, replacement.id); + + let error = + check_and_continue_goal_at_path_for_goal(&path, "session", &original.id, 123, 7) + .unwrap_err(); + + assert!(error.to_string().contains("replaced")); + let current = GoalStore::open(&path) + .unwrap() + .try_get_goal("session") + .unwrap() + .unwrap(); + assert_eq!(current.id, replacement.id); + assert_eq!(current.tokens_used, 0); + assert_eq!(current.time_used_secs, 0); + assert_eq!(current.turns_used, 0); + } +} diff --git a/src-rust/crates/query/src/lib.rs b/src-rust/crates/query/src/lib.rs index 894af30..e1fd51a 100644 --- a/src-rust/crates/query/src/lib.rs +++ b/src-rust/crates/query/src/lib.rs @@ -30,7 +30,11 @@ pub use compact::{ CompactResult, CompactTrigger, MessageGroup, MicroCompactConfig, TokenWarningState, }; pub use cron_scheduler::start_cron_scheduler; -pub use goal_loop::{check_and_continue_goal, mark_goal_complete, GoalContinuation, StopReason}; +pub use goal_loop::{ + check_and_continue_goal, check_and_continue_goal_at_path, + check_and_continue_goal_at_path_for_goal, check_and_continue_goal_for_goal, mark_goal_complete, + GoalContinuation, StopReason, +}; pub use session_memory::{ ExtractedMemory, MemoryCandidate, MemoryCandidateStatus, MemoryCandidateStore, MemoryCategory, MemoryPersistenceOutcome, SessionMemoryExtractor, SessionMemoryState, @@ -614,6 +618,83 @@ pub enum QueryEvent { state: TokenWarningState, pct_used: f64, }, + /// A finalized assistant or tool-result transcript fact that must survive + /// context compaction and other working-history rewrites. + DurableMessage { message: Message }, +} + +fn emit_durable_message(event_tx: Option<&mpsc::UnboundedSender>, message: &Message) { + if let Some(tx) = event_tx { + let _ = tx.send(QueryEvent::DurableMessage { + message: message.clone(), + }); + } +} + +fn emit_turn_complete( + event_tx: Option<&mpsc::UnboundedSender>, + turn: u32, + stop_reason: &str, + usage: &UsageInfo, +) { + if let Some(tx) = event_tx { + let _ = tx.send(QueryEvent::TurnComplete { + turn, + stop_reason: stop_reason.to_string(), + usage: Some(usage.clone()), + }); + } +} + +fn sync_finalized_assistant_message(messages: &mut Vec, finalized: &Message) { + if let Some(uuid) = finalized.uuid.as_deref() { + if let Some(stored) = messages + .iter_mut() + .rev() + .find(|message| message.uuid.as_deref() == Some(uuid)) + { + *stored = finalized.clone(); + return; + } + } + messages.push(finalized.clone()); +} + +async fn finalize_terminal_assistant_message( + messages: &mut Vec, + assistant_msg: &mut Message, + shadow_snap: Option<&Arc>, + initial_snapshot: Option<&str>, + event_tx: Option<&mpsc::UnboundedSender>, +) { + if let (Some(snap), Some(hash)) = (shadow_snap, initial_snapshot) { + let patch = snap.patch(hash).await; + if !patch.files.is_empty() { + assistant_msg.snapshot_patch = Some(patch); + } + } + sync_finalized_assistant_message(messages, assistant_msg); + emit_durable_message(event_tx, assistant_msg); +} + +async fn finish_anthropic_stall( + messages: &mut Vec, + assistant_msg: &mut Message, + shadow_snap: Option<&Arc>, + initial_snapshot: Option<&str>, + event_tx: Option<&mpsc::UnboundedSender>, +) -> QueryOutcome { + finalize_terminal_assistant_message( + messages, + assistant_msg, + shadow_snap, + initial_snapshot, + event_tx, + ) + .await; + QueryOutcome::Error(ClaudeError::Api( + "Anthropic stream stalled after all retry attempts".to_string(), + )) } // --------------------------------------------------------------------------- @@ -1241,15 +1322,10 @@ pub async fn run_query_loop( effective_max_turns ))); } - // Return the last assistant message if any - let last_msg = messages - .last() - .cloned() - .unwrap_or_else(|| Message::assistant("Max turns reached.")); - return QueryOutcome::EndTurn { - message: last_msg, - usage: UsageInfo::default(), - }; + return QueryOutcome::Error(ClaudeError::Api(format!( + "Maximum turn limit ({}) reached", + effective_max_turns + ))); } // Check for cancellation @@ -1710,11 +1786,14 @@ pub async fn run_query_loop( let provider_stall = tokio::time::sleep(provider_stall_timeout); tokio::pin!(provider_stall); let mut provider_stream_stalled = false; + let mut provider_stream_cancelled = false; + let mut terminal_stream_error = None; loop { tokio::select! { _ = cancel_token.cancelled() => { - return QueryOutcome::Cancelled; + provider_stream_cancelled = true; + break; } _ = &mut provider_stall => { provider_stream_stalled = true; @@ -1726,6 +1805,13 @@ pub async fn run_query_loop( None => break, Some(Err(e)) => { error!(provider = %provider_id_str, error = %e, "Provider stream error"); + let claude_error: claurst_core::error::ClaudeError = e.into(); + terminal_stream_error = Some(enrich_error_notice( + claude_error, + &provider_id_str, + &model_id_str, + active_account_notice(&provider_id_str).as_deref(), + )); break; } Some(Ok(evt)) => { @@ -1793,6 +1879,25 @@ pub async fn run_query_loop( // If the stream stalled (no data for 45s), retry. if provider_stream_stalled && retries_left > 0 { + let partial_text = text_chunks.join(""); + cost_tracker.add_usage( + usage.input_tokens, + usage.output_tokens, + usage.cache_creation_input_tokens, + usage.cache_read_input_tokens, + ); + if !partial_text.is_empty() { + let mut partial_assistant = Message::assistant(partial_text); + partial_assistant.uuid = Some(msg_id.clone()); + finalize_terminal_assistant_message( + messages, + &mut partial_assistant, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + } retries_left -= 1; warn!(provider = %provider_id_str, model = %model_id_str, retries_left, "Provider stream stalled — retrying"); if let Some(ref tx) = event_tx { @@ -1804,6 +1909,12 @@ pub async fn run_query_loop( turn -= 1; continue; } + if provider_stream_stalled { + terminal_stream_error = Some(ClaudeError::Api(format!( + "Provider '{}' stream stalled after all retry attempts", + provider_id_str + ))); + } // Build the content blocks from accumulated stream data. let mut content_blocks: Vec = Vec::new(); @@ -1846,6 +1957,18 @@ pub async fn run_query_loop( snapshot_patch: None, }; + if provider_stream_cancelled { + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + return QueryOutcome::Cancelled; + } + cost_tracker.add_usage( usage.input_tokens, usage.output_tokens, @@ -1853,8 +1976,48 @@ pub async fn run_query_loop( usage.cache_read_input_tokens, ); + // Apply terminal guards before any tool execution. Provider + // responses use the same finalization path as Anthropic so + // durable history is complete before reporting a terminal + // outcome. + if let Some(limit) = config.max_budget_usd { + let spent = cost_tracker.total_cost_usd(); + if spent >= limit { + if let Some(ref tx) = event_tx { + let _ = tx.send(QueryEvent::Status(format!( + "Budget limit ${:.4} exceeded (spent ${:.4}) — stopping.", + limit, spent + ))); + } + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + return QueryOutcome::BudgetExceeded { + cost_usd: spent, + limit_usd: limit, + }; + } + } + messages.push(assistant_msg.clone()); + if let Some(error) = terminal_stream_error { + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + return QueryOutcome::Error(error); + } + // Handle tool-use turn: execute tools and loop. let tool_use_blocks: Vec<_> = content_blocks .iter() @@ -1867,12 +2030,106 @@ pub async fn run_query_loop( }) .collect(); + // PostModelTurn hooks run before tool execution, matching + // the Anthropic stream path. A hard veto makes the sampled + // response terminal and must not be converted into a + // successful end turn. + let hook_result = fire_post_sampling_hooks(&assistant_msg, &tool_ctx.config); + if !hook_result.blocking_errors.is_empty() { + if hook_result.prevent_continuation { + let hook_reason = hook_result + .blocking_errors + .iter() + .map(Message::get_all_text) + .collect::>() + .join("\n"); + for err_msg in hook_result.blocking_errors { + messages.push(err_msg); + } + if let Some(ref tx) = event_tx { + let _ = tx.send(QueryEvent::Status( + "PostModelTurn hook vetoed continuation.".to_string(), + )); + } + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + emit_turn_complete(event_tx.as_ref(), turn, &stop_str, &usage); + return QueryOutcome::Error(ClaudeError::Api(format!( + "PostModelTurn hook vetoed continuation: {}", + hook_reason + ))); + } + for err_msg in hook_result.blocking_errors { + debug!("PostModelTurn hook injecting error message"); + messages.push(err_msg); + } + } + + if stop_str == "max_tokens" { + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + emit_turn_complete(event_tx.as_ref(), turn, &stop_str, &usage); + return QueryOutcome::MaxTokens { + partial_message: assistant_msg, + usage, + }; + } + + if stop_str == "content_filtered" + || !matches!( + stop_str.as_str(), + "end_turn" | "tool_use" | "max_tokens" | "stop_sequence" + ) + { + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + emit_turn_complete(event_tx.as_ref(), turn, &stop_str, &usage); + return QueryOutcome::Error(ClaudeError::Api(format!( + "Provider '{}' returned unsupported stop reason '{}'", + provider_id_str, stop_str + ))); + } + + if stop_str == "tool_use" && tool_use_blocks.is_empty() { + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + emit_turn_complete(event_tx.as_ref(), turn, &stop_str, &usage); + return QueryOutcome::Error(ClaudeError::Api( + "Provider returned tool_use without tool blocks".to_string(), + )); + } + // Execute tool-use blocks through the same hook-aware path used by // Anthropic streaming. Some OpenAI-compatible providers return // finish_reason "stop" even when tool calls are present, so the // block presence is authoritative here, but execution must still // enforce PreToolUse/plugin policy and emit post-hook events. if !tool_use_blocks.is_empty() { + emit_durable_message(event_tx.as_ref(), &assistant_msg); let tool_results = execute_tool_blocks_with_hooks( tool_use_blocks, tools, @@ -1880,13 +2137,15 @@ pub async fn run_query_loop( event_tx.as_ref(), ) .await; - messages.push(Message { + let tool_result_message = Message { role: claurst_core::types::Role::User, content: claurst_core::types::MessageContent::Blocks(tool_results), uuid: None, cost: None, snapshot_patch: None, - }); + }; + messages.push(tool_result_message.clone()); + emit_durable_message(event_tx.as_ref(), &tool_result_message); // Real tool work happened; give the next // announce-then-stop stall a fresh recovery budget. stall_recovery_count = 0; @@ -1898,10 +2157,9 @@ pub async fn run_query_loop( // that executed no tools and whose text is empty or ends // on an imminent-action announcement gets a bounded // continuation nudge instead of ending the turn. - if stall_recovery_enabled() - && stall_recovery_count < STALL_RECOVERY_LIMIT - && is_stalled_announcement(&combined_text) - { + let stalled_announcement = + stall_recovery_enabled() && is_stalled_announcement(&combined_text); + if stalled_announcement && stall_recovery_count < STALL_RECOVERY_LIMIT { stall_recovery_count += 1; warn!( attempt = stall_recovery_count, @@ -1916,6 +2174,7 @@ pub async fn run_query_loop( stall_recovery_count, STALL_RECOVERY_LIMIT ))); } + emit_durable_message(event_tx.as_ref(), &assistant_msg); messages.push(Message::user(STALL_RECOVERY_MSG)); continue; } @@ -1954,22 +2213,22 @@ pub async fn run_query_loop( } } - if let Some(ref tx) = event_tx { - let _ = tx.send(QueryEvent::TurnComplete { - stop_reason: stop_str.clone(), - turn, - usage: Some(usage.clone()), - }); - } - - // Attach snapshot patch covering all file changes this query. - if let (Some(ref snap), Some(ref hash)) = (&shadow_snap, &initial_snapshot) { - let patch = snap.patch(hash).await; - if !patch.files.is_empty() { - assistant_msg.snapshot_patch = Some(patch); - } + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + + emit_turn_complete(event_tx.as_ref(), turn, &stop_str, &usage); + if stalled_announcement && stall_recovery_count >= STALL_RECOVERY_LIMIT { + return QueryOutcome::Error(ClaudeError::Api(format!( + "Provider '{}' repeatedly ended without completing the turn", + provider_id_str + ))); } - return QueryOutcome::EndTurn { message: assistant_msg, usage, @@ -2040,18 +2299,14 @@ pub async fn run_query_loop( // (some providers are slow; we don't want to give up too early). const STALL_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(45); let mut accumulator = StreamAccumulator::new(); + let mut terminal_stream_error = None; let stall_deadline = tokio::time::sleep(STALL_TIMEOUT); tokio::pin!(stall_deadline); + let mut stream_cancelled = false; let stream_stalled = loop { tokio::select! { - _ = cancel_token.cancelled() => { - return QueryOutcome::Cancelled; - } - _ = &mut stall_deadline => { - // No data for 45s — stall detected - break true; - } + biased; event = stream_rx.recv() => { // Reset stall timer on every received event. stall_deadline.as_mut().reset(tokio::time::Instant::now() + STALL_TIMEOUT); @@ -2064,6 +2319,11 @@ pub async fn run_query_loop( warn!(model = %effective_model, "API overloaded"); } error!(error_type, message, "Stream error"); + terminal_stream_error = Some(ClaudeError::Api(format!( + "Anthropic stream error ({}): {}", + error_type, message + ))); + break false; } AnthropicStreamEvent::MessageStop => break false, _ => {} @@ -2072,10 +2332,45 @@ pub async fn run_query_loop( None => break false, // Stream ended } } + _ = cancel_token.cancelled() => { + // The stream handler can publish a delta before this loop + // has accumulated it. Drain buffered events so cancellation + // finalizes the assistant the user actually saw. + stream_rx.close(); + while let Ok(event) = stream_rx.try_recv() { + accumulator.on_event(&event); + } + stream_cancelled = true; + break false; + } + _ = &mut stall_deadline => { + // No data for 45s — stall detected + break true; + } } }; if stream_stalled && retries_left > 0 { + let (mut assistant_msg, usage, _) = accumulator.finish(); + cost_tracker.add_usage( + usage.input_tokens, + usage.output_tokens, + usage.cache_creation_input_tokens, + usage.cache_read_input_tokens, + ); + if !assistant_msg.get_all_text().is_empty() { + if assistant_msg.uuid.is_none() { + assistant_msg.uuid = Some(uuid::Uuid::new_v4().to_string()); + } + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + } retries_left -= 1; warn!(model = %effective_model, retries_left, "Stream stalled — retrying request"); if let Some(ref tx) = event_tx { @@ -2087,8 +2382,33 @@ pub async fn run_query_loop( turn -= 1; // don't count this stalled attempt continue; } - let (mut assistant_msg, usage, stop_reason) = accumulator.finish(); + if assistant_msg.uuid.is_none() { + assistant_msg.uuid = Some(uuid::Uuid::new_v4().to_string()); + } + + if stream_cancelled { + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + return QueryOutcome::Cancelled; + } + + if stream_stalled { + return finish_anthropic_stall( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + } // Track costs cost_tracker.add_usage( @@ -2108,6 +2428,14 @@ pub async fn run_query_loop( limit, spent ))); } + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; return QueryOutcome::BudgetExceeded { cost_usd: spent, limit_usd: limit, @@ -2118,22 +2446,26 @@ pub async fn run_query_loop( // Append assistant message to conversation messages.push(assistant_msg.clone()); - // If the provider returned an unknown stop reason but the assistant - // message contains tool_use blocks, treat it as tool_use so we don't - // silently end the turn (issue #149: agent stops after tool call for - // providers that emit non-standard finish reasons). + if let Some(error) = terminal_stream_error { + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + return QueryOutcome::Error(error); + } + + // Unknown stop reasons are terminal failures. In particular, do not + // infer tool continuation from blocks when the provider could not + // report a recognized completion reason. let raw_stop = stop_reason.as_deref().unwrap_or("end_turn"); let stop = match raw_stop { "end_turn" | "tool_use" | "max_tokens" | "stop_sequence" | "content_filtered" => { raw_stop } - _ if !assistant_msg.get_tool_use_blocks().is_empty() => { - warn!( - stop_reason = raw_stop, - "Unknown stop reason with tool_use blocks present; treating as tool_use" - ); - "tool_use" - } _ => raw_stop, }; @@ -2143,6 +2475,12 @@ pub async fn run_query_loop( let hook_result = fire_post_sampling_hooks(&assistant_msg, &tool_ctx.config); if !hook_result.blocking_errors.is_empty() { if hook_result.prevent_continuation { + let hook_reason = hook_result + .blocking_errors + .iter() + .map(Message::get_all_text) + .collect::>() + .join("\n"); // Hard veto: push the errors into the conversation and abort. for err_msg in hook_result.blocking_errors { messages.push(err_msg); @@ -2152,14 +2490,19 @@ pub async fn run_query_loop( "PostModelTurn hook vetoed continuation.".to_string(), )); } - let last = messages - .last() - .cloned() - .unwrap_or_else(|| Message::assistant("Hook blocked continuation.")); - return QueryOutcome::EndTurn { - message: last, - usage, - }; + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + emit_turn_complete(event_tx.as_ref(), turn, stop, &usage); + return QueryOutcome::Error(ClaudeError::Api(format!( + "PostModelTurn hook vetoed continuation: {}", + hook_reason + ))); } // Soft errors: inject them so the model can react next turn. for err_msg in hook_result.blocking_errors { @@ -2207,50 +2550,109 @@ pub async fn run_query_loop( "Compacting context... (emergency collapse)".to_string(), )); } - match compact::context_collapse(std::mem::take(messages), client, config).await { - Ok(result) => { + let original_messages = std::mem::take(messages); + let collapse_result = { + let collapse = + compact::context_collapse(original_messages.clone(), client, config); + tokio::pin!(collapse); + tokio::select! { + _ = cancel_token.cancelled() => None, + result = &mut collapse => Some(result), + } + }; + match collapse_result { + None => { + *messages = original_messages; + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + return QueryOutcome::Cancelled; + } + Some(Ok(result)) => { *messages = result.messages; info!( tokens_freed = result.tokens_freed, "Context-collapse complete" ); } - Err(e) => { + Some(Err(e)) => { warn!(error = %e, "Context-collapse failed"); - // Put messages back on failure (mem::take drained them). - // We can't recover them here — re-run auto-compact as fallback. + *messages = original_messages; + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + return QueryOutcome::Error(e); } } } else if compact::should_compact(usage.input_tokens, context_limit) { if let Some(ref tx) = event_tx { let _ = tx.send(QueryEvent::Status("Compacting context...".to_string())); } - match compact::reactive_compact( - std::mem::take(messages), - client, - config, - cancel_token.clone(), - &[], - ) - .await - { - Ok(result) => { + let original_messages = std::mem::take(messages); + let compact_result = { + let compact = compact::reactive_compact( + original_messages.clone(), + client, + config, + cancel_token.clone(), + &[], + ); + tokio::pin!(compact); + tokio::select! { + _ = cancel_token.cancelled() => None, + result = &mut compact => Some(result), + } + }; + match compact_result { + Some(Ok(result)) => { *messages = result.messages; info!( tokens_freed = result.tokens_freed, "Reactive compact complete" ); } - Err(claurst_core::error::ClaudeError::Cancelled) => { + None | Some(Err(claurst_core::error::ClaudeError::Cancelled)) => { warn!("Reactive compact was cancelled"); + *messages = original_messages; + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + return QueryOutcome::Cancelled; } - Err(e) => { + Some(Err(e)) => { warn!(error = %e, "Reactive compact failed"); + *messages = original_messages; + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + return QueryOutcome::Error(e); } } } } else if stop == "end_turn" || stop == "tool_use" { // Proactive auto-compact (original path, used when reactive compact is off). + let compaction_required = + compact::should_auto_compact(usage.input_tokens, &config.model, &compact_state); if let Some(new_msgs) = compact::auto_compact_if_needed( client, messages, @@ -2266,6 +2668,18 @@ pub async fn run_query_loop( "Context compacted to stay within limits.".to_string(), )); } + } else if compaction_required { + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + return QueryOutcome::Error(ClaudeError::Other( + "Auto-compaction failed after a completed response".to_string(), + )); } } @@ -2274,21 +2688,14 @@ pub async fn run_query_loop( // imminent-action announcement gets a bounded continuation nudge // instead of ending the turn. Decided before TurnComplete emission // so recovered rounds stay invisible, like max_tokens recovery. - let stall_recovery_pending = stop == "end_turn" + let stalled_announcement = stop == "end_turn" && stall_recovery_enabled() - && stall_recovery_count < STALL_RECOVERY_LIMIT && assistant_msg.get_tool_use_blocks().is_empty() && is_stalled_announcement(&assistant_msg.get_all_text()); - - if should_emit_turn_complete(stop, max_tokens_recovery_count, stall_recovery_pending) { - if let Some(ref tx) = event_tx { - let _ = tx.send(QueryEvent::TurnComplete { - turn, - stop_reason: stop.to_string(), - usage: Some(usage.clone()), - }); - } - } + let stall_recovery_pending = + stalled_announcement && stall_recovery_count < STALL_RECOVERY_LIMIT; + let stall_recovery_exhausted = + stalled_announcement && stall_recovery_count >= STALL_RECOVERY_LIMIT; // Helper closure for firing the Stop hook. macro_rules! fire_stop_hook { @@ -2330,6 +2737,8 @@ pub async fn run_query_loop( // The stalled assistant message is already in the history, // so the nudge reads in context. Stop hooks are not fired: // the turn is not over. + sync_finalized_assistant_message(messages, &assistant_msg); + emit_durable_message(event_tx.as_ref(), &assistant_msg); messages.push(Message::user(STALL_RECOVERY_MSG)); continue; } @@ -2457,14 +2866,23 @@ pub async fn run_query_loop( } } - // Attach snapshot patch covering all file changes this query. - if let (Some(ref snap), Some(ref hash)) = (&shadow_snap, &initial_snapshot) { - let patch = snap.patch(hash).await; - if !patch.files.is_empty() { - assistant_msg.snapshot_patch = Some(patch); - } + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + if should_emit_turn_complete(stop, max_tokens_recovery_count, false) { + emit_turn_complete(event_tx.as_ref(), turn, stop, &usage); } + if stall_recovery_exhausted { + return QueryOutcome::Error(ClaudeError::Api( + "Model repeatedly ended without completing the turn".to_string(), + )); + } return QueryOutcome::EndTurn { message: assistant_msg, usage, @@ -2491,6 +2909,8 @@ pub async fn run_query_loop( } // The partial assistant message must be in the history so // the continuation makes sense to the model. + sync_finalized_assistant_message(messages, &assistant_msg); + emit_durable_message(event_tx.as_ref(), &assistant_msg); messages.push(Message::user(MAX_TOKENS_RECOVERY_MSG)); continue; } @@ -2499,6 +2919,17 @@ pub async fn run_query_loop( "max_tokens recovery exhausted after {} attempts", MAX_TOKENS_RECOVERY_LIMIT ); + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + if should_emit_turn_complete(stop, max_tokens_recovery_count, false) { + emit_turn_complete(event_tx.as_ref(), turn, stop, &usage); + } return QueryOutcome::MaxTokens { partial_message: assistant_msg, usage, @@ -2515,10 +2946,18 @@ pub async fn run_query_loop( let tool_blocks = assistant_msg.get_tool_use_blocks(); if tool_blocks.is_empty() { // Shouldn't happen but treat as end_turn - return QueryOutcome::EndTurn { - message: assistant_msg, - usage, - }; + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + emit_turn_complete(event_tx.as_ref(), turn, stop, &usage); + return QueryOutcome::Error(ClaudeError::Api( + "Provider returned tool_use without tool blocks".to_string(), + )); } let tool_invocations: Vec<_> = tool_blocks @@ -2531,6 +2970,8 @@ pub async fn run_query_loop( } }) .collect(); + sync_finalized_assistant_message(messages, &assistant_msg); + emit_durable_message(event_tx.as_ref(), &assistant_msg); let result_blocks = execute_tool_blocks_with_hooks( tool_invocations, tools, @@ -2540,7 +2981,9 @@ pub async fn run_query_loop( .await; // Append tool results as a user message - messages.push(Message::user_blocks(result_blocks)); + let tool_result_message = Message::user_blocks(result_blocks); + messages.push(tool_result_message.clone()); + emit_durable_message(event_tx.as_ref(), &tool_result_message); // Continue the loop to send results back to the model continue; @@ -2552,38 +2995,60 @@ pub async fn run_query_loop( &tool_ctx.config, tool_ctx.working_dir.clone(), ); - if let (Some(ref snap), Some(ref hash)) = (&shadow_snap, &initial_snapshot) { - let patch = snap.patch(hash).await; - if !patch.files.is_empty() { - assistant_msg.snapshot_patch = Some(patch); - } + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + if should_emit_turn_complete(stop, max_tokens_recovery_count, false) { + emit_turn_complete(event_tx.as_ref(), turn, stop, &usage); } return QueryOutcome::EndTurn { message: assistant_msg, usage, }; } + "content_filtered" => { + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + if should_emit_turn_complete(stop, max_tokens_recovery_count, false) { + emit_turn_complete(event_tx.as_ref(), turn, stop, &usage); + } + return QueryOutcome::Error(ClaudeError::Api( + "Provider content filter blocked the response".to_string(), + )); + } other => { - warn!( - stop_reason = other, - "Unknown stop reason, treating as end_turn" - ); + warn!(stop_reason = other, "Unknown stop reason"); fire_stop_hook!(assistant_msg); let _bg = stop_hooks_with_full_behavior( &assistant_msg, &tool_ctx.config, tool_ctx.working_dir.clone(), ); - if let (Some(ref snap), Some(ref hash)) = (&shadow_snap, &initial_snapshot) { - let patch = snap.patch(hash).await; - if !patch.files.is_empty() { - assistant_msg.snapshot_patch = Some(patch); - } + finalize_terminal_assistant_message( + messages, + &mut assistant_msg, + shadow_snap.as_ref(), + initial_snapshot.as_deref(), + event_tx.as_ref(), + ) + .await; + if should_emit_turn_complete(stop, max_tokens_recovery_count, false) { + emit_turn_complete(event_tx.as_ref(), turn, stop, &usage); } - return QueryOutcome::EndTurn { - message: assistant_msg, - usage, - }; + return QueryOutcome::Error(ClaudeError::Api(format!( + "Provider returned unsupported stop reason '{other}'" + ))); } } } @@ -2941,36 +3406,1868 @@ pub async fn run_single_query( #[cfg(test)] mod tests { + use std::collections::{HashSet, VecDeque}; + use std::pin::Pin; + use std::process::Command; + use std::sync::{Arc, Mutex}; + use super::*; - use claurst_api::SystemPrompt; + use async_trait::async_trait; + use claurst_api::provider_types::{ + ProviderCapabilities, ProviderRequest, ProviderResponse, ProviderStatus, + StopReason as ProviderStopReason, StreamEvent, SystemPromptStyle, + }; + use claurst_api::{LlmProvider, ModelInfo, ProviderError, SystemPrompt}; + use claurst_core::config::{Config, HookEntry, HookEvent, PermissionMode}; + use claurst_core::permissions::AutoPermissionHandler; + use claurst_core::provider_id::ProviderId; + use claurst_core::types::{MessageContent, Role}; + use claurst_tools::{PermissionLevel, ToolResult}; + use futures::{Stream, StreamExt}; + + const SCRIPTED_PROVIDER_ID: &str = "scripted-goal-journal"; + + struct ScriptedProvider { + id: ProviderId, + rounds: Mutex>, + } - fn make_config(sys: Option<&str>, append: Option<&str>) -> QueryConfig { - QueryConfig { - model: "claude-sonnet-4-6".to_string(), - max_tokens: 4096, - max_turns: 10, - system_prompt: sys.map(String::from), - append_system_prompt: append.map(String::from), - output_style: claurst_core::system_prompt::OutputStyle::Default, - output_style_prompt: None, - working_directory: None, - thinking_budget: None, - temperature: None, - tool_result_budget: 50_000, - effort_level: None, - command_queue: None, - skill_index: None, - max_budget_usd: None, - fallback_model: None, - provider_registry: None, - agent_name: None, - agent_definition: None, - model_registry: None, - managed_agents: None, - preserve_selected_model: false, + enum ScriptedRound { + Events(Vec>), + PendingAfter(Vec>), + } + + impl ScriptedProvider { + fn new(rounds: Vec>) -> Self { + Self { + id: ProviderId::new(SCRIPTED_PROVIDER_ID), + rounds: Mutex::new( + rounds + .into_iter() + .map(|round| ScriptedRound::Events(round.into_iter().map(Ok).collect())) + .collect(), + ), + } + } + + fn with_results(rounds: Vec>>) -> Self { + Self { + id: ProviderId::new(SCRIPTED_PROVIDER_ID), + rounds: Mutex::new(rounds.into_iter().map(ScriptedRound::Events).collect()), + } + } + + fn partial_then_pending(events: Vec) -> Self { + Self { + id: ProviderId::new(SCRIPTED_PROVIDER_ID), + rounds: Mutex::new( + vec![ScriptedRound::PendingAfter( + events.into_iter().map(Ok).collect(), + )] + .into(), + ), + } + } + } + + #[async_trait] + impl LlmProvider for ScriptedProvider { + fn id(&self) -> &ProviderId { + &self.id + } + + fn name(&self) -> &str { + "Scripted goal journal provider" + } + + async fn create_message( + &self, + _request: ProviderRequest, + ) -> Result { + Ok(ProviderResponse { + id: "unused".to_string(), + content: Vec::new(), + stop_reason: ProviderStopReason::EndTurn, + usage: UsageInfo::default(), + model: "scripted".to_string(), + }) + } + + async fn create_message_stream( + &self, + _request: ProviderRequest, + ) -> Result< + Pin> + Send>>, + ProviderError, + > { + let round = self + .rounds + .lock() + .unwrap() + .pop_front() + .expect("scripted provider ran out of rounds"); + let stream: Pin> + Send>> = + match round { + ScriptedRound::Events(events) => Box::pin(futures::stream::iter(events)), + ScriptedRound::PendingAfter(events) => { + Box::pin(futures::stream::iter(events).chain(futures::stream::pending())) + } + }; + Ok(stream) + } + + async fn list_models(&self) -> Result, ProviderError> { + Ok(Vec::new()) + } + + async fn health_check(&self) -> Result { + Ok(ProviderStatus::Healthy) + } + + fn capabilities(&self) -> ProviderCapabilities { + ProviderCapabilities { + streaming: true, + tool_calling: true, + thinking: false, + image_input: false, + pdf_input: false, + audio_input: false, + video_input: false, + caching: false, + structured_output: false, + system_prompt_style: SystemPromptStyle::TopLevel, + } } } + struct WriteFixtureTool; + + #[async_trait] + impl Tool for WriteFixtureTool { + fn name(&self) -> &str { + "WriteFixture" + } + + fn description(&self) -> &str { + "Writes the durable-message journal test fixture" + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::None + } + + fn input_schema(&self) -> Value { + serde_json::json!({"type": "object"}) + } + + async fn execute(&self, _input: Value, ctx: &ToolContext) -> ToolResult { + match tokio::fs::write(ctx.working_dir.join("fixture.txt"), "written\n").await { + Ok(()) => ToolResult::success("wrote fixture.txt"), + Err(error) => ToolResult::error(error.to_string()), + } + } + } + + fn scripted_round( + id: &str, + blocks: Vec, + stop_reason: ProviderStopReason, + ) -> Vec { + let mut events = vec![StreamEvent::MessageStart { + id: id.to_string(), + model: "scripted".to_string(), + usage: UsageInfo { + input_tokens: 1, + ..UsageInfo::default() + }, + }]; + for (index, block) in blocks.into_iter().enumerate() { + events.push(StreamEvent::ContentBlockStart { + index, + content_block: block.clone(), + }); + match block { + ContentBlock::Text { text } => { + events.push(StreamEvent::TextDelta { index, text }); + } + ContentBlock::ToolUse { input, .. } => { + events.push(StreamEvent::InputJsonDelta { + index, + partial_json: input.to_string(), + }); + } + _ => {} + } + events.push(StreamEvent::ContentBlockStop { index }); + } + events.extend([ + StreamEvent::MessageDelta { + stop_reason: Some(stop_reason), + usage: Some(UsageInfo { + output_tokens: 1, + ..UsageInfo::default() + }), + }, + StreamEvent::MessageStop, + ]); + events + } + + fn init_test_repository(path: &std::path::Path) { + let run = |args: &[&str]| { + let output = Command::new("git") + .args(args) + .current_dir(path) + .output() + .unwrap(); + assert!( + output.status.success(), + "git {args:?} failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + }; + run(&["init", "-b", "main"]); + run(&["config", "user.email", "test@example.invalid"]); + run(&["config", "user.name", "Test"]); + std::fs::write(path.join("seed.txt"), "seed\n").unwrap(); + run(&["add", "seed.txt"]); + run(&["commit", "-m", "seed"]); + } + + fn is_tool_result_carrier(message: &Message) -> bool { + matches!(message.role, Role::User) + && matches!( + &message.content, + MessageContent::Blocks(blocks) + if blocks.iter().any(|block| matches!(block, ContentBlock::ToolResult { .. })) + ) + } + + fn unique_message_fingerprints(messages: &[Message]) -> HashSet { + messages + .iter() + .map(|message| serde_json::to_string(message).unwrap()) + .collect() + } + + async fn run_scripted_provider( + provider: ScriptedProvider, + ) -> (QueryOutcome, Vec, Vec) { + run_scripted_provider_with( + provider, + make_config(None, None), + Config { + provider: Some(SCRIPTED_PROVIDER_ID.to_string()), + permission_mode: PermissionMode::BypassPermissions, + ..Config::default() + }, + &[], + ) + .await + } + + async fn run_scripted_provider_with( + provider: ScriptedProvider, + mut config: QueryConfig, + tool_config: Config, + tools: &[Box], + ) -> (QueryOutcome, Vec, Vec) { + let temp = tempfile::tempdir().unwrap(); + let mut registry = claurst_api::ProviderRegistry::new(); + registry.register(Arc::new(provider)); + config.model = "scripted".to_string(); + config.provider_registry = Some(Arc::new(registry)); + let cost_tracker = claurst_core::cost::CostTracker::new(); + let tool_ctx = ToolContext { + working_dir: temp.path().to_path_buf(), + permission_mode: PermissionMode::BypassPermissions, + permission_handler: Arc::new(AutoPermissionHandler { + mode: PermissionMode::BypassPermissions, + }), + cost_tracker: cost_tracker.clone(), + session_id: "provider-outcome".to_string(), + file_history: Default::default(), + current_turn: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + non_interactive: true, + mcp_manager: None, + config: tool_config, + managed_agent_config: None, + completion_notifier: None, + pending_permissions: None, + permission_manager: None, + user_question_tx: None, + }; + let client = claurst_api::AnthropicClient::new(claurst_api::client::ClientConfig { + api_key: "unused-test-key".to_string(), + ..Default::default() + }) + .unwrap(); + let mut messages = vec![Message::user("test provider outcome")]; + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + let outcome = run_query_loop( + &client, + &mut messages, + tools, + &tool_ctx, + &config, + cost_tracker, + Some(event_tx), + tokio_util::sync::CancellationToken::new(), + None, + ) + .await; + let mut events = Vec::new(); + while let Ok(event) = event_rx.try_recv() { + events.push(event); + } + (outcome, events, messages) + } + + async fn run_scripted_provider_until_partial_then_cancel( + provider: ScriptedProvider, + ) -> (QueryOutcome, Vec, Vec) { + let temp = tempfile::tempdir().unwrap(); + let mut registry = claurst_api::ProviderRegistry::new(); + registry.register(Arc::new(provider)); + let mut config = make_config(None, None); + config.model = "scripted".to_string(); + config.provider_registry = Some(Arc::new(registry)); + let cost_tracker = claurst_core::cost::CostTracker::new(); + let tool_ctx = ToolContext { + working_dir: temp.path().to_path_buf(), + permission_mode: PermissionMode::BypassPermissions, + permission_handler: Arc::new(AutoPermissionHandler { + mode: PermissionMode::BypassPermissions, + }), + cost_tracker: cost_tracker.clone(), + session_id: "provider-cancellation".to_string(), + file_history: Default::default(), + current_turn: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + non_interactive: true, + mcp_manager: None, + config: Config { + provider: Some(SCRIPTED_PROVIDER_ID.to_string()), + permission_mode: PermissionMode::BypassPermissions, + ..Config::default() + }, + managed_agent_config: None, + completion_notifier: None, + pending_permissions: None, + permission_manager: None, + user_question_tx: None, + }; + let client = claurst_api::AnthropicClient::new(claurst_api::client::ClientConfig { + api_key: "unused-test-key".to_string(), + ..Default::default() + }) + .unwrap(); + let mut messages = vec![Message::user("test provider cancellation")]; + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + let cancel_token = tokio_util::sync::CancellationToken::new(); + let mut events = Vec::new(); + let outcome = { + let query = run_query_loop( + &client, + &mut messages, + &[], + &tool_ctx, + &config, + cost_tracker, + Some(event_tx), + cancel_token.clone(), + None, + ); + tokio::pin!(query); + loop { + tokio::select! { + outcome = &mut query => break outcome, + event = event_rx.recv() => { + let event = event.expect("query event channel must stay open before cancellation"); + let received_partial = matches!( + &event, + QueryEvent::Stream(AnthropicStreamEvent::ContentBlockDelta { + delta: claurst_api::streaming::ContentDelta::TextDelta { text }, + .. + }) if text == "partial before cancellation" + ); + events.push(event); + if received_partial { + cancel_token.cancel(); + } + } + } + } + }; + while let Ok(event) = event_rx.try_recv() { + events.push(event); + } + (outcome, events, messages) + } + + async fn serve_anthropic_sse_sequence(bodies: Vec<&'static str>) -> String { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind Anthropic test listener"); + let addr = listener.local_addr().expect("Anthropic test listener addr"); + tokio::spawn(async move { + for body in bodies { + let (mut stream, _) = listener + .accept() + .await + .expect("accept Anthropic test connection"); + let mut request = [0u8; 8192]; + let _ = stream.read(&mut request).await; + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ); + stream + .write_all(response.as_bytes()) + .await + .expect("write Anthropic test response"); + let _ = stream.shutdown().await; + } + }); + format!("http://{addr}") + } + + async fn serve_anthropic_sse(body: &'static str) -> String { + serve_anthropic_sse_sequence(vec![body]).await + } + + async fn serve_anthropic_partial_sse(body: &'static str) -> String { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind partial Anthropic test listener"); + let addr = listener + .local_addr() + .expect("partial Anthropic test listener addr"); + tokio::spawn(async move { + let (mut stream, _) = listener + .accept() + .await + .expect("accept partial Anthropic test connection"); + let mut request = [0u8; 8192]; + let _ = stream.read(&mut request).await; + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: keep-alive\r\n\r\n{}", + body + ); + stream + .write_all(response.as_bytes()) + .await + .expect("write partial Anthropic test response"); + stream + .flush() + .await + .expect("flush partial Anthropic response"); + std::future::pending::<()>().await; + }); + format!("http://{addr}") + } + + async fn serve_anthropic_sse_then_wait_for_compaction( + body: &'static str, + ) -> (String, tokio::sync::oneshot::Receiver<()>) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind compaction Anthropic test listener"); + let addr = listener + .local_addr() + .expect("compaction Anthropic test listener addr"); + let (compaction_started_tx, compaction_started_rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let (mut first_stream, _) = listener + .accept() + .await + .expect("accept initial Anthropic test connection"); + let mut request = [0u8; 8192]; + let _ = first_stream.read(&mut request).await; + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ); + first_stream + .write_all(response.as_bytes()) + .await + .expect("write initial Anthropic test response"); + let _ = first_stream.shutdown().await; + + let (mut compact_stream, _) = listener + .accept() + .await + .expect("accept compaction Anthropic test connection"); + let _ = compact_stream.read(&mut request).await; + let _ = compaction_started_tx.send(()); + std::future::pending::<()>().await; + }); + (format!("http://{addr}"), compaction_started_rx) + } + + async fn run_anthropic_sse( + body: &'static str, + tool_config: Config, + ) -> (QueryOutcome, Vec, Vec) { + let temp = tempfile::tempdir().unwrap(); + let api_base = serve_anthropic_sse(body).await; + let client = claurst_api::AnthropicClient::new(claurst_api::client::ClientConfig { + api_key: "test-key".to_string(), + api_base, + max_retries: 0, + ..Default::default() + }) + .unwrap(); + let cost_tracker = claurst_core::cost::CostTracker::new(); + let tool_ctx = ToolContext { + working_dir: temp.path().to_path_buf(), + permission_mode: PermissionMode::BypassPermissions, + permission_handler: Arc::new(AutoPermissionHandler { + mode: PermissionMode::BypassPermissions, + }), + cost_tracker: cost_tracker.clone(), + session_id: "anthropic-outcome".to_string(), + file_history: Default::default(), + current_turn: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + non_interactive: true, + mcp_manager: None, + config: tool_config, + managed_agent_config: None, + completion_notifier: None, + pending_permissions: None, + permission_manager: None, + user_question_tx: None, + }; + let mut messages = vec![Message::user("test Anthropic outcome")]; + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + let outcome = run_query_loop( + &client, + &mut messages, + &[], + &tool_ctx, + &make_config(None, None), + cost_tracker, + Some(event_tx), + tokio_util::sync::CancellationToken::new(), + None, + ) + .await; + let mut events = Vec::new(); + while let Ok(event) = event_rx.try_recv() { + events.push(event); + } + (outcome, events, messages) + } + + async fn run_anthropic_until_partial_then_cancel( + body: &'static str, + ) -> (QueryOutcome, Vec, Vec) { + let temp = tempfile::tempdir().unwrap(); + let api_base = serve_anthropic_partial_sse(body).await; + let client = claurst_api::AnthropicClient::new(claurst_api::client::ClientConfig { + api_key: "test-key".to_string(), + api_base, + max_retries: 0, + ..Default::default() + }) + .unwrap(); + let cost_tracker = claurst_core::cost::CostTracker::new(); + let tool_ctx = ToolContext { + working_dir: temp.path().to_path_buf(), + permission_mode: PermissionMode::BypassPermissions, + permission_handler: Arc::new(AutoPermissionHandler { + mode: PermissionMode::BypassPermissions, + }), + cost_tracker: cost_tracker.clone(), + session_id: "anthropic-cancellation".to_string(), + file_history: Default::default(), + current_turn: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + non_interactive: true, + mcp_manager: None, + config: Config::default(), + managed_agent_config: None, + completion_notifier: None, + pending_permissions: None, + permission_manager: None, + user_question_tx: None, + }; + let mut messages = vec![Message::user("test Anthropic cancellation")]; + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + let cancel_token = tokio_util::sync::CancellationToken::new(); + let query_config = make_config(None, None); + let mut events = Vec::new(); + let outcome = { + let query = run_query_loop( + &client, + &mut messages, + &[], + &tool_ctx, + &query_config, + cost_tracker, + Some(event_tx), + cancel_token.clone(), + None, + ); + tokio::pin!(query); + loop { + tokio::select! { + outcome = &mut query => break outcome, + event = event_rx.recv() => { + let event = event.expect("query event channel must stay open before cancellation"); + let completed_partial_block = matches!( + &event, + QueryEvent::Stream(AnthropicStreamEvent::ContentBlockStop { index: 0 }) + ); + events.push(event); + if completed_partial_block { + cancel_token.cancel(); + } + } + } + } + }; + while let Ok(event) = event_rx.try_recv() { + events.push(event); + } + (outcome, events, messages) + } + + async fn run_anthropic_sse_sequence_in_worktree( + bodies: Vec<&'static str>, + working_dir: &std::path::Path, + tool_config: Config, + query_config: QueryConfig, + tools: &[Box], + ) -> (QueryOutcome, Vec, Vec) { + let api_base = serve_anthropic_sse_sequence(bodies).await; + let client = claurst_api::AnthropicClient::new(claurst_api::client::ClientConfig { + api_key: "test-key".to_string(), + api_base, + max_retries: 0, + ..Default::default() + }) + .unwrap(); + let cost_tracker = claurst_core::cost::CostTracker::new(); + let tool_ctx = ToolContext { + working_dir: working_dir.to_path_buf(), + permission_mode: PermissionMode::BypassPermissions, + permission_handler: Arc::new(AutoPermissionHandler { + mode: PermissionMode::BypassPermissions, + }), + cost_tracker: cost_tracker.clone(), + session_id: "anthropic-snapshot".to_string(), + file_history: Default::default(), + current_turn: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + non_interactive: true, + mcp_manager: None, + config: tool_config, + managed_agent_config: None, + completion_notifier: None, + pending_permissions: None, + permission_manager: None, + user_question_tx: None, + }; + let mut messages = vec![Message::user("exercise terminal snapshot")]; + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + let outcome = run_query_loop( + &client, + &mut messages, + tools, + &tool_ctx, + &query_config, + cost_tracker, + Some(event_tx), + tokio_util::sync::CancellationToken::new(), + None, + ) + .await; + let mut events = Vec::new(); + while let Ok(event) = event_rx.try_recv() { + events.push(event); + } + (outcome, events, messages) + } + + fn make_config(sys: Option<&str>, append: Option<&str>) -> QueryConfig { + QueryConfig { + model: "claude-sonnet-4-6".to_string(), + max_tokens: 4096, + max_turns: 10, + system_prompt: sys.map(String::from), + append_system_prompt: append.map(String::from), + output_style: claurst_core::system_prompt::OutputStyle::Default, + output_style_prompt: None, + working_directory: None, + thinking_budget: None, + temperature: None, + tool_result_budget: 50_000, + effort_level: None, + command_queue: None, + skill_index: None, + max_budget_usd: None, + fallback_model: None, + provider_registry: None, + agent_name: None, + agent_definition: None, + model_registry: None, + managed_agents: None, + preserve_selected_model: false, + } + } + + #[tokio::test] + async fn durable_message_event_survives_context_rewrite() { + let (tx, mut rx) = mpsc::unbounded_channel(); + let mut assistant = Message::assistant(""); + let mut messages = vec![assistant.clone()]; + assistant = Message::assistant("final placeholder"); + messages[0] = assistant.clone(); + emit_durable_message(Some(&tx), &assistant); + messages = vec![Message::user("compacted summary")]; + + let event = rx.recv().await.unwrap(); + match event { + QueryEvent::DurableMessage { message } => { + assert_eq!(message.get_all_text(), "final placeholder") + } + _ => panic!("expected durable message event"), + } + assert_eq!(messages.len(), 1); + } + + #[tokio::test] + async fn exhausted_anthropic_stall_finalizes_partial_assistant_before_error() { + let (tx, mut rx) = mpsc::unbounded_channel(); + let mut messages = vec![Message::user("continue")]; + let mut assistant = Message::assistant("partial response"); + assistant.uuid = Some("stalled-response".to_string()); + + let outcome = + finish_anthropic_stall(&mut messages, &mut assistant, None, None, Some(&tx)).await; + + assert!(matches!(outcome, QueryOutcome::Error(_))); + let event = rx.recv().await.unwrap(); + assert!(matches!( + event, + QueryEvent::DurableMessage { message } + if message.uuid.as_deref() == Some("stalled-response") + && message.get_all_text() == "partial response" + )); + assert!(messages.iter().any(|message| { + message.uuid.as_deref() == Some("stalled-response") + && message.get_all_text() == "partial response" + })); + } + + #[test] + fn finalized_assistant_is_restored_once_after_context_rewrite() { + let mut finalized = Message::assistant("final response"); + finalized.uuid = Some("current-response".to_string()); + let mut unrelated = Message::assistant("compacted summary"); + unrelated.uuid = Some("summary".to_string()); + let mut messages = vec![unrelated.clone()]; + + sync_finalized_assistant_message(&mut messages, &finalized); + sync_finalized_assistant_message(&mut messages, &finalized); + + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].uuid, unrelated.uuid); + assert_eq!(messages[0].get_all_text(), "compacted summary"); + assert_eq!(messages[1].uuid, finalized.uuid); + assert_eq!(messages[1].get_all_text(), "final response"); + } + + #[tokio::test] + async fn provider_stream_error_finalizes_partial_response_before_returning_error() { + let error = ProviderError::StreamError { + provider: ProviderId::new(SCRIPTED_PROVIDER_ID), + message: "scripted stream failed".to_string(), + partial_response: Some("partial".to_string()), + }; + let provider = ScriptedProvider::with_results(vec![vec![ + Ok(StreamEvent::MessageStart { + id: "stream-error".to_string(), + model: "scripted".to_string(), + usage: UsageInfo::default(), + }), + Ok(StreamEvent::TextDelta { + index: 0, + text: "partial".to_string(), + }), + Err(error), + ]]); + + let (outcome, events, messages) = run_scripted_provider(provider).await; + + assert!(matches!(&outcome, QueryOutcome::Error(_))); + assert!(!matches!(&outcome, QueryOutcome::EndTurn { .. })); + let durable = events + .iter() + .find_map(|event| match event { + QueryEvent::DurableMessage { message } => Some(message), + _ => None, + }) + .expect("stream error must publish its finalized partial assistant"); + assert_eq!(durable.uuid.as_deref(), Some("stream-error")); + assert_eq!(durable.get_all_text(), "partial"); + assert!(messages.iter().any(|message| { + message.uuid == durable.uuid && message.get_all_text() == "partial" + })); + } + + #[tokio::test] + async fn provider_cancellation_after_partial_sampling_finalizes_exactly_once() { + let provider = ScriptedProvider::partial_then_pending(vec![ + StreamEvent::MessageStart { + id: "provider-cancelled-partial".to_string(), + model: "scripted".to_string(), + usage: UsageInfo::default(), + }, + StreamEvent::TextDelta { + index: 0, + text: "partial before cancellation".to_string(), + }, + ]); + + let (outcome, events, messages) = + run_scripted_provider_until_partial_then_cancel(provider).await; + + assert!(matches!(outcome, QueryOutcome::Cancelled)); + let durable_index = events + .iter() + .position(|event| { + matches!(event, QueryEvent::DurableMessage { message } + if message.uuid.as_deref() == Some("provider-cancelled-partial") + && message.get_all_text() == "partial before cancellation") + }) + .expect("cancelled provider response must be durable"); + assert_eq!( + events + .iter() + .filter(|event| { + matches!(event, QueryEvent::DurableMessage { message } + if message.uuid.as_deref() == Some("provider-cancelled-partial")) + }) + .count(), + 1 + ); + assert!(messages.iter().any(|message| { + message.uuid.as_deref() == Some("provider-cancelled-partial") + && message.get_all_text() == "partial before cancellation" + })); + assert!(events[durable_index + 1..] + .iter() + .all(|event| !matches!(event, QueryEvent::TurnComplete { .. }))); + } + + #[tokio::test] + async fn anthropic_stream_error_finalizes_partial_response_before_returning_error() { + const STREAM_ERROR_SSE: &str = concat!( + "event: message_start\n", + "data: {\"message\":{\"id\":\"anthropic-stream-error\",\"model\":\"claude-test\",\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n", + "event: content_block_start\n", + "data: {\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "event: content_block_delta\n", + "data: {\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"partial\"}}\n\n", + "event: content_block_stop\n", + "data: {\"index\":0}\n\n", + "event: error\n", + "data: {\"error\":{\"type\":\"api_error\",\"message\":\"scripted stream failed\"}}\n\n", + ); + + let (outcome, events, messages) = + run_anthropic_sse(STREAM_ERROR_SSE, Config::default()).await; + + assert!(matches!(outcome, QueryOutcome::Error(_))); + let durable = events + .iter() + .find_map(|event| match event { + QueryEvent::DurableMessage { message } => Some(message), + _ => None, + }) + .expect("stream error must publish its finalized partial assistant"); + assert_eq!(durable.get_all_text(), "partial"); + assert!(messages.iter().any(|message| { + message.uuid == durable.uuid && message.get_all_text() == "partial" + })); + } + + #[tokio::test] + async fn anthropic_cancellation_after_partial_sampling_finalizes_exactly_once() { + const PARTIAL_SSE: &str = concat!( + "event: message_start\n", + "data: {\"message\":{\"id\":\"anthropic-cancelled-partial\",\"model\":\"claude-test\",\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n", + "event: content_block_start\n", + "data: {\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "event: content_block_delta\n", + "data: {\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"partial before cancellation\"}}\n\n", + "event: content_block_stop\n", + "data: {\"index\":0}\n\n", + ); + + let (outcome, events, messages) = + run_anthropic_until_partial_then_cancel(PARTIAL_SSE).await; + + assert!(matches!(outcome, QueryOutcome::Cancelled)); + let durable_index = events + .iter() + .position(|event| { + matches!(event, QueryEvent::DurableMessage { message } + if message.get_all_text() == "partial before cancellation") + }) + .unwrap_or_else(|| { + panic!("cancelled Anthropic response must be durable; events: {events:?}") + }); + let durable_uuid = match &events[durable_index] { + QueryEvent::DurableMessage { message } => message.uuid.clone(), + _ => unreachable!("durable index points to durable event"), + }; + assert_eq!( + events + .iter() + .filter(|event| { + matches!(event, QueryEvent::DurableMessage { message } + if message.uuid == durable_uuid) + }) + .count(), + 1 + ); + assert_eq!( + messages + .iter() + .filter(|message| message.uuid == durable_uuid) + .count(), + 1 + ); + assert!(events[durable_index + 1..] + .iter() + .all(|event| !matches!(event, QueryEvent::TurnComplete { .. }))); + } + + #[tokio::test] + async fn reactive_compaction_cancellation_restores_and_finalizes_sampled_assistant() { + struct FeatureGateGuard { + prior: Option, + } + + impl FeatureGateGuard { + fn enable_reactive_compact() -> Self { + let prior = std::env::var_os("COVEN_CODE_FEATURE_REACTIVE_COMPACT"); + std::env::set_var("COVEN_CODE_FEATURE_REACTIVE_COMPACT", "1"); + Self { prior } + } + } + + impl Drop for FeatureGateGuard { + fn drop(&mut self) { + match self.prior.take() { + Some(value) => std::env::set_var("COVEN_CODE_FEATURE_REACTIVE_COMPACT", value), + None => std::env::remove_var("COVEN_CODE_FEATURE_REACTIVE_COMPACT"), + } + } + } + + const COMPACTION_TRIGGER_SSE: &str = concat!( + "event: message_start\n", + "data: {\"message\":{\"id\":\"compaction-cancelled\",\"model\":\"claude-test\",\"usage\":{\"input_tokens\":180000,\"output_tokens\":0}}}\n\n", + "event: content_block_start\n", + "data: {\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "event: content_block_delta\n", + "data: {\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"sampled before compaction cancellation\"}}\n\n", + "event: content_block_stop\n", + "data: {\"index\":0}\n\n", + "event: message_delta\n", + "data: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":1}}\n\n", + "event: message_stop\n", + "data: {}\n\n", + ); + + let _feature_gate = FeatureGateGuard::enable_reactive_compact(); + let temp = tempfile::tempdir().unwrap(); + let (api_base, compaction_started) = + serve_anthropic_sse_then_wait_for_compaction(COMPACTION_TRIGGER_SSE).await; + let client = claurst_api::AnthropicClient::new(claurst_api::client::ClientConfig { + api_key: "test-key".to_string(), + api_base, + max_retries: 0, + ..Default::default() + }) + .unwrap(); + let cost_tracker = claurst_core::cost::CostTracker::new(); + let tool_ctx = ToolContext { + working_dir: temp.path().to_path_buf(), + permission_mode: PermissionMode::BypassPermissions, + permission_handler: Arc::new(AutoPermissionHandler { + mode: PermissionMode::BypassPermissions, + }), + cost_tracker: cost_tracker.clone(), + session_id: "compaction-cancellation".to_string(), + file_history: Default::default(), + current_turn: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + non_interactive: true, + mcp_manager: None, + config: Config::default(), + managed_agent_config: None, + completion_notifier: None, + pending_permissions: None, + permission_manager: None, + user_question_tx: None, + }; + let query_config = make_config(None, None); + let mut messages = (0..12) + .map(|index| Message::user(format!("history {index}"))) + .collect::>(); + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + let cancel_token = tokio_util::sync::CancellationToken::new(); + let mut events = Vec::new(); + let outcome = { + tokio::pin!(compaction_started); + let query = run_query_loop( + &client, + &mut messages, + &[], + &tool_ctx, + &query_config, + cost_tracker, + Some(event_tx), + cancel_token.clone(), + None, + ); + tokio::pin!(query); + tokio::select! { + outcome = &mut query => outcome, + result = &mut compaction_started => { + result.expect("compaction request must start"); + cancel_token.cancel(); + tokio::time::timeout(std::time::Duration::from_secs(1), &mut query) + .await + .expect("reactive compaction must observe cancellation") + } + } + }; + while let Ok(event) = event_rx.try_recv() { + events.push(event); + } + + assert!(matches!(outcome, QueryOutcome::Cancelled)); + let durable_index = events + .iter() + .position(|event| { + matches!(event, QueryEvent::DurableMessage { message } + if message.get_all_text() == "sampled before compaction cancellation") + }) + .expect("cancelled compaction response must be durable"); + let durable_uuid = match &events[durable_index] { + QueryEvent::DurableMessage { message } => message.uuid.clone(), + _ => unreachable!("durable index points to durable event"), + }; + assert_eq!( + messages + .iter() + .filter(|message| { + message.uuid == durable_uuid + && message.get_all_text() == "sampled before compaction cancellation" + }) + .count(), + 1 + ); + } + + #[tokio::test] + async fn provider_max_tokens_is_not_goal_continuation_eligible() { + let provider = ScriptedProvider::new(vec![scripted_round( + "max-tokens", + vec![ContentBlock::Text { + text: "partial response".to_string(), + }], + ProviderStopReason::MaxTokens, + )]); + + let (outcome, events, _) = run_scripted_provider(provider).await; + + assert!(matches!(&outcome, QueryOutcome::MaxTokens { .. })); + assert!(!matches!(&outcome, QueryOutcome::EndTurn { .. })); + let durable: Vec<_> = events + .iter() + .filter_map(|event| match event { + QueryEvent::DurableMessage { message } => Some(message), + _ => None, + }) + .collect(); + assert_eq!(durable.len(), 1); + assert_eq!(durable[0].uuid.as_deref(), Some("max-tokens")); + let durable_index = events + .iter() + .position(|event| matches!(event, QueryEvent::DurableMessage { .. })) + .unwrap(); + let complete_index = events + .iter() + .position(|event| matches!(event, QueryEvent::TurnComplete { .. })) + .unwrap(); + assert!(durable_index < complete_index); + } + + #[tokio::test] + async fn provider_max_tokens_with_tool_blocks_does_not_execute_or_continue_tools() { + let provider = ScriptedProvider::new(vec![ + scripted_round( + "max-tokens-tool", + vec![ContentBlock::ToolUse { + id: "write-fixture".to_string(), + name: "WriteFixture".to_string(), + input: serde_json::json!({}), + }], + ProviderStopReason::MaxTokens, + ), + scripted_round( + "must-not-request-second-round", + vec![ContentBlock::Text { + text: "unexpected continuation".to_string(), + }], + ProviderStopReason::EndTurn, + ), + ]); + let tools: Vec> = vec![Box::new(WriteFixtureTool)]; + + let (outcome, events, _) = run_scripted_provider_with( + provider, + make_config(None, None), + Config { + provider: Some(SCRIPTED_PROVIDER_ID.to_string()), + permission_mode: PermissionMode::BypassPermissions, + ..Config::default() + }, + &tools, + ) + .await; + + assert!(matches!(&outcome, QueryOutcome::MaxTokens { .. })); + assert!(!matches!(&outcome, QueryOutcome::EndTurn { .. })); + assert!(!events.iter().any(|event| { + matches!( + event, + QueryEvent::ToolStart { tool_id, .. } | QueryEvent::ToolEnd { tool_id, .. } + if tool_id == "write-fixture" + ) + })); + assert_eq!( + events + .iter() + .filter(|event| matches!(event, QueryEvent::DurableMessage { message } if message.uuid.as_deref() == Some("max-tokens-tool"))) + .count(), + 1 + ); + } + + #[tokio::test] + async fn provider_max_tokens_runs_hard_post_model_veto_before_classification() { + let provider = ScriptedProvider::new(vec![scripted_round( + "max-tokens-veto", + vec![ContentBlock::ToolUse { + id: "write-fixture".to_string(), + name: "WriteFixture".to_string(), + input: serde_json::json!({}), + }], + ProviderStopReason::MaxTokens, + )]); + let tools: Vec> = vec![Box::new(WriteFixtureTool)]; + let mut tool_config = Config { + provider: Some(SCRIPTED_PROVIDER_ID.to_string()), + permission_mode: PermissionMode::BypassPermissions, + ..Config::default() + }; + tool_config.hooks.insert( + HookEvent::PostModelTurn, + vec![HookEntry { + command: "printf 'max-token veto ran' >&2; exit 2".to_string(), + tool_filter: None, + blocking: true, + }], + ); + + let (outcome, events, _) = + run_scripted_provider_with(provider, make_config(None, None), tool_config, &tools) + .await; + + assert!( + matches!(&outcome, QueryOutcome::Error(error) if error.to_string().contains("max-token veto ran")) + ); + assert!(!matches!(&outcome, QueryOutcome::EndTurn { .. })); + assert!(!matches!(&outcome, QueryOutcome::MaxTokens { .. })); + assert!(!events.iter().any(|event| { + matches!( + event, + QueryEvent::ToolStart { tool_id, .. } | QueryEvent::ToolEnd { tool_id, .. } + if tool_id == "write-fixture" + ) + })); + let durable_index = events + .iter() + .position(|event| { + matches!(event, QueryEvent::DurableMessage { message } if message.uuid.as_deref() == Some("max-tokens-veto")) + }) + .expect("vetoed max_tokens response must be durable"); + let complete_index = events + .iter() + .position(|event| matches!(event, QueryEvent::TurnComplete { .. })) + .expect("vetoed max_tokens response must publish completion after durability"); + assert!(durable_index < complete_index); + } + + #[tokio::test] + async fn provider_content_filtered_and_unknown_stops_are_non_successful_and_durable() { + for (id, stop_reason) in [ + ( + "provider-content-filtered", + ProviderStopReason::ContentFiltered, + ), + ( + "provider-unknown-stop", + ProviderStopReason::Other("provider_extension".to_string()), + ), + ] { + let provider = ScriptedProvider::new(vec![scripted_round( + id, + vec![ContentBlock::Text { + text: "sampled terminal text".to_string(), + }], + stop_reason, + )]); + + let (outcome, events, messages) = run_scripted_provider(provider).await; + + assert!(matches!(&outcome, QueryOutcome::Error(_))); + assert!(!matches!(&outcome, QueryOutcome::EndTurn { .. })); + let durable_index = events + .iter() + .position(|event| { + matches!(event, QueryEvent::DurableMessage { message } if message.uuid.as_deref() == Some(id)) + }) + .expect("terminal provider response must be durable"); + let complete_index = events + .iter() + .position(|event| matches!(event, QueryEvent::TurnComplete { .. })) + .expect("terminal provider response must publish completion after durability"); + assert!(durable_index < complete_index); + assert_eq!( + messages + .iter() + .filter(|message| message.uuid.as_deref() == Some(id)) + .count(), + 1 + ); + } + } + + #[tokio::test] + async fn provider_budget_exhaustion_with_tool_blocks_does_not_execute_or_continue_tools() { + let provider = ScriptedProvider::new(vec![ + scripted_round( + "over-budget-tool", + vec![ContentBlock::ToolUse { + id: "write-fixture".to_string(), + name: "WriteFixture".to_string(), + input: serde_json::json!({}), + }], + ProviderStopReason::ToolUse, + ), + scripted_round( + "must-not-request-second-round", + vec![ContentBlock::Text { + text: "unexpected continuation".to_string(), + }], + ProviderStopReason::EndTurn, + ), + ]); + let mut query_config = make_config(None, None); + query_config.max_budget_usd = Some(0.0); + let tools: Vec> = vec![Box::new(WriteFixtureTool)]; + + let (outcome, events, _) = run_scripted_provider_with( + provider, + query_config, + Config { + provider: Some(SCRIPTED_PROVIDER_ID.to_string()), + permission_mode: PermissionMode::BypassPermissions, + ..Config::default() + }, + &tools, + ) + .await; + + assert!(matches!(&outcome, QueryOutcome::BudgetExceeded { .. })); + assert!(!matches!(&outcome, QueryOutcome::EndTurn { .. })); + assert!(!events.iter().any(|event| { + matches!( + event, + QueryEvent::ToolStart { tool_id, .. } | QueryEvent::ToolEnd { tool_id, .. } + if tool_id == "write-fixture" + ) + })); + assert_eq!( + events + .iter() + .filter(|event| matches!(event, QueryEvent::DurableMessage { message } if message.uuid.as_deref() == Some("over-budget-tool"))) + .count(), + 1 + ); + } + + #[tokio::test] + async fn provider_tool_use_without_tool_blocks_is_a_non_successful_terminal() { + let provider = ScriptedProvider::new(vec![scripted_round( + "malformed-tool-use", + Vec::new(), + ProviderStopReason::ToolUse, + )]); + + let (outcome, events, _) = run_scripted_provider(provider).await; + + assert!(matches!(&outcome, QueryOutcome::Error(_))); + assert!(!matches!(&outcome, QueryOutcome::EndTurn { .. })); + let durable_index = events + .iter() + .position(|event| { + matches!(event, QueryEvent::DurableMessage { message } if message.uuid.as_deref() == Some("malformed-tool-use")) + }) + .expect("malformed provider terminal must be durable"); + let complete_index = events + .iter() + .position(|event| matches!(event, QueryEvent::TurnComplete { .. })) + .expect("malformed provider terminal must publish completion after durability"); + assert!(durable_index < complete_index); + } + + #[tokio::test] + async fn max_turn_exhaustion_after_a_tool_round_is_not_goal_continuation_eligible() { + let provider = ScriptedProvider::new(vec![scripted_round( + "max-turn-tool-round", + vec![ContentBlock::ToolUse { + id: "write-fixture".to_string(), + name: "WriteFixture".to_string(), + input: serde_json::json!({}), + }], + ProviderStopReason::ToolUse, + )]); + let tools: Vec> = vec![Box::new(WriteFixtureTool)]; + let mut config = make_config(None, None); + config.max_turns = 1; + + let (outcome, events, messages) = run_scripted_provider_with( + provider, + config, + Config { + provider: Some(SCRIPTED_PROVIDER_ID.to_string()), + permission_mode: PermissionMode::BypassPermissions, + ..Config::default() + }, + &tools, + ) + .await; + + assert!(matches!(&outcome, QueryOutcome::Error(_))); + assert!(!matches!(&outcome, QueryOutcome::EndTurn { .. })); + assert_eq!( + events + .iter() + .filter(|event| { + matches!(event, QueryEvent::DurableMessage { message } if message.uuid.as_deref() == Some("max-turn-tool-round")) + }) + .count(), + 1 + ); + assert_eq!( + events + .iter() + .filter(|event| match event { + QueryEvent::DurableMessage { message } => is_tool_result_carrier(message), + _ => false, + }) + .count(), + 1 + ); + assert_eq!( + messages + .iter() + .filter(|message| message.uuid.as_deref() == Some("max-turn-tool-round")) + .count(), + 1 + ); + } + + #[tokio::test] + async fn provider_hard_post_model_veto_prevents_tool_execution_and_end_turn() { + let provider = ScriptedProvider::new(vec![ + scripted_round( + "vetoed-tool-round", + vec![ContentBlock::ToolUse { + id: "write-fixture".to_string(), + name: "WriteFixture".to_string(), + input: serde_json::json!({}), + }], + ProviderStopReason::ToolUse, + ), + scripted_round( + "must-not-request-second-round", + vec![ContentBlock::Text { + text: "unexpected continuation".to_string(), + }], + ProviderStopReason::EndTurn, + ), + ]); + let tools: Vec> = vec![Box::new(WriteFixtureTool)]; + let mut tool_config = Config { + provider: Some(SCRIPTED_PROVIDER_ID.to_string()), + permission_mode: PermissionMode::BypassPermissions, + ..Config::default() + }; + tool_config.hooks.insert( + HookEvent::PostModelTurn, + vec![HookEntry { + command: "printf 'blocked by policy' >&2; exit 2".to_string(), + tool_filter: None, + blocking: true, + }], + ); + + let (outcome, events, _) = + run_scripted_provider_with(provider, make_config(None, None), tool_config, &tools) + .await; + + assert!( + matches!(&outcome, QueryOutcome::Error(error) if error.to_string().contains("blocked by policy")) + ); + assert!(!matches!(&outcome, QueryOutcome::EndTurn { .. })); + assert!(!events.iter().any(|event| { + matches!( + event, + QueryEvent::ToolStart { tool_id, .. } | QueryEvent::ToolEnd { tool_id, .. } + if tool_id == "write-fixture" + ) + })); + let durable_index = events + .iter() + .position(|event| { + matches!(event, QueryEvent::DurableMessage { message } if message.uuid.as_deref() == Some("vetoed-tool-round")) + }) + .expect("vetoed provider response must be durable"); + let complete_index = events + .iter() + .position(|event| matches!(event, QueryEvent::TurnComplete { .. })) + .expect("vetoed provider response must publish completion after durability"); + assert!(durable_index < complete_index); + } + + #[tokio::test] + async fn provider_stall_recovery_exhaustion_is_not_a_successful_end_turn() { + let provider = ScriptedProvider::new(vec![ + scripted_round("stall-1", Vec::new(), ProviderStopReason::EndTurn), + scripted_round("stall-2", Vec::new(), ProviderStopReason::EndTurn), + scripted_round("stall-exhausted", Vec::new(), ProviderStopReason::EndTurn), + ]); + + let (outcome, events, _) = run_scripted_provider(provider).await; + + assert!(matches!(&outcome, QueryOutcome::Error(_))); + assert!(!matches!(&outcome, QueryOutcome::EndTurn { .. })); + let final_durable = events.iter().rev().find_map(|event| match event { + QueryEvent::DurableMessage { message } => Some(message), + _ => None, + }); + assert_eq!( + final_durable.map(Message::get_all_text).as_deref(), + Some("(no response — model ended the turn with stop_reason \"end_turn\")") + ); + } + + #[tokio::test] + async fn hard_post_model_hook_veto_is_not_goal_continuation_eligible() { + const END_TURN_SSE: &str = concat!( + "event: message_start\n", + "data: {\"message\":{\"id\":\"hard-veto\",\"model\":\"claude-test\",\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n", + "event: content_block_start\n", + "data: {\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "event: content_block_delta\n", + "data: {\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"sampled answer\"}}\n\n", + "event: content_block_stop\n", + "data: {\"index\":0}\n\n", + "event: message_delta\n", + "data: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":2}}\n\n", + "event: message_stop\n", + "data: {}\n\n", + ); + let mut tool_config = Config::default(); + tool_config.hooks.insert( + HookEvent::PostModelTurn, + vec![HookEntry { + command: "printf 'blocked by policy' >&2; exit 2".to_string(), + tool_filter: None, + blocking: true, + }], + ); + + let (outcome, events, _) = run_anthropic_sse(END_TURN_SSE, tool_config).await; + + match outcome { + QueryOutcome::Error(error) => { + assert!(error.to_string().contains("blocked by policy")); + } + other => panic!("hard hook veto must be non-success, got {other:?}"), + } + let durable_index = events + .iter() + .position(|event| matches!(event, QueryEvent::DurableMessage { .. })) + .unwrap(); + let terminal_index = events + .iter() + .position(|event| matches!(event, QueryEvent::TurnComplete { .. })) + .unwrap(); + assert!(durable_index < terminal_index); + } + + #[tokio::test] + async fn anthropic_content_filtered_and_unknown_stops_are_non_successful_and_durable() { + const CONTENT_FILTERED_SSE: &str = concat!( + "event: message_start\n", + "data: {\"message\":{\"id\":\"anthropic-content-filtered\",\"model\":\"claude-test\",\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n", + "event: content_block_start\n", + "data: {\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "event: content_block_delta\n", + "data: {\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"filtered terminal text\"}}\n\n", + "event: content_block_stop\n", + "data: {\"index\":0}\n\n", + "event: message_delta\n", + "data: {\"delta\":{\"stop_reason\":\"content_filtered\"},\"usage\":{\"output_tokens\":1}}\n\n", + "event: message_stop\n", + "data: {}\n\n", + ); + const UNKNOWN_STOP_SSE: &str = concat!( + "event: message_start\n", + "data: {\"message\":{\"id\":\"anthropic-unknown-stop\",\"model\":\"claude-test\",\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n", + "event: content_block_start\n", + "data: {\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "event: content_block_delta\n", + "data: {\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"unknown terminal text\"}}\n\n", + "event: content_block_stop\n", + "data: {\"index\":0}\n\n", + "event: message_delta\n", + "data: {\"delta\":{\"stop_reason\":\"provider_extension\"},\"usage\":{\"output_tokens\":1}}\n\n", + "event: message_stop\n", + "data: {}\n\n", + ); + + for (text, body) in [ + ("filtered terminal text", CONTENT_FILTERED_SSE), + ("unknown terminal text", UNKNOWN_STOP_SSE), + ] { + let (outcome, events, messages) = run_anthropic_sse(body, Config::default()).await; + + assert!(matches!(&outcome, QueryOutcome::Error(_))); + assert!(!matches!(&outcome, QueryOutcome::EndTurn { .. })); + let durable_index = events + .iter() + .position(|event| { + matches!(event, QueryEvent::DurableMessage { message } if message.get_all_text() == text) + }) + .expect("terminal Anthropic response must be durable"); + let complete_index = events + .iter() + .position(|event| matches!(event, QueryEvent::TurnComplete { .. })) + .expect("terminal Anthropic response must publish completion after durability"); + assert!(durable_index < complete_index); + let durable_uuid = match &events[durable_index] { + QueryEvent::DurableMessage { message } => message.uuid.clone(), + _ => unreachable!("durable index points to durable event"), + }; + assert_eq!( + messages + .iter() + .filter(|message| message.uuid == durable_uuid) + .count(), + 1 + ); + } + } + + #[tokio::test] + async fn malformed_tool_use_terminal_keeps_cumulative_snapshot_durable() { + const TOOL_USE_SSE: &str = concat!( + "event: message_start\n", + "data: {\"message\":{\"id\":\"write-round\",\"model\":\"claude-test\",\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n", + "event: content_block_start\n", + "data: {\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"write-fixture\",\"name\":\"WriteFixture\",\"input\":{}}}\n\n", + "event: content_block_stop\n", + "data: {\"index\":0}\n\n", + "event: message_delta\n", + "data: {\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":1}}\n\n", + "event: message_stop\n", + "data: {}\n\n", + ); + const MALFORMED_TOOL_USE_SSE: &str = concat!( + "event: message_start\n", + "data: {\"message\":{\"id\":\"malformed-tool-use\",\"model\":\"claude-test\",\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n", + "event: message_delta\n", + "data: {\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":1}}\n\n", + "event: message_stop\n", + "data: {}\n\n", + ); + let temp = tempfile::tempdir().unwrap(); + init_test_repository(temp.path()); + let mut tool_config = Config { + auto_commits: Some(true), + permission_mode: PermissionMode::BypassPermissions, + ..Config::default() + }; + tool_config.project_dir = Some(temp.path().to_path_buf()); + let tools: Vec> = vec![Box::new(WriteFixtureTool)]; + + let (outcome, events, messages) = run_anthropic_sse_sequence_in_worktree( + vec![TOOL_USE_SSE, MALFORMED_TOOL_USE_SSE], + temp.path(), + tool_config, + make_config(None, None), + &tools, + ) + .await; + + assert!(matches!(outcome, QueryOutcome::Error(_))); + let final_durable = events.iter().rev().find_map(|event| match event { + QueryEvent::DurableMessage { message } if message.role == Role::Assistant => { + Some(message) + } + _ => None, + }); + let final_durable = final_durable.expect("malformed terminal must be durable"); + assert!(final_durable.snapshot_patch.is_some()); + let final_uuid = final_durable.uuid.clone().expect("stable assistant UUID"); + assert!(messages.iter().any(|message| { + message.uuid.as_ref() == Some(&final_uuid) && message.snapshot_patch.is_some() + })); + let durable_index = events + .iter() + .position(|event| matches!(event, QueryEvent::DurableMessage { message } if message.uuid.as_ref() == Some(&final_uuid))) + .unwrap(); + let terminal_index = events + .iter() + .position(|event| matches!(event, QueryEvent::TurnComplete { .. })) + .unwrap(); + assert!(durable_index < terminal_index); + } + + #[tokio::test] + async fn budget_exhaustion_finalizes_sampled_response_with_cumulative_snapshot() { + const ZERO_COST_TOOL_USE_SSE: &str = concat!( + "event: message_start\n", + "data: {\"message\":{\"id\":\"budget-write-round\",\"model\":\"claude-test\",\"usage\":{\"input_tokens\":0,\"output_tokens\":0}}}\n\n", + "event: content_block_start\n", + "data: {\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"write-fixture\",\"name\":\"WriteFixture\",\"input\":{}}}\n\n", + "event: content_block_stop\n", + "data: {\"index\":0}\n\n", + "event: message_delta\n", + "data: {\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":0}}\n\n", + "event: message_stop\n", + "data: {}\n\n", + ); + const EXPENSIVE_END_TURN_SSE: &str = concat!( + "event: message_start\n", + "data: {\"message\":{\"id\":\"budget-terminal\",\"model\":\"claude-test\",\"usage\":{\"input_tokens\":1000000,\"output_tokens\":0}}}\n\n", + "event: content_block_start\n", + "data: {\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "event: content_block_delta\n", + "data: {\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"sampled before budget stop\"}}\n\n", + "event: content_block_stop\n", + "data: {\"index\":0}\n\n", + "event: message_delta\n", + "data: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":1000000}}\n\n", + "event: message_stop\n", + "data: {}\n\n", + ); + let temp = tempfile::tempdir().unwrap(); + init_test_repository(temp.path()); + let mut tool_config = Config { + auto_commits: Some(true), + permission_mode: PermissionMode::BypassPermissions, + ..Config::default() + }; + tool_config.project_dir = Some(temp.path().to_path_buf()); + let tools: Vec> = vec![Box::new(WriteFixtureTool)]; + let mut query_config = make_config(None, None); + query_config.max_budget_usd = Some(0.000_001); + + let (outcome, events, messages) = run_anthropic_sse_sequence_in_worktree( + vec![ZERO_COST_TOOL_USE_SSE, EXPENSIVE_END_TURN_SSE], + temp.path(), + tool_config, + query_config, + &tools, + ) + .await; + + assert!(matches!(outcome, QueryOutcome::BudgetExceeded { .. })); + let final_durable = events.iter().rev().find_map(|event| match event { + QueryEvent::DurableMessage { message } + if message.get_all_text() == "sampled before budget stop" => + { + Some(message) + } + _ => None, + }); + let final_durable = final_durable.expect("budget terminal assistant must be durable"); + assert!(final_durable.snapshot_patch.is_some()); + let final_uuid = final_durable.uuid.as_ref().expect("stable assistant UUID"); + assert_eq!( + messages + .iter() + .filter(|message| message.uuid.as_ref() == Some(final_uuid)) + .count(), + 1 + ); + assert!(messages.iter().any(|message| { + message.uuid.as_ref() == Some(final_uuid) + && message.get_all_text() == "sampled before budget stop" + && message.snapshot_patch.is_some() + })); + assert!(matches!( + events.last(), + Some(QueryEvent::DurableMessage { message }) + if message.uuid.as_ref() == Some(final_uuid) + )); + assert!(!events + .iter() + .any(|event| matches!(event, QueryEvent::TurnComplete { .. }))); + } + + #[tokio::test] + async fn provider_dispatch_emits_finalized_durable_messages_exactly_once() { + let temp = tempfile::tempdir().unwrap(); + init_test_repository(temp.path()); + + let rounds = vec![ + scripted_round( + "tool-round", + vec![ContentBlock::ToolUse { + id: "write-fixture".to_string(), + name: "WriteFixture".to_string(), + input: serde_json::json!({}), + }], + ProviderStopReason::ToolUse, + ), + scripted_round("stall-round-1", Vec::new(), ProviderStopReason::EndTurn), + scripted_round("stall-round-2", Vec::new(), ProviderStopReason::EndTurn), + scripted_round("final-round", Vec::new(), ProviderStopReason::EndTurn), + ]; + let mut registry = claurst_api::ProviderRegistry::new(); + registry.register(Arc::new(ScriptedProvider::new(rounds))); + + let mut config = make_config(None, None); + config.model = "scripted".to_string(); + config.provider_registry = Some(Arc::new(registry)); + + let cost_tracker = claurst_core::cost::CostTracker::new(); + let mut tool_config = Config { + provider: Some(SCRIPTED_PROVIDER_ID.to_string()), + auto_commits: Some(true), + permission_mode: PermissionMode::BypassPermissions, + ..Config::default() + }; + tool_config.project_dir = Some(temp.path().to_path_buf()); + let tool_ctx = ToolContext { + working_dir: temp.path().to_path_buf(), + permission_mode: PermissionMode::BypassPermissions, + permission_handler: Arc::new(AutoPermissionHandler { + mode: PermissionMode::BypassPermissions, + }), + cost_tracker: cost_tracker.clone(), + session_id: "durable-journal".to_string(), + file_history: Default::default(), + current_turn: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + non_interactive: true, + mcp_manager: None, + config: tool_config, + managed_agent_config: None, + completion_notifier: None, + pending_permissions: None, + permission_manager: None, + user_question_tx: None, + }; + let client = claurst_api::AnthropicClient::new(claurst_api::client::ClientConfig { + api_key: "unused-test-key".to_string(), + ..Default::default() + }) + .unwrap(); + let tools: Vec> = vec![Box::new(WriteFixtureTool)]; + let mut messages = vec![Message::user("write the fixture")]; + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + + let outcome = run_query_loop( + &client, + &mut messages, + &tools, + &tool_ctx, + &config, + cost_tracker, + Some(event_tx), + tokio_util::sync::CancellationToken::new(), + None, + ) + .await; + assert!(matches!(outcome, QueryOutcome::Error(_))); + + let mut events = Vec::new(); + while let Ok(event) = event_rx.try_recv() { + events.push(event); + } + let durable: Vec<_> = events + .iter() + .filter_map(|event| match event { + QueryEvent::DurableMessage { message } => Some(message.clone()), + _ => None, + }) + .collect(); + + assert_eq!(durable.len(), 5); + let assistant_ids: Vec<_> = durable + .iter() + .filter(|message| matches!(message.role, Role::Assistant)) + .map(|message| message.uuid.as_deref().unwrap()) + .collect(); + assert_eq!( + assistant_ids, + [ + "tool-round", + "stall-round-1", + "stall-round-2", + "final-round" + ] + ); + assert!(!durable[0].get_tool_use_blocks().is_empty()); + assert!(is_tool_result_carrier(&durable[1])); + assert_eq!(durable[2].get_all_text(), ""); + assert_eq!(durable[3].get_all_text(), ""); + assert_eq!( + durable[4].get_all_text(), + "(no response — model ended the turn with stop_reason \"end_turn\")" + ); + assert!(durable[4].snapshot_patch.is_some()); + assert_eq!(unique_message_fingerprints(&durable).len(), durable.len()); + let final_durable_index = events + .iter() + .rposition(|event| matches!(event, QueryEvent::DurableMessage { .. })) + .unwrap(); + let terminal_index = events + .iter() + .rposition(|event| matches!(event, QueryEvent::TurnComplete { .. })) + .unwrap(); + assert!(final_durable_index < terminal_index); + } + #[test] fn parses_github_origin_repo_slug_without_credentials() { assert_eq!( diff --git a/src-rust/crates/tools/src/goal_complete.rs b/src-rust/crates/tools/src/goal_complete.rs index fa1e904..c0971a6 100644 --- a/src-rust/crates/tools/src/goal_complete.rs +++ b/src-rust/crates/tools/src/goal_complete.rs @@ -8,9 +8,14 @@ use crate::{PermissionLevel, Tool, ToolContext, ToolResult}; use async_trait::async_trait; use serde::Deserialize; use serde_json::{json, Value}; +use std::path::PathBuf; pub struct GoalCompleteTool; +pub struct PathScopedGoalCompleteTool { + goal_store_path: PathBuf, +} + #[derive(Debug, Deserialize)] struct GoalCompleteInput { /// A concise summary of what was accomplished (the audit). @@ -19,6 +24,39 @@ struct GoalCompleteInput { evidence: String, } +impl GoalCompleteTool { + pub fn at_path(goal_store_path: PathBuf) -> PathScopedGoalCompleteTool { + PathScopedGoalCompleteTool { goal_store_path } + } +} + +fn complete_goal(input: Value, ctx: &ToolContext, store: claurst_core::GoalStore) -> ToolResult { + let params: GoalCompleteInput = match serde_json::from_value(input) { + Ok(params) => params, + Err(error) => return ToolResult::error(format!("Invalid input: {error}")), + }; + + if params.audit_summary.trim().is_empty() { + return ToolResult::error( + "audit_summary cannot be empty. Provide a concise description of what was completed." + .to_string(), + ); + } + if params.evidence.trim().is_empty() { + return ToolResult::error( + "evidence cannot be empty. Provide test output, diffs, or command results.".to_string(), + ); + } + + match store.complete_active_goal(&ctx.session_id) { + Ok(()) => ToolResult::success(format!( + "Goal marked complete.\n\nAudit summary: {}\n\nEvidence: {}", + params.audit_summary, params.evidence, + )), + Err(error) => ToolResult::error(format!("Failed to mark goal complete: {error}")), + } +} + #[async_trait] impl Tool for GoalCompleteTool { fn name(&self) -> &str { @@ -55,39 +93,139 @@ impl Tool for GoalCompleteTool { } async fn execute(&self, input: Value, ctx: &ToolContext) -> ToolResult { - let params: GoalCompleteInput = match serde_json::from_value(input) { - Ok(p) => p, - Err(e) => return ToolResult::error(format!("Invalid input: {}", e)), - }; - - if params.audit_summary.trim().is_empty() { - return ToolResult::error( - "audit_summary cannot be empty. Provide a concise description of what was completed." - .to_string(), - ); + match claurst_core::GoalStore::open_default() { + None => ToolResult::error("Could not open goal store.".to_string()), + Some(store) => complete_goal(input, ctx, store), } - if params.evidence.trim().is_empty() { - return ToolResult::error( - "evidence cannot be empty. Provide test output, diffs, or command results." - .to_string(), - ); + } +} + +impl PathScopedGoalCompleteTool { + fn open_store(&self) -> Result { + claurst_core::GoalStore::open(&self.goal_store_path).map_err(|error| error.to_string()) + } +} + +#[async_trait] +impl Tool for PathScopedGoalCompleteTool { + fn name(&self) -> &str { + "GoalComplete" + } + + fn description(&self) -> &str { + GoalCompleteTool.description() + } + + fn permission_level(&self) -> PermissionLevel { + PermissionLevel::None + } + + fn input_schema(&self) -> Value { + GoalCompleteTool.input_schema() + } + + async fn execute(&self, input: Value, ctx: &ToolContext) -> ToolResult { + match self.open_store() { + Ok(store) => complete_goal(input, ctx, store), + Err(error) => ToolResult::error(format!("Could not open goal store: {error}")), } + } +} - let session_id = &ctx.session_id; +#[cfg(test)] +mod tests { + use std::path::PathBuf; + use std::sync::atomic::AtomicUsize; + use std::sync::Arc; - match claurst_core::GoalStore::open_default() { - None => ToolResult::error("Could not open goal store.".to_string()), - Some(store) => match store.set_status(session_id, claurst_core::GoalStatus::Complete) { - Ok(()) => ToolResult::success(format!( - "Goal marked complete.\n\nAudit summary: {}\n\nEvidence: {}", - params.audit_summary, params.evidence, - )), - Err(e) => ToolResult::error(format!( - "Failed to mark goal complete: {}. \ - There may be no active goal for this session.", - e - )), - }, + use super::*; + use claurst_core::config::{Config, PermissionMode}; + use claurst_core::file_history::FileHistory; + use claurst_core::permissions::AutoPermissionHandler; + + fn test_tool_context(session_id: &str) -> ToolContext { + ToolContext { + working_dir: PathBuf::from("/workspace"), + permission_mode: PermissionMode::Default, + permission_handler: Arc::new(AutoPermissionHandler { + mode: PermissionMode::Default, + }), + cost_tracker: claurst_core::cost::CostTracker::new(), + session_id: session_id.to_string(), + file_history: Arc::new(parking_lot::Mutex::new(FileHistory::new())), + current_turn: Arc::new(AtomicUsize::new(0)), + non_interactive: true, + mcp_manager: None, + config: Config::default(), + managed_agent_config: None, + completion_notifier: None, + pending_permissions: None, + permission_manager: None, + user_question_tx: None, } } + + #[test] + fn default_tool_remains_unit_constructible() { + let tool = GoalCompleteTool; + assert_eq!(tool.name(), "GoalComplete"); + } + + #[tokio::test] + async fn explicit_path_completes_only_the_matching_active_goal() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("goals.sqlite"); + let store = claurst_core::GoalStore::open(&path).unwrap(); + store.set_goal("target", "finish", None).unwrap(); + store.set_goal("other", "stay active", None).unwrap(); + + let result = GoalCompleteTool::at_path(path.clone()) + .execute( + serde_json::json!({ + "audit_summary": "finished", + "evidence": "tests passed" + }), + &test_tool_context("target"), + ) + .await; + + assert!(!result.is_error); + let reopened = claurst_core::GoalStore::open(&path).unwrap(); + assert_eq!( + reopened.try_get_goal("target").unwrap().unwrap().status, + claurst_core::GoalStatus::Complete + ); + assert_eq!( + reopened.try_get_goal("other").unwrap().unwrap().status, + claurst_core::GoalStatus::Active + ); + } + + #[tokio::test] + async fn explicit_path_rejects_paused_and_missing_goals() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("goals.sqlite"); + let store = claurst_core::GoalStore::open(&path).unwrap(); + store.set_goal("paused", "finish", None).unwrap(); + store + .set_status("paused", claurst_core::GoalStatus::Paused) + .unwrap(); + let input = serde_json::json!({ + "audit_summary": "finished", + "evidence": "tests passed" + }); + + assert!( + GoalCompleteTool::at_path(path.clone()) + .execute(input.clone(), &test_tool_context("paused")) + .await + .is_error + ); + assert!( + GoalCompleteTool::at_path(path) + .execute(input, &test_tool_context("missing")) + .await + .is_error + ); + } } diff --git a/src-rust/crates/tools/src/lib.rs b/src-rust/crates/tools/src/lib.rs index 88bc30e..59d9da3 100644 --- a/src-rust/crates/tools/src/lib.rs +++ b/src-rust/crates/tools/src/lib.rs @@ -70,7 +70,7 @@ pub use file_read::FileReadTool; pub use file_write::FileWriteTool; pub use formatter::try_format_file; pub use glob_tool::GlobTool; -pub use goal_complete::GoalCompleteTool; +pub use goal_complete::{GoalCompleteTool, PathScopedGoalCompleteTool}; pub use grep_tool::GrepTool; pub use lsp_tool::LspTool; pub use mcp_auth_tool::McpAuthTool; diff --git a/src-rust/crates/tui/src/app.rs b/src-rust/crates/tui/src/app.rs index bf74a6e..33aee9b 100644 --- a/src-rust/crates/tui/src/app.rs +++ b/src-rust/crates/tui/src/app.rs @@ -6850,6 +6850,7 @@ impl App { _ => {} } } + QueryEvent::DurableMessage { .. } => {} } // Update token count from tracker.