Mujoco_WASM/decision_server/tests/test_codex.py

import asyncio
import json
import os
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, patch

from aiohttp import web
from aiohttp.test_utils import TestServer

from decision_server.credentials import deepseek_llm, openrouter_jev
from decision_server.protocol import DecisionError
from decision_server.providers.codex import CodexAccount
from decision_server.providers.codex_gate import verify_no_tools
from decision_server.providers.jev import answer
from decision_server.tests.test_service import plan_value, request_value


class CredentialTests(unittest.TestCase):
    def test_explicit_single_variable_no_eval(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "env"
            path.write_text(
                "OTHER_SECRET=do-not-import\n"
                'export OPENROUTER_API_KEY="fixture-key-at-least-16" # note\n'
            )
            conn = openrouter_jev(path)
            self.assertEqual(conn.protocol, "openrouter-decisions")
            self.assertEqual(conn.model, "typesafe/jev-1.13")
            self.assertNotIn("OTHER_SECRET", os.environ)
            path.write_text("DEEPSEEK_API_KEY=another-fixture-key-16\n")
            self.assertEqual(deepseek_llm(path).model, "deepseek-flash")
            self.assertEqual(deepseek_llm(path).base_url, "https://api.deepseek.com")
            for value in ("$(touch PWNED)", "`some-command`", "short"):
                path.write_text("OPENROUTER_API_KEY=" + value)
                with self.assertRaises(DecisionError):
                    openrouter_jev(path)
            path.write_text(
                "OPENROUTER_API_KEY=fixture-key-at-least-16\nOPENROUTER_API_KEY=duplicate-fixture-16"
            )
            with self.assertRaises(DecisionError):
                openrouter_jev(path)

    def test_official_choice_not_reordered(self):
        result = {"answers": {"test": {"choice": "a", "probabilities": {"a": 0.1, "b": 0.9}}}}
        self.assertEqual(answer(result, "test", ["a", "b"]), "a")
        for value in (float("nan"), True, -1, 1.01):
            result["answers"]["test"]["probabilities"]["a"] = value
            with self.assertRaises(DecisionError):
                answer(result, "test", ["a", "b"])


class CodexTests(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.client = CodexAccount(Path(self.temp.name))
        self.client.status = AsyncMock(return_value={"loggedIn": True})
        self.client.check_model = AsyncMock(return_value={"toolGatePassed": True})
        self.client.close = AsyncMock()
        self.calls = []
        self.mode = "success"
        self.started = asyncio.Event()

        async def rpc(method, params):
            self.calls.append((method, params))
            if method == "thread/start":
                return {"thread": {"id": "thread"}}
            if method == "turn/start":
                queue = self.client.queues["thread"]
                self.started.set()
                if self.mode != "wait":
                    item = {"type": "agentMessage", "text": json.dumps(plan_value())}
                    if self.mode == "tool":
                        item = {"type": "commandExecution"}
                    queue.put_nowait(
                        {
                            "method": "item/completed",
                            "params": {
                                "threadId": "thread",
                                "turnId": "turn",
                                "item": item,
                            },
                        }
                    )
                    queue.put_nowait(
                        {
                            "method": "turn/completed",
                            "params": {
                                "threadId": "thread",
                                "turn": {
                                    "id": "turn",
                                    "status": self.mode if self.mode == "failed" else "completed",
                                },
                            },
                        }
                    )
                return {"turn": {"id": "turn"}}
            return {}

        self.client.rpc = AsyncMock(side_effect=rpc)
        self.client.process = SimpleNamespace(returncode=None)

    async def asyncTearDown(self):
        self.temp.cleanup()

    async def test_plan_is_structured_and_ephemeral(self):
        value, _ = await self.client.plan(request_value(), "allowed")
        self.assertEqual(value, plan_value())
        self.assertTrue(self.calls[0][1]["ephemeral"])
        self.assertIn("outputSchema", self.calls[1][1])
        self.assertEqual(self.calls[-1][0], "thread/unsubscribe")
        self.assertFalse(self.client.queues)

    async def test_tool_error_and_failed_turn_interrupt(self):
        for mode in ("tool", "failed"):
            self.mode = mode
            self.calls.clear()
            with self.assertRaises(DecisionError):
                await self.client.plan(request_value(), "allowed")
            self.assertIn("turn/interrupt", [method for method, _ in self.calls])

    async def test_cancel_uses_exact_turn_interrupt(self):
        self.mode = "wait"
        task = asyncio.create_task(self.client.plan(request_value(), "allowed"))
        await self.started.wait()
        task.cancel()
        with self.assertRaises(asyncio.CancelledError):
            await task
        self.assertIn(("turn/interrupt", {"threadId": "thread", "turnId": "turn"}), self.calls)
        self.assertFalse(self.client.queues)

    async def test_failed_gate_never_starts_turn(self):
        self.client.check_model.side_effect = DecisionError("codex_tool_gate_failed")
        with self.assertRaises(DecisionError):
            await self.client.plan(request_value(), "allowed")
        self.assertFalse(self.calls)

    async def test_no_login_no_turn(self):
        self.client.status.return_value = {"loggedIn": False}
        with self.assertRaises(DecisionError):
            await self.client.plan(request_value(), "allowed")
        self.assertFalse(self.calls)

    async def test_login_cancel_and_logout_are_named_operations(self):
        self.client.start = AsyncMock()
        self.client.rpc = AsyncMock(
            return_value={
                "authUrl": "https://auth.openai.com/oauth/authorize?state=fixture",
                "loginId": "login",
            }
        )
        result = await self.client.login()
        self.assertEqual(result["storage"], "session-only")
        self.client.rpc.assert_awaited_with("account/login/start", {"type": "chatgpt"})
        await self.client.cancel_login()
        self.client.rpc.assert_awaited_with("account/login/cancel", {"loginId": "login"})
        self.assertIsNone(self.client.login_id)
        await self.client.logout()
        self.client.rpc.assert_awaited_with("account/logout", {})
        self.client.close.assert_awaited()

    async def test_device_login_uses_official_protocol(self):
        self.client.start = AsyncMock()
        self.client.rpc = AsyncMock(
            return_value={
                "verificationUrl": "https://auth.openai.com/codex/device",
                "userCode": "ABCD-1234",
                "loginId": "device-login",
            }
        )
        result = await self.client.login(device=True)
        self.client.rpc.assert_awaited_with("account/login/start", {"type": "chatgptDeviceCode"})
        self.assertEqual(result["userCode"], "ABCD-1234")
        self.assertNotIn("authUrl", result)
        self.assertEqual(self.client.login_id, "device-login")

    async def test_unexpected_auth_url_rejected(self):
        self.client.start = AsyncMock()
        self.client.rpc = AsyncMock(return_value={"authUrl": "https://evil.test/login"})
        with self.assertRaisesRegex(DecisionError, "codex_unexpected_login_url"):
            await self.client.login()
        self.client.close.assert_awaited()

    async def test_hidden_or_unknown_model_rejected_before_gate(self):
        native = CodexAccount(Path(self.temp.name))
        native.models = AsyncMock(return_value={"models": [{"id": "current"}]})
        with self.assertRaisesRegex(DecisionError, "codex_model_unavailable"):
            await native.check_model("gpt-5.4")
        self.assertFalse(native.checked)

    async def test_unknown_rpc_forbidden(self):
        native = CodexAccount(Path(self.temp.name))
        with self.assertRaises(DecisionError):
            await native.rpc("command/exec", {})
        with (
            patch("decision_server.providers.codex.shutil.which", return_value=None),
            self.assertRaisesRegex(DecisionError, "codex_not_installed"),
        ):
            await native.start()


class NativeGates(unittest.IsolatedAsyncioTestCase):
    @unittest.skipUnless(
        os.environ.get("DECISION_CODEX_SMOKE") == "1", "native fake inference is opt-in"
    )
    async def test_native_structured_turn_and_unsubscribe_without_login(self):
        async def respond(request):
            incoming = await request.json()
            self.assertEqual(incoming.get("tools", []), [])
            item = {
                "type": "message",
                "role": "assistant",
                "id": "msg_plan",
                "status": "completed",
                "phase": "final_answer",
                "content": [{"type": "output_text", "text": json.dumps(plan_value())}],
            }
            events = [
                {
                    "type": "response.created",
                    "response": {"id": "resp_plan", "status": "in_progress"},
                },
                {"type": "response.output_item.done", "output_index": 0, "item": item},
                {
                    "type": "response.completed",
                    "response": {
                        "id": "resp_plan",
                        "status": "completed",
                        "output": [item],
                        "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
                    },
                },
            ]
            return web.Response(
                content_type="text/event-stream",
                text="".join(
                    "event: " + e["type"] + "\ndata: " + json.dumps(e) + "\n\n" for e in events
                ),
            )

        app = web.Application()
        app.router.add_post("/v1/responses", respond)
        server = TestServer(app)
        await server.start_server()
        try:
            with tempfile.TemporaryDirectory() as directory:
                client = CodexAccount(Path(directory), probe_url=str(server.make_url("/v1")))
                try:
                    await client.start()
                    # Fake-inference fixture only: no OAuth or cloud inference.
                    with (
                        patch.object(client, "status", AsyncMock(return_value={"loggedIn": True})),
                        patch.object(client, "check_model", AsyncMock()),
                    ):
                        result, _ = await client.plan(request_value(), "gpt-5.6-terra")
                    self.assertEqual(result, plan_value())
                    self.assertIsNotNone(client.process)
                    self.assertIsNone(client.process.returncode)
                finally:
                    await client.close()
        finally:
            await server.close()

    @unittest.skipUnless(
        os.environ.get("DECISION_CODEX_SMOKE") == "1", "native offline gate is opt-in"
    )
    async def test_all_visible_models_no_tools_and_injected_calls_rejected(self):
        catalog = Path(__file__).parents[1] / "providers/codex_models_0_147.json"
        models = json.loads(catalog.read_text())["models"]
        with tempfile.TemporaryDirectory() as directory:
            for model in models:
                if model["visibility"] != "list":
                    continue
                with self.subTest(model=model["slug"]):
                    evidence = {}
                    await verify_no_tools(Path(directory), model["slug"], evidence)
                    self.assertTrue(evidence["passed"])
                    self.assertEqual(evidence["tools"], [])