diff --git a/default.nix b/default.nix index baf5908b..5306b49b 100644 --- a/default.nix +++ b/default.nix @@ -1,4 +1,5 @@ -{ lib, stdenv, rustPlatform, buildGoModule, fetchurl, gleam, beamPackages, python3, bash, coreutils, makeWrapper }: +{ lib, stdenv, rustPlatform, buildGoModule, fetchurl, gleam, beamPackages, python3, bash, coreutils +, makeWrapper, zig, usage-core }: let erlang = beamPackages.erlang; # esqlite loads the pc plugin during its Rebar3 build, which cannot fetch @@ -44,6 +45,42 @@ let platforms = lib.platforms.unix; }; }; + # The provide-usage CLI: a pure Zig core the daemon drives one round at a + # time (see robot-docs/usage-feed.md). Pinned by the flake input, built with + # the toolchain its build.zig.zon demands. + usageCore = stdenv.mkDerivation { + pname = "usage-core"; + version = "0.1.0"; + src = usage-core; + # The knot serves the archive with a query string, so the store path has + # no extension for unpackPhase to sniff: one explicit tar it is. + unpackPhase = '' + runHook preUnpack + tar -xzf $src + runHook postUnpack + ''; + sourceRoot = "provide-usage-main"; + nativeBuildInputs = [ zig ]; + # The build.zig ranlib step hardcodes zig-out, so the default prefix it is + # (no --prefix): artifacts are copied out in installPhase instead. + buildPhase = '' + runHook preBuild + export HOME="$TMPDIR/home" + zig build --cache-dir "$TMPDIR/zig-cache" --global-cache-dir "$TMPDIR/zig-global-cache" + runHook postBuild + ''; + installPhase = '' + runHook preInstall + mkdir -p $out/bin + cp zig-out/bin/usage $out/bin/usage + runHook postInstall + ''; + meta = { + description = "stateless provider usage and quota core with a CLI"; + mainProgram = "usage"; + platforms = lib.platforms.unix; + }; + }; daemon = stdenv.mkDerivation { pname = "albedo-daemon"; version = "1.0.0"; @@ -69,6 +106,7 @@ let cp -R build/erlang-shipment/. $out/lib/albedo/ mkdir -p $out/lib/albedo/albedo/priv/bin ln -s ${render}/bin/albedo-render $out/lib/albedo/albedo/priv/bin/albedo-render + ln -s ${usageCore}/bin/usage $out/lib/albedo/albedo/priv/bin/usage makeWrapper $out/lib/albedo/entrypoint.sh $out/bin/albedo-daemon \ --add-flags run \ --prefix PATH : ${lib.makeBinPath [ erlang python3 bash coreutils render ]} @@ -88,7 +126,7 @@ in buildGoModule { wrapProgram $out/bin/albedo \ --set-default ALBEDO_DAEMON ${daemon}/bin/albedo-daemon ''; - passthru = { inherit daemon render; }; + passthru = { inherit daemon render usageCore; }; meta = { description = "coding agent daemon with a Charm terminal client"; license = lib.licenses.wtfpl; diff --git a/flake.lock b/flake.lock index f1a94352..8a6115c0 100644 --- a/flake.lock +++ b/flake.lock @@ -18,7 +18,20 @@ }, "root": { "inputs": { - "nixpkgs": "nixpkgs" + "nixpkgs": "nixpkgs", + "usage-core": "usage-core" + } + }, + "usage-core": { + "flake": false, + "locked": { + "narHash": "sha256-UVsepzVfYIDXUsoAH4KRdAWkeFPylUurWWXFw2Ereqc=", + "type": "file", + "url": "https://api.next.tangled.org/xrpc/org.tangled.temp.git.getArchive?format=tar.gz&prefix=provide-usage-main&ref=1a4bff9&repo=did:plc:2lf7buutfnfcucljnmaypf7u" + }, + "original": { + "type": "file", + "url": "https://api.next.tangled.org/xrpc/org.tangled.temp.git.getArchive?format=tar.gz&prefix=provide-usage-main&ref=1a4bff9&repo=did:plc:2lf7buutfnfcucljnmaypf7u" } } }, diff --git a/flake.nix b/flake.nix index e3ea0600..2f2d2982 100644 --- a/flake.nix +++ b/flake.nix @@ -2,8 +2,15 @@ description = "albedo coding agent and terminal client"; inputs.nixpkgs.url = "github:NixOS/nixpkgs/nixpkgs-unstable"; + # usage-core, pinned: a pure Zig provider-usage core whose CLI the daemon + # drives one round at a time (see robot-docs/usage-feed.md). A tarball off + # the knot, not a git input, so no fetchGit machinery enters the build. + inputs.usage-core = { + url = "https://api.next.tangled.org/xrpc/org.tangled.temp.git.getArchive?repo=did%3Aplc%3A2lf7buutfnfcucljnmaypf7u&ref=1a4bff9&format=tar.gz&prefix=provide-usage-main"; + flake = false; + }; - outputs = { self, nixpkgs }: + outputs = { self, nixpkgs, usage-core }: let systems = [ "aarch64-darwin" "x86_64-darwin" "aarch64-linux" "x86_64-linux" ]; eachSystem = nixpkgs.lib.genAttrs systems; @@ -11,8 +18,8 @@ packages = eachSystem (system: let pkgs = import nixpkgs { inherit system; }; in rec { - albedo = pkgs.callPackage ./default.nix { }; - inherit (albedo) daemon render; + albedo = pkgs.callPackage ./default.nix { inherit usage-core; }; + inherit (albedo) daemon render usageCore; default = albedo; }); devShells = eachSystem (system: diff --git a/robot-docs/quota.md b/robot-docs/quota.md new file mode 100644 index 00000000..c3081167 --- /dev/null +++ b/robot-docs/quota.md @@ -0,0 +1,55 @@ +# quota + +Provider quota is polled per account, once per daemon, and recorded as raw readings. Nothing is derived at record time: percentages, windows, reset times, statuses, and errors are stored exactly as the usage feed reported them, so later analysis can recompute anything from the rows. This is the phase-1 half of cache- and quota-aware compaction; a meter estimator reads these rows on demand. + +## accounts + +One poller, started with the daemon, maps every account albedo can see onto a [provide-usage](https://next.tangled.org/ptr.pet/provide-usage) provider id and credential: + +- Each OAuth account in a provider's rotation pool: Claude accounts under `auth.json`'s `anthropic` key poll as `anthropic`, ChatGPT accounts under `openai-codex` as `openai-codex`, Google accounts under `google-antigravity` as `google-antigravity` (with `extra.projectId`). Credentials come from the same readers the request path uses, and a stale token is refreshed through the provider extension's own refresh path under `auth.lock`, so polling refreshes tokens exactly the way requests do. Only the access token crosses into the core: the refresh token stays home, since albedo does all the refreshing and the core never reads one. A refresh that fails keeps the stored credential: the poll then records the auth failure as a reading instead of dropping the account. +- Profiles in `config.json`: an `alibaba` profile polls as `alibaba` with a `cli` credential — the feed runs the signed-in `bl` CLI. A generic `openai` profile maps by its `baseUrl` host: `hyper.charm.land` polls as `hyper`, `api.deepseek.com` as `deepseek`, both with the profile's API key. A profile whose host has no feed is skipped. + +The account label is always non-secret: an email, an account id, a short hash of a refresh token, or a profile name. + +## cadence + +Each account is polled on its own schedule; two polls of one account never overlap. The default cadence is 10 minutes. An account moves onto the 2-minute busy cadence while any limit is at or above 80 percent or resets within 15 minutes. A failed poll (a driver failure or a report carrying an error) backs off, doubling each time up to an hour. Account discovery re-runs on the ordinary cadence, off the poller, so a refresh never blocks scheduling; new accounts poll immediately, and an account that is mid-poll when it disappears finishes first. + +Settings live in `$ALBEDO_HOME/extensions.json` under `"quota"`, read like any extension's settings and re-read every round, so changes apply without a restart: + +```json +{ "quota": { "pollSeconds": 600, "busyPollSeconds": 120, "enabled": true } } +``` + +`enabled: false` stops polling and drops the account list; the readings already stored stay readable. + +## readings + +`quota_sample` in `albedo.sqlite` holds one row per limit a report carried: the account label, provider id, the report's plan (what 100 percent means changes with it), limit id and label, used percent, the window label and window seconds, reset time (epoch milliseconds), scope, status, the report's error when it failed, the observation time, and the source (`poll`). A report with an error and no limits stores one row carrying the error; a missing percentage stays null rather than becoming zero, so unknown is never mistaken for empty. + +The feed itself is `albedo/harness/usage_feed`: a stateless `usage advance -` round loop whose HTTP and command I/O runs through albedo. `Error` from the feed means the driver failed (missing binary, unparseable output); a provider failure arrives as a report with an error, which is data. + +## reading them back + +`GET /quota` serves the latest reading per account and limit, read-only: + +```json +{ "readings": [ { "id": 41, "account": "charm-hyper", "provider": "hyper", + "plan": null, "limitId": "primary", "label": "Primary", + "usedPercent": 85.0, "windowLabel": "1 hour", "windowSeconds": 3600, + "resetsAt": 1760000000000, "scope": null, "status": "ok", + "error": null, "observedAt": 1759999400000, + "source": "poll" } ] } +``` + +`GET /quota?history=` pages the raw samples newest-first by row id (`history=0` starts at the newest, `limit` caps the page at 200), answering `{ "items", "nextCursor", "hasMore" }` like the other paged routes. The route queries the store directly and writes nothing. + +## hermetic tests + +The provider URLs live inside the usage core, so `test/e2e/quota_test.py` points `ALBEDO_USAGE_CORE` at a fake `usage` executable that answers the envelope with scripted reports (anthropic ordinary, hyper busy, deepseek failing), with short poll settings; it asserts readings land for all three account shapes and that history pages. The override is authoritative: a nonexistent value errors rather than falling back, so the fake is the whole feed for that daemon. + +Each cadence is asserted as the interval between one account's consecutive readings — the median hyper interval at `busyPollSeconds`, the ordinary account at `pollSeconds`, the failing account's intervals growing past the ordinary cadence — not as poll counts over wall time, which race the machine rather than measure the schedule. The fake answers even a bad envelope or its own vanished directory with an error report instead of crashing: a crashed fake would exit 1 and read as a driver failure, so the fake failing loudly and the test failing on any reading whose error names the driver keeps the two halves distinguishable. + +## source references + +`albedo/daemon/quota.gleam` owns the poller, the table, and the read path. The poller takes its fetch function as a parameter, so tests can drive it against a fake. The per-provider account listing lives with each provider's credentials: `albedo_claude_auth:accounts/1`, `albedo_openai_auth:accounts/1`, and `albedo_antigravity:accounts/1`. Startup wiring is in `albedo/daemon/server.gleam`: the table in `prepare_storage`, the poller under the daemon supervisor, and the route among the daemon's own. diff --git a/robot-docs/usage-feed.md b/robot-docs/usage-feed.md new file mode 100644 index 00000000..272a95ec --- /dev/null +++ b/robot-docs/usage-feed.md @@ -0,0 +1,36 @@ +# usage feed + +provider usage and quota, through [provide-usage](https://next.tangled.org/ptr.pet/provide-usage): a pure Zig core (`usage-core`) that knows each provider's rules, and an albedo driver that owns every byte of I/O. the daemon never grows provider-specific quota code; a new provider is a `usage-core` bump plus the pinned flake input moving. + +## the split + +`usage-core` never opens a socket, reads a clock, or touches the filesystem. a provider is a state machine: given a credential and the responses so far, `usage advance` says either which requests to send next or the finished report. albedo sends those requests, hands the bodies back, and repeats. one feed is a handful of round trips at most. + +the driver is two files: + +- [`src/albedo/harness/usage_feed.gleam`](../src/albedo/harness/usage_feed.gleam) — the loop, the envelope, and the report types. `fetch(provider, credential, now_ms)` runs a feed to completion. a provider failure (a 4xx, an unknown provider, a flaky `bl`) is data: `Ok(Report(error: Some(..)))`. `Error` means the driver itself failed: the binary is missing, its output is unparseable, or the feed did not finish in 6 rounds. +- [`src/albedo/harness/albedo_usage_core.erl`](../src/albedo/harness/albedo_usage_core.erl) — the I/O half: one `usage advance -` process per round, `http` requests through [`albedo_http`](extensions.md), and `command` execution. + +each round is a fresh process fed one JSON envelope line on stdin: `{"provider", "credential", "responses": [round, ...], "nowMs"}`. the envelope carries the credential, so it never reaches argv where other local processes could read it. the host keeps the state (it resends every round's responses) and the core keeps the rules; neither holds anything the other could corrupt. an erlang port cannot half-close stdin, and the CLI stops reading at the first newline, so the driver writes the line and reads stdout to exit: one step line per replayed round, the last line is the answer. + +the binary resolves like `albedo-render`: `priv/bin/usage` first, then `PATH`. `ALBEDO_USAGE_CORE` overrides both — tests point it at a fake, and a value that does not exist is an error, never a silent fallback. + +## requests and responses + +a step asks for requests; each is answered with one `{"status", "body"}` object: + +- `http` — sent through albedo's one httpc front door, verified TLS and all. a transport failure is data: status `0` with the reason in the body, so a feed reports it instead of the call failing. +- `command` (alibaba's `bl`) — run with only `PATH`, `HOME`, `LANG`, `LC_ALL`, `LC_CTYPE`, `TZ` and the proxy variables in the environment, a timeout (the request's `timeoutMs`, 15 s by default), and a 64 KiB stdout cap. `127` means "this host will not run commands" and is reserved; a command that cannot be found or timed out answers `126`, and the feed routes around either the same way. + +## packaging + +`usage-core` is a pinned tarball flake input (`flake = false`) on the knot, built in `default.nix` with nixpkgs' zig (its `build.zig.zon` demands 0.16.0, which nixpkgs-unstable carries). the CLI is linked into the daemon next to `albedo-render`: + +``` +nix build .#albedo.usageCore # the derivation alone +nix build .#daemon # the whole daemon, priv/bin/usage included +``` + +## consumers + +the daemon-side poller (per-account quota) reads `Report` and `Limit` and records what arrives; nothing else should parse the core's JSON. `Limit.resets_at` is epoch milliseconds; `used_percent` is clamped 0–100; absent optionals mean the provider did not state the value, never "zero". diff --git a/robot-docs/usage-ledger.md b/robot-docs/usage-ledger.md new file mode 100644 index 00000000..3b4b0fb8 --- /dev/null +++ b/robot-docs/usage-ledger.md @@ -0,0 +1,36 @@ +# Usage ledger + +Every provider call the daemon makes leaves one row in `request_ledger`, a durable account of what the transcript cannot say later: what the request cost, which account served it, how it ended, and what its prefix looked like. This is the recording half of cache- and quota-aware compaction; the analysis (a meter estimator and λ) is derived from these rows on demand, so nothing beyond them is stored. + +Rows are written by the worker that streams the call, one per upstream attempt: a turn retried after a gateway hiccup leaves the failed attempt and the successful one as separate rows. Recording must never fail a turn — a ledger write error is logged and swallowed, and the call proceeds. + +## What a row holds + +- `session`, `seq` — the conversation, and the transcript row the call produced: the first assistant row of the response that turn committed (for a tool-call turn, the row carrying the call itself, not the tool output after it). `seq` is null when the call produced no transcript row: a failed call, or a compaction summary, which writes strategy state instead of the transcript. +- `kind` — `turn` for the model/tool loop, `summarizer` for the compaction summary call. +- `profile`, `provider`, `account`, `model` — the saved profile the request went through, `protocol:endpoint` as it was reached, the non-secret label of the account that served (an account id, email, or short key hash; null when the provider has no account pool), and the model. +- `startedMs`, `finishedMs` — when the attempt ran. +- `outcome`, `status`, `error` — `ok`, or `error` with the HTTP status and a truncated error body, so a 429 is a quota reading. +- token counts as the provider reported them: `inputTokens`, `cachedInputTokens`, `cacheCreationTokens`, `cacheWrite5mTokens`, `cacheWrite1hTokens`, `outputTokens`, `reasoningTokens`. Unreported counts are null, not zero. + +Providers report these differently and each adapter fills what it can: Anthropic's `cache_creation.ephemeral_5m_input_tokens` / `ephemeral_1h_input_tokens` become the TTL-split writes; OpenAI's `completion_tokens_details.reasoning_tokens` and the Responses API's `output_tokens_details.reasoning_tokens` become reasoning; Antigravity's `thoughtsTokenCount` is reported as reasoning while output keeps folding thoughts in, as before. + +## Prefix identity + +What makes cache analysis possible without storing requests. Each row records: + +- `headHash` — SHA-256 of the request head: instructions plus tool schemas. One head, one hash. +- `inputs` — how many projected inputs the request carried. +- `replaced` and `projectionHash` — the projection identity. The loop knows how many original inputs the projection replaced (`original length − shared suffix with the projection`); the row carries that count and a hash of the replacement part that stands in for them. +- `strategy` — the active compaction strategy name, when one prepared the request. + +Within one projection the request is append-only, so equal head hash plus equal projection identity means a warm prefix; a changed projection identity means a rewrite that invalidated the cache. A summarizer call carries its own head hash with `replaced` and `projectionHash` null: its history is exactly what it says. + +## Local API + +`GET /sessions/:id/ledger` answers rows oldest first, read-only straight from the store: + +- `?after=` continues from the last row id of the previous page. +- `?limit=` bounds the page (1–500, default 100). + +The response is `{session, rows, after}`; `after` is the id to continue from. Like every session route it needs the daemon token. The later `/quota` inspector reads the same rows; nothing is stored twice. diff --git a/src/albedo/daemon/conversation.gleam b/src/albedo/daemon/conversation.gleam index a2c35620..b6b4724c 100644 --- a/src/albedo/daemon/conversation.gleam +++ b/src/albedo/daemon/conversation.gleam @@ -1,5 +1,6 @@ import albedo/daemon/events import albedo/daemon/images +import albedo/daemon/ledger import albedo/daemon/mail import albedo/daemon/note import albedo/daemon/notice @@ -69,7 +70,8 @@ pub fn initialise(store: store.Store) -> Result(Nil, String) { use _ <- result.try(store.exec( db, "CREATE TABLE IF NOT EXISTS sessions(id TEXT PRIMARY KEY,title TEXT NOT NULL DEFAULT 'new session',cwd TEXT NOT NULL,model TEXT NOT NULL,protocol TEXT NOT NULL,stage TEXT NOT NULL DEFAULT 'idle',provider TEXT,activity_seq INTEGER,last_assistant_at INTEGER,usage_model TEXT,usage_recorded_at INTEGER,usage_prompt_tokens INTEGER,usage_completion_tokens INTEGER,usage_cached_prompt_tokens INTEGER,usage_cache_creation_tokens INTEGER,effort TEXT); CREATE TABLE IF NOT EXISTS transcript(seq INTEGER PRIMARY KEY AUTOINCREMENT,session TEXT NOT NULL REFERENCES sessions(id),payload BLOB NOT NULL,timestamp INTEGER,provider TEXT,thought_ms INTEGER); CREATE INDEX IF NOT EXISTS transcript_session ON transcript(session,seq);" - <> images.schema, + <> images.schema + <> ledger.schema, )) use _ <- result.try( store.add_columns(db, "sessions", [ @@ -83,6 +85,9 @@ pub fn initialise(store: store.Store) -> Result(Nil, String) { #("usage_completion_tokens", "INTEGER"), #("usage_cached_prompt_tokens", "INTEGER"), #("usage_cache_creation_tokens", "INTEGER"), + #("usage_cache_write_5m_tokens", "INTEGER"), + #("usage_cache_write_1h_tokens", "INTEGER"), + #("usage_reasoning_tokens", "INTEGER"), #("effort", "TEXT"), #("pinned_instructions", "TEXT"), #("pinned_context", "BLOB"), @@ -291,6 +296,7 @@ pub fn delete(store: store.Store, id: String) -> Result(Nil, String) { list.try_each( [ "transcript", + "request_ledger", "schedules", "session_extensions", "rolling_compaction_state", @@ -396,6 +402,16 @@ fn has_visible_assistant(inputs: List(types.Input)) -> Bool { }) } +/// Whether one committed input is a row the model itself produced: an +/// assistant message or a replayed provider item, including a tool call. The +/// first such row of a commit is the transcript row that commit links to. +fn is_assistant_row(input: types.Input) -> Bool { + case input { + types.Replay(_) | types.Assistant(_) -> True + _ -> False + } +} + /// The newest message a person wrote. Notes are the daemon talking and mail /// comes from other agents or outside, so neither names a session. pub fn latest_user(inputs: List(types.Input)) -> Option(String) { @@ -803,10 +819,14 @@ pub fn commit_from( provider: Option(String), ) -> Result(Int, String) { commit_with_letters(store, id, inputs, stage, provider, None, []) + |> result.map(fn(committed) { committed.0 }) } /// Appends a model response's inputs; `thought_ms`, how long the model -/// thought before it, is kept on the first that shows the thinking. +/// thought before it, is kept on the first that shows the thinking. Answers +/// the shared daemon time and the seq of the response's first assistant row, +/// so a recorded provider call can link to the transcript row it produced; +/// a commit with no assistant row links nothing. pub fn commit_response( store: store.Store, id: String, @@ -814,7 +834,7 @@ pub fn commit_response( stage: Stage, provider: Option(String), thought_ms: Option(Int), -) -> Result(Int, String) { +) -> Result(#(Int, Option(Int)), String) { commit_with_letters(store, id, inputs, stage, provider, thought_ms, []) } @@ -829,6 +849,7 @@ pub fn commit_letters( letters: List(String), ) -> Result(Int, String) { commit_with_letters(store, id, inputs, stage, provider, None, letters) + |> result.map(fn(committed) { committed.0 }) } fn commit_with_letters( @@ -839,7 +860,7 @@ fn commit_with_letters( provider: Option(String), thought_ms: Option(Int), letters: List(String), -) -> Result(Int, String) { +) -> Result(#(Int, Option(Int)), String) { let timestamp = usage.now() let read = images.reader(store) store.query(store, fn(db) { @@ -849,21 +870,37 @@ fn commit_with_letters( True -> 1 False -> 0 } - use _ <- result.try( + // The first assistant row of the commit, so a caller that recorded the + // provider call behind it can link to the transcript row it produced. + use linked <- result.try( entries(inputs, Some(timestamp), provider, thought_ms) - |> list.try_each(fn(entry) { + |> list.try_fold(None, fn(linked, entry) { use input <- result.try(images.externalize(db, entry.input, read)) - store.run( - db, - "INSERT INTO transcript(session,payload,timestamp,provider,thought_ms) VALUES(?,?,?,?,?)", - [ - sqlight.text(id), - sqlight.blob(pack(input)), - sqlight.int(timestamp), - sqlight.nullable(sqlight.text, provider), - sqlight.nullable(sqlight.int, entry.thought_ms), - ], + use _ <- result.try( + store.run( + db, + "INSERT INTO transcript(session,payload,timestamp,provider,thought_ms) VALUES(?,?,?,?,?)", + [ + sqlight.text(id), + sqlight.blob(pack(input)), + sqlight.int(timestamp), + sqlight.nullable(sqlight.text, provider), + sqlight.nullable(sqlight.int, entry.thought_ms), + ], + ), ) + case linked, is_assistant_row(input) { + None, True -> + store.one( + db, + "SELECT last_insert_rowid()", + [], + decode.field(0, decode.int, decode.success), + "committed row sequence", + ) + |> result.map(Some) + _, _ -> Ok(linked) + } }), ) // COALESCE keeps the old title when no new user message suggests one. @@ -877,8 +914,9 @@ fn commit_with_letters( sqlight.text(id), ], ) + |> result.replace(linked) }) - |> result.replace(timestamp) + |> result.map(fn(linked) { #(timestamp, linked) }) }) } @@ -889,6 +927,9 @@ fn usage_decoder() { use completion <- decode.field(3, decode.optional(decode.int)) use cached <- decode.field(4, decode.optional(decode.int)) use creation <- decode.field(5, decode.optional(decode.int)) + use write_5m <- decode.field(6, decode.optional(decode.int)) + use write_1h <- decode.field(7, decode.optional(decode.int)) + use reasoning <- decode.field(8, decode.optional(decode.int)) case model, recorded_at, prompt, completion { None, None, None, None -> decode.success(None) Some(model), Some(recorded_at), None, None -> @@ -898,7 +939,15 @@ fn usage_decoder() { Some(usage.Metadata( model, recorded_at, - Some(usage.Tokens(prompt, completion, cached, creation)), + Some(usage.Tokens( + prompt, + completion, + cached, + creation, + write_5m, + write_1h, + reasoning, + )), )), ) _, _, _, _ -> decode.failure(None, "consistent saved usage metadata") @@ -912,7 +961,7 @@ pub fn load_usage( store.query(store, fn(db) { store.one( db, - "SELECT usage_model,usage_recorded_at,usage_prompt_tokens,usage_completion_tokens,usage_cached_prompt_tokens,usage_cache_creation_tokens FROM sessions WHERE id=?", + "SELECT usage_model,usage_recorded_at,usage_prompt_tokens,usage_completion_tokens,usage_cached_prompt_tokens,usage_cache_creation_tokens,usage_cache_write_5m_tokens,usage_cache_write_1h_tokens,usage_reasoning_tokens FROM sessions WHERE id=?", [sqlight.text(id)], usage_decoder(), "session not found", @@ -923,7 +972,7 @@ pub fn load_usage( pub fn clear_usage(store: store.Store, id: String) -> Result(Nil, String) { store.write( store, - "UPDATE sessions SET usage_model=NULL,usage_recorded_at=NULL,usage_prompt_tokens=NULL,usage_completion_tokens=NULL,usage_cached_prompt_tokens=NULL,usage_cache_creation_tokens=NULL WHERE id=?", + "UPDATE sessions SET usage_model=NULL,usage_recorded_at=NULL,usage_prompt_tokens=NULL,usage_completion_tokens=NULL,usage_cached_prompt_tokens=NULL,usage_cache_creation_tokens=NULL,usage_cache_write_5m_tokens=NULL,usage_cache_write_1h_tokens=NULL,usage_reasoning_tokens=NULL WHERE id=?", [sqlight.text(id)], ) } @@ -934,18 +983,31 @@ pub fn record_usage( metadata: usage.Metadata, ) -> Result(Nil, String) { let usage.Metadata(model, recorded_at, tokens) = metadata - let #(prompt, completion, cached, creation) = case tokens { - Some(usage.Tokens(prompt, completion, cached, creation)) -> #( + let #(prompt, completion, cached, creation, write_5m, write_1h, reasoning) = case + tokens + { + Some(usage.Tokens( + prompt, + completion, + cached, + creation, + write_5m, + write_1h, + reasoning, + )) -> #( Some(prompt), Some(completion), cached, creation, + write_5m, + write_1h, + reasoning, ) - None -> #(None, None, None, None) + None -> #(None, None, None, None, None, None, None) } store.write( store, - "UPDATE sessions SET usage_model=?,usage_recorded_at=?,usage_prompt_tokens=?,usage_completion_tokens=?,usage_cached_prompt_tokens=?,usage_cache_creation_tokens=? WHERE id=?", + "UPDATE sessions SET usage_model=?,usage_recorded_at=?,usage_prompt_tokens=?,usage_completion_tokens=?,usage_cached_prompt_tokens=?,usage_cache_creation_tokens=?,usage_cache_write_5m_tokens=?,usage_cache_write_1h_tokens=?,usage_reasoning_tokens=? WHERE id=?", [ sqlight.text(model), sqlight.int(recorded_at), @@ -953,6 +1015,9 @@ pub fn record_usage( sqlight.nullable(sqlight.int, completion), sqlight.nullable(sqlight.int, cached), sqlight.nullable(sqlight.int, creation), + sqlight.nullable(sqlight.int, write_5m), + sqlight.nullable(sqlight.int, write_1h), + sqlight.nullable(sqlight.int, reasoning), sqlight.text(id), ], ) diff --git a/src/albedo/daemon/ledger.gleam b/src/albedo/daemon/ledger.gleam new file mode 100644 index 00000000..a543d9f7 --- /dev/null +++ b/src/albedo/daemon/ledger.gleam @@ -0,0 +1,470 @@ +//// One durable row per provider call. +//// +//// The transcript keeps what a turn said; it cannot say what the request +//// looked like when the provider saw it. Each row records what later cache +//// and quota analysis cannot recover: which account served, how long the +//// call took, what it cost in tokens (cached reads, cache writes split by +//// TTL, reasoning), how it ended (an HTTP status turns a 429 into a quota +//// reading), and a prefix identity — a hash of the request head plus the +//// projection that stands in for replaced history — so two rows can be told +//// apart as a warm append or a cache-busting rewrite without storing any +//// request body. Requests are never stored, only hashed. + +import albedo/daemon/store +import albedo/harness/compaction +import albedo/openai_api/types +import gleam/bit_array +import gleam/crypto +import gleam/dynamic/decode +import gleam/json.{type Json} +import gleam/list +import gleam/option.{type Option, None, Some} +import gleam/result +import gleam/string +import sqlight + +pub const schema = "CREATE TABLE IF NOT EXISTS request_ledger(id INTEGER PRIMARY KEY AUTOINCREMENT,session TEXT NOT NULL REFERENCES sessions(id),seq INTEGER,kind TEXT NOT NULL,profile TEXT NOT NULL,provider TEXT NOT NULL,account TEXT,model TEXT NOT NULL,started_ms INTEGER NOT NULL,finished_ms INTEGER NOT NULL,outcome TEXT NOT NULL,status INTEGER,error TEXT,input_tokens INTEGER,cached_input_tokens INTEGER,cache_creation_tokens INTEGER,cache_write_5m_tokens INTEGER,cache_write_1h_tokens INTEGER,output_tokens INTEGER,reasoning_tokens INTEGER,head_hash TEXT NOT NULL,inputs INTEGER NOT NULL,replaced INTEGER,projection_hash TEXT,strategy TEXT); CREATE INDEX IF NOT EXISTS request_ledger_session ON request_ledger(session,id);" + +/// What the call was for. A turn is the model loop; a summarizer is the +/// compaction summary call `loop.summarize` makes. +pub type Kind { + Turn + Summarizer +} + +fn kind_name(kind: Kind) -> String { + case kind { + Turn -> "turn" + Summarizer -> "summarizer" + } +} + +/// How a call ended: completed, or failed with the HTTP status when there was +/// one and a bounded detail that keeps an error body readable, not whole. +pub type Outcome { + Completed + Failed(status: Option(Int), detail: String) +} + +pub fn outcome(result: Result(a, types.Error)) -> Outcome { + case result { + Ok(_) -> Completed + Error(types.HttpError(status, body)) -> Failed(Some(status), bounded(body)) + Error(types.ProviderError(message)) -> Failed(None, bounded(message)) + Error(error) -> Failed(None, string.inspect(error)) + } +} + +/// Error bodies are truncated so one huge response cannot bloat the ledger. +const detail_limit = 2000 + +fn bounded(text: String) -> String { + case string.length(text) > detail_limit { + True -> string.slice(text, 0, detail_limit) <> "…[truncated]" + False -> text + } +} + +/// What makes cache analysis possible without storing requests: the request +/// head (instructions and tools) as a hash, how many projected inputs the +/// request carries, and a projection identity that changes exactly when +/// compaction changes what precedes the verbatim tail — the count of original +/// inputs the projection replaced and a hash of the replacement part. Within +/// one projection the request is append-only, so these tell a warm prefix +/// from a rewrite. +pub type Prefix { + Prefix( + head_hash: String, + inputs: Int, + replaced: Option(Int), + projection_hash: Option(String), + strategy: Option(String), + ) +} + +/// The prefix identity of a projected request: `projected` is the history the +/// request carries, `original` the durable history it stands in for, and +/// `strategy` the compaction strategy that prepared it, when one did. +pub fn prefix( + instructions: String, + tools: List(types.Tool), + projected: List(types.Input), + original: List(types.Input), + strategy: Option(String), +) -> Prefix { + let suffix = compaction.common_suffix(original, projected) + let replaced = list.length(original) - suffix + Prefix( + head_hash: head_hash(instructions, tools), + inputs: list.length(projected), + replaced: Some(replaced), + projection_hash: case replaced { + 0 -> None + _ -> + Some(inputs_hash(list.take(projected, list.length(projected) - suffix))) + }, + strategy: strategy, + ) +} + +/// The prefix identity of a request outside the session projection, whose +/// history is exactly what it says: nothing replaced anything. +pub fn direct_prefix( + instructions: String, + tools: List(types.Tool), + inputs: List(types.Input), +) -> Prefix { + Prefix( + head_hash: head_hash(instructions, tools), + inputs: list.length(inputs), + replaced: None, + projection_hash: None, + strategy: None, + ) +} + +/// The hash of a request head: its instructions and the tools it offers. +pub fn head_hash(instructions: String, tools: List(types.Tool)) -> String { + json.object([ + #("instructions", json.string(instructions)), + #( + "tools", + json.array(tools, fn(tool) { + json.object([ + #("name", json.string(tool.name)), + #("description", json.string(tool.description)), + #("parameters", tool.parameters), + #("strict", json.bool(tool.strict)), + ]) + }), + ), + ]) + |> json.to_string + |> hash +} + +/// The hash of a run of inputs, as the ledger identities them. +pub fn inputs_hash(inputs: List(types.Input)) -> String { + inputs + |> list.map(input_identity) + |> string.join("\n") + |> hash +} + +fn input_identity(input: types.Input) -> String { + case input { + types.User(text) -> "user\n" <> text + types.Assistant(text) -> "assistant\n" <> text + types.UserImage(text, image) -> + "user_image\n" <> image_identity(image) <> "\n" <> text + types.ToolOutput(id, output, images) -> + "tool_output\n" + <> id + <> "\n" + <> output + <> string.concat( + list.map(images, fn(image) { "\n" <> image_identity(image) }), + ) + types.Replay(item) -> "replay\n" <> json.to_string(types.replay_json(item)) + } +} + +/// An image by the content identity its payload already carries. +fn image_identity(image: types.Image) -> String { + case types.image_data(image) { + types.InlineData(data) -> hash(data) + types.StoredData(hash, _, _) -> hash + } +} + +fn hash(text: String) -> String { + text + |> bit_array.from_string + |> crypto.hash(crypto.Sha256, _) + |> bit_array.base16_encode + |> string.lowercase +} + +/// One provider call, as it is about to be recorded. +pub type Call { + Call( + session: String, + kind: Kind, + /// The saved profile the request went through. + profile: String, + /// How the request reached the provider: protocol and endpoint. + provider: String, + /// Non-secret label of the account that served; None without a pool. + account: Option(String), + model: String, + started_ms: Int, + finished_ms: Int, + outcome: Outcome, + usage: Option(types.Usage), + prefix: Prefix, + ) +} + +/// Writes one row and answers its id, so the caller can attach the transcript +/// seq the call produced once it is committed. +pub fn record(ledger: store.Store, call: Call) -> Result(Int, String) { + store.query(ledger, fn(db) { + use _ <- result.try(store.run(db, insert, values(call))) + store.one( + db, + "SELECT last_insert_rowid()", + [], + decode.field(0, decode.int, decode.success), + "request ledger row id", + ) + }) +} + +// `seq` starts null and is attached after the call's transcript row commits. +const insert = "INSERT INTO request_ledger(session,seq,kind,profile,provider,account,model,started_ms,finished_ms,outcome,status,error,input_tokens,cached_input_tokens,cache_creation_tokens,cache_write_5m_tokens,cache_write_1h_tokens,output_tokens,reasoning_tokens,head_hash,inputs,replaced,projection_hash,strategy) VALUES(?,NULL,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)" + +fn values(call: Call) -> List(sqlight.Value) { + let Call( + session, + kind, + profile, + provider, + account, + model, + started_ms, + finished_ms, + outcome, + usage, + prefix, + ) = call + let Prefix(head_hash, inputs, replaced, projection_hash, strategy) = prefix + let #(status, error) = case outcome { + Completed -> #(None, None) + Failed(status, detail) -> #(status, Some(bounded(detail))) + } + let #(input, output, cached, creation, write_5m, write_1h, reasoning) = case + usage + { + Some(types.Usage( + input, + output, + cached, + creation, + write_5m, + write_1h, + reasoning, + )) -> #( + Some(input), + Some(output), + cached, + creation, + write_5m, + write_1h, + reasoning, + ) + None -> #(None, None, None, None, None, None, None) + } + [ + sqlight.text(session), + sqlight.text(kind_name(kind)), + sqlight.text(profile), + sqlight.text(provider), + sqlight.nullable(sqlight.text, account), + sqlight.text(model), + sqlight.int(started_ms), + sqlight.int(finished_ms), + sqlight.text(case outcome { + Completed -> "ok" + Failed(..) -> "error" + }), + sqlight.nullable(sqlight.int, status), + sqlight.nullable(sqlight.text, error), + sqlight.nullable(sqlight.int, input), + sqlight.nullable(sqlight.int, cached), + sqlight.nullable(sqlight.int, creation), + sqlight.nullable(sqlight.int, write_5m), + sqlight.nullable(sqlight.int, write_1h), + sqlight.nullable(sqlight.int, output), + sqlight.nullable(sqlight.int, reasoning), + sqlight.text(head_hash), + sqlight.int(inputs), + sqlight.nullable(sqlight.int, replaced), + sqlight.nullable(sqlight.text, projection_hash), + sqlight.nullable(sqlight.text, strategy), + ] +} + +/// Attaches the transcript row a recorded call produced. +pub fn attach( + ledger: store.Store, + session: String, + row: Int, + seq: Int, +) -> Result(Nil, String) { + store.write( + ledger, + "UPDATE request_ledger SET seq=? WHERE id=? AND session=?", + [sqlight.int(seq), sqlight.int(row), sqlight.text(session)], + ) +} + +/// One stored row, as the read path answers it. +pub type Row { + Row( + id: Int, + session: String, + seq: Option(Int), + kind: String, + profile: String, + provider: String, + account: Option(String), + model: String, + started_ms: Int, + finished_ms: Int, + outcome: String, + status: Option(Int), + error: Option(String), + input_tokens: Option(Int), + cached_input_tokens: Option(Int), + cache_creation_tokens: Option(Int), + cache_write_5m_tokens: Option(Int), + cache_write_1h_tokens: Option(Int), + output_tokens: Option(Int), + reasoning_tokens: Option(Int), + head_hash: String, + inputs: Int, + replaced: Option(Int), + projection_hash: Option(String), + strategy: Option(String), + ) +} + +/// Rows after `after`, oldest first, at most `limit`; the id to continue from. +pub fn page( + ledger: store.Store, + session: String, + after: Int, + limit: Int, +) -> Result(#(List(Row), Int), String) { + store.query(ledger, fn(db) { + use rows <- result.try(store.rows( + db, + "SELECT " + <> columns + <> " FROM request_ledger WHERE session=? AND id>? ORDER BY id LIMIT ?", + [sqlight.text(session), sqlight.int(after), sqlight.int(limit)], + row_decoder(), + )) + let next = + rows + |> list.last + |> result.map(fn(row) { row.id }) + |> result.unwrap(after) + Ok(#(rows, next)) + }) +} + +// Named, in decoder order, so a column added later cannot shift a read. +const columns = "id,session,seq,kind,profile,provider,account,model,started_ms,finished_ms,outcome,status,error,input_tokens,cached_input_tokens,cache_creation_tokens,cache_write_5m_tokens,cache_write_1h_tokens,output_tokens,reasoning_tokens,head_hash,inputs,replaced,projection_hash,strategy" + +fn row_decoder() -> decode.Decoder(Row) { + use id <- decode.field(0, decode.int) + use session <- decode.field(1, decode.string) + use seq <- decode.field(2, decode.optional(decode.int)) + use kind <- decode.field(3, decode.string) + use profile <- decode.field(4, decode.string) + use provider <- decode.field(5, decode.string) + use account <- decode.field(6, decode.optional(decode.string)) + use model <- decode.field(7, decode.string) + use started <- decode.field(8, decode.int) + use finished <- decode.field(9, decode.int) + use outcome <- decode.field(10, decode.string) + use status <- decode.field(11, decode.optional(decode.int)) + use error <- decode.field(12, decode.optional(decode.string)) + use input <- decode.field(13, decode.optional(decode.int)) + use cached <- decode.field(14, decode.optional(decode.int)) + use creation <- decode.field(15, decode.optional(decode.int)) + use write_5m <- decode.field(16, decode.optional(decode.int)) + use write_1h <- decode.field(17, decode.optional(decode.int)) + use output <- decode.field(18, decode.optional(decode.int)) + use reasoning <- decode.field(19, decode.optional(decode.int)) + use head_hash <- decode.field(20, decode.string) + use inputs <- decode.field(21, decode.int) + use replaced <- decode.field(22, decode.optional(decode.int)) + use projection_hash <- decode.field(23, decode.optional(decode.string)) + use strategy <- decode.field(24, decode.optional(decode.string)) + decode.success(Row( + id, + session, + seq, + kind, + profile, + provider, + account, + model, + started, + finished, + outcome, + status, + error, + input, + cached, + creation, + write_5m, + write_1h, + output, + reasoning, + head_hash, + inputs, + replaced, + projection_hash, + strategy, + )) +} + +/// One row as the read route serves it. Unreported counts are null, so a +/// reader distinguishes unknown from zero. +pub fn row_json(row: Row) -> Json { + json.object([ + #("id", json.int(row.id)), + #("session", json.string(row.session)), + #("seq", json.nullable(row.seq, json.int)), + #("kind", json.string(row.kind)), + #("profile", json.string(row.profile)), + #("provider", json.string(row.provider)), + #("account", json.nullable(row.account, json.string)), + #("model", json.string(row.model)), + #("startedMs", json.int(row.started_ms)), + #("finishedMs", json.int(row.finished_ms)), + #("outcome", json.string(row.outcome)), + #("status", json.nullable(row.status, json.int)), + #("error", json.nullable(row.error, json.string)), + #("inputTokens", json.nullable(row.input_tokens, json.int)), + #("cachedInputTokens", json.nullable(row.cached_input_tokens, json.int)), + #("cacheCreationTokens", json.nullable(row.cache_creation_tokens, json.int)), + #("cacheWrite5mTokens", json.nullable(row.cache_write_5m_tokens, json.int)), + #("cacheWrite1hTokens", json.nullable(row.cache_write_1h_tokens, json.int)), + #("outputTokens", json.nullable(row.output_tokens, json.int)), + #("reasoningTokens", json.nullable(row.reasoning_tokens, json.int)), + #("headHash", json.string(row.head_hash)), + #("inputs", json.int(row.inputs)), + #("replaced", json.nullable(row.replaced, json.int)), + #("projectionHash", json.nullable(row.projection_hash, json.string)), + #("strategy", json.nullable(row.strategy, json.string)), + ]) +} + +/// The ms clock the ledger timestamps rows with. +@external(erlang, "erlang", "system_time") +fn system_time(unit: TimeUnit) -> Int + +type TimeUnit { + Millisecond +} + +pub fn now() -> Int { + system_time(Millisecond) +} + +/// Rows a page asks for when it does not say. +pub const page_rows = 100 diff --git a/src/albedo/daemon/quota.gleam b/src/albedo/daemon/quota.gleam new file mode 100644 index 00000000..bfd6a502 --- /dev/null +++ b/src/albedo/daemon/quota.gleam @@ -0,0 +1,810 @@ +//// Provider quota readings, polled per account and recorded raw. +//// +//// One poller for the whole daemon owns the schedule. Every account albedo +//// can see is polled on its own cadence: each OAuth account in a provider's +//// rotation pool, and each API-key profile whose base url names a feed. Never +//// two polls of one account at once; a limit at or above 80 percent, or one +//// resetting within fifteen minutes, moves that account onto the busy +//// cadence; a failure backs off. Credentials come from the same readers and +//// refreshers the request path uses, so polling refreshes tokens the way +//// requests do. +//// +//// Only readings are stored — percentages, windows, reset times, statuses and +//// errors exactly as the feed reported them, never anything derived. Phase 2 +//// derives what it needs from these rows. + +import albedo/daemon/configuration +import albedo/daemon/store +import albedo/harness/settings +import albedo/harness/usage_feed.{type Limit, type Report} +import gleam/dict.{type Dict} +import gleam/dynamic.{type Dynamic} +import gleam/dynamic/decode +import gleam/erlang/process.{type Subject} +import gleam/int +import gleam/io +import gleam/json +import gleam/list +import gleam/option.{type Option, None, Some} +import gleam/otp/actor +import gleam/result +import gleam/string +import gleam/uri +import sqlight + +/// The feed the poller polls through: the real one, or a fake in tests. +pub type Fetch = + fn(String, json.Json, Int) -> Result(Report, String) + +/// One polled account: a provide-usage provider id, the account's non-secret +/// label, and the credential to poll it with. +pub type Target { + Target(provider: String, label: String, credential: json.Json) +} + +pub type Settings { + Settings(poll_seconds: Int, busy_poll_seconds: Int, enabled: Bool) +} + +const default_settings = Settings(600, 120, True) + +/// `$ALBEDO_HOME/extensions.json` under `"quota"`, read like any extension's +/// settings, and re-read every round so changes apply without a restart. +fn load_settings() -> Settings { + let decoder = { + use poll_seconds <- decode.optional_field("pollSeconds", 600, decode.int) + use busy_poll_seconds <- decode.optional_field( + "busyPollSeconds", + 120, + decode.int, + ) + use enabled <- decode.optional_field("enabled", True, decode.bool) + decode.success(Settings(poll_seconds, busy_poll_seconds, enabled)) + } + settings.load("quota", decoder, default_settings) + |> result.unwrap(default_settings) + |> clamped +} + +fn clamped(settings: Settings) -> Settings { + let poll_seconds = int.clamp(settings.poll_seconds, 1, 86_400) + Settings( + poll_seconds, + int.clamp(settings.busy_poll_seconds, 1, poll_seconds), + settings.enabled, + ) +} + +/// The accounts this daemon can poll, with credentials refreshed the way +/// requests refresh them. +fn enumerate(home: String) -> List(Target) { + list.flatten([ + oauth_targets("anthropic", claude_accounts(home)), + oauth_targets("openai-codex", codex_accounts(home)), + oauth_targets("google-antigravity", antigravity_accounts(home)), + profile_targets(home), + ]) +} + +@external(erlang, "albedo_claude_auth", "accounts") +fn claude_accounts(home: String) -> List(Dynamic) + +@external(erlang, "albedo_openai_auth", "accounts") +fn codex_accounts(home: String) -> List(Dynamic) + +@external(erlang, "albedo_antigravity", "accounts") +fn antigravity_accounts(home: String) -> List(Dynamic) + +/// A stored OAuth account: the label the provider's own identity logic gives +/// it, plus the provide-usage credential fields. +type Stored { + Stored( + label: String, + access: String, + account_id: String, + email: String, + project_id: String, + ) +} + +fn stored_decoder() -> decode.Decoder(Stored) { + use label <- decode.field("label", decode.string) + use access <- decode.optional_field("access", "", decode.string) + use account_id <- decode.optional_field("accountId", "", decode.string) + use email <- decode.optional_field("email", "", decode.string) + use project_id <- decode.optional_field("projectId", "", decode.string) + decode.success(Stored(label, access, account_id, email, project_id)) +} + +fn oauth_targets(provider: String, stored: List(Dynamic)) -> List(Target) { + stored + |> list.filter_map(fn(value) { + decode.run(value, stored_decoder()) |> result.replace_error(Nil) + }) + |> list.filter_map(fn(account) { + case account.access, account.label { + "", _ | _, "" -> Error(Nil) + _, _ -> Ok(Target(provider, account.label, oauth_credential(account))) + } + }) +} + +/// Only the access token crosses into the core: albedo refreshes it, and the +/// core never reads a refresh token. +fn oauth_credential(account: Stored) -> json.Json { + let fields = [ + #("kind", json.string("oauth")), + #("access", json.string(account.access)), + ..text_fields([ + #("accountId", account.account_id), + #("email", account.email), + ]) + ] + // Antigravity's quota is per Cloud Code Assist project. + case account.project_id { + "" -> json.object(fields) + project_id -> + json.object( + list.append(fields, [ + #("extra", json.object([#("projectId", json.string(project_id))])), + ]), + ) + } +} + +/// Config profiles: an alibaba profile polls through the signed-in `bl` CLI, +/// and a generic `openai` profile maps to a feed by its base url host. A +/// profile whose host has no feed is skipped. +fn profile_targets(home: String) -> List(Target) { + configuration.providers(home) + |> result.unwrap([]) + |> list.filter_map(fn(profile) { + case profile.extension { + "alibaba" -> + Ok(Target( + "alibaba", + "alibaba", + json.object([ + #("kind", json.string("cli")), + ]), + )) + "openai" -> key_target(home, profile.name) + _ -> Error(Nil) + } + }) +} + +fn key_target(home: String, name: String) -> Result(Target, Nil) { + use #(base_url, api_key) <- result.try( + configuration.settings(home, name, profile_decoder()) + |> result.replace_error(Nil), + ) + use provider <- result.try(feed_provider(base_url)) + case string.trim(api_key) { + "" -> Error(Nil) + key -> + Ok(Target( + provider, + name, + json.object([ + #("kind", json.string("apiKey")), + #("apiKey", json.string(key)), + ]), + )) + } +} + +fn profile_decoder() -> decode.Decoder(#(String, String)) { + use base_url <- decode.optional_field("baseUrl", "", decode.string) + use api_key <- decode.optional_field("apiKey", "", decode.string) + decode.success(#(base_url, api_key)) +} + +fn feed_provider(base_url: String) -> Result(String, Nil) { + use parsed <- result.try(uri.parse(base_url) |> result.replace_error(Nil)) + use host <- result.try(parsed.host |> option.to_result(Nil)) + case string.lowercase(host) { + "hyper.charm.land" -> Ok("hyper") + "api.deepseek.com" -> Ok("deepseek") + _ -> Error(Nil) + } +} + +fn text_fields(fields: List(#(String, String))) -> List(#(String, json.Json)) { + list.filter_map(fields, fn(field) { + case field.1 { + "" -> Error(Nil) + value -> Ok(#(field.0, json.string(value))) + } + }) +} + +// ---- the poller --------------------------------------------------------- + +pub type Message { + Tick + /// The account list, as a spawn saw it. + Enumerated(List(Target)) + /// One poll finished: the account key, whether it failed, and whether it + /// is busy. + Polled(key: String, failed: Bool, busy: Bool) + Down(process.Down) +} + +type Account { + Account( + /// provider <> "/" <> label: one polled account, never two polls at once. + key: String, + target: Target, + /// Wall-clock milliseconds for this account's next poll. + next_due: Int, + /// Consecutive failed polls, for backoff. + failures: Int, + /// The process polling now, so a crash frees the slot. + polling: Option(process.Pid), + ) +} + +type State { + State( + home: String, + database: store.Store, + fetch: Fetch, + accounts: Dict(String, Account), + self: Subject(Message), + /// A tick is already scheduled. + ticking: Bool, + /// The process listing accounts now, so a crash cannot wedge it. + enumerating: Option(process.Pid), + last_enumerated: Int, + ) +} + +/// One poller for the whole daemon. It starts with the daemon, restarts under +/// its supervisor, and polls each account it can see on its own cadence. +pub fn start( + home: String, + database: store.Store, + fetch: Fetch, +) -> Result(actor.Started(Subject(Message)), actor.StartError) { + actor.new_with_initialiser(30_000, fn(self) { + label("albedo_quota", home) + let _ = process.send_after(self, 0, Tick) + Ok( + actor.initialised(State( + home, + database, + fetch, + dict.new(), + self, + True, + None, + 0, + )) + |> actor.returning(self) + |> actor.selecting( + process.new_selector() + |> process.select(self) + |> process.select_monitors(Down), + ), + ) + }) + |> actor.on_message(handle) + |> actor.start +} + +fn handle(state: State, message: Message) -> actor.Next(State, Message) { + case message { + Tick -> retick(State(..on_tick(state), ticking: False)) + Enumerated(targets) -> + retick(merged(State(..state, enumerating: None), targets, now_ms())) + Polled(key, failed, busy) -> retick(finished(state, key, failed, busy)) + Down(process.ProcessDown(_, pid, _)) -> retick(fell(state, pid)) + // The poller monitors processes only; a port is never one of them. + Down(process.PortDown(_, _, _)) -> actor.continue(state) + } +} + +fn on_tick(state: State) -> State { + let config = load_settings() + let now = now_ms() + case config.enabled { + False -> State(..state, accounts: dict.new()) + True -> { + let state = case + state.enumerating, + now - state.last_enumerated >= config.poll_seconds * 1000 + { + None, True -> enumerate_now(state) + _, _ -> state + } + start_due(state, now) + } + } +} + +/// Account discovery runs off the poller: a refresh may wait on the network. +fn enumerate_now(state: State) -> State { + let self = state.self + let home = state.home + let pid = + process.spawn_unlinked(fn() { + process.send(self, Enumerated(enumerate(home))) + }) + let _ = process.monitor(pid) + State(..state, enumerating: Some(pid)) +} + +/// New accounts poll at once; an account that is polling now finishes first, +/// so its readings stay even if the account has since gone. +fn merged(state: State, targets: List(Target), now: Int) -> State { + let fresh = + list.fold(targets, dict.new(), fn(accounts, target) { + let key = key_of(target) + let account = case dict.get(state.accounts, key) { + Ok(existing) -> Account(..existing, target: target) + Error(_) -> Account(key, target, now, 0, None) + } + dict.insert(accounts, key, account) + }) + let accounts = + dict.fold(state.accounts, fresh, fn(accounts, key, account) { + case dict.has_key(accounts, key), account.polling { + False, Some(_) -> dict.insert(accounts, key, account) + _, _ -> accounts + } + }) + State(..state, accounts: accounts) +} + +fn start_due(state: State, now: Int) -> State { + let fetch = state.fetch + let database = state.database + let self = state.self + dict.fold(state.accounts, state, fn(state, key, account) { + case account.next_due <= now, account.polling { + True, None -> { + let target = account.target + let pid = + process.spawn_unlinked(fn() { + let #(failed, busy) = run(fetch, database, target) + process.send(self, Polled(key, failed, busy)) + }) + let _ = process.monitor(pid) + State( + ..state, + accounts: dict.insert( + state.accounts, + key, + Account(..account, polling: Some(pid)), + ), + ) + } + _, _ -> state + } + }) +} + +/// The next wait for an account: the busy cadence while any limit is at or +/// above 80 percent or resets within fifteen minutes, the ordinary cadence +/// otherwise, and a doubling backoff on failure. +fn finished(state: State, key: String, failed: Bool, busy: Bool) -> State { + case dict.get(state.accounts, key) { + Error(_) -> state + Ok(account) -> { + let config = load_settings() + let failures = case failed { + True -> account.failures + 1 + False -> 0 + } + let wait = case failed { + True -> + config.poll_seconds * int.bitwise_shift_left(1, int.min(failures, 6)) + |> int.min(3600) + False -> + case busy { + True -> config.busy_poll_seconds + False -> config.poll_seconds + } + } + State( + ..state, + accounts: dict.insert( + state.accounts, + key, + Account( + ..account, + failures: failures, + polling: None, + next_due: now_ms() + wait * 1000, + ), + ), + ) + } + } +} + +/// A spawned poll that crashed never answers `Polled`: free its slot and +/// back the account off, so a crashing feed cannot wedge the schedule. +fn fell(state: State, pid: process.Pid) -> State { + case state.enumerating { + Some(enumerating) if enumerating == pid -> State(..state, enumerating: None) + _ -> + case + dict.to_list(state.accounts) + |> list.find(fn(entry) { entry.1.polling == Some(pid) }) + { + Error(_) -> state + Ok(#(key, _)) -> finished(state, key, True, False) + } + } +} + +/// One tick in flight at a time: the next is scheduled for whichever account +/// or enumeration is soonest, so the busy cadence is honoured exactly. +fn retick(state: State) -> actor.Next(State, Message) { + case state.ticking { + True -> actor.continue(state) + False -> { + let _ = process.send_after(state.self, tick_delay(state), Tick) + actor.continue(State(..state, ticking: True)) + } + } +} + +fn tick_delay(state: State) -> Int { + let config = load_settings() + case config.enabled { + False -> config.poll_seconds * 1000 + True -> { + let now = now_ms() + let due = + [ + list.map(dict.values(state.accounts), fn(account) { account.next_due }), + [state.last_enumerated + config.poll_seconds * 1000], + ] + |> list.flatten + let delay = case minimum(due) { + None -> config.poll_seconds * 1000 + Some(when) -> when - now + } + int.clamp(delay, 250, config.poll_seconds * 1000) + } + } +} + +fn minimum(values: List(Int)) -> Option(Int) { + list.fold(values, None, fn(best, value) { + case best { + Some(current) -> Some(int.min(current, value)) + None -> Some(value) + } + }) +} + +/// One poll of one account, off the poller: whatever comes back is recorded +/// exactly as it came, then the schedule hears whether it failed and whether +/// the account is busy. +fn run(fetch: Fetch, database: store.Store, target: Target) -> #(Bool, Bool) { + let now = now_ms() + case fetch(target.provider, target.credential, now) { + Ok(report) -> { + record(database, target, report.plan, report.limits, report.error, now) + #(report.error != None, is_busy(report, now)) + } + Error(message) -> { + record(database, target, None, [], Some(message), now) + #(True, False) + } + } +} + +fn is_busy(report: Report, now: Int) -> Bool { + list.any(report.limits, fn(limit) { + case limit.used_percent { + Some(used) if used >=. 80.0 -> True + _ -> + case limit.resets_at { + Some(resets_at) -> + resets_at > now && resets_at - now <= 15 * 60 * 1000 + None -> False + } + } + }) +} + +// ---- the readings ------------------------------------------------------- + +/// One stored reading, exactly as the feed reported it. +pub type Sample { + Sample( + id: Int, + /// The account's non-secret label. + account: String, + provider: String, + /// The report plan: what 100 percent means can change with it. + plan: Option(String), + limit_id: String, + label: String, + used_percent: Option(Float), + window_label: Option(String), + window_seconds: Option(Int), + resets_at: Option(Int), + scope: Option(String), + status: String, + error: Option(String), + observed_at: Int, + source: String, + ) +} + +const schema = " +CREATE TABLE IF NOT EXISTS quota_sample ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + account TEXT NOT NULL, + provider TEXT NOT NULL, + plan TEXT, + limit_id TEXT NOT NULL DEFAULT '', + label TEXT NOT NULL DEFAULT '', + used_percent REAL, + window_label TEXT, + window_seconds INTEGER, + resets_at INTEGER, + scope TEXT, + status TEXT NOT NULL DEFAULT '', + error TEXT, + observed_at INTEGER NOT NULL, + source TEXT NOT NULL DEFAULT 'poll' +); +CREATE INDEX IF NOT EXISTS quota_sample_latest + ON quota_sample(account, provider, limit_id, id); +" + +/// The table, before any session or poller starts. +pub fn initialise(database: store.Store) -> Result(Nil, String) { + store.query(database, fn(db) { store.exec(db, schema) }) +} + +/// Only readings: one row per limit the report carried, or one row carrying +/// the report's error when it carried none. +fn record( + database: store.Store, + target: Target, + plan: Option(String), + limits: List(Limit), + error: Option(String), + now: Int, +) -> Nil { + let rows = case limits { + [] -> [ + Row( + target.label, + target.provider, + plan, + "", + "", + None, + None, + None, + None, + None, + "", + error, + now, + ), + ] + _ -> + list.map(limits, fn(limit) { + Row( + target.label, + target.provider, + plan, + limit.id, + limit.label, + limit.used_percent, + limit.window_label, + limit.window_seconds, + limit.resets_at, + limit.scope, + limit.status, + error, + now, + ) + }) + } + let outcome = + store.query(database, fn(db) { + list.try_each(rows, fn(row) { insert(db, row) }) + }) + case outcome { + Ok(_) -> Nil + Error(message) -> io.println_error("quota sample write failed: " <> message) + } +} + +type Row { + Row( + account: String, + provider: String, + plan: Option(String), + limit_id: String, + label: String, + used_percent: Option(Float), + window_label: Option(String), + window_seconds: Option(Int), + resets_at: Option(Int), + scope: Option(String), + status: String, + error: Option(String), + observed_at: Int, + ) +} + +fn insert(db: sqlight.Connection, row: Row) -> Result(Nil, String) { + store.run( + db, + "INSERT INTO quota_sample(account, provider, plan, limit_id, label, used_percent, window_label, window_seconds, resets_at, scope, status, error, observed_at, source) + VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'poll')", + [ + sqlight.text(row.account), + sqlight.text(row.provider), + sqlight.nullable(sqlight.text, row.plan), + sqlight.text(row.limit_id), + sqlight.text(row.label), + sqlight.nullable(sqlight.float, row.used_percent), + sqlight.nullable(sqlight.text, row.window_label), + sqlight.nullable(sqlight.int, row.window_seconds), + sqlight.nullable(sqlight.int, row.resets_at), + sqlight.nullable(sqlight.text, row.scope), + sqlight.text(row.status), + sqlight.nullable(sqlight.text, row.error), + sqlight.int(row.observed_at), + ], + ) +} + +const sample_columns = "id, account, provider, plan, limit_id, label, used_percent, window_label, window_seconds, resets_at, scope, status, error, observed_at, source" + +fn sample_decoder() -> decode.Decoder(Sample) { + use id <- decode.field(0, decode.int) + use account <- decode.field(1, decode.string) + use provider <- decode.field(2, decode.string) + use plan <- decode.field(3, decode.optional(decode.string)) + use limit_id <- decode.field(4, decode.string) + use label <- decode.field(5, decode.string) + use used_percent <- decode.field(6, decode.optional(decode.float)) + use window_label <- decode.field(7, decode.optional(decode.string)) + use window_seconds <- decode.field(8, decode.optional(decode.int)) + use resets_at <- decode.field(9, decode.optional(decode.int)) + use scope <- decode.field(10, decode.optional(decode.string)) + use status <- decode.field(11, decode.string) + use error <- decode.field(12, decode.optional(decode.string)) + use observed_at <- decode.field(13, decode.int) + use source <- decode.field(14, decode.string) + decode.success(Sample( + id, + account, + provider, + plan, + limit_id, + label, + used_percent, + window_label, + window_seconds, + resets_at, + scope, + status, + error, + observed_at, + source, + )) +} + +/// The latest reading per account and limit, read-only. +pub fn latest(database: store.Store) -> Result(List(Sample), String) { + store.read(database, "SELECT " <> sample_columns <> " FROM quota_sample s + WHERE s.id = (SELECT MAX(t.id) FROM quota_sample t + WHERE t.account = s.account AND t.provider = s.provider + AND t.limit_id = s.limit_id) + ORDER BY s.account, s.provider, s.limit_id", [], sample_decoder()) +} + +/// Readings newest-first with id keyset paging: `before` 0 starts at the +/// newest, any other value continues after that row id. +pub fn history( + database: store.Store, + before: Int, + limit: Int, +) -> Result(List(Sample), String) { + case before <= 0 { + True -> + store.read( + database, + "SELECT " + <> sample_columns + <> " FROM quota_sample ORDER BY id DESC LIMIT ?", + [sqlight.int(limit)], + sample_decoder(), + ) + False -> + store.read( + database, + "SELECT " + <> sample_columns + <> " FROM quota_sample WHERE id < ? ORDER BY id DESC LIMIT ?", + [sqlight.int(before), sqlight.int(limit)], + sample_decoder(), + ) + } +} + +/// The latest readings, as the `/quota` route serves them. +pub fn latest_json(database: store.Store) -> Result(json.Json, String) { + use samples <- result.try(latest(database)) + Ok( + json.object([ + #("readings", json.array(samples, sample_json)), + ]), + ) +} + +/// One history page: the rows and where the next page continues. +pub fn page_json( + database: store.Store, + before: Int, + limit: Int, +) -> Result(json.Json, String) { + use samples <- result.try(history(database, before, limit)) + let next = case list.length(samples) < limit, list.last(samples) { + True, _ -> None + False, Ok(last) -> Some(last.id) + False, Error(_) -> None + } + Ok( + json.object([ + #("items", json.array(samples, sample_json)), + #("nextCursor", json.nullable(next, json.int)), + #("hasMore", json.bool(next != None)), + ]), + ) +} + +/// A reading, exactly as it was recorded: unknown stays unknown, so a +/// missing percentage or reset time is never mistaken for zero. +pub fn sample_json(sample: Sample) -> json.Json { + json.object([ + #("id", json.int(sample.id)), + #("account", json.string(sample.account)), + #("provider", json.string(sample.provider)), + #("plan", json.nullable(sample.plan, json.string)), + #("limitId", json.string(sample.limit_id)), + #("label", json.string(sample.label)), + #("usedPercent", json.nullable(sample.used_percent, json.float)), + #("windowLabel", json.nullable(sample.window_label, json.string)), + #("windowSeconds", json.nullable(sample.window_seconds, json.int)), + #("resetsAt", json.nullable(sample.resets_at, json.int)), + #("scope", json.nullable(sample.scope, json.string)), + #("status", json.string(sample.status)), + #("error", json.nullable(sample.error, json.string)), + #("observedAt", json.int(sample.observed_at)), + #("source", json.string(sample.source)), + ]) +} + +fn key_of(target: Target) -> String { + target.provider <> "/" <> target.label +} + +type TimeUnit { + Millisecond +} + +@external(erlang, "erlang", "system_time") +fn system_time(unit: TimeUnit) -> Int + +fn now_ms() -> Int { + system_time(Millisecond) +} + +@external(erlang, "albedo_inspect", "label") +fn label(kind: String, id: String) -> Nil diff --git a/src/albedo/daemon/server.gleam b/src/albedo/daemon/server.gleam index 794d3a2d..02fecedf 100644 --- a/src/albedo/daemon/server.gleam +++ b/src/albedo/daemon/server.gleam @@ -6,7 +6,9 @@ import albedo/daemon/family import albedo/daemon/history import albedo/daemon/image import albedo/daemon/images +import albedo/daemon/ledger import albedo/daemon/mail +import albedo/daemon/quota import albedo/daemon/reaper import albedo/daemon/session import albedo/daemon/session_provider @@ -17,6 +19,7 @@ import albedo/harness/extension import albedo/harness/extensions/schedule/ledger as schedule import albedo/harness/oauth import albedo/harness/runtime +import albedo/harness/usage_feed import albedo/openai_api/types import gleam/bytes_tree import gleam/dict.{type Dict} @@ -169,6 +172,13 @@ pub fn start(config: Config, port: Int) -> Result(Int, String) { |> supervisor.add( supervision.worker(fn() { start_registry(config, host, name) }), ) + // One daemon-wide quota poller: every account it can see, on its own + // cadence, recorded as raw readings. + |> supervisor.add( + supervision.worker(fn() { + quota.start(config.home, runtime.ledger(host), usage_feed.fetch) + }), + ) |> supervisor.start |> result.map_error(string.inspect), ) @@ -526,6 +536,7 @@ fn prepare_storage( host: runtime.Runtime, ) -> Result(Nil, String) { use _ <- result.try(conversation.initialise(runtime.ledger(host))) + use _ <- result.try(quota.initialise(runtime.ledger(host))) use _ <- result.try(mail.initialise(runtime.ledger(host))) use _ <- result.try(family.initialise(runtime.ledger(host))) use moved <- result.try(images.migrate( @@ -1549,7 +1560,7 @@ fn uri_decode(segment: String) -> String { /// The daemon's own top-level routes; a service never shadows them. const daemon_routes = [ - "health", "sessions", "models", "auth", "shutdown", "agents", + "health", "sessions", "models", "auth", "shutdown", "agents", "quota", ] fn route( @@ -1673,6 +1684,24 @@ fn daemon_route( ) } Get, ["agents", "stream"] -> agents_stream(req) + // Quota readings, read-only: the latest per account and limit, or + // with `?history=` the raw samples after that row, newest first. + Get, ["quota"] -> { + use host <- with_host(registry) + let database = runtime.ledger(host) + case list.key_find(query(req), "history") { + Ok(cursor) -> + quota.page_json( + database, + int.parse(cursor) |> result.unwrap(0), + query_int(req, "limit", 50) |> int.clamp(1, 200), + ) + |> answered(200, fn(value) { value }, 400) + Error(_) -> + quota.latest_json(database) + |> answered(200, fn(value) { value }, 400) + } + } // `details` lists objects with catalog facts. Without it, plain ids. // Efforts come from where sessions of this extension look them up, // and facts from the endpoint, or there when none is given. @@ -1776,6 +1805,27 @@ fn daemon_route( 400, ) } + // The session's recorded provider calls, oldest first, for tests and + // the quota inspector: `after` is the last row id to continue from. + Get, ["sessions", id, "ledger"] -> { + let after = query_int(req, "after", 0) + let limit = + query_int(req, "limit", ledger.page_rows) + |> int.clamp(1, 500) + use host <- with_host(registry) + ledger.page(runtime.ledger(host), id, after, limit) + |> answered( + 200, + fn(page) { + json.object([ + #("session", json.string(id)), + #("rows", json.array(page.0, ledger.row_json)), + #("after", json.int(page.1)), + ]) + }, + 400, + ) + } // Older transcript as rendered events, for a client scrolled past what // its reset carried: `before` is the `before` of the page it holds. Get, ["sessions", id, "history"] -> { diff --git a/src/albedo/daemon/session.gleam b/src/albedo/daemon/session.gleam index 6e226615..9c369191 100644 --- a/src/albedo/daemon/session.gleam +++ b/src/albedo/daemon/session.gleam @@ -125,7 +125,7 @@ pub type Message { List(types.Input), conversation.Stage, Option(Int), - Subject(Result(Int, String)), + Subject(Result(#(Int, Option(Int)), String)), ) RecordContext(String, context_snapshot.Snapshot, Bool, Subject(Nil)) RecordUsage(String, usage.Metadata, Subject(Result(Nil, String))) @@ -658,7 +658,7 @@ fn handle(state: State, message: Message) { thought_ms, ) case written { - Ok(timestamp) -> + Ok(#(timestamp, _)) -> answer( session_state.State( ..session_history.remember_response( @@ -1533,7 +1533,7 @@ fn finish_run( usage.Metadata( state.info.model, usage.now(), - Some(usage.Tokens(tokens, 0, None, None)), + Some(usage.Tokens(tokens, 0, None, None, None, None, None)), ) session_state.emit( session_state.State(..state, latest_usage: Some(metadata)), diff --git a/src/albedo/daemon/session_run.gleam b/src/albedo/daemon/session_run.gleam index 32e62d2c..6df7aa90 100644 --- a/src/albedo/daemon/session_run.gleam +++ b/src/albedo/daemon/session_run.gleam @@ -20,7 +20,7 @@ pub type Messages(message) { List(types.Input), conversation.Stage, Option(Int), - Subject(Result(Int, String)), + Subject(Result(#(Int, Option(Int)), String)), ) -> message, context: fn(String, context_snapshot.Snapshot, Bool, Subject(Nil)) -> message, @@ -112,13 +112,16 @@ fn confirm( } } +/// The worker's commit answers the daemon time and the seq of the commit's +/// first assistant row, so the loop can link a recorded provider call to the +/// transcript row it produced. pub fn commit_fn( owner: Subject(message), run_id: String, messages: Messages(message), waiting timeout: Int, ) -> fn(List(types.Input), conversation.Stage, Option(Int)) -> - Result(Int, String) { + Result(#(Int, Option(Int)), String) { fn(inputs, stage, thought_ms) { confirm(owner, "commit", timeout, messages.commit( run_id, @@ -206,6 +209,9 @@ pub fn start( usage_fn(owner, run_id, messages, 10_000), drain_fn(owner, run_id, messages, 10_000), fn(head) { report(owner, messages.pin(run_id, head, _)) }, + // The request ledger's identity for this session's provider calls. + state.info.id, + state.info.provider, ) let pid = process.spawn_unlinked(fn() { diff --git a/src/albedo/daemon/usage.gleam b/src/albedo/daemon/usage.gleam index d080d5ff..6292b35d 100644 --- a/src/albedo/daemon/usage.gleam +++ b/src/albedo/daemon/usage.gleam @@ -16,6 +16,9 @@ pub type Tokens { completion_tokens: Int, cached_prompt_tokens: Option(Int), cache_creation_tokens: Option(Int), + cache_write_5m_tokens: Option(Int), + cache_write_1h_tokens: Option(Int), + reasoning_tokens: Option(Int), ) } @@ -25,8 +28,24 @@ pub fn from_completion( recorded_at: Int, ) -> Metadata { let tokens = case provider_usage { - Some(types.Usage(prompt, completion, cached, creation)) -> - Some(Tokens(prompt, completion, cached, creation)) + Some(types.Usage( + prompt, + completion, + cached, + creation, + write_5m, + write_1h, + reasoning, + )) -> + Some(Tokens( + prompt, + completion, + cached, + creation, + write_5m, + write_1h, + reasoning, + )) None -> None } Metadata(model, recorded_at, tokens) @@ -37,7 +56,15 @@ pub fn from_completion( pub fn event(metadata: Metadata) -> String { let Metadata(model, recorded_at, tokens) = metadata let token_fields = case tokens { - Some(Tokens(prompt, completion, cached, creation)) -> + Some(Tokens( + prompt, + completion, + cached, + creation, + write_5m, + write_1h, + reasoning, + )) -> [ #("promptTokens", json.int(prompt)), #("completionTokens", json.int(completion)), @@ -45,6 +72,9 @@ pub fn event(metadata: Metadata) -> String { ] |> opt_field("cacheCreationTokens", creation) |> opt_field("cachedPromptTokens", cached) + |> opt_field("cacheWrite5mTokens", write_5m) + |> opt_field("cacheWrite1hTokens", write_1h) + |> opt_field("reasoningTokens", reasoning) None -> [] } json.object([ diff --git a/src/albedo/harness/albedo_usage_core.erl b/src/albedo/harness/albedo_usage_core.erl new file mode 100644 index 00000000..7a5a39a6 --- /dev/null +++ b/src/albedo/harness/albedo_usage_core.erl @@ -0,0 +1,218 @@ +%% The I/O half of albedo/harness/usage_feed: usage-core is a pure Zig core, +%% so every byte of I/O happens here. One `usage advance -` process per round, +%% fed the envelope on stdin because it carries the credential; `http` +%% requests go through albedo_http; `command` requests run with an environment +%% allowlist, a timeout and a stdout cap. Nothing the envelope or a response +%% carries is ever logged: both can hold credentials. +-module(albedo_usage_core). + +-export([advance/1, http_request/4, run_command/4]). + +-define(ADVANCE_TIMEOUT_MS, 15000). +-define(COMMAND_TIMEOUT_MS, 15000). +-define(COMMAND_TIMEOUT_CAP_MS, 60000). +-define(STDOUT_CAP, 65536). + +%% PATH, HOME, locale, timezone and proxies: what a CLI needs to find and +%% speak to its own backend, and nothing that could leak this session. +-define(ALLOWED_ENV, + ["PATH", "HOME", "LANG", "LC_ALL", "LC_CTYPE", "TZ", + "HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY", + "http_proxy", "https_proxy", "all_proxy", "no_proxy"]). + +%% One round. The envelope line is the whole message: the CLI stops reading +%% at the first newline, which is what an Erlang port needs (it cannot +%% half-close stdin), and stdout is read to exit. The last output line is +%% the step; the caller parses it, this side only moves bytes. +advance(Envelope) when is_binary(Envelope) -> + case resolve() of + false -> + {error, <<"usage is not built; run nix build or put it on PATH">>}; + Bin -> + Port = open_port({spawn_executable, Bin}, + [{args, ["advance", "-"]}, exit_status, binary, hide]), + Port ! {self(), {command, [Envelope, $\n]}}, + drain(Port, [], deadline(?ADVANCE_TIMEOUT_MS)) + end. + +%% An `http` request through albedo's one front door. A transport failure is +%% data for the feed, not a driver failure: the protocol says "no response" +%% with status 0, and the reason travels in the body. +http_request(Method, Url, Headers, Body) -> + RequestBody = case Body of + none -> none; + {some, Payload} -> {content_type(Headers), Payload} + end, + HeaderList = without_content_type(Headers), + case albedo_http:request(method(Method), Url, HeaderList, RequestBody, 30000, 10000) of + {ok, {Status, _ResponseHeaders, ResponseBody}} -> + {Status, ResponseBody}; + {error, Reason} -> + ReasonText = unicode:characters_to_binary(io_lib:format("~0p", [Reason])), + {0, ReasonText} + end. + +%% A `command` request, run with only the allowed environment, a timeout and +%% a stdout cap. 127 is reserved for "this host will not run commands", so a +%% command that cannot run or timed out answers 126: the feed routes around +%% either the same way. +run_command(Command, Args, Stdin, TimeoutMs) -> + Timeout = case TimeoutMs of + {some, Ms} when is_integer(Ms), Ms > 0 -> min(Ms, ?COMMAND_TIMEOUT_CAP_MS); + _ -> ?COMMAND_TIMEOUT_MS + end, + case executable(Command) of + false -> + Body = iolist_to_binary([Command, <<" is not on PATH">>]), + {126, Body}; + Exe -> + Port = open_port({spawn_executable, Exe}, + [{args, Args}, exit_status, binary, hide, {env, command_env()}]), + case Stdin of + {some, Data} -> Port ! {self(), {command, Data}}; + none -> ok + end, + case drain_capped(Port, [], 0, deadline(Timeout)) of + {ok, Result} -> Result; + timeout -> {126, <<"the command timed out">>} + end + end. + +%% ALBEDO_USAGE_CORE overrides resolution, so tests can point at a fake; the +%% shipped binary sits in priv/bin beside albedo-render, and PATH is the last +%% resort. +resolve() -> + case os:getenv("ALBEDO_USAGE_CORE") of + Override when Override =/= false, Override =/= "" -> + case filelib:is_file(Override) of + true -> Override; + false -> on_path(Override) + end; + _ -> + Bin = case code:priv_dir(albedo) of + {error, _} -> false; + Priv -> filename:join(Priv, "bin/usage") + end, + case is_file(Bin) of + true -> Bin; + false -> on_path("usage") + end + end. + +%% os:find_executable wants a list, and the commands arrive as binaries. +on_path(Name) -> + os:find_executable(unicode:characters_to_list(Name)). + +is_file(false) -> false; +is_file(Path) -> filelib:is_file(Path). + +%% A path in the command is used as given; a bare name is looked up on PATH. +executable(Command) -> + case binary:match(Command, <<"/">>) of + nomatch -> on_path(Command); + _ -> case filelib:is_file(Command) of + true -> Command; + false -> on_path(Command) + end + end. + +%% open_port's {env, ...} adds to the inherited environment rather than +%% replacing it, so every other variable is removed by name: the allowlist is +%% the whole environment a command sees. +command_env() -> + lists:map( + fun({K, V}) -> + case lists:member(K, ?ALLOWED_ENV) of + true -> {K, V}; + false -> {K, false} + end + end, + os:env()). + +method(Method) -> + binary_to_atom(string:lowercase(Method), utf8). + +%% httpc wants the body's content type in the request tuple rather than the +%% header list, so it moves there and the rest travel as headers. +content_type(Headers) -> + case [V || {K, V} <- Headers, is_content_type(K)] of + [Type | _] -> unicode:characters_to_list(Type); + [] -> "application/json" + end. + +without_content_type(Headers) -> + [{K, V} || {K, V} <- Headers, not is_content_type(K)]. + +is_content_type(K) -> + string:equal(string:lowercase(K), "content-type"). + +%% A receive's `after` is an idle timeout: a child that trickles a byte now +%% and then never trips it. The deadline is computed once and counts down +%% across every chunk. +deadline(Timeout) -> + erlang:monotonic_time(millisecond) + Timeout. + +remaining(Deadline) -> + max(0, Deadline - erlang:monotonic_time(millisecond)). + +drain(Port, Acc, Deadline) -> + receive + {Port, {data, Data}} -> drain(Port, [Data | Acc], Deadline); + {Port, {exit_status, 0}} -> + catch erlang:port_close(Port), + {ok, iolist_to_binary(lists:reverse(Acc))}; + {Port, {exit_status, Status}} -> + catch erlang:port_close(Port), + Text = integer_to_binary(Status), + {error, <<"the usage CLI exited with status ", Text/binary>>} + after remaining(Deadline) -> + kill(Port), + {error, <<"the usage CLI timed out">>} + end. + +%% Stdout up to the cap, then discarded, so a runaway command cannot fill +%% memory. Stderr is not merged: the body the feed sees stays stdout only. +drain_capped(Port, Acc, Size, Deadline) -> + receive + {Port, {data, Data}} -> + {Kept, Size2} = cap(Size, Data), + drain_capped(Port, [Kept | Acc], Size2, Deadline); + {Port, {exit_status, Status}} -> + catch erlang:port_close(Port), + {ok, {Status, iolist_to_binary(lists:reverse(Acc))}} + after remaining(Deadline) -> + kill(Port), + timeout + end. + +%% port_close closes the pipes; a child blocked on the network would keep +%% running. The child is ours, so the timeout path takes its pid and sends +%% KILL. A pid that already exited and was reaped makes kill itself fail, +%% which is the outcome we wanted anyway. +kill(Port) -> + Pid = case erlang:port_info(Port, os_pid) of + {os_pid, OsPid} when is_integer(OsPid) -> OsPid; + _ -> false + end, + catch erlang:port_close(Port), + case Pid of + false -> ok; + _ -> + Killer = open_port({spawn_executable, "/bin/kill"}, + [{args, ["-KILL", integer_to_list(Pid)]}, + exit_status, binary, hide]), + receive + {Killer, {exit_status, _}} -> ok + after 1000 -> + catch erlang:port_close(Killer) + end + end. + +cap(Size, _Data) when Size >= ?STDOUT_CAP -> + {<<>>, Size}; +cap(Size, Data) -> + Room = ?STDOUT_CAP - Size, + case byte_size(Data) =< Room of + true -> {Data, Size + byte_size(Data)}; + false -> {binary:part(Data, 0, Room), ?STDOUT_CAP} + end. diff --git a/src/albedo/harness/compaction.gleam b/src/albedo/harness/compaction.gleam index bab3b22b..8a7fe672 100644 --- a/src/albedo/harness/compaction.gleam +++ b/src/albedo/harness/compaction.gleam @@ -419,5 +419,18 @@ pub fn is_user(input: types.Input) -> Bool { } } +/// How many inputs two histories share at their end, oldest first: the tail a +/// projection kept verbatim, so what precedes it is what it replaced. +pub fn common_suffix(a: List(types.Input), b: List(types.Input)) -> Int { + suffix_length(list.reverse(a), list.reverse(b), 0) +} + +fn suffix_length(a: List(types.Input), b: List(types.Input), n: Int) -> Int { + case a, b { + [x, ..xs], [y, ..ys] if x == y -> suffix_length(xs, ys, n + 1) + _, _ -> n + } +} + @external(erlang, "albedo_compaction", "fingerprint") pub fn fingerprint(value: a) -> String diff --git a/src/albedo/harness/extension.gleam b/src/albedo/harness/extension.gleam index 4fdb2600..e0f7f02d 100644 --- a/src/albedo/harness/extension.gleam +++ b/src/albedo/harness/extension.gleam @@ -138,6 +138,9 @@ pub type Upstream { stream: fn(types.Request, fn(types.Event) -> types.Control) -> Result(types.Turn, types.Error), explain: fn(types.Error) -> Option(String), + /// The non-secret label of the account that served the most recent + /// request, when the provider has an account pool to name one from. + account: fn() -> Option(String), ) } diff --git a/src/albedo/harness/extensions/alibaba/extension.gleam b/src/albedo/harness/extensions/alibaba/extension.gleam index 9cecc9ab..886546de 100644 --- a/src/albedo/harness/extensions/alibaba/extension.gleam +++ b/src/albedo/harness/extensions/alibaba/extension.gleam @@ -41,9 +41,14 @@ fn resolve( ) let pool = pool(context.home, context.profile, context.session) use client <- result.map(pool.current()) - rotation.client_upstream(pool, client, fn(client, error) { - explain_key(context.home, context.profile, client, error) - }) + rotation.client_upstream( + pool, + client, + fn(client, error) { + explain_key(context.home, context.profile, client, error) + }, + fn(client) { rotation.key_label(client.api_key) }, + ) } /// Every Alibaba key albedo can see, as a rotation pool: the profile's own key diff --git a/src/albedo/harness/extensions/antigravity/albedo_antigravity.erl b/src/albedo/harness/extensions/antigravity/albedo_antigravity.erl index 7a3d5673..86cc4929 100644 --- a/src/albedo/harness/extensions/antigravity/albedo_antigravity.erl +++ b/src/albedo/harness/extensions/antigravity/albedo_antigravity.erl @@ -4,7 +4,7 @@ -include_lib("kernel/include/file.hrl"). -export([access/1, access/2, limited/3, encode/1, normalize_schema/1, session_number/1, uuid/1, call_id/0, now_ms/0, user_agent/1, discovered/1, refresh/1, reload/1, expire/2, - exchange/4, discover/3, account/1, client_id/0]). + exchange/4, discover/3, account/1, client_id/0, accounts/1]). -define(KEY, <<"google-antigravity">>). -define(CLIENT_ID, <<"MTA3MTAwNjA2MDU5MS10bWhzc2luMmgyMWxjcmUyMzV2dG9sb2poNGc0MDNlcC5hcHBzLmdvb2dsZXVzZXJjb250ZW50LmNvbQ==">>). @@ -65,6 +65,46 @@ first_usable([Credential | Rest], Path, Session) -> credentials(Data) -> albedo_credentials:oauth(Data, ?KEY). +%% Every Google account for the quota poller, each refreshed the way a request +%% refreshes it: `label` is the account's non-secret identity, the rest are +%% the provide-usage credential fields, projectId included. The refresh +%% token stays home: albedo refreshes, and the core never reads it. A refresh +%% that fails keeps the stored credential, so the poll records the auth failure as +%% a reading instead of silently dropping the account. +accounts(Home) -> + Path = albedo_credentials:auth_path(Home), + case albedo_credentials:read(Path) of + {ok, Data} -> + [entry(refreshed_account(Path, V)) || V <- credentials(Data), account_id(V) =/= <<>>]; + _ -> [] + end. + +refreshed_account(Path, Credential) -> + case usable(Credential) of + fresh -> Credential; + stale -> + case albedo_credentials:with_lock(Path, + fun() -> refresh_current(Path, identity(Credential)) end, + fun() -> {error, busy} end) of + {ok, Updated} -> Updated; + _ -> Credential + end; + invalid -> Credential + end. + +entry(Value) -> + maps:merge(#{<<"label">> => account_id(Value)}, fields(Value)). + +fields(Value) -> + maps:from_list([{K, field(Value, K)} + || K <- [<<"access">>, <<"email">>, <<"projectId">>]]). + +field(Value, Key) -> + case maps:get(Key, Value, <<>>) of + Text when is_binary(Text) -> Text; + _ -> <<>> + end. + usable(Credential) -> Valid = fun(Key) -> case maps:get(Key, Credential, <<>>) of <<_, _/binary>> -> true; _ -> false end end, Expires = maps:get(<<"expires">>, Credential, 0), diff --git a/src/albedo/harness/extensions/antigravity/extension.gleam b/src/albedo/harness/extensions/antigravity/extension.gleam index 28488656..d3ee0b00 100644 --- a/src/albedo/harness/extensions/antigravity/extension.gleam +++ b/src/albedo/harness/extensions/antigravity/extension.gleam @@ -139,6 +139,8 @@ pub fn upstream( }), access, fn(access, error) { explain(home, access, error) }, + // The Google account's email; blank credentials report no account. + fn(access) { access.email }, ) } diff --git a/src/albedo/harness/extensions/antigravity/stream.gleam b/src/albedo/harness/extensions/antigravity/stream.gleam index 77b5b55e..8f43eee7 100644 --- a/src/albedo/harness/extensions/antigravity/stream.gleam +++ b/src/albedo/harness/extensions/antigravity/stream.gleam @@ -357,7 +357,17 @@ fn chunk_decoder() -> decode.Decoder(Chunk) { None, decode.optional(decode.int), ) - decode.success(types.Usage(input, output + thoughts, cached, None)) + // Output keeps folding thoughts in; the thought count is also reported + // as reasoning so usage can tell the two apart. + decode.success(types.Usage( + input, + output + thoughts, + cached, + None, + None, + None, + Some(thoughts), + )) } let response = { use response_id <- decode.optional_field( diff --git a/src/albedo/harness/extensions/claude/albedo_claude_auth.erl b/src/albedo/harness/extensions/claude/albedo_claude_auth.erl index 80e738dd..5a2f6681 100644 --- a/src/albedo/harness/extensions/claude/albedo_claude_auth.erl +++ b/src/albedo/harness/extensions/claude/albedo_claude_auth.erl @@ -1,6 +1,6 @@ -module(albedo_claude_auth). --export([exchange/4, account/1, access/2, profile/2, expire/2, token/1]). +-export([exchange/4, account/1, access/2, profile/2, expire/2, token/1, accounts/1]). -define(KEY, <<"anthropic">>). -define(CLIENT_ID, <<"9d1c250a-e61b-44d9-88ed-5944d1962f5e">>). @@ -103,6 +103,47 @@ session_uuid(Session) -> credentials(Data) -> albedo_credentials:oauth(Data, ?KEY). +%% Every Claude account for the quota poller, each with a fresh access token: +%% `label` is the account's non-secret identity, the rest are the +%% provide-usage credential fields. The refresh token stays home: albedo +%% refreshes, and the core never reads it. A refresh that fails keeps the stored +%% credential, so the poll records the auth failure as a reading instead of +%% silently dropping the account. +accounts(Home0) -> + Path = albedo_credentials:auth_path(Home0), + case albedo_credentials:read(Path) of + {ok, Data} -> + [entry(refreshed(Path, V)) || V <- credentials(Data), identity(V) =/= <<>>]; + _ -> [] + end. + +refreshed(Path, Value) -> + Expires = maps:get(<<"expires">>, Value, 0), + case is_integer(Expires) andalso + Expires > erlang:system_time(millisecond) + ?REFRESH_SKEW_MS of + true -> Value; + false -> + case albedo_credentials:with_lock(Path, + fun() -> refresh(Path, identity(Value)) end, + fun() -> {error, busy} end) of + {ok, Updated} -> Updated; + _ -> Value + end + end. + +entry(Value) -> + maps:merge(#{<<"label">> => identity(Value)}, fields(Value)). + +fields(Value) -> + maps:from_list([{K, field(Value, K)} + || K <- [<<"access">>, <<"accountId">>, <<"email">>]]). + +field(Value, Key) -> + case maps:get(Key, Value, <<>>) of + Text when is_binary(Text) -> Text; + _ -> <<>> + end. + first_access([], _, _) -> {error, <<"Claude is not authenticated; run /login and add a Claude account">>}; first_access([C | Rest], Path, Session) -> Now = erlang:system_time(millisecond), diff --git a/src/albedo/harness/extensions/claude/extension.gleam b/src/albedo/harness/extensions/claude/extension.gleam index e9a38fcc..2895b879 100644 --- a/src/albedo/harness/extensions/claude/extension.gleam +++ b/src/albedo/harness/extensions/claude/extension.gleam @@ -125,6 +125,8 @@ fn resolve( _ -> None } }, + // The OAuth profile's account UUID, not a credential. + fn() { Some(account) }, )) } diff --git a/src/albedo/harness/extensions/claude/stream.gleam b/src/albedo/harness/extensions/claude/stream.gleam index 4235e8d4..f586f33e 100644 --- a/src/albedo/harness/extensions/claude/stream.gleam +++ b/src/albedo/harness/extensions/claude/stream.gleam @@ -91,7 +91,7 @@ fn apply( |> option.from_result let usage = case state.usage, output { Some(old), Some(n) -> Some(types.Usage(..old, output_tokens: n)) - None, Some(n) -> Some(types.Usage(0, n, None, None)) + None, Some(n) -> Some(types.Usage(0, n, None, None, None, None, None)) _, None -> state.usage } emit( @@ -297,6 +297,12 @@ fn usage_decoder() -> decode.Decoder(types.Usage) { None, decode.optional(decode.int), ) + // The same writes split by how long the entry lives. + use writes <- decode.optional_field( + "cache_creation", + None, + decode.optional(cache_creation_decoder()), + ) // Anthropic reports input_tokens without cached reads or writes; the // harness counts the whole input context so cached is a subset of input. decode.success(types.Usage( @@ -304,9 +310,30 @@ fn usage_decoder() -> decode.Decoder(types.Usage) { output, cache, creation, + option.then(writes, fn(split) { split.write_5m }), + option.then(writes, fn(split) { split.write_1h }), + None, )) } +type CacheWrites { + CacheWrites(write_5m: Option(Int), write_1h: Option(Int)) +} + +fn cache_creation_decoder() -> decode.Decoder(CacheWrites) { + use write_5m <- decode.optional_field( + "ephemeral_5m_input_tokens", + None, + decode.optional(decode.int), + ) + use write_1h <- decode.optional_field( + "ephemeral_1h_input_tokens", + None, + decode.optional(decode.int), + ) + decode.success(CacheWrites(write_5m, write_1h)) +} + fn original_name(name: String, tools: List(types.Tool)) -> String { tools |> list.find(fn(tool) { diff --git a/src/albedo/harness/extensions/codex/albedo_openai_auth.erl b/src/albedo/harness/extensions/codex/albedo_openai_auth.erl index 9c99ec9b..1d08ce8e 100644 --- a/src/albedo/harness/extensions/codex/albedo_openai_auth.erl +++ b/src/albedo/harness/extensions/codex/albedo_openai_auth.erl @@ -1,7 +1,7 @@ -module(albedo_openai_auth). %% Codex sign-in, multi-account credential selection, and refresh. --export([codex_access/2, codex_revoke/2, codex_limited/3, codex_exchange/3, codex_account/1]). +-export([codex_access/2, codex_revoke/2, codex_limited/3, codex_exchange/3, codex_account/1, accounts/1]). -define(SCOPE, <<"codex">>). -define(STORE, <<"openai-codex">>). @@ -167,6 +167,51 @@ usable(_) -> error. credentials(Data) when is_map(Data) -> albedo_credentials:oauth(Data, ?STORE); credentials(_) -> []. +%% Every stored ChatGPT account for the quota poller, each refreshed the way a +%% request refreshes it: `label` is the account's non-secret identity, the rest +%% are the provide-usage credential fields. The refresh token stays home: +%% albedo refreshes, and the core never reads it. A refresh that fails keeps the +%% stored credential, so the poll records the auth failure as a reading +%% instead of silently dropping the account. +accounts(Home0) -> + Path = albedo_credentials:auth_path(text(Home0)), + case albedo_credentials:read(Path) of + {ok, Data} -> + [entry(V) || V <- [refreshed(Path, C) || C <- credentials(Data)], label(V) =/= <<>>]; + _ -> [] + end. + +refreshed(Path, Credential) -> + case usable(Credential) of + {ok, Ready} -> Ready; + refresh -> + case refresh_locked(Path, identity(Credential)) of + {ok, Updated} -> Updated; + _ -> Credential + end; + error -> Credential + end. + +label(Credential) -> + first([email(Credential), account_id(Credential), short_hash(maps:get(<<"refresh">>, Credential, <<>>))]). + +entry(Value) -> + maps:merge(#{<<"label">> => label(Value)}, fields(Value)). + +fields(Value) -> + maps:from_list([{K, field(Value, K)} + || K <- [<<"access">>, <<"accountId">>, <<"email">>]]). + +field(Value, Key) -> + case maps:get(Key, Value, <<>>) of + Text when is_binary(Text) -> Text; + _ -> <<>> + end. + +short_hash(Token) when is_binary(Token), byte_size(Token) > 0 -> + binary:encode_hex(binary:part(crypto:hash(sha256, Token), 0, 8), lowercase); +short_hash(_) -> <<>>. + selected(Credential) -> maps:get(<<"selected">>, Credential, false) =:= true. %% Records a Codex limit response against the account that received it, so diff --git a/src/albedo/harness/extensions/codex/extension.gleam b/src/albedo/harness/extensions/codex/extension.gleam index 7ab71051..d743be48 100644 --- a/src/albedo/harness/extensions/codex/extension.gleam +++ b/src/albedo/harness/extensions/codex/extension.gleam @@ -105,9 +105,18 @@ fn resolve( pool(context.home, context.session, openai_api.stream), client, fn(client, error) { account_failure(context.home, client, error) }, + account_label, ) } +/// The ChatGPT account a client speaks for, or a hash of its key. +fn account_label(client: types.Client) -> String { + case client.policy { + types.Codex(account_id, _) -> account_id + _ -> rotation.key_label(client.api_key) + } +} + /// The session's current account as a client. Its model list refreshes in /// the background, so a new model shows up without a picker visit. fn connect(home: String, session: String) -> Result(types.Client, String) { diff --git a/src/albedo/harness/extensions/openai/extension.gleam b/src/albedo/harness/extensions/openai/extension.gleam index 72b2fbb9..44280736 100644 --- a/src/albedo/harness/extensions/openai/extension.gleam +++ b/src/albedo/harness/extensions/openai/extension.gleam @@ -59,6 +59,7 @@ pub fn upstream( client.protocol, fn(request, on_event) { openai_api.stream(client, request, on_event) }, explain, + fn() { None }, ) } diff --git a/src/albedo/harness/extensions/proxy/chat.gleam b/src/albedo/harness/extensions/proxy/chat.gleam index c271f7ef..57a66ec2 100644 --- a/src/albedo/harness/extensions/proxy/chat.gleam +++ b/src/albedo/harness/extensions/proxy/chat.gleam @@ -587,17 +587,25 @@ fn finish(finish: types.Finish) -> String { } fn usage(usage: Option(types.Usage)) -> Json { - let types.Usage(input, output, cached, ..) = - option.unwrap(usage, types.Usage(0, 0, None, None)) - json.object([ - #("prompt_tokens", json.int(input)), - #("completion_tokens", json.int(output)), - #("total_tokens", json.int(input + output)), - #( - "prompt_tokens_details", - json.object([#("cached_tokens", json.int(option.unwrap(cached, 0)))]), - ), - ]) + let types.Usage(input, output, cached, reasoning, ..) = + option.unwrap(usage, types.Usage(0, 0, None, None, None, None, None)) + json.object( + [ + #("prompt_tokens", json.int(input)), + #("completion_tokens", json.int(output)), + #("total_tokens", json.int(input + output)), + #( + "prompt_tokens_details", + json.object([#("cached_tokens", json.int(option.unwrap(cached, 0)))]), + ), + ] + |> when(option.is_some(reasoning), #( + "completion_tokens_details", + json.object([ + #("reasoning_tokens", json.int(option.unwrap(reasoning, 0))), + ]), + )), + ) } fn when( diff --git a/src/albedo/harness/loop.gleam b/src/albedo/harness/loop.gleam index 528196a3..bc15be0b 100644 --- a/src/albedo/harness/loop.gleam +++ b/src/albedo/harness/loop.gleam @@ -2,12 +2,14 @@ import albedo/daemon/conversation import albedo/daemon/events as view +import albedo/daemon/ledger import albedo/daemon/usage import albedo/harness/compaction import albedo/harness/extension import albedo/harness/runtime import albedo/openai_api/types import gleam/int +import gleam/io import gleam/json import gleam/list import gleam/option.{type Option, None, Some} @@ -24,9 +26,11 @@ pub type Loop { upstream: extension.Upstream, publish: fn(String) -> Bool, /// Commits inputs at a stage; a model response passes how long the - /// model thought before it. + /// model thought before it. Answers the shared daemon time and the seq + /// of the commit's first assistant row, which the call's ledger row + /// links to. commit: fn(List(types.Input), conversation.Stage, Option(Int)) -> - Result(Int, String), + Result(#(Int, Option(Int)), String), record_context: fn(types.Request, Option(compaction.Observation), Bool) -> Nil, record_usage: fn(usage.Metadata) -> Result(Nil, String), @@ -34,6 +38,10 @@ pub type Loop { /// Reports the pin's compaction baseline (`Some`) or that compaction /// replaced the history the pinned prompt was cached with (`None`). report_pin: fn(Option(Int)) -> Nil, + /// The session whose provider calls the request ledger records. + session: String, + /// The saved profile those calls went through. + profile: String, ) } @@ -60,30 +68,37 @@ pub fn run( case state.pin { Unpinned -> state.report_pin(Some( - list.length(original) - common_suffix(original, history), + list.length(original) - compaction.common_suffix(original, history), )) Pinned(..) -> Nil } let current_instructions = request_instructions(state) let request = request(state, current_instructions, history) state.record_context(request, prepared.observation, prepared.compacted) - use turn <- result.try( - stream_with_retries(state.upstream, request, state.publish, fn(event) { - case view.stream_event(id, step, event) { - Some(serialized) -> - case state.publish(serialized) { - True -> types.Continue - False -> types.Stop - } - None -> types.Continue - } - }) + use #(row, turn) <- result.try( + call( + state, + ledger.Turn, + request, + request_prefix(request, history, original, prepared.observation), + state.publish, + fn(event) { + case view.stream_event(id, step, event) { + Some(serialized) -> + case state.publish(serialized) { + True -> types.Continue + False -> types.Stop + } + None -> types.Continue + } + }, + ) |> result.map_error(describe(state.upstream, _)), ) let completed_usage = usage.from_completion(state.model, turn.usage, usage.now()) let replay = list.map(turn.output, types.Replay) - use timestamp <- result.try(state.commit( + use #(timestamp, seq) <- result.try(state.commit( replay, case turn.tool_calls { [] -> conversation.Idle @@ -91,6 +106,7 @@ pub fn run( }, turn.thought_ms, )) + attach(state, row, seq) list.each(replay, fn(input) { let _ = list.each(view.assistant_message(input, Some(timestamp)), state.publish) @@ -128,7 +144,7 @@ pub fn compact(state: Loop, inputs: List(types.Input)) -> Result(Nil, String) { let history = prepared.inputs // The projection keeps a verbatim tail, so everything before the shared // suffix is what the strategy's replacement stands in for. - let suffix = common_suffix(original, history) + let suffix = compaction.common_suffix(original, history) let evicted = list.length(inputs) - suffix let _ = case evicted > 0 { True -> { @@ -236,17 +252,6 @@ fn unshift( }) } -fn common_suffix(a: List(types.Input), b: List(types.Input)) -> Int { - suffix_length(list.reverse(a), list.reverse(b), 0) -} - -fn suffix_length(a: List(types.Input), b: List(types.Input), n: Int) -> Int { - case a, b { - [x, ..xs], [y, ..ys] if x == y -> suffix_length(xs, ys, n + 1) - _, _ -> n - } -} - fn display_text(items: List(types.Input)) -> String { items |> list.map(fn(input) { @@ -277,22 +282,117 @@ fn bounded(text: String, limit: Int, note: String) -> String { } } -/// Reissue transient transport and gateway failures. A failed attempt has no committed output -/// or tool effects; discard its live previews before forwarding the next attempt. -fn stream_with_retries( - upstream: extension.Upstream, +/// Streams one provider call, writing a request-ledger row per attempt: what +/// it cost, how it ended, and the prefix identity it went out with. The row +/// id of the attempt that succeeded is answered so its transcript seq can be +/// attached once committed. A row that cannot be written is logged and +/// skipped, never a failed turn. +fn call( + state: Loop, + kind: ledger.Kind, request: types.Request, + prefix: ledger.Prefix, publish: fn(String) -> Bool, on_event: fn(types.Event) -> types.Control, -) -> Result(types.Turn, types.Error) { - retry_stream(fn() { upstream.stream(request, on_event) }, publish, 1) +) -> Result(#(Option(Int), types.Turn), types.Error) { + retry_stream( + fn() { + let started = ledger.now() + let outcome = state.upstream.stream(request, on_event) + let usage = case outcome { + Ok(turn) -> turn.usage + Error(_) -> None + } + // The account label is read after the attempt: rotation records which + // account served while the request streams. + let row = + ledger.record( + runtime.ledger(state.host), + ledger.Call( + state.session, + kind, + state.profile, + provider_label(state), + state.upstream.account(), + request.model, + started, + ledger.now(), + ledger.outcome(outcome), + usage, + prefix, + ), + ) + let row = case row { + Ok(id) -> Some(id) + Error(error) -> { + io.println_error( + "request ledger write failed for session " + <> state.session + <> ": " + <> error, + ) + None + } + } + case outcome { + Ok(turn) -> Ok(#(row, turn)) + Error(error) -> Error(error) + } + }, + publish, + 1, + ) +} + +/// Attaches the transcript row a recorded call produced — the first +/// assistant row of its response commit. A commit with no assistant row, or +/// a failed update, is logged and skipped, never a failed turn. +fn attach(state: Loop, row: Option(Int), seq: Option(Int)) -> Nil { + case row, seq { + Some(id), Some(seq) -> + case ledger.attach(runtime.ledger(state.host), state.session, id, seq) { + Ok(_) -> Nil + Error(error) -> + io.println_error( + "request ledger seq attach failed for session " + <> state.session + <> ": " + <> error, + ) + } + _, _ -> Nil + } +} + +/// How a call reached its provider: the protocol over the endpoint. +fn provider_label(state: Loop) -> String { + types.protocol_name(state.upstream.protocol) <> ":" <> state.upstream.endpoint +} + +/// The request's prefix identity: its head, and the projection that stands in +/// for the history compaction replaced. +fn request_prefix( + request: types.Request, + history: List(types.Input), + original: List(types.Input), + observation: Option(compaction.Observation), +) -> ledger.Prefix { + ledger.prefix( + option.unwrap(request.instructions, ""), + request.tools, + history, + original, + option.map(observation, fn(observation) { observation.strategy }), + ) } +/// Reissue transient transport and gateway failures. A failed attempt has no committed output +/// or tool effects; discard its live previews before forwarding the next attempt. fn retry_stream( - run: fn() -> Result(types.Turn, types.Error), + run: fn() -> Result(a, types.Error), publish: fn(String) -> Bool, attempt: Int, -) -> Result(types.Turn, types.Error) { +) -> Result(a, types.Error) { case run() { Error(error) -> case attempt < 3 && retryable(error) { @@ -369,7 +469,8 @@ fn settle_pin( Pinned(prompt, baseline) -> { // Strategies may rebuild recap text every request, but the count of // original inputs they stand in for only grows when compaction runs. - let head = list.length(original) - common_suffix(original, history) + let head = + list.length(original) - compaction.common_suffix(original, history) case baseline { None -> { state.report_pin(Some(head)) @@ -426,10 +527,16 @@ fn summarize( Some(max_output_tokens), types.defaults, ) - use turn <- result.try( - stream_with_retries(state.upstream, summary_request, state.publish, fn(_) { - types.Continue - }) + use #(_, turn) <- result.try( + call( + state, + ledger.Summarizer, + summary_request, + // The summarizer's history is exactly what it says; nothing replaced. + ledger.direct_prefix(summary_instructions, [], [types.User(prompt)]), + state.publish, + fn(_) { types.Continue }, + ) |> result.map_error(fn(error) { "summarizer provider request failed: " <> describe(state.upstream, error) }), diff --git a/src/albedo/harness/rotation.gleam b/src/albedo/harness/rotation.gleam index f6a97cd2..ab818f89 100644 --- a/src/albedo/harness/rotation.gleam +++ b/src/albedo/harness/rotation.gleam @@ -10,11 +10,14 @@ import albedo/harness/extension import albedo/openai_api/types +import gleam/bit_array +import gleam/crypto import gleam/dynamic/decode import gleam/json import gleam/list import gleam/option.{type Option, None, Some} import gleam/result +import gleam/string /// What a 429 said about the account that received it. pub type Limit { @@ -113,12 +116,25 @@ pub fn same_client(a: types.Client, b: types.Client) -> Bool { } /// An upstream for an OpenAI-compatible client rotating through pool. +/// `label` names an account without naming its credential. pub fn client_upstream( pool: Pool(types.Client), first: types.Client, explain: fn(types.Client, types.Error) -> Option(String), + label: fn(types.Client) -> String, ) -> extension.Upstream { - upstream(first.base_url, first.protocol, pool, first, explain) + upstream(first.base_url, first.protocol, pool, first, explain, label) +} + +/// A non-secret account label for a key-authenticated client: the first +/// eight hex characters of the key's SHA-256, enough to tell siblings apart. +pub fn key_label(key: String) -> String { + key + |> bit_array.from_string + |> crypto.hash(crypto.Sha256, _) + |> bit_array.base16_encode + |> string.slice(0, 8) + |> string.lowercase } pub type Pool(account) { @@ -151,13 +167,15 @@ pub fn budget() -> Budget { const waits = [4000, 8000, 15_000, 30_000] /// An upstream that rotates through `pool` starting from `first`. `explain` -/// is given the account the failing attempt went to, not necessarily `first`. +/// is given the account the failing attempt went to, not necessarily `first`; +/// `label` names the account that served, without naming its credential. pub fn upstream( endpoint: String, protocol: types.Protocol, pool: Pool(account), first: account, explain: fn(account, types.Error) -> Option(String), + label: fn(account) -> String, ) -> extension.Upstream { let slot = new_slot() extension.Upstream( @@ -170,6 +188,12 @@ pub fn upstream( }) }, fn(error) { explain(last_served(slot, first), error) }, + fn() { + case string.trim(label(last_served(slot, first))) { + "" -> None + account -> Some(account) + } + }, ) } diff --git a/src/albedo/harness/usage_feed.gleam b/src/albedo/harness/usage_feed.gleam new file mode 100644 index 00000000..324a5aee --- /dev/null +++ b/src/albedo/harness/usage_feed.gleam @@ -0,0 +1,291 @@ +//// provide-usage as a stateless one-shot CLI, driven from Erlang. The core +//// (usage-core, Zig) decides which requests to send and nothing else: each +//// round is one fresh `usage advance -` process fed the whole history on +//// stdin, so a credential never reaches argv, and every byte of I/O is +//// albedo's - `http` requests through albedo's HTTP client, `command` +//// requests through a PATH, HOME, locale, timezone and proxy environment +//// with a timeout and a stdout cap. The round loop lives here and the I/O +//// lives in albedo_usage_core, so neither half can drift from the other. +//// +//// `Error` from `fetch` means the driver itself failed: the binary is +//// missing, its output is unparseable, or the feed did not finish in +//// `max_rounds` rounds. A provider failure (a 4xx, an unknown provider, a +//// flaky CLI) is data: it comes back as `Ok(Report(error: Some(..)))`. + +import gleam/dynamic/decode +import gleam/int +import gleam/json +import gleam/list +import gleam/option.{type Option, None, Some} +import gleam/result +import gleam/string + +pub type Limit { + Limit( + id: String, + label: String, + used_percent: Option(Float), + window_label: Option(String), + window_seconds: Option(Int), + resets_at: Option(Int), + scope: Option(String), + status: String, + ) +} + +pub type Report { + Report( + account_id: Option(String), + email: Option(String), + plan: Option(String), + limits: List(Limit), + error: Option(String), + ) +} + +/// A feed that needs more round trips than this is stuck, not slow: the +/// deepest feed today (hyper) is three rounds. +const max_rounds = 6 + +@external(erlang, "albedo_usage_core", "advance") +fn advance(envelope: String) -> Result(String, String) + +@external(erlang, "albedo_usage_core", "http_request") +fn http_request( + method: String, + url: String, + headers: List(#(String, String)), + body: Option(String), +) -> #(Int, String) + +@external(erlang, "albedo_usage_core", "run_command") +fn run_command( + command: String, + args: List(String), + stdin: Option(String), + timeout_ms: Option(Int), +) -> #(Int, String) + +/// Runs a provider's feed to completion. Each round is one `usage advance -` +/// process fed one envelope line on stdin (never argv: it carries the +/// credential); `http` requests go through albedo's HTTP client, `command` +/// requests (alibaba's `bl`) run with only PATH, HOME, locale, timezone and +/// proxy variables, a timeout, and a stdout cap. +pub fn fetch( + provider: String, + credential: json.Json, + now_ms: Int, +) -> Result(Report, String) { + loop(provider, credential, now_ms, []) +} + +/// `rounds` is the whole history so far, one entry per answered round, so +/// every `usage advance -` process replays the feed from scratch and the core +/// stays stateless: the host keeps the state, the core keeps the rules. +fn loop( + provider: String, + credential: json.Json, + now_ms: Int, + rounds: List(json.Json), +) -> Result(Report, String) { + case list.length(rounds) >= max_rounds { + True -> + Error( + "usage feed for " + <> provider + <> " did not finish in " + <> int.to_string(max_rounds) + <> " rounds", + ) + False -> { + let envelope = + json.object([ + #("provider", json.string(provider)), + #("credential", credential), + #("responses", json.array(rounds, fn(round) { round })), + #("nowMs", json.int(now_ms)), + ]) + |> json.to_string + use output <- result.try(advance(envelope)) + use step <- result.try(step(output)) + case step { + Reported(report) -> Ok(report) + // The core answers an unknown provider (or its own parse failure) + // with an error step, which is still the provider's answer, not a + // driver failure. + Failed(message) -> Ok(Report(None, None, None, [], Some(message))) + Requests(requests) -> { + let round = + json.array(list.map(requests, answer), fn(response) { response }) + loop(provider, credential, now_ms, list.append(rounds, [round])) + } + } + } + } +} + +fn answer(request: Request) -> json.Json { + // A transport failure or a command that could not run is already folded + // into the status by the Erlang half, so both request kinds answer the + // same way: one `{"status", "body"}` object per request. + let #(status, body) = case request { + HttpRequest(method, url, headers, body) -> + http_request(method, url, headers, body) + CommandRequest(command, args, stdin, timeout_ms) -> + run_command(command, args, stdin, timeout_ms) + } + json.object([ + #("status", json.int(status)), + #("body", json.string(body)), + ]) +} + +/// One step is printed per replayed round; the last line is the answer. +fn step(output: String) -> Result(Step, String) { + let last = + output + |> string.split("\n") + |> list.map(string.trim) + |> list.filter(fn(line) { line != "" }) + |> list.last + case last { + Error(Nil) -> Error("the usage CLI printed no step") + Ok(last) -> + json.parse(last, step_decoder()) + |> result.replace_error("the usage CLI printed an unparseable step") + } +} + +type Request { + HttpRequest( + method: String, + url: String, + headers: List(#(String, String)), + body: Option(String), + ) + CommandRequest( + command: String, + args: List(String), + stdin: Option(String), + timeout_ms: Option(Int), + ) +} + +type Step { + Requests(List(Request)) + Reported(Report) + Failed(String) +} + +fn step_decoder() -> decode.Decoder(Step) { + decode.one_of( + decode.field("report", report_decoder(), fn(report) { + decode.success(Reported(report)) + }), + or: [ + decode.field("requests", decode.list(request_decoder()), fn(requests) { + decode.success(Requests(requests)) + }), + decode.field("error", decode.string, fn(message) { + decode.success(Failed(message)) + }), + ], + ) +} + +fn request_decoder() -> decode.Decoder(Request) { + use kind <- decode.field("kind", decode.string) + case kind { + "command" -> command_decoder() + _ -> http_decoder() + } +} + +fn command_decoder() -> decode.Decoder(Request) { + use command <- decode.field("command", decode.string) + use args <- decode.optional_field("args", [], decode.list(decode.string)) + use stdin <- decode.optional_field("stdin", None, optional_string()) + use timeout_ms <- decode.optional_field("timeoutMs", None, optional_int()) + decode.success(CommandRequest(command, args, stdin, timeout_ms)) +} + +fn http_decoder() -> decode.Decoder(Request) { + use method <- decode.optional_field("method", "GET", decode.string) + use url <- decode.field("url", decode.string) + use headers <- decode.optional_field("headers", [], decode.list(pair())) + use body <- decode.optional_field("body", None, optional_string()) + decode.success(HttpRequest(method, url, headers, body)) +} + +fn report_decoder() -> decode.Decoder(Report) { + use account_id <- decode.optional_field("accountId", None, optional_string()) + use email <- decode.optional_field("email", None, optional_string()) + use plan <- decode.optional_field("planType", None, optional_string()) + use limits <- decode.optional_field( + "limits", + [], + decode.list(limit_decoder()), + ) + use error <- decode.optional_field("error", None, optional_string()) + decode.success(Report(account_id, email, plan, limits, error)) +} + +fn limit_decoder() -> decode.Decoder(Limit) { + use id <- decode.field("id", decode.string) + use label <- decode.field("label", decode.string) + use used_percent <- decode.optional_field( + "usedPercent", + None, + optional_float(), + ) + use window_label <- decode.optional_field( + "windowLabel", + None, + optional_string(), + ) + use window_seconds <- decode.optional_field( + "windowSeconds", + None, + optional_int(), + ) + use resets_at <- decode.optional_field("resetsAt", None, optional_int()) + use scope <- decode.optional_field("scope", None, optional_string()) + use status <- decode.field("status", decode.string) + decode.success(Limit( + id, + label, + used_percent, + window_label, + window_seconds, + resets_at, + scope, + status, + )) +} + +fn optional_string() -> decode.Decoder(Option(String)) { + decode.optional(decode.string) +} + +fn optional_int() -> decode.Decoder(Option(Int)) { + decode.optional(decode.int) +} + +/// The core clamps percentages but emits whole numbers without a decimal +/// point, which arrive as integers; a percentage is a float either way. +fn optional_float() -> decode.Decoder(Option(Float)) { + decode.optional( + decode.one_of(decode.float, or: [ + decode.map(decode.int, int.to_float), + ]), + ) +} + +fn pair() -> decode.Decoder(#(String, String)) { + decode.then(decode.list(decode.string), fn(parts) { + case parts { + [name, value] -> decode.success(#(name, value)) + _ -> decode.failure(#("", ""), "a two-element header pair") + } + }) +} diff --git a/src/albedo/openai_api/chat.gleam b/src/albedo/openai_api/chat.gleam index 92858cef..95f3403b 100644 --- a/src/albedo/openai_api/chat.gleam +++ b/src/albedo/openai_api/chat.gleam @@ -102,6 +102,7 @@ fn chunk_decoder() -> decode.Decoder(Chunk) { "prompt_tokens", "completion_tokens", "prompt_tokens_details", + "completion_tokens_details", )), ) decode.success(Chunk(id, choices, usage)) diff --git a/src/albedo/openai_api/responses.gleam b/src/albedo/openai_api/responses.gleam index 93638fd6..b204d555 100644 --- a/src/albedo/openai_api/responses.gleam +++ b/src/albedo/openai_api/responses.gleam @@ -316,6 +316,7 @@ fn response_decoder(incomplete: Bool) -> decode.Decoder(Response) { "input_tokens", "output_tokens", "input_tokens_details", + "output_tokens_details", )), ) use details <- decode.optional_field( diff --git a/src/albedo/openai_api/types.gleam b/src/albedo/openai_api/types.gleam index fe37c2b3..d5943455 100644 --- a/src/albedo/openai_api/types.gleam +++ b/src/albedo/openai_api/types.gleam @@ -263,6 +263,13 @@ pub type Usage { output_tokens: Int, cached_input_tokens: Option(Int), cache_creation_tokens: Option(Int), + /// Cache writes the provider splits by TTL, when it reports the split. + cache_write_5m_tokens: Option(Int), + cache_write_1h_tokens: Option(Int), + /// Tokens the provider names as reasoning, separate from its output + /// count when it does. Antigravity folds thoughts into output and also + /// reports the thought count here. + reasoning_tokens: Option(Int), ) } @@ -272,6 +279,7 @@ pub fn usage_decoder( input_tokens: String, output_tokens: String, input_tokens_details: String, + output_tokens_details: String, ) -> decode.Decoder(Usage) { use input <- decode.field(input_tokens, decode.int) use output <- decode.field(output_tokens, decode.int) @@ -280,7 +288,20 @@ pub fn usage_decoder( None, decode.optional(cached_tokens_decoder()), ) - decode.success(Usage(input, output, option.flatten(details), None)) + use reasoning <- decode.optional_field( + output_tokens_details, + None, + decode.optional(reasoning_tokens_decoder()), + ) + decode.success(Usage( + input, + output, + option.flatten(details), + None, + None, + None, + option.flatten(reasoning), + )) } fn cached_tokens_decoder() -> decode.Decoder(Option(Int)) { @@ -292,6 +313,15 @@ fn cached_tokens_decoder() -> decode.Decoder(Option(Int)) { ) } +fn reasoning_tokens_decoder() -> decode.Decoder(Option(Int)) { + decode.optional_field( + "reasoning_tokens", + None, + decode.optional(decode.int), + decode.success, + ) +} + pub type Finish { Complete ToolCalls diff --git a/test/daemon/session_run_test.gleam b/test/daemon/session_run_test.gleam index 52514cd7..ddb61bfa 100644 --- a/test/daemon/session_run_test.gleam +++ b/test/daemon/session_run_test.gleam @@ -80,7 +80,7 @@ pub fn a_reply_that_races_the_callee_exit_still_arrives_test() { /// A stand-in for the session actor's protocol. type Owner { Publish(reply: Subject(Bool)) - Commit(reply: Subject(Result(Int, String))) + Commit(reply: Subject(Result(#(Int, option.Option(Int)), String))) Commits(reply: Subject(Int)) Unused } @@ -151,7 +151,7 @@ pub fn publish_stops_the_stream_when_the_owner_dies_test() { |> should.be_false } -fn commit(owner: Subject(Owner)) -> Result(Int, String) { +fn commit(owner: Subject(Owner)) -> Result(#(Int, option.Option(Int)), String) { session_run.commit_fn(owner, "run", messages(), waiting: 20)( [], conversation.Tool, diff --git a/test/e2e/ledger_test.py b/test/e2e/ledger_test.py new file mode 100644 index 00000000..46b617b8 --- /dev/null +++ b/test/e2e/ledger_test.py @@ -0,0 +1,180 @@ +"""The request ledger: one durable row per provider call, recording what cache +and quota analysis cannot recover from the transcript. + +Catches: usage counts that never reach the ledger (cache reads, TTL-split +writes, reasoning), the account/provider identity of each call, the prefix +identity that distinguishes a warm append from a compaction rewrite, an HTTP +failure recorded as a quota reading, and the seq link to the transcript. None +of this is observable from events or history alone. +""" +import json +import unittest + +from harness import Albedo, Provider, Reply, text + +FIRST = {"input_tokens": 100, "output_tokens": 20, + "input_tokens_details": {"cached_tokens": 40}, + "output_tokens_details": {"reasoning_tokens": 7}} +TOOL = {"input_tokens": 200, "output_tokens": 30, + "input_tokens_details": {"cached_tokens": 80}, + "output_tokens_details": {"reasoning_tokens": 11}} +AFTER = {"input_tokens": 300, "output_tokens": 40, + "input_tokens_details": {"cached_tokens": 120}, + "output_tokens_details": {"reasoning_tokens": 3}} +THIRD = {"input_tokens": 400, "output_tokens": 50, + "input_tokens_details": {"cached_tokens": 160}, + "output_tokens_details": {"reasoning_tokens": 5}} +SUMMARY = {"input_tokens": 60, "output_tokens": 10} +POST = {"input_tokens": 70, "output_tokens": 15, + "input_tokens_details": {"cached_tokens": 30}, + "output_tokens_details": {"reasoning_tokens": 2}} + + +def reply(request): + inputs = request["input"] + # The tool result arrives while the newest user message is still the one + # that asked for the tool, so it is told apart by input shape. + if inputs[-1].get("type") == "function_call_output": + return text("tool done", usage=AFTER) + user = next((str(item.get("content", "")) for item in reversed(inputs) + if item.get("role") == "user"), "") + if "" in user: + return text("older conversation summary", usage=SUMMARY) + if user == "use the tool": + return Reply("python", "2 + 2", usage=TOOL) + if user == "hit the limit": + return Reply("error", "limit exceeded", status=429) + if user == "after compaction": + return text("post compaction reply", usage=POST) + if user == "third": + return text("third reply", usage=THIRD) + return text("first reply", usage=FIRST) + + +class LedgerTest(unittest.TestCase): + def rows(self, app, session, query=""): + with app.api(f"/sessions/{session}/ledger{query}") as response: + return json.load(response) + + def tree(self, app, session): + """Every transcript row through the tree route, oldest first.""" + items, after = [], 0 + while True: + with app.api(f"/sessions/{session}/tree?after={after}&limit=100") \ + as response: + page = json.load(response) + items.extend(page["items"]) + if not page["hasMore"]: + return items + after = page["nextCursor"] + + def test_rows_record_usage_identity_prefix_and_projection(self): + provider = Provider(reply) + try: + with Albedo(provider, protocol="responses") as app: + session = app.session() + for prompt in ("first", "use the tool", "third", "hit the limit"): + app.prompt(session, prompt).close() + app.idle(session) + with app.api(f"/sessions/{session}/commands", + {"name": "/compact"}) as response: + self.assertTrue(json.load(response)["result"]["started"]) + app.idle(session) + app.prompt(session, "after compaction").close() + app.idle(session) + + page = self.rows(app, session) + rows = page["rows"] + self.assertEqual([row["kind"] for row in rows], + ["turn", "turn", "turn", "turn", "turn", + "summarizer", "turn"]) + # Every call carries its own reported counts, split as sent. + self.assertEqual([(row["inputTokens"], row["outputTokens"], + row["cachedInputTokens"], + row["reasoningTokens"]) for row in rows], + [(100, 20, 40, 7), (200, 30, 80, 11), + (300, 40, 120, 3), (400, 50, 160, 5), + (None, None, None, None), + (60, 10, None, None), (70, 15, 30, 2)]) + # The head (instructions and tools) never changed, so one hash. + heads = {row["headHash"] for row in rows if row["kind"] == "turn"} + self.assertEqual(len(heads), 1) + self.assertNotIn(rows[5]["headHash"], heads) + + for row in rows: + self.assertIsNone(row["account"]) + # The shared harness names the fixture profile per route. + self.assertTrue(row["profile"].startswith("fixture")) + self.assertEqual(row["model"], "fixture-model") + self.assertTrue(row["provider"].startswith("responses:")) + self.assertLessEqual(row["startedMs"], row["finishedMs"]) + + for row in rows[:4] + [rows[6]]: + self.assertEqual(row["outcome"], "ok") + self.assertIsNone(row["status"]) + self.assertIsNone(row["error"]) + + for row in rows[:4]: + # Nothing replaced yet: the identity is a plain append. + self.assertEqual(row["replaced"], 0) + self.assertIsNone(row["projectionHash"]) + + # A 429 is a quota reading, with no usage and no transcript row. + limited = rows[4] + self.assertEqual(limited["outcome"], "error") + self.assertEqual(limited["status"], 429) + self.assertIn("limit exceeded", limited["error"]) + self.assertIsNone(limited["seq"]) + + # Each turn row's seq is the transcript row its response + # committed, exactly: the assistant row that turn produced, + # found again through the tree route straight from the + # transcript. + tree = self.tree(app, session) + assistant = {item["preview"]: item["id"] for item in tree + if item["type"] == "assistant"} + self.assertEqual(rows[0]["seq"], assistant["first reply"]) + self.assertEqual(rows[2]["seq"], assistant["tool done"]) + self.assertEqual(rows[3]["seq"], assistant["third reply"]) + self.assertEqual(rows[6]["seq"], + assistant["post compaction reply"]) + # The tool-call turn's row points at its own assistant row, + # the call it produced, not the tool output after it. + calls = [item for item in tree + if item["preview"] == "call python"] + self.assertEqual(len(calls), 1) + self.assertEqual(rows[1]["seq"], calls[0]["id"]) + tool_output = tree[tree.index(calls[0]) + 1] + self.assertEqual(tool_output["type"], "tool") + self.assertNotEqual(tool_output["id"], rows[1]["seq"]) + # Linked rows sit in transcript order. + linked = [row["seq"] for row in rows if row["seq"] is not None] + self.assertEqual(linked, sorted(linked)) + + # The compaction summary call is its own kind, outside the + # session projection. + summary = rows[5] + self.assertIsNone(summary["seq"]) + self.assertIsNone(summary["replaced"]) + self.assertIsNone(summary["projectionHash"]) + self.assertEqual(summary["inputs"], 1) + + # The projection identity changes once compaction rewrites the + # head the verbatim tail hangs from. + after = rows[6] + self.assertGreater(after["replaced"], 0) + self.assertIsNotNone(after["projectionHash"]) + self.assertTrue(after["strategy"]) + + # Paging: continue after the first row, bounded by limit. + second = self.rows(app, session, + f"?after={rows[0]['id']}&limit=2") + self.assertEqual([row["id"] for row in second["rows"]], + [row["id"] for row in rows[1:3]]) + self.assertEqual(second["after"], rows[2]["id"]) + finally: + provider.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/e2e/quota_test.py b/test/e2e/quota_test.py new file mode 100644 index 00000000..d61a97a0 --- /dev/null +++ b/test/e2e/quota_test.py @@ -0,0 +1,262 @@ +"""Quota readings land raw, each account keeps its own cadence, and failures back off. + +The daemon-wide poller maps every account it can see — an OAuth account from +auth.json, an API-key profile whose base url names a feed — onto provide-usage +credentials and records each report as it came. This is the only place that +wiring is visible from outside, so it is E2E: the usage core itself is a fake +executable answering the `usage advance -` envelope (the real one needs +accounts), pointed at by ALBEDO_USAGE_CORE with short poll settings. + +The cadence is asserted through the interval between one account's consecutive +readings, not through poll counts over wall time: counts race the machine, and +a hyper poll arriving a second late was read as "the busy cadence did not kick +in" when the schedule was exactly as configured. The intervals are what the +poller actually kept. +""" +import json +import os +import stat +import tempfile +import time +import unittest + +from harness import Albedo, Provider, exclusive, text + +POLL_SECONDS = 3 +BUSY_SECONDS = 1 + +FAKE_USAGE = """#!/usr/bin/env python3 +# Fake usage-core: one `advance -` envelope line in, one report line out. Each +# account is answered immediately, so one poll is one process. The provider in +# the envelope picks the script: anthropic is ordinary (5 percent), hyper is +# busy (85 percent, resetting soon), deepseek fails (an error report is data, +# not a driver failure). The happy path touches no file, so nothing can make +# it exit 1; anything unexpected is logged to fake-error.txt and answered as +# an error report rather than crashing, which the driver would report as the +# CLI itself failing. +import json, os, sys, time, traceback + +ERRORS = os.path.join(os.path.dirname(os.path.abspath(__file__)), "fake-error.txt") + +def report_for(provider): + if provider == "anthropic": + return {"planType": "pro", + "limits": [{"id": "weekly", "label": "Weekly usage", + "usedPercent": 5, "windowLabel": "7 days", + "status": "ok"}]} + if provider == "deepseek": + return {"error": "the fixture deepseek key has no quota"} + return {"limits": [{"id": "primary", "label": "Primary", "usedPercent": 85, + "windowLabel": "1 hour", "windowSeconds": 3600, + "resetsAt": int(time.time() * 1000) + 600000, + "status": "ok"}]} + +try: + line = sys.stdin.readline() + provider = json.loads(line)["provider"] + print(json.dumps({"report": report_for(provider)})) +except Exception: + # A bad envelope lands here too: record it and answer an error report, so + # the failure shows up as a reading rather than as an exit-1 crash. + try: + with open(ERRORS, "a") as errors: + errors.write(traceback.format_exc()) + except Exception: + pass + print(json.dumps({"report": {"error": "the fake usage core failed"}})) +""" + + +@exclusive +class QuotaTests(unittest.TestCase): + def setUp(self): + self.provider = Provider(lambda _request: text("ok")) + self.addCleanup(self.provider.close) + self.core = tempfile.TemporaryDirectory(prefix="quota-core-") + self.addCleanup(self.core.cleanup) + self.usage = os.path.join(self.core.name, "usage") + with open(self.usage, "w") as out: + out.write(FAKE_USAGE) + os.chmod(self.usage, os.stat(self.usage).st_mode | stat.S_IEXEC) + + def prepare(app): + # The poller resolves the core through the daemon's environment. + app.env["ALBEDO_USAGE_CORE"] = self.usage + settings = {"models": {"refreshHours": 0}, "quota": { + "pollSeconds": POLL_SECONDS, "busyPollSeconds": BUSY_SECONDS}} + (app.home / "extensions.json").write_text(json.dumps(settings)) + (app.home / "auth.json").write_text(json.dumps({"anthropic": [ + {"type": "oauth", "access": "fixture-anthropic-access", + "refresh": "fixture-anthropic-refresh", + "expires": int(time.time() * 1000) + 3600000, + "accountId": "fixture-account"}]})) + + self.app = Albedo(self.provider, providers={ + "fixture": {"baseUrl": self.provider.url, "apiKey": "fixture-key", + "model": "fixture-model", "protocol": "chat_completions"}, + "charm-hyper": {"baseUrl": "https://hyper.charm.land/v1", + "apiKey": "fixture-hyper-key", "model": "fixture-model", + "protocol": "chat_completions"}, + "deepseek-fixture": {"baseUrl": "https://api.deepseek.com/v1", + "apiKey": "fixture-deepseek-key", + "model": "fixture-model", + "protocol": "chat_completions"}, + }, prepare=prepare) + self.app.__enter__() + self.addCleanup(self.restore_home) + self.addCleanup(self.app.__exit__, None, None, None) + # A daemon booted before this fixture (the whole suite shares one) did + # not see ALBEDO_USAGE_CORE; one booted now does. Restarting makes both + # cases identical, and this test runs exclusively, after the rest. + self.app.restart() + # Readings from before the restart belong to a daemon that may not have + # had the fake; the cadence starts from the daemon under test. + self.since = time.time() * 1000 + + def restore_home(self): + self.app.env.pop("ALBEDO_USAGE_CORE", None) + (self.app.home / "auth.json").unlink(missing_ok=True) + + def readings(self): + with self.app.api("/quota") as response: + return json.load(response)["readings"] + + def samples(self): + """Every stored sample, newest first, through the history pages.""" + rows, cursor = [], 0 + while len(rows) < 1000: + with self.app.api(f"/quota?history={cursor}&limit=200") as response: + page = json.load(response) + rows.extend(page["items"]) + if not page["hasMore"]: + break + cursor = page["nextCursor"] + return rows + + def observed(self, account): + """One timestamp per poll of one account since the restart, oldest first.""" + return sorted({row["observedAt"] for row in self.samples() + if row["account"] == account and row["observedAt"] > self.since}) + + def intervals(self, account): + times = self.observed(account) + return [later - earlier for earlier, later in zip(times, times[1:])] + + def fake_errors(self): + try: + with open(os.path.join(self.core.name, "fake-error.txt")) as errors: + return errors.read() + except FileNotFoundError: + return "" + + def wait_for(self, condition, timeout, message): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if condition(): + return + time.sleep(0.25) + self.fail(message) + + def test_readings_land_and_each_account_keeps_its_cadence(self): + # Readings land for every account: the OAuth account from auth.json, + # the hyper-charm profile, and the deepseek profile. + wanted = {("anthropic", "weekly"), ("hyper", "primary"), ("deepseek", "")} + reading = {} + deadline = time.monotonic() + 40 + while time.monotonic() < deadline: + reading = {(r["provider"], r["limitId"]): r for r in self.readings()} + if wanted.issubset(reading): + break + time.sleep(0.25) + else: + self.fail("quota readings did not land: " + + json.dumps(self.readings()) + self.fake_errors()) + hyper = reading[("hyper", "primary")] + self.assertEqual(hyper["account"], "charm-hyper") + self.assertEqual(hyper["usedPercent"], 85.0) + self.assertEqual(hyper["windowSeconds"], 3600) + self.assertGreater(hyper["resetsAt"], int(time.time() * 1000)) + self.assertEqual(hyper["status"], "ok") + self.assertEqual(hyper["source"], "poll") + anthropic = reading[("anthropic", "weekly")] + self.assertEqual(anthropic["account"], "fixture-account") + self.assertEqual(anthropic["usedPercent"], 5.0) + # The report plan and each limit's window label are readings too, kept + # exactly as reported; where the report had neither, they stay null. + self.assertEqual(anthropic["plan"], "pro") + self.assertEqual(anthropic["windowLabel"], "7 days") + self.assertIsNone(hyper["plan"]) + self.assertEqual(hyper["windowLabel"], "1 hour") + # A failed report is recorded as a reading too, error and all; a + # missing percentage stays unknown rather than becoming zero. + deepseek = reading[("deepseek", "")] + self.assertIn("no quota", deepseek["error"]) + self.assertIsNone(deepseek["usedPercent"]) + self.assertIsNone(deepseek["plan"]) + + # Enough polls of each account to read a cadence off the intervals. + self.wait_for(lambda: len(self.observed("charm-hyper")) >= 6, 60, + "the busy cadence never produced six polls: " + + json.dumps(self.intervals("charm-hyper")) + self.fake_errors()) + self.wait_for(lambda: len(self.observed("fixture-account")) >= 4, 60, + "the ordinary cadence never produced four polls: " + + json.dumps(self.intervals("fixture-account")) + self.fake_errors()) + self.wait_for(lambda: len(self.observed("deepseek-fixture")) >= 3, 60, + "the failure backoff never produced three polls: " + + json.dumps(self.intervals("deepseek-fixture")) + self.fake_errors()) + + # The schedule as the poller kept it: one account's consecutive + # readings are one poll duration plus its wait, so each cadence shows + # up as an interval with generous bounds either side of its setting. + # The busy cadence (85 percent, resetting soon) is busyPollSeconds, + # the ordinary one is pollSeconds, and the failing account backs off + # past the ordinary cadence, each wait longer than the last. + busy = self.intervals("charm-hyper") + ordinary = self.intervals("fixture-account") + backed_off = self.intervals("deepseek-fixture") + + def median(values): + return sorted(values)[len(values) // 2] + + message = (f"busy {busy}, ordinary {ordinary}, backed off {backed_off}, " + f"fake errors {self.fake_errors()!r}") + self.assertGreaterEqual(median(busy), BUSY_SECONDS * 900, message) + self.assertLess(median(busy), BUSY_SECONDS * 2200, + message + " — the busy account polled at the ordinary cadence") + self.assertGreaterEqual(median(ordinary), POLL_SECONDS * 900, message) + self.assertLess(median(ordinary), POLL_SECONDS * 1700, message) + self.assertGreaterEqual(min(backed_off), POLL_SECONDS * 1800, + message + " — failures did not back off") + self.assertTrue(all(later >= earlier for earlier, later in zip(backed_off, backed_off[1:])), + message + " — the backoff did not grow") + + # The feed is the fake, so a driver failure ("usage is not built", the + # CLI crashing, the feed not finishing) is a bug in the host half, not + # weather. An exit-1 burst would surface here as readings carrying it. + failures = [row for row in self.samples() + if row["observedAt"] > self.since + and row["error"] and "usage" in row["error"]] + self.assertEqual(failures, [], "the usage driver failed against the fake") + self.assertEqual(self.fake_errors(), "") + + # The latest reading per account and limit is the newest sample, and + # history pages from it by row id, newest first. + latest = {(r["provider"], r["limitId"]): r for r in self.readings()} + with self.app.api("/quota?history=0") as response: + page = json.load(response) + samples = page["items"] + ids = [sample["id"] for sample in samples] + self.assertTrue(all(later < earlier for earlier, later in zip(ids, ids[1:])), + f"history is not newest first: {ids}") + newest = max((s for s in samples if s["provider"] == "hyper"), key=lambda s: s["id"]) + self.assertEqual(newest["observedAt"], latest[("hyper", "primary")]["observedAt"]) + cursor = ids[len(ids) // 2] + with self.app.api(f"/quota?history={cursor}") as response: + older = json.load(response) + self.assertTrue(all(sample["id"] < cursor for sample in older["items"]), + f"history after {cursor} reached newer rows: " + f"{[s['id'] for s in older['items']]}") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/harness/albedo_usage_test_support.erl b/test/harness/albedo_usage_test_support.erl new file mode 100644 index 00000000..8df7eadf --- /dev/null +++ b/test/harness/albedo_usage_test_support.erl @@ -0,0 +1,233 @@ +%% A fake usage CLI, a local HTTP endpoint and a command whose whole job is to +%% dump the environment it ran with: the three things usage_feed's loop has to +%% get right (envelope on stdin, responses replayed, allowlist on commands) +%% without a real provider, a real credential or the network. +%% +%% ALBEDO_USAGE_CORE points the driver's binary resolution at the fake, so the +%% daemon is never involved: no route exercises this loop end to end yet. +-module(albedo_usage_test_support). + +-export([start/0, stop/1, envelope/1, env_dump/1, request_line/1, + break_binary/1, mono_ms/0, trickle_pid_gone/1]). + +%% eunit runs a module's tests in parallel ({inparallel, _} reaches into the +%% module), and ALBEDO_USAGE_CORE is one global environment variable, so the +%% fakes must not overlap. A registered name is the lock: it is held from +%% start to stop, and the death of a crashed test process releases it. +-define(LOCK, albedo_usage_fake). + +start() -> + await_lock(), + Dir = filename:join(tmp(), "albedo-usage-" ++ os:getpid() ++ "-" + ++ integer_to_list(erlang:unique_integer([positive]))), + ok = file:make_dir(Dir), + {ok, Listen} = gen_tcp:listen(0, [binary, {active, false}, {reuseaddr, true}, {packet, raw}]), + {ok, {_Address, Port}} = inet:sockname(Listen), + Server = spawn(fun() -> serve(Listen, []) end), + Url = io_lib:format("http://127.0.0.1:~B/usage", [Port]), + Usage = filename:join(Dir, "usage"), + Dump = filename:join(Dir, "env-dump.sh"), + Trickle = filename:join(Dir, "trickle.sh"), + write(Usage, fake_usage(Url, Dump, Trickle, filename:join(Dir, "envelope.txt"))), + write(Dump, ["#!/bin/sh\nenv >", filename:join(Dir, "env.txt"), "\necho dumped\n"]), + %% Prints its own pid, then a line forever: a command that always has + %% output but never exits, for the deadline-kill test. + write(Trickle, ["#!/bin/sh\necho $$ >", filename:join(Dir, "trickle.pid"), + "\nwhile true; do\n echo tick\n sleep 0.2\ndone\n"]), + ok = file:change_mode(Usage, 8#755), + ok = file:change_mode(Dump, 8#755), + ok = file:change_mode(Trickle, 8#755), + os:putenv("ALBEDO_USAGE_CORE", Usage), + os:putenv("ALBEDO_USAGE_TEST_CANARY", "leak"), + {ok, #{dir => Dir, listen => Listen, server => Server, + envelope => filename:join(Dir, "envelope.txt"), + env => filename:join(Dir, "env.txt"), + trickle_pid => filename:join(Dir, "trickle.pid")}}. + +stop(Ref) -> + exit(maps:get(server, Ref), kill), + gen_tcp:close(maps:get(listen, Ref)), + os:unsetenv("ALBEDO_USAGE_CORE"), + os:unsetenv("ALBEDO_USAGE_TEST_CANARY"), + _ = file:del_dir_r(maps:get(dir, Ref)), + catch unregister(?LOCK), + nil. + +await_lock() -> + case whereis(?LOCK) of + Pid when is_pid(Pid) -> + Ref = erlang:monitor(process, Pid), + receive + {'DOWN', Ref, process, Pid, _} -> ok + after 15000 -> + exit(usage_fake_lock_timeout) + end, + await_lock(); + undefined -> + try register(?LOCK, self()) of + true -> ok + catch + _:_ -> await_lock() + end + end. + +envelope(Ref) -> + read(maps:get(envelope, Ref)). + +env_dump(Ref) -> + read(maps:get(env, Ref)). + +request_line(Ref) -> + Server = maps:get(server, Ref), + Server ! {requests, self()}, + receive + {requests, Server, Lines} -> Lines + after 2000 -> [] + end. + +%% Points resolution at a path that does not exist: the override must be +%% authoritative, never silently falling back to PATH. +break_binary(Ref) -> + os:putenv("ALBEDO_USAGE_CORE", filename:join(maps:get(dir, Ref), "not-there")), + nil. + +mono_ms() -> + erlang:monotonic_time(millisecond). + +%% The trickle command is killed on timeout, not orphaned: its pid stops +%% answering signal 0. The beam reaps its port children asynchronously, so +%% the check polls briefly rather than asserting the very first probe. +trickle_pid_gone(Ref) -> + Path = maps:get(trickle_pid, Ref), + case file:read_file(Path) of + {ok, Bin} -> + Pid = binary_to_integer(string:trim(Bin)), + gone(Pid, 20); + _ -> + false + end. + +gone(Pid, 0) -> + signal(Pid) =/= 0; +gone(Pid, N) -> + case signal(Pid) of + 0 -> timer:sleep(50), gone(Pid, N - 1); + _ -> true + end. + +%% Exit status of `kill -0 Pid`: 0 while the process exists. +signal(Pid) -> + Port = open_port({spawn_executable, "/bin/kill"}, + [{args, ["-0", integer_to_list(Pid)]}, exit_status, binary, hide]), + receive + {Port, {exit_status, Status}} -> Status + after 1000 -> + catch erlang:port_close(Port), + 0 + end. + +read(Path) -> + case file:read_file(Path) of + {ok, Bin} -> Bin; + _ -> <<>> + end. + +write(Path, Content) -> + ok = file:write_file(Path, Content). + +tmp() -> + case os:getenv("TMPDIR") of + false -> "/tmp"; + T -> T + end. + +%% The fake core speaks the real protocol: one envelope line in (the whole +%% history, on stdin), one step line out, and the provider name decides what +%% kind of feed it is. The default feed answers round one with an http request +%% and a command, and round two with a report built from the http body - so +%% the loop, the envelope and the command environment are all load-bearing. +fake_usage(Url, Dump, Trickle, EnvelopeLog) -> + ["#!/usr/bin/env python3\n" + "import json, sys\n" + "line = sys.stdin.readline()\n" + "open(", quote(EnvelopeLog), ", 'a').write(line)\n" + "envelope = json.loads(line)\n" + "provider = envelope.get('provider')\n" + "responses = envelope.get('responses', [])\n" + "if provider == 'endless':\n" + " step = {'requests': [{'kind': 'command', 'command': 'true', 'args': []}]}\n" + "elif provider == 'broken':\n" + " step = {'report': {'limits': [], 'error': 'the provider is unreachable'}}\n" + "elif provider == 'trickle':\n" + " if responses:\n" + " status = responses[-1][0]['status']\n" + " step = {'report': {'limits': [], 'error': 'command answered %d' % status}}\n" + " else:\n" + " step = {'requests': [\n" + " {'kind': 'command', 'command': ", quote(Trickle), ", 'args': [],\n" + " 'timeoutMs': 1000}]}\n" + "elif provider == 'missing-command' and responses:\n" + " status = responses[-1][0]['status']\n" + " step = {'report': {'limits': [], 'error': 'command answered %d' % status}}\n" + "elif responses:\n" + " body = json.loads(responses[-1][0]['body'])\n" + " step = {'report': {\n" + " 'accountId': 'acct-1', 'email': 'dawn@example.com', 'planType': 'token',\n" + " 'limits': [{'id': 'five_hour', 'label': 'Session',\n" + " 'usedPercent': body['pct'], 'windowLabel': '5h window',\n" + " 'windowSeconds': 18000, 'resetsAt': 1893456000000,\n" + " 'status': 'ok'}],\n" + " 'error': None}}\n" + "elif provider == 'missing-command':\n" + " step = {'requests': [\n" + " {'kind': 'command', 'command': 'no-such-binary-xyz', 'args': []}]}\n" + "else:\n" + " step = {'requests': [\n" + " {'kind': 'http', 'method': 'GET', 'url': ", quote(Url), ",\n" + " 'headers': [['accept', 'application/json']]},\n" + " {'kind': 'command', 'command': ", quote(Dump), ", 'args': [],\n" + " 'timeoutMs': 5000}]}\n" + "print(json.dumps(step))\n"]. + +quote(Text) -> + [$', Text, $']. + +serve(Listen, Acc) -> + receive + {requests, From} -> + From ! {requests, self(), lists:reverse(Acc)} + after 0 -> + ok + end, + case gen_tcp:accept(Listen, 200) of + {ok, Socket} -> + case read_head(Socket, <<>>) of + {ok, Head} -> + [RequestLine | _] = binary:split(Head, <<"\r\n">>), + Body = <<"{\"pct\":37}">>, + ok = gen_tcp:send(Socket, [ + <<"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\n">>, + <<"content-length: ">>, integer_to_binary(byte_size(Body)), + <<"\r\n\r\n">>, Body]), + gen_tcp:close(Socket), + serve(Listen, [RequestLine | Acc]); + _ -> + gen_tcp:close(Socket), + serve(Listen, Acc) + end; + {error, timeout} -> + serve(Listen, Acc); + _ -> + ok + end. + +read_head(Socket, Acc) -> + case binary:match(Acc, <<"\r\n\r\n">>) of + nomatch -> + case gen_tcp:recv(Socket, 0, 5000) of + {ok, Bytes} -> read_head(Socket, <>); + Error -> Error + end; + _ -> + {ok, Acc} + end. diff --git a/test/harness/antigravity_test.gleam b/test/harness/antigravity_test.gleam index a64e6b85..4f1c8e48 100644 --- a/test/harness/antigravity_test.gleam +++ b/test/harness/antigravity_test.gleam @@ -103,7 +103,9 @@ pub fn stream_reduces_thoughts_text_and_calls_test() { ] = events assert arguments == "{\"command\":\"ls\"}" assert turn.finish == types.ToolCalls - assert turn.usage == Some(types.Usage(10, 5, Some(4), None)) + // Output folds the 2 thought tokens in; reasoning carries the thought count. + assert turn.usage + == Some(types.Usage(10, 5, Some(4), None, None, None, Some(2))) let assert [types.ToolCall(id, "bash", "{\"command\":\"ls\"}")] = turn.tool_calls assert string.starts_with(id, "call_") diff --git a/test/harness/usage_feed_test.gleam b/test/harness/usage_feed_test.gleam new file mode 100644 index 00000000..f0ccee55 --- /dev/null +++ b/test/harness/usage_feed_test.gleam @@ -0,0 +1,131 @@ +//// The usage-feed driver against a fake usage CLI (ALBEDO_USAGE_CORE), a +//// local HTTP endpoint and a command that dumps its own environment. The +//// daemon has no route for quota yet, so no e2e scenario can reach this loop; +//// the bugs it catches are the host half drifting from the core's protocol: +//// the credential landing in argv instead of stdin, responses not replayed, +//// or the command environment growing beyond the allowlist. + +import albedo/harness/usage_feed +import gleam/json +import gleam/option.{None, Some} +import gleam/string +import gleeunit/should + +pub type Fake + +@external(erlang, "albedo_usage_test_support", "start") +fn start_fake() -> Result(Fake, Nil) + +@external(erlang, "albedo_usage_test_support", "stop") +fn stop_fake(fake: Fake) -> Nil + +@external(erlang, "albedo_usage_test_support", "envelope") +fn fake_envelope(fake: Fake) -> String + +@external(erlang, "albedo_usage_test_support", "env_dump") +fn fake_env_dump(fake: Fake) -> String + +@external(erlang, "albedo_usage_test_support", "request_line") +fn fake_request_line(fake: Fake) -> List(String) + +@external(erlang, "albedo_usage_test_support", "break_binary") +fn break_binary(fake: Fake) -> Nil + +@external(erlang, "albedo_usage_test_support", "mono_ms") +fn mono_ms() -> Int + +@external(erlang, "albedo_usage_test_support", "trickle_pid_gone") +fn trickle_pid_gone(fake: Fake) -> Bool + +fn credential() -> json.Json { + json.object([ + #("kind", json.string("oauth")), + #("access", json.string("tok_n0t_a_real_secret")), + ]) +} + +pub fn fetch_runs_http_and_command_rounds_test() { + let assert Ok(fake) = start_fake() + let assert Ok(report) = usage_feed.fetch("fake", credential(), 1_234_567_890) + // The fake core built the report from the real HTTP response body, which + // only the loop could have delivered. + report.account_id |> should.equal(Some("acct-1")) + report.email |> should.equal(Some("dawn@example.com")) + report.plan |> should.equal(Some("token")) + report.error |> should.equal(None) + let assert [limit] = report.limits + limit.id |> should.equal("five_hour") + limit.label |> should.equal("Session") + limit.used_percent |> should.equal(Some(37.0)) + limit.window_label |> should.equal(Some("5h window")) + limit.window_seconds |> should.equal(Some(18_000)) + limit.resets_at |> should.equal(Some(1_893_456_000_000)) + limit.scope |> should.equal(None) + limit.status |> should.equal("ok") + // The HTTP request went through albedo's HTTP client to the local endpoint. + fake_request_line(fake) |> should.equal(["GET /usage HTTP/1.1"]) + // The envelope crossed on stdin, never argv: both logged rounds carry the + // provider, the credential and nowMs, and the second replays the responses. + let envelope = fake_envelope(fake) + envelope |> string.contains("\"provider\":\"fake\"") |> should.be_true + envelope |> string.contains("tok_n0t_a_real_secret") |> should.be_true + envelope |> string.contains("\"nowMs\":1234567890") |> should.be_true + envelope |> string.contains("\"responses\":[]") |> should.be_true + envelope |> string.contains("\"status\":200") |> should.be_true + // The command ran with only the allowed environment: no canary, and the + // resolver override itself never reached the child. + let dumped = fake_env_dump(fake) + dumped |> string.contains("ALBEDO_USAGE_TEST_CANARY") |> should.be_false + dumped |> string.contains("ALBEDO_USAGE_CORE") |> should.be_false + dumped |> string.contains("PATH=") |> should.be_true + dumped |> string.contains("HOME=") |> should.be_true + stop_fake(fake) +} + +pub fn a_provider_error_is_data_not_a_driver_failure_test() { + let assert Ok(fake) = start_fake() + let assert Ok(report) = usage_feed.fetch("broken", credential(), 0) + report.error |> should.equal(Some("the provider is unreachable")) + report.limits |> should.equal([]) + report.account_id |> should.equal(None) + stop_fake(fake) +} + +pub fn a_feed_that_never_reports_is_a_driver_failure_test() { + let assert Ok(fake) = start_fake() + let assert Error(message) = usage_feed.fetch("endless", credential(), 0) + message |> string.contains("did not finish in 6 rounds") |> should.be_true + stop_fake(fake) +} + +pub fn a_missing_binary_is_a_driver_failure_test() { + let assert Ok(fake) = start_fake() + break_binary(fake) + let assert Error(message) = usage_feed.fetch("fake", credential(), 0) + message |> string.contains("usage is not built") |> should.be_true + stop_fake(fake) +} + +pub fn a_command_off_path_answers_126_test() { + let assert Ok(fake) = start_fake() + // The fake core asks for a command that does not exist; the feed routes + // around a non-zero status, so the report still arrives. + let assert Ok(report) = usage_feed.fetch("missing-command", credential(), 0) + report.error |> should.equal(Some("command answered 126")) + stop_fake(fake) +} + +pub fn a_trickling_command_hits_its_deadline_and_is_killed_test() { + let assert Ok(fake) = start_fake() + // The trickle command always has output but never exits: an idle timeout + // would let it run forever, so only the deadline across chunks ends it. + let started = mono_ms() + let assert Ok(report) = usage_feed.fetch("trickle", credential(), 0) + let elapsed = mono_ms() - started + report.error |> should.equal(Some("command answered 126")) + // The request's 1s timeoutMs is honoured, not the 15s default. + should.be_true(elapsed < 4000) + // And the child was killed, not orphaned with closed pipes. + trickle_pid_gone(fake) |> should.be_true + stop_fake(fake) +} diff --git a/test/openai_api/stream_test.gleam b/test/openai_api/stream_test.gleam index 91d074b3..9b1c0e01 100644 --- a/test/openai_api/stream_test.gleam +++ b/test/openai_api/stream_test.gleam @@ -93,7 +93,7 @@ pub fn chat_accumulates_interleaved_tools_usage_and_native_replay_test() { types.ToolCall("call_a", "alpha", "{\"a\":1}"), types.ToolCall("call_b", "beta", "{\"b\":2}"), ], - Some(types.Usage(5, 8, Some(0), None)), + Some(types.Usage(5, 8, Some(0), None, None, None, None)), types.ToolCalls, None, )), @@ -121,7 +121,14 @@ pub fn cache_usage_details_tolerate_null_and_reject_malformed_values_test() { let assert Ok(#( _, [], - Some(types.Turn(_, [], [], Some(types.Usage(3, 2, None, None)), _, None)), + Some(types.Turn( + _, + [], + [], + Some(types.Usage(3, 2, None, None, None, None, None)), + _, + None, + )), )) = send(stream.new(types.Responses), "", responses_null) let chat_null = @@ -131,7 +138,14 @@ pub fn cache_usage_details_tolerate_null_and_reject_malformed_values_test() { let assert Ok(#( _, [], - Some(types.Turn(_, _, [], Some(types.Usage(3, 2, None, None)), _, None)), + Some(types.Turn( + _, + _, + [], + Some(types.Usage(3, 2, None, None, None, None, None)), + _, + None, + )), )) = send(state, "", "[DONE]") let responses_malformed =