okstra 0.173.0 → 0.175.1

package/runtime/python/okstra_ctl/usage_identity.py

@@ -0,0 +1,86 @@
+"""Project usage rows from stored execution-identity refs."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from collections.abc import Mapping
+from typing import Any
+
+from .execution_identity import (
+    ExecutionManifest,
+    execution_identity_for_ref,
+    stored_execution_label,
+    stored_identity,
+)
+from .execution_manifest import read_execution_manifest
+
+
+def run_execution_manifest(run_root: Path) -> ExecutionManifest:
+    """Load the first run-manifest under ``run_root/manifests``."""
+    return read_execution_manifest(_manifest_path(run_root))
+
+
+def identity_rows_from_state(state: Mapping[str, Any] | None) -> list[dict[str, Any]]:
+    """Collect stored identity rows from a team-state document."""
+    if not isinstance(state, Mapping):
+        return []
+    rows: list[dict[str, Any]] = []
+    seen: set[str] = set()
+    for collection in (
+        state.get("workers") or [],
+        state.get("workerDispatches") or [],
+        state.get("agentDispatches") or [],
+    ):
+        if not isinstance(collection, list):
+            continue
+        for item in collection:
+            ident = stored_identity(item if isinstance(item, Mapping) else None)
+            if not ident:
+                continue
+            key = str(ident.get("roleExecutionRef") or ident.get("executionLabel"))
+            if key in seen:
+                continue
+            seen.add(key)
+            rows.append(ident)
+    return rows
+
+
+def collect_usage_rows(run_root: Path) -> list[dict[str, Any]]:
+    """Project usage rows with stored execution identity refs."""
+    manifest = run_execution_manifest(run_root)
+    team_state = _load_json(run_root / "state" / "team-state.json")
+    rows: list[dict[str, Any]] = []
+    for collection in (
+        team_state.get("workerDispatches") or [],
+        team_state.get("agentDispatches") or [],
+    ):
+        for item in collection:
+            if not isinstance(item, Mapping):
+                continue
+            rows.append(_usage_row(manifest, item))
+    return rows
+
+
+def _usage_row(manifest: ExecutionManifest, item: Mapping[str, Any]) -> dict[str, Any]:
+    ref = str(item.get("roleExecutionRef") or "")
+    role = execution_identity_for_ref(manifest, ref)
+    return {
+        "participantRef": item.get("participantRef") or role.participant_ref,
+        "roleExecutionRef": role.role_execution_ref,
+        "invocationRef": item.get("invocationRef"),
+        "attempt": item.get("attempt", 1),
+        "executionLabel": stored_execution_label(
+            {"executionLabel": item.get("executionLabel") or role.execution_label}
+        ),
+    }
+
+
+def _manifest_path(run_root: Path) -> Path:
+    manifests = run_root / "manifests"
+    for path in sorted(manifests.glob("*.json")):
+        return path
+    raise FileNotFoundError(f"no run manifest under {run_root}")
+
+
+def _load_json(path: Path) -> dict[str, Any]:
+    return json.loads(path.read_text(encoding="utf-8"))

package/runtime/python/okstra_ctl/usage_report.py

@@ -147,9 +147,13 @@ def _collect_usage(
     end_at: dt.datetime,
     unavailable: dict[str, int],
     warnings: dict[str, int],
-) -> tuple[dict[str, dict], set[str], int, int]:
+) -> tuple[dict[str, dict], set[str], int, int, list[dict]]:
+    from okstra_ctl.usage_identity import identity_rows_from_state
+
     buckets: dict[str, dict] = {}
     unmatched_models: set[str] = set()
+    identity_rows: list[dict] = []
+    seen_identity: set[str] = set()
     total_runs = 0
     collected_runs = 0
     for entry in list_project_tasks(project_root):
@@ -173,6 +177,13 @@ def _collect_usage(
             total_runs += 1
             bucket["runs"] += 1
             metrics, reason = _run_metrics(run, project_root)
+            state = _read_state(project_root, str(run.get("teamStatePath") or ""))
+            for row in identity_rows_from_state(state):
+                key = str(row.get("roleExecutionRef") or row.get("executionLabel"))
+                if key in seen_identity:
+                    continue
+                seen_identity.add(key)
+                identity_rows.append(row)
             if reason is not None:
                 unavailable[reason] += 1
                 continue
@@ -182,7 +193,7 @@ def _collect_usage(
                 bucket[key] += metrics[key]
             bucket["costUsd"] += metrics["costUsd"]
             unmatched_models.update(metrics["unmatchedModels"])
-    return buckets, unmatched_models, total_runs, collected_runs
+    return buckets, unmatched_models, total_runs, collected_runs, identity_rows
 
 
 def build_usage_report(project_root: Path, days: int, now: dt.datetime) -> dict:
@@ -196,7 +207,7 @@ def build_usage_report(project_root: Path, days: int, now: dt.datetime) -> dict:
         start_at = dt.datetime.min.replace(tzinfo=UTC)
     unavailable = {key: 0 for key in UNAVAILABLE_KEYS}
     warnings = {"missingOrInvalidTimelines": 0}
-    buckets, unmatched_models, total_runs, collected_runs = _collect_usage(
+    buckets, unmatched_models, total_runs, collected_runs, identity_rows = _collect_usage(
         project_root, start_at, end_at, unavailable, warnings
     )
 
@@ -226,6 +237,7 @@ def build_usage_report(project_root: Path, days: int, now: dt.datetime) -> dict:
         "unavailable": unavailable,
         "warnings": warnings,
         "unmatchedModels": sorted(unmatched_models),
+        "identityRows": identity_rows,
     }