mirror of
https://github.com/codewhale-hq/Codewhale.git
synced 2026-11-11 10:01:30 -10:00
fix(runtime): preserve history, controls, and partial MCP frames
Bind saved sessions to a verified durable turn checkpoint and replay newer turns after Engine eviction or Runtime reopen. Preserve signed/media blocks and exact fork boundaries; reject ambiguous or changed saved histories. Bind cancellation and steering to admitted turns before Engine startup. Keep incomplete MCP stdio frames across receive cancellation and enforce the frame cap across resumed reads, without replaying tool requests. Validation: scoped product tests: 603 passed; 0 failed; 3 ignored. Focused regressions: 11 passed; 0 failed. Negative controls restoring former behaviors: 0 passed; 6 failed at the expected assertions, then fixes restored. Production library cargo check and git diff --check pass. Earlier sandbox run: 545 passed; 22 localhost-listener permission failures; all 22 passed when allowed to bind their local fixtures. Full workspace/web gate and hosted CI remain integration-owner checks.
This commit is contained in:
1 parent
d6698f193d
commit
800474b2a9
10 files changed
+1029
-159
No files matched your search
@@ -684,7 +684,8 @@ pub struct EngineHandle {
|
||||
/// Send user input responses to the engine
|
||||
tx_user_input: mpsc::Sender<UserInputDecision>,
|
||||
/// Send steer input for an in-flight turn.
|
||||
tx_steer: mpsc::Sender<String>,
|
||||
tx_steer: mpsc::Sender<handle::SteerInput>,
|
||||
turn_controls: Arc<StdMutex<handle::TurnControls>>,
|
||||
/// Shared pause flag set by the TUI and read by the turn loop.
|
||||
shared_paused: Arc<StdMutex<bool>>,
|
||||
/// Whether the host must construct the route's concrete provider client
|
||||
@@ -875,7 +876,9 @@ pub struct Engine {
|
||||
/// approval gate still fails closed.
|
||||
approval_receipt_store: Result<ApprovalReceiptStore, String>,
|
||||
rx_user_input: mpsc::Receiver<UserInputDecision>,
|
||||
rx_steer: mpsc::Receiver<String>,
|
||||
rx_steer: mpsc::Receiver<handle::SteerInput>,
|
||||
turn_controls: Arc<StdMutex<handle::TurnControls>>,
|
||||
admitted_turn_control: Option<handle::TurnControl>,
|
||||
tx_event: mpsc::Sender<Event>,
|
||||
/// Wakeup channel for the parent turn loop when a direct child sub-agent
|
||||
/// terminates (issue #756). Cloned into `SubAgentRuntime` so the runtime
|
||||
@@ -1291,28 +1294,46 @@ impl Engine {
|
||||
.finish(id);
|
||||
}
|
||||
|
||||
fn reset_cancel_token(&mut self) {
|
||||
let token = CancellationToken::new();
|
||||
self.cancel_token = token.clone();
|
||||
match self.shared_cancel_token.lock() {
|
||||
Ok(mut shared) => {
|
||||
*shared = token;
|
||||
}
|
||||
Err(poisoned) => {
|
||||
*poisoned.into_inner() = token;
|
||||
fn begin_turn_control(&mut self) -> handle::TurnControlGuard {
|
||||
let mut controls = self
|
||||
.turn_controls
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
let control = self
|
||||
.admitted_turn_control
|
||||
.take()
|
||||
.unwrap_or_else(|| controls.fresh());
|
||||
self.cancel_token = control.cancel.clone();
|
||||
self.cancel_reason = Arc::clone(&control.reason);
|
||||
*self
|
||||
.shared_cancel_token
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = control.cancel.clone();
|
||||
*self
|
||||
.shared_paused
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = false;
|
||||
controls.active = Some(control.clone());
|
||||
handle::TurnControlGuard {
|
||||
controls: Arc::clone(&self.turn_controls),
|
||||
id: control.id,
|
||||
}
|
||||
}
|
||||
|
||||
fn next_turn_steer(&mut self) -> Option<String> {
|
||||
let active_id = self
|
||||
.turn_controls
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.active
|
||||
.as_ref()
|
||||
.map(|control| control.id);
|
||||
while let Ok(steer) = self.rx_steer.try_recv() {
|
||||
if steer.turn_id == active_id {
|
||||
return Some(steer.content);
|
||||
}
|
||||
}
|
||||
// Fresh turn → clear any latched cancellation reason from the
|
||||
// previous turn so a downstream "request cancelled" message
|
||||
// doesn't inherit a stale cause.
|
||||
match self.cancel_reason.lock() {
|
||||
Ok(mut slot) => *slot = None,
|
||||
Err(poisoned) => *poisoned.into_inner() = None,
|
||||
}
|
||||
match self.shared_paused.lock() {
|
||||
Ok(mut paused) => *paused = false,
|
||||
Err(poisoned) => *poisoned.into_inner() = false,
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn env_only_api_key_recovery_hint(api_config: &Config) -> Option<String> {
|
||||
@@ -1472,6 +1493,7 @@ impl Engine {
|
||||
let (tx_approval, rx_approval) = mpsc::channel(64);
|
||||
let (tx_user_input, rx_user_input) = mpsc::channel(32);
|
||||
let (tx_steer, rx_steer) = mpsc::channel(64);
|
||||
let turn_controls = Arc::new(StdMutex::new(handle::TurnControls::default()));
|
||||
let (tx_subagent_completion, rx_subagent_completion) = mpsc::unbounded_channel();
|
||||
let cancel_token = CancellationToken::new();
|
||||
let shared_cancel_token = Arc::new(StdMutex::new(cancel_token.clone()));
|
||||
@@ -1726,6 +1748,8 @@ impl Engine {
|
||||
approval_receipt_store,
|
||||
rx_user_input,
|
||||
rx_steer,
|
||||
turn_controls: Arc::clone(&turn_controls),
|
||||
admitted_turn_control: None,
|
||||
tx_event,
|
||||
tx_subagent_completion,
|
||||
rx_subagent_completion,
|
||||
@@ -1757,6 +1781,7 @@ impl Engine {
|
||||
tx_approval,
|
||||
tx_user_input,
|
||||
tx_steer,
|
||||
turn_controls,
|
||||
shared_paused,
|
||||
client_preflight_required: true,
|
||||
live_runtime_authority,
|
||||
@@ -1793,7 +1818,7 @@ impl Engine {
|
||||
auto_approve: bool,
|
||||
approval_mode: crate::tui::approval::ApprovalMode,
|
||||
) {
|
||||
self.reset_cancel_token();
|
||||
let turn_control = self.begin_turn_control();
|
||||
self.turn_counter = self.turn_counter.saturating_add(1);
|
||||
|
||||
let turn_id = format!(
|
||||
@@ -1947,6 +1972,7 @@ impl Engine {
|
||||
if status == TurnOutcomeStatus::Interrupted {
|
||||
self.emit_interrupted_survivor_status().await;
|
||||
}
|
||||
drop(turn_control);
|
||||
let _ = self
|
||||
.tx_event
|
||||
.send(Event::TurnComplete {
|
||||
@@ -2593,6 +2619,15 @@ impl Engine {
|
||||
verbosity,
|
||||
provenance,
|
||||
} => {
|
||||
self.admitted_turn_control = {
|
||||
let mut controls = self
|
||||
.turn_controls
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
let control = controls.pending.pop_front();
|
||||
controls.active = control.clone();
|
||||
control
|
||||
};
|
||||
self.handle_send_message(
|
||||
content,
|
||||
mode,
|
||||
@@ -4692,6 +4727,7 @@ impl Engine {
|
||||
verbosity: Option<String>,
|
||||
provenance: UserInputProvenance,
|
||||
) -> SendMessageOutcome {
|
||||
let turn_control = self.begin_turn_control();
|
||||
let mut goal_objective = goal_objective;
|
||||
let mut goal_token_budget = goal_token_budget;
|
||||
let mut goal_status = goal_status;
|
||||
@@ -4838,17 +4874,12 @@ impl Engine {
|
||||
if let Some(status) = input_policy.status() {
|
||||
let _ = self.tx_event.send(Event::status(status)).await;
|
||||
}
|
||||
// Reset cancel token for fresh turn (in case previous was cancelled)
|
||||
self.reset_cancel_token();
|
||||
|
||||
// Track the complete effective mode policy so mid-turn metadata, `/edit`,
|
||||
// idle worker resumptions, and approval gates cannot read a stale policy
|
||||
// after the UI changed modes (#3568).
|
||||
self.apply_runtime_mode_policy(&input_policy);
|
||||
|
||||
// Drain stale steer messages from previous turns.
|
||||
while self.rx_steer.try_recv().is_ok() {}
|
||||
|
||||
// Create turn context first so start event includes a stable turn id.
|
||||
// An active goal gets the host's goal allowance (#5994); turns with
|
||||
// an explicit per-invocation ceiling (exec --max-turns, child
|
||||
@@ -5262,6 +5293,7 @@ impl Engine {
|
||||
if status == TurnOutcomeStatus::Interrupted {
|
||||
self.emit_interrupted_survivor_status().await;
|
||||
}
|
||||
drop(turn_control);
|
||||
let turn_complete_delivered = self
|
||||
.tx_event
|
||||
.send(Event::TurnComplete {
|
||||
@@ -7432,7 +7464,7 @@ pub(crate) struct MockEngineHandle {
|
||||
pub rx_op: mpsc::Receiver<Op>,
|
||||
rx_approval: mpsc::Receiver<ApprovalDecision>,
|
||||
rx_user_input: mpsc::Receiver<UserInputDecision>,
|
||||
pub rx_steer: mpsc::Receiver<String>,
|
||||
pub rx_steer: mpsc::Receiver<handle::SteerInput>,
|
||||
pub tx_event: mpsc::Sender<Event>,
|
||||
pub cancel_token: CancellationToken,
|
||||
}
|
||||
@@ -7519,6 +7551,7 @@ pub(crate) fn mock_engine_handle() -> MockEngineHandle {
|
||||
tx_approval,
|
||||
tx_user_input,
|
||||
tx_steer,
|
||||
turn_controls: Arc::new(StdMutex::new(handle::TurnControls::default())),
|
||||
shared_paused,
|
||||
client_preflight_required: false,
|
||||
live_runtime_authority,
|
||||
|
||||
@@ -9,7 +9,10 @@
|
||||
//! so the agent loop's mailbox API is reviewable on its own.
|
||||
|
||||
use anyhow::Result;
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::{Arc, Mutex as StdMutex};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use super::approval::{ApprovalDecision, UserInputDecision};
|
||||
use super::{
|
||||
@@ -17,6 +20,92 @@ use super::{
|
||||
UserInputResponse,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(super) struct TurnControl {
|
||||
pub id: u64,
|
||||
pub cancel: CancellationToken,
|
||||
pub reason: Arc<StdMutex<Option<CancelReason>>>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct TurnControls {
|
||||
next_id: u64,
|
||||
pub active: Option<TurnControl>,
|
||||
pub pending: VecDeque<TurnControl>,
|
||||
}
|
||||
|
||||
impl TurnControls {
|
||||
pub fn fresh(&mut self) -> TurnControl {
|
||||
self.next_id = self
|
||||
.next_id
|
||||
.checked_add(1)
|
||||
.expect("turn control id exhausted");
|
||||
TurnControl {
|
||||
id: self.next_id,
|
||||
cancel: CancellationToken::new(),
|
||||
reason: Arc::new(StdMutex::new(None)),
|
||||
}
|
||||
}
|
||||
|
||||
fn target(&self) -> Option<&TurnControl> {
|
||||
self.active.as_ref().or_else(|| self.pending.front())
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct TurnControlGuard {
|
||||
pub controls: Arc<StdMutex<TurnControls>>,
|
||||
pub id: u64,
|
||||
}
|
||||
|
||||
impl Drop for TurnControlGuard {
|
||||
fn drop(&mut self) {
|
||||
let mut controls = self
|
||||
.controls
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if controls
|
||||
.active
|
||||
.as_ref()
|
||||
.is_some_and(|active| active.id == self.id)
|
||||
{
|
||||
controls.active = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct SteerInput {
|
||||
pub(super) turn_id: Option<u64>,
|
||||
pub(crate) content: String,
|
||||
}
|
||||
|
||||
impl std::ops::Deref for SteerInput {
|
||||
type Target = str;
|
||||
fn deref(&self) -> &str {
|
||||
&self.content
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for SteerInput {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
self.content.fmt(f)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct SteerPermit {
|
||||
permit: mpsc::OwnedPermit<SteerInput>,
|
||||
turn_id: Option<u64>,
|
||||
}
|
||||
|
||||
impl SteerPermit {
|
||||
pub(crate) fn send(self, content: String) {
|
||||
self.permit.send(SteerInput {
|
||||
turn_id: self.turn_id,
|
||||
content,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
impl EngineHandle {
|
||||
/// Called only while Runtime holds the idle turn admission claim. The
|
||||
/// following SendMessage refreshes the existing prompt/config projection.
|
||||
@@ -50,11 +139,11 @@ impl EngineHandle {
|
||||
/// Send an operation to the engine
|
||||
pub async fn send(&self, op: Op) -> Result<()> {
|
||||
let authority = Self::change_mode_authority(&op);
|
||||
let permit = self.tx_op.reserve().await?;
|
||||
let permit = self.tx_op.clone().reserve_owned().await?;
|
||||
if let Some(authority) = authority {
|
||||
self.publish_runtime_authority(authority);
|
||||
}
|
||||
permit.send(op);
|
||||
self.send_reserved_op(permit, op);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -65,7 +154,7 @@ impl EngineHandle {
|
||||
/// safely be dropped and re-requested on the next drain cycle.
|
||||
pub fn try_send(&self, op: Op) -> Result<()> {
|
||||
let authority = Self::change_mode_authority(&op);
|
||||
let result = self.tx_op.try_send(op);
|
||||
let result = self.tx_op.clone().try_reserve_owned();
|
||||
// A full channel already guarantees that the engine will wake and
|
||||
// drain an operation. Publish the typed authority anyway: the drain
|
||||
// applies pending authority before handling that queued operation, so
|
||||
@@ -76,10 +165,24 @@ impl EngineHandle {
|
||||
{
|
||||
self.publish_runtime_authority(authority);
|
||||
}
|
||||
result?;
|
||||
self.send_reserved_op(result?, op);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Bind controls and enqueue under one lock, preserving the same FIFO as
|
||||
/// the operation mailbox even when several senders hold reserved slots.
|
||||
pub(crate) fn send_reserved_op(&self, permit: mpsc::OwnedPermit<Op>, op: Op) {
|
||||
let mut controls = self
|
||||
.turn_controls
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if matches!(&op, Op::SendMessage { .. }) {
|
||||
let control = controls.fresh();
|
||||
controls.pending.push_back(control);
|
||||
}
|
||||
permit.send(op);
|
||||
}
|
||||
|
||||
fn change_mode_authority(op: &Op) -> Option<LiveRuntimeAuthority> {
|
||||
let Op::ChangeMode {
|
||||
mode,
|
||||
@@ -144,8 +247,15 @@ impl EngineHandle {
|
||||
/// Reserve capacity for a runtime steer before it mutates durable state.
|
||||
/// The owned permit lets the caller persist and dispatch synchronously,
|
||||
/// without a cancellation point between those two operations.
|
||||
pub(crate) async fn reserve_steer(&self) -> Result<mpsc::OwnedPermit<String>> {
|
||||
Ok(self.tx_steer.clone().reserve_owned().await?)
|
||||
pub(crate) async fn reserve_steer(&self) -> Result<SteerPermit> {
|
||||
let permit = self.tx_steer.clone().reserve_owned().await?;
|
||||
let turn_id = self
|
||||
.turn_controls
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.target()
|
||||
.map(|control| control.id);
|
||||
Ok(SteerPermit { permit, turn_id })
|
||||
}
|
||||
|
||||
/// Cancel the current request (user-initiated path — keeps the
|
||||
@@ -158,6 +268,19 @@ impl EngineHandle {
|
||||
/// Cancel the current request and latch the reason so downstream
|
||||
/// "request cancelled" error messages can name a cause.
|
||||
pub fn cancel_with_reason(&self, reason: CancelReason) {
|
||||
// Keep turn activation excluded until both the admitted control and
|
||||
// the legacy shared token have been canceled.
|
||||
let controls = self
|
||||
.turn_controls
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some(control) = controls.target() {
|
||||
*control
|
||||
.reason
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(reason);
|
||||
control.cancel.cancel();
|
||||
}
|
||||
match self.cancel_reason.lock() {
|
||||
Ok(mut slot) => *slot = Some(reason),
|
||||
Err(poisoned) => *poisoned.into_inner() = Some(reason),
|
||||
@@ -173,6 +296,14 @@ impl EngineHandle {
|
||||
#[must_use]
|
||||
#[allow(dead_code)]
|
||||
pub fn is_cancelled(&self) -> bool {
|
||||
if let Some(control) = self
|
||||
.turn_controls
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.target()
|
||||
{
|
||||
return control.cancel.is_cancelled();
|
||||
}
|
||||
match self.cancel_token.lock() {
|
||||
Ok(token) => token.is_cancelled(),
|
||||
Err(poisoned) => poisoned.into_inner().is_cancelled(),
|
||||
@@ -253,7 +384,7 @@ impl EngineHandle {
|
||||
|
||||
/// Steer an in-flight turn with additional user input.
|
||||
pub async fn steer(&self, content: impl Into<String>) -> Result<()> {
|
||||
self.tx_steer.send(content.into()).await?;
|
||||
self.reserve_steer().await?.send(content.into());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -4955,10 +4955,21 @@ async fn steer_during_final_coordination_response_gets_its_own_provider_reply()
|
||||
};
|
||||
let (mut engine, handle) =
|
||||
Engine::new_with_model_client(engine_config, &Config::default(), client);
|
||||
let tx_steer = handle.tx_steer.clone();
|
||||
let steer_handle = handle.clone();
|
||||
mock.push_factory(move |_request| {
|
||||
tx_steer
|
||||
.try_send("Include this user steer in the final answer.".to_string())
|
||||
let turn_id = steer_handle
|
||||
.turn_controls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.active
|
||||
.as_ref()
|
||||
.map(|control| control.id);
|
||||
steer_handle
|
||||
.tx_steer
|
||||
.try_send(handle::SteerInput {
|
||||
turn_id,
|
||||
content: "Include this user steer in the final answer.".to_string(),
|
||||
})
|
||||
.expect("test steer channel remains open");
|
||||
canned::tool_call_turn(
|
||||
"call-coordination-read-2",
|
||||
@@ -9282,7 +9293,7 @@ fn engine_handle_cancel_tracks_latest_turn_token() {
|
||||
let (mut engine, handle) = Engine::new(EngineConfig::default(), &Config::default());
|
||||
let stale_token = engine.cancel_token.clone();
|
||||
|
||||
engine.reset_cancel_token();
|
||||
let _turn_control = engine.begin_turn_control();
|
||||
handle.cancel();
|
||||
|
||||
assert!(engine.cancel_token.is_cancelled());
|
||||
@@ -20911,6 +20922,7 @@ fn engine_handle_try_send_does_not_block_when_op_channel_is_full() {
|
||||
tx_approval: mpsc::channel(1).0,
|
||||
tx_user_input: mpsc::channel(1).0,
|
||||
tx_steer: mpsc::channel(1).0,
|
||||
turn_controls: Arc::new(StdMutex::new(handle::TurnControls::default())),
|
||||
shared_paused: Arc::new(StdMutex::new(false)),
|
||||
client_preflight_required: true,
|
||||
live_runtime_authority: Arc::new(StdMutex::new(LiveRuntimeAuthorityState::new(
|
||||
@@ -21182,7 +21194,7 @@ readline.createInterface({ input: process.stdin }).on('line', async line => {
|
||||
)
|
||||
.await
|
||||
.expect("stop interrupts explicit schema wait");
|
||||
engine.reset_cancel_token();
|
||||
let _turn_control = engine.begin_turn_control();
|
||||
fs::write(&release, "release").unwrap();
|
||||
let build = build.expect("explicit fast/failed selections must not wait for slow");
|
||||
let active = build.surface.active.unwrap_or_default();
|
||||
|
||||
@@ -805,7 +805,7 @@ impl Engine {
|
||||
}
|
||||
|
||||
let mut accepted_steer = false;
|
||||
while let Ok(steer) = self.rx_steer.try_recv() {
|
||||
while let Some(steer) = self.next_turn_steer() {
|
||||
let steer = steer.trim().to_string();
|
||||
if steer.is_empty() {
|
||||
continue;
|
||||
@@ -4386,7 +4386,7 @@ impl Engine {
|
||||
let Some(event_result) = poll_outcome else {
|
||||
break;
|
||||
};
|
||||
while let Ok(steer) = self.rx_steer.try_recv() {
|
||||
while let Some(steer) = self.next_turn_steer() {
|
||||
let steer = steer.trim().to_string();
|
||||
if steer.is_empty() {
|
||||
continue;
|
||||
|
||||
+86
-21
@@ -15,6 +15,8 @@ pub(super) struct StdioTransport {
|
||||
pub(super) child: Arc<TokioMutex<Child>>,
|
||||
pub(super) stdin: ChildStdin,
|
||||
pub(super) reader: tokio::io::BufReader<ChildStdout>,
|
||||
/// Partial frame bytes survive cancellation of the receive future.
|
||||
pub(super) pending_line: Vec<u8>,
|
||||
/// Tail of stderr lines from the spawned MCP server. A background task
|
||||
/// drains the child's stderr into this buffer so a mid-run crash leaves
|
||||
/// some context behind instead of `Stdio::null` swallowing it.
|
||||
@@ -222,6 +224,7 @@ impl StdioTransport {
|
||||
child,
|
||||
stdin,
|
||||
reader: tokio::io::BufReader::new(stdout),
|
||||
pending_line: Vec::new(),
|
||||
stderr_tail,
|
||||
authority_cancel_watch,
|
||||
_reviewed_launch: reviewed_launch,
|
||||
@@ -310,22 +313,24 @@ impl McpTransport for StdioTransport {
|
||||
}
|
||||
|
||||
async fn recv(&mut self) -> Result<Vec<u8>> {
|
||||
let mut line_bytes: Vec<u8> = Vec::new();
|
||||
loop {
|
||||
// Bounded read: a server emitting a newline-free multi-GB "line"
|
||||
// must not OOM us (read_line is unbounded).
|
||||
let bytes =
|
||||
match read_line_capped(&mut self.reader, &mut line_bytes, MAX_MCP_RESPONSE_BYTES)
|
||||
.await
|
||||
{
|
||||
Ok(b) => b,
|
||||
Err(err) => {
|
||||
if let Some(stderr) = format_stderr_context(&self.stderr_tail).await {
|
||||
anyhow::bail!("Stdio transport read error: {err}\n{stderr}");
|
||||
}
|
||||
return Err(err.into());
|
||||
let bytes = match read_line_capped(
|
||||
&mut self.reader,
|
||||
&mut self.pending_line,
|
||||
MAX_MCP_RESPONSE_BYTES,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(b) => b,
|
||||
Err(err) => {
|
||||
if let Some(stderr) = format_stderr_context(&self.stderr_tail).await {
|
||||
anyhow::bail!("Stdio transport read error: {err}\n{stderr}");
|
||||
}
|
||||
};
|
||||
return Err(err.into());
|
||||
}
|
||||
};
|
||||
if bytes == 0 {
|
||||
// Let the stderr drain task catch up before snapshotting, and
|
||||
// name the exit status: a reviewed plugin's stderr is never
|
||||
@@ -340,6 +345,7 @@ impl McpTransport for StdioTransport {
|
||||
anyhow::bail!("Stdio transport closed{exit}");
|
||||
}
|
||||
|
||||
let line_bytes = std::mem::take(&mut self.pending_line);
|
||||
let line = String::from_utf8_lossy(&line_bytes);
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
@@ -385,10 +391,10 @@ impl Drop for StdioTransport {
|
||||
}
|
||||
}
|
||||
|
||||
/// Read one newline-terminated line into `out` (cleared first), aborting if it
|
||||
/// exceeds `max` bytes without a newline. Bounds an otherwise-unbounded
|
||||
/// `read_line` so a misbehaving MCP server cannot OOM the client. Returns the
|
||||
/// number of bytes accumulated; 0 means EOF.
|
||||
/// Continue one newline-terminated line in caller-owned `out`, aborting if it
|
||||
/// exceeds `max` bytes. Cancellation retains consumed bytes; the caller clears
|
||||
/// the buffer only after receiving a complete frame. Returns the total bytes
|
||||
/// accumulated; 0 means EOF.
|
||||
async fn read_line_capped<R>(
|
||||
reader: &mut R,
|
||||
out: &mut Vec<u8>,
|
||||
@@ -398,7 +404,6 @@ where
|
||||
R: tokio::io::AsyncBufRead + Unpin,
|
||||
{
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
out.clear();
|
||||
loop {
|
||||
let (chunk, consumed, done) = {
|
||||
let available = reader.fill_buf().await?;
|
||||
@@ -414,15 +419,15 @@ where
|
||||
reader.consume(consumed);
|
||||
}
|
||||
out.extend_from_slice(&chunk);
|
||||
if done {
|
||||
break;
|
||||
}
|
||||
if out.len() > max {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
format!("MCP stdio line exceeded {max} bytes without a newline"),
|
||||
format!("MCP stdio line exceeded {max} bytes"),
|
||||
));
|
||||
}
|
||||
if done {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(out.len())
|
||||
}
|
||||
@@ -431,6 +436,64 @@ where
|
||||
mod read_cap_tests {
|
||||
use super::read_line_capped;
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancelled_partial_read_preserves_next_frame() {
|
||||
use futures_util::FutureExt;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
let (mut writer, reader) = tokio::io::duplex(4096);
|
||||
let mut reader = tokio::io::BufReader::new(reader);
|
||||
let prefix = br#"{"jsonrpc":"2.0","id":"1","result":"#;
|
||||
writer.write_all(prefix).await.unwrap();
|
||||
let mut pending = Vec::new();
|
||||
// Poll through the consumed prefix to Pending, then drop the future.
|
||||
assert!(
|
||||
read_line_capped(&mut reader, &mut pending, 1024)
|
||||
.now_or_never()
|
||||
.is_none()
|
||||
);
|
||||
assert_eq!(pending, prefix);
|
||||
writer.write_all(b"null}\n").await.unwrap();
|
||||
read_line_capped(&mut reader, &mut pending, 1024)
|
||||
.await
|
||||
.unwrap();
|
||||
let first: serde_json::Value =
|
||||
serde_json::from_slice(&std::mem::take(&mut pending)).unwrap();
|
||||
assert_eq!(first["id"], "1");
|
||||
writer
|
||||
.write_all(b"{\"id\":\"2\",\"result\":true}\n")
|
||||
.await
|
||||
.unwrap();
|
||||
read_line_capped(&mut reader, &mut pending, 1024)
|
||||
.await
|
||||
.unwrap();
|
||||
let second: serde_json::Value = serde_json::from_slice(&pending).unwrap();
|
||||
assert_eq!(second["id"], "2");
|
||||
assert_eq!(second["result"], true);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resumed_frame_still_enforces_cap_at_newline() {
|
||||
use futures_util::FutureExt;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
let (mut writer, reader) = tokio::io::duplex(4096);
|
||||
let mut reader = tokio::io::BufReader::new(reader);
|
||||
let mut pending = Vec::new();
|
||||
writer.write_all(b"1234").await.unwrap();
|
||||
assert!(
|
||||
read_line_capped(&mut reader, &mut pending, 6)
|
||||
.now_or_never()
|
||||
.is_none()
|
||||
);
|
||||
writer.write_all(b"567\n").await.unwrap();
|
||||
assert_eq!(
|
||||
read_line_capped(&mut reader, &mut pending, 6)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
std::io::ErrorKind::InvalidData
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_a_line_and_reports_eof() {
|
||||
let data = b"hello\nworld\n".to_vec();
|
||||
@@ -441,11 +504,13 @@ mod read_cap_tests {
|
||||
6
|
||||
);
|
||||
assert_eq!(out, b"hello\n");
|
||||
out.clear();
|
||||
assert_eq!(
|
||||
read_line_capped(&mut reader, &mut out, 1024).await.unwrap(),
|
||||
6
|
||||
);
|
||||
assert_eq!(out, b"world\n");
|
||||
out.clear();
|
||||
// EOF.
|
||||
assert_eq!(
|
||||
read_line_capped(&mut reader, &mut out, 1024).await.unwrap(),
|
||||
|
||||
@@ -2776,6 +2776,80 @@ fn test_connection(transport: Box<dyn McpTransport>) -> McpConnection {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn execute_timeout_after_partial_stdio_response_does_not_corrupt_next_call() -> Result<()> {
|
||||
use serde_json::{Value, json};
|
||||
struct SharedStdio(Arc<tokio::sync::Mutex<StdioTransport>>);
|
||||
#[async_trait::async_trait]
|
||||
impl McpTransport for SharedStdio {
|
||||
async fn send(&mut self, bytes: Vec<u8>) -> Result<()> {
|
||||
self.0.lock().await.send(bytes).await
|
||||
}
|
||||
async fn recv(&mut self) -> Result<Vec<u8>> {
|
||||
self.0.lock().await.recv().await
|
||||
}
|
||||
}
|
||||
let dir = tempfile::tempdir()?;
|
||||
let requests = dir.path().join("requests.jsonl");
|
||||
let ready = dir.path().join("ready");
|
||||
let script = r#"
|
||||
printf '%s' '{"jsonrpc":"2.0","id":"1","result":'
|
||||
: > "$2"
|
||||
IFS= read -r first
|
||||
printf '%s\n' "$first" >> "$1"
|
||||
IFS= read -r second
|
||||
printf '%s\n' "$second" >> "$1"
|
||||
printf '%s\n' 'null}' '{"jsonrpc":"2.0","id":"2","result":{"ok":true}}'
|
||||
IFS= read -r keep_open
|
||||
"#;
|
||||
let mut config = test_server_config();
|
||||
config.args = vec![
|
||||
"-c".into(),
|
||||
script.into(),
|
||||
"cw-partial-frame-fixture".into(),
|
||||
requests.display().to_string(),
|
||||
ready.display().to_string(),
|
||||
];
|
||||
let transport = Arc::new(tokio::sync::Mutex::new(StdioTransport::spawn(
|
||||
"partial-frame",
|
||||
"sh",
|
||||
&config,
|
||||
tokio_util::sync::CancellationToken::new(),
|
||||
)?));
|
||||
tokio::time::timeout(Duration::from_secs(10), async {
|
||||
while !ready.exists() {
|
||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
||||
}
|
||||
})
|
||||
.await?;
|
||||
let mut connection = test_connection(Box::new(SharedStdio(Arc::clone(&transport))));
|
||||
let error = connection
|
||||
.call_tool("first", json!({}), 1)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(error.to_string().contains("timed out"), "{error:#}");
|
||||
assert_eq!(
|
||||
transport.lock().await.pending_line,
|
||||
br#"{"jsonrpc":"2.0","id":"1","result":"#
|
||||
);
|
||||
assert!(connection.is_ready());
|
||||
assert_eq!(
|
||||
connection.call_tool("second", json!({}), 5).await?,
|
||||
json!({"ok": true})
|
||||
);
|
||||
let sent = fs::read_to_string(requests)?;
|
||||
let sent = sent
|
||||
.lines()
|
||||
.map(serde_json::from_str::<Value>)
|
||||
.collect::<std::result::Result<Vec<_>, _>>()?;
|
||||
assert_eq!(sent.len(), 2, "neither request may be replayed");
|
||||
assert_eq!(sent[0]["id"], "1");
|
||||
assert_eq!(sent[1]["id"], "2");
|
||||
transport.lock().await.shutdown().await;
|
||||
Ok::<_, anyhow::Error>(())
|
||||
}
|
||||
|
||||
fn json_frame(value: serde_json::Value) -> Vec<u8> {
|
||||
serde_json::to_vec(&value).unwrap()
|
||||
}
|
||||
@@ -4688,6 +4762,7 @@ async fn stdio_transport_shutdown_terminates_child() {
|
||||
child: Arc::new(tokio::sync::Mutex::new(child)),
|
||||
stdin,
|
||||
reader: tokio::io::BufReader::new(stdout),
|
||||
pending_line: Vec::new(),
|
||||
stderr_tail: StderrTail::new(),
|
||||
authority_cancel_watch: None,
|
||||
_reviewed_launch: None,
|
||||
@@ -4816,6 +4891,7 @@ async fn stdio_transport_recv_error_includes_stderr_tail() {
|
||||
child: Arc::new(tokio::sync::Mutex::new(child)),
|
||||
stdin,
|
||||
reader: tokio::io::BufReader::new(stdout),
|
||||
pending_line: Vec::new(),
|
||||
stderr_tail,
|
||||
authority_cancel_watch: None,
|
||||
_reviewed_launch: None,
|
||||
|
||||
@@ -290,6 +290,7 @@ pub(super) async fn resume_session_thread(
|
||||
Path(id): Path<String>,
|
||||
Json(req): Json<ResumeSessionRequest>,
|
||||
) -> Result<(StatusCode, Json<ResumeSessionResponse>), ApiError> {
|
||||
let _checkpoint_admission = state.runtime_threads.session_checkpoint_guard().await;
|
||||
let manager = SessionManager::new(state.sessions_dir.clone())
|
||||
.map_err(|e| ApiError::internal(format!("Failed to open sessions dir: {e}")))?;
|
||||
let session = manager
|
||||
@@ -333,19 +334,15 @@ pub(super) async fn resume_session_thread(
|
||||
|
||||
// Link the session to the new thread so that `ensure_engine_loaded`
|
||||
// can restore the full message history from the session file.
|
||||
if let Err(e) = state
|
||||
state
|
||||
.runtime_threads
|
||||
.set_thread_session_id(&thread.id, &id)
|
||||
.set_thread_session_checkpoint(&thread.id, &session)
|
||||
.await
|
||||
{
|
||||
let session_ref = crate::utils::redacted_identifier_for_log(&id);
|
||||
tracing::warn!(
|
||||
session = %session_ref,
|
||||
thread_id = %thread.id,
|
||||
error = %e,
|
||||
"Failed to link session to thread"
|
||||
);
|
||||
}
|
||||
.map_err(|e| {
|
||||
ApiError::internal(format!(
|
||||
"Saved session was read but its Runtime checkpoint could not be bound: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
let summary = format!(
|
||||
"Resumed session '{}' ({} messages) into thread {}",
|
||||
@@ -367,6 +364,7 @@ pub(super) async fn create_session_from_thread(
|
||||
State(state): State<RuntimeApiState>,
|
||||
Json(req): Json<CreateSessionRequest>,
|
||||
) -> Result<(StatusCode, Json<CreateSessionResponse>), ApiError> {
|
||||
let _checkpoint_admission = state.runtime_threads.session_checkpoint_guard().await;
|
||||
let thread_id = req.thread_id.trim();
|
||||
if thread_id.is_empty() {
|
||||
return Err(ApiError::bad_request("thread_id is required"));
|
||||
@@ -435,19 +433,15 @@ pub(super) async fn create_session_from_thread(
|
||||
|
||||
// Link the session to the thread so that `ensure_engine_loaded` can
|
||||
// restore the full message history from the session file.
|
||||
if let Err(e) = state
|
||||
state
|
||||
.runtime_threads
|
||||
.set_thread_session_id(&detail.thread.id, &session_handle)
|
||||
.set_thread_session_checkpoint(&detail.thread.id, &session)
|
||||
.await
|
||||
{
|
||||
let session_ref = crate::utils::redacted_identifier_for_log(&session_handle);
|
||||
tracing::warn!(
|
||||
session = %session_ref,
|
||||
thread_id = %detail.thread.id,
|
||||
error = %e,
|
||||
"Failed to link session to thread"
|
||||
);
|
||||
}
|
||||
.map_err(|e| {
|
||||
ApiError::internal(format!(
|
||||
"Session was saved but its Runtime checkpoint could not be bound: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok((
|
||||
StatusCode::CREATED,
|
||||
@@ -719,6 +713,7 @@ pub(super) async fn save_current_session(
|
||||
State(state): State<RuntimeApiState>,
|
||||
Json(req): Json<SaveSessionRequest>,
|
||||
) -> Result<Json<SaveSessionResponse>, ApiError> {
|
||||
let _checkpoint_admission = state.runtime_threads.session_checkpoint_guard().await;
|
||||
// Find the thread to save.
|
||||
let thread_id = match req.thread_id {
|
||||
Some(id) => id,
|
||||
@@ -826,19 +821,15 @@ pub(super) async fn save_current_session(
|
||||
// restore the full message history (including thinking/tool blocks)
|
||||
// from the session file instead of reconstructing from turns.
|
||||
let session_handle = session.metadata.id.clone();
|
||||
if let Err(e) = state
|
||||
state
|
||||
.runtime_threads
|
||||
.set_thread_session_id(&thread_id, &session_handle)
|
||||
.set_thread_session_checkpoint(&thread_id, &session)
|
||||
.await
|
||||
{
|
||||
let session_ref = crate::utils::redacted_identifier_for_log(&session_handle);
|
||||
tracing::warn!(
|
||||
session = %session_ref,
|
||||
thread_id = %thread_id,
|
||||
error = %e,
|
||||
"Failed to link session to thread"
|
||||
);
|
||||
}
|
||||
.map_err(|e| {
|
||||
ApiError::internal(format!(
|
||||
"Session was saved but its Runtime checkpoint could not be bound: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(Json(SaveSessionResponse {
|
||||
session_id: session_handle,
|
||||
|
||||
@@ -432,6 +432,7 @@ fn messages_from_thread_detail_batches_tool_results() {
|
||||
task_id: None,
|
||||
title: None,
|
||||
session_id: None,
|
||||
saved_session_checkpoint: None,
|
||||
};
|
||||
let turn = TurnRecord {
|
||||
schema_version: 2,
|
||||
@@ -622,6 +623,7 @@ fn legacy_exact_thread_export_normalizes_provider_kind_and_id() {
|
||||
task_id: None,
|
||||
title: None,
|
||||
session_id: None,
|
||||
saved_session_checkpoint: None,
|
||||
},
|
||||
turns: Vec::new(),
|
||||
items: Vec::new(),
|
||||
|
||||
@@ -686,11 +686,88 @@ pub struct ThreadRecord {
|
||||
/// additive metadata — older readers ignore it without misinterpretation.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub title: Option<String>,
|
||||
/// The session ID associated with this thread. When set, `ensure_engine_loaded`
|
||||
/// loads the full message history (including thinking/tool blocks) from the
|
||||
/// session file instead of reconstructing from turns (which loses process info).
|
||||
/// Full-fidelity saved history prefix. Engine recovery verifies its checkpoint
|
||||
/// and appends later durable Runtime turns without duplicating saved messages.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub session_id: Option<String>,
|
||||
/// Exact saved transcript and the last Runtime turn it already contains.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub saved_session_checkpoint: Option<SavedSessionCheckpoint>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct SavedSessionCheckpoint {
|
||||
pub covered_turn_id: Option<String>,
|
||||
pub messages_sha256: String,
|
||||
/// A backtracked fork may retain only a prefix of the verified snapshot.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub retained_messages: Option<usize>,
|
||||
}
|
||||
|
||||
fn session_messages_sha256(messages: &[Message]) -> Result<String> {
|
||||
Ok(Sha256::digest(serde_json::to_vec(messages)?)
|
||||
.iter()
|
||||
.map(|byte| format!("{byte:02x}"))
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Compare only fields represented by legacy seeding. A match identifies a
|
||||
/// prefix boundary; the saved messages themselves retain every raw block.
|
||||
fn session_recovery_projection(messages: &[Message]) -> Vec<Value> {
|
||||
let mut projection = Vec::new();
|
||||
for message in messages {
|
||||
let role = message.role.as_str();
|
||||
if role == "user" {
|
||||
let text = message
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|block| match block {
|
||||
ContentBlock::Text { text, .. } if !text.trim().is_empty() => {
|
||||
Some(text.as_str())
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
if !text.is_empty() {
|
||||
projection.push(json!(["user", text]));
|
||||
}
|
||||
}
|
||||
for block in &message.content {
|
||||
match block {
|
||||
ContentBlock::Text { text, .. }
|
||||
if role == "assistant" && !text.trim().is_empty() =>
|
||||
{
|
||||
projection.push(json!(["assistant", text]))
|
||||
}
|
||||
ContentBlock::Thinking { thinking, .. }
|
||||
if role == "assistant" && !thinking.trim().is_empty() =>
|
||||
{
|
||||
projection.push(json!(["thinking", thinking]))
|
||||
}
|
||||
ContentBlock::ToolUse {
|
||||
id, name, input, ..
|
||||
}
|
||||
| ContentBlock::ServerToolUse {
|
||||
id, name, input, ..
|
||||
} if role == "assistant" => projection.push(json!(["tool_use", id, name, input])),
|
||||
ContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
is_error,
|
||||
content_blocks,
|
||||
} if role == "user" => projection.push(json!([
|
||||
"tool_result",
|
||||
tool_use_id,
|
||||
content,
|
||||
is_error.unwrap_or(false),
|
||||
content_blocks
|
||||
])),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
projection
|
||||
}
|
||||
|
||||
fn thread_execution_state_matches(left: &ThreadRecord, right: &ThreadRecord) -> bool {
|
||||
@@ -713,6 +790,7 @@ fn thread_execution_state_matches(left: &ThreadRecord, right: &ThreadRecord) ->
|
||||
&& left.system_prompt == right.system_prompt
|
||||
&& left.task_id == right.task_id
|
||||
&& left.session_id == right.session_id
|
||||
&& left.saved_session_checkpoint == right.saved_session_checkpoint
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -3411,9 +3489,9 @@ pub struct RuntimeThreadManager {
|
||||
store: RuntimeThreadStore,
|
||||
_process_owner_lock: Arc<RuntimeProcessOwnerLock>,
|
||||
/// Concurrent turn admissions share a read lease; config reload owns the
|
||||
/// write lease from validation through publication. A turn therefore
|
||||
/// cannot snapshot old credentials/routes and dispatch after reload has
|
||||
/// returned.
|
||||
/// write lease from validation through publication. Saved-session binding
|
||||
/// also owns the write lease across snapshot and checkpoint persistence.
|
||||
/// This orders both route changes and saved-history boundaries with dispatch.
|
||||
config_admission: Arc<AsyncRwLock<()>>,
|
||||
engine_load: Arc<Mutex<()>>,
|
||||
active: Arc<Mutex<ActiveThreads>>,
|
||||
@@ -5906,6 +5984,7 @@ impl RuntimeThreadManager {
|
||||
task_id: req.task_id,
|
||||
title: None,
|
||||
session_id: None,
|
||||
saved_session_checkpoint: None,
|
||||
};
|
||||
self.store.save_thread(&thread)?;
|
||||
if let Err(error) = self
|
||||
@@ -6367,20 +6446,30 @@ impl RuntimeThreadManager {
|
||||
Ok(thread)
|
||||
}
|
||||
|
||||
/// Link a session to a thread so that `ensure_engine_loaded` can restore
|
||||
/// the full message history (including thinking/tool blocks) from the
|
||||
/// session file instead of reconstructing from turns.
|
||||
pub async fn set_thread_session_id(&self, thread_id: &str, session_id: &str) -> Result<()> {
|
||||
/// Save/resume holds this existing admission lease from snapshot through
|
||||
/// binding, so a newer turn cannot be mistaken for part of the snapshot.
|
||||
pub(crate) async fn session_checkpoint_guard(&self) -> tokio::sync::OwnedRwLockWriteGuard<()> {
|
||||
Arc::clone(&self.config_admission).write_owned().await
|
||||
}
|
||||
|
||||
/// Bind full-fidelity saved history to the exact prefix it covers. Callers
|
||||
/// hold session_checkpoint_guard across the snapshot, save and this write.
|
||||
pub(crate) async fn set_thread_session_checkpoint(
|
||||
&self,
|
||||
thread_id: &str,
|
||||
session: &crate::session_manager::SavedSession,
|
||||
) -> Result<()> {
|
||||
let session_id = &session.metadata.id;
|
||||
let messages_sha256 = session_messages_sha256(&session.messages)?;
|
||||
let thread = {
|
||||
let _thread_mutation = self.store.thread_mutation.lock();
|
||||
let mut thread = self
|
||||
.store
|
||||
.load_thread(thread_id)
|
||||
.with_context(|| format!("Thread not found: {thread_id}"))?;
|
||||
if thread.session_id.as_deref() == Some(session_id) {
|
||||
return Ok(());
|
||||
}
|
||||
thread.session_id = Some(session_id.to_string());
|
||||
let mut thread = self.store.load_thread(thread_id)?;
|
||||
thread.session_id = Some(session_id.clone());
|
||||
thread.saved_session_checkpoint = Some(SavedSessionCheckpoint {
|
||||
covered_turn_id: thread.latest_turn_id.clone(),
|
||||
messages_sha256,
|
||||
retained_messages: None,
|
||||
});
|
||||
thread.updated_at = Utc::now();
|
||||
self.store.save_thread(&thread)?;
|
||||
thread
|
||||
@@ -6480,6 +6569,11 @@ impl RuntimeThreadManager {
|
||||
let mut cloned_turn = source_turn.clone();
|
||||
cloned_turn.id = format!("turn_{}", &Uuid::new_v4().to_string()[..8]);
|
||||
cloned_turn.thread_id = forked.id.clone();
|
||||
if let Some(checkpoint) = forked.saved_session_checkpoint.as_mut()
|
||||
&& checkpoint.covered_turn_id.as_deref() == Some(source_turn.id.as_str())
|
||||
{
|
||||
checkpoint.covered_turn_id = Some(cloned_turn.id.clone());
|
||||
}
|
||||
cloned_turn.item_ids.clear();
|
||||
|
||||
let items = self.store.list_items_for_turn(&source_turn.id)?;
|
||||
@@ -6592,11 +6686,40 @@ impl RuntimeThreadManager {
|
||||
forked.latest_turn_id = None;
|
||||
forked.archived = false;
|
||||
|
||||
if let Some((messages, covered)) = self.saved_session_prefix(&source, &source_turns)? {
|
||||
let kept_turns = covered.min(target_turn_idx);
|
||||
let retained_messages = if covered <= target_turn_idx {
|
||||
messages.len()
|
||||
} else {
|
||||
let expected = session_recovery_projection(
|
||||
&self.reconstruct_messages_from_turns(&source_turns[..target_turn_idx])?,
|
||||
);
|
||||
(0..=messages.len())
|
||||
.find(|count| session_recovery_projection(&messages[..*count]) == expected)
|
||||
.context("Cannot identify an exact saved-history boundary for this backtrack; the source thread was preserved")?
|
||||
};
|
||||
forked.saved_session_checkpoint = Some(SavedSessionCheckpoint {
|
||||
covered_turn_id: kept_turns
|
||||
.checked_sub(1)
|
||||
.map(|index| source_turns[index].id.clone()),
|
||||
messages_sha256: match &source.saved_session_checkpoint {
|
||||
Some(checkpoint) => checkpoint.messages_sha256.clone(),
|
||||
None => session_messages_sha256(&messages)?,
|
||||
},
|
||||
retained_messages: Some(retained_messages),
|
||||
});
|
||||
}
|
||||
|
||||
let mut cloned_records = Vec::with_capacity(target_turn_idx);
|
||||
for source_turn in source_turns.iter().take(target_turn_idx) {
|
||||
let mut cloned_turn = source_turn.clone();
|
||||
cloned_turn.id = format!("turn_{}", &Uuid::new_v4().to_string()[..8]);
|
||||
cloned_turn.thread_id = forked.id.clone();
|
||||
if let Some(checkpoint) = forked.saved_session_checkpoint.as_mut()
|
||||
&& checkpoint.covered_turn_id.as_deref() == Some(source_turn.id.as_str())
|
||||
{
|
||||
checkpoint.covered_turn_id = Some(cloned_turn.id.clone());
|
||||
}
|
||||
cloned_turn.item_ids.clear();
|
||||
|
||||
let items = self.store.list_items_for_turn(&source_turn.id)?;
|
||||
@@ -8002,7 +8125,7 @@ impl RuntimeThreadManager {
|
||||
policy.permission,
|
||||
configured_sandbox_mode,
|
||||
);
|
||||
let _sender = permit.send(op);
|
||||
engine.send_reserved_op(permit, op);
|
||||
touch_lru(&mut active.lru, thread_id);
|
||||
self.spawn_claimed_turn_monitor(
|
||||
turn.clone(),
|
||||
@@ -8320,7 +8443,7 @@ impl RuntimeThreadManager {
|
||||
policy.permission,
|
||||
configured_sandbox_mode,
|
||||
);
|
||||
let _sender = permit.send(op);
|
||||
engine.send_reserved_op(permit, op);
|
||||
touch_lru(&mut active.lru, thread_id);
|
||||
self.spawn_claimed_turn_monitor(
|
||||
turn.clone(),
|
||||
@@ -8665,52 +8788,14 @@ impl RuntimeThreadManager {
|
||||
.unwrap_or_else(crate::tools::subagent::AdvisorConfig::disabled),
|
||||
};
|
||||
|
||||
// Verify the persisted history before spawning an Engine task.
|
||||
let session_messages = self.restore_thread_messages(&thread)?;
|
||||
let engine = spawn_engine_with_authoritative_route_config(
|
||||
engine_cfg,
|
||||
&cfg,
|
||||
Arc::clone(&self.config),
|
||||
);
|
||||
|
||||
// When the thread has an associated session, load the full message history
|
||||
// (including thinking/tool blocks) from the session file. This preserves
|
||||
// process information that `reconstruct_messages_from_turns` would lose.
|
||||
let session_messages = if let Some(ref sid) = thread.session_id {
|
||||
match crate::session_manager::default_sessions_dir() {
|
||||
Ok(sessions_dir) => {
|
||||
match crate::session_manager::SessionManager::new(sessions_dir) {
|
||||
Ok(manager) => match manager.load_session(sid) {
|
||||
Ok(session) => session.messages,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Failed to load session {} for thread {}: {e}; falling back to turn reconstruction",
|
||||
sid,
|
||||
thread.id
|
||||
);
|
||||
let turns = self.store.list_turns_for_thread(&thread.id)?;
|
||||
self.reconstruct_messages_from_turns(&turns)?
|
||||
}
|
||||
},
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Failed to open sessions dir: {e}; falling back to turn reconstruction"
|
||||
);
|
||||
let turns = self.store.list_turns_for_thread(&thread.id)?;
|
||||
self.reconstruct_messages_from_turns(&turns)?
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Failed to resolve sessions dir: {e}; falling back to turn reconstruction"
|
||||
);
|
||||
let turns = self.store.list_turns_for_thread(&thread.id)?;
|
||||
self.reconstruct_messages_from_turns(&turns)?
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let turns = self.store.list_turns_for_thread(&thread.id)?;
|
||||
self.reconstruct_messages_from_turns(&turns)?
|
||||
};
|
||||
let sys_prompt = thread
|
||||
.system_prompt
|
||||
.as_ref()
|
||||
@@ -8790,6 +8875,77 @@ impl RuntimeThreadManager {
|
||||
self.ensure_engine_loaded(&thread).await
|
||||
}
|
||||
|
||||
fn restore_thread_messages(&self, thread: &ThreadRecord) -> Result<Vec<Message>> {
|
||||
let turns = self.store.list_turns_for_thread(&thread.id)?;
|
||||
let (mut messages, covered) = self
|
||||
.saved_session_prefix(thread, &turns)?
|
||||
.unwrap_or_default();
|
||||
messages.extend(self.reconstruct_messages_from_turns(&turns[covered..])?);
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
fn saved_session_prefix(
|
||||
&self,
|
||||
thread: &ThreadRecord,
|
||||
turns: &[TurnRecord],
|
||||
) -> Result<Option<(Vec<Message>, usize)>> {
|
||||
let Some(session_id) = thread.session_id.as_deref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let session = crate::session_manager::default_sessions_dir()
|
||||
.and_then(crate::session_manager::SessionManager::new)
|
||||
.and_then(|manager| manager.load_session(session_id))
|
||||
.with_context(|| format!("Cannot read saved session {session_id}; restore that session file before resuming thread {}", thread.id))?;
|
||||
let covered = if let Some(checkpoint) = &thread.saved_session_checkpoint {
|
||||
if checkpoint.messages_sha256 != session_messages_sha256(&session.messages)? {
|
||||
bail!(
|
||||
"Saved session {session_id} changed after this thread's checkpoint; re-import it into a separate thread to preserve both histories"
|
||||
);
|
||||
}
|
||||
match checkpoint.covered_turn_id.as_deref() {
|
||||
Some(id) => turns.iter().position(|turn| turn.id == id)
|
||||
.map(|index| index + 1)
|
||||
.with_context(|| format!("Saved session checkpoint turn {id} is missing; restore the Runtime turn records before resuming"))?,
|
||||
None => 0,
|
||||
}
|
||||
} else {
|
||||
// Legacy links have no cursor. Establish it only from an exact
|
||||
// match to the existing seeder's projection; never use mtime or
|
||||
// silently let a saved file replace a newer Runtime transcript.
|
||||
let expected = session_recovery_projection(&session.messages);
|
||||
let mut prefix = Vec::new();
|
||||
let mut covered = (expected.is_empty()).then_some(0);
|
||||
for (index, turn) in turns.iter().enumerate() {
|
||||
if covered.is_some() {
|
||||
break;
|
||||
}
|
||||
prefix.extend(session_recovery_projection(
|
||||
&self.reconstruct_messages_from_turns(std::slice::from_ref(turn))?,
|
||||
));
|
||||
if prefix == expected {
|
||||
covered = Some(index + 1);
|
||||
} else if !expected.starts_with(&prefix) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
covered.with_context(|| format!("Saved session {session_id} has no verifiable Runtime checkpoint; keep both histories and re-import the saved session into a separate thread"))?
|
||||
};
|
||||
let mut messages = session.messages;
|
||||
if let Some(retained) = thread
|
||||
.saved_session_checkpoint
|
||||
.as_ref()
|
||||
.and_then(|checkpoint| checkpoint.retained_messages)
|
||||
{
|
||||
if retained > messages.len() {
|
||||
bail!(
|
||||
"Saved session checkpoint exceeds its transcript; restore the saved session before resuming"
|
||||
);
|
||||
}
|
||||
messages.truncate(retained);
|
||||
}
|
||||
Ok(Some((messages, covered)))
|
||||
}
|
||||
|
||||
fn reconstruct_messages_from_turns(&self, turns: &[TurnRecord]) -> Result<Vec<Message>> {
|
||||
let mut messages = Vec::new();
|
||||
for turn in turns {
|
||||
|
||||
@@ -7,6 +7,406 @@ use tokio::sync::{mpsc, oneshot};
|
||||
use tokio::time::sleep;
|
||||
use uuid::Uuid;
|
||||
|
||||
mod recovery {
|
||||
use super::*;
|
||||
use crate::core::engine::Engine;
|
||||
use crate::llm_client::mock::{MockLlmClient, canned};
|
||||
|
||||
fn config() -> Config {
|
||||
let mut config = Config {
|
||||
provider: Some("recovery-fixture".into()),
|
||||
providers: Some(crate::config::ProvidersConfig {
|
||||
custom: HashMap::from([(
|
||||
"recovery-fixture".into(),
|
||||
crate::config::ProviderConfig {
|
||||
kind: Some("openai-compatible".into()),
|
||||
base_url: Some("http://127.0.0.1:18181/v1".into()),
|
||||
model: Some("fixture-model".into()),
|
||||
api_key: Some("local-test-key".into()),
|
||||
..Default::default()
|
||||
},
|
||||
)]),
|
||||
..Default::default()
|
||||
}),
|
||||
..Config::default()
|
||||
};
|
||||
config.set_feature("mcp", false).unwrap();
|
||||
config.set_feature("subagents", false).unwrap();
|
||||
config
|
||||
}
|
||||
|
||||
async fn close_engines(manager: &RuntimeThreadManager) -> Result<()> {
|
||||
let handles = manager
|
||||
.active
|
||||
.lock()
|
||||
.await
|
||||
.engines
|
||||
.drain()
|
||||
.map(|(_, state)| state.engine)
|
||||
.collect::<Vec<_>>();
|
||||
for handle in handles {
|
||||
handle.send(Op::Shutdown).await?;
|
||||
tokio::time::timeout(Duration::from_secs(10), async {
|
||||
while handle.rx_event.write().await.recv().await.is_some() {}
|
||||
})
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn history_case(restart: bool) -> Result<()> {
|
||||
let _env = crate::test_support::lock_test_env();
|
||||
let dir = tempfile::tempdir()?;
|
||||
let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", dir.path());
|
||||
let workspace = dir.path().join("workspace");
|
||||
fs::create_dir(&workspace)?;
|
||||
let runtime_dir = dir.path().join("runtime");
|
||||
let manager_cfg = RuntimeThreadManagerConfig {
|
||||
max_active_threads: 1,
|
||||
..test_manager_config(runtime_dir)
|
||||
};
|
||||
let mut manager =
|
||||
RuntimeThreadManager::open(config(), workspace.clone(), manager_cfg.clone())?;
|
||||
let thread = manager
|
||||
.create_thread(CreateThreadRequest::default())
|
||||
.await?;
|
||||
let old: Vec<Message> = serde_json::from_value(json!([
|
||||
{"role":"user","content":[{"type":"text","text":"OLD"},{"type":"image_url","image_url":{"url":"data:image/png;base64,AA=="}}]},
|
||||
{"role":"assistant","content":[{"type":"thinking","thinking":"prior reasoning","signature":"fixture-signed-thinking"},{"type":"tool_use","id":"fixture-call","name":"read_file","input":{"path":"fixture.txt"},"thought_signature":"fixture-tool-signature"}]},
|
||||
{"role":"user","content":[{"type":"tool_result","tool_use_id":"fixture-call","content":"old result"}]},
|
||||
{"role":"assistant","content":[{"type":"text","text":"OLD ANSWER"}]}
|
||||
]))?;
|
||||
manager.seed_thread_from_messages(&thread.id, &old).await?;
|
||||
let mut saved = crate::session_manager::create_saved_session_with_id_and_mode(
|
||||
Uuid::new_v4().to_string(),
|
||||
&old,
|
||||
&thread.model,
|
||||
&workspace,
|
||||
0,
|
||||
None,
|
||||
Some("agent"),
|
||||
);
|
||||
saved.metadata.set_model_provider_route(
|
||||
thread.model_provider.as_deref().unwrap(),
|
||||
thread.model_provider_id.as_deref(),
|
||||
);
|
||||
let sessions = crate::session_manager::SessionManager::new(
|
||||
crate::session_manager::default_sessions_dir()?,
|
||||
)?;
|
||||
{
|
||||
let _admission = manager.session_checkpoint_guard().await;
|
||||
sessions.save_session(&saved)?;
|
||||
manager
|
||||
.set_thread_session_checkpoint(&thread.id, &saved)
|
||||
.await?;
|
||||
}
|
||||
let original_engine = manager.get_engine(&thread.id).await?;
|
||||
assert_eq!(original_engine.get_session_snapshot().await?.messages, old);
|
||||
// Simulate a later durable turn using the real store writer. The saved
|
||||
// full-fidelity snapshot deliberately remains at the OLD boundary.
|
||||
let tail: Vec<Message> = serde_json::from_value(json!([
|
||||
{"role":"user","content":[{"type":"text","text":"NEW"}]},
|
||||
{"role":"assistant","content":[{"type":"text","text":"NEW ANSWER"}]}
|
||||
]))?;
|
||||
manager.seed_thread_from_messages(&thread.id, &tail).await?;
|
||||
assert_eq!(sessions.load_session(&saved.metadata.id)?.messages, old);
|
||||
if restart {
|
||||
close_engines(&manager).await?;
|
||||
drop(original_engine);
|
||||
drop(manager);
|
||||
manager = RuntimeThreadManager::open(config(), workspace, manager_cfg)?;
|
||||
} else {
|
||||
let other = manager
|
||||
.create_thread(CreateThreadRequest::default())
|
||||
.await?;
|
||||
manager
|
||||
.get_engine(&other.id)
|
||||
.await?
|
||||
.get_session_snapshot()
|
||||
.await?;
|
||||
assert!(!manager.active.lock().await.engines.contains_key(&thread.id));
|
||||
}
|
||||
let restored = manager
|
||||
.get_engine(&thread.id)
|
||||
.await?
|
||||
.get_session_snapshot()
|
||||
.await?
|
||||
.messages;
|
||||
let mut expected = old;
|
||||
expected.extend(tail);
|
||||
assert_eq!(
|
||||
restored, expected,
|
||||
"snapshot plus durable tail must be exact, including raw signatures/media"
|
||||
);
|
||||
let detail = manager.get_thread_detail(&thread.id).await?;
|
||||
assert!(
|
||||
detail
|
||||
.items
|
||||
.iter()
|
||||
.any(|item| item.detail.as_deref() == Some("NEW"))
|
||||
);
|
||||
close_engines(&manager).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn linked_saved_session_replays_runtime_tail_after_eviction() -> Result<()> {
|
||||
history_case(false).await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn linked_saved_session_replays_runtime_tail_after_restart() -> Result<()> {
|
||||
history_case(true).await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn saved_checkpoint_validation_and_forks_preserve_history() -> Result<()> {
|
||||
let _env = crate::test_support::lock_test_env();
|
||||
let dir = tempfile::tempdir()?;
|
||||
let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", dir.path());
|
||||
let manager = RuntimeThreadManager::open(
|
||||
config(),
|
||||
dir.path().to_path_buf(),
|
||||
test_manager_config(dir.path().join("runtime")),
|
||||
)?;
|
||||
let thread = manager
|
||||
.create_thread(CreateThreadRequest::default())
|
||||
.await?;
|
||||
let messages: Vec<Message> = serde_json::from_value(json!([
|
||||
{"role":"user","content":[{"type":"text","text":"keep"},{"type":"image_url","image_url":{"url":"data:image/png;base64,AA=="}}]},
|
||||
{"role":"assistant","content":[{"type":"thinking","thinking":"kept reasoning","signature":"kept signature"},{"type":"text","text":"kept answer"}]},
|
||||
{"role":"user","content":[{"type":"text","text":"drop"}]},
|
||||
{"role":"assistant","content":[{"type":"text","text":"dropped answer"}]}
|
||||
]))?;
|
||||
manager
|
||||
.seed_thread_from_messages(&thread.id, &messages)
|
||||
.await?;
|
||||
let saved = crate::session_manager::create_saved_session_with_id_and_mode(
|
||||
Uuid::new_v4().to_string(),
|
||||
&messages,
|
||||
&thread.model,
|
||||
dir.path(),
|
||||
0,
|
||||
None,
|
||||
Some("agent"),
|
||||
);
|
||||
let sessions = crate::session_manager::SessionManager::new(
|
||||
crate::session_manager::default_sessions_dir()?,
|
||||
)?;
|
||||
let _admission = manager.session_checkpoint_guard().await;
|
||||
sessions.save_session(&saved)?;
|
||||
manager
|
||||
.set_thread_session_checkpoint(&thread.id, &saved)
|
||||
.await?;
|
||||
let thread = manager.get_thread(&thread.id).await?;
|
||||
let fork = manager.fork_thread(&thread.id).await?;
|
||||
assert_eq!(manager.restore_thread_messages(&fork)?, messages);
|
||||
let (backtrack, _) = manager.fork_at_user_message(&thread.id, 0).await?;
|
||||
assert_eq!(manager.restore_thread_messages(&backtrack)?, messages[..2]);
|
||||
let (empty, _) = manager.fork_at_user_message(&thread.id, 1).await?;
|
||||
assert!(manager.restore_thread_messages(&empty)?.is_empty());
|
||||
assert_eq!(manager.restore_thread_messages(&thread)?, messages);
|
||||
let mut legacy = thread.clone();
|
||||
legacy.saved_session_checkpoint = None;
|
||||
assert_eq!(manager.restore_thread_messages(&legacy)?, messages);
|
||||
let tail = vec![Message {
|
||||
role: Role::User,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "new tail".into(),
|
||||
cache_control: None,
|
||||
}],
|
||||
}];
|
||||
manager.seed_thread_from_messages(&thread.id, &tail).await?;
|
||||
let mut expected = messages.clone();
|
||||
expected.extend(tail);
|
||||
assert_eq!(manager.restore_thread_messages(&legacy)?, expected);
|
||||
let mut missing = thread.clone();
|
||||
missing
|
||||
.saved_session_checkpoint
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.covered_turn_id = Some("turn_missing".into());
|
||||
assert!(
|
||||
manager
|
||||
.restore_thread_messages(&missing)
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("checkpoint turn turn_missing is missing")
|
||||
);
|
||||
let mut changed = saved.clone();
|
||||
changed.messages[0] = Message {
|
||||
role: Role::User,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "conflicting history".into(),
|
||||
cache_control: None,
|
||||
}],
|
||||
};
|
||||
sessions.save_session(&changed)?;
|
||||
assert!(
|
||||
manager
|
||||
.restore_thread_messages(&thread)
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("changed after this thread's checkpoint")
|
||||
);
|
||||
assert!(
|
||||
manager
|
||||
.restore_thread_messages(&legacy)
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("no verifiable Runtime checkpoint")
|
||||
);
|
||||
assert!(manager.get_engine(&thread.id).await.is_err());
|
||||
assert!(manager.active.lock().await.engines.is_empty());
|
||||
sessions.save_session(&saved)?;
|
||||
assert_eq!(manager.restore_thread_messages(&thread)?, expected);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn control_case(interrupt: bool, follow_up: bool) -> Result<()> {
|
||||
let _env = crate::test_support::lock_test_env();
|
||||
let dir = tempfile::tempdir()?;
|
||||
let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", dir.path());
|
||||
let workspace = dir.path().join("workspace");
|
||||
fs::create_dir(&workspace)?;
|
||||
let config = config();
|
||||
let manager = RuntimeThreadManager::open(
|
||||
config.clone(),
|
||||
workspace.clone(),
|
||||
test_manager_config(dir.path().join("runtime")),
|
||||
)?;
|
||||
let thread = manager
|
||||
.create_thread(CreateThreadRequest::default())
|
||||
.await?;
|
||||
let mock = Arc::new(MockLlmClient::new(vec![canned::simple_text_turn(
|
||||
"fixture response",
|
||||
)]));
|
||||
let (engine, handle) = Engine::new_with_model_client(
|
||||
EngineConfig {
|
||||
workspace,
|
||||
model: thread.model.clone(),
|
||||
features: crate::features::Features::default(),
|
||||
subagents_enabled: false,
|
||||
snapshots_enabled: false,
|
||||
memory_enabled: false,
|
||||
terminal_chrome_enabled: false,
|
||||
runtime_services: crate::tools::spec::RuntimeToolServices {
|
||||
active_thread_id: Some(thread.id.clone()),
|
||||
..Default::default()
|
||||
},
|
||||
..EngineConfig::default()
|
||||
},
|
||||
&config,
|
||||
mock.clone(),
|
||||
);
|
||||
manager
|
||||
.install_test_engine(&thread.id, handle.clone())
|
||||
.await?;
|
||||
let (release, begin) = oneshot::channel();
|
||||
let run = tokio::spawn(async move {
|
||||
begin.await.unwrap();
|
||||
engine.run().await;
|
||||
});
|
||||
let first = manager
|
||||
.start_turn(
|
||||
&thread.id,
|
||||
StartTurnRequest {
|
||||
prompt: "first request".into(),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(mock.call_count(), 0);
|
||||
if interrupt {
|
||||
manager.interrupt_turn(&thread.id, &first.id).await?;
|
||||
} else {
|
||||
manager
|
||||
.steer_turn(
|
||||
&thread.id,
|
||||
&first.id,
|
||||
SteerTurnRequest {
|
||||
prompt: "NEW CORRECTION".into(),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
release.send(()).unwrap();
|
||||
let terminal = wait_for_terminal_turn(&manager, &first.id, Duration::from_secs(10)).await?;
|
||||
if interrupt {
|
||||
assert_eq!(terminal.status, RuntimeTurnStatus::Interrupted);
|
||||
assert_eq!(
|
||||
mock.call_count(),
|
||||
0,
|
||||
"pre-dispatch interrupt must prevent the first provider request"
|
||||
);
|
||||
assert!(
|
||||
manager
|
||||
.store
|
||||
.list_items_for_turn(&first.id)?
|
||||
.iter()
|
||||
.all(|item| item.kind != TurnItemKind::ToolCall)
|
||||
);
|
||||
} else {
|
||||
assert_eq!(terminal.status, RuntimeTurnStatus::Completed);
|
||||
let request = mock
|
||||
.last_request()
|
||||
.expect("real Engine reached mock provider");
|
||||
let text = serde_json::to_string(&request.messages)?;
|
||||
assert_eq!(text.matches("NEW CORRECTION").count(), 1);
|
||||
}
|
||||
assert_eq!(
|
||||
manager
|
||||
.events_since(&thread.id, None)?
|
||||
.iter()
|
||||
.filter(|event| event.event == "turn.completed"
|
||||
&& event.turn_id.as_deref() == Some(first.id.as_str()))
|
||||
.count(),
|
||||
1
|
||||
);
|
||||
if follow_up {
|
||||
handle.steer("stale correction after completion").await?;
|
||||
let second = manager
|
||||
.start_turn(
|
||||
&thread.id,
|
||||
StartTurnRequest {
|
||||
prompt: "second request".into(),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(
|
||||
wait_for_terminal_turn(&manager, &second.id, Duration::from_secs(10))
|
||||
.await?
|
||||
.status,
|
||||
RuntimeTurnStatus::Completed
|
||||
);
|
||||
assert_eq!(mock.call_count(), 1);
|
||||
assert!(
|
||||
!serde_json::to_string(&mock.last_request().unwrap().messages)?
|
||||
.contains("stale correction after completion")
|
||||
);
|
||||
}
|
||||
handle.send(Op::Shutdown).await?;
|
||||
tokio::time::timeout(Duration::from_secs(10), run).await??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn queued_turn_interrupt_survives_engine_startup() -> Result<()> {
|
||||
control_case(true, false).await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn queued_turn_steer_reaches_first_request() -> Result<()> {
|
||||
control_case(false, false).await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn new_turn_does_not_inherit_previous_turn_controls() -> Result<()> {
|
||||
control_case(true, true).await
|
||||
}
|
||||
}
|
||||
|
||||
fn test_runtime_dir() -> PathBuf {
|
||||
std::env::temp_dir().join(format!("deepseek-runtime-threads-{}", Uuid::new_v4()))
|
||||
}
|
||||
@@ -410,6 +810,7 @@ fn sample_thread(thread_id: &str) -> ThreadRecord {
|
||||
task_id: None,
|
||||
title: None,
|
||||
session_id: None,
|
||||
saved_session_checkpoint: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10450,7 +10851,7 @@ async fn steer_turn_on_active_turn_records_item_and_event() -> Result<()> {
|
||||
})
|
||||
.await;
|
||||
if let Some(steer) = rx_steer.recv().await {
|
||||
let _ = steer_seen_tx.send(steer);
|
||||
let _ = steer_seen_tx.send(steer.content);
|
||||
}
|
||||
let _ = tx_event
|
||||
.send(EngineEvent::MessageStarted { index: 0 })
|
||||
@@ -10570,7 +10971,9 @@ async fn steer_receipts_outlive_caller_cancellation_after_engine_acceptance() ->
|
||||
.await
|
||||
});
|
||||
assert_eq!(
|
||||
tokio::time::timeout(Duration::from_secs(2), rx_steer.recv()).await?,
|
||||
tokio::time::timeout(Duration::from_secs(2), rx_steer.recv())
|
||||
.await?
|
||||
.map(|steer| steer.content),
|
||||
Some("keep the accepted steer".to_string())
|
||||
);
|
||||
steer_task.abort();
|
||||
@@ -11154,6 +11557,7 @@ fn opening_manager_recovers_stale_queued_and_in_progress_work() -> Result<()> {
|
||||
task_id: None,
|
||||
title: None,
|
||||
session_id: None,
|
||||
saved_session_checkpoint: None,
|
||||
};
|
||||
manager.store.save_thread(&thread)?;
|
||||
|
||||
|
||||
Reference in new issue
Block a user