Skip to content

feat: add StreamReservation context manager for streaming DX#39

Merged
amavashev merged 2 commits intomainfrom
feat/streaming-context-manager
Apr 8, 2026
Merged

feat: add StreamReservation context manager for streaming DX#39
amavashev merged 2 commits intomainfrom
feat/streaming-context-manager

Conversation

@amavashev
Copy link
Copy Markdown
Contributor

Summary

  • Add StreamReservation and AsyncStreamReservation context managers that automate the reserve → commit/release lifecycle for streaming use cases, reducing boilerplate from ~50 lines to ~15
  • Add StreamUsage dataclass for accumulating tokens/cost during streaming
  • Add CyclesClient.stream_reservation() / AsyncCyclesClient.stream_reservation() convenience factory methods
  • Bump version 0.2.0 → 0.3.0

Before (~50 lines manual boilerplate)

# Manual: create reservation, extract ID, try/except stream,
# build commit request, handle release on error, manage idempotency keys...

After (~15 lines)

with client.stream_reservation(
    action=Action(kind="llm.completion", name="gpt-4o"),
    estimate=Amount(unit=Unit.USD_MICROCENTS, amount=max_tokens * 1000),
    cost_fn=lambda u: u.tokens_input * 250 + u.tokens_output * 1000,
) as reservation:
    for chunk in stream:
        if chunk.usage:
            reservation.usage.tokens_input = chunk.usage.prompt_tokens
            reservation.usage.tokens_output = chunk.usage.completion_tokens
# Auto-committed on success, auto-released on exception

Key details

  • Cost resolution: explicit usage.actual_cost > cost_fn(usage) > estimate fallback
  • Heartbeat: automatic TTL extension, same formula as decorator lifecycle
  • Commit retry: uses existing CommitRetryEngine / AsyncCommitRetryEngine
  • Context propagation: sets/clears CyclesContext via ContextVar; respects user-set ctx.metrics
  • Spec validation: validate_ttl_ms, validate_grace_period_ms, validate_subject — matches lifecycle.py
  • Error handling: RESERVATION_FINALIZED, RESERVATION_EXPIRED, IDEMPOTENCY_MISMATCH do not trigger release — matches lifecycle.py exactly

Files changed

File Change
runcycles/streaming.py New — core module
runcycles/client.py Add stream_reservation() to both clients
runcycles/__init__.py Export new public symbols
tests/test_streaming.py New — 64 tests
examples/streaming_usage.py Rewritten to use context manager API
README.md Add streaming section with complete example
AUDIT.md Document streaming module
examples/README.md Update description
pyproject.toml Version 0.3.0

Test plan

  • pytest --cov=runcycles — 364 passed, 99.38% coverage
  • ruff check — all clean
  • ruff format --check — all clean
  • mypy — no errors
  • 64 streaming-specific tests covering: success, deny, error codes, retry, heartbeat, cost resolution, context propagation, spec validation, IDEMPOTENCY_MISMATCH

Add StreamReservation and AsyncStreamReservation context managers that
automate the reserve → commit/release lifecycle for streaming use cases,
reducing boilerplate from ~50 lines to ~15.

- StreamUsage dataclass for accumulating tokens/cost during streaming
- Auto-commit on successful exit, auto-release on exception
- Heartbeat-based TTL extension for long-running streams
- Commit retry via existing CommitRetryEngine
- Cost resolution: explicit actual_cost > cost_fn > estimate fallback
- Respects user-set ctx.metrics during streaming
- Full spec validation (TTL, grace_period, subject constraints)
- Handles IDEMPOTENCY_MISMATCH correctly (no release)
- Client convenience: CyclesClient.stream_reservation()
- 64 tests, 97% module coverage, 99.38% total coverage
- Version bump: 0.2.0 → 0.3.0
- Fix weak async cost_fn test: use distinct value (1500) vs estimate
  (1000) so test fails if cost_fn is ignored
- Fix README streaming example: add missing variable definitions
  (max_tokens, openai_client, stream creation) so example is runnable
- Update AUDIT.md: correct test count (64), add validation and
  IDEMPOTENCY_MISMATCH details
- Update examples/README.md: streaming_usage.py description
@amavashev amavashev merged commit f29fd3a into main Apr 8, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant