From 310b0ee379e5d8b913024ce84fbf3f6fac0d46f7 Mon Sep 17 00:00:00 2001
From: morluto <76467478+morluto@users.noreply.github.com>
Date: Wed, 7 Oct 2026 03:20:10 +0800
Subject: [PATCH] perf(task): avoid quadratic response history diffs

---
 src/agent_env/task_step/context_ops.py     |  3 ++-
 tst/unit/env/task_step/test_context_ops.py | 27 +++++++++++++++++++++-
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/agent_env/task_step/context_ops.py b/src/agent_env/task_step/context_ops.py
index a80cd265..9835dcca 100644
--- a/src/agent_env/task_step/context_ops.py
+++ b/src/agent_env/task_step/context_ops.py
@@ -142,7 +142,8 @@ def build_context_update_ops(
     for name in _FIELDS.lists:
         pre_items = getattr(pre_ctx, name)
         post_items = getattr(post, name)
-        new_items = [x for x in post_items if x not in pre_items]
+        candidates = post_items[len(pre_items):] if _is_prefix(pre_items, post_items) else post_items
+        new_items = [x for x in candidates if x not in pre_items]
         if new_items:
             ops.add_to_sets[f"context.{name}"] = [
                 dataclasses.asdict(x) if dataclasses.is_dataclass(x) else x
diff --git a/tst/unit/env/task_step/test_context_ops.py b/tst/unit/env/task_step/test_context_ops.py
index 051a1145..64c49e10 100644
--- a/tst/unit/env/task_step/test_context_ops.py
+++ b/tst/unit/env/task_step/test_context_ops.py
@@ -1,5 +1,6 @@
 """Unit tests for context_ops — path-level diff + UpdateSpec compilation."""
 
+import copy
 import dataclasses
 
 import pytest
@@ -107,6 +108,31 @@ def test_agents_and_responses_lists_use_add_to_set():
     assert "context.prompt_responses" in ops.add_to_sets
 
 
+@pytest.mark.parametrize("append_new", [False, True])
+def test_response_history_diff_uses_linear_comparisons(monkeypatch, append_new):
+    pre = TaskStepContext(prompt_responses=[_response(f"p{i}") for i in range(100)])
+    post = copy.deepcopy(pre)
+    if append_new:
+        post.prompt_responses.append(_response("new"))
+
+    comparisons = 0
+    original_eq = PromptResponse.__eq__
+
+    def count_equal(self, other):
+        nonlocal comparisons
+        comparisons += 1
+        return original_eq(self, other)
+
+    monkeypatch.setattr(PromptResponse, "__eq__", count_equal)
+    ops = build_context_update_ops(pre, post)
+
+    if append_new:
+        assert ops.add_to_sets["context.prompt_responses"] == [dataclasses.asdict(_response("new"))]
+    else:
+        assert ops.is_empty()
+    assert comparisons <= 2 * len(pre.prompt_responses)
+
+
 def test_scalar_nil_to_value_emits_set():
     pre = TaskStepContext()
     post = TaskStepContext(agent_model="gpt-4")
@@ -325,7 +351,6 @@ def test_filename_keyed_output_urls_falls_back_to_wholesale_set():
 
 def _apply_ops_to_mirror(doc: dict, ops: ContextUpdateOps) -> dict:
     """Apply ops to an in-memory dict, modeling Mongo's $set / $unset / $addToSet."""
-    import copy
     out = copy.deepcopy(doc)
 
     def descend(d: dict, path: list[str]) -> dict: