openai--openai-agents-python/tests/test_example_workflows.py

from __future__ import annotations

import asyncio
import json
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal, cast
from unittest.mock import AsyncMock

import pytest
from openai.types.responses import ResponseTextDeltaEvent
from pydantic import BaseModel

from agents import (
    Agent,
    AgentBase,
    AgentToolStreamEvent,
    AgentUpdatedStreamEvent,
    GuardrailFunctionOutput,
    InputGuardrailTripwireTriggered,
    ItemHelpers,
    ModelSettings,
    OutputGuardrailTripwireTriggered,
    RawResponsesStreamEvent,
    RunContextWrapper,
    Runner,
    input_guardrail,
    output_guardrail,
)
from agents.agent import ToolsToFinalOutputResult
from agents.items import TResponseInputItem
from agents.tool import FunctionToolResult, function_tool
from examples.financial_research_agent.agents.verifier_agent import (
    VerificationIssue,
    VerificationResult,
)
from examples.financial_research_agent.agents.writer_agent import FinancialReportData
from examples.financial_research_agent.manager import (
    FinancialResearchManager,
    FinancialSearchEvidence,
    FinancialSource,
    _extract_financial_sources,
)
from examples.sandbox.basic import _import_docker_from_env
from examples.sandbox.docker.docker_runner import (
    _format_tool_call,
    _format_tool_output,
)
from examples.sandbox.sandbox_agents_as_tools import (
    PricingPacketReview,
    RolloutRiskReview,
    _structured_tool_output_extractor,
)
from examples.tools.web_search_filters import _normalized_source_urls
from examples.web_search_utils import extract_url_citations, extract_web_search_source_urls

from .fake_model import FakeModel
from .test_responses import (
    get_final_output_message,
    get_function_tool_call,
    get_handoff_tool_call,
    get_text_input_item,
    get_text_message,
)


def test_web_search_source_urls_reject_decoded_reserved_delimiters() -> None:
    assert (
        _normalized_source_urls(
            ["https://developers.openai.com/api/docs/models/finding-the-right-model%3F.pls"]
        )
        == []
    )


def test_web_search_source_urls_are_canonical_and_domain_scoped() -> None:
    assert _normalized_source_urls(
        [
            "https://developers.openai.com/api/docs/models/gpt-5.6-sol?utm_source=openai",
            "https://developers.openai.com/api/docs/models/gpt-5.6-sol#pricing",
            "https://subdomain.developers.openai.com/api/docs/models/gpt-5.6-terra/",
            "https://developers.openai.com/assets/logo.svg",
            "https://user@developers.openai.com/api/docs/models/gpt-5.6-sol",
            "https://example.com/api/docs/models/gpt-5.6-sol",
        ]
    ) == [
        "https://developers.openai.com/api/docs/models/gpt-5.6-sol",
        "https://subdomain.developers.openai.com/api/docs/models/gpt-5.6-terra",
    ]


def test_web_search_metadata_distinguishes_citations_from_retrieved_sources() -> None:
    items = [
        {
            "raw_item": {
                "type": "web_search_call",
                "action": {
                    "type": "search",
                    "sources": [
                        {
                            "type": "url",
                            "url": "https://developers.openai.com/api/docs/models/gpt-5.6-sol",
                        },
                        {
                            "type": "url",
                            "url": "https://developers.openai.com/api/docs/models/gpt-5.6-terra",
                        },
                    ],
                },
            }
        },
        {
            "raw_item": {
                "type": "message",
                "content": [
                    {
                        "type": "output_text",
                        "text": "Use Sol for the most demanding work.",
                        "annotations": [
                            {
                                "type": "url_citation",
                                "title": "GPT-5.6 Sol",
                                "url": (
                                    "https://developers.openai.com/api/docs/models/gpt-5.6-sol"
                                ),
                            }
                        ],
                    }
                ],
            }
        },
    ]

    assert extract_web_search_source_urls(items) == [
        "https://developers.openai.com/api/docs/models/gpt-5.6-sol",
        "https://developers.openai.com/api/docs/models/gpt-5.6-terra",
    ]
    assert [(citation.title, citation.url) for citation in extract_url_citations(items)] == [
        (
            "GPT-5.6 Sol",
            "https://developers.openai.com/api/docs/models/gpt-5.6-sol",
        )
    ]


def test_financial_search_evidence_preserves_citations_and_retrieved_sources() -> None:
    sources = _extract_financial_sources(
        [
            {
                "raw_item": {
                    "type": "message",
                    "content": [
                        {
                            "type": "output_text",
                            "annotations": [
                                {
                                    "type": "url_citation",
                                    "title": "Annual report",
                                    "url": "https://example.com/annual-report",
                                }
                            ],
                        }
                    ],
                }
            },
            {
                "raw_item": {
                    "type": "web_search_call",
                    "action": {
                        "sources": [
                            {"type": "url", "url": "https://example.com/annual-report"},
                            {"type": "url", "url": "https://example.com/earnings"},
                        ]
                    },
                }
            },
        ]
    )

    assert sources == [
        FinancialSource(title="Annual report", url="https://example.com/annual-report"),
        FinancialSource(
            title="https://example.com/earnings",
            url="https://example.com/earnings",
        ),
    ]


@pytest.mark.asyncio
async def test_financial_report_revises_once_after_failed_verification(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    manager = object.__new__(FinancialResearchManager)
    original_report = FinancialReportData(
        short_summary="Original",
        markdown_report="Unsupported claim",
        follow_up_questions=[],
    )
    revised_report = FinancialReportData(
        short_summary="Revised",
        markdown_report="Supported claim",
        follow_up_questions=[],
    )
    rejected = VerificationResult(
        verified=False,
        issues=[
            VerificationIssue(
                claim="Unsupported claim",
                category="unsupported",
                explanation="No supplied evidence supports it.",
                source_urls=[],
            )
        ],
    )
    accepted = VerificationResult(verified=True, issues=[])
    write_report = AsyncMock(return_value=original_report)
    verify_report = AsyncMock(side_effect=[rejected, accepted])
    revise_report = AsyncMock(return_value=revised_report)
    monkeypatch.setattr(manager, "_write_report", write_report)
    monkeypatch.setattr(manager, "_verify_report", verify_report)
    monkeypatch.setattr(manager, "_revise_report", revise_report)

    report, verification = await manager._produce_verified_report("query", [])

    assert report == revised_report
    assert verification == accepted
    write_report.assert_awaited_once_with("query", [])
    revise_report.assert_awaited_once_with("query", original_report, [], rejected)
    assert verify_report.await_count == 2


@pytest.mark.asyncio
async def test_financial_report_fails_after_second_rejected_verification(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    manager = object.__new__(FinancialResearchManager)
    report = FinancialReportData(
        short_summary="Summary",
        markdown_report="Unsupported claim",
        follow_up_questions=[],
    )
    rejected = VerificationResult(
        verified=False,
        issues=[
            VerificationIssue(
                claim="Unsupported claim",
                category="unsupported",
                explanation="No supplied evidence supports it.",
                source_urls=[],
            )
        ],
    )
    monkeypatch.setattr(manager, "_write_report", AsyncMock(return_value=report))
    monkeypatch.setattr(manager, "_verify_report", AsyncMock(return_value=rejected))
    monkeypatch.setattr(manager, "_revise_report", AsyncMock(return_value=report))

    with pytest.raises(RuntimeError, match="failed evidence verification after one revision"):
        await manager._produce_verified_report("query", [])


def test_financial_report_input_includes_cutoff_and_evidence() -> None:
    manager = object.__new__(FinancialResearchManager)
    manager.research_cutoff = "2026-07-11"
    evidence = FinancialSearchEvidence(
        query="company annual report",
        reason="Ground annual metrics",
        summary="Revenue increased.",
        sources=[FinancialSource(title="Annual report", url="https://example.com/report")],
        retrieved_at="2026-07-11",
    )

    payload = json.loads(manager._report_input("Analyze the company", [evidence]))

    assert payload == {
        "original_query": "Analyze the company",
        "research_cutoff": "2026-07-11",
        "evidence": [evidence.model_dump(mode="json")],
    }


def test_sandbox_basic_direct_run_imports_external_docker_sdk(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    sdk_dir = tmp_path / "sdk"
    docker_package = sdk_dir / "docker"
    docker_package.mkdir(parents=True)
    docker_package.joinpath("__init__.py").write_text(
        "def from_env():\n    return 'external docker sdk'\n"
    )

    script_dir = Path("examples/sandbox").resolve()
    monkeypatch.setattr(sys, "path", [str(script_dir), str(sdk_dir)])
    for module_name in list(sys.modules):
        if module_name == "docker" or module_name.startswith("docker."):
            monkeypatch.delitem(sys.modules, module_name, raising=False)

    docker_from_env = _import_docker_from_env()

    assert docker_from_env() == "external docker sdk"
    assert sys.path == [str(script_dir), str(sdk_dir)]


@dataclass
class EvaluationFeedback:
    feedback: str
    score: Literal["pass", "needs_improvement"]


@dataclass
class OutlineCheckerOutput:
    good_quality: bool
    is_scifi: bool


@pytest.mark.asyncio
async def test_llm_as_judge_loop_handles_dataclass_feedback() -> None:
    """Mimics the llm_as_a_judge example: loop until the evaluator passes the outline."""
    outline_model = FakeModel()
    outline_model.add_multiple_turn_outputs(
        [
            [get_text_message("Outline v1")],
            [get_text_message("Outline v2")],
        ]
    )

    judge_model = FakeModel()
    judge_model.add_multiple_turn_outputs(
        [
            [
                get_final_output_message(
                    json.dumps(
                        {
                            "response": {
                                "feedback": "Add more suspense",
                                "score": "needs_improvement",
                            }
                        }
                    )
                )
            ],
            [
                get_final_output_message(
                    json.dumps({"response": {"feedback": "Looks good", "score": "pass"}})
                )
            ],
        ]
    )

    outline_agent = Agent(name="outline", model=outline_model)
    judge_agent = Agent(name="judge", model=judge_model, output_type=EvaluationFeedback)

    conversation: list[TResponseInputItem] = [get_text_input_item("Tell me a space story")]
    latest_outline: str | None = None

    for expected_outline, expected_score in [
        ("Outline v1", "needs_improvement"),
        ("Outline v2", "pass"),
    ]:
        outline_result = await Runner.run(outline_agent, conversation)
        latest_outline = ItemHelpers.text_message_outputs(outline_result.new_items)
        assert latest_outline == expected_outline

        conversation = outline_result.to_input_list()

        judge_result = await Runner.run(judge_agent, conversation)
        feedback = judge_result.final_output
        assert isinstance(feedback, EvaluationFeedback)
        assert feedback.score == expected_score

        if feedback.score == "pass":
            break

        conversation.append({"content": f"Feedback: {feedback.feedback}", "role": "user"})

    assert latest_outline == "Outline v2"
    assert len(conversation) == 4
    assert judge_model.last_turn_args["input"] == conversation


@pytest.mark.asyncio
async def test_parallel_translation_flow_reuses_runner_outputs() -> None:
    """Covers the parallelization example by feeding multiple translations into a picker agent."""
    translation_model = FakeModel()
    translation_model.add_multiple_turn_outputs(
        [
            [get_text_message("Uno")],
            [get_text_message("Dos")],
            [get_text_message("Tres")],
        ]
    )
    spanish_agent = Agent(name="spanish_agent", model=translation_model)

    picker_model = FakeModel()
    picker_model.set_next_output([get_text_message("Pick: Dos")])
    picker_agent = Agent(name="picker", model=picker_model)

    translations: list[str] = []
    for _ in range(3):
        result = await Runner.run(spanish_agent, input="Hello")
        translations.append(ItemHelpers.text_message_outputs(result.new_items))

    combined = "\n\n".join(translations)
    picker_result = await Runner.run(
        picker_agent,
        input=f"Input: Hello\n\nTranslations:\n{combined}",
    )

    assert translations == ["Uno", "Dos", "Tres"]
    assert picker_result.final_output == "Pick: Dos"
    assert picker_model.last_turn_args["input"] == [
        {"content": f"Input: Hello\n\nTranslations:\n{combined}", "role": "user"}
    ]


@pytest.mark.asyncio
async def test_deterministic_story_flow_stops_when_checker_blocks() -> None:
    """Mimics deterministic flow: stop early when quality gate fails."""
    outline_model = FakeModel()
    outline_model.set_next_output([get_text_message("Outline v1")])
    checker_model = FakeModel()
    checker_model.set_next_output(
        [
            get_final_output_message(
                json.dumps({"response": {"good_quality": False, "is_scifi": True}})
            )
        ]
    )
    story_model = FakeModel()
    story_model.set_next_output(RuntimeError("story should not run"))

    outline_agent = Agent(name="outline", model=outline_model)
    checker_agent = Agent(
        name="checker",
        model=checker_model,
        output_type=OutlineCheckerOutput,
    )
    story_agent = Agent(name="story", model=story_model)

    inputs: list[TResponseInputItem] = [get_text_input_item("Sci-fi please")]
    outline_result = await Runner.run(outline_agent, inputs)
    inputs = outline_result.to_input_list()

    checker_result = await Runner.run(checker_agent, inputs)
    decision = checker_result.final_output

    assert isinstance(decision, OutlineCheckerOutput)
    assert decision.good_quality is False
    assert decision.is_scifi is True
    if decision.good_quality and decision.is_scifi:
        await Runner.run(story_agent, outline_result.final_output)
    assert story_model.first_turn_args is None, "story agent should never be invoked when gated"


@pytest.mark.asyncio
async def test_deterministic_story_flow_runs_story_on_pass() -> None:
    """Mimics deterministic flow: run full path when checker approves."""
    outline_model = FakeModel()
    outline_model.set_next_output([get_text_message("Outline ready")])
    checker_model = FakeModel()
    checker_model.set_next_output(
        [
            get_final_output_message(
                json.dumps({"response": {"good_quality": True, "is_scifi": True}})
            )
        ]
    )
    story_model = FakeModel()
    story_model.set_next_output([get_text_message("Final story")])

    outline_agent = Agent(name="outline", model=outline_model)
    checker_agent = Agent(
        name="checker",
        model=checker_model,
        output_type=OutlineCheckerOutput,
    )
    story_agent = Agent(name="story", model=story_model)

    inputs: list[TResponseInputItem] = [get_text_input_item("Sci-fi please")]
    outline_result = await Runner.run(outline_agent, inputs)
    inputs = outline_result.to_input_list()

    checker_result = await Runner.run(checker_agent, inputs)
    decision = checker_result.final_output
    assert isinstance(decision, OutlineCheckerOutput)
    assert decision.good_quality is True
    assert decision.is_scifi is True

    story_result = await Runner.run(story_agent, outline_result.final_output)
    assert story_result.final_output == "Final story"
    assert story_model.last_turn_args["input"] == [{"content": "Outline ready", "role": "user"}]


@pytest.mark.asyncio
async def test_routing_stream_emits_text_and_updates_inputs() -> None:
    """Mimics routing example stream: text deltas flow through and input history updates."""
    model = FakeModel()
    model.set_next_output([get_text_message("Bonjour")])
    triage_agent = Agent(name="triage_agent", model=model)

    streamed = Runner.run_streamed(triage_agent, input="Salut")

    deltas: list[str] = []
    async for event in streamed.stream_events():
        if isinstance(event, RawResponsesStreamEvent) and isinstance(
            event.data, ResponseTextDeltaEvent
        ):
            deltas.append(event.data.delta)

    assert "".join(deltas) == "Bonjour"
    assert streamed.final_output == "Bonjour"
    assert len(streamed.new_items) == 1
    input_list = streamed.to_input_list()
    assert len(input_list) == 2
    assert input_list[0] == {"content": "Salut", "role": "user"}
    assistant_item = input_list[1]
    assert isinstance(assistant_item, dict)
    assert assistant_item.get("role") == "assistant"
    assert assistant_item.get("type") == "message"
    content: Any = assistant_item.get("content")
    assert isinstance(content, list)
    first_content = content[0]
    assert isinstance(first_content, dict)
    assert first_content.get("text") == "Bonjour"


class MathHomeworkOutput(BaseModel):
    reasoning: str
    is_math_homework: bool


@pytest.mark.asyncio
async def test_input_guardrail_agent_trips_and_returns_info() -> None:
    """Mimics math guardrail example: guardrail agent runs and trips before main agent completes."""
    guardrail_model = FakeModel()
    guardrail_model.set_next_output(
        [
            get_final_output_message(
                json.dumps({"reasoning": "math detected", "is_math_homework": True})
            )
        ]
    )
    guardrail_agent = Agent(name="guardrail", model=guardrail_model, output_type=MathHomeworkOutput)

    @input_guardrail
    async def math_guardrail(
        context: RunContextWrapper[None], agent: Agent, input: str | list[TResponseInputItem]
    ) -> GuardrailFunctionOutput:
        result = await Runner.run(guardrail_agent, input, context=context.context)
        output = result.final_output_as(MathHomeworkOutput)
        return GuardrailFunctionOutput(
            output_info=output, tripwire_triggered=output.is_math_homework
        )

    main_model = FakeModel()
    main_model.set_next_output([get_text_message("Should not run")])
    main_agent = Agent(name="main", model=main_model, input_guardrails=[math_guardrail])

    with pytest.raises(InputGuardrailTripwireTriggered) as excinfo:
        await Runner.run(main_agent, "Solve 2x+5=11")

    guardrail_result = excinfo.value.guardrail_result
    assert isinstance(guardrail_result.output.output_info, MathHomeworkOutput)
    assert guardrail_result.output.output_info.is_math_homework is True
    assert guardrail_result.output.output_info.reasoning == "math detected"


class MessageOutput(BaseModel):
    reasoning: str
    response: str
    user_name: str | None


@pytest.mark.asyncio
async def test_output_guardrail_blocks_sensitive_data() -> None:
    """Mimics sensitive data guardrail example: trips when phone number is present."""

    @output_guardrail
    async def sensitive_data_check(
        context: RunContextWrapper, agent: Agent, output: MessageOutput
    ) -> GuardrailFunctionOutput:
        contains_phone = "650" in output.response or "650" in output.reasoning
        return GuardrailFunctionOutput(
            output_info={"contains_phone": contains_phone},
            tripwire_triggered=contains_phone,
        )

    model = FakeModel()
    model.set_next_output(
        [
            get_final_output_message(
                json.dumps(
                    {
                        "reasoning": "User shared phone 650-123-4567",
                        "response": "Thanks!",
                        "user_name": None,
                    }
                )
            )
        ]
    )
    agent = Agent(
        name="Assistant",
        model=model,
        output_type=MessageOutput,
        output_guardrails=[sensitive_data_check],
    )

    with pytest.raises(OutputGuardrailTripwireTriggered) as excinfo:
        await Runner.run(agent, "My phone number is 650-123-4567.")

    guardrail_output = excinfo.value.guardrail_result.output.output_info
    assert isinstance(guardrail_output, dict)
    assert guardrail_output["contains_phone"] is True


@pytest.mark.asyncio
async def test_streaming_guardrail_style_cancel_after_threshold() -> None:
    """Mimics streaming guardrail example: stop streaming once threshold is reached."""
    model = FakeModel()
    model.set_next_output(
        [
            get_text_message("Chunk1 "),
            get_text_message("Chunk2 "),
            get_text_message("Chunk3"),
        ]
    )
    agent = Agent(name="talkative", model=model)

    streamed = Runner.run_streamed(agent, input="Start")

    deltas: list[str] = []
    async for event in streamed.stream_events():
        if isinstance(event, RawResponsesStreamEvent) and isinstance(
            event.data, ResponseTextDeltaEvent
        ):
            deltas.append(event.data.delta)
            if len("".join(deltas)) >= len("Chunk1 Chunk2 "):
                streamed.cancel(mode="immediate")

    collected = "".join(deltas)
    assert "Chunk1" in collected
    assert "Chunk3" not in collected
    assert streamed.final_output is None
    assert streamed.is_complete is True


@pytest.mark.asyncio
async def test_streaming_cancel_after_turn_allows_turn_completion() -> None:
    """Ensure cancel(after_turn) lets the current turn finish and final_output is populated."""
    model = FakeModel()
    model.set_next_output([get_text_message("Hello"), get_text_message("World")])
    agent = Agent(name="talkative", model=model)

    streamed = Runner.run_streamed(agent, input="Hi")

    deltas: list[str] = []
    async for event in streamed.stream_events():
        if isinstance(event, RawResponsesStreamEvent) and isinstance(
            event.data, ResponseTextDeltaEvent
        ):
            deltas.append(event.data.delta)
            streamed.cancel(mode="after_turn")

    assert "".join(deltas).startswith("Hello")
    assert streamed.final_output == "World"
    assert streamed.is_complete is True
    assert len(streamed.new_items) == 2


@pytest.mark.asyncio
async def test_streaming_handoff_emits_agent_updated_event() -> None:
    """Mimics routing handoff stream: emits AgentUpdatedStreamEvent and switches agent."""
    delegate_model = FakeModel()
    delegate_model.set_next_output([get_text_message("delegate reply")])
    delegate_agent = Agent(name="delegate", model=delegate_model)

    triage_model = FakeModel()
    triage_model.set_next_output(
        [
            get_text_message("triage summary"),
            get_handoff_tool_call(delegate_agent),
        ]
    )
    triage_agent = Agent(name="triage", model=triage_model, handoffs=[delegate_agent])

    streamed = Runner.run_streamed(triage_agent, input="Help me")

    agent_updates: list[AgentUpdatedStreamEvent] = []
    async for event in streamed.stream_events():
        if isinstance(event, AgentUpdatedStreamEvent):
            agent_updates.append(event)

    assert streamed.final_output == "delegate reply"
    assert streamed.last_agent == delegate_agent
    assert len(agent_updates) >= 1
    assert any(update.new_agent == delegate_agent for update in agent_updates)


@pytest.mark.asyncio
async def test_agent_as_tool_streaming_example_collects_events() -> None:
    """Mimics agents_as_tools_streaming example: on_stream receives nested streaming events."""
    billing_agent = Agent(name="billing")

    received: list[AgentToolStreamEvent] = []

    async def on_stream(event: AgentToolStreamEvent) -> None:
        received.append(event)

    billing_tool = billing_agent.as_tool(
        tool_name="billing_agent",
        tool_description="Answer billing questions",
        on_stream=on_stream,
    )

    async def fake_invoke(ctx, input: str) -> str:
        event_payload: AgentToolStreamEvent = {
            "event": RawResponsesStreamEvent(data=cast(Any, {"type": "output_text_delta"})),
            "agent": billing_agent,
            "tool_call": ctx.tool_call,
        }
        await on_stream(event_payload)
        return "Billing: $100"

    billing_tool.on_invoke_tool = fake_invoke

    main_model = FakeModel()
    main_model.add_multiple_turn_outputs(
        [
            [get_function_tool_call("billing_agent", json.dumps({"input": "Need bill"}))],
            [get_text_message("Final answer")],
        ]
    )

    main_agent = Agent(
        name="support",
        model=main_model,
        tools=[billing_tool],
        model_settings=ModelSettings(tool_choice="required"),
    )

    result = await Runner.run(main_agent, "How much is my bill?")

    assert result.final_output == "Final answer"
    assert received, "on_stream should capture nested streaming events"
    assert all(event["agent"] == billing_agent for event in received)
    assert all(
        event["tool_call"] and event["tool_call"].name == "billing_agent" for event in received
    )


@pytest.mark.asyncio
async def test_sandbox_agents_as_tools_example_serializes_structured_reviews() -> None:
    pricing_model = FakeModel()
    pricing_model.set_next_output(
        [
            get_final_output_message(
                json.dumps(
                    {
                        "requested_discount_percent": 15,
                        "requested_term_months": 24,
                        "pricing_risk": "medium",
                        "summary": "Discount ask is above target band.",
                        "recommended_next_step": "Trade discount for a stronger give-get.",
                        "evidence_files": ["pricing_summary.md", "commercial_notes.md"],
                    }
                )
            )
        ]
    )
    rollout_model = FakeModel()
    rollout_model.set_next_output(
        [
            get_final_output_message(
                json.dumps(
                    {
                        "rollout_risk": "medium",
                        "summary": "Launch timing is compressed.",
                        "blockers": [
                            "Regional admin training is incomplete.",
                            "SSO migration lands in week 2.",
                        ],
                        "recommended_next_step": "Require a phased rollout plan.",
                        "evidence_files": ["rollout_plan.md", "support_history.md"],
                    }
                )
            )
        ]
    )
    orchestrator_model = FakeModel()
    orchestrator_model.add_multiple_turn_outputs(
        [
            [
                get_function_tool_call(
                    "review_pricing_packet",
                    json.dumps({"input": "Review pricing"}),
                    call_id="outer_pricing",
                ),
                get_function_tool_call(
                    "review_rollout_risk",
                    json.dumps({"input": "Review rollout"}),
                    call_id="outer_rollout",
                ),
                get_function_tool_call(
                    "get_discount_approval_rule",
                    json.dumps({"discount_percent": 15}),
                    call_id="outer_approval",
                ),
            ],
            [get_text_message("Recommendation complete")],
        ]
    )

    @function_tool
    def get_discount_approval_rule(discount_percent: int) -> str:
        if discount_percent <= 10:
            return "AE"
        if discount_percent <= 15:
            return "RSD"
        return "Finance + RSD"

    pricing_agent = Agent(
        name="pricing",
        model=pricing_model,
        output_type=PricingPacketReview,
    )
    rollout_agent = Agent(
        name="rollout",
        model=rollout_model,
        output_type=RolloutRiskReview,
    )
    orchestrator = Agent(
        name="orchestrator",
        model=orchestrator_model,
        tools=[
            pricing_agent.as_tool(
                "review_pricing_packet",
                "Pricing review",
                custom_output_extractor=_structured_tool_output_extractor,
            ),
            rollout_agent.as_tool(
                "review_rollout_risk",
                "Rollout review",
                custom_output_extractor=_structured_tool_output_extractor,
            ),
            get_discount_approval_rule,
        ],
        model_settings=ModelSettings(tool_choice="required"),
    )

    result = await Runner.run(orchestrator, "Review the renewal")

    assert result.final_output == "Recommendation complete"
    outer_second_turn_input = cast(
        list[dict[str, Any]],
        orchestrator_model.last_turn_args["input"],
    )
    outer_tool_outputs = [
        item for item in outer_second_turn_input if item.get("type") == "function_call_output"
    ]
    assert outer_tool_outputs == [
        {
            "call_id": "outer_pricing",
            "output": json.dumps(
                {
                    "evidence_files": ["pricing_summary.md", "commercial_notes.md"],
                    "pricing_risk": "medium",
                    "recommended_next_step": "Trade discount for a stronger give-get.",
                    "requested_discount_percent": 15,
                    "requested_term_months": 24,
                    "summary": "Discount ask is above target band.",
                },
                sort_keys=True,
            ),
            "type": "function_call_output",
        },
        {
            "call_id": "outer_rollout",
            "output": json.dumps(
                {
                    "blockers": [
                        "Regional admin training is incomplete.",
                        "SSO migration lands in week 2.",
                    ],
                    "evidence_files": ["rollout_plan.md", "support_history.md"],
                    "recommended_next_step": "Require a phased rollout plan.",
                    "rollout_risk": "medium",
                    "summary": "Launch timing is compressed.",
                },
                sort_keys=True,
            ),
            "type": "function_call_output",
        },
        {
            "call_id": "outer_approval",
            "output": "RSD",
            "type": "function_call_output",
        },
    ]


def test_docker_runner_formats_tool_calls_without_dumping_run_item() -> None:
    assert (
        _format_tool_call(
            {
                "type": "function_call",
                "name": "read_file",
                "arguments": json.dumps({"path": "README.md"}),
            }
        )
        == '[tool call] read_file: {"path": "README.md"}'
    )

    assert (
        _format_tool_call(
            {
                "type": "shell_call",
                "action": {
                    "commands": ["find . -maxdepth 2 -type f", "cat README.md"],
                },
            }
        )
        == "[tool call] shell: find . -maxdepth 2 -type f; cat README.md"
    )


def test_docker_runner_formats_tool_output_as_readable_block() -> None:
    assert _format_tool_output("$ ls\nREADME.md\nsrc\n") == "[tool output]\n$ ls\nREADME.md\nsrc\n"


@pytest.mark.asyncio
async def test_forcing_tool_use_behaviors_align_with_example() -> None:
    """Mimics forcing_tool_use example: default vs first_tool vs custom behaviors."""

    @function_tool
    def get_weather(city: str) -> str:
        return f"{city}: Sunny"

    # default: run_llm_again -> model responds after tool call
    default_model = FakeModel()
    default_model.add_multiple_turn_outputs(
        [
            [
                get_text_message("Tool call coming"),
                get_function_tool_call("get_weather", json.dumps({"city": "Tokyo"})),
            ],
            [get_text_message("Done after tool")],
        ]
    )

    default_agent = Agent(
        name="default",
        model=default_model,
        tools=[get_weather],
        tool_use_behavior="run_llm_again",
        model_settings=ModelSettings(tool_choice=None),
    )

    default_result = await Runner.run(default_agent, "Weather?")
    assert default_result.final_output == "Done after tool"
    assert len(default_result.raw_responses) == 2

    # first_tool: stop_on_first_tool -> final output from first tool result
    first_model = FakeModel()
    first_model.set_next_output(
        [
            get_text_message("Tool call coming"),
            get_function_tool_call("get_weather", json.dumps({"city": "Paris"})),
        ]
    )

    first_agent = Agent(
        name="first",
        model=first_model,
        tools=[get_weather],
        tool_use_behavior="stop_on_first_tool",
        model_settings=ModelSettings(tool_choice="required"),
    )

    first_result = await Runner.run(first_agent, "Weather?")
    assert first_result.final_output == "Paris: Sunny"
    assert len(first_result.raw_responses) == 1

    # custom: uses custom tool_use_behavior to format output, still with required tool choice
    async def custom_tool_use_behavior(
        context: RunContextWrapper[Any], results: list[FunctionToolResult]
    ) -> ToolsToFinalOutputResult:
        return ToolsToFinalOutputResult(
            is_final_output=True, final_output=f"Custom:{results[0].output}"
        )

    custom_model = FakeModel()
    custom_model.set_next_output(
        [
            get_text_message("Tool call coming"),
            get_function_tool_call("get_weather", json.dumps({"city": "Berlin"})),
        ]
    )

    custom_agent = Agent(
        name="custom",
        model=custom_model,
        tools=[get_weather],
        tool_use_behavior=custom_tool_use_behavior,
        model_settings=ModelSettings(tool_choice="required"),
    )

    custom_result = await Runner.run(custom_agent, "Weather?")
    assert custom_result.final_output == "Custom:Berlin: Sunny"


@pytest.mark.asyncio
async def test_routing_multi_turn_continues_with_handoff_agent() -> None:
    """Mimics routing example multi-turn: first handoff, then continue with delegated agent."""
    delegate_model = FakeModel()
    delegate_model.set_next_output([get_text_message("Bonjour")])
    delegate_agent = Agent(name="delegate", model=delegate_model)

    triage_model = FakeModel()
    triage_model.add_multiple_turn_outputs(
        [
            [get_handoff_tool_call(delegate_agent)],
            [get_text_message("handoff completed")],
        ]
    )
    triage_agent = Agent(name="triage", model=triage_model, handoffs=[delegate_agent])

    first_result = await Runner.run(triage_agent, "Help me in French")
    assert first_result.final_output == "Bonjour"
    assert first_result.last_agent == delegate_agent

    # Next user turn continues with delegate.
    delegate_model.set_next_output([get_text_message("Encore?")])
    follow_up_input = first_result.to_input_list()
    follow_up_input.append({"role": "user", "content": "Encore!"})

    second_result = await Runner.run(delegate_agent, follow_up_input)
    assert second_result.final_output == "Encore?"
    assert delegate_model.last_turn_args["input"] == follow_up_input


@pytest.mark.asyncio
async def test_agents_as_tools_conditional_enabling_matches_preference() -> None:
    """Mimics agents_as_tools_conditional example: only enabled tools are invoked per preference."""

    class AppContext(BaseModel):
        language_preference: str

    def french_spanish_enabled(ctx: RunContextWrapper[AppContext], _agent: AgentBase) -> bool:
        return ctx.context.language_preference in ["french_spanish", "european"]

    def european_enabled(ctx: RunContextWrapper[AppContext], _agent: AgentBase) -> bool:
        return ctx.context.language_preference == "european"

    scenarios = [
        ("spanish_only", {"respond_spanish"}),
        ("french_spanish", {"respond_spanish", "respond_french"}),
        ("european", {"respond_spanish", "respond_french", "respond_italian"}),
    ]

    for preference, expected_tools in scenarios:
        spanish_model = FakeModel()
        spanish_model.set_next_output([get_text_message("ES hola")])
        spanish_agent = Agent(name="spanish", model=spanish_model)

        french_model = FakeModel()
        french_model.set_next_output([get_text_message("FR bonjour")])
        french_agent = Agent(name="french", model=french_model)

        italian_model = FakeModel()
        italian_model.set_next_output([get_text_message("IT ciao")])
        italian_agent = Agent(name="italian", model=italian_model)

        orchestrator_model = FakeModel()
        # Build tool calls only for expected tools to avoid missing-tool errors.
        tool_calls = [
            get_function_tool_call(tool_name, json.dumps({"input": "Hi"}))
            for tool_name in sorted(expected_tools)
        ]
        orchestrator_model.add_multiple_turn_outputs([tool_calls, [get_text_message("Done")]])

        context = AppContext(language_preference=preference)

        orchestrator = Agent(
            name="orchestrator",
            model=orchestrator_model,
            tools=[
                spanish_agent.as_tool(
                    tool_name="respond_spanish",
                    tool_description="Spanish",
                    is_enabled=True,
                ),
                french_agent.as_tool(
                    tool_name="respond_french",
                    tool_description="French",
                    is_enabled=french_spanish_enabled,
                ),
                italian_agent.as_tool(
                    tool_name="respond_italian",
                    tool_description="Italian",
                    is_enabled=european_enabled,
                ),
            ],
            model_settings=ModelSettings(tool_choice="required"),
        )

        result = await Runner.run(orchestrator, "Hello", context=context)

        assert result.final_output == "Done"
        assert (
            spanish_model.first_turn_args is not None
            if "respond_spanish" in expected_tools
            else spanish_model.first_turn_args is None
        )
        assert (
            french_model.first_turn_args is not None
            if "respond_french" in expected_tools
            else french_model.first_turn_args is None
        )
        assert (
            italian_model.first_turn_args is not None
            if "respond_italian" in expected_tools
            else italian_model.first_turn_args is None
        )


@pytest.mark.asyncio
async def test_agents_as_tools_orchestrator_runs_multiple_translations() -> None:
    """Orchestrator calls multiple translation agent tools then summarizes."""
    spanish_model = FakeModel()
    spanish_model.set_next_output([get_text_message("ES hola")])
    spanish_agent = Agent(name="spanish", model=spanish_model)

    french_model = FakeModel()
    french_model.set_next_output([get_text_message("FR bonjour")])
    french_agent = Agent(name="french", model=french_model)

    orchestrator_model = FakeModel()
    orchestrator_model.add_multiple_turn_outputs(
        [
            [get_function_tool_call("translate_to_spanish", json.dumps({"input": "Hi"}))],
            [get_function_tool_call("translate_to_french", json.dumps({"input": "Hi"}))],
            [get_text_message("Summary complete")],
        ]
    )

    orchestrator = Agent(
        name="orchestrator",
        model=orchestrator_model,
        tools=[
            spanish_agent.as_tool("translate_to_spanish", "Spanish"),
            french_agent.as_tool("translate_to_french", "French"),
        ],
    )

    result = await Runner.run(orchestrator, "Hi")

    assert result.final_output == "Summary complete"
    assert spanish_model.last_turn_args["input"] == [{"content": "Hi", "role": "user"}]
    assert french_model.last_turn_args["input"] == [{"content": "Hi", "role": "user"}]
    assert len(result.raw_responses) == 3


@pytest.mark.asyncio
async def test_agents_as_tools_subagent_cancellation_preserves_parent_final_output() -> None:
    """A cancelled nested subagent should not drop sibling outputs from the parent turn."""

    async def _cancel_tool() -> str:
        raise asyncio.CancelledError("tool-cancelled")

    success_model = FakeModel()
    success_model.set_next_output([get_text_message("Status: ok")])
    success_agent = Agent(name="status", model=success_model)

    observability_model = FakeModel()
    observability_model.set_next_output(
        [get_function_tool_call("cancel_tool", "{}", call_id="inner_cancel")]
    )
    observability_agent = Agent(
        name="observability",
        model=observability_model,
        tools=[function_tool(_cancel_tool, name_override="cancel_tool")],
        model_settings=ModelSettings(tool_choice="required"),
    )

    orchestrator_model = FakeModel()
    orchestrator_model.add_multiple_turn_outputs(
        [
            [
                get_function_tool_call(
                    "status_agent",
                    json.dumps({"input": "Hi"}),
                    call_id="outer_status",
                ),
                get_function_tool_call(
                    "observability_agent",
                    json.dumps({"input": "Hi"}),
                    call_id="outer_observability",
                ),
            ],
            [get_text_message("Summary complete")],
        ]
    )

    orchestrator = Agent(
        name="orchestrator",
        model=orchestrator_model,
        tools=[
            success_agent.as_tool("status_agent", "Status"),
            observability_agent.as_tool("observability_agent", "Observability"),
        ],
        model_settings=ModelSettings(tool_choice="required"),
    )

    result = await Runner.run(orchestrator, "Hi")

    assert result.final_output == "Summary complete"
    assert len(result.raw_responses) == 2
    assert success_model.last_turn_args["input"] == [{"content": "Hi", "role": "user"}]
    assert observability_model.first_turn_args is not None
    assert observability_model.first_turn_args["input"] == [{"content": "Hi", "role": "user"}]

    second_turn_input = cast(list[dict[str, Any]], orchestrator_model.last_turn_args["input"])
    tool_outputs = [
        item for item in second_turn_input if item.get("type") == "function_call_output"
    ]
    assert len(tool_outputs) == 2
    assert tool_outputs[0] == {
        "call_id": "outer_status",
        "output": "Status: ok",
        "type": "function_call_output",
    }
    assert tool_outputs[1]["call_id"] == "outer_observability"
    assert tool_outputs[1]["type"] == "function_call_output"
    assert tool_outputs[1]["output"].startswith(
        "An error occurred while running the tool. Please try again. Error:"
    )
    assert "cancel" in tool_outputs[1]["output"].lower()


@pytest.mark.asyncio
async def test_agents_as_tools_streaming_subagent_cancellation_preserves_parent_output() -> None:
    """A streaming nested subagent should retain sibling outputs after cancellation."""

    async def _ok_tool() -> str:
        return "Investigation: ok"

    async def _cancel_tool() -> str:
        raise asyncio.CancelledError("tool-cancelled")

    received_events: list[AgentToolStreamEvent] = []

    async def on_stream(event: AgentToolStreamEvent) -> None:
        received_events.append(event)

    status_model = FakeModel()
    status_model.set_next_output([get_text_message("Status: ok")])
    status_agent = Agent(name="status", model=status_model)

    observability_model = FakeModel()
    observability_model.add_multiple_turn_outputs(
        [
            [
                get_function_tool_call("ok_tool", "{}", call_id="inner_ok"),
                get_function_tool_call("cancel_tool", "{}", call_id="inner_cancel"),
            ],
            [get_text_message("Nested summary")],
        ]
    )
    observability_agent = Agent(
        name="observability",
        model=observability_model,
        tools=[
            function_tool(_ok_tool, name_override="ok_tool"),
            function_tool(_cancel_tool, name_override="cancel_tool"),
        ],
        model_settings=ModelSettings(tool_choice="required"),
    )

    orchestrator_model = FakeModel()
    orchestrator_model.add_multiple_turn_outputs(
        [
            [
                get_function_tool_call(
                    "status_agent",
                    json.dumps({"input": "Hi"}),
                    call_id="outer_status",
                ),
                get_function_tool_call(
                    "observability_agent",
                    json.dumps({"input": "Hi"}),
                    call_id="outer_observability",
                ),
            ],
            [get_text_message("Summary complete")],
        ]
    )

    orchestrator = Agent(
        name="orchestrator",
        model=orchestrator_model,
        tools=[
            status_agent.as_tool("status_agent", "Status"),
            observability_agent.as_tool(
                "observability_agent",
                "Observability",
                on_stream=on_stream,
            ),
        ],
        model_settings=ModelSettings(tool_choice="required"),
    )

    result = await Runner.run(orchestrator, "Hi")

    assert result.final_output == "Summary complete"
    assert len(result.raw_responses) == 2
    assert received_events, "on_stream should confirm the nested streaming path ran"
    assert status_model.last_turn_args["input"] == [{"content": "Hi", "role": "user"}]
    assert observability_model.last_turn_args is not None

    nested_second_turn_input = cast(
        list[dict[str, Any]],
        observability_model.last_turn_args["input"],
    )
    nested_tool_outputs = [
        item for item in nested_second_turn_input if item.get("type") == "function_call_output"
    ]
    assert nested_tool_outputs == [
        {
            "call_id": "inner_ok",
            "output": "Investigation: ok",
            "type": "function_call_output",
        },
        {
            "call_id": "inner_cancel",
            "output": (
                "An error occurred while running the tool. Please try again. Error: tool-cancelled"
            ),
            "type": "function_call_output",
        },
    ]

    outer_second_turn_input = cast(
        list[dict[str, Any]],
        orchestrator_model.last_turn_args["input"],
    )
    outer_tool_outputs = [
        item for item in outer_second_turn_input if item.get("type") == "function_call_output"
    ]
    assert outer_tool_outputs == [
        {
            "call_id": "outer_status",
            "output": "Status: ok",
            "type": "function_call_output",
        },
        {
            "call_id": "outer_observability",
            "output": "Nested summary",
            "type": "function_call_output",
        },
    ]


@pytest.mark.asyncio
async def test_agents_as_tools_failure_error_function_none_reraises_cancelled_error() -> None:
    """Explicit None should preserve cancellation semantics for nested agent tools."""

    async def _cancel_tool() -> str:
        raise asyncio.CancelledError("tool-cancelled")

    status_model = FakeModel()
    status_model.set_next_output([get_text_message("Status: ok")])
    status_agent = Agent(name="status", model=status_model)

    observability_model = FakeModel()
    observability_model.set_next_output(
        [get_function_tool_call("cancel_tool", "{}", call_id="inner_cancel")]
    )
    observability_agent = Agent(
        name="observability",
        model=observability_model,
        tools=[
            function_tool(_cancel_tool, name_override="cancel_tool", failure_error_function=None)
        ],
        model_settings=ModelSettings(tool_choice="required"),
    )

    orchestrator_model = FakeModel()
    orchestrator_model.set_next_output(
        [
            get_function_tool_call(
                "status_agent",
                json.dumps({"input": "Hi"}),
                call_id="outer_status",
            ),
            get_function_tool_call(
                "observability_agent",
                json.dumps({"input": "Hi"}),
                call_id="outer_observability",
            ),
        ]
    )

    orchestrator = Agent(
        name="orchestrator",
        model=orchestrator_model,
        tools=[
            status_agent.as_tool("status_agent", "Status"),
            observability_agent.as_tool(
                "observability_agent",
                "Observability",
                failure_error_function=None,
            ),
        ],
        model_settings=ModelSettings(tool_choice="required"),
    )

    with pytest.raises(asyncio.CancelledError):
        await Runner.run(orchestrator, "Hi")