From 07f5a85b23a8b575f78e9fe094ed449f888f6e6d Mon Sep 17 00:00:00 2001
From: Pravin Kamble <iampbkamble@gmail.com>
Date: Tue, 10 Nov 2026 20:03:10 +0530
Subject: [PATCH] DRU-1147 -- Count an opencode or pi call's cached input and
 reasoning tokens in its token totals

---
 backend/druks/harnesses/artifacts.py | 10 +++++++
 backend/tests/test_cost_capture.py   | 43 ++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/backend/druks/harnesses/artifacts.py b/backend/druks/harnesses/artifacts.py
index fb8117d9..069aa7ef 100644
--- a/backend/druks/harnesses/artifacts.py
+++ b/backend/druks/harnesses/artifacts.py
@@ -67,6 +67,16 @@ def normalize_token_usage(metadata: dict[str, Any] | None) -> dict[str, int] | N
         # Codex's output_tokens already include its reasoning tokens.
         output = _int(metadata.get("output_tokens"))
         reasoning = _int(metadata.get("reasoning_output_tokens"))
+    elif provider in ("opencode", "pi"):
+        # Both CLIs report input without its cache tokens. opencode also
+        # reports output without its reasoning tokens; pi includes them.
+        cached_input = _int(metadata.get("cached_input_tokens"))
+        cache_creation = _int(metadata.get("cache_creation_tokens"))
+        input_total = _int(metadata.get("input_tokens")) + cached_input + cache_creation
+        reasoning = _int(metadata.get("reasoning_tokens"))
+        output = _int(metadata.get("output_tokens"))
+        if provider == "opencode":
+            output += reasoning
     else:
         # Unknown provider — try the fields generously so a future
         # adapter that writes the canonical names just works.
diff --git a/backend/tests/test_cost_capture.py b/backend/tests/test_cost_capture.py
index d867a10e..4ae80526 100644
--- a/backend/tests/test_cost_capture.py
+++ b/backend/tests/test_cost_capture.py
@@ -373,6 +373,49 @@ def test_normalize_token_usage_openai_keeps_input_and_output_inclusive():
     }
 
 
+def test_normalize_token_usage_opencode_adds_cache_to_input_and_reasoning_to_output():
+    metadata = {
+        "provider": "opencode",
+        "input_tokens": 8_000,
+        "cached_input_tokens": 90_000,
+        "cache_creation_tokens": 2_000,
+        "output_tokens": 900,
+        "reasoning_tokens": 600,
+    }
+
+    canonical = normalize_token_usage(metadata)
+
+    assert canonical == {
+        "input_tokens": 100_000,
+        "output_tokens": 1_500,
+        "cached_input_tokens": 90_000,
+        "cache_creation_tokens": 2_000,
+        "reasoning_tokens": 600,
+        "total_tokens": 101_500,
+    }
+
+
+def test_normalize_token_usage_pi_adds_cache_to_input():
+    metadata = {
+        "provider": "pi",
+        "input_tokens": 8_000,
+        "cached_input_tokens": 90_000,
+        "cache_creation_tokens": 2_000,
+        "output_tokens": 1_500,
+    }
+
+    canonical = normalize_token_usage(metadata)
+
+    assert canonical == {
+        "input_tokens": 100_000,
+        "output_tokens": 1_500,
+        "cached_input_tokens": 90_000,
+        "cache_creation_tokens": 2_000,
+        "reasoning_tokens": 0,
+        "total_tokens": 101_500,
+    }
+
+
 def test_normalize_token_usage_returns_none_for_empty_metadata():
     assert normalize_token_usage({"provider": "anthropic"}) is None
     assert normalize_token_usage({}) is None