llm_stack/tool/loop_resumable.rs
1//! Resumable tool loop with caller-controlled iteration.
2//!
3//! Unlike [`tool_loop`](super::tool_loop) which runs to completion autonomously,
4//! the resumable loop yields control to the caller after each tool execution
5//! round. The caller can then decide to continue, inject messages, or stop.
6//!
7//! This enables orchestration patterns like:
8//! - Multi-agent systems where a master inspects tool results between iterations
9//! - External event injection (user follow-ups, worker completions)
10//! - Context compaction between iterations
11//! - Custom routing logic based on which tools were called
12//!
13//! # Example
14//!
15//! ```rust,no_run
16//! use llm_stack::tool::{ToolLoopConfig, ToolRegistry, ToolLoopHandle, TurnResult, LoopCommand};
17//! use llm_stack::{ChatParams, ChatMessage};
18//!
19//! # async fn example(provider: &dyn llm_stack::DynProvider) -> Result<(), llm_stack::LlmError> {
20//! let registry: ToolRegistry<()> = ToolRegistry::new();
21//! let params = ChatParams {
22//! messages: vec![ChatMessage::user("Hello")],
23//! ..Default::default()
24//! };
25//!
26//! let mut handle = ToolLoopHandle::new(
27//! provider,
28//! ®istry,
29//! params,
30//! ToolLoopConfig::default(),
31//! &(),
32//! );
33//!
34//! loop {
35//! match handle.next_turn().await {
36//! TurnResult::Yielded(turn) => {
37//! // Text from this turn is directly available
38//! if let Some(text) = turn.assistant_text() {
39//! println!("LLM said: {text}");
40//! }
41//! // Inspect results, decide what to do
42//! turn.continue_loop();
43//! }
44//! TurnResult::Completed(done) => {
45//! println!("Done: {:?}", done.response.text());
46//! break;
47//! }
48//! TurnResult::Error(err) => {
49//! eprintln!("Error: {}", err.error);
50//! break;
51//! }
52//! }
53//! }
54//! # Ok(())
55//! # }
56//! ```
57
58use crate::chat::{ChatMessage, ChatResponse, ContentBlock, ToolCall, ToolResult};
59use crate::error::LlmError;
60use crate::provider::{ChatParams, DynProvider};
61use crate::usage::Usage;
62
63use super::LoopDepth;
64use super::ToolRegistry;
65use super::config::{TerminationReason, ToolLoopConfig, ToolLoopResult};
66use super::loop_core::{CompletedData, ErrorData, IterationOutcome, LoopCore};
67
68// ── Shared macros for Yielded-like types ─────────────────────────────
69
70/// Implements the common methods on a `Yielded`-like struct.
71///
72/// Both `Yielded` and `OwnedYielded` have identical field layouts and
73/// method bodies. The only difference is the handle type they borrow.
74/// This macro eliminates the duplication.
75///
76/// Expects the struct to have fields: `handle`, `assistant_content`,
77/// `tool_calls`, `results`, `iteration`, `total_usage`.
78macro_rules! impl_yielded_methods {
79 ($yielded:ident < $($lt:lifetime),* >) => {
80 impl<$($lt,)* Ctx: LoopDepth + Send + Sync + 'static> $yielded<$($lt,)* Ctx> {
81 /// Continue with the given command.
82 pub fn resume(self, command: LoopCommand) {
83 self.handle.resume(command);
84 }
85
86 /// Convenience: continue to the next LLM iteration with no injected messages.
87 pub fn continue_loop(self) {
88 self.resume(LoopCommand::Continue);
89 }
90
91 /// Convenience: inject messages and continue.
92 pub fn inject_and_continue(self, messages: Vec<ChatMessage>) {
93 self.resume(LoopCommand::InjectMessages(messages));
94 }
95
96 /// Convenience: stop the loop.
97 pub fn stop(self, reason: Option<String>) {
98 self.resume(LoopCommand::Stop(reason));
99 }
100
101 /// Extract text from `assistant_content` blocks.
102 ///
103 /// Returns the LLM's "thinking aloud" text emitted alongside tool calls,
104 /// or `None` if there were no text blocks.
105 pub fn assistant_text(&self) -> Option<String> {
106 let text: String = self
107 .assistant_content
108 .iter()
109 .filter_map(|block| match block {
110 ContentBlock::Text(t) => Some(t.as_str()),
111 _ => None,
112 })
113 .collect::<Vec<_>>()
114 .join("\n");
115 if text.is_empty() { None } else { Some(text) }
116 }
117
118 /// Access the full message history (read-only).
119 pub fn messages(&self) -> &[ChatMessage] {
120 self.handle.messages()
121 }
122
123 /// Access the full message history (mutable, for context compaction).
124 pub fn messages_mut(&mut self) -> &mut Vec<ChatMessage> {
125 self.handle.messages_mut()
126 }
127 }
128 };
129}
130
131pub(crate) use impl_yielded_methods;
132
133/// Converts an [`IterationOutcome`] into a `TurnResult`-like enum.
134///
135/// Both `TurnResult` and `OwnedTurnResult` need identical destructuring
136/// of `IterationOutcome` into their respective `Yielded`/`Completed`/`Error`
137/// variants. This macro generates that conversion.
138macro_rules! outcome_to_turn_result {
139 ($outcome:expr, $handle:expr, $turn_ty:ident, $yielded_ty:ident) => {
140 match $outcome {
141 IterationOutcome::ToolsExecuted {
142 tool_calls,
143 results,
144 assistant_content,
145 iteration,
146 total_usage,
147 } => $turn_ty::Yielded($yielded_ty {
148 handle: $handle,
149 tool_calls,
150 results,
151 assistant_content,
152 iteration,
153 total_usage,
154 }),
155 IterationOutcome::Completed(CompletedData {
156 response,
157 termination_reason,
158 iterations,
159 total_usage,
160 }) => $turn_ty::Completed(Completed {
161 response,
162 termination_reason,
163 iterations,
164 total_usage,
165 }),
166 IterationOutcome::Error(ErrorData {
167 error,
168 iterations,
169 total_usage,
170 }) => $turn_ty::Error(TurnError {
171 error,
172 iterations,
173 total_usage,
174 }),
175 }
176 };
177}
178
179pub(crate) use outcome_to_turn_result;
180
181/// Result of one turn of the tool loop.
182///
183/// Match on this to determine what happened and what you can do next.
184/// Each variant carries the data from the turn AND (for `Yielded`) a handle
185/// scoped to valid operations for that state.
186///
187/// This follows the same pattern as [`std::collections::hash_map::Entry`] —
188/// the variant gives you exactly the methods that make sense for that state.
189#[must_use = "a TurnResult must be matched — Yielded requires resume() to continue"]
190pub enum TurnResult<'a, 'h, Ctx: LoopDepth + Send + Sync + 'static> {
191 /// Tools were executed. The caller MUST consume this via `resume()`,
192 /// `continue_loop()`, `inject_and_continue()`, or `stop()`.
193 ///
194 /// While this variant exists, the `ToolLoopHandle` is mutably borrowed
195 /// and cannot be used directly. Consuming the `Yielded` releases the
196 /// borrow.
197 Yielded(Yielded<'a, 'h, Ctx>),
198
199 /// The loop completed (no tool calls, stop condition, max iterations, or timeout).
200 Completed(Completed),
201
202 /// An unrecoverable error occurred.
203 Error(TurnError),
204}
205
206/// Handle returned when tools were executed. Borrows the [`ToolLoopHandle`]
207/// mutably, so the caller cannot call `next_turn()` again until this is
208/// consumed via `resume()`, `continue_loop()`, `inject_and_continue()`, or
209/// `stop()`.
210///
211/// The text content the LLM produced alongside tool calls is available
212/// directly via [`assistant_content`](Self::assistant_content) and
213/// [`assistant_text()`](Self::assistant_text) — no need to scan
214/// `messages()`.
215#[must_use = "must call .resume(), .continue_loop(), .inject_and_continue(), or .stop() to continue"]
216pub struct Yielded<'a, 'h, Ctx: LoopDepth + Send + Sync + 'static> {
217 handle: &'h mut ToolLoopHandle<'a, Ctx>,
218
219 /// The tool calls the LLM requested.
220 pub tool_calls: Vec<ToolCall>,
221
222 /// Results from executing those tool calls.
223 pub results: Vec<ToolResult>,
224
225 /// Text content from the LLM's response alongside the tool calls.
226 ///
227 /// This is the `other_content` from `partition_content()` — `Text`,
228 /// `Reasoning`, `Image`, etc. — everything that isn't a `ToolCall` or
229 /// `ToolResult`. Previously only accessible by scanning `messages()`.
230 pub assistant_content: Vec<ContentBlock>,
231
232 /// Current iteration number (1-indexed).
233 pub iteration: u32,
234
235 /// Accumulated usage across all iterations so far.
236 pub total_usage: Usage,
237}
238
239impl_yielded_methods!(Yielded<'a, 'h>);
240
241/// Terminal: the loop completed successfully.
242pub struct Completed {
243 /// The final LLM response. Use `.text()` to get the response text.
244 pub response: ChatResponse,
245 /// Why the loop terminated.
246 pub termination_reason: TerminationReason,
247 /// Total iterations performed.
248 pub iterations: u32,
249 /// Accumulated usage.
250 pub total_usage: Usage,
251}
252
253/// Terminal: the loop errored.
254pub struct TurnError {
255 /// The error.
256 pub error: LlmError,
257 /// Iterations completed before the error.
258 pub iterations: u32,
259 /// Usage accumulated before the error.
260 pub total_usage: Usage,
261}
262
263/// Commands sent by the caller to control the resumable loop.
264///
265/// Passed to [`Yielded::resume()`] after receiving a
266/// [`TurnResult::Yielded`].
267#[derive(Debug)]
268pub enum LoopCommand {
269 /// Continue to the next LLM iteration normally.
270 Continue,
271
272 /// Inject additional messages before the next LLM call.
273 ///
274 /// The injected messages are appended after the tool results from
275 /// the current round. Use this to provide additional context
276 /// (e.g., worker agent results, user follow-ups).
277 InjectMessages(Vec<ChatMessage>),
278
279 /// Stop the loop immediately.
280 ///
281 /// Returns a `Completed` event with `TerminationReason::StopCondition`.
282 Stop(Option<String>),
283}
284
285// ── ToolLoopHandle ──────────────────────────────────────────────────
286
287/// Caller-driven resumable tool loop.
288///
289/// Unlike [`tool_loop`](super::tool_loop) which runs autonomously, this struct
290/// gives the caller control between each tool execution round. Call
291/// [`next_turn()`](Self::next_turn) to advance the loop, inspect the result,
292/// then consume the [`Yielded`] handle to control what happens next.
293///
294/// # No spawning required
295///
296/// This is a direct state machine — no background tasks, no channels. The
297/// caller drives it by calling `next_turn()` which performs one iteration
298/// (LLM call + tool execution) and returns.
299///
300/// # Lifecycle
301///
302/// 1. Create with [`new()`](Self::new)
303/// 2. Call [`next_turn()`](Self::next_turn) to get the first result
304/// 3. If `Yielded`, consume via `resume()` / `continue_loop()` / etc., then
305/// call `next_turn()` again
306/// 4. Repeat until `Completed` or `Error`
307/// 5. Optionally call [`into_result()`](Self::into_result) for a `ToolLoopResult`
308pub struct ToolLoopHandle<'a, Ctx: LoopDepth + Send + Sync + 'static> {
309 provider: &'a dyn DynProvider,
310 registry: &'a ToolRegistry<Ctx>,
311 core: LoopCore<Ctx>,
312}
313
314impl<'a, Ctx: LoopDepth + Send + Sync + 'static> ToolLoopHandle<'a, Ctx> {
315 /// Create a new resumable tool loop.
316 ///
317 /// Does not start execution — call [`next_turn()`](Self::next_turn) to
318 /// begin the first iteration.
319 ///
320 /// # Depth Tracking
321 ///
322 /// Same as [`tool_loop`](super::tool_loop) — if `Ctx` implements [`LoopDepth`],
323 /// nested calls are tracked and `max_depth` is enforced. If the depth limit
324 /// is already exceeded, the first call to `next_turn()` returns `Error`.
325 pub fn new(
326 provider: &'a dyn DynProvider,
327 registry: &'a ToolRegistry<Ctx>,
328 params: ChatParams,
329 config: ToolLoopConfig,
330 ctx: &Ctx,
331 ) -> Self {
332 Self {
333 provider,
334 registry,
335 core: LoopCore::new(params, config, ctx),
336 }
337 }
338
339 /// Advance the loop and return the result of this turn.
340 ///
341 /// Each call performs one iteration: LLM generation, tool execution (if
342 /// applicable), and returns the result.
343 ///
344 /// Returns a [`TurnResult`] that must be matched:
345 /// - [`TurnResult::Yielded`] — tools ran, consume via `resume()` /
346 /// `continue_loop()` / `inject_and_continue()` / `stop()` to continue
347 /// - [`TurnResult::Completed`] — loop is done, read `.response`
348 /// - [`TurnResult::Error`] — loop failed, read `.error`
349 ///
350 /// After `Completed` or `Error`, all subsequent calls return the same
351 /// terminal result.
352 pub async fn next_turn(&mut self) -> TurnResult<'a, '_, Ctx> {
353 let outcome = self.core.do_iteration(self.provider, self.registry).await;
354 outcome_to_turn_result!(outcome, self, TurnResult, Yielded)
355 }
356
357 /// Tell the loop how to proceed before the next [`next_turn()`](Self::next_turn) call.
358 ///
359 /// When using [`TurnResult::Yielded`], prefer the convenience methods on
360 /// [`Yielded`] (`continue_loop()`, `inject_and_continue()`, `stop()`),
361 /// which consume the yielded handle and call this internally.
362 ///
363 /// This method is useful when you need to set a command on the handle
364 /// directly — for example, when driving the handle from an external
365 /// event loop that receives the command asynchronously after the
366 /// `Yielded` has already been consumed.
367 ///
368 /// Has no effect after `Completed` or `Error`.
369 pub fn resume(&mut self, command: LoopCommand) {
370 self.core.resume(command);
371 }
372
373 /// Get a snapshot of the current conversation messages.
374 ///
375 /// Useful for context window management or debugging.
376 pub fn messages(&self) -> &[ChatMessage] {
377 self.core.messages()
378 }
379
380 /// Get a mutable reference to the conversation messages.
381 ///
382 /// Allows direct manipulation of the message history between iterations
383 /// (e.g., for context compaction/summarization).
384 pub fn messages_mut(&mut self) -> &mut Vec<ChatMessage> {
385 self.core.messages_mut()
386 }
387
388 /// Get the accumulated usage across all iterations so far.
389 pub fn total_usage(&self) -> &Usage {
390 self.core.total_usage()
391 }
392
393 /// Get the current iteration count.
394 pub fn iterations(&self) -> u32 {
395 self.core.iterations()
396 }
397
398 /// Whether the loop has finished (returned Completed or Error).
399 pub fn is_finished(&self) -> bool {
400 self.core.is_finished()
401 }
402
403 /// Consume the handle and return a `ToolLoopResult`.
404 ///
405 /// If the loop hasn't completed yet, returns a result with current
406 /// iteration count and `TerminationReason::Complete`.
407 pub fn into_result(self) -> ToolLoopResult {
408 self.core.into_result()
409 }
410
411 /// Convert this borrowed handle into an owned handle.
412 ///
413 /// The provider and registry must be provided as `Arc` since this
414 /// handle only holds references. The loop state (iterations, messages,
415 /// usage, etc.) is transferred as-is.
416 pub fn into_owned(
417 self,
418 provider: std::sync::Arc<dyn DynProvider>,
419 registry: std::sync::Arc<ToolRegistry<Ctx>>,
420 ) -> super::OwnedToolLoopHandle<Ctx> {
421 super::OwnedToolLoopHandle::from_core(provider, registry, self.core)
422 }
423}
424
425impl<Ctx: LoopDepth + Send + Sync + 'static> std::fmt::Debug for ToolLoopHandle<'_, Ctx> {
426 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
427 f.debug_struct("ToolLoopHandle")
428 .field("core", &self.core)
429 .finish_non_exhaustive()
430 }
431}