perf(diagnosis): O(n) context truncation instead of O(n^2) (#70)

sentinel/diagnosis/truncation.py

@@ -26,6 +26,10 @@
 
 _SAFETY_MARGIN = 0.10
 _CHARS_PER_TOKEN = 4
+# Default json.dumps separator between list items (separators=(", ", ": ")).
+# Popping the last element of a list removes exactly this separator plus the
+# element's own serialization — the basis for the O(n) running-total below.
+_ITEM_SEP = ", "
 
 _DROP_ORDER: tuple[str, ...] = (
     "recent_logs",
@@ -49,7 +53,11 @@ def to_dict(self) -> dict[str, int]:
 
 
 def _estimate_tokens(blob_json: str) -> int:
-    return int(len(blob_json) / _CHARS_PER_TOKEN * (1 + _SAFETY_MARGIN))
+    return _estimate_tokens_for_chars(len(blob_json))
+
+
+def _estimate_tokens_for_chars(n_chars: int) -> int:
+    return int(n_chars / _CHARS_PER_TOKEN * (1 + _SAFETY_MARGIN))
 
 
 def truncate_for_budget(
@@ -71,18 +79,24 @@ def truncate_for_budget(
     blob: dict[str, Any] = json.loads(raw_json)
     _sort_sections_in_place(ctx, blob)
 
+    # Track the serialized length with a running total rather than re-dumping the
+    # whole blob on every pop (issue #70: that re-serialization was O(n^2) on
+    # exactly the large-context case truncation exists for). Popping the last
+    # element of a JSON list removes its leading ", " separator plus its own
+    # serialization, so the total stays byte-exact against json.dumps(blob).
+    current_chars = len(json.dumps(blob))
     dropped: dict[str, int] = {}
     while True:
         progress = False
         for section in _DROP_ORDER:
             items: list[Any] = blob[section]["data"]
             if len(items) <= 1:
                 continue
-            items.pop()
+            popped = items.pop()
             dropped[section] = dropped.get(section, 0) + 1
             progress = True
-            current_json = json.dumps(blob)
-            if _estimate_tokens(current_json) <= max_input_tokens:
+            current_chars -= len(_ITEM_SEP) + len(json.dumps(popped))
+            if _estimate_tokens_for_chars(current_chars) <= max_input_tokens:
                 return IncidentContext.model_validate(blob), TruncationStats(dropped=dropped)
         if not progress:
             # Every section is already at ≤1 item; nothing left to drop.

tests/unit/diagnosis/test_truncation.py

@@ -1,8 +1,10 @@
 # tests/unit/diagnosis/test_truncation.py
 from __future__ import annotations
 
+import json
 from datetime import UTC, datetime, timedelta
 
+import pytest
 from sentinel.diagnosis.truncation import truncate_for_budget
 
 from tests.unit.diagnosis.fakes import (
@@ -60,6 +62,38 @@ def test_deploys_section_always_keeps_newest() -> None:
     assert "deploy:new" in deploy_ids
 
 
+def test_whole_blob_serialized_at_most_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Truncation must not re-serialize the entire context once per dropped item
+    (issue #70 — that was O(n^2)). Assert the full blob is dumped at most once,
+    regardless of how many items get dropped. Item-sized dumps are fine."""
+    ctx = make_context(
+        deploys=[make_deploy()],
+        logs=[make_log(i) for i in range(120)],
+        runbooks=[make_runbook(), make_runbook()],
+        similar=[make_similar(), make_similar()],
+    )
+    real_dumps = json.dumps
+    blob_dumps = 0
+
+    def counting_dumps(obj: object) -> str:
+        nonlocal blob_dumps
+        # The whole-context blob is the dict carrying the section keys. Require two
+        # distinct sections so a future rename of one can't silently turn this
+        # discriminator into a no-op (which would pass <= 1 as a false green).
+        if isinstance(obj, dict) and "recent_logs" in obj and "recent_deploys" in obj:
+            blob_dumps += 1
+        return real_dumps(obj)
+
+    # truncate_for_budget calls json.dumps with a single positional arg, so a
+    # one-arg wrapper covers every call made during the patched window.
+    monkeypatch.setattr(json, "dumps", counting_dumps)
+    out, stats = truncate_for_budget(ctx, max_input_tokens=500)
+
+    assert stats.dropped.get("recent_logs", 0) > 50  # genuinely dropped many
+    assert len(out.recent_logs.data) >= 1
+    assert blob_dumps <= 1, f"whole blob serialized {blob_dumps}x (expected <=1)"
+
+
 def test_stats_to_dict_carries_per_section_counts() -> None:
     ctx = make_context(
         logs=[make_log(i) for i in range(100)],