ReMe/tests/unit/test_codex_agent_wrapper.py

"""Unit tests for the Codex agent wrapper and its FastMCP bridge."""

# pylint: disable=missing-class-docstring,missing-function-docstring,protected-access,too-many-lines

import asyncio
from dataclasses import dataclass
import json
from pathlib import Path
import stat
import sys
from types import SimpleNamespace

import pytest
from openai_codex.generated.v2_all import TokenUsageBreakdown
from pydantic import BaseModel

from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper
from reme.components.agent_wrapper.codex_mcp_server import _prepare_config
from reme.components.job import BackgroundJob
from reme.components.outbound_proxy import FixedHttpOutboundProxy
from reme.config import resolve_app_config
from reme.enumeration import ChunkEnum, ComponentEnum
from reme.schema import ApplicationConfig, Response


class _Job:
    def __init__(self, name="search"):
        self.name = name
        self.description = "Search memory"
        self.parameters = {
            "type": "object",
            "properties": {"query": {"type": "string"}},
            "required": ["query"],
        }
        self.calls = []

    async def __call__(self, **kwargs):
        self.calls.append(kwargs)
        return Response(answer=f"found:{kwargs['query']}")


def _wrapper(tmp_path, **kwargs):
    job = _Job()
    config = SimpleNamespace(
        workspace_dir=str(tmp_path),
        mem_session_dir="mem_session",
        environment={},
        components={ComponentEnum.AS_LLM: {}},
        model_dump=lambda **_kwargs: {
            "workspace_dir": str(tmp_path),
            "enable_logo": False,
            "log_to_console": False,
            "log_to_file": False,
            "jobs": {},
            "components": {},
        },
    )
    context = SimpleNamespace(app_config=config, components={}, jobs={job.name: job})
    return CodexAgentWrapper(app_context=context, **kwargs), job


def test_mcp_config_uses_stdio_bridge_and_selected_jobs(tmp_path):
    wrapper, _job = _wrapper(tmp_path, mcp_config="custom.yaml")

    config = wrapper._mcp_server_config(  # pylint: disable=protected-access
        {"job_tools": ["search", "search"], "tool_context_id": "ctx-1"},
    )

    assert config["command"]
    assert config["enabled_tools"] == ["search"]
    assert config["args"].count("--job") == 1
    assert config["args"][config["args"].index("--job") + 1] == "search"
    assert "reme.components.agent_wrapper.codex_mcp_server" in config["args"]
    assert config["args"][config["args"].index("--config") + 1] == str(tmp_path / "custom.yaml")
    assert config["args"][config["args"].index("--tool-context-id") + 1] == "ctx-1"


def test_mcp_config_serializes_injected_job_kwargs(tmp_path):
    wrapper, _job = _wrapper(tmp_path, mcp_config="custom.yaml")

    config = wrapper._mcp_server_config(  # pylint: disable=protected-access
        {
            "job_tools": ["search"],
            "tool_context_id": "ctx-1",
            "injected_job_kwargs": {"_allowed_paths": ["daily/2025-06-01/note.md"]},
        },
    )

    raw = config["args"][config["args"].index("--injected-job-kwargs") + 1]
    assert json.loads(raw) == {"_allowed_paths": ["daily/2025-06-01/note.md"]}

    plain = wrapper._mcp_server_config({"job_tools": ["search"]})  # pylint: disable=protected-access
    assert "--injected-job-kwargs" not in plain["args"]


def test_prepare_config_merges_injected_job_kwargs_with_tool_context():
    prepared = _prepare_config(
        {"jobs": {"selected": {"backend": "base"}}},
        ["selected"],
        "ctx-1",
        {"_allowed_paths": ["daily/2025-06-01/note.md"]},
    )

    assert prepared["service"]["injected_job_kwargs"] == {
        "_allowed_paths": ["daily/2025-06-01/note.md"],
        "tool_context_id": "ctx-1",
    }


def test_thread_config_preserves_other_mcp_servers(tmp_path):
    wrapper, _job = _wrapper(tmp_path)
    config = wrapper._thread_config(  # pylint: disable=protected-access
        {
            "job_tools": ["search"],
            "config": {"mcp_servers": {"docs": {"url": "https://example.test/mcp"}}},
        },
    )

    assert "docs" in config["mcp_servers"]
    assert len(config["mcp_servers"]) == 2
    assert next(name for name in config["mcp_servers"] if name != "docs").startswith("reme_jobs_")


@pytest.mark.asyncio
async def test_thread_config_injects_proxy_only_into_codex_shell_policy(tmp_path):
    wrapper, _job = _wrapper(tmp_path)
    proxy = FixedHttpOutboundProxy(url="http://127.0.0.1:18080")
    await proxy.start()
    wrapper.app_context.components = {ComponentEnum.OUTBOUND_PROXY: {"default": proxy}}
    await wrapper.start()

    config = wrapper._thread_config(  # pylint: disable=protected-access
        {
            "config": {
                "shell_environment_policy": {
                    "inherit": "core",
                    "set": {
                        "CUSTOM": "preserved",
                        "HTTP_PROXY": "http://user-proxy.example:8080",
                    },
                },
            },
        },
    )
    environment = config["shell_environment_policy"]["set"]

    assert config["shell_environment_policy"]["inherit"] == "core"
    assert environment["CUSTOM"] == "preserved"
    for key in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"):
        assert environment[key] == proxy.http_url

    auth = wrapper._resolve_auth_config("oauth")  # pylint: disable=protected-access
    client_config = wrapper._build_client_config(auth)  # pylint: disable=protected-access
    assert "HTTP_PROXY" not in client_config.env
    assert "HTTPS_PROXY" not in client_config.env

    await wrapper.close()
    await proxy.close()


def test_mcp_config_rejects_background_jobs(tmp_path):
    wrapper, _job = _wrapper(tmp_path)
    wrapper.app_context.jobs["watch"] = BackgroundJob(name="watch", app_context=wrapper.app_context)

    with pytest.raises(TypeError, match="non-stream request jobs"):
        wrapper._mcp_server_config({"job_tools": ["watch"]})


def test_prepare_config_reuses_selected_stdio_mcp_service():
    prepared = _prepare_config(
        {
            "service": {"backend": "http"},
            "jobs": {
                "selected": {"backend": "base", "enable_serve": False},
                "helper": {"backend": "base"},
                "watch": {"backend": "background"},
            },
        },
        ["selected"],
        "ctx-1",
    )

    assert prepared["service"] == {
        "backend": "mcp",
        "transport": "stdio",
        "jobs": ["selected"],
        "tool_error_on_failure": True,
        "injected_job_kwargs": {"tool_context_id": "ctx-1"},
    }
    assert set(prepared["jobs"]) == {"selected", "helper"}
    assert prepared["jobs"]["selected"]["enable_serve"] is True


def test_prepare_config_rejects_missing_or_background_selected_jobs():
    with pytest.raises(KeyError, match="missing"):
        _prepare_config({"jobs": {}}, ["missing"])
    with pytest.raises(KeyError, match="watch"):
        _prepare_config({"jobs": {"watch": {"backend": "background"}}}, ["watch"])


@pytest.mark.asyncio
async def test_stdio_bridge_starts_and_lists_selected_job(tmp_path):
    from fastmcp import Client
    from fastmcp.client import StdioTransport

    config_path = tmp_path / "bridge.json"
    config_path.write_text(
        json.dumps(
            {
                "service": {"backend": "mcp"},
                "workspace_dir": str(tmp_path / "workspace"),
                "jobs": {
                    "empty": {
                        "backend": "base",
                        "description": "Return an empty response",
                        "parameters": {"type": "object", "properties": {}},
                        "steps": [],
                    },
                },
            },
        ),
        encoding="utf-8",
    )
    transport = StdioTransport(
        command=sys.executable,
        args=[
            "-m",
            "reme.components.agent_wrapper.codex_mcp_server",
            "--config",
            str(config_path),
            "--workspace",
            str(tmp_path / "workspace"),
            "--job",
            "empty",
        ],
        cwd=str(Path(__file__).resolve().parents[2]),
    )

    async with Client(transport, timeout=10) as client:
        tools = await client.list_tools()

    assert [tool.name for tool in tools] == ["empty"]


@pytest.mark.asyncio
async def test_stdio_bridge_stdout_is_protocol_clean(tmp_path):
    config_path = tmp_path / "bridge.json"
    config_path.write_text(
        json.dumps(
            {
                "service": {"backend": "mcp"},
                "workspace_dir": str(tmp_path / "workspace"),
                "jobs": {
                    "empty": {
                        "backend": "base",
                        "description": "Empty",
                        "parameters": {"type": "object", "properties": {}},
                        "steps": [],
                    },
                },
            },
        ),
        encoding="utf-8",
    )
    proc = await asyncio.create_subprocess_exec(
        sys.executable,
        "-m",
        "reme.components.agent_wrapper.codex_mcp_server",
        "--config",
        str(config_path),
        "--workspace",
        str(tmp_path / "workspace"),
        "--job",
        "empty",
        cwd=str(Path(__file__).resolve().parents[2]),
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    request = {
        "jsonrpc": "2.0",
        "id": 1,
        "method": "initialize",
        "params": {
            "protocolVersion": "2025-06-18",
            "capabilities": {},
            "clientInfo": {"name": "test", "version": "1"},
        },
    }
    assert proc.stdin is not None and proc.stdout is not None and proc.stderr is not None
    proc.stdin.write((json.dumps(request) + "\n").encode())
    await proc.stdin.drain()
    first_line = await asyncio.wait_for(proc.stdout.readline(), timeout=10)
    message = json.loads(first_line)
    assert message["jsonrpc"] == "2.0"
    assert message["id"] == 1
    proc.terminate()
    await asyncio.wait_for(proc.wait(), timeout=10)
    stdout = first_line + await proc.stdout.read()
    stderr = (await proc.stderr.read()).decode()
    assert b"Loading config" not in stdout
    assert b"INFO" not in stdout
    assert b"WARNING" not in stdout
    assert b"2026-" not in stdout
    assert "Failed to parse JSONRPC message" not in stderr
    assert "Invalid JSON" not in stderr


def test_sdk_responds_to_interactive_approval_server_requests():
    from openai_codex.client import CodexClient

    client = CodexClient()
    assert client._default_approval_handler(
        "item/commandExecution/requestApproval",
        {},
    ) == {  # pylint: disable=protected-access
        "decision": "accept",
    }
    assert client._default_approval_handler(
        "item/fileChange/requestApproval",
        {},
    ) == {  # pylint: disable=protected-access
        "decision": "accept",
    }


def test_event_to_chunks_maps_content_usage_and_completion():
    content_event = SimpleNamespace(
        method="item/agentMessage/delta",
        payload=SimpleNamespace(item_id="item-1", delta="hello"),
    )
    usage = TokenUsageBreakdown(
        cachedInputTokens=1,
        inputTokens=3,
        outputTokens=5,
        reasoningOutputTokens=2,
        totalTokens=8,
    )
    usage_event = SimpleNamespace(
        method="thread/tokenUsage/updated",
        payload=SimpleNamespace(token_usage=SimpleNamespace(last=usage)),
    )
    completed_event = SimpleNamespace(
        method="turn/completed",
        payload=SimpleNamespace(
            turn=SimpleNamespace(id="turn-1", status=SimpleNamespace(value="completed"), duration_ms=10, error=None),
        ),
    )

    content = CodexAgentWrapper._event_to_chunks(content_event, "thread-1")  # pylint: disable=protected-access
    usage_chunks = CodexAgentWrapper._event_to_chunks(usage_event, "thread-1")  # pylint: disable=protected-access
    completed = CodexAgentWrapper._event_to_chunks(completed_event, "thread-1")  # pylint: disable=protected-access

    assert content[0].chunk_type == ChunkEnum.CONTENT
    assert content[0].chunk == "hello"
    assert usage_chunks[0].chunk_type == ChunkEnum.USAGE
    assert usage_chunks[0].input_tokens == 3
    assert usage_chunks[0].output_tokens == 5
    assert completed[0].chunk_type == ChunkEnum.REPLY_END
    assert completed[0].metadata["status"] == "completed"


def test_event_to_chunks_preserves_new_turn_scoped_notifications():
    from openai_codex.types import Notification
    from openai_codex.generated.v2_all import TurnDiffUpdatedNotification

    event = Notification(
        method="turn/diff/updated",
        payload=TurnDiffUpdatedNotification(
            threadId="thread-1",
            turnId="turn-1",
            diff="diff --git a/a b/a",
        ),
    )

    chunk = CodexAgentWrapper._event_to_chunks(event, "thread-1")[0]  # pylint: disable=protected-access

    assert chunk.chunk_type == ChunkEnum.DATA
    assert chunk.chunk == {
        "threadId": "thread-1",
        "turnId": "turn-1",
        "diff": "diff --git a/a b/a",
    }
    assert chunk.metadata == {"codex_method": "turn/diff/updated"}


@dataclass
class _TurnResult:
    final_response: str
    status: str = "completed"


@pytest.mark.asyncio
async def test_reply_returns_thread_id_and_structured_output(tmp_path, monkeypatch):
    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")

    class FakeThread:
        id = "thread-1"

        async def run(self, inputs, **kwargs):
            assert inputs == "answer"
            assert kwargs["output_schema"] == {"type": "object"}
            return _TurnResult(final_response=json.dumps({"ok": True}))

    class FakeCodex:
        def __init__(self, _config):
            pass

        async def __aenter__(self):
            return self

        async def __aexit__(self, _exc_type, _exc, _tb):
            return None

        async def close(self):
            return None

        async def account(self):
            return SimpleNamespace(account=SimpleNamespace())

        async def thread_start(self, **_kwargs):
            return FakeThread()

    monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex)

    result = await wrapper.reply("answer", output_schema={"type": "object"})
    await wrapper.close()

    assert result["session_id"] == "thread-1"
    assert result["structured_output"] == {"ok": True}


@pytest.mark.asyncio
async def test_reply_accepts_latest_sdk_run_input(tmp_path, monkeypatch):
    from openai_codex import LocalImageInput, TextInput

    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
    inputs = [TextInput("describe this image"), LocalImageInput("image.png")]
    observed = {}

    class FakeThread:
        id = "thread-1"

        async def run(self, run_input, **_kwargs):
            observed["input"] = run_input
            return _TurnResult(final_response="done")

    class FakeCodex:
        def __init__(self, _config):
            pass

        async def account(self):
            return SimpleNamespace(account=SimpleNamespace())

        async def close(self):
            return None

        async def thread_start(self, **_kwargs):
            return FakeThread()

    monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex)

    result = await wrapper.reply(inputs)
    await wrapper.close()

    assert observed["input"] is inputs
    assert result["last_message"] == "done"


def test_codex_skills_add_all_without_deleting_existing_content(tmp_path):
    wrapper, _job = _wrapper(tmp_path)
    for name in ("reme_memory", "qwenpaw_memory"):
        source = tmp_path / "skills" / name
        source.mkdir(parents=True)
        (source / "SKILL.md").write_text(f"# {name}", encoding="utf-8")
    existing = tmp_path / ".agents" / "skills" / "user_skill"
    existing.mkdir(parents=True)
    marker = existing / "marker"
    marker.write_text("keep", encoding="utf-8")

    wrapper._ensure_skills("all")  # pylint: disable=protected-access

    assert marker.read_text(encoding="utf-8") == "keep"
    for name in ("reme_memory", "qwenpaw_memory"):
        target = tmp_path / ".agents" / "skills" / name
        assert target.is_symlink()
        assert target.resolve() == (tmp_path / "skills" / name).resolve()


def test_codex_skills_support_single_name_and_are_idempotent(tmp_path):
    wrapper, _job = _wrapper(tmp_path)
    source = tmp_path / "skills" / "one"
    source.mkdir(parents=True)
    (source / "SKILL.md").write_text("# one", encoding="utf-8")

    wrapper._ensure_skills("one")  # pylint: disable=protected-access
    wrapper._ensure_skills(["one"])  # pylint: disable=protected-access

    assert (tmp_path / ".agents" / "skills" / "one").resolve() == source.resolve()


@pytest.mark.parametrize("kind", ["directory", "external_link"])
def test_codex_skills_preserve_conflicts(tmp_path, kind):
    wrapper, _job = _wrapper(tmp_path)
    source = tmp_path / "skills" / "one"
    source.mkdir(parents=True)
    (source / "SKILL.md").write_text("# one", encoding="utf-8")
    target = tmp_path / ".agents" / "skills" / "one"
    target.parent.mkdir(parents=True)
    if kind == "directory":
        target.mkdir()
        (target / "marker").write_text("keep", encoding="utf-8")
    else:
        external = tmp_path / "external"
        external.mkdir()
        target.symlink_to(external, target_is_directory=True)

    with pytest.raises(FileExistsError, match="Codex skill conflict"):
        wrapper._ensure_skills("one")  # pylint: disable=protected-access
    assert target.exists()


@pytest.mark.parametrize("create_dir", [False, True])
def test_codex_skills_reject_missing_or_invalid_skill(tmp_path, create_dir):
    wrapper, _job = _wrapper(tmp_path)
    if create_dir:
        (tmp_path / "skills" / "missing_manifest").mkdir(parents=True)
        name = "missing_manifest"
    else:
        name = "missing"
    with pytest.raises(FileNotFoundError):
        wrapper._ensure_skills(name)  # pylint: disable=protected-access


def test_codex_skills_do_not_modify_codex_home(tmp_path):
    codex_home = tmp_path / "codex-home"
    marker = codex_home / "skills" / "marker"
    marker.parent.mkdir(parents=True)
    marker.write_text("keep", encoding="utf-8")
    wrapper, _job = _wrapper(tmp_path, codex_home=codex_home)
    source = tmp_path / "skills" / "one"
    source.mkdir(parents=True)
    (source / "SKILL.md").write_text("# one", encoding="utf-8")

    wrapper._ensure_skills("one")  # pylint: disable=protected-access

    assert marker.read_text(encoding="utf-8") == "keep"


def test_effective_mcp_config_snapshot_is_private_and_removed_on_close(tmp_path):
    wrapper, _job = _wrapper(tmp_path)
    config = wrapper._mcp_server_config({"job_tools": ["search"]})  # pylint: disable=protected-access
    snapshot = Path(config["args"][config["args"].index("--config") + 1])
    assert snapshot.exists()
    assert stat.S_IMODE(snapshot.stat().st_mode) == 0o600
    assert json.loads(snapshot.read_text(encoding="utf-8"))["workspace_dir"] == str(tmp_path)

    async def close_started_wrapper():
        await wrapper.start()
        await wrapper.close()

    asyncio.run(close_started_wrapper())
    assert not snapshot.exists()


@pytest.mark.asyncio
async def test_effective_snapshot_exposes_parent_only_custom_job(tmp_path):
    from fastmcp import Client
    from fastmcp.client import StdioTransport

    app_config = ApplicationConfig(
        workspace_dir=str(tmp_path),
        enable_logo=False,
        log_to_console=False,
        log_to_file=False,
        service={"backend": "mcp"},
        jobs={
            "only_custom": {
                "backend": "base",
                "description": "Parent-only inline job",
                "parameters": {"type": "object", "properties": {}},
                "steps": [],
            },
            "referenced_helper": {
                "backend": "base",
                "description": "A normal job that custom jobs may reference.",
                "parameters": {"type": "object", "properties": {}},
                "steps": [],
            },
        },
    )
    job = _Job("only_custom")
    context = SimpleNamespace(app_config=app_config, components={}, jobs={"only_custom": job})
    wrapper = CodexAgentWrapper(app_context=context)
    server_config = wrapper._mcp_server_config({"job_tools": ["only_custom"]})  # pylint: disable=protected-access
    transport = StdioTransport(
        command=server_config["command"],
        args=server_config["args"],
        cwd=server_config["cwd"],
    )

    await wrapper.start()
    async with Client(transport, timeout=10) as client:
        tools = await client.list_tools()
        snapshot = Path(server_config["args"][server_config["args"].index("--config") + 1])
        assert snapshot.exists()
        assert set(json.loads(snapshot.read_text(encoding="utf-8"))["jobs"]) == {
            "only_custom",
            "referenced_helper",
        }
    await wrapper.close()

    assert [tool.name for tool in tools] == ["only_custom"]
    assert not snapshot.exists()


@pytest.mark.asyncio
async def test_open_thread_defaults_to_full_access(tmp_path):
    from openai_codex import ApprovalMode, Sandbox

    wrapper, _job = _wrapper(tmp_path)
    observed = {}

    class FakeCodex:
        async def thread_start(self, **kwargs):
            observed.update(kwargs)
            return SimpleNamespace(id="thread-1")

    await wrapper._open_thread(FakeCodex(), {})  # pylint: disable=protected-access

    assert observed["approval_mode"] == ApprovalMode.auto_review
    assert observed["sandbox"] == Sandbox.full_access


@pytest.mark.asyncio
async def test_compact_session_uses_native_thread_operation(tmp_path, monkeypatch):
    wrapper, _job = _wrapper(tmp_path)
    observed = {}

    class FakeThread:
        async def compact(self):
            observed["compacted"] = True

    async def start():
        observed["started"] = True

    async def resume(session_id):
        observed["resume"] = session_id
        return FakeThread()

    async def get_codex():
        return SimpleNamespace(thread_resume=resume)

    monkeypatch.setattr(wrapper, "start", start)
    monkeypatch.setattr(wrapper, "_get_codex", get_codex)

    await wrapper.compact_session("thread-1")

    assert observed["started"] is True
    assert observed["resume"] == "thread-1"
    assert observed["compacted"] is True


@pytest.mark.asyncio
async def test_open_thread_forwards_latest_sdk_options(tmp_path):
    from openai_codex.types import Personality, ThreadSource, ThreadStartSource

    wrapper, _job = _wrapper(tmp_path)
    observed = {}

    class FakeCodex:
        async def thread_start(self, **kwargs):
            observed["start"] = kwargs
            return SimpleNamespace(id="thread-1")

        async def thread_resume(self, _thread_id, **kwargs):
            observed["resume"] = kwargs
            return SimpleNamespace(id="thread-1")

        async def thread_fork(self, _thread_id, **kwargs):
            observed["fork"] = kwargs
            return SimpleNamespace(id="thread-2")

    codex = FakeCodex()
    await wrapper._open_thread(  # pylint: disable=protected-access
        codex,
        {
            "personality": "friendly",
            "service_name": "reme",
            "session_start_source": "startup",
            "thread_source": "user",
        },
    )
    await wrapper._open_thread(  # pylint: disable=protected-access
        codex,
        {"session_id": "thread-1", "personality": "pragmatic"},
    )
    await wrapper._open_thread(  # pylint: disable=protected-access
        codex,
        {"session_id": "thread-1", "fork_session": True, "thread_source": "subagent"},
    )

    assert observed["start"]["personality"] == Personality.friendly
    assert observed["start"]["service_name"] == "reme"
    assert observed["start"]["session_start_source"] == ThreadStartSource.startup
    assert observed["start"]["thread_source"] == ThreadSource.user
    assert observed["resume"]["personality"] == Personality.pragmatic
    assert observed["fork"]["thread_source"] == ThreadSource.subagent


@pytest.mark.asyncio
async def test_resume_reuses_tool_context_and_rejects_context_change(tmp_path):
    wrapper, _job = _wrapper(tmp_path)

    class FakeCodex:
        async def thread_start(self, **_kwargs):
            return SimpleNamespace(id="thread-1")

        async def thread_resume(self, _thread_id, **_kwargs):
            return SimpleNamespace(id="thread-1")

    await wrapper._open_thread(FakeCodex(), {"tool_context_id": "ctx-a"})  # pylint: disable=protected-access
    await wrapper._open_thread(  # pylint: disable=protected-access
        FakeCodex(),
        {"resume": "thread-1", "tool_context_id": "ctx-a"},
    )
    with pytest.raises(ValueError, match="cannot change"):
        await wrapper._open_thread(  # pylint: disable=protected-access
            FakeCodex(),
            {"resume": "thread-1", "tool_context_id": "ctx-b"},
        )


class _StructuredModel(BaseModel):
    ok: bool


@pytest.mark.parametrize("schema", [_StructuredModel(ok=True), str])
def test_output_schema_rejects_instances_and_arbitrary_classes(tmp_path, schema):
    wrapper, _job = _wrapper(tmp_path)
    with pytest.raises(TypeError, match="JSON schema dict or BaseModel class"):
        wrapper._merged_kwargs({"output_schema": schema})  # pylint: disable=protected-access


def test_output_schema_normalizes_model_class_and_preserves_dict(tmp_path):
    wrapper, _job = _wrapper(tmp_path)
    schema = {"type": "object", "properties": {"ok": {"type": "boolean"}}}

    assert (
        wrapper._merged_kwargs({"output_schema": _StructuredModel})["output_schema"]  # pylint: disable=protected-access
        == _StructuredModel.model_json_schema()
    )
    assert (
        wrapper._merged_kwargs({"output_schema": schema})["output_schema"] is schema
    )  # pylint: disable=protected-access


@pytest.mark.asyncio
async def test_reply_normalizes_schema_and_reuses_persistent_client(tmp_path, monkeypatch):
    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
    clients = []
    observed_schemas = []
    close_count = 0

    class FakeThread:
        id = "thread-1"

        async def run(self, _inputs, **kwargs):
            observed_schemas.append(kwargs["output_schema"])
            return _TurnResult(final_response=json.dumps({"ok": True}))

    class FakeCodex:
        def __init__(self, _config):
            clients.append(self)

        async def __aenter__(self):
            return self

        async def close(self):
            nonlocal close_count
            close_count += 1

        async def account(self):
            return SimpleNamespace(account=SimpleNamespace())

        async def thread_start(self, **_kwargs):
            return FakeThread()

    monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex)

    await wrapper.start()
    result = await wrapper.reply("first", output_schema=_StructuredModel)
    await wrapper.close()

    assert result["structured_output"] == {"ok": True}
    assert observed_schemas == [_StructuredModel.model_json_schema()]
    assert len(clients) == 1
    assert close_count == 1


@pytest.mark.asyncio
@pytest.mark.parametrize("schema", [{}, _StructuredModel])
async def test_reply_stream_rejects_output_schema(tmp_path, schema):
    wrapper, _job = _wrapper(tmp_path)

    with pytest.raises(NotImplementedError, match="Structured output is not supported"):
        await anext(wrapper.reply_stream("answer", output_schema=schema))


@pytest.mark.asyncio
async def test_reply_stream_interrupts_turn_when_consumer_closes_early(tmp_path, monkeypatch):
    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
    stream_closed = False
    interrupt_count = 0

    class FakeTurn:
        id = "turn-1"

        async def stream(self):
            nonlocal stream_closed
            try:
                yield SimpleNamespace(
                    method="turn/started",
                    payload=SimpleNamespace(turn=SimpleNamespace(id=self.id)),
                )
                await asyncio.Event().wait()
            finally:
                stream_closed = True

        async def interrupt(self):
            nonlocal interrupt_count
            interrupt_count += 1

    class FakeThread:
        id = "thread-1"

        async def turn(self, _inputs, **_kwargs):
            return FakeTurn()

    class FakeCodex:
        def __init__(self, _config):
            pass

        async def account(self):
            return SimpleNamespace(account=SimpleNamespace())

        async def close(self):
            return None

        async def thread_start(self, **_kwargs):
            return FakeThread()

    monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex)

    stream = wrapper.reply_stream("answer")
    first = await anext(stream)
    assert first.chunk_type == ChunkEnum.REPLY_START
    await stream.aclose()
    await wrapper.close()

    assert interrupt_count == 1
    assert stream_closed


@pytest.mark.asyncio
async def test_reply_stream_does_not_record_token_usage(tmp_path, monkeypatch):
    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
    recorded_usages = []
    usage = TokenUsageBreakdown(
        cachedInputTokens=0,
        inputTokens=3,
        outputTokens=5,
        reasoningOutputTokens=0,
        totalTokens=8,
    )

    class FakeTurn:
        id = "turn-1"

        async def stream(self):
            yield SimpleNamespace(
                method="thread/tokenUsage/updated",
                payload=SimpleNamespace(token_usage=SimpleNamespace(last=usage)),
            )
            yield SimpleNamespace(
                method="turn/completed",
                payload=SimpleNamespace(
                    turn=SimpleNamespace(
                        id=self.id,
                        status=SimpleNamespace(value="completed"),
                        duration_ms=1,
                        error=None,
                    ),
                ),
            )

        async def interrupt(self):
            raise AssertionError("Completed turns must not be interrupted")

    class FakeThread:
        id = "thread-1"

        async def turn(self, _inputs, **_kwargs):
            return FakeTurn()

    class FakeCodex:
        def __init__(self, _config):
            pass

        async def account(self):
            return SimpleNamespace(account=SimpleNamespace())

        async def close(self):
            return None

        async def thread_start(self, **_kwargs):
            return FakeThread()

    monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex)
    monkeypatch.setattr(wrapper, "_record_token_usage", recorded_usages.append)

    chunks = [chunk async for chunk in wrapper.reply_stream("answer")]
    await wrapper.close()

    assert any(chunk.chunk_type == ChunkEnum.USAGE for chunk in chunks)
    assert not recorded_usages


@pytest.mark.asyncio
async def test_close_waits_for_active_turn(tmp_path, monkeypatch):
    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
    turn_started = asyncio.Event()
    release_turn = asyncio.Event()
    client_closed = asyncio.Event()

    class FakeThread:
        id = "thread-1"

        async def run(self, _inputs, **_kwargs):
            turn_started.set()
            await release_turn.wait()
            return _TurnResult(final_response="done")

    class FakeCodex:
        def __init__(self, _config):
            pass

        async def account(self):
            return SimpleNamespace(account=SimpleNamespace())

        async def close(self):
            client_closed.set()

        async def thread_start(self, **_kwargs):
            return FakeThread()

    monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex)

    await wrapper.start()
    reply_task = asyncio.create_task(wrapper.reply("answer"))
    await turn_started.wait()
    close_task = asyncio.create_task(wrapper.close())
    await asyncio.sleep(0)

    assert not client_closed.is_set()
    assert not close_task.done()

    release_turn.set()
    result = await reply_task
    await close_task

    assert result["last_message"] == "done"
    assert client_closed.is_set()


@pytest.mark.asyncio
async def test_component_start_keeps_optional_client_lazy(tmp_path, monkeypatch):
    wrapper, _job = _wrapper(tmp_path, auth_mode="api_key", api_key="")

    def fail_if_constructed(_config):
        raise AssertionError("Codex client should be lazy")

    monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", fail_if_constructed)

    await wrapper.start()
    await wrapper.close()

    assert wrapper._codex is None


@pytest.mark.asyncio
async def test_client_config_is_fixed_for_component_lifetime(tmp_path, monkeypatch):
    wrapper, _job = _wrapper(tmp_path, auth_mode="api_key", api_key="one")
    clients = []

    class FakeCodex:
        def __init__(self, _config):
            clients.append(self)

        async def login_api_key(self, _api_key):
            return None

        async def close(self):
            return None

    monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex)

    await wrapper.start()
    assert await wrapper._get_codex() is clients[0]  # pylint: disable=protected-access
    with pytest.raises(TypeError, match="configured on the wrapper: api_key"):
        await wrapper.reply("answer", api_key="two")
    await wrapper.close()

    assert len(clients) == 1


@pytest.mark.parametrize(
    ("name", "value"),
    [
        ("auth_mode", "oauth"),
        ("base_url", "https://example.test/v1"),
        ("codex_bin", "/tmp/codex"),
        ("codex_home", "/tmp/codex-home"),
        ("config_overrides", ['model="test"']),
        ("cwd", "/tmp"),
        ("experimental_api", False),
        ("launch_args_override", ["codex", "app-server"]),
    ],
)
@pytest.mark.asyncio
async def test_reply_rejects_call_time_client_options(tmp_path, name, value):
    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")

    with pytest.raises(TypeError, match=f"configured on the wrapper: {name}"):
        await wrapper.reply("answer", **{name: value})

    assert wrapper._codex is None


def test_constructor_rejects_launch_args_override(tmp_path):
    with pytest.raises(TypeError, match="configure codex_bin instead"):
        _wrapper(tmp_path, launch_args_override=["codex", "app-server"])


def test_oauth_mode_ignores_api_credentials_and_forces_chatgpt(tmp_path, monkeypatch):
    wrapper, _job = _wrapper(
        tmp_path,
        auth_mode="oauth",
        api_key="explicit-key",
        base_url="https://explicit.example.test/v1",
    )
    wrapper.app_context.app_config.environment = {"TOOL_ENV": "preserved"}
    monkeypatch.setenv("CODEX_API_KEY", "ambient-key")
    monkeypatch.setenv("CODEX_BASE_URL", "https://ambient.example.test/v1")

    auth = wrapper._resolve_auth_config(  # pylint: disable=protected-access
        "oauth",
        "explicit-key",
        "https://explicit.example.test/v1",
    )
    config = wrapper._build_client_config(auth)  # pylint: disable=protected-access

    assert auth.mode == "oauth"
    assert auth.api_key == ""
    assert auth.base_url == ""
    assert config.env["TOOL_ENV"] == "preserved"
    assert "CODEX_HOME" not in wrapper.app_context.app_config.environment
    assert 'forced_login_method="chatgpt"' in config.config_overrides
    assert not any(value.startswith("openai_base_url=") for value in config.config_overrides)


@pytest.mark.asyncio
async def test_api_key_mode_logs_in_app_server_explicitly(tmp_path, monkeypatch):
    wrapper, _job = _wrapper(
        tmp_path,
        auth_mode="api_key",
        api_key="explicit-key",
        base_url="https://proxy.example.test/v1",
    )
    observed = {}

    class FakeCodex:
        def __init__(self, config):
            observed["config"] = config

        async def login_api_key(self, api_key):
            observed["api_key"] = api_key

        async def close(self):
            observed["closed"] = True

    monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex)

    await wrapper.start()
    await wrapper._get_codex()  # pylint: disable=protected-access
    await wrapper.close()

    config = observed["config"]
    assert observed["api_key"] == "explicit-key"
    assert 'openai_base_url="https://proxy.example.test/v1"' in config.config_overrides
    assert 'forced_login_method="api"' in config.config_overrides
    assert observed["closed"] is True


def test_auth_selection_uses_explicit_wrapper_options(tmp_path, monkeypatch):
    wrapper, _job = _wrapper(tmp_path)
    wrapper.app_context.app_config.components[ComponentEnum.AS_LLM]["default"] = SimpleNamespace(
        credential={"api_key": "default-key", "base_url": "https://default.example.test/v1"},
    )
    for name in ("CODEX_API_KEY", "OPENAI_API_KEY", "LLM_API_KEY"):
        monkeypatch.setenv(name, "ambient-key")
    for name in ("CODEX_BASE_URL", "OPENAI_BASE_URL", "LLM_BASE_URL"):
        monkeypatch.setenv(name, "https://ambient.example.test/v1")

    auth = wrapper._resolve_auth_config(  # pylint: disable=protected-access
        "api_key",
        "configured-key",
    )

    assert auth.mode == "api_key"
    assert auth.api_key == "configured-key"
    assert auth.base_url == ""

    with pytest.raises(ValueError, match="requires a non-empty API key"):
        wrapper._resolve_auth_config(  # pylint: disable=protected-access
            "api_key",
        )


@pytest.mark.parametrize("review_status", ["approved", "denied"])
def test_event_to_chunks_maps_approval_started_and_completed(review_status):
    action = {"type": "futureApprovalAction", "value": "preserved"}
    review = {"status": review_status, "rationale": "policy"}
    started = SimpleNamespace(
        method="item/autoApprovalReview/started",
        payload=SimpleNamespace(
            action=action,
            review=review,
            review_id="review-1",
            target_item_id="item-1",
            turn_id="turn-1",
        ),
    )
    completed = SimpleNamespace(
        method="item/autoApprovalReview/completed",
        payload=SimpleNamespace(
            action=action,
            review=review,
            review_id="review-1",
            target_item_id="item-1",
            turn_id="turn-1",
            decision_source={"type": "guardian"},
        ),
    )

    started_chunk = CodexAgentWrapper._event_to_chunks(started, "thread-1")[0]  # pylint: disable=protected-access
    completed_chunk = CodexAgentWrapper._event_to_chunks(completed, "thread-1")[0]  # pylint: disable=protected-access

    assert started_chunk.chunk_type == ChunkEnum.APPROVAL
    assert started_chunk.chunk == action
    assert started_chunk.metadata["status"] == "started"
    assert completed_chunk.metadata["status"] == "completed"
    assert completed_chunk.metadata["review"]["status"] == review_status
    assert completed_chunk.metadata["decision_source"] == {"type": "guardian"}


def test_codex_home_expands_user_directory(tmp_path):
    wrapper, _job = _wrapper(tmp_path, codex_home="~/.codex")
    assert wrapper.session_path == Path.home() / ".codex"


def test_named_default_mcp_config_remains_supported(tmp_path):
    wrapper, _job = _wrapper(tmp_path, mcp_config="default")
    assert wrapper._mcp_config_source({}) == "default"  # pylint: disable=protected-access


def test_default_config_provides_codex_oauth_wrapper(monkeypatch):
    monkeypatch.delenv("CODEX_HOME", raising=False)
    config = resolve_app_config(log_config=False)
    oauth = config["components"]["agent_wrapper"]["codex_oauth"]
    codex = config["components"]["agent_wrapper"]["codex"]
    assert oauth["backend"] == "codex"
    assert oauth["auth_mode"] == "oauth"
    assert oauth["codex_home"] == "~/.codex"
    assert oauth["sandbox"] == "full-access"
    assert "api_key" not in oauth
    assert codex["auth_mode"] == "api_key"
    assert codex["sandbox"] == "full-access"