feat: session-metrics実装

2026-05-03 15:10:43 +09:00
parent 702ed79517
commit 70c4f1930e
15 changed files with 982 additions and 28 deletions
@@ -0,0 +1,367 @@
+//! End-to-end coverage for the prune-projection metrics path.
+//!
+//! Drives a Pod with a scripted mock LLM client and a custom tool that
+//! returns a long `ToolOutput.content`, then inspects the persisted
+//! session log to verify:
+//!
+//! - `prune.skip { reason: "no_candidates" }` lands when the protected-turn
+//!   window covers the entire history.
+//! - `prune.fire` lands once enough turns + usage measurements exist for
+//!   the projection to actually apply.
+//! - The fire metric and the immediately-following `prune.post_request`
+//!   metric share the same `correlation_id`, so cache_read / cache_write
+//!   from the LlmUsage that triggered the projection can be joined back
+//!   to the originating event.
+//! - `prune.skip { reason: "below_min_savings" }` lands when candidates
+//!   exist but their estimated savings are below the configured floor.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use async_trait::async_trait;
+use futures::Stream;
+use llm_worker::Worker;
+use llm_worker::llm_client::event::{
+    Event as LlmEvent, ResponseStatus, StatusEvent, UsageEvent,
+};
+use llm_worker::llm_client::{ClientError, LlmClient, Request};
+use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
+use session_metrics::{DOMAIN, Metric, metrics_from_extensions};
+use session_store::FsStore;
+
+use pod::{Pod, PodManifest};
+
+#[derive(Clone)]
+struct MockClient {
+    responses: Arc<Vec<Vec<LlmEvent>>>,
+    call_count: Arc<AtomicUsize>,
+}
+
+impl MockClient {
+    fn new(responses: Vec<Vec<LlmEvent>>) -> Self {
+        Self {
+            responses: Arc::new(responses),
+            call_count: Arc::new(AtomicUsize::new(0)),
+        }
+    }
+}
+
+#[async_trait]
+impl LlmClient for MockClient {
+    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);
+        if count >= self.responses.len() {
+            return Err(ClientError::Config("mock client exhausted".into()));
+        }
+        let events = self.responses[count].clone();
+        let stream = futures::stream::iter(events.into_iter().map(Ok));
+        Ok(Box::pin(stream))
+    }
+}
+
+/// Tool that returns a fixed `ToolOutput { summary, content: Some(big) }`.
+/// `content` is long enough for prune savings to comfortably clear small
+/// `min_savings` thresholds.
+struct BigContentTool {
+    summary: &'static str,
+    content: String,
+}
+
+#[async_trait]
+impl Tool for BigContentTool {
+    async fn execute(&self, _input: &str) -> Result<ToolOutput, ToolError> {
+        Ok(ToolOutput {
+            summary: self.summary.into(),
+            content: Some(self.content.clone()),
+        })
+    }
+}
+
+fn big_content_tool_definition(name: &'static str) -> ToolDefinition {
+    Arc::new(move || {
+        let summary = "tool result summary";
+        let content = "x".repeat(2048);
+        (
+            ToolMeta::new(name)
+                .description("test tool that returns a long content")
+                .input_schema(serde_json::json!({"type": "object"})),
+            Arc::new(BigContentTool { summary, content }) as Arc<dyn Tool>,
+        )
+    })
+}
+
+fn usage_event(input_total: u64, cache_read: u64, cache_write: u64, output: u64) -> LlmEvent {
+    LlmEvent::Usage(UsageEvent {
+        input_tokens: Some(input_total),
+        output_tokens: Some(output),
+        total_tokens: Some(input_total + output),
+        cache_read_input_tokens: Some(cache_read),
+        cache_creation_input_tokens: Some(cache_write),
+    })
+}
+
+/// Tool-call response from the assistant: emits a `tool_use` block then a
+/// usage event so usage_history gains a measurement on this turn.
+fn tool_use_response(call_id: &str, tool_name: &str) -> Vec<LlmEvent> {
+    vec![
+        LlmEvent::tool_use_start(0, call_id, tool_name),
+        LlmEvent::tool_input_delta(0, "{}"),
+        LlmEvent::tool_use_stop(0),
+        usage_event(500, 0, 0, 10),
+        LlmEvent::Status(StatusEvent {
+            status: ResponseStatus::Completed,
+        }),
+    ]
+}
+
+/// Plain text response with explicit cache_read/cache_write so that
+/// `prune.post_request` can carry meaningful values when this is the
+/// LLM call that follows a `prune.fire` event.
+fn text_response_with_cache(text: &str, cache_read: u64, cache_write: u64) -> Vec<LlmEvent> {
+    vec![
+        LlmEvent::text_block_start(0),
+        LlmEvent::text_delta(0, text),
+        LlmEvent::text_block_stop(0, None),
+        usage_event(800, cache_read, cache_write, 5),
+        LlmEvent::Status(StatusEvent {
+            status: ResponseStatus::Completed,
+        }),
+    ]
+}
+
+fn manifest_toml(prune_protected_turns: usize, prune_min_savings: u64) -> String {
+    format!(
+        r#"
+[pod]
+name = "test-pod"
+pwd = "./"
+
+[model]
+scheme = "anthropic"
+model_id = "test-model"
+
+[worker]
+max_tokens = 100
+
+[compaction]
+prune_protected_turns = {prune_protected_turns}
+prune_min_savings = {prune_min_savings}
+
+[[scope.allow]]
+target = "./"
+permission = "write"
+"#
+    )
+}
+
+async fn make_pod(
+    manifest_toml: String,
+    client: MockClient,
+    tool_name: &'static str,
+) -> (Pod<MockClient, FsStore>, tempfile::TempDir, tempfile::TempDir) {
+    let manifest = PodManifest::from_toml(&manifest_toml).unwrap();
+    let store_tmp = tempfile::tempdir().unwrap();
+    let store = FsStore::new(store_tmp.path()).await.unwrap();
+    let pwd_tmp = tempfile::tempdir().unwrap();
+    let pwd = pwd_tmp.path().to_path_buf();
+    let scope = pod::Scope::writable(&pwd).unwrap();
+
+    let mut worker = Worker::new(client);
+    worker.register_tool(big_content_tool_definition(tool_name));
+
+    let pod = Pod::new(manifest, worker, store, pwd, scope).await.unwrap();
+    (pod, store_tmp, pwd_tmp)
+}
+
+/// Drive Pod through enough runs to exercise both skip-no_candidates and
+/// fire branches, then read the session log back and assert the metric
+/// stream.
+#[tokio::test]
+async fn prune_metrics_emit_skip_then_fire_with_post_request_join() {
+    // Run 1 (request 0): tool_use → triggers tool execution → request 1
+    //   on the second iteration to produce the assistant reply.
+    // Run 2 (request 2): plain assistant text. Prune evaluation here
+    //   sees user1's tool_result outside the 1-protected-turn window and
+    //   should fire.
+    let client = MockClient::new(vec![
+        tool_use_response("call-1", "big_tool"),
+        text_response_with_cache("ok", 0, 200),
+        text_response_with_cache("done", 1234, 50),
+    ]);
+    let (mut pod, _store_tmp, _pwd_tmp) =
+        make_pod(manifest_toml(1, 1), client, "big_tool").await;
+    let session_id = pod.session_id();
+    // Cloning the store handle to read the session log back after the
+    // runs complete — the Pod retains its own copy.
+    let store = pod.store().clone();
+
+    pod.run_text("first").await.unwrap();
+    pod.run_text("second").await.unwrap();
+
+    let state = session_store::restore(&store, session_id).await.unwrap();
+    let metrics = metrics_from_extensions(&state.extensions);
+
+    // Run 1 has 2 LLM iterations (tool loop), each evaluates prune with
+    // only one user-message turn → 2x skip{no_candidates}.
+    // Run 2 has 1 LLM iteration with enough turns → 1x fire +
+    // 1x post_request paired by correlation_id.
+    let names: Vec<&str> = metrics.iter().map(|m| m.name.as_str()).collect();
+    assert!(
+        names.contains(&"prune.skip"),
+        "expected prune.skip in {names:?}"
+    );
+    assert!(
+        names.contains(&"prune.fire"),
+        "expected prune.fire in {names:?}"
+    );
+    assert!(
+        names.contains(&"prune.post_request"),
+        "expected prune.post_request in {names:?}"
+    );
+
+    // All skips in run 1 must record reason=no_candidates.
+    for m in metrics.iter().filter(|m| m.name == "prune.skip") {
+        assert_eq!(
+            m.dimensions.get("reason").map(String::as_str),
+            Some("no_candidates"),
+            "skip metric should be no_candidates here, got {m:?}"
+        );
+        assert!(m.correlation_id.is_none());
+    }
+
+    // The fire metric carries dimensions and correlation_id.
+    let fire = metrics
+        .iter()
+        .find(|m| m.name == "prune.fire")
+        .expect("prune.fire missing");
+    assert!(
+        fire.dimensions.contains_key("candidate_count"),
+        "fire missing candidate_count: {fire:?}"
+    );
+    assert!(
+        fire.dimensions.contains_key("border_turn"),
+        "fire missing border_turn: {fire:?}"
+    );
+    assert!(
+        fire.value.is_some(),
+        "fire missing estimated_savings value"
+    );
+    let fire_id = fire
+        .correlation_id
+        .as_ref()
+        .expect("fire metric missing correlation_id");
+
+    // Exactly one post_request metric should exist with the same id, and
+    // its value/dimension should reflect the cache numbers from the
+    // text_response_with_cache call (cache_read=1234, cache_write=50).
+    let post = metrics
+        .iter()
+        .find(|m| m.name == "prune.post_request")
+        .expect("prune.post_request missing");
+    assert_eq!(post.correlation_id.as_ref(), Some(fire_id));
+    assert_eq!(post.value, Some(1234.0));
+    assert_eq!(
+        post.dimensions.get("cache_write_tokens").map(String::as_str),
+        Some("50")
+    );
+    assert!(post.dimensions.contains_key("history_len"));
+}
+
+/// `min_savings` set high enough that candidates exist but the estimated
+/// savings always fall short → the second run should record
+/// `prune.skip { reason: "below_min_savings" }`.
+#[tokio::test]
+async fn prune_metrics_record_below_min_savings_skip() {
+    let client = MockClient::new(vec![
+        tool_use_response("call-1", "big_tool"),
+        text_response_with_cache("ok", 0, 100),
+        text_response_with_cache("done", 0, 0),
+    ]);
+    let (mut pod, _store_tmp, _pwd_tmp) =
+        make_pod(manifest_toml(1, u64::MAX), client, "big_tool").await;
+    let session_id = pod.session_id();
+    let store = pod.store().clone();
+
+    pod.run_text("first").await.unwrap();
+    pod.run_text("second").await.unwrap();
+
+    let state = session_store::restore(&store, session_id).await.unwrap();
+    let metrics = metrics_from_extensions(&state.extensions);
+    let below = metrics
+        .iter()
+        .find(|m| {
+            m.name == "prune.skip"
+                && m.dimensions.get("reason").map(String::as_str) == Some("below_min_savings")
+        })
+        .expect("expected prune.skip with reason=below_min_savings");
+    assert!(
+        below.dimensions.contains_key("candidate_count"),
+        "below_min_savings skip should report candidate_count: {below:?}"
+    );
+    assert!(
+        below.value.is_some(),
+        "below_min_savings skip should report estimated savings as value: {below:?}"
+    );
+    // No prune.fire for this scenario.
+    assert!(metrics.iter().all(|m| m.name != "prune.fire"));
+    // No prune.post_request either (no fire to join with).
+    assert!(metrics.iter().all(|m| m.name != "prune.post_request"));
+}
+
+/// Sessions that have no metrics in the log restore cleanly: the
+/// `RestoredState.extensions` simply contains no `metrics` domain
+/// payloads, and `metrics_from_extensions` returns an empty Vec.
+/// Backward-compatibility check for old logs predating this feature.
+#[tokio::test]
+async fn old_sessions_without_metrics_replay_cleanly() {
+    // Manifest without any `[compaction]` section → prune (and therefore
+    // the prune observer) is never installed, so no metrics get written.
+    let manifest_toml = r#"
+[pod]
+name = "test-pod"
+pwd = "./"
+
+[model]
+scheme = "anthropic"
+model_id = "test-model"
+
+[worker]
+max_tokens = 100
+
+[[scope.allow]]
+target = "./"
+permission = "write"
+"#;
+    let client = MockClient::new(vec![text_response_with_cache("hi", 0, 0)]);
+    let manifest = PodManifest::from_toml(manifest_toml).unwrap();
+    let store_tmp = tempfile::tempdir().unwrap();
+    let store = FsStore::new(store_tmp.path()).await.unwrap();
+    let pwd_tmp = tempfile::tempdir().unwrap();
+    let pwd = pwd_tmp.path().to_path_buf();
+    let scope = pod::Scope::writable(&pwd).unwrap();
+    let worker = Worker::new(client);
+    let mut pod = Pod::new(manifest, worker, store.clone(), pwd, scope)
+        .await
+        .unwrap();
+    let session_id = pod.session_id();
+    pod.run_text("hello").await.unwrap();
+
+    let state = session_store::restore(&store, session_id).await.unwrap();
+    let metrics = metrics_from_extensions(&state.extensions);
+    assert!(metrics.is_empty(), "no metrics should be recorded: {metrics:?}");
+    // And no extension entries at all in the metrics domain.
+    assert!(state.extensions.iter().all(|(d, _)| d != DOMAIN));
+
+    // Smoke check that fold helper is robust on a sentinel Metric value:
+    let m = Metric::now("smoke");
+    assert_eq!(m.name, "smoke");
+}