perf(task): reuse keys during journal union

src/agent_env/task/step_journal.py

@@ -29,11 +29,15 @@ def _union(existing: list, items: list, *, move_to_end: bool = False, path: str
     """Union by value (deployed envs by instance), order-stable. ``move_to_end`` is the re-record case: that step's
     completion moves to the end of the replay order, so its items must move with it."""
     key = _item_key(path)
-    fresh = [x for i, x in enumerate(items) if key(x) not in [key(y) for y in items[:i]]]
-    fresh_keys = [key(x) for x in fresh]
+    fresh, fresh_keys = [], []
+    for item in items:
+        item_key = key(item)
+        if item_key not in fresh_keys:
+            fresh.append(item)
+            fresh_keys.append(item_key)
+    existing_keys = [key(item) for item in existing]
     if move_to_end:
-        return [x for x in existing if key(x) not in fresh_keys] + fresh
-    existing_keys = [key(x) for x in existing]
+        return [item for item, item_key in zip(existing, existing_keys) if item_key not in fresh_keys] + fresh
     return list(existing) + [x for x, k in zip(fresh, fresh_keys) if k not in existing_keys]
 
 

tst/benchmarks/journal_union.py

@@ -0,0 +1,80 @@
+"""Compare journal list union with a git revision using the production function bodies.
+
+Run from the repository root, for example:
+    python tst/benchmarks/journal_union.py --base origin/main --repeats 5
+"""
+
+from __future__ import annotations
+
+import argparse
+import ast
+import statistics
+import subprocess
+import time
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[2]
+TARGET = "src/agent_env/task/step_journal.py"
+
+
+def positive_int(value: str) -> int:
+    try:
+        result = int(value)
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError("must be a positive integer") from exc
+    if result <= 0:
+        raise argparse.ArgumentTypeError("must be a positive integer")
+    return result
+
+
+def load_union(source: str):
+    tree = ast.parse(source)
+    wanted = {"_union", "_item_key"}
+    nodes = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in wanted]
+    if {node.name for node in nodes} != wanted:
+        raise ValueError("source does not define both _union and _item_key")
+    namespace = {}
+    module = ast.Module(body=nodes, type_ignores=[])
+    exec(compile(ast.fix_missing_locations(module), "<step_journal>", "exec"), namespace)
+    return namespace["_union"]
+
+
+def median_ms(fn, repeats: int) -> float:
+    samples = []
+    for _ in range(repeats):
+        start = time.perf_counter()
+        fn()
+        samples.append((time.perf_counter() - start) * 1000)
+    return statistics.median(samples)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--base", default="origin/main", help="git revision to compare against (default: origin/main)")
+    parser.add_argument("--repeats", type=positive_int, default=5, help="timing samples per implementation and size (default: 5)")
+    args = parser.parse_args()
+
+    baseline_source = subprocess.check_output(
+        ["git", "show", f"{args.base}:{TARGET}"], cwd=ROOT, text=True
+    )
+    candidate_source = (ROOT / TARGET).read_text()
+    baseline, candidate = load_union(baseline_source), load_union(candidate_source)
+
+    print(f"Median runtime; production _union/_item_key function bodies extracted from {args.base} and the worktree")
+    print("items  baseline_ms  worktree_ms  speedup")
+    for size in (100, 1000, 3000):
+        items = [
+            {"instance_id": str(i), "sandbox_id": str(i), "card": {"tools": ["a", "b"]}}
+            for i in range(size)
+        ]
+        expected = baseline([], items, path="context.deployed_envs")
+        actual = candidate([], items, path="context.deployed_envs")
+        if actual != expected:
+            raise AssertionError(f"union output differs at {size} items")
+        before = median_ms(lambda: baseline([], items, path="context.deployed_envs"), args.repeats)
+        after = median_ms(lambda: candidate([], items, path="context.deployed_envs"), args.repeats)
+        print(f"{size:>5}  {before:>11.3f}  {after:>11.3f}  {before / after:>6.2f}x")
+
+
+if __name__ == "__main__":
+    main()

tst/unit/task/step_journal_test.py

@@ -9,6 +9,7 @@
 
 import pytest
 
+import agent_env.task.step_journal as journal_mod
 import agent_env.task.store as store_mod
 from agent_env.store import Filter, LocalSqliteDocumentStore, Sort
 from agent_env.task.step_journal import _union, commit_ordered, replay_context
@@ -574,6 +575,30 @@ def test_a_deployed_env_written_by_two_kernels_stays_one_entry():
     assert _union([old], [new], path="context.prompt_responses") == [old, new]  # every other list stays by value
 
 
+def test_union_projects_each_item_key_once(monkeypatch):
+    calls = 0
+
+    def item_key(path):
+        def project(item):
+            nonlocal calls
+            calls += 1
+            return item
+
+        return project
+
+    monkeypatch.setattr(journal_mod, "_item_key", item_key)
+    existing = [["old"], ["shared"]]
+    items = [["shared"], ["new"], ["new"]]
+
+    assert _union(existing, items) == [["old"], ["shared"], ["new"]]
+    assert calls == len(existing) + len(items)
+
+    calls = 0
+    moved_items = [["old"]]
+    assert _union(existing, moved_items, move_to_end=True) == [["shared"], ["old"]]
+    assert calls == len(existing) + len(moved_items)
+
+
 def test_a_completion_that_re_adds_a_deployed_env_keeps_one_entry(store):
     iid = _new(store)
     old = {"env_id": "e", "instance_id": "i1", "mcp_server_name": None}