fix: #162 — budget-overflow no longer corrupts session state in submit_message

Previously, QueryEnginePort.submit_message() checked the token budget AFTER
appending the prompt to mutable_messages, transcript_store, and permission_denials,
and AFTER calling compact_messages_if_needed(). On overflow it set
stop_reason='max_budget_reached' but the overflow turn was already committed.
Any caller that persisted the session afterwards wrote the rejected prompt to
disk — the session was silently poisoned even though the TurnResult said the
turn never completed.

Fix:
- Restructure submit_message so the budget check early-returns BEFORE any
  mutation of mutable_messages, transcript_store, permission_denials, or
  total_usage.
- The returned TurnResult.usage reflects pre-call state (overflow never
  advanced the usage counter).
- Normal (in-budget) path unchanged: mutation happens exactly once, at the
  end, only on 'completed' results.

This closes the atomicity gap: submit_message is now either 'turn committed'
(stop_reason='completed') or 'turn rejected, state untouched'
(stop_reason in {'max_budget_reached', 'max_turns_reached'}). Callers can
safely retry with a fresh budget or a smaller prompt without worrying about
phantom committed turns from prior rejections.

Tests (tests/test_submit_message_budget.py, 10 tests):
- TestBudgetOverflowDoesNotMutate (5): mutable_messages / transcript /
  permission_denials / total_usage / TurnResult.usage all pre-mutation after overflow
- TestOverflowPersistence (2): first-turn overflow persists empty session;
  successful-turn-then-overflow persists only the successful turn
- TestEngineUsableAfterOverflow (2): subsequent in-budget call still works
  with no residue; repeated overflows don't accumulate hidden state
- TestNormalPathStillCommits (1): regression guard — non-overflow path still
  commits mutable_messages/transcript/usage as expected

Full suite: 59/59 passing, zero regression.

Blocker: none. Closes ROADMAP #162.
This commit is contained in:
YeonGyu-Kim
2026-04-22 17:29:55 +09:00
parent 3f4d46d7b4
commit 4813a2b351
2 changed files with 186 additions and 3 deletions

View File

@@ -85,9 +85,25 @@ class QueryEnginePort:
]
output = self._format_output(summary_lines)
projected_usage = self.total_usage.add_turn(prompt, output)
stop_reason = 'completed'
# #162: budget check must precede mutation. Previously this block set
# stop_reason='max_budget_reached' but still appended the overflow turn
# to mutable_messages / transcript_store / permission_denials, corrupting
# the session for any caller that persisted it afterwards. The overflow
# prompt was effectively committed even though the TurnResult signalled
# rejection. Now we early-return with pre-mutation state intact so
# callers can safely retry with a smaller prompt or a fresh budget.
if projected_usage.input_tokens + projected_usage.output_tokens > self.config.max_budget_tokens:
stop_reason = 'max_budget_reached'
return TurnResult(
prompt=prompt,
output=output,
matched_commands=matched_commands,
matched_tools=matched_tools,
permission_denials=denied_tools,
usage=self.total_usage, # unchanged — overflow turn was rejected
stop_reason='max_budget_reached',
)
self.mutable_messages.append(prompt)
self.transcript_store.append(prompt)
self.permission_denials.extend(denied_tools)
@@ -100,7 +116,7 @@ class QueryEnginePort:
matched_tools=matched_tools,
permission_denials=denied_tools,
usage=self.total_usage,
stop_reason=stop_reason,
stop_reason='completed',
)
def stream_submit_message(