Skip to content

feat(plugins): carry token usage on the on_stream_end observer payload - #90077

Open
CocaKova wants to merge 1 commit into
NousResearch:mainfrom
CocaKova:feat/stream-end-usage
Open

feat(plugins): carry token usage on the on_stream_end observer payload#90077
CocaKova wants to merge 1 commit into
NousResearch:mainfrom
CocaKova:feat/stream-end-usage

Conversation

@CocaKova

Copy link
Copy Markdown
Contributor

What

on_stream_end observer callbacks now receive a usage kwarg alongside final_text/finished/error — the turn's token counters in turn_finalizer's existing field vocabulary:

usage = {
  prompt_tokens, completion_tokens, total_tokens, reasoning_tokens, api_calls,
  last_prompt_tokens,   # final API call's prompt size = current context occupancy
  context_length,       # the window that occupancy sits in
}

Why

Stream observer hooks (added with the on_stream_start/delta/end family) let a plugin watch every token of a turn, but none of the numbers behind it — so anything usage-shaped (a context meter in a client, a token dashboard) still requires core patches, which the plugin policy forbids. last_prompt_tokens + context_length are exactly a context gauge; the session counters come along at zero extra cost since they're already on the agent.

Real consumer exists: a standalone gateway side-channel plugin that renders a per-turn context ring in a Matrix client. No behavior change for existing callbacks — the payload is additive, observer-only, and every read is a defensive getattr with a zero default (a partially initialized agent yields zeroed usage, never a raise).

Changes

  • run_agent.py: _stream_hook_usage_payload() + wired into _emit_stream_end
  • hermes_cli/plugins.py: payload documented in the VALID_HOOKS comment block
  • Tests: 2 added, 14 passing in tests/run_agent/test_plugin_stream_hooks.py

@CocaKova

Copy link
Copy Markdown
Contributor Author

The Python tests / slice 12 failure is pre-existing on main, not from this PR:

  • Failing test: tests/tools/test_image_generation.py::TestFalCatalog::test_upscale_defaults_are_all_off
  • This PR's diff is 3 files, 79 insertions, all in the stream-observer hook path (run_agent.py, hermes_cli/plugins.py, tests/run_agent/test_plugin_stream_hooks.py) — the image catalog and its tests are byte-identical to main here.
  • Reproduces locally on bare 13ce0c5c67 (current main): the xai/grok-imagine-image/v2.0 entry added in ceabb030fb ships "upscale": True, violating the all-off invariant f06c41522e established.
  • One-line fix opened as fix(image-gen): grok-imagine-image 2.0 must not default upscale on #90080; happy to rebase this branch once it lands.

Hook tests on this branch: 14/14 passing (tests/run_agent/test_plugin_stream_hooks.py).

@alt-glitch alt-glitch added type/feature New feature or request P3 Low — cosmetic, nice to have comp/agent Core agent runtime: loop, agent_init, prompt builder, context-compression, responses endpoint comp/plugins Plugin system and bundled plugins labels Aug 19, 2026
Stream observer plugins can see every token of a turn but none of the
numbers behind it, so anything usage-shaped (a context meter, a token
dashboard) still needs core patches. on_stream_end already carries the
turn's lifecycle (final_text/finished/error); this adds a usage dict in
turn_finalizer's existing field vocabulary: prompt/completion/total/
reasoning token counters, api_calls, plus last_prompt_tokens (the final
call's prompt size, i.e. the model's current context occupancy) and
context_length. Together the last two are exactly a context gauge.

All reads are defensive getattr with zero defaults, so a partially
initialized agent yields zeroed usage instead of breaking the hook path.
Payload documented in the VALID_HOOKS comment block. Consumed by a
standalone gateway-side-channel plugin rendering a context ring.

Tests: 2 added, 14 passing in tests/run_agent/test_plugin_stream_hooks.py
@CocaKova
CocaKova force-pushed the feat/stream-end-usage branch from 8d5c92a to 95c528a Compare August 19, 2026 23:57
@CocaKova

Copy link
Copy Markdown
Contributor Author

Correction to my comment above, and a rebase.

The slice-12 failure is resolved — but by #89933 (merged), not by #90080, which I opened for the same one-line change before finding it was a duplicate. #90080 is now closed; #89933 is the fix. Sorry for the misdirection.

Rebased onto current main (7b25941b0e), which includes that fix. The branch was 83 commits behind; it's now zero, and the diff is unchanged — 3 files, all in the stream-observer hook path. tests/run_agent/test_plugin_stream_hooks.py is 14 passing on the rebased branch, and test_upscale_defaults_are_all_off passes on this base.

Also relevant: #90310 adds the delivery target (chat_id/chat_type/thread_id) to the shared _stream_hook_base_payload. It touches the same file and test module as this one, so whichever lands second wants a trivial rebase. They're independent — usage on stream end vs. chat identity on every observer payload — and together they're what lets the consuming side-channel plugin stop being a set of core patches.

@Enough1122

Copy link
Copy Markdown
Contributor

AI code review — automated review for reference, author can ignore or act on any point.

Additive and contract-clean: usage rides on_stream_end as a new kwarg in turn_finalizer's existing field vocabulary, the payload is snapshotted synchronously at emit time (so async dispatch observes immutable values even as counters advance), and defensive getattr reads mean a partially initialized agent yields zeros instead of raising — with a test asserting precisely that.

Points:

  1. Pin the narrow-signature compatibility path: per the plugin compatibility contract, old-style callbacks declaring only (final_text, finished, error) must keep receiving just those fields while **kwargs observers get everything. If agent/plugin_stream_hooks signature-inspects before invoking, add one test with a narrow-signature callback proving no TypeError and no unexpected kwarg — that's the regression a future refactor of the dispatcher is most likely to introduce.
  2. Document the context_length == 0 sentinel ("unknown") for meter plugins in the hook docs — otherwise the first plugin will divide by it.
  3. Minor: _stream_hook_usage_payload reads six session_* attributes plus compressor state; if turn_finalizer already centralizes counter naming, importing its vocabulary (or a shared constant list) keeps the two from drifting when a new counter is added there.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

comp/agent Core agent runtime: loop, agent_init, prompt builder, context-compression, responses endpoint comp/plugins Plugin system and bundled plugins P3 Low — cosmetic, nice to have type/feature New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants