oclaw/interfaces/gateway/server_methods/usage.py

from __future__ import annotations

import time
from datetime import datetime, timedelta, timezone
from typing import Any

from .shared_types import GatewayRequestHandlers
from .validation import error_shape

COST_USAGE_CACHE_TTL_MS = 30_000
DAY_MS = 24 * 60 * 60 * 1000
_cost_usage_cache: dict[str, dict[str, Any]] = {}


def _ok(respond, payload: Any, meta: dict[str, Any] | None = None) -> None:
    if callable(respond):
        respond(True, payload, None, meta or None)


def _bad(respond, message: str) -> None:
    if callable(respond):
        respond(False, None, error_shape("INVALID_REQUEST", message), None)


def _unavailable(respond, message: str) -> None:
    if callable(respond):
        respond(False, None, error_shape("UNAVAILABLE", message), None)


def _parse_date_parts(raw: Any) -> tuple[int, int, int] | None:
    if not isinstance(raw, str) or not raw.strip():
        return None
    try:
        dt = datetime.strptime(raw.strip(), "%Y-%m-%d")
    except ValueError:
        return None
    return dt.year, dt.month, dt.day


def _parse_utc_offset_minutes(raw: Any) -> int | None:
    if not isinstance(raw, str) or not raw.strip():
        return None
    text = raw.strip()
    if not text.startswith("UTC"):
        return None
    sign_part = text[3:4]
    if sign_part not in {"+", "-"}:
        return None
    rest = text[4:]
    if ":" in rest:
        hh_s, mm_s = rest.split(":", 1)
    else:
        hh_s, mm_s = rest, "0"
    try:
        hh = int(hh_s)
        mm = int(mm_s)
    except ValueError:
        return None
    if hh > 14 or mm < 0 or mm >= 60:
        return None
    total = hh * 60 + mm
    if sign_part == "-":
        total = -total
    if total < -12 * 60 or total > 14 * 60:
        return None
    return total


def _resolve_date_mode(params: dict[str, Any]) -> dict[str, Any]:
    mode = params.get("mode")
    if mode == "gateway":
        return {"mode": "gateway"}
    if mode == "specific":
        offset = _parse_utc_offset_minutes(params.get("utcOffset"))
        if offset is not None:
            return {"mode": "specific", "utcOffsetMinutes": offset}
    return {"mode": "utc"}


def _parse_date_to_ms(raw: Any, interpretation: dict[str, Any]) -> int | None:
    parts = _parse_date_parts(raw)
    if not parts:
        return None
    y, m, d = parts
    if interpretation["mode"] == "gateway":
        return int(datetime(y, m, d).timestamp() * 1000)
    if interpretation["mode"] == "specific":
        offset = interpretation["utcOffsetMinutes"]
        base = datetime(y, m, d, tzinfo=timezone.utc).timestamp() * 1000
        return int(base - offset * 60 * 1000)
    return int(datetime(y, m, d, tzinfo=timezone.utc).timestamp() * 1000)


def _today_start_ms(now: datetime, interpretation: dict[str, Any]) -> int:
    if interpretation["mode"] == "gateway":
        local = datetime(now.year, now.month, now.day)
        return int(local.timestamp() * 1000)
    if interpretation["mode"] == "specific":
        offset = interpretation["utcOffsetMinutes"]
        shifted = now + timedelta(minutes=offset)
        start = datetime(shifted.year, shifted.month, shifted.day, tzinfo=timezone.utc)
        return int(start.timestamp() * 1000) - offset * 60 * 1000
    start = datetime(now.year, now.month, now.day, tzinfo=timezone.utc)
    return int(start.timestamp() * 1000)


def _parse_days(raw: Any) -> int | None:
    if isinstance(raw, (int, float)) and not isinstance(raw, bool):
        return int(raw)
    if isinstance(raw, str) and raw.strip():
        try:
            return int(float(raw.strip()))
        except ValueError:
            return None
    return None


def _parse_date_range(params: dict[str, Any]) -> dict[str, int]:
    now = datetime.now(timezone.utc)
    interpretation = _resolve_date_mode(params)
    today_start = _today_start_ms(now, interpretation)
    today_end = today_start + DAY_MS - 1

    start_ms = _parse_date_to_ms(params.get("startDate"), interpretation)
    end_ms = _parse_date_to_ms(params.get("endDate"), interpretation)
    if start_ms is not None and end_ms is not None:
        return {"startMs": start_ms, "endMs": end_ms + DAY_MS - 1}

    days = _parse_days(params.get("days"))
    if days is not None:
        clamped = max(1, days)
        return {"startMs": today_start - (clamped - 1) * DAY_MS, "endMs": today_end}

    return {"startMs": today_start - 29 * DAY_MS, "endMs": today_end}


def _load_cost_usage_summary_cached(start_ms: int, end_ms: int, context: Any) -> tuple[dict[str, Any], bool]:
    cache_key = f"{start_ms}-{end_ms}"
    now_ms = int(time.time() * 1000)
    cached = _cost_usage_cache.get(cache_key)
    if cached and now_ms - int(cached.get("updatedAt", 0)) < COST_USAGE_CACHE_TTL_MS:
        return dict(cached.get("summary") or {}), True

    hook = context.get("load_cost_usage_summary") if isinstance(context, dict) else None
    if callable(hook):
        summary = hook({"startMs": start_ms, "endMs": end_ms})
        if not isinstance(summary, dict):
            summary = {}
    else:
        summary = {
            "startMs": start_ms,
            "endMs": end_ms,
            "totals": {
                "input": 0,
                "output": 0,
                "cacheRead": 0,
                "cacheWrite": 0,
                "totalTokens": 0,
                "totalCost": 0,
            },
        }
    _cost_usage_cache[cache_key] = {"summary": summary, "updatedAt": now_ms}
    return summary, False


def _usage_status_handler(opts: dict[str, Any]) -> None:
    respond = opts.get("respond")
    context = opts.get("context")
    hook = context.get("load_provider_usage_summary") if isinstance(context, dict) else None
    try:
        summary = hook() if callable(hook) else {"providers": [], "generatedAt": int(time.time() * 1000)}
        _ok(respond, summary if isinstance(summary, dict) else {})
    except Exception as exc:
        _unavailable(respond, str(exc))


def _usage_cost_handler(opts: dict[str, Any]) -> None:
    respond = opts.get("respond")
    params = opts.get("params") or {}
    context = opts.get("context")
    if params is not None and not isinstance(params, dict):
        _bad(respond, "invalid usage.cost params")
        return
    date_range = _parse_date_range(dict(params))
    try:
        summary, was_cached = _load_cost_usage_summary_cached(date_range["startMs"], date_range["endMs"], context)
        _ok(respond, summary, {"cached": True} if was_cached else None)
    except Exception as exc:
        _unavailable(respond, str(exc))


def _sessions_usage_handler(opts: dict[str, Any]) -> None:
    respond = opts.get("respond")
    params = opts.get("params") or {}
    context = opts.get("context")
    if not isinstance(params, dict):
        _bad(respond, "invalid sessions.usage params")
        return
    limit = params.get("limit")
    limit = int(limit) if isinstance(limit, (int, float)) and not isinstance(limit, bool) else 50
    limit = max(1, min(limit, 500))
    specific_key = params.get("key")
    if specific_key is not None and not isinstance(specific_key, str):
        _bad(respond, "invalid sessions.usage params")
        return
    date_range = _parse_date_range(params)

    hook = context.get("load_sessions_usage") if isinstance(context, dict) else None
    if callable(hook):
        try:
            out = hook(
                {
                    "startMs": date_range["startMs"],
                    "endMs": date_range["endMs"],
                    "limit": limit,
                    "key": specific_key,
                    "includeContextWeight": bool(params.get("includeContextWeight", False)),
                }
            )
            _ok(respond, out if isinstance(out, dict) else {"sessions": [], "aggregates": {}, "range": date_range})
            return
        except Exception as exc:
            _unavailable(respond, str(exc))
            return

    payload = {
        "sessions": ([] if not specific_key else [{"key": specific_key, "sessionId": specific_key, "updatedAt": date_range["endMs"]}])[:limit],
        "aggregates": {
            "totals": {
                "input": 0,
                "output": 0,
                "cacheRead": 0,
                "cacheWrite": 0,
                "totalTokens": 0,
                "totalCost": 0,
                "missingCostEntries": 0,
            },
            "messages": {
                "total": 0,
                "user": 0,
                "assistant": 0,
                "toolCalls": 0,
                "toolResults": 0,
                "errors": 0,
            },
            "tools": [],
            "byModel": [],
            "byProvider": [],
            "byAgent": [],
            "byChannel": [],
            "daily": [],
            "latency": None,
            "dailyLatency": [],
            "modelDaily": [],
            "aggregateTail": [],
        },
        "range": date_range,
        "limit": limit,
    }
    _ok(respond, payload)


usage_handlers: GatewayRequestHandlers = {
    "usage.status": _usage_status_handler,
    "usage.cost": _usage_cost_handler,
    "sessions.usage": _sessions_usage_handler,
}