Python: exclude encrypted reasoning payloads from compaction token counts

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"
@@ -496,7 +502,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
@@ -508,6 +520,7 @@ def _write_group_annotation(
                 GROUP_INDEX_KEY,
                 GROUP_HAS_REASONING_KEY,
                 GROUP_TOKEN_COUNT_KEY,
+                GROUP_TOKEN_COUNT_BASIS_KEY,
             }
         }
 
@@ -517,6 +530,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
@@ -550,6 +564,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
 
 
@@ -558,6 +576,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
 
 
@@ -732,12 +751,68 @@ 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):
+        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):
+        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
+
+
 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`` 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:
+        typed_properties = cast("dict[str, Any]", additional_properties)
+        payload["additional_properties"] = {
+            key: value for key, value in typed_properties.items() if key != "encrypted_content"
+        }
     return payload
 
 

python/packages/core/tests/core/test_compaction.py

@@ -435,7 +435,106 @@ 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_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.
+    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:
+
+    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)
@@ -445,6 +544,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)