diff --git a/src/powercontext/builtin/persistence/migrations/atomic_memory_references.py b/src/powercontext/builtin/persistence/migrations/atomic_memory_references.py index 54ed0b6cac..8434718193 100644 --- a/src/powercontext/builtin/persistence/migrations/atomic_memory_references.py +++ b/src/powercontext/builtin/persistence/migrations/atomic_memory_references.py @@ -85,6 +85,24 @@ _TERMINAL_DREAM = frozenset({"succeeded", "failed"}) +def _stored_work_payload(payload: bytes, /) -> tuple[type[BaseModel], Mapping[str, Any], Mapping[str, Any]] | None: + """Return one stored Source's typed Work payload with its stored root, or None when it holds none. + + The stored payload envelope arrived after structured Work Sources did, so a Work + payload can sit at the top level of a row written by an earlier release. The root + lets a rewrite keep the layout the row already uses. + """ + + stored = json.loads(payload) + if not isinstance(stored, dict): + return None + value = stored.get("value") + if not isinstance(value, dict): + value = stored + model = _WORK_MODELS.get(str((value.get("metadata") or {}).get("kind"))) + return None if model is None else (model, value, stored) + + class _LegacyValue(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True, strict=True) @@ -710,11 +728,10 @@ async def convert(connection: AsyncConnection, key: Mapping[str, Any]) -> None: ), dict(key), ) - envelope = json.loads(bytes(stored)) - value = envelope["value"] - model = _WORK_MODELS.get(str((value.get("metadata") or {}).get("kind"))) - if model is None: + work = _stored_work_payload(bytes(stored)) + if work is None: return + model, value, root = work original = json.loads(value["content"]) # A receipt's unavailable evidence records what was missing; it asserts no support to drop. dropped: list[tuple[str, dict[str, Any]]] | None = None if model is HandoffReceipt else [] @@ -749,13 +766,15 @@ async def convert(connection: AsyncConnection, key: Mapping[str, Any]) -> None: context.count("work_sources") if not context.write: return - value = {**value, "content": after.model_dump_json(by_alias=True, exclude_none=False, indent=2)} + updated = {**value, "content": after.model_dump_json(by_alias=True, exclude_none=False, indent=2)} + # A legacy row keeps its own layout; only its content changes. + rewritten = {**root, "content": updated["content"]} if value is root else {**root, "value": updated} await connection.execute( text( "UPDATE pc_sources SET payload = :payload WHERE scope_id = :scope_id AND source_type = :source_type " "AND source_id = :source_id" ), - {**key, "payload": rfc8785.dumps({**envelope, "value": value})}, + {**key, "payload": rfc8785.dumps(rewritten)}, ) if isinstance(after, TaskOutcome): changed = _outcome_digests(_legacy_item_digests(original), after) @@ -932,10 +951,10 @@ async def _dream_inputs_change( # noqa: C901 - One check over each kind of pinn ) if stored is None: continue - value = json.loads(bytes(stored)).get("value") or {} - model = _WORK_MODELS.get(str((value.get("metadata") or {}).get("kind"))) - if model is None: + work = _stored_work_payload(bytes(stored)) + if work is None: continue + model, value, _root = work content = json.loads(value["content"]) if model is HandoffReceipt: # These unavailable addresses survive unchanged, so the Source digest stays pinned. @@ -1295,12 +1314,11 @@ async def count(sql: str) -> int: ("payload",), "WHERE source_type = 'content' AND CAST(payload AS CHAR) LIKE '%memory%'", ): - value = json.loads(bytes(row["payload"]))["value"] - model = _WORK_MODELS.get(str((value.get("metadata") or {}).get("kind"))) - if model is None: + work = _stored_work_payload(bytes(row["payload"])) + if work is None: continue - content = value["content"] - current = json.loads(content) + model, value, _root = work + current = json.loads(value["content"]) if model is HandoffReceipt: # Only unavailable addresses may describe an archived collection. current["unavailable_evidence"] = [ diff --git a/tests/e2e/test_atomic_memory_migration_evidence.py b/tests/e2e/test_atomic_memory_migration_evidence.py index 1f2f5bdd07..e64eb10309 100644 --- a/tests/e2e/test_atomic_memory_migration_evidence.py +++ b/tests/e2e/test_atomic_memory_migration_evidence.py @@ -25,7 +25,8 @@ from typing import Any import pytest -from sqlalchemy import Engine, ForeignKeyConstraint, MetaData, event, insert, inspect, select, text, tuple_ +import rfc8785 +from sqlalchemy import Engine, ForeignKeyConstraint, MetaData, event, insert, inspect, select, text, tuple_, update from powercontext.artifacts import ArtifactRef from powercontext.builtin.artifacts.atomic_memory import AtomicMemoryContent @@ -997,6 +998,84 @@ async def scenario() -> None: asyncio.run(scenario()) +async def _unenvelope(connection, ref: SourceRef) -> None: + """Rewrite one stored Source row into the payload layout used before the envelope.""" + + where = (SOURCES_TABLE.c.scope_id == SCOPE, SOURCES_TABLE.c.source_id == ref.source_id) + payload = await connection.scalar(select(SOURCES_TABLE.c.payload).where(*where)) + assert isinstance(payload, bytes) + await connection.execute( + update(SOURCES_TABLE).where(*where).values(payload=rfc8785.dumps(json.loads(bytes(payload))["value"])) + ) + + +def test_unenveloped_source_mentioning_memory_migrates_unchanged(migration_backend: MigrationBackend) -> None: + config = migration_backend.config + mention = "Restore durable memory context before accepting this Handoff." + capture = SourceRef(source_type="content", source_id="unenveloped-capture") + + async def seeded(contexts, connection) -> None: + await contexts.repositories.sources.add( + connection, + SCOPE, + ContentSource(name=capture.source_id, materialization=SourceMaterialization.CAPTURED, content=mention), + ) + await _unenvelope(connection, capture) + + async def scenario() -> None: + await _seed_and_migrate(config, seeded) + async with ( + open_builtin_contexts(BuiltinConfig(database=config)) as contexts, + contexts.database.transaction() as connection, + ): + stored = (await contexts.repositories.sources.get(connection, SCOPE, capture)).value + assert isinstance(stored, ContentSource) and stored.content == mention + report = await verify_atomic_memory_migration(connection, index=contexts.atomic_memory.index) + assert report.ready, report.errors + + asyncio.run(scenario()) + + +def test_unenveloped_work_source_citations_convert_and_keep_their_layout( + migration_backend: MigrationBackend, +) -> None: + config = migration_backend.config + alpha = _atomic("alpha", 2) + outcome = SourceRef(source_type="content", source_id="legacy-layout-outcome") + legacy = {**LEGACY_OUTCOME, "produced_artifacts": [_collection(2)]} + + async def seeded(contexts, connection) -> None: + await _scope_row(connection) + await _work_source(contexts, connection, outcome, "task-outcome", "powercontext.task-outcome.v1", legacy) + await _unenvelope(connection, outcome) + + async def scenario() -> None: + await _seed_and_migrate(config, seeded) + async with ( + open_builtin_contexts(BuiltinConfig(database=config)) as contexts, + contexts.database.transaction() as connection, + ): + payload = await connection.scalar( + select(SOURCES_TABLE.c.payload).where( + SOURCES_TABLE.c.scope_id == SCOPE, SOURCES_TABLE.c.source_id == outcome.source_id + ) + ) + assert isinstance(payload, bytes) + stored = json.loads(bytes(payload)) + assert "value" not in stored, "a converted legacy row keeps the layout it was written with" + converted = TaskOutcome.model_validate_json(stored["content"]) + assert converted.observations[0].evidence == (HandoffArtifactCitation(artifact_ref=alpha),) + assert converted.produced_artifacts == () + incoming = (await _archive_metadata(connection, 2))["incoming_references"] + assert [(item["carrier"], item["value"]) for item in incoming] == [("work_source", _collection(2))] + source = (await contexts.repositories.sources.get(connection, SCOPE, outcome)).value + assert isinstance(source, ContentSource) + report = await verify_atomic_memory_migration(connection, index=contexts.atomic_memory.index) + assert report.ready, report.errors + + asyncio.run(scenario()) + + def test_collections_citing_each_other_are_removed_regardless_of_identifier_order( migration_backend: MigrationBackend, ) -> None: