From 97c30f51a16c50873e3b21a8e47ef4b423f1af17 Mon Sep 17 00:00:00 2001 From: Brian Krabach Date: Wed, 7 Oct 2026 22:01:16 -0700 Subject: [PATCH] Preserve checkpoint continuity across token counting outages --- .../boundary.py | 8 ++ .../tests/test_count_failure_continuity.py | 83 +++++++++++++++++++ 2 files changed, 91 insertions(+) create mode 100644 modules/context-managed/tests/test_count_failure_continuity.py diff --git a/modules/context-managed/amplifier_module_context_managed/boundary.py b/modules/context-managed/amplifier_module_context_managed/boundary.py index c6bed26..e6ca6c2 100644 --- a/modules/context-managed/amplifier_module_context_managed/boundary.py +++ b/modules/context-managed/amplifier_module_context_managed/boundary.py @@ -745,6 +745,14 @@ async def prepare_note(): if not self._same_prefix(self._summary_progress, self.messages): self._summary_progress = None raise + # A counting outage says nothing about the validity of this + # history or checkpoint. Preserve the typed diagnostic and do + # not install a summary-failure latch that blocks continuation. + if (type(exc).__module__, type(exc).__name__) == ( + "amplifier_module_provider_openai._token_count", "TokenCountError"): + outcome = "failed" + stats["failure"] = {"type": "TokenCountError", "code": "token_count_failed"} + raise # Stop this request instead of hiding compaction failures with # trimming. Keep canonical history and the last valid checkpoint. outcome = "failed" diff --git a/modules/context-managed/tests/test_count_failure_continuity.py b/modules/context-managed/tests/test_count_failure_continuity.py new file mode 100644 index 0000000..d6e5b11 --- /dev/null +++ b/modules/context-managed/tests/test_count_failure_continuity.py @@ -0,0 +1,83 @@ +"""A count outage never poisons a valid checkpoint or replays history.""" +import copy +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from amplifier_core import ChatRequest, Message +from amplifier_module_context_managed.boundary import BoundaryContextManager + + +def checkpoint(label): + return {'role': 'user', 'content': 'Native checkpoint', 'metadata': { + 'source': 'context-managed', 'ephemeral': True, 'persisted': True, 'opaque': label}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize('phase', ['initial_count', 'compact_count', 'result_count']) +async def test_count_failure_preserves_checkpoint_then_continues(phase): + error_type = type('TokenCountError', (RuntimeError,), { + '__module__': 'amplifier_module_provider_openai._token_count'}) + failure = error_type('Could not check conversation size') + ctx = BoundaryContextManager({'durable_checkpoints': True, 'max_tokens': 5000, + 'summarize_trigger': .2, 'native_min_new_tokens': 0}) + identity = {'provider': 'fixture', 'model': 'native'} + ctx.checkpoint_identity = lambda: identity + rows = [{'role': 'user', 'content': 'Produce a report'}, + {'role': 'assistant', 'content': 'Earlier report saved'}, + {'role': 'user', 'content': 'Check the remaining evidence'}, + {'role': 'assistant', 'content': '', 'tool_calls': [{'id': 'done', 'name': 'read_file', 'arguments': {}}]}, + {'role': 'tool', 'tool_call_id': 'done', 'content': 'Already completed evidence'}, + {'role': 'user', 'content': 'Continue'}, + {'role': 'assistant', 'content': 'Progress saved'}, + {'role': 'user', 'content': 'Keep going'}] + await ctx.set_messages(rows) + ctx.summary = (2, checkpoint('old')) + ctx.summary_identity = identity + before = json.dumps(ctx.export_checkpoint(identity), sort_keys=True) + recovered = False + compacted = False + + async def budget(request, **kwargs): + if not recovered and phase == 'compact_count': + raise failure + return {'estimated_input_tokens': 2000, 'input_limit_tokens': 5000, + 'measurement': {'kind': 'provider_count', 'source': 'fixture', 'input_tokens': 2000}} + + async def compact(request): + nonlocal compacted + compacted = True + return {'kind': 'native', 'message': checkpoint('new')} + + provider = SimpleNamespace(name='fixture', default_model='native', + supports_native_compaction=lambda: True, + validate_compacted_context=lambda row: bool(row.get('metadata', {}).get('opaque')), + request_budget=budget, compact_context=AsyncMock(side_effect=compact), complete=AsyncMock()) + + async def count_view(view): + if not recovered and ((phase == 'initial_count' and not compacted) or + (phase == 'result_count' and compacted)): + raise failure + tokens = 100 if compacted else 4500 + return {'dispatch': ChatRequest(messages=[Message(**row) for row in view]), + 'budget_decision': {'estimated_input_tokens': tokens, 'input_limit_tokens': 5000, + 'measurement': {'kind': 'provider_count', 'source': 'fixture', 'input_tokens': tokens}}} + + with pytest.raises(error_type) as caught: + await ctx.get_measured_request_view(provider=provider, retain_contents=[], count_view=count_view) + assert caught.value is failure + assert json.dumps(ctx.export_checkpoint(identity), sort_keys=True) == before + assert ctx.summary_failure is None + assert ctx.messages == rows + provider.complete.assert_not_called() + recovered, compacted = True, False + result = await ctx.get_measured_request_view(provider=provider, retain_contents=[], count_view=count_view) + assert result['final_attempt']['budget_decision']['measurement']['input_tokens'] == 100 + assert ctx.summary[0] > 2 + assert ctx.messages == rows + assert sum(row.get('tool_call_id') == 'done' for row in ctx.messages) == 1 + provider.complete.assert_not_called() + restored = BoundaryContextManager(copy.deepcopy(ctx.config)) + await restored.set_messages(rows) + assert restored.restore_checkpoint(ctx.export_checkpoint(identity), identity)['status'] == 'restored'