Cognitive_Agent/tests/fixtures/fake_codex_app_server.py

"""Deterministic JSONL App Server used by protocol adapter tests."""

from __future__ import annotations

import json
import os
import sys
import time
import tomllib
from pathlib import Path

TRACE_PATH = Path.cwd() / "wire.jsonl"
ENVIRONMENT_PATH = Path.cwd() / "environment.json"
EXTRA_MCP_NAME = next(
    (
        argument.removeprefix("--fake-extra-mcp=")
        for argument in sys.argv[1:]
        if argument.startswith("--fake-extra-mcp=")
    ),
    None,
)
INVALID_MCP_CONFIG = "--fake-invalid-mcp-config" in sys.argv[1:]
UNAUTHENTICATED = "--fake-unauthenticated" in sys.argv[1:]
LOGIN_FAILURE = "--fake-login-failure" in sys.argv[1:]
NO_OPENAI_AUTH_REQUIRED = "--fake-no-openai-auth-required" in sys.argv[1:]
thread_number = 0
turn_number = 0
waiting_turns: dict[str, str] = {}
account_present = not UNAUTHENTICATED

ENVIRONMENT_PATH.write_text(
    json.dumps(sorted(os.environ), ensure_ascii=True),
    encoding="utf-8",
)


def emit(message: dict[str, object]) -> None:
    print(json.dumps(message, ensure_ascii=False, separators=(",", ":")), flush=True)


def trace(message: dict[str, object]) -> None:
    recorded = message
    if message.get("method") == "account/login/start":
        recorded = dict(message)
        params = message.get("params")
        if isinstance(params, dict):
            recorded["params"] = {**params, "apiKey": "<redacted>"}
    with TRACE_PATH.open("a", encoding="utf-8") as stream:
        stream.write(json.dumps(recorded, ensure_ascii=False) + "\n")


for raw_line in sys.stdin:
    message = json.loads(raw_line)
    trace(message)
    method = message.get("method")
    request_id = message.get("id")
    params = message.get("params") if isinstance(message.get("params"), dict) else {}
    if method == "initialize":
        emit(
            {
                "id": request_id,
                "result": {
                    "userAgent": "fake-codex/1.0",
                    "platformFamily": "test",
                    "platformOs": "test",
                },
            }
        )
    elif method == "initialized":
        continue
    elif method == "account/read":
        emit(
            {
                "id": request_id,
                "result": {
                    "account": ({"type": "apiKey"} if account_present else None),
                    "requiresOpenaiAuth": not NO_OPENAI_AUTH_REQUIRED,
                },
            }
        )
    elif method == "account/login/start":
        if LOGIN_FAILURE:
            emit(
                {
                    "id": request_id,
                    "error": {
                        "code": -32000,
                        "message": f"authentication rejected: {params.get('apiKey', '')}",
                    },
                }
            )
            continue
        if params.get("type") != "apiKey" or not isinstance(params.get("apiKey"), str):
            emit(
                {
                    "id": request_id,
                    "error": {"code": -32602, "message": "invalid login request"},
                }
            )
            continue
        account_present = True
        emit({"id": request_id, "result": {"type": "apiKey"}})
    elif method == "config/read":
        mcp_override = next(
            (
                argument.removeprefix("mcp_servers=")
                for argument in sys.argv
                if argument.startswith("mcp_servers=")
            ),
            "{}",
        )
        mcp_servers = tomllib.loads(f"mcp_servers = {mcp_override}")["mcp_servers"]
        if EXTRA_MCP_NAME and EXTRA_MCP_NAME not in mcp_servers:
            mcp_servers[EXTRA_MCP_NAME] = {"enabled": True}
        if INVALID_MCP_CONFIG:
            mcp_servers = {}
        emit(
            {
                "id": request_id,
                "result": {
                    "config": {
                        "model_provider": "fake-provider",
                        "mcp_servers": mcp_servers,
                    },
                    "origins": {},
                    "layers": None,
                },
            }
        )
    elif method == "thread/start":
        thread_number += 1
        thread_id = f"thread-{thread_number}"
        emit(
            {
                "id": request_id,
                "result": {
                    "thread": {"id": thread_id},
                    "model": params.get("model"),
                    "modelProvider": params.get("modelProvider"),
                },
            }
        )
    elif method == "turn/start":
        turn_number += 1
        turn_id = f"turn-{turn_number}"
        thread_id = str(params.get("threadId"))
        text = str((params.get("input") or [{}])[0].get("text") or "")
        # Emit before the RPC response to exercise out-of-order event backlog.
        emit(
            {
                "method": "turn/started",
                "params": {
                    "threadId": thread_id,
                    "turn": {"id": turn_id, "status": "inProgress", "items": []},
                },
            }
        )
        emit(
            {
                "id": request_id,
                "result": {"turn": {"id": turn_id, "status": "inProgress", "items": []}},
            }
        )
        if "WAIT_FOR_INTERRUPT" in text:
            waiting_turns[turn_id] = thread_id
            continue
        if "EXIT_DURING_TURN" in text:
            # Let the client register the returned turn handle before stdout
            # closes, then exercise its unexpected-process-exit path.
            time.sleep(0.1)
            raise SystemExit(23)
        if "FAIL_DURING_TURN" in text:
            emit(
                {
                    "method": "error",
                    "params": {
                        "threadId": thread_id,
                        "turnId": turn_id,
                        "willRetry": False,
                        "error": {"message": "upstream failed without a terminal event"},
                    },
                }
            )
            continue
        item_id = f"answer-{turn_number}"
        emit(
            {
                "method": "item/started",
                "params": {
                    "threadId": thread_id,
                    "turnId": turn_id,
                    "startedAtMs": 1,
                    "item": {
                        "id": item_id,
                        "type": "agentMessage",
                        "text": "",
                        "phase": "final_answer",
                    },
                },
            }
        )
        for delta in ("协议", "正常"):
            emit(
                {
                    "method": "item/agentMessage/delta",
                    "params": {
                        "threadId": thread_id,
                        "turnId": turn_id,
                        "itemId": item_id,
                        "delta": delta,
                    },
                }
            )
        emit(
            {
                "method": "item/completed",
                "params": {
                    "threadId": thread_id,
                    "turnId": turn_id,
                    "completedAtMs": 2,
                    "item": {
                        "id": item_id,
                        "type": "agentMessage",
                        "text": "协议正常",
                        "phase": "final_answer",
                    },
                },
            }
        )
        emit(
            {
                "method": "turn/completed",
                "params": {
                    "threadId": thread_id,
                    "turn": {"id": turn_id, "status": "completed", "items": []},
                },
            }
        )
    elif method == "turn/interrupt":
        turn_id = str(params.get("turnId"))
        thread_id = waiting_turns.pop(turn_id, str(params.get("threadId")))
        emit({"id": request_id, "result": {}})
        emit(
            {
                "method": "turn/completed",
                "params": {
                    "threadId": thread_id,
                    "turn": {"id": turn_id, "status": "interrupted", "items": []},
                },
            }
        )
    elif method == "thread/unsubscribe":
        emit({"id": request_id, "result": {"status": "unsubscribed"}})
    else:
        emit(
            {
                "id": request_id,
                "error": {"code": -32601, "message": f"unsupported: {method}"},
            }
        )