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

fix: recover partial session log writes

Browse Source
This commit is contained in:
Hare
2026-08-05 18:15:49 +09:00
parent 36df79e561
commit fd391ef705
10 changed files with 488 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/llm-engine/src/engine.rs
+45 -27
View File
@@ -50,6 +50,9 @@ pub enum EngineError {
/// Config warnings (unsupported options)
#[error("Config warnings: {}", .0.iter().map(|w| w.to_string()).collect::<Vec<_>>().join(", "))]
ConfigWarnings(Vec<ConfigWarning>),
/// A durable-history observer rejected an item before it entered history.
#[error("History append failed: {0}")]
HistoryAppend(String),
}
/// Tool registration error
@@ -222,10 +225,10 @@ pub struct Engine<C: LlmClient, S: EngineState = Mutable> {
/// truncation have been applied — i.e. on the same data that
/// enters history.
tool_result_cbs: Vec<Box<dyn Fn(&ToolResult) + Send + Sync>>,
/// History-append callbacks. Invoked for non-streamed items when they
/// are appended to persistent engine history, so upper layers can
/// broadcast those items using history itself as the source of truth.
history_append_cbs: Vec<Box<dyn Fn(&Item) + Send + Sync>>,
/// History-append callbacks. Invoked before non-streamed items enter
/// engine history. An error rejects the item and aborts the turn, allowing
/// upper layers to make durable storage the commit gate.
history_append_cbs: Vec<Box<dyn Fn(&Item) -> Result<(), String> + Send + Sync>>,
/// Request configuration (max_tokens, temperature, etc.)
request_config: RequestConfig,
/// Whether the previous run was interrupted
@@ -498,23 +501,31 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
}
}
/// Register a callback invoked for items appended directly to engine
/// history outside streaming timeline callbacks.
pub fn on_history_append(&mut self, callback: impl Fn(&Item) + Send + Sync + 'static) {
/// Register a fallible callback invoked before an item enters engine
/// history. Returning an error rejects that item and aborts the turn.
pub fn on_history_append(
&mut self,
callback: impl Fn(&Item) -> Result<(), String> + Send + Sync + 'static,
) {
self.history_append_cbs.push(Box::new(callback));
}
fn emit_history_append(&self, item: &Item) {
fn emit_history_append(&self, item: &Item) -> Result<(), EngineError> {
for cb in &self.history_append_cbs {
cb(item);
cb(item).map_err(EngineError::HistoryAppend)?;
}
Ok(())
}
fn append_history_items(&mut self, items: impl IntoIterator<Item = Item>) {
fn append_history_items(
&mut self,
items: impl IntoIterator<Item = Item>,
) -> Result<(), EngineError> {
for item in items {
self.emit_history_append(&item);
self.emit_history_append(&item)?;
self.history.push(item);
}
Ok(())
}
fn request_trace_payload(&self, request: &Request) -> Value {
@@ -1125,9 +1136,13 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
// These are committed *before* the per-request clone so they
// participate in the LLM request below and get persisted by
// the caller that owns durable history.
let pending = self.interceptor.pending_history_appends().await;
let pending = self
.interceptor
.pending_history_appends()
.await
.map_err(EngineError::HistoryAppend)?;
if !pending.is_empty() {
self.append_history_items(pending);
self.append_history_items(pending)?;
}
// Clone the history into a per-request context. Everything
@@ -1202,7 +1217,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(items.clone())?;
request_context.extend(items);
info!("Yielded by interceptor after pre-request history append");
for cb in &self.turn_end_cbs {
@@ -1220,7 +1235,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(items.clone())?;
request_context.extend(items);
}
PreRequestAction::Continue => {}
@@ -1280,7 +1295,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(assistant_items)?;
}
self.emit_llm_continuation(
current_llm_call,
@@ -1307,7 +1322,7 @@ 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(assistant_items)?;
if tool_calls.is_empty() {
match self.interceptor.on_turn_end(&self.history).await {
@@ -1316,7 +1331,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
return Ok(EngineResult::Finished);
}
TurnEndAction::ContinueWithMessages(additional) => {
self.append_history_items(additional);
self.append_history_items(additional)?;
continue;
}
TurnEndAction::Pause => {
@@ -1610,7 +1625,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
result.is_error,
)
});
self.append_history_items(items);
self.append_history_items(items)?;
Ok(None)
}
Err(err) => {
@@ -1815,12 +1830,15 @@ impl<C: LlmClient> Engine<C, Mutable> {
self.history = items;
}
/// 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 engine 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);
/// 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(
&mut self,
items: impl IntoIterator<Item = Item>,
) -> Result<(), EngineError> {
self.append_history_items(items)
}
/// Truncate history without emitting append callbacks.
@@ -1969,9 +1987,9 @@ 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(std::iter::once(user_item))?;
if !extras.is_empty() {
self.append_history_items(extras);
self.append_history_items(extras)?;
}
let result = self.run_turn_loop().await;
self.finalize_interruption(result).await
crates/llm-engine/src/interceptor.rs
+2 -2
View File
@@ -158,8 +158,8 @@ pub trait Interceptor: Send + Sync {
/// reproducible per-request transformations (pruning, content
/// trimming, cache anchors) that depend only on the existing
/// history.
async fn pending_history_appends(&self) -> Vec<Item> {
Vec::new()
async fn pending_history_appends(&self) -> Result<Vec<Item>, String> {
Ok(Vec::new())
}
/// Called before each LLM request. The context starts as a clone
crates/llm-engine/tests/engine_state_test.rs
+81 -20
View File
@@ -44,12 +44,18 @@ fn test_mutable_history_manipulation() {
assert!(engine.history().is_empty());
// Add to history
engine.append_history(vec![Item::user_message("Hello")]);
engine.append_history(vec![Item::assistant_message("Hi there!")]);
engine
.append_history(vec![Item::user_message("Hello")])
.unwrap();
engine
.append_history(vec![Item::assistant_message("Hi there!")])
.unwrap();
assert_eq!(engine.history().len(), 2);
// Append to history via the callback-aware API.
engine.append_history(vec![Item::user_message("How are you?")]);
engine
.append_history(vec![Item::user_message("How are you?")])
.unwrap();
assert_eq!(engine.history().len(), 3);
// Clear history
@@ -86,15 +92,20 @@ fn test_mutable_append_history() {
if let Some(text) = item.as_text() {
observed_for_callback.lock().unwrap().push(text.to_string());
}
Ok(())
});
engine.append_history(vec![Item::user_message("First")]);
engine
.append_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"),
]);
engine
.append_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!(
@@ -157,6 +168,40 @@ fn test_mutable_can_register_tool() {
engine.register_tool(tool.definition());
}
/// A durable-history failure on a tool call must stop the turn before the
/// tool can produce an external side effect.
#[tokio::test]
async fn history_append_failure_stops_before_tool_execution() {
let client = MockLlmClient::new(vec![
Event::tool_use_start(0, "call_1", "count_tool"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]);
let tool = CountingTool::new("count_tool");
let mut engine = Engine::new(client);
engine.register_tool(tool.definition());
engine.on_history_append(|item| {
if item.is_tool_call() {
Err("simulated ENOSPC".to_string())
} else {
Ok(())
}
});
let mut engine = engine.lock();
let error = engine.run("use the tool").await.unwrap_err();
assert!(
matches!(error, EngineError::HistoryAppend(ref message) 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"));
}
// =============================================================================
// State Transition Tests
// =============================================================================
@@ -168,8 +213,12 @@ fn test_lock_transition() {
let mut engine = Engine::new(client);
engine.set_system_prompt("System");
engine.append_history(vec![Item::user_message("Hello")]);
engine.append_history(vec![Item::assistant_message("Hi")]);
engine
.append_history(vec![Item::user_message("Hello")])
.unwrap();
engine
.append_history(vec![Item::assistant_message("Hi")])
.unwrap();
// Lock
let locked_engine = engine.lock();
@@ -186,14 +235,18 @@ fn test_unlock_transition() {
let client = MockLlmClient::new(vec![]);
let mut engine = Engine::new(client);
engine.append_history(vec![Item::user_message("Hello")]);
engine
.append_history(vec![Item::user_message("Hello")])
.unwrap();
let locked_engine = engine.lock();
// Unlock
let mut engine = locked_engine.unlock();
// History operations are available again in Mutable state
engine.append_history(vec![Item::assistant_message("Hi")]);
engine
.append_history(vec![Item::assistant_message("Hi")])
.unwrap();
engine.clear_history();
assert!(engine.history().is_empty());
}
@@ -316,8 +369,12 @@ async fn test_locked_prefix_len_tracking() {
let mut engine = Engine::new(client);
// Add items beforehand
engine.append_history(vec![Item::user_message("Pre-existing message 1")]);
engine.append_history(vec![Item::assistant_message("Pre-existing response 1")]);
engine
.append_history(vec![Item::user_message("Pre-existing message 1")])
.unwrap();
engine
.append_history(vec![Item::assistant_message("Pre-existing response 1")])
.unwrap();
assert_eq!(engine.history().len(), 2);
@@ -387,10 +444,12 @@ async fn test_unlock_edit_relock() {
]]);
let mut engine = Engine::new(client);
engine.append_history(vec![
Item::user_message("Hello"),
Item::assistant_message("Hi"),
]);
engine
.append_history(vec![
Item::user_message("Hello"),
Item::assistant_message("Hi"),
])
.unwrap();
// Lock -> Unlock
let locked = engine.lock();
@@ -400,7 +459,9 @@ async fn test_unlock_edit_relock() {
// Edit history
unlocked.clear_history();
unlocked.append_history(vec![Item::user_message("Fresh start")]);
unlocked
.append_history(vec![Item::user_message("Fresh start")])
.unwrap();
// Re-lock
let relocked = unlocked.lock();