fix: integrate post-compact usage boundary

2026-09-20 19:36:48 +09:00
parent c492cd9a94 b2ddab6d7e
commit 4e7f48a389
8 changed files with 1209 additions and 108 deletions
@@ -7,15 +7,16 @@
 //!   (driven by `compact_request_threshold` → `PreRequestAction::Yield`)

 use std::pin::Pin;
-use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};

 use agen::Engine;
 use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent, UsageEvent};
 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,
@@ -125,10 +126,13 @@ fn annotated(item: Item) -> session_store::LoggedHistoryEntry {
    }
 }

+type CallObserver = Arc<dyn Fn(usize) + Send + Sync>;
+
 #[derive(Clone)]
 struct MockClient {
    responses: Arc<Vec<Vec<LlmEvent>>>,
    call_count: Arc<AtomicUsize>,
+    observer: Arc<Mutex<Option<CallObserver>>>,
 }

 impl MockClient {
@@ -136,8 +140,13 @@ impl MockClient {
        Self {
            responses: Arc::new(responses),
            call_count: Arc::new(AtomicUsize::new(0)),
+            observer: Arc::new(Mutex::new(None)),
        }
    }
+
+    fn set_observer(&self, observer: impl Fn(usize) + Send + Sync + 'static) {
+        *self.observer.lock().unwrap() = Some(Arc::new(observer));
+    }
 }

 #[async_trait]
@@ -152,6 +161,9 @@ impl LlmClient for MockClient {
    ) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + Send>>, ClientError>
    {
        let count = self.call_count.fetch_add(1, Ordering::SeqCst);
+        if let Some(observer) = self.observer.lock().unwrap().clone() {
+            observer(count);
+        }
        if count >= self.responses.len() {
            return Err(ClientError::Config("mock client exhausted".into()));
        }
@@ -161,6 +173,90 @@ 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;
+
+#[async_trait]
+impl Tool for GrowTool {
+    async fn execute(
+        &self,
+        _input: &str,
+        _ctx: agen::tool::ToolExecutionContext,
+    ) -> Result<ToolOutput, ToolError> {
+        Ok(ToolOutput {
+            summary: "grew context".into(),
+            content: Some("x".repeat(400_000)),
+            attachments: Vec::new(),
+        })
+    }
+}
+
+fn grow_tool_definition() -> ToolDefinition {
+    Arc::new(|| {
+        (
+            ToolMeta::new("grow")
+                .description("test tool that grows request context")
+                .input_schema(serde_json::json!({"type": "object"})),
+            Arc::new(GrowTool) as Arc<dyn Tool>,
+        )
+    })
+}
+
+fn grow_tool_use_events(call_id: &str, input_tokens: u64) -> Vec<LlmEvent> {
+    vec![
+        LlmEvent::tool_use_start(0, call_id, "grow"),
+        LlmEvent::tool_input_delta(0, "{}"),
+        LlmEvent::tool_use_stop(0),
+        LlmEvent::usage(input_tokens, 1),
+        LlmEvent::Status(StatusEvent {
+            status: ResponseStatus::Completed,
+        }),
+    ]
+}
+
 #[derive(Clone)]
 struct BlockingCompactClient {
    calls: Arc<AtomicUsize>,
@@ -365,12 +461,13 @@ where
    let scope = worker::Scope::writable(&pwd).unwrap();
    std::mem::forget(pwd_tmp);

-    let worker =
+    let mut engine =
        Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client);
+    engine.register_tool(grow_tool_definition());
    let observed_store = store.clone();
    let mut worker = Worker::new(
        manifest,
-        worker,
+        engine,
        store,
        worker::WorkerWorkspaceContext::local_filesystem(None),
        worker::WorkerFilesystemAuthority::local(pwd.clone(), pwd.clone()),
@@ -379,6 +476,7 @@ where
    .await
    .unwrap();
    worker.enable_worker_metadata_write_through().unwrap();
+    worker.attach_log_writer(Arc::new(worker.log_writer_handle()));
    (worker, observed_store)
 }

@@ -1057,6 +1155,289 @@ async fn request_threshold_compact_publishes_runtime_progress() {
    assert_eq!(post.metric.correlation_id.as_deref(), Some(correlation_id));
 }

+#[tokio::test]
+async fn completed_post_compact_usage_rearms_within_one_tool_loop() {
+    let client = MockClient::new(vec![
+        text_events_with_usage("seed", 100_000),
+        write_summary_tool_use_events("compact-1", "first summary"),
+        single_text_events("done"),
+        grow_tool_use_events("grow-1", 1_000),
+        write_summary_tool_use_events("compact-2", "second summary"),
+        single_text_events("done"),
+        text_events_with_usage("finished", 800),
+    ]);
+    let call_count = Arc::clone(&client.call_count);
+    let observed_client = client.clone();
+    let manifest = MID_TURN_MANIFEST_TOML.replace(
+        "compact_request_threshold = 100",
+        "compact_threshold = 80000\ncompact_request_threshold = 50000",
+    );
+    let (mut worker, store) = make_worker_with_manifest_and_store(&manifest, client).await;
+    let session_id = worker.session_id();
+    let durable_at_next_request = Arc::new(AtomicBool::new(false));
+    let observed_barrier = Arc::clone(&durable_at_next_request);
+    let observed_store = store.clone();
+    observed_client.set_observer(move |call_index| {
+        if call_index != 4 {
+            return;
+        }
+        let usage_is_durable = observed_store
+            .list_segments(session_id)
+            .unwrap()
+            .into_iter()
+            .flat_map(|segment_id| observed_store.read_all(session_id, segment_id).unwrap())
+            .any(|entry| {
+                matches!(
+                    entry,
+                    LogEntry::LlmUsage {
+                        input_total_tokens: 1_000,
+                        ..
+                    }
+                )
+            });
+        assert!(
+            usage_is_durable,
+            "the first post-compact usage must be durable before the next provider request"
+        );
+        observed_barrier.store(true, Ordering::SeqCst);
+    });
+
+    worker.run_text("seed input").await.unwrap();
+    worker
+        .run_text("continue through two compactions")
+        .await
+        .expect("a successful post-compact request must rearm before the tool-loop threshold");
+
+    assert_eq!(
+        call_count.load(Ordering::SeqCst),
+        7,
+        "both compactor cycles and both normal post-compact requests must run"
+    );
+    assert!(durable_at_next_request.load(Ordering::SeqCst));
+
+    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);
+    assert!(compact_starts.iter().any(|record| {
+        record.metric.dimensions.get("trigger").map(String::as_str) == Some("pre_run")
+    }));
+    assert!(compact_starts.iter().any(|record| {
+        record.metric.dimensions.get("trigger").map(String::as_str) == Some("request_threshold")
+    }));
+    let compact_posts = metrics
+        .iter()
+        .filter(|record| record.metric.name == "compact.post_request")
+        .collect::<Vec<_>>();
+    assert_eq!(compact_posts.len(), 2);
+    assert_ne!(
+        compact_posts[0].metric.correlation_id,
+        compact_posts[1].metric.correlation_id
+    );
+
+    let 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!(usage.iter().filter(|tokens| **tokens == 1_000).count(), 1);
+    assert_eq!(usage.iter().filter(|tokens| **tokens == 800).count(), 1);
+
+    let restored = session_store::restore(&store, session_id, worker.segment_id()).unwrap();
+    assert_eq!(restored.usage_history.len(), 1);
+    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![
@@ -3,7 +3,9 @@ use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex};

 use agen::Engine;
-use agen::llm_client::event::{ErrorEvent, Event as LlmEvent, ResponseStatus, StatusEvent};
+use agen::llm_client::event::{
+    ErrorEvent, Event as LlmEvent, ResponseStatus, StatusEvent, UsageEvent,
+};
 use agen::llm_client::types::Item;
 use agen::llm_client::{ClientError, LlmClient, Request};
 use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
@@ -2580,18 +2582,32 @@ async fn drain_until<F: FnMut(&Event) -> bool>(
 /// the interrupted turn's history consistency.
 #[tokio::test]
 async fn pause_then_resume_preserves_notifications_and_history_consistency() {
-    // Response 1: hang after opening a text block (no stop / completed),
-    // so the Engine is parked inside the stream read and `cancel_rx`
-    // races it cleanly on Method::Pause.
+    // Response 1: report billable partial usage, then hang after opening a text
+    // block (no stop / completed). Pause must flush that usage without treating
+    // the provider request as normally completed.
    let hang = MockResponse::Hang(vec![
        LlmEvent::text_block_start(0),
        LlmEvent::text_delta(0, "partial..."),
+        LlmEvent::Usage(UsageEvent {
+            input_tokens: Some(90),
+            output_tokens: Some(2),
+            total_tokens: Some(92),
+            cache_read_input_tokens: Some(0),
+            cache_creation_input_tokens: Some(0),
+        }),
    ]);
    // Response 2: a clean assistant reply delivered on Resume.
    let ok = MockResponse::Complete(vec![
        LlmEvent::text_block_start(0),
        LlmEvent::text_delta(0, "resumed output"),
        LlmEvent::text_block_stop(0, None),
+        LlmEvent::Usage(UsageEvent {
+            input_tokens: Some(100),
+            output_tokens: Some(3),
+            total_tokens: Some(103),
+            cache_read_input_tokens: Some(0),
+            cache_creation_input_tokens: Some(0),
+        }),
        LlmEvent::Status(StatusEvent {
            status: ResponseStatus::Completed,
        }),
@@ -2645,6 +2661,23 @@ async fn pause_then_resume_preserves_notifications_and_history_consistency() {

    tokio::time::sleep(std::time::Duration::from_millis(50)).await;
    assert_eq!(handle.shared_state.catalog_status(), WorkerStatus::Paused);
+    let (paused_entries, _) = handle.sink.subscribe_with_snapshot();
+    let paused_usage = paused_entries
+        .iter()
+        .filter_map(|entry| match entry {
+            LogEntry::LlmUsage {
+                input_total_tokens,
+                output_tokens,
+                ..
+            } => Some((*input_total_tokens, *output_tokens)),
+            _ => None,
+        })
+        .collect::<Vec<_>>();
+    assert_eq!(
+        paused_usage,
+        vec![(90, 2)],
+        "Pause must durably account for usage already observed on the interrupted request"
+    );

    handle
        .send(Method::Notify {
@@ -2729,6 +2762,23 @@ async fn pause_then_resume_preserves_notifications_and_history_consistency() {
        })
        .unwrap_or_default();
    assert_eq!(assistant_text, "resumed output");
+    let (final_entries, _) = handle.sink.subscribe_with_snapshot();
+    let final_usage = final_entries
+        .iter()
+        .filter_map(|entry| match entry {
+            LogEntry::LlmUsage {
+                input_total_tokens,
+                output_tokens,
+                ..
+            } => Some((*input_total_tokens, *output_tokens)),
+            _ => None,
+        })
+        .collect::<Vec<_>>();
+    assert_eq!(
+        final_usage,
+        vec![(90, 2), (100, 3)],
+        "the paused request must remain a distinct billing record rather than merging into resume"
+    );
    let has_tool_call = history.iter().any(|i| i.is_tool_call());
    assert!(!has_tool_call, "no orphan tool_call in history");
 }