Something went wrong. Try again.
small gleam coding and (not yet) persistent agent daemon with a detachable cli forked from okami.mom/albedo
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450"""Usage summary: what the provider calls used, totalled by day, profile and model.
Catches: a call whose counts never reach the hourly totals, a failed attemptcounted as tokens or not counted at all, a total that shrinks when a session isdeleted (the request rows go with it), a store from before the totals existedthat starts with an empty summary instead of its history, a trigger that is notinstalled again after the backfill, an hour stamped after today counted inthe lines but not the days, and days that begin at UTC midnight insteadof the client's. None of this is visible from one session's request records.
The selector half catches what a narrowed read gets wrong: scopes that followthe filter instead of the window, so a client cannot offer the scope it is notshowing; an account matched without its provider, so one provider's accountselects another's; a value that a percent-decoded query turns into a differentvalue (a `+` arriving as a space); a list cut short without saying so; and daytotals taken from the hundred lines shown rather than from every hour matched."""
import contextlibimport jsonimport sqlite3import timeimport unittestimport urllib.errorimport urllib.parsefrom datetime import datetime
from harness import Albedo, Provider, Reply, exclusive, text, wait_until
FIRST = { "input_tokens": 100, "output_tokens": 20, "input_tokens_details": {"cached_tokens": 40}, "output_tokens_details": {"reasoning_tokens": 7},}SECOND = { "input_tokens": 300, "output_tokens": 40, "input_tokens_details": {"cached_tokens": 120}, "output_tokens_details": {"reasoning_tokens": 3},}DAY_MS = 86_400_000
ANTHROPIC = "anthropic:https://api.anthropic.com"OPENAI = "responses:https://api.openai.com/v1"GEMINI = "gemini:https://generativelanguage.googleapis.com"
def encode(parameters): """A query the daemon reads back unchanged: `+` is a space to its query decoder, so a literal plus travels as `%2B`, never as a form-encoded `+`.""" return "?" + urllib.parse.urlencode(parameters, quote_via=urllib.parse.quote)
def reply(request): user = next( ( str(item.get("content", "")) for item in reversed(request["input"]) if item.get("role") == "user" ), "", ) # A non-retryable error isolates one failed attempt. Upstream now treats # 429 as account rotation/backoff, not an immediate terminal response. if user == "fail this request": return Reply("error", "invalid request", status=400) return text("done", usage=SECOND if user == "second" else FIRST)
def instant(timestamp): return datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
def millis(timestamp): return int(instant(timestamp).timestamp() * 1000)
# exclusive: it reads totals across the whole daemon and restarts it on an older store@exclusiveclass UsageSummaryTests(unittest.TestCase): def setUp(self): self.provider = Provider(reply) self.addCleanup(self.provider.close) self.app = Albedo(self.provider, protocol="responses") self.app.__enter__() self.addCleanup(self.app.__exit__, None, None, None)
def usage(self, query=""): with self.app.api(f"/usage{query}") as response: return json.load(response)
def run_prompts(self, *prompts): sessions = [] for prompt in prompts: session = self.app.session(name="usage fixture") self.app.prompt(session, prompt).close() self.app.idle(session) sessions.append(session) return sessions
def delete_session(self, session): resource = f"/sessions/{session}?view=configuration" with self.app.api(resource) as response: validator = response.getheader("ETag") response.read() with self.app.api( resource, method="DELETE", headers={"If-Match": validator} ) as response: self.assertEqual(json.load(response)["state"], "complete")
def seed(self, *rows, hour=None): """Writes one hour of call totals, each row being `(provider, account, model, calls, input_tokens, output_tokens)`.
The route reads these totals, never the request rows, and only a provider with an account pool labels a call with an account, which a fixture would need real credentials for. `account` is `""` for a provider without one, which is how the trigger stores it. The hour is the current one unless `hour` says otherwise, so every window covers it. """ if hour is None: hour = int(time.time() * 1000) // 3_600_000 * 3_600_000 with contextlib.closing(sqlite3.connect(self.app.home / "albedo.sqlite")) as db: db.execute("PRAGMA busy_timeout=5000") db.executemany( "INSERT OR REPLACE INTO usage_hours(hour_ms,profile,provider,account," "model,kind,calls,errors,input_tokens,cached_input_tokens," "cache_creation_tokens,output_tokens,reasoning_tokens)" " VALUES(?,'fixture',?,?,?,'turn',?,0,?,0,0,?,0)", [ (hour, provider, account, model, calls, input_tokens, output_tokens) for provider, account, model, calls, input_tokens, output_tokens in rows ], ) db.commit()
def totals(self, summary): (line,) = summary["lines"] return { key: line[key] for key in ( "calls", "errors", "input_tokens", "cached_input_tokens", "cache_creation_tokens", "output_tokens", "reasoning_tokens", ) }
def test_calls_are_totalled_by_day_and_by_model_and_a_failed_attempt_is_an_error( self, ): self.run_prompts("first", "second", "fail this request") summary = self.usage("?days=3") expected = { "calls": 3, "errors": 1, "input_tokens": 400, "cached_input_tokens": 160, "cache_creation_tokens": 0, "output_tokens": 60, "reasoning_tokens": 10, } self.assertEqual(self.totals(summary), expected) (line,) = summary["lines"] self.assertEqual(line["model"], "fixture-model") # Upstream migrates the fixture key into a labelled account. self.assertEqual(line["account"], self.app.profile) self.assertTrue(line["profile"].startswith("fixture")) # Three days, oldest first, the earlier two empty, and the day of use # holds the same totals as the one line. days = summary["days"] self.assertEqual(len(days), 3) for day in days[:2]: self.assertEqual(day["calls"], 0) self.assertEqual({key: days[2][key] for key in expected}, expected) self.assertEqual( [millis(day["start"]) - millis(days[0]["start"]) for day in days], [0, DAY_MS, 2 * DAY_MS], )
def test_title_generation_counts_alongside_the_main_turn(self): self.provider.title_script = lambda request: text( "Generated title", usage=SECOND ) session = self.app.session() self.app.prompt(session, "first").close() self.app.idle(session) wait_until( lambda: self.totals(self.usage())["calls"] == 2, 10, "the title request was not counted", ) self.assertEqual( self.totals(self.usage()), { "calls": 2, "errors": 0, "input_tokens": 400, "cached_input_tokens": 160, "cache_creation_tokens": 0, "output_tokens": 60, "reasoning_tokens": 10, }, )
def test_a_deleted_session_keeps_what_it_used(self): first, second = self.run_prompts("first", "second") before = self.totals(self.usage()) self.assertEqual(before["calls"], 2) self.delete_session(first) # The premise: the session's own request rows went with it. with sqlite3.connect(self.app.home / "albedo.sqlite") as db: self.assertEqual( db.execute( "SELECT COUNT(*) FROM provider_requests WHERE session=?", (first,) ).fetchone()[0], 0, ) self.assertEqual(self.totals(self.usage()), before) self.delete_session(second) self.assertEqual(self.totals(self.usage()), before)
def test_a_store_from_before_the_totals_is_backfilled_and_counts_on_afterwards( self, ): self.run_prompts("first", "second") before = self.totals(self.usage()) database = self.app.home / "albedo.sqlite"
def older_store(_app=None): with sqlite3.connect(database) as db: db.execute("DROP TRIGGER provider_requests_usage") db.execute("DROP TABLE usage_hours")
self.app.restart(prepare=older_store) self.assertEqual(self.totals(self.usage()), before) # The trigger is back: a call after the upgrade is counted once. self.run_prompts("first") after = self.totals(self.usage()) self.assertEqual(after["calls"], before["calls"] + 1) self.assertEqual(after["input_tokens"], before["input_tokens"] + 100)
def test_an_hour_stamped_after_today_is_in_no_list(self): # A clock set back leaves hours stamped later than now. The days end at # today, so the lines and scopes must end there too, or a model would # be counted in a list the days cannot show. now = int(time.time() * 1000) // 3_600_000 * 3_600_000 self.seed((ANTHROPIC, "", "opus", 1, 100, 10)) self.seed((OPENAI, "", "gpt", 1, 400, 40), hour=now + 2 * DAY_MS) summary = self.usage(encode({"days": 2})) self.assertEqual([line["model"] for line in summary["lines"]], ["opus"]) self.assertEqual(summary["scopes"], [{"provider": ANTHROPIC, "account": None}]) self.assertEqual(sum(day["calls"] for day in summary["days"]), 1)
def test_days_begin_at_the_clients_midnight(self): utc = self.usage("?days=2") ahead = self.usage("?days=2&utc_offset_minutes=300") behind = self.usage("?days=2&utc_offset_minutes=-480") for summary, offset in ((utc, 0), (ahead, 300), (behind, -480)): self.assertEqual(summary["utc_offset_minutes"], offset) now = millis(summary["generated_at"]) today = (now + offset * 60_000) // DAY_MS * DAY_MS - offset * 60_000 self.assertEqual( [millis(day["start"]) for day in summary["days"]], [today - DAY_MS, today], )
def test_the_window_is_bounded_and_checked(self): self.assertEqual(len(self.usage()["days"]), 7) self.assertEqual(len(self.usage("?days=90")["days"]), 90) for query in ( "?days=0", "?days=91", "?days=x", "?utc_offset_minutes=900", "?utc_offset_minutes=-721", "?hours=1", ): with self.subTest(query=query): with self.assertRaises(urllib.error.HTTPError) as caught: self.usage(query) self.assertEqual(caught.exception.code, 400) caught.exception.close()
def test_scopes_list_the_whole_window_while_the_filter_narrows_the_days_and_lines( self, ): self.seed( (ANTHROPIC, "ada@example.com", "opus", 1, 100, 10), (ANTHROPIC, "grace@example.com", "opus", 1, 200, 20), (OPENAI, "ada@example.com", "gpt", 1, 400, 40), (GEMINI, "", "gemini-pro", 1, 800, 80), ) whole = self.usage(encode({"days": 1})) self.assertEqual( whole["scopes"], [ {"provider": ANTHROPIC, "account": "ada@example.com"}, {"provider": ANTHROPIC, "account": "grace@example.com"}, {"provider": GEMINI, "account": None}, {"provider": OPENAI, "account": "ada@example.com"}, ], ) self.assertFalse(whole["scopes_truncated"]) self.assertFalse(whole["lines_truncated"]) self.assertEqual(whole["days"][0]["calls"], 4) self.assertEqual(whole["days"][0]["input_tokens"], 1500) self.assertEqual( [ (line["provider"], line["account"], line["model"]) for line in whole["lines"] ], [ (GEMINI, None, "gemini-pro"), (OPENAI, "ada@example.com", "gpt"), (ANTHROPIC, "grace@example.com", "opus"), (ANTHROPIC, "ada@example.com", "opus"), ], )
# One provider: the days and the lines count only its rows, and every # scope stays on offer. anthropic = self.usage(encode({"days": 1, "provider": ANTHROPIC})) self.assertEqual(anthropic["scopes"], whole["scopes"]) self.assertEqual(anthropic["days"][0]["calls"], 2) self.assertEqual(anthropic["days"][0]["input_tokens"], 300) self.assertEqual( [line["account"] for line in anthropic["lines"]], ["grace@example.com", "ada@example.com"], )
# One account of one provider: the same label under another provider is # another scope, and only the named provider's rows are counted. ada = self.usage( encode({"days": 1, "provider": OPENAI, "account": "ada@example.com"}) ) self.assertEqual([line["model"] for line in ada["lines"]], ["gpt"]) self.assertEqual(ada["days"][0]["input_tokens"], 400) self.assertEqual(ada["scopes"], whole["scopes"])
# A provider with no use in the window: zeroed days and no lines, and # still every scope, because the scopes are read before the filter. unused = self.usage( encode({"days": 1, "provider": "nobody:https://example.invalid"}) ) self.assertEqual(unused["lines"], []) self.assertEqual(unused["days"][0]["calls"], 0) self.assertEqual(unused["days"][0]["input_tokens"], 0) self.assertEqual(unused["scopes"], whole["scopes"])
def test_scope_values_are_matched_exactly_even_with_an_ampersand_a_quote_or_a_plus( self, ): provider = "custom:https://api.example.com/v1?model='a'&flag=x+y" account = "ada+albedo@example.com" self.seed( (provider, account, "odd", 1, 7, 1), (provider, "ada albedo@example.com", "spaced", 1, 11, 1), (provider, "", "none", 1, 13, 1), ) found = self.usage( encode({"days": 1, "provider": provider, "account": account}) ) self.assertEqual( [(line["model"], line["account"]) for line in found["lines"]], [("odd", account)], ) self.assertEqual(found["days"][0]["input_tokens"], 7)
# A raw `+` in a query arrives as a space, so a client that escaped its # values reaches the plus account and the space account separately. quoted = urllib.parse.quote(provider, safe="") raw = f"?days=1&provider={quoted}&account=ada+albedo%40example.com" self.assertEqual( [line["model"] for line in self.usage(raw)["lines"]], ["spaced"] )
def test_an_empty_account_selects_the_rows_without_one_and_needs_its_provider(self): self.seed( (GEMINI, "", "gemini-pro", 1, 800, 80), (GEMINI, "ada@example.com", "gemini-pro", 1, 100, 10), ) # The provider alone is every account of it, which is not the same as # its rows without an account at all. pooled = self.usage(encode({"days": 1, "provider": GEMINI})) self.assertEqual(len(pooled["lines"]), 2) self.assertEqual(pooled["days"][0]["input_tokens"], 900) without = self.usage(encode({"days": 1, "provider": GEMINI, "account": ""})) self.assertEqual(len(without["lines"]), 1) self.assertIsNone(without["lines"][0]["account"]) self.assertEqual(without["days"][0]["input_tokens"], 800)
for query in ( "?account=ada%40example.com", "?account=", "?provider=", "?provider=&account=", ): with self.subTest(query=query): with self.assertRaises(urllib.error.HTTPError) as caught: self.usage(query) self.assertEqual(caught.exception.code, 400) caught.exception.close()
def test_the_lines_are_cut_at_a_hundred_while_the_days_total_every_row(self): provider = "cap:https://example.com/v1" self.seed( *[(provider, "", f"model-{index:03}", 1, 5, 1) for index in range(100)] ) exact = self.usage(encode({"days": 7})) self.assertEqual(len(exact["lines"]), 100) self.assertFalse(exact["lines_truncated"]) self.assertEqual(exact["days"][-1]["calls"], 100)
self.seed( *[(provider, "", f"model-{index:03}", 1, 5, 1) for index in range(100, 120)] ) cut = self.usage(encode({"days": 7})) self.assertEqual(len(cut["lines"]), 100) self.assertTrue(cut["lines_truncated"]) # The totals come from the hours, not from the hundred lines shown. self.assertEqual(cut["days"][-1]["calls"], 120) self.assertEqual(cut["days"][-1]["input_tokens"], 120 * 5) self.assertEqual(sum(line["input_tokens"] for line in cut["lines"]), 100 * 5)
def test_the_scopes_are_cut_at_two_hundred_and_say_so(self): self.seed(*[(f"scope:{index:03}", "", "one", 1, 1, 1) for index in range(200)]) exact = self.usage(encode({"days": 7})) self.assertEqual(len(exact["scopes"]), 200) self.assertFalse(exact["scopes_truncated"])
self.seed( *[(f"scope:{index:03}", "", "one", 1, 1, 1) for index in range(200, 210)] ) cut = self.usage(encode({"days": 7})) self.assertEqual(len(cut["scopes"]), 200) self.assertTrue(cut["scopes_truncated"]) # Sorted by provider, so the first two hundred are still the ones shown, # and the day still totals every row. self.assertEqual(cut["scopes"], exact["scopes"]) self.assertEqual(cut["days"][-1]["calls"], 210)
if __name__ == "__main__": unittest.main()