harness-rs-loop 0.0.67

//! `TelemetryHook` — maps the agent's lifecycle [`Event`] stream onto
//! structured `tracing` spans and events, so a run becomes observable in any
//! `tracing` subscriber.
//!
//! Why `tracing` and not a hard OpenTelemetry dependency? Because `tracing` is
//! the idiomatic Rust instrumentation seam: the library emits spans + events,
//! and the *binary* chooses the exporter. Attach
//! [`tracing-opentelemetry`](https://docs.rs/tracing-opentelemetry) with an
//! OTLP pipeline and every span below is exported to Jaeger / Tempo / any OTLP
//! backend with **zero changes here**; attach `tracing_subscriber::fmt().json()`
//! and you get newline-delimited JSON for log pipelines. One instrumentation,
//! many backends.
//!
//! Field names follow the OpenTelemetry **GenAI semantic conventions**
//! (`gen_ai.*`), so any OTel-aware backend — Logfire, SigNoz, Langfuse via OTLP,
//! Grafana — recognizes token counts, model, and finish reason automatically and
//! computes cost/latency with zero mapping. The pre-convention flat names
//! (`input_tokens`, `tool`, …) are kept alongside as aliases for existing
//! consumers.
//!
//! The spans come from [`AgentLoop`](crate::AgentLoop) itself — it knows when
//! a call really starts and ends, which a hook does not (parallel read-only
//! tools are dispatched before their `PreToolUse` fires). This hook adds the
//! events and the run totals. Span/event shape (target `harness.telemetry`):
//!
//! ```text
//! agent_run (span, from the loop: gen_ai.operation.name=invoke_agent,
//! │          gen_ai.request.model, gen_ai.provider.name, outcome)
//!   ├─ run.start         (source)
//!   ├─ iter              (iter)
//!   ├─ chat              (span, from the loop, one per model call:
//!   │   │                 otel.name="chat {model}", gen_ai.usage.*,
//!   │   │                 gen_ai.response.finish_reasons, tool_calls)
//!   │   └─ model.first_token (ttft_ms)
//!   ├─ model.complete    (gen_ai.operation.name=chat,
//!   │                     gen_ai.usage.input_tokens, gen_ai.usage.output_tokens,
//!   │                     gen_ai.usage.cached_input_tokens,
//!   │                     gen_ai.response.finish_reasons
//!   │                     + aliases: input_tokens, output_tokens,
//!   │                       cached_input_tokens, tool_calls, stop)
//!   ├─ execute_tool      (span, from the loop, one per tool call:
//!   │                     otel.name="execute_tool {tool}", gen_ai.tool.name,
//!   │                     gen_ai.tool.call.id, ok, otel.status_code=error on failure)
//!   ├─ tool.call         (gen_ai.operation.name=execute_tool, gen_ai.tool.name,
//!   │                     ok, duration_ms + alias: tool)
//!   ├─ sensor            (sensor, signals)
//!   ├─ compact           (stage, tokens_before, tokens_after, tokens_saved)
//!   ├─ budget.warning    (ratio)
//!   ├─ run.cancelled     (warn: the caller stopped the run)
//!   └─ run.end           (gen_ai.usage.*, total_tokens, model_calls, tool_calls,
//!                         tool_failures, compactions, tokens_saved, duration_ms)
//! ```
//!
//! To export over OTLP, enable the crate's `otel` feature and call
//! [`crate::otel::init_tracing_with_otlp`] from your binary; see that module.
//!
//! Wire it like any hook:
//! ```ignore
//! let loop_ = AgentLoop::new(model).with_hook(std::sync::Arc::new(TelemetryHook::new()));
//! ```

use harness_core::{Event, Hook, HookOutcome, World};
use std::collections::HashMap;
use std::sync::Mutex;
use std::time::Instant;

/// Emits a span per run and a structured event per model call, tool call,
/// sensor, compaction, and budget warning. See the module docs for the OTLP
/// bridge.
pub struct TelemetryHook {
    /// `call_id -> PreToolUse time`: which calls are in flight, for a cancel
    /// to settle, and a duration fallback for a loop that sends no `ToolTimed`.
    tool_starts: Mutex<HashMap<String, Instant>>,
    /// `call_id -> how long it really ran`, from `ToolTimed`. Parallel calls
    /// finish before their `PreToolUse` fires, so the fallback reads ~0 for them.
    tool_ran_ms: Mutex<HashMap<String, u64>>,
    /// When the current model call was handed off, so the run summary can say
    /// how much of the wall clock was spent waiting on the provider.
    model_start: Mutex<Option<Instant>>,
    /// True until the current streamed call produces its first fragment — the
    /// latency a person actually experiences, as distinct from how long the
    /// whole answer took.
    awaiting_first_token: Mutex<bool>,
    /// Running totals for the whole run, so `run.end` can answer "what did this
    /// cost?" without the reader summing per-turn lines by hand.
    totals: Mutex<RunTotals>,
}

/// What a run added up to, accumulated across its turns.
#[derive(Default)]
struct RunTotals {
    started: Option<Instant>,
    input_tokens: u64,
    output_tokens: u64,
    cached_input_tokens: u64,
    cache_write_input_tokens: u64,
    /// `model_calls` counts calls that *completed* — a cancelled one has no
    /// usage to attribute, so it adds only its wait to `model_ms`. `tool_calls`
    /// counts *dispatches*, including one cut short by a cancel, so it
    /// reconciles with `Outcome::Cancelled.tools_called`, which is incremented
    /// before the cancel check.
    model_calls: u64,
    tool_calls: u64,
    tool_failures: u64,
    compactions: u64,
    tokens_saved: u64,
    /// Wall-clock spent inside model calls, and the summed duration of tool
    /// calls. Tools dispatched in parallel overlap, so `tool_ms` can exceed the
    /// wall clock it occupied — it is a cost, not a span.
    model_ms: u64,
    tool_ms: u64,
    /// Read-only calls that exactly repeated an earlier one — wasted rounds the
    /// stuck-detector cannot see, because they are not consecutive.
    repeats: u64,
    /// Time to the first streamed fragment of the run's first model call: what
    /// the person watching waited before anything appeared. Zero when the run
    /// did not stream.
    first_token_ms: u64,
}

impl TelemetryHook {
    pub fn new() -> Self {
        Self {
            tool_starts: Mutex::new(HashMap::new()),
            tool_ran_ms: Mutex::new(HashMap::new()),
            model_start: Mutex::new(None),
            awaiting_first_token: Mutex::new(false),
            totals: Mutex::new(RunTotals::default()),
        }
    }
}

impl Default for TelemetryHook {
    fn default() -> Self {
        Self::new()
    }
}

impl Hook for TelemetryHook {
    fn name(&self) -> &str {
        "telemetry"
    }
    fn matches(&self, _ev: &Event<'_>) -> bool {
        true
    }

    fn fire(&self, ev: &Event<'_>, _world: &mut World) -> HookOutcome {
        match ev {
            Event::SessionStart { source } => {
                tracing::info!(
                    target: "harness.telemetry",
                    event = "run.start",
                    source = format!("{source:?}"),
                );
                *self.totals.lock().unwrap() = RunTotals {
                    started: Some(Instant::now()),
                    ..Default::default()
                };
                // Per-call state, too. A tool call a hook denied fired
                // PreToolUse and never PostToolUse, so its entry would sit here
                // for the hook's lifetime and a later cancel would count it as
                // a dispatch of *that* run. Every run starts from a clean slate.
                self.tool_starts.lock().unwrap().clear();
                self.tool_ran_ms.lock().unwrap().clear();
                *self.model_start.lock().unwrap() = None;
                *self.awaiting_first_token.lock().unwrap() = false;
            }
            Event::PreModel { .. } => {
                *self.model_start.lock().unwrap() = Some(Instant::now());
                *self.awaiting_first_token.lock().unwrap() = true;
            }
            Event::ModelTokenDelta { .. } => {
                // Only the first fragment of each call; the rest are noise.
                let mut awaiting = self.awaiting_first_token.lock().unwrap();
                if !*awaiting {
                    return HookOutcome::Allow;
                }
                *awaiting = false;
                drop(awaiting);
                let ttft = self
                    .model_start
                    .lock()
                    .unwrap()
                    .map(|s| s.elapsed().as_millis() as u64)
                    .unwrap_or(0);
                {
                    let mut t = self.totals.lock().unwrap();
                    // The first call's figure is the one a person felt; later
                    // calls are the agent thinking, not the answer starting.
                    if t.first_token_ms == 0 {
                        t.first_token_ms = ttft;
                    }
                }
                {
                    tracing::info!(
                        target: "harness.telemetry",
                        event = "model.first_token",
                        ttft_ms = ttft,
                    );
                };
            }
            Event::Heartbeat { iter } => {
                tracing::info!(target: "harness.telemetry", event = "iter", iter = *iter);
            }
            Event::PostModel { out } => {
                let waited = self
                    .model_start
                    .lock()
                    .unwrap()
                    .take()
                    .map(|s| s.elapsed().as_millis() as u64)
                    .unwrap_or(0);
                {
                    let mut t = self.totals.lock().unwrap();
                    t.model_calls += 1;
                    t.model_ms += waited;
                    t.input_tokens += out.usage.input_tokens as u64;
                    t.output_tokens += out.usage.output_tokens as u64;
                    t.cached_input_tokens += out.usage.cached_input_tokens as u64;
                    t.cache_write_input_tokens += out.usage.cache_write_input_tokens as u64;
                }
                let stop = format!("{:?}", out.stop_reason);
                tracing::info!(
                    target: "harness.telemetry",
                    event = "model.complete",
                    // OTel GenAI semantic conventions:
                    "gen_ai.operation.name" = "chat",
                    "gen_ai.usage.input_tokens" = out.usage.input_tokens,
                    "gen_ai.usage.output_tokens" = out.usage.output_tokens,
                    "gen_ai.usage.cached_input_tokens" = out.usage.cached_input_tokens,
                    "gen_ai.usage.cache_write_input_tokens" = out.usage.cache_write_input_tokens,
                    "gen_ai.response.finish_reasons" = %stop,
                    // pre-convention aliases:
                    input_tokens = out.usage.input_tokens,
                    output_tokens = out.usage.output_tokens,
                    cached_input_tokens = out.usage.cached_input_tokens,
                    cache_write_input_tokens = out.usage.cache_write_input_tokens,
                    tool_calls = out.tool_calls.len(),
                    stop = %stop,
                    duration_ms = waited,
                );
            }
            Event::PreToolUse { action } => {
                self.tool_starts
                    .lock()
                    .unwrap()
                    .insert(action.call_id.clone(), Instant::now());
            }
            Event::ToolTimed {
                action,
                started_ms,
                finished_ms,
            } => {
                self.tool_ran_ms.lock().unwrap().insert(
                    action.call_id.clone(),
                    finished_ms.saturating_sub(*started_ms).max(0) as u64,
                );
            }
            Event::PostToolUse { action, result } => {
                let since_pre = self
                    .tool_starts
                    .lock()
                    .unwrap()
                    .remove(&action.call_id)
                    .map(|s| s.elapsed().as_millis() as u64);
                let duration_ms = self
                    .tool_ran_ms
                    .lock()
                    .unwrap()
                    .remove(&action.call_id)
                    .or(since_pre)
                    .unwrap_or(0);
                let repeat = result
                    .content
                    .get("repeat_of_earlier_call")
                    .and_then(|v| v.as_bool())
                    .unwrap_or(false);
                {
                    let mut t = self.totals.lock().unwrap();
                    t.tool_calls += 1;
                    t.tool_ms += duration_ms;
                    if repeat {
                        t.repeats += 1;
                    }
                    if !result.ok {
                        t.tool_failures += 1;
                    }
                }
                {
                    tracing::info!(
                        target: "harness.telemetry",
                        event = "tool.call",
                        "gen_ai.operation.name" = "execute_tool",
                        "gen_ai.tool.name" = %action.tool,
                        ok = result.ok,
                        duration_ms,
                        tool = %action.tool, // alias
                    );
                };
            }
            Event::PostSensor { sensor, signals } => {
                tracing::debug!(
                    target: "harness.telemetry",
                    event = "sensor",
                    sensor = %sensor,
                    signals = signals.len(),
                );
            }
            Event::PostCompact {
                stage,
                before,
                after,
            } => {
                let saved = before.saturating_sub(*after);
                {
                    let mut t = self.totals.lock().unwrap();
                    t.compactions += 1;
                    t.tokens_saved += saved as u64;
                }
                // At info, not debug: compaction is what keeps a long run
                // affordable, and "it ran" without "it saved 12k" is not an
                // observation anyone can act on.
                tracing::info!(
                    target: "harness.telemetry",
                    event = "compact",
                    stage = format!("{stage:?}"),
                    tokens_before = *before,
                    tokens_after = *after,
                    tokens_saved = saved,
                );
            }
            Event::BudgetWarning { ratio } => {
                tracing::warn!(
                    target: "harness.telemetry",
                    event = "budget.warning",
                    ratio = *ratio,
                );
            }
            Event::Cancelled => {
                // A cancelled dispatch fired PreToolUse but will never fire
                // PostToolUse: the loop returns before it. Settle its entry
                // here, for two reasons. `run.end` must agree with
                // `Outcome::Cancelled.tools_called`, which counts that dispatch;
                // and `tool_starts` must not grow by one per cancel across a
                // long `Session` that reuses this hook.
                let dangling: Vec<Instant> = self
                    .tool_starts
                    .lock()
                    .unwrap()
                    .drain()
                    .map(|(_, started)| started)
                    .collect();
                // The same for a model call cut off between PreModel and
                // PostModel: the wait was real and belongs in `model_ms`, so
                // `duration_ms` does not exceed `model_ms + tool_ms` by an
                // unexplained gap; the call itself never completed, so
                // `model_calls` is left alone. Clearing the two fields is
                // hygiene — the next PreModel would overwrite them anyway.
                let cut_off_model = self.model_start.lock().unwrap().take();
                *self.awaiting_first_token.lock().unwrap() = false;
                {
                    let mut t = self.totals.lock().unwrap();
                    for started in dangling {
                        t.tool_calls += 1;
                        t.tool_ms += started.elapsed().as_millis() as u64;
                    }
                    if let Some(started) = cut_off_model {
                        t.model_ms += started.elapsed().as_millis() as u64;
                    }
                }
                {
                    // Warn, not info: a cancel is the person deciding the run was
                    // not worth finishing, which is worth seeing in a trace that
                    // would otherwise look like any other run.end.
                    tracing::warn!(target: "harness.telemetry", event = "run.cancelled");
                };
            }
            Event::SessionEnd => {
                let t = std::mem::take(&mut *self.totals.lock().unwrap());
                {
                    // One line with the whole bill. Per-turn events answer "what
                    // happened"; this answers "what did it cost", which is the
                    // question asked after every run and previously required
                    // adding the turns up by hand.
                    tracing::info!(
                        target: "harness.telemetry",
                        event = "run.end",
                        "gen_ai.usage.input_tokens" = t.input_tokens,
                        "gen_ai.usage.output_tokens" = t.output_tokens,
                        "gen_ai.usage.cached_input_tokens" = t.cached_input_tokens,
                        "gen_ai.usage.cache_write_input_tokens" = t.cache_write_input_tokens,
                        total_tokens = t.input_tokens + t.output_tokens,
                        model_calls = t.model_calls,
                        tool_calls = t.tool_calls,
                        tool_failures = t.tool_failures,
                        compactions = t.compactions,
                        tokens_saved = t.tokens_saved,
                        model_ms = t.model_ms,
                        tool_ms = t.tool_ms,
                        repeat_calls = t.repeats,
                        first_token_ms = t.first_token_ms,
                        duration_ms = t.started.map(|s| s.elapsed().as_millis() as u64).unwrap_or(0),
                    );
                };
            }
            _ => {}
        }
        HookOutcome::Allow
    }
}