diff --git a/deploy/beacon/package.json b/deploy/beacon/package.json new file mode 100644 index 000000000..b1913f3bd --- /dev/null +++ b/deploy/beacon/package.json @@ -0,0 +1,5 @@ +{ + "name": "headroom-beacon", + "private": true, + "type": "module" +} diff --git a/deploy/beacon/test-rollup.mjs b/deploy/beacon/test-rollup.mjs new file mode 100644 index 000000000..3878ee2d4 --- /dev/null +++ b/deploy/beacon/test-rollup.mjs @@ -0,0 +1,214 @@ +/** + * Self-check for scheduled()'s hourly compaction. node test-rollup.mjs [dir] + * + * The one thing that must never drift: rollupHour() and the QUALIFY in + * headroom-beacon-stats/beacon.sh have to agree on which heartbeat wins. If + * they disagree the reports get quietly wrong rather than loudly broken, so + * this asserts the JS picks exactly the max-seq row per (install, session). + * + * Point it at a directory of real beacon objects to check against the corpus: + * aws s3 sync s3://headroom-telemetry/sessions/dt=.../hh=.../ /tmp/hr/ ... + * node test-rollup.mjs /tmp/hr + * With no argument it runs on a small fixture and needs no network. + */ +import { readdirSync, readFileSync } from 'node:fs'; +import assert from 'node:assert/strict'; +import { oldestRawDay, rollupHour } from './worker.js'; + +// R2 returns at most 1000 keys per list page, so on a real hour (~4,000 +// objects) the cursor loop in rollupHour is load-bearing. The stub paginates at +// a deliberately tiny size so that loop is exercised by every case below: with +// a single-page stub, a regression that dropped the cursor would still print +// "ok" while silently rolling up only the first page of every hour. +const PAGE = 3; + +/** The slice of the R2 binding rollupHour uses, backed by a plain object. */ +function stubBucket(files, { failKeys = new Set() } = {}) { + const written = {}; + const reads = []; + return { + written, + reads, + list: async ({ prefix, cursor, delimiter }) => { + const keys = Object.keys(files) + .filter((k) => k.startsWith(prefix)) + .sort(); + if (delimiter) { + const seen = new Set(); + for (const k of keys) { + const cut = k.indexOf(delimiter, prefix.length); + if (cut >= 0) seen.add(k.slice(0, cut + 1)); + } + return { objects: [], delimitedPrefixes: [...seen], truncated: false }; + } + const start = cursor ? keys.indexOf(cursor) : 0; + const page = keys.slice(start, start + PAGE); + const next = start + PAGE; + return { + objects: page.map((key) => ({ key })), + truncated: next < keys.length, + cursor: next < keys.length ? keys[next] : undefined, + }; + }, + get: async (key) => { + reads.push(key); + if (failKeys.has(key)) throw new Error(`simulated R2 failure: ${key}`); + if (!(key in files)) return null; + return { text: async () => files[key] }; + }, + put: async (key, body) => { + written[key] = body; + }, + }; +} + +const beacon = (install, id, seq) => + JSON.stringify({ resource: { 'headroom.install_id': install }, session: { id, seq } }); + +const PART = 'dt=2026-08-06/hh=14'; + +/** Run rollupHour against a stub bucket and decode whatever it wrote. */ +async function run(files, opts = {}) { + const CORPUS = stubBucket(files, opts); + const spend = { read: 0 }; + let threw = null; + let out = null; + try { + out = await rollupHour({ CORPUS }, PART, spend); + } catch (err) { + threw = err; + } + const body = CORPUS.written[`rollup/${PART}/data.ndjson`]; + return { + threw, + spend, + wrote: out ? out.wrote : 0, + empty: `rollup/${PART}/empty` in CORPUS.written, + keys: Object.keys(CORPUS.written), + rows: body ? body.split('\n').map((l) => JSON.parse(l)) : [], + }; +} + +// 1. Highest seq wins, out-of-order input, one row per (install, session). +// More objects than PAGE, so the list cursor loop runs. +{ + const files = { + [`sessions/${PART}/a.json`]: [beacon('i1', 's1', 3), beacon('i1', 's2', 1)].join('\n'), + [`sessions/${PART}/b.json`]: beacon('i1', 's1', 9), + [`sessions/${PART}/c.json`]: beacon('i1', 's1', 7), + // Same session id under a different install must not collapse together. + [`sessions/${PART}/d.json`]: beacon('i2', 's1', 2), + [`sessions/${PART}/e.json`]: beacon('i1', 's1', 5), + }; + const { rows, spend, threw } = await run(files); + assert.equal(threw, null); + // 5 objects at PAGE=3 is two pages: proves the cursor loop, which is + // load-bearing at the real ~4,000 objects/hour. + assert.ok(Object.keys(files).length > PAGE, 'fixture must span pages'); + assert.equal(spend.read, 5, 'reads every object across every page'); + assert.equal(rows.length, 3, 'one row per (install, session)'); + const seq = Object.fromEntries( + rows.map((r) => [`${r.resource['headroom.install_id']} ${r.session.id}`, r.session.seq]) + ); + assert.deepEqual(seq, { 'i1 s1': 9, 'i1 s2': 1, 'i2 s1': 2 }); +} + +// 2. An unparseable record loses only itself. Content this Worker wrote with +// JSON.stringify never becomes valid later, so blocking the hour on it would +// strand the hour rather than one record. +{ + const files = { + [`sessions/${PART}/a.json`]: '{ this is not json', + [`sessions/${PART}/b.json`]: `\n${beacon('i1', 's1', 4)}\n`, + }; + const { rows, threw } = await run(files); + assert.equal(threw, null, 'corrupt content does not abandon the hour'); + assert.deepEqual(rows.map((r) => r.session.seq), [4], 'survives a corrupt object'); +} + +// 3. A failed get is transient, so the hour must NOT be written — a rollup is +// built once and then trusted forever, so a short read would silently become +// the permanent record. +{ + const files = { + [`sessions/${PART}/a.json`]: beacon('i1', 's1', 1), + [`sessions/${PART}/b.json`]: beacon('i1', 's2', 1), + }; + const { threw, keys } = await run(files, { + failKeys: new Set([`sessions/${PART}/b.json`]), + }); + assert.ok(threw, 'a failed get throws so the hour is retried'); + assert.deepEqual(keys, [], 'nothing written on a partial read'); +} + +// 4. Spend is reported even when the hour throws. Charging a flat guess instead +// lets a run that failed late overshoot the subrequest ceiling. +{ + const files = Object.fromEntries( + Array.from({ length: 7 }, (_, i) => [`sessions/${PART}/o${i}.json`, beacon('i1', `s${i}`, 1)]) + ); + const { threw, spend } = await run(files, { + failKeys: new Set([`sessions/${PART}/o6.json`]), + }); + assert.ok(threw); + assert.equal(spend.read, 7, 'caller sees real spend, not a guess'); +} + +// 5. An empty hour writes a marker, not a zero-byte NDJSON. Without it the hour +// stays "missing" and is re-listed on every run forever. +{ + const { rows, empty, keys } = await run({}); + assert.deepEqual(rows, []); + assert.ok(empty, 'empty hour leaves a marker'); + assert.ok( + keys.every((k) => !k.endsWith('.ndjson')), + 'no zero-byte ndjson for readers to special-case' + ); +} + +// 6. oldestRawDay floors the backfill. A fixed lookback window silently strands +// every hour older than it once analysis stopped reading sessions/. +{ + const CORPUS = stubBucket({ + 'sessions/dt=2026-08-03/hh=01/a.json': beacon('i1', 's1', 1), + 'sessions/dt=2026-08-06/hh=14/b.json': beacon('i1', 's2', 1), + 'sessions/dt=2026-08-07/hh=00/c.json': beacon('i1', 's3', 1), + }); + assert.equal(await oldestRawDay({ CORPUS }), '2026-08-03'); + assert.equal(await oldestRawDay({ CORPUS: stubBucket({}) }), null, 'empty bucket -> null'); +} + +// 7. Against real objects, if a directory was given: same answer as the QUALIFY +// in beacon.sh, which is `count(DISTINCT install||session)` rows, each +// carrying that pair's max seq. +const dir = process.argv[2]; +if (dir) { + const files = {}; + for (const f of readdirSync(dir).filter((f) => f.endsWith('.json'))) { + files[`sessions/${PART}/${f}`] = readFileSync(`${dir}/${f}`, 'utf8'); + } + const { rows, spend, threw } = await run(files); + assert.equal(threw, null); + + const expected = new Map(); + for (const text of Object.values(files)) { + for (const line of text.split('\n')) { + if (!line.trim()) continue; + const r = JSON.parse(line); + const k = `${r.resource?.['headroom.install_id']} ${r.session?.id}`; + expected.set(k, Math.max(expected.get(k) ?? -1, r.session?.seq ?? 0)); + } + } + assert.equal(spend.read, Object.keys(files).length); + assert.equal(rows.length, expected.size, 'row count matches DISTINCT sessions'); + for (const r of rows) { + const k = `${r.resource['headroom.install_id']} ${r.session.id}`; + assert.equal(r.session.seq, expected.get(k), `max seq for ${k}`); + } + console.log( + `real corpus: ${spend.read} objects -> ${rows.length} sessions in 1 object` + + ` (${Math.ceil(spend.read / PAGE)} list pages)` + ); +} + +console.log('ok'); diff --git a/deploy/beacon/worker.js b/deploy/beacon/worker.js index 15bc22d4f..7fea03001 100644 --- a/deploy/beacon/worker.js +++ b/deploy/beacon/worker.js @@ -121,7 +121,184 @@ function extract(payload) { return records; } +// ----------------------------------------------------------------- rollup -- +// +// The corpus is one object per heartbeat, ~1KB each — 65k on 2026-08-06 and +// climbing. DuckDB reads them correctly, but a full `pull` is ~100k HTTPS round +// trips for 95MB: minutes of pure per-object latency, no real bytes or compute. +// Listing the bucket alone took 88 seconds. +// +// This job collapses each COMPLETE hour into one object under rollup/, keeping +// only the highest-seq heartbeat per (install, session). One measured hour +// (dt=2026-08-06/hh=14): 3,938 objects and 3,938 rows in, 1 object and 1,061 +// rows out. Analysis reads rollup/**, never sessions/**. Raw is left exactly as +// written, so any rollup can be rebuilt by deleting it. +// +// Hourly rather than daily because every R2 binding call is a subrequest: a day +// is ~65k of them against a 10k-per-invocation ceiling, an hour is ~4k. + +const READ_BUDGET = 60000; // objects per run; see [limits] in wrangler.toml +// A get costs ~45ms of round trip and almost no CPU, so this is what decides +// whether a run finishes: at 20 an hour took ~3 minutes, against a 15-minute +// wall clock for a cron invocation. Raise it if an hour ever stops fitting. +const FANOUT = 100; // concurrent R2 gets + +const partition = (d) => + `dt=${d.toISOString().slice(0, 10)}/hh=${d.toISOString().slice(11, 13)}`; + +/** + * One hour of heartbeats -> one deduped NDJSON object. + * + * Returns `{ read, wrote }`. Spend is reported through the mutable `spend` + * accumulator so the caller still knows it even when this throws: the budget + * has to track real spend, and a flat guess lets a run that failed late + * overshoot the subrequest ceiling and get killed inside an hour that would + * otherwise have succeeded. + * + * Writes nothing unless the whole hour read cleanly. A rollup is built once and + * then treated as done forever, so a partial read would silently become the + * permanent record — better to write nothing and let the next run retry. + */ +export async function rollupHour(env, part, spend = { read: 0 }) { + const best = new Map(); + let failed = 0; // transient: retry the hour + let corrupt = 0; // permanent: record and move on + let cursor; + do { + const page = await env.CORPUS.list({ prefix: `sessions/${part}/`, cursor }); + for (let i = 0; i < page.objects.length; i += FANOUT) { + // allSettled, not all: one transient R2 error among the ~4,000 gets in a + // real hour would otherwise reject the batch and discard the whole hour. + const settled = await Promise.allSettled( + page.objects + .slice(i, i + FANOUT) + .map((o) => env.CORPUS.get(o.key).then((r) => (r ? r.text() : null))) + ); + for (const outcome of settled) { + spend.read++; + // A miss counts as a failure too. The key came from a LIST, so the + // object existed; treating it as empty would quietly shrink the rollup. + if (outcome.status !== 'fulfilled' || outcome.value === null) { + failed++; + continue; + } + for (const line of outcome.value.split('\n')) { + if (!line) continue; + let rec; + try { + rec = JSON.parse(line); + } catch { + // Counted and logged, but NOT a reason to abandon the hour. A + // failed get is transient and worth retrying; content this Worker + // itself wrote with JSON.stringify does not become valid later, so + // blocking on it would strand the hour until its raw objects + // expire and then lose the whole hour instead of one record. + corrupt++; + continue; + } + // A session heartbeats every 5 minutes carrying CUMULATIVE totals, so + // the highest seq IS the whole session and every earlier row is a + // strict subset. Sessions straddle hours, so readers still dedupe + // across rollups on this same key — this only shrinks each hour. + const id = `${rec.resource?.['headroom.install_id']} ${rec.session?.id}`; + const prev = best.get(id); + if (!prev || (rec.session?.seq ?? 0) > (prev.session?.seq ?? 0)) { + best.set(id, rec); + } + } + } + } + cursor = page.truncated ? page.cursor : undefined; + } while (cursor); + + if (failed) { + throw new Error(`${part}: ${failed} of ${spend.read} objects unreadable`); + } + if (corrupt) { + console.error(`rollup ${part}: skipped ${corrupt} unparseable record(s)`); + } + + // A genuinely empty hour gets a marker rather than a zero-byte NDJSON that + // every reader would have to special-case. Without it the hour stays + // "missing" and is re-listed on every run for the life of the bucket. + if (best.size === 0) { + await env.CORPUS.put(`rollup/${part}/empty`, ''); + return { read: spend.read, wrote: 0 }; + } + await env.CORPUS.put( + `rollup/${part}/data.ndjson`, + [...best.values()].map((r) => JSON.stringify(r)).join('\n'), + { httpMetadata: { contentType: 'application/x-ndjson' } } + ); + return { read: spend.read, wrote: best.size }; +} + +/** Oldest `dt=` day still under sessions/, or null. One delimited LIST. */ +export async function oldestRawDay(env) { + const page = await env.CORPUS.list({ prefix: 'sessions/', delimiter: '/' }); + const days = (page.delimitedPrefixes || []) + .map((p) => p.slice('sessions/dt='.length).replace(/\/$/, '')) + .filter((d) => /^\d{4}-\d{2}-\d{2}$/.test(d)) + .sort(); + return days.length ? days[0] : null; +} + export default { + /** Hourly cron. Builds every complete hour back to the oldest raw data. */ + async scheduled(event, env) { + // Backfill reaches all the way to the oldest surviving raw day, NOT a fixed + // window. A fixed window silently strands everything older than it the + // moment analysis stopped reading sessions/ — the raw objects are still + // there, but nothing would ever compact them, so they vanish from every + // report. Bounding by real data instead means the floor rises only when a + // lifecycle rule actually expires the raw objects. + const oldest = await oldestRawDay(env); + if (!oldest) return; + const floorMs = Date.parse(`${oldest}T00:00:00Z`); + if (Number.isNaN(floorMs)) return; + + // Only list from the floor forward. Rollups older than the oldest raw day + // can never be rebuilt, so enumerating them answers nothing — this is what + // keeps the listing bounded by retention rather than by total history. + const done = new Set(); + let cursor; + do { + const page = await env.CORPUS.list({ + prefix: 'rollup/', + startAfter: `rollup/dt=${oldest}`, + cursor, + }); + for (const o of page.objects) { + // Tolerates both `/data.ndjson` and the `/empty` marker. + const rel = o.key.slice('rollup/'.length); + const cut = rel.lastIndexOf('/'); + if (cut > 0) done.add(rel.slice(0, cut)); + } + cursor = page.truncated ? page.cursor : undefined; + } while (cursor); + + // Newest first, so a backlog drains from the present backwards and the + // freshest hour is never the one starved by the budget. Starts one hour + // back: the current hour is still being written to. + let budget = READ_BUDGET; + for (let t = event.scheduledTime - 3600_000; t >= floorMs && budget > 0; t -= 3600_000) { + const part = partition(new Date(t)); + if (done.has(part)) continue; + // Shared with rollupHour so a throw still reports what it spent. + const spend = { read: 0 }; + try { + await rollupHour(env, part, spend); + } catch (err) { + // Newest-first means an hour that always throws — one grown past the + // subrequest ceiling, say — would otherwise block every older hour + // behind it forever. Skip it and keep draining; it has no marker, so + // the next run retries it. + console.error(`rollup ${part} failed after ${spend.read} objects: ${err}`); + } + budget -= spend.read; + } + }, + async fetch(request, env, ctx) { if (request.method !== 'POST') { return new Response('beacon: POST OTLP logs to /v1/logs', { status: 405 }); @@ -145,14 +322,13 @@ export default { } if (records.length === 0) return new Response(null, { status: 204 }); - const now = new Date(); - const day = now.toISOString().slice(0, 10); - const hour = now.toISOString().slice(11, 13); // Hive-style partitioning so DuckDB can prune by date without a catalog. - // ponytail: one object per request. At beacon volume that is a few hundred - // thousand objects a month, which globs fine. Add a daily compaction job - // when the file count starts to slow queries, not before. - const key = `sessions/dt=${day}/hh=${hour}/${crypto.randomUUID()}.json`; + // Shares partition() with the rollup: the cron lists `sessions//`, so + // two independent spellings of this scheme would mean the writer and the + // compactor could drift apart and silently match zero objects. + // ponytail: one object per request. Compacted hourly into rollup/ by + // scheduled() above — analysis reads that, never this. + const key = `sessions/${partition(new Date())}/${crypto.randomUUID()}.json`; const ndjson = records.map((r) => JSON.stringify(r)).join('\n'); // Respond immediately; durability work continues after the response. diff --git a/deploy/beacon/wrangler.toml b/deploy/beacon/wrangler.toml index 412b22045..826dba6bb 100644 --- a/deploy/beacon/wrangler.toml +++ b/deploy/beacon/wrangler.toml @@ -32,6 +32,24 @@ bucket_name = "headroom-telemetry" # npx wrangler secret put METRICS_OTLP_AUTH # Absent = R2 only, which is the right place to start. +# Hourly compaction of sessions/ into rollup/ — see scheduled() in worker.js. +# At :05 so the hour being rolled up is definitely closed. A >=1h interval also +# buys the 15-minute CPU limit instead of 30s, which the backfill run needs. +[triggers] +crons = ["5 * * * *"] + +# Every R2 binding call is a subrequest, and one hour is already ~4k objects. +# The paid default of 10k would cap a run at two hours and stall the backfill +# behind live traffic forever. This only raises a ceiling; a normal run spends +# ~4k. READ_BUDGET in worker.js is what actually bounds the work. +# +# Workers Paid only — on the Free plan this key is rejected outright ("CPU +# limits are not supported for the Free plan"), and the cron could not run +# anyway: Free gives a scheduled handler 10ms of CPU, and parsing an hour of +# heartbeats is tens of ms. +[limits] +subrequests = 100000 + [observability] enabled = true diff --git a/headroom/integrations/langchain/langgraph.py b/headroom/integrations/langchain/langgraph.py index 2e6734795..83aad7560 100644 --- a/headroom/integrations/langchain/langgraph.py +++ b/headroom/integrations/langchain/langgraph.py @@ -49,6 +49,7 @@ except ImportError: from headroom.ccr.tool_injection import CCR_TOOL_NAME from headroom.config import is_tool_excluded +from headroom.telemetry.session import BeaconCompressionObserver from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig logger = logging.getLogger(__name__) @@ -136,7 +137,11 @@ class _CrusherSingleton: config = SmartCrusherConfig( min_tokens_to_crush=self._min_tokens, ) - self._crusher = SmartCrusher(config=config) + # observer: no proxy here, so nothing else reports these + # compressions to the beacon. See BeaconCompressionObserver. + self._crusher = SmartCrusher( + config=config, observer=BeaconCompressionObserver() + ) return self._crusher diff --git a/headroom/integrations/mcp/server.py b/headroom/integrations/mcp/server.py index 23fdf3e8f..888ad6115 100644 --- a/headroom/integrations/mcp/server.py +++ b/headroom/integrations/mcp/server.py @@ -56,6 +56,7 @@ from typing import Any from headroom.config import HeadroomConfig, SmartCrusherConfig from headroom.providers.openai import OpenAIProvider +from headroom.telemetry.session import BeaconCompressionObserver from headroom.transforms.smart_crusher import SmartCrusher @@ -263,7 +264,18 @@ class HeadroomMCPCompressor: min_tokens_to_crush=profile.min_tokens_to_compress, max_items_after_crush=profile.max_items, ) - crusher = SmartCrusher(config=smart_config, with_compaction=False) # type: ignore[arg-type] + # observer: MCP runs outside the proxy, so PrometheusMetrics (the + # proxy's observer, which forwards to the beacon) never sees these + # compressions. Without one, an MCP install reports real tokens.saved + # with an empty compression.by_strategy. + crusher = SmartCrusher( + # headroom.config.SmartCrusherConfig vs the transform's own + # same-named dataclass; the ignore has to sit on the argument line + # because that is where mypy reports a multi-line call's arg-type. + config=smart_config, # type: ignore[arg-type] + with_compaction=False, + observer=BeaconCompressionObserver(), + ) # Build messages for SmartCrusher (it expects conversation format) messages = [ diff --git a/headroom/integrations/strands/hooks.py b/headroom/integrations/strands/hooks.py index 275b6f64d..07ef31887 100644 --- a/headroom/integrations/strands/hooks.py +++ b/headroom/integrations/strands/hooks.py @@ -50,6 +50,7 @@ except ImportError: from headroom import HeadroomConfig from headroom.ccr.tool_injection import CCR_TOOL_NAME from headroom.config import is_tool_excluded +from headroom.telemetry.session import BeaconCompressionObserver from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig logger = logging.getLogger(__name__) @@ -173,7 +174,11 @@ class HeadroomHookProvider(HookProvider): # type: ignore[misc] crusher_config = SmartCrusherConfig( min_tokens_to_crush=self.min_tokens_to_compress ) - self._crusher = SmartCrusher(config=crusher_config) + # observer: no proxy here, so nothing else reports these + # compressions to the beacon. See BeaconCompressionObserver. + self._crusher = SmartCrusher( + config=crusher_config, observer=BeaconCompressionObserver() + ) logger.debug( "SmartCrusher initialized with min_tokens=%d", self.min_tokens_to_compress ) diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index 3bf4b4103..fe477290b 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -417,14 +417,14 @@ class PrometheusMetrics: return total_input_tokens, total_input_cost_usd try: - # totals() rather than stats(): identical numbers, without the - # 31-day cost-record walk that stats()["budget_basis"] performs and - # this caller throws away. See CostTracker.totals. - tracked_input_tokens, tracked_input_cost_usd = self.cost_tracker.totals() + cost_stats = self.cost_tracker.stats() except Exception: logger.debug("Failed to read cost tracker totals for savings history", exc_info=True) return total_input_tokens, total_input_cost_usd + tracked_input_tokens = cost_stats.get("total_input_tokens") + tracked_input_cost_usd = cost_stats.get("total_input_cost_usd") + if tracked_input_tokens is not None: try: total_input_tokens = self._savings_tracker_input_tokens_offset + max( @@ -467,6 +467,14 @@ class PrometheusMetrics: self.requests_by_stack[slug] += 1 self.savings_tracker.record_lifetime_stack(slug) + # Same fan-out as record_compression. This header is the only signal + # that names the harness when an agent is pointed at a persistent proxy + # rather than launched by `headroom wrap`, and the beacon cannot import + # headroom.proxy to read requests_by_stack itself. + from headroom.telemetry.session import record_stack as _beacon_stack + + _beacon_stack(slug) + def record_compression( self, strategy: str, @@ -497,6 +505,22 @@ class PrometheusMetrics: if saved > 0: self.tokens_saved_by_strategy[strategy] += saved + # Fan out to the beacon. This object is the configured + # CompressionObserver for the proxy's pipelines, so it is where those + # events already arrive with both token counts — a second observer here + # would mean a second measurement pass for numbers in hand. (The paths + # that have no observer at all pass telemetry's + # BeaconCompressionObserver directly instead.) + # + # The beacon is ON by default, so this does not short-circuit in + # practice and must stay off the aggregator's lock: it stages into a + # dedicated mutex that the request path never takes, which is what + # keeps this method's "synchronous + lock-free" contract honest with + # respect to everything else in the process. + from headroom.telemetry.session import record_compression as _beacon_compression + + _beacon_compression(strategy, original_tokens, compressed_tokens) + def record_extension_savings(self, key: str, saved: int) -> None: """Accumulate tokens saved by a proxy extension, keyed by ``key``. diff --git a/headroom/telemetry/session.py b/headroom/telemetry/session.py index cfa77f44c..9aa287cb2 100644 --- a/headroom/telemetry/session.py +++ b/headroom/telemetry/session.py @@ -94,6 +94,39 @@ _SLUG_RE = re.compile(r"^[a-z][a-z0-9_]{0,31}$") # vocabulary. Values are slug-validated before they are counted. _REASON_TAGS = ("passthrough_reason", "image_skip_reason", "memory_skip_reason") +# Cardinality cap on `by_strategy`. The real vocabulary is CompressionStrategy +# plus a couple of literals — under a dozen — but `record_compression` takes a +# free string, so an extension or a future caller could invent keys per request. +# Matches the same guard on `requests_by_stack` (MAX_DISTINCT_STACKS). +MAX_STRATEGIES = 32 + +# Compression events arrive on the compression executor thread, mid-request, +# before that request's outcome ever reaches `SessionAggregator.record`. They +# are staged here rather than written straight into the live session, which +# keeps three things true at once: +# +# * The executor thread never takes the aggregator's lock, so compression +# cannot serialise against the request path. The beacon is on by default, +# and ContentRouter observes once per routing decision — once per content +# section per request — so that contention would be real. +# * A compression event cannot CREATE a session. Sessions are started only by +# an outcome, which preserves the invariant that every emitted session has +# turns >= 1; otherwise a request abandoned between compression and its +# outcome (Claude Code users interrupt streaming routinely) would emit a +# phantom all-zero row that inflates fleet session and install counts. +# * The first turn's numbers still survive, because the outcome that follows +# milliseconds later drains this into the session it opens. +# +# A request that dies before its outcome leaves its events staged, and they are +# attributed to the next session instead. That is a rounding error against +# inventing a session that never happened. +_staged_lock = threading.Lock() +_staged_strategies: dict[str, list[int]] = {} +# Per-request stack slugs, for `detect_stack`'s by_stack branch. Same staging +# and the same reason: the proxy sees the X-Headroom-Stack header per request, +# and this is the only place the beacon can learn it without importing proxy. +_staged_stacks: dict[str, int] = {} + def _pct(numerator: float, denominator: float) -> float: """Percentage to 2dp, or 0.0 when undefined. @@ -218,6 +251,28 @@ def resource_attributes( } if install_mode: attrs["headroom.install_mode"] = install_mode + # Detect when the caller did not supply one. Every caller so far supplies + # nothing, so `headroom.stack` was absent from the entire corpus while the + # detector sat unused — which made the fleet unsegmentable by agent, the + # question the corpus is most often asked ("what does this look like under + # Claude Code?"). + # + # The env vars detect_stack checks first are only set by `headroom wrap`. + # The common deployment points an agent at a persistent proxy through + # ANTHROPIC_BASE_URL and sets neither, so environment-only detection would + # answer the literal "proxy" for almost the whole fleet — a populated, + # authoritative-looking column that cannot answer the question it exists + # for. The slugs staged by `record_stack` are that fleet's only real + # signal, so they are fed to detect_stack's by_stack branch. + if stack is None: + try: + from headroom.telemetry.context import detect_stack + + with _staged_lock: + by_stack = dict(_staged_stacks) + stack = detect_stack({"requests": {"by_stack": by_stack}} if by_stack else None) + except Exception: # a broken detector must not silence telemetry + logger.debug("telemetry: stack detection failed", exc_info=True) if stack: attrs["headroom.stack"] = stack return attrs @@ -252,6 +307,10 @@ class _Session: overhead_ms: float = 0.0 latency_ms: float = 0.0 transforms: dict[str, int] = field(default_factory=dict) + # strategy slug -> [events, tokens_in, tokens_out]. `transforms` says which + # compressors ran; this says whether they were worth running. A list rather + # than three parallel dicts so the three numbers cannot drift apart. + strategies: dict[str, list[int]] = field(default_factory=dict) skips: dict[str, int] = field(default_factory=dict) sources: dict[str, int] = field(default_factory=dict) providers: set[str] = field(default_factory=set) @@ -363,6 +422,37 @@ class _Session: }, "compression": { "transforms": dict(self.transforms), + # Per-strategy effectiveness. `transforms` counts invocations, + # which cannot tell a compressor that saved 60% from one that + # ran constantly and saved nothing — the fleet's top transform + # by count contributes an unknown share of `tokens.saved`. + # + # These do NOT sum to `tokens.saved`, and must not be presented + # as if they do: strategies compose (the router routes, a + # strategy runs inside it) so the same text is measured by more + # than one, and tool-schema savings never appear here at all. + # Read a row as "of what this strategy was handed, it removed + # this much" — a per-strategy yield, not a share of the total. + # + # A LIST of uniform records, not a {strategy: {...}} object, + # and that shape is deliberate. DuckDB infers a JSON object as + # a STRUCT while its keys are few and consistent and as a MAP + # once they are not, so an object keyed by strategy would + # change COLUMN TYPE as the fleet adopts new compressors — + # exactly the break that silently took out the `transforms` + # report. A list of records has fixed field names, so the type + # is the same on day one and after the 30th strategy ships, and + # a new field inside a record is absorbed by union_by_name. + # Sorted so a payload is byte-comparable between heartbeats. + "by_strategy": [ + { + "strategy": name, + "n": counts[0], + "tokens_in": counts[1], + "tokens_out": counts[2], + } + for name, counts in sorted(self.strategies.items()) + ], "overhead_ms_total": round(self.overhead_ms, 1), # Sum of per-request durations, NOT elapsed time: concurrent turns # make this exceed `session.duration_s`. Kept under the original @@ -505,6 +595,20 @@ def _fold(sess: _Session, outcome: Any, now: float, source: str = "proxy") -> No sess.last_seen = now sess.turns += 1 sess.sources[source] = sess.sources.get(source, 0) + 1 + + # Compression ran on the executor thread before this outcome arrived; take + # what it staged. Done here rather than in the observer so the executor + # thread never touches the aggregator lock — see the note on _staged_lock. + for name, staged in _drain_staged_strategies().items(): + counts = sess.strategies.get(name) + if counts is None: + if len(sess.strategies) >= MAX_STRATEGIES: + continue + counts = [0, 0, 0] + sess.strategies[name] = counts + counts[0] += staged[0] + counts[1] += staged[1] + counts[2] += staged[2] sess.original_tokens += int(get("original_tokens") or 0) sess.attempted_tokens += int(get("attempted_input_tokens") or 0) # Billed/volume figure, so prefer the provider's own count and fall back to @@ -776,6 +880,116 @@ def record_mcp_compression( ) +def record_compression(strategy: str, original_tokens: int, compressed_tokens: int) -> None: + """Beacon entry point for one compression event. + + Signature-compatible with + :class:`headroom.transforms.observability.CompressionObserver`, so the + proxy's existing observer can forward here without a second measurement + pass — the numbers are already computed on the hot path for Prometheus + (``PrometheusMetrics.tokens_saved_by_strategy``); they just never left the + process. + + Same discipline as the rest of this module: off by default and cheap when + off, never raises. This runs once per routing decision, so it must not do + anything a request would notice. + """ + from headroom.telemetry.beacon import is_beacon_enabled + + if not is_beacon_enabled(): + return + # A slug, not the raw string. The real values are CompressionStrategy enum + # tags, but the observer protocol takes a free string, and anything that is + # not already a bounded lowercase identifier collapses to "other" rather + # than reaching the wire. + slug = _safe_slug(strategy) + try: + before = int(original_tokens or 0) + after = int(compressed_tokens or 0) + except (TypeError, ValueError): + return + if before <= 0: + return + # Clamped at the input: a compressor that emits more than it received is a + # bug, and letting `out` exceed `in` would surface downstream as negative + # savings rather than as the bug it is. Prometheus clamps the same way. + after = min(max(after, 0), before) + with _staged_lock: + counts = _staged_strategies.get(slug) + if counts is None: + if len(_staged_strategies) >= MAX_STRATEGIES: + return + counts = [0, 0, 0] + _staged_strategies[slug] = counts + counts[0] += 1 + counts[1] += before + counts[2] += after + + +def record_stack(slug: str) -> None: + """Beacon entry point for one request's stack slug. + + The harness identity lives in the ``X-Headroom-Stack`` header, which only + the proxy sees, and per request rather than per process. Counting slugs + here lets :func:`resource_attributes` answer ``detect_stack``'s by_stack + branch without the telemetry package importing ``headroom.proxy``. + + Without this the beacon can only read the two environment variables, so + every install that points an agent at a persistent proxy — the common + deployment for Claude Code, Cursor, Codex and the adapters — reports the + literal ``"proxy"`` and the fleet is unsegmentable by agent. + """ + from headroom.telemetry.beacon import is_beacon_enabled + + if not is_beacon_enabled(): + return + # normalize_stack is the same chokepoint the proxy applies at ingress; an + # unbounded header value must not reach the wire or grow this dict. + from headroom.telemetry.context import normalize_stack + + clean = normalize_stack(slug) + if not clean: + return + with _staged_lock: + if clean not in _staged_stacks and len(_staged_stacks) >= MAX_STRATEGIES: + return + _staged_stacks[clean] = _staged_stacks.get(clean, 0) + 1 + + +class BeaconCompressionObserver: + """A `CompressionObserver` that forwards to the beacon and nothing else. + + The proxy's `PrometheusMetrics` is already an observer and forwards from + there, so this is for the paths that never had one: the MCP servers, the + bare transform pipeline, and the LangChain/Strands integrations. Those + processes report `tokens.saved` either way, so without this they emit + sessions with real token totals and an empty `by_strategy` — a silently + biased subset that cannot be reconciled with the fleet totals. + + Only `record_compression` is implemented. ContentRouter's two other + observer hooks (`record_kompress_size_gate`, `record_router_route_counts`) + are each individually guarded at the call site, and both feed `/stats` + rather than the beacon. + """ + + __slots__ = () + + def record_compression( + self, strategy: str, original_tokens: int, compressed_tokens: int + ) -> None: + record_compression(strategy, original_tokens, compressed_tokens) + + +def _drain_staged_strategies() -> dict[str, list[int]]: + """Take everything staged since the last drain. Caller merges it.""" + with _staged_lock: + if not _staged_strategies: + return {} + drained = {name: counts[:] for name, counts in _staged_strategies.items()} + _staged_strategies.clear() + return drained + + def record_outcome(outcome: Any) -> None: """Beacon entry point, called from the proxy's outcome funnel. @@ -874,6 +1088,97 @@ def demo() -> None: agg.flush_all() assert len(emitted) == 2, emitted assert emitted[1]["session"]["turns"] == 1 + + # --- per-strategy compression ----------------------------------------- + # Compression runs on the executor thread before its request's outcome + # arrives, so events are staged and drained by the next outcome. That is + # what keeps the first turn's numbers while letting only an outcome open a + # session. `_staged_*` is module state, so clear it between cases. + _staged_strategies.clear() + _staged_stacks.clear() + + strat: list[dict[str, Any]] = [] + sa = SessionAggregator(strat.append, idle_s=10.0) + record_compression("smart_crusher", 1000, 400) + record_compression("smart_crusher", 500, 300) + record_compression("code_aware", 800, 800) + assert sa._current is None, "a compression event must not open a session" + sa.record(FakeOutcome(), now=2000.0) + sa.flush_all() + by = {row["strategy"]: row for row in strat[-1]["compression"]["by_strategy"]} + assert by["smart_crusher"] == { + "strategy": "smart_crusher", + "n": 2, + "tokens_in": 1500, + "tokens_out": 700, + }, by + # A strategy that ran and saved nothing must still appear: "ran 800 tokens + # through and removed none" is the finding, and dropping it would make + # every strategy look effective. + assert by["code_aware"]["tokens_in"] == by["code_aware"]["tokens_out"] == 800, by + assert strat[-1]["session"]["turns"] == 1, "compression events are not turns" + # A list of records, not an object keyed by strategy: the type must not + # change as strategies are added. See the note in payload(). + assert isinstance(strat[-1]["compression"]["by_strategy"], list) + assert [r["strategy"] for r in strat[-1]["compression"]["by_strategy"]] == sorted( + r["strategy"] for r in strat[-1]["compression"]["by_strategy"] + ), "sorted so heartbeats are byte-comparable" + + # Draining is exhaustive: a second session must not re-count the first + # session's events. + assert not _staged_strategies, "record() drains everything it staged" + again: list[dict[str, Any]] = [] + sb = SessionAggregator(again.append, idle_s=10.0) + sb.record(FakeOutcome(), now=3000.0) + sb.flush_all() + assert again[-1]["compression"]["by_strategy"] == [], again[-1] + + # An abandoned request — compression ran, the outcome never arrived — must + # not invent a session. Before staging, this emitted a phantom turns=0 row + # with all-zero tokens that inflated fleet session and install counts. + ghost: list[dict[str, Any]] = [] + sc = SessionAggregator(ghost.append, idle_s=10.0) + record_compression("smart_crusher", 900, 100) + sc.flush_all() + assert ghost == [], "no outcome, no session" + _staged_strategies.clear() + + # Over the cardinality cap, extra strategies are dropped rather than + # allowed to grow the payload without bound. + cap: list[dict[str, Any]] = [] + cc = SessionAggregator(cap.append, idle_s=10.0) + for i in range(MAX_STRATEGIES + 5): + record_compression(f"s{i}", 100, 50) + cc.record(FakeOutcome(), now=4000.0) + cc.flush_all() + assert len(cap[-1]["compression"]["by_strategy"]) == MAX_STRATEGIES, cap[-1] + _staged_strategies.clear() + + # Strategy names are slugged, never passed through: the observer protocol + # takes a free string and this is the only chokepoint before the wire. + assert _safe_slug("smart_crusher") == "smart_crusher" + assert _safe_slug("../../etc/passwd") == "other" + + # --- stack detection --------------------------------------------------- + # Environment-only detection answers "proxy" for every install that points + # an agent at a persistent proxy instead of using `headroom wrap` — i.e. + # most of the fleet. The per-request slugs are the only real signal. + _staged_stacks.clear() + for _ in range(9): + record_stack("wrap_claude") + record_stack("wrap_cursor") + assert resource_attributes()["headroom.stack"] == "wrap_claude", "dominant stack wins" + _staged_stacks.clear() + for _ in range(5): + record_stack("wrap_claude") + for _ in range(5): + record_stack("wrap_cursor") + assert resource_attributes()["headroom.stack"] == "mixed", "no dominant stack" + _staged_stacks.clear() + record_stack("../../etc/passwd") + assert not _staged_stacks, "junk slugs never reach the wire" + assert resource_attributes()["headroom.stack"] == "proxy", "falls back with no signal" + assert emitted[1]["session"]["id"] != emitted[0]["session"]["id"] assert emitted[1]["session"]["ended"] == "shutdown" diff --git a/headroom/transforms/pipeline.py b/headroom/transforms/pipeline.py index c7061dee1..769580d42 100644 --- a/headroom/transforms/pipeline.py +++ b/headroom/transforms/pipeline.py @@ -163,7 +163,13 @@ class TransformPipeline: # - Logs -> LogCompressor # - Search results -> SearchCompressor # - HTML -> HTMLExtractor - transforms.append(ContentRouter()) + # observer: the proxy passes PrometheusMetrics; this bare pipeline is + # used by the library/adapter paths, which would otherwise report + # tokens.saved with an empty by_strategy. Imported here rather than at + # module scope — transforms sits below telemetry in the import graph. + from headroom.telemetry.session import BeaconCompressionObserver + + transforms.append(ContentRouter(observer=BeaconCompressionObserver())) logger.info("Pipeline using ContentRouter for intelligent content-aware compression") return transforms