Mujoco_WASM/decision_server/connections.py

"""Non-secret metadata outside the repository; API credentials are session-memory only."""

import ipaddress
import json
import os
import re
from dataclasses import dataclass, field
from pathlib import Path
from urllib.parse import urlsplit

from .protocol import DecisionError, fields


def endpoint(value):
    if not isinstance(value, str) or len(value) > 512 or any(c.isspace() for c in value):
        raise DecisionError("invalid_endpoint")
    try:
        url = urlsplit(value)
        port = url.port
        host = url.hostname
        local = host == "localhost" or ipaddress.ip_address(host).is_loopback
    except ValueError:
        local = False
        try:
            port, host = url.port, url.hostname
        except (ValueError, UnboundLocalError) as exc:
            raise DecisionError("invalid_endpoint") from exc
    if (
        not host
        or url.username
        or url.password
        or url.query
        or url.fragment
        or "?" in value
        or "#" in value
        or "\\" in value
        or (port is not None and not 1 <= port <= 65535)
        or url.scheme not in ("https", "http")
        or (url.scheme == "http" and not local)
    ):
        raise DecisionError("https_or_loopback_required")
    return value.rstrip("/")


@dataclass(frozen=True)
class Connection:
    protocol: str
    base_url: str
    model: str
    key: str = field(default="", repr=False)

    def public(self):
        return {
            "protocol": self.protocol,
            "baseUrl": self.base_url,
            "model": self.model,
            "hasKey": bool(self.key),
            "keyStorage": "memory-only",
        }


class Connections:
    def __init__(self, directory: Path):
        repo = Path(__file__).resolve().parents[1]
        directory = directory.expanduser().resolve()
        if directory.is_relative_to(repo):
            raise DecisionError("state_directory_must_be_outside_repository")
        self.path = directory / "connections.json"
        self.values = {}
        if self.path.is_symlink():
            raise DecisionError("saved_metadata_symlink_forbidden")
        if self.path.exists():
            try:
                data = json.loads(self.path.read_text())
                for role, config in data.items():
                    self.values[role] = self.parse({"role": role, **config, "apiKey": ""})
            except (OSError, ValueError, DecisionError, TypeError, AttributeError):
                raise DecisionError("invalid_saved_metadata") from None

    @staticmethod
    def parse(data):
        fields(data, ["role", "protocol", "baseUrl", "model", "apiKey"])
        role, protocol = data["role"], data["protocol"]
        allowed = {
            "llm": ("responses", "chat-completions", "codex"),
            "jev": ("typesafe", "openrouter-decisions"),
        }
        if not isinstance(role, str) or role not in allowed or protocol not in allowed[role]:
            raise DecisionError("invalid_provider")
        if not isinstance(data["model"], str) or not re.fullmatch(
            r"[A-Za-z0-9_./:-]{1,128}", data["model"]
        ):
            raise DecisionError("invalid_model")
        key = data["apiKey"]
        if not isinstance(key, str) or len(key) > 4096 or any(c.isspace() for c in key):
            raise DecisionError("invalid_key")
        if key and any(key in str(data[k]) for k in ("model", "baseUrl")):
            raise DecisionError("credential_in_metadata")
        if protocol == "codex":
            if data["baseUrl"] != "" or key:
                raise DecisionError("codex_does_not_accept_api_keys_or_urls")
            return Connection(protocol, "", data["model"])
        return Connection(protocol, endpoint(data["baseUrl"]), data["model"], key)

    def set(self, data):
        conn = self.parse(data)
        # Every update supplies a key anew: never send an old host's credentials to a new host.
        values = {**self.values, data["role"]: conn}
        metadata = {
            role: {"protocol": c.protocol, "baseUrl": c.base_url, "model": c.model}
            for role, c in values.items()
        }
        encoded = json.dumps(metadata)
        if any(c.key and c.key in encoded for c in [*self.values.values(), *values.values()]):
            raise DecisionError("credential_in_metadata")
        self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
        temporary = self.path.with_suffix(".tmp")
        flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC | os.O_NOFOLLOW
        fd = os.open(temporary, flags, 0o600)
        with os.fdopen(fd, "w") as stream:
            os.fchmod(stream.fileno(), 0o600)
            json.dump(metadata, stream)
        os.replace(temporary, self.path)
        self.values = values
        return conn.public()

    def get(self, role):
        if role not in self.values:
            raise DecisionError("connection_not_configured", 409)
        conn = self.values[role]
        if conn.protocol != "codex" and not conn.key:
            raise DecisionError("api_key_required", 409)
        return conn

    def redact(self, value):
        if isinstance(value, str):
            for conn in self.values.values():
                if conn.key:
                    value = value.replace(conn.key, "[redacted]")
        elif isinstance(value, dict):
            return {key: self.redact(item) for key, item in value.items()}
        elif isinstance(value, list):
            return [self.redact(item) for item in value]
        return value