feat: add provenance-aware worker history
This commit is contained in:
@@ -21,20 +21,22 @@ agen = { version = "0.2.1", features = ["codex"] }
|
||||
|
||||
## Quick start
|
||||
|
||||
Supply an implementation of [`LlmClient`](https://docs.rs/agen/latest/agen/llm_client/trait.LlmClient.html), then run a turn. The first call consumes the mutable engine and returns a cache-locked engine for later turns.
|
||||
Supply an implementation of [`LlmClient`](https://docs.rs/agen/latest/agen/llm_client/trait.LlmClient.html), keep conversation history in your application, then run a turn. The first call consumes the mutable engine and returns a cache-locked engine for later turns.
|
||||
|
||||
```no_run
|
||||
use agen::Engine;
|
||||
use agen::{Engine, EngineError, History};
|
||||
use agen::llm_client::LlmClient;
|
||||
|
||||
async fn conversation<C: LlmClient>(client: C) {
|
||||
async fn conversation<C: LlmClient>(client: C) -> Result<(), EngineError> {
|
||||
let mut history = History::new();
|
||||
let output = Engine::new(client)
|
||||
.system_prompt("You are a concise assistant.")
|
||||
.run("Explain typed state in one sentence.")
|
||||
.await;
|
||||
.run(&mut history, "Explain typed state in one sentence.")
|
||||
.await?;
|
||||
|
||||
let mut engine = output.engine;
|
||||
let _exit = engine.run("Give a Rust example.").await;
|
||||
let _result = engine.run(&mut history, "Give a Rust example.").await?;
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
|
||||
use agen::llm_client::scheme::{Scheme, anthropic::AnthropicScheme};
|
||||
use agen::llm_client::transport::{HttpTransport, ResolvedAuth};
|
||||
use agen::{Engine, EngineRunExit, StopReason};
|
||||
use agen::{Engine, EngineResult, History};
|
||||
use std::time::Duration;
|
||||
|
||||
#[tokio::main]
|
||||
@@ -29,6 +29,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let base_url = scheme.default_base_url().to_string();
|
||||
let client = HttpTransport::new(scheme, model, base_url, ResolvedAuth::ApiKey(api_key), cap);
|
||||
let engine = Engine::new(client);
|
||||
let mut history = History::new();
|
||||
|
||||
println!("🚀 Starting Engine...");
|
||||
println!("💡 Will cancel after 2 seconds\n");
|
||||
@@ -45,13 +46,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
|
||||
println!("📡 Sending request to LLM...");
|
||||
|
||||
let output = engine.run("Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await;
|
||||
match output.result {
|
||||
EngineRunExit::Finished => println!("✅ Task completed normally"),
|
||||
EngineRunExit::Paused => println!("⏸️ Task paused"),
|
||||
EngineRunExit::Yielded => println!("↩️ Task yielded"),
|
||||
EngineRunExit::Interrupted(StopReason::LimitReached) => {
|
||||
println!("🔒 Turn limit reached")
|
||||
match engine.run(&mut history, "Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await {
|
||||
Ok(out) => match out.result {
|
||||
EngineResult::Finished => println!("✅ Task completed normally"),
|
||||
EngineResult::Paused => println!("⏸️ Task paused"),
|
||||
EngineResult::LimitReached => println!("🔒 Turn limit reached"),
|
||||
EngineResult::Yielded => println!("↩️ Task yielded"),
|
||||
},
|
||||
Err(e) => {
|
||||
println!("❌ Task error: {}", e);
|
||||
}
|
||||
EngineRunExit::Interrupted(reason) => println!("❌ Task interrupted: {reason:?}"),
|
||||
}
|
||||
|
||||
@@ -39,7 +39,7 @@ use tracing::info;
|
||||
use tracing_subscriber::EnvFilter;
|
||||
|
||||
use agen::{
|
||||
Engine, EngineRunExit, StopReason,
|
||||
Engine, History,
|
||||
interceptor::{Interceptor, PostToolAction, ToolResultInfo},
|
||||
llm_client::{
|
||||
LlmClient,
|
||||
@@ -474,11 +474,16 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
|
||||
engine.set_interceptor(ToolResultPrinterPolicy::new(tool_call_names));
|
||||
|
||||
let mut history = History::new();
|
||||
|
||||
// One-shot mode
|
||||
if let Some(prompt) = args.prompt {
|
||||
let output = engine.run(&prompt).await;
|
||||
if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) = output.result {
|
||||
eprintln!("\n❌ Error: {error}");
|
||||
match engine.run(&mut history, &prompt).await {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
eprintln!("\n❌ Error: {}", e);
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
return Ok(());
|
||||
@@ -497,8 +502,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let output = engine.run(first_input).await;
|
||||
let mut locked = output.engine;
|
||||
let mut locked = match engine.run(&mut history, first_input).await {
|
||||
Ok(out) => out.engine,
|
||||
Err(e) => {
|
||||
eprintln!("\n❌ Error: {}", e);
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
|
||||
loop {
|
||||
print!("\n👤 You: ");
|
||||
@@ -517,8 +527,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
break;
|
||||
}
|
||||
|
||||
if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) = locked.run(input).await {
|
||||
eprintln!("\n❌ Error: {error}");
|
||||
match locked.run(&mut history, input).await {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
eprintln!("\n❌ Error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+200
-94
@@ -7,7 +7,7 @@ use tokio::sync::mpsc;
|
||||
use tracing::{debug, info, trace, warn};
|
||||
|
||||
use crate::{
|
||||
Item,
|
||||
History, HistoryEntry, Item,
|
||||
callback::{
|
||||
ClosureMetaHandler, ClosureTextBlockHandler, ClosureThinkingBlockHandler,
|
||||
ClosureToolUseBlockHandler, TextBlockScope, ThinkingBlockScope, ToolUseBlockScope,
|
||||
@@ -117,9 +117,9 @@ impl From<Result<EngineResult, EngineError>> for EngineRunExit {
|
||||
/// Result of [`Engine::run`] or [`Engine::resume`].
|
||||
///
|
||||
/// Contains the `Locked` Engine (ready for subsequent runs) and the outcome.
|
||||
pub struct EngineRunOutput<C: LlmClient> {
|
||||
pub struct EngineRunOutput<C: LlmClient, A = ()> {
|
||||
/// The Engine, now in Locked state.
|
||||
pub engine: Engine<C, Locked>,
|
||||
pub engine: Engine<C, Locked, A>,
|
||||
/// Outcome of the turn.
|
||||
pub result: EngineRunExit,
|
||||
}
|
||||
@@ -139,29 +139,31 @@ const MAX_STREAM_CONTINUATIONS: u32 = 3;
|
||||
///
|
||||
/// # State Transitions (Type-state)
|
||||
///
|
||||
/// - [`Mutable`]: Initial state. System prompt, history, and tools can be freely edited.
|
||||
/// - [`Mutable`]: Initial state. System prompt and tools can be edited; history is caller-owned.
|
||||
/// - [`Locked`]: Cache-protected state. Prefix context is immutable; only `run()` / `resume()` are available.
|
||||
///
|
||||
/// Calling `run()` on a `Mutable` Engine consumes it and returns a
|
||||
/// `Locked` Engine together with the result. This ensures the
|
||||
/// cache prefix is fixed for optimal KV cache hit rate.
|
||||
/// `Locked` Engine together with the result. The engine borrows the caller's
|
||||
/// [`History`](crate::History) only while running, so host annotations stay with
|
||||
/// the host-owned history and are never projected to providers.
|
||||
///
|
||||
/// ```ignore
|
||||
/// let mut history = History::new();
|
||||
/// let mut engine = Engine::new(client)
|
||||
/// .system_prompt("You are a helpful assistant.");
|
||||
/// engine.register_tool(my_tool);
|
||||
///
|
||||
/// // Mutable::run() consumes self → EngineRunOutput { engine: Locked, result }
|
||||
/// let out = engine.run("Hello").await;
|
||||
/// let out = engine.run(&mut history, "Hello").await?;
|
||||
/// let mut engine = out.engine;
|
||||
///
|
||||
/// // Locked::run() borrows &mut self
|
||||
/// let _exit = engine.run("Follow-up").await;
|
||||
/// engine.run(&mut history, "Follow-up").await?;
|
||||
///
|
||||
/// // To edit between turns, unlock back to Mutable
|
||||
/// let mut engine = engine.unlock();
|
||||
/// engine.truncate_history(5);
|
||||
/// let out = engine.run("Continue").await;
|
||||
/// history.truncate(5);
|
||||
/// let out = engine.run(&mut history, "Continue").await?;
|
||||
/// let mut engine = out.engine;
|
||||
/// ```
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -181,7 +183,7 @@ enum StreamCompletion {
|
||||
Interrupted { reason: String },
|
||||
}
|
||||
|
||||
pub struct Engine<C: LlmClient, S: EngineState = Mutable> {
|
||||
pub struct Engine<C: LlmClient, S: EngineState = Mutable, A = ()> {
|
||||
/// LLM client
|
||||
client: C,
|
||||
/// Retry policy for opening an LLM response stream.
|
||||
@@ -201,8 +203,6 @@ pub struct Engine<C: LlmClient, S: EngineState = Mutable> {
|
||||
interceptor: Box<dyn Interceptor>,
|
||||
/// System prompt
|
||||
system_prompt: Option<String>,
|
||||
/// Item history (owned by Engine)
|
||||
history: Vec<Item>,
|
||||
/// History length at lock time (only meaningful in Locked state)
|
||||
locked_prefix_len: usize,
|
||||
/// AgentTurn count across the lifetime of this Engine.
|
||||
@@ -290,10 +290,10 @@ pub struct Engine<C: LlmClient, S: EngineState = Mutable> {
|
||||
/// stable conversation identifier when the backend benefits from one.
|
||||
cache_key: Option<String>,
|
||||
/// State marker
|
||||
_state: PhantomData<S>,
|
||||
_state: PhantomData<(S, A)>,
|
||||
}
|
||||
|
||||
impl<C: LlmClient, S: EngineState> Engine<C, S> {
|
||||
impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
|
||||
fn start_logical_run(&mut self) {
|
||||
self.active_run_turn_count = Some(0);
|
||||
}
|
||||
@@ -559,11 +559,15 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
|
||||
|
||||
fn append_history_items(
|
||||
&mut self,
|
||||
history: &mut History<A>,
|
||||
items: impl IntoIterator<Item = Item>,
|
||||
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||
) -> Result<(), EngineError> {
|
||||
for item in items {
|
||||
self.emit_history_append(&item)?;
|
||||
self.history.push(item);
|
||||
history
|
||||
.append_with(item, annotate)
|
||||
.map_err(EngineError::HistoryAppend)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -670,9 +674,9 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
|
||||
&self.client
|
||||
}
|
||||
|
||||
/// Get a reference to the history
|
||||
pub fn history(&self) -> &[Item] {
|
||||
&self.history
|
||||
/// Borrow caller-owned annotated history entries.
|
||||
pub fn history<'h>(&self, history: &'h History<A>) -> &'h [HistoryEntry<A>] {
|
||||
history.entries()
|
||||
}
|
||||
|
||||
/// Get a reference to the system prompt
|
||||
@@ -929,20 +933,20 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
|
||||
}
|
||||
|
||||
/// Check for pending tool calls (for resuming from Pause)
|
||||
fn get_pending_tool_calls(&self) -> Option<Vec<ToolCall>> {
|
||||
fn get_pending_tool_calls(&self, history: &History<A>) -> Option<Vec<ToolCall>> {
|
||||
// Find the last ToolCall items that don't have corresponding ToolResult
|
||||
let mut pending_calls = Vec::new();
|
||||
let mut answered_call_ids = std::collections::HashSet::new();
|
||||
|
||||
// First pass: collect all answered call IDs
|
||||
for item in &self.history {
|
||||
for item in history.items() {
|
||||
if let Item::ToolResult { call_id, .. } = item {
|
||||
answered_call_ids.insert(call_id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// Second pass: find unanswered tool calls
|
||||
for item in &self.history {
|
||||
for item in history.items() {
|
||||
if let Item::ToolCall {
|
||||
call_id,
|
||||
name,
|
||||
@@ -1142,19 +1146,26 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
|
||||
}
|
||||
|
||||
/// Internal turn execution logic
|
||||
async fn run_turn_loop(&mut self) -> Result<EngineResult, EngineError> {
|
||||
async fn run_turn_loop(
|
||||
&mut self,
|
||||
history: &mut History<A>,
|
||||
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||
) -> Result<EngineResult, EngineError> {
|
||||
let tool_definitions = self.build_tool_definitions();
|
||||
|
||||
info!(
|
||||
item_count = self.history.len(),
|
||||
item_count = history.len(),
|
||||
tool_count = tool_definitions.len(),
|
||||
"Starting engine run"
|
||||
);
|
||||
|
||||
// Resume pending tool calls from a previous Pause
|
||||
if let Some(tool_calls) = self.get_pending_tool_calls() {
|
||||
if let Some(tool_calls) = self.get_pending_tool_calls(history) {
|
||||
info!("Resuming pending tool calls");
|
||||
if let Some(result) = self.execute_and_commit_tools(tool_calls).await? {
|
||||
if let Some(result) = self
|
||||
.execute_and_commit_tools(history, annotate, tool_calls)
|
||||
.await?
|
||||
{
|
||||
return Ok(result);
|
||||
}
|
||||
}
|
||||
@@ -1199,13 +1210,13 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
|
||||
.await
|
||||
.map_err(EngineError::HistoryAppend)?;
|
||||
if !pending.is_empty() {
|
||||
self.append_history_items(pending)?;
|
||||
self.append_history_items(history, pending, annotate)?;
|
||||
}
|
||||
|
||||
// Clone the history into a per-request context. Everything
|
||||
// below (prune projection, interceptor hooks) mutates only
|
||||
// this clone, so the persistent `self.history` stays intact.
|
||||
let mut request_context = self.history.clone();
|
||||
// this clone, so the caller-owned `history` stays intact.
|
||||
let mut request_context = history.items_cloned();
|
||||
|
||||
// Prune projection: if both the config and the savings
|
||||
// estimator are configured, drop ToolResult.content from
|
||||
@@ -1273,7 +1284,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
|
||||
return Err(EngineError::Aborted(reason));
|
||||
}
|
||||
PreRequestAction::YieldWith(items) => {
|
||||
self.append_history_items(items.clone())?;
|
||||
self.append_history_items(history, items.clone(), annotate)?;
|
||||
request_context.extend(items);
|
||||
info!("Yielded by interceptor after pre-request history append");
|
||||
for cb in &self.turn_end_cbs {
|
||||
@@ -1289,7 +1300,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
|
||||
return Ok(EngineResult::Yielded);
|
||||
}
|
||||
PreRequestAction::ContinueWith(items) => {
|
||||
self.append_history_items(items.clone())?;
|
||||
self.append_history_items(history, items.clone(), annotate)?;
|
||||
request_context.extend(items);
|
||||
}
|
||||
PreRequestAction::Continue => {}
|
||||
@@ -1348,7 +1359,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
|
||||
let assistant_items =
|
||||
self.build_assistant_items(&reasoning_items, &text_blocks, &[]);
|
||||
if !assistant_items.is_empty() {
|
||||
self.append_history_items(assistant_items)?;
|
||||
self.append_history_items(history, assistant_items, annotate)?;
|
||||
}
|
||||
self.emit_llm_continuation(
|
||||
current_llm_call,
|
||||
@@ -1376,15 +1387,16 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
|
||||
let tool_calls = self.tool_call_collector.take_collected();
|
||||
let assistant_items =
|
||||
self.build_assistant_items(&reasoning_items, &text_blocks, &tool_calls);
|
||||
self.append_history_items(assistant_items)?;
|
||||
self.append_history_items(history, assistant_items, annotate)?;
|
||||
|
||||
if tool_calls.is_empty() {
|
||||
match self.interceptor.on_turn_end(&self.history).await {
|
||||
let turn_end_context = history.items_cloned();
|
||||
match self.interceptor.on_turn_end(&turn_end_context).await {
|
||||
TurnEndAction::Finish => {
|
||||
return Ok(EngineResult::Finished);
|
||||
}
|
||||
TurnEndAction::ContinueWithMessages(additional) => {
|
||||
self.append_history_items(additional)?;
|
||||
self.append_history_items(history, additional, annotate)?;
|
||||
continue;
|
||||
}
|
||||
TurnEndAction::Pause => {
|
||||
@@ -1393,7 +1405,10 @@ 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(history, annotate, tool_calls)
|
||||
.await?
|
||||
{
|
||||
return Ok(result);
|
||||
}
|
||||
}
|
||||
@@ -1639,6 +1654,8 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
|
||||
/// `None` if the turn loop should continue.
|
||||
async fn execute_and_commit_tools(
|
||||
&mut self,
|
||||
history: &mut History<A>,
|
||||
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||
tool_calls: Vec<ToolCall>,
|
||||
) -> Result<Option<EngineResult>, EngineError> {
|
||||
match self.execute_tools(tool_calls).await {
|
||||
@@ -1655,7 +1672,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
|
||||
result.attachments,
|
||||
)
|
||||
});
|
||||
self.append_history_items(items)?;
|
||||
self.append_history_items(history, items, annotate)?;
|
||||
Ok(None)
|
||||
}
|
||||
Err(err) => Err(err),
|
||||
@@ -1663,9 +1680,9 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: LlmClient> Engine<C, Mutable> {
|
||||
/// Create a new Engine (in Mutable state)
|
||||
pub fn new(client: C) -> Self {
|
||||
impl<C: LlmClient, A> Engine<C, Mutable, A> {
|
||||
/// Create a new annotated Engine (in Mutable state).
|
||||
pub fn new_annotated(client: C) -> Self {
|
||||
let text_block_collector = TextBlockCollector::new();
|
||||
let tool_call_collector = ToolCallCollector::new();
|
||||
let thinking_block_collector = ThinkingBlockCollector::new();
|
||||
@@ -1687,7 +1704,6 @@ impl<C: LlmClient> Engine<C, Mutable> {
|
||||
tool_server: ToolServer::new().handle(),
|
||||
interceptor: Box::new(DefaultInterceptor),
|
||||
system_prompt: None,
|
||||
history: Vec::new(),
|
||||
locked_prefix_len: 0,
|
||||
turn_count: 0,
|
||||
active_run_turn_count: None,
|
||||
@@ -1847,36 +1863,38 @@ impl<C: LlmClient> Engine<C, Mutable> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Replace history during restore/rebuild without emitting append callbacks.
|
||||
/// Replace caller-owned history during restore/rebuild without emitting append callbacks.
|
||||
///
|
||||
/// This is not a history-growth API. Live append paths must use
|
||||
/// [`append_history`](Self::append_history) so `on_history_append` observers
|
||||
/// see every inserted item.
|
||||
pub fn set_history(&mut self, items: Vec<Item>) {
|
||||
self.history = items;
|
||||
}
|
||||
|
||||
/// Append items to history after every history-append observer accepts the
|
||||
/// item. This is the only public Mutable-state API for growing engine
|
||||
/// history; callers that need session-log persistence must install
|
||||
/// [`on_history_append`](Self::on_history_append) before calling it.
|
||||
pub fn append_history(
|
||||
/// [`append_history_with`](Self::append_history_with) so observers and the
|
||||
/// trusted annotation callback see every inserted item.
|
||||
pub fn replace_history_entries(
|
||||
&mut self,
|
||||
history: &mut History<A>,
|
||||
entries: Vec<HistoryEntry<A>>,
|
||||
) -> Vec<HistoryEntry<A>> {
|
||||
history.replace_entries(entries)
|
||||
}
|
||||
|
||||
/// Append items to caller-owned history after every observer and the trusted
|
||||
/// annotation callback accepts the item.
|
||||
pub fn append_history_with(
|
||||
&mut self,
|
||||
history: &mut History<A>,
|
||||
items: impl IntoIterator<Item = Item>,
|
||||
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||
) -> Result<(), EngineError> {
|
||||
self.append_history_items(items)
|
||||
self.append_history_items(history, items, annotate)
|
||||
}
|
||||
|
||||
/// Truncate history without emitting append callbacks.
|
||||
///
|
||||
/// This is an edit operation, not a history-growth path.
|
||||
pub fn truncate_history(&mut self, len: usize) {
|
||||
self.history.truncate(len);
|
||||
/// Truncate caller-owned history without emitting append callbacks.
|
||||
pub fn truncate_history(&mut self, history: &mut History<A>, len: usize) {
|
||||
history.truncate(len);
|
||||
}
|
||||
|
||||
/// Clear history
|
||||
pub fn clear_history(&mut self) {
|
||||
self.history.clear();
|
||||
/// Clear caller-owned history.
|
||||
pub fn clear_history(&mut self, history: &mut History<A>) {
|
||||
history.clear();
|
||||
}
|
||||
|
||||
/// Set the turn count (for session restoration)
|
||||
@@ -1895,16 +1913,21 @@ impl<C: LlmClient> Engine<C, Mutable> {
|
||||
self
|
||||
}
|
||||
|
||||
/// Execute a turn, consuming self and transitioning to Locked.
|
||||
/// Run the engine with one user input, appending to caller-owned history.
|
||||
///
|
||||
/// This is the primary entry point for first use. Equivalent to
|
||||
/// `self.lock()` followed by `locked.run(user_input)`.
|
||||
///
|
||||
/// Subsequent runs can call [`Engine::run`] directly.
|
||||
/// To edit state between turns, call [`unlock()`](Engine::unlock) first.
|
||||
pub async fn run(self, user_input: impl Into<String>) -> EngineRunOutput<C> {
|
||||
let mut locked = self.lock();
|
||||
let result = locked.run(user_input).await;
|
||||
/// The trusted `annotate` callback is invoked after append observers and before
|
||||
/// each new item becomes live in `history`. Providers, token counters, pruners,
|
||||
/// and interceptors receive only the `Item` projection.
|
||||
pub async fn run_with_annotation(
|
||||
self,
|
||||
history: &mut History<A>,
|
||||
user_input: impl Into<String>,
|
||||
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||
) -> EngineRunOutput<C, A> {
|
||||
let mut locked = self.lock(history);
|
||||
let result = locked
|
||||
.run_with_annotation(history, user_input, annotate)
|
||||
.await;
|
||||
EngineRunOutput {
|
||||
engine: locked,
|
||||
result,
|
||||
@@ -1914,9 +1937,13 @@ impl<C: LlmClient> Engine<C, Mutable> {
|
||||
/// Resume from Paused, consuming self and transitioning to Locked.
|
||||
///
|
||||
/// Used after `unlock()` → edit → resume.
|
||||
pub async fn resume(self) -> EngineRunOutput<C> {
|
||||
let mut locked = self.lock();
|
||||
let result = locked.resume().await;
|
||||
pub async fn resume_with_annotation(
|
||||
self,
|
||||
history: &mut History<A>,
|
||||
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||
) -> EngineRunOutput<C, A> {
|
||||
let mut locked = self.lock(history);
|
||||
let result = locked.resume_with_annotation(history, annotate).await;
|
||||
EngineRunOutput {
|
||||
engine: locked,
|
||||
result,
|
||||
@@ -1936,9 +1963,9 @@ impl<C: LlmClient> Engine<C, Mutable> {
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if a pending tool factory produces a duplicate name.
|
||||
pub fn lock(self) -> Engine<C, Locked> {
|
||||
pub fn lock(self, history: &History<A>) -> Engine<C, Locked, A> {
|
||||
self.tool_server.flush_pending();
|
||||
let locked_prefix_len = self.history.len();
|
||||
let locked_prefix_len = history.len();
|
||||
Engine {
|
||||
client: self.client,
|
||||
retry_policy: self.retry_policy,
|
||||
@@ -1949,7 +1976,6 @@ impl<C: LlmClient> Engine<C, Mutable> {
|
||||
tool_server: self.tool_server,
|
||||
interceptor: self.interceptor,
|
||||
system_prompt: self.system_prompt,
|
||||
history: self.history,
|
||||
locked_prefix_len,
|
||||
turn_count: self.turn_count,
|
||||
active_run_turn_count: self.active_run_turn_count,
|
||||
@@ -1983,19 +2009,73 @@ impl<C: LlmClient> Engine<C, Mutable> {
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: LlmClient> Engine<C, Locked> {
|
||||
fn unit_history_annotation(_: &Item) -> Result<(), String> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
impl<C: LlmClient> Engine<C, Mutable, ()> {
|
||||
/// Create a new Engine (in Mutable state) using unit history annotations.
|
||||
pub fn new(client: C) -> Self {
|
||||
Self::new_annotated(client)
|
||||
}
|
||||
|
||||
/// Append unit-annotated items to caller-owned history.
|
||||
pub fn append_history(
|
||||
&mut self,
|
||||
history: &mut History<()>,
|
||||
items: impl IntoIterator<Item = Item>,
|
||||
) -> Result<(), EngineError> {
|
||||
let mut annotate = unit_history_annotation;
|
||||
self.append_history_items(history, items, &mut annotate)
|
||||
}
|
||||
|
||||
/// Replace unit-annotated history from plain items.
|
||||
pub fn set_history(&mut self, history: &mut History<()>, items: Vec<Item>) {
|
||||
history.replace_items(items);
|
||||
}
|
||||
|
||||
/// Run using unit annotations.
|
||||
pub async fn run(
|
||||
self,
|
||||
history: &mut History<()>,
|
||||
user_input: impl Into<String>,
|
||||
) -> EngineRunOutput<C> {
|
||||
let mut annotate = unit_history_annotation;
|
||||
self.run_with_annotation(history, user_input, &mut annotate)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Resume using unit annotations.
|
||||
pub async fn resume(self, history: &mut History<()>) -> EngineRunOutput<C> {
|
||||
let mut annotate = unit_history_annotation;
|
||||
self.resume_with_annotation(history, &mut annotate).await
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: LlmClient, A> Engine<C, Locked, A> {
|
||||
/// Execute a turn
|
||||
///
|
||||
/// Adds a new user message to history and sends a request to the LLM.
|
||||
/// Automatically loops if there are tool calls.
|
||||
pub async fn run(&mut self, user_input: impl Into<String>) -> EngineRunExit {
|
||||
self.run_result(user_input.into()).await.into()
|
||||
pub async fn run_with_annotation(
|
||||
&mut self,
|
||||
history: &mut History<A>,
|
||||
user_input: impl Into<String>,
|
||||
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||
) -> EngineRunExit {
|
||||
self.run_result_with_annotation(history, user_input.into(), annotate)
|
||||
.await
|
||||
.into()
|
||||
}
|
||||
|
||||
async fn run_result(&mut self, user_input: String) -> Result<EngineResult, EngineError> {
|
||||
async fn run_result_with_annotation(
|
||||
&mut self,
|
||||
history: &mut History<A>,
|
||||
user_input: String,
|
||||
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||
) -> Result<EngineResult, EngineError> {
|
||||
// Supplying new user input abandons any paused/yielded logical run.
|
||||
self.active_run_turn_count = None;
|
||||
// Interceptor: on_prompt_submit
|
||||
let mut user_item = Item::user_message(user_input);
|
||||
let extras = match self.interceptor.on_prompt_submit(&mut user_item).await {
|
||||
PromptAction::Cancel(reason) => {
|
||||
@@ -2006,27 +2086,35 @@ impl<C: LlmClient> Engine<C, Locked> {
|
||||
PromptAction::Continue => Vec::new(),
|
||||
PromptAction::ContinueWith(items) => items,
|
||||
};
|
||||
self.append_history_items(std::iter::once(user_item))?;
|
||||
self.append_history_items(history, std::iter::once(user_item), annotate)?;
|
||||
if !extras.is_empty() {
|
||||
self.append_history_items(extras)?;
|
||||
self.append_history_items(history, extras, annotate)?;
|
||||
}
|
||||
self.start_logical_run();
|
||||
let result = self.run_turn_loop().await;
|
||||
let result = self.run_turn_loop(history, annotate).await;
|
||||
let result = self.finalize_interruption(result).await;
|
||||
self.finish_logical_run(&result);
|
||||
result
|
||||
}
|
||||
|
||||
/// Resume execution (from Paused state)
|
||||
///
|
||||
/// Resumes turn processing from current state without adding a new user message.
|
||||
pub async fn resume(&mut self) -> EngineRunExit {
|
||||
self.resume_result().await.into()
|
||||
/// Resume execution (from Paused state).
|
||||
pub async fn resume_with_annotation(
|
||||
&mut self,
|
||||
history: &mut History<A>,
|
||||
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||
) -> EngineRunExit {
|
||||
self.resume_result_with_annotation(history, annotate)
|
||||
.await
|
||||
.into()
|
||||
}
|
||||
|
||||
async fn resume_result(&mut self) -> Result<EngineResult, EngineError> {
|
||||
async fn resume_result_with_annotation(
|
||||
&mut self,
|
||||
history: &mut History<A>,
|
||||
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||
) -> Result<EngineResult, EngineError> {
|
||||
self.ensure_logical_run();
|
||||
let result = self.run_turn_loop().await;
|
||||
let result = self.run_turn_loop(history, annotate).await;
|
||||
let result = self.finalize_interruption(result).await;
|
||||
self.finish_logical_run(&result);
|
||||
result
|
||||
@@ -2041,7 +2129,7 @@ impl<C: LlmClient> Engine<C, Locked> {
|
||||
///
|
||||
/// Note: After this operation, subsequent requests may not hit the cache.
|
||||
/// Use only when you need to edit history.
|
||||
pub fn unlock(self) -> Engine<C, Mutable> {
|
||||
pub fn unlock(self) -> Engine<C, Mutable, A> {
|
||||
Engine {
|
||||
client: self.client,
|
||||
retry_policy: self.retry_policy,
|
||||
@@ -2052,7 +2140,6 @@ impl<C: LlmClient> Engine<C, Locked> {
|
||||
tool_server: self.tool_server,
|
||||
interceptor: self.interceptor,
|
||||
system_prompt: self.system_prompt,
|
||||
history: self.history,
|
||||
locked_prefix_len: 0,
|
||||
turn_count: self.turn_count,
|
||||
active_run_turn_count: self.active_run_turn_count,
|
||||
@@ -2086,6 +2173,25 @@ impl<C: LlmClient> Engine<C, Locked> {
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: LlmClient> Engine<C, Locked, ()> {
|
||||
/// Run another turn using unit annotations.
|
||||
pub async fn run(
|
||||
&mut self,
|
||||
history: &mut History<()>,
|
||||
user_input: impl Into<String>,
|
||||
) -> EngineRunExit {
|
||||
let mut annotate = unit_history_annotation;
|
||||
self.run_with_annotation(history, user_input, &mut annotate)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Resume using unit annotations.
|
||||
pub async fn resume(&mut self, history: &mut History<()>) -> EngineRunExit {
|
||||
let mut annotate = unit_history_annotation;
|
||||
self.resume_with_annotation(history, &mut annotate).await
|
||||
}
|
||||
}
|
||||
|
||||
enum FirstStreamEvent {
|
||||
Ready(ResponseStream),
|
||||
Empty(ResponseStream),
|
||||
|
||||
@@ -0,0 +1,199 @@
|
||||
//! Typed conversation history containers.
|
||||
//!
|
||||
//! Agen keeps provider-visible [`Item`](crate::Item) values separate from any
|
||||
//! host-domain provenance. The host chooses the annotation type `A`, while Agen
|
||||
//! preserves each item and annotation as one entry for clone/truncate/restore
|
||||
//! style history operations.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::Item;
|
||||
|
||||
/// One conversation-history entry with host-owned annotation.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct HistoryEntry<A = ()> {
|
||||
/// Provider/model-visible conversation item.
|
||||
pub item: Item,
|
||||
/// Host-domain metadata kept with the item and never projected to providers.
|
||||
pub annotation: A,
|
||||
}
|
||||
|
||||
impl<A> HistoryEntry<A> {
|
||||
/// Build an entry from an item and its annotation.
|
||||
pub fn new(item: Item, annotation: A) -> Self {
|
||||
Self { item, annotation }
|
||||
}
|
||||
|
||||
/// Split the entry into its item and annotation.
|
||||
pub fn into_parts(self) -> (Item, A) {
|
||||
(self.item, self.annotation)
|
||||
}
|
||||
}
|
||||
|
||||
impl HistoryEntry<()> {
|
||||
/// Build a unit-annotated entry.
|
||||
pub fn from_item(item: Item) -> Self {
|
||||
Self {
|
||||
item,
|
||||
annotation: (),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Conversation history with one annotation per item.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
|
||||
pub struct History<A = ()> {
|
||||
entries: Vec<HistoryEntry<A>>,
|
||||
}
|
||||
|
||||
impl<A> History<A> {
|
||||
/// Create an empty history.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
entries: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Build history from already annotated entries, preserving order.
|
||||
pub fn from_entries(entries: Vec<HistoryEntry<A>>) -> Self {
|
||||
Self { entries }
|
||||
}
|
||||
|
||||
/// Replace all entries as one restore/rebuild operation and return the old entries.
|
||||
pub fn replace_entries(&mut self, entries: Vec<HistoryEntry<A>>) -> Vec<HistoryEntry<A>> {
|
||||
std::mem::replace(&mut self.entries, entries)
|
||||
}
|
||||
|
||||
/// Borrow annotated entries.
|
||||
pub fn entries(&self) -> &[HistoryEntry<A>] {
|
||||
&self.entries
|
||||
}
|
||||
|
||||
/// Mutably borrow annotated entries for host-owned rebuild operations.
|
||||
pub fn entries_mut(&mut self) -> &mut [HistoryEntry<A>] {
|
||||
&mut self.entries
|
||||
}
|
||||
|
||||
/// Consume the history into annotated entries.
|
||||
pub fn into_entries(self) -> Vec<HistoryEntry<A>> {
|
||||
self.entries
|
||||
}
|
||||
|
||||
/// Number of entries.
|
||||
pub fn len(&self) -> usize {
|
||||
self.entries.len()
|
||||
}
|
||||
|
||||
/// Whether the history is empty.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.entries.is_empty()
|
||||
}
|
||||
|
||||
/// Iterate over annotated entries.
|
||||
pub fn iter(&self) -> impl ExactSizeIterator<Item = &HistoryEntry<A>> {
|
||||
self.entries.iter()
|
||||
}
|
||||
|
||||
/// Iterate over provider-visible items only.
|
||||
pub fn items(&self) -> impl ExactSizeIterator<Item = &Item> {
|
||||
self.entries.iter().map(|entry| &entry.item)
|
||||
}
|
||||
|
||||
/// Clone provider-visible items into a request-local projection.
|
||||
pub fn items_cloned(&self) -> Vec<Item> {
|
||||
self.items().cloned().collect()
|
||||
}
|
||||
|
||||
/// Append an already annotated entry.
|
||||
pub fn push_entry(&mut self, entry: HistoryEntry<A>) {
|
||||
self.entries.push(entry);
|
||||
}
|
||||
|
||||
/// Append many already annotated entries.
|
||||
pub fn extend_entries(&mut self, entries: impl IntoIterator<Item = HistoryEntry<A>>) {
|
||||
self.entries.extend(entries);
|
||||
}
|
||||
|
||||
/// Commit one item through a trusted annotation callback before it becomes live.
|
||||
///
|
||||
/// The callback may durably persist the item and returns the annotation that
|
||||
/// must be stored with it. If the callback fails, the history is left unchanged.
|
||||
pub fn append_with(
|
||||
&mut self,
|
||||
item: Item,
|
||||
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||
) -> Result<(), String> {
|
||||
let annotation = annotate(&item)?;
|
||||
self.entries.push(HistoryEntry { item, annotation });
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Commit items through a trusted annotation callback before they become live.
|
||||
///
|
||||
/// Items before a failure remain appended; the failing item and later items do
|
||||
/// not enter history. This mirrors append-only durable logs where each accepted
|
||||
/// item is already committed before the next item is attempted.
|
||||
pub fn extend_with(
|
||||
&mut self,
|
||||
items: impl IntoIterator<Item = Item>,
|
||||
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||
) -> Result<(), String> {
|
||||
for item in items {
|
||||
self.append_with(item, annotate)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Truncate entries, preserving item+annotation pairing for retained entries.
|
||||
pub fn truncate(&mut self, len: usize) {
|
||||
self.entries.truncate(len);
|
||||
}
|
||||
|
||||
/// Clear all entries.
|
||||
pub fn clear(&mut self) {
|
||||
self.entries.clear();
|
||||
}
|
||||
}
|
||||
|
||||
impl History<()> {
|
||||
/// Build unit-annotated history from provider-visible items.
|
||||
pub fn from_items(items: Vec<Item>) -> Self {
|
||||
Self {
|
||||
entries: items.into_iter().map(HistoryEntry::from_item).collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Replace history from provider-visible items using unit annotations.
|
||||
pub fn replace_items(&mut self, items: Vec<Item>) -> Vec<HistoryEntry<()>> {
|
||||
self.replace_entries(items.into_iter().map(HistoryEntry::from_item).collect())
|
||||
}
|
||||
|
||||
/// Append one item with unit annotation.
|
||||
pub fn push(&mut self, item: Item) {
|
||||
self.entries.push(HistoryEntry::from_item(item));
|
||||
}
|
||||
|
||||
/// Append items with unit annotations.
|
||||
pub fn extend_items(&mut self, items: impl IntoIterator<Item = Item>) {
|
||||
self.entries
|
||||
.extend(items.into_iter().map(HistoryEntry::from_item));
|
||||
}
|
||||
}
|
||||
|
||||
impl<A> IntoIterator for History<A> {
|
||||
type Item = HistoryEntry<A>;
|
||||
type IntoIter = std::vec::IntoIter<HistoryEntry<A>>;
|
||||
|
||||
fn into_iter(self) -> Self::IntoIter {
|
||||
self.entries.into_iter()
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, A> IntoIterator for &'a History<A> {
|
||||
type Item = &'a HistoryEntry<A>;
|
||||
type IntoIter = std::slice::Iter<'a, HistoryEntry<A>>;
|
||||
|
||||
fn into_iter(self) -> Self::IntoIter {
|
||||
self.entries.iter()
|
||||
}
|
||||
}
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
mod engine;
|
||||
mod handler;
|
||||
mod history;
|
||||
mod message;
|
||||
|
||||
pub(crate) mod callback;
|
||||
@@ -24,6 +25,7 @@ pub use engine::{
|
||||
LlmRetryNotice, StopReason, ToolRegistryError,
|
||||
};
|
||||
pub use handler::ToolUseBlockStart;
|
||||
pub use history::{History, HistoryEntry};
|
||||
pub use interceptor::Interceptor;
|
||||
pub use message::{ContentPart, Item, Message, Role};
|
||||
pub use tool::{ToolCall, ToolExecutionContext, ToolOutputLimits, ToolResult};
|
||||
|
||||
@@ -19,7 +19,7 @@ mod private {
|
||||
/// - Editing message history (add, delete, clear)
|
||||
/// - Registering tools and hooks
|
||||
///
|
||||
/// Can transition to [`Locked`] state via `Engine::lock()`.
|
||||
/// Can transition to [`Locked`] state via `Engine::lock(&history)`.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
mod common;
|
||||
|
||||
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
|
||||
use agen::{Engine, EngineError, History, HistoryEntry, Item, Role};
|
||||
use common::MockLlmClient;
|
||||
|
||||
fn completed_text_events(text: &str) -> Vec<Event> {
|
||||
vec![
|
||||
Event::text_block_start(0),
|
||||
Event::text_delta(0, text),
|
||||
Event::text_block_stop(0, None),
|
||||
Event::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}),
|
||||
]
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn run_preserves_item_annotations_without_projecting_them() {
|
||||
let client = MockLlmClient::new(completed_text_events("assistant reply"));
|
||||
let engine = Engine::<_, agen::state::Mutable, String>::new_annotated(client);
|
||||
let mut history = History::<String>::new();
|
||||
let mut next = 0usize;
|
||||
let mut annotate = |item: &Item| {
|
||||
next += 1;
|
||||
let kind = match item {
|
||||
Item::Message { role, .. } => match role {
|
||||
Role::User => "user",
|
||||
Role::Assistant => "assistant",
|
||||
Role::System => "system",
|
||||
},
|
||||
Item::ToolCall { .. } => "tool_call",
|
||||
Item::ToolResult { .. } => "tool_result",
|
||||
Item::Reasoning { .. } => "reasoning",
|
||||
};
|
||||
Ok(format!("{next}:{kind}"))
|
||||
};
|
||||
|
||||
let output = engine
|
||||
.run_with_annotation(&mut history, "hello", &mut annotate)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(output.result, agen::EngineResult::Finished));
|
||||
assert_eq!(history.len(), 2);
|
||||
assert_eq!(history.entries()[0].annotation, "1:user");
|
||||
assert_eq!(history.entries()[1].annotation, "2:assistant");
|
||||
assert_eq!(history.items_cloned().len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn append_failure_does_not_make_item_live() {
|
||||
let client = MockLlmClient::new(vec![]);
|
||||
let mut engine = Engine::<_, agen::state::Mutable, usize>::new_annotated(client);
|
||||
let mut history = History::<usize>::new();
|
||||
let mut fail = |_item: &Item| Err("commit failed".to_string());
|
||||
|
||||
let err = engine
|
||||
.append_history_with(&mut history, [Item::user_message("uncommitted")], &mut fail)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(err, EngineError::HistoryAppend(message) if message == "commit failed"));
|
||||
assert!(history.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replacement_keeps_items_and_annotations_together() {
|
||||
let mut history = History::from_entries(vec![
|
||||
HistoryEntry::new(Item::user_message("old"), "old-ann".to_string()),
|
||||
HistoryEntry::new(Item::user_message("second"), "second-ann".to_string()),
|
||||
]);
|
||||
|
||||
history.truncate(1);
|
||||
assert_eq!(history.entries()[0].item.as_text(), Some("old"));
|
||||
assert_eq!(history.entries()[0].annotation, "old-ann");
|
||||
|
||||
let previous = history.replace_entries(vec![HistoryEntry::new(
|
||||
Item::user_message("restored"),
|
||||
"restored-ann".to_string(),
|
||||
)]);
|
||||
|
||||
assert_eq!(previous.len(), 1);
|
||||
assert_eq!(history.entries()[0].item.as_text(), Some("restored"));
|
||||
assert_eq!(history.entries()[0].annotation, "restored-ann");
|
||||
}
|
||||
@@ -8,11 +8,11 @@ use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use agen::Engine;
|
||||
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent as ClientStatusEvent};
|
||||
use agen::llm_client::retry::RetryPolicy;
|
||||
use agen::llm_client::{ClientError, LlmClient, Request, ResponseStream};
|
||||
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
||||
use agen::{Engine, History};
|
||||
use async_trait::async_trait;
|
||||
use common::MockLlmClient;
|
||||
|
||||
@@ -58,6 +58,7 @@ async fn test_callback_llm_retry_event() {
|
||||
max_attempts: 2,
|
||||
total_timeout: Duration::from_secs(1),
|
||||
});
|
||||
let mut history: History = History::new();
|
||||
|
||||
let notices = Arc::new(Mutex::new(Vec::new()));
|
||||
let sink = notices.clone();
|
||||
@@ -65,11 +66,8 @@ async fn test_callback_llm_retry_event() {
|
||||
sink.lock().unwrap().push((llm_call, notice.clone()));
|
||||
});
|
||||
|
||||
let result = engine.run("retry once").await;
|
||||
assert!(
|
||||
matches!(result.result, agen::EngineRunExit::Finished),
|
||||
"engine should succeed after one retry"
|
||||
);
|
||||
let result = engine.run(&mut history, "retry once").await;
|
||||
assert!(result.is_ok(), "engine should succeed after one retry");
|
||||
|
||||
let notices = notices.lock().unwrap();
|
||||
assert_eq!(notices.len(), 1);
|
||||
@@ -94,6 +92,7 @@ async fn test_callback_text_block_events() {
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
let text_deltas = Arc::new(Mutex::new(Vec::new()));
|
||||
let text_completes = Arc::new(Mutex::new(Vec::new()));
|
||||
@@ -111,12 +110,9 @@ async fn test_callback_text_block_events() {
|
||||
});
|
||||
});
|
||||
|
||||
// Mutable::run consumes self, returns (Locked, EngineRunExit)
|
||||
let result = engine.run("Greet me").await;
|
||||
assert!(
|
||||
matches!(result.result, agen::EngineRunExit::Finished),
|
||||
"Engine should complete"
|
||||
);
|
||||
// Mutable::run consumes self, returns (Locked, EngineResult)
|
||||
let result = engine.run(&mut history, "Greet me").await;
|
||||
assert!(result.is_ok(), "Engine should complete");
|
||||
|
||||
let deltas = text_deltas.lock().unwrap();
|
||||
assert_eq!(deltas.len(), 2);
|
||||
@@ -143,6 +139,7 @@ async fn test_callback_tool_call_complete() {
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
let tool_starts = Arc::new(Mutex::new(Vec::<(String, String)>::new()));
|
||||
let tool_completes = Arc::new(Mutex::new(Vec::new()));
|
||||
@@ -160,8 +157,8 @@ async fn test_callback_tool_call_complete() {
|
||||
});
|
||||
});
|
||||
|
||||
// Mutable::run consumes self, returns (Locked, EngineRunExit)
|
||||
let _ = engine.run("Weather please").await;
|
||||
// Mutable::run consumes self, returns (Locked, EngineResult)
|
||||
let _ = engine.run(&mut history, "Weather please").await;
|
||||
|
||||
let starts = tool_starts.lock().unwrap();
|
||||
assert_eq!(starts.len(), 1);
|
||||
@@ -189,6 +186,7 @@ async fn test_callback_turn_events() {
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
let turn_starts = Arc::new(Mutex::new(Vec::new()));
|
||||
let turn_ends = Arc::new(Mutex::new(Vec::new()));
|
||||
@@ -203,9 +201,9 @@ async fn test_callback_turn_events() {
|
||||
ends.lock().unwrap().push(turn);
|
||||
});
|
||||
|
||||
// Mutable::run consumes self, returns (Locked, EngineRunExit)
|
||||
let result = engine.run("Do something").await;
|
||||
assert!(matches!(result.result, agen::EngineRunExit::Finished));
|
||||
// Mutable::run consumes self, returns (Locked, EngineResult)
|
||||
let result = engine.run(&mut history, "Do something").await;
|
||||
assert!(result.is_ok());
|
||||
|
||||
let starts = turn_starts.lock().unwrap();
|
||||
let ends = turn_ends.lock().unwrap();
|
||||
@@ -260,6 +258,7 @@ async fn test_callback_tool_result_events() {
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
engine.register_tool(fixed_tool(
|
||||
"fixed",
|
||||
@@ -282,7 +281,7 @@ async fn test_callback_tool_result_events() {
|
||||
));
|
||||
});
|
||||
|
||||
let _ = engine.run("call it").await;
|
||||
let _ = engine.run(&mut history, "call it").await;
|
||||
|
||||
let observed = captured.lock().unwrap();
|
||||
assert_eq!(observed.len(), 1);
|
||||
@@ -336,6 +335,7 @@ async fn test_callback_tool_result_error_path() {
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
engine.register_tool(erroring_tool("erroring", "boom"));
|
||||
|
||||
@@ -351,7 +351,7 @@ async fn test_callback_tool_result_error_path() {
|
||||
));
|
||||
});
|
||||
|
||||
let _ = engine.run("fail it").await;
|
||||
let _ = engine.run(&mut history, "fail it").await;
|
||||
|
||||
let observed = captured.lock().unwrap();
|
||||
assert_eq!(observed.len(), 1);
|
||||
@@ -380,6 +380,7 @@ async fn test_callback_usage_events() {
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
let usage_events = Arc::new(Mutex::new(Vec::new()));
|
||||
|
||||
@@ -388,8 +389,8 @@ async fn test_callback_usage_events() {
|
||||
usages.lock().unwrap().push(event.clone());
|
||||
});
|
||||
|
||||
// Mutable::run consumes self, returns (Locked, EngineRunExit)
|
||||
let _ = engine.run("Hello").await;
|
||||
// Mutable::run consumes self, returns (Locked, EngineResult)
|
||||
let _ = engine.run(&mut history, "Hello").await;
|
||||
|
||||
let usages = usage_events.lock().unwrap();
|
||||
assert_eq!(usages.len(), 1);
|
||||
|
||||
@@ -9,8 +9,8 @@ use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use agen::Engine;
|
||||
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
||||
use agen::{Engine, History};
|
||||
use async_trait::async_trait;
|
||||
use common::MockLlmClient;
|
||||
|
||||
@@ -134,9 +134,10 @@ async fn test_engine_simple_text_response() {
|
||||
|
||||
let client = MockLlmClient::from_fixture(&fixture_path).unwrap();
|
||||
let engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
// Send a simple message (Mutable::run consumes self, returns tuple)
|
||||
let result = engine.run("Hello").await;
|
||||
let result = engine.run(&mut history, "Hello").await;
|
||||
|
||||
assert!(
|
||||
matches!(result.result, agen::EngineRunExit::Finished),
|
||||
@@ -159,6 +160,7 @@ async fn test_engine_tool_call() {
|
||||
|
||||
let client = MockLlmClient::from_fixture(&fixture_path).unwrap();
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
// Register tool
|
||||
let weather_tool = MockWeatherTool::new();
|
||||
@@ -166,7 +168,9 @@ async fn test_engine_tool_call() {
|
||||
engine.register_tool(weather_tool.definition());
|
||||
|
||||
// Send message (Mutable::run consumes self, returns tuple)
|
||||
let _result = engine.run("What's the weather in Tokyo?").await;
|
||||
let _result = engine
|
||||
.run(&mut history, "What's the weather in Tokyo?")
|
||||
.await;
|
||||
|
||||
// Verify tool was called
|
||||
// Note: max_turns=1 so no request is sent after tool result
|
||||
@@ -198,9 +202,10 @@ async fn test_engine_with_programmatic_events() {
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
// Mutable::run consumes self, returns tuple
|
||||
let result = engine.run("Greet me").await;
|
||||
let result = engine.run(&mut history, "Greet me").await;
|
||||
|
||||
assert!(
|
||||
matches!(result.result, agen::EngineRunExit::Finished),
|
||||
|
||||
@@ -15,7 +15,7 @@ use agen::interceptor::{
|
||||
use agen::llm_client::ClientError;
|
||||
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
|
||||
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
||||
use agen::{Engine, EngineError, EngineRunExit, StopReason};
|
||||
use agen::{Engine, EngineError, EngineResult, History};
|
||||
use async_trait::async_trait;
|
||||
use common::MockLlmClient;
|
||||
|
||||
@@ -75,36 +75,37 @@ fn test_mutable_set_system_prompt() {
|
||||
fn test_mutable_history_manipulation() {
|
||||
let client = MockLlmClient::new(vec![]);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
// Initial state is empty
|
||||
assert!(engine.history().is_empty());
|
||||
assert!(history.is_empty());
|
||||
|
||||
// Add to history
|
||||
engine
|
||||
.append_history(vec![Item::user_message("Hello")])
|
||||
.append_history(&mut history, vec![Item::user_message("Hello")])
|
||||
.unwrap();
|
||||
engine
|
||||
.append_history(vec![Item::assistant_message("Hi there!")])
|
||||
.append_history(&mut history, vec![Item::assistant_message("Hi there!")])
|
||||
.unwrap();
|
||||
assert_eq!(engine.history().len(), 2);
|
||||
assert_eq!(history.len(), 2);
|
||||
|
||||
// Append to history via the callback-aware API.
|
||||
engine
|
||||
.append_history(vec![Item::user_message("How are you?")])
|
||||
.append_history(&mut history, vec![Item::user_message("How are you?")])
|
||||
.unwrap();
|
||||
assert_eq!(engine.history().len(), 3);
|
||||
assert_eq!(history.len(), 3);
|
||||
|
||||
// Clear history
|
||||
engine.clear_history();
|
||||
assert!(engine.history().is_empty());
|
||||
engine.clear_history(&mut history);
|
||||
assert!(history.is_empty());
|
||||
|
||||
// Set history
|
||||
let items = vec![
|
||||
Item::user_message("Test"),
|
||||
Item::assistant_message("Response"),
|
||||
];
|
||||
engine.set_history(items);
|
||||
assert_eq!(engine.history().len(), 2);
|
||||
engine.set_history(&mut history, items);
|
||||
assert_eq!(history.len(), 2);
|
||||
}
|
||||
|
||||
/// Verify that Engine can be constructed using builder pattern
|
||||
@@ -112,9 +113,10 @@ fn test_mutable_history_manipulation() {
|
||||
fn test_mutable_builder_pattern() {
|
||||
let client = MockLlmClient::new(vec![]);
|
||||
let engine = Engine::new(client).system_prompt("System prompt");
|
||||
let history: History = History::new();
|
||||
|
||||
assert_eq!(engine.get_system_prompt(), Some("System prompt"));
|
||||
assert!(engine.history().is_empty());
|
||||
assert!(history.is_empty());
|
||||
}
|
||||
|
||||
/// Verify that multiple items can be added with append_history and callbacks fire.
|
||||
@@ -124,6 +126,7 @@ fn test_mutable_append_history() {
|
||||
let observed = Arc::new(Mutex::new(Vec::new()));
|
||||
let observed_for_callback = Arc::clone(&observed);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
engine.on_history_append(move |item| {
|
||||
if let Some(text) = item.as_text() {
|
||||
observed_for_callback.lock().unwrap().push(text.to_string());
|
||||
@@ -132,18 +135,21 @@ fn test_mutable_append_history() {
|
||||
});
|
||||
|
||||
engine
|
||||
.append_history(vec![Item::user_message("First")])
|
||||
.append_history(&mut history, vec![Item::user_message("First")])
|
||||
.unwrap();
|
||||
|
||||
engine
|
||||
.append_history(vec![
|
||||
Item::assistant_message("Response 1"),
|
||||
Item::user_message("Second"),
|
||||
Item::assistant_message("Response 2"),
|
||||
])
|
||||
.append_history(
|
||||
&mut history,
|
||||
vec![
|
||||
Item::assistant_message("Response 1"),
|
||||
Item::user_message("Second"),
|
||||
Item::assistant_message("Response 2"),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(engine.history().len(), 4);
|
||||
assert_eq!(history.len(), 4);
|
||||
assert_eq!(
|
||||
observed.lock().unwrap().as_slice(),
|
||||
["First", "Response 1", "Second", "Response 2"]
|
||||
@@ -218,6 +224,7 @@ async fn history_append_failure_stops_before_tool_execution() {
|
||||
]);
|
||||
let tool = CountingTool::new("count_tool");
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
engine.register_tool(tool.definition());
|
||||
engine.on_history_append(|item| {
|
||||
if item.is_tool_call() {
|
||||
@@ -227,8 +234,8 @@ async fn history_append_failure_stops_before_tool_execution() {
|
||||
}
|
||||
});
|
||||
|
||||
let mut engine = engine.lock();
|
||||
let exit = engine.run("use the tool").await;
|
||||
let mut engine = engine.lock(&history);
|
||||
let error = engine.run(&mut history, "use the tool").await.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
exit,
|
||||
@@ -236,8 +243,8 @@ async fn history_append_failure_stops_before_tool_execution() {
|
||||
if message == "simulated ENOSPC"
|
||||
));
|
||||
assert_eq!(tool.call_count(), 0);
|
||||
assert_eq!(engine.history().len(), 1);
|
||||
assert_eq!(engine.history()[0].as_text(), Some("use the tool"));
|
||||
assert_eq!(history.len(), 1);
|
||||
assert_eq!(history.entries()[0].item.as_text(), Some("use the tool"));
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
@@ -249,21 +256,22 @@ async fn history_append_failure_stops_before_tool_execution() {
|
||||
fn test_lock_transition() {
|
||||
let client = MockLlmClient::new(vec![]);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
engine.set_system_prompt("System");
|
||||
engine
|
||||
.append_history(vec![Item::user_message("Hello")])
|
||||
.append_history(&mut history, vec![Item::user_message("Hello")])
|
||||
.unwrap();
|
||||
engine
|
||||
.append_history(vec![Item::assistant_message("Hi")])
|
||||
.append_history(&mut history, vec![Item::assistant_message("Hi")])
|
||||
.unwrap();
|
||||
|
||||
// Lock
|
||||
let locked_engine = engine.lock();
|
||||
let locked_engine = engine.lock(&history);
|
||||
|
||||
// History and system prompt are still accessible in Locked state
|
||||
assert_eq!(locked_engine.get_system_prompt(), Some("System"));
|
||||
assert_eq!(locked_engine.history().len(), 2);
|
||||
assert_eq!(history.len(), 2);
|
||||
assert_eq!(locked_engine.locked_prefix_len(), 2);
|
||||
}
|
||||
|
||||
@@ -272,21 +280,22 @@ fn test_lock_transition() {
|
||||
fn test_unlock_transition() {
|
||||
let client = MockLlmClient::new(vec![]);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
engine
|
||||
.append_history(vec![Item::user_message("Hello")])
|
||||
.append_history(&mut history, vec![Item::user_message("Hello")])
|
||||
.unwrap();
|
||||
let locked_engine = engine.lock();
|
||||
let locked_engine = engine.lock(&history);
|
||||
|
||||
// Unlock
|
||||
let mut engine = locked_engine.unlock();
|
||||
|
||||
// History operations are available again in Mutable state
|
||||
engine
|
||||
.append_history(vec![Item::assistant_message("Hi")])
|
||||
.append_history(&mut history, vec![Item::assistant_message("Hi")])
|
||||
.unwrap();
|
||||
engine.clear_history();
|
||||
assert!(engine.history().is_empty());
|
||||
engine.clear_history(&mut history);
|
||||
assert!(history.is_empty());
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
@@ -307,20 +316,20 @@ async fn test_mutable_run_updates_history() -> Result<(), EngineError> {
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
// Execute (Mutable::run consumes self, returns EngineRunOutput)
|
||||
let out = engine.run("Hi there").await;
|
||||
let engine = out.engine;
|
||||
let _out = engine.run(&mut history, "Hi there").await?;
|
||||
|
||||
// History is updated
|
||||
let history = engine.history();
|
||||
let entries = history.entries();
|
||||
assert_eq!(history.len(), 2); // user + assistant
|
||||
|
||||
// User message
|
||||
assert_eq!(history[0].as_text(), Some("Hi there"));
|
||||
assert_eq!(entries[0].item.as_text(), Some("Hi there"));
|
||||
|
||||
// Assistant message
|
||||
assert_eq!(history[1].as_text(), Some("Hello, I'm an assistant!"));
|
||||
assert_eq!(entries[1].item.as_text(), Some("Hello, I'm an assistant!"));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -351,35 +360,36 @@ async fn test_locked_multi_turn_history_accumulation() {
|
||||
]);
|
||||
|
||||
let engine = Engine::new(client).system_prompt("You are helpful.");
|
||||
let mut history: History = History::new();
|
||||
|
||||
// Lock (after setting system prompt)
|
||||
let mut locked_engine = engine.lock();
|
||||
let mut locked_engine = engine.lock(&history);
|
||||
assert_eq!(locked_engine.locked_prefix_len(), 0); // No items yet
|
||||
|
||||
// Turn 1
|
||||
let result1 = locked_engine.run("Hello!").await;
|
||||
assert!(matches!(result1, EngineRunExit::Finished));
|
||||
assert_eq!(locked_engine.history().len(), 2); // user + assistant
|
||||
let result1 = locked_engine.run(&mut history, "Hello!").await;
|
||||
assert!(result1.is_ok());
|
||||
assert_eq!(history.len(), 2); // user + assistant
|
||||
|
||||
// Turn 2
|
||||
let result2 = locked_engine.run("Can you help me?").await;
|
||||
assert!(matches!(result2, EngineRunExit::Finished));
|
||||
assert_eq!(locked_engine.history().len(), 4); // 2 * (user + assistant)
|
||||
let result2 = locked_engine.run(&mut history, "Can you help me?").await;
|
||||
assert!(result2.is_ok());
|
||||
assert_eq!(history.len(), 4); // 2 * (user + assistant)
|
||||
|
||||
// Verify history contents
|
||||
let history = locked_engine.history();
|
||||
let entries = history.entries();
|
||||
|
||||
// Turn 1 user message
|
||||
assert_eq!(history[0].as_text(), Some("Hello!"));
|
||||
assert_eq!(entries[0].item.as_text(), Some("Hello!"));
|
||||
|
||||
// Turn 1 assistant message
|
||||
assert_eq!(history[1].as_text(), Some("Nice to meet you!"));
|
||||
assert_eq!(entries[1].item.as_text(), Some("Nice to meet you!"));
|
||||
|
||||
// Turn 2 user message
|
||||
assert_eq!(history[2].as_text(), Some("Can you help me?"));
|
||||
assert_eq!(entries[2].item.as_text(), Some("Can you help me?"));
|
||||
|
||||
// Turn 2 assistant message
|
||||
assert_eq!(history[3].as_text(), Some("I can help with that."));
|
||||
assert_eq!(entries[3].item.as_text(), Some("I can help with that."));
|
||||
}
|
||||
|
||||
/// Verify that locked_prefix_len correctly records history length at lock time
|
||||
@@ -405,26 +415,36 @@ async fn test_locked_prefix_len_tracking() {
|
||||
]);
|
||||
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
// Add items beforehand
|
||||
engine
|
||||
.append_history(vec![Item::user_message("Pre-existing message 1")])
|
||||
.append_history(
|
||||
&mut history,
|
||||
vec![Item::user_message("Pre-existing message 1")],
|
||||
)
|
||||
.unwrap();
|
||||
engine
|
||||
.append_history(vec![Item::assistant_message("Pre-existing response 1")])
|
||||
.append_history(
|
||||
&mut history,
|
||||
vec![Item::assistant_message("Pre-existing response 1")],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(engine.history().len(), 2);
|
||||
assert_eq!(history.len(), 2);
|
||||
|
||||
// Lock
|
||||
let mut locked_engine = engine.lock();
|
||||
let mut locked_engine = engine.lock(&history);
|
||||
assert_eq!(locked_engine.locked_prefix_len(), 2); // 2 items at lock time
|
||||
|
||||
// Execute turn
|
||||
locked_engine.run("New message").await;
|
||||
locked_engine
|
||||
.run(&mut history, "New message")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// History grows but locked_prefix_len remains unchanged
|
||||
assert_eq!(locked_engine.history().len(), 4); // 2 + 2
|
||||
assert_eq!(history.len(), 4); // 2 + 2
|
||||
assert_eq!(locked_engine.locked_prefix_len(), 2); // Unchanged
|
||||
}
|
||||
|
||||
@@ -451,18 +471,19 @@ async fn test_turn_count_increment() -> Result<(), EngineError> {
|
||||
]);
|
||||
|
||||
let engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
assert_eq!(engine.turn_count(), 0);
|
||||
assert_eq!(engine.llm_call_count(), 0);
|
||||
|
||||
// First run consumes Mutable, returns EngineRunOutput
|
||||
let mut engine = engine.run("First").await.engine;
|
||||
let mut engine = engine.run(&mut history, "First").await?.engine;
|
||||
assert_eq!(engine.turn_count(), 1);
|
||||
// Retry not yet implemented → AgentTurn:LlmCall is 1:1.
|
||||
assert_eq!(engine.llm_call_count(), 1);
|
||||
|
||||
// Subsequent runs on Locked take &mut self
|
||||
engine.run("Second").await;
|
||||
engine.run(&mut history, "Second").await?;
|
||||
assert_eq!(engine.turn_count(), 2);
|
||||
assert_eq!(engine.llm_call_count(), 2);
|
||||
|
||||
@@ -482,28 +503,29 @@ async fn test_unlock_edit_relock() {
|
||||
]]);
|
||||
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
engine
|
||||
.append_history(vec![
|
||||
Item::user_message("Hello"),
|
||||
Item::assistant_message("Hi"),
|
||||
])
|
||||
.append_history(
|
||||
&mut history,
|
||||
vec![Item::user_message("Hello"), Item::assistant_message("Hi")],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Lock -> Unlock
|
||||
let locked = engine.lock();
|
||||
let locked = engine.lock(&history);
|
||||
assert_eq!(locked.locked_prefix_len(), 2);
|
||||
|
||||
let mut unlocked = locked.unlock();
|
||||
|
||||
// Edit history
|
||||
unlocked.clear_history();
|
||||
unlocked.clear_history(&mut history);
|
||||
unlocked
|
||||
.append_history(vec![Item::user_message("Fresh start")])
|
||||
.append_history(&mut history, vec![Item::user_message("Fresh start")])
|
||||
.unwrap();
|
||||
|
||||
// Re-lock
|
||||
let relocked = unlocked.lock();
|
||||
assert_eq!(relocked.history().len(), 1);
|
||||
let relocked = unlocked.lock(&history);
|
||||
assert_eq!(history.len(), 1);
|
||||
assert_eq!(relocked.locked_prefix_len(), 1);
|
||||
}
|
||||
|
||||
@@ -546,19 +568,23 @@ async fn test_lock_unlock_relock_tools_remain_effective() {
|
||||
]);
|
||||
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
let tool_a = CountingTool::new("tool_a");
|
||||
engine.register_tool(tool_a.definition());
|
||||
|
||||
let mut locked = engine.lock();
|
||||
locked.run("first").await;
|
||||
let mut locked = engine.lock(&history);
|
||||
locked.run(&mut history, "first").await.expect("first run");
|
||||
assert_eq!(tool_a.call_count(), 1, "tool_a should be called once");
|
||||
|
||||
let mut unlocked = locked.unlock();
|
||||
let tool_b = CountingTool::new("tool_b");
|
||||
unlocked.register_tool(tool_b.definition());
|
||||
|
||||
let mut relocked = unlocked.lock();
|
||||
relocked.run("second").await;
|
||||
let mut relocked = unlocked.lock(&history);
|
||||
relocked
|
||||
.run(&mut history, "second")
|
||||
.await
|
||||
.expect("second run");
|
||||
|
||||
assert_eq!(tool_a.call_count(), 1, "tool_a should not be called again");
|
||||
assert_eq!(tool_b.call_count(), 1, "tool_b should be called once");
|
||||
@@ -573,8 +599,9 @@ async fn test_lock_unlock_relock_tools_remain_effective() {
|
||||
fn test_system_prompt_preserved_in_locked_state() {
|
||||
let client = MockLlmClient::new(vec![]);
|
||||
let engine = Engine::new(client).system_prompt("Important system prompt");
|
||||
let history: History = History::new();
|
||||
|
||||
let locked = engine.lock();
|
||||
let locked = engine.lock(&history);
|
||||
assert_eq!(locked.get_system_prompt(), Some("Important system prompt"));
|
||||
|
||||
let unlocked = locked.unlock();
|
||||
@@ -589,14 +616,15 @@ fn test_system_prompt_preserved_in_locked_state() {
|
||||
fn test_system_prompt_change_after_unlock() {
|
||||
let client = MockLlmClient::new(vec![]);
|
||||
let engine = Engine::new(client).system_prompt("Original prompt");
|
||||
let history: History = History::new();
|
||||
|
||||
let locked = engine.lock();
|
||||
let locked = engine.lock(&history);
|
||||
let mut unlocked = locked.unlock();
|
||||
|
||||
unlocked.set_system_prompt("New prompt");
|
||||
assert_eq!(unlocked.get_system_prompt(), Some("New prompt"));
|
||||
|
||||
let relocked = unlocked.lock();
|
||||
let relocked = unlocked.lock(&history);
|
||||
assert_eq!(relocked.get_system_prompt(), Some("New prompt"));
|
||||
}
|
||||
|
||||
@@ -660,17 +688,21 @@ impl Interceptor for ContinueTurnOnce {
|
||||
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));
|
||||
let mut history: History = History::new();
|
||||
engine.set_max_turns(Some(1));
|
||||
let mut engine = engine.lock();
|
||||
let mut engine = engine.lock(&history);
|
||||
|
||||
assert!(matches!(engine.run("first").await, EngineRunExit::Finished));
|
||||
assert_eq!(
|
||||
engine.run(&mut history, "first").await.unwrap(),
|
||||
EngineResult::Finished
|
||||
);
|
||||
assert_eq!(engine.turn_count(), 1);
|
||||
assert_eq!(engine.active_run_turn_count(), None);
|
||||
|
||||
assert!(matches!(
|
||||
engine.run("second").await,
|
||||
EngineRunExit::Finished
|
||||
));
|
||||
assert_eq!(
|
||||
engine.run(&mut history, "second").await.unwrap(),
|
||||
EngineResult::Finished
|
||||
);
|
||||
assert_eq!(engine.turn_count(), 2);
|
||||
assert_eq!(engine.active_run_turn_count(), None);
|
||||
}
|
||||
@@ -678,17 +710,24 @@ async fn max_turns_is_scoped_to_each_fresh_run() {
|
||||
#[tokio::test]
|
||||
async fn yielded_resume_keeps_the_same_unspent_turn_budget() {
|
||||
let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
|
||||
let mut history: History = History::new();
|
||||
engine.set_max_turns(Some(1));
|
||||
engine.set_interceptor(YieldOnce {
|
||||
calls: AtomicUsize::new(0),
|
||||
});
|
||||
let mut engine = engine.lock();
|
||||
let mut engine = engine.lock(&history);
|
||||
|
||||
assert!(matches!(engine.run("start").await, EngineRunExit::Yielded));
|
||||
assert_eq!(
|
||||
engine.run(&mut history, "start").await.unwrap(),
|
||||
EngineResult::Yielded
|
||||
);
|
||||
assert_eq!(engine.turn_count(), 0);
|
||||
assert_eq!(engine.active_run_turn_count(), Some(0));
|
||||
|
||||
assert!(matches!(engine.resume().await, EngineRunExit::Finished));
|
||||
assert_eq!(
|
||||
engine.resume(&mut history).await.unwrap(),
|
||||
EngineResult::Finished
|
||||
);
|
||||
assert_eq!(engine.turn_count(), 1);
|
||||
assert_eq!(engine.active_run_turn_count(), None);
|
||||
}
|
||||
@@ -705,22 +744,26 @@ async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() {
|
||||
];
|
||||
let tool = CountingTool::new("count_tool");
|
||||
let mut engine = Engine::new(MockLlmClient::new(events));
|
||||
let mut history: History = History::new();
|
||||
engine.set_max_turns(Some(1));
|
||||
engine.register_tool(tool.definition());
|
||||
engine.set_interceptor(PauseToolOnce {
|
||||
calls: AtomicUsize::new(0),
|
||||
});
|
||||
let mut engine = engine.lock();
|
||||
let mut engine = engine.lock(&history);
|
||||
|
||||
assert!(matches!(engine.run("call it").await, EngineRunExit::Paused));
|
||||
assert_eq!(
|
||||
engine.run(&mut history, "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!(matches!(
|
||||
engine.resume().await,
|
||||
EngineRunExit::Interrupted(StopReason::LimitReached)
|
||||
));
|
||||
assert_eq!(
|
||||
engine.resume(&mut history).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");
|
||||
@@ -739,20 +782,24 @@ async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() {
|
||||
let client = MockLlmClient::with_responses(vec![tool_events, completed_text_events()]);
|
||||
let tool = CountingTool::new("count_tool");
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
engine.set_max_turns(Some(1));
|
||||
engine.register_tool(tool.definition());
|
||||
engine.set_interceptor(PauseToolOnce {
|
||||
calls: AtomicUsize::new(0),
|
||||
});
|
||||
let mut engine = engine.lock();
|
||||
let mut engine = engine.lock(&history);
|
||||
|
||||
assert!(matches!(engine.run("pause").await, EngineRunExit::Paused));
|
||||
assert_eq!(
|
||||
engine.run(&mut history, "pause").await.unwrap(),
|
||||
EngineResult::Paused
|
||||
);
|
||||
assert_eq!(engine.active_run_turn_count(), Some(1));
|
||||
|
||||
assert!(matches!(
|
||||
engine.run("replace").await,
|
||||
EngineRunExit::Finished
|
||||
));
|
||||
assert_eq!(
|
||||
engine.run(&mut history, "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");
|
||||
@@ -761,16 +808,17 @@ async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() {
|
||||
#[tokio::test]
|
||||
async fn interceptor_continuation_consumes_the_logical_run_budget() {
|
||||
let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
|
||||
let mut history: History = History::new();
|
||||
engine.set_max_turns(Some(1));
|
||||
engine.set_interceptor(ContinueTurnOnce {
|
||||
calls: AtomicUsize::new(0),
|
||||
});
|
||||
let mut engine = engine.lock();
|
||||
let mut engine = engine.lock(&history);
|
||||
|
||||
assert!(matches!(
|
||||
engine.run("start").await,
|
||||
EngineRunExit::Interrupted(StopReason::LimitReached)
|
||||
));
|
||||
assert_eq!(
|
||||
engine.run(&mut history, "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);
|
||||
@@ -779,15 +827,16 @@ async fn interceptor_continuation_consumes_the_logical_run_budget() {
|
||||
#[tokio::test]
|
||||
async fn restored_active_run_budget_is_enforced_before_another_llm_call() {
|
||||
let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
|
||||
let mut history: History = History::new();
|
||||
engine.set_max_turns(Some(1));
|
||||
engine.set_turn_count(7);
|
||||
engine.set_active_run_turn_count(Some(1));
|
||||
let mut engine = engine.lock();
|
||||
let mut engine = engine.lock(&history);
|
||||
|
||||
assert!(matches!(
|
||||
engine.resume().await,
|
||||
EngineRunExit::Interrupted(StopReason::LimitReached)
|
||||
));
|
||||
assert_eq!(
|
||||
engine.resume(&mut history).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);
|
||||
|
||||
@@ -6,12 +6,12 @@ use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use agen::Engine;
|
||||
use agen::interceptor::{Interceptor, PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo};
|
||||
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
|
||||
use agen::tool::{
|
||||
Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult,
|
||||
};
|
||||
use agen::{Engine, History};
|
||||
use async_trait::async_trait;
|
||||
|
||||
mod common;
|
||||
@@ -145,6 +145,7 @@ async fn test_parallel_tool_execution() {
|
||||
],
|
||||
]);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
let tool1 = SlowTool::new("slow_tool_1", 100);
|
||||
let tool2 = SlowTool::new("slow_tool_2", 100);
|
||||
let tool3 = SlowTool::new("slow_tool_3", 100);
|
||||
@@ -158,8 +159,8 @@ async fn test_parallel_tool_execution() {
|
||||
engine.register_tool(tool3.definition());
|
||||
|
||||
let start = Instant::now();
|
||||
// Mutable::run consumes self, returns (Locked, EngineRunExit)
|
||||
let _result = engine.run("Run all tools").await;
|
||||
// Mutable::run consumes self, returns (Locked, EngineResult)
|
||||
let _result = engine.run(&mut history, "Run all tools").await;
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
// Verify all tools were called
|
||||
@@ -205,13 +206,14 @@ async fn test_tool_execution_context_order_and_batch_id() {
|
||||
],
|
||||
]);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
let contexts = Arc::new(Mutex::new(Vec::new()));
|
||||
|
||||
engine.register_tool(ContextRecordingTool::new("record_a", contexts.clone()).definition());
|
||||
engine.register_tool(ContextRecordingTool::new("record_b", contexts.clone()).definition());
|
||||
engine.register_tool(ContextRecordingTool::new("record_c", contexts.clone()).definition());
|
||||
|
||||
let _ = engine.run("record contexts").await;
|
||||
let _ = engine.run(&mut history, "record contexts").await;
|
||||
|
||||
let mut contexts = contexts.lock().unwrap().clone();
|
||||
contexts.sort_by_key(|ctx| ctx.call_index);
|
||||
@@ -256,11 +258,12 @@ async fn test_tool_execution_context_batch_id_changes_between_batches() {
|
||||
],
|
||||
]);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
let contexts = Arc::new(Mutex::new(Vec::new()));
|
||||
|
||||
engine.register_tool(ContextRecordingTool::new("record", contexts.clone()).definition());
|
||||
|
||||
let _ = engine.run("record batches").await;
|
||||
let _ = engine.run(&mut history, "record batches").await;
|
||||
|
||||
let contexts = contexts.lock().unwrap().clone();
|
||||
assert_eq!(contexts.len(), 2);
|
||||
@@ -298,6 +301,7 @@ async fn test_tool_execution_context_for_skipped_and_synthetic_paths() {
|
||||
],
|
||||
]);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
let executed_contexts = Arc::new(Mutex::new(Vec::new()));
|
||||
let pre_contexts = Arc::new(Mutex::new(Vec::new()));
|
||||
let post_contexts = Arc::new(Mutex::new(Vec::new()));
|
||||
@@ -344,7 +348,9 @@ async fn test_tool_execution_context_for_skipped_and_synthetic_paths() {
|
||||
post_contexts: post_contexts.clone(),
|
||||
});
|
||||
|
||||
let _ = engine.run("record skipped and synthetic contexts").await;
|
||||
let _ = engine
|
||||
.run(&mut history, "record skipped and synthetic contexts")
|
||||
.await;
|
||||
|
||||
let mut pre_contexts = pre_contexts.lock().unwrap().clone();
|
||||
pre_contexts.sort_by_key(|ctx| ctx.call_index);
|
||||
@@ -389,6 +395,7 @@ async fn test_before_tool_call_skip() {
|
||||
|
||||
let client = MockLlmClient::new(events);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
let allowed_tool = SlowTool::new("allowed_tool", 10);
|
||||
let blocked_tool = SlowTool::new("blocked_tool", 10);
|
||||
@@ -415,8 +422,8 @@ async fn test_before_tool_call_skip() {
|
||||
|
||||
engine.set_interceptor(BlockingPolicy);
|
||||
|
||||
// Mutable::run consumes self, returns (Locked, EngineRunExit)
|
||||
let _result = engine.run("Test hook").await;
|
||||
// Mutable::run consumes self, returns (Locked, EngineResult)
|
||||
let _result = engine.run(&mut history, "Test hook").await;
|
||||
|
||||
// allowed_tool is called, but blocked_tool is not
|
||||
assert_eq!(
|
||||
@@ -457,6 +464,7 @@ async fn test_post_tool_call_modification() {
|
||||
]);
|
||||
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
|
||||
#[derive(Clone)]
|
||||
struct SimpleTool;
|
||||
@@ -502,8 +510,8 @@ async fn test_post_tool_call_modification() {
|
||||
modified_content: modified_content.clone(),
|
||||
});
|
||||
|
||||
// Mutable::run consumes self, returns (Locked, EngineRunExit)
|
||||
let result = engine.run("Test modification").await;
|
||||
// Mutable::run consumes self, returns (Locked, EngineResult)
|
||||
let result = engine.run(&mut history, "Test modification").await;
|
||||
|
||||
assert!(
|
||||
matches!(result.result, agen::EngineRunExit::Finished),
|
||||
@@ -543,6 +551,7 @@ async fn test_before_tool_call_synthetic_result_committed() {
|
||||
],
|
||||
]);
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
let blocked_tool = SlowTool::new("blocked_tool", 10);
|
||||
let blocked_clone = blocked_tool.clone();
|
||||
engine.register_tool(blocked_tool.definition());
|
||||
@@ -561,10 +570,13 @@ async fn test_before_tool_call_synthetic_result_committed() {
|
||||
|
||||
engine.set_interceptor(SyntheticPolicy);
|
||||
|
||||
let result = engine.run("Test synthetic result").await;
|
||||
let _result = engine
|
||||
.run(&mut history, "Test synthetic result")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(blocked_clone.call_count(), 0, "Blocked tool should not run");
|
||||
assert!(result.engine.history().iter().any(|item| matches!(
|
||||
assert!(history.items().any(|item| matches!(
|
||||
item,
|
||||
agen::Item::ToolResult {
|
||||
call_id,
|
||||
|
||||
@@ -13,12 +13,12 @@
|
||||
|
||||
mod common;
|
||||
|
||||
use agen::Engine;
|
||||
use agen::Item;
|
||||
use agen::llm_client::event::{
|
||||
BlockMetadata, BlockStart, BlockStop, BlockType, Event, ReasoningBlockData, ResponseStatus,
|
||||
StatusEvent,
|
||||
};
|
||||
use agen::{Engine, History};
|
||||
use common::MockLlmClient;
|
||||
|
||||
fn reasoning_block(text: impl Into<String>, data: ReasoningBlockData) -> Vec<Event> {
|
||||
@@ -65,15 +65,15 @@ async fn anthropic_thinking_round_trips_signature_into_history() {
|
||||
]);
|
||||
let client = MockLlmClient::new(events);
|
||||
let engine = Engine::new(client);
|
||||
let out = engine.run("question?").await;
|
||||
let engine = out.engine;
|
||||
let mut history: History = History::new();
|
||||
let _out = engine.run(&mut history, "question?").await.expect("run ok");
|
||||
|
||||
let history = engine.history();
|
||||
let entries = history.entries();
|
||||
// user / reasoning / assistant_message
|
||||
assert_eq!(history.len(), 3, "history: {history:?}");
|
||||
|
||||
assert!(matches!(history[0], Item::Message { .. }));
|
||||
match &history[1] {
|
||||
assert!(matches!(entries[0].item, Item::Message { .. }));
|
||||
match &entries[1].item {
|
||||
Item::Reasoning {
|
||||
text, signature, ..
|
||||
} => {
|
||||
@@ -82,7 +82,7 @@ async fn anthropic_thinking_round_trips_signature_into_history() {
|
||||
}
|
||||
other => panic!("expected Reasoning, got {other:?}"),
|
||||
}
|
||||
assert_eq!(history[2].as_text(), Some("Here's the answer"));
|
||||
assert_eq!(entries[2].item.as_text(), Some("Here's the answer"));
|
||||
}
|
||||
|
||||
/// OpenAI Responses 風: encrypted_content + summary を持った reasoning が
|
||||
@@ -109,11 +109,11 @@ async fn openai_reasoning_round_trips_encrypted_and_summary() {
|
||||
]);
|
||||
let client = MockLlmClient::new(events);
|
||||
let engine = Engine::new(client);
|
||||
let out = engine.run("q").await;
|
||||
let engine = out.engine;
|
||||
let mut history: History = History::new();
|
||||
let _out = engine.run(&mut history, "q").await.expect("run ok");
|
||||
|
||||
let history = engine.history();
|
||||
match &history[1] {
|
||||
let entries = history.entries();
|
||||
match &entries[1].item {
|
||||
Item::Reasoning {
|
||||
text,
|
||||
summary,
|
||||
@@ -155,13 +155,13 @@ async fn reasoning_precedes_text_in_assistant_burst() {
|
||||
}));
|
||||
let client = MockLlmClient::new(events);
|
||||
let engine = Engine::new(client);
|
||||
let out = engine.run("q").await;
|
||||
let engine = out.engine;
|
||||
let mut history: History = History::new();
|
||||
let _out = engine.run(&mut history, "q").await.expect("run ok");
|
||||
|
||||
let history = engine.history();
|
||||
let entries = history.entries();
|
||||
// user / reasoning(先頭) / assistant_message
|
||||
assert!(matches!(history[1], Item::Reasoning { .. }));
|
||||
assert_eq!(history[2].as_text(), Some("intermediate"));
|
||||
assert!(matches!(entries[1].item, Item::Reasoning { .. }));
|
||||
assert_eq!(entries[2].item.as_text(), Some("intermediate"));
|
||||
}
|
||||
|
||||
/// resume シナリオ: history.json 由来の Item::Reasoning(signature) を Engine に
|
||||
@@ -207,14 +207,18 @@ async fn injected_reasoning_survives_into_outgoing_request() {
|
||||
};
|
||||
|
||||
let mut engine = Engine::new(client);
|
||||
let mut history: History = History::new();
|
||||
// resume: 既存 history を流し込む
|
||||
engine.set_history(vec![
|
||||
Item::user_message("prior question"),
|
||||
Item::reasoning("prior thinking").with_signature("SIG-PRIOR"),
|
||||
Item::assistant_message("prior answer"),
|
||||
]);
|
||||
engine.set_history(
|
||||
&mut history,
|
||||
vec![
|
||||
Item::user_message("prior question"),
|
||||
Item::reasoning("prior thinking").with_signature("SIG-PRIOR"),
|
||||
Item::assistant_message("prior answer"),
|
||||
],
|
||||
);
|
||||
|
||||
let _ = engine.run("follow up").await;
|
||||
let _ = engine.run(&mut history, "follow up").await.expect("run ok");
|
||||
|
||||
let req = captured
|
||||
.lock()
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use agen::Engine;
|
||||
use agen::{Engine, History};
|
||||
use agen::llm_client::capability::{
|
||||
CacheStrategy, ModelCapability, StructuredOutput, ToolCallingSupport,
|
||||
};
|
||||
@@ -22,7 +22,8 @@ fn main() {
|
||||
cap,
|
||||
);
|
||||
let engine = Engine::new(client);
|
||||
let mut locked = engine.lock();
|
||||
let history = History::new();
|
||||
let mut locked = engine.lock(&history);
|
||||
let def: agen::tool::ToolDefinition = Arc::new(|| panic!("unused"));
|
||||
let _ = locked.register_tool(def);
|
||||
}
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
error[E0599]: no method named `register_tool` found for struct `Engine<HttpTransport<AnthropicScheme>, Locked>` in the current scope
|
||||
--> tests/ui/locked_register_tool.rs:27:20
|
||||
--> tests/ui/locked_register_tool.rs:28:20
|
||||
|
|
||||
27 | let _ = locked.register_tool(def);
|
||||
28 | let _ = locked.register_tool(def);
|
||||
| ^^^^^^^^^^^^^ method not found in `Engine<HttpTransport<AnthropicScheme>, Locked>`
|
||||
|
|
||||
= note: the method was found for
|
||||
- `Engine<C>`
|
||||
- `Engine<C, Mutable, A>`
|
||||
|
||||
Reference in New Issue
Block a user