codewhale_telemetry/

counters.rs

//! Process-wide session accumulators.
//!
//! Atomics rather than a lock, because every bump happens on a hot path and
//! must cost nothing measurable. Every bump also happens at a **call site**,
//! never inside a conditionally-entered handler: a counter sited above an early
//! return is a landmine, and the natural future optimization — hoisting the
//! guard to the caller — silently zeroes it for exactly the users who do not
//! use that feature.

use std::sync::Mutex;
use std::sync::atomic::{AtomicU32, Ordering};

use crate::event::{Counters, Errors, TurnWall};

/// One session's accumulated counts.
#[derive(Debug, Default)]
pub struct SessionCounters {
    turns: AtomicU32,
    tool_calls: AtomicU32,
    fleet_dispatch: AtomicU32,
    workflow_run: AtomicU32,
    subagent_spawn: AtomicU32,
    mcp_server_connected: AtomicU32,
    memory_search: AtomicU32,
    approval_modal_shown: AtomicU32,
    approval_auto_allowed: AtomicU32,
    command_palette_open: AtomicU32,

    auth_preflight_failed: AtomicU32,
    provider_http_4xx: AtomicU32,
    provider_http_5xx: AtomicU32,
    tool_denied_by_policy: AtomicU32,
    tool_timeout: AtomicU32,
    network_error: AtomicU32,

    turn_wall: Mutex<TurnWall>,
    /// Sorted, deduplicated `ProviderKind::as_str()` values. `&'static str`
    /// only: a `String` here would be the seam through which a customer's
    /// `[providers.<name>]` table key reaches the wire.
    providers: Mutex<Vec<&'static str>>,
}

/// Which counter to bump. A closed enum so a call site cannot invent a key.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Counter {
    /// A model turn completed.
    Turns,
    /// A tool call was executed.
    ToolCalls,
    /// A fleet dispatch started.
    FleetDispatch,
    /// A `workflow_run` action was parsed and executed.
    WorkflowRun,
    /// A sub-agent was spawned.
    SubagentSpawn,
    /// An MCP server reached `connected`.
    McpServerConnected,
    /// A native-memory search ran.
    MemorySearch,
    /// An approval modal was shown.
    ApprovalModalShown,
    /// An approval was granted by an auto-allow rule.
    ApprovalAutoAllowed,
    /// The command palette was opened.
    CommandPaletteOpen,
}

/// Which error counter to bump.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ErrorCounter {
    /// Credential preflight rejected the route.
    AuthPreflightFailed,
    /// Provider responded 4xx.
    ProviderHttp4xx,
    /// Provider responded 5xx.
    ProviderHttp5xx,
    /// A tool call was denied by policy.
    ToolDeniedByPolicy,
    /// A tool call timed out.
    ToolTimeout,
    /// A request failed below HTTP.
    NetworkError,
}

impl SessionCounters {
    /// Bump a feature counter by one.
    pub fn bump(&self, counter: Counter) {
        self.slot(counter).fetch_add(1, Ordering::Relaxed);
    }

    /// Bump an error counter by one.
    pub fn bump_error(&self, counter: ErrorCounter) {
        self.error_slot(counter).fetch_add(1, Ordering::Relaxed);
    }

    /// Record one turn's wall-clock time in the histogram.
    pub fn observe_turn_secs(&self, secs: u64) {
        if let Ok(mut wall) = self.turn_wall.lock() {
            wall.observe_secs(secs);
        }
    }

    /// Record that a provider was routed to.
    ///
    /// Takes a `ProviderKind` **by value**, never a `&str`: the four
    /// persistence and label accessors that look like the natural seam all
    /// return the customer's own `[providers.<name>]` table key when the route
    /// is custom, and `/status` already prints it. `ProviderKind::Custom`
    /// yields the literal `"custom"` and nothing else.
    pub fn record_provider(&self, provider: codewhale_config::ProviderKind) {
        let name = provider.as_str();
        if let Ok(mut providers) = self.providers.lock()
            && !providers.contains(&name)
        {
            providers.push(name);
            providers.sort_unstable();
        }
    }

    /// Snapshot the feature counters.
    #[must_use]
    pub fn counters(&self) -> Counters {
        Counters {
            turns: self.turns.load(Ordering::Relaxed),
            tool_calls: self.tool_calls.load(Ordering::Relaxed),
            fleet_dispatch: self.fleet_dispatch.load(Ordering::Relaxed),
            workflow_run: self.workflow_run.load(Ordering::Relaxed),
            subagent_spawn: self.subagent_spawn.load(Ordering::Relaxed),
            mcp_server_connected: self.mcp_server_connected.load(Ordering::Relaxed),
            memory_search: self.memory_search.load(Ordering::Relaxed),
            approval_modal_shown: self.approval_modal_shown.load(Ordering::Relaxed),
            approval_auto_allowed: self.approval_auto_allowed.load(Ordering::Relaxed),
            command_palette_open: self.command_palette_open.load(Ordering::Relaxed),
        }
    }

    /// Snapshot the error counters.
    #[must_use]
    pub fn errors(&self) -> Errors {
        Errors {
            auth_preflight_failed: self.auth_preflight_failed.load(Ordering::Relaxed),
            provider_http_4xx: self.provider_http_4xx.load(Ordering::Relaxed),
            provider_http_5xx: self.provider_http_5xx.load(Ordering::Relaxed),
            tool_denied_by_policy: self.tool_denied_by_policy.load(Ordering::Relaxed),
            tool_timeout: self.tool_timeout.load(Ordering::Relaxed),
            network_error: self.network_error.load(Ordering::Relaxed),
        }
    }

    /// Snapshot the turn wall-clock histogram.
    #[must_use]
    pub fn turn_wall(&self) -> TurnWall {
        self.turn_wall.lock().map(|wall| *wall).unwrap_or_default()
    }

    /// Snapshot the sorted, deduplicated provider set.
    #[must_use]
    pub fn providers(&self) -> Vec<String> {
        self.providers
            .lock()
            .map(|providers| providers.iter().map(|name| (*name).to_string()).collect())
            .unwrap_or_default()
    }

    fn slot(&self, counter: Counter) -> &AtomicU32 {
        match counter {
            Counter::Turns => &self.turns,
            Counter::ToolCalls => &self.tool_calls,
            Counter::FleetDispatch => &self.fleet_dispatch,
            Counter::WorkflowRun => &self.workflow_run,
            Counter::SubagentSpawn => &self.subagent_spawn,
            Counter::McpServerConnected => &self.mcp_server_connected,
            Counter::MemorySearch => &self.memory_search,
            Counter::ApprovalModalShown => &self.approval_modal_shown,
            Counter::ApprovalAutoAllowed => &self.approval_auto_allowed,
            Counter::CommandPaletteOpen => &self.command_palette_open,
        }
    }

    fn error_slot(&self, counter: ErrorCounter) -> &AtomicU32 {
        match counter {
            ErrorCounter::AuthPreflightFailed => &self.auth_preflight_failed,
            ErrorCounter::ProviderHttp4xx => &self.provider_http_4xx,
            ErrorCounter::ProviderHttp5xx => &self.provider_http_5xx,
            ErrorCounter::ToolDeniedByPolicy => &self.tool_denied_by_policy,
            ErrorCounter::ToolTimeout => &self.tool_timeout,
            ErrorCounter::NetworkError => &self.network_error,
        }
    }
}

/// Classify a provider HTTP status into the 4xx or 5xx counter, if either
/// applies. Captured from `status.as_u16() / 100` **before** the error is
/// built — every `LlmError` variant carries the raw provider body verbatim.
#[must_use]
pub fn http_status_counter(status: u16) -> Option<ErrorCounter> {
    match status / 100 {
        4 => Some(ErrorCounter::ProviderHttp4xx),
        5 => Some(ErrorCounter::ProviderHttp5xx),
        _ => None,
    }
}