Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: route worker history appends through callbacks

Browse Source
This commit is contained in:
Hare
2026-05-24 06:44:19 +09:00
parent b64e098b5b
commit 65c399e6d9
5 changed files with 178 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/llm-worker/src/worker.rs
+24 -34
View File
@@ -127,7 +127,7 @@ enum ToolExecutionResult {
///
/// // To edit between turns, unlock back to Mutable
/// let mut worker = worker.unlock();
/// worker.history_mut().truncate(5);
/// worker.truncate_history(5);
/// let out = worker.run("Continue").await?;
/// let mut worker = out.worker;
/// ```
@@ -400,7 +400,7 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
}
fn extend_history_with_callbacks(&mut self, items: impl IntoIterator<Item = Item>) {
fn append_history_items(&mut self, items: impl IntoIterator<Item = Item>) {
for item in items {
self.emit_history_append(&item);
self.history.push(item);
@@ -985,7 +985,7 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
// get persisted by the upper layer that owns history.json.
let pending = self.interceptor.pending_history_appends().await;
if !pending.is_empty() {
self.extend_history_with_callbacks(pending);
self.append_history_items(pending);
}
// Clone the history into a per-request context. Everything
@@ -1093,7 +1093,7 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
self.turn_count += 1;
// Collect and commit assistant items. Routed through
// `extend_history_with_callbacks` so observers (e.g. the
// `append_history_items` so observers (e.g. the
// Pod-side per-item session-log committer) see each item
// as it lands.
let reasoning_items = self.reasoning_item_collector.take_collected();
@@ -1101,7 +1101,7 @@ impl<C: LlmClient, S: WorkerState> Worker<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.extend_history_with_callbacks(assistant_items);
self.append_history_items(assistant_items);
if tool_calls.is_empty() {
match self.interceptor.on_turn_end(&self.history).await {
@@ -1110,7 +1110,7 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
return Ok(WorkerResult::Finished);
}
TurnEndAction::ContinueWithMessages(additional) => {
self.extend_history_with_callbacks(additional);
self.append_history_items(additional);
continue;
}
TurnEndAction::Pause => {
@@ -1227,7 +1227,7 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
result.is_error,
)
});
self.extend_history_with_callbacks(items);
self.append_history_items(items);
Ok(None)
}
Err(err) => {
@@ -1411,38 +1411,28 @@ impl<C: LlmClient> Worker<C, Mutable> {
}
}
/// Get a mutable reference to history
/// Replace history during restore/rebuild without emitting append callbacks.
///
/// Available only in Mutable state.
pub fn history_mut(&mut self) -> &mut Vec<Item> {
&mut self.history
}
/// Set history
/// 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;
}
/// Add an item to history (builder pattern)
pub fn with_item(mut self, item: Item) -> Self {
self.history.push(item);
self
/// Append items to history and notify history-append observers for each
/// item before it lands. This is the only public Mutable-state API for
/// growing worker history; callers that need session-log persistence must
/// install [`on_history_append`](Self::on_history_append) before calling it.
pub fn append_history(&mut self, items: impl IntoIterator<Item = Item>) {
self.append_history_items(items);
}
/// Add an item to history
pub fn push_item(&mut self, item: Item) {
self.history.push(item);
}
/// Add multiple items to history (builder pattern)
pub fn with_items(mut self, items: impl IntoIterator<Item = Item>) -> Self {
self.history.extend(items);
self
}
/// Add multiple items to history
pub fn extend_history(&mut self, items: impl IntoIterator<Item = Item>) {
self.history.extend(items);
/// 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);
}
/// Clear history
@@ -1575,9 +1565,9 @@ impl<C: LlmClient> Worker<C, Locked> {
PromptAction::Continue => Vec::new(),
PromptAction::ContinueWith(items) => items,
};
self.history.push(user_item);
self.append_history_items(std::iter::once(user_item));
if !extras.is_empty() {
self.extend_history_with_callbacks(extras);
self.append_history_items(extras);
}
let result = self.run_turn_loop().await;
self.finalize_interruption(result).await
crates/llm-worker/tests/worker_state_test.rs
+34 -30
View File
@@ -5,8 +5,8 @@
mod common;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use common::MockLlmClient;
@@ -44,14 +44,12 @@ fn test_mutable_history_manipulation() {
assert!(worker.history().is_empty());
// Add to history
worker.push_item(Item::user_message("Hello"));
worker.push_item(Item::assistant_message("Hi there!"));
worker.append_history(vec![Item::user_message("Hello")]);
worker.append_history(vec![Item::assistant_message("Hi there!")]);
assert_eq!(worker.history().len(), 2);
// Mutable access to history
worker
.history_mut()
.push(Item::user_message("How are you?"));
// Append to history via the callback-aware API.
worker.append_history(vec![Item::user_message("How are you?")]);
assert_eq!(worker.history().len(), 3);
// Clear history
@@ -71,34 +69,38 @@ fn test_mutable_history_manipulation() {
#[test]
fn test_mutable_builder_pattern() {
let client = MockLlmClient::new(vec![]);
let worker = Worker::new(client)
.system_prompt("System prompt")
.with_item(Item::user_message("Hello"))
.with_item(Item::assistant_message("Hi!"))
.with_items(vec![
Item::user_message("How are you?"),
Item::assistant_message("I'm fine!"),
]);
let worker = Worker::new(client).system_prompt("System prompt");
assert_eq!(worker.get_system_prompt(), Some("System prompt"));
assert_eq!(worker.history().len(), 4);
assert!(worker.history().is_empty());
}
/// Verify that multiple items can be added with extend_history
/// Verify that multiple items can be added with append_history and callbacks fire.
#[test]
fn test_mutable_extend_history() {
fn test_mutable_append_history() {
let client = MockLlmClient::new(vec![]);
let observed = Arc::new(Mutex::new(Vec::new()));
let observed_for_callback = Arc::clone(&observed);
let mut worker = Worker::new(client);
worker.on_history_append(move |item| {
if let Some(text) = item.as_text() {
observed_for_callback.lock().unwrap().push(text.to_string());
}
});
worker.push_item(Item::user_message("First"));
worker.append_history(vec![Item::user_message("First")]);
worker.extend_history(vec![
worker.append_history(vec![
Item::assistant_message("Response 1"),
Item::user_message("Second"),
Item::assistant_message("Response 2"),
]);
assert_eq!(worker.history().len(), 4);
assert_eq!(
observed.lock().unwrap().as_slice(),
["First", "Response 1", "Second", "Response 2"]
);
}
#[derive(Clone)]
@@ -162,8 +164,8 @@ fn test_lock_transition() {
let mut worker = Worker::new(client);
worker.set_system_prompt("System");
worker.push_item(Item::user_message("Hello"));
worker.push_item(Item::assistant_message("Hi"));
worker.append_history(vec![Item::user_message("Hello")]);
worker.append_history(vec![Item::assistant_message("Hi")]);
// Lock
let locked_worker = worker.lock();
@@ -180,14 +182,14 @@ fn test_unlock_transition() {
let client = MockLlmClient::new(vec![]);
let mut worker = Worker::new(client);
worker.push_item(Item::user_message("Hello"));
worker.append_history(vec![Item::user_message("Hello")]);
let locked_worker = worker.lock();
// Unlock
let mut worker = locked_worker.unlock();
// History operations are available again in Mutable state
worker.push_item(Item::assistant_message("Hi"));
worker.append_history(vec![Item::assistant_message("Hi")]);
worker.clear_history();
assert!(worker.history().is_empty());
}
@@ -310,8 +312,8 @@ async fn test_locked_prefix_len_tracking() {
let mut worker = Worker::new(client);
// Add items beforehand
worker.push_item(Item::user_message("Pre-existing message 1"));
worker.push_item(Item::assistant_message("Pre-existing response 1"));
worker.append_history(vec![Item::user_message("Pre-existing message 1")]);
worker.append_history(vec![Item::assistant_message("Pre-existing response 1")]);
assert_eq!(worker.history().len(), 2);
@@ -380,9 +382,11 @@ async fn test_unlock_edit_relock() {
}),
]]);
let worker = Worker::new(client)
.with_item(Item::user_message("Hello"))
.with_item(Item::assistant_message("Hi"));
let mut worker = Worker::new(client);
worker.append_history(vec![
Item::user_message("Hello"),
Item::assistant_message("Hi"),
]);
// Lock -> Unlock
let locked = worker.lock();
@@ -392,7 +396,7 @@ async fn test_unlock_edit_relock() {
// Edit history
unlocked.clear_history();
unlocked.push_item(Item::user_message("Fresh start"));
unlocked.append_history(vec![Item::user_message("Fresh start")]);
// Re-lock
let relocked = unlocked.lock();