diff --git a/polylogue/schemas/runtime_registry.py b/polylogue/schemas/runtime_registry.py index 4f4ae6276e..5765740cff 100644 --- a/polylogue/schemas/runtime_registry.py +++ b/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( diff --git a/tests/unit/schemas/test_promotion_monotonicity.py b/tests/unit/schemas/test_promotion_monotonicity.py index 4d4ba52506..caac007f83 100644 --- a/tests/unit/schemas/test_promotion_monotonicity.py +++ b/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=)`` 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