diff --git a/src/agent_comms/context_store.py b/src/agent_comms/context_store.py index 8aa26590c7..c561396af4 100644 --- a/src/agent_comms/context_store.py +++ b/src/agent_comms/context_store.py @@ -146,9 +146,11 @@ def page( 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 @@ def latest( ) -> 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 @@ def original_segments(self) -> Iterator[SegmentManifest]: 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 @@ def require_retained( 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 @@ def resolve( 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, diff --git a/tests/test_context_manifest_deltas.py b/tests/test_context_manifest_deltas.py index b4d2d827b1..4a0fbeb17f 100644 --- a/tests/test_context_manifest_deltas.py +++ b/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: