Skip to content

Close the streaming feature gaps: handle storage and mid-stream budget accounting #187

Description

@dgenio

Summary

Bring the streaming path to parity with the standard invoke path: store a Handle for
streamed results (so they can be expanded later) and charge the BudgetManager as
chunks flow (so a stream cannot exceed the budget it reserved).

Why this matters

The module's own docstring acknowledges both gaps. Without a Handle, streamed data is
ephemeral — the "summary now, full data via handle later" contract silently does not
apply to streams. Without mid-stream charging, a long stream consumes context far
beyond its reservation, undermining the budget guarantees the firewall provides on
the buffered path. As streaming becomes the dominant agent UX, parity stops being
optional.

Current evidence

  • kernel/_stream.py:183-185 (docstring): "The effective response mode is resolved once up front — no budget is [re-]escalat[ed]… honouring apply_stream's stateless contract."
  • kernel/_stream.py:100: handle: Handle | None = None — the accumulator exists but no handle is stored on the true streaming path.
  • firewall/budget_manager.py supports reserve/release but is not consulted per chunk.

External context

Not required for this issue.

Proposed implementation

  1. Accumulate streamed rows/chunks (bounded by existing budgets) and store a Handle
    on stream completion via kernel._handles.store(...), attaching it to the final
    chunk's Frame; respect HandleStore caps.
  2. Add per-chunk budget charging: estimate chunk cost (reuse the token-counting
    seam, firewall/token_counting.py), decrement, and on exhaustion truncate the
    stream with a stable warning code and a trace note.
  3. Coordinate with the cross-chunk redaction buffer (ISSUE 3) since both add bounded
    state to the stream loop.
  4. Record final totals in the single ActionTrace the stream already produces.

AI-agent execution notes

  • Inspect first: kernel/_stream.py (whole flow), firewall/transform.py (apply_stream), firewall/budget_manager.py, handles.py, tests/test_firewall_stream.py.
  • Determinism and the documented "single ActionTrace per stream" behavior must hold.
  • Edge cases: consumer abandons the stream mid-way (handle should reflect data received; budget released); empty stream; exhaustion exactly at a chunk boundary.
  • Do not change non-streaming paths.

Acceptance criteria

  • A completed stream yields a Frame with a usable Handle; expanding it returns the streamed rows under normal handle constraints.
  • A stream exceeding its budget is truncated with a stable warning and an accurate trace.
  • Abandoned streams release reservations (no leak).

Test plan

Extend tests/test_firewall_stream.py: completion handle round-trip, exhaustion
truncation, abandonment cleanup. Run make ci.

Documentation plan

Update streaming sections in docs/architecture.md/docs/context_firewall.md
(remove the documented limitation); CHANGELOG Added/Fixed.

Migration and compatibility notes

Streams gain a handle on the final chunk (additive field already exists). Budget
truncation is new behavior under exhaustion — previously the stream continued;
document as a fix toward the stated budget contract.

Risks and tradeoffs

Buffering streamed rows for the handle costs memory — bound it with the existing
budget machinery and document the cap. Chunk cost estimation is approximate; err
conservative.

Suggested labels

product, reliability, architecture

Metadata

Metadata

Assignees

No one assigned

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions