okstra 0.146.1 → 0.148.0

package/runtime/python/okstra_token_usage/collect.py

@@ -20,7 +20,9 @@ from .antigravity import (
     find_antigravity_sessions,
 )
 from .paths import claude_project_dir, utc_now
-from .pricing import codex_cost_usd, antigravity_cost_usd
+from .pricing import antigravity_cost_usd, provider_cost_usd
+from okstra_ctl.models import provider_wrappers
+from okstra_ctl.lead_runtime import ARTIFACT_ONLY_LEAD_RUNTIMES
 from okstra_ctl.wrapper_status import read_wrapper_status, status_path_for_prompt
 
 
@@ -446,14 +448,16 @@ def _codex_worker_windows(project_root: Path, state: dict) -> dict[str, list[tup
     return windows
 
 
-def _codex_runtime_provider(worker: dict) -> str:
+def _cli_assignment_provider(worker: dict) -> str:
     worker_id = str(worker.get("workerId") or "").strip()
-    agent = str(worker.get("agent") or "").strip()
-    if worker_id == "report-writer":
-        return "codex"
-    if agent in {"codex", "antigravity"}:
-        return agent
-    return ""
+    explicit_provider = str(worker.get("provider") or "").strip()
+    provider = explicit_provider or str(worker.get("agent") or "").strip()
+    runner = str(worker.get("runner") or "").strip()
+    if runner and runner != "cli-wrapper":
+        return ""
+    if not explicit_provider and worker_id == "report-writer":
+        provider = "codex"
+    return provider if provider in provider_wrappers() else ""
 
 
 def _cli_sessions_for_windows(
@@ -482,8 +486,10 @@ def _cli_session_totals(provider: str, session_paths: list[Path]) -> list[dict]:
     for session_path in session_paths:
         if provider == "codex":
             total = codex_session_total(session_path)
-        else:
+        elif provider == "antigravity":
             total = antigravity_session_total(session_path)
+        else:
+            continue
         if total.get("available"):
             totals.append(total)
     return totals
@@ -496,8 +502,9 @@ def _cli_usage_block(provider: str, totals: dict, session_paths: list[Path]) ->
     if totals.get("model"):
         block["model"] = totals["model"]
         block["cliModel"] = totals["model"]
-    if provider == "codex":
-        cost = codex_cost_usd(
+    if provider in {"codex", "grok", "kimi"}:
+        cost = provider_cost_usd(
+            provider,
             totals.get("model"),
             totals.get("inputTokens", 0) or 0,
             totals.get("cachedInputTokens", 0) or 0,
@@ -627,7 +634,7 @@ def _attach_cli_usage(
 ) -> None:
     """Layer aggregated CLI token/cost onto a Claude-side worker ``block``.
 
-    A codex/antigravity worker re-dispatched within one aggregated wrapper
+    A CLI-wrapper worker re-dispatched within one aggregated wrapper
     window produces several rollout jsonls; pricing only the latest (the old
     ``find_codex_session`` -> ``sessions[-1]`` behavior) dropped the earlier
     attempts' tokens. We sum every in-window session here so the redispatched
@@ -653,18 +660,19 @@ def _attach_cli_usage(
             block[key] = cli[key]
 
 
-def _collect_codex_runtime_usage(state: dict, project_root: Path) -> dict:
+def _collect_cli_runtime_usage(state: dict, project_root: Path) -> dict:
     windows_by_worker = _codex_worker_windows(project_root, state)
     state["leadUsage"] = na_block(
-        "Codex lead token accounting is not available yet; worker CLI usage is collected from provider logs."
+        "Host lead token accounting is unavailable; CLI worker usage is collected only from attributable provider logs."
     )
     for worker in state.get("workers", []):
         if not isinstance(worker, dict):
             continue
-        provider = _codex_runtime_provider(worker)
+        provider = _cli_assignment_provider(worker)
         if not provider:
             worker["usage"] = na_block(
-                f"worker provider is not CLI-backed in Codex runtime: {worker.get('agent') or worker.get('workerId')}"
+                "worker is host-native or has no registered CLI provider: "
+                f"{worker.get('provider') or worker.get('agent') or worker.get('workerId')}"
             )
             continue
         worker_id = str(worker.get("workerId") or "").strip()
@@ -745,7 +753,7 @@ def _populate_usage_summary(
         "definitions": {
             "totalTokens": "Sum of input + output + cache_creation + cache_read tokens (raw processed volume; matches Anthropic API breakdown). Cache reads are 95%+ in long sessions.",
             "billableEquivalentTokens": "Tokens normalized to base-input-price units (cache_creation_5m x1.25, cache_creation_1h x2.0, cache_read x0.1, output x5). 5m vs 1h is split from usage.cache_creation when the API breakdown is present; otherwise all cache_creation falls into 5m.",
-            "estimatedCostUsd": "USD cost using public list pricing for the model recorded in the session. cliWorkers covers Codex/Antigravity CLI calls billed under those providers.",
+            "estimatedCostUsd": "USD cost using public list pricing for the model recorded in the session. cliWorkers covers attributable registered-provider CLI calls.",
         },
     }
 
@@ -761,8 +769,8 @@ def collect(team_state_path: Path, project_root: Path | None = None, *,
     task_key = state.get("taskKey", "")
     team_name = resolve_team_name(state)
     lead_sid = (state.get("lead") or {}).get("sessionId")
-    if state.get("leadRuntime") == "codex":
-        return _collect_codex_runtime_usage(state, cwd)
+    if state.get("leadRuntime") in ARTIFACT_ONLY_LEAD_RUNTIMES:
+        return _collect_cli_runtime_usage(state, cwd)
 
     # 1) Claude sessions (lead + claude-side workers). Cache totals at scan
     # time so we don't re-read the jsonl when a worker matches multiple
@@ -878,9 +886,10 @@ def collect(team_state_path: Path, project_root: Path | None = None, *,
             block["additionalSessionIds"] = [sid for sid, _, _ in matched[1:]]
             block["matchedAgentNames"] = sorted({t.get("agentName") for _, _, t in matched if t.get("agentName")})
 
-        # For codex/antigravity workers, sum every CLI session that fell inside
-        # the aggregated wrapper window (re-dispatches leave one rollout each).
-        if agent in ("codex", "antigravity"):
+        # For CLI-wrapper workers, retain execution evidence and attach any
+        # attributable provider transcript usage without treating absence as zero.
+        provider = _cli_assignment_provider(worker)
+        if provider:
             windows, status_path, used_fallback = _worker_cli_windows(
                 cwd,
                 worker,
@@ -888,7 +897,7 @@ def collect(team_state_path: Path, project_root: Path | None = None, *,
             )
             _attach_cli_usage(
                 block,
-                agent,
+                provider,
                 cwd,
                 windows,
                 status_path,

package/runtime/python/okstra_token_usage/pricing.py

@@ -4,18 +4,18 @@ Pricing is matched by substring against the model id recorded in the session
 transcript, so keys must reflect the *actual* model id form emitted by each
 provider:
 
-  * Anthropic — `claude-fable-5*`, `claude-opus-5*`, `claude-opus-4-*`, `claude-sonnet-4-*`,
+  * Anthropic — `claude-fable-5*`, `claude-opus-5*`, `claude-opus-4-*`, `claude-sonnet-5*`, `claude-sonnet-4-*`,
     `claude-haiku-4-5-*`, `claude-3-5-sonnet-*`, `claude-3-5-haiku-*`,
     `claude-3-opus-*`, `claude-3-haiku-*`.
   * OpenAI / Codex — `gpt-5*`, `gpt-4o*`, `gpt-4*`.
-  * Google / Gemini — `gemini-3.5-flash*`, `gemini-3.1-pro*`, `gemini-3-flash*`,
+  * Google / Gemini — `gemini-3.6-flash*`, `gemini-3.5-flash*`, `gemini-3.1-pro*`, `gemini-3-flash*`,
     `gemini-2.5-pro*`, `gemini-2.5-flash*`, `gemini-2.0-flash*`.
 
 Matching prefers the longest key that is a substring of the model id (most
 specific wins), so table ordering does not affect the result. Update when
 providers change list pricing.
 
-Sources (last verified 2026-05-17, public list prices, USD per 1M tokens):
+Sources (last verified 2026-08-03, public list prices, USD per 1M tokens):
   * Anthropic: https://www.anthropic.com/pricing
   * OpenAI:    https://openai.com/api/pricing
   * Google:    https://ai.google.dev/gemini-api/docs/pricing
@@ -55,6 +55,9 @@ CLAUDE_PRICING = {
     # Claude Opus 5.
     "opus-5":     (5.0, 6.25, 0.50, 25.0),     # Opus 5 (cache prices derived from ratios)
 
+    # Claude Sonnet 5.
+    "sonnet-5":   (3.0, 3.75, 0.30, 15.0),     # Sonnet 5 (cache prices derived from ratios)
+
     # Claude 4 point releases (explicit so future divergence is easy to see).
     "opus-4-8":   (5.0, 6.25, 0.50, 25.0),     # Opus 4.8 (cache prices derived from ratios)
     "opus-4-7":   (5.0, 6.25, 0.50, 25.0),     # Opus 4.7 (cache prices derived from ratios)
@@ -107,6 +110,8 @@ _LEGACY_CODEX_PRICING = {
 # catalog + legacy: catalog owns selectable prices, legacy covers the rest. No
 # key overlaps today; if one ever collided, the legacy (right) side would win.
 CODEX_PRICING = {**catalog_pricing("codex"), **_LEGACY_CODEX_PRICING}
+GROK_PRICING = catalog_pricing("grok")
+KIMI_PRICING = catalog_pricing("kimi")
 
 GEMINI_PRICING = {
     # model substring -> (input USD/1M, output USD/1M).
@@ -121,6 +126,8 @@ GEMINI_PRICING = {
     # forms appear in the wild, so include both for the new 3.x families.
 
     # Gemini 3 series.
+    "3.6-flash":       (1.50, 7.50),
+    "3-6-flash":       (1.50, 7.50),
     "3.5-flash":       (1.50, 9.00),
     "3-5-flash":       (1.50, 9.00),
     "3.1-pro":         (2.00, 12.0),
@@ -213,6 +220,31 @@ def codex_cost_usd(model: str | None, input_t: int, cached_input_t: int, output_
     return round((fresh * pi + cached_input_t * pci + output_t * po) / 1_000_000, 4)
 
 
+def provider_cost_usd(
+    provider: str,
+    model: str | None,
+    input_t: int,
+    cached_input_t: int,
+    output_t: int,
+) -> float | None:
+    """Price a three-rate provider without fabricating an unknown model rate."""
+    tables = {
+        "codex": CODEX_PRICING,
+        "grok": GROK_PRICING,
+        "kimi": KIMI_PRICING,
+    }
+    pricing = _match_pricing(model, tables.get(provider, {}))
+    if pricing is None:
+        return None
+    input_price, cached_price, output_price = pricing
+    fresh_input = max(0, input_t - cached_input_t)
+    return round((
+        fresh_input * input_price
+        + cached_input_t * cached_price
+        + output_t * output_price
+    ) / 1_000_000, 4)
+
+
 def antigravity_cost_usd(model: str | None, input_t: int, output_t: int) -> float | None:
     p = _match_pricing(model, GEMINI_PRICING)
     if p is None: