fix(codegen): preserve regenerated response aliases (#854)

* fix(codegen): preserve regenerated response aliases

* test: pin semantic response aliases

* fix: satisfy mypy for codegen aliases

* fix: wrap response alias fallbacks

---------

Co-authored-by: sangilish <56685007+sangilish@users.noreply.github.com>

Author: bokelley

scripts/consolidate_exports.py

@@ -28,6 +28,10 @@ def extract_exports_from_module(module_path: Path) -> set[str]:
 
     exports = set()
 
+    def _add_public_type_name(name: str) -> None:
+        if name and not name.startswith("_") and name[0].isupper():
+            exports.add(name)
+
     # Only look at module-level nodes (not inside classes)
     for node in tree.body:
         # Class definitions
@@ -39,8 +43,11 @@ def extract_exports_from_module(module_path: Path) -> set[str]:
             for target in node.targets:
                 if isinstance(target, ast.Name) and not target.id.startswith("_"):
                     # Only export if it looks like a type name (starts with capital)
-                    if target.id and target.id[0].isupper():
-                        exports.add(target.id)
+                    _add_public_type_name(target.id)
+        elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
+            # ``Foo: TypeAlias = ...`` is common in post-generated response
+            # unions; keep those public aliases in the consolidated namespace.
+            _add_public_type_name(node.target.id)
 
     return exports
 
@@ -484,11 +491,11 @@ def _stem_matches_export(module_stem: str, export_name: str) -> bool:
     content = "\n".join(lines)
     # Product is generated as a RootModel union, but the SDK's public Product
     # export intentionally points at the concrete first arm for subclassing.
-    # Avoid importing the RootModel under the public name so mypy sees the
-    # later compatibility assignment as a plain class alias, not a class
-    # redefinition.
+    # Avoid importing the RootModel so mypy sees the later compatibility
+    # assignment as the first public Product binding, not a class redefinition.
     content = content.replace(
-        "    Product,\n    Product1,\n", "    Product as ProductUnion,\n    Product1,\n"
+        "    Product,\n    Product1,\n",
+        "    Product1,\n",
     )
     return content
 

scripts/post_generate_fixes.py

@@ -46,6 +46,28 @@ def _load_resolve_bundle_key():
 OUTPUT_DIR = REPO_ROOT / "src" / "adcp" / "types" / "generated_poc"
 SCHEMA_DIR = REPO_ROOT / "schemas" / "cache" / _BUNDLE_KEY
 
+_PROTOCOL_ENVELOPE_IMPORT = "from ..core.protocol_envelope import ProtocolEnvelope\n"
+_VERSION_ENVELOPE_IMPORT = "from ..core.version_envelope import AdcpVersionEnvelope\n"
+
+
+def _sync_protocol_envelope_import(source: str) -> str:
+    """Keep the ProtocolEnvelope import aligned with restored response arms."""
+    uses_protocol_envelope = "ProtocolEnvelope" in source.replace(_PROTOCOL_ENVELOPE_IMPORT, "")
+    if not uses_protocol_envelope:
+        return source.replace(_PROTOCOL_ENVELOPE_IMPORT, "")
+    if _PROTOCOL_ENVELOPE_IMPORT in source:
+        return source
+    if _VERSION_ENVELOPE_IMPORT in source:
+        return source.replace(
+            _VERSION_ENVELOPE_IMPORT,
+            _PROTOCOL_ENVELOPE_IMPORT + _VERSION_ENVELOPE_IMPORT,
+            1,
+        )
+    future_import = "from __future__ import annotations\n\n"
+    if future_import in source:
+        return source.replace(future_import, future_import + _PROTOCOL_ENVELOPE_IMPORT, 1)
+    return _PROTOCOL_ENVELOPE_IMPORT + source
+
 
 def add_model_validator_to_product():
     """Add model_validators to Product class.
@@ -1593,15 +1615,12 @@ def restore_response_variant_aliases() -> None:
 
     def _remove_original_response_class(source: str, base: str) -> str:
         """Remove the generator's envelope-only class before restoring a union alias."""
-        protocol_import = "from ..core.protocol_envelope import ProtocolEnvelope\n"
         source = re.sub(
             rf"\n\nclass {re.escape(base)}\(AdcpVersionEnvelope, ProtocolEnvelope\):\n    pass\n",
             "\n",
             source,
         )
-        if "ProtocolEnvelope" not in source.replace(protocol_import, ""):
-            source = source.replace(protocol_import, "")
-        return source
+        return _sync_protocol_envelope_import(source)
 
     def _normalize_existing_arms(target: Path, base: str) -> None:
         """Keep compatibility arms payload-shaped and expose final names as aliases."""
@@ -1621,6 +1640,7 @@ def _normalize_existing_arms(target: Path, base: str) -> None:
             f"\n{base}: TypeAlias = ",
             new_source,
         )
+        new_source = _sync_protocol_envelope_import(new_source)
         if new_source != original:
             target.write_text(new_source)
 

src/adcp/types/_ergonomic.py

@@ -51,13 +51,15 @@
 from adcp.types.generated_poc.core.error import Error
 from adcp.types.generated_poc.core.ext import ExtensionObject
 from adcp.types.generated_poc.core.format import Format
+from adcp.types.generated_poc.core.package import Package
 from adcp.types.generated_poc.core.product import Product
 from adcp.types.generated_poc.enums.advertiser_industry import AdvertiserIndustry
 from adcp.types.generated_poc.enums.asset_content_type import AssetContentType
 from adcp.types.generated_poc.enums.creative_sort_field import CreativeSortField
 from adcp.types.generated_poc.enums.delivery_type import DeliveryType
 from adcp.types.generated_poc.enums.disclosure_persistence import DisclosurePersistence
 from adcp.types.generated_poc.enums.disclosure_position import DisclosurePosition
+from adcp.types.generated_poc.enums.media_buy_status import MediaBuyStatus
 from adcp.types.generated_poc.enums.pacing import Pacing
 from adcp.types.generated_poc.enums.sort_direction import SortDirection
 from adcp.types.generated_poc.enums.task_status import TaskStatus
@@ -79,6 +81,9 @@
 )
 from adcp.types.generated_poc.media_buy.package_request import PackageRequest
 from adcp.types.generated_poc.media_buy.package_update import PackageUpdate
+from adcp.types.generated_poc.media_buy.create_media_buy_response import (
+    CreateMediaBuyResponse1,
+)
 from adcp.types.generated_poc.media_buy.get_media_buy_delivery_response import (
     GetMediaBuyDeliveryResponse,
     MediaBuyDelivery,
@@ -476,6 +481,24 @@ def _apply_coercion() -> None:
     )
     ListCreativeFormatsResponse.model_rebuild(force=True)
 
+    # Apply coercion to CreateMediaBuyResponse1
+    # - packages: list[Package] (accepts subclass instances)
+    # - media_buy_status: MediaBuyStatus | str | None
+    _patch_field_annotation(
+        CreateMediaBuyResponse1,
+        "packages",
+        Annotated[
+            list[Package],
+            BeforeValidator(coerce_subclass_list(Package)),
+        ],
+    )
+    _patch_field_annotation(
+        CreateMediaBuyResponse1,
+        "media_buy_status",
+        Annotated[MediaBuyStatus | None, BeforeValidator(coerce_to_enum(MediaBuyStatus))],
+    )
+    CreateMediaBuyResponse1.model_rebuild(force=True)
+
     # Apply coercion to GetMediaBuyDeliveryResponse
     # - context: ContextObject | dict | None
     # - status: TaskStatus | str | None

src/adcp/types/_generated.py

@@ -10,7 +10,7 @@
 DO NOT EDIT MANUALLY.
 
 Generated from: https://github.com/adcontextprotocol/adcp/tree/main/schemas
-Generation date: 2026-05-22 14:10:57 UTC
+Generation date: 2026-05-24 21:33:52 UTC
 """
 
 # ruff: noqa: E501, I001
@@ -260,7 +260,13 @@
     Tags,
 )
 from adcp.types.generated_poc.brand.acquire_rights_request import AcquireRightsRequest, Campaign
-from adcp.types.generated_poc.brand.acquire_rights_response import AcquireRightsResponse
+from adcp.types.generated_poc.brand.acquire_rights_response import (
+    AcquireRightsResponse,
+    AcquireRightsResponse1,
+    AcquireRightsResponse2,
+    AcquireRightsResponse3,
+    AcquireRightsResponse4,
+)
 from adcp.types.generated_poc.brand.creative_approval_request import CreativeApprovalRequest
 from adcp.types.generated_poc.brand.creative_approval_response import CreativeApprovalResponse
 from adcp.types.generated_poc.brand.get_brand_identity_request import (
@@ -273,6 +279,8 @@
     FontRole,
     FontRole2,
     GetBrandIdentityResponse,
+    GetBrandIdentityResponse1,
+    GetBrandIdentityResponse2,
     OpentypeFeature,
     Style,
     WeightRangeItem,
@@ -305,8 +313,15 @@
     UpdateRightsResponse2,
 )
 from adcp.types.generated_poc.brand.verification_status import VerificationStatus
-from adcp.types.generated_poc.brand.verify_brand_claim_request import VerifyBrandClaimRequest
-from adcp.types.generated_poc.brand.verify_brand_claim_response import VerifyBrandClaimResponse
+from adcp.types.generated_poc.brand.verify_brand_claim_request import (
+    ClaimType,
+    VerifyBrandClaimRequest,
+)
+from adcp.types.generated_poc.brand.verify_brand_claim_response import (
+    VerifyBrandClaimErrorResponse,
+    VerifyBrandClaimResponse,
+    VerifyBrandClaimSuccessResponse,
+)
 from adcp.types.generated_poc.brand.verify_brand_claims_request import (
     Claim,
     Claim1,
@@ -322,7 +337,6 @@
     VerifyBrandClaimsRequestBulk,
 )
 from adcp.types.generated_poc.brand.verify_brand_claims_response import (
-    ClaimType,
     ResultEntry,
     ResultEntry1,
     ResultEntry2,
@@ -435,13 +449,18 @@
 )
 from adcp.types.generated_poc.content_standards.get_content_standards_response import (
     GetContentStandardsResponse,
+    GetContentStandardsResponse1,
+    GetContentStandardsResponse2,
 )
 from adcp.types.generated_poc.content_standards.get_media_buy_artifacts_request import (
     GetMediaBuyArtifactsRequest,
     TimeRange,
 )
 from adcp.types.generated_poc.content_standards.get_media_buy_artifacts_response import (
+    ArtifactRecord,
     GetMediaBuyArtifactsResponse,
+    GetMediaBuyArtifactsResponse1,
+    GetMediaBuyArtifactsResponse2,
 )
 from adcp.types.generated_poc.content_standards.list_content_standards_request import (
     ListContentStandardsRequest,
@@ -753,7 +772,6 @@
     MaterialSubmission,
     MetricOptimization,
     MetricOptimization1,
-    Product as ProductUnion,
     Product1,
     Product2,
     ProductCardDetailed,
@@ -1043,6 +1061,8 @@
 )
 from adcp.types.generated_poc.creative.get_creative_features_response import (
     GetCreativeFeaturesResponse,
+    GetCreativeFeaturesResponse1,
+    GetCreativeFeaturesResponse2,
 )
 from adcp.types.generated_poc.creative.list_creative_formats_request import (
     AssetType,
@@ -1719,24 +1739,12 @@
 # Backward compatibility aliases for renamed types
 MediaBuyPackage = _PackageFromGetMediaBuysResponse
 DeliveryStatus = _DeliveryStatusFromGetMediaBuysResponse
-AcquireRightsResponse1 = AcquireRightsResponse
-AcquireRightsResponse2 = AcquireRightsResponse
-AcquireRightsResponse3 = AcquireRightsResponse
-AcquireRightsResponse4 = AcquireRightsResponse
 ComplyTestControllerResponse1 = ComplyTestControllerResponse
 ComplyTestControllerResponse2 = ComplyTestControllerResponse
 ComplyTestControllerResponse3 = ComplyTestControllerResponse
 ComplyTestControllerResponse4 = ComplyTestControllerResponse
 CreateContentStandardsResponse1 = CreateContentStandardsResponse
 CreateContentStandardsResponse2 = CreateContentStandardsResponse
-GetBrandIdentityResponse1 = GetBrandIdentityResponse
-GetBrandIdentityResponse2 = GetBrandIdentityResponse
-GetContentStandardsResponse1 = GetContentStandardsResponse
-GetContentStandardsResponse2 = GetContentStandardsResponse
-GetCreativeFeaturesResponse1 = GetCreativeFeaturesResponse
-GetCreativeFeaturesResponse2 = GetCreativeFeaturesResponse
-GetMediaBuyArtifactsResponse1 = GetMediaBuyArtifactsResponse
-GetMediaBuyArtifactsResponse2 = GetMediaBuyArtifactsResponse
 ListContentStandardsResponse1 = ListContentStandardsResponse
 ListContentStandardsResponse2 = ListContentStandardsResponse
 UpdateContentStandardsResponse1 = UpdateContentStandardsResponse
@@ -1831,6 +1839,7 @@
     "Area",
     "Arm",
     "Artifact",
+    "ArtifactRecord",
     "ArtifactRef",
     "ArtifactWebhook",
     "ArtifactWebhookPayload",
@@ -3056,8 +3065,10 @@
     "VerificationStatus",
     "VerifyAgent",
     "VerifyAgent759",
+    "VerifyBrandClaimErrorResponse",
     "VerifyBrandClaimRequest",
     "VerifyBrandClaimResponse",
+    "VerifyBrandClaimSuccessResponse",
     "VerifyBrandClaimsRequest",
     "VerifyBrandClaimsRequestBulk",
     "VerifyBrandClaimsResponseBulk",