litellm/tests/harness_e2e/test_harness_e2e.py

"""End-to-end: every harness, real runtime, real LiteLLM AI Gateway via litellm_proxy/."""

from pathlib import Path

import pytest
from pydantic import BaseModel

import litellm
from litellm import Harness, sandbox
from litellm.harness import (
    CapabilityUnsupported,
    Done,
    FileChange,
    State,
    Text,
    ToolCall,
)

from .conftest import harness_available, harness_params, model_for, requires_gateway

pytestmark = [requires_gateway]

TURN_TIMEOUT = 300


class Answer(BaseModel):
    city: str
    country: str


@pytest.mark.parametrize("harness", harness_params())
def test_agent_creates_file_and_reports_cost(harness: Harness, workspace: Path) -> None:
    result = litellm.agent(
        harness,
        "Create a file named hello.txt whose entire content is the single word: hi",
        sandbox=sandbox.local(workspace),
        model=model_for(harness),
        timeout=TURN_TIMEOUT,
    )

    assert result.stop_reason == "done", result.text
    assert (workspace / "hello.txt").read_text().strip().lower() == "hi"
    assert [f.path for f in result.files if f.kind == "created"] == ["hello.txt"]
    assert result.usage.calls >= 1
    assert result.usage.input_tokens > 0
    assert result.cost >= 0


@pytest.mark.parametrize("harness", harness_params())
def test_agent_stream_event_order(harness: Harness, workspace: Path) -> None:
    (workspace / "secret.txt").write_text("The secret word is ZEBRA.\n")
    events = list(
        litellm.agent(
            harness,
            "Read secret.txt and reply with just the secret word in it.",
            sandbox=sandbox.local(workspace),
            model=model_for(harness),
            timeout=TURN_TIMEOUT,
            stream=True,
        )
    )

    assert isinstance(events[-1], Done)
    assert sum(isinstance(e, Done) for e in events) == 1
    assert any(isinstance(e, Text) for e in events)
    assert any(isinstance(e, ToolCall) for e in events)
    assert "zebra" in events[-1].result.text.lower()


@pytest.mark.parametrize("harness", harness_params())
def test_agent_structured_output(harness: Harness, workspace: Path) -> None:
    result = litellm.agent(
        harness,
        "What is the capital of France? Do not use any tools.",
        sandbox=sandbox.local(workspace),
        model=model_for(harness),
        output=Answer,
        permissions="read-only",
        timeout=TURN_TIMEOUT,
    )

    assert isinstance(result.output, Answer)
    assert result.output.city.lower() == "paris"


@pytest.mark.parametrize("harness", harness_params())
def test_agent_session_remembers_previous_turn(
    harness: Harness, workspace: Path
) -> None:
    with litellm.agent_session(
        harness,
        sandbox=sandbox.local(workspace),
        model=model_for(harness),
        timeout=TURN_TIMEOUT,
    ) as s:
        s.run("Remember this code word: PELICAN. Reply with just OK.")
        second = s.run(
            "What code word did I ask you to remember? Reply with just the word."
        )
        assert "pelican" in second.text.lower()
        assert s.cost >= second.cost


@pytest.mark.parametrize("harness", harness_params())
def test_agent_detach_and_resume(harness: Harness, workspace: Path) -> None:
    box = sandbox.local(workspace)
    s = litellm.agent_session(
        harness, sandbox=box, model=model_for(harness), timeout=TURN_TIMEOUT
    )
    s.run("Remember this number: 4817. Reply with just OK.")
    raw = s.detach().dumps()

    with litellm.agent_resume(
        State.loads(raw), sandbox=box, model=model_for(harness)
    ) as resumed:
        r = resumed.run(
            "What number did I ask you to remember? Reply with just the number."
        )
    assert "4817" in r.text


requires_opencode = pytest.mark.skipif(
    not harness_available(Harness.OPENCODE), reason="opencode runtime not installed"
)


@requires_opencode
def test_opencode_session_resumed_in_another_sandbox_keeps_working_in_both(
    workspace: Path, tmp_path: Path
) -> None:
    model = model_for(Harness.OPENCODE)
    s = litellm.agent_session(
        Harness.OPENCODE,
        sandbox=sandbox.local(workspace),
        model=model,
        timeout=TURN_TIMEOUT,
    )
    s.run("Remember this number: 4817. Reply with just OK.")
    raw = s.detach().dumps()
    moved = tmp_path / "moved"
    moved.mkdir()

    with litellm.agent_resume(
        State.loads(raw), sandbox=sandbox.local(workspace), model=model, timeout=TURN_TIMEOUT
    ) as original:
        original.run("Reply with just OK.")
        with litellm.agent_resume(
            State.loads(raw), sandbox=sandbox.local(moved), model=model, timeout=TURN_TIMEOUT
        ) as resumed:
            resumed.run(
                "Create a file named hello.txt whose only content is the number I asked "
                "you to remember."
            )
        original.run(
            "Create a file named again.txt whose only content is the number I asked "
            "you to remember."
        )

    assert (moved / "hello.txt").read_text().strip() == "4817"
    assert (workspace / "again.txt").read_text().strip() == "4817"
    assert not (workspace / "hello.txt").exists()


@requires_opencode
def test_opencode_resume_of_a_missing_session_ends_with_a_runtime_error(
    workspace: Path,
) -> None:
    model = model_for(Harness.OPENCODE)
    missing = State(
        harness=Harness.OPENCODE,
        native_session_id="ses_0000000000000000000missing",
        workdir=str(workspace),
        model=model,
    )

    with litellm.agent_resume(
        missing, sandbox=sandbox.local(workspace), model=model, timeout=TURN_TIMEOUT
    ) as resumed:
        result = resumed.run("Reply with just OK.")

    assert result.stop_reason == "runtime_error"
    assert "ses_0000000000000000000missing" in result.text


@pytest.mark.parametrize("harness", harness_params())
def test_agent_read_only_blocks_writes(harness: Harness, workspace: Path) -> None:
    result = litellm.agent(
        harness,
        "Create a file named blocked.txt containing x. If you cannot, just say you cannot.",
        sandbox=sandbox.local(workspace),
        model=model_for(harness),
        permissions="read-only",
        timeout=TURN_TIMEOUT,
    )

    assert not (workspace / "blocked.txt").exists()
    assert not [f for f in result.files if isinstance(f, FileChange)]


def test_string_harness_rejected(workspace: Path) -> None:
    with pytest.raises(TypeError, match=r"Harness\.CODEX"):
        litellm.agent("codex", "hi", sandbox=sandbox.local(workspace))  # type: ignore[arg-type]


def test_capability_checked_before_start(workspace: Path) -> None:
    with pytest.raises(CapabilityUnsupported):
        litellm.agent(
            Harness.CODEX,
            "hi",
            sandbox=sandbox.local(workspace),
            model=model_for(Harness.CODEX),
            disable_tools=["bash"],
        )