Resolve shared context entries once per read batch

src/agent_comms/context_store.py

@@ -1,3 +1,3 @@
 """Append-only original turn contexts, independent of message delivery bytes."""
 
 from __future__ import annotations
@@ -146,9 +146,11 @@
             older = BeforeContextManifest(self.source, rows[0].id).available(db, incarnation)
             newer = AfterContextManifest(self.source, rows[-1].id).available(db, incarnation)
         retained = {item.record_id: item for item in previous if item.source == self.source}
+        resolved = {item.record_id: item for item in store.resolve(
+            db, self.source, (row for row in rows if row.id not in retained),
+        )}
         resources = tuple(
-            retained[row.id] if row.id in retained else store.resolve(db, self.source, row)
-            for row in rows
+            retained[row.id] if row.id in retained else resolved[row.id] for row in rows
         )
         return ContextManifestPage(resources, older, newer)
 
@@ -504,7 +506,7 @@
     ) -> RecordedContextState:
         rows = ContextRequest.newest(db, owner)
         if rows:
-            return self.resolve(db, source, rows[0])
+            return next(self.resolve(db, source, rows))
         return AbsentContext()
 
     def entries(
@@ -526,10 +528,12 @@
         with self.snapshot() as db:
             for entry in ContextEntry.stream(db):
                 yield from entry.manifest.original_values()
-            for request in ContextRequest.stream(db):
+            requests = tuple(ContextRequest.stream(db))
+            if requests:
                 source = ContextStoreInfo.require(db, singleton=1).source
-                for root in self.resolve(db, source, request).value.segments:
-                    yield from root.original_values()
+                for request in self.resolve(db, source, requests):
+                    for root in request.value.segments:
+                        yield from root.original_values()
 
     def retain_contexts(
         self, destination: ContextStore, identities: Iterable[NativeSessionIdentity], *,
@@ -584,14 +588,18 @@
                 if stored != reference.source:
                     raise RelationViolationError("Retained context reference names another source")
                 ContextRequest.require(retained, id=reference.record_id)
-            for row in ContextRequest.stream(retained):
+            rows = tuple(ContextRequest.stream(retained))
+            if rows:
                 original_source = ContextStoreInfo.require(source, singleton=1).source
                 retained_source = ContextStoreInfo.require(retained, singleton=1).source
-                original_row = ContextRequest.require(source, id=row.id)
-                expected = self.resolve(source, original_source, original_row)
-                actual = destination.resolve(retained, retained_source, row)
-                if row != original_row or actual != expected:
+                originals = tuple(ContextRequest.require(source, id=row.id) for row in rows)
+                if rows != originals:
                     raise RelationViolationError("Retained context reader differs from its original")
+                expected = self.resolve(source, original_source, originals)
+                actual = destination.resolve(retained, retained_source, rows)
+                for original, copied in zip(expected, actual, strict=True):
+                    if copied != original:
+                        raise RelationViolationError("Retained context reader differs from its original")
             for version in ContextSegmentVersion.stream(retained):
                 if version != ContextSegmentVersion.require(source, id=version.id):
                     raise RelationViolationError("Retained context version differs from its original")
@@ -619,25 +627,36 @@
         self,
         db: sqlite3.Connection,
         source: ContextStoreIdentity,
-        row: ContextRequest,
-    ) -> CapturedContext:
-        records = {row.id: row.recorded(db)}
+        rows: Iterable[ContextRequest],
+    ) -> Iterator[CapturedContext]:
+        """Resolve one snapshot's requested rows, decoding shared originals once.
+
+        Only selected rows and their direct referenced segments participate; the
+        maps die with this iterator and never cross a snapshot or retain history.
+        """
+        selected = tuple(rows)
+        originals = {row.id: row for row in selected}
+        records = {key: row.recorded(db) for key, row in originals.items()}
         references: set[ContextEntryReference] = set()
-        for position, part in enumerate(records[row.id].segments):
-            references.update(part.entries)
-            for predecessor in part.predecessors(row.id):
-                if predecessor not in records:
-                    prior = ContextRequest.require(db, id=predecessor)
+        for row in selected:
+            for position, part in enumerate(records[row.id].segments):
+                references.update(part.entries)
+                for predecessor in part.predecessors(row.id):
+                    if predecessor not in originals:
+                        originals[predecessor] = ContextRequest.require(db, id=predecessor)
+                    prior = originals[predecessor]
                     if prior.identity.thread != row.identity.thread:
                         raise RelationViolationError("Context predecessor is missing or foreign")
-                    records[predecessor] = prior.recorded(db)
-                # Only the referenced segment is needed, not the predecessor
-                # request's other segments and their unrelated dependencies.
-                prior_part = part.predecessor(records, position)
-                references.update(prior_part.entries)
+                    if predecessor not in records:
+                        records[predecessor] = prior.recorded(db)
+                    # Only the referenced segment is needed, not the predecessor
+                    # request's other segments and their unrelated dependencies.
+                    prior_part = part.predecessor(records, position)
+                    references.update(prior_part.entries)
         entries = {ref: ContextEntry.original(db, ref) for ref in references}
-        record = records[row.id]
-        return CapturedContext(source, row.id, record, record.restore(entries, records))
+        for row in selected:
+            record = records[row.id]
+            yield CapturedContext(source, row.id, record, record.restore(entries, records))
 
     def page(
         self,

tests/test_context_manifest_deltas.py

@@ -163,7 +163,7 @@ def test_forty_turns_use_direct_sources_and_fixed_transcript_keyframes(tmp_path)
         rows = ContextRequest.select(db, order_by=('id',))
         for turn in (1, 32, 33, 40):
             row = rows[turn - 1]
-            result = comms.contexts.resolve(db, source, row)
+            result, = comms.contexts.resolve(db, source, (row,))
             assert result.value.segments[:2] == expected[turn - 1].segments[:2]
             assert result.value.segments[2].contributors == expected[turn - 1].segments[2].contributors
             recorded = result.recorded
@@ -178,7 +178,7 @@ def observe(frame, event, argument):
                 decoded.append(frame.f_locals['self'].id)
         try:
             sys.setprofile(observe)
-            comms.contexts.resolve(db, source, rows[-1])
+            tuple(comms.contexts.resolve(db, source, (rows[-1],)))
         finally:
             sys.setprofile(previous_profile)
         assert len(decoded) == len(set(decoded)) == 3
@@ -188,6 +188,67 @@ def observe(frame, event, argument):
                 {}, {rows[32].id: delta}, position=2)
 
 
+def test_context_batch_decodes_shared_entries_once(tmp_path: Path) -> None:
+    comms = context_comms(tmp_path)
+    owner = comms.registry.require('Alice')
+    expected = tuple(context(owner, tmp_path, count, turn)
+                     for turn, count in enumerate((3, 4, 5), start=1))
+    prior = UNREAD_CONTEXT
+    for manifest in expected:
+        prior = comms.contexts.record_after(prior, manifest)
+    decoded = []
+    with comms.contexts.snapshot() as db:
+        source = ContextStoreInfo.require(db, singleton=1).source
+        rows = ContextRequest.select(db, order_by=('id',))
+        previous_profile = sys.getprofile()
+        def observe(frame, event, argument):
+            if event == 'call' and frame.f_code is ContextEntry.original.__func__.__code__:
+                decoded.append(frame.f_locals['reference'])
+        try:
+            sys.setprofile(observe)
+            actual = tuple(comms.contexts.resolve(db, source, rows))
+        finally:
+            sys.setprofile(previous_profile)
+    assert len(decoded) == len(set(decoded)) == 5
+    for original, restored in zip(expected, actual, strict=True):
+        assert restored.value.segments == original.segments
+
+
+@pytest.mark.parametrize('preloaded', (False, True))
+def test_context_batch_refuses_foreign_predecessor_even_when_preloaded(
+    tmp_path: Path, preloaded: bool,
+) -> None:
+    comms = context_comms(tmp_path)
+    alice, bob = (comms.registry.require(name) for name in ('Alice', 'Bob'))
+    foreign = comms.contexts.record_after(UNREAD_CONTEXT, context(bob, tmp_path, 1, 1))
+    original = comms.contexts.record_after(UNREAD_CONTEXT, context(alice, tmp_path, 1, 1))
+    record = replace(original.recorded, segments=(
+        replace(original.recorded.segments[0], previous=RecordedContextPredecessor(foreign.record_id)),
+        *original.recorded.segments[1:],
+    ))
+    with comms.contexts.writing() as db:
+        forged = ContextRequest.append(db, record)
+    with comms.contexts.snapshot() as db:
+        selected = ContextRequest.require(db, id=forged)
+        rows = ((ContextRequest.require(db, id=foreign.record_id), selected)
+                if preloaded else (selected,))
+        with pytest.raises(RelationViolationError, match='predecessor is missing or foreign'):
+            tuple(comms.contexts.resolve(db, original.source, rows))
+
+
+def test_context_batch_keeps_same_session_and_entry_on_different_paths(tmp_path: Path) -> None:
+    comms = context_comms(tmp_path)
+    owner = comms.registry.require('Alice')
+    originals = tuple(comms.contexts.record_after(
+        UNREAD_CONTEXT, context(owner, tmp_path / path, 1, turn),
+    ) for turn, path in enumerate(('one', 'two'), start=1))
+    with comms.contexts.snapshot() as db:
+        rows = ContextRequest.select(db, order_by=('id',))
+        actual = tuple(comms.contexts.resolve(db, originals[0].source, rows))
+    for original, restored in zip(originals, actual, strict=True):
+        assert restored == original
+
+
 def test_missing_store_identity_refuses_reads_and_further_records(tmp_path):
 
     comms = context_comms(tmp_path)
@@ -230,8 +291,8 @@ def test_selected_context_copy_preserves_closure_and_original_sequences(tmp_path
     selected = originals[1]
     reference = ContextManifestReference(selected.source, selected.record_id)
     with comms.contexts.snapshot() as db:
-        expected = comms.contexts.resolve(
-            db, reference.source, ContextRequest.require(db, id=reference.record_id),
+        expected, = comms.contexts.resolve(
+            db, reference.source, (ContextRequest.require(db, id=reference.record_id),),
         )
     before = ContextStore.digest(comms.contexts.path)
     copy = PracticeCopy(comms.root, tmp_path / 'retained')
@@ -244,7 +305,8 @@ def test_selected_context_copy_preserves_closure_and_original_sequences(tmp_path
         assert info.source == selected.source
         rows = ContextRequest.select(db, order_by=('id',))
         assert tuple(row.id for row in rows) == tuple(row.record_id for row in originals[:2])
-        assert retained.resolve(db, info.source, rows[-1]) == expected
+        actual, = retained.resolve(db, info.source, (rows[-1],))
+        assert actual == expected
         assert ContextEntry.count(db) == 2
         retained_version = max(row.id for row in ContextSegmentVersion.select(db))
     with comms.contexts.snapshot() as db: