codex_harness.py

from __future__ import annotations

import json
import os
import urllib.error
import urllib.request
import uuid
from dataclasses import dataclass


@dataclass
class CodexResult:
    final_response: str
    finish_reason: str = "completed"
    tool_names: tuple[str, ...] = ()
    tool_events: tuple[str, ...] = ()


class CodexHarness:
    """OpenCode Go adapter with a restricted Marine function-call bridge.

    OpenCode Go is only the model endpoint.  It does not inherit the MCP
    server from the old Codex CLI process, so the adapter must expose the
    allowed Marine functions and feed their results back into the Responses
    loop.  Without this loop ordinary chat works but Ocean exports can only
    produce a prose answer (the failure seen in the Space UI).
    """

    def __init__(self, model=None, api_key=None, base_url=None,
                 session_root=None, system_prompt=None, **kwargs):
        self.model = model or os.environ.get("OPENCODE_GO_MODEL", "gpt-5.6-luna")
        self.api_key = (
            api_key
            or os.environ.get("OPENCODE_GO_API_KEY")
            or os.environ.get("OPENAI_API_KEY")
            or ""
        ).strip()
        self.base_url = (
            base_url
            or os.environ.get("OPENCODE_GO_BASE_URL")
            or os.environ.get("OPENAI_BASE_URL")
            or "https://opencode.ai/zen/go/v1"
        ).rstrip("/")
        self.system_prompt = system_prompt or ""

    @staticmethod
    def _text(payload: dict) -> str:
        text = str(payload.get("output_text") or "").strip()
        if text:
            return text
        parts = []
        for item in payload.get("output") or []:
            for content in item.get("content") or []:
                if content.get("type") in {"output_text", "text"}:
                    value = content.get("text")
                    if isinstance(value, str):
                        parts.append(value)
                    elif isinstance(value, dict) and isinstance(value.get("value"), str):
                        parts.append(value["value"])
        return "".join(parts).strip()

    @staticmethod
    def _tools() -> list[dict]:
        """OpenAI Responses-compatible schemas for the supported MCP tools."""
        def fn(name, description, properties=None, required=None):
            return {
                "type": "function",
                "name": name,
                "description": description,
                "parameters": {
                    "type": "object",
                    "properties": properties or {},
                    "required": required or [],
                    "additionalProperties": False,
                },
            }

        number = {"type": "number"}
        string = {"type": "string"}
        domain = {"type": "string", "enum": ["ocean", "tuna", "squid"]}
        tools = [
            fn("mcp_marine_marine_health", "Check the live school Marine API."),
            fn("mcp_marine_marine_domains", "Return live Ocean, Tuna and Squid data-center status."),
            fn("mcp_marine_marine_status", "Return detailed status for one data center.", {"domain": domain}),
            fn("mcp_marine_marine_catalog", "Return the live Ocean catalog."),
            fn("mcp_marine_marine_query", "Check whether an Ocean source/variable/date exists.", {
                "date": string, "variable": string, "source": string, "domain": domain,
            }, ["date", "variable", "source"]),
            fn("mcp_marine_marine_subset", "Create a NetCDF Ocean subset and return its download URL.", {
                "date": string, "lon_min": number, "lon_max": number,
                "lat_min": number, "lat_max": number, "variable": string,
                "source": string, "domain": domain, "depth": number,
            }, ["date", "lon_min", "lon_max", "lat_min", "lat_max", "variable", "source"]),
            fn("mcp_marine_marine_export", "Export Ocean data as netcdf, csv, xlsx, json, geotiff or png.", {
                "date": string, "lon_min": number, "lon_max": number,
                "lat_min": number, "lat_max": number, "variable": string,
                "source": string, "format": string, "domain": domain, "depth": number,
            }, ["date", "lon_min", "lon_max", "lat_min", "lat_max", "variable", "source", "format"]),
            fn("mcp_marine_marine_export_range", "Export an inclusive Ocean date range (maximum 31 days) and return one download URL per day.", {
                "start_date": string, "end_date": string, "lon_min": number,
                "lon_max": number, "lat_min": number, "lat_max": number,
                "variable": string, "source": string, "format": string,
                "domain": domain, "depth": number,
            }, ["start_date", "end_date", "lon_min", "lon_max", "lat_min", "lat_max", "variable", "source", "format"]),
            fn("mcp_marine_marine_download", "Turn a valid export token into a download URL.", {"token": string}, ["token"]),
            fn("mcp_marine_fisheries_catalog", "Return the live HF fisheries catalog.", {"domain": domain}),
            fn("mcp_marine_fisheries_inventory", "Search the live HF fisheries file inventory.", {
                "domain": domain, "keyword": string, "max_results": {"type": "integer"},
            }),
            fn("mcp_marine_fisheries_search", "Search live HF fisheries files.", {
                "query": string, "max_results": {"type": "integer"},
            }, ["query"]),
            fn("mcp_marine_fisheries_data_rules", "Return fisheries data interpretation rules."),
            fn("mcp_marine_fisheries_analyze_export", "Read and analyze a specific HF fisheries file, optionally exporting CSV.", {
                "repository": string, "path": string, "operation": string,
                "filters": {"type": "object"}, "metric": string,
            }, ["repository", "path"]),
        ]
        # The backend functions accept omitted optional values.  Responses
        # strict mode is intentionally not enabled so older OpenCode models
        # can omit optional keys.
        return tools

    @staticmethod
    def _jsonable(value):
        try:
            json.dumps(value, ensure_ascii=False)
            return value
        except TypeError:
            return {"status": "ok", "result": str(value)}

    @staticmethod
    def _run_tool(name: str, arguments: dict) -> dict:
        """Execute only the allow-listed Marine/HF functions.

        Importing the existing bridge keeps one source of truth for request
        validation and download URL construction.  No shell, filesystem or
        arbitrary Python tool is exposed to the model.
        """
        import marine_mcp

        short = str(name or "")
        if short.startswith("mcp_marine_"):
            short = short[len("mcp_marine_"):]
        aliases = {
            "marine_fisheries_catalog": "marine_fisheries_catalog",
            "fisheries_catalog": "fisheries_catalog",
        }
        short = aliases.get(short, short)
        allowed = {
            "marine_health", "marine_domains", "marine_status", "marine_catalog",
            "marine_query", "marine_subset", "marine_export", "marine_export_range", "marine_download",
            "marine_fisheries_catalog", "fisheries_catalog", "fisheries_inventory",
            "fisheries_search", "fisheries_data_rules", "fisheries_analyze_export",
        }
        if short not in allowed or not hasattr(marine_mcp, short):
            raise ValueError("unsupported Marine tool: " + str(name))
        args = dict(arguments or {})
        # OpenCode models sometimes emit depth=0 for every Ocean request.
        # ERA5, ERA5 accumulated fields, OISST and OC-CCI are 2-D products;
        # forwarding that synthetic depth makes the school API reject a valid
        # request with "has no depth dimension".  Preserve non-zero explicit
        # depths so an actually invalid user request still gets the server's
        # honest validation error.
        if short in {"marine_export", "marine_export_range", "marine_subset"}:
            source = str(args.get("source") or "").strip().lower()
            if source in {"era5", "era5_accum", "oisst", "occci", "oc-cci"}:
                try:
                    if float(args.get("depth")) == 0.0:
                        args.pop("depth", None)
                except (TypeError, ValueError):
                    pass
        result = getattr(marine_mcp, short)(**args)
        return CodexHarness._jsonable(result)

    def _request(self, body: dict, session_id: str) -> dict:
        request = urllib.request.Request(
            self.base_url + "/responses",
            data=json.dumps(body, ensure_ascii=False).encode("utf-8"),
            method="POST",
            headers={
                "Authorization": "Bearer " + self.api_key,
                "Content-Type": "application/json",
                "Accept": "application/json",
                "User-Agent": "Global-Marine-Foundation/4.5",
                "x-opencode-session": str(session_id),
            },
        )
        try:
            with urllib.request.urlopen(request, timeout=600) as response:
                return json.loads(response.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")[:1200]
            raise RuntimeError(f"OpenCode Go API {exc.code}: {detail}") from exc
        except urllib.error.URLError as exc:
            raise RuntimeError(f"OpenCode Go connection failed: {exc.reason}") from exc

    @staticmethod
    def _function_calls(payload: dict) -> list[dict]:
        calls = []
        for item in payload.get("output") or []:
            if not isinstance(item, dict):
                continue
            if item.get("type") not in {"function_call", "tool_call"}:
                continue
            name = item.get("name") or item.get("function", {}).get("name")
            raw_args = item.get("arguments")
            if raw_args is None and isinstance(item.get("function"), dict):
                raw_args = item["function"].get("arguments")
            if isinstance(raw_args, str):
                try:
                    args = json.loads(raw_args or "{}")
                except json.JSONDecodeError as exc:
                    raise RuntimeError(f"Marine tool arguments are invalid: {exc}") from exc
            else:
                args = raw_args or {}
            calls.append({
                "name": str(name or ""),
                "arguments": args if isinstance(args, dict) else {},
                "call_id": item.get("call_id") or item.get("id") or str(uuid.uuid4()),
            })
        return calls

    def run(self, prompt, session_id=None):
        if not self.api_key:
            raise RuntimeError("OPENCODE_GO_API_KEY is not configured")
        request_session_id = str(session_id or uuid.uuid4())
        full_prompt = (self.system_prompt + "\n\n" + prompt).strip()
        tools = self._tools()
        payload = self._request({
            "model": self.model,
            "input": full_prompt,
            "tools": tools,
            "tool_choice": "auto",
        }, request_session_id)
        tool_names = []
        tool_events = []

        # Responses can contain multiple tool calls.  Execute them locally,
        # then continue the same Responses conversation until text is final.
        for _ in range(8):
            calls = self._function_calls(payload)
            if not calls:
                break
            outputs = []
            for call in calls:
                if call["name"] not in tool_names:
                    tool_names.append(call["name"])
                try:
                    value = self._run_tool(call["name"], call["arguments"])
                except Exception as exc:
                    value = {"status": "error", "detail": str(exc)[:1000]}
                if len(tool_events)<40:
                    tool_events.append(json.dumps({
                        "name":call["name"],
                        "arguments":call["arguments"],
                        "result":value,
                    },ensure_ascii=False,default=str)[:12000])
                outputs.append({
                    "type": "function_call_output",
                    "call_id": call["call_id"],
                    "output": json.dumps(value, ensure_ascii=False, default=str),
                })
            next_body = {
                "model": self.model,
                "previous_response_id": payload.get("id"),
                "input": outputs,
                "tools": tools,
                "tool_choice": "auto",
            }
            payload = self._request(next_body, request_session_id)
        else:
            raise RuntimeError("OpenCode Go tool loop exceeded 8 rounds")

        answer = self._text(payload)
        if not answer:
            raise RuntimeError("OpenCode Go returned no text: " + json.dumps(payload, ensure_ascii=False)[:1200])
        return CodexResult(final_response=answer, tool_names=tuple(tool_names), tool_events=tuple(tool_events))