diff --git a/companion/codex_display/main.py b/companion/codex_display/main.py index 8622b7c..fe3b284 100644 --- a/companion/codex_display/main.py +++ b/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: diff --git a/companion/tests/test_app_server.py b/companion/tests/test_app_server.py index cebf830..5114f42 100644 --- a/companion/tests/test_app_server.py +++ b/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)