Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/agent_env/task_step/context_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
27 changes: 26 additions & 1 deletion tst/unit/env/task_step/test_context_ops.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
"""Unit tests for context_ops — path-level diff + UpdateSpec compilation."""

import copy
import dataclasses

import pytest
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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:
Expand Down
Loading