Compaction follow-up: JSON-serialise shrunk tool returns, strengthen guard test, DRY threshold (#1824) (#1853)

* Compaction follow-up: JSON-serialise shrunk tool returns, strengthen guard test, DRY threshold (#1824)

Code-review follow-up for the in-run history compaction processor (#1817).

- Serialise non-string ToolReturnPart content with json.dumps (not str())
  when shrinking, via a shared _stringify_tool_content helper that is reused
  by the token estimator so the pre-shrink estimate and post-shrink payload
  are serialised identically. Matches the existing
  json.dumps(..., default=str) convention used for tool-result content
  elsewhere in pydantic_ai_agents.py. (finding #2)
- Strengthen test_thinking_only_modelresponse_is_not_emptied: place a large
  ToolReturnPart in the older prefix so the empty-parts guard -- not the
  no-op early return -- is the verified reason the thinking-only ModelResponse
  survives. Add test_non_string_tool_return_serialized_as_json to pin the
  JSON-serialisation fix. (finding #1)
- DRY: new context_window_and_threshold() is the single definition of the
  compaction kick-in point, shared verbatim by should_compact,
  compact_message_history, and the in-run shrink processor (replaces three
  copies of int(get_context_window_for_model(...) * ratio)).
- Document that IN_RUN_TOOL_RETURN_TARGET_CHARS / _shrink_tool_return_part
  bound the preserved prefix, not the final string (the trim notice adds a
  small fixed overhead) (finding #3), and add a comment confirming the
  _on_in_run_shrink closure's captured message IDs are stable for the turn
  (finding #5).

Findings #4 (test reaching into a pydantic-ai internal) intentionally left as
documented in the issue -- no change needed.

* Address review: guard RecursionError in tool-content stringify, drop stale line refs

- _stringify_tool_content now also catches RecursionError so a (pathological)
  self-referential tool return falls back to str() instead of crashing the
  shrink/estimate — the stated 'never crash on exotic content' contract.
- Drop the hardcoded 'lines ~866-889' from the closure-stability comment in
  pydantic_ai_agents.py (would drift); say 'earlier in this function'.
- Add the symmetric positive assertion (json null present) to
  test_non_string_tool_return_serialized_as_json.

* Address PR #1853 review: clarify _stringify_tool_content except + 'None' assertion comments

- Keep the broad except in _stringify_tool_content but correct the
  misleading comment: TypeError/ValueError/RecursionError are all genuinely
  reachable (default=str does not rescue non-string dict keys, NaN/Inf with a
  strict encoder, or circular refs), so narrowing to RecursionError would
  turn a graceful fallback into a crash (review nit #1).
- Document why 'None' not in part.content is safe: the only non-JSON text
  spliced in is the fixed trim notice, which contains no literal 'None'
  (review nit #2).

---------

Co-authored-by: Claude <noreply@anthropic.com>

Author: JSv4

CHANGELOG.md

@@ -111,6 +111,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- **In-run history compaction follow-up (issue #1824).** Code-review fixes for
+  the pydantic-ai `shrink_old_artifacts_processor` shipped in #1817.
+  - **Structured tool returns now shrink to valid JSON**
+    (`opencontractserver/llms/history_processors.py`). `ToolReturnPart.content`
+    can be a dict/list; the shrink path stringified it with `str()` (Python
+    repr — single-quoted keys, `True`/`None`), which the model can misparse. A
+    new `_stringify_tool_content` helper serialises non-strings with
+    `json.dumps(content, default=str)` (falling back to `str()` only if JSON
+    serialisation fails) and is shared by both the token estimator and the
+    shrink, so the pre-shrink estimate matches the post-shrink payload.
+  - **Strengthened `test_thinking_only_modelresponse_is_not_emptied`**
+    (`opencontractserver/tests/test_history_processors.py`). The old fixture had
+    no shrinkable content in the older prefix, so the test passed via the no-op
+    early-return rather than the empty-parts guard it names. It now places a
+    large `ToolReturnPart` alongside the thinking-only response so the shrink
+    actually runs, and asserts (via telemetry) that the guard — not an early
+    return — is why the `ThinkingPart` survives. Added
+    `test_non_string_tool_return_serialized_as_json` to pin the JSON fix.
 - **MCP `WWW-Authenticate` base-URL hardening** (`opencontractserver/mcp/server.py`)
   — the 401 challenge now prefers the trusted `MCP_PUBLIC_BASE_URL` over the
   request `Host` header (MCP bypasses `ALLOWED_HOSTS`), falling back to the
@@ -234,6 +252,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- **DRY: single definition of the compaction kick-in point** (issue #1824)
+  (`opencontractserver/llms/context_guardrails.py`). New
+  `context_window_and_threshold(model_name, threshold_ratio)` replaces the
+  `int(get_context_window_for_model(...) * ratio)` triplication across
+  `should_compact`, `compact_message_history`, and the in-run
+  `shrink_old_artifacts_processor`, so turn-level and in-run compaction derive
+  the threshold identically from one place. Documented (finding #3) that
+  `IN_RUN_TOOL_RETURN_TARGET_CHARS` bounds the preserved prefix, not the final
+  string — the appended trim notice adds a small fixed overhead — and added a
+  comment confirming the `_on_in_run_shrink` closure's captured message IDs are
+  stable for the turn (finding #5).
 - **`BaseService.filter_visible_qs` now fails closed** (`opencontractserver/shared/services/base.py`)
   — an input lacking a `visible_to_user` method previously passed through
   **unfiltered** (fail-open, a latent row-leak); it now raises `TypeError`.

opencontractserver/constants/context_guardrails.py

@@ -168,6 +168,12 @@
 # trim notice is appended. Chosen so a shrunk return still carries enough
 # signal for the model to recall what the tool produced without dominating
 # the budget.
+#
+# This bounds the preserved *prefix*, not the final string: the appended
+# trim notice (see ``history_processors._TRIM_NOTICE_TEMPLATE``, ~40 chars
+# plus the digits of the elided count) makes the post-shrink content
+# slightly longer than this value. Negligible at the default 4K, but
+# callers tuning this down to a tight ceiling should budget for the notice.
 IN_RUN_TOOL_RETURN_TARGET_CHARS: int = 4_000
 
 # Whether the in-run processor strips ThinkingPart instances from messages

opencontractserver/llms/agents/pydantic_ai_agents.py

@@ -991,6 +991,12 @@ async def _stream_core(
         # compaction fires. Non-streaming (_chat_raw) never sets this
         # callback, so those chats get log-only telemetry from the
         # processor itself.
+        #
+        # The closure captures ``user_msg_id`` / ``llm_msg_id`` by Python
+        # late-binding, but both are resolved earlier in this function and
+        # never reassigned afterwards: ``_stream_core`` handles a single
+        # user turn, so the IDs are stable for the closure's lifetime and
+        # every in-loop shrink in this turn tags the same message pair.
         def _on_in_run_shrink(event: Any) -> None:
             try:
                 thought_evt = ThoughtEvent(

opencontractserver/llms/context_guardrails.py

@@ -44,6 +44,7 @@
 __all__ = [
     "estimate_token_count",
     "get_context_window_for_model",
+    "context_window_and_threshold",
     "truncate_tool_output",
     "cap_summary_length",
     "strip_compaction_prefix",
@@ -122,6 +123,24 @@ def get_context_window_for_model(model_name: str) -> int:
     return DEFAULT_CONTEXT_WINDOW
 
 
+def context_window_and_threshold(
+    model_name: str,
+    threshold_ratio: float = COMPACTION_THRESHOLD_RATIO,
+) -> tuple[int, int]:
+    """Return ``(context_window, threshold)`` for *model_name*.
+
+    ``threshold`` is ``int(context_window * threshold_ratio)`` — the single
+    definition of "the estimated-token count at which compaction kicks in".
+    Shared verbatim by turn-level compaction (:func:`should_compact`,
+    :func:`compact_message_history`) and the in-run history processor
+    (``opencontractserver.llms.history_processors.shrink_old_artifacts_processor``)
+    so both layers derive the kick-in point identically from the same
+    ``threshold_ratio``.
+    """
+    window = get_context_window_for_model(model_name)
+    return window, int(window * threshold_ratio)
+
+
 # ---------------------------------------------------------------------------
 # Tool output truncation
 # ---------------------------------------------------------------------------
@@ -220,13 +239,12 @@ def should_compact(
     prompt + stored summary + all messages) exceeds *threshold_ratio*
     of the model's context window.
     """
-    context_window = get_context_window_for_model(model_name)
     total_tokens = (
         system_prompt_tokens
         + stored_summary_tokens
         + sum(m.token_estimate for m in messages)
     )
-    threshold = int(context_window * threshold_ratio)
+    _, threshold = context_window_and_threshold(model_name, threshold_ratio)
     return total_tokens > threshold
 
 
@@ -276,8 +294,7 @@ def compact_message_history(
     """
     message_tokens = sum(m.token_estimate for m in messages)
     total_before = system_prompt_tokens + stored_summary_tokens + message_tokens
-    context_window = get_context_window_for_model(model_name)
-    threshold = int(context_window * threshold_ratio)
+    _, threshold = context_window_and_threshold(model_name, threshold_ratio)
 
     if total_before <= threshold:
         return CompactionResult(

opencontractserver/llms/history_processors.py

@@ -22,6 +22,7 @@
 
 from __future__ import annotations
 
+import json
 import logging
 from dataclasses import dataclass, replace
 from typing import Any
@@ -40,8 +41,8 @@
 )
 from opencontractserver.llms.context_guardrails import (
     CompactionConfig,
+    context_window_and_threshold,
     estimate_token_count,
-    get_context_window_for_model,
 )
 
 logger = logging.getLogger(__name__)
@@ -84,18 +85,44 @@ class InRunShrinkEvent:
 )
 
 
+def _stringify_tool_content(content: Any) -> str:
+    """Serialise (possibly non-string) tool-return content to a string.
+
+    ``ToolReturnPart.content`` is typed ``Any`` — tools may return dicts or
+    lists. We serialise those with ``json.dumps`` rather than ``str`` so the
+    model receives valid JSON (double-quoted keys, ``true``/``null`` literals)
+    instead of Python ``repr`` output it can misparse. ``default=str`` absorbs
+    stray non-JSON-serialisable values (datetimes, Decimals); if ``json.dumps``
+    still fails we fall back to ``str`` so the shrink/estimate never crashes on
+    exotic content.
+
+    Used by both the token estimator and the actual shrink, so the
+    pre-shrink estimate and the post-shrink payload are serialised
+    identically.
+    """
+    if isinstance(content, str):
+        return content
+    try:
+        return json.dumps(content, default=str)
+    except (TypeError, ValueError, RecursionError):
+        # All three are genuinely reachable, not defensive padding:
+        #   - TypeError: ``default=str`` only rescues non-serialisable *values*;
+        #     a dict with a non-string/number key (e.g. a tuple) still raises.
+        #   - ValueError: out-of-range floats (NaN/Infinity) with a strict
+        #     encoder, etc.
+        #   - RecursionError: a self-referential (circular) tool return.
+        # The fallback intent is "never crash on exotic content", so str() it.
+        return str(content)
+
+
 def _estimate_total_tokens(messages: list[ModelMessage], system_prompt: str) -> int:
     """Cheap upper-bound token estimate for the messages + system prompt."""
     total = estimate_token_count(system_prompt or "")
     for msg in messages:
         for part in getattr(msg, "parts", []):
             content = getattr(part, "content", None)
-            if isinstance(content, str):
-                total += estimate_token_count(content)
-            elif content is not None:
-                # Tool returns can be non-string (lists/dicts); stringify
-                # cheaply for estimation.
-                total += estimate_token_count(str(content))
+            if content is not None:
+                total += estimate_token_count(_stringify_tool_content(content))
     return total
 
 
@@ -157,13 +184,18 @@ def _shrink_tool_return_part(
 ) -> tuple[ToolReturnPart, bool]:
     """Return ``(maybe_shrunk_part, was_shrunk)``.
 
-    For non-string contents (dict/list), stringifies first so the
-    truncation arithmetic is uniform; the shrunk copy carries the
-    stringified, truncated form so the model sees a string.
+    For non-string contents (dict/list), serialises first (see
+    :func:`_stringify_tool_content`) so the truncation arithmetic is
+    uniform; the shrunk copy carries the serialised, truncated form so the
+    model sees a string.
+
+    ``target_chars`` bounds the preserved *prefix*, not the final string:
+    the returned content is up to ``target_chars`` plus the length of the
+    appended trim notice (``_TRIM_NOTICE_TEMPLATE`` — ~40 chars plus the
+    digits of the elided count). Callers setting a tight ceiling should
+    budget for that overhead.
     """
-    content = part.content
-    if not isinstance(content, str):
-        content = str(content)
+    content = _stringify_tool_content(part.content)
 
     if len(content) <= target_chars:
         return part, False
@@ -240,15 +272,18 @@ async def shrink_old_artifacts_processor(
             return messages
 
         tokens_before = _estimate_total_tokens(messages, system_prompt)
-        context_window = get_context_window_for_model(model_name)
         # ``threshold_ratio`` is shared by design with the outer
         # turn-level compaction (in ``MessageHistoryService``). A single
         # ratio keeps the kick-in point coherent across both layers — if
         # the outer pass already wants to compact persisted history,
         # the in-loop pass should also be willing to shrink older
         # tool returns. Do not split into two ratios without rethinking
-        # both call sites together.
-        threshold = int(context_window * cfg.threshold_ratio)
+        # both call sites together. ``context_window_and_threshold`` is
+        # the single definition of that kick-in point, shared verbatim
+        # with the turn-level compaction functions.
+        context_window, threshold = context_window_and_threshold(
+            model_name, cfg.threshold_ratio
+        )
 
         if tokens_before <= threshold:
             return messages

opencontractserver/tests/test_history_processors.py

@@ -491,30 +491,143 @@ def _raises(_event: InRunShrinkEvent) -> None:
 
 
 def test_thinking_only_modelresponse_is_not_emptied():
-    """A ModelResponse with only ThinkingPart is left intact (no empty parts)."""
+    """A ModelResponse whose only part is a ThinkingPart survives the
+    drop-thinking pass because of the empty-parts guard.
+
+    The fixture deliberately places a *second*, shrinkable artifact in the
+    older prefix (a large ``ToolReturnPart``). Without it the processor
+    would hit the ``tool_returns_shrunk == 0 and thinking_parts_dropped ==
+    0`` early-return and leave the whole list untouched — so the
+    thinking-only response would survive simply because *nothing* ran, not
+    because the guard fired. The big tool return forces the real shrink
+    path: the processor does modify the older prefix this pass, and the
+    guard is the reason the thinking-only ModelResponse alone is spared.
+    """
     # An old ModelResponse with only a giant ThinkingPart.
     only_thinking = ModelResponse(parts=[ThinkingPart(content="t" * 600_000)])
-    # Pair it with a ModelRequest so the structure is valid.
-    old_req = ModelRequest(
-        parts=[UserPromptPart(content="paired with thinking-only response")]
+    # A large, shrinkable tool return alongside it in the older prefix so
+    # the processor has something to actually shrink this pass.
+    big_return_req = ModelRequest(
+        parts=[
+            ToolReturnPart(
+                tool_name="load_document_text",
+                content="r" * 600_000,
+                tool_call_id="BIG",
+            )
+        ]
     )
     recent_pairs: list[ModelMessage] = []
     for i in range(5):
         r1, r2 = _make_pair(tool_call_id=f"R{i}", tool_name="ping", return_chars=50)
         recent_pairs.extend([r1, r2])
     messages: list[ModelMessage] = [
         only_thinking,
-        old_req,
+        big_return_req,
         *recent_pairs,
         ModelRequest(parts=[UserPromptPart(content="continue")]),
     ]
 
     deps = _FakeDeps()
     result = _run(messages, deps)
 
-    # The thinking-only ModelResponse is left intact (we didn't strip the
-    # only thing it contained).
+    # (a) The thinking-only ModelResponse is left intact — the guard refused
+    #     to strip its only part even though in_run_drop_thinking is on.
     first = result[0]
     assert isinstance(first, ModelResponse)
     assert len(first.parts) == 1
     assert isinstance(first.parts[0], ThinkingPart)
+
+    # (b) The shrink genuinely ran this pass (not the no-op early return):
+    #     the large sibling tool return WAS truncated.
+    big_return = result[1]
+    assert isinstance(big_return, ModelRequest)
+    big_part = big_return.parts[0]
+    assert isinstance(big_part, ToolReturnPart)
+    assert isinstance(big_part.content, str)
+    assert len(big_part.content) < 5_000
+    assert IN_RUN_TRIM_NOTICE_MARKER in big_part.content
+    assert big_part.tool_call_id == "BIG"
+
+    # (c) Telemetry confirms the guard — not an early return — is why the
+    #     ThinkingPart survived: a shrink event fired with the tool return
+    #     counted but zero thinking parts dropped (the only droppable
+    #     ThinkingPart was the one the guard protected).
+    assert len(deps.events) == 1
+    evt = deps.events[0]
+    assert evt.tool_returns_shrunk == 1
+    assert evt.thinking_parts_dropped == 0
+
+
+def test_non_string_tool_return_serialized_as_json():
+    """A structured (dict) ToolReturnPart is serialised with ``json.dumps``,
+    not ``str()``, when shrunk.
+
+    Tools may return dicts/lists. ``str()`` on those yields Python repr
+    (single-quoted keys, ``True``/``None`` literals) the model can
+    misparse; ``json.dumps`` yields valid JSON. This pins that the shrunk
+    content is JSON-shaped.
+    """
+    # A structured return whose JSON serialisation alone exceeds the
+    # threshold, so it is both the over-threshold trigger and the thing
+    # that gets shrunk. ``True``/``None`` round-trip differently under
+    # json.dumps ("true"/"null") vs str() ("True"/"None"), giving the
+    # assertion a clean discriminator.
+    big_payload = {
+        "flag": True,
+        "missing": None,
+        "blob": "v" * 700_000,
+    }
+    old_resp = ModelResponse(
+        parts=[
+            TextPart(content="step"),
+            ToolCallPart(tool_name="search", args="{}", tool_call_id="JSON"),
+        ]
+    )
+    old_req = ModelRequest(
+        parts=[
+            ToolReturnPart(
+                tool_name="search",
+                content=big_payload,
+                tool_call_id="JSON",
+            )
+        ]
+    )
+    recent_pairs: list[ModelMessage] = []
+    for i in range(5):
+        r1, r2 = _make_pair(tool_call_id=f"R{i}", tool_name="ping", return_chars=50)
+        recent_pairs.extend([r1, r2])
+    messages: list[ModelMessage] = [
+        ModelRequest(parts=[UserPromptPart(content="start")]),
+        old_resp,
+        old_req,
+        *recent_pairs,
+        ModelRequest(parts=[UserPromptPart(content="continue")]),
+    ]
+
+    deps = _FakeDeps()
+    result = _run(messages, deps)
+
+    shrunk = result[2]
+    assert isinstance(shrunk, ModelRequest)
+    part = shrunk.parts[0]
+    assert isinstance(part, ToolReturnPart)
+    # Serialised to a string and truncated with the trim notice.
+    assert isinstance(part.content, str)
+    assert IN_RUN_TRIM_NOTICE_MARKER in part.content
+    # JSON serialisation, not Python repr: object opens with ``{``, keys are
+    # double-quoted, and there are no single-quoted Python-repr keys.
+    assert part.content.startswith("{")
+    assert '"flag"' in part.content
+    assert "'flag'" not in part.content
+    # The booleans/null are JSON tokens, never Python ``True``/``None``.
+    assert "true" in part.content
+    assert "True" not in part.content
+    assert "null" in part.content  # JSON null, not Python None
+    # Safe because the only non-JSON text spliced into the payload is the trim
+    # notice (IN_RUN_TRIM_NOTICE_MARKER), whose template is fixed numeric/text
+    # and contains no literal "None"; if that template ever changes to include
+    # the word, tighten this to assert against the JSON prefix only.
+    assert "None" not in part.content
+    # tool_call_id correlation preserved across the serialise+shrink.
+    assert part.tool_call_id == "JSON"
+    assert deps.events and deps.events[0].tool_returns_shrunk == 1