airflow.providers.common.ai.mixins.cancellable_run¶

Mixin that cancels an operator’s in-flight pydantic-ai run when the task is killed.

Classes¶

CancellableAgentRunMixin

Run a pydantic-ai agent synchronously with kill-time cancellation wired in.

Module Contents¶

class airflow.providers.common.ai.mixins.cancellable_run.CancellableAgentRunMixin[source]¶

Run a pydantic-ai agent synchronously with kill-time cancellation wired in.

The wrapper holds the in-flight run’s CancellationToken so on_kill() can cancel it. Cancelling makes run_sync raise RunCancelled and unwind, giving the agent’s toolsets a chance to exit (tearing down a provisioned sandbox, for one) before SIGKILL rather than leaving the run to die mid-flight.

This needs the Task SDK to call on_kill on SIGTERM, which lands in Airflow 3.0.4 and 3.1.0. On 3.0.0 to 3.0.3 the handler is absent, so a kill runs the pre-existing path (the run continues until SIGKILL).

log: Any[source]¶

run_agent_sync(agent, user_prompt, **run_kwargs)[source]¶

Call agent.run_sync under a fresh cancellation token held for on_kill().

on_kill()[source]¶