diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 298067a13..7a08f05ef 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -4478,6 +4478,15 @@ impl Worker { 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 Worker { 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 Worker { mut cancel: Option>, trigger: CompactionTrigger, ) -> Result { - 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 diff --git a/crates/worker/tests/compact_events_test.rs b/crates/worker/tests/compact_events_test.rs index 51ec551cb..7e086533f 100644 --- a/crates/worker/tests/compact_events_test.rs +++ b/crates/worker/tests/compact_events_test.rs @@ -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>>, + call_count: Arc, + hanging_call: usize, +} + +impl PauseAfterUsageClient { + fn new(responses: Vec>, 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 { + Box::new(self.clone()) + } + + async fn stream( + &self, + _request: Request, + ) -> Result> + 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::>(); + 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::>(); + 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::>(); + 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![