feat(engine): warn on max_tokens_per_story at session boundaries - #346
Conversation
The advisory per-story cap was read once, inside _finalize_commit_phase and past advance(task, Phase.DONE): an overrun surfaced only after every token was spent, and a story that deferred or escalated was never checked at all. Re-check the story-cumulative weighted spend at every session boundary in _run_session instead, where task.tokens is fresh and story-cumulative and every run type funnels through. The first crossing journals (now with a `budget` field) and raises an ATTENTION + desktop notice; a persisted StoryTask.token_budget_warned latch keeps it one notice per story across resumes. Still advisory - nothing is terminated, and the #158 per-session guard is untouched. Closes #336
WalkthroughChangesThe per-story token budget is now checked after every session boundary using cumulative weighted spend. The first over-budget crossing emits an attention notice and journal entry, while a persisted story latch prevents duplicate warnings after later sessions or resumes. Story token budget
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Session as Session runner
participant Engine
participant Story as StoryTask
participant Journal
participant Notice as Attention notice
Session->>Engine: finish session and save state
Engine->>Story: check cumulative weighted spend
Engine->>Journal: record token-budget-exceeded
Engine->>Notice: emit one attention notification
Engine->>Story: persist token_budget_warned
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/engine.py`:
- Around line 2792-2809: The token-budget warning latch in the task handling
flow must be persisted immediately after setting task.token_budget_warned,
before self.journal.append or gates.notify executes. Update the relevant method
around the token-budget warning logic, and add a regression test covering the
save-failure/process-death boundary after notification delivery while preserving
single-warning behavior across resume.
🪄 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: 9e60938b-5f14-4aa8-b39d-9210012d3b34
📒 Files selected for processing (10)
CHANGELOG.mdREADME.mddocs/FEATURES.mddocs/tui-guide.mdsrc/bmad_loop/data/settings/core.tomlsrc/bmad_loop/engine.pysrc/bmad_loop/model.pysrc/bmad_loop/policy.pytests/test_engine.pytests/test_model.py
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| task.token_budget_warned = True | ||
| self.journal.append( | ||
| "token-budget-exceeded", | ||
| story_key=task.story_key, | ||
| weighted=weighted, | ||
| total=task.tokens.total, | ||
| budget=budget, | ||
| ) | ||
| gates.notify( | ||
| self.policy, | ||
| self.run_dir, | ||
| f"story token budget exceeded: {task.story_key}", | ||
| f"{weighted} weighted tokens ({task.tokens.total} raw incl. cache reads) " | ||
| f"past the {budget} `limits.max_tokens_per_story` cap — advisory: the " | ||
| f"story continues. Stop the run with `bmad-loop stop` if the spend is " | ||
| f"no longer worth it.", | ||
| ) | ||
| self._save() |
There was a problem hiding this comment.
task.token_budget_warned = True is committed to state.json only at the very end of this method, after both the journal write and gates.notify. If the process is killed between those side effects and the final self._save() (e.g., host OOM, SIGKILL on the post-session window), the state on disk still has token_budget_warned = False. On resume the budget check fires a second time: a duplicate token-budget-exceeded entry lands in the journal and the user gets a second ATTENTION + desktop notification.
The existing test_token_budget_latch_survives_resume test validates the crash-after-_note_story_token_budget case but not a crash inside it; moving the save to the top of the "triggered" branch closes that window with a one-line reorder.
| task.token_budget_warned = True | |
| self.journal.append( | |
| "token-budget-exceeded", | |
| story_key=task.story_key, | |
| weighted=weighted, | |
| total=task.tokens.total, | |
| budget=budget, | |
| ) | |
| gates.notify( | |
| self.policy, | |
| self.run_dir, | |
| f"story token budget exceeded: {task.story_key}", | |
| f"{weighted} weighted tokens ({task.tokens.total} raw incl. cache reads) " | |
| f"past the {budget} `limits.max_tokens_per_story` cap — advisory: the " | |
| f"story continues. Stop the run with `bmad-loop stop` if the spend is " | |
| f"no longer worth it.", | |
| ) | |
| self._save() | |
| task.token_budget_warned = True | |
| self._save() | |
| self.journal.append( | |
| "token-budget-exceeded", | |
| story_key=task.story_key, | |
| weighted=weighted, | |
| total=task.tokens.total, | |
| budget=budget, | |
| ) | |
| gates.notify( | |
| self.policy, | |
| self.run_dir, | |
| f"story token budget exceeded: {task.story_key}", | |
| f"{weighted} weighted tokens ({task.tokens.total} raw incl. cache reads) " | |
| f"past the {budget} `limits.max_tokens_per_story` cap — advisory: the " | |
| f"story continues. Stop the run with `bmad-loop stop` if the spend is " | |
| f"no longer worth it.", | |
| ) |
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/bmad_loop/engine.py
Line: 2792-2809
Comment:
**Latch not persisted before journal/notify side effects**
`task.token_budget_warned = True` is committed to `state.json` only at the very end of this method, after both the journal write and `gates.notify`. If the process is killed between those side effects and the final `self._save()` (e.g., host OOM, SIGKILL on the post-session window), the state on disk still has `token_budget_warned = False`. On resume the budget check fires a second time: a duplicate `token-budget-exceeded` entry lands in the journal and the user gets a second ATTENTION + desktop notification.
The existing `test_token_budget_latch_survives_resume` test validates the crash-after-`_note_story_token_budget` case but not a crash inside it; moving the save to the top of the "triggered" branch closes that window with a one-line reorder.
```suggestion
task.token_budget_warned = True
self._save()
self.journal.append(
"token-budget-exceeded",
story_key=task.story_key,
weighted=weighted,
total=task.tokens.total,
budget=budget,
)
gates.notify(
self.policy,
self.run_dir,
f"story token budget exceeded: {task.story_key}",
f"{weighted} weighted tokens ({task.tokens.total} raw incl. cache reads) "
f"past the {budget} `limits.max_tokens_per_story` cap — advisory: the "
f"story continues. Stop the run with `bmad-loop stop` if the spend is "
f"no longer worth it.",
)
```
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.There was a problem hiding this comment.
Same finding as the CodeRabbit thread — full reasoning there: #346 (comment)
Short version: the direction is right but the _save() hoist is the wrong end of the block.
_run_inner'sfinallyatengine.py:491already_save()s on every abort arm (including the genericexcept Exceptionat 441-481 — see the comment at 451), so only a SIGKILL reaches the window you describe, and a duplicate breaks nothing downstream (diagnostics.py:378-389's histogram is descriptive; the derived counters at 407-409 don't include this kind).- The hoist converts at-least-once into at-most-once: a death between
_save()andjournal.appendwould leave the latch durable with no record and no notice — the max_tokens_per_story is checked only post-done and only journals — warn while the story runs #336 silent-overrun bug, reintroduced. journal → notify → saveis also the uniform house order (_record_defer3135/3141/3147,_escalate3242/3243/3249,_vetoed998-1014,recovery_flow.py:635/641/647).
The actual defect was one line above: task.token_budget_warned = True sat before the journal.append. Journal.append is unguarded IO (journal.py:41-42), so an OSError there let that same finally persist the latch with nothing written anywhere — permanent silence, no SIGKILL needed. Unlike the phase its siblings set pre-write, this bit is rendered on no surface, so a lost record is unrecoverable. Moved the latch below the append (emit order unchanged) in 2771189, with test_token_budget_latch_waits_for_the_durable_record — which fails on 9a15718.
Your point about the existing resume test not covering a crash inside the method was exactly right, and that's the gap the new test closes.
There was a problem hiding this comment.
You're right, and the reasoning is solid on all three counts.
The OSError path is the stronger argument — Journal.append is unguarded IO and the finally in _run_inner will persist whatever is in memory at that point. My hoist protects against SIGKILL (rare, duplicate harmless) at the cost of reintroducing silent loss on a plain OSError (no SIGKILL required, and unrecoverable because the latch surfaces nowhere). That's a worse trade.
The house order (journal → notify → save) also makes the intent legible: the record must exist before the latch is durable. My suggestion inverted that.
The fix in 2771189 — latch below journal.append, emit order unchanged — is the right call. The new test covering the OSError path is exactly the gap I noted, closed correctly. Nothing further to add here.
Should I remember the house order (journal → notify → save) as a rule for future reviews in this repo?
Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.
Greptile SummaryThis PR moves the
Confidence Score: 5/5Safe to merge — the change is surgical (new method + moved check), all session paths are covered, and every ordering invariant has a dedicated test with a proven ablation. The logic is confined to a single new method called at a well-defined point in Files Needing Attention: No files require special attention.
|
| Filename | Overview |
|---|---|
| src/bmad_loop/engine.py | New _note_story_token_budget method placed after the tail _save() in _run_session; covers all session paths including defer/escalate. Ordering invariants (journal → latch → notify → save) are clearly documented and validated by dedicated tests. |
| src/bmad_loop/model.py | Adds token_budget_warned: bool = False to StoryTask with correct to_dict/from_dict round-trip and legacy-state default; test coverage for both cases. |
| src/bmad_loop/policy.py | Adds isinstance(bool) + non-int guard and < 1 range check for max_tokens_per_story, matching the existing validation for max_tokens_per_session. Previously a bare int() silently accepted true → 1-token cap. |
| src/bmad_loop/gates.py | File sink for ATTENTION wrapped in try/except OSError; honours the documented "never raises" contract that previously only the desktop half kept. |
| tests/test_engine.py | Adds four new budget tests covering: the deferring-story case (#336), once-per-story latch, latch survival across resume, and latch-ordering invariant (journal write must precede latch set). Ablation table in PR description validates negative assertions. |
| tests/test_policy.py | New parametrised tests cover default parse, <= 0 rejection, and non-integer rejection for max_tokens_per_story. |
| tests/test_model.py | Adds round-trip and legacy-default tests for token_budget_warned. |
| tests/test_gates.py | New test verifies the ATTENTION-path OSError is silently swallowed by placing a directory at the file path (portable across OS/root). |
Reviews (2): Last reviewed commit: "fix(policy): validate max_tokens_per_sto..." | Re-trigger Greptile
…ning `token_budget_warned` was set before the journal.append it suppresses. Journal.append is unguarded IO, and _run_inner's finally persists the task on every abort arm, so an OSError there made the latch durable with no journal entry and no notice - the story's overrun suppressed forever, which is #336 itself. Unlike the phase _record_defer/_escalate set before their own writes, this bit is rendered on no surface, so a lost record is unrecoverable. Latch after the append instead; the journal -> notify -> save emit order is unchanged. _save() deliberately stays last: at-least-once is the right bias for an anti-silence warning, and the run-level finally already covers every in-process abort.
notify() documented "never raises", but only the desktop half was guarded - the ATTENTION append was bare IO. An unwritable run dir therefore turned an advisory notice into a run crash at every record-a-decision site (defer, escalate, plugin veto, manual-recovery pause, budget warnings), several of which notify after already journaling a decision they cannot unwind. Degrade the file sink like the desktop one, per the observe-degrade rule the adapter budget guards already apply to this same call. The journal entry each caller writes first stays the durable record.
The field was parsed with a bare int(): `true` coerced to a 1-token story cap and `0` was accepted, both against the documented int >= 1 (core.toml already sets minimum = 1, so the settings TUI could not produce either). max_tokens_per_session has carried the bool/int and >= 1 guards since #158; mirror them. Pre-existing, but the session-boundary check makes it loud: a 0/1 cap now warns at the first session boundary of every story instead of once post-done. Rejects a policy.toml that previously loaded - there is no 0 = off semantic.
1 participant
max_tokens_per_story(2M weighted, advisory) had exactly one read site: inside_finalize_commit_phase, afteradvance(task, Phase.DONE). Two consequences,both hit by the 0.9.0 field report behind #337:
reported run burned 8.35M weighted against the 2M cap and produced no record
of any kind. Total silence, not late silence.
What changed
The story's cumulative weighted spend is re-checked at every session
boundary: a new
_note_story_token_budget(task)called from_run_sessionright after its tail
_save(). That is the point wheretask.tokensis bothfresh (
attach_session_usagefolded this session in just above) andstory-cumulative, and every session of every run type funnels through it — so
deferring and escalating stories are judged like committing ones, and
SweepEngine/StoriesEngineinherit it unchanged. A session that abortednever reaches the site (the exception propagates past the
finally), so nostory is judged on a usage read that did not happen.
The first crossing:
token-budget-exceededkind, now carryingbudgetalongside
weighted/total;story token budget exceeded: <key>,body naming both units, the cap, and that this is advisory and the story
continues.
It fires once per story, latched on a new persisted
StoryTask.token_budget_warned— every session after the crossing is over thecap too, and the latch is in
to_dict/from_dictprecisely so a resumed rundoes not re-notify what the pre-pause process already did.
Enforcement weight is the live
self.policy.limits.cache_read_weight, theenforcement side of the split documented at
summary()(displays mustreproduce from
state.json; budgets bind at the policy the running processenforces).
Still advisory — nothing is terminated. No new policy knob. Zero changes to
the #158 per-session guard:
SessionSpec,adapters/generic.pyandadapters/opencode_http.pyare untouched.Also:
core.toml'smax_tokens_per_storyentry was the only[limits]fieldwith no
label/description(presentation-only; the settings sync tests areunchanged and pass as-is).
Tests
test_token_budget_exceeded_journals_weightedamended — still exactly oneentry at
weighted=160000, plus thebudgetfield, the ATTENTION content andthe latch.
test_token_budget_discounts_cache_readsuntouched — the no-false-positivecase still holds through the new site.
test_token_budget_warns_on_a_story_that_never_commits(the max_tokens_per_story is checked only post-done and only journals — warn while the story runs #336 case),test_token_budget_warns_once_per_story,test_token_budget_latch_survives_resume.test_model.pyround-trip + legacy-state default.Per the repo's ablation rule, each negative assertion was proven by deleting its
gate first:
task.phase != Phase.DONE: return..._warns_on_a_story_that_never_commitstoken_budget_warnedearly return..._warns_once_per_storytoken_budget_warnedfromto_dict..._latch_survives_resumetoken_budget_warned = Trueback abovejournal.append..._latch_waits_for_the_durable_recordTrueon disk, record absent) — and only this testexcept OSErroraround the ATTENTION writetest_notify_file_swallows_an_unwritable_attention_pathIsADirectoryErrorescapes)max_tokens_per_storybool/int guard..._rejects_non_integersmax_tokens_per_story >= 1guard..._must_be_positiveuv run pytest -q -n auto: 3416 passed, 24 skipped. Fulltrunk checkclean;pyright@1.1.411(CI pin) 0 errors.Review round 1 — two bots, one finding, plus two defects they surfaced
Both reviewers flagged the same thing: the latch
task.token_budget_warned = Truewas only madedurable by the trailing
self._save(), afterjournal.appendandgates.notify. Both proposedhoisting the
_save()up.Declined the hoist, fixed the actual defect at the other end of the block (2771189). The
window they describe is narrower than stated —
_run_inner'sfinallyatengine.py:491already_save()s on every abort arm, so only SIGKILL reaches it, and a duplicate breaks nothingdownstream. The hoist would convert at-least-once into at-most-once: a death between the save and
the append leaves the latch durable with no record and no notice, i.e. the #336 silent overrun
reintroduced.
The real bug was one line above: the latch was set before the
journal.append. That appendis unguarded IO (
journal.py:41-42), so anOSErrorlet the samefinallypersist the latch withnothing written anywhere — permanent silence, no SIGKILL needed. Unlike the phase
_record_deferand
_escalateset before their own writes, this bit is rendered on no surface, so a lost recordis unrecoverable. The latch now sits between the append and the notify; the
journal → notify → saveemit order is unchanged, so it still matches every other record-a-decision site.test_token_budget_latch_waits_for_the_durable_recordfails on 9a15718.Two further defects surfaced during that review and are fixed here too (maintainer's call to
handle in-PR rather than as follow-ups) — this makes the PR three logical changes rather than
one, kept as three separate commits:
gates.notifydidn't honor its "never raises" contract. The ATTENTION append(
gates.py:94) was bare IO while only the desktop half was guarded, so an unwritable run dirturned an advisory notice into a run crash at all ~11 record-a-decision sites — several of which
notify after journaling a decision they cannot unwind. Guarded inside
notifyso every callsite is covered at once, matching the observe-degrade rule already applied to this same call at
adapters/generic.py:633-637.limits.max_tokens_per_storyhad no validation. Bareint()atpolicy.py:743,so
truebecame a 1-token cap and0was accepted, against the documentedint >= 1(
core.tomlalready setsminimum = 1, so the settings TUI could never produce either).Mirrors
max_tokens_per_session's bool/int +>= 1guards.maintainer: this rejects a
policy.tomlthat previously loaded. There is no0 = offsemantic — set the cap high instead, it only warns.
Explicit non-goals (follow-ups)
Both were considered and deliberately left out of this PR:
SessionSpecso the running session itself samples against it (~30s cadence,like the No mid-session guardrail: a productive-but-looping session (265 turns, 49.7M tokens) is bounded only by the clock #158 per-session guard). Needs a latch separate from No mid-session guardrail: a productive-but-looping session (265 turns, 49.7M tokens) is bounded only by the clock #158's single
budget_trippedboolean. This PR's session-boundary check would already havefired after session 1–2 of the reported run.
status/ TUI — both surface totals only, with nobudget or remaining-headroom column (
cli.pystatus text,tui/widgets.pytoken columns).
Closes #336
Summary by CodeRabbit