Upload folder using huggingface_hub

README.md

@@ -1,12 +1,6 @@
 ---
-title: Deep Research-personal
-emoji: 🔥
-colorFrom: gray
-colorTo: blue
-sdk: gradio
-sdk_version: 6.1.0
+title: deep_research-personal
 app_file: app.py
-pinned: false
+sdk: gradio
+sdk_version: 6.0.2
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

agents/__init__.py

@@ -0,0 +1,319 @@
+import logging
+import sys
+from typing import Literal
+
+from openai import AsyncOpenAI
+
+from . import _config
+from .agent import (
+    Agent,
+    AgentBase,
+    StopAtTools,
+    ToolsToFinalOutputFunction,
+    ToolsToFinalOutputResult,
+)
+from .agent_output import AgentOutputSchema, AgentOutputSchemaBase
+from .computer import AsyncComputer, Button, Computer, Environment
+from .exceptions import (
+    AgentsException,
+    InputGuardrailTripwireTriggered,
+    MaxTurnsExceeded,
+    ModelBehaviorError,
+    OutputGuardrailTripwireTriggered,
+    RunErrorDetails,
+    ToolInputGuardrailTripwireTriggered,
+    ToolOutputGuardrailTripwireTriggered,
+    UserError,
+)
+from .guardrail import (
+    GuardrailFunctionOutput,
+    InputGuardrail,
+    InputGuardrailResult,
+    OutputGuardrail,
+    OutputGuardrailResult,
+    input_guardrail,
+    output_guardrail,
+)
+from .handoffs import Handoff, HandoffInputData, HandoffInputFilter, handoff
+from .items import (
+    HandoffCallItem,
+    HandoffOutputItem,
+    ItemHelpers,
+    MessageOutputItem,
+    ModelResponse,
+    ReasoningItem,
+    RunItem,
+    ToolCallItem,
+    ToolCallOutputItem,
+    TResponseInputItem,
+)
+from .lifecycle import AgentHooks, RunHooks
+from .memory import OpenAIConversationsSession, Session, SessionABC, SQLiteSession
+from .model_settings import ModelSettings
+from .models.interface import Model, ModelProvider, ModelTracing
+from .models.multi_provider import MultiProvider
+from .models.openai_chatcompletions import OpenAIChatCompletionsModel
+from .models.openai_provider import OpenAIProvider
+from .models.openai_responses import OpenAIResponsesModel
+from .prompts import DynamicPromptFunction, GenerateDynamicPromptData, Prompt
+from .repl import run_demo_loop
+from .result import RunResult, RunResultStreaming
+from .run import RunConfig, Runner
+from .run_context import RunContextWrapper, TContext
+from .stream_events import (
+    AgentUpdatedStreamEvent,
+    RawResponsesStreamEvent,
+    RunItemStreamEvent,
+    StreamEvent,
+)
+from .tool import (
+    CodeInterpreterTool,
+    ComputerTool,
+    FileSearchTool,
+    FunctionTool,
+    FunctionToolResult,
+    HostedMCPTool,
+    ImageGenerationTool,
+    LocalShellCommandRequest,
+    LocalShellExecutor,
+    LocalShellTool,
+    MCPToolApprovalFunction,
+    MCPToolApprovalFunctionResult,
+    MCPToolApprovalRequest,
+    Tool,
+    WebSearchTool,
+    default_tool_error_function,
+    function_tool,
+)
+from .tool_guardrails import (
+    ToolGuardrailFunctionOutput,
+    ToolInputGuardrail,
+    ToolInputGuardrailData,
+    ToolInputGuardrailResult,
+    ToolOutputGuardrail,
+    ToolOutputGuardrailData,
+    ToolOutputGuardrailResult,
+    tool_input_guardrail,
+    tool_output_guardrail,
+)
+from .tracing import (
+    AgentSpanData,
+    CustomSpanData,
+    FunctionSpanData,
+    GenerationSpanData,
+    GuardrailSpanData,
+    HandoffSpanData,
+    MCPListToolsSpanData,
+    Span,
+    SpanData,
+    SpanError,
+    SpeechGroupSpanData,
+    SpeechSpanData,
+    Trace,
+    TracingProcessor,
+    TranscriptionSpanData,
+    add_trace_processor,
+    agent_span,
+    custom_span,
+    function_span,
+    gen_span_id,
+    gen_trace_id,
+    generation_span,
+    get_current_span,
+    get_current_trace,
+    guardrail_span,
+    handoff_span,
+    mcp_tools_span,
+    set_trace_processors,
+    set_trace_provider,
+    set_tracing_disabled,
+    set_tracing_export_api_key,
+    speech_group_span,
+    speech_span,
+    trace,
+    transcription_span,
+)
+from .usage import Usage
+from .version import __version__
+
+
+def set_default_openai_key(key: str, use_for_tracing: bool = True) -> None:
+    """Set the default OpenAI API key to use for LLM requests (and optionally tracing()). This is
+    only necessary if the OPENAI_API_KEY environment variable is not already set.
+
+    If provided, this key will be used instead of the OPENAI_API_KEY environment variable.
+
+    Args:
+        key: The OpenAI key to use.
+        use_for_tracing: Whether to also use this key to send traces to OpenAI. Defaults to True
+            If False, you'll either need to set the OPENAI_API_KEY environment variable or call
+            set_tracing_export_api_key() with the API key you want to use for tracing.
+    """
+    _config.set_default_openai_key(key, use_for_tracing)
+
+
+def set_default_openai_client(client: AsyncOpenAI, use_for_tracing: bool = True) -> None:
+    """Set the default OpenAI client to use for LLM requests and/or tracing. If provided, this
+    client will be used instead of the default OpenAI client.
+
+    Args:
+        client: The OpenAI client to use.
+        use_for_tracing: Whether to use the API key from this client for uploading traces. If False,
+            you'll either need to set the OPENAI_API_KEY environment variable or call
+            set_tracing_export_api_key() with the API key you want to use for tracing.
+    """
+    _config.set_default_openai_client(client, use_for_tracing)
+
+
+def set_default_openai_api(api: Literal["chat_completions", "responses"]) -> None:
+    """Set the default API to use for OpenAI LLM requests. By default, we will use the responses API
+    but you can set this to use the chat completions API instead.
+    """
+    _config.set_default_openai_api(api)
+
+
+def enable_verbose_stdout_logging():
+    """Enables verbose logging to stdout. This is useful for debugging."""
+    logger = logging.getLogger("openai.agents")
+    logger.setLevel(logging.DEBUG)
+    logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+__all__ = [
+    "Agent",
+    "AgentBase",
+    "StopAtTools",
+    "ToolsToFinalOutputFunction",
+    "ToolsToFinalOutputResult",
+    "Runner",
+    "run_demo_loop",
+    "Model",
+    "ModelProvider",
+    "ModelTracing",
+    "ModelSettings",
+    "OpenAIChatCompletionsModel",
+    "MultiProvider",
+    "OpenAIProvider",
+    "OpenAIResponsesModel",
+    "AgentOutputSchema",
+    "AgentOutputSchemaBase",
+    "Computer",
+    "AsyncComputer",
+    "Environment",
+    "Button",
+    "AgentsException",
+    "InputGuardrailTripwireTriggered",
+    "OutputGuardrailTripwireTriggered",
+    "ToolInputGuardrailTripwireTriggered",
+    "ToolOutputGuardrailTripwireTriggered",
+    "DynamicPromptFunction",
+    "GenerateDynamicPromptData",
+    "Prompt",
+    "MaxTurnsExceeded",
+    "ModelBehaviorError",
+    "UserError",
+    "InputGuardrail",
+    "InputGuardrailResult",
+    "OutputGuardrail",
+    "OutputGuardrailResult",
+    "GuardrailFunctionOutput",
+    "input_guardrail",
+    "output_guardrail",
+    "ToolInputGuardrail",
+    "ToolOutputGuardrail",
+    "ToolGuardrailFunctionOutput",
+    "ToolInputGuardrailData",
+    "ToolInputGuardrailResult",
+    "ToolOutputGuardrailData",
+    "ToolOutputGuardrailResult",
+    "tool_input_guardrail",
+    "tool_output_guardrail",
+    "handoff",
+    "Handoff",
+    "HandoffInputData",
+    "HandoffInputFilter",
+    "TResponseInputItem",
+    "MessageOutputItem",
+    "ModelResponse",
+    "RunItem",
+    "HandoffCallItem",
+    "HandoffOutputItem",
+    "ToolCallItem",
+    "ToolCallOutputItem",
+    "ReasoningItem",
+    "ItemHelpers",
+    "RunHooks",
+    "AgentHooks",
+    "Session",
+    "SessionABC",
+    "SQLiteSession",
+    "OpenAIConversationsSession",
+    "RunContextWrapper",
+    "TContext",
+    "RunErrorDetails",
+    "RunResult",
+    "RunResultStreaming",
+    "RunConfig",
+    "RawResponsesStreamEvent",
+    "RunItemStreamEvent",
+    "AgentUpdatedStreamEvent",
+    "StreamEvent",
+    "FunctionTool",
+    "FunctionToolResult",
+    "ComputerTool",
+    "FileSearchTool",
+    "CodeInterpreterTool",
+    "ImageGenerationTool",
+    "LocalShellCommandRequest",
+    "LocalShellExecutor",
+    "LocalShellTool",
+    "Tool",
+    "WebSearchTool",
+    "HostedMCPTool",
+    "MCPToolApprovalFunction",
+    "MCPToolApprovalRequest",
+    "MCPToolApprovalFunctionResult",
+    "function_tool",
+    "Usage",
+    "add_trace_processor",
+    "agent_span",
+    "custom_span",
+    "function_span",
+    "generation_span",
+    "get_current_span",
+    "get_current_trace",
+    "guardrail_span",
+    "handoff_span",
+    "set_trace_processors",
+    "set_trace_provider",
+    "set_tracing_disabled",
+    "speech_group_span",
+    "transcription_span",
+    "speech_span",
+    "mcp_tools_span",
+    "trace",
+    "Trace",
+    "TracingProcessor",
+    "SpanError",
+    "Span",
+    "SpanData",
+    "AgentSpanData",
+    "CustomSpanData",
+    "FunctionSpanData",
+    "GenerationSpanData",
+    "GuardrailSpanData",
+    "HandoffSpanData",
+    "SpeechGroupSpanData",
+    "SpeechSpanData",
+    "MCPListToolsSpanData",
+    "TranscriptionSpanData",
+    "set_default_openai_key",
+    "set_default_openai_client",
+    "set_default_openai_api",
+    "set_tracing_export_api_key",
+    "enable_verbose_stdout_logging",
+    "gen_trace_id",
+    "gen_span_id",
+    "default_tool_error_function",
+    "__version__",
+]

agents/_config.py

@@ -0,0 +1,26 @@
+from openai import AsyncOpenAI
+from typing_extensions import Literal
+
+from .models import _openai_shared
+from .tracing import set_tracing_export_api_key
+
+
+def set_default_openai_key(key: str, use_for_tracing: bool) -> None:
+    _openai_shared.set_default_openai_key(key)
+
+    if use_for_tracing:
+        set_tracing_export_api_key(key)
+
+
+def set_default_openai_client(client: AsyncOpenAI, use_for_tracing: bool) -> None:
+    _openai_shared.set_default_openai_client(client)
+
+    if use_for_tracing:
+        set_tracing_export_api_key(client.api_key)
+
+
+def set_default_openai_api(api: Literal["chat_completions", "responses"]) -> None:
+    if api == "chat_completions":
+        _openai_shared.set_use_responses_by_default(False)
+    else:
+        _openai_shared.set_use_responses_by_default(True)

agents/_debug.py

@@ -0,0 +1,28 @@
+import os
+
+
+def _debug_flag_enabled(flag: str, default: bool = False) -> bool:
+    flag_value = os.getenv(flag)
+    if flag_value is None:
+        return default
+    else:
+        return flag_value == "1" or flag_value.lower() == "true"
+
+
+def _load_dont_log_model_data() -> bool:
+    return _debug_flag_enabled("OPENAI_AGENTS_DONT_LOG_MODEL_DATA", default=True)
+
+
+def _load_dont_log_tool_data() -> bool:
+    return _debug_flag_enabled("OPENAI_AGENTS_DONT_LOG_TOOL_DATA", default=True)
+
+
+DONT_LOG_MODEL_DATA = _load_dont_log_model_data()
+"""By default we don't log LLM inputs/outputs, to prevent exposing sensitive information. Set this
+flag to enable logging them.
+"""
+
+DONT_LOG_TOOL_DATA = _load_dont_log_tool_data()
+"""By default we don't log tool call inputs/outputs, to prevent exposing sensitive information. Set
+this flag to enable logging them.
+"""

agents/_run_impl.py

@@ -0,0 +1,1442 @@
+from __future__ import annotations
+
+import asyncio
+import dataclasses
+import inspect
+from collections.abc import Awaitable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, cast
+
+from openai.types.responses import (
+    ResponseComputerToolCall,
+    ResponseFileSearchToolCall,
+    ResponseFunctionToolCall,
+    ResponseFunctionWebSearch,
+    ResponseOutputMessage,
+)
+from openai.types.responses.response_code_interpreter_tool_call import (
+    ResponseCodeInterpreterToolCall,
+)
+from openai.types.responses.response_computer_tool_call import (
+    ActionClick,
+    ActionDoubleClick,
+    ActionDrag,
+    ActionKeypress,
+    ActionMove,
+    ActionScreenshot,
+    ActionScroll,
+    ActionType,
+    ActionWait,
+)
+from openai.types.responses.response_input_item_param import (
+    ComputerCallOutputAcknowledgedSafetyCheck,
+)
+from openai.types.responses.response_input_param import ComputerCallOutput, McpApprovalResponse
+from openai.types.responses.response_output_item import (
+    ImageGenerationCall,
+    LocalShellCall,
+    McpApprovalRequest,
+    McpCall,
+    McpListTools,
+)
+from openai.types.responses.response_reasoning_item import ResponseReasoningItem
+
+from .agent import Agent, ToolsToFinalOutputResult
+from .agent_output import AgentOutputSchemaBase
+from .computer import AsyncComputer, Computer
+from .exceptions import (
+    AgentsException,
+    ModelBehaviorError,
+    ToolInputGuardrailTripwireTriggered,
+    ToolOutputGuardrailTripwireTriggered,
+    UserError,
+)
+from .guardrail import InputGuardrail, InputGuardrailResult, OutputGuardrail, OutputGuardrailResult
+from .handoffs import Handoff, HandoffInputData
+from .items import (
+    HandoffCallItem,
+    HandoffOutputItem,
+    ItemHelpers,
+    MCPApprovalRequestItem,
+    MCPApprovalResponseItem,
+    MCPListToolsItem,
+    MessageOutputItem,
+    ModelResponse,
+    ReasoningItem,
+    RunItem,
+    ToolCallItem,
+    ToolCallOutputItem,
+    TResponseInputItem,
+)
+from .lifecycle import RunHooks
+from .logger import logger
+from .model_settings import ModelSettings
+from .models.interface import ModelTracing
+from .run_context import RunContextWrapper, TContext
+from .stream_events import RunItemStreamEvent, StreamEvent
+from .tool import (
+    ComputerTool,
+    ComputerToolSafetyCheckData,
+    FunctionTool,
+    FunctionToolResult,
+    HostedMCPTool,
+    LocalShellCommandRequest,
+    LocalShellTool,
+    MCPToolApprovalRequest,
+    Tool,
+)
+from .tool_context import ToolContext
+from .tool_guardrails import (
+    ToolInputGuardrailData,
+    ToolInputGuardrailResult,
+    ToolOutputGuardrailData,
+    ToolOutputGuardrailResult,
+)
+from .tracing import (
+    SpanError,
+    Trace,
+    function_span,
+    get_current_trace,
+    guardrail_span,
+    handoff_span,
+    trace,
+)
+from .util import _coro, _error_tracing
+
+if TYPE_CHECKING:
+    from .run import RunConfig
+
+
+class QueueCompleteSentinel:
+    pass
+
+
+QUEUE_COMPLETE_SENTINEL = QueueCompleteSentinel()
+
+_NOT_FINAL_OUTPUT = ToolsToFinalOutputResult(is_final_output=False, final_output=None)
+
+
+@dataclass
+class AgentToolUseTracker:
+    agent_to_tools: list[tuple[Agent, list[str]]] = field(default_factory=list)
+    """Tuple of (agent, list of tools used). Can't use a dict because agents aren't hashable."""
+
+    def add_tool_use(self, agent: Agent[Any], tool_names: list[str]) -> None:
+        existing_data = next((item for item in self.agent_to_tools if item[0] == agent), None)
+        if existing_data:
+            existing_data[1].extend(tool_names)
+        else:
+            self.agent_to_tools.append((agent, tool_names))
+
+    def has_used_tools(self, agent: Agent[Any]) -> bool:
+        existing_data = next((item for item in self.agent_to_tools if item[0] == agent), None)
+        return existing_data is not None and len(existing_data[1]) > 0
+
+
+@dataclass
+class ToolRunHandoff:
+    handoff: Handoff
+    tool_call: ResponseFunctionToolCall
+
+
+@dataclass
+class ToolRunFunction:
+    tool_call: ResponseFunctionToolCall
+    function_tool: FunctionTool
+
+
+@dataclass
+class ToolRunComputerAction:
+    tool_call: ResponseComputerToolCall
+    computer_tool: ComputerTool
+
+
+@dataclass
+class ToolRunMCPApprovalRequest:
+    request_item: McpApprovalRequest
+    mcp_tool: HostedMCPTool
+
+
+@dataclass
+class ToolRunLocalShellCall:
+    tool_call: LocalShellCall
+    local_shell_tool: LocalShellTool
+
+
+@dataclass
+class ProcessedResponse:
+    new_items: list[RunItem]
+    handoffs: list[ToolRunHandoff]
+    functions: list[ToolRunFunction]
+    computer_actions: list[ToolRunComputerAction]
+    local_shell_calls: list[ToolRunLocalShellCall]
+    tools_used: list[str]  # Names of all tools used, including hosted tools
+    mcp_approval_requests: list[ToolRunMCPApprovalRequest]  # Only requests with callbacks
+
+    def has_tools_or_approvals_to_run(self) -> bool:
+        # Handoffs, functions and computer actions need local processing
+        # Hosted tools have already run, so there's nothing to do.
+        return any(
+            [
+                self.handoffs,
+                self.functions,
+                self.computer_actions,
+                self.local_shell_calls,
+                self.mcp_approval_requests,
+            ]
+        )
+
+
+@dataclass
+class NextStepHandoff:
+    new_agent: Agent[Any]
+
+
+@dataclass
+class NextStepFinalOutput:
+    output: Any
+
+
+@dataclass
+class NextStepRunAgain:
+    pass
+
+
+@dataclass
+class SingleStepResult:
+    original_input: str | list[TResponseInputItem]
+    """The input items i.e. the items before run() was called. May be mutated by handoff input
+    filters."""
+
+    model_response: ModelResponse
+    """The model response for the current step."""
+
+    pre_step_items: list[RunItem]
+    """Items generated before the current step."""
+
+    new_step_items: list[RunItem]
+    """Items generated during this current step."""
+
+    next_step: NextStepHandoff | NextStepFinalOutput | NextStepRunAgain
+    """The next step to take."""
+
+    tool_input_guardrail_results: list[ToolInputGuardrailResult]
+    """Tool input guardrail results from this step."""
+
+    tool_output_guardrail_results: list[ToolOutputGuardrailResult]
+    """Tool output guardrail results from this step."""
+
+    @property
+    def generated_items(self) -> list[RunItem]:
+        """Items generated during the agent run (i.e. everything generated after
+        `original_input`)."""
+        return self.pre_step_items + self.new_step_items
+
+
+def get_model_tracing_impl(
+    tracing_disabled: bool, trace_include_sensitive_data: bool
+) -> ModelTracing:
+    if tracing_disabled:
+        return ModelTracing.DISABLED
+    elif trace_include_sensitive_data:
+        return ModelTracing.ENABLED
+    else:
+        return ModelTracing.ENABLED_WITHOUT_DATA
+
+
+class RunImpl:
+    @classmethod
+    async def execute_tools_and_side_effects(
+        cls,
+        *,
+        agent: Agent[TContext],
+        # The original input to the Runner
+        original_input: str | list[TResponseInputItem],
+        # Everything generated by Runner since the original input, but before the current step
+        pre_step_items: list[RunItem],
+        new_response: ModelResponse,
+        processed_response: ProcessedResponse,
+        output_schema: AgentOutputSchemaBase | None,
+        hooks: RunHooks[TContext],
+        context_wrapper: RunContextWrapper[TContext],
+        run_config: RunConfig,
+    ) -> SingleStepResult:
+        # Make a copy of the generated items
+        pre_step_items = list(pre_step_items)
+
+        new_step_items: list[RunItem] = []
+        new_step_items.extend(processed_response.new_items)
+
+        # First, lets run the tool calls - function tools and computer actions
+        (
+            (function_results, tool_input_guardrail_results, tool_output_guardrail_results),
+            computer_results,
+        ) = await asyncio.gather(
+            cls.execute_function_tool_calls(
+                agent=agent,
+                tool_runs=processed_response.functions,
+                hooks=hooks,
+                context_wrapper=context_wrapper,
+                config=run_config,
+            ),
+            cls.execute_computer_actions(
+                agent=agent,
+                actions=processed_response.computer_actions,
+                hooks=hooks,
+                context_wrapper=context_wrapper,
+                config=run_config,
+            ),
+        )
+        new_step_items.extend([result.run_item for result in function_results])
+        new_step_items.extend(computer_results)
+
+        # Next, run the MCP approval requests
+        if processed_response.mcp_approval_requests:
+            approval_results = await cls.execute_mcp_approval_requests(
+                agent=agent,
+                approval_requests=processed_response.mcp_approval_requests,
+                context_wrapper=context_wrapper,
+            )
+            new_step_items.extend(approval_results)
+
+        # Next, check if there are any handoffs
+        if run_handoffs := processed_response.handoffs:
+            return await cls.execute_handoffs(
+                agent=agent,
+                original_input=original_input,
+                pre_step_items=pre_step_items,
+                new_step_items=new_step_items,
+                new_response=new_response,
+                run_handoffs=run_handoffs,
+                hooks=hooks,
+                context_wrapper=context_wrapper,
+                run_config=run_config,
+            )
+
+        # Next, we'll check if the tool use should result in a final output
+        check_tool_use = await cls._check_for_final_output_from_tools(
+            agent=agent,
+            tool_results=function_results,
+            context_wrapper=context_wrapper,
+            config=run_config,
+        )
+
+        if check_tool_use.is_final_output:
+            # If the output type is str, then let's just stringify it
+            if not agent.output_type or agent.output_type is str:
+                check_tool_use.final_output = str(check_tool_use.final_output)
+
+            if check_tool_use.final_output is None:
+                logger.error(
+                    "Model returned a final output of None. Not raising an error because we assume"
+                    "you know what you're doing."
+                )
+
+            return await cls.execute_final_output(
+                agent=agent,
+                original_input=original_input,
+                new_response=new_response,
+                pre_step_items=pre_step_items,
+                new_step_items=new_step_items,
+                final_output=check_tool_use.final_output,
+                hooks=hooks,
+                context_wrapper=context_wrapper,
+                tool_input_guardrail_results=tool_input_guardrail_results,
+                tool_output_guardrail_results=tool_output_guardrail_results,
+            )
+
+        # Now we can check if the model also produced a final output
+        message_items = [item for item in new_step_items if isinstance(item, MessageOutputItem)]
+
+        # We'll use the last content output as the final output
+        potential_final_output_text = (
+            ItemHelpers.extract_last_text(message_items[-1].raw_item) if message_items else None
+        )
+
+        # Generate final output only when there are no pending tool calls or approval requests.
+        if not processed_response.has_tools_or_approvals_to_run():
+            if output_schema and not output_schema.is_plain_text() and potential_final_output_text:
+                final_output = output_schema.validate_json(potential_final_output_text)
+                return await cls.execute_final_output(
+                    agent=agent,
+                    original_input=original_input,
+                    new_response=new_response,
+                    pre_step_items=pre_step_items,
+                    new_step_items=new_step_items,
+                    final_output=final_output,
+                    hooks=hooks,
+                    context_wrapper=context_wrapper,
+                    tool_input_guardrail_results=tool_input_guardrail_results,
+                    tool_output_guardrail_results=tool_output_guardrail_results,
+                )
+            elif not output_schema or output_schema.is_plain_text():
+                return await cls.execute_final_output(
+                    agent=agent,
+                    original_input=original_input,
+                    new_response=new_response,
+                    pre_step_items=pre_step_items,
+                    new_step_items=new_step_items,
+                    final_output=potential_final_output_text or "",
+                    hooks=hooks,
+                    context_wrapper=context_wrapper,
+                    tool_input_guardrail_results=tool_input_guardrail_results,
+                    tool_output_guardrail_results=tool_output_guardrail_results,
+                )
+
+        # If there's no final output, we can just run again
+        return SingleStepResult(
+            original_input=original_input,
+            model_response=new_response,
+            pre_step_items=pre_step_items,
+            new_step_items=new_step_items,
+            next_step=NextStepRunAgain(),
+            tool_input_guardrail_results=tool_input_guardrail_results,
+            tool_output_guardrail_results=tool_output_guardrail_results,
+        )
+
+    @classmethod
+    def maybe_reset_tool_choice(
+        cls, agent: Agent[Any], tool_use_tracker: AgentToolUseTracker, model_settings: ModelSettings
+    ) -> ModelSettings:
+        """Resets tool choice to None if the agent has used tools and the agent's reset_tool_choice
+        flag is True."""
+
+        if agent.reset_tool_choice is True and tool_use_tracker.has_used_tools(agent):
+            return dataclasses.replace(model_settings, tool_choice=None)
+
+        return model_settings
+
+    @classmethod
+    def process_model_response(
+        cls,
+        *,
+        agent: Agent[Any],
+        all_tools: list[Tool],
+        response: ModelResponse,
+        output_schema: AgentOutputSchemaBase | None,
+        handoffs: list[Handoff],
+    ) -> ProcessedResponse:
+        items: list[RunItem] = []
+
+        run_handoffs = []
+        functions = []
+        computer_actions = []
+        local_shell_calls = []
+        mcp_approval_requests = []
+        tools_used: list[str] = []
+        handoff_map = {handoff.tool_name: handoff for handoff in handoffs}
+        function_map = {tool.name: tool for tool in all_tools if isinstance(tool, FunctionTool)}
+        computer_tool = next((tool for tool in all_tools if isinstance(tool, ComputerTool)), None)
+        local_shell_tool = next(
+            (tool for tool in all_tools if isinstance(tool, LocalShellTool)), None
+        )
+        hosted_mcp_server_map = {
+            tool.tool_config["server_label"]: tool
+            for tool in all_tools
+            if isinstance(tool, HostedMCPTool)
+        }
+
+        for output in response.output:
+            if isinstance(output, ResponseOutputMessage):
+                items.append(MessageOutputItem(raw_item=output, agent=agent))
+            elif isinstance(output, ResponseFileSearchToolCall):
+                items.append(ToolCallItem(raw_item=output, agent=agent))
+                tools_used.append("file_search")
+            elif isinstance(output, ResponseFunctionWebSearch):
+                items.append(ToolCallItem(raw_item=output, agent=agent))
+                tools_used.append("web_search")
+            elif isinstance(output, ResponseReasoningItem):
+                items.append(ReasoningItem(raw_item=output, agent=agent))
+            elif isinstance(output, ResponseComputerToolCall):
+                items.append(ToolCallItem(raw_item=output, agent=agent))
+                tools_used.append("computer_use")
+                if not computer_tool:
+                    _error_tracing.attach_error_to_current_span(
+                        SpanError(
+                            message="Computer tool not found",
+                            data={},
+                        )
+                    )
+                    raise ModelBehaviorError(
+                        "Model produced computer action without a computer tool."
+                    )
+                computer_actions.append(
+                    ToolRunComputerAction(tool_call=output, computer_tool=computer_tool)
+                )
+            elif isinstance(output, McpApprovalRequest):
+                items.append(MCPApprovalRequestItem(raw_item=output, agent=agent))
+                if output.server_label not in hosted_mcp_server_map:
+                    _error_tracing.attach_error_to_current_span(
+                        SpanError(
+                            message="MCP server label not found",
+                            data={"server_label": output.server_label},
+                        )
+                    )
+                    raise ModelBehaviorError(f"MCP server label {output.server_label} not found")
+                else:
+                    server = hosted_mcp_server_map[output.server_label]
+                    if server.on_approval_request:
+                        mcp_approval_requests.append(
+                            ToolRunMCPApprovalRequest(
+                                request_item=output,
+                                mcp_tool=server,
+                            )
+                        )
+                    else:
+                        logger.warning(
+                            f"MCP server {output.server_label} has no on_approval_request hook"
+                        )
+            elif isinstance(output, McpListTools):
+                items.append(MCPListToolsItem(raw_item=output, agent=agent))
+            elif isinstance(output, McpCall):
+                items.append(ToolCallItem(raw_item=output, agent=agent))
+                tools_used.append("mcp")
+            elif isinstance(output, ImageGenerationCall):
+                items.append(ToolCallItem(raw_item=output, agent=agent))
+                tools_used.append("image_generation")
+            elif isinstance(output, ResponseCodeInterpreterToolCall):
+                items.append(ToolCallItem(raw_item=output, agent=agent))
+                tools_used.append("code_interpreter")
+            elif isinstance(output, LocalShellCall):
+                items.append(ToolCallItem(raw_item=output, agent=agent))
+                tools_used.append("local_shell")
+                if not local_shell_tool:
+                    _error_tracing.attach_error_to_current_span(
+                        SpanError(
+                            message="Local shell tool not found",
+                            data={},
+                        )
+                    )
+                    raise ModelBehaviorError(
+                        "Model produced local shell call without a local shell tool."
+                    )
+                local_shell_calls.append(
+                    ToolRunLocalShellCall(tool_call=output, local_shell_tool=local_shell_tool)
+                )
+
+            elif not isinstance(output, ResponseFunctionToolCall):
+                logger.warning(f"Unexpected output type, ignoring: {type(output)}")
+                continue
+
+            # At this point we know it's a function tool call
+            if not isinstance(output, ResponseFunctionToolCall):
+                continue
+
+            tools_used.append(output.name)
+
+            # Handoffs
+            if output.name in handoff_map:
+                items.append(HandoffCallItem(raw_item=output, agent=agent))
+                handoff = ToolRunHandoff(
+                    tool_call=output,
+                    handoff=handoff_map[output.name],
+                )
+                run_handoffs.append(handoff)
+            # Regular function tool call
+            else:
+                if output.name not in function_map:
+                    if output_schema is not None and output.name == "json_tool_call":
+                        # LiteLLM could generate non-existent tool calls for structured outputs
+                        items.append(ToolCallItem(raw_item=output, agent=agent))
+                        functions.append(
+                            ToolRunFunction(
+                                tool_call=output,
+                                # this tool does not exist in function_map, so generate ad-hoc one,
+                                # which just parses the input if it's a string, and returns the
+                                # value otherwise
+                                function_tool=_build_litellm_json_tool_call(output),
+                            )
+                        )
+                        continue
+                    else:
+                        _error_tracing.attach_error_to_current_span(
+                            SpanError(
+                                message="Tool not found",
+                                data={"tool_name": output.name},
+                            )
+                        )
+                        error = f"Tool {output.name} not found in agent {agent.name}"
+                        raise ModelBehaviorError(error)
+
+                items.append(ToolCallItem(raw_item=output, agent=agent))
+                functions.append(
+                    ToolRunFunction(
+                        tool_call=output,
+                        function_tool=function_map[output.name],
+                    )
+                )
+
+        return ProcessedResponse(
+            new_items=items,
+            handoffs=run_handoffs,
+            functions=functions,
+            computer_actions=computer_actions,
+            local_shell_calls=local_shell_calls,
+            tools_used=tools_used,
+            mcp_approval_requests=mcp_approval_requests,
+        )
+
+    @classmethod
+    async def _execute_input_guardrails(
+        cls,
+        *,
+        func_tool: FunctionTool,
+        tool_context: ToolContext[TContext],
+        agent: Agent[TContext],
+        tool_input_guardrail_results: list[ToolInputGuardrailResult],
+    ) -> str | None:
+        """Execute input guardrails for a tool.
+
+        Args:
+            func_tool: The function tool being executed.
+            tool_context: The tool execution context.
+            agent: The agent executing the tool.
+            tool_input_guardrail_results: List to append guardrail results to.
+
+        Returns:
+            None if tool execution should proceed, or a message string if execution should be
+            skipped.
+
+        Raises:
+            ToolInputGuardrailTripwireTriggered: If a guardrail triggers an exception.
+        """
+        if not func_tool.tool_input_guardrails:
+            return None
+
+        for guardrail in func_tool.tool_input_guardrails:
+            gr_out = await guardrail.run(
+                ToolInputGuardrailData(
+                    context=tool_context,
+                    agent=agent,
+                )
+            )
+
+            # Store the guardrail result
+            tool_input_guardrail_results.append(
+                ToolInputGuardrailResult(
+                    guardrail=guardrail,
+                    output=gr_out,
+                )
+            )
+
+            # Handle different behavior types
+            if gr_out.behavior["type"] == "raise_exception":
+                raise ToolInputGuardrailTripwireTriggered(guardrail=guardrail, output=gr_out)
+            elif gr_out.behavior["type"] == "reject_content":
+                # Set final_result to the message and skip tool execution
+                return gr_out.behavior["message"]
+            elif gr_out.behavior["type"] == "allow":
+                # Continue to next guardrail or tool execution
+                continue
+
+        return None
+
+    @classmethod
+    async def _execute_output_guardrails(
+        cls,
+        *,
+        func_tool: FunctionTool,
+        tool_context: ToolContext[TContext],
+        agent: Agent[TContext],
+        real_result: Any,
+        tool_output_guardrail_results: list[ToolOutputGuardrailResult],
+    ) -> Any:
+        """Execute output guardrails for a tool.
+
+        Args:
+            func_tool: The function tool being executed.
+            tool_context: The tool execution context.
+            agent: The agent executing the tool.
+            real_result: The actual result from the tool execution.
+            tool_output_guardrail_results: List to append guardrail results to.
+
+        Returns:
+            The final result after guardrail processing (may be modified).
+
+        Raises:
+            ToolOutputGuardrailTripwireTriggered: If a guardrail triggers an exception.
+        """
+        if not func_tool.tool_output_guardrails:
+            return real_result
+
+        final_result = real_result
+        for output_guardrail in func_tool.tool_output_guardrails:
+            gr_out = await output_guardrail.run(
+                ToolOutputGuardrailData(
+                    context=tool_context,
+                    agent=agent,
+                    output=real_result,
+                )
+            )
+
+            # Store the guardrail result
+            tool_output_guardrail_results.append(
+                ToolOutputGuardrailResult(
+                    guardrail=output_guardrail,
+                    output=gr_out,
+                )
+            )
+
+            # Handle different behavior types
+            if gr_out.behavior["type"] == "raise_exception":
+                raise ToolOutputGuardrailTripwireTriggered(
+                    guardrail=output_guardrail, output=gr_out
+                )
+            elif gr_out.behavior["type"] == "reject_content":
+                # Override the result with the guardrail message
+                final_result = gr_out.behavior["message"]
+                break
+            elif gr_out.behavior["type"] == "allow":
+                # Continue to next guardrail
+                continue
+
+        return final_result
+
+    @classmethod
+    async def _execute_tool_with_hooks(
+        cls,
+        *,
+        func_tool: FunctionTool,
+        tool_context: ToolContext[TContext],
+        agent: Agent[TContext],
+        hooks: RunHooks[TContext],
+        tool_call: ResponseFunctionToolCall,
+    ) -> Any:
+        """Execute the core tool function with before/after hooks.
+
+        Args:
+            func_tool: The function tool being executed.
+            tool_context: The tool execution context.
+            agent: The agent executing the tool.
+            hooks: The run hooks to execute.
+            tool_call: The tool call details.
+
+        Returns:
+            The result from the tool execution.
+        """
+        await asyncio.gather(
+            hooks.on_tool_start(tool_context, agent, func_tool),
+            (
+                agent.hooks.on_tool_start(tool_context, agent, func_tool)
+                if agent.hooks
+                else _coro.noop_coroutine()
+            ),
+        )
+
+        return await func_tool.on_invoke_tool(tool_context, tool_call.arguments)
+
+    @classmethod
+    async def execute_function_tool_calls(
+        cls,
+        *,
+        agent: Agent[TContext],
+        tool_runs: list[ToolRunFunction],
+        hooks: RunHooks[TContext],
+        context_wrapper: RunContextWrapper[TContext],
+        config: RunConfig,
+    ) -> tuple[
+        list[FunctionToolResult], list[ToolInputGuardrailResult], list[ToolOutputGuardrailResult]
+    ]:
+        # Collect guardrail results
+        tool_input_guardrail_results: list[ToolInputGuardrailResult] = []
+        tool_output_guardrail_results: list[ToolOutputGuardrailResult] = []
+
+        async def run_single_tool(
+            func_tool: FunctionTool, tool_call: ResponseFunctionToolCall
+        ) -> Any:
+            with function_span(func_tool.name) as span_fn:
+                tool_context = ToolContext.from_agent_context(
+                    context_wrapper,
+                    tool_call.call_id,
+                    tool_call=tool_call,
+                )
+                if config.trace_include_sensitive_data:
+                    span_fn.span_data.input = tool_call.arguments
+                try:
+                    # 1) Run input tool guardrails, if any
+                    rejected_message = await cls._execute_input_guardrails(
+                        func_tool=func_tool,
+                        tool_context=tool_context,
+                        agent=agent,
+                        tool_input_guardrail_results=tool_input_guardrail_results,
+                    )
+
+                    if rejected_message is not None:
+                        # Input guardrail rejected the tool call
+                        final_result = rejected_message
+                    else:
+                        # 2) Actually run the tool
+                        real_result = await cls._execute_tool_with_hooks(
+                            func_tool=func_tool,
+                            tool_context=tool_context,
+                            agent=agent,
+                            hooks=hooks,
+                            tool_call=tool_call,
+                        )
+
+                        # 3) Run output tool guardrails, if any
+                        final_result = await cls._execute_output_guardrails(
+                            func_tool=func_tool,
+                            tool_context=tool_context,
+                            agent=agent,
+                            real_result=real_result,
+                            tool_output_guardrail_results=tool_output_guardrail_results,
+                        )
+
+                        # 4) Tool end hooks (with final result, which may have been overridden)
+                        await asyncio.gather(
+                            hooks.on_tool_end(tool_context, agent, func_tool, final_result),
+                            (
+                                agent.hooks.on_tool_end(
+                                    tool_context, agent, func_tool, final_result
+                                )
+                                if agent.hooks
+                                else _coro.noop_coroutine()
+                            ),
+                        )
+                    result = final_result
+                except Exception as e:
+                    _error_tracing.attach_error_to_current_span(
+                        SpanError(
+                            message="Error running tool",
+                            data={"tool_name": func_tool.name, "error": str(e)},
+                        )
+                    )
+                    if isinstance(e, AgentsException):
+                        raise e
+                    raise UserError(f"Error running tool {func_tool.name}: {e}") from e
+
+                if config.trace_include_sensitive_data:
+                    span_fn.span_data.output = result
+            return result
+
+        tasks = []
+        for tool_run in tool_runs:
+            function_tool = tool_run.function_tool
+            tasks.append(run_single_tool(function_tool, tool_run.tool_call))
+
+        results = await asyncio.gather(*tasks)
+
+        function_tool_results = [
+            FunctionToolResult(
+                tool=tool_run.function_tool,
+                output=result,
+                run_item=ToolCallOutputItem(
+                    output=result,
+                    raw_item=ItemHelpers.tool_call_output_item(tool_run.tool_call, str(result)),
+                    agent=agent,
+                ),
+            )
+            for tool_run, result in zip(tool_runs, results)
+        ]
+
+        return function_tool_results, tool_input_guardrail_results, tool_output_guardrail_results
+
+    @classmethod
+    async def execute_local_shell_calls(
+        cls,
+        *,
+        agent: Agent[TContext],
+        calls: list[ToolRunLocalShellCall],
+        context_wrapper: RunContextWrapper[TContext],
+        hooks: RunHooks[TContext],
+        config: RunConfig,
+    ) -> list[RunItem]:
+        results: list[RunItem] = []
+        # Need to run these serially, because each call can affect the local shell state
+        for call in calls:
+            results.append(
+                await LocalShellAction.execute(
+                    agent=agent,
+                    call=call,
+                    hooks=hooks,
+                    context_wrapper=context_wrapper,
+                    config=config,
+                )
+            )
+        return results
+
+    @classmethod
+    async def execute_computer_actions(
+        cls,
+        *,
+        agent: Agent[TContext],
+        actions: list[ToolRunComputerAction],
+        hooks: RunHooks[TContext],
+        context_wrapper: RunContextWrapper[TContext],
+        config: RunConfig,
+    ) -> list[RunItem]:
+        results: list[RunItem] = []
+        # Need to run these serially, because each action can affect the computer state
+        for action in actions:
+            acknowledged: list[ComputerCallOutputAcknowledgedSafetyCheck] | None = None
+            if action.tool_call.pending_safety_checks and action.computer_tool.on_safety_check:
+                acknowledged = []
+                for check in action.tool_call.pending_safety_checks:
+                    data = ComputerToolSafetyCheckData(
+                        ctx_wrapper=context_wrapper,
+                        agent=agent,
+                        tool_call=action.tool_call,
+                        safety_check=check,
+                    )
+                    maybe = action.computer_tool.on_safety_check(data)
+                    ack = await maybe if inspect.isawaitable(maybe) else maybe
+                    if ack:
+                        acknowledged.append(
+                            ComputerCallOutputAcknowledgedSafetyCheck(
+                                id=check.id,
+                                code=check.code,
+                                message=check.message,
+                            )
+                        )
+                    else:
+                        raise UserError("Computer tool safety check was not acknowledged")
+
+            results.append(
+                await ComputerAction.execute(
+                    agent=agent,
+                    action=action,
+                    hooks=hooks,
+                    context_wrapper=context_wrapper,
+                    config=config,
+                    acknowledged_safety_checks=acknowledged,
+                )
+            )
+
+        return results
+
+    @classmethod
+    async def execute_handoffs(
+        cls,
+        *,
+        agent: Agent[TContext],
+        original_input: str | list[TResponseInputItem],
+        pre_step_items: list[RunItem],
+        new_step_items: list[RunItem],
+        new_response: ModelResponse,
+        run_handoffs: list[ToolRunHandoff],
+        hooks: RunHooks[TContext],
+        context_wrapper: RunContextWrapper[TContext],
+        run_config: RunConfig,
+    ) -> SingleStepResult:
+        # If there is more than one handoff, add tool responses that reject those handoffs
+        multiple_handoffs = len(run_handoffs) > 1
+        if multiple_handoffs:
+            output_message = "Multiple handoffs detected, ignoring this one."
+            new_step_items.extend(
+                [
+                    ToolCallOutputItem(
+                        output=output_message,
+                        raw_item=ItemHelpers.tool_call_output_item(
+                            handoff.tool_call, output_message
+                        ),
+                        agent=agent,
+                    )
+                    for handoff in run_handoffs[1:]
+                ]
+            )
+
+        actual_handoff = run_handoffs[0]
+        with handoff_span(from_agent=agent.name) as span_handoff:
+            handoff = actual_handoff.handoff
+            new_agent: Agent[Any] = await handoff.on_invoke_handoff(
+                context_wrapper, actual_handoff.tool_call.arguments
+            )
+            span_handoff.span_data.to_agent = new_agent.name
+            if multiple_handoffs:
+                requested_agents = [handoff.handoff.agent_name for handoff in run_handoffs]
+                span_handoff.set_error(
+                    SpanError(
+                        message="Multiple handoffs requested",
+                        data={
+                            "requested_agents": requested_agents,
+                        },
+                    )
+                )
+
+            # Append a tool output item for the handoff
+            new_step_items.append(
+                HandoffOutputItem(
+                    agent=agent,
+                    raw_item=ItemHelpers.tool_call_output_item(
+                        actual_handoff.tool_call,
+                        handoff.get_transfer_message(new_agent),
+                    ),
+                    source_agent=agent,
+                    target_agent=new_agent,
+                )
+            )
+
+            # Execute handoff hooks
+            await asyncio.gather(
+                hooks.on_handoff(
+                    context=context_wrapper,
+                    from_agent=agent,
+                    to_agent=new_agent,
+                ),
+                (
+                    agent.hooks.on_handoff(
+                        context_wrapper,
+                        agent=new_agent,
+                        source=agent,
+                    )
+                    if agent.hooks
+                    else _coro.noop_coroutine()
+                ),
+            )
+
+            # If there's an input filter, filter the input for the next agent
+            input_filter = handoff.input_filter or (
+                run_config.handoff_input_filter if run_config else None
+            )
+            if input_filter:
+                logger.debug("Filtering inputs for handoff")
+                handoff_input_data = HandoffInputData(
+                    input_history=tuple(original_input)
+                    if isinstance(original_input, list)
+                    else original_input,
+                    pre_handoff_items=tuple(pre_step_items),
+                    new_items=tuple(new_step_items),
+                    run_context=context_wrapper,
+                )
+                if not callable(input_filter):
+                    _error_tracing.attach_error_to_span(
+                        span_handoff,
+                        SpanError(
+                            message="Invalid input filter",
+                            data={"details": "not callable()"},
+                        ),
+                    )
+                    raise UserError(f"Invalid input filter: {input_filter}")
+                filtered = input_filter(handoff_input_data)
+                if inspect.isawaitable(filtered):
+                    filtered = await filtered
+                if not isinstance(filtered, HandoffInputData):
+                    _error_tracing.attach_error_to_span(
+                        span_handoff,
+                        SpanError(
+                            message="Invalid input filter result",
+                            data={"details": "not a HandoffInputData"},
+                        ),
+                    )
+                    raise UserError(f"Invalid input filter result: {filtered}")
+
+                original_input = (
+                    filtered.input_history
+                    if isinstance(filtered.input_history, str)
+                    else list(filtered.input_history)
+                )
+                pre_step_items = list(filtered.pre_handoff_items)
+                new_step_items = list(filtered.new_items)
+
+        return SingleStepResult(
+            original_input=original_input,
+            model_response=new_response,
+            pre_step_items=pre_step_items,
+            new_step_items=new_step_items,
+            next_step=NextStepHandoff(new_agent),
+            tool_input_guardrail_results=[],
+            tool_output_guardrail_results=[],
+        )
+
+    @classmethod
+    async def execute_mcp_approval_requests(
+        cls,
+        *,
+        agent: Agent[TContext],
+        approval_requests: list[ToolRunMCPApprovalRequest],
+        context_wrapper: RunContextWrapper[TContext],
+    ) -> list[RunItem]:
+        async def run_single_approval(approval_request: ToolRunMCPApprovalRequest) -> RunItem:
+            callback = approval_request.mcp_tool.on_approval_request
+            assert callback is not None, "Callback is required for MCP approval requests"
+            maybe_awaitable_result = callback(
+                MCPToolApprovalRequest(context_wrapper, approval_request.request_item)
+            )
+            if inspect.isawaitable(maybe_awaitable_result):
+                result = await maybe_awaitable_result
+            else:
+                result = maybe_awaitable_result
+            reason = result.get("reason", None)
+            raw_item: McpApprovalResponse = {
+                "approval_request_id": approval_request.request_item.id,
+                "approve": result["approve"],
+                "type": "mcp_approval_response",
+            }
+            if not result["approve"] and reason:
+                raw_item["reason"] = reason
+            return MCPApprovalResponseItem(
+                raw_item=raw_item,
+                agent=agent,
+            )
+
+        tasks = [run_single_approval(approval_request) for approval_request in approval_requests]
+        return await asyncio.gather(*tasks)
+
+    @classmethod
+    async def execute_final_output(
+        cls,
+        *,
+        agent: Agent[TContext],
+        original_input: str | list[TResponseInputItem],
+        new_response: ModelResponse,
+        pre_step_items: list[RunItem],
+        new_step_items: list[RunItem],
+        final_output: Any,
+        hooks: RunHooks[TContext],
+        context_wrapper: RunContextWrapper[TContext],
+        tool_input_guardrail_results: list[ToolInputGuardrailResult],
+        tool_output_guardrail_results: list[ToolOutputGuardrailResult],
+    ) -> SingleStepResult:
+        # Run the on_end hooks
+        await cls.run_final_output_hooks(agent, hooks, context_wrapper, final_output)
+
+        return SingleStepResult(
+            original_input=original_input,
+            model_response=new_response,
+            pre_step_items=pre_step_items,
+            new_step_items=new_step_items,
+            next_step=NextStepFinalOutput(final_output),
+            tool_input_guardrail_results=tool_input_guardrail_results,
+            tool_output_guardrail_results=tool_output_guardrail_results,
+        )
+
+    @classmethod
+    async def run_final_output_hooks(
+        cls,
+        agent: Agent[TContext],
+        hooks: RunHooks[TContext],
+        context_wrapper: RunContextWrapper[TContext],
+        final_output: Any,
+    ):
+        await asyncio.gather(
+            hooks.on_agent_end(context_wrapper, agent, final_output),
+            agent.hooks.on_end(context_wrapper, agent, final_output)
+            if agent.hooks
+            else _coro.noop_coroutine(),
+        )
+
+    @classmethod
+    async def run_single_input_guardrail(
+        cls,
+        agent: Agent[Any],
+        guardrail: InputGuardrail[TContext],
+        input: str | list[TResponseInputItem],
+        context: RunContextWrapper[TContext],
+    ) -> InputGuardrailResult:
+        with guardrail_span(guardrail.get_name()) as span_guardrail:
+            result = await guardrail.run(agent, input, context)
+            span_guardrail.span_data.triggered = result.output.tripwire_triggered
+            return result
+
+    @classmethod
+    async def run_single_output_guardrail(
+        cls,
+        guardrail: OutputGuardrail[TContext],
+        agent: Agent[Any],
+        agent_output: Any,
+        context: RunContextWrapper[TContext],
+    ) -> OutputGuardrailResult:
+        with guardrail_span(guardrail.get_name()) as span_guardrail:
+            result = await guardrail.run(agent=agent, agent_output=agent_output, context=context)
+            span_guardrail.span_data.triggered = result.output.tripwire_triggered
+            return result
+
+    @classmethod
+    def stream_step_items_to_queue(
+        cls,
+        new_step_items: list[RunItem],
+        queue: asyncio.Queue[StreamEvent | QueueCompleteSentinel],
+    ):
+        for item in new_step_items:
+            if isinstance(item, MessageOutputItem):
+                event = RunItemStreamEvent(item=item, name="message_output_created")
+            elif isinstance(item, HandoffCallItem):
+                event = RunItemStreamEvent(item=item, name="handoff_requested")
+            elif isinstance(item, HandoffOutputItem):
+                event = RunItemStreamEvent(item=item, name="handoff_occured")
+            elif isinstance(item, ToolCallItem):
+                event = RunItemStreamEvent(item=item, name="tool_called")
+            elif isinstance(item, ToolCallOutputItem):
+                event = RunItemStreamEvent(item=item, name="tool_output")
+            elif isinstance(item, ReasoningItem):
+                event = RunItemStreamEvent(item=item, name="reasoning_item_created")
+            elif isinstance(item, MCPApprovalRequestItem):
+                event = RunItemStreamEvent(item=item, name="mcp_approval_requested")
+            elif isinstance(item, MCPListToolsItem):
+                event = RunItemStreamEvent(item=item, name="mcp_list_tools")
+
+            else:
+                logger.warning(f"Unexpected item type: {type(item)}")
+                event = None
+
+            if event:
+                queue.put_nowait(event)
+
+    @classmethod
+    def stream_step_result_to_queue(
+        cls,
+        step_result: SingleStepResult,
+        queue: asyncio.Queue[StreamEvent | QueueCompleteSentinel],
+    ):
+        cls.stream_step_items_to_queue(step_result.new_step_items, queue)
+
+    @classmethod
+    async def _check_for_final_output_from_tools(
+        cls,
+        *,
+        agent: Agent[TContext],
+        tool_results: list[FunctionToolResult],
+        context_wrapper: RunContextWrapper[TContext],
+        config: RunConfig,
+    ) -> ToolsToFinalOutputResult:
+        """Determine if tool results should produce a final output.
+        Returns:
+            ToolsToFinalOutputResult: Indicates whether final output is ready, and the output value.
+        """
+        if not tool_results:
+            return _NOT_FINAL_OUTPUT
+
+        if agent.tool_use_behavior == "run_llm_again":
+            return _NOT_FINAL_OUTPUT
+        elif agent.tool_use_behavior == "stop_on_first_tool":
+            return ToolsToFinalOutputResult(
+                is_final_output=True, final_output=tool_results[0].output
+            )
+        elif isinstance(agent.tool_use_behavior, dict):
+            names = agent.tool_use_behavior.get("stop_at_tool_names", [])
+            for tool_result in tool_results:
+                if tool_result.tool.name in names:
+                    return ToolsToFinalOutputResult(
+                        is_final_output=True, final_output=tool_result.output
+                    )
+            return ToolsToFinalOutputResult(is_final_output=False, final_output=None)
+        elif callable(agent.tool_use_behavior):
+            if inspect.iscoroutinefunction(agent.tool_use_behavior):
+                return await cast(
+                    Awaitable[ToolsToFinalOutputResult],
+                    agent.tool_use_behavior(context_wrapper, tool_results),
+                )
+            else:
+                return cast(
+                    ToolsToFinalOutputResult, agent.tool_use_behavior(context_wrapper, tool_results)
+                )
+
+        logger.error(f"Invalid tool_use_behavior: {agent.tool_use_behavior}")
+        raise UserError(f"Invalid tool_use_behavior: {agent.tool_use_behavior}")
+
+
+class TraceCtxManager:
+    """Creates a trace only if there is no current trace, and manages the trace lifecycle."""
+
+    def __init__(
+        self,
+        workflow_name: str,
+        trace_id: str | None,
+        group_id: str | None,
+        metadata: dict[str, Any] | None,
+        disabled: bool,
+    ):
+        self.trace: Trace | None = None
+        self.workflow_name = workflow_name
+        self.trace_id = trace_id
+        self.group_id = group_id
+        self.metadata = metadata
+        self.disabled = disabled
+
+    def __enter__(self) -> TraceCtxManager:
+        current_trace = get_current_trace()
+        if not current_trace:
+            self.trace = trace(
+                workflow_name=self.workflow_name,
+                trace_id=self.trace_id,
+                group_id=self.group_id,
+                metadata=self.metadata,
+                disabled=self.disabled,
+            )
+            self.trace.start(mark_as_current=True)
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.trace:
+            self.trace.finish(reset_current=True)
+
+
+class ComputerAction:
+    @classmethod
+    async def execute(
+        cls,
+        *,
+        agent: Agent[TContext],
+        action: ToolRunComputerAction,
+        hooks: RunHooks[TContext],
+        context_wrapper: RunContextWrapper[TContext],
+        config: RunConfig,
+        acknowledged_safety_checks: list[ComputerCallOutputAcknowledgedSafetyCheck] | None = None,
+    ) -> RunItem:
+        output_func = (
+            cls._get_screenshot_async(action.computer_tool.computer, action.tool_call)
+            if isinstance(action.computer_tool.computer, AsyncComputer)
+            else cls._get_screenshot_sync(action.computer_tool.computer, action.tool_call)
+        )
+
+        _, _, output = await asyncio.gather(
+            hooks.on_tool_start(context_wrapper, agent, action.computer_tool),
+            (
+                agent.hooks.on_tool_start(context_wrapper, agent, action.computer_tool)
+                if agent.hooks
+                else _coro.noop_coroutine()
+            ),
+            output_func,
+        )
+
+        await asyncio.gather(
+            hooks.on_tool_end(context_wrapper, agent, action.computer_tool, output),
+            (
+                agent.hooks.on_tool_end(context_wrapper, agent, action.computer_tool, output)
+                if agent.hooks
+                else _coro.noop_coroutine()
+            ),
+        )
+
+        # TODO: don't send a screenshot every single time, use references
+        image_url = f"data:image/png;base64,{output}"
+        return ToolCallOutputItem(
+            agent=agent,
+            output=image_url,
+            raw_item=ComputerCallOutput(
+                call_id=action.tool_call.call_id,
+                output={
+                    "type": "computer_screenshot",
+                    "image_url": image_url,
+                },
+                type="computer_call_output",
+                acknowledged_safety_checks=acknowledged_safety_checks,
+            ),
+        )
+
+    @classmethod
+    async def _get_screenshot_sync(
+        cls,
+        computer: Computer,
+        tool_call: ResponseComputerToolCall,
+    ) -> str:
+        action = tool_call.action
+        if isinstance(action, ActionClick):
+            computer.click(action.x, action.y, action.button)
+        elif isinstance(action, ActionDoubleClick):
+            computer.double_click(action.x, action.y)
+        elif isinstance(action, ActionDrag):
+            computer.drag([(p.x, p.y) for p in action.path])
+        elif isinstance(action, ActionKeypress):
+            computer.keypress(action.keys)
+        elif isinstance(action, ActionMove):
+            computer.move(action.x, action.y)
+        elif isinstance(action, ActionScreenshot):
+            computer.screenshot()
+        elif isinstance(action, ActionScroll):
+            computer.scroll(action.x, action.y, action.scroll_x, action.scroll_y)
+        elif isinstance(action, ActionType):
+            computer.type(action.text)
+        elif isinstance(action, ActionWait):
+            computer.wait()
+
+        return computer.screenshot()
+
+    @classmethod
+    async def _get_screenshot_async(
+        cls,
+        computer: AsyncComputer,
+        tool_call: ResponseComputerToolCall,
+    ) -> str:
+        action = tool_call.action
+        if isinstance(action, ActionClick):
+            await computer.click(action.x, action.y, action.button)
+        elif isinstance(action, ActionDoubleClick):
+            await computer.double_click(action.x, action.y)
+        elif isinstance(action, ActionDrag):
+            await computer.drag([(p.x, p.y) for p in action.path])
+        elif isinstance(action, ActionKeypress):
+            await computer.keypress(action.keys)
+        elif isinstance(action, ActionMove):
+            await computer.move(action.x, action.y)
+        elif isinstance(action, ActionScreenshot):
+            await computer.screenshot()
+        elif isinstance(action, ActionScroll):
+            await computer.scroll(action.x, action.y, action.scroll_x, action.scroll_y)
+        elif isinstance(action, ActionType):
+            await computer.type(action.text)
+        elif isinstance(action, ActionWait):
+            await computer.wait()
+
+        return await computer.screenshot()
+
+
+class LocalShellAction:
+    @classmethod
+    async def execute(
+        cls,
+        *,
+        agent: Agent[TContext],
+        call: ToolRunLocalShellCall,
+        hooks: RunHooks[TContext],
+        context_wrapper: RunContextWrapper[TContext],
+        config: RunConfig,
+    ) -> RunItem:
+        await asyncio.gather(
+            hooks.on_tool_start(context_wrapper, agent, call.local_shell_tool),
+            (
+                agent.hooks.on_tool_start(context_wrapper, agent, call.local_shell_tool)
+                if agent.hooks
+                else _coro.noop_coroutine()
+            ),
+        )
+
+        request = LocalShellCommandRequest(
+            ctx_wrapper=context_wrapper,
+            data=call.tool_call,
+        )
+        output = call.local_shell_tool.executor(request)
+        if inspect.isawaitable(output):
+            result = await output
+        else:
+            result = output
+
+        await asyncio.gather(
+            hooks.on_tool_end(context_wrapper, agent, call.local_shell_tool, result),
+            (
+                agent.hooks.on_tool_end(context_wrapper, agent, call.local_shell_tool, result)
+                if agent.hooks
+                else _coro.noop_coroutine()
+            ),
+        )
+
+        return ToolCallOutputItem(
+            agent=agent,
+            output=output,
+            raw_item={
+                "type": "local_shell_call_output",
+                "id": call.tool_call.call_id,
+                "output": result,
+                # "id": "out" + call.tool_call.id,  # TODO remove this, it should be optional
+            },
+        )
+
+
+def _build_litellm_json_tool_call(output: ResponseFunctionToolCall) -> FunctionTool:
+    async def on_invoke_tool(_ctx: ToolContext[Any], value: Any) -> Any:
+        if isinstance(value, str):
+            import json
+
+            return json.loads(value)
+        return value
+
+    return FunctionTool(
+        name=output.name,
+        description=output.name,
+        params_json_schema={},
+        on_invoke_tool=on_invoke_tool,
+        strict_json_schema=True,
+        is_enabled=True,
+    )

agents/agent.py

@@ -0,0 +1,476 @@
+from __future__ import annotations
+
+import asyncio
+import dataclasses
+import inspect
+from collections.abc import Awaitable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, cast
+
+from openai.types.responses.response_prompt_param import ResponsePromptParam
+from typing_extensions import NotRequired, TypeAlias, TypedDict
+
+from .agent_output import AgentOutputSchemaBase
+from .guardrail import InputGuardrail, OutputGuardrail
+from .handoffs import Handoff
+from .items import ItemHelpers
+from .logger import logger
+from .mcp import MCPUtil
+from .model_settings import ModelSettings
+from .models.default_models import (
+    get_default_model_settings,
+    gpt_5_reasoning_settings_required,
+    is_gpt_5_default,
+)
+from .models.interface import Model
+from .prompts import DynamicPromptFunction, Prompt, PromptUtil
+from .run_context import RunContextWrapper, TContext
+from .tool import FunctionTool, FunctionToolResult, Tool, function_tool
+from .util import _transforms
+from .util._types import MaybeAwaitable
+
+if TYPE_CHECKING:
+    from .lifecycle import AgentHooks, RunHooks
+    from .mcp import MCPServer
+    from .memory.session import Session
+    from .result import RunResult
+    from .run import RunConfig
+
+
+@dataclass
+class ToolsToFinalOutputResult:
+    is_final_output: bool
+    """Whether this is the final output. If False, the LLM will run again and receive the tool call
+    output.
+    """
+
+    final_output: Any | None = None
+    """The final output. Can be None if `is_final_output` is False, otherwise must match the
+    `output_type` of the agent.
+    """
+
+
+ToolsToFinalOutputFunction: TypeAlias = Callable[
+    [RunContextWrapper[TContext], list[FunctionToolResult]],
+    MaybeAwaitable[ToolsToFinalOutputResult],
+]
+"""A function that takes a run context and a list of tool results, and returns a
+`ToolsToFinalOutputResult`.
+"""
+
+
+class StopAtTools(TypedDict):
+    stop_at_tool_names: list[str]
+    """A list of tool names, any of which will stop the agent from running further."""
+
+
+class MCPConfig(TypedDict):
+    """Configuration for MCP servers."""
+
+    convert_schemas_to_strict: NotRequired[bool]
+    """If True, we will attempt to convert the MCP schemas to strict-mode schemas. This is a
+    best-effort conversion, so some schemas may not be convertible. Defaults to False.
+    """
+
+
+@dataclass
+class AgentBase(Generic[TContext]):
+    """Base class for `Agent` and `RealtimeAgent`."""
+
+    name: str
+    """The name of the agent."""
+
+    handoff_description: str | None = None
+    """A description of the agent. This is used when the agent is used as a handoff, so that an
+    LLM knows what it does and when to invoke it.
+    """
+
+    tools: list[Tool] = field(default_factory=list)
+    """A list of tools that the agent can use."""
+
+    mcp_servers: list[MCPServer] = field(default_factory=list)
+    """A list of [Model Context Protocol](https://modelcontextprotocol.io/) servers that
+    the agent can use. Every time the agent runs, it will include tools from these servers in the
+    list of available tools.
+
+    NOTE: You are expected to manage the lifecycle of these servers. Specifically, you must call
+    `server.connect()` before passing it to the agent, and `server.cleanup()` when the server is no
+    longer needed.
+    """
+
+    mcp_config: MCPConfig = field(default_factory=lambda: MCPConfig())
+    """Configuration for MCP servers."""
+
+    async def get_mcp_tools(self, run_context: RunContextWrapper[TContext]) -> list[Tool]:
+        """Fetches the available tools from the MCP servers."""
+        convert_schemas_to_strict = self.mcp_config.get("convert_schemas_to_strict", False)
+        return await MCPUtil.get_all_function_tools(
+            self.mcp_servers, convert_schemas_to_strict, run_context, self
+        )
+
+    async def get_all_tools(self, run_context: RunContextWrapper[TContext]) -> list[Tool]:
+        """All agent tools, including MCP tools and function tools."""
+        mcp_tools = await self.get_mcp_tools(run_context)
+
+        async def _check_tool_enabled(tool: Tool) -> bool:
+            if not isinstance(tool, FunctionTool):
+                return True
+
+            attr = tool.is_enabled
+            if isinstance(attr, bool):
+                return attr
+            res = attr(run_context, self)
+            if inspect.isawaitable(res):
+                return bool(await res)
+            return bool(res)
+
+        results = await asyncio.gather(*(_check_tool_enabled(t) for t in self.tools))
+        enabled: list[Tool] = [t for t, ok in zip(self.tools, results) if ok]
+        return [*mcp_tools, *enabled]
+
+
+@dataclass
+class Agent(AgentBase, Generic[TContext]):
+    """An agent is an AI model configured with instructions, tools, guardrails, handoffs and more.
+
+    We strongly recommend passing `instructions`, which is the "system prompt" for the agent. In
+    addition, you can pass `handoff_description`, which is a human-readable description of the
+    agent, used when the agent is used inside tools/handoffs.
+
+    Agents are generic on the context type. The context is a (mutable) object you create. It is
+    passed to tool functions, handoffs, guardrails, etc.
+
+    See `AgentBase` for base parameters that are shared with `RealtimeAgent`s.
+    """
+
+    instructions: (
+        str
+        | Callable[
+            [RunContextWrapper[TContext], Agent[TContext]],
+            MaybeAwaitable[str],
+        ]
+        | None
+    ) = None
+    """The instructions for the agent. Will be used as the "system prompt" when this agent is
+    invoked. Describes what the agent should do, and how it responds.
+
+    Can either be a string, or a function that dynamically generates instructions for the agent. If
+    you provide a function, it will be called with the context and the agent instance. It must
+    return a string.
+    """
+
+    prompt: Prompt | DynamicPromptFunction | None = None
+    """A prompt object (or a function that returns a Prompt). Prompts allow you to dynamically
+    configure the instructions, tools and other config for an agent outside of your code. Only
+    usable with OpenAI models, using the Responses API.
+    """
+
+    handoffs: list[Agent[Any] | Handoff[TContext, Any]] = field(default_factory=list)
+    """Handoffs are sub-agents that the agent can delegate to. You can provide a list of handoffs,
+    and the agent can choose to delegate to them if relevant. Allows for separation of concerns and
+    modularity.
+    """
+
+    model: str | Model | None = None
+    """The model implementation to use when invoking the LLM.
+
+    By default, if not set, the agent will use the default model configured in
+    `agents.models.get_default_model()` (currently "gpt-4.1").
+    """
+
+    model_settings: ModelSettings = field(default_factory=get_default_model_settings)
+    """Configures model-specific tuning parameters (e.g. temperature, top_p).
+    """
+
+    input_guardrails: list[InputGuardrail[TContext]] = field(default_factory=list)
+    """A list of checks that run in parallel to the agent's execution, before generating a
+    response. Runs only if the agent is the first agent in the chain.
+    """
+
+    output_guardrails: list[OutputGuardrail[TContext]] = field(default_factory=list)
+    """A list of checks that run on the final output of the agent, after generating a response.
+    Runs only if the agent produces a final output.
+    """
+
+    output_type: type[Any] | AgentOutputSchemaBase | None = None
+    """The type of the output object. If not provided, the output will be `str`. In most cases,
+    you should pass a regular Python type (e.g. a dataclass, Pydantic model, TypedDict, etc).
+    You can customize this in two ways:
+    1. If you want non-strict schemas, pass `AgentOutputSchema(MyClass, strict_json_schema=False)`.
+    2. If you want to use a custom JSON schema (i.e. without using the SDK's automatic schema)
+       creation, subclass and pass an `AgentOutputSchemaBase` subclass.
+    """
+
+    hooks: AgentHooks[TContext] | None = None
+    """A class that receives callbacks on various lifecycle events for this agent.
+    """
+
+    tool_use_behavior: (
+        Literal["run_llm_again", "stop_on_first_tool"] | StopAtTools | ToolsToFinalOutputFunction
+    ) = "run_llm_again"
+    """
+    This lets you configure how tool use is handled.
+    - "run_llm_again": The default behavior. Tools are run, and then the LLM receives the results
+        and gets to respond.
+    - "stop_on_first_tool": The output from the first tool call is treated as the final result.
+        In other words, it isn’t sent back to the LLM for further processing but is used directly
+        as the final output.
+    - A StopAtTools object: The agent will stop running if any of the tools listed in
+        `stop_at_tool_names` is called.
+        The final output will be the output of the first matching tool call.
+        The LLM does not process the result of the tool call.
+    - A function: If you pass a function, it will be called with the run context and the list of
+      tool results. It must return a `ToolsToFinalOutputResult`, which determines whether the tool
+      calls result in a final output.
+
+      NOTE: This configuration is specific to FunctionTools. Hosted tools, such as file search,
+      web search, etc. are always processed by the LLM.
+    """
+
+    reset_tool_choice: bool = True
+    """Whether to reset the tool choice to the default value after a tool has been called. Defaults
+    to True. This ensures that the agent doesn't enter an infinite loop of tool usage."""
+
+    def __post_init__(self):
+        from typing import get_origin
+
+        if not isinstance(self.name, str):
+            raise TypeError(f"Agent name must be a string, got {type(self.name).__name__}")
+
+        if self.handoff_description is not None and not isinstance(self.handoff_description, str):
+            raise TypeError(
+                f"Agent handoff_description must be a string or None, "
+                f"got {type(self.handoff_description).__name__}"
+            )
+
+        if not isinstance(self.tools, list):
+            raise TypeError(f"Agent tools must be a list, got {type(self.tools).__name__}")
+
+        if not isinstance(self.mcp_servers, list):
+            raise TypeError(
+                f"Agent mcp_servers must be a list, got {type(self.mcp_servers).__name__}"
+            )
+
+        if not isinstance(self.mcp_config, dict):
+            raise TypeError(
+                f"Agent mcp_config must be a dict, got {type(self.mcp_config).__name__}"
+            )
+
+        if (
+            self.instructions is not None
+            and not isinstance(self.instructions, str)
+            and not callable(self.instructions)
+        ):
+            raise TypeError(
+                f"Agent instructions must be a string, callable, or None, "
+                f"got {type(self.instructions).__name__}"
+            )
+
+        if (
+            self.prompt is not None
+            and not callable(self.prompt)
+            and not hasattr(self.prompt, "get")
+        ):
+            raise TypeError(
+                f"Agent prompt must be a Prompt, DynamicPromptFunction, or None, "
+                f"got {type(self.prompt).__name__}"
+            )
+
+        if not isinstance(self.handoffs, list):
+            raise TypeError(f"Agent handoffs must be a list, got {type(self.handoffs).__name__}")
+
+        if self.model is not None and not isinstance(self.model, str):
+            from .models.interface import Model
+
+            if not isinstance(self.model, Model):
+                raise TypeError(
+                    f"Agent model must be a string, Model, or None, got {type(self.model).__name__}"
+                )
+
+        if not isinstance(self.model_settings, ModelSettings):
+            raise TypeError(
+                f"Agent model_settings must be a ModelSettings instance, "
+                f"got {type(self.model_settings).__name__}"
+            )
+
+        if (
+            # The user sets a non-default model
+            self.model is not None
+            and (
+                # The default model is gpt-5
+                is_gpt_5_default() is True
+                # However, the specified model is not a gpt-5 model
+                and (
+                    isinstance(self.model, str) is False
+                    or gpt_5_reasoning_settings_required(self.model) is False  # type: ignore
+                )
+                # The model settings are not customized for the specified model
+                and self.model_settings == get_default_model_settings()
+            )
+        ):
+            # In this scenario, we should use a generic model settings
+            # because non-gpt-5 models are not compatible with the default gpt-5 model settings.
+            # This is a best-effort attempt to make the agent work with non-gpt-5 models.
+            self.model_settings = ModelSettings()
+
+        if not isinstance(self.input_guardrails, list):
+            raise TypeError(
+                f"Agent input_guardrails must be a list, got {type(self.input_guardrails).__name__}"
+            )
+
+        if not isinstance(self.output_guardrails, list):
+            raise TypeError(
+                f"Agent output_guardrails must be a list, "
+                f"got {type(self.output_guardrails).__name__}"
+            )
+
+        if self.output_type is not None:
+            from .agent_output import AgentOutputSchemaBase
+
+            if not (
+                isinstance(self.output_type, (type, AgentOutputSchemaBase))
+                or get_origin(self.output_type) is not None
+            ):
+                raise TypeError(
+                    f"Agent output_type must be a type, AgentOutputSchemaBase, or None, "
+                    f"got {type(self.output_type).__name__}"
+                )
+
+        if self.hooks is not None:
+            from .lifecycle import AgentHooksBase
+
+            if not isinstance(self.hooks, AgentHooksBase):
+                raise TypeError(
+                    f"Agent hooks must be an AgentHooks instance or None, "
+                    f"got {type(self.hooks).__name__}"
+                )
+
+        if (
+            not (
+                isinstance(self.tool_use_behavior, str)
+                and self.tool_use_behavior in ["run_llm_again", "stop_on_first_tool"]
+            )
+            and not isinstance(self.tool_use_behavior, dict)
+            and not callable(self.tool_use_behavior)
+        ):
+            raise TypeError(
+                f"Agent tool_use_behavior must be 'run_llm_again', 'stop_on_first_tool', "
+                f"StopAtTools dict, or callable, got {type(self.tool_use_behavior).__name__}"
+            )
+
+        if not isinstance(self.reset_tool_choice, bool):
+            raise TypeError(
+                f"Agent reset_tool_choice must be a boolean, "
+                f"got {type(self.reset_tool_choice).__name__}"
+            )
+
+    def clone(self, **kwargs: Any) -> Agent[TContext]:
+        """Make a copy of the agent, with the given arguments changed.
+        Notes:
+            - Uses `dataclasses.replace`, which performs a **shallow copy**.
+            - Mutable attributes like `tools` and `handoffs` are shallow-copied:
+              new list objects are created only if overridden, but their contents
+              (tool functions and handoff objects) are shared with the original.
+            - To modify these independently, pass new lists when calling `clone()`.
+        Example:
+            ```python
+            new_agent = agent.clone(instructions="New instructions")
+            ```
+        """
+        return dataclasses.replace(self, **kwargs)
+
+    def as_tool(
+        self,
+        tool_name: str | None,
+        tool_description: str | None,
+        custom_output_extractor: Callable[[RunResult], Awaitable[str]] | None = None,
+        is_enabled: bool
+        | Callable[[RunContextWrapper[Any], AgentBase[Any]], MaybeAwaitable[bool]] = True,
+        run_config: RunConfig | None = None,
+        max_turns: int | None = None,
+        hooks: RunHooks[TContext] | None = None,
+        previous_response_id: str | None = None,
+        conversation_id: str | None = None,
+        session: Session | None = None,
+    ) -> Tool:
+        """Transform this agent into a tool, callable by other agents.
+
+        This is different from handoffs in two ways:
+        1. In handoffs, the new agent receives the conversation history. In this tool, the new agent
+           receives generated input.
+        2. In handoffs, the new agent takes over the conversation. In this tool, the new agent is
+           called as a tool, and the conversation is continued by the original agent.
+
+        Args:
+            tool_name: The name of the tool. If not provided, the agent's name will be used.
+            tool_description: The description of the tool, which should indicate what it does and
+                when to use it.
+            custom_output_extractor: A function that extracts the output from the agent. If not
+                provided, the last message from the agent will be used.
+            is_enabled: Whether the tool is enabled. Can be a bool or a callable that takes the run
+                context and agent and returns whether the tool is enabled. Disabled tools are hidden
+                from the LLM at runtime.
+        """
+
+        @function_tool(
+            name_override=tool_name or _transforms.transform_string_function_style(self.name),
+            description_override=tool_description or "",
+            is_enabled=is_enabled,
+        )
+        async def run_agent(context: RunContextWrapper, input: str) -> str:
+            from .run import DEFAULT_MAX_TURNS, Runner
+
+            resolved_max_turns = max_turns if max_turns is not None else DEFAULT_MAX_TURNS
+
+            output = await Runner.run(
+                starting_agent=self,
+                input=input,
+                context=context.context,
+                run_config=run_config,
+                max_turns=resolved_max_turns,
+                hooks=hooks,
+                previous_response_id=previous_response_id,
+                conversation_id=conversation_id,
+                session=session,
+            )
+            if custom_output_extractor:
+                return await custom_output_extractor(output)
+
+            return ItemHelpers.text_message_outputs(output.new_items)
+
+        return run_agent
+
+    async def get_system_prompt(self, run_context: RunContextWrapper[TContext]) -> str | None:
+        if isinstance(self.instructions, str):
+            return self.instructions
+        elif callable(self.instructions):
+            # Inspect the signature of the instructions function
+            sig = inspect.signature(self.instructions)
+            params = list(sig.parameters.values())
+
+            # Enforce exactly 2 parameters
+            if len(params) != 2:
+                raise TypeError(
+                    f"'instructions' callable must accept exactly 2 arguments (context, agent), "
+                    f"but got {len(params)}: {[p.name for p in params]}"
+                )
+
+            # Call the instructions function properly
+            if inspect.iscoroutinefunction(self.instructions):
+                return await cast(Awaitable[str], self.instructions(run_context, self))
+            else:
+                return cast(str, self.instructions(run_context, self))
+
+        elif self.instructions is not None:
+            logger.error(
+                f"Instructions must be a string or a callable function, "
+                f"got {type(self.instructions).__name__}"
+            )
+
+        return None
+
+    async def get_prompt(
+        self, run_context: RunContextWrapper[TContext]
+    ) -> ResponsePromptParam | None:
+        """Get the prompt for the agent."""
+        return await PromptUtil.to_model_input(self.prompt, run_context, self)

agents/agent_output.py

@@ -0,0 +1,194 @@
+import abc
+from dataclasses import dataclass
+from typing import Any
+
+from pydantic import BaseModel, TypeAdapter
+from typing_extensions import TypedDict, get_args, get_origin
+
+from .exceptions import ModelBehaviorError, UserError
+from .strict_schema import ensure_strict_json_schema
+from .tracing import SpanError
+from .util import _error_tracing, _json
+
+_WRAPPER_DICT_KEY = "response"
+
+
+class AgentOutputSchemaBase(abc.ABC):
+    """An object that captures the JSON schema of the output, as well as validating/parsing JSON
+    produced by the LLM into the output type.
+    """
+
+    @abc.abstractmethod
+    def is_plain_text(self) -> bool:
+        """Whether the output type is plain text (versus a JSON object)."""
+        pass
+
+    @abc.abstractmethod
+    def name(self) -> str:
+        """The name of the output type."""
+        pass
+
+    @abc.abstractmethod
+    def json_schema(self) -> dict[str, Any]:
+        """Returns the JSON schema of the output. Will only be called if the output type is not
+        plain text.
+        """
+        pass
+
+    @abc.abstractmethod
+    def is_strict_json_schema(self) -> bool:
+        """Whether the JSON schema is in strict mode. Strict mode constrains the JSON schema
+        features, but guarantees valid JSON. See here for details:
+        https://platform.openai.com/docs/guides/structured-outputs#supported-schemas
+        """
+        pass
+
+    @abc.abstractmethod
+    def validate_json(self, json_str: str) -> Any:
+        """Validate a JSON string against the output type. You must return the validated object,
+        or raise a `ModelBehaviorError` if the JSON is invalid.
+        """
+        pass
+
+
+@dataclass(init=False)
+class AgentOutputSchema(AgentOutputSchemaBase):
+    """An object that captures the JSON schema of the output, as well as validating/parsing JSON
+    produced by the LLM into the output type.
+    """
+
+    output_type: type[Any]
+    """The type of the output."""
+
+    _type_adapter: TypeAdapter[Any]
+    """A type adapter that wraps the output type, so that we can validate JSON."""
+
+    _is_wrapped: bool
+    """Whether the output type is wrapped in a dictionary. This is generally done if the base
+    output type cannot be represented as a JSON Schema object.
+    """
+
+    _output_schema: dict[str, Any]
+    """The JSON schema of the output."""
+
+    _strict_json_schema: bool
+    """Whether the JSON schema is in strict mode. We **strongly** recommend setting this to True,
+    as it increases the likelihood of correct JSON input.
+    """
+
+    def __init__(self, output_type: type[Any], strict_json_schema: bool = True):
+        """
+        Args:
+            output_type: The type of the output.
+            strict_json_schema: Whether the JSON schema is in strict mode. We **strongly** recommend
+                setting this to True, as it increases the likelihood of correct JSON input.
+        """
+        self.output_type = output_type
+        self._strict_json_schema = strict_json_schema
+
+        if output_type is None or output_type is str:
+            self._is_wrapped = False
+            self._type_adapter = TypeAdapter(output_type)
+            self._output_schema = self._type_adapter.json_schema()
+            return
+
+        # We should wrap for things that are not plain text, and for things that would definitely
+        # not be a JSON Schema object.
+        self._is_wrapped = not _is_subclass_of_base_model_or_dict(output_type)
+
+        if self._is_wrapped:
+            OutputType = TypedDict(
+                "OutputType",
+                {
+                    _WRAPPER_DICT_KEY: output_type,  # type: ignore
+                },
+            )
+            self._type_adapter = TypeAdapter(OutputType)
+            self._output_schema = self._type_adapter.json_schema()
+        else:
+            self._type_adapter = TypeAdapter(output_type)
+            self._output_schema = self._type_adapter.json_schema()
+
+        if self._strict_json_schema:
+            try:
+                self._output_schema = ensure_strict_json_schema(self._output_schema)
+            except UserError as e:
+                raise UserError(
+                    "Strict JSON schema is enabled, but the output type is not valid. "
+                    "Either make the output type strict, "
+                    "or wrap your type with AgentOutputSchema(YourType, strict_json_schema=False)"
+                ) from e
+
+    def is_plain_text(self) -> bool:
+        """Whether the output type is plain text (versus a JSON object)."""
+        return self.output_type is None or self.output_type is str
+
+    def is_strict_json_schema(self) -> bool:
+        """Whether the JSON schema is in strict mode."""
+        return self._strict_json_schema
+
+    def json_schema(self) -> dict[str, Any]:
+        """The JSON schema of the output type."""
+        if self.is_plain_text():
+            raise UserError("Output type is plain text, so no JSON schema is available")
+        return self._output_schema
+
+    def validate_json(self, json_str: str) -> Any:
+        """Validate a JSON string against the output type. Returns the validated object, or raises
+        a `ModelBehaviorError` if the JSON is invalid.
+        """
+        validated = _json.validate_json(json_str, self._type_adapter, partial=False)
+        if self._is_wrapped:
+            if not isinstance(validated, dict):
+                _error_tracing.attach_error_to_current_span(
+                    SpanError(
+                        message="Invalid JSON",
+                        data={"details": f"Expected a dict, got {type(validated)}"},
+                    )
+                )
+                raise ModelBehaviorError(
+                    f"Expected a dict, got {type(validated)} for JSON: {json_str}"
+                )
+
+            if _WRAPPER_DICT_KEY not in validated:
+                _error_tracing.attach_error_to_current_span(
+                    SpanError(
+                        message="Invalid JSON",
+                        data={"details": f"Could not find key {_WRAPPER_DICT_KEY} in JSON"},
+                    )
+                )
+                raise ModelBehaviorError(
+                    f"Could not find key {_WRAPPER_DICT_KEY} in JSON: {json_str}"
+                )
+            return validated[_WRAPPER_DICT_KEY]
+        return validated
+
+    def name(self) -> str:
+        """The name of the output type."""
+        return _type_to_str(self.output_type)
+
+
+def _is_subclass_of_base_model_or_dict(t: Any) -> bool:
+    if not isinstance(t, type):
+        return False
+
+    # If it's a generic alias, 'origin' will be the actual type, e.g. 'list'
+    origin = get_origin(t)
+
+    allowed_types = (BaseModel, dict)
+    # If it's a generic alias e.g. list[str], then we should check the origin type i.e. list
+    return issubclass(origin or t, allowed_types)
+
+
+def _type_to_str(t: type[Any]) -> str:
+    origin = get_origin(t)
+    args = get_args(t)
+
+    if origin is None:
+        # It's a simple type like `str`, `int`, etc.
+        return t.__name__
+    elif args:
+        args_str = ", ".join(_type_to_str(arg) for arg in args)
+        return f"{origin.__name__}[{args_str}]"
+    else:
+        return str(t)

agents/computer.py

@@ -0,0 +1,107 @@
+import abc
+from typing import Literal
+
+Environment = Literal["mac", "windows", "ubuntu", "browser"]
+Button = Literal["left", "right", "wheel", "back", "forward"]
+
+
+class Computer(abc.ABC):
+    """A computer implemented with sync operations. The Computer interface abstracts the
+    operations needed to control a computer or browser."""
+
+    @property
+    @abc.abstractmethod
+    def environment(self) -> Environment:
+        pass
+
+    @property
+    @abc.abstractmethod
+    def dimensions(self) -> tuple[int, int]:
+        pass
+
+    @abc.abstractmethod
+    def screenshot(self) -> str:
+        pass
+
+    @abc.abstractmethod
+    def click(self, x: int, y: int, button: Button) -> None:
+        pass
+
+    @abc.abstractmethod
+    def double_click(self, x: int, y: int) -> None:
+        pass
+
+    @abc.abstractmethod
+    def scroll(self, x: int, y: int, scroll_x: int, scroll_y: int) -> None:
+        pass
+
+    @abc.abstractmethod
+    def type(self, text: str) -> None:
+        pass
+
+    @abc.abstractmethod
+    def wait(self) -> None:
+        pass
+
+    @abc.abstractmethod
+    def move(self, x: int, y: int) -> None:
+        pass
+
+    @abc.abstractmethod
+    def keypress(self, keys: list[str]) -> None:
+        pass
+
+    @abc.abstractmethod
+    def drag(self, path: list[tuple[int, int]]) -> None:
+        pass
+
+
+class AsyncComputer(abc.ABC):
+    """A computer implemented with async operations. The Computer interface abstracts the
+    operations needed to control a computer or browser."""
+
+    @property
+    @abc.abstractmethod
+    def environment(self) -> Environment:
+        pass
+
+    @property
+    @abc.abstractmethod
+    def dimensions(self) -> tuple[int, int]:
+        pass
+
+    @abc.abstractmethod
+    async def screenshot(self) -> str:
+        pass
+
+    @abc.abstractmethod
+    async def click(self, x: int, y: int, button: Button) -> None:
+        pass
+
+    @abc.abstractmethod
+    async def double_click(self, x: int, y: int) -> None:
+        pass
+
+    @abc.abstractmethod
+    async def scroll(self, x: int, y: int, scroll_x: int, scroll_y: int) -> None:
+        pass
+
+    @abc.abstractmethod
+    async def type(self, text: str) -> None:
+        pass
+
+    @abc.abstractmethod
+    async def wait(self) -> None:
+        pass
+
+    @abc.abstractmethod
+    async def move(self, x: int, y: int) -> None:
+        pass
+
+    @abc.abstractmethod
+    async def keypress(self, keys: list[str]) -> None:
+        pass
+
+    @abc.abstractmethod
+    async def drag(self, path: list[tuple[int, int]]) -> None:
+        pass

agents/exceptions.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from .agent import Agent
+    from .guardrail import InputGuardrailResult, OutputGuardrailResult
+    from .items import ModelResponse, RunItem, TResponseInputItem
+    from .run_context import RunContextWrapper
+    from .tool_guardrails import (
+        ToolGuardrailFunctionOutput,
+        ToolInputGuardrail,
+        ToolOutputGuardrail,
+    )
+
+from .util._pretty_print import pretty_print_run_error_details
+
+
+@dataclass
+class RunErrorDetails:
+    """Data collected from an agent run when an exception occurs."""
+
+    input: str | list[TResponseInputItem]
+    new_items: list[RunItem]
+    raw_responses: list[ModelResponse]
+    last_agent: Agent[Any]
+    context_wrapper: RunContextWrapper[Any]
+    input_guardrail_results: list[InputGuardrailResult]
+    output_guardrail_results: list[OutputGuardrailResult]
+
+    def __str__(self) -> str:
+        return pretty_print_run_error_details(self)
+
+
+class AgentsException(Exception):
+    """Base class for all exceptions in the Agents SDK."""
+
+    run_data: RunErrorDetails | None
+
+    def __init__(self, *args: object) -> None:
+        super().__init__(*args)
+        self.run_data = None
+
+
+class MaxTurnsExceeded(AgentsException):
+    """Exception raised when the maximum number of turns is exceeded."""
+
+    message: str
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(message)
+
+
+class ModelBehaviorError(AgentsException):
+    """Exception raised when the model does something unexpected, e.g. calling a tool that doesn't
+    exist, or providing malformed JSON.
+    """
+
+    message: str
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(message)
+
+
+class UserError(AgentsException):
+    """Exception raised when the user makes an error using the SDK."""
+
+    message: str
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(message)
+
+
+class InputGuardrailTripwireTriggered(AgentsException):
+    """Exception raised when a guardrail tripwire is triggered."""
+
+    guardrail_result: InputGuardrailResult
+    """The result data of the guardrail that was triggered."""
+
+    def __init__(self, guardrail_result: InputGuardrailResult):
+        self.guardrail_result = guardrail_result
+        super().__init__(
+            f"Guardrail {guardrail_result.guardrail.__class__.__name__} triggered tripwire"
+        )
+
+
+class OutputGuardrailTripwireTriggered(AgentsException):
+    """Exception raised when a guardrail tripwire is triggered."""
+
+    guardrail_result: OutputGuardrailResult
+    """The result data of the guardrail that was triggered."""
+
+    def __init__(self, guardrail_result: OutputGuardrailResult):
+        self.guardrail_result = guardrail_result
+        super().__init__(
+            f"Guardrail {guardrail_result.guardrail.__class__.__name__} triggered tripwire"
+        )
+
+
+class ToolInputGuardrailTripwireTriggered(AgentsException):
+    """Exception raised when a tool input guardrail tripwire is triggered."""
+
+    guardrail: ToolInputGuardrail[Any]
+    """The guardrail that was triggered."""
+
+    output: ToolGuardrailFunctionOutput
+    """The output from the guardrail function."""
+
+    def __init__(self, guardrail: ToolInputGuardrail[Any], output: ToolGuardrailFunctionOutput):
+        self.guardrail = guardrail
+        self.output = output
+        super().__init__(f"Tool input guardrail {guardrail.__class__.__name__} triggered tripwire")
+
+
+class ToolOutputGuardrailTripwireTriggered(AgentsException):
+    """Exception raised when a tool output guardrail tripwire is triggered."""
+
+    guardrail: ToolOutputGuardrail[Any]
+    """The guardrail that was triggered."""
+
+    output: ToolGuardrailFunctionOutput
+    """The output from the guardrail function."""
+
+    def __init__(self, guardrail: ToolOutputGuardrail[Any], output: ToolGuardrailFunctionOutput):
+        self.guardrail = guardrail
+        self.output = output
+        super().__init__(f"Tool output guardrail {guardrail.__class__.__name__} triggered tripwire")

agents/extensions/handoff_filters.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from ..handoffs import HandoffInputData
+from ..items import (
+    HandoffCallItem,
+    HandoffOutputItem,
+    ReasoningItem,
+    RunItem,
+    ToolCallItem,
+    ToolCallOutputItem,
+    TResponseInputItem,
+)
+
+"""Contains common handoff input filters, for convenience. """
+
+
+def remove_all_tools(handoff_input_data: HandoffInputData) -> HandoffInputData:
+    """Filters out all tool items: file search, web search and function calls+output."""
+
+    history = handoff_input_data.input_history
+    new_items = handoff_input_data.new_items
+
+    filtered_history = (
+        _remove_tool_types_from_input(history) if isinstance(history, tuple) else history
+    )
+    filtered_pre_handoff_items = _remove_tools_from_items(handoff_input_data.pre_handoff_items)
+    filtered_new_items = _remove_tools_from_items(new_items)
+
+    return HandoffInputData(
+        input_history=filtered_history,
+        pre_handoff_items=filtered_pre_handoff_items,
+        new_items=filtered_new_items,
+        run_context=handoff_input_data.run_context,
+    )
+
+
+def _remove_tools_from_items(items: tuple[RunItem, ...]) -> tuple[RunItem, ...]:
+    filtered_items = []
+    for item in items:
+        if (
+            isinstance(item, HandoffCallItem)
+            or isinstance(item, HandoffOutputItem)
+            or isinstance(item, ToolCallItem)
+            or isinstance(item, ToolCallOutputItem)
+            or isinstance(item, ReasoningItem)
+        ):
+            continue
+        filtered_items.append(item)
+    return tuple(filtered_items)
+
+
+def _remove_tool_types_from_input(
+    items: tuple[TResponseInputItem, ...],
+) -> tuple[TResponseInputItem, ...]:
+    tool_types = [
+        "function_call",
+        "function_call_output",
+        "computer_call",
+        "computer_call_output",
+        "file_search_call",
+        "web_search_call",
+    ]
+
+    filtered_items: list[TResponseInputItem] = []
+    for item in items:
+        itype = item.get("type")
+        if itype in tool_types:
+            continue
+        filtered_items.append(item)
+    return tuple(filtered_items)

agents/extensions/handoff_prompt.py

@@ -0,0 +1,19 @@
+# A recommended prompt prefix for agents that use handoffs. We recommend including this or
+# similar instructions in any agents that use handoffs.
+RECOMMENDED_PROMPT_PREFIX = (
+    "# System context\n"
+    "You are part of a multi-agent system called the Agents SDK, designed to make agent "
+    "coordination and execution easy. Agents uses two primary abstraction: **Agents** and "
+    "**Handoffs**. An agent encompasses instructions and tools and can hand off a "
+    "conversation to another agent when appropriate. "
+    "Handoffs are achieved by calling a handoff function, generally named "
+    "`transfer_to_<agent_name>`. Transfers between agents are handled seamlessly in the background;"
+    " do not mention or draw attention to these transfers in your conversation with the user.\n"
+)
+
+
+def prompt_with_handoff_instructions(prompt: str) -> str:
+    """
+    Add recommended instructions to the prompt for agents that use handoffs.
+    """
+    return f"{RECOMMENDED_PROMPT_PREFIX}\n\n{prompt}"

agents/extensions/memory/__init__.py

@@ -0,0 +1,65 @@
+"""Session memory backends living in the extensions namespace.
+
+This package contains optional, production-grade session implementations that
+introduce extra third-party dependencies (database drivers, ORMs, etc.). They
+conform to the :class:`agents.memory.session.Session` protocol so they can be
+used as a drop-in replacement for :class:`agents.memory.session.SQLiteSession`.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+__all__: list[str] = [
+    "EncryptedSession",
+    "RedisSession",
+    "SQLAlchemySession",
+    "AdvancedSQLiteSession",
+]
+
+
+def __getattr__(name: str) -> Any:
+    if name == "EncryptedSession":
+        try:
+            from .encrypt_session import EncryptedSession  # noqa: F401
+
+            return EncryptedSession
+        except ModuleNotFoundError as e:
+            raise ImportError(
+                "EncryptedSession requires the 'cryptography' extra. "
+                "Install it with: pip install openai-agents[encrypt]"
+            ) from e
+
+    if name == "RedisSession":
+        try:
+            from .redis_session import RedisSession  # noqa: F401
+
+            return RedisSession
+        except ModuleNotFoundError as e:
+            raise ImportError(
+                "RedisSession requires the 'redis' extra. "
+                "Install it with: pip install openai-agents[redis]"
+            ) from e
+
+    if name == "SQLAlchemySession":
+        try:
+            from .sqlalchemy_session import SQLAlchemySession  # noqa: F401
+
+            return SQLAlchemySession
+        except ModuleNotFoundError as e:
+            raise ImportError(
+                "SQLAlchemySession requires the 'sqlalchemy' extra. "
+                "Install it with: pip install openai-agents[sqlalchemy]"
+            ) from e
+
+    if name == "AdvancedSQLiteSession":
+        try:
+            from .advanced_sqlite_session import AdvancedSQLiteSession  # noqa: F401
+
+            return AdvancedSQLiteSession
+        except ModuleNotFoundError as e:
+            raise ImportError(
+                f"Failed to import AdvancedSQLiteSession: {e}"
+            ) from e
+
+    raise AttributeError(f"module {__name__} has no attribute {name}")

agents/extensions/memory/advanced_sqlite_session.py

@@ -0,0 +1,1285 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import threading
+from contextlib import closing
+from pathlib import Path
+from typing import Any, Union, cast
+
+from agents.result import RunResult
+from agents.usage import Usage
+
+from ...items import TResponseInputItem
+from ...memory import SQLiteSession
+
+
+class AdvancedSQLiteSession(SQLiteSession):
+    """Enhanced SQLite session with conversation branching and usage analytics."""
+
+    def __init__(
+        self,
+        *,
+        session_id: str,
+        db_path: str | Path = ":memory:",
+        create_tables: bool = False,
+        logger: logging.Logger | None = None,
+        **kwargs,
+    ):
+        """Initialize the AdvancedSQLiteSession.
+
+        Args:
+            session_id: The ID of the session
+            db_path: The path to the SQLite database file. Defaults to `:memory:` for in-memory storage
+            create_tables: Whether to create the structure tables
+            logger: The logger to use. Defaults to the module logger
+            **kwargs: Additional keyword arguments to pass to the superclass
+        """  # noqa: E501
+        super().__init__(session_id, db_path, **kwargs)
+        if create_tables:
+            self._init_structure_tables()
+        self._current_branch_id = "main"
+        self._logger = logger or logging.getLogger(__name__)
+
+    def _init_structure_tables(self):
+        """Add structure and usage tracking tables.
+
+        Creates the message_structure and turn_usage tables with appropriate
+        indexes for conversation branching and usage analytics.
+        """
+        conn = self._get_connection()
+
+        # Message structure with branch support
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS message_structure (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                session_id TEXT NOT NULL,
+                message_id INTEGER NOT NULL,
+                branch_id TEXT NOT NULL DEFAULT 'main',
+                message_type TEXT NOT NULL,
+                sequence_number INTEGER NOT NULL,
+                user_turn_number INTEGER,
+                branch_turn_number INTEGER,
+                tool_name TEXT,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                FOREIGN KEY (session_id) REFERENCES agent_sessions(session_id) ON DELETE CASCADE,
+                FOREIGN KEY (message_id) REFERENCES agent_messages(id) ON DELETE CASCADE
+            )
+        """)
+
+        # Turn-level usage tracking with branch support and full JSON details
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS turn_usage (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                session_id TEXT NOT NULL,
+                branch_id TEXT NOT NULL DEFAULT 'main',
+                user_turn_number INTEGER NOT NULL,
+                requests INTEGER DEFAULT 0,
+                input_tokens INTEGER DEFAULT 0,
+                output_tokens INTEGER DEFAULT 0,
+                total_tokens INTEGER DEFAULT 0,
+                input_tokens_details JSON,
+                output_tokens_details JSON,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                FOREIGN KEY (session_id) REFERENCES agent_sessions(session_id) ON DELETE CASCADE,
+                UNIQUE(session_id, branch_id, user_turn_number)
+            )
+        """)
+
+        # Indexes
+        conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_structure_session_seq
+            ON message_structure(session_id, sequence_number)
+        """)
+        conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_structure_branch
+            ON message_structure(session_id, branch_id)
+        """)
+        conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_structure_turn
+            ON message_structure(session_id, branch_id, user_turn_number)
+        """)
+        conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_structure_branch_seq
+            ON message_structure(session_id, branch_id, sequence_number)
+        """)
+        conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_turn_usage_session_turn
+            ON turn_usage(session_id, branch_id, user_turn_number)
+        """)
+
+        conn.commit()
+
+    async def add_items(self, items: list[TResponseInputItem]) -> None:
+        """Add items to the session.
+
+        Args:
+            items: The items to add to the session
+        """
+        # Add to base table first
+        await super().add_items(items)
+
+        # Extract structure metadata with precise sequencing
+        if items:
+            await self._add_structure_metadata(items)
+
+    async def get_items(
+        self,
+        limit: int | None = None,
+        branch_id: str | None = None,
+    ) -> list[TResponseInputItem]:
+        """Get items from current or specified branch.
+
+        Args:
+            limit: Maximum number of items to return. If None, returns all items.
+            branch_id: Branch to get items from. If None, uses current branch.
+
+        Returns:
+            List of conversation items from the specified branch.
+        """
+        if branch_id is None:
+            branch_id = self._current_branch_id
+
+            # Get all items for this branch
+            def _get_all_items_sync():
+                """Synchronous helper to get all items for a branch."""
+                conn = self._get_connection()
+                # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+                with self._lock if self._is_memory_db else threading.Lock():
+                    with closing(conn.cursor()) as cursor:
+                        if limit is None:
+                            cursor.execute(
+                                """
+                                SELECT m.message_data
+                                FROM agent_messages m
+                                JOIN message_structure s ON m.id = s.message_id
+                                WHERE m.session_id = ? AND s.branch_id = ?
+                                ORDER BY s.sequence_number ASC
+                            """,
+                                (self.session_id, branch_id),
+                            )
+                        else:
+                            cursor.execute(
+                                """
+                                SELECT m.message_data
+                                FROM agent_messages m
+                                JOIN message_structure s ON m.id = s.message_id
+                                WHERE m.session_id = ? AND s.branch_id = ?
+                                ORDER BY s.sequence_number DESC
+                                LIMIT ?
+                            """,
+                                (self.session_id, branch_id, limit),
+                            )
+
+                        rows = cursor.fetchall()
+                        if limit is not None:
+                            rows = list(reversed(rows))
+
+                    items = []
+                    for (message_data,) in rows:
+                        try:
+                            item = json.loads(message_data)
+                            items.append(item)
+                        except json.JSONDecodeError:
+                            continue
+                    return items
+
+            return await asyncio.to_thread(_get_all_items_sync)
+
+        def _get_items_sync():
+            """Synchronous helper to get items for a specific branch."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                with closing(conn.cursor()) as cursor:
+                    # Get message IDs in correct order for this branch
+                    if limit is None:
+                        cursor.execute(
+                            """
+                            SELECT m.message_data
+                            FROM agent_messages m
+                            JOIN message_structure s ON m.id = s.message_id
+                            WHERE m.session_id = ? AND s.branch_id = ?
+                            ORDER BY s.sequence_number ASC
+                        """,
+                            (self.session_id, branch_id),
+                        )
+                    else:
+                        cursor.execute(
+                            """
+                            SELECT m.message_data
+                            FROM agent_messages m
+                            JOIN message_structure s ON m.id = s.message_id
+                            WHERE m.session_id = ? AND s.branch_id = ?
+                            ORDER BY s.sequence_number DESC
+                            LIMIT ?
+                        """,
+                            (self.session_id, branch_id, limit),
+                        )
+
+                    rows = cursor.fetchall()
+                    if limit is not None:
+                        rows = list(reversed(rows))
+
+                items = []
+                for (message_data,) in rows:
+                    try:
+                        item = json.loads(message_data)
+                        items.append(item)
+                    except json.JSONDecodeError:
+                        continue
+                return items
+
+        return await asyncio.to_thread(_get_items_sync)
+
+    async def store_run_usage(self, result: RunResult) -> None:
+        """Store usage data for the current conversation turn.
+
+        This is designed to be called after `Runner.run()` completes.
+        Session-level usage can be aggregated from turn data when needed.
+
+        Args:
+            result: The result from the run
+        """
+        try:
+            if result.context_wrapper.usage is not None:
+                # Get the current turn number for this branch
+                current_turn = self._get_current_turn_number()
+                # Only update turn-level usage - session usage is aggregated on demand
+                await self._update_turn_usage_internal(current_turn, result.context_wrapper.usage)
+        except Exception as e:
+            self._logger.error(f"Failed to store usage for session {self.session_id}: {e}")
+
+    def _get_next_turn_number(self, branch_id: str) -> int:
+        """Get the next turn number for a specific branch.
+
+        Args:
+            branch_id: The branch ID to get the next turn number for.
+
+        Returns:
+            The next available turn number for the specified branch.
+        """
+        conn = self._get_connection()
+        with closing(conn.cursor()) as cursor:
+            cursor.execute(
+                """
+                SELECT COALESCE(MAX(user_turn_number), 0)
+                FROM message_structure
+                WHERE session_id = ? AND branch_id = ?
+            """,
+                (self.session_id, branch_id),
+            )
+            result = cursor.fetchone()
+            max_turn = result[0] if result else 0
+            return max_turn + 1
+
+    def _get_next_branch_turn_number(self, branch_id: str) -> int:
+        """Get the next branch turn number for a specific branch.
+
+        Args:
+            branch_id: The branch ID to get the next branch turn number for.
+
+        Returns:
+            The next available branch turn number for the specified branch.
+        """
+        conn = self._get_connection()
+        with closing(conn.cursor()) as cursor:
+            cursor.execute(
+                """
+                SELECT COALESCE(MAX(branch_turn_number), 0)
+                FROM message_structure
+                WHERE session_id = ? AND branch_id = ?
+            """,
+                (self.session_id, branch_id),
+            )
+            result = cursor.fetchone()
+            max_turn = result[0] if result else 0
+            return max_turn + 1
+
+    def _get_current_turn_number(self) -> int:
+        """Get the current turn number for the current branch.
+
+        Returns:
+            The current turn number for the active branch.
+        """
+        conn = self._get_connection()
+        with closing(conn.cursor()) as cursor:
+            cursor.execute(
+                """
+                SELECT COALESCE(MAX(user_turn_number), 0)
+                FROM message_structure
+                WHERE session_id = ? AND branch_id = ?
+                """,
+                (self.session_id, self._current_branch_id),
+            )
+            result = cursor.fetchone()
+            return result[0] if result else 0
+
+    async def _add_structure_metadata(self, items: list[TResponseInputItem]) -> None:
+        """Extract structure metadata with branch-aware turn tracking.
+
+        This method:
+        - Assigns turn numbers per branch (not globally)
+        - Assigns explicit sequence numbers for precise ordering
+        - Links messages to their database IDs for structure tracking
+        - Handles multiple user messages in a single batch correctly
+
+        Args:
+            items: The items to add to the session
+        """
+
+        def _add_structure_sync():
+            """Synchronous helper to add structure metadata to database."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                # Get the IDs of messages we just inserted, in order
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(
+                        f"SELECT id FROM {self.messages_table} "
+                        f"WHERE session_id = ? ORDER BY id DESC LIMIT ?",
+                        (self.session_id, len(items)),
+                    )
+                    message_ids = [row[0] for row in cursor.fetchall()]
+                    message_ids.reverse()  # Match order of items
+
+                # Get current max sequence number (global)
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(
+                        """
+                        SELECT COALESCE(MAX(sequence_number), 0)
+                        FROM message_structure
+                        WHERE session_id = ?
+                    """,
+                        (self.session_id,),
+                    )
+                    seq_start = cursor.fetchone()[0]
+
+                # Get current turn numbers atomically with a single query
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(
+                        """
+                        SELECT
+                            COALESCE(MAX(user_turn_number), 0) as max_global_turn,
+                            COALESCE(MAX(branch_turn_number), 0) as max_branch_turn
+                        FROM message_structure
+                        WHERE session_id = ? AND branch_id = ?
+                    """,
+                        (self.session_id, self._current_branch_id),
+                    )
+                    result = cursor.fetchone()
+                    current_turn = result[0] if result else 0
+                    current_branch_turn = result[1] if result else 0
+
+                # Process items and assign turn numbers correctly
+                structure_data = []
+                user_message_count = 0
+
+                for i, (item, msg_id) in enumerate(zip(items, message_ids)):
+                    msg_type = self._classify_message_type(item)
+                    tool_name = self._extract_tool_name(item)
+
+                    # If this is a user message, increment turn counters
+                    if self._is_user_message(item):
+                        user_message_count += 1
+                        item_turn = current_turn + user_message_count
+                        item_branch_turn = current_branch_turn + user_message_count
+                    else:
+                        # Non-user messages inherit the turn number of the most recent user message
+                        item_turn = current_turn + user_message_count
+                        item_branch_turn = current_branch_turn + user_message_count
+
+                    structure_data.append(
+                        (
+                            self.session_id,
+                            msg_id,
+                            self._current_branch_id,
+                            msg_type,
+                            seq_start + i + 1,  # Global sequence
+                            item_turn,  # Global turn number
+                            item_branch_turn,  # Branch-specific turn number
+                            tool_name,
+                        )
+                    )
+
+                with closing(conn.cursor()) as cursor:
+                    cursor.executemany(
+                        """
+                        INSERT INTO message_structure
+                        (session_id, message_id, branch_id, message_type, sequence_number,
+                         user_turn_number, branch_turn_number, tool_name)
+                        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                    """,
+                        structure_data,
+                    )
+                    conn.commit()
+
+        try:
+            await asyncio.to_thread(_add_structure_sync)
+        except Exception as e:
+            self._logger.error(
+                f"Failed to add structure metadata for session {self.session_id}: {e}"
+            )
+            # Try to clean up any orphaned messages to maintain consistency
+            try:
+                await self._cleanup_orphaned_messages()
+            except Exception as cleanup_error:
+                self._logger.error(f"Failed to cleanup orphaned messages: {cleanup_error}")
+            # Don't re-raise - structure metadata is supplementary
+
+    async def _cleanup_orphaned_messages(self) -> None:
+        """Remove messages that exist in agent_messages but not in message_structure.
+
+        This can happen if _add_structure_metadata fails after super().add_items() succeeds.
+        Used for maintaining data consistency.
+        """
+
+        def _cleanup_sync():
+            """Synchronous helper to cleanup orphaned messages."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                with closing(conn.cursor()) as cursor:
+                    # Find messages without structure metadata
+                    cursor.execute(
+                        """
+                        SELECT am.id
+                        FROM agent_messages am
+                        LEFT JOIN message_structure ms ON am.id = ms.message_id
+                        WHERE am.session_id = ? AND ms.message_id IS NULL
+                    """,
+                        (self.session_id,),
+                    )
+
+                    orphaned_ids = [row[0] for row in cursor.fetchall()]
+
+                    if orphaned_ids:
+                        # Delete orphaned messages
+                        placeholders = ",".join("?" * len(orphaned_ids))
+                        cursor.execute(
+                            f"DELETE FROM agent_messages WHERE id IN ({placeholders})", orphaned_ids
+                        )
+
+                        deleted_count = cursor.rowcount
+                        conn.commit()
+
+                        self._logger.info(f"Cleaned up {deleted_count} orphaned messages")
+                        return deleted_count
+
+                    return 0
+
+        return await asyncio.to_thread(_cleanup_sync)
+
+    def _classify_message_type(self, item: TResponseInputItem) -> str:
+        """Classify the type of a message item.
+
+        Args:
+            item: The message item to classify.
+
+        Returns:
+            String representing the message type (user, assistant, etc.).
+        """
+        if isinstance(item, dict):
+            if item.get("role") == "user":
+                return "user"
+            elif item.get("role") == "assistant":
+                return "assistant"
+            elif item.get("type"):
+                return str(item.get("type"))
+        return "other"
+
+    def _extract_tool_name(self, item: TResponseInputItem) -> str | None:
+        """Extract tool name if this is a tool call/output.
+
+        Args:
+            item: The message item to extract tool name from.
+
+        Returns:
+            Tool name if item is a tool call, None otherwise.
+        """
+        if isinstance(item, dict):
+            item_type = item.get("type")
+
+            # For MCP tools, try to extract from server_label if available
+            if item_type in {"mcp_call", "mcp_approval_request"} and "server_label" in item:
+                server_label = item.get("server_label")
+                tool_name = item.get("name")
+                if tool_name and server_label:
+                    return f"{server_label}.{tool_name}"
+                elif server_label:
+                    return str(server_label)
+                elif tool_name:
+                    return str(tool_name)
+
+            # For tool types without a 'name' field, derive from the type
+            elif item_type in {
+                "computer_call",
+                "file_search_call",
+                "web_search_call",
+                "code_interpreter_call",
+            }:
+                return item_type
+
+            # Most other tool calls have a 'name' field
+            elif "name" in item:
+                name = item.get("name")
+                return str(name) if name is not None else None
+
+        return None
+
+    def _is_user_message(self, item: TResponseInputItem) -> bool:
+        """Check if this is a user message.
+
+        Args:
+            item: The message item to check.
+
+        Returns:
+            True if the item is a user message, False otherwise.
+        """
+        return isinstance(item, dict) and item.get("role") == "user"
+
+    async def create_branch_from_turn(
+        self, turn_number: int, branch_name: str | None = None
+    ) -> str:
+        """Create a new branch starting from a specific user message turn.
+
+        Args:
+            turn_number: The branch turn number of the user message to branch from
+            branch_name: Optional name for the branch (auto-generated if None)
+
+        Returns:
+            The branch_id of the newly created branch
+
+        Raises:
+            ValueError: If turn doesn't exist or doesn't contain a user message
+        """
+        import time
+
+        # Validate the turn exists and contains a user message
+        def _validate_turn():
+            """Synchronous helper to validate turn exists and contains user message."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT am.message_data
+                    FROM message_structure ms
+                    JOIN agent_messages am ON ms.message_id = am.id
+                    WHERE ms.session_id = ? AND ms.branch_id = ?
+                    AND ms.branch_turn_number = ? AND ms.message_type = 'user'
+                    """,
+                    (self.session_id, self._current_branch_id, turn_number),
+                )
+
+                result = cursor.fetchone()
+                if not result:
+                    raise ValueError(
+                        f"Turn {turn_number} does not contain a user message "
+                        f"in branch '{self._current_branch_id}'"
+                    )
+
+                message_data = result[0]
+                try:
+                    content = json.loads(message_data).get("content", "")
+                    return content[:50] + "..." if len(content) > 50 else content
+                except Exception:
+                    return "Unable to parse content"
+
+        turn_content = await asyncio.to_thread(_validate_turn)
+
+        # Generate branch name if not provided
+        if branch_name is None:
+            timestamp = int(time.time())
+            branch_name = f"branch_from_turn_{turn_number}_{timestamp}"
+
+        # Copy messages before the branch point to the new branch
+        await self._copy_messages_to_new_branch(branch_name, turn_number)
+
+        # Switch to new branch
+        old_branch = self._current_branch_id
+        self._current_branch_id = branch_name
+
+        self._logger.debug(
+            f"Created branch '{branch_name}' from turn {turn_number} ('{turn_content}') in '{old_branch}'"  # noqa: E501
+        )
+        return branch_name
+
+    async def create_branch_from_content(
+        self, search_term: str, branch_name: str | None = None
+    ) -> str:
+        """Create branch from the first user turn matching the search term.
+
+        Args:
+            search_term: Text to search for in user messages.
+            branch_name: Optional name for the branch (auto-generated if None).
+
+        Returns:
+            The branch_id of the newly created branch.
+
+        Raises:
+            ValueError: If no matching turns are found.
+        """
+        matching_turns = await self.find_turns_by_content(search_term)
+        if not matching_turns:
+            raise ValueError(f"No user turns found containing '{search_term}'")
+
+        # Use the first (earliest) match
+        turn_number = matching_turns[0]["turn"]
+        return await self.create_branch_from_turn(turn_number, branch_name)
+
+    async def switch_to_branch(self, branch_id: str) -> None:
+        """Switch to a different branch.
+
+        Args:
+            branch_id: The branch to switch to.
+
+        Raises:
+            ValueError: If the branch doesn't exist.
+        """
+
+        # Validate branch exists
+        def _validate_branch():
+            """Synchronous helper to validate branch exists."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT COUNT(*) FROM message_structure
+                    WHERE session_id = ? AND branch_id = ?
+                """,
+                    (self.session_id, branch_id),
+                )
+
+                count = cursor.fetchone()[0]
+                if count == 0:
+                    raise ValueError(f"Branch '{branch_id}' does not exist")
+
+        await asyncio.to_thread(_validate_branch)
+
+        old_branch = self._current_branch_id
+        self._current_branch_id = branch_id
+        self._logger.info(f"Switched from branch '{old_branch}' to '{branch_id}'")
+
+    async def delete_branch(self, branch_id: str, force: bool = False) -> None:
+        """Delete a branch and all its associated data.
+
+        Args:
+            branch_id: The branch to delete.
+            force: If True, allows deleting the current branch (will switch to 'main').
+
+        Raises:
+            ValueError: If branch doesn't exist, is 'main', or is current branch without force.
+        """
+        if not branch_id or not branch_id.strip():
+            raise ValueError("Branch ID cannot be empty")
+
+        branch_id = branch_id.strip()
+
+        # Protect main branch
+        if branch_id == "main":
+            raise ValueError("Cannot delete the 'main' branch")
+
+        # Check if trying to delete current branch
+        if branch_id == self._current_branch_id:
+            if not force:
+                raise ValueError(
+                    f"Cannot delete current branch '{branch_id}'. Use force=True or switch branches first"  # noqa: E501
+                )
+            else:
+                # Switch to main before deleting
+                await self.switch_to_branch("main")
+
+        def _delete_sync():
+            """Synchronous helper to delete branch and associated data."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                with closing(conn.cursor()) as cursor:
+                    # First verify the branch exists
+                    cursor.execute(
+                        """
+                        SELECT COUNT(*) FROM message_structure
+                        WHERE session_id = ? AND branch_id = ?
+                    """,
+                        (self.session_id, branch_id),
+                    )
+
+                    count = cursor.fetchone()[0]
+                    if count == 0:
+                        raise ValueError(f"Branch '{branch_id}' does not exist")
+
+                    # Delete from turn_usage first (foreign key constraint)
+                    cursor.execute(
+                        """
+                        DELETE FROM turn_usage
+                        WHERE session_id = ? AND branch_id = ?
+                    """,
+                        (self.session_id, branch_id),
+                    )
+
+                    usage_deleted = cursor.rowcount
+
+                    # Delete from message_structure
+                    cursor.execute(
+                        """
+                        DELETE FROM message_structure
+                        WHERE session_id = ? AND branch_id = ?
+                    """,
+                        (self.session_id, branch_id),
+                    )
+
+                    structure_deleted = cursor.rowcount
+
+                    conn.commit()
+
+                    return usage_deleted, structure_deleted
+
+        usage_deleted, structure_deleted = await asyncio.to_thread(_delete_sync)
+
+        self._logger.info(
+            f"Deleted branch '{branch_id}': {structure_deleted} message entries, {usage_deleted} usage entries"  # noqa: E501
+        )
+
+    async def list_branches(self) -> list[dict[str, Any]]:
+        """List all branches in this session.
+
+        Returns:
+            List of dicts with branch info containing:
+                - 'branch_id': Branch identifier
+                - 'message_count': Number of messages in branch
+                - 'user_turns': Number of user turns in branch
+                - 'is_current': Whether this is the current branch
+                - 'created_at': When the branch was first created
+        """
+
+        def _list_branches_sync():
+            """Synchronous helper to list all branches."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT
+                        ms.branch_id,
+                        COUNT(*) as message_count,
+                        COUNT(CASE WHEN ms.message_type = 'user' THEN 1 END) as user_turns,
+                        MIN(ms.created_at) as created_at
+                    FROM message_structure ms
+                    WHERE ms.session_id = ?
+                    GROUP BY ms.branch_id
+                    ORDER BY created_at
+                """,
+                    (self.session_id,),
+                )
+
+                branches = []
+                for row in cursor.fetchall():
+                    branch_id, msg_count, user_turns, created_at = row
+                    branches.append(
+                        {
+                            "branch_id": branch_id,
+                            "message_count": msg_count,
+                            "user_turns": user_turns,
+                            "is_current": branch_id == self._current_branch_id,
+                            "created_at": created_at,
+                        }
+                    )
+
+                return branches
+
+        return await asyncio.to_thread(_list_branches_sync)
+
+    async def _copy_messages_to_new_branch(self, new_branch_id: str, from_turn_number: int) -> None:
+        """Copy messages before the branch point to the new branch.
+
+        Args:
+            new_branch_id: The ID of the new branch to copy messages to.
+            from_turn_number: The turn number to copy messages up to (exclusive).
+        """
+
+        def _copy_sync():
+            """Synchronous helper to copy messages to new branch."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                with closing(conn.cursor()) as cursor:
+                    # Get all messages before the branch point
+                    cursor.execute(
+                        """
+                        SELECT
+                            ms.message_id,
+                            ms.message_type,
+                            ms.sequence_number,
+                            ms.user_turn_number,
+                            ms.branch_turn_number,
+                            ms.tool_name
+                        FROM message_structure ms
+                        WHERE ms.session_id = ? AND ms.branch_id = ?
+                        AND ms.branch_turn_number < ?
+                        ORDER BY ms.sequence_number
+                    """,
+                        (self.session_id, self._current_branch_id, from_turn_number),
+                    )
+
+                    messages_to_copy = cursor.fetchall()
+
+                    if messages_to_copy:
+                        # Get the max sequence number for the new inserts
+                        cursor.execute(
+                            """
+                            SELECT COALESCE(MAX(sequence_number), 0)
+                            FROM message_structure
+                            WHERE session_id = ?
+                        """,
+                            (self.session_id,),
+                        )
+
+                        seq_start = cursor.fetchone()[0]
+
+                        # Insert copied messages with new branch_id
+                        new_structure_data = []
+                        for i, (
+                            msg_id,
+                            msg_type,
+                            _,
+                            user_turn,
+                            branch_turn,
+                            tool_name,
+                        ) in enumerate(messages_to_copy):
+                            new_structure_data.append(
+                                (
+                                    self.session_id,
+                                    msg_id,  # Same message_id (sharing the actual message data)
+                                    new_branch_id,
+                                    msg_type,
+                                    seq_start + i + 1,  # New sequence number
+                                    user_turn,  # Keep same global turn number
+                                    branch_turn,  # Keep same branch turn number
+                                    tool_name,
+                                )
+                            )
+
+                        cursor.executemany(
+                            """
+                            INSERT INTO message_structure
+                            (session_id, message_id, branch_id, message_type, sequence_number,
+                             user_turn_number, branch_turn_number, tool_name)
+                            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                        """,
+                            new_structure_data,
+                        )
+
+                    conn.commit()
+
+        await asyncio.to_thread(_copy_sync)
+
+    async def get_conversation_turns(self, branch_id: str | None = None) -> list[dict[str, Any]]:
+        """Get user turns with content for easy browsing and branching decisions.
+
+        Args:
+            branch_id: Branch to get turns from (current branch if None).
+
+        Returns:
+            List of dicts with turn info containing:
+                - 'turn': Branch turn number
+                - 'content': User message content (truncated)
+                - 'full_content': Full user message content
+                - 'timestamp': When the turn was created
+                - 'can_branch': Always True (all user messages can branch)
+        """
+        if branch_id is None:
+            branch_id = self._current_branch_id
+
+        def _get_turns_sync():
+            """Synchronous helper to get conversation turns."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT
+                        ms.branch_turn_number,
+                        am.message_data,
+                        ms.created_at
+                    FROM message_structure ms
+                    JOIN agent_messages am ON ms.message_id = am.id
+                    WHERE ms.session_id = ? AND ms.branch_id = ?
+                    AND ms.message_type = 'user'
+                    ORDER BY ms.branch_turn_number
+                """,
+                    (self.session_id, branch_id),
+                )
+
+                turns = []
+                for row in cursor.fetchall():
+                    turn_num, message_data, created_at = row
+                    try:
+                        content = json.loads(message_data).get("content", "")
+                        turns.append(
+                            {
+                                "turn": turn_num,
+                                "content": content[:100] + "..." if len(content) > 100 else content,
+                                "full_content": content,
+                                "timestamp": created_at,
+                                "can_branch": True,
+                            }
+                        )
+                    except (json.JSONDecodeError, AttributeError):
+                        continue
+
+                return turns
+
+        return await asyncio.to_thread(_get_turns_sync)
+
+    async def find_turns_by_content(
+        self, search_term: str, branch_id: str | None = None
+    ) -> list[dict[str, Any]]:
+        """Find user turns containing specific content.
+
+        Args:
+            search_term: Text to search for in user messages.
+            branch_id: Branch to search in (current branch if None).
+
+        Returns:
+            List of matching turns with same format as get_conversation_turns().
+        """
+        if branch_id is None:
+            branch_id = self._current_branch_id
+
+        def _search_sync():
+            """Synchronous helper to search turns by content."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT
+                        ms.branch_turn_number,
+                        am.message_data,
+                        ms.created_at
+                    FROM message_structure ms
+                    JOIN agent_messages am ON ms.message_id = am.id
+                    WHERE ms.session_id = ? AND ms.branch_id = ?
+                    AND ms.message_type = 'user'
+                    AND am.message_data LIKE ?
+                    ORDER BY ms.branch_turn_number
+                """,
+                    (self.session_id, branch_id, f"%{search_term}%"),
+                )
+
+                matches = []
+                for row in cursor.fetchall():
+                    turn_num, message_data, created_at = row
+                    try:
+                        content = json.loads(message_data).get("content", "")
+                        matches.append(
+                            {
+                                "turn": turn_num,
+                                "content": content,
+                                "full_content": content,
+                                "timestamp": created_at,
+                                "can_branch": True,
+                            }
+                        )
+                    except (json.JSONDecodeError, AttributeError):
+                        continue
+
+                return matches
+
+        return await asyncio.to_thread(_search_sync)
+
+    async def get_conversation_by_turns(
+        self, branch_id: str | None = None
+    ) -> dict[int, list[dict[str, str | None]]]:
+        """Get conversation grouped by user turns for specified branch.
+
+        Args:
+            branch_id: Branch to get conversation from (current branch if None).
+
+        Returns:
+            Dictionary mapping turn numbers to lists of message metadata.
+        """
+        if branch_id is None:
+            branch_id = self._current_branch_id
+
+        def _get_conversation_sync():
+            """Synchronous helper to get conversation by turns."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT user_turn_number, message_type, tool_name
+                    FROM message_structure
+                    WHERE session_id = ? AND branch_id = ?
+                    ORDER BY sequence_number
+                """,
+                    (self.session_id, branch_id),
+                )
+
+                turns: dict[int, list[dict[str, str | None]]] = {}
+                for row in cursor.fetchall():
+                    turn_num, msg_type, tool_name = row
+                    if turn_num not in turns:
+                        turns[turn_num] = []
+                    turns[turn_num].append({"type": msg_type, "tool_name": tool_name})
+                return turns
+
+        return await asyncio.to_thread(_get_conversation_sync)
+
+    async def get_tool_usage(self, branch_id: str | None = None) -> list[tuple[str, int, int]]:
+        """Get all tool usage by turn for specified branch.
+
+        Args:
+            branch_id: Branch to get tool usage from (current branch if None).
+
+        Returns:
+            List of tuples containing (tool_name, usage_count, turn_number).
+        """
+        if branch_id is None:
+            branch_id = self._current_branch_id
+
+        def _get_tool_usage_sync():
+            """Synchronous helper to get tool usage statistics."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT tool_name, COUNT(*), user_turn_number
+                    FROM message_structure
+                    WHERE session_id = ? AND branch_id = ? AND message_type IN (
+                        'tool_call', 'function_call', 'computer_call', 'file_search_call',
+                        'web_search_call', 'code_interpreter_call', 'custom_tool_call',
+                        'mcp_call', 'mcp_approval_request'
+                    )
+                    GROUP BY tool_name, user_turn_number
+                    ORDER BY user_turn_number
+                """,
+                    (self.session_id, branch_id),
+                )
+                return cursor.fetchall()
+
+        return await asyncio.to_thread(_get_tool_usage_sync)
+
+    async def get_session_usage(self, branch_id: str | None = None) -> dict[str, int] | None:
+        """Get cumulative usage for session or specific branch.
+
+        Args:
+            branch_id: If provided, only get usage for that branch. If None, get all branches.
+
+        Returns:
+            Dictionary with usage statistics or None if no usage data found.
+        """
+
+        def _get_usage_sync():
+            """Synchronous helper to get session usage data."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                if branch_id:
+                    # Branch-specific usage
+                    query = """
+                        SELECT
+                            SUM(requests) as total_requests,
+                            SUM(input_tokens) as total_input_tokens,
+                            SUM(output_tokens) as total_output_tokens,
+                            SUM(total_tokens) as total_total_tokens,
+                            COUNT(*) as total_turns
+                        FROM turn_usage
+                        WHERE session_id = ? AND branch_id = ?
+                    """
+                    params: tuple[str, ...] = (self.session_id, branch_id)
+                else:
+                    # All branches
+                    query = """
+                        SELECT
+                            SUM(requests) as total_requests,
+                            SUM(input_tokens) as total_input_tokens,
+                            SUM(output_tokens) as total_output_tokens,
+                            SUM(total_tokens) as total_total_tokens,
+                            COUNT(*) as total_turns
+                        FROM turn_usage
+                        WHERE session_id = ?
+                    """
+                    params = (self.session_id,)
+
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(query, params)
+                    row = cursor.fetchone()
+
+                    if row and row[0] is not None:
+                        return {
+                            "requests": row[0] or 0,
+                            "input_tokens": row[1] or 0,
+                            "output_tokens": row[2] or 0,
+                            "total_tokens": row[3] or 0,
+                            "total_turns": row[4] or 0,
+                        }
+                    return None
+
+        result = await asyncio.to_thread(_get_usage_sync)
+
+        return cast(Union[dict[str, int], None], result)
+
+    async def get_turn_usage(
+        self,
+        user_turn_number: int | None = None,
+        branch_id: str | None = None,
+    ) -> list[dict[str, Any]] | dict[str, Any]:
+        """Get usage statistics by turn with full JSON token details.
+
+        Args:
+            user_turn_number: Specific turn to get usage for. If None, returns all turns.
+            branch_id: Branch to get usage from (current branch if None).
+
+        Returns:
+            Dictionary with usage data for specific turn, or list of dictionaries for all turns.
+        """
+
+        if branch_id is None:
+            branch_id = self._current_branch_id
+
+        def _get_turn_usage_sync():
+            """Synchronous helper to get turn usage statistics."""
+            conn = self._get_connection()
+
+            if user_turn_number is not None:
+                query = """
+                    SELECT requests, input_tokens, output_tokens, total_tokens,
+                           input_tokens_details, output_tokens_details
+                    FROM turn_usage
+                    WHERE session_id = ? AND branch_id = ? AND user_turn_number = ?
+                """
+
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(query, (self.session_id, branch_id, user_turn_number))
+                    row = cursor.fetchone()
+
+                    if row:
+                        # Parse JSON details if present
+                        input_details = None
+                        output_details = None
+
+                        if row[4]:  # input_tokens_details
+                            try:
+                                input_details = json.loads(row[4])
+                            except json.JSONDecodeError:
+                                pass
+
+                        if row[5]:  # output_tokens_details
+                            try:
+                                output_details = json.loads(row[5])
+                            except json.JSONDecodeError:
+                                pass
+
+                        return {
+                            "requests": row[0],
+                            "input_tokens": row[1],
+                            "output_tokens": row[2],
+                            "total_tokens": row[3],
+                            "input_tokens_details": input_details,
+                            "output_tokens_details": output_details,
+                        }
+                    return {}
+            else:
+                query = """
+                    SELECT user_turn_number, requests, input_tokens, output_tokens,
+                           total_tokens, input_tokens_details, output_tokens_details
+                    FROM turn_usage
+                    WHERE session_id = ? AND branch_id = ?
+                    ORDER BY user_turn_number
+                """
+
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(query, (self.session_id, branch_id))
+                    results = []
+                    for row in cursor.fetchall():
+                        # Parse JSON details if present
+                        input_details = None
+                        output_details = None
+
+                        if row[5]:  # input_tokens_details
+                            try:
+                                input_details = json.loads(row[5])
+                            except json.JSONDecodeError:
+                                pass
+
+                        if row[6]:  # output_tokens_details
+                            try:
+                                output_details = json.loads(row[6])
+                            except json.JSONDecodeError:
+                                pass
+
+                        results.append(
+                            {
+                                "user_turn_number": row[0],
+                                "requests": row[1],
+                                "input_tokens": row[2],
+                                "output_tokens": row[3],
+                                "total_tokens": row[4],
+                                "input_tokens_details": input_details,
+                                "output_tokens_details": output_details,
+                            }
+                        )
+                    return results
+
+        result = await asyncio.to_thread(_get_turn_usage_sync)
+
+        return cast(Union[list[dict[str, Any]], dict[str, Any]], result)
+
+    async def _update_turn_usage_internal(self, user_turn_number: int, usage_data: Usage) -> None:
+        """Internal method to update usage for a specific turn with full JSON details.
+
+        Args:
+            user_turn_number: The turn number to update usage for.
+            usage_data: The usage data to store.
+        """
+
+        def _update_sync():
+            """Synchronous helper to update turn usage data."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                # Serialize token details as JSON
+                input_details_json = None
+                output_details_json = None
+
+                if hasattr(usage_data, "input_tokens_details") and usage_data.input_tokens_details:
+                    try:
+                        input_details_json = json.dumps(usage_data.input_tokens_details.__dict__)
+                    except (TypeError, ValueError) as e:
+                        self._logger.warning(f"Failed to serialize input tokens details: {e}")
+                        input_details_json = None
+
+                    if (
+                        hasattr(usage_data, "output_tokens_details")
+                        and usage_data.output_tokens_details
+                    ):
+                        try:
+                            output_details_json = json.dumps(
+                                usage_data.output_tokens_details.__dict__
+                            )
+                        except (TypeError, ValueError) as e:
+                            self._logger.warning(f"Failed to serialize output tokens details: {e}")
+                            output_details_json = None
+
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(
+                        """
+                        INSERT OR REPLACE INTO turn_usage
+                        (session_id, branch_id, user_turn_number, requests, input_tokens, output_tokens,
+                         total_tokens, input_tokens_details, output_tokens_details)
+                        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    """,  # noqa: E501
+                        (
+                            self.session_id,
+                            self._current_branch_id,
+                            user_turn_number,
+                            usage_data.requests or 0,
+                            usage_data.input_tokens or 0,
+                            usage_data.output_tokens or 0,
+                            usage_data.total_tokens or 0,
+                            input_details_json,
+                            output_details_json,
+                        ),
+                    )
+                    conn.commit()
+
+        await asyncio.to_thread(_update_sync)

agents/extensions/memory/encrypt_session.py

@@ -0,0 +1,185 @@
+"""Encrypted Session wrapper for secure conversation storage.
+
+This module provides transparent encryption for session storage with automatic
+expiration of old data. When TTL expires, expired items are silently skipped.
+
+Usage::
+
+    from agents.extensions.memory import EncryptedSession, SQLAlchemySession
+
+    # Create underlying session (e.g. SQLAlchemySession)
+    underlying_session = SQLAlchemySession.from_url(
+        session_id="user-123",
+        url="postgresql+asyncpg://app:[email protected]/agents",
+        create_tables=True,
+    )
+
+    # Wrap with encryption and TTL-based expiration
+    session = EncryptedSession(
+        session_id="user-123",
+        underlying_session=underlying_session,
+        encryption_key="your-encryption-key",
+        ttl=600,  # 10 minutes
+    )
+
+    await Runner.run(agent, "Hello", session=session)
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+from typing import Any, cast
+
+from cryptography.fernet import Fernet, InvalidToken
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+from typing_extensions import Literal, TypedDict, TypeGuard
+
+from ...items import TResponseInputItem
+from ...memory.session import SessionABC
+
+
+class EncryptedEnvelope(TypedDict):
+    """TypedDict for encrypted message envelopes stored in the underlying session."""
+
+    __enc__: Literal[1]
+    v: int
+    kid: str
+    payload: str
+
+
+def _ensure_fernet_key_bytes(master_key: str) -> bytes:
+    """
+    Accept either a Fernet key (urlsafe-b64, 32 bytes after decode) or a raw string.
+    Returns raw bytes suitable for HKDF input.
+    """
+    if not master_key:
+        raise ValueError("encryption_key not set; required for EncryptedSession.")
+    try:
+        key_bytes = base64.urlsafe_b64decode(master_key)
+        if len(key_bytes) == 32:
+            return key_bytes
+    except Exception:
+        pass
+    return master_key.encode("utf-8")
+
+
+def _derive_session_fernet_key(master_key_bytes: bytes, session_id: str) -> Fernet:
+    hkdf = HKDF(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=session_id.encode("utf-8"),
+        info=b"agents.session-store.hkdf.v1",
+    )
+    derived = hkdf.derive(master_key_bytes)
+    return Fernet(base64.urlsafe_b64encode(derived))
+
+
+def _to_json_bytes(obj: Any) -> bytes:
+    return json.dumps(obj, ensure_ascii=False, separators=(",", ":"), default=str).encode("utf-8")
+
+
+def _from_json_bytes(data: bytes) -> Any:
+    return json.loads(data.decode("utf-8"))
+
+
+def _is_encrypted_envelope(item: object) -> TypeGuard[EncryptedEnvelope]:
+    """Type guard to check if an item is an encrypted envelope."""
+    return (
+        isinstance(item, dict)
+        and item.get("__enc__") == 1
+        and "payload" in item
+        and "kid" in item
+        and "v" in item
+    )
+
+
+class EncryptedSession(SessionABC):
+    """Encrypted wrapper for Session implementations with TTL-based expiration.
+
+    This class wraps any SessionABC implementation to provide transparent
+    encryption/decryption of stored items using Fernet encryption with
+    per-session key derivation and automatic expiration of old data.
+
+    When items expire (exceed TTL), they are silently skipped during retrieval.
+
+    Note: Expired tokens are rejected based on the system clock of the application server.
+    To avoid valid tokens being rejected due to clock drift, ensure all servers in
+    your environment are synchronized using NTP.
+    """
+
+    def __init__(
+        self,
+        session_id: str,
+        underlying_session: SessionABC,
+        encryption_key: str,
+        ttl: int = 600,
+    ):
+        """
+        Args:
+            session_id: ID for this session
+            underlying_session: The real session store (e.g. SQLiteSession, SQLAlchemySession)
+            encryption_key: Master key (Fernet key or raw secret)
+            ttl: Token time-to-live in seconds (default 10 min)
+        """
+        self.session_id = session_id
+        self.underlying_session = underlying_session
+        self.ttl = ttl
+
+        master = _ensure_fernet_key_bytes(encryption_key)
+        self.cipher = _derive_session_fernet_key(master, session_id)
+        self._kid = "hkdf-v1"
+        self._ver = 1
+
+    def __getattr__(self, name):
+        return getattr(self.underlying_session, name)
+
+    def _wrap(self, item: TResponseInputItem) -> EncryptedEnvelope:
+        if isinstance(item, dict):
+            payload = item
+        elif hasattr(item, "model_dump"):
+            payload = item.model_dump()
+        elif hasattr(item, "__dict__"):
+            payload = item.__dict__
+        else:
+            payload = dict(item)
+
+        token = self.cipher.encrypt(_to_json_bytes(payload)).decode("utf-8")
+        return {"__enc__": 1, "v": self._ver, "kid": self._kid, "payload": token}
+
+    def _unwrap(self, item: TResponseInputItem | EncryptedEnvelope) -> TResponseInputItem | None:
+        if not _is_encrypted_envelope(item):
+            return cast(TResponseInputItem, item)
+
+        try:
+            token = item["payload"].encode("utf-8")
+            plaintext = self.cipher.decrypt(token, ttl=self.ttl)
+            return cast(TResponseInputItem, _from_json_bytes(plaintext))
+        except (InvalidToken, KeyError):
+            return None
+
+    async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
+        encrypted_items = await self.underlying_session.get_items(limit)
+        valid_items: list[TResponseInputItem] = []
+        for enc in encrypted_items:
+            item = self._unwrap(enc)
+            if item is not None:
+                valid_items.append(item)
+        return valid_items
+
+    async def add_items(self, items: list[TResponseInputItem]) -> None:
+        wrapped: list[EncryptedEnvelope] = [self._wrap(it) for it in items]
+        await self.underlying_session.add_items(cast(list[TResponseInputItem], wrapped))
+
+    async def pop_item(self) -> TResponseInputItem | None:
+        while True:
+            enc = await self.underlying_session.pop_item()
+            if not enc:
+                return None
+            item = self._unwrap(enc)
+            if item is not None:
+                return item
+
+    async def clear_session(self) -> None:
+        await self.underlying_session.clear_session()

agents/extensions/memory/redis_session.py

@@ -0,0 +1,267 @@
+"""Redis-powered Session backend.
+
+Usage::
+
+    from agents.extensions.memory import RedisSession
+
+    # Create from Redis URL
+    session = RedisSession.from_url(
+        session_id="user-123",
+        url="redis://localhost:6379/0",
+    )
+
+    # Or pass an existing Redis client that your application already manages
+    session = RedisSession(
+        session_id="user-123",
+        redis_client=my_redis_client,
+    )
+
+    await Runner.run(agent, "Hello", session=session)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+from typing import Any
+from urllib.parse import urlparse
+
+try:
+    import redis.asyncio as redis
+    from redis.asyncio import Redis
+except ImportError as e:
+    raise ImportError(
+        "RedisSession requires the 'redis' package. Install it with: pip install redis"
+    ) from e
+
+from ...items import TResponseInputItem
+from ...memory.session import SessionABC
+
+
+class RedisSession(SessionABC):
+    """Redis implementation of :pyclass:`agents.memory.session.Session`."""
+
+    def __init__(
+        self,
+        session_id: str,
+        *,
+        redis_client: Redis,
+        key_prefix: str = "agents:session",
+        ttl: int | None = None,
+    ):
+        """Initializes a new RedisSession.
+
+        Args:
+            session_id (str): Unique identifier for the conversation.
+            redis_client (Redis[bytes]): A pre-configured Redis async client.
+            key_prefix (str, optional): Prefix for Redis keys to avoid collisions.
+                Defaults to "agents:session".
+            ttl (int | None, optional): Time-to-live in seconds for session data.
+                If None, data persists indefinitely. Defaults to None.
+        """
+        self.session_id = session_id
+        self._redis = redis_client
+        self._key_prefix = key_prefix
+        self._ttl = ttl
+        self._lock = asyncio.Lock()
+        self._owns_client = False  # Track if we own the Redis client
+
+        # Redis key patterns
+        self._session_key = f"{self._key_prefix}:{self.session_id}"
+        self._messages_key = f"{self._session_key}:messages"
+        self._counter_key = f"{self._session_key}:counter"
+
+    @classmethod
+    def from_url(
+        cls,
+        session_id: str,
+        *,
+        url: str,
+        redis_kwargs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> RedisSession:
+        """Create a session from a Redis URL string.
+
+        Args:
+            session_id (str): Conversation ID.
+            url (str): Redis URL, e.g. "redis://localhost:6379/0" or "rediss://host:6380".
+            redis_kwargs (dict[str, Any] | None): Additional keyword arguments forwarded to
+                redis.asyncio.from_url.
+            **kwargs: Additional keyword arguments forwarded to the main constructor
+                (e.g., key_prefix, ttl, etc.).
+
+        Returns:
+            RedisSession: An instance of RedisSession connected to the specified Redis server.
+        """
+        redis_kwargs = redis_kwargs or {}
+
+        # Parse URL to determine if we need SSL
+        parsed = urlparse(url)
+        if parsed.scheme == "rediss":
+            redis_kwargs.setdefault("ssl", True)
+
+        redis_client = redis.from_url(url, **redis_kwargs)
+        session = cls(session_id, redis_client=redis_client, **kwargs)
+        session._owns_client = True  # We created the client, so we own it
+        return session
+
+    async def _serialize_item(self, item: TResponseInputItem) -> str:
+        """Serialize an item to JSON string. Can be overridden by subclasses."""
+        return json.dumps(item, separators=(",", ":"))
+
+    async def _deserialize_item(self, item: str) -> TResponseInputItem:
+        """Deserialize a JSON string to an item. Can be overridden by subclasses."""
+        return json.loads(item)  # type: ignore[no-any-return]  # json.loads returns Any but we know the structure
+
+    async def _get_next_id(self) -> int:
+        """Get the next message ID using Redis INCR for atomic increment."""
+        result = await self._redis.incr(self._counter_key)
+        return int(result)
+
+    async def _set_ttl_if_configured(self, *keys: str) -> None:
+        """Set TTL on keys if configured."""
+        if self._ttl is not None:
+            pipe = self._redis.pipeline()
+            for key in keys:
+                pipe.expire(key, self._ttl)
+            await pipe.execute()
+
+    # ------------------------------------------------------------------
+    # Session protocol implementation
+    # ------------------------------------------------------------------
+
+    async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
+        """Retrieve the conversation history for this session.
+
+        Args:
+            limit: Maximum number of items to retrieve. If None, retrieves all items.
+                   When specified, returns the latest N items in chronological order.
+
+        Returns:
+            List of input items representing the conversation history
+        """
+        async with self._lock:
+            if limit is None:
+                # Get all messages in chronological order
+                raw_messages = await self._redis.lrange(self._messages_key, 0, -1)  # type: ignore[misc]  # Redis library returns Union[Awaitable[T], T] in async context
+            else:
+                if limit <= 0:
+                    return []
+                # Get the latest N messages (Redis list is ordered chronologically)
+                # Use negative indices to get from the end - Redis uses -N to -1 for last N items
+                raw_messages = await self._redis.lrange(self._messages_key, -limit, -1)  # type: ignore[misc]  # Redis library returns Union[Awaitable[T], T] in async context
+
+            items: list[TResponseInputItem] = []
+            for raw_msg in raw_messages:
+                try:
+                    # Handle both bytes (default) and str (decode_responses=True) Redis clients
+                    if isinstance(raw_msg, bytes):
+                        msg_str = raw_msg.decode("utf-8")
+                    else:
+                        msg_str = raw_msg  # Already a string
+                    item = await self._deserialize_item(msg_str)
+                    items.append(item)
+                except (json.JSONDecodeError, UnicodeDecodeError):
+                    # Skip corrupted messages
+                    continue
+
+            return items
+
+    async def add_items(self, items: list[TResponseInputItem]) -> None:
+        """Add new items to the conversation history.
+
+        Args:
+            items: List of input items to add to the history
+        """
+        if not items:
+            return
+
+        async with self._lock:
+            pipe = self._redis.pipeline()
+
+            # Set session metadata with current timestamp
+            pipe.hset(
+                self._session_key,
+                mapping={
+                    "session_id": self.session_id,
+                    "created_at": str(int(time.time())),
+                    "updated_at": str(int(time.time())),
+                },
+            )
+
+            # Add all items to the messages list
+            serialized_items = []
+            for item in items:
+                serialized = await self._serialize_item(item)
+                serialized_items.append(serialized)
+
+            if serialized_items:
+                pipe.rpush(self._messages_key, *serialized_items)
+
+            # Update the session timestamp
+            pipe.hset(self._session_key, "updated_at", str(int(time.time())))
+
+            # Execute all commands
+            await pipe.execute()
+
+            # Set TTL if configured
+            await self._set_ttl_if_configured(
+                self._session_key, self._messages_key, self._counter_key
+            )
+
+    async def pop_item(self) -> TResponseInputItem | None:
+        """Remove and return the most recent item from the session.
+
+        Returns:
+            The most recent item if it exists, None if the session is empty
+        """
+        async with self._lock:
+            # Use RPOP to atomically remove and return the rightmost (most recent) item
+            raw_msg = await self._redis.rpop(self._messages_key)  # type: ignore[misc]  # Redis library returns Union[Awaitable[T], T] in async context
+
+            if raw_msg is None:
+                return None
+
+            try:
+                # Handle both bytes (default) and str (decode_responses=True) Redis clients
+                if isinstance(raw_msg, bytes):
+                    msg_str = raw_msg.decode("utf-8")
+                else:
+                    msg_str = raw_msg  # Already a string
+                return await self._deserialize_item(msg_str)
+            except (json.JSONDecodeError, UnicodeDecodeError):
+                # Return None for corrupted messages (already removed)
+                return None
+
+    async def clear_session(self) -> None:
+        """Clear all items for this session."""
+        async with self._lock:
+            # Delete all keys associated with this session
+            await self._redis.delete(
+                self._session_key,
+                self._messages_key,
+                self._counter_key,
+            )
+
+    async def close(self) -> None:
+        """Close the Redis connection.
+
+        Only closes the connection if this session owns the Redis client
+        (i.e., created via from_url). If the client was injected externally,
+        the caller is responsible for managing its lifecycle.
+        """
+        if self._owns_client:
+            await self._redis.aclose()
+
+    async def ping(self) -> bool:
+        """Test Redis connectivity.
+
+        Returns:
+            True if Redis is reachable, False otherwise.
+        """
+        try:
+            await self._redis.ping()
+            return True
+        except Exception:
+            return False

agents/extensions/memory/sqlalchemy_session.py

@@ -0,0 +1,312 @@
+"""SQLAlchemy-powered Session backend.
+
+Usage::
+
+    from agents.extensions.memory import SQLAlchemySession
+
+    # Create from SQLAlchemy URL (uses asyncpg driver under the hood for Postgres)
+    session = SQLAlchemySession.from_url(
+        session_id="user-123",
+        url="postgresql+asyncpg://app:[email protected]/agents",
+        create_tables=True, # If you want to auto-create tables, set to True.
+    )
+
+    # Or pass an existing AsyncEngine that your application already manages
+    session = SQLAlchemySession(
+        session_id="user-123",
+        engine=my_async_engine,
+        create_tables=True, # If you want to auto-create tables, set to True.
+    )
+
+    await Runner.run(agent, "Hello", session=session)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Any
+
+from sqlalchemy import (
+    TIMESTAMP,
+    Column,
+    ForeignKey,
+    Index,
+    Integer,
+    MetaData,
+    String,
+    Table,
+    Text,
+    delete,
+    insert,
+    select,
+    text as sql_text,
+    update,
+)
+from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine
+
+from ...items import TResponseInputItem
+from ...memory.session import SessionABC
+
+
+class SQLAlchemySession(SessionABC):
+    """SQLAlchemy implementation of :pyclass:`agents.memory.session.Session`."""
+
+    _metadata: MetaData
+    _sessions: Table
+    _messages: Table
+
+    def __init__(
+        self,
+        session_id: str,
+        *,
+        engine: AsyncEngine,
+        create_tables: bool = False,
+        sessions_table: str = "agent_sessions",
+        messages_table: str = "agent_messages",
+    ):
+        """Initializes a new SQLAlchemySession.
+
+        Args:
+            session_id (str): Unique identifier for the conversation.
+            engine (AsyncEngine): A pre-configured SQLAlchemy async engine. The engine
+                must be created with an async driver (e.g., 'postgresql+asyncpg://',
+                'mysql+aiomysql://', or 'sqlite+aiosqlite://').
+            create_tables (bool, optional): Whether to automatically create the required
+                tables and indexes. Defaults to False for production use. Set to True for
+                development and testing when migrations aren't used.
+            sessions_table (str, optional): Override the default table name for sessions if needed.
+            messages_table (str, optional): Override the default table name for messages if needed.
+        """
+        self.session_id = session_id
+        self._engine = engine
+        self._lock = asyncio.Lock()
+
+        self._metadata = MetaData()
+        self._sessions = Table(
+            sessions_table,
+            self._metadata,
+            Column("session_id", String, primary_key=True),
+            Column(
+                "created_at",
+                TIMESTAMP(timezone=False),
+                server_default=sql_text("CURRENT_TIMESTAMP"),
+                nullable=False,
+            ),
+            Column(
+                "updated_at",
+                TIMESTAMP(timezone=False),
+                server_default=sql_text("CURRENT_TIMESTAMP"),
+                onupdate=sql_text("CURRENT_TIMESTAMP"),
+                nullable=False,
+            ),
+        )
+
+        self._messages = Table(
+            messages_table,
+            self._metadata,
+            Column("id", Integer, primary_key=True, autoincrement=True),
+            Column(
+                "session_id",
+                String,
+                ForeignKey(f"{sessions_table}.session_id", ondelete="CASCADE"),
+                nullable=False,
+            ),
+            Column("message_data", Text, nullable=False),
+            Column(
+                "created_at",
+                TIMESTAMP(timezone=False),
+                server_default=sql_text("CURRENT_TIMESTAMP"),
+                nullable=False,
+            ),
+            Index(
+                f"idx_{messages_table}_session_time",
+                "session_id",
+                "created_at",
+            ),
+            sqlite_autoincrement=True,
+        )
+
+        # Async session factory
+        self._session_factory = async_sessionmaker(self._engine, expire_on_commit=False)
+
+        self._create_tables = create_tables
+
+    # ---------------------------------------------------------------------
+    # Convenience constructors
+    # ---------------------------------------------------------------------
+    @classmethod
+    def from_url(
+        cls,
+        session_id: str,
+        *,
+        url: str,
+        engine_kwargs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> SQLAlchemySession:
+        """Create a session from a database URL string.
+
+        Args:
+            session_id (str): Conversation ID.
+            url (str): Any SQLAlchemy async URL, e.g. "postgresql+asyncpg://user:pass@host/db".
+            engine_kwargs (dict[str, Any] | None): Additional keyword arguments forwarded to
+                sqlalchemy.ext.asyncio.create_async_engine.
+            **kwargs: Additional keyword arguments forwarded to the main constructor
+                (e.g., create_tables, custom table names, etc.).
+
+        Returns:
+            SQLAlchemySession: An instance of SQLAlchemySession connected to the specified database.
+        """
+        engine_kwargs = engine_kwargs or {}
+        engine = create_async_engine(url, **engine_kwargs)
+        return cls(session_id, engine=engine, **kwargs)
+
+    async def _serialize_item(self, item: TResponseInputItem) -> str:
+        """Serialize an item to JSON string. Can be overridden by subclasses."""
+        return json.dumps(item, separators=(",", ":"))
+
+    async def _deserialize_item(self, item: str) -> TResponseInputItem:
+        """Deserialize a JSON string to an item. Can be overridden by subclasses."""
+        return json.loads(item)  # type: ignore[no-any-return]
+
+    # ------------------------------------------------------------------
+    # Session protocol implementation
+    # ------------------------------------------------------------------
+    async def _ensure_tables(self) -> None:
+        """Ensure tables are created before any database operations."""
+        if self._create_tables:
+            async with self._engine.begin() as conn:
+                await conn.run_sync(self._metadata.create_all)
+            self._create_tables = False  # Only create once
+
+    async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
+        """Retrieve the conversation history for this session.
+
+        Args:
+            limit: Maximum number of items to retrieve. If None, retrieves all items.
+                   When specified, returns the latest N items in chronological order.
+
+        Returns:
+            List of input items representing the conversation history
+        """
+        await self._ensure_tables()
+        async with self._session_factory() as sess:
+            if limit is None:
+                stmt = (
+                    select(self._messages.c.message_data)
+                    .where(self._messages.c.session_id == self.session_id)
+                    .order_by(self._messages.c.created_at.asc())
+                )
+            else:
+                stmt = (
+                    select(self._messages.c.message_data)
+                    .where(self._messages.c.session_id == self.session_id)
+                    # Use DESC + LIMIT to get the latest N
+                    # then reverse later for chronological order.
+                    .order_by(self._messages.c.created_at.desc())
+                    .limit(limit)
+                )
+
+            result = await sess.execute(stmt)
+            rows: list[str] = [row[0] for row in result.all()]
+
+            if limit is not None:
+                rows.reverse()
+
+            items: list[TResponseInputItem] = []
+            for raw in rows:
+                try:
+                    items.append(await self._deserialize_item(raw))
+                except json.JSONDecodeError:
+                    # Skip corrupted rows
+                    continue
+            return items
+
+    async def add_items(self, items: list[TResponseInputItem]) -> None:
+        """Add new items to the conversation history.
+
+        Args:
+            items: List of input items to add to the history
+        """
+        if not items:
+            return
+
+        await self._ensure_tables()
+        payload = [
+            {
+                "session_id": self.session_id,
+                "message_data": await self._serialize_item(item),
+            }
+            for item in items
+        ]
+
+        async with self._session_factory() as sess:
+            async with sess.begin():
+                # Ensure the parent session row exists - use merge for cross-DB compatibility
+                # Check if session exists
+                existing = await sess.execute(
+                    select(self._sessions.c.session_id).where(
+                        self._sessions.c.session_id == self.session_id
+                    )
+                )
+                if not existing.scalar_one_or_none():
+                    # Session doesn't exist, create it
+                    await sess.execute(
+                        insert(self._sessions).values({"session_id": self.session_id})
+                    )
+
+                # Insert messages in bulk
+                await sess.execute(insert(self._messages), payload)
+
+                # Touch updated_at column
+                await sess.execute(
+                    update(self._sessions)
+                    .where(self._sessions.c.session_id == self.session_id)
+                    .values(updated_at=sql_text("CURRENT_TIMESTAMP"))
+                )
+
+    async def pop_item(self) -> TResponseInputItem | None:
+        """Remove and return the most recent item from the session.
+
+        Returns:
+            The most recent item if it exists, None if the session is empty
+        """
+        await self._ensure_tables()
+        async with self._session_factory() as sess:
+            async with sess.begin():
+                # Fallback for all dialects - get ID first, then delete
+                subq = (
+                    select(self._messages.c.id)
+                    .where(self._messages.c.session_id == self.session_id)
+                    .order_by(self._messages.c.created_at.desc())
+                    .limit(1)
+                )
+                res = await sess.execute(subq)
+                row_id = res.scalar_one_or_none()
+                if row_id is None:
+                    return None
+                # Fetch data before deleting
+                res_data = await sess.execute(
+                    select(self._messages.c.message_data).where(self._messages.c.id == row_id)
+                )
+                row = res_data.scalar_one_or_none()
+                await sess.execute(delete(self._messages).where(self._messages.c.id == row_id))
+
+                if row is None:
+                    return None
+                try:
+                    return await self._deserialize_item(row)
+                except json.JSONDecodeError:
+                    return None
+
+    async def clear_session(self) -> None:
+        """Clear all items for this session."""
+        await self._ensure_tables()
+        async with self._session_factory() as sess:
+            async with sess.begin():
+                await sess.execute(
+                    delete(self._messages).where(self._messages.c.session_id == self.session_id)
+                )
+                await sess.execute(
+                    delete(self._sessions).where(self._sessions.c.session_id == self.session_id)
+                )

agents/extensions/models/litellm_model.py

@@ -0,0 +1,601 @@
+from __future__ import annotations
+
+import json
+import time
+from collections.abc import AsyncIterator
+from copy import copy
+from typing import Any, Literal, cast, overload
+
+from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails
+
+from agents.exceptions import ModelBehaviorError
+
+try:
+    import litellm
+except ImportError as _e:
+    raise ImportError(
+        "`litellm` is required to use the LitellmModel. You can install it via the optional "
+        "dependency group: `pip install 'openai-agents[litellm]'`."
+    ) from _e
+
+from openai import NOT_GIVEN, AsyncStream, NotGiven
+from openai.types.chat import (
+    ChatCompletionChunk,
+    ChatCompletionMessageCustomToolCall,
+    ChatCompletionMessageFunctionToolCall,
+    ChatCompletionMessageParam,
+)
+from openai.types.chat.chat_completion_message import (
+    Annotation,
+    AnnotationURLCitation,
+    ChatCompletionMessage,
+)
+from openai.types.chat.chat_completion_message_function_tool_call import Function
+from openai.types.responses import Response
+
+from ... import _debug
+from ...agent_output import AgentOutputSchemaBase
+from ...handoffs import Handoff
+from ...items import ModelResponse, TResponseInputItem, TResponseStreamEvent
+from ...logger import logger
+from ...model_settings import ModelSettings
+from ...models.chatcmpl_converter import Converter
+from ...models.chatcmpl_helpers import HEADERS, HEADERS_OVERRIDE
+from ...models.chatcmpl_stream_handler import ChatCmplStreamHandler
+from ...models.fake_id import FAKE_RESPONSES_ID
+from ...models.interface import Model, ModelTracing
+from ...tool import Tool
+from ...tracing import generation_span
+from ...tracing.span_data import GenerationSpanData
+from ...tracing.spans import Span
+from ...usage import Usage
+from ...util._json import _to_dump_compatible
+
+
+class InternalChatCompletionMessage(ChatCompletionMessage):
+    """
+    An internal subclass to carry reasoning_content and thinking_blocks without modifying the original model.
+    """  # noqa: E501
+
+    reasoning_content: str
+    thinking_blocks: list[dict[str, Any]] | None = None
+
+
+class LitellmModel(Model):
+    """This class enables using any model via LiteLLM. LiteLLM allows you to acess OpenAPI,
+    Anthropic, Gemini, Mistral, and many other models.
+    See supported models here: [litellm models](https://docs.litellm.ai/docs/providers).
+    """
+
+    def __init__(
+        self,
+        model: str,
+        base_url: str | None = None,
+        api_key: str | None = None,
+    ):
+        self.model = model
+        self.base_url = base_url
+        self.api_key = api_key
+
+    async def get_response(
+        self,
+        system_instructions: str | None,
+        input: str | list[TResponseInputItem],
+        model_settings: ModelSettings,
+        tools: list[Tool],
+        output_schema: AgentOutputSchemaBase | None,
+        handoffs: list[Handoff],
+        tracing: ModelTracing,
+        previous_response_id: str | None = None,  # unused
+        conversation_id: str | None = None,  # unused
+        prompt: Any | None = None,
+    ) -> ModelResponse:
+        with generation_span(
+            model=str(self.model),
+            model_config=model_settings.to_json_dict()
+            | {"base_url": str(self.base_url or ""), "model_impl": "litellm"},
+            disabled=tracing.is_disabled(),
+        ) as span_generation:
+            response = await self._fetch_response(
+                system_instructions,
+                input,
+                model_settings,
+                tools,
+                output_schema,
+                handoffs,
+                span_generation,
+                tracing,
+                stream=False,
+                prompt=prompt,
+            )
+
+            assert isinstance(response.choices[0], litellm.types.utils.Choices)
+
+            if _debug.DONT_LOG_MODEL_DATA:
+                logger.debug("Received model response")
+            else:
+                logger.debug(
+                    f"""LLM resp:\n{
+                        json.dumps(
+                            response.choices[0].message.model_dump(), indent=2, ensure_ascii=False
+                        )
+                    }\n"""
+                )
+
+            if hasattr(response, "usage"):
+                response_usage = response.usage
+                usage = (
+                    Usage(
+                        requests=1,
+                        input_tokens=response_usage.prompt_tokens,
+                        output_tokens=response_usage.completion_tokens,
+                        total_tokens=response_usage.total_tokens,
+                        input_tokens_details=InputTokensDetails(
+                            cached_tokens=getattr(
+                                response_usage.prompt_tokens_details, "cached_tokens", 0
+                            )
+                            or 0
+                        ),
+                        output_tokens_details=OutputTokensDetails(
+                            reasoning_tokens=getattr(
+                                response_usage.completion_tokens_details, "reasoning_tokens", 0
+                            )
+                            or 0
+                        ),
+                    )
+                    if response.usage
+                    else Usage()
+                )
+            else:
+                usage = Usage()
+                logger.warning("No usage information returned from Litellm")
+
+            if tracing.include_data():
+                span_generation.span_data.output = [response.choices[0].message.model_dump()]
+            span_generation.span_data.usage = {
+                "input_tokens": usage.input_tokens,
+                "output_tokens": usage.output_tokens,
+            }
+
+            items = Converter.message_to_output_items(
+                LitellmConverter.convert_message_to_openai(response.choices[0].message)
+            )
+
+            return ModelResponse(
+                output=items,
+                usage=usage,
+                response_id=None,
+            )
+
+    async def stream_response(
+        self,
+        system_instructions: str | None,
+        input: str | list[TResponseInputItem],
+        model_settings: ModelSettings,
+        tools: list[Tool],
+        output_schema: AgentOutputSchemaBase | None,
+        handoffs: list[Handoff],
+        tracing: ModelTracing,
+        previous_response_id: str | None = None,  # unused
+        conversation_id: str | None = None,  # unused
+        prompt: Any | None = None,
+    ) -> AsyncIterator[TResponseStreamEvent]:
+        with generation_span(
+            model=str(self.model),
+            model_config=model_settings.to_json_dict()
+            | {"base_url": str(self.base_url or ""), "model_impl": "litellm"},
+            disabled=tracing.is_disabled(),
+        ) as span_generation:
+            response, stream = await self._fetch_response(
+                system_instructions,
+                input,
+                model_settings,
+                tools,
+                output_schema,
+                handoffs,
+                span_generation,
+                tracing,
+                stream=True,
+                prompt=prompt,
+            )
+
+            final_response: Response | None = None
+            async for chunk in ChatCmplStreamHandler.handle_stream(response, stream):
+                yield chunk
+
+                if chunk.type == "response.completed":
+                    final_response = chunk.response
+
+            if tracing.include_data() and final_response:
+                span_generation.span_data.output = [final_response.model_dump()]
+
+            if final_response and final_response.usage:
+                span_generation.span_data.usage = {
+                    "input_tokens": final_response.usage.input_tokens,
+                    "output_tokens": final_response.usage.output_tokens,
+                }
+
+    @overload
+    async def _fetch_response(
+        self,
+        system_instructions: str | None,
+        input: str | list[TResponseInputItem],
+        model_settings: ModelSettings,
+        tools: list[Tool],
+        output_schema: AgentOutputSchemaBase | None,
+        handoffs: list[Handoff],
+        span: Span[GenerationSpanData],
+        tracing: ModelTracing,
+        stream: Literal[True],
+        prompt: Any | None = None,
+    ) -> tuple[Response, AsyncStream[ChatCompletionChunk]]: ...
+
+    @overload
+    async def _fetch_response(
+        self,
+        system_instructions: str | None,
+        input: str | list[TResponseInputItem],
+        model_settings: ModelSettings,
+        tools: list[Tool],
+        output_schema: AgentOutputSchemaBase | None,
+        handoffs: list[Handoff],
+        span: Span[GenerationSpanData],
+        tracing: ModelTracing,
+        stream: Literal[False],
+        prompt: Any | None = None,
+    ) -> litellm.types.utils.ModelResponse: ...
+
+    async def _fetch_response(
+        self,
+        system_instructions: str | None,
+        input: str | list[TResponseInputItem],
+        model_settings: ModelSettings,
+        tools: list[Tool],
+        output_schema: AgentOutputSchemaBase | None,
+        handoffs: list[Handoff],
+        span: Span[GenerationSpanData],
+        tracing: ModelTracing,
+        stream: bool = False,
+        prompt: Any | None = None,
+    ) -> litellm.types.utils.ModelResponse | tuple[Response, AsyncStream[ChatCompletionChunk]]:
+        # Preserve reasoning messages for tool calls when reasoning is on
+        # This is needed for models like Claude 4 Sonnet/Opus which support interleaved thinking
+        preserve_thinking_blocks = (
+            model_settings.reasoning is not None and model_settings.reasoning.effort is not None
+        )
+
+        converted_messages = Converter.items_to_messages(
+            input, preserve_thinking_blocks=preserve_thinking_blocks
+        )
+
+        # Fix for interleaved thinking bug: reorder messages to ensure tool_use comes before tool_result  # noqa: E501
+        if preserve_thinking_blocks:
+            converted_messages = self._fix_tool_message_ordering(converted_messages)
+
+        if system_instructions:
+            converted_messages.insert(
+                0,
+                {
+                    "content": system_instructions,
+                    "role": "system",
+                },
+            )
+        converted_messages = _to_dump_compatible(converted_messages)
+
+        if tracing.include_data():
+            span.span_data.input = converted_messages
+
+        parallel_tool_calls = (
+            True
+            if model_settings.parallel_tool_calls and tools and len(tools) > 0
+            else False
+            if model_settings.parallel_tool_calls is False
+            else None
+        )
+        tool_choice = Converter.convert_tool_choice(model_settings.tool_choice)
+        response_format = Converter.convert_response_format(output_schema)
+
+        converted_tools = [Converter.tool_to_openai(tool) for tool in tools] if tools else []
+
+        for handoff in handoffs:
+            converted_tools.append(Converter.convert_handoff_tool(handoff))
+
+        converted_tools = _to_dump_compatible(converted_tools)
+
+        if _debug.DONT_LOG_MODEL_DATA:
+            logger.debug("Calling LLM")
+        else:
+            messages_json = json.dumps(
+                converted_messages,
+                indent=2,
+                ensure_ascii=False,
+            )
+            tools_json = json.dumps(
+                converted_tools,
+                indent=2,
+                ensure_ascii=False,
+            )
+            logger.debug(
+                f"Calling Litellm model: {self.model}\n"
+                f"{messages_json}\n"
+                f"Tools:\n{tools_json}\n"
+                f"Stream: {stream}\n"
+                f"Tool choice: {tool_choice}\n"
+                f"Response format: {response_format}\n"
+            )
+
+        reasoning_effort = model_settings.reasoning.effort if model_settings.reasoning else None
+
+        stream_options = None
+        if stream and model_settings.include_usage is not None:
+            stream_options = {"include_usage": model_settings.include_usage}
+
+        extra_kwargs = {}
+        if model_settings.extra_query:
+            extra_kwargs["extra_query"] = copy(model_settings.extra_query)
+        if model_settings.metadata:
+            extra_kwargs["metadata"] = copy(model_settings.metadata)
+        if model_settings.extra_body and isinstance(model_settings.extra_body, dict):
+            extra_kwargs.update(model_settings.extra_body)
+
+        # Add kwargs from model_settings.extra_args, filtering out None values
+        if model_settings.extra_args:
+            extra_kwargs.update(model_settings.extra_args)
+
+        ret = await litellm.acompletion(
+            model=self.model,
+            messages=converted_messages,
+            tools=converted_tools or None,
+            temperature=model_settings.temperature,
+            top_p=model_settings.top_p,
+            frequency_penalty=model_settings.frequency_penalty,
+            presence_penalty=model_settings.presence_penalty,
+            max_tokens=model_settings.max_tokens,
+            tool_choice=self._remove_not_given(tool_choice),
+            response_format=self._remove_not_given(response_format),
+            parallel_tool_calls=parallel_tool_calls,
+            stream=stream,
+            stream_options=stream_options,
+            reasoning_effort=reasoning_effort,
+            top_logprobs=model_settings.top_logprobs,
+            extra_headers=self._merge_headers(model_settings),
+            api_key=self.api_key,
+            base_url=self.base_url,
+            **extra_kwargs,
+        )
+
+        if isinstance(ret, litellm.types.utils.ModelResponse):
+            return ret
+
+        response = Response(
+            id=FAKE_RESPONSES_ID,
+            created_at=time.time(),
+            model=self.model,
+            object="response",
+            output=[],
+            tool_choice=cast(Literal["auto", "required", "none"], tool_choice)
+            if tool_choice != NOT_GIVEN
+            else "auto",
+            top_p=model_settings.top_p,
+            temperature=model_settings.temperature,
+            tools=[],
+            parallel_tool_calls=parallel_tool_calls or False,
+            reasoning=model_settings.reasoning,
+        )
+        return response, ret
+
+    def _fix_tool_message_ordering(
+        self, messages: list[ChatCompletionMessageParam]
+    ) -> list[ChatCompletionMessageParam]:
+        """
+        Fix the ordering of tool messages to ensure tool_use messages come before tool_result messages.
+
+        This addresses the interleaved thinking bug where conversation histories may contain
+        tool results before their corresponding tool calls, causing Anthropic API to reject the request.
+        """  # noqa: E501
+        if not messages:
+            return messages
+
+        # Collect all tool calls and tool results
+        tool_call_messages = {}  # tool_id -> (index, message)
+        tool_result_messages = {}  # tool_id -> (index, message)
+        other_messages = []  # (index, message) for non-tool messages
+
+        for i, message in enumerate(messages):
+            if not isinstance(message, dict):
+                other_messages.append((i, message))
+                continue
+
+            role = message.get("role")
+
+            if role == "assistant" and message.get("tool_calls"):
+                # Extract tool calls from this assistant message
+                tool_calls = message.get("tool_calls", [])
+                if isinstance(tool_calls, list):
+                    for tool_call in tool_calls:
+                        if isinstance(tool_call, dict):
+                            tool_id = tool_call.get("id")
+                            if tool_id:
+                                # Create a separate assistant message for each tool call
+                                single_tool_msg = cast(dict[str, Any], message.copy())
+                                single_tool_msg["tool_calls"] = [tool_call]
+                                tool_call_messages[tool_id] = (
+                                    i,
+                                    cast(ChatCompletionMessageParam, single_tool_msg),
+                                )
+
+            elif role == "tool":
+                tool_call_id = message.get("tool_call_id")
+                if tool_call_id:
+                    tool_result_messages[tool_call_id] = (i, message)
+                else:
+                    other_messages.append((i, message))
+            else:
+                other_messages.append((i, message))
+
+        # First, identify which tool results will be paired to avoid duplicates
+        paired_tool_result_indices = set()
+        for tool_id in tool_call_messages:
+            if tool_id in tool_result_messages:
+                tool_result_idx, _ = tool_result_messages[tool_id]
+                paired_tool_result_indices.add(tool_result_idx)
+
+        # Create the fixed message sequence
+        fixed_messages: list[ChatCompletionMessageParam] = []
+        used_indices = set()
+
+        # Add messages in their original order, but ensure tool_use → tool_result pairing
+        for i, original_message in enumerate(messages):
+            if i in used_indices:
+                continue
+
+            if not isinstance(original_message, dict):
+                fixed_messages.append(original_message)
+                used_indices.add(i)
+                continue
+
+            role = original_message.get("role")
+
+            if role == "assistant" and original_message.get("tool_calls"):
+                # Process each tool call in this assistant message
+                tool_calls = original_message.get("tool_calls", [])
+                if isinstance(tool_calls, list):
+                    for tool_call in tool_calls:
+                        if isinstance(tool_call, dict):
+                            tool_id = tool_call.get("id")
+                            if (
+                                tool_id
+                                and tool_id in tool_call_messages
+                                and tool_id in tool_result_messages
+                            ):
+                                # Add tool_use → tool_result pair
+                                _, tool_call_msg = tool_call_messages[tool_id]
+                                tool_result_idx, tool_result_msg = tool_result_messages[tool_id]
+
+                                fixed_messages.append(tool_call_msg)
+                                fixed_messages.append(tool_result_msg)
+
+                                # Mark both as used
+                                used_indices.add(tool_call_messages[tool_id][0])
+                                used_indices.add(tool_result_idx)
+                            elif tool_id and tool_id in tool_call_messages:
+                                # Tool call without result - add just the tool call
+                                _, tool_call_msg = tool_call_messages[tool_id]
+                                fixed_messages.append(tool_call_msg)
+                                used_indices.add(tool_call_messages[tool_id][0])
+
+                used_indices.add(i)  # Mark original multi-tool message as used
+
+            elif role == "tool":
+                # Only preserve unmatched tool results to avoid duplicates
+                if i not in paired_tool_result_indices:
+                    fixed_messages.append(original_message)
+                used_indices.add(i)
+
+            else:
+                # Regular message - add it normally
+                fixed_messages.append(original_message)
+                used_indices.add(i)
+
+        return fixed_messages
+
+    def _remove_not_given(self, value: Any) -> Any:
+        if isinstance(value, NotGiven):
+            return None
+        return value
+
+    def _merge_headers(self, model_settings: ModelSettings):
+        return {**HEADERS, **(model_settings.extra_headers or {}), **(HEADERS_OVERRIDE.get() or {})}
+
+
+class LitellmConverter:
+    @classmethod
+    def convert_message_to_openai(
+        cls, message: litellm.types.utils.Message
+    ) -> ChatCompletionMessage:
+        if message.role != "assistant":
+            raise ModelBehaviorError(f"Unsupported role: {message.role}")
+
+        tool_calls: (
+            list[ChatCompletionMessageFunctionToolCall | ChatCompletionMessageCustomToolCall] | None
+        ) = (
+            [LitellmConverter.convert_tool_call_to_openai(tool) for tool in message.tool_calls]
+            if message.tool_calls
+            else None
+        )
+
+        provider_specific_fields = message.get("provider_specific_fields", None)
+        refusal = (
+            provider_specific_fields.get("refusal", None) if provider_specific_fields else None
+        )
+
+        reasoning_content = ""
+        if hasattr(message, "reasoning_content") and message.reasoning_content:
+            reasoning_content = message.reasoning_content
+
+        # Extract full thinking blocks including signatures (for Anthropic)
+        thinking_blocks: list[dict[str, Any]] | None = None
+        if hasattr(message, "thinking_blocks") and message.thinking_blocks:
+            # Convert thinking blocks to dict format for compatibility
+            thinking_blocks = []
+            for block in message.thinking_blocks:
+                if isinstance(block, dict):
+                    thinking_blocks.append(cast(dict[str, Any], block))
+                else:
+                    # Convert object to dict by accessing its attributes
+                    block_dict: dict[str, Any] = {}
+                    if hasattr(block, "__dict__"):
+                        block_dict = dict(block.__dict__.items())
+                    elif hasattr(block, "model_dump"):
+                        block_dict = block.model_dump()
+                    else:
+                        # Last resort: convert to string representation
+                        block_dict = {"thinking": str(block)}
+                    thinking_blocks.append(block_dict)
+
+        return InternalChatCompletionMessage(
+            content=message.content,
+            refusal=refusal,
+            role="assistant",
+            annotations=cls.convert_annotations_to_openai(message),
+            audio=message.get("audio", None),  # litellm deletes audio if not present
+            tool_calls=tool_calls,
+            reasoning_content=reasoning_content,
+            thinking_blocks=thinking_blocks,
+        )
+
+    @classmethod
+    def convert_annotations_to_openai(
+        cls, message: litellm.types.utils.Message
+    ) -> list[Annotation] | None:
+        annotations: list[litellm.types.llms.openai.ChatCompletionAnnotation] | None = message.get(
+            "annotations", None
+        )
+        if not annotations:
+            return None
+
+        return [
+            Annotation(
+                type="url_citation",
+                url_citation=AnnotationURLCitation(
+                    start_index=annotation["url_citation"]["start_index"],
+                    end_index=annotation["url_citation"]["end_index"],
+                    url=annotation["url_citation"]["url"],
+                    title=annotation["url_citation"]["title"],
+                ),
+            )
+            for annotation in annotations
+        ]
+
+    @classmethod
+    def convert_tool_call_to_openai(
+        cls, tool_call: litellm.types.utils.ChatCompletionMessageToolCall
+    ) -> ChatCompletionMessageFunctionToolCall:
+        return ChatCompletionMessageFunctionToolCall(
+            id=tool_call.id,
+            type="function",
+            function=Function(
+                name=tool_call.function.name or "",
+                arguments=tool_call.function.arguments,
+            ),
+        )

agents/extensions/models/litellm_provider.py

@@ -0,0 +1,23 @@
+from ...models.default_models import get_default_model
+from ...models.interface import Model, ModelProvider
+from .litellm_model import LitellmModel
+
+# This is kept for backward compatiblity but using get_default_model() method is recommended.
+DEFAULT_MODEL: str = "gpt-4.1"
+
+
+class LitellmProvider(ModelProvider):
+    """A ModelProvider that uses LiteLLM to route to any model provider. You can use it via:
+    ```python
+    Runner.run(agent, input, run_config=RunConfig(model_provider=LitellmProvider()))
+    ```
+    See supported models here: [litellm models](https://docs.litellm.ai/docs/providers).
+
+    NOTE: API keys must be set via environment variables. If you're using models that require
+    additional configuration (e.g. Azure API base or version), those must also be set via the
+    environment variables that LiteLLM expects. If you have more advanced needs, we recommend
+    copy-pasting this class and making any modifications you need.
+    """
+
+    def get_model(self, model_name: str | None) -> Model:
+        return LitellmModel(model_name or get_default_model())

agents/extensions/visualization.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import graphviz  # type: ignore
+
+from agents import Agent
+from agents.handoffs import Handoff
+from agents.tool import Tool
+
+
+def get_main_graph(agent: Agent) -> str:
+    """
+    Generates the main graph structure in DOT format for the given agent.
+
+    Args:
+        agent (Agent): The agent for which the graph is to be generated.
+
+    Returns:
+        str: The DOT format string representing the graph.
+    """
+    parts = [
+        """
+    digraph G {
+        graph [splines=true];
+        node [fontname="Arial"];
+        edge [penwidth=1.5];
+    """
+    ]
+    parts.append(get_all_nodes(agent))
+    parts.append(get_all_edges(agent))
+    parts.append("}")
+    return "".join(parts)
+
+
+def get_all_nodes(
+    agent: Agent, parent: Agent | None = None, visited: set[str] | None = None
+) -> str:
+    """
+    Recursively generates the nodes for the given agent and its handoffs in DOT format.
+
+    Args:
+        agent (Agent): The agent for which the nodes are to be generated.
+
+    Returns:
+        str: The DOT format string representing the nodes.
+    """
+    if visited is None:
+        visited = set()
+    if agent.name in visited:
+        return ""
+    visited.add(agent.name)
+
+    parts = []
+
+    # Start and end the graph
+    if not parent:
+        parts.append(
+            '"__start__" [label="__start__", shape=ellipse, style=filled, '
+            "fillcolor=lightblue, width=0.5, height=0.3];"
+            '"__end__" [label="__end__", shape=ellipse, style=filled, '
+            "fillcolor=lightblue, width=0.5, height=0.3];"
+        )
+        # Ensure parent agent node is colored
+        parts.append(
+            f'"{agent.name}" [label="{agent.name}", shape=box, style=filled, '
+            "fillcolor=lightyellow, width=1.5, height=0.8];"
+        )
+
+    for tool in agent.tools:
+        parts.append(
+            f'"{tool.name}" [label="{tool.name}", shape=ellipse, style=filled, '
+            f"fillcolor=lightgreen, width=0.5, height=0.3];"
+        )
+
+    for mcp_server in agent.mcp_servers:
+        parts.append(
+            f'"{mcp_server.name}" [label="{mcp_server.name}", shape=box, style=filled, '
+            f"fillcolor=lightgrey, width=1, height=0.5];"
+        )
+
+    for handoff in agent.handoffs:
+        if isinstance(handoff, Handoff):
+            parts.append(
+                f'"{handoff.agent_name}" [label="{handoff.agent_name}", '
+                f"shape=box, style=filled, style=rounded, "
+                f"fillcolor=lightyellow, width=1.5, height=0.8];"
+            )
+        if isinstance(handoff, Agent):
+            if handoff.name not in visited:
+                parts.append(
+                    f'"{handoff.name}" [label="{handoff.name}", '
+                    f"shape=box, style=filled, style=rounded, "
+                    f"fillcolor=lightyellow, width=1.5, height=0.8];"
+                )
+            parts.append(get_all_nodes(handoff, agent, visited))
+
+    return "".join(parts)
+
+
+def get_all_edges(
+    agent: Agent, parent: Agent | None = None, visited: set[str] | None = None
+) -> str:
+    """
+    Recursively generates the edges for the given agent and its handoffs in DOT format.
+
+    Args:
+        agent (Agent): The agent for which the edges are to be generated.
+        parent (Agent, optional): The parent agent. Defaults to None.
+
+    Returns:
+        str: The DOT format string representing the edges.
+    """
+    if visited is None:
+        visited = set()
+    if agent.name in visited:
+        return ""
+    visited.add(agent.name)
+
+    parts = []
+
+    if not parent:
+        parts.append(f'"__start__" -> "{agent.name}";')
+
+    for tool in agent.tools:
+        parts.append(f"""
+        "{agent.name}" -> "{tool.name}" [style=dotted, penwidth=1.5];
+        "{tool.name}" -> "{agent.name}" [style=dotted, penwidth=1.5];""")
+
+    for mcp_server in agent.mcp_servers:
+        parts.append(f"""
+        "{agent.name}" -> "{mcp_server.name}" [style=dashed, penwidth=1.5];
+        "{mcp_server.name}" -> "{agent.name}" [style=dashed, penwidth=1.5];""")
+
+    for handoff in agent.handoffs:
+        if isinstance(handoff, Handoff):
+            parts.append(f"""
+            "{agent.name}" -> "{handoff.agent_name}";""")
+        if isinstance(handoff, Agent):
+            parts.append(f"""
+            "{agent.name}" -> "{handoff.name}";""")
+            parts.append(get_all_edges(handoff, agent, visited))
+
+    if not agent.handoffs and not isinstance(agent, Tool):  # type: ignore
+        parts.append(f'"{agent.name}" -> "__end__";')
+
+    return "".join(parts)
+
+
+def draw_graph(agent: Agent, filename: str | None = None) -> graphviz.Source:
+    """
+    Draws the graph for the given agent and optionally saves it as a PNG file.
+
+    Args:
+        agent (Agent): The agent for which the graph is to be drawn.
+        filename (str): The name of the file to save the graph as a PNG.
+
+    Returns:
+        graphviz.Source: The graphviz Source object representing the graph.
+    """
+    dot_code = get_main_graph(agent)
+    graph = graphviz.Source(dot_code)
+
+    if filename:
+        graph.render(filename, format="png", cleanup=True)
+
+    return graph

agents/function_schema.py

@@ -0,0 +1,398 @@
+from __future__ import annotations
+
+import contextlib
+import inspect
+import logging
+import re
+from dataclasses import dataclass
+from typing import Annotated, Any, Callable, Literal, get_args, get_origin, get_type_hints
+
+from griffe import Docstring, DocstringSectionKind
+from pydantic import BaseModel, Field, create_model
+from pydantic.fields import FieldInfo
+
+from .exceptions import UserError
+from .run_context import RunContextWrapper
+from .strict_schema import ensure_strict_json_schema
+from .tool_context import ToolContext
+
+
+@dataclass
+class FuncSchema:
+    """
+    Captures the schema for a python function, in preparation for sending it to an LLM as a tool.
+    """
+
+    name: str
+    """The name of the function."""
+    description: str | None
+    """The description of the function."""
+    params_pydantic_model: type[BaseModel]
+    """A Pydantic model that represents the function's parameters."""
+    params_json_schema: dict[str, Any]
+    """The JSON schema for the function's parameters, derived from the Pydantic model."""
+    signature: inspect.Signature
+    """The signature of the function."""
+    takes_context: bool = False
+    """Whether the function takes a RunContextWrapper argument (must be the first argument)."""
+    strict_json_schema: bool = True
+    """Whether the JSON schema is in strict mode. We **strongly** recommend setting this to True,
+    as it increases the likelihood of correct JSON input."""
+
+    def to_call_args(self, data: BaseModel) -> tuple[list[Any], dict[str, Any]]:
+        """
+        Converts validated data from the Pydantic model into (args, kwargs), suitable for calling
+        the original function.
+        """
+        positional_args: list[Any] = []
+        keyword_args: dict[str, Any] = {}
+        seen_var_positional = False
+
+        # Use enumerate() so we can skip the first parameter if it's context.
+        for idx, (name, param) in enumerate(self.signature.parameters.items()):
+            # If the function takes a RunContextWrapper and this is the first parameter, skip it.
+            if self.takes_context and idx == 0:
+                continue
+
+            value = getattr(data, name, None)
+            if param.kind == param.VAR_POSITIONAL:
+                # e.g. *args: extend positional args and mark that *args is now seen
+                positional_args.extend(value or [])
+                seen_var_positional = True
+            elif param.kind == param.VAR_KEYWORD:
+                # e.g. **kwargs handling
+                keyword_args.update(value or {})
+            elif param.kind in (param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD):
+                # Before *args, add to positional args. After *args, add to keyword args.
+                if not seen_var_positional:
+                    positional_args.append(value)
+                else:
+                    keyword_args[name] = value
+            else:
+                # For KEYWORD_ONLY parameters, always use keyword args.
+                keyword_args[name] = value
+        return positional_args, keyword_args
+
+
+@dataclass
+class FuncDocumentation:
+    """Contains metadata about a Python function, extracted from its docstring."""
+
+    name: str
+    """The name of the function, via `__name__`."""
+    description: str | None
+    """The description of the function, derived from the docstring."""
+    param_descriptions: dict[str, str] | None
+    """The parameter descriptions of the function, derived from the docstring."""
+
+
+DocstringStyle = Literal["google", "numpy", "sphinx"]
+
+
+# As of Feb 2025, the automatic style detection in griffe is an Insiders feature. This
+# code approximates it.
+def _detect_docstring_style(doc: str) -> DocstringStyle:
+    scores: dict[DocstringStyle, int] = {"sphinx": 0, "numpy": 0, "google": 0}
+
+    # Sphinx style detection: look for :param, :type, :return:, and :rtype:
+    sphinx_patterns = [r"^:param\s", r"^:type\s", r"^:return:", r"^:rtype:"]
+    for pattern in sphinx_patterns:
+        if re.search(pattern, doc, re.MULTILINE):
+            scores["sphinx"] += 1
+
+    # Numpy style detection: look for headers like 'Parameters', 'Returns', or 'Yields' followed by
+    # a dashed underline
+    numpy_patterns = [
+        r"^Parameters\s*\n\s*-{3,}",
+        r"^Returns\s*\n\s*-{3,}",
+        r"^Yields\s*\n\s*-{3,}",
+    ]
+    for pattern in numpy_patterns:
+        if re.search(pattern, doc, re.MULTILINE):
+            scores["numpy"] += 1
+
+    # Google style detection: look for section headers with a trailing colon
+    google_patterns = [r"^(Args|Arguments):", r"^(Returns):", r"^(Raises):"]
+    for pattern in google_patterns:
+        if re.search(pattern, doc, re.MULTILINE):
+            scores["google"] += 1
+
+    max_score = max(scores.values())
+    if max_score == 0:
+        return "google"
+
+    # Priority order: sphinx > numpy > google in case of tie
+    styles: list[DocstringStyle] = ["sphinx", "numpy", "google"]
+
+    for style in styles:
+        if scores[style] == max_score:
+            return style
+
+    return "google"
+
+
+@contextlib.contextmanager
+def _suppress_griffe_logging():
+    # Suppresses warnings about missing annotations for params
+    logger = logging.getLogger("griffe")
+    previous_level = logger.getEffectiveLevel()
+    logger.setLevel(logging.ERROR)
+    try:
+        yield
+    finally:
+        logger.setLevel(previous_level)
+
+
+def generate_func_documentation(
+    func: Callable[..., Any], style: DocstringStyle | None = None
+) -> FuncDocumentation:
+    """
+    Extracts metadata from a function docstring, in preparation for sending it to an LLM as a tool.
+
+    Args:
+        func: The function to extract documentation from.
+        style: The style of the docstring to use for parsing. If not provided, we will attempt to
+            auto-detect the style.
+
+    Returns:
+        A FuncDocumentation object containing the function's name, description, and parameter
+        descriptions.
+    """
+    name = func.__name__
+    doc = inspect.getdoc(func)
+    if not doc:
+        return FuncDocumentation(name=name, description=None, param_descriptions=None)
+
+    with _suppress_griffe_logging():
+        docstring = Docstring(doc, lineno=1, parser=style or _detect_docstring_style(doc))
+        parsed = docstring.parse()
+
+    description: str | None = next(
+        (section.value for section in parsed if section.kind == DocstringSectionKind.text), None
+    )
+
+    param_descriptions: dict[str, str] = {
+        param.name: param.description
+        for section in parsed
+        if section.kind == DocstringSectionKind.parameters
+        for param in section.value
+    }
+
+    return FuncDocumentation(
+        name=func.__name__,
+        description=description,
+        param_descriptions=param_descriptions or None,
+    )
+
+
+def _strip_annotated(annotation: Any) -> tuple[Any, tuple[Any, ...]]:
+    """Returns the underlying annotation and any metadata from typing.Annotated."""
+
+    metadata: tuple[Any, ...] = ()
+    ann = annotation
+
+    while get_origin(ann) is Annotated:
+        args = get_args(ann)
+        if not args:
+            break
+        ann = args[0]
+        metadata = (*metadata, *args[1:])
+
+    return ann, metadata
+
+
+def _extract_description_from_metadata(metadata: tuple[Any, ...]) -> str | None:
+    """Extracts a human readable description from Annotated metadata if present."""
+
+    for item in metadata:
+        if isinstance(item, str):
+            return item
+    return None
+
+
+def function_schema(
+    func: Callable[..., Any],
+    docstring_style: DocstringStyle | None = None,
+    name_override: str | None = None,
+    description_override: str | None = None,
+    use_docstring_info: bool = True,
+    strict_json_schema: bool = True,
+) -> FuncSchema:
+    """
+    Given a Python function, extracts a `FuncSchema` from it, capturing the name, description,
+    parameter descriptions, and other metadata.
+
+    Args:
+        func: The function to extract the schema from.
+        docstring_style: The style of the docstring to use for parsing. If not provided, we will
+            attempt to auto-detect the style.
+        name_override: If provided, use this name instead of the function's `__name__`.
+        description_override: If provided, use this description instead of the one derived from the
+            docstring.
+        use_docstring_info: If True, uses the docstring to generate the description and parameter
+            descriptions.
+        strict_json_schema: Whether the JSON schema is in strict mode. If True, we'll ensure that
+            the schema adheres to the "strict" standard the OpenAI API expects. We **strongly**
+            recommend setting this to True, as it increases the likelihood of the LLM producing
+            correct JSON input.
+
+    Returns:
+        A `FuncSchema` object containing the function's name, description, parameter descriptions,
+        and other metadata.
+    """
+
+    # 1. Grab docstring info
+    if use_docstring_info:
+        doc_info = generate_func_documentation(func, docstring_style)
+        param_descs = dict(doc_info.param_descriptions or {})
+    else:
+        doc_info = None
+        param_descs = {}
+
+    type_hints_with_extras = get_type_hints(func, include_extras=True)
+    type_hints: dict[str, Any] = {}
+    annotated_param_descs: dict[str, str] = {}
+
+    for name, annotation in type_hints_with_extras.items():
+        if name == "return":
+            continue
+
+        stripped_ann, metadata = _strip_annotated(annotation)
+        type_hints[name] = stripped_ann
+
+        description = _extract_description_from_metadata(metadata)
+        if description is not None:
+            annotated_param_descs[name] = description
+
+    for name, description in annotated_param_descs.items():
+        param_descs.setdefault(name, description)
+
+    # Ensure name_override takes precedence even if docstring info is disabled.
+    func_name = name_override or (doc_info.name if doc_info else func.__name__)
+
+    # 2. Inspect function signature and get type hints
+    sig = inspect.signature(func)
+    params = list(sig.parameters.items())
+    takes_context = False
+    filtered_params = []
+
+    if params:
+        first_name, first_param = params[0]
+        # Prefer the evaluated type hint if available
+        ann = type_hints.get(first_name, first_param.annotation)
+        if ann != inspect._empty:
+            origin = get_origin(ann) or ann
+            if origin is RunContextWrapper or origin is ToolContext:
+                takes_context = True  # Mark that the function takes context
+            else:
+                filtered_params.append((first_name, first_param))
+        else:
+            filtered_params.append((first_name, first_param))
+
+    # For parameters other than the first, raise error if any use RunContextWrapper or ToolContext.
+    for name, param in params[1:]:
+        ann = type_hints.get(name, param.annotation)
+        if ann != inspect._empty:
+            origin = get_origin(ann) or ann
+            if origin is RunContextWrapper or origin is ToolContext:
+                raise UserError(
+                    f"RunContextWrapper/ToolContext param found at non-first position in function"
+                    f" {func.__name__}"
+                )
+        filtered_params.append((name, param))
+
+    # We will collect field definitions for create_model as a dict:
+    #   field_name -> (type_annotation, default_value_or_Field(...))
+    fields: dict[str, Any] = {}
+
+    for name, param in filtered_params:
+        ann = type_hints.get(name, param.annotation)
+        default = param.default
+
+        # If there's no type hint, assume `Any`
+        if ann == inspect._empty:
+            ann = Any
+
+        # If a docstring param description exists, use it
+        field_description = param_descs.get(name, None)
+
+        # Handle different parameter kinds
+        if param.kind == param.VAR_POSITIONAL:
+            # e.g. *args: extend positional args
+            if get_origin(ann) is tuple:
+                # e.g. def foo(*args: tuple[int, ...]) -> treat as List[int]
+                args_of_tuple = get_args(ann)
+                if len(args_of_tuple) == 2 and args_of_tuple[1] is Ellipsis:
+                    ann = list[args_of_tuple[0]]  # type: ignore
+                else:
+                    ann = list[Any]
+            else:
+                # If user wrote *args: int, treat as List[int]
+                ann = list[ann]  # type: ignore
+
+            # Default factory to empty list
+            fields[name] = (
+                ann,
+                Field(default_factory=list, description=field_description),
+            )
+
+        elif param.kind == param.VAR_KEYWORD:
+            # **kwargs handling
+            if get_origin(ann) is dict:
+                # e.g. def foo(**kwargs: dict[str, int])
+                dict_args = get_args(ann)
+                if len(dict_args) == 2:
+                    ann = dict[dict_args[0], dict_args[1]]  # type: ignore
+                else:
+                    ann = dict[str, Any]
+            else:
+                # e.g. def foo(**kwargs: int) -> Dict[str, int]
+                ann = dict[str, ann]  # type: ignore
+
+            fields[name] = (
+                ann,
+                Field(default_factory=dict, description=field_description),
+            )
+
+        else:
+            # Normal parameter
+            if default == inspect._empty:
+                # Required field
+                fields[name] = (
+                    ann,
+                    Field(..., description=field_description),
+                )
+            elif isinstance(default, FieldInfo):
+                # Parameter with a default value that is a Field(...)
+                fields[name] = (
+                    ann,
+                    FieldInfo.merge_field_infos(
+                        default, description=field_description or default.description
+                    ),
+                )
+            else:
+                # Parameter with a default value
+                fields[name] = (
+                    ann,
+                    Field(default=default, description=field_description),
+                )
+
+    # 3. Dynamically build a Pydantic model
+    dynamic_model = create_model(f"{func_name}_args", __base__=BaseModel, **fields)
+
+    # 4. Build JSON schema from that model
+    json_schema = dynamic_model.model_json_schema()
+    if strict_json_schema:
+        json_schema = ensure_strict_json_schema(json_schema)
+
+    # 5. Return as a FuncSchema dataclass
+    return FuncSchema(
+        name=func_name,
+        # Ensure description_override takes precedence even if docstring info is disabled.
+        description=description_override or (doc_info.description if doc_info else None),
+        params_pydantic_model=dynamic_model,
+        params_json_schema=json_schema,
+        signature=sig,
+        takes_context=takes_context,
+        strict_json_schema=strict_json_schema,
+    )

agents/guardrail.py

@@ -0,0 +1,329 @@
+from __future__ import annotations
+
+import inspect
+from collections.abc import Awaitable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Generic, Union, overload
+
+from typing_extensions import TypeVar
+
+from .exceptions import UserError
+from .items import TResponseInputItem
+from .run_context import RunContextWrapper, TContext
+from .util._types import MaybeAwaitable
+
+if TYPE_CHECKING:
+    from .agent import Agent
+
+
+@dataclass
+class GuardrailFunctionOutput:
+    """The output of a guardrail function."""
+
+    output_info: Any
+    """
+    Optional information about the guardrail's output. For example, the guardrail could include
+    information about the checks it performed and granular results.
+    """
+
+    tripwire_triggered: bool
+    """
+    Whether the tripwire was triggered. If triggered, the agent's execution will be halted.
+    """
+
+
+@dataclass
+class InputGuardrailResult:
+    """The result of a guardrail run."""
+
+    guardrail: InputGuardrail[Any]
+    """
+    The guardrail that was run.
+    """
+
+    output: GuardrailFunctionOutput
+    """The output of the guardrail function."""
+
+
+@dataclass
+class OutputGuardrailResult:
+    """The result of a guardrail run."""
+
+    guardrail: OutputGuardrail[Any]
+    """
+    The guardrail that was run.
+    """
+
+    agent_output: Any
+    """
+    The output of the agent that was checked by the guardrail.
+    """
+
+    agent: Agent[Any]
+    """
+    The agent that was checked by the guardrail.
+    """
+
+    output: GuardrailFunctionOutput
+    """The output of the guardrail function."""
+
+
+@dataclass
+class InputGuardrail(Generic[TContext]):
+    """Input guardrails are checks that run in parallel to the agent's execution.
+    They can be used to do things like:
+    - Check if input messages are off-topic
+    - Take over control of the agent's execution if an unexpected input is detected
+
+    You can use the `@input_guardrail()` decorator to turn a function into an `InputGuardrail`, or
+    create an `InputGuardrail` manually.
+
+    Guardrails return a `GuardrailResult`. If `result.tripwire_triggered` is `True`,
+    the agent's execution will immediately stop, and
+    an `InputGuardrailTripwireTriggered` exception will be raised
+    """
+
+    guardrail_function: Callable[
+        [RunContextWrapper[TContext], Agent[Any], str | list[TResponseInputItem]],
+        MaybeAwaitable[GuardrailFunctionOutput],
+    ]
+    """A function that receives the agent input and the context, and returns a
+     `GuardrailResult`. The result marks whether the tripwire was triggered, and can optionally
+     include information about the guardrail's output.
+    """
+
+    name: str | None = None
+    """The name of the guardrail, used for tracing. If not provided, we'll use the guardrail
+    function's name.
+    """
+
+    def get_name(self) -> str:
+        if self.name:
+            return self.name
+
+        return self.guardrail_function.__name__
+
+    async def run(
+        self,
+        agent: Agent[Any],
+        input: str | list[TResponseInputItem],
+        context: RunContextWrapper[TContext],
+    ) -> InputGuardrailResult:
+        if not callable(self.guardrail_function):
+            raise UserError(f"Guardrail function must be callable, got {self.guardrail_function}")
+
+        output = self.guardrail_function(context, agent, input)
+        if inspect.isawaitable(output):
+            return InputGuardrailResult(
+                guardrail=self,
+                output=await output,
+            )
+
+        return InputGuardrailResult(
+            guardrail=self,
+            output=output,
+        )
+
+
+@dataclass
+class OutputGuardrail(Generic[TContext]):
+    """Output guardrails are checks that run on the final output of an agent.
+    They can be used to do check if the output passes certain validation criteria
+
+    You can use the `@output_guardrail()` decorator to turn a function into an `OutputGuardrail`,
+    or create an `OutputGuardrail` manually.
+
+    Guardrails return a `GuardrailResult`. If `result.tripwire_triggered` is `True`, an
+    `OutputGuardrailTripwireTriggered` exception will be raised.
+    """
+
+    guardrail_function: Callable[
+        [RunContextWrapper[TContext], Agent[Any], Any],
+        MaybeAwaitable[GuardrailFunctionOutput],
+    ]
+    """A function that receives the final agent, its output, and the context, and returns a
+     `GuardrailResult`. The result marks whether the tripwire was triggered, and can optionally
+     include information about the guardrail's output.
+    """
+
+    name: str | None = None
+    """The name of the guardrail, used for tracing. If not provided, we'll use the guardrail
+    function's name.
+    """
+
+    def get_name(self) -> str:
+        if self.name:
+            return self.name
+
+        return self.guardrail_function.__name__
+
+    async def run(
+        self, context: RunContextWrapper[TContext], agent: Agent[Any], agent_output: Any
+    ) -> OutputGuardrailResult:
+        if not callable(self.guardrail_function):
+            raise UserError(f"Guardrail function must be callable, got {self.guardrail_function}")
+
+        output = self.guardrail_function(context, agent, agent_output)
+        if inspect.isawaitable(output):
+            return OutputGuardrailResult(
+                guardrail=self,
+                agent=agent,
+                agent_output=agent_output,
+                output=await output,
+            )
+
+        return OutputGuardrailResult(
+            guardrail=self,
+            agent=agent,
+            agent_output=agent_output,
+            output=output,
+        )
+
+
+TContext_co = TypeVar("TContext_co", bound=Any, covariant=True)
+
+# For InputGuardrail
+_InputGuardrailFuncSync = Callable[
+    [RunContextWrapper[TContext_co], "Agent[Any]", Union[str, list[TResponseInputItem]]],
+    GuardrailFunctionOutput,
+]
+_InputGuardrailFuncAsync = Callable[
+    [RunContextWrapper[TContext_co], "Agent[Any]", Union[str, list[TResponseInputItem]]],
+    Awaitable[GuardrailFunctionOutput],
+]
+
+
+@overload
+def input_guardrail(
+    func: _InputGuardrailFuncSync[TContext_co],
+) -> InputGuardrail[TContext_co]: ...
+
+
+@overload
+def input_guardrail(
+    func: _InputGuardrailFuncAsync[TContext_co],
+) -> InputGuardrail[TContext_co]: ...
+
+
+@overload
+def input_guardrail(
+    *,
+    name: str | None = None,
+) -> Callable[
+    [_InputGuardrailFuncSync[TContext_co] | _InputGuardrailFuncAsync[TContext_co]],
+    InputGuardrail[TContext_co],
+]: ...
+
+
+def input_guardrail(
+    func: _InputGuardrailFuncSync[TContext_co]
+    | _InputGuardrailFuncAsync[TContext_co]
+    | None = None,
+    *,
+    name: str | None = None,
+) -> (
+    InputGuardrail[TContext_co]
+    | Callable[
+        [_InputGuardrailFuncSync[TContext_co] | _InputGuardrailFuncAsync[TContext_co]],
+        InputGuardrail[TContext_co],
+    ]
+):
+    """
+    Decorator that transforms a sync or async function into an `InputGuardrail`.
+    It can be used directly (no parentheses) or with keyword args, e.g.:
+
+        @input_guardrail
+        def my_sync_guardrail(...): ...
+
+        @input_guardrail(name="guardrail_name")
+        async def my_async_guardrail(...): ...
+    """
+
+    def decorator(
+        f: _InputGuardrailFuncSync[TContext_co] | _InputGuardrailFuncAsync[TContext_co],
+    ) -> InputGuardrail[TContext_co]:
+        return InputGuardrail(
+            guardrail_function=f,
+            # If not set, guardrail name uses the function’s name by default.
+            name=name if name else f.__name__,
+        )
+
+    if func is not None:
+        # Decorator was used without parentheses
+        return decorator(func)
+
+    # Decorator used with keyword arguments
+    return decorator
+
+
+_OutputGuardrailFuncSync = Callable[
+    [RunContextWrapper[TContext_co], "Agent[Any]", Any],
+    GuardrailFunctionOutput,
+]
+_OutputGuardrailFuncAsync = Callable[
+    [RunContextWrapper[TContext_co], "Agent[Any]", Any],
+    Awaitable[GuardrailFunctionOutput],
+]
+
+
+@overload
+def output_guardrail(
+    func: _OutputGuardrailFuncSync[TContext_co],
+) -> OutputGuardrail[TContext_co]: ...
+
+
+@overload
+def output_guardrail(
+    func: _OutputGuardrailFuncAsync[TContext_co],
+) -> OutputGuardrail[TContext_co]: ...
+
+
+@overload
+def output_guardrail(
+    *,
+    name: str | None = None,
+) -> Callable[
+    [_OutputGuardrailFuncSync[TContext_co] | _OutputGuardrailFuncAsync[TContext_co]],
+    OutputGuardrail[TContext_co],
+]: ...
+
+
+def output_guardrail(
+    func: _OutputGuardrailFuncSync[TContext_co]
+    | _OutputGuardrailFuncAsync[TContext_co]
+    | None = None,
+    *,
+    name: str | None = None,
+) -> (
+    OutputGuardrail[TContext_co]
+    | Callable[
+        [_OutputGuardrailFuncSync[TContext_co] | _OutputGuardrailFuncAsync[TContext_co]],
+        OutputGuardrail[TContext_co],
+    ]
+):
+    """
+    Decorator that transforms a sync or async function into an `OutputGuardrail`.
+    It can be used directly (no parentheses) or with keyword args, e.g.:
+
+        @output_guardrail
+        def my_sync_guardrail(...): ...
+
+        @output_guardrail(name="guardrail_name")
+        async def my_async_guardrail(...): ...
+    """
+
+    def decorator(
+        f: _OutputGuardrailFuncSync[TContext_co] | _OutputGuardrailFuncAsync[TContext_co],
+    ) -> OutputGuardrail[TContext_co]:
+        return OutputGuardrail(
+            guardrail_function=f,
+            # Guardrail name defaults to function's name when not specified (None).
+            name=name if name else f.__name__,
+        )
+
+    if func is not None:
+        # Decorator was used without parentheses
+        return decorator(func)
+
+    # Decorator used with keyword arguments
+    return decorator