Add OpenAIAgentSessionOperator for OpenAI Managed Agents

providers/openai/docs/operators/openai.rst

@@ -224,3 +224,76 @@ An example of using the operator:
     :language: python
     :start-after: [START howto_operator_openai_trigger_operator]
     :end-before: [END howto_operator_openai_trigger_operator]
+
+.. _howto/operator:OpenAIAgentSessionOperator:
+
+Managed Agents sessions
+=======================
+
+Use :class:`~airflow.providers.openai.operators.agent.OpenAIAgentSessionOperator`
+to submit a message to OpenAI's Managed Agents service. The service runs the agent
+loop. Airflow waits for the first turn to complete, optionally releasing the worker
+with ``deferrable=True``. This requires OpenAI Python SDK 3.13.0 or newer and access
+to the beta Agents API on your configured endpoint.
+
+The provider's base dependency still permits older SDKs for other OpenAI APIs.
+Install ``openai>=3.13.0`` on both workers and triggerers to use Managed Agents.
+Libraries that require ``openai<3`` (including current LlamaIndex OpenAI LLM
+integrations) cannot share that environment.
+
+.. exampleinclude:: /../../openai/tests/system/openai/example_openai_agent.py
+    :language: python
+    :start-after: [START howto_operator_openai_agent]
+    :end-before: [END howto_operator_openai_agent]
+
+Parameters
+^^^^^^^^^^
+
+* ``input``: Initial user message.
+* ``environment``: SDK environment configuration, such as ``{"type": "none"}``,
+  or an environment template reference for a hosted sandbox.
+* ``agent_id``: An existing saved agent. Alternatively, supply an inline agent
+  with a model in ``session_kwargs["agent"]``.
+* ``session_kwargs``: SDK session creation options, including agent overrides,
+  ``vault_ids`` and ``metadata``. The keys ``input``, ``environment``, ``agent_id``
+  and ``stream`` are reserved.
+* ``conn_id``: OpenAI connection, defaulting to ``openai_default``.
+* ``deferrable``: Whether to release the worker while waiting. Defaults to the
+  Airflow ``operators.default_deferrable`` setting.
+* ``poll_interval``: Seconds between checks, defaulting to 10.
+* ``timeout``: Seconds to wait for completion, defaulting to 3600. A shorter
+  ``execution_timeout`` still applies to a deferred task and preempts the
+  cancel-on-timeout path below.
+
+Transient polling failures are retried; three consecutive failures fail the task.
+
+The operator returns the session ID. When XCom pushing is enabled, it also writes
+``session_id``, ``turn_id`` and the turn's available token ``usage``. Usage includes
+the Airflow ``try_number``; it represents the current attempt, not cumulative spend
+across retries. Full message histories and artifacts are not stored in XCom.
+Retrieve them with ``OpenAIHook().get_conn().beta.agents.sessions.items`` and
+``.artifacts`` using the returned session ID.
+
+Each attempt creates a fresh session. Do not submit additional turns to it while
+this task is running. An idle session without a visible turn is not treated as
+success. Failed or cancelled turns fail the task. Client-side function tools are
+not executed by the operator and fail the task when requested; use service-side
+tools instead. A self-hosted environment must have an independently managed worker.
+
+On timeout or polling failure, the operator requests cancellation of its session's
+active turn. It retains the session and artifacts for inspection. Cancellation does
+not delete the environment or guarantee that its resources have been released.
+Killing a synchronous task also requests cancellation. Cancellation of a killed
+deferred task requires Airflow 3.3 or newer; on older versions, cancel it manually.
+A hard worker termination or Airflow execution timeout can bypass cleanup. Retrying
+the task creates another session and can repeat external side effects.
+
+Hook methods
+^^^^^^^^^^^^
+
+:class:`~airflow.providers.openai.hooks.openai.OpenAIHook` provides
+``create_agent``, ``create_agent_session``, ``get_agent_session`` and
+``cancel_agent_session``. ``poll_agent_session`` checks the first turn of a fresh,
+exclusively owned session; it is not a general waiter for reused sessions.
+For other resources, use the SDK client returned by ``get_conn()``. See the
+`OpenAI Agents API reference <https://developers.openai.com/api/reference/python/resources/beta/subresources/agents>`__.

providers/openai/provider.yaml

@@ -87,11 +87,13 @@ operators:
   - integration-name: OpenAI
     python-modules:
       - airflow.providers.openai.operators.openai
+      - airflow.providers.openai.operators.agent
 
 triggers:
   - integration-name: OpenAI
     python-modules:
       - airflow.providers.openai.triggers.openai
+      - airflow.providers.openai.triggers.agent
 
 connection-types:
   - hook-class-name: airflow.providers.openai.hooks.openai.OpenAIHook

providers/openai/src/airflow/providers/openai/exceptions.py

@@ -30,3 +30,7 @@ class OpenAIBatchTimeout(AirflowException):
 
 class OpenAITriggerEventError(AirflowException):
     """Raise when a deferred task resumes with a missing or malformed trigger event."""
+
+
+class OpenAIAgentSessionError(AirflowException):
+    """Raise when a Managed Agents session fails or cannot run."""

providers/openai/src/airflow/providers/openai/get_provider_info.py

@@ -48,10 +48,22 @@ def get_provider_info():
             {"integration-name": "OpenAI", "python-modules": ["airflow.providers.openai.hooks.openai"]}
         ],
         "operators": [
-            {"integration-name": "OpenAI", "python-modules": ["airflow.providers.openai.operators.openai"]}
+            {
+                "integration-name": "OpenAI",
+                "python-modules": [
+                    "airflow.providers.openai.operators.openai",
+                    "airflow.providers.openai.operators.agent",
+                ],
+            }
         ],
         "triggers": [
-            {"integration-name": "OpenAI", "python-modules": ["airflow.providers.openai.triggers.openai"]}
+            {
+                "integration-name": "OpenAI",
+                "python-modules": [
+                    "airflow.providers.openai.triggers.openai",
+                    "airflow.providers.openai.triggers.agent",
+                ],
+            }
         ],
         "connection-types": [
             {

providers/openai/src/airflow/providers/openai/hooks/openai.py

@@ -56,6 +56,7 @@
 from airflow.providers.common.compat.module_loading import import_string
 from airflow.providers.common.compat.sdk import BaseHook
 from airflow.providers.openai.exceptions import (
+    OpenAIAgentSessionError,
     OpenAIBatchJobException,
     OpenAIBatchTimeout,
     OpenAITriggerEventError,
@@ -714,3 +715,69 @@ def cancel_batch(self, batch_id: str) -> Batch:
         """
         batch = self.conn.batches.cancel(batch_id=batch_id)
         return batch
+
+    #: Consecutive ``poll_agent_session`` failures tolerated before an agent wait gives up.
+    MAX_CONSECUTIVE_POLL_FAILURES = 3
+
+    @cached_property
+    def _managed_agents(self) -> Any:
+        agents = getattr(self.conn.beta, "agents", None)
+        if agents is None:
+            raise OpenAIAgentSessionError(
+                "Managed Agents requires openai>=3.13.0. Upgrade the OpenAI SDK on workers and triggerers."
+            )
+        return agents
+
+    def create_agent(self, **kwargs: Any) -> Any:
+        """Create a reusable Managed Agent using the SDK's agent configuration arguments."""
+        return self._managed_agents.create(**kwargs)
+
+    def create_agent_session(self, *, input: str, environment: dict[str, Any], **kwargs: Any) -> Any:
+        """Create a fresh Managed Agents session and submit its initial turn."""
+        if "stream" in kwargs:
+            raise ValueError("create_agent_session does not support streaming")
+        return self._managed_agents.sessions.create(
+            input=input, environment=environment, stream=False, **kwargs
+        )
+
+    def get_agent_session(self, session_id: str) -> Any:
+        """Retrieve a Managed Agents session, including required actions and usage."""
+        return self._managed_agents.sessions.retrieve(session_id)
+
+    def cancel_agent_session(self, session_id: str) -> None:
+        """Request cancellation of the session's active turn, preserving its history and artifacts."""
+        self._managed_agents.sessions.events.create(
+            session_id, events=[{"type": "agent.session.input.cancel"}]
+        )
+
+    def poll_agent_session(self, session_id: str) -> dict[str, Any] | None:
+        """
+        Check the first turn of a fresh, exclusively owned session.
+
+        Return a terminal result, or ``None`` while waiting. An idle session without
+        a visible turn is not completion: the submitted input may still be queued.
+        This helper must not be used to wait for subsequent turns of a reused session.
+        """
+        session = self.get_agent_session(session_id)
+        turns = self._managed_agents.sessions.turns.list(session_id, order="asc", limit=1)
+        turn = turns.data[0] if turns.data else None
+        result: dict[str, Any] = {"session_id": session_id}
+        if turn is not None:
+            result["turn_id"] = turn.id
+            result["usage"] = turn.usage.model_dump(mode="json") if turn.usage is not None else None
+            if turn.status == "completed":
+                return {**result, "status": "success"}
+            if turn.status in {"failed", "cancelled"}:
+                message = f"Agent turn {turn.id} {turn.status}"
+                if turn.error is not None:
+                    message += f" ({turn.error.code}): {turn.error.message}"
+                return {**result, "status": "error", "message": message}
+        if session.status == "failed":
+            return {**result, "status": "error", "message": f"Agent session failed: {session.error}"}
+        if any(action.type == "function_call" for action in session.required_actions):
+            return {
+                **result,
+                "status": "error",
+                "message": "The agent requested a client-side function tool. Use server-side tools with this operator.",
+            }
+        return None

providers/openai/src/airflow/providers/openai/operators/agent.py

@@ -0,0 +1,179 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import math
+import time
+from collections.abc import Sequence
+from datetime import timedelta
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.compat.sdk import BaseOperator, conf
+from airflow.providers.openai.exceptions import OpenAIAgentSessionError, OpenAITriggerEventError
+from airflow.providers.openai.hooks.openai import OpenAIHook
+from airflow.providers.openai.triggers.agent import OpenAIAgentSessionTrigger
+
+if TYPE_CHECKING:
+    from airflow.providers.common.compat.sdk import Context
+
+
+class OpenAIAgentSessionOperator(BaseOperator):
+    """
+    Run one turn in a fresh OpenAI Managed Agents session and return its session ID.
+
+    The session is retained for downstream retrieval of items and artifacts. A retry
+    creates a new session and can repeat external side effects.
+
+    :param input: Initial user message. (templated)
+    :param environment: SDK environment configuration or template reference. (templated)
+    :param agent_id: Saved agent ID. Alternatively supply agent.model in session_kwargs. (templated)
+    :param session_kwargs: Additional SDK session creation arguments, such as agent,
+        vault_ids and metadata. Must not contain input, environment, agent_id or stream. (templated)
+    :param conn_id: OpenAI connection ID. (templated)
+    :param deferrable: Release the worker while waiting for completion.
+    :param poll_interval: Seconds between polls.
+    :param timeout: Maximum seconds to wait for the initial turn. A shorter
+        ``execution_timeout`` still applies and preempts the cancel-on-timeout path.
+    """
+
+    template_fields: Sequence[str] = ("input", "environment", "agent_id", "session_kwargs", "conn_id")
+    template_fields_renderers = {"environment": "json", "session_kwargs": "json"}
+
+    def __init__(
+        self,
+        *,
+        input: str,
+        environment: dict[str, Any],
+        agent_id: str | None = None,
+        session_kwargs: dict[str, Any] | None = None,
+        conn_id: str = OpenAIHook.default_conn_name,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        poll_interval: float = 10,
+        timeout: float = 3600,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        for name, value in (("poll_interval", poll_interval), ("timeout", timeout)):
+            if not math.isfinite(value) or value <= 0:
+                raise ValueError(f"{name} must be finite and positive")
+        self.input = input
+        self.environment = environment
+        self.agent_id = agent_id
+        self.session_kwargs = session_kwargs or {}
+        self.conn_id = conn_id
+        self.deferrable = deferrable
+        self.poll_interval = poll_interval
+        self.timeout = timeout
+        self.session_id: str | None = None
+
+    @cached_property
+    def hook(self) -> OpenAIHook:
+        """Return the connection's OpenAI hook."""
+        return OpenAIHook(conn_id=self.conn_id)
+
+    def execute(self, context: Context) -> str:
+        reserved = {"input", "environment", "agent_id", "stream"} & self.session_kwargs.keys()
+        if reserved:
+            raise ValueError(f"Reserved session_kwargs: {sorted(reserved)}")
+        agent = self.session_kwargs.get("agent")
+        if agent is not None and not isinstance(agent, dict):
+            raise ValueError("session_kwargs['agent'] must be a dict of SDK agent fields")
+        if not self.agent_id and not (agent or {}).get("model"):
+            raise ValueError("Supply agent_id or session_kwargs['agent']['model']")
+        if not self.input:
+            raise ValueError("input must not be empty")
+        create_kwargs = dict(self.session_kwargs)
+        if self.agent_id:
+            create_kwargs["agent_id"] = self.agent_id
+        session = self.hook.create_agent_session(
+            input=self.input, environment=self.environment, **create_kwargs
+        )
+        self.session_id = session.id
+        try:
+            if self.do_xcom_push:
+                context["ti"].xcom_push(key="session_id", value=session.id)
+            if self.deferrable:
+                self.defer(
+                    trigger=OpenAIAgentSessionTrigger(
+                        conn_id=self.conn_id,
+                        session_id=session.id,
+                        poll_interval=self.poll_interval,
+                        end_time=time.time() + self.timeout,
+                    ),
+                    method_name="execute_complete",
+                    kwargs={"session_id": session.id},
+                    timeout=self.execution_timeout
+                    or timedelta(seconds=self.timeout + self.poll_interval + 60),
+                )
+            deadline = time.monotonic() + self.timeout
+            consecutive_failures = 0
+            while time.monotonic() < deadline:
+                try:
+                    result = self.hook.poll_agent_session(session.id)
+                except Exception as exc:
+                    consecutive_failures += 1
+                    if consecutive_failures >= OpenAIHook.MAX_CONSECUTIVE_POLL_FAILURES:
+                        raise
+                    self.log.warning("Polling agent session %s failed (%s); retrying.", session.id, exc)
+                else:
+                    consecutive_failures = 0
+                    if result is not None:
+                        break
+                time.sleep(min(self.poll_interval, max(0, deadline - time.monotonic())))
+            else:
+                raise OpenAIAgentSessionError(f"Agent session {session.id} timed out")
+        except Exception:
+            self.on_kill()
+            raise
+        return self.execute_complete(context, result, session_id=session.id)
+
+    def execute_complete(self, context: Context, event: Any = None, session_id: str | None = None) -> str:
+        """Validate completion, record usage, and return the owned session ID."""
+        self.session_id = session_id or self.session_id
+        if (
+            not isinstance(event, dict)
+            or event.get("status") not in ("success", "error", "timeout")
+            or not self.session_id
+            or event.get("session_id") != self.session_id
+        ):
+            self.on_kill()
+            raise OpenAITriggerEventError("Invalid Managed Agents trigger event")
+        if event["status"] != "success":
+            self.on_kill()
+        if self.do_xcom_push:
+            try:
+                if event.get("turn_id"):
+                    context["ti"].xcom_push(key="turn_id", value=event["turn_id"])
+                usage = event.get("usage")
+                if usage is not None:
+                    context["ti"].xcom_push(
+                        key="usage", value={**usage, "try_number": context["ti"].try_number}
+                    )
+            except Exception:
+                self.log.exception("Could not record agent turn usage for session %s", self.session_id)
+        if event["status"] != "success":
+            raise OpenAIAgentSessionError(event.get("message", "Agent session failed"))
+        return self.session_id
+
+    def on_kill(self) -> None:
+        """Request cancellation without deleting session history or artifacts."""
+        if self.session_id:
+            try:
+                self.hook.cancel_agent_session(self.session_id)
+            except Exception:
+                self.log.exception("Could not cancel agent session %s", self.session_id)