Keep the larger today usage value

companion/codex_display/main.py

@@ -45,10 +45,13 @@ def _active_count(self) -> int:
     def _current_snapshot(self) -> Snapshot:
         assert self._snapshot is not None
         snapshot = self._snapshot
-        if snapshot.tokens_today_estimated:
-            local_tokens = read_local_tokens()
-            if local_tokens is not None:
-                snapshot = replace(snapshot, tokens_today=local_tokens)
+        local_tokens = read_local_tokens()
+        if local_tokens is not None and local_tokens > snapshot.tokens_today:
+            snapshot = replace(
+                snapshot,
+                tokens_today=local_tokens,
+                tokens_today_estimated=True,
+            )
         return replace(snapshot, active_threads=self._active_count())
 
     async def encoded(self) -> bytes:

companion/tests/test_app_server.py

@@ -1,6 +1,8 @@
 import asyncio
 import json
 import unittest
+from datetime import datetime, timezone
+from unittest import mock
 
 from companion.codex_display.app_server import AppServerClient, AppServerError
 from companion.codex_display.main import SnapshotCache, reconcile_active_loop
@@ -10,6 +12,7 @@
 class FakeClient:
     def __init__(self):
         self.threads = [{"status": {"type": "active"}}]
+        self.usage = {"dailyUsageBuckets": []}
 
     async def request(self, method, params=None):
         if method == "account/rateLimits/read":
@@ -22,7 +25,7 @@ async def request(self, method, params=None):
                 }
             }
         if method == "account/usage/read":
-            return {"dailyUsageBuckets": []}
+            return self.usage
         if method == "thread/list":
             self.thread_params = params
             return {"data": self.threads}
@@ -40,6 +43,42 @@ def test_collects_all_three_sources(self):
 
 
 class LifecycleTests(unittest.IsolatedAsyncioTestCase):
+    async def test_today_usage_uses_larger_local_value(self):
+        client = FakeClient()
+        cache = SnapshotCache(client, refresh_seconds=60)
+        await cache.refresh()
+
+        with mock.patch(
+            "companion.codex_display.main.read_local_tokens",
+            return_value=250,
+        ):
+            payload = json.loads(await cache.encoded())
+
+        self.assertEqual(payload["d"], 250)
+        self.assertEqual(payload["e"], 1)
+
+    async def test_today_usage_keeps_larger_official_value(self):
+        client = FakeClient()
+        client.usage = {
+            "dailyUsageBuckets": [
+                {
+                    "startDate": datetime.now(timezone.utc).date().isoformat(),
+                    "tokens": 250,
+                }
+            ]
+        }
+        cache = SnapshotCache(client, refresh_seconds=60)
+        await cache.refresh()
+
+        with mock.patch(
+            "companion.codex_display.main.read_local_tokens",
+            return_value=100,
+        ):
+            payload = json.loads(await cache.encoded())
+
+        self.assertEqual(payload["d"], 250)
+        self.assertEqual(payload["e"], 0)
+
     async def test_hook_change_reconciles_cached_active_count(self):
         client = FakeClient()
         cache = SnapshotCache(client, refresh_seconds=60)