stats.py

Go to the documentation of this file.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# oc_ui/terminal/stats.py
 
"""
Real-time statistics tracking for the OC Chat enhanced terminal UI.
 
Purpose
-------
Provides the presentation-side statistics tracker that backs the
enhanced terminal header with real-time turn / session / context-usage
counters. Relocated here from ``oc_chat.core.stats`` by plan 153: the
tracker's sole owner and mutator is :class:`~oc_ui.terminal.enhanced_presenter.EnhancedPresenter`,
so it belongs in the ``oc_ui`` adapter layer, not in the harness core.
Keeping it here removes the ``oc_ui -> oc_chat.core.stats`` import edge
named as a boundary residual in plan 151 §8.
 
Responsibilities
-------------
- StatisticsTracker: Main class for tracking all metrics
- Token formatting utilities for K/MB/GB display
- Color-coded statistics for different metric types
- Real-time updates during sessions
 
Diagnostics
-----------
Domain: OC_UI.STATS
Levels:
    L2 — lifecycle events
    L3 — semantic details
    L4 — deep tracing
 
Contracts
---------
- All statistics are thread-safe and atomic
- Token counts are formatted with proper units
- Percentages are always displayed with % symbol
- Time values are formatted with appropriate units
- This module imports stdlib + oc_diagnostics only; it MUST NOT import
  ``oc_chat`` (the plan-153 boundary the DAG guard asserts).
 
Dependencies
------------
- oc_diagnostics (structured debug logging via _dbg)
"""
 
from __future__ import annotations
import threading
from typing import Dict, List, NamedTuple
from dataclasses import dataclass, field
 
from oc_diagnostics import _dbg
 
DEFAULT_CONTEXT_WINDOW_SIZE: int = 128_000
# OCT-LINT: disable=no_hardcoded_secrets reason="Docstring prose (context-window cap description) mis-flagged as a high-entropy string; not a secret"
"""Default context-window token cap for the enhanced UI's context meter.
 
UI-local fallback that mirrors the harness default of the same name in
``oc_chat.core.stats`` (consumed by ``cli.py``). The presenter always
overrides it via :meth:`StatisticsTracker.set_context_window_size` right
after construction — this constant only governs the meter before the
first bind. Kept UI-local so this module carries no ``oc_chat`` import.
"""
 
 
@dataclass
class TokenStats:
    """Token usage statistics."""
    prompt: int = 0
    completion: int = 0
    total: int = 0
    cache_hits: int = 0
    cache_hit_rate: float = 0.0
    duration_ms: int = 0
 
    def add_usage(
        self,
        usage_data: dict | int,
        completion: int = 0,
        cache_hits: int = 0,
        duration_ms: int = 0,
    ) -> None:
        """Add token usage from an API response dictionary or counters."""
        if isinstance(usage_data, dict):
            prompt = usage_data.get('prompt_tokens', 0)
            completion = usage_data.get('completion_tokens', 0)
            cache_hits = usage_data.get('cache_hit_tokens', 0)
            duration = usage_data.get('dur_ms', 0)
        else:
            prompt = usage_data
            duration = duration_ms
 
        self.prompt += prompt
        self.completion += completion
        self.total += prompt + completion
        self.cache_hits += cache_hits
        self.duration_ms = duration
 
        if self.total > 0:
            self.cache_hit_rate = (self.cache_hits / self.total) * 100
 
    def format_tokens(self, count: int) -> str:
        """Format token count with appropriate units."""
        if count < 1024:
            return f"{count}"
        elif count < 1_048_576:
            return f"{count // 1024}K"
        elif count < 1_073_741_824:
            return f"{count / 1_048_576:.1f}M"
        else:
            return f"{count / 1_073_741_824:.1f}G"
 
 
@dataclass
class TimingStats:
    """Timing and performance statistics."""
    duration_ms: int = 0
    tool_calls: int = 0
    iteration_count: int = 0
 
    def add_timing(self, duration_ms: int, tool_calls: int = 0) -> None:
        """Add timing information from a turn."""
        self.duration_ms = duration_ms
        self.tool_calls += tool_calls
 
 
@dataclass
class SessionStats:
    """Session-level statistics.
 
    Wall-clock session duration is NOT tracked here (plan 147 D1) — the
    old ``update_session_totals`` accumulator wrote the same
    ``timing.duration_ms`` field ``TimingStats.add_timing`` overwrites,
    a genuine overwrite-vs-accumulate conflict with zero readers. Real
    session-elapsed timing now lives on
    :class:`~oc_ui.terminal.input.TurnStats` (monotonic, process-local)
    and reaches the enhanced header via :class:`~oc_ui.terminal.statusline.StatusBarState`.
    """
    tokens: TokenStats = field(default_factory=TokenStats)
    timing: TimingStats = field(default_factory=TimingStats)
 
 
@dataclass
class AgentStats:
    """Agent/sub-loop tracking statistics."""
    active_agents: int = 0
    total_agents_launched: int = 0
    sub_loops_running: int = 0
    total_sub_loops_launched: int = 0
 
    def add_agent(self) -> None:
        """Add a new agent to tracking."""
        self.active_agents += 1
        self.total_agents_launched += 1
 
    def add_sub_loop(self) -> None:
        """Add a new sub-loop to tracking."""
        self.sub_loops_running += 1
        self.total_sub_loops_launched += 1
 
    def finish_agent(self) -> None:
        """Mark an agent as finished."""
        if self.active_agents > 0:
            self.active_agents -= 1
 
    def finish_sub_loop(self) -> None:
        """Mark a sub-loop as finished."""
        if self.sub_loops_running > 0:
            self.sub_loops_running -= 1
 
 
class StatisticsTracker:
    """Main statistics tracker for OC Chat enhanced UI."""
 
    def __init__(self) -> None:
        func = "StatisticsTracker.__init__"
        self._lock = threading.Lock()
        self._session_stats = SessionStats()
        self._turn_stats = TokenStats()
        self._agent_stats = AgentStats()
        self._context_usage = 0.0  # Percentage of context window used
        self._max_context_size = DEFAULT_CONTEXT_WINDOW_SIZE
        self._last_turn_tool_calls = 0
        _dbg("OC_UI.STATS", func, "initialized", 2)
 
    def update_turn_stats(
        self,
        prompt: int = 0,
        completion: int = 0,
        cache_hits: int = 0,
        duration_ms: int = 0,
        tool_calls: int = 0,
        **aliases,
    ) -> None:
        """Update statistics for a completed turn."""
        func = "update_turn_stats"
        prompt = aliases.pop("prompt_tokens", prompt)
        completion = aliases.pop("completion_tokens", completion)
        cache_hits = aliases.pop("cache_hit_tokens", cache_hits)
        cache_hits = aliases.pop("prompt_cached", cache_hits)
        with self._lock:
            self._turn_stats.add_usage(prompt, completion, cache_hits, duration_ms)
            self._session_stats.tokens.add_usage(prompt, completion, cache_hits)
            self._session_stats.timing.add_timing(duration_ms, tool_calls)
            self._session_stats.timing.iteration_count += 1
            self._last_turn_tool_calls = tool_calls
            _dbg("OC_UI.STATS", func,
                   f"turn: prompt={prompt}, completion={completion}, "
                   f"cache_hits={cache_hits}, duration={duration_ms}ms, tools={tool_calls}", 3)
 
    def update_context_usage(self, used_tokens: int) -> None:
        """Update context window usage percentage."""
        func = "update_context_usage"
        with self._lock:
            if self._max_context_size > 0:
                self._context_usage = min(100.0, (used_tokens / self._max_context_size) * 100)
            _dbg("OC_UI.STATS", func, f"context: {used_tokens}/{self._max_context_size} = {self._context_usage:.1f}%", 3)
 
    def add_agent(self) -> None:
        """Add a new agent to tracking."""
        func = "add_agent"
        with self._lock:
            self._agent_stats.add_agent()
            _dbg("OC_UI.STATS", func, f"active_agents={self._agent_stats.active_agents}", 3)
 
    def finish_agent(self) -> None:
        """Mark an agent as finished."""
        func = "finish_agent"
        with self._lock:
            self._agent_stats.finish_agent()
            _dbg("OC_UI.STATS", func, f"active_agents={self._agent_stats.active_agents}", 3)
 
    def add_sub_loop(self) -> None:
        """Add a new sub-loop to tracking."""
        func = "add_sub_loop"
        with self._lock:
            self._agent_stats.add_sub_loop()
            _dbg("OC_UI.STATS", func, f"sub_loops_running={self._agent_stats.sub_loops_running}", 3)
 
    def finish_sub_loop(self) -> None:
        """Mark a sub-loop as finished."""
        func = "finish_sub_loop"
        with self._lock:
            self._agent_stats.finish_sub_loop()
            _dbg("OC_UI.STATS", func, f"sub_loops_running={self._agent_stats.sub_loops_running}", 3)
 
    def get_stats_snapshot(self) -> Dict[str, any]:
        """Get a snapshot of all current statistics."""
        with self._lock:
            agents = {
                "active": self._agent_stats.active_agents,
                "total_launched": self._agent_stats.total_agents_launched,
                "sub_loops_running": self._agent_stats.sub_loops_running,
                "total_sub_loops": self._agent_stats.total_sub_loops_launched,
            }
            context = {
                "usage_percent": self._context_usage,
                "max_size": self._max_context_size,
            }
            last_turn = {
                "prompt": self._turn_stats.prompt,
                "completion": self._turn_stats.completion,
                "total": self._turn_stats.total,
                "cache_hit_rate": self._turn_stats.cache_hit_rate,
                "duration_ms": self._turn_stats.duration_ms,
                "tool_calls": self._last_turn_tool_calls,
            }
            return {
                "agents": agents,
                "context": context,
                "last_turn": last_turn,
                "session": {
                    "tokens": {
                        "prompt": self._session_stats.tokens.prompt,
                        "completion": self._session_stats.tokens.completion,
                        "total": self._session_stats.tokens.total,
                        "cache_hit_rate": self._session_stats.tokens.cache_hit_rate,
                    },
                    "timing": {
                        "duration_ms": self._session_stats.timing.duration_ms,
                        "tool_calls": self._session_stats.timing.tool_calls,
                        "iteration_count": self._session_stats.timing.iteration_count,
                    },
                    "agents": agents,
                    "context": context,
                    "last_turn": last_turn,
                }
            }
 
    def set_context_window_size(self, max_size: int) -> None:
        """Set the maximum context window size."""
        func = "set_context_window_size"
        with self._lock:
            self._max_context_size = max_size
            _dbg("OC_UI.STATS", func, f"max_context_size={max_size}", 3)