From 66b45352f3f7ac23cf419be26a3bbd8ddd5fb3cb Mon Sep 17 00:00:00 2001 From: Wanjun Gu Date: Sat, 17 Oct 2026 23:11:08 -0700 Subject: [PATCH] fix(local): keep a continued reply's joining space, retire an unanswered retry note on Stop, and drop 'try again' from the GPU out-of-memory notice Three defects the final regression run found in the stream retry of #403/#405/#407, each reproduced independently first. - A completion that continues a cut-off reply (a retry after a partial reply, or a continuation past the output limit) no longer loses its leading whitespace in the local providers' stream recovery, so 'brown fox' + ' jumps' no longer reads 'foxjumps'. The agent marks the provider call (CONTINUES_A_REPLY) from the shape of the history: the model's own text followed by an agent-only user row. That shape also matches a few fresh replies after another agent-only note (self-critique or done-gate feedback, a supervision prompt); those keep a leading space too, which is harmless. Whitespace alone is still no reply (R14). - A Stop, or a failed stream, while a provider retry waits hides that retry's note from the model in place (a compare-and-set of its metadata; nothing is deleted), so the person's next message is not sent behind a 'continue the reply' order the retry never acted on. The daemon runner and Crew do it in settle_interrupted_turn and record_turn_stopped; the CLI and ACP call Agent::retire_unanswered_retry_note once the turn's token is cancelled, because there a Stop usually ends the stream on its own rather than through the cancel arm. A note the retry answered stays, as does a note followed by an agent-visible row such as a steer. - Llama Server's GPU out-of-memory notice, shown right before 'Trying again (1/1)', no longer asks the person to try again or says the server was restarted, matching the compute-failure notice. --- crates/biorouter-acp/src/server.rs | 11 + crates/biorouter-cli/src/session/mod.rs | 12 + crates/biorouter/src/agents/agent.rs | 181 ++++++++++++- .../src/providers/formats/text_tool_calls.rs | 125 ++++++++- crates/biorouter/src/providers/llamacpp.rs | 3 +- .../src/providers/llamacpp_sidecar.rs | 39 ++- crates/biorouter/src/providers/ollama.rs | 9 +- .../biorouter/src/session/session_manager.rs | 21 ++ .../tests/output_recovery_agent_loop.rs | 253 +++++++++++++++++- 9 files changed, 625 insertions(+), 29 deletions(-) diff --git a/crates/biorouter-acp/src/server.rs b/crates/biorouter-acp/src/server.rs index 081dfa60b..40069ef61 100644 --- a/crates/biorouter-acp/src/server.rs +++ b/crates/biorouter-acp/src/server.rs @@ -1024,12 +1024,23 @@ impl BioRouterAcpAgent { } Ok(_) => {} Err(e) => { + drop(stream); + self.agent.retire_unanswered_retry_note(&session_id).await; return Err(sacp::Error::internal_error() .data(format!("Error in agent response stream: {}", e))); } } } + // A cancel while a provider retry waits leaves the retry's note the last + // row the model reads, and the next prompt would be sent behind it. The + // token, not `was_cancelled`: the reply loop sees the cancel too, and its + // stream then usually ends on its own before another event arrives. + drop(stream); + if cancel_token.is_cancelled() { + self.agent.retire_unanswered_retry_note(&session_id).await; + } + let mut sessions = self.sessions.lock().await; if let Some(session) = sessions.get_mut(&session_id) { session.cancel_token = None; diff --git a/crates/biorouter-cli/src/session/mod.rs b/crates/biorouter-cli/src/session/mod.rs index 232fe7510..419787b7f 100644 --- a/crates/biorouter-cli/src/session/mod.rs +++ b/crates/biorouter-cli/src/session/mod.rs @@ -1741,6 +1741,18 @@ impl CliSession { } } + // A Stop (or a failed stream) while a provider retry waits leaves the + // retry's note the last row the model reads, and the next message would + // be sent behind it. The daemon settles that in `settle_interrupted_turn`; + // the CLI settles no turn, so it retires the note itself. Every exit that + // abandons the turn has cancelled the token, including the commonest: the + // reply loop sees the Stop and its stream ends on its own (`None`). + if cancel_token_clone.is_cancelled() { + self.agent + .retire_unanswered_retry_note(&self.session_id) + .await; + } + // Whatever is still held (a call that never answered) belongs to this // turn and prints before it ends. print_shell_output_lines(shell_order.finish_all(), &mut progress_bars, interactive); diff --git a/crates/biorouter/src/agents/agent.rs b/crates/biorouter/src/agents/agent.rs index 4e9f5fe0d..d2215ff91 100644 --- a/crates/biorouter/src/agents/agent.rs +++ b/crates/biorouter/src/agents/agent.rs @@ -2652,6 +2652,83 @@ fn model_only_user_text_with_new_id>(text: S) -> Message { .with_visibility(false, true) } +/// Whether the next completion goes on from the model's own text: the last +/// rows the model reads are its reply and then a note the person never sees, +/// as when a cut-off reply is retried or continued past its output limit. Its +/// first text then joins that reply, so a local provider keeps the leading +/// whitespace that separates the two +/// ([`crate::providers::formats::text_tool_calls::RecoveryContext::continues_reply`]). +fn request_continues_a_reply(messages: &[Message]) -> bool { + let mut seen = messages + .iter() + .rev() + .filter(|message| message.is_agent_visible()); + let (Some(note), Some(reply)) = (seen.next(), seen.next()) else { + return false; + }; + note.role == rmcp::model::Role::User + && !note.is_user_visible() + && reply.role == rmcp::model::Role::Assistant + && matches!(reply.content.last(), Some(MessageContent::Text(_))) +} + +#[cfg(test)] +mod request_continues_a_reply_tests { + use super::*; + + /// Both continuations the loop writes, a retry of a cut-off reply and one + /// past the output limit, go on from the model's text; a new message from + /// the person, or a note after tool results, does not. + #[test] + fn only_a_note_right_after_the_models_text_continues_it() { + let asked = Message::user().with_text("Write one sentence about foxes."); + let partial = Message::assistant().with_text("The quick brown fox"); + let stopped = Message::assistant() + .with_text("Stopped.") + .with_visibility(true, false); + for note in [ + "Continue the reply", + TRUNCATION_CONTINUATION_MESSAGE, + ] { + assert!(request_continues_a_reply(&[ + asked.clone(), + partial.clone(), + model_only_user_text(note), + ])); + assert!(request_continues_a_reply(&[ + asked.clone(), + partial.clone(), + model_only_user_text(note), + stopped.clone(), + ])); + } + assert!(!request_continues_a_reply(&[ + asked.clone(), + partial.clone(), + stopped, + Message::user().with_text("test"), + ])); + let call = Message::assistant().with_tool_request( + "c1", + Ok(CallToolRequestParams { + task: None, + name: "developer__shell".into(), + arguments: None, + meta: None, + }), + ); + let result = Message::user() + .with_tool_response("c1", Ok(rmcp::model::CallToolResult::success(vec![]))); + assert!(!request_continues_a_reply(&[ + asked.clone(), + call, + result, + model_only_user_text("retry"), + ])); + assert!(!request_continues_a_reply(&[asked])); + } +} + /// Persist one loop iteration's messages and hand them back carrying the uids /// the store actually used. /// @@ -4084,6 +4161,9 @@ pub struct Agent { pub(super) steer_arrivals: tokio::sync::watch::Sender, /// The newest steer a `SteerWaiting` frame has already been sent for. pub(super) announced_steer_seq: std::sync::atomic::AtomicU64, + /// The note the last provider retry was stored behind, as (session id, + /// message id). See [`Agent::retire_unanswered_retry_note`]. + pub(super) pending_retry_note: std::sync::Mutex>, } /// What an accepted steer is waiting behind. See [`AgentEvent::SteerWaiting`]. @@ -5140,6 +5220,7 @@ impl Agent { turn_exit: std::sync::Mutex::new(None), steer_arrivals: tokio::sync::watch::Sender::new(0), announced_steer_seq: std::sync::atomic::AtomicU64::new(0), + pending_retry_note: std::sync::Mutex::new(None), } } @@ -10754,7 +10835,12 @@ impl Agent { /// /// `None` when the write failed: a notice the store does not hold must not be /// handed to a client as history, because a reload would not reproduce it. + /// + /// A retry note nothing answered is retired first: a reply stream that ran + /// out on its own after the Stop reaches here without + /// [`Self::settle_interrupted_turn`]. pub async fn record_turn_stopped(&self, session_id: &str) -> Option { + self.retire_unanswered_retry_note(session_id).await; let notice = self .durable_notice( session_id, @@ -10764,6 +10850,67 @@ impl Agent { notice.id.is_some().then_some(notice) } + /// Hide from the model a retry note nothing answered. A provider retry is + /// stored behind a note telling the model to continue its cut-off reply, or + /// to go on after the tool results; when the turn is settled with that note + /// still the last row the model reads (a Stop while the retry was waiting), + /// the person's next message would be sent behind it, in one user message, + /// and the model told to continue a reply the person stopped. The row is + /// updated in place, not deleted, so nothing is rewritten; a note the retry + /// already answered is followed by that answer and stays. + /// + /// [`Self::settle_interrupted_turn`] and [`Self::record_turn_stopped`] call + /// it. A client that settles neither, the CLI and ACP, calls it itself after + /// a cancelled or failed reply stream has ended or been dropped. + pub async fn retire_unanswered_retry_note(&self, session_id: &str) { + let note_id = { + let mut pending = self + .pending_retry_note + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match pending.take() { + Some((session, id)) if session == session_id => id, + other => { + *pending = other; + return; + } + } + }; + let last = match self + .config + .session_manager + .get_session(session_id, true) + .await + { + Ok(session) => session.conversation.and_then(|conversation| { + conversation + .messages() + .iter() + .rev() + .find(|message| message.is_agent_visible()) + .cloned() + }), + Err(e) => { + warn!( + "could not read the interrupted turn's conversation; its retry note stays: {e}" + ); + return; + } + }; + let Some(note) = last.filter(|message| message.id.as_deref() == Some(note_id.as_str())) + else { + return; + }; + if let Err(e) = self + .config + .session_manager + .hide_from_agent(session_id, ¬e) + .await + { + warn!("could not retire the interrupted turn's retry note: {e}"); + } + } + /// Settle a turn whose reply stream is being abandoned — a Stop, or a stream /// that failed — writing, in the order things happened (D19): /// @@ -10785,6 +10932,7 @@ impl Agent { stopped: bool, ) -> Result { let prose = self.settle_stopped_reply(session_id, in_flight).await; + self.retire_unanswered_retry_note(session_id).await; let steers = self.settle_carried_over_soft_interrupts(session_id).await?; let notice = if stopped { self.record_turn_stopped(session_id).await @@ -11521,13 +11669,16 @@ impl Agent { Some(session_config.id.clone()), coding_agent_bridge::ACTIVE_BRIDGE_URL.scope( bridge_url, - Self::stream_response_from_provider( - iteration_provider, - &system_prompt, - conversation_with_moim.messages(), - &tools, - &toolshim_tools, - live_steer_receiver, + crate::providers::formats::text_tool_calls::CONTINUES_A_REPLY.scope( + request_continues_a_reply(conversation.messages()), + Self::stream_response_from_provider( + iteration_provider, + &system_prompt, + conversation_with_moim.messages(), + &tools, + &toolshim_tools, + live_steer_receiver, + ), ), ), ), @@ -12702,13 +12853,17 @@ impl Agent { // then the same request, so the model answers what // it was asked rather than the note (`RetryPoint`). if let Some(notice) = notice { - messages_to_add.push( - model_only_user_text_with_new_id( - crate::tool_inspection::frame_loop_warnings( - std::slice::from_ref(¬ice), - ) - ), + let note = model_only_user_text_with_new_id( + crate::tool_inspection::frame_loop_warnings( + std::slice::from_ref(¬ice), + ) ); + *self + .pending_retry_note + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + note.id.clone().map(|id| (session_config.id.clone(), id)); + messages_to_add.push(note); } did_recover_provider_error_this_iteration = true; break; diff --git a/crates/biorouter/src/providers/formats/text_tool_calls.rs b/crates/biorouter/src/providers/formats/text_tool_calls.rs index cb832222e..db84bf3c1 100644 --- a/crates/biorouter/src/providers/formats/text_tool_calls.rs +++ b/crates/biorouter/src/providers/formats/text_tool_calls.rs @@ -1060,6 +1060,37 @@ pub struct RecoveryContext { /// For the log line. pub provider: &'static str, pub model: String, + /// This completion continues a reply already shown: a cut-off reply being + /// retried, or continued past its output limit. Its first text joins that + /// reply, so its leading whitespace is the separator between the two and + /// is kept; a reply that is whitespace alone is still dropped (R14). + pub continues_reply: bool, +} + +tokio::task_local! { + /// Set by the agent around a provider call that continues a reply + /// ([`RecoveryContext::continues_reply`]); a local provider reads it where + /// it opens its stream, inside that call. + pub static CONTINUES_A_REPLY: bool; +} + +/// Whether the provider call in progress continues a reply already shown +/// ([`CONTINUES_A_REPLY`]); false outside the agent's scope. +pub fn continues_a_reply() -> bool { + CONTINUES_A_REPLY + .try_with(|continues| *continues) + .unwrap_or(false) +} + +/// Held text as its release shows it: without its leading whitespace (R14), +/// unless the completion continues a reply, where that whitespace separates +/// the continuation from the text it joins. +fn released(held: &str, continues_reply: bool) -> &str { + if continues_reply { + held + } else { + held.trim_start() + } } #[derive(PartialEq)] @@ -1104,7 +1135,7 @@ pub fn recover_leaked_tool_calls( None => break, Some(Err(error)) => { if !held.trim().is_empty() { - let text = held.trim_start().to_string(); + let text = released(&held, context.continues_reply).to_string(); yield (Some(released_at_error(text_message(text, message_id.as_deref()))), None, None); } Err(error)?; @@ -1142,7 +1173,7 @@ pub fn recover_leaked_tool_calls( } else { mode = Mode::Passthrough; // R14: leading whitespace is not an answer. - release.push_str(held.trim_start()); + release.push_str(released(&held, context.continues_reply)); held.clear(); } } @@ -1152,7 +1183,7 @@ pub fn recover_leaked_tool_calls( // G2: a structured call. What was held is text, shown as // text, ahead of the call. if !held.trim().is_empty() { - release.insert_str(0, held.trim_start()); + release.insert_str(0, released(&held, context.continues_reply)); } held.clear(); mode = Mode::Passthrough; @@ -1216,10 +1247,10 @@ pub fn recover_leaked_tool_calls( .iter() .position(|c| matches!(c, MessageContent::ToolRequest(_))) .unwrap_or(message.content.len()); - message.content.insert(at, MessageContent::text(text.trim_start().to_string())); + message.content.insert(at, MessageContent::text(released(&text, context.continues_reply).to_string())); message } - None => text_message(text.trim_start().to_string(), message_id.as_deref()), + None => text_message(released(&text, context.continues_reply).to_string(), message_id.as_deref()), }; yield (Some(message), Some(usage), pending); continue; @@ -1236,7 +1267,7 @@ pub fn recover_leaked_tool_calls( if !held.trim().is_empty() { // The stream ended without a usage item: not a confirmed complete // turn (G3), so whatever was held is shown as the text it is. - yield text_item(held.trim_start().to_string(), message_id.as_deref()); + yield text_item(released(&held, context.continues_reply).to_string(), message_id.as_deref()); } }) } @@ -1717,6 +1748,7 @@ mod tests { tools: tools(), provider: "test", model: "m".to_string(), + continues_reply: false, }, ) .try_collect() @@ -1859,6 +1891,7 @@ mod tests { tools: tools(), provider: "test", model: "m".to_string(), + continues_reply: false, }, ); let mut out = Vec::new(); @@ -1906,6 +1939,86 @@ mod tests { .all(|(message, _, _)| !message.as_ref().is_some_and(is_released_at_error))); } + /// A completion that continues a reply already shown keeps the whitespace + /// it starts with. llama-server cuts between tokens and a token carries its + /// leading space, so trimming it joined a retried reply as "brown + /// foxjumps over" (live, against a scripted llama-server). A fresh reply + /// still loses it, and whitespace alone is still no reply (R14). + #[tokio::test] + async fn a_continued_reply_keeps_the_whitespace_that_joins_it() { + async fn shown(items: Vec, continues_reply: bool) -> String { + let stream: MessageStream = Box::pin(futures::stream::iter( + items.into_iter().map(Ok::<_, ProviderError>), + )); + let out: Vec = recover_leaked_tool_calls( + stream, + RecoveryContext { + tools: tools(), + provider: "test", + model: "m".to_string(), + continues_reply, + }, + ) + .try_collect() + .await + .unwrap(); + out.iter() + .filter_map(|(message, _, _)| message.as_ref()) + .map(Message::as_concat_text) + .collect() + } + let text = |text: &str| { + ( + Some(assistant(vec![MessageContent::text(text)])), + None, + None, + ) + }; + let cases: Vec<(&str, Vec)> = vec![ + // Shown as it arrives. + (" jumps over the lazy dog.", { + let mut items = chunks(" jumps over the lazy dog."); + items.push((None, Some(usage("stop")), None)); + items + }), + ("\n\nThey are clever.", { + let mut items = chunks("\n\nThey are clever."); + items.push((None, Some(usage("stop")), None)); + items + }), + // Held as a possible call, then released as text. + ( + " {\"x\": 1} is JSON.", + vec![ + text(" {"), + text("\"x\": 1} is JSON."), + (None, Some(usage("stop")), None), + ], + ), + // Held until a usage cut at its limit. + ( + " {\"name\": \"x\"", + vec![ + text(" {\"name\": \"x\""), + (None, Some(usage("length")), None), + ], + ), + // Held until the stream ended without a usage. + (" {\"name\": \"x\"", vec![text(" {\"name\": \"x\"")]), + ]; + for (expected, items) in cases { + assert_eq!(shown(items.clone(), true).await, expected, "{expected:?}"); + assert_eq!( + shown(items, false).await, + expected.trim_start(), + "a fresh reply: {expected:?}" + ); + } + let mut whitespace = chunks(" \n"); + whitespace.push((None, Some(usage("stop")), None)); + assert_eq!(shown(whitespace, true).await, "", "R14"); + } + /// L8: a whole-message payload in the template's CALL tag is recovered, /// and the tag is held back while it streams. #[test] diff --git a/crates/biorouter/src/providers/llamacpp.rs b/crates/biorouter/src/providers/llamacpp.rs index 09956570b..9537f3006 100644 --- a/crates/biorouter/src/providers/llamacpp.rs +++ b/crates/biorouter/src/providers/llamacpp.rs @@ -43,7 +43,7 @@ use crate::model::ModelConfig; use crate::privacy::ProviderTier; use crate::providers::formats::openai::{create_request, get_usage, response_to_message}; use crate::providers::formats::text_tool_calls::{ - recover_in_message, recover_leaked_tool_calls, RecoveryContext, + continues_a_reply, recover_in_message, recover_leaked_tool_calls, RecoveryContext, }; use crate::utils::safe_truncate; @@ -1394,6 +1394,7 @@ impl Provider for LlamaCppProvider { tools: filtered_tools.to_vec(), provider: "llamacpp", model: self.model.model_name.clone(), + continues_reply: continues_a_reply(), }, ); Ok(hold_lease_while_generating_with_recovery( diff --git a/crates/biorouter/src/providers/llamacpp_sidecar.rs b/crates/biorouter/src/providers/llamacpp_sidecar.rs index c7caba3f5..49e746ce4 100644 --- a/crates/biorouter/src/providers/llamacpp_sidecar.rs +++ b/crates/biorouter/src/providers/llamacpp_sidecar.rs @@ -1311,11 +1311,10 @@ fn backend_failure_message(model: &str, ctx_size: Option, tail: &str) -> Some(window) if window > LLAMACPP_AGENT_CONTEXT_SIZE => format!( "Quit other apps or local models that hold GPU memory, or set \ LLAMACPP_CONTEXT_SIZE to {LLAMACPP_AGENT_CONTEXT_SIZE} (not lower: Biorouter's \ - instructions and tool definitions alone take about 31,000-40,000 tokens), and \ - try again." + instructions and tool definitions alone take about 31,000-40,000 tokens)." ), _ => format!( - "Quit other apps or local models that hold GPU memory and try again. A smaller \ + "Quit other apps or local models that hold GPU memory. A smaller \ context window will not help: {LLAMACPP_AGENT_CONTEXT_SIZE} tokens is the \ smallest Biorouter runs a local model with, because its instructions and tool \ definitions alone take about 31,000-40,000 tokens. If this keeps happening, use \ @@ -1323,8 +1322,8 @@ fn backend_failure_message(model: &str, ctx_size: Option, tail: &str) -> ), }; format!( - "Llama Server ran out of GPU memory while running {model} and has been restarted. \ - {advice}" + "Llama Server ran out of GPU memory while running {model}. Biorouter stopped that \ + server, and the next request starts a fresh one. {advice}" ) } else { // No "try again": the chat shows this right before "Trying again @@ -3658,7 +3657,10 @@ mod tests { assert!(looks_like_memory_failure(tail)); let message = backend_failure_message("gemma4-31b", Some(131_072), tail); assert!(message.contains("ran out of GPU memory"), "{message}"); - assert!(message.contains("restarted"), "{message}"); + assert!( + message.contains("Biorouter stopped that server"), + "{message}" + ); let other = backend_failure_message("gemma4-31b", Some(131_072), "E decode: failed"); assert!(other.contains("E decode: failed"), "{other}"); @@ -3690,6 +3692,31 @@ mod tests { ); } + /// The out-of-memory notice is a compute-failure notice too, shown right + /// before "Trying again (1/1)…": it keeps its advice but must not ask the + /// person to try again, nor say a server was restarted that was stopped. + #[test] + fn an_out_of_memory_failure_does_not_ask_the_person_to_try_again() { + let tail = + "E error: Insufficient Memory (00000008:kIOGPUCommandBufferCallbackErrorOutOfMemory)"; + for window in [Some(131_072), Some(65_536), None] { + let message = backend_failure_message("qwen3.8-27b", window, tail); + assert!( + message.starts_with( + "Llama Server ran out of GPU memory while running qwen3.8-27b. Biorouter \ + stopped that server, and the next request starts a fresh one. Quit other \ + apps or local models that hold GPU memory" + ), + "{window:?}: {message}" + ); + assert!( + !message.to_lowercase().contains("try again"), + "{window:?}: {message}" + ); + assert!(!message.contains("restarted"), "{window:?}: {message}"); + } + } + /// A GPU out-of-memory mid-request offers a smaller window only while there /// is one above the floor, and names the floor when it does. The old advice /// ("set a smaller LLAMACPP_CONTEXT_SIZE") sent a machine already at 64k to diff --git a/crates/biorouter/src/providers/ollama.rs b/crates/biorouter/src/providers/ollama.rs index 6d063af2f..6e187e2df 100644 --- a/crates/biorouter/src/providers/ollama.rs +++ b/crates/biorouter/src/providers/ollama.rs @@ -44,7 +44,7 @@ use super::formats::ollama::{ OLLAMA_CHAT_PATH, }; use super::formats::text_tool_calls::{ - recover_in_message, recover_leaked_tool_calls, RecoveryContext, + continues_a_reply, recover_in_message, recover_leaked_tool_calls, RecoveryContext, }; use super::retry::{ProviderRetry, TimeoutRetry}; use super::utils::{ @@ -1595,6 +1595,9 @@ struct NativeTurn { /// Whether the turn's one corrective retry is already spent. retried: bool, memory: Option, + /// Read when the turn opened, inside the agent's provider call: the + /// stream is read later, outside it. + continues_reply: bool, } /// A native turn: decode, recover leaked calls (TIER1), guard, and — at most @@ -1627,6 +1630,7 @@ fn native_turn(turn: NativeTurn, first: Response, first_log: RequestLog) -> Mess tools, mut retried, memory, + continues_reply, } = turn; let model = model_config.model_name.clone(); let mut next_response = Some((first, first_log)); @@ -1669,6 +1673,7 @@ fn native_turn(turn: NativeTurn, first: Response, first_log: RequestLog) -> Mess tools: tools.clone(), provider: "ollama", model: model.clone(), + continues_reply, }, ); let mut guard = TurnGuard::new(check.clone()); @@ -2477,6 +2482,7 @@ impl OllamaProvider { tools: tools.to_vec(), retried: false, memory: self.memory_observer(model_config), + continues_reply: continues_a_reply(), }; match outcome { NativeOutcome::Ok(response) => Ok(Some(native_turn(turn, response, log))), @@ -2659,6 +2665,7 @@ impl OllamaProvider { tools: tools.to_vec(), provider: "ollama", model: self.model.model_name.clone(), + continues_reply: continues_a_reply(), }, ); Ok(guard_turn( diff --git a/crates/biorouter/src/session/session_manager.rs b/crates/biorouter/src/session/session_manager.rs index d425b714b..93f6fbe04 100644 --- a/crates/biorouter/src/session/session_manager.rs +++ b/crates/biorouter/src/session/session_manager.rs @@ -2200,6 +2200,27 @@ impl SessionManager { self.add_message_adopting_uid(id, message).await } + /// Hide a stored row from the model, keeping it stored: an update in place, + /// compared against the metadata it was read with, so a row changed since + /// is left alone. + pub(crate) async fn hide_from_agent(&self, id: &str, message: &Message) -> Result<()> { + let Some(uid) = message.id.as_deref() else { + return Ok(()); + }; + sqlx::query( + "UPDATE messages SET metadata_json = ? WHERE session_id = ? AND msg_uid = ? AND metadata_json = ?", + ) + .bind(serde_json::to_string( + &message.metadata.clone().with_agent_invisible(), + )?) + .bind(id) + .bind(uid) + .bind(serde_json::to_string(&message.metadata)?) + .execute(self.storage.pool().await?) + .await?; + Ok(()) + } + /// Unconditional whole-history rewrite: DELETE every message of the session /// and re-INSERT the supplied ones. /// diff --git a/crates/biorouter/tests/output_recovery_agent_loop.rs b/crates/biorouter/tests/output_recovery_agent_loop.rs index cdfa80821..ad7288260 100644 --- a/crates/biorouter/tests/output_recovery_agent_loop.rs +++ b/crates/biorouter/tests/output_recovery_agent_loop.rs @@ -1275,6 +1275,11 @@ struct CutStreamProvider { failures: usize, first: &'static str, recover: bool, + /// What a call after the failures answers. + answer: &'static str, + /// A call that never answers, not even with headers: a retry held by a + /// busy server, which the person stops while it waits. + held_call: Option, /// The agent-visible text of every request, `` block stripped. sent: Mutex>>, /// Whether each request carried the `` block. @@ -1288,11 +1293,23 @@ impl CutStreamProvider { failures, first, recover, + answer: CUT_ANSWER, + held_call: None, sent: Mutex::default(), with_block: Mutex::default(), } } + fn answering(mut self, answer: &'static str) -> Self { + self.answer = answer; + self + } + + fn holding_call(mut self, call: usize) -> Self { + self.held_call = Some(call); + self + } + fn record(&self, messages: &[Message]) { let visible: Vec<&Message> = messages .iter() @@ -1352,10 +1369,13 @@ impl Provider for CutStreamProvider { tools: &[Tool], ) -> Result { use biorouter::providers::formats::text_tool_calls::{ - recover_leaked_tool_calls, RecoveryContext, + continues_a_reply, recover_leaked_tool_calls, RecoveryContext, }; self.record(messages); let call = self.calls.fetch_add(1, Ordering::SeqCst); + if self.held_call == Some(call) { + std::future::pending::<()>().await; + } let id = format!("reply-{call}"); let items: Vec> = if call < self.failures { let mut items: Vec> = self @@ -1382,7 +1402,7 @@ impl Provider for CutStreamProvider { ); usage.finish_reason = Some("stop".to_string()); vec![Ok(( - Some(Message::assistant().with_id(id).with_text(CUT_ANSWER)), + Some(Message::assistant().with_id(id).with_text(self.answer)), Some(usage), None, ))] @@ -1395,6 +1415,7 @@ impl Provider for CutStreamProvider { tools: tools.to_vec(), provider: "llamacpp", model: "mock-model".to_string(), + continues_reply: continues_a_reply(), }, ) } else { @@ -1637,3 +1658,231 @@ async fn a_partial_reply_ending_in_whitespace_keeps_the_block() { ); assert!(reply.shown.iter().any(|text| text == CUT_ANSWER)); } + +/// Live, against a scripted llama-server: a reply cut after "The quick brown +/// fox" was continued with " jumps over the lazy dog.", and recovery dropped +/// the continuation's leading space, so the CLI printed "foxjumps" and the +/// stored rows read "fox" and "jumps over the lazy dog.". The retry continues +/// the reply, so the space that joins the two is kept. +#[tokio::test] +async fn a_retried_partial_reply_keeps_the_space_that_joins_it() { + const CONTINUED: &str = " jumps over the lazy dog."; + let provider = + Arc::new(CutStreamProvider::new("The quick brown fox", 1, true).answering(CONTINUED)); + let reply = reply_through_a_cut_stream(provider.clone()).await; + + assert_eq!(provider.calls.load(Ordering::SeqCst), 2); + assert!(!reply.aborted); + assert!( + reply.shown.iter().any(|text| text == CONTINUED), + "{:?}", + reply.shown + ); + let replies: String = reply + .stored + .iter() + .filter_map(|row| row.strip_prefix("Assistant: ")) + .collect(); + assert_eq!( + replies, "The quick brown fox jumps over the lazy dog.", + "{:?}", + reply.stored + ); + assert!( + reply + .stored + .iter() + .any(|row| row.contains("biorouter-loop-guard")), + "the continuation still follows its note: {:?}", + reply.stored + ); +} + +/// Live, against a scripted llama-server: a Stop while the retry of a cut-off +/// reply was waiting left its "Continue the reply from exactly where it +/// stopped" note in the history, and the person's next message was sent behind +/// it in one user message, so the model continued the stopped sentence instead +/// of answering. The settled turn hides the note from the model, on each of the +/// runner's paths: it drops the stream and settles the turn (a Stop, or a +/// failed stream), or the stream runs out on its own after the cancel and only +/// the "Stopped." notice is written. +#[tokio::test] +async fn a_stop_while_a_retry_waits_retires_its_note() { + stop_while_a_retry_waits(Abandoned::Stopped).await; +} + +#[tokio::test] +async fn a_stop_the_reply_loop_sees_retires_the_retry_note_too() { + stop_while_a_retry_waits(Abandoned::StoppedInTheLoop).await; +} + +#[tokio::test] +async fn a_stream_that_fails_while_a_retry_waits_retires_its_note() { + stop_while_a_retry_waits(Abandoned::Failed).await; +} + +/// Live, with the CLI: a Ctrl-C while the retry waited settled nothing, so +/// `biorouter run -r -t test` sent "[note]\ntest" in one user message and the +/// model echoed the note. The CLI and ACP settle no turn and write no "Stopped." +/// notice; they retire the note themselves, whether the reply loop saw the +/// cancel and its stream ran out (what the CLI measured) or the client dropped +/// the stream first. +#[tokio::test] +async fn a_client_that_settles_no_turn_retires_the_retry_note_itself() { + stop_while_a_retry_waits(Abandoned::ByTheClientAfterTheLoop).await; +} + +#[tokio::test] +async fn a_client_that_drops_the_stream_retires_the_retry_note_itself() { + stop_while_a_retry_waits(Abandoned::DroppedByTheClient).await; +} + +#[derive(PartialEq)] +enum Abandoned { + Stopped, + StoppedInTheLoop, + Failed, + ByTheClientAfterTheLoop, + DroppedByTheClient, +} + +async fn stop_while_a_retry_waits(how: Abandoned) { + let provider = Arc::new(CutStreamProvider::new("The quick brown fox", 1, true).holding_call(1)); + let overrides = HashMap::from([ + ( + "BIOROUTER_MISTAKE_STREAK_DETECTION".to_string(), + "true".to_string(), + ), + ( + "BIOROUTER_PROVIDER_ERROR_RETRIES".to_string(), + "1".to_string(), + ), + ]); + with_config_overrides(overrides, async { + let (agent, sessions, session_id, _work, _data, _permissions) = + test_agent(provider.clone()).await; + let session_config = || SessionConfig { + id: session_id.clone(), + schedule_id: None, + max_turns: Some(4), + max_tool_calls: None, + budget: None, + retry_config: None, + reasoning_effort: None, + }; + + // The runner's in-flight rows: every message since the last persisted + // batch. A Stop drops the stream and hands them to the settlement. + let mut in_flight: Vec = Vec::new(); + let cancel = tokio_util::sync::CancellationToken::new(); + { + let stream = agent + .reply( + Message::user().with_text("Write one sentence about foxes."), + session_config(), + Some(cancel.clone()), + ) + .await + .unwrap(); + tokio::pin!(stream); + tokio::time::timeout(Duration::from_secs(10), async { + while provider.calls.load(Ordering::SeqCst) < 2 { + match tokio::time::timeout(Duration::from_millis(50), stream.next()).await { + Ok(Some(event)) => match event.unwrap() { + AgentEvent::Message(message) => in_flight.push(message), + AgentEvent::MessagesPersisted(_) => in_flight.clear(), + _ => {} + }, + Ok(None) => panic!("the turn ended before its retry was sent"), + Err(_) => {} + } + } + }) + .await + .expect("the retry is sent"); + if matches!( + how, + Abandoned::StoppedInTheLoop | Abandoned::ByTheClientAfterTheLoop + ) { + cancel.cancel(); + tokio::time::timeout(Duration::from_secs(10), async { + while let Some(event) = stream.next().await { + event.unwrap(); + } + }) + .await + .expect("the stopped turn ends"); + } + } + match how { + Abandoned::StoppedInTheLoop => { + agent.record_turn_stopped(&session_id).await.unwrap(); + } + Abandoned::ByTheClientAfterTheLoop => { + agent.retire_unanswered_retry_note(&session_id).await; + } + Abandoned::DroppedByTheClient => { + cancel.cancel(); + agent.retire_unanswered_retry_note(&session_id).await; + } + Abandoned::Stopped | Abandoned::Failed => { + agent + .settle_interrupted_turn(&session_id, &in_flight, how == Abandoned::Stopped) + .await + .unwrap(); + } + } + + let stream = agent + .reply(Message::user().with_text("test"), session_config(), None) + .await + .unwrap(); + tokio::pin!(stream); + tokio::time::timeout(Duration::from_secs(10), async { + while let Some(event) = stream.next().await { + event.unwrap(); + } + }) + .await + .expect("the next reply settles"); + + let sent = provider.sent.lock().unwrap().clone(); + assert_eq!(sent.len(), 3, "{sent:?}"); + assert!( + sent[1] + .iter() + .any(|row| row.contains("biorouter-loop-guard")), + "the retry was sent behind its note: {sent:?}" + ); + assert_eq!( + sent[2].last().map(String::as_str), + Some("User: test"), + "{sent:?}" + ); + assert!( + sent[2] + .iter() + .all(|row| !row.contains("biorouter-loop-guard")), + "{sent:?}" + ); + let stored = sessions + .get_session(&session_id, true) + .await + .unwrap() + .conversation + .unwrap(); + let notes: Vec<&Message> = stored + .messages() + .iter() + .filter(|message| message.as_concat_text().contains("biorouter-loop-guard")) + .collect(); + assert_eq!(notes.len(), 1, "the note is kept, not deleted"); + assert!(!notes[0].is_agent_visible(), "{:?}", notes[0]); + assert!(stored + .messages() + .iter() + .any(|message| message.is_agent_visible() + && message.as_concat_text() == "The quick brown fox")); + }) + .await; +}