From 5c5568b9000ca7c881ce6aa92e4cf4bff260cbec Mon Sep 17 00:00:00 2001 From: Lubaoshuai <128781758+Lubaoshuai@users.noreply.github.com> Date: Mon, 14 Sep 2026 00:59:39 +0800 Subject: [PATCH 1/3] Python: exclude encrypted reasoning payloads from compaction token counts The compaction token counter serialises each message and tokenises the JSON. Content.protected_data (OpenAI encrypted_content, Anthropic thinking signature, Gemini thought_signature, OpenAI reasoning_details) is opaque ciphertext the model replays but never tokenises, so counting it inflated every token-aware compaction decision. The same blob can also ride in additional_properties as encrypted_content. Both are now excluded from serialised content, matching the existing exclusions for raw_representation and mirrored items. Fixes #8348 --- .../core/agent_framework/_compaction.py | 10 +++++ .../core/tests/core/test_compaction.py | 38 +++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index b9080d0a814..64e62e086f1 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -700,6 +700,16 @@ def _serialize_content(content: Content) -> dict[str, Any]: # ``items`` mirrors ``result`` for function_result content; exclude it # to avoid double-counting tokens during estimation. payload.pop("items", None) + # ``protected_data`` holds provider-encrypted reasoning payloads (e.g. + # OpenAI ``encrypted_content``, Anthropic thinking ``signature``) that the + # model replays opaquely and never tokenises; exclude them so estimation + # measures the text the model actually sees. + payload.pop("protected_data", None) + additional_properties = payload.get("additional_properties") + if isinstance(additional_properties, dict) and "encrypted_content" in additional_properties: + payload["additional_properties"] = { + key: value for key, value in additional_properties.items() if key != "encrypted_content" + } return payload diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 2cfadff942b..9d3be5097f0 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -434,6 +434,44 @@ def test_annotate_message_groups_with_tokenizer_adds_token_counts() -> None: assert isinstance(_token_count(messages[1]), int) +def test_annotate_message_groups_token_counts_ignore_protected_data() -> None: + encrypted_blob = "A" * 4096 + with_payload = Message( + role="assistant", + contents=[ + Content.from_text_reasoning(id="rs_1", text="", protected_data=encrypted_blob), + Content.from_text("final answer"), + ], + ) + without_payload = Message( + role="assistant", + contents=[ + Content.from_text_reasoning(id="rs_1", text=""), + Content.from_text("final answer"), + ], + ) + + annotate_message_groups([with_payload], tokenizer=CharacterEstimatorTokenizer()) + annotate_message_groups([without_payload], tokenizer=CharacterEstimatorTokenizer()) + + assert _token_count(with_payload) == _token_count(without_payload) + + +def test_annotate_message_groups_token_counts_ignore_encrypted_content_in_additional_properties() -> None: + encrypted_blob = "A" * 4096 + with_blob = Content.from_text_reasoning(id="rs_1", text="") + with_blob.additional_properties["encrypted_content"] = encrypted_blob + without_blob = Content.from_text_reasoning(id="rs_1", text="") + + with_message = Message(role="assistant", contents=[with_blob, Content.from_text("final answer")]) + without_message = Message(role="assistant", contents=[without_blob, Content.from_text("final answer")]) + + annotate_message_groups([with_message], tokenizer=CharacterEstimatorTokenizer()) + annotate_message_groups([without_message], tokenizer=CharacterEstimatorTokenizer()) + + assert _token_count(with_message) == _token_count(without_message) + + def test_extend_compaction_messages_preserves_existing_annotations_and_tokens() -> None: tokenizer = CharacterEstimatorTokenizer() messages = [_assistant_function_call("c3")] From efb44d8893a7be32f91597699650fc344f91682f Mon Sep 17 00:00:00 2001 From: Lubaoshuai <128781758+Lubaoshuai@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:50:38 +0800 Subject: [PATCH 2/3] fix(compaction): count clear-text reasoning details and invalidate stale token caches Review follow-up for the protected-data token exclusion: - reasoning_details stored in protected_data are JSON-serialised and replayed to the provider as clear text (summary, reasoning_text and nested reasoning.text), so only opaque members (encrypted_content) are excluded from the estimate; non-JSON payloads (Anthropic signature, Responses API encrypted blobs) are still dropped entirely. - group token counts are stamped with a serialization basis version; counts cached under an older basis (e.g. sessions annotated before a serialization change) are recomputed instead of reused. --- .../core/agent_framework/_compaction.py | 68 +++++++++++++++++-- .../core/tests/core/test_compaction.py | 59 ++++++++++++++++ 2 files changed, 121 insertions(+), 6 deletions(-) diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index 64e62e086f1..67d49356c1b 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -30,6 +30,12 @@ GROUP_INDEX_KEY = "index" GROUP_HAS_REASONING_KEY = "has_reasoning" GROUP_TOKEN_COUNT_KEY = "token_count" # ruff:ignore[hardcoded-password-string] # nosec B105 - compaction metadata key, not a credential +GROUP_TOKEN_COUNT_BASIS_KEY = "token_count_basis" # ruff:ignore[hardcoded-password-string] # nosec B105 - compaction metadata key, not a credential +# Bumped whenever the token-estimation serialization changes. Cached counts +# stamped with an older basis are recomputed instead of reused: annotations +# survive Message.to_dict()/from_dict() round-trips, so without the basis a +# transcript annotated before a serialization change keeps its stale counts. +TOKEN_COUNT_BASIS_VERSION = 2 EXCLUDED_KEY = "_excluded" EXCLUDE_REASON_KEY = "_exclude_reason" SUMMARY_OF_MESSAGE_IDS_KEY = "_summary_of_message_ids" @@ -458,7 +464,13 @@ def _write_group_annotation( token_count: int | None = None if existing_raw_annotation is not None: raw_token_count = existing_raw_annotation.get(GROUP_TOKEN_COUNT_KEY) - if isinstance(raw_token_count, int) or raw_token_count is None: + raw_token_count_basis = existing_raw_annotation.get(GROUP_TOKEN_COUNT_BASIS_KEY) + # Preserve a cached count only when it was computed with the current + # serialization basis; counts from older bases are stale and get + # recomputed (see TOKEN_COUNT_BASIS_VERSION). + if raw_token_count_basis == TOKEN_COUNT_BASIS_VERSION and ( + isinstance(raw_token_count, int) or raw_token_count is None + ): token_count = raw_token_count unknown_fields = { key: value @@ -470,6 +482,7 @@ def _write_group_annotation( GROUP_INDEX_KEY, GROUP_HAS_REASONING_KEY, GROUP_TOKEN_COUNT_KEY, + GROUP_TOKEN_COUNT_BASIS_KEY, } } @@ -479,6 +492,7 @@ def _write_group_annotation( GROUP_INDEX_KEY: index, GROUP_HAS_REASONING_KEY: has_reasoning, GROUP_TOKEN_COUNT_KEY: token_count, + GROUP_TOKEN_COUNT_BASIS_KEY: TOKEN_COUNT_BASIS_VERSION if token_count is not None else None, } annotation.update(unknown_fields) message.additional_properties[GROUP_ANNOTATION_KEY] = annotation @@ -512,6 +526,10 @@ def _token_count(message: Message) -> int | None: if annotation is None: return None token_count = annotation.get(GROUP_TOKEN_COUNT_KEY) + # Counts stamped with an older serialization basis are stale: they were + # computed from a different serialized form and must be recomputed. + if annotation.get(GROUP_TOKEN_COUNT_BASIS_KEY) != TOKEN_COUNT_BASIS_VERSION: + return None return token_count if isinstance(token_count, int) else None @@ -520,6 +538,7 @@ def _write_token_count(message: Message, token_count: int) -> None: if annotation is None: return annotation[GROUP_TOKEN_COUNT_KEY] = token_count + annotation[GROUP_TOKEN_COUNT_BASIS_KEY] = TOKEN_COUNT_BASIS_VERSION message.additional_properties[GROUP_ANNOTATION_KEY] = annotation @@ -694,17 +713,54 @@ def annotate_message_groups( return _ordered_group_ids_from_annotations(messages) +_OPAQUE_REASONING_KEYS: Final[frozenset[str]] = frozenset({"encrypted_content"}) + + +def _strip_opaque_reasoning_payload(value: Any) -> Any: + """Recursively drop opaque reasoning members, keeping clear-text payloads.""" + if isinstance(value, dict): + filtered = { + key: _strip_opaque_reasoning_payload(item) + for key, item in value.items() + if key not in _OPAQUE_REASONING_KEYS + } + return filtered or None + if isinstance(value, list): + filtered = [item for item in (_strip_opaque_reasoning_payload(item) for item in value) if item is not None] + return filtered or None + return value + + def _serialize_content(content: Content) -> dict[str, Any]: payload = content.to_dict(exclude_none=True) payload.pop("raw_representation", None) # ``items`` mirrors ``result`` for function_result content; exclude it # to avoid double-counting tokens during estimation. payload.pop("items", None) - # ``protected_data`` holds provider-encrypted reasoning payloads (e.g. - # OpenAI ``encrypted_content``, Anthropic thinking ``signature``) that the - # model replays opaquely and never tokenises; exclude them so estimation - # measures the text the model actually sees. - payload.pop("protected_data", None) + # ``protected_data`` carries provider reasoning payloads. JSON-serialised + # reasoning_details (Chat Completions) are replayed to the provider as + # clear text -- ``summary``, ``reasoning_text`` and nested ``reasoning.text`` + # are part of the context the provider receives, so only their opaque + # members (``encrypted_content``) are excluded and the clear text stays + # counted. Anything that is not such a JSON structure (Anthropic thinking + # ``signature``, Responses API ``encrypted_content`` blobs) is replayed + # opaquely and never tokenised; exclude it so estimation measures the text + # the model actually sees. + protected_data = payload.get("protected_data") + if isinstance(protected_data, str) and protected_data: + try: + reasoning_payload = json.loads(protected_data) + except ValueError: + reasoning_payload = None + filtered = ( + _strip_opaque_reasoning_payload(reasoning_payload) if isinstance(reasoning_payload, (dict, list)) else None + ) + if filtered is None: + payload.pop("protected_data", None) + else: + payload["protected_data"] = json.dumps(filtered, ensure_ascii=False) + else: + payload.pop("protected_data", None) additional_properties = payload.get("additional_properties") if isinstance(additional_properties, dict) and "encrypted_content" in additional_properties: payload["additional_properties"] = { diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 9d3be5097f0..2e164e1cee1 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -472,7 +472,62 @@ def test_annotate_message_groups_token_counts_ignore_encrypted_content_in_additi assert _token_count(with_message) == _token_count(without_message) +def test_annotate_message_groups_token_counts_keep_clear_text_reasoning_details() -> None: + import json + + clear_text = "visible reasoning " * 128 + with_details = Content.from_text_reasoning( + id="rs_1", + text="", + protected_data=json.dumps([{"type": "reasoning.text", "text": clear_text, "encrypted_content": "A" * 4096}]), + ) + summary_only = Content.from_text_reasoning( + id="rs_1", + text="", + protected_data=json.dumps([{"type": "reasoning.text", "text": clear_text}]), + ) + without_details = Content.from_text_reasoning(id="rs_1", text="") + final_answer = Content.from_text("final answer") + + with_message = Message(role="assistant", contents=[with_details, final_answer]) + summary_only_message = Message(role="assistant", contents=[summary_only, final_answer]) + bare_message = Message(role="assistant", contents=[without_details, final_answer]) + + annotate_message_groups([with_message], tokenizer=CharacterEstimatorTokenizer()) + annotate_message_groups([summary_only_message], tokenizer=CharacterEstimatorTokenizer()) + annotate_message_groups([bare_message], tokenizer=CharacterEstimatorTokenizer()) + + # The opaque encrypted member is excluded while the clear-text reasoning + # the provider replays stays counted. + assert _token_count(with_message) == _token_count(summary_only_message) + assert _token_count(with_message) > _token_count(bare_message) + + +def test_annotate_token_counts_recomputes_counts_from_stale_serialization_basis() -> None: + + from agent_framework._compaction import TOKEN_COUNT_BASIS_VERSION, _write_group_annotation, annotate_token_counts + + message = Message(role="assistant", contents=[Content.from_text("final answer " * 64)]) + _write_group_annotation(message, group_id="g0", kind="assistant_text", index=0, has_reasoning=False) + annotation = message.additional_properties[GROUP_ANNOTATION_KEY] + assert isinstance(annotation, dict) + # Simulate a count cached by a pre-fix serialization basis: value present, + # no basis stamp (annotations survive to_dict/from_dict round-trips). + annotation[GROUP_TOKEN_COUNT_KEY] = 1 + annotation.pop("token_count_basis", None) + assert _token_count(message) == 1 + + annotate_token_counts([message], tokenizer=CharacterEstimatorTokenizer()) + + annotation = message.additional_properties[GROUP_ANNOTATION_KEY] + assert isinstance(annotation, dict) + assert annotation[GROUP_TOKEN_COUNT_KEY] != 1 + assert annotation["token_count_basis"] == TOKEN_COUNT_BASIS_VERSION + + def test_extend_compaction_messages_preserves_existing_annotations_and_tokens() -> None: + from agent_framework._compaction import TOKEN_COUNT_BASIS_VERSION + tokenizer = CharacterEstimatorTokenizer() messages = [_assistant_function_call("c3")] annotate_message_groups(messages) @@ -482,6 +537,10 @@ def test_extend_compaction_messages_preserves_existing_annotations_and_tokens() annotation = messages[0].additional_properties.get(GROUP_ANNOTATION_KEY) if isinstance(annotation, dict): annotation[GROUP_TOKEN_COUNT_KEY] = old_token_count + # Counts cached under the current serialization basis are preserved + # across extensions; counts without the current basis stamp are stale + # and get recomputed (see test_annotate_token_counts_recomputes_...). + annotation["token_count_basis"] = TOKEN_COUNT_BASIS_VERSION extend_compaction_messages(messages, [_tool_result("c3", "ok")], tokenizer=tokenizer) From 48b4e325c1d1b538acc08738944c39bfa6fbd127 Mon Sep 17 00:00:00 2001 From: Lubaoshuai <128781758+Lubaoshuai@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:04:41 +0800 Subject: [PATCH 3/3] fix(compaction): resolve pyright/mypy typing errors in token estimation - annotate the recursive reasoning-payload filter with explicit dict[Any, Any] / list[Any] casts so iterating the narrowed Any dict does not leak Unknown types under pyright strict - cast additional_properties before the encrypted_content filter - assert non-None token counts in the clear-text regression test so the ordering comparison does not involve optional integers --- .../core/agent_framework/_compaction.py | 25 +++++++++++++------ .../core/tests/core/test_compaction.py | 10 ++++++-- 2 files changed, 25 insertions(+), 10 deletions(-) diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index b69b7156775..6b292183717 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -757,15 +757,23 @@ def annotate_message_groups( def _strip_opaque_reasoning_payload(value: Any) -> Any: """Recursively drop opaque reasoning members, keeping clear-text payloads.""" if isinstance(value, dict): - filtered = { - key: _strip_opaque_reasoning_payload(item) - for key, item in value.items() - if key not in _OPAQUE_REASONING_KEYS - } + entries = cast("dict[Any, Any]", value) + filtered: dict[Any, Any] = {} + for key, item in entries.items(): + if key in _OPAQUE_REASONING_KEYS: + continue + stripped = _strip_opaque_reasoning_payload(item) + if stripped is not None: + filtered[key] = stripped return filtered or None if isinstance(value, list): - filtered = [item for item in (_strip_opaque_reasoning_payload(item) for item in value) if item is not None] - return filtered or None + entries = cast("list[Any]", value) + kept = [ + stripped + for stripped in (_strip_opaque_reasoning_payload(entry) for entry in entries) + if stripped is not None + ] + return kept or None return value @@ -801,8 +809,9 @@ def _serialize_content(content: Content) -> dict[str, Any]: payload.pop("protected_data", None) additional_properties = payload.get("additional_properties") if isinstance(additional_properties, dict) and "encrypted_content" in additional_properties: + typed_properties = cast("dict[str, Any]", additional_properties) payload["additional_properties"] = { - key: value for key, value in additional_properties.items() if key != "encrypted_content" + key: value for key, value in typed_properties.items() if key != "encrypted_content" } return payload diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 853cea3532e..c02eb5364b4 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -500,8 +500,14 @@ def test_annotate_message_groups_token_counts_keep_clear_text_reasoning_details( # The opaque encrypted member is excluded while the clear-text reasoning # the provider replays stays counted. - assert _token_count(with_message) == _token_count(summary_only_message) - assert _token_count(with_message) > _token_count(bare_message) + with_count = _token_count(with_message) + summary_count = _token_count(summary_only_message) + bare_count = _token_count(bare_message) + assert with_count is not None + assert summary_count is not None + assert bare_count is not None + assert with_count == summary_count + assert with_count > bare_count def test_annotate_token_counts_recomputes_counts_from_stale_serialization_basis() -> None: