fix: persist post-compact usage at request boundary

crates/worker/tests/compact_events_test.rs

@@ -7,13 +7,14 @@
 //!   (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 protocol::{Event, Method, RunResult};
@@ -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,47 @@ impl LlmClient for MockClient {
     }
 }
 
+#[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 +418,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 +433,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 +1112,108 @@ 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 compacted_context_above_request_threshold_fails_before_provider_request() {
     let client = MockClient::new(vec![