fix: accurate per-call token usage and oversized-context guards

Token accounting:
- Drain each tool round's provider stream to exhaustion before running
  tools and recursing, so the usage decorator persists exactly one
  token_usage row per LLM call, at call end. Previously every round's
  generator was abandoned mid-iteration and flushed together at request
  teardown, writing N near-identical rows stamped with the final
  round's provider counts (duplicate billing).
- Consume the Chat Completions include_usage terminal chunk (it arrives
  after finish_reason and was never read) so streamed calls record
  provider-exact token counts instead of tiktoken estimates.
- Claim provider-reported usage once per call (_last_usage_claimed) so
  a late-finalized generator can never adopt another call's counts.

Oversized-context guards:
- Enforce Responses API function_call/function_call_output pairing in
  the input builder (drop unpaired items; bypassed for store-mode
  previous_response_id chaining where calls are matched server-side).
- Hard pre-send context gate: shrink oversized tool results and refuse
  payloads that cannot fit the model's window before dispatch, so a
  hopeless request is never sent or billed.
- Cap a single tool result entering the LLM context
  (TOOL_RESULT_MAX_TOKENS, default 20000); the tool journal and
  persistence keep the full result. Applied on the resume/continuation
  path too.
- Skip the fallback attempt when the payload cannot fit the fallback
  model's context window (10% estimation slack).
- Compression: never save a compression point that does not reduce
  tokens; bound oversized verbatim fields kept after a compression
  point (COMPRESSION_RECENT_FIELD_MAX_TOKENS, default 8000).

Robustness fixes from review:
- Google parallel function calls: complete index-less ToolCalls are no
  longer merged into one another (dict arguments raised TypeError on
  +=; second call could execute with the first call's arguments).
- Trailing-frame failures after a delivered answer no longer error the
  stream or restream the whole answer from the fallback
  (_stream_reached_finish).
- In-memory compression falls back to minimal pruning when the summary
  is not smaller than the original.
- keep<=0 guard in the middle-truncation helpers (a tiny cap returned
  marker + full text).

Frontend: tooltip on the Analytics tokens stat card explaining that
agent tool loops re-send conversation context on every step.
This commit is contained in:
Alex committed 2026-07-15 21:20:39 +01:00
1 parent a7fd6d7c02
commit ca5f80995d
16 files changed
+1626 -119

No files matched your search

+121
View File
@@ -811,3 +811,124 @@ class TestRespondingProviderTracking:
primary = _Google(fail_at=0, backup_models=["backup-model"])
primary.gen(**CALL_ARGS)
assert primary._responding_provider == "openai"
# Tests — fallback payload size gate
@pytest.mark.integration
class TestFallbackPayloadSizeGate:
"""A payload that cannot fit the fallback's context window must skip the
fallback attempt (it would be a guaranteed second rejection) and
propagate the primary's error instead."""
def _primary_with_backup(self, patch_model_utils, backup):
patch_model_utils(
get_provider=lambda mid, **_kw: "openai",
get_api_key=lambda p: "k",
create_llm=lambda type, **kw: backup,
)
return FakeLLM(fail_at=0, backup_models=["backup-model"])
BIG_ARGS = dict(
model="test-model",
messages=[{"role": "user", "content": "word " * 300}],
)
def test_gen_skips_fallback_when_payload_cannot_fit(
self, monkeypatch, patch_model_utils
):
backup = FakeLLM(responses=["backup ok"])
primary = self._primary_with_backup(patch_model_utils, backup)
monkeypatch.setattr(
"application.core.model_utils.get_token_limit",
lambda mid, user_id=None: 10,
)
with pytest.raises(RuntimeError, match="primary model unavailable"):
primary.gen(**self.BIG_ARGS)
assert backup.gen_called is False
def test_stream_skips_fallback_when_payload_cannot_fit(
self, monkeypatch, patch_model_utils
):
backup = FakeLLM(stream_chunks=["backup chunk"])
primary = self._primary_with_backup(patch_model_utils, backup)
monkeypatch.setattr(
"application.core.model_utils.get_token_limit",
lambda mid, user_id=None: 10,
)
with pytest.raises(RuntimeError, match="mid-stream failure"):
list(primary.gen_stream(**self.BIG_ARGS))
assert backup.gen_stream_called is False
def test_fallback_proceeds_when_payload_fits(
self, monkeypatch, patch_model_utils
):
backup = FakeLLM(responses=["backup ok"])
primary = self._primary_with_backup(patch_model_utils, backup)
monkeypatch.setattr(
"application.core.model_utils.get_token_limit",
lambda mid, user_id=None: 100000,
)
assert primary.gen(**self.BIG_ARGS) == "backup ok"
assert backup.gen_called is True
def test_estimation_failure_never_blocks_fallback(
self, monkeypatch, patch_model_utils
):
backup = FakeLLM(responses=["backup ok"])
primary = self._primary_with_backup(patch_model_utils, backup)
def boom(*a, **kw):
raise ValueError("estimator broken")
monkeypatch.setattr("application.usage._count_prompt_tokens", boom)
assert primary.gen(**self.BIG_ARGS) == "backup ok"
assert backup.gen_called is True
# Tests — no fallback restream after the primary delivered its finish signal
@pytest.mark.integration
class TestNoRestreamAfterFinish:
"""A trailing-frame failure (between the finish chunk and stream end)
must NOT restream the already-delivered answer from the fallback."""
class FinishThenFailLLM(FakeLLM):
def _raw_gen_stream(self, baseself, model, messages, stream, tools=None, **kwargs):
self.gen_stream_called = True
yield "the full answer"
# Provider marked the stream finished (finish_reason arrived)...
self._stream_reached_finish = True
# ...then the trailing usage/[DONE] frame dies.
raise RuntimeError("connection reset in trailing frame")
def test_trailing_frame_failure_skips_fallback(self, patch_model_utils):
backup = FakeLLM(stream_chunks=["fallback chunk"])
patch_model_utils(
get_provider=lambda mid, **_kw: "openai",
get_api_key=lambda p: "k",
create_llm=lambda type, **kw: backup,
)
primary = self.FinishThenFailLLM(backup_models=["backup-model"])
received = []
with pytest.raises(RuntimeError, match="trailing frame"):
for chunk in primary.gen_stream(**CALL_ARGS):
received.append(chunk)
assert received == ["the full answer"]
assert backup.gen_stream_called is False
def test_pre_finish_failure_still_falls_back(self, patch_model_utils):
backup = FakeLLM(stream_chunks=["fallback chunk"])
patch_model_utils(
get_provider=lambda mid, **_kw: "openai",
get_api_key=lambda p: "k",
create_llm=lambda type, **kw: backup,
)
primary = FakeLLM(fail_at=0, backup_models=["backup-model"])
out = list(primary.gen_stream(**CALL_ARGS))
assert out == ["fallback chunk"]
assert backup.gen_stream_called is True