Skip to content

feat(adapters,engine): mid-session token-budget guard with off/warn/enforce modes (#158) - #176

Merged
pbean merged 2 commits into
mainfrom
feat/mid-session-budget-guard-158
Jul 18, 2026
Merged

feat(adapters,engine): mid-session token-budget guard with off/warn/enforce modes (#158)#176
pbean merged 2 commits into
mainfrom
feat/mid-session-budget-guard-158

Conversation

@pbean

@pbean pbean commented Jul 17, 2026

Copy link
Copy Markdown
Collaborator

Closes #158.

What

A runaway-but-busy session (265 turns / 49.7M raw tokens in the reference incident) was bounded only by session_timeout_min: the stall detector fires on silence (pane output deliberately re-arms it), and max_tokens_per_story is advisory and evaluated only post-done. Nothing observed how much a running session was spending.

Both adapter wait loops (generic tmux + opencode-http) now sample cumulative session usage on the existing ~30s heartbeat cadence and compare the weighted total (cache_read_weight-discounted, same metric as the story cap) against a new per-session cap. Crossing it trips once per session, per limits.session_budget_mode:

  • warn (default) — one ATTENTION line + one budget-tripped lifecycle breadcrumb; the session is never interfered with. Session-end journal entries carry budget_weighted / budget / budget_mode.
  • enforce — warn actions + a wrap-up nudge, then session_budget_grace_s (default 240s, monotonic and wall-clock co-bound per the session-end journaled 2h19 after session_timeout_min fires when the session is wedged inside a tool call #157 lesson) to finish; a session that completes inside the grace is a normal completed, otherwise it terminates with the new over_budget status.
  • off — no sampling; behavior identical to today.

over_budget deliberately rides the existing != "completed" routing (retry with a fresh context → plateau-defer when attempts are exhausted) — zero escalation.py changes — and joins the post-kill-reconcile rescue set so an artifact flushed at kill time is still honored.

New policy keys

[limits]
session_budget_mode = "warn"     # off | warn | enforce
max_tokens_per_session = 4000000 # weighted; healthy sessions ran ~1-2.5M on the reference run
session_budget_grace_s = 240     # enforce: wrap-up window; 0 = terminate at trip, no nudge

Defaulting to warn: upgrades gain visibility (the #158 ask — see the 49.7M session as it happens) with no behavior change to session outcomes; enforce is the opt-in hard bound.

Design constraints held

Testing

  • 2400 passed, 1 pre-existing skip (live opencode binary); trunk clean.
  • Guard behavior matrix on the generic adapter (warn latch, enforce nudge→grace→over_budget, wrap-up-in-grace, zero grace, dead-window-at-expiry/at-trip, torn transcript, wall-clock-frozen-monotonic grace expiry, parser-none/mode-off/OSError inert).
  • opencode E2Es against the fake server: enforce trip→abort, zero-grace, warn-mode (no abort, no nudge, one ATTENTION across multiple samples — pins the trip latch on its separate copy of the logic), under-cap inert.
  • Engine: over_budget rides retry→defer for dev and review sessions; budget extras on session-end; all four fields threaded into every SessionSpec.

Summary by CodeRabbit

  • New Features
    • Added configurable mid-session token-budget guard with off, warn (attention + breadcrumbs, no termination), and enforce (optional wrap-up grace then over-budget termination).
    • Introduced cost-weighted budget tracking and cache-read weighting, with weighted budget details recorded in session history and lifecycle logs.
    • Added new limits settings: session_budget_mode, max_tokens_per_session, and session_budget_grace_s.
  • Documentation
    • Expanded guidance and TUI/Journaling documentation to cover mid-session budget trip and over-budget outcomes.
  • Bug Fixes
    • Improved robustness when sampling usage, sending nudges, notifying events, or during post-kill reconciliation with over-budget results.

…nforce modes (#158)

Sample cumulative session usage on the 30s heartbeat inside both adapter
wait loops; a weighted spend crossing limits.max_tokens_per_session trips
once per session — warn mirrors to ATTENTION, enforce nudges a wrap-up,
arms a wall-co-bound grace window, then terminates over_budget (rides the
existing retry→defer arm untouched).

Claude-Session: https://claude.ai/code/session_01K1McZgyLxJumxC2U8KBtoP
@coderabbitai

coderabbitai Bot commented Jul 17, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 6cb971cb-92fa-4709-9ca8-149ff47da9e1

📥 Commits

Reviewing files that changed from the base of the PR and between 85cb395 and 631ff8c.

📒 Files selected for processing (3)
  • src/bmad_loop/policy.py
  • tests/test_opencode_http.py
  • tests/test_policy.py
🚧 Files skipped from review as they are similar to previous changes (3)
  • src/bmad_loop/policy.py
  • tests/test_policy.py
  • tests/test_opencode_http.py

Walkthrough

Adds a configurable mid-session weighted token-budget guard to generic and OpenCode adapters. Sessions can warn or enforce a cap, record budget metadata, terminate as over_budget, and follow existing retry/defer handling.

Changes

Mid-session token-budget guard

Layer / File(s) Summary
Budget policy and session contracts
src/bmad_loop/policy.py, src/bmad_loop/adapters/base.py, src/bmad_loop/engine.py, src/bmad_loop/data/settings/core.toml, tests/test_policy.py, tests/test_engine.py, tests/test_settings_schema.py
Adds budget settings, validation, session fields, engine wiring, lifecycle metadata, and schema coverage.
Generic adapter enforcement
src/bmad_loop/adapters/generic.py, tests/test_generic_tmux.py
Samples transcript usage, supports warn/enforce modes and grace periods, emits lifecycle notifications, propagates weighted usage, and reconciles over-budget results.
OpenCode HTTP enforcement
src/bmad_loop/adapters/opencode_http.py, tests/test_opencode_http.py
Samples usage through HTTP, applies budget handling in the SSE wait loop, and tests enforcement, warning, grace, failure, and crash paths.
Budget behavior documentation
README.md, CHANGELOG.md, docs/FEATURES.md, docs/tui-guide.md
Documents configuration, weighted accounting, lifecycle events, enforcement behavior, and adapter limitations.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Engine
  participant Adapter
  participant UsageSource
  participant Lifecycle
  Engine->>Adapter: start session with budget settings
  Adapter->>UsageSource: sample cumulative weighted usage
  Adapter->>Lifecycle: record budget trip
  Adapter-->>Engine: return completed or over_budget result
Loading

Possibly related PRs

Poem

I’m a rabbit with a token tally,
Guarding each busy session’s valley.
Warn, then nudge, then hop away,
When budgets bloom beyond their stay.
Journals sparkle: “over budget!” hooray!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 44.94% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title is concise and accurately summarizes the main feature: a mid-session token-budget guard with enforcement modes.
Linked Issues check ✅ Passed The PR implements the requested mid-session usage sampling, weighted budget enforcement, over_budget handling, lifecycle visibility, and retry/defer routing.
Out of Scope Changes check ✅ Passed The documented changes are all supportive of the budget-guard feature, with no clear unrelated scope creep identified.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch feat/mid-session-budget-guard-158

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🧹 Nitpick comments (1)
tests/test_opencode_http.py (1)

737-757: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Make the under-cap test observe nonzero usage.

The completed scenario finishes before the next normal heartbeat, so the only sample sees []; this test passes without exercising weighted usage below the cap. Use big-usage-then-complete with a cap above 5M and a shortened heartbeat.

🤖 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 `@tests/test_opencode_http.py` around lines 737 - 757, Update
test_e2e_budget_inert_under_cap to use the big-usage-then-complete scenario, set
token_budget above 5M, and shorten token_budget_grace_s so a heartbeat samples
nonzero usage before completion. Preserve the assertions that the run completes
without ATTENTION, a breadcrumb, or budget_weighted.
🤖 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 `@src/bmad_loop/adapters/opencode_http.py`:
- Line 647: Replace the Bandit-only suppression on the budget mode condition in
src/bmad_loop/adapters/opencode_http.py:647 with the targeted Ruff suppression #
noqa: S105. Add # noqa: S106 to the corresponding budget mode enum usages in
tests/test_opencode_http.py at lines 684, 725, 747, 773, and 844; make no other
changes.

In `@src/bmad_loop/policy.py`:
- Around line 609-614: Update the LimitsPolicy construction for
max_tokens_per_session and session_budget_grace_s to validate the raw
configuration values as non-boolean integers before conversion, matching the
existing scm.preserve_keep validation pattern. Reject booleans and floats rather
than coercing or truncating them, while preserving the configured/default
integer values for valid inputs.

---

Nitpick comments:
In `@tests/test_opencode_http.py`:
- Around line 737-757: Update test_e2e_budget_inert_under_cap to use the
big-usage-then-complete scenario, set token_budget above 5M, and shorten
token_budget_grace_s so a heartbeat samples nonzero usage before completion.
Preserve the assertions that the run completes without ATTENTION, a breadcrumb,
or budget_weighted.
🪄 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

Run ID: 0573df78-8474-4b8d-8e8e-164f699e0776

📥 Commits

Reviewing files that changed from the base of the PR and between d524a78 and 85cb395.

📒 Files selected for processing (15)
  • CHANGELOG.md
  • README.md
  • docs/FEATURES.md
  • docs/tui-guide.md
  • src/bmad_loop/adapters/base.py
  • src/bmad_loop/adapters/generic.py
  • src/bmad_loop/adapters/opencode_http.py
  • src/bmad_loop/data/settings/core.toml
  • src/bmad_loop/engine.py
  • src/bmad_loop/policy.py
  • tests/test_engine.py
  • tests/test_generic_tmux.py
  • tests/test_opencode_http.py
  • tests/test_policy.py
  • tests/test_settings_schema.py

Comment thread src/bmad_loop/adapters/opencode_http.py
Comment thread src/bmad_loop/policy.py Outdated
- policy: reject bool/float/string for max_tokens_per_session and
  session_budget_grace_s instead of int()-coercing (true -> 1 token would
  terminate every enforce-mode session; strings crashed with a raw
  ValueError) — same strict pattern as scm.preserve_keep
- tests: rework test_e2e_budget_inert_under_cap onto big-usage-then-complete
  with a shrunk heartbeat so samples actually observe the 5M weighted spend
  under the cap (the completed scenario finished before ever reporting
  usage, leaving the comparison untested); assert no budget-tripped
  breadcrumb as the docstring claims
- rejected the S105/S106 noqa suggestion: repo Ruff runs default rules
  (no lint.select), S rules never fire; Bandit is the security gate and
  the existing nosec B105 covers it

Claude-Session: https://claude.ai/code/session_01K1McZgyLxJumxC2U8KBtoP
@pbean

pbean commented Jul 18, 2026

Copy link
Copy Markdown
Collaborator Author

Review round addressed in 631ff8c: strict non-bool-integer validation for the two new budget knobs (+ rejection tests), and the under-cap nitpick — test_e2e_budget_inert_under_cap now uses big-usage-then-complete with a shrunk heartbeat so samples observe the 5M weighted spend under the cap, and asserts no budget-tripped breadcrumb. The S105/S106 noqa suggestion was rejected (repo Ruff runs default rules with no lint.select; Bandit is the security gate and the existing # nosec B105 covers it — see the inline reply).

@pbean

pbean commented Jul 18, 2026

Copy link
Copy Markdown
Collaborator Author

@CodeRabbit review

@coderabbitai

coderabbitai Bot commented Jul 18, 2026

Copy link
Copy Markdown
✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

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.

No mid-session guardrail: a productive-but-looping session (265 turns, 49.7M tokens) is bounded only by the clock

1 participant