Add WebSocketSensor and WebSocketTrigger to the standard provider

[ColtenOuO]

Summary

Deferrable operators that hand a long-lived request off to a remote server currently have no built-in way to resume once that server replies over a WebSocket connection. This closes #21139, which asked for exactly this: a trigger that opens a WebSocket connection, hands off the wait to the triggerer, and fires once the remote server replies.

Why we need this

Without this, the only options today are a plain HTTP-based deferrable sensor/trigger that polls, or a non-deferrable operator that blocks synchronously, and both have real costs for a long-running, unknown-duration remote job. Polling means repeated API calls against the remote server for the entire duration of the job — on a poke_interval of a few seconds that adds up fast over a job that can run for hours, and it puts avoidable load on a server that has no way to tell Airflow "I'm still working" other than answering yet another request. Polling also caps how quickly Airflow finds out the job is done: the result only surfaces on the next poll, so the effective latency is bounded by poke_interval, not by when the server actually finished. A synchronous, non-deferrable wait avoids the polling cost but instead ties up a full worker slot for the whole duration, which does not scale when many such jobs run concurrently. A WebSocket trigger avoids all three: the remote server pushes the reply the moment it is ready (no poll delay, no repeated calls), and the wait happens in the triggerer rather than on a worker.

This Provider Design (Work flow)

sequenceDiagram
    autonumber
    participant DAG as DAG task<br/>(WebSocketSensor, deferrable=True)
    participant Worker
    participant Triggerer
    participant Server as Remote WebSocket server

    DAG->>Worker: execute()
    Worker->>Worker: self.defer(trigger=WebSocketTrigger(...))
    Note over Worker: TaskDeferred raised —<br/>state becomes DEFERRED,<br/>worker slot released immediately
    Worker->>Triggerer: hand off serialized WebSocketTrigger
    Triggerer->>Server: connect() + send(message_to_send)
    Note over Triggerer,Server: await recv() — no polling loop,<br/>no worker occupied while waiting
    Server-->>Triggerer: WebSocket message (e.g. job done)
    Triggerer-->>Worker: TriggerEvent(payload)
    Note over Worker: A worker is scheduled again<br/>only once the event fires
    Worker->>DAG: execute_complete(event)

Loading

The same WebSocketTrigger also backs the non-deferrable path (deferrable=False, the default): WebSocketSensor.poke() opens one connection directly on the worker and blocks on recv() up to the sensor's own timeout — no repeated reconnects, since a WebSocket message can only be read once and re-sending message_to_send could restart the remote job. mode="reschedule" is rejected for the same reason: a rescheduled attempt would need a brand-new connection anyway.

Because Airflow does not guarantee a trigger's run() executes only once (a triggerer restart or redistribution re-runs it from scratch), a reconnect there re-sends message_to_send too. If that message has a side effect on the remote server, such as starting a job, the server must treat a resend as safe — for example by deduplicating on a request id embedded in the message. This is documented on the trigger and sensor, and covered by a test that reconstructs a trigger from its own serialize() output and runs it twice.

Usage

Wait for a message, non-deferrable (holds a worker for the duration):

WebSocketSensor(
    task_id="wait_for_websocket_message",
    url="wss://example.com/socket",
    timeout=3,
)

Same thing, but deferrable — releases the worker while waiting:

WebSocketSensor(
    task_id="wait_for_websocket_message_async",
    url="wss://example.com/socket",
    deferrable=True,
    timeout=3,
)

Send a request over the connection right after it opens, then wait for the async reply — the common case where the remote server needs to be told what to do before it has anything to report back. url, header, and message_to_send are all templated, so a run id or auth token can come from the Dag context instead of being hard-coded:

WebSocketSensor(
    task_id="request_and_wait_for_websocket_reply",
    url="wss://example.com/socket",
    message_to_send='{"action": "start_job"}',
    header={"Authorization": "Bearer my-token"},
    deferrable=True,
    timeout=3,
)

Changes

• Added WebSocketTrigger in airflow/providers/standard/triggers/websocket.py — a BaseTrigger (not BaseEventTrigger, which is the event-driven-scheduling marker and not meant for resuming a deferred task) that opens a ws:///wss:// connection with websockets.asyncio.client, optionally sends an initial message, and fires a TriggerEvent with the payload once a message is received.
• Added WebSocketSensor in airflow/providers/standard/sensors/websocket.py, following the deferrable / poke() / execute() / execute_complete() pattern used by FileSensor. poke() shares one deadline between the connection handshake (open_timeout) and the message wait (recv(timeout=...)), and the sensor is @poke_mode_only since a rescheduled attempt can't reuse the connection anyway. url, header, and message_to_send are template fields, with a JSON renderer for header.
• Registered both in provider.yaml (triggers and sensors sections); get_provider_info.py is regenerated from it via the update-providers-build-files prek hook rather than hand-edited, since that file is release-managed.
• Added the websocket optional extra (websockets>=14.0) to pyproject.toml, plus the same dependency in the provider's dev group so tests run without installing the extra separately. uv.lock updated to match via update-providers-dependencies.
• Added docs/sensors/websocket.rst (three usage examples plus the idempotency caveat above) and linked it from provider.yaml's how-to-guide list, following docs/sensors/file.rst.
• Added the three examples above to example_dags/example_sensors.py, referenced by the docs page via exampleinclude.
• Added unit tests for the trigger (serialization, message received, no-message-to-send, trigger-reconstruction resend) and the sensor (poke success/timeout, handshake timeout, remaining-time budgeting, reschedule-mode rejection, defer-without-poking-first, single-poke-per-timeout, template rendering) in tests/unit/standard/triggers/test_websocket.py and tests/unit/standard/sensors/test_websocket.py, using spec/autospec on all WebSocket mocks per the testing guidelines.

closes: #21139

---

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

• Yes — Claude Code (Sonnet 5)

Generated-by: Claude Code (Sonnet 5) following the guidelines

[ColtenOuO]

self review.

[ColtenOuO]

self review.

[ColtenOuO]

self review again.

[ColtenOuO]

self review again, needs fixed PR description and CI

[ColtenOuO]

One Nit, final review