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

fix: fence pending usage across segment forks

Browse Source
This commit is contained in:
Hare
2026-09-20 19:31:45 +09:00
parent 2dda4b90b4
commit b2ddab6d7e
2 changed files with 263 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/worker.rs
+38 -4
View File
@@ -4478,6 +4478,15 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
self.handle_worker_result(result, history_before).await
}
fn ensure_no_pending_usage_for_segment_replacement(&self) -> Result<(), WorkerError> {
let pending_records = self.usage_tracker.pending_record_count();
if pending_records == 0 {
Ok(())
} else {
Err(WorkerError::PendingUsagePersistence { pending_records })
}
}
/// Ensure the session exists and the writer's tally still matches
/// the on-disk entry count.
///
@@ -4519,6 +4528,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
if store_count == entries_written {
return Ok(());
}
self.ensure_no_pending_usage_for_segment_replacement()?;
// Auto-fork within the same Session: mint a fresh Segment and
// switch to it. The source segment is left immutable (no terminal
// marker is written back); the fork relationship is recorded
@@ -5146,10 +5156,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
mut cancel: Option<tokio::sync::watch::Receiver<bool>>,
trigger: CompactionTrigger,
) -> Result<SegmentId, WorkerError> {
let pending_records = self.usage_tracker.pending_record_count();
if pending_records > 0 {
return Err(WorkerError::PendingUsagePersistence { pending_records });
}
self.ensure_no_pending_usage_for_segment_replacement()?;
self.release_pending_compaction_service().await?;
let _rewrite_guard = self
.prepare_session_rewrite(SessionRewriteKind::Compact)
@@ -8626,6 +8633,33 @@ mod build_summary_prompt_tests {
assert!(append_error.to_string().contains("synthetic failure"));
assert_eq!(worker.usage_tracker.pending_record_count(), 1);
// Simulate another writer extending the active Segment. The normal run
// prelude must not auto-fork while the failed record still belongs to
// this Segment.
store
.append(
session_id,
source_segment,
&LogEntry::TurnEnd {
ts: segment_log::now_millis(),
turn_count: 99,
},
)
.unwrap();
let run_error = worker
.run_text("must stop before provider ownership")
.await
.expect_err("head drift must not fork past pending usage");
assert!(matches!(
run_error,
WorkerError::PendingUsagePersistence { pending_records: 1 }
));
assert_eq!(worker.segment_id(), source_segment);
assert_eq!(
store.list_segments(session_id).unwrap(),
vec![source_segment]
);
let error = worker
.compact(0)
.await
crates/worker/tests/compact_events_test.rs
+225 -1
View File
@@ -16,7 +16,7 @@ use agen::llm_client::types::Item;
use agen::llm_client::{ClientError, LlmClient, Request};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use async_trait::async_trait;
use futures::Stream;
use futures::{Stream, StreamExt};
use protocol::{Event, Method, RunResult};
use session_store::{
CombinedStore, FsStore, FsWorkerStore, LogEntry, Store, WorkerMetadata, WorkerMetadataStore,
@@ -173,6 +173,49 @@ impl LlmClient for MockClient {
}
}
#[derive(Clone)]
struct PauseAfterUsageClient {
responses: Arc<Vec<Vec<LlmEvent>>>,
call_count: Arc<AtomicUsize>,
hanging_call: usize,
}
impl PauseAfterUsageClient {
fn new(responses: Vec<Vec<LlmEvent>>, hanging_call: usize) -> Self {
Self {
responses: Arc::new(responses),
call_count: Arc::new(AtomicUsize::new(0)),
hanging_call,
}
}
}
#[async_trait]
impl LlmClient for PauseAfterUsageClient {
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
async fn stream(
&self,
_request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + Send>>, ClientError>
{
let count = self.call_count.fetch_add(1, Ordering::SeqCst);
let events = self
.responses
.get(count)
.cloned()
.ok_or_else(|| ClientError::Config("pause mock client exhausted".into()))?;
let finite = futures::stream::iter(events.into_iter().map(Ok));
if count == self.hanging_call {
Ok(Box::pin(finite.chain(futures::stream::pending())))
} else {
Ok(Box::pin(finite))
}
}
}
#[derive(Clone)]
struct GrowTool;
@@ -1214,6 +1257,187 @@ async fn completed_post_compact_usage_rearms_within_one_tool_loop() {
assert_eq!(restored.usage_history[0].input_total_tokens, 800);
}
#[tokio::test]
async fn post_compact_pause_bills_partial_without_rearming_or_consuming_correlation() {
let client = PauseAfterUsageClient::new(
vec![
text_events_with_usage("seed", 100_000),
write_summary_tool_use_events("manual-compact", "pause test summary"),
single_text_events("done"),
vec![
LlmEvent::text_block_start(0),
LlmEvent::text_delta(0, "partial"),
LlmEvent::usage(90, 2),
],
grow_tool_use_events("grow-after-resume", 80),
write_summary_tool_use_events("automatic-compact", "rearmed summary"),
single_text_events("done"),
text_events_with_usage("finished", 50),
],
3,
);
let call_count = Arc::clone(&client.call_count);
let manifest = MID_TURN_MANIFEST_TOML.replace(
"compact_request_threshold = 100",
"compact_threshold = 80000\ncompact_request_threshold = 50000",
);
let (worker, store) = make_worker_with_manifest_and_store(&manifest, client).await;
let session_id = worker.session_id();
let runtime_tmp = tempfile::tempdir().unwrap();
let bash_output_dir = runtime_tmp.path().join("bash-output");
let (handle, shutdown_receiver) =
WorkerController::spawn(worker, runtime_tmp.path(), &bash_output_dir)
.await
.unwrap();
let mut rx = handle.subscribe();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"seed history",
))
.await
.unwrap();
loop {
if matches!(
tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for seed")
.expect("event"),
Event::RunEnd {
result: RunResult::Finished
}
) {
break;
}
}
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"pause the first post-compact request",
))
.await
.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(2), async {
while call_count.load(Ordering::SeqCst) < 4 {
tokio::task::yield_now().await;
}
})
.await
.expect("post-compact provider request should start");
handle
.send(Method::Pause {
command: protocol::WorkerCommandEnvelope::new(1),
})
.await
.unwrap();
loop {
let event = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for pause")
.expect("event");
if matches!(
event,
Event::RunEnd {
result: RunResult::Paused
}
) {
break;
}
}
let paused_usage = store
.list_segments(session_id)
.unwrap()
.into_iter()
.flat_map(|segment_id| store.read_all(session_id, segment_id).unwrap())
.filter_map(|entry| match entry {
LogEntry::LlmUsage {
input_total_tokens, ..
} => Some(input_total_tokens),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(
paused_usage.iter().filter(|tokens| **tokens == 90).count(),
1
);
let paused_metrics = session_metrics::read_session_metrics(&store, session_id).unwrap();
assert!(
paused_metrics
.iter()
.all(|record| record.metric.name != "compact.post_request"),
"partial Pause usage must not consume compact correlation or rearm"
);
handle
.send(Method::Resume {
command: protocol::WorkerCommandEnvelope::new(2),
})
.await
.unwrap();
loop {
if matches!(
tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for resume completion")
.expect("event"),
Event::RunEnd {
result: RunResult::Finished
}
) {
break;
}
}
assert_eq!(
call_count.load(Ordering::SeqCst),
8,
"the completed resumed request must rearm and permit request-threshold compaction"
);
let metrics = session_metrics::read_session_metrics(&store, session_id).unwrap();
let compact_starts = metrics
.iter()
.filter(|record| record.metric.name == "compact.start")
.collect::<Vec<_>>();
assert_eq!(compact_starts.len(), 2);
let pre_run_correlation = compact_starts
.iter()
.find(|record| {
record.metric.dimensions.get("trigger").map(String::as_str) == Some("pre_run")
})
.and_then(|record| record.metric.correlation_id.as_deref())
.expect("pre-run automatic compact correlation");
let pre_run_posts = metrics
.iter()
.filter(|record| {
record.metric.name == "compact.post_request"
&& record.metric.correlation_id.as_deref() == Some(pre_run_correlation)
})
.collect::<Vec<_>>();
assert_eq!(pre_run_posts.len(), 1);
assert_eq!(
pre_run_posts[0].metric.dimensions["input_total_tokens"],
"80"
);
assert!(metrics.iter().all(|record| {
record.metric.name != "compact.post_request"
|| record.metric.dimensions["input_total_tokens"] != "90"
}));
handle
.send(Method::Shutdown {
command: protocol::WorkerCommandEnvelope::new(3),
})
.await
.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(2), shutdown_receiver)
.await
.expect("controller shutdown timeout")
.expect("shutdown confirmation");
}
#[tokio::test]
async fn compacted_context_above_request_threshold_fails_before_provider_request() {
let client = MockClient::new(vec![