3 Commits
Author SHA1 Message Date
Hare ccf7de1a55 fix: preserve run budget across segment forks 2026-08-26 14:13:43 +09:00
Hare ccf3c80d29 fix: clear abandoned run budget before compaction 2026-08-26 14:03:07 +09:00
Hare 17c629136a fix: scope max turns to logical runs 2026-08-26 13:47:35 +09:00
7 changed files with 545 additions and 31 deletions
+69 -16
View File
@@ -179,14 +179,20 @@ pub struct Engine<C: LlmClient, S: EngineState = Mutable> {
history: Vec<Item>, history: Vec<Item>,
/// History length at lock time (only meaningful in Locked state) /// History length at lock time (only meaningful in Locked state)
locked_prefix_len: usize, locked_prefix_len: usize,
/// AgentTurn count. /// AgentTurn count across the lifetime of this Engine.
/// ///
/// Once retry (`agen-stream-continuation`) is implemented, an /// Once retry (`agen-stream-continuation`) is implemented, an
/// AgentTurn collapses N retried `LlmCall`s with identical input; /// AgentTurn collapses N retried `LlmCall`s with identical input;
/// today retry is not implemented so AgentTurn and LlmCall fire 1:1 /// today retry is not implemented so AgentTurn and LlmCall fire 1:1
/// and the increment site (the LLM-call loop) is shared. /// and the increment site (the LLM-call loop) is shared.
/// `max_turns` is interpreted as a per-`run()` AgentTurn cap.
turn_count: usize, turn_count: usize,
/// AgentTurns consumed by the currently active logical run.
///
/// A fresh [`run`](Self::run) starts at zero. Pause and Yield retain the
/// count for [`resume`](Self::resume), while terminal outcomes clear it.
/// `max_turns` is enforced against this run-scoped count rather than the
/// cumulative `turn_count` above.
active_run_turn_count: Option<usize>,
/// LlmCall count (per-Engine running counter, monotonic). Unlike /// LlmCall count (per-Engine running counter, monotonic). Unlike
/// `turn_count` this never collapses retries. /// `turn_count` this never collapses retries.
llm_call_count: usize, llm_call_count: usize,
@@ -268,6 +274,20 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
self.last_run_interrupted = false; self.last_run_interrupted = false;
} }
fn start_logical_run(&mut self) {
self.active_run_turn_count = Some(0);
}
fn ensure_logical_run(&mut self) {
self.active_run_turn_count.get_or_insert(0);
}
fn finish_logical_run(&mut self, result: &Result<EngineResult, EngineError>) {
if !matches!(result, Ok(EngineResult::Paused) | Ok(EngineResult::Yielded)) {
self.active_run_turn_count = None;
}
}
fn drain_cancel_queue(&mut self) { fn drain_cancel_queue(&mut self) {
while self.cancel_rx.try_recv().is_ok() {} while self.cancel_rx.try_recv().is_ok() {}
} }
@@ -650,6 +670,23 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
self.turn_count self.turn_count
} }
/// Get the AgentTurns consumed by an interrupted logical run.
///
/// `Some` is retained only while Pause or Yield permits a later
/// [`resume`](Self::resume). Terminal outcomes return this to `None`.
pub fn active_run_turn_count(&self) -> Option<usize> {
self.active_run_turn_count
}
/// Restore the persisted turn budget of an interrupted logical run.
///
/// Session owners restore this together with the cumulative turn count and
/// history. `None` means there is no resumable logical run and the next
/// [`resume`](Self::resume) starts a fresh budget.
pub fn set_active_run_turn_count(&mut self, turn_count: Option<usize>) {
self.active_run_turn_count = turn_count;
}
/// Get the current LlmCall count (per-Engine running counter, never /// Get the current LlmCall count (per-Engine running counter, never
/// collapsed by retry). /// collapsed by retry).
pub fn llm_call_count(&self) -> usize { pub fn llm_call_count(&self) -> usize {
@@ -1123,6 +1160,19 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
return Err(EngineError::Cancelled); return Err(EngineError::Cancelled);
} }
if let Some(max) = self.max_turns
&& self.active_run_turn_count.unwrap_or(0) >= max as usize
{
info!(
active_run_turn_count = self.active_run_turn_count.unwrap_or(0),
total_turn_count = self.turn_count,
max_turns = max,
"Logical run turn limit reached"
);
self.last_run_interrupted = false;
return Ok(EngineResult::LimitReached);
}
let current_turn = self.turn_count; let current_turn = self.turn_count;
if !continuing_stream { if !continuing_stream {
debug!(turn = current_turn, "Turn start"); debug!(turn = current_turn, "Turn start");
@@ -1314,6 +1364,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
cb(current_turn); cb(current_turn);
} }
self.turn_count += 1; self.turn_count += 1;
*self.active_run_turn_count.get_or_insert(0) += 1;
// Collect and commit assistant items. Routed through // Collect and commit assistant items. Routed through
// `append_history_items` so observers see each item as it lands. // `append_history_items` so observers see each item as it lands.
@@ -1344,18 +1395,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
if let Some(result) = self.execute_and_commit_tools(tool_calls).await? { if let Some(result) = self.execute_and_commit_tools(tool_calls).await? {
return Ok(result); return Ok(result);
} }
if let Some(max) = self.max_turns {
if self.turn_count >= max as usize {
info!(
turn_count = self.turn_count,
max_turns = max,
"Turn limit reached"
);
self.last_run_interrupted = false;
return Ok(EngineResult::LimitReached);
}
}
} }
} }
@@ -1664,6 +1703,7 @@ impl<C: LlmClient> Engine<C, Mutable> {
history: Vec::new(), history: Vec::new(),
locked_prefix_len: 0, locked_prefix_len: 0,
turn_count: 0, turn_count: 0,
active_run_turn_count: None,
llm_call_count: 0, llm_call_count: 0,
tool_execution_batch_count: 0, tool_execution_batch_count: 0,
max_turns: None, max_turns: None,
@@ -1866,6 +1906,9 @@ impl<C: LlmClient> Engine<C, Mutable> {
/// Set the last_run_interrupted flag (for session restoration) /// Set the last_run_interrupted flag (for session restoration)
pub fn set_last_run_interrupted(&mut self, interrupted: bool) { pub fn set_last_run_interrupted(&mut self, interrupted: bool) {
self.last_run_interrupted = interrupted; self.last_run_interrupted = interrupted;
if !interrupted {
self.active_run_turn_count = None;
}
} }
/// Apply configuration (reserved for future extensions) /// Apply configuration (reserved for future extensions)
@@ -1934,6 +1977,7 @@ impl<C: LlmClient> Engine<C, Mutable> {
history: self.history, history: self.history,
locked_prefix_len, locked_prefix_len,
turn_count: self.turn_count, turn_count: self.turn_count,
active_run_turn_count: self.active_run_turn_count,
llm_call_count: self.llm_call_count, llm_call_count: self.llm_call_count,
tool_execution_batch_count: self.tool_execution_batch_count, tool_execution_batch_count: self.tool_execution_batch_count,
max_turns: self.max_turns, max_turns: self.max_turns,
@@ -1974,6 +2018,8 @@ impl<C: LlmClient> Engine<C, Locked> {
&mut self, &mut self,
user_input: impl Into<String>, user_input: impl Into<String>,
) -> Result<EngineResult, EngineError> { ) -> Result<EngineResult, EngineError> {
// Supplying new user input abandons any paused/yielded logical run.
self.active_run_turn_count = None;
self.reset_interruption_state(); self.reset_interruption_state();
// Interceptor: on_prompt_submit // Interceptor: on_prompt_submit
let mut user_item = Item::user_message(user_input); let mut user_item = Item::user_message(user_input);
@@ -1991,8 +2037,11 @@ impl<C: LlmClient> Engine<C, Locked> {
if !extras.is_empty() { if !extras.is_empty() {
self.append_history_items(extras)?; self.append_history_items(extras)?;
} }
self.start_logical_run();
let result = self.run_turn_loop().await; let result = self.run_turn_loop().await;
self.finalize_interruption(result).await let result = self.finalize_interruption(result).await;
self.finish_logical_run(&result);
result
} }
/// Resume execution (from Paused state) /// Resume execution (from Paused state)
@@ -2000,8 +2049,11 @@ impl<C: LlmClient> Engine<C, Locked> {
/// Resumes turn processing from current state without adding a new user message. /// Resumes turn processing from current state without adding a new user message.
pub async fn resume(&mut self) -> Result<EngineResult, EngineError> { pub async fn resume(&mut self) -> Result<EngineResult, EngineError> {
self.reset_interruption_state(); self.reset_interruption_state();
self.ensure_logical_run();
let result = self.run_turn_loop().await; let result = self.run_turn_loop().await;
self.finalize_interruption(result).await let result = self.finalize_interruption(result).await;
self.finish_logical_run(&result);
result
} }
/// Get the prefix length at lock time /// Get the prefix length at lock time
@@ -2027,6 +2079,7 @@ impl<C: LlmClient> Engine<C, Locked> {
history: self.history, history: self.history,
locked_prefix_len: 0, locked_prefix_len: 0,
turn_count: self.turn_count, turn_count: self.turn_count,
active_run_turn_count: self.active_run_turn_count,
llm_call_count: self.llm_call_count, llm_call_count: self.llm_call_count,
tool_execution_batch_count: self.tool_execution_batch_count, tool_execution_batch_count: self.tool_execution_batch_count,
max_turns: self.max_turns, max_turns: self.max_turns,
+186 -1
View File
@@ -9,9 +9,12 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use agen::Item; use agen::Item;
use agen::interceptor::{
Interceptor, PreRequestAction, PreToolAction, ToolCallInfo, TurnEndAction,
};
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use agen::{Engine, EngineError}; use agen::{Engine, EngineError, EngineResult};
use async_trait::async_trait; use async_trait::async_trait;
use common::MockLlmClient; use common::MockLlmClient;
@@ -561,3 +564,185 @@ fn test_system_prompt_change_after_unlock() {
let relocked = unlocked.lock(); let relocked = unlocked.lock();
assert_eq!(relocked.get_system_prompt(), Some("New prompt")); assert_eq!(relocked.get_system_prompt(), Some("New prompt"));
} }
fn completed_text_events() -> Vec<Event> {
vec![
Event::text_block_start(0),
Event::text_delta(0, "done"),
Event::text_block_stop(0, None),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
struct YieldOnce {
calls: AtomicUsize,
}
#[async_trait]
impl Interceptor for YieldOnce {
async fn pre_llm_request(&self, _context: &mut Vec<Item>) -> PreRequestAction {
if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
PreRequestAction::Yield
} else {
PreRequestAction::Continue
}
}
}
struct PauseToolOnce {
calls: AtomicUsize,
}
#[async_trait]
impl Interceptor for PauseToolOnce {
async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> PreToolAction {
if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
PreToolAction::Pause
} else {
PreToolAction::Continue
}
}
}
struct ContinueTurnOnce {
calls: AtomicUsize,
}
#[async_trait]
impl Interceptor for ContinueTurnOnce {
async fn on_turn_end(&self, _history: &[Item]) -> TurnEndAction {
if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
TurnEndAction::ContinueWithMessages(vec![Item::system_message("continue")])
} else {
TurnEndAction::Finish
}
}
}
#[tokio::test]
async fn max_turns_is_scoped_to_each_fresh_run() {
let responses = vec![completed_text_events(), completed_text_events()];
let mut engine = Engine::new(MockLlmClient::with_responses(responses));
engine.set_max_turns(Some(1));
let mut engine = engine.lock();
assert_eq!(engine.run("first").await.unwrap(), EngineResult::Finished);
assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), None);
assert_eq!(engine.run("second").await.unwrap(), EngineResult::Finished);
assert_eq!(engine.turn_count(), 2);
assert_eq!(engine.active_run_turn_count(), None);
}
#[tokio::test]
async fn yielded_resume_keeps_the_same_unspent_turn_budget() {
let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
engine.set_max_turns(Some(1));
engine.set_interceptor(YieldOnce {
calls: AtomicUsize::new(0),
});
let mut engine = engine.lock();
assert_eq!(engine.run("start").await.unwrap(), EngineResult::Yielded);
assert_eq!(engine.turn_count(), 0);
assert_eq!(engine.active_run_turn_count(), Some(0));
assert_eq!(engine.resume().await.unwrap(), EngineResult::Finished);
assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), None);
}
#[tokio::test]
async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() {
let events = vec![
Event::tool_use_start(0, "call_1", "count_tool"),
Event::tool_input_delta(0, "{}"),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
];
let tool = CountingTool::new("count_tool");
let mut engine = Engine::new(MockLlmClient::new(events));
engine.set_max_turns(Some(1));
engine.register_tool(tool.definition());
engine.set_interceptor(PauseToolOnce {
calls: AtomicUsize::new(0),
});
let mut engine = engine.lock();
assert_eq!(engine.run("call it").await.unwrap(), EngineResult::Paused);
assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), Some(1));
assert_eq!(tool.call_count(), 0);
assert_eq!(engine.resume().await.unwrap(), EngineResult::LimitReached);
assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), None);
assert_eq!(tool.call_count(), 1, "the consumed turn's tool still runs");
}
#[tokio::test]
async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() {
let tool_events = vec![
Event::tool_use_start(0, "call_1", "count_tool"),
Event::tool_input_delta(0, "{}"),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
];
let client = MockLlmClient::with_responses(vec![tool_events, completed_text_events()]);
let tool = CountingTool::new("count_tool");
let mut engine = Engine::new(client);
engine.set_max_turns(Some(1));
engine.register_tool(tool.definition());
engine.set_interceptor(PauseToolOnce {
calls: AtomicUsize::new(0),
});
let mut engine = engine.lock();
assert_eq!(engine.run("pause").await.unwrap(), EngineResult::Paused);
assert_eq!(engine.active_run_turn_count(), Some(1));
assert_eq!(engine.run("replace").await.unwrap(), EngineResult::Finished);
assert_eq!(engine.turn_count(), 2);
assert_eq!(engine.active_run_turn_count(), None);
assert_eq!(tool.call_count(), 1, "pending-tool semantics are unchanged");
}
#[tokio::test]
async fn interceptor_continuation_consumes_the_logical_run_budget() {
let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
engine.set_max_turns(Some(1));
engine.set_interceptor(ContinueTurnOnce {
calls: AtomicUsize::new(0),
});
let mut engine = engine.lock();
assert_eq!(
engine.run("start").await.unwrap(),
EngineResult::LimitReached
);
assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.llm_call_count(), 1);
assert_eq!(engine.active_run_turn_count(), None);
}
#[tokio::test]
async fn restored_active_run_budget_is_enforced_before_another_llm_call() {
let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
engine.set_max_turns(Some(1));
engine.set_turn_count(7);
engine.set_last_run_interrupted(true);
engine.set_active_run_turn_count(Some(1));
let mut engine = engine.lock();
assert_eq!(engine.resume().await.unwrap(), EngineResult::LimitReached);
assert_eq!(engine.turn_count(), 7);
assert_eq!(engine.llm_call_count(), 0);
assert_eq!(engine.active_run_turn_count(), None);
}
+2
View File
@@ -307,6 +307,7 @@ pub fn save_run_completed(
segment_id: SegmentId, segment_id: SegmentId,
result: EngineResult, result: EngineResult,
interrupted: bool, interrupted: bool,
active_run_turn_count: Option<usize>,
) -> Result<(), StoreError> { ) -> Result<(), StoreError> {
append_entry( append_entry(
store, store,
@@ -316,6 +317,7 @@ pub fn save_run_completed(
ts: segment_log::now_millis(), ts: segment_log::now_millis(),
interrupted, interrupted,
result, result,
active_run_turn_count,
}, },
) )
} }
+132 -2
View File
@@ -125,11 +125,16 @@ pub enum LogEntry {
TurnEnd { ts: u64, turn_count: usize }, TurnEnd { ts: u64, turn_count: usize },
/// `run()` / `resume()` が `EngineResult` で正常終了した。 /// `run()` / `resume()` が `EngineResult` で正常終了した。
/// Audit-only metadata: replay は `interrupted` のみ反映する。 /// Replay restores both interruption state and any resumable logical-run
/// turn budget.
RunCompleted { RunCompleted {
ts: u64, ts: u64,
interrupted: bool, interrupted: bool,
result: EngineResult, result: EngineResult,
/// AgentTurns consumed by a paused/yielded logical run. Terminal
/// outcomes persist `None`.
#[serde(default, skip_serializing_if = "Option::is_none")]
active_run_turn_count: Option<usize>,
}, },
/// `run()` / `resume()` が `EngineError` で終了した。 /// `run()` / `resume()` が `EngineError` で終了した。
@@ -141,6 +146,15 @@ pub enum LogEntry {
message: String, message: String,
}, },
/// Restores an active logical-run budget at a segment boundary, notably
/// after compaction replaced the segment that held the original Invoke and
/// RunCompleted entries.
ActiveRunCheckpoint {
ts: u64,
active_turn_count: usize,
total_turn_count: usize,
},
/// A paused interrupted turn was explicitly abandoned without calling /// A paused interrupted turn was explicitly abandoned without calling
/// `run()` or `resume()` again. Replay clears the interrupted marker so /// `run()` or `resume()` again. Replay clears the interrupted marker so
/// the restored Worker is idle and future user input starts a normal new turn. /// the restored Worker is idle and future user input starts a normal new turn.
@@ -209,6 +223,8 @@ pub struct RestoredState {
pub config: RequestConfig, pub config: RequestConfig,
pub history: Vec<Item>, pub history: Vec<Item>,
pub turn_count: usize, pub turn_count: usize,
/// AgentTurns consumed by the active paused/yielded logical run.
pub active_run_turn_count: Option<usize>,
pub last_run_interrupted: bool, pub last_run_interrupted: bool,
/// Number of entries replayed. `0` means the segment log was empty. /// Number of entries replayed. `0` means the segment log was empty.
/// Writers track their own append count via the same counter so /// Writers track their own append count via the same counter so
@@ -238,6 +254,7 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState {
config: RequestConfig::default(), config: RequestConfig::default(),
history: Vec::new(), history: Vec::new(),
turn_count: 0, turn_count: 0,
active_run_turn_count: None,
last_run_interrupted: false, last_run_interrupted: false,
entries_count: 0, entries_count: 0,
usage_history: Vec::new(), usage_history: Vec::new(),
@@ -265,6 +282,7 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState {
// A terminal run record below clears or refines this. If the // A terminal run record below clears or refines this. If the
// log ends first, restore must treat the turn as interrupted. // log ends first, restore must treat the turn as interrupted.
state.last_run_interrupted = true; state.last_run_interrupted = true;
state.active_run_turn_count = Some(0);
} }
LogEntry::UserInput { LogEntry::UserInput {
segments, segments,
@@ -290,16 +308,44 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState {
state.history.push(item.to_history_item()); state.history.push(item.to_history_item());
} }
LogEntry::TurnEnd { turn_count, .. } => { LogEntry::TurnEnd { turn_count, .. } => {
if let Some(active_turn_count) = &mut state.active_run_turn_count {
*active_turn_count += turn_count.saturating_sub(state.turn_count);
}
state.turn_count = *turn_count; state.turn_count = *turn_count;
} }
LogEntry::RunCompleted { interrupted, .. } => { LogEntry::RunCompleted {
interrupted,
result,
active_run_turn_count,
..
} => {
state.last_run_interrupted = *interrupted; state.last_run_interrupted = *interrupted;
if *interrupted && matches!(result, EngineResult::Paused | EngineResult::Yielded) {
// Legacy entries omit the explicit field; retain the
// Invoke/TurnEnd-derived count in that case.
if let Some(turn_count) = active_run_turn_count {
state.active_run_turn_count = Some(*turn_count);
}
} else {
state.active_run_turn_count = None;
}
} }
LogEntry::RunErrored { interrupted, .. } => { LogEntry::RunErrored { interrupted, .. } => {
state.last_run_interrupted = *interrupted; state.last_run_interrupted = *interrupted;
state.active_run_turn_count = None;
}
LogEntry::ActiveRunCheckpoint {
active_turn_count,
total_turn_count,
..
} => {
state.active_run_turn_count = Some(*active_turn_count);
state.turn_count = *total_turn_count;
state.last_run_interrupted = true;
} }
LogEntry::PausedTurnAbandoned { .. } => { LogEntry::PausedTurnAbandoned { .. } => {
state.last_run_interrupted = false; state.last_run_interrupted = false;
state.active_run_turn_count = None;
} }
LogEntry::ConfigChanged { config, .. } => { LogEntry::ConfigChanged { config, .. } => {
state.config = config.clone(); state.config = config.clone();
@@ -397,6 +443,7 @@ mod tests {
ts: 3200, ts: 3200,
interrupted: false, interrupted: false,
result: EngineResult::Finished, result: EngineResult::Finished,
active_run_turn_count: None,
}, },
]); ]);
assert_eq!(state.history.len(), 2); assert_eq!(state.history.len(), 2);
@@ -695,10 +742,93 @@ mod tests {
ts: 100, ts: 100,
interrupted: true, interrupted: true,
result: EngineResult::Paused, result: EngineResult::Paused,
active_run_turn_count: Some(1),
}, },
LogEntry::PausedTurnAbandoned { ts: 200 }, LogEntry::PausedTurnAbandoned { ts: 200 },
]); ]);
assert!(!state.last_run_interrupted); assert!(!state.last_run_interrupted);
assert_eq!(state.active_run_turn_count, None);
}
#[test]
fn replay_restores_active_run_budget_across_compaction_checkpoint() {
let state = collect_state(&[
LogEntry::SegmentStart {
ts: 0,
session_id: uuid::Uuid::nil(),
system_prompt: None,
config: RequestConfig::default(),
history: vec![],
forked_from: None,
compacted_from: None,
},
LogEntry::ActiveRunCheckpoint {
ts: 100,
active_turn_count: 3,
total_turn_count: 9,
},
]);
assert_eq!(state.turn_count, 9);
assert_eq!(state.active_run_turn_count, Some(3));
assert!(state.last_run_interrupted);
}
#[test]
fn legacy_interrupted_run_derives_budget_from_invoke_and_turn_end() {
let entry: LogEntry = serde_json::from_value(serde_json::json!({
"kind": "run_completed",
"ts": 300,
"interrupted": true,
"result": "paused"
}))
.expect("legacy run-completed entry");
let state = collect_state(&[
LogEntry::SegmentStart {
ts: 0,
session_id: uuid::Uuid::nil(),
system_prompt: None,
config: RequestConfig::default(),
history: vec![],
forked_from: None,
compacted_from: None,
},
LogEntry::Invoke {
ts: 100,
trigger: InvokeKind::UserSend,
},
LogEntry::TurnEnd {
ts: 200,
turn_count: 2,
},
entry,
]);
assert_eq!(state.active_run_turn_count, Some(2));
assert!(state.last_run_interrupted);
}
#[test]
fn non_resumable_interruption_clears_the_active_run_budget() {
let state = collect_state(&[
LogEntry::Invoke {
ts: 100,
trigger: InvokeKind::UserSend,
},
LogEntry::TurnEnd {
ts: 200,
turn_count: 2,
},
LogEntry::RunCompleted {
ts: 300,
interrupted: true,
result: EngineResult::LimitReached,
active_run_turn_count: None,
},
]);
assert!(state.last_run_interrupted);
assert_eq!(state.active_run_turn_count, None);
} }
#[test] #[test]
@@ -51,6 +51,7 @@ fn round_trip_write_and_read() {
ts: 3200, ts: 3200,
interrupted: false, interrupted: false,
result: EngineResult::Finished, result: EngineResult::Finished,
active_run_turn_count: None,
}, },
]; ];
@@ -132,6 +132,7 @@ async fn run_and_persist(
segment_id, segment_id,
r.clone(), r.clone(),
worker.last_run_interrupted(), worker.last_run_interrupted(),
worker.active_run_turn_count(),
) )
.unwrap(); .unwrap();
} }
@@ -309,6 +310,7 @@ async fn session_resume_after_pause() {
// Restore state and verify // Restore state and verify
let state = session_store::restore(&store, sid, segid).unwrap(); let state = session_store::restore(&store, sid, segid).unwrap();
assert!(state.last_run_interrupted); assert!(state.last_run_interrupted);
assert_eq!(state.active_run_turn_count, Some(2));
} }
#[tokio::test] #[tokio::test]
+153 -12
View File
@@ -781,9 +781,21 @@ struct EmptyTurnRollbackSnapshot {
usage_history_len: usize, usage_history_len: usize,
ai_activity_count: usize, ai_activity_count: usize,
last_run_interrupted: bool, last_run_interrupted: bool,
active_run_turn_count: Option<usize>,
flow_runtime_state: Option<flow::FlowRuntimeState>, flow_runtime_state: Option<flow::FlowRuntimeState>,
} }
fn active_run_checkpoint_entry(
active_run_turn_count: Option<usize>,
total_turn_count: usize,
) -> Option<LogEntry> {
active_run_turn_count.map(|active_turn_count| LogEntry::ActiveRunCheckpoint {
ts: segment_log::now_millis(),
active_turn_count,
total_turn_count,
})
}
fn is_ai_materialized_item(item: &Item) -> bool { fn is_ai_materialized_item(item: &Item) -> bool {
match item { match item {
Item::Message { role, .. } => *role == Role::Assistant, Item::Message { role, .. } => *role == Role::Assistant,
@@ -1777,6 +1789,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
self.engine_mut().set_turn_count(state.turn_count); self.engine_mut().set_turn_count(state.turn_count);
self.engine_mut() self.engine_mut()
.set_last_run_interrupted(state.last_run_interrupted); .set_last_run_interrupted(state.last_run_interrupted);
self.engine_mut()
.set_active_run_turn_count(state.active_run_turn_count);
self.user_segments = state.user_segments; self.user_segments = state.user_segments;
*self.usage_history.lock().expect("usage_history poisoned") = state.usage_history; *self.usage_history.lock().expect("usage_history poisoned") = state.usage_history;
*self *self
@@ -2350,6 +2364,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
usage_history_len, usage_history_len,
ai_activity_count: self.ai_activity_counter.load(Ordering::SeqCst), ai_activity_count: self.ai_activity_counter.load(Ordering::SeqCst),
last_run_interrupted: self.engine().last_run_interrupted(), last_run_interrupted: self.engine().last_run_interrupted(),
active_run_turn_count: self.engine().active_run_turn_count(),
flow_runtime_state: self flow_runtime_state: self
.flow_runtime_state .flow_runtime_state
.lock() .lock()
@@ -2381,6 +2396,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
self.engine_mut().truncate_history(snapshot.history_len); self.engine_mut().truncate_history(snapshot.history_len);
self.engine_mut() self.engine_mut()
.set_last_run_interrupted(snapshot.last_run_interrupted); .set_last_run_interrupted(snapshot.last_run_interrupted);
self.engine_mut()
.set_active_run_turn_count(snapshot.active_run_turn_count);
*self *self
.flow_runtime_state .flow_runtime_state
.lock() .lock()
@@ -2541,9 +2558,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
// `last_run_interrupted` flag; `Worker::resume` reuses the prior // `last_run_interrupted` flag; `Worker::resume` reuses the prior
// context via a different entry point and never triggers this // context via a different entry point and never triggers this
// path. // path.
if self.engine.as_ref().unwrap().last_run_interrupted() { self.prepare_interrupted_history_for_fresh_run()?;
self.apply_interrupt_prep()?;
}
self.prepare_for_run().await?; self.prepare_for_run().await?;
@@ -2654,6 +2669,19 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
out out
} }
/// Close interrupted history before a fresh user/notification run.
///
/// Clearing the interrupted flag also ends the old logical-run budget and
/// must happen before `prepare_for_run`: proactive compaction checkpoints
/// only resumable runs, never the run this invocation is abandoning.
fn prepare_interrupted_history_for_fresh_run(&mut self) -> Result<(), WorkerError> {
if self.engine().last_run_interrupted() {
self.apply_interrupt_prep()?;
self.engine_mut().set_last_run_interrupted(false);
}
Ok(())
}
/// Stage the post-interruption cleanup at the front of worker /// Stage the post-interruption cleanup at the front of worker
/// history: close every unanswered `Item::ToolCall` with a synthetic /// history: close every unanswered `Item::ToolCall` with a synthetic
/// `Item::ToolResult` (Anthropic wire-validity), then append a /// `Item::ToolResult` (Anthropic wire-validity), then append a
@@ -2769,12 +2797,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
), ),
"run_for_notification expects a non-UserSend InvokeKind; got {kind:?}" "run_for_notification expects a non-UserSend InvokeKind; got {kind:?}"
); );
// This is a fresh Invoke, not an explicit resume of the interrupted self.prepare_interrupted_history_for_fresh_run()?;
// turn. Close any dangling tool calls before an auto-run notification
// can enter `Engine::resume` and execute them again after a crash.
if self.engine.as_ref().unwrap().last_run_interrupted() {
self.apply_interrupt_prep()?;
}
self.prepare_for_run().await?; self.prepare_for_run().await?;
// IDLE → active marker for the buffered notification / worker-event // IDLE → active marker for the buffered notification / worker-event
@@ -2866,15 +2889,23 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
}), }),
compacted_from: None, compacted_from: None,
}; };
let mut initial_entries = vec![entry.clone()];
if let Some(checkpoint) =
active_run_checkpoint_entry(w.active_run_turn_count(), w.turn_count())
{
initial_entries.push(checkpoint);
}
self.store self.store
.create_segment(loc.session_id, fork_segment_id, &[entry.clone()]) .create_segment(loc.session_id, fork_segment_id, &initial_entries)
.map_err(WorkerError::from)?; .map_err(WorkerError::from)?;
self.segment_state.set_location(SegmentLocation { self.segment_state.set_location(SegmentLocation {
session_id: loc.session_id, session_id: loc.session_id,
segment_id: fork_segment_id, segment_id: fork_segment_id,
}); });
self.segment_state.set_entries_written(1); self.segment_state
self.sink.reset_with_initial(entry); .set_entries_written(initial_entries.len());
self.sink
.reset_with_initial_entries(initial_entries.clone());
if self.scope_allocation.is_some() { if self.scope_allocation.is_some() {
worker_allocation::update_segment(&self.manifest.worker.name, fork_segment_id)?; worker_allocation::update_segment(&self.manifest.worker.name, fork_segment_id)?;
} }
@@ -3218,12 +3249,14 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
} }
let interrupted = self.engine.as_ref().unwrap().last_run_interrupted(); let interrupted = self.engine.as_ref().unwrap().last_run_interrupted();
let active_run_turn_count = self.engine.as_ref().unwrap().active_run_turn_count();
match result { match result {
Ok(r) => { Ok(r) => {
self.commit_entry(LogEntry::RunCompleted { self.commit_entry(LogEntry::RunCompleted {
ts: segment_log::now_millis(), ts: segment_log::now_millis(),
interrupted, interrupted,
result: r.clone(), result: r.clone(),
active_run_turn_count,
})?; })?;
} }
Err(e) => { Err(e) => {
@@ -3754,6 +3787,11 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
}), }),
}; };
let mut initial_entries = vec![entry.clone()]; let mut initial_entries = vec![entry.clone()];
if let Some(checkpoint) =
active_run_checkpoint_entry(w.active_run_turn_count(), source_turn_count)
{
initial_entries.push(checkpoint);
}
if let Some(flow_state) = self if let Some(flow_state) = self
.flow_runtime_state .flow_runtime_state
.lock() .lock()
@@ -5179,6 +5217,7 @@ where
worker.set_request_config(state.config.clone()); worker.set_request_config(state.config.clone());
worker.set_turn_count(state.turn_count); worker.set_turn_count(state.turn_count);
worker.set_last_run_interrupted(state.last_run_interrupted); worker.set_last_run_interrupted(state.last_run_interrupted);
worker.set_active_run_turn_count(state.active_run_turn_count);
if anchored_on_summary { if anchored_on_summary {
worker.set_cache_anchor(Some(0)); worker.set_cache_anchor(Some(0));
} }
@@ -6957,6 +6996,108 @@ mod build_summary_prompt_tests {
} }
} }
#[tokio::test]
async fn fresh_run_clears_interrupted_budget_before_pre_run_compaction() {
let dir = tempfile::tempdir().unwrap();
let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap();
let mut worker = Worker::new(
minimal_manifest(),
Engine::new(NoopClient),
store,
WorkerWorkspaceContext::no_workspace(),
WorkerFilesystemAuthority::None,
Scope::empty(),
)
.await
.unwrap();
worker.ensure_segment_head().unwrap();
worker.engine_mut().set_last_run_interrupted(true);
worker.engine_mut().set_active_run_turn_count(Some(3));
worker.prepare_interrupted_history_for_fresh_run().unwrap();
assert!(!worker.engine().last_run_interrupted());
assert_eq!(worker.engine().active_run_turn_count(), None);
let checkpoint = active_run_checkpoint_entry(
worker.engine().active_run_turn_count(),
worker.engine().turn_count(),
);
assert!(checkpoint.is_none());
let mut replacement_entries = vec![LogEntry::SegmentStart {
ts: segment_log::now_millis(),
session_id: uuid::Uuid::nil(),
system_prompt: None,
config: RequestConfig::default(),
history: vec![],
forked_from: None,
compacted_from: None,
}];
replacement_entries.extend(checkpoint);
let restored = session_store::collect_state(&replacement_entries);
assert!(!restored.last_run_interrupted);
assert_eq!(restored.active_run_turn_count, None);
}
#[tokio::test]
async fn auto_fork_checkpoints_interrupted_run_budget_for_restore() {
let dir = tempfile::tempdir().unwrap();
let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap();
let mut worker = Worker::new(
minimal_manifest(),
Engine::new(NoopClient),
store,
WorkerWorkspaceContext::no_workspace(),
WorkerFilesystemAuthority::None,
Scope::empty(),
)
.await
.unwrap();
worker.ensure_segment_head().unwrap();
worker.engine_mut().set_turn_count(7);
worker.engine_mut().set_last_run_interrupted(true);
worker.engine_mut().set_active_run_turn_count(Some(3));
let session_id = worker.session_id();
let source_segment_id = worker.segment_id();
worker
.store()
.append(
session_id,
source_segment_id,
&LogEntry::Extension {
ts: segment_log::now_millis(),
domain: "test.auto_fork_drift".into(),
payload: serde_json::json!({}),
},
)
.unwrap();
worker.ensure_segment_head().unwrap();
let fork_segment_id = worker.segment_id();
assert_ne!(fork_segment_id, source_segment_id);
let fork_entries = worker
.store()
.read_all(session_id, fork_segment_id)
.unwrap();
assert!(matches!(
fork_entries.as_slice(),
[
LogEntry::SegmentStart { .. },
LogEntry::ActiveRunCheckpoint {
active_turn_count: 3,
total_turn_count: 7,
..
}
]
));
let restored = session_store::collect_state(&fork_entries);
assert!(restored.last_run_interrupted);
assert_eq!(restored.turn_count, 7);
assert_eq!(restored.active_run_turn_count, Some(3));
}
#[tokio::test] #[tokio::test]
async fn flow_transition_feature_installs_runtime_local_coordinator() { async fn flow_transition_feature_installs_runtime_local_coordinator() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();