Repository: apache/airflow
File: airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py (Python)

Code before:
class TIDeferredStatePayload(StrictBaseModel):
    """Schema for updating TaskInstance to a deferred state."""

    state: Annotated[
        Literal[IntermediateTIState.DEFERRED],
        # Specify a default in the schema, but not in code, so Pydantic marks it as required.
        WithJsonSchema(
            {
                "type": "string",
                "enum": [IntermediateTIState.DEFERRED],
                "default": IntermediateTIState.DEFERRED,
            }
        ),
    ]
    classpath: str
    trigger_kwargs: Annotated[JsonValue | str, Field(default_factory=dict)]
    """
    Kwargs to pass to the trigger constructor, either a plain dict or an encrypted string.

    Both forms will be passed along to the trigger, the server will not handle either.
    """

    trigger_timeout: timedelta | None = None
    next_method: str
    """The name of the method on the operator to call in the worker after the trigger has fired."""
    next_kwargs: Annotated[JsonValue, Field(default_factory=dict)]
    """
    Kwargs to pass to the above method, either a plain dict or an encrypted string.

    Both forms will be passed along to the TaskSDK upon resume, the server will not handle either.
    """
    rendered_map_index: str | None = None


class TIRescheduleStatePayload(StrictBaseModel):
    """Schema for updating TaskInstance to a up_for_reschedule state."""

    state: Annotated[
        Literal[IntermediateTIState.UP_FOR_RESCHEDULE],
        # Specify a default in the schema, but not in code, so Pydantic marks it as required.
        WithJsonSchema(
            {
                "type": "string",
                "enum": [IntermediateTIState.UP_FOR_RESCHEDULE],
                "default": IntermediateTIState.UP_FOR_RESCHEDULE,
            }
        ),
    ]
    reschedule_date: UtcDateTime
    end_date: UtcDateTime


Review comment (suggestion):
Gave it some more thought and I think that we should make it as `dict[str, JsonValue]` instead.

Worker -> API Server (Deferral):
```python
{  "state": "deferred",  "trigger_kwargs": {"moment": {"__classname__": "datetime.datetime", "__version__": 1, "datetime": "2024-01-15T10:30:00+00:00"}},  "next_kwargs": {"retry": {"__classname__": "bool", "__version__": 1, "bool": true}}}
```
Which will always be a `dict[str, JsonValue]`

API Server -> Worker (Resumption):
When worker resumes task, API server sends TaskInstance context with:
```python
{  "next_kwargs": {"retry": {"__classname__": "bool", "__version__": 1, "bool": true}}}
```
Which will always be a `dict[str, JsonValue]`

Because in both directions, we send `dict[str, JsonValue]`, we should have `next_kwargs` be `dict[str, JsonValue]`

Code after:
class TIDeferredStatePayload(StrictBaseModel):
    """Schema for updating TaskInstance to a deferred state."""

    state: Annotated[
        Literal[IntermediateTIState.DEFERRED],
        # Specify a default in the schema, but not in code, so Pydantic marks it as required.
        WithJsonSchema(
            {
                "type": "string",
                "enum": [IntermediateTIState.DEFERRED],
                "default": IntermediateTIState.DEFERRED,
            }
        ),
    ]
    classpath: str
    trigger_kwargs: Annotated[dict[str, JsonValue] | str, Field(default_factory=dict)]
    """
    Kwargs to pass to the trigger constructor, either a plain dict or an encrypted string.

    Both forms will be passed along to the trigger, the server will not handle either.
    """

    trigger_timeout: timedelta | None = None
    next_method: str
    """The name of the method on the operator to call in the worker after the trigger has fired."""
    next_kwargs: Annotated[dict[str, JsonValue], Field(default_factory=dict)]
    """
    Kwargs to pass to the above method, either a plain dict or an encrypted string.

    Both forms will be passed along to the TaskSDK upon resume, the server will not handle either.
    """
    rendered_map_index: str | None = None


class TIRescheduleStatePayload(StrictBaseModel):
    """Schema for updating TaskInstance to a up_for_reschedule state."""

    state: Annotated[
        Literal[IntermediateTIState.UP_FOR_RESCHEDULE],
        # Specify a default in the schema, but not in code, so Pydantic marks it as required.
        WithJsonSchema(
            {
                "type": "string",
                "enum": [IntermediateTIState.UP_FOR_RESCHEDULE],
                "default": IntermediateTIState.UP_FOR_RESCHEDULE,
            }
        ),
    ]
    reschedule_date: UtcDateTime
    end_date: UtcDateTime