Enforce supervisor schema class name matches its type literal

[jason810496]

Why

The six supervisor discriminated unions (ToTask, ToSupervisor, ToManager, ToDagProcessor, ToTriggerRunner, ToTriggerSupervisor) need every member's class __name__ to equal its type: Literal[...] value so CommsDecoder routes wire frames to the right class — but nothing enforced that invariant, and two members had silently drifted on main.

What

• Add task-sdk/tests/task_sdk/execution_time/test_supervisor_schemas_name_type_sync.py — a parametrized unit test (one case per union) that walks every member and asserts class __name__ equals its single-value type Literal. Catches drift, a missing type field, and multi-value Literals. Runs as part of the task-sdk test suite.
• Fix the two pre-existing mismatches the new test surfaces in task-sdk/src/airflow/sdk/execution_time/comms.py:
  • XComCountResponse: Literal["XComLengthResponse"] → Literal["XComCountResponse"]
  • GetXComCount: Literal["GetNumberXComs"] → Literal["GetXComCount"]
• Update the one assertion in task-sdk/tests/task_sdk/execution_time/test_supervisor.py that pinned the old wire string.

Verfication

---

Was generative AI tooling used to co-author this PR?

• Yes, with help of Claude Code Opus 4.7 following the guidelines

[jscheffl]

Not sure if considered but besides making it static check at coding time, have we considered making this a runtime property handled by Pydantic?

class BaseCommClass(BaseModel):
    """Base type class."""

    type: str = "__undefined__"

    def __init__(self):
        super().__init__()
        self.type = self.__class__.__name__

class MyType1(BaseCommClass):
    """Another type."""

    pass

class MyType2(BaseCommClass):
    """Another type."""

    pass


print(MyType1().type)
print(MyType2().type)

[jason810496]

Not sure if considered but besides making it static check at coding time, have we considered making this a runtime property handled by Pydantic?

I will follow-up on this one. BTW, not sure if there's any reason to explictly specify the typ literal field (e.g. binding the field at runtime might change the generated schema?) I will check it out.

[jason810496]

Tried it. Runtime-binding changes the schema and also breaks the tagged union itself, because Pydantic v2 needs the discriminator field to be a Literal.

from typing import Annotated, Literal, Union
from pydantic import BaseModel, Field, TypeAdapter

# --- Current pattern: Literal type field ---
class A1(BaseModel):
    type: Literal["A1"] = "A1"

class A2(BaseModel):
    type: Literal["A2"] = "A2"

print(A1.model_json_schema())
#. {'properties': {'type': {'const': 'A1', 'default': 'A1', 'title': 'Type', 'type': 'string'}}, 'title': 'A1', 'type': 'object'}
TypeAdapter(Annotated[Union[A1, A2], Field(discriminator="type")])  # OK

# --- Proposed pattern: runtime-bind in __init__ ---
class B(BaseModel):
    type: str = "__undefined__"
    def __init__(self, **kw):
        super().__init__(**kw)
        self.type = type(self).__name__

class B1(B): pass
class B2(B): pass

print(B1.model_json_schema())
# {'properties': {'type': {'default': '__undefined__', 'title': 'Type', 'type': 'string'}}, 'title': 'B1', 'type': 'object'}
print(B1().type)  # 'B1' (instance state only; not reflected in schema)

TypeAdapter(Annotated[Union[B1, B2], Field(discriminator="type")])
# PydanticUserError: Model 'B1' needs field 'type' to be of type `Literal`

So CommsDecoder's TypeAdapter-based routing relies on the Literal annotation, and any downstream schema consumer relies on the const it produces.