Add OpenAIAgentSessionOperator for OpenAI Managed Agents

[kaxil]

Summary

OpenAI shipped a Managed Agents API in openai SDK 3.13 (client.beta.agents, header OpenAI-Beta: agents=v1): persisted agents, sessions with turns, hosted or self-hosted environments, vaults for MCP credentials, and subagents. It is the counterpart of the Anthropic Managed Agents surface this provider's sibling already integrates.

This adds OpenAIAgentSessionOperator, which creates a fresh session with an initial message, waits for the first turn to finish, and returns the session ID. deferrable=True hands the wait to OpenAIAgentSessionTrigger. When XCom pushing is on it also records session_id, turn_id and the turn's token usage. Five OpenAIHook helpers back it (create agent, create session, get session, cancel session, poll session).

Usage from the example DAG:

OpenAIAgentSessionOperator(
    task_id="run_agent",
    input="Explain how Airflow retries affect a task that calls an external API.",
    environment={"type": "none"},
    session_kwargs={"agent": {"model": "gpt-6-astra", "instructions": "..."}},
    deferrable=True,
)

Pass agent_id instead of an inline agent to run a saved agent, and vault_ids or an environment template reference through session_kwargs or environment.

Design rationale

The SDK floor stays at openai>=2.37.0; the Agents API is feature-detected. Raising the floor to 3.13 fails the provider dependency resolution because llama-index-llms-openai (pulled in through the Common AI provider's LlamaIndex extra) still pins openai<3. The hook checks for client.beta.agents at first use and raises an error naming the required version, so every other OpenAI operator keeps working on older SDKs and the new test module skips itself when the installed SDK predates the API.

Each task attempt owns one fresh session and waits only for its first turn. A retry creates a new session. Reusing the previous one would mean picking the right turn out of a session with history the operator did not write. The poll lists turns in ascending order with limit=1 and reads that turn's status. An idle session with no turn yet counts as still waiting, because the initial input can sit queued before a turn object exists, and treating idle as done would return success for work that has not started.

Client-side function tools fail the task. The worker is deferred or asleep while the agent runs, and a function_call required action would need Airflow to run arbitrary tool code mid-wait and post results back. The error message points at service-side tools, MCP servers or hosted environments. Self-hosted environments are supported but their worker (codex exec-server) has to be run independently; an environment_connection required action is treated as still waiting.

Failure paths cancel the active turn but keep the session. Timeout, polling failure, a bad trigger event, a failed turn and a killed task all send agent.session.input.cancel, so a runaway turn stops consuming quota while the items and artifacts remain retrievable by session ID. Up to three consecutive polling errors are retried before the session is cancelled, the same tolerance the Anthropic operators use. A shorter execution_timeout caps the deferral timeout.

A mock-transport test against the real SDK covers the wire format, including the beta header, request payloads, turn listing parameters and the cancel event. The failure path was also exercised against the live API: a turn that failed server-side was cancelled and surfaced as the task error.

Live run

The example DAG run through a local scheduler and triggerer against the live Managed Agents API with deferrable=True. The task deferred, the trigger polled the session and its first turn, and the task resumed and succeeded. Session ID, turn ID and token usage landed in XCom.

Gotchas

• Managed Agents needs openai>=3.13.0 on workers and triggerers. Libraries pinning openai<3, including current LlamaIndex OpenAI integrations, cannot share that environment.
• Cancellation of a killed deferred task relies on trigger on_kill cleanup, which needs Airflow 3.3 or newer. On older versions a killed deferred task leaves the turn running until it finishes or times out server-side.
• Usage pushed to XCom is for the current attempt only and carries try_number; it is not cumulative across retries.

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

[Kunal8954]

Nice work overall — the feature-detection of beta.agents (so the floor stays compatible with openai<3), the session ownership/XCom handling, and the trigger/hook split all look clean. I checked the whole diff including tests and docs.

One thing I verified explicitly since it's easy to get wrong: the defer() path is safe. TaskDeferred subclasses BaseException (not Exception) on every supported Airflow version (>=2.11), so the except Exception: raise in execute() doesn't intercept the deferral and won't call on_kill() / cancel the fresh session at defer time. Good.

A few minor points:

1. test_agent.py — unguarded import httpx2. httpx2 is the vendored HTTP client that only ships in openai 3.x. Because it's imported at module top level (not behind pytest.importorskip), in an env that pins openai<3 (which this PR explicitly claims to keep supported) the entire test module fails at import with ModuleNotFoundError, instead of skipping like the importorskip lines below. Consider guarding it the same way.

2. hooks/openai.py poll_agent_session — session.required_actions may be None. any(action.type == "function_call" for action in session.required_actions) raises TypeError if a response ever omits the field. for action in (session.required_actions or []) is a cheap hardening.

3. operators/agent.py execute — session_kwargs["agent"] assumed to be a dict. self.session_kwargs.get("agent", {}).get("model") raises AttributeError (instead of the intended ValueError) if a user passes an actual Agent pydantic object — which the SDK accepts — or None. getattr(...)/isinstance handling would give a clearer error.

4. hooks/openai.py — turn.usage.model_dump(mode="json") assumes a pydantic model; if an endpoint ever returns usage as a plain dict this breaks. Minor robustness nit.

5. Nit: f"Agent turn {turn.id} {turn.status}: {turn.error}" prints None when turn.error has no message.

Everything else (reserved-key enforcement, max-consecutive-failure retry, cancel-on-timeout/kill, end_time persisted across triggerer restarts, docs exampleinclude path) checks out.

[kaxil]

Thanks for the careful read. I pushed 4785899 for the error message and the agent validation.

A failed or cancelled turn now reports the SDK error's code and message. When a cancelled turn carries no error the message is just the turn and its status, so the None you noticed is gone. A non-dict session_kwargs["agent"] now raises ValueError before any session is created. The SDK's agent parameter is a TypedDict, so a pydantic Agent object was never valid there, and the check mainly catches None and typos.

The httpx2 import is already behind pytest.importorskip at this commit, next to the AgentSession and Turn skips. The bare import was in the first push and the lowest-dependency CI job caught it, which is why the module now skips as a whole on openai<3.

I left required_actions and usage as they are. In the SDK, AgentSession.required_actions is a required List[RequiredAction] and Turn.usage is Optional[TokenUsage], and the None usage case is handled. The client cannot produce a None list or a plain-dict usage, so I would rather not add guards for shapes that cannot occur.

[Kunal8954]

Approved — the fixes look good and I verified the remaining items:

• Deferral is safe: TaskDeferred subclasses BaseException (not Exception) on all supported Airflow versions (>=2.11), so the except-Exception on_kill() path in execute() cannot cancel the fresh session at defer time.
• template_fields_renderers={'environment': 'json', ...} only affects RTIF/UI display (airflow-core renderedtifields.py) and does not mutate session_kwargs/environment into strings at execution, so the reserved-keys and .get('model') checks stay correct after templating.
• httpx2 is behind pytest.importorskip so the module skips cleanly on openai<3.
• The required_actions / usage guards you declined are reasonable given the SDK models guarantee those shapes.

No blockers.

[Kunal8954]

Great work