fix: scope max turns to logical runs

crates/agen/src/engine.rs

@@ -179,14 +179,20 @@ pub struct Engine<C: LlmClient, S: EngineState = Mutable> {
     history: Vec<Item>,
     /// History length at lock time (only meaningful in Locked state)
     locked_prefix_len: usize,
-    /// AgentTurn count.
+    /// AgentTurn count across the lifetime of this Engine.
     ///
     /// Once retry (`agen-stream-continuation`) is implemented, an
     /// AgentTurn collapses N retried `LlmCall`s with identical input;
     /// today retry is not implemented so AgentTurn and LlmCall fire 1:1
     /// and the increment site (the LLM-call loop) is shared.
-    /// `max_turns` is interpreted as a per-`run()` AgentTurn cap.
     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
     /// `turn_count` this never collapses retries.
     llm_call_count: usize,
@@ -268,6 +274,20 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
         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) {
         while self.cancel_rx.try_recv().is_ok() {}
     }
@@ -650,6 +670,23 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
         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
     /// collapsed by retry).
     pub fn llm_call_count(&self) -> usize {
@@ -1123,6 +1160,19 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
                 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;
             if !continuing_stream {
                 debug!(turn = current_turn, "Turn start");
@@ -1314,6 +1364,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
                 cb(current_turn);
             }
             self.turn_count += 1;
+            *self.active_run_turn_count.get_or_insert(0) += 1;
 
             // Collect and commit assistant items. Routed through
             // `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? {
                 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(),
             locked_prefix_len: 0,
             turn_count: 0,
+            active_run_turn_count: None,
             llm_call_count: 0,
             tool_execution_batch_count: 0,
             max_turns: None,
@@ -1866,6 +1906,9 @@ impl<C: LlmClient> Engine<C, Mutable> {
     /// Set the last_run_interrupted flag (for session restoration)
     pub fn set_last_run_interrupted(&mut self, interrupted: bool) {
         self.last_run_interrupted = interrupted;
+        if !interrupted {
+            self.active_run_turn_count = None;
+        }
     }
 
     /// Apply configuration (reserved for future extensions)
@@ -1934,6 +1977,7 @@ impl<C: LlmClient> Engine<C, Mutable> {
             history: self.history,
             locked_prefix_len,
             turn_count: self.turn_count,
+            active_run_turn_count: self.active_run_turn_count,
             llm_call_count: self.llm_call_count,
             tool_execution_batch_count: self.tool_execution_batch_count,
             max_turns: self.max_turns,
@@ -1974,6 +2018,8 @@ impl<C: LlmClient> Engine<C, Locked> {
         &mut self,
         user_input: impl Into<String>,
     ) -> Result<EngineResult, EngineError> {
+        // Supplying new user input abandons any paused/yielded logical run.
+        self.active_run_turn_count = None;
         self.reset_interruption_state();
         // Interceptor: on_prompt_submit
         let mut user_item = Item::user_message(user_input);
@@ -1991,8 +2037,11 @@ impl<C: LlmClient> Engine<C, Locked> {
         if !extras.is_empty() {
             self.append_history_items(extras)?;
         }
+        self.start_logical_run();
         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)
@@ -2000,8 +2049,11 @@ impl<C: LlmClient> Engine<C, Locked> {
     /// Resumes turn processing from current state without adding a new user message.
     pub async fn resume(&mut self) -> Result<EngineResult, EngineError> {
         self.reset_interruption_state();
+        self.ensure_logical_run();
         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
@@ -2027,6 +2079,7 @@ impl<C: LlmClient> Engine<C, Locked> {
             history: self.history,
             locked_prefix_len: 0,
             turn_count: self.turn_count,
+            active_run_turn_count: self.active_run_turn_count,
             llm_call_count: self.llm_call_count,
             tool_execution_batch_count: self.tool_execution_batch_count,
             max_turns: self.max_turns,

crates/agen/tests/engine_state_test.rs

@@ -9,9 +9,12 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex};
 
 use agen::Item;
+use agen::interceptor::{
+    Interceptor, PreRequestAction, PreToolAction, ToolCallInfo, TurnEndAction,
+};
 use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
 use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
-use agen::{Engine, EngineError};
+use agen::{Engine, EngineError, EngineResult};
 use async_trait::async_trait;
 use common::MockLlmClient;
 
@@ -561,3 +564,185 @@ fn test_system_prompt_change_after_unlock() {
     let relocked = unlocked.lock();
     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);
+}