openswarm/backend/apps/agents/anthropic_proxy.py

"""Lightweight Anthropic-format HTTP proxy.

When a user is on openswarm-pro with a non-Claude primary (GPT/Gemini/etc.),
the Claude Code CLI needs a single `ANTHROPIC_BASE_URL` that can serve BOTH:

    1. the primary model calls (e.g. `cx/gpt-5` → must go to 9Router)
    2. auxiliary Claude calls for subagents, WebSearch delegation
       (e.g. `claude-haiku-4-5` → must go to OpenSwarm Pro's cloud proxy)

9Router doesn't know about OpenSwarm Pro, and we don't want to maintain a
custom 9Router provider-node for that. This proxy splits requests by the
`model` field in the body and forwards each to the correct upstream.

Mounted at `/api/anthropic-proxy`. Set `ANTHROPIC_BASE_URL` to
`http://127.0.0.1:<backend-port>/api/anthropic-proxy` in the CLI env for
Pro users with non-Claude primaries.
"""

import json
import logging
from contextlib import asynccontextmanager

import httpx
from fastapi import Request
from fastapi.responses import JSONResponse, StreamingResponse

from backend.config.Apps import SubApp

logger = logging.getLogger(__name__)


@asynccontextmanager
async def anthropic_proxy_lifespan():
    yield


anthropic_proxy = SubApp("anthropic-proxy", anthropic_proxy_lifespan)


_CLAUDE_MODEL_PREFIXES = (
    "claude-",
    "claude/",
    "sonnet",
    "opus",
    "haiku",
    "cc/",
)

_GEMINI_MODEL_PREFIXES = ("gemini/", "gc/", "ag/")

# Bare-model patterns that resolve to Gemini's native API (gemini-3-flash-api,
# gemini-3.1-pro-api, gemini-3.1-flash-lite-api, etc. — when user supplies own
# Google API key in Settings → Models). These bypass our `gemini/` prefix so
# the prefix-only check above misses them; we match on the bare-name shape
# here too so $schema scrubbing fires for own-key Gemini sessions.
# Pre-fix: 8/8 own-key Gemini sessions in production failed with 400 because
# JSON Schema's $schema field leaked into Google's tools[].function_declarations
# payload. (See raw_payloads where status=error on every gemini-*-api session.)
_GEMINI_BARE_MODEL_PATTERNS = ("gemini-",)

# Fields Gemini's function_declarations validator rejects. 9Router 0.3.60's
# translator strips allOf/anyOf/oneOf/const-toplevel/required but misses
# these. Each one we've seen Gemini 400 on in production with "Unknown
# name 'X' at request.tools[N].function_declarations[N].parameters.…"
_GEMINI_FORBIDDEN_SCHEMA_KEYS = {
    # JSON-Schema metadata fields Gemini's stricter validator doesn't accept.
    "$schema",
    "$id",                  # ag/gemini-3.1-pro-high session, 2026-05-08
    "$ref",                 # JSON-Schema reference; Gemini wants inlined types
    "$defs",                # ditto
    "definitions",          # legacy alias for $defs
    # Constraint fields Gemini doesn't implement.
    "additionalProperties",
    "propertyNames",
    "patternProperties",
    "exclusiveMinimum",
    "exclusiveMaximum",
    "const",                # nested const leaks through 9Router's top-level-only strip.
    # Anthropic-specific tool-call hints not part of vanilla JSON Schema.
    # Anthropic's CLI emits these on tools that benefit from response
    # priming; Gemini's validator rejects all unknown keys.
    "prefill",              # ag/gemini-3.1-pro-high session, 2026-05-08
    "enumTitles",           # human-readable enum labels; OpenAI-only convention
    "title",                # safe to keep usually but Gemini sometimes rejects under nested arrays
    "examples",             # JSON-Schema 2019-09 keyword Gemini doesn't honor
    "default",              # often allowed but rejected in nested array.items
    "readOnly",
    "writeOnly",
    "deprecated",
}


def _scrub_gemini_schema(node):
    """Recursive in-place strip of Gemini-rejected JSON Schema fields."""
    if isinstance(node, dict):
        for k in list(node.keys()):
            if k in _GEMINI_FORBIDDEN_SCHEMA_KEYS:
                node.pop(k, None)
                continue
            node[k] = _scrub_gemini_schema(node[k])
        return node
    if isinstance(node, list):
        for i, v in enumerate(node):
            node[i] = _scrub_gemini_schema(v)
        return node
    return node


# Models that REQUIRE max_completion_tokens instead of max_tokens.
# OpenAI's GPT-5.x family (gpt-5.4, gpt-5.4-mini, gpt-5.5, gpt-5.3-codex,
# etc.) introduced this in late 2025 — the legacy `max_tokens` field returns
# a 400 "Unsupported parameter: 'max_tokens' is not supported with this
# model. Use 'max_completion_tokens' instead." Anthropic's CLI / SDK still
# emits `max_tokens` because that's the Anthropic-format wire shape; we
# rename it on the way out for OpenAI-routed GPT-5 models.
_OPENAI_MAX_COMPLETION_TOKENS_MODELS = ("gpt-5",)


def _is_openai_max_completion_tokens_model(model: str) -> bool:
    """Match every shape a GPT-5 model name might arrive in. Includes:
      - bare:               "gpt-5", "gpt-5.5", "gpt-5.4-mini"
      - api-suffixed:       "gpt-5.5-api"  (desktop's pinned-api naming)
      - 9router-prefixed:   "openai/gpt-5.5"  (post-translation name)
      - codex-routed:       "cx/gpt-5.3-codex"  (CLI subscription)
    Anything WITHOUT "gpt-5" in the (lowercased) string is rejected.
    """
    m = (model or "").strip().lower()
    if not m:
        return False
    # Strip common routing prefixes so we can match the bare model body.
    for prefix in ("openai/", "cx/", "openrouter/", "or:openai/", "cp/", "cp-"):
        if m.startswith(prefix):
            m = m[len(prefix):]
            break
    return any(m.startswith(p) for p in _OPENAI_MAX_COMPLETION_TOKENS_MODELS)


def _scrub_request_for_openai_gpt5(body: bytes) -> bytes:
    """Rename `max_tokens` → `max_completion_tokens` for GPT-5 models.

    Bytes-in/out, never raises. No-op if the body isn't JSON or doesn't
    contain `max_tokens`. Drops the legacy field if BOTH are present so
    the API doesn't reject for "both fields specified".
    """
    if not body:
        return body
    try:
        parsed = json.loads(body)
    except Exception:
        return body
    if not isinstance(parsed, dict):
        return body
    if "max_tokens" in parsed and "max_completion_tokens" not in parsed:
        parsed["max_completion_tokens"] = parsed.pop("max_tokens")
        return json.dumps(parsed).encode("utf-8")
    if "max_tokens" in parsed and "max_completion_tokens" in parsed:
        parsed.pop("max_tokens", None)
        return json.dumps(parsed).encode("utf-8")
    return body


def _scrub_request_for_gemini(body: bytes) -> bytes:
    """Strip Gemini-incompatible schema keys from request tools. Bytes-in/out, never raises."""
    if not body:
        return body
    try:
        parsed = json.loads(body)
    except Exception:
        return body
    tools = parsed.get("tools") if isinstance(parsed, dict) else None
    if isinstance(tools, list):
        for t in tools:
            if not isinstance(t, dict):
                continue
            if isinstance(t.get("input_schema"), (dict, list)):
                _scrub_gemini_schema(t["input_schema"])
            if isinstance(t.get("parameters"), (dict, list)):
                _scrub_gemini_schema(t["parameters"])
    return json.dumps(parsed).encode("utf-8")


# Headers we strip before forwarding — these change hop-by-hop or we
# replace them with upstream-specific auth.
_HOP_HEADERS = {
    "host",
    "content-length",
    "authorization",
    "x-api-key",
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
}


def _is_claude_model(model: str) -> bool:
    m = (model or "").strip().lower()
    return m.startswith(_CLAUDE_MODEL_PREFIXES)


def _is_gemini_model(model: str) -> bool:
    m = (model or "").strip().lower()
    if m.startswith(_GEMINI_MODEL_PREFIXES):
        return True
    # Bare-name match: "gemini-3-flash-api", "gemini-3.1-pro-api", etc.
    # Excludes anthropic-routed gemini models (those carry "/" or other
    # routing prefixes via the registry).
    if "/" in m:
        return False
    return any(m.startswith(p) for p in _GEMINI_BARE_MODEL_PATTERNS)


def _pick_upstream(model: str) -> tuple[str, dict[str, str]]:
    """Return (base_url_without_v1, auth_headers) for this model."""
    from backend.apps.settings.settings import load_settings
    s = load_settings()

    if _is_claude_model(model):
        # Prefer Pro cloud proxy when configured.
        if getattr(s, "connection_mode", "own_key") == "openswarm-pro":
            bearer = getattr(s, "openswarm_bearer_token", "") or ""
            proxy = (getattr(s, "openswarm_proxy_url", "") or "https://api.openswarm.com").rstrip("/")
            if bearer and proxy:
                return (proxy, {"Authorization": f"Bearer {bearer}"})
        # Fall through — let 9Router handle it (maybe user has a real Claude sub).

    # Default: 9Router for everything else (cx/, gc/, gh/, apikey-routed models).
    return ("http://127.0.0.1:20128", {"x-api-key": "9router"})


@anthropic_proxy.router.api_route(
    "",
    methods=["GET", "HEAD", "OPTIONS"],
    include_in_schema=False,
)
@anthropic_proxy.router.api_route(
    "/",
    methods=["GET", "HEAD", "OPTIONS"],
    include_in_schema=False,
)
async def _healthcheck():
    """CLI healthchecks the proxy root — return 200 so it doesn't 404."""
    return {"ok": True}


@anthropic_proxy.router.api_route(
    "/v1/{rest:path}",
    methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"],
)
async def proxy(rest: str, request: Request):
    body = await request.body()
    model = ""
    if body:
        try:
            parsed = json.loads(body)
            model = str(parsed.get("model") or "")
        except Exception:
            pass

    if _is_gemini_model(model):
        body = _scrub_request_for_gemini(body)
    if _is_openai_max_completion_tokens_model(model):
        body = _scrub_request_for_openai_gpt5(body)

    base_url, auth_headers = _pick_upstream(model)

    forward_headers: dict[str, str] = {}
    for k, v in request.headers.items():
        if k.lower() in _HOP_HEADERS:
            continue
        # The CLI we spawn carries our per-install auth token via
        # `x-api-key` (we set `ANTHROPIC_API_KEY=<our_token>` on the
        # spawn env, and the CLI forwards that value as x-api-key). We
        # must NOT forward that header to the real upstream — it would
        # leak our local token to api.openswarm.com / 9Router, AND it
        # would shadow the real upstream auth (bearer or `9router`
        # literal) that `_pick_upstream` wants to set. Strip it here.
        if k.lower() == "x-api-key":
            continue
        forward_headers[k] = v
    forward_headers.update(auth_headers)

    url = f"{base_url}/v1/{rest}"
    wants_stream = False
    if body:
        try:
            wants_stream = bool(json.loads(body).get("stream"))
        except Exception:
            pass

    try:
        if wants_stream:
            client = httpx.AsyncClient(timeout=httpx.Timeout(600.0, connect=30.0))
            req = client.build_request(
                request.method, url, content=body, headers=forward_headers,
                params=dict(request.query_params),
            )
            upstream = await client.send(req, stream=True)

            async def streamer():
                try:
                    async for chunk in upstream.aiter_raw():
                        if chunk:
                            yield chunk
                finally:
                    await upstream.aclose()
                    await client.aclose()

            return StreamingResponse(
                streamer(),
                status_code=upstream.status_code,
                headers={k: v for k, v in upstream.headers.items()
                         if k.lower() not in _HOP_HEADERS},
                media_type=upstream.headers.get("content-type", "text/event-stream"),
            )
        else:
            async with httpx.AsyncClient(timeout=httpx.Timeout(600.0, connect=30.0)) as client:
                r = await client.request(
                    request.method, url, content=body, headers=forward_headers,
                    params=dict(request.query_params),
                )
                return JSONResponse(
                    content=r.json() if r.headers.get("content-type", "").startswith("application/json") else {"raw": r.text},
                    status_code=r.status_code,
                    headers={k: v for k, v in r.headers.items() if k.lower() not in _HOP_HEADERS},
                )
    except httpx.TimeoutException:
        return JSONResponse({"error": "upstream timeout"}, status_code=504)
    except Exception as e:
        logger.warning(f"anthropic-proxy error: {e}")
        return JSONResponse({"error": str(e)[:300]}, status_code=502)