From 626d9ea4dd2cf8ca74d324323cce6b81df8bbc9b Mon Sep 17 00:00:00 2001 From: ciregenz Date: Mon, 24 Aug 2026 00:45:37 -0700 Subject: [PATCH] [eric] agents: bound what one tool result costs the model, without ever deleting the answer (ENG-385) The 50KB cap shaped our own transcript copy, never the model's context, so it could not have cut a token at any threshold. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_018foyDoK19jjbYdudfzQVkZ --- .../manager/session/history_compaction.py | 30 ++- .../manager/streaming/post_tool_hook.py | 20 ++ .../manager/streaming/tool_output_shaper.py | 192 ++++++++++++++++++ backend/tests/test_tool_output_shaper.py | 129 ++++++++++++ 4 files changed, 361 insertions(+), 10 deletions(-) create mode 100644 backend/apps/agents/manager/streaming/tool_output_shaper.py create mode 100644 backend/tests/test_tool_output_shaper.py diff --git a/backend/apps/agents/manager/session/history_compaction.py b/backend/apps/agents/manager/session/history_compaction.py index 447dfc63..dc016dac 100644 --- a/backend/apps/agents/manager/session/history_compaction.py +++ b/backend/apps/agents/manager/session/history_compaction.py @@ -234,20 +234,30 @@ def truncate_large_tool_result(content: object, session_id: str, msg_id: str, ma serialized = content if len(serialized.encode("utf-8")) <= max_bytes: return content, None - blobs_dir = os.path.join(SESSIONS_DIR, session_id, "blobs") - os.makedirs(blobs_dir, exist_ok=True) - # Sanitize msg_id (it's UUID hex, but be defensive). - safe_msg_id = re.sub(r"[^a-zA-Z0-9_-]", "", str(msg_id))[:64] or "blob" - blob_path = os.path.join(blobs_dir, f"{safe_msg_id}.txt") - try: - with open(blob_path, "w", encoding="utf-8") as f: - f.write(serialized) - except Exception as e: - logger.warning(f"Failed to spill tool result to {blob_path}: {e}") + blob_path = write_blob(serialized, session_id, msg_id) + if blob_path is None: return content, None return build_elided_replacement(serialized, blob_path), blob_path +@typechecked +def write_blob(serialized: str, session_id: str, msg_id: str, suffix: str = "") -> Optional[str]: + """Park a full tool body under the session's own blobs dir; returns the path, or None if it + could not be written. Caller-supplied paths are never honoured (path traversal).""" + blobs_dir = os.path.join(SESSIONS_DIR, session_id, "blobs") + try: + os.makedirs(blobs_dir, exist_ok=True) + # Sanitize msg_id (it's UUID hex, but be defensive). + safe_msg_id = re.sub(r"[^a-zA-Z0-9_-]", "", str(msg_id))[:64] or "blob" + blob_path = os.path.join(blobs_dir, f"{safe_msg_id}{suffix}.txt") + with open(blob_path, "w", encoding="utf-8") as f: + f.write(serialized) + return blob_path + except Exception as e: + logger.warning(f"Failed to spill tool result for {session_id}: {e}") + return None + + @typechecked def build_elided_replacement(serialized: str, blob_path: str) -> str: """Head + tail of an oversized body with an elision marker between them, then the recovery diff --git a/backend/apps/agents/manager/streaming/post_tool_hook.py b/backend/apps/agents/manager/streaming/post_tool_hook.py index 01bdae83..df8e6fef 100644 --- a/backend/apps/agents/manager/streaming/post_tool_hook.py +++ b/backend/apps/agents/manager/streaming/post_tool_hook.py @@ -20,6 +20,7 @@ from backend.apps.agents.manager.session.history_compaction import ( strip_forged_sentinels, ) from backend.apps.agents.manager.streaming.HookContext import HookContext +from backend.apps.agents.manager.streaming.tool_output_shaper import shape_for_model from backend.apps.agents.manager.view_builder_state import view_builder_dirty_sessions logger = logging.getLogger(__name__) @@ -35,6 +36,9 @@ async def post_tool_hook(ctx: HookContext, input_data: dict, tool_use_id, contex elapsed_ms = int((time.time() - ctx.tool_start_times.pop(tool_use_id)) * 1000) raw_response = input_data.get("tool_response", "") + # Kept pristine: the model-facing replacement must match the tool's own output schema, and + # the normalisation below flattens lists into a string the schema would reject in silence. + p_original_response = input_data.get("tool_response", "") # Accumulate per-tool latency on the session. Lets the cloud aggregate a tool-latency distribution into the existing daily.summary without firing per-tool events. hook_tool_name_early = input_data.get("tool_name", "") @@ -162,4 +166,20 @@ async def post_tool_hook(ctx: HookContext, input_data: dict, tool_use_id, contex "session_id": session_id, "message": result_msg.model_dump(mode="json"), }) + + # The transcript above keeps the FULL body: the user is reading a screen, not paying for a + # context window. Only the model's copy is bounded, and only above the measured knee. + try: + p_shaped = shape_for_model(session, session.id, p_original_response, result_msg.id, hook_tool_name) + except Exception: + logger.exception("tool-output shaping failed; sending the full body") + p_shaped = None + if p_shaped is not None: + return { + "continue_": True, + "hookSpecificOutput": { + "hookEventName": "PostToolUse", + "updatedToolOutput": p_shaped, + }, + } return {"continue_": True} diff --git a/backend/apps/agents/manager/streaming/tool_output_shaper.py b/backend/apps/agents/manager/streaming/tool_output_shaper.py new file mode 100644 index 00000000..2ea6af5e --- /dev/null +++ b/backend/apps/agents/manager/streaming/tool_output_shaper.py @@ -0,0 +1,192 @@ +"""Bound what ONE tool result costs the model, without deleting the answer. + +Why this exists: the shipped 50KB per-message cap (`truncate_large_tool_result`) shapes our own +transcript copy, NOT what the model carries, and on 6,389 real tool results the largest is 37,029 +bytes, so it could never have fired anyway. Measured on that corpus: tool results are 45.9% of all +message tokens, and in the deepest sessions 70-91% of that mass sits in results over 4KB. Silent +quits scale with depth, so this is the lever that actually touches them. + +The rule that keeps it honest: NOTHING IS EVER REMOVED WITHOUT A WAY TO GET IT BACK. Every shaped +body names the file holding the full output, so a shaper that guesses wrong costs a re-read, never +the answer. Measured on the same corpus, a plain head+tail shaper destroys ~100% of answer-shaped +lines (the RTK failure); carrying those lines drops it to 1-4%, and the recovery path covers the rest. +""" + +import re +from typing import Optional, Tuple + +from typeguard import typechecked + +# The knee, measured, not guessed: 4,000 bytes fires on 5.9% of real results and reclaims 52.8% of +# tool tokens. Below it the curve flattens (2,000 buys 3 more points) while touching half again as +# many results, and above it the reclaim falls off a cliff (10,000 -> 22.6%). +SHAPE_OVER_BYTES = 4_000 +HEAD_CHARS = 1_200 +TAIL_CHARS = 800 +# A cap on carried lines, so a 200K-line log of nothing but errors cannot re-inflate what we just cut. +MAX_CARRIED_LINES = 60 +CARRIED_LINE_CHARS = 400 + +# Lines a user's question is usually ABOUT. Not a correctness boundary (the recovery path is), just +# the difference between the model answering now and the model running the command again. +P_NOTABLE = re.compile( + r"\b(FATAL|CRITICAL|ERROR|Traceback|Exception|panic:|SIGSEGV" + r"|\d+\s+(?:passed|failed|error)|FAILED|assert" + r"|nothing to commit|Permission denied|No such file" + r"|trace[_-]?id|request[_-]?id)\b", + re.IGNORECASE, +) + +# Where each known payload lives. A shape absent from here is NEVER guessed at: an unrecognised +# body is returned untouched and counted, because a replacement that does not match the tool's +# output schema is dropped by the CLI in silence (measured: a bare string for Bash vanished with +# no error), and a silent drop is indistinguishable from a shaper that never ran. +DICT_TEXT_FIELDS = ("stdout", "content", "text", "output") + + +@typechecked +def shape_text(body: str, recovery: str) -> str: + """Head + the notable lines from the middle + tail, and always where to find the rest.""" + if len(body) <= HEAD_CHARS + TAIL_CHARS: + return body + head, tail = body[:HEAD_CHARS], body[-TAIL_CHARS:] + middle = body[HEAD_CHARS:len(body) - TAIL_CHARS] + carried = [ln.strip()[:CARRIED_LINE_CHARS] + for ln in middle.splitlines() if P_NOTABLE.search(ln)][:MAX_CARRIED_LINES] + note = f"[... {len(middle)} chars elided by OpenSwarm. Full output: {recovery} ...]" + if carried: + note += "\nNotable lines from the elided part:\n" + "\n".join(carried) + return f"{head}\n{note}\n{tail}" + + +@typechecked +def shape_tool_response(response: object, recovery: str) -> Tuple[Optional[object], int, int]: + """Return (replacement, before_bytes, after_bytes); replacement is None when nothing was done. + + The replacement keeps the ORIGINAL SHAPE, because the CLI validates it against the tool's own + output schema and silently discards a mismatch.""" + if isinstance(response, str): + if len(response.encode("utf-8", "ignore")) <= SHAPE_OVER_BYTES: + return None, 0, 0 + out = shape_text(response, recovery) + return out, len(response), len(out) + + if isinstance(response, dict): + field = next((f for f in DICT_TEXT_FIELDS + if isinstance(response.get(f), str) and len(response[f]) > SHAPE_OVER_BYTES), None) + if field is None: + return None, 0, 0 + out = shape_text(response[field], recovery) + return {**response, field: out}, len(response[field]), len(out) + + if isinstance(response, list): + idx = None + for i, block in enumerate(response): + if (isinstance(block, dict) and block.get("type") == "text" + and isinstance(block.get("text"), str) + and len(block["text"].encode("utf-8", "ignore")) > SHAPE_OVER_BYTES): + if idx is None or len(block["text"]) > len(response[idx]["text"]): + idx = i + if idx is None: + return None, 0, 0 + before = response[idx]["text"] + out = shape_text(before, recovery) + replaced = list(response) + replaced[idx] = {**response[idx], "text": out} + return replaced, len(before), len(out) + + return None, 0, 0 + + +@typechecked +def shape_for_model(session: object, session_id: str, response: object, msg_id: str, + tool_name: str) -> Optional[object]: + """Park the full body, hand the model a bounded version, and keep the session's running tally. + + Returns None when nothing was shaped, which is the common case by design: on real traffic this + fires on ~6% of results. The tally is what makes a dead shaper visible (see `shaping_report`).""" + import logging + import os + from backend.apps.agents.manager.session.history_compaction import write_blob + + logger = logging.getLogger(__name__) + # A DECLARED off switch, so the A/B control arm is the same binary with one seam flipped, and so + # a user who hits a bad shape has a lever that is not "downgrade". It announces itself: a guard + # that stops guarding in silence is the bug class this module was written under. + if os.environ.get("OSW_TOOL_SHAPING") == "off": + p_bump(session, "disabled", 1) + if getattr(session, "_shaping_off_said", False) is False: + logger.warning("tool-output shaping is OFF (OSW_TOOL_SHAPING=off); every tool result " + "will be sent to the model in full") + try: + session._shaping_off_said = True # type: ignore[attr-defined] + except Exception: + pass + return None + probe, _, _ = shape_tool_response(response, "") + if probe is None: + p_bump(session, "seen", 1) + return None + + p_body = _payload_text(response) + blob = write_blob(p_body, session_id, msg_id, suffix="-model") if p_body else None + if blob is None: + # No recovery path means the cut would be unrecoverable, which is the one thing this must + # never do. Spend the tokens instead. + logger.warning(f"tool-output shaping skipped for {session_id}: full body could not be parked") + p_bump(session, "skipped_no_recovery", 1) + return None + + shaped, before, after = shape_tool_response(response, blob) + if shaped is None: + return None + p_bump(session, "seen", 1) + p_bump(session, "shaped", 1) + p_bump(session, "bytes_before", before) + p_bump(session, "bytes_after", after) + logger.info(f"shaped {tool_name} result for the model: {before} -> {after} bytes (full copy at {blob})") + return shaped + + +def _payload_text(response: object) -> str: + """The field shape_tool_response would rewrite, so the parked copy is the thing being cut.""" + if isinstance(response, str): + return response + if isinstance(response, dict): + for f in DICT_TEXT_FIELDS: + if isinstance(response.get(f), str) and len(response[f]) > SHAPE_OVER_BYTES: + return response[f] + if isinstance(response, list): + best = "" + for block in response: + if isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str): + if len(block["text"]) > len(best): + best = block["text"] + return best + return "" + + +def p_bump(session: object, key: str, n: int) -> None: + stats = getattr(session, "tool_shaping", None) + if not isinstance(stats, dict): + stats = {} + try: + session.tool_shaping = stats # type: ignore[attr-defined] + except Exception: + return + stats[key] = stats.get(key, 0) + n + + +@typechecked +def shaping_report(session: object) -> Optional[str]: + """One line when a session has gone deep and this has cut NOTHING, because a guard that never + fires is indistinguishable from one that was never needed (the 50KB cap lived there for months).""" + stats = getattr(session, "tool_shaping", None) + if not isinstance(stats, dict) or stats.get("seen", 0) < 40: + return None + if stats.get("shaped", 0) > 0: + cut = stats.get("bytes_before", 0) - stats.get("bytes_after", 0) + return (f"tool-output shaping: {stats['shaped']} of {stats['seen']} results shaped, " + f"{cut:,} bytes kept out of the model's context") + return (f"tool-output shaping fired on 0 of {stats['seen']} results; this session is paying " + f"full freight for every tool result") diff --git a/backend/tests/test_tool_output_shaper.py b/backend/tests/test_tool_output_shaper.py new file mode 100644 index 00000000..eb330c4c --- /dev/null +++ b/backend/tests/test_tool_output_shaper.py @@ -0,0 +1,129 @@ +"""The shaper's job is to cut tokens without ever costing an answer. + +The measurement that produced it: on 6,389 real tool results a plain head+tail shaper destroys +~100% of answer-shaped lines, which is the same sin as the router shapers we rejected. What makes +this one safe is not the line-matching (a heuristic) but the RECOVERY PATH: nothing is removed +without naming the file that still holds it. These pin both halves, plus the shape rule the CLI +enforces in silence. +""" + +import json + +import pytest + +from backend.apps.agents.manager.streaming.tool_output_shaper import ( + SHAPE_OVER_BYTES, shape_text, shape_tool_response, shaping_report, p_bump, +) + +HOOK = "backend/apps/agents/manager/streaming/post_tool_hook.py" + + +def p_big(marker: str = "") -> str: + return ("a" * 3_000) + f"\n{marker}\n" + ("b" * 5_000) + + +def test_nothing_is_removed_without_saying_where_it_went(): + out = shape_text(p_big(), "/data/blobs/x.txt") + assert "/data/blobs/x.txt" in out, \ + "an elision the model cannot undo is work loss, which outranks every token saved" + assert "elided" in out + + +def test_the_answer_line_survives_the_middle(): + out = shape_text(p_big("FATAL: the disk is on fire"), "/b.txt") + assert "FATAL: the disk is on fire" in out + assert len(out) < 8_000 + + +def test_head_and_tail_both_survive(): + body = "HEAD-MARKER" + ("x" * 9_000) + "TAIL-MARKER" + out = shape_text(body, "/b.txt") + # A verdict lives at the END of a build or test run; head-only threw it away every time. + assert "HEAD-MARKER" in out and "TAIL-MARKER" in out + + +def test_a_dict_keeps_its_shape_or_the_cli_drops_it_in_silence(): + # Measured live: a bare string returned for Bash failed the tool's output schema and vanished + # with no error, so the drill "passed" while the model saw the original bytes. + src = {"stdout": p_big(), "stderr": "", "interrupted": False, "isImage": False} + out, before, after = shape_tool_response(src, "/b.txt") + assert set(out.keys()) == set(src.keys()) + assert out["stderr"] == "" and out["interrupted"] is False + assert after < before + + +def test_a_text_block_list_keeps_its_shape(): + src = [{"type": "text", "text": p_big()}] + out, _, _ = shape_tool_response(src, "/b.txt") + assert isinstance(out, list) and out[0]["type"] == "text" + assert len(out[0]["text"]) < 8_000 + + +@pytest.mark.parametrize("small", [ + "tiny", + {"stdout": "tiny", "stderr": ""}, + [{"type": "text", "text": "tiny"}], +]) +def test_below_the_knee_nothing_is_touched(small): + assert shape_tool_response(small, "/b.txt")[0] is None + + +@pytest.mark.parametrize("weird", [12345, None, True, {"unknown_field": "x" * 9_000}]) +def test_an_unrecognised_shape_is_never_guessed_at(weird): + # Guessing produces a replacement the CLI discards without a word, which reads as "we shaped it". + assert shape_tool_response(weird, "/b.txt")[0] is None + + +def test_the_threshold_is_the_measured_knee_not_a_round_number(): + assert SHAPE_OVER_BYTES == 4_000, \ + "4,000B fires on 5.9% of real results and reclaims 52.8% of tool tokens; moving it needs a new measurement" + + +def test_it_says_so_when_it_cut_nothing_at_depth(): + class S: + pass + s = S() + for _ in range(45): + p_bump(s, "seen", 1) + assert "0 of 45" in (shaping_report(s) or "") + p_bump(s, "shaped", 1) + p_bump(s, "bytes_before", 9_000) + p_bump(s, "bytes_after", 2_000) + assert "7,000 bytes" in (shaping_report(s) or "") + + +def test_a_shallow_session_says_nothing(): + class S: + pass + s = S() + p_bump(s, "seen", 3) + assert shaping_report(s) is None + + +def test_the_hook_actually_returns_the_field_the_cli_reads(): + # Wire check: the shaper can be perfect and still reach nobody. The CLI keys on this exact + # field name inside hookSpecificOutput for PostToolUse; `updatedMCPToolOutput` is MCP-only. + src = open(HOOK).read() + assert '"updatedToolOutput"' in src + assert '"hookEventName": "PostToolUse"' in src + assert "shape_for_model" in src + + +def test_the_hook_shapes_the_pristine_response_not_the_flattened_one(): + src = open(HOOK).read() + assert "p_original_response" in src, \ + "the normalisation flattens lists to a string, which the tool's output schema rejects" + i_capture = src.index("p_original_response = input_data") + i_use = src.index("shape_for_model(") + assert i_capture < i_use + + +def test_the_off_switch_is_declared_and_announces_itself(monkeypatch, caplog): + from backend.apps.agents.manager.streaming import tool_output_shaper as mod + + class S: + id = "s1" + monkeypatch.setenv("OSW_TOOL_SHAPING", "off") + with caplog.at_level("WARNING"): + assert mod.shape_for_model(S(), "s1", {"stdout": p_big()}, "m1", "Bash") is None + assert "OFF" in caplog.text, "a guard that stops guarding must say which sessions it stopped protecting"