Repository: apache/airflow
File: task-sdk/src/airflow/sdk/api/datamodels/_generated.py (Python)

Code before:

class PrevSuccessfulDagRunResponse(BaseModel):
    """
    Schema for response with previous successful DagRun information for Task Template Context.
    """

    data_interval_start: Annotated[AwareDatetime | None, Field(title="Data Interval Start")] = None
    data_interval_end: Annotated[AwareDatetime | None, Field(title="Data Interval End")] = None
    start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] = None
    end_date: Annotated[AwareDatetime | None, Field(title="End Date")] = None


class TIDeferredStatePayload(BaseModel):
    """
    Schema for updating TaskInstance to a deferred state.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    state: Annotated[Literal["deferred"] | None, Field(title="State")] = "deferred"
    classpath: Annotated[str, Field(title="Classpath")]
    trigger_kwargs: Annotated[JsonValue | str | None, Field(title="Trigger Kwargs")] = None
    trigger_timeout: Annotated[timedelta | None, Field(title="Trigger Timeout")] = None
    next_method: Annotated[str, Field(title="Next Method")]
    next_kwargs: JsonValue | None = None
    rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None


class TIEnterRunningPayload(BaseModel):
    """
    Schema for updating TaskInstance to 'RUNNING' state with minimal required fields.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    state: Annotated[Literal["running"] | None, Field(title="State")] = "running"
    hostname: Annotated[str, Field(title="Hostname")]
    unixname: Annotated[str, Field(title="Unixname")]
    pid: Annotated[int, Field(title="Pid")]
    start_date: Annotated[AwareDatetime, Field(title="Start Date")]


class TIHeartbeatInfo(BaseModel):
    """
    Schema for TaskInstance heartbeat endpoint.
    """

    model_config = ConfigDict(
        extra="forbid",

Review comment (bug):
`JsonValue` is `dict | list | str | int | float | bool | None` in Pydantic. But `next_kwargs` is semantically always a dict - it represents keyword arguments that get unpacked as `**next_kwargs` in `execute_callable(context, **next_kwargs)`. Changing the type to `JsonValue` is incorrect because:

serde output is still a dict: `serde.serialize({"foo": datetime})` returns `{"foo": {"__classname__": ...}}` - the outer container is still a dict

Code after:

class PrevSuccessfulDagRunResponse(BaseModel):
    """
    Schema for response with previous successful DagRun information for Task Template Context.
    """

    data_interval_start: Annotated[AwareDatetime | None, Field(title="Data Interval Start")] = None
    data_interval_end: Annotated[AwareDatetime | None, Field(title="Data Interval End")] = None
    start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] = None
    end_date: Annotated[AwareDatetime | None, Field(title="End Date")] = None


class PreviousTIResponse(BaseModel):
    """
    Schema for response with previous TaskInstance information.
    """

    task_id: Annotated[str, Field(title="Task Id")]
    dag_id: Annotated[str, Field(title="Dag Id")]
    run_id: Annotated[str, Field(title="Run Id")]
    logical_date: Annotated[AwareDatetime | None, Field(title="Logical Date")] = None
    start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] = None
    end_date: Annotated[AwareDatetime | None, Field(title="End Date")] = None
    state: Annotated[str | None, Field(title="State")] = None
    try_number: Annotated[int, Field(title="Try Number")]
    map_index: Annotated[int | None, Field(title="Map Index")] = -1
    duration: Annotated[float | None, Field(title="Duration")] = None


class TIDeferredStatePayload(BaseModel):
    """
    Schema for updating TaskInstance to a deferred state.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    state: Annotated[Literal["deferred"] | None, Field(title="State")] = "deferred"
    classpath: Annotated[str, Field(title="Classpath")]
    trigger_kwargs: Annotated[dict[str, JsonValue] | str | None, Field(title="Trigger Kwargs")] = None
    trigger_timeout: Annotated[timedelta | None, Field(title="Trigger Timeout")] = None
    next_method: Annotated[str, Field(title="Next Method")]
    next_kwargs: Annotated[dict[str, JsonValue] | None, Field(title="Next Kwargs")] = None
    rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None


class TIEnterRunningPayload(BaseModel):
    """
    Schema for updating TaskInstance to 'RUNNING' state with minimal required fields.
    """