(MOT-4327) feat(harness): live context accounting with a console chip - #686
Conversation
…nd serve it through metrics Each generate step folds the assembly it already performed into a ContextSnapshotV1: category token estimates (system prompt, function schemas, messages by role, overhead, post-assembly hook growth), the usable budget, the final request total, compaction state, and the estimator that produced the numbers. No extra counting round trips. After the terminal frame the snapshot is stamped with the generation's actual provider usage and stored under harness_context/<session_id>. harness::turn-completed now carries the snapshot as context, and harness::metrics by_session entries return each session's latest snapshot beside the existing usage sums. context::assemble's new breakdown response feeds the categories; an older context-manager degrades to totals-only snapshots.
…hip and metrics renderer harness ships console assets: a session chip (id context) rendering the real per-turn context meter with a breakdown popover (stacked category bar, legend, estimator provenance, actual usage and cost), fed by harness::metrics on mount and harness::turn-completed live; and a function-trigger renderer for harness::metrics results in chat. The chip registers through the console's session chip slot and skips silently on consoles that predate it. Registration implements the console asset wire contract directly because the shared console-ui crate pins iii-sdk =0.21.6 while harness pins =0.21.8, which cargo cannot co-resolve; relaxing the crate pin is a follow-up outside this worker.
…rn and ticks mid-turn Before the first generation the chip now renders 0% of the model's catalog context window instead of a bare dash. And because snapshots are written after every generate step while the turn-completed push only fires at turn end, the chip additionally polls the session's harness_context state row every five seconds so long multi-step turns update live instead of staying frozen until the turn finishes.
…te trigger Drop the interval poll. The chip now binds the state worker's `state` trigger with an engine-side scope and key filter on the session's harness_context row, so every per-step snapshot write streams to the chip in real time. One state::get on mount hydrates; everything after is push.
…token counting When a provider token counter is available (router::count_tokens), the snapshot's categories become real numbers: the generation's billed usage is the exact total, the system prompt and tool schemas are counted through the provider's tokenizer with a probe-delta and cached per content hash (one wire count per session, not per step), and the message window is the remainder with the estimated by-role proportions rescaled onto it. The estimator field reports what produced the numbers; rigs without a counter keep the heuristic snapshot untouched. Counting never runs the model, never enters the session context, and costs nothing.
The popover foot renders exact plus the counting source when the snapshot numbers came from a provider tokenizer, keeping est. only for the chars heuristic fallback.
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe harness now records and persists context snapshots with provider token counts, exposes them through events and metrics, and renders context and metrics data in an embedded console UI with a pnpm/esbuild build pipeline. ChangesContext accounting and console UI
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant TurnLoop
participant RouterClient
participant HarnessState
participant TurnEvents
participant SessionMetrics
participant ConsoleUI
TurnLoop->>RouterClient: Count prompt and tool tokens
RouterClient-->>TurnLoop: Return token count and estimator
TurnLoop->>HarnessState: Store context snapshot
TurnLoop->>TurnEvents: Emit snapshot with turn event
TurnLoop->>SessionMetrics: Include snapshot in usage
ConsoleUI->>HarnessState: Hydrate session snapshot
HarnessState-->>ConsoleUI: Return context data
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
skill-check — worker0 verified, 54 skipped (no docs/).
Four for four. Nicely done. |
…lim the accounting paths Review cleanups, wire-identical (goldens unchanged). The console-ui crate's iii-sdk pin relaxes to a 0.21 range so harness can link it, deleting the 314-line wire-contract fork; harness keeps a small stub naming its two assets. exactify loses the probe-per-part duplication (base counted once and cached, the two category deltas run concurrently) and the empty-string sentinel; a rig without a counter now keeps its heuristic snapshot instead of stamping exact totals. The snapshot build reads the assemble breakdown through Default and a From impl, Assembled carries Applied whole, and metrics reads the snapshot off the turn record it already fetched instead of a second state round trip per session. Chip UI: the turn-completed subscription goes away (mount get plus the state trigger already cover every write, earlier and finer), snapshot category types match what the wire always sends, one categories() list feeds both the stacked bar and the legend, and the tone thresholds live once in lib/tone shared with the metrics renderer.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@harness/src/context_snapshot.rs`:
- Around line 165-171: Update cache_key and its callers in the snapshot
token-counting flow to include the resolved provider identity alongside model,
kind, and content. Ensure both calls to router::count_tokens use provider-scoped
keys, and when the provider is None, resolve a stable default before key
creation or bypass caching if that identity is mutable.
- Around line 48-55: Update the scaling logic around the apply closure in the
snapshot scaling method to distribute any residual tokens after independently
flooring the category values, so the resulting category total exactly equals the
requested numerator. Preserve proportional scaling as closely as possible, and
add a test asserting that the scaled category total matches snapshot.total.
In `@harness/src/functions/send.rs`:
- Line 690: Preserve the latest per-session context snapshot when replacing or
creating a turn record: in harness/src/functions/send.rs lines 690-690,
initialize the new TurnRecord from the persisted session snapshot or retain the
prior snapshot; in harness/src/functions/metrics.rs lines 162-164, make the
metrics path fall back to that persisted session snapshot when the current
record’s context_snapshot is absent.
In `@harness/src/turn_loop.rs`:
- Around line 2205-2213: Rebuild the token-category breakdown from the final,
post-hook request before constructing SnapshotCategoriesV1, rather than reusing
b.system_prompt_tokens and the pre-mutation counts. Ensure the heuristic
fallback remains available when provider exactification is unavailable, and make
the category values sum consistently with final_request_tokens, including
hook_guidance.
- Around line 510-517: Update the context snapshot construction around
build_context_snapshot so the final request total excludes reservation-only
extra_overhead_tokens, including when the hook leaves the request unchanged. Use
the actual final request token count rather than assembled.token_count plus
reserved headroom, while preserving the existing snapshot assignment flow.
In `@harness/ui/src/context-chip/index.tsx`:
- Around line 279-294: Update the hydration logic in the useEffect keyed by host
and sessionId to compare the fetched snapshot’s timestamp with the current
snapshot before calling setSnapshot. Accept the hydrated value only when it is
not older than the already-streamed snapshot, while preserving the
cancelled/session_id guards and existing state reset behavior.
In `@harness/ui/src/lib/metrics.ts`:
- Around line 107-117: Update isSnapshot to validate that
snap.categories.messages exists and is a non-array object before accepting the
value as a ContextSnapshot. Preserve the existing total, usable, and categories
checks so malformed records are rejected and rendered as empty data.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: e727b737-44c2-41c9-b3a8-acae32fc40c7
⛔ Files ignored due to path filters (2)
harness/Cargo.lockis excluded by!**/*.lockpnpm-lock.yamlis excluded by!**/pnpm-lock.yaml
📒 Files selected for processing (27)
crates/console-ui/Cargo.tomlharness/Cargo.tomlharness/build.rsharness/src/clients/context.rsharness/src/clients/router.rsharness/src/context_snapshot.rsharness/src/events.rsharness/src/functions/metrics.rsharness/src/functions/send.rsharness/src/lib.rsharness/src/main.rsharness/src/subagent.rsharness/src/turn_loop.rsharness/src/types/turn.rsharness/src/ui.rsharness/tests/golden/schemas/harness.metrics.jsonharness/ui/build.mjsharness/ui/package.jsonharness/ui/page.tsxharness/ui/src/context-chip/index.tsxharness/ui/src/function-trigger-message/index.tsxharness/ui/src/lib/format.tsharness/ui/src/lib/metrics.tsharness/ui/src/lib/tone.tsharness/ui/styles.cssharness/ui/tsconfig.jsonpnpm-workspace.yaml
…snapshot # Conflicts: # harness/src/events.rs # harness/src/main.rs # harness/src/subagent.rs # harness/src/turn_loop.rs
CodeRabbit findings. The count cache key now includes the provider so two providers serving one model id cannot share entries. A freshly seeded turn no longer hides the session's snapshot from harness::metrics: the record copy is preferred and the durable harness_context row is the fallback. An unchanged request subtracts the unused pre-generate hook reservation from its total instead of reporting reserved headroom as consumed. Scaled by-role buckets absorb the flooring remainder into the largest bucket so categories keep summing to the exact total. The chip keeps whichever snapshot is newest when the hydration read races the streamed trigger, and isSnapshot validates the messages object it later reads.
eval path-depends on harness, whose exact 0.21.8 pins can never co-resolve with eval's 0.21.6 ones in a single graph; the eval CI lane trips on it for any harness-touching change. Pins aligned, the SendOptions literal gains the max_validation_retries field harness grew, and the eval goldens pick up the context snapshot schema that now flows through the embedded harness metrics response.
…t field The integration probe's LifecycleEventV1 is an exact wire shape and the scenario floor pins the payload key set, so every turn-completed event carrying the new context snapshot was rejected and scenarios timed out waiting for a terminal status they had already received. Both contracts learn the optional context field; the snapshot's consumer-facing shape stays pinned by the harness::metrics schema golden. Full scenario suite green locally against the installed engine.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
eval/Cargo.toml (1)
19-20: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueAlign the eval dependency pins to the workspace SDK policy.
evalandharnessnow both useiii-sdk/iii-helpers0.21.8, but the rest of the workspace pins these crates to 0.21.6. Shared SDK types and the sharedopentelemetryversion used by the SDK can still split across the workspace unless the root lock resolves a single SDK version.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@eval/Cargo.toml` around lines 19 - 20, Update the iii-sdk and iii-helpers dependency pins in eval/Cargo.toml to match the workspace policy at version 0.21.6, aligning them with the rest of the workspace so the root lock resolves one shared SDK and opentelemetry dependency version.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@eval/Cargo.toml`:
- Around line 19-20: Update the iii-sdk and iii-helpers dependency pins in
eval/Cargo.toml to match the workspace policy at version 0.21.6, aligning them
with the rest of the workspace so the root lock resolves one shared SDK and
opentelemetry dependency version.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 6ce9619e-772f-4586-8218-78588fb4dc50
⛔ Files ignored due to path filters (1)
eval/Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (8)
eval/Cargo.tomleval/src/runtime.rseval/tests/golden/schemas/eval.assert.exact.jsoneval/tests/golden/schemas/eval.assert.normalized_text.jsoneval/tests/golden/schemas/eval.result.jsonharness/tests/integration/src/probe.rsharness/tests/integration/src/scenario/floor.rsharness/tests/integration/src/types/probe.rs
… in the context panel The panel summed cached and fresh input into one actual figure, which hid the number that drives cost on a long session. The footer now reports fresh input separately and adds a cache line with tokens read, tokens written, and the cached share of the prompt. The line is absent when the provider reported no cache activity.
The provenance, actuals, cache and cost lines sat in the ghost tint used for decoration, which made the numbers hard to read. They are findings, so they move to normal ink, and the cache hit rate carries a verdict colour: green above seventy percent, warn below thirty, plain between. A cold prefix re-bills the whole prompt at the write premium every turn, so it is worth flagging rather than dimming.
| Ok(serde_json::from_value(v).ok()) | ||
| } | ||
|
|
||
| pub async fn delete( |
There was a problem hiding this comment.
delete() has no callers in this PR. on_session_deleted::handle purges filesystem grants and budget state, but not the new harness_context/<session_id> row, so deleted sessions leak their snapshot rows in the state worker. Should this be wired into harness::on-session-deleted alongside the other purges (or dropped, if cleanup is deliberately deferred to a follow-up)?
There was a problem hiding this comment.
Good catch, it was dead code. Wired into on-session-deleted in 085fac7, right next to the grants and budget purges.
| let estimator = counted_system | ||
| .estimator | ||
| .or(counted_tools.estimator) | ||
| .unwrap_or_else(|| "provider".into()); |
There was a problem hiding this comment.
Cache hits lose the estimator provenance: counted_delta returns estimator: None when served from cache, so once both parts are cached (every step after the first with a stable prompt and tools) this falls back to "provider". On a tiktoken-backed rig — the Codex example in the PR description — step 2+ snapshots would read exact · provider tokenizer in the popover footer. Consider caching the estimator string alongside the token count (e.g. store (u64, String) in the cache).
There was a problem hiding this comment.
Yep. The cache holds (tokens, estimator) now, so a hit keeps the provenance. Probe base stores None deliberately, its estimator never reaches a snapshot anyway.
| snapshot.categories.messages = messages; | ||
| snapshot.categories.overhead = 0; | ||
| snapshot.categories.hook_guidance = 0; | ||
| snapshot.total = billed; |
There was a problem hiding this comment.
billed can exceed usable (actual billed usage vs a budget derived from the heuristic estimate — e.g. the 32% system-prompt underestimate quoted in the PR description), at which point total > usable contradicts the documented invariant on ContextSnapshotV1 ("total <= usable always holds for a generation that ran"), which also ships in the schema goldens' descriptions. free saturates and the UI clamps, so nothing breaks today, but schema consumers may rely on the stated invariant. Suggest softening the doc (or clamping total to usable).
There was a problem hiding this comment.
Right, billed can exceed usable once usage lands. Went with softening the doc rather than clamping, since clamping would hide the overage and that is the one number worth seeing. Golden regenerated.
| let context = match turn.context_snapshot.clone() { | ||
| Some(snapshot) => Some(snapshot), | ||
| None => { | ||
| crate::context_snapshot::get(&deps.iii, &node.session_id, cfg.session_timeout_ms) |
There was a problem hiding this comment.
This fallback runs one awaited state::get per session, serially, inside the tree walk. Any session whose latest turn record carries no snapshot pays the round trip — which today is every pre-existing session and every freshly seeded turn — so a large session tree adds sessions × RTT to each harness::metrics call, and partial snapshots are explicitly polled as progress signals (per the comment just below). Consider fetching these concurrently (join_all) or skipping the fallback on the polling/partial path.
There was a problem hiding this comment.
Fair. Took your second option: the indices are collected during the walk and read once for the final complete response, so polled partials pay nothing.
| // must never fail a turn that generated successfully. | ||
| if let Some(snapshot) = record.context_snapshot.as_mut() { | ||
| snapshot.usage = outcome.message.usage.clone(); | ||
| crate::context_snapshot::exactify(snapshot, &router, gen_system_prompt.as_deref(), &tools) |
There was a problem hiding this comment.
exactify sits on the generation hot path: on a cache miss it is up to 3 counting round trips (probe base, then the two part deltas), each bounded only by the router client timeout, before the turn can finalize and emit turn-completed. With a stable prompt/tools that's first-step-only, but any per-step system-prompt change re-counts, and a slow or flaky counting endpoint delays every turn's completion. Worth a tighter dedicated timeout for count_tokens, or a comment on why the shared router timeout is acceptable here.
| const offTrigger = host.iii.registerTrigger({ | ||
| type: 'state', | ||
| function_id: `${STATE_FN}::${host.iii.browserId}`, | ||
| config: { scope: 'harness_context', key: sessionId }, |
There was a problem hiding this comment.
The trigger's function_id is namespaced per tab (browserId) but not per session, and a handler/trigger pair is registered per mounted chip. If two chips are ever mounted at once (two sessions visible, or a mount/unmount race), both register the same function_id with different key configs — the second registration can clobber the first's engine-side filter, and one unmount's offTrigger() may tear down the other chip's stream. Including sessionId in the function_id would make the registrations independent.
There was a problem hiding this comment.
Fixed, the id carries the session now. One wrinkle worth noting: browserId has to stay the last segment, since on() namespaces the handler with it internally, so it ended up as STATE_FN::sessionId::browserId.
|
|
||
| let pnpm = locate_pnpm(); | ||
|
|
||
| let status = Command::new(&pnpm) |
There was a problem hiding this comment.
Heads-up: a plain cargo build/cargo test of the harness now hard-requires Node + pnpm (plus network for the first pnpm install), and the failure mode is a panic mid-build for someone touching only Rust. This matches the console worker's web/ precedent per the module doc, so mainly: please confirm the Rust CI lanes and dev-setup docs already guarantee pnpm is present.
There was a problem hiding this comment.
Checked, both lanes are already covered. ci.yml and release.yml key on /ui/package.json, and _rust-binary.yml pre-builds the bundle on Linux so the cross-compile shards never need a node toolchain. Same wiring as console, editor, database and pdf. The real gap is that nothing tells a rust-only contributor a harness build now needs pnpm, and that is equally true of the other four, so I left it out of this PR.
…ce and hot-path cost Six findings from review. `count_tokens` inherited the router's generation-sized timeout, which defaults to 320s. It runs inside `exactify`, on the path that finalizes a turn, so one slow counting endpoint could hold `turn-completed` open for as long as a whole generation. It now has its own 10s ceiling and still yields to a smaller configured timeout. The count cache stored a bare token count, so a cache hit reported no estimator and the snapshot fell back to naming the provider. Every step after the first with a stable prompt and tool set therefore mislabelled where its numbers came from. The estimator is cached alongside the count. `context_snapshot::delete` had no caller: `harness::on-session-deleted` purged filesystem grants and budget state but left the session's `harness_context` row behind. It is purged with the rest now. `ContextSnapshotV1` documented `total <= usable` as an invariant, and that sentence ships in the wire schema. Once provider usage lands, `total` is what was billed, which can exceed a budget derived beforehand from an estimate. The doc says that instead, and the golden is regenerated. `harness::metrics` read the durable snapshot inline for every session whose turn record carried none, adding a serial round trip per session to a walk that is polled as a progress signal. Those reads now happen once, for the final complete response. The context chip's state trigger was namespaced per tab but not per session, so two mounted chips registered conflicting engine-side filters under one id and either teardown took the other's stream with it. The id carries the session too.
`eval` depends on `harness` by path, so the `ContextSnapshotV1` doc rewrite in 085fac7 landed in eval's schema descriptions too.
…turn A snapshot is written after every generate step and carries that step's usage, so on a multi-step turn the row moved while the turn was still running: cost fell from one step to the next as the prompt cache warmed and more of the window was billed as cache reads. The numbers were right and the word above them was wrong.
2 participants
What
The harness accounts for every generation's context window and surfaces it live in the console.
Data plane:
ContextSnapshotV1: category token counts (system prompt, function schemas, messages by role, request overhead, post-assembly hook growth), the usable budget, the final request total, compaction state, and the estimator. Zero extra counting round trips; the categories come from context-manager's new assemblebreakdownand degrade to totals-only against an older context-manager.harness_context/<session_id>(one row per session, latest wins).router::count_tokens, (MOT-4329) feat(llm-router,providers): provider-backed token counting #685), the snapshot becomes exact: billed usage is the total, the system prompt and tool schemas are counted through the provider's tokenizer with a probe delta and cached per model, provider and content, and the message window is the remainder. Rigs without a counter keep the heuristic numbers, and theestimatorfield always says which produced them.harness::turn-completedcarries the snapshot ascontext;harness::metrics.by_session[]returns each session's latest snapshot, preferring the copy already on the turn record and falling back to the durable row for a freshly seeded turn.Console surface (injectable UI shipped by this worker):
contextsession chip in the chat header: real percent of the usable window, live after every step. Onestate::geton mount, then streaming via the state worker'sstatetrigger with an engine-side scope and key filter. No polling anywhere. Before the first turn it shows the model's catalog window.exact · provider tokenizerorest. heuristic), fresh input and output, prompt cache read, write and hit rate, and cost. The hit rate is coloured by verdict, since a cold prefix re-bills the whole prompt at the write premium every turn.harness::metricsresults in chat.Live verification (running rig)
Verified across three providers. Anthropic haiku: estimator
provider, total 13,072 equal to billed to the token, system prompt 12,230 exact where the heuristic had said 8,262 (32 percent under). Codex gpt-5.4-mini: estimatortiktoken, total 11,481 equal to billed. Chip, popover, and mid-turn streaming verified in the console against live sessions.The cache readout immediately earned itself: a three-turn probe showed a session pays two full cache writes before its first read, because the first generation's system prompt is 14 tokens larger than every turn after it. Filed as MOT-4347; not addressed here.
Testing
316 tests including the integration scenario suite, wire-schema goldens regenerated, cargo fmt and clippy with
-D warningsclean, UI bundle built and scope-asserted.Notes for review
crates/console-uirelaxes itsiii-sdkpin from an exact=0.21.6to a0.21range so harness can link it. That deleted a 314-line fork of the asset wire contract this PR originally carried, and each consumer still resolves its own exact pin in its own lockfile.evalfollows harness ontoiii-sdk0.21.8. It path-depends on harness, so the two exact pins could never co-resolve once this branch moved; the eval CI lane fails on any harness-touching change without it.LifecycleEventV1is an exact wire shape and the scenario floor pins the payload key set, so both learn the optionalcontextfield. Without that, every scenario timed out waiting for a terminal status it had already received./simplifyare in their own commits:exactifycounts the probe base once and runs the two category deltas concurrently,AssembledcarriesAppliedwhole, and the chip drops a redundantturn-completedsubscription that the state trigger already covers earlier and finer.Part of MOT-4324 (MOT-4327 and the harness half of MOT-4329).
Summary by CodeRabbit
New Features
Bug Fixes