fix(schemas): preserve nested annotations and carry forward element kinds

polylogue/schemas/runtime_registry.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import copy
+import dataclasses
 import gzip
 import json
 import threading
@@ -568,7 +569,23 @@ def _load_local_element_schema(
         element = package.element(element_kind)
         if element is None or element.schema_file is None:
             return None
-        path = self._provider_dir(provider_token) / "versions" / package.version / "elements" / element.schema_file
+        return self._read_local_element_schema_file(provider_token, package.version, element.schema_file)
+
+    def _read_local_element_schema_file(
+        self,
+        provider_token: str,
+        package_version: str,
+        schema_file: str,
+    ) -> PublicSchemaDocument | None:
+        """Read one element schema file given an already-resolved package/element.
+
+        Split out of ``_load_local_element_schema`` so a caller that already
+        holds the catalog and the resolved package/element (e.g.
+        ``_existing_provider_element_schemas``'s loop) reads the schema file
+        directly instead of re-loading and re-parsing the whole catalog JSON
+        per element/version.
+        """
+        path = self._provider_dir(provider_token) / "versions" / package_version / "elements" / schema_file
         if not path.exists():
             return None
         return _read_gzip_json_dict(path)
@@ -593,11 +610,7 @@ def _existing_provider_element_schemas(self, provider_token: str) -> dict[str, P
             for element in package.elements:
                 if element.schema_file is None:
                     continue
-                existing = self._load_local_element_schema(
-                    provider_token,
-                    version=package.version,
-                    element_kind=element.element_kind,
-                )
+                existing = self._read_local_element_schema_file(provider_token, package.version, element.schema_file)
                 if existing is None:
                     continue
                 prior = merged_by_kind.get(element.element_kind)
@@ -608,6 +621,67 @@ def _existing_provider_element_schemas(self, provider_token: str) -> dict[str, P
                 )
         return merged_by_kind
 
+    @staticmethod
+    def _annotate_merged_schema_node(
+        merged: PublicSchemaDocument,
+        *,
+        existing: PublicSchemaDocument | None,
+        candidate: PublicSchemaDocument | None,
+    ) -> PublicSchemaDocument:
+        """Recursively reattach ``x-polylogue-*`` annotations onto a structurally merged node.
+
+        ``merge_observed_structure_schemas`` merges only structural keywords
+        (``type``/``properties``/``items``/``additionalProperties`` --
+        its own docstring says "without retaining property history"), so
+        every node of the merged tree comes back with no annotation overlay
+        at all, not just the root. This walks ``merged``/``existing``/
+        ``candidate`` in parallel by matching property name / items /
+        additionalProperties position, preferring the candidate's freshly
+        computed annotations at each node and falling back to the existing
+        package's annotations for that same node when the candidate didn't
+        recompute one (e.g. this run observed the field's type again but
+        didn't rerun semantic-role/format/frequency/distribution inference).
+        """
+        result: PublicSchemaDocument = dict(merged)
+        existing = existing or {}
+        candidate = candidate or {}
+        for key, value in existing.items():
+            if key.startswith("x-polylogue-") and key not in candidate:
+                result[key] = value
+        for key, value in candidate.items():
+            if key.startswith("x-polylogue-"):
+                result[key] = value
+
+        merged_properties = json_document(merged.get("properties"))
+        if merged_properties:
+            existing_properties = json_document(existing.get("properties"))
+            candidate_properties = json_document(candidate.get("properties"))
+            result["properties"] = {
+                name: SchemaRegistry._annotate_merged_schema_node(
+                    json_document(child),
+                    existing=json_document(existing_properties.get(name)) or None,
+                    candidate=json_document(candidate_properties.get(name)) or None,
+                )
+                for name, child in merged_properties.items()
+            }
+
+        merged_items = json_document(merged.get("items"))
+        if merged_items:
+            result["items"] = SchemaRegistry._annotate_merged_schema_node(
+                merged_items,
+                existing=json_document(existing.get("items")) or None,
+                candidate=json_document(candidate.get("items")) or None,
+            )
+
+        merged_additional = json_document(merged.get("additionalProperties"))
+        if merged_additional:
+            result["additionalProperties"] = SchemaRegistry._annotate_merged_schema_node(
+                merged_additional,
+                existing=json_document(existing.get("additionalProperties")) or None,
+                candidate=json_document(candidate.get("additionalProperties")) or None,
+            )
+        return result
+
     @staticmethod
     def _merge_element_schema_with_existing(
         existing: PublicSchemaDocument | None,
@@ -633,16 +707,14 @@ def _merge_element_schema_with_existing(
         from polylogue.schemas.generation.dynamic_keys import merge_observed_structure_schemas
 
         merged = json_document(merge_observed_structure_schemas([json_document(existing), candidate]))
-        # Structural merge owns type/properties/items only; provenance and
-        # the x-polylogue-* annotation overlay come from the candidate (this
-        # run's fresh observation), falling back to the existing package so
-        # the merge never drops an annotation the candidate simply didn't
-        # recompute.
-        for key, value in existing.items():
-            if key.startswith("x-polylogue-") and key not in candidate:
-                merged[key] = value
+        # Structural merge owns type/properties/items/additionalProperties
+        # only; the x-polylogue-* annotation overlay is reattached node by
+        # node (not just at the document root) by
+        # _annotate_merged_schema_node, preferring the candidate's fresh
+        # annotations with the existing package's as fallback.
+        merged = SchemaRegistry._annotate_merged_schema_node(merged, existing=existing, candidate=candidate)
         for key, value in candidate.items():
-            if key.startswith("x-polylogue-") or key in ("$schema", "title"):
+            if key in ("$schema", "title"):
                 merged[key] = value
         return merged
 
@@ -656,6 +728,14 @@ def replace_provider_packages(
     ) -> None:
         provider_token = _provider_token(provider)
         existing_element_schemas = self._existing_provider_element_schemas(provider_token)
+        existing_catalog = self._load_local_catalog(provider_token)
+        existing_elements_by_version: dict[str, dict[str, SchemaElementManifest]] = {}
+        if existing_catalog is not None:
+            for existing_package in existing_catalog.packages:
+                existing_elements_by_version[existing_package.version] = {
+                    element.element_kind: element for element in existing_package.elements
+                }
+
         prepared_packages: list[tuple[SchemaVersionPackage, ElementSchemaMap, Mapping[str, object] | None]] = []
         for package in catalog.packages:
             element_schemas = package_schemas.get(package.version)
@@ -667,6 +747,32 @@ def replace_provider_packages(
                 )
                 for element_kind, schema in element_schemas.items()
             }
+
+            # A thinner regeneration window can observe zero samples for an
+            # element kind this same version previously committed, so that
+            # kind is absent from `element_schemas` entirely and the merge
+            # loop above never runs for it. Carry it forward unmerged
+            # (pass-through) instead of letting it vanish from the
+            # destructive versions/-tree rewrite below -- the same
+            # destructive-loss bug class ov5r fixed, narrower: whole missing
+            # element kinds rather than narrowed types within an observed
+            # kind. Scoped to kinds this SAME version previously carried
+            # (matching get_element_schema/package.element()'s per-version
+            # lookup) -- a kind that only ever lived on a version this
+            # regeneration dropped entirely is a separate, out-of-scope loss
+            # class (a whole retired version, not a kind within a version).
+            carried_elements = list(package.elements)
+            for element_kind, prior_manifest in existing_elements_by_version.get(package.version, {}).items():
+                if element_kind in element_schemas:
+                    continue
+                carried_schema = existing_element_schemas.get(element_kind)
+                if carried_schema is None:
+                    continue
+                merged_element_schemas[element_kind] = carried_schema
+                carried_elements.append(prior_manifest)
+            if len(carried_elements) != len(package.elements):
+                package = dataclasses.replace(package, elements=carried_elements)
+
             workload_profile = (
                 package_workload_profiles.get(package.version) if package_workload_profiles is not None else None
             )
@@ -695,6 +801,12 @@ def replace_provider_packages(
                 element_schemas=element_schemas,
                 workload_profile=workload_profile,
             )
+        # Persist the (possibly element-carry-forward-augmented) packages,
+        # not the caller's original `catalog.packages` -- otherwise a carried
+        # forward element's schema file is written to disk but the saved
+        # manifest never lists it, so get_element_schema/package.element()
+        # still can't find it.
+        catalog = dataclasses.replace(catalog, packages=[package for package, _, _ in prepared_packages])
         self.save_package_catalog(catalog)
 
     def _single_element_package(

tests/unit/schemas/test_promotion_monotonicity.py

@@ -194,6 +194,43 @@ def _single_package_catalog(provider: str, version: str, element_kind: str) -> S
     )
 
 
+def _two_element_package_catalog(
+    provider: str, version: str, anchor_kind: str, other_kind: str
+) -> SchemaPackageCatalog:
+    return SchemaPackageCatalog(
+        provider=provider,
+        packages=[
+            SchemaVersionPackage(
+                provider=provider,
+                version=version,
+                anchor_kind=anchor_kind,
+                default_element_kind=anchor_kind,
+                first_seen="2026-08-01T00:00:00Z",
+                last_seen="2026-08-01T00:00:00Z",
+                bundle_scope_count=0,
+                sample_count=1,
+                elements=[
+                    SchemaElementManifest(
+                        element_kind=anchor_kind,
+                        schema_file=f"{anchor_kind}.schema.json.gz",
+                        sample_count=1,
+                        artifact_count=1,
+                    ),
+                    SchemaElementManifest(
+                        element_kind=other_kind,
+                        schema_file=f"{other_kind}.schema.json.gz",
+                        sample_count=1,
+                        artifact_count=1,
+                    ),
+                ],
+            )
+        ],
+        default_version=version,
+        latest_version=version,
+        recommended_version=version,
+    )
+
+
 class TestReplaceProviderPackagesMonotonicity:
     """Guards ``SchemaRegistry.replace_provider_packages`` -- the code path every
     full-corpus ``devtools schema-generate`` run writes through
@@ -299,3 +336,151 @@ def test_regen_type_unions_only_grow_never_narrow(self, tmp_registry: SchemaRegi
         declared_type = timestamp["type"]
         observed = set(declared_type) if isinstance(declared_type, list) else {declared_type}
         assert observed == {"string", "number"}
+
+
+class TestMergeElementSchemaAnnotationPreservation:
+    """Guards nested ``x-polylogue-*`` annotation survival across
+    ``SchemaRegistry.replace_provider_packages`` regenerations (polylogue-46kg
+    P1, found by automated review on PR #3502 minutes after merge).
+
+    ``merge_observed_structure_schemas`` (``schemas/generation/dynamic_keys.py``)
+    merges only structural keywords (``type``/``properties``/``items``/
+    ``additionalProperties``) -- its own docstring says "without retaining
+    property history". The prior ``_merge_element_schema_with_existing``
+    restored ``x-polylogue-*`` annotations only at the document root after
+    merging, so every regeneration stripped freshly-computed *nested*
+    annotations (semantic role, format, frequency, observed distribution)
+    from property-level nodes.
+
+    Anti-vacuity: reverting
+    ``SchemaRegistry._merge_element_schema_with_existing`` to restore
+    ``x-polylogue-*`` keys only at the document root (dropping the recursive
+    ``_annotate_merged_schema_node`` walk into ``properties``) makes both
+    tests below fail -- the nested ``user_id``/``value`` property's
+    annotation is stripped by the structural merge and never restored.
+    """
+
+    def test_nested_annotation_survives_when_fresh_pass_recomputes_no_annotation(
+        self, tmp_registry: SchemaRegistry
+    ) -> None:
+        annotated_schema: JSONDocument = {
+            "type": "object",
+            "properties": {
+                "user_id": {
+                    "type": "string",
+                    "x-polylogue-semantic-role": "identifier",
+                }
+            },
+        }
+        tmp_registry.replace_provider_packages(
+            "regen-annotation",
+            _single_package_catalog("regen-annotation", "v1", "session_record_stream"),
+            {"v1": {"session_record_stream": annotated_schema}},
+        )
+
+        # A subsequent regeneration re-observes the same property's type but
+        # its annotation pass didn't recompute a semantic role this time --
+        # e.g. a narrower sample window that never re-derived it.
+        unannotated_schema: JSONDocument = {
+            "type": "object",
+            "properties": {"user_id": {"type": "string"}},
+        }
+        tmp_registry.replace_provider_packages(
+            "regen-annotation",
+            _single_package_catalog("regen-annotation", "v1", "session_record_stream"),
+            {"v1": {"session_record_stream": unannotated_schema}},
+        )
+
+        merged = tmp_registry.get_schema("regen-annotation", version="v1")
+        assert merged is not None
+        user_id = cast("dict[str, Any]", cast("dict[str, Any]", merged["properties"])["user_id"])
+        assert user_id.get("x-polylogue-semantic-role") == "identifier"
+
+    def test_nested_annotation_fresh_value_wins_over_stale_existing(self, tmp_registry: SchemaRegistry) -> None:
+        stale_schema: JSONDocument = {
+            "type": "object",
+            "properties": {
+                "value": {
+                    "type": "string",
+                    "x-polylogue-semantic-role": "identifier",
+                }
+            },
+        }
+        tmp_registry.replace_provider_packages(
+            "regen-annotation-update",
+            _single_package_catalog("regen-annotation-update", "v1", "session_record_stream"),
+            {"v1": {"session_record_stream": stale_schema}},
+        )
+
+        # This regeneration DOES carry a fresh, different annotation for the
+        # same nested property -- the fresh value must win, not the stale one.
+        updated_schema: JSONDocument = {
+            "type": "object",
+            "properties": {
+                "value": {
+                    "type": "string",
+                    "x-polylogue-semantic-role": "timestamp",
+                }
+            },
+        }
+        tmp_registry.replace_provider_packages(
+            "regen-annotation-update",
+            _single_package_catalog("regen-annotation-update", "v1", "session_record_stream"),
+            {"v1": {"session_record_stream": updated_schema}},
+        )
+
+        merged = tmp_registry.get_schema("regen-annotation-update", version="v1")
+        assert merged is not None
+        value = cast("dict[str, Any]", cast("dict[str, Any]", merged["properties"])["value"])
+        assert value.get("x-polylogue-semantic-role") == "timestamp"
+
+
+class TestReplaceProviderPackagesCarriesForwardUnobservedElementKinds:
+    """Guards ``SchemaRegistry.replace_provider_packages`` against dropping a
+    whole element kind when a thinner regeneration observes zero samples for
+    it (polylogue-46kg P2, found by automated review on PR #3502 minutes
+    after merge).
+
+    This is the same destructive-loss bug class ``ov5r``
+    (``TestReplaceProviderPackagesMonotonicity`` above) fixed, just narrower:
+    a whole element kind absent from ``element_schemas.items()`` for a
+    version never enters the merge loop at all, so the destructive
+    ``versions/`` delete-and-rewrite drops it entirely -- afterward
+    ``get_element_schema(..., element_kind=<old kind>)`` silently returns
+    ``None``.
+
+    Anti-vacuity: removing the ``existing_elements_by_version`` carry-forward
+    loop in ``replace_provider_packages`` (or the fix that persists the
+    carry-forward-augmented packages via ``dataclasses.replace(catalog, ...)``
+    instead of the caller's original ``catalog.packages``) makes this test
+    fail -- ``"tool_result"`` vanishes from the second run's manifest/elements
+    directory.
+    """
+
+    def test_kind_absent_from_fresh_observation_window_is_carried_forward(self, tmp_registry: SchemaRegistry) -> None:
+        message_schema: JSONDocument = {"type": "object", "properties": {"text": {"type": "string"}}}
+        tool_result_schema: JSONDocument = {"type": "object", "properties": {"exit_code": {"type": "integer"}}}
+        tmp_registry.replace_provider_packages(
+            "regen-kind-drop",
+            _two_element_package_catalog("regen-kind-drop", "v1", "message", "tool_result"),
+            {"v1": {"message": message_schema, "tool_result": tool_result_schema}},
+        )
+        assert tmp_registry.get_element_schema("regen-kind-drop", version="v1", element_kind="tool_result") is not None
+
+        # A thinner regeneration window observes zero tool_result samples
+        # this pass (e.g. a corpus subset with no tool calls) -- the fresh
+        # catalog and package_schemas mapping for this version now mention
+        # only "message".
+        tmp_registry.replace_provider_packages(
+            "regen-kind-drop",
+            _single_package_catalog("regen-kind-drop", "v1", "message"),
+            {"v1": {"message": message_schema}},
+        )
+
+        carried = tmp_registry.get_element_schema("regen-kind-drop", version="v1", element_kind="tool_result")
+        assert carried is not None
+        properties = cast("dict[str, Any]", carried["properties"])
+        assert properties["exit_code"]["type"] == "integer"
+
+        # "message" itself is untouched and still resolvable.
+        assert tmp_registry.get_element_schema("regen-kind-drop", version="v1", element_kind="message") is not None