proxy.py

import os
import sys
import json
import time
import hmac
import secrets
import asyncio
from pathlib import Path

import httpx
import yaml
from fastapi import FastAPI, Request, Response, BackgroundTasks, Depends, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
from pydantic import BaseModel
from langfuse import Langfuse, propagate_attributes
from opentelemetry.sdk.trace import TracerProvider
try:
    # Private to the SDK -- only used to drop its per-public-key cache when a
    # team's Langfuse keys change. If a future SDK moves it, key changes just
    # need a restart again, as before.
    from langfuse._client.resource_manager import LangfuseResourceManager
except Exception:
    LangfuseResourceManager = None
from contextlib import asynccontextmanager
import threading
import uvicorn

# ── Config ────────────────────────────────────────────────────────────────────
PROXY_PORT       = int(os.getenv("PROXY_PORT",       "4000"))
PROXY_TIMEOUT    = int(os.getenv("PROXY_TIMEOUT",    "600"))
ADMIN_TOKEN      = os.getenv("ADMIN_TOKEN",           "")
CONFIG_DATA_PATH = os.getenv("CONFIG_DATA_PATH",      "data/config.json")
LEGACY_YAML_PATH = os.getenv("CONFIG_PATH",           "backends.yaml")
# The domain teams should hit, e.g. https://llm.company.com — used verbatim
# in /whoami's curl examples instead of guessing from request headers, which
# only works if a reverse proxy is correctly forwarding X-Forwarded-Proto/Host.
# Leave unset to fall back to that header-based guess (fine for local/direct use).
PUBLIC_BASE_URL  = os.getenv("PUBLIC_BASE_URL",       "").rstrip("/")
# Default cap on concurrent in-flight requests to a single backend, used
# whenever a backend doesn't specify its own max_concurrency. The shared
# httpx connection pool (see http_client below) has no per-backend
# awareness on its own, so without this a burst to one backend can starve
# every other backend and team of connections too.
DEFAULT_BACKEND_CONCURRENCY = 20

# Starting value for the editable "supported endpoints" list (managed live
# at /admin, persisted to data/config.json as "llm_endpoints") — only used
# to seed a config.json that doesn't have that key yet, e.g. on first boot
# or when upgrading from a version of the proxy older than this setting.
DEFAULT_LLM_ENDPOINTS = [
    "v1/chat/completions", "v1/completions",
    "chat/completions",    "completions",
    "v1/responses",        "responses",
    "api/chat",            "api/generate",
    "v1/rerank",           "rerank",
    "v1/messages",         "messages",
]

# Anthropic's API-format version, sent as the required `anthropic-version`
# header on every request (not a model version -- it hasn't changed since
# the Messages API's initial release, and covers every Claude model).
ANTHROPIC_VERSION = "2023-06-01"

# Legacy env vars — read once, only to seed data/config.json on first boot.
# After that file exists, these are ignored; edit backends/teams/tracing via
# the admin UI (or the JSON file) instead.
_LEGACY_LANGFUSE_HOST       = os.getenv("LANGFUSE_HOST",       "")
_LEGACY_LANGFUSE_PUBLIC_KEY = os.getenv("LANGFUSE_PUBLIC_KEY", "")
_LEGACY_LANGFUSE_SECRET_KEY = os.getenv("LANGFUSE_SECRET_KEY", "")

if not ADMIN_TOKEN:
    print("Missing required env var: ADMIN_TOKEN (generate one: openssl rand -hex 24)", flush=True)
    sys.exit(1)


# ── Runtime config store ─────────────────────────────────────────────────────
# Single source of truth for backends, teams, and tracing settings, editable
# live via the admin UI/API — no restart needed. Persisted as JSON (not the
# old backends.yaml + .env-var-indirection scheme) because the whole point
# of the UI is that secrets get typed into a form and saved, not hand-edited
# in files. Treat this file like .env: never commit it.
class ConfigStore:
    def __init__(self, path: str, legacy_yaml_path: str):
        self.path = Path(path)
        self._lock = asyncio.Lock()
        self.data = self._load(legacy_yaml_path)

    def _load(self, legacy_yaml_path: str) -> dict:
        if self.path.exists():
            with open(self.path) as f:
                data = json.load(f)
        else:
            data = self._migrate_from_yaml(legacy_yaml_path)
            self._write(data)
        data.setdefault("backends", {})
        data.setdefault("teams", {})
        data.setdefault("llm_endpoints", list(DEFAULT_LLM_ENDPOINTS))
        for cfg in data["teams"].values():
            cfg.setdefault("langfuse", {"enabled": False, "host": "", "public_key": "", "secret_key": ""})
        # Upgrade a config.json written before per-model pricing existed,
        # where "models" was just a list of id strings.
        for cfg in data["backends"].values():
            cfg["models"] = [
                m if isinstance(m, dict) else {"id": m, "input_price": 0.0, "output_price": 0.0}
                for m in cfg.get("models", [])
            ]
            cfg.setdefault("max_concurrency", DEFAULT_BACKEND_CONCURRENCY)
        return data

    def _migrate_from_yaml(self, legacy_yaml_path: str) -> dict:
        data = {"backends": {}, "teams": {}}
        # Every team gets its own Langfuse project (host/public/secret key),
        # entered on that team's admin form. The legacy LANGFUSE_* env vars
        # are only used here, once, as the starting value for teams imported
        # from backends.yaml — new teams get their keys typed in at /admin.
        seed_langfuse = {
            "enabled":    bool(_LEGACY_LANGFUSE_HOST and _LEGACY_LANGFUSE_PUBLIC_KEY
                                and _LEGACY_LANGFUSE_SECRET_KEY),
            "host":       _LEGACY_LANGFUSE_HOST,
            "public_key": _LEGACY_LANGFUSE_PUBLIC_KEY,
            "secret_key": _LEGACY_LANGFUSE_SECRET_KEY,
        }
        yp = Path(legacy_yaml_path)
        if not yp.exists():
            return data
        raw = yaml.safe_load(yp.read_text()) or {}
        for name, cfg in (raw.get("backends") or {}).items():
            api_key_env = cfg.get("api_key_env") or ""
            data["backends"][name] = {
                "type":            cfg.get("type", "openai"),
                "base_url":        cfg["base_url"].rstrip("/"),
                "api_key":         os.getenv(api_key_env, "") if api_key_env else "",
                "models":          [{"id": m, "input_price": 0.0, "output_price": 0.0}
                                    for m in (cfg.get("models") or [])],
                "max_concurrency": cfg.get("max_concurrency") or DEFAULT_BACKEND_CONCURRENCY,
            }
        for name, cfg in (raw.get("teams") or {}).items():
            token_env = cfg.get("token_env")
            data["teams"][name] = {
                "token":    os.getenv(token_env, "") if token_env else "",
                "backends": cfg.get("backends") or [],
                "langfuse": dict(seed_langfuse),
            }
        print(f"One-time migration: imported {legacy_yaml_path} -> {self.path}", flush=True)
        return data

    def _write(self, data=None):
        data = self.data if data is None else data
        self.path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self.path.with_suffix(".tmp")
        with open(tmp, "w") as f:
            json.dump(data, f, indent=2, sort_keys=True)
        tmp.replace(self.path)
        try:
            os.chmod(self.path, 0o600)
        except Exception:
            pass

    async def save(self):
        async with self._lock:
            self._write()


store = ConfigStore(CONFIG_DATA_PATH, LEGACY_YAML_PATH)

BACKENDS:      dict = {}
TEAMS:         dict = {}
TOKENS:        dict = {}
LLM_ENDPOINTS: set  = set()


def _rebuild_indexes():
    global BACKENDS, TEAMS, TOKENS, LLM_ENDPOINTS
    BACKENDS      = store.data["backends"]
    TEAMS         = {name: {"backends": cfg.get("backends", [])} for name, cfg in store.data["teams"].items()}
    TOKENS        = {cfg["token"]: name for name, cfg in store.data["teams"].items() if cfg.get("token")}
    LLM_ENDPOINTS = set(store.data.get("llm_endpoints") or DEFAULT_LLM_ENDPOINTS)


_rebuild_indexes()

# ── Per-backend concurrency limiting ─────────────────────────────────────────
# Each backend gets its own asyncio.Semaphore capping how many requests can
# be actively in flight to it at once; requests beyond that limit simply
# wait their turn rather than being rejected. A semaphore's limit is fixed
# at construction, so a backend gets a fresh one any time its config is
# saved (including just to change the limit) — see _rebuild_backend_semaphore.
_BACKEND_SEMAPHORES: dict = {}


def _rebuild_backend_semaphore(name: str):
    limit = BACKENDS.get(name, {}).get("max_concurrency") or DEFAULT_BACKEND_CONCURRENCY
    _BACKEND_SEMAPHORES[name] = asyncio.Semaphore(limit)


def _backend_semaphore(name: str) -> asyncio.Semaphore:
    return _BACKEND_SEMAPHORES.setdefault(name, asyncio.Semaphore(DEFAULT_BACKEND_CONCURRENCY))


for _backend_name in BACKENDS:
    _rebuild_backend_semaphore(_backend_name)


def _presented_token(request: Request) -> str:
    auth = request.headers.get("authorization", "")
    return auth[7:].strip() if auth.lower().startswith("bearer ") else auth.strip()


def _authenticate(request: Request):
    """Returns (team_name, None) on success, or (None, JSONResponse) on failure."""
    presented = _presented_token(request)

    # Constant-time compare against every token so a wrong guess leaks no timing.
    matched = None
    for tok, team in TOKENS.items():
        if hmac.compare_digest(presented, tok):
            matched = team
    if matched:
        return matched, None

    return None, JSONResponse(
        status_code=401,
        content={"error": {
            "message": "Incorrect API key provided.",
            "type":    "invalid_request_error",
            "code":    "invalid_api_key",
        }},
    )


def require_admin(request: Request):
    presented = _presented_token(request)
    if not presented or not hmac.compare_digest(presented, ADMIN_TOKEN):
        raise HTTPException(status_code=401, detail="unauthorized")


# ── Tracing (Langfuse) ────────────────────────────────────────────────────────
# One Langfuse client per team — each team traces into its own project. Built
# directly (not get_client()) so a team's client can be swapped out the
# moment an admin edits its keys, with no restart. A team with no (or
# incomplete) Langfuse keys gets tracing_enabled=False, which makes every
# start_as_current_observation() call a real no-op instead of needing a
# separate code path in the proxy handler below.
_LANGFUSE_CLIENTS: dict = {}

# The Langfuse SDK caches its exporter per public_key for the life of the
# process: a second Langfuse(...) with the same public_key silently reuses
# the first one's host/secret/enabled state. So an admin fixing a team's
# secret key or host (same public key) would otherwise change nothing until
# a restart. Each public key therefore gets its own TracerProvider here,
# tracked with the host/secret it was built for; when those change, the old
# provider is shut down, the SDK's cached entry is dropped, and a fresh one
# is built. A per-key provider also means one team's flush() only waits on
# its own Langfuse project, not on every team's.
_LANGFUSE_PROVIDERS: dict = {}   # public_key -> {"sig": (host, secret_key), "provider": TracerProvider}


def _langfuse_host(cfg: dict) -> str:
    # A trailing "/" makes the SDK post to "<host>//api/public/otel/...",
    # which Langfuse answers with a 308 redirect on every export.
    return (cfg.get("host") or "").strip().rstrip("/")


def _evict_sdk_cache(public_key: str) -> bool:
    """Drops the SDK's cached exporter for this public key, shutting it down
    in the background. False if the SDK internals aren't reachable -- the
    caller then keeps the existing exporter (the pre-fix behaviour: new keys
    apply after a restart) rather than wiring a team to a dead one."""
    try:
        if LangfuseResourceManager is None:
            raise RuntimeError("langfuse SDK internals not available")
        with LangfuseResourceManager._lock:
            stale = LangfuseResourceManager._instances.pop(public_key, None)
    except Exception as e:
        print(f"Langfuse error: could not reset cached client for {public_key} "
              f"(restart to apply new keys): {e}", flush=True)
        return False
    if stale is not None:
        threading.Thread(target=_shutdown_quietly, args=(stale.shutdown,), daemon=True).start()
    return True


def _shutdown_quietly(fn):
    try:
        fn()
    except Exception as e:
        print(f"Langfuse error: {e}", flush=True)


def _build_langfuse(cfg: dict, replace: bool = False) -> Langfuse:
    """replace=True only for an admin save of this team: its host/secret then
    take over the exporter for its public key. Every other build (startup,
    other teams sharing the key) reuses the existing exporter as-is -- the SDK
    can only hold one per public key, so two teams saved with different
    secrets for the same key must not keep replacing each other's."""
    cfg = cfg or {}
    host = _langfuse_host(cfg)
    public_key = cfg.get("public_key") or ""
    secret_key = cfg.get("secret_key") or ""
    enabled = bool(cfg.get("enabled") and host and public_key and secret_key)
    if not enabled:
        # Never under the team's real public key: the SDK would cache a
        # tracing-disabled exporter for it, and turning tracing on later
        # would keep silently reusing that.
        return Langfuse(public_key="disabled", secret_key="disabled",
                        host="http://localhost", tracing_enabled=False)
    sig = (host, secret_key)
    state = _LANGFUSE_PROVIDERS.get(public_key)
    if state is None or (replace and state["sig"] != sig):
        if _evict_sdk_cache(public_key):
            if state is not None:
                # Flush what the old keys still had queued, then stop that
                # exporter so it can't keep re-sending with stale credentials.
                threading.Thread(target=_shutdown_quietly, args=(state["provider"].shutdown,),
                                 daemon=True).start()
            state = {"sig": sig, "provider": TracerProvider()}
            _LANGFUSE_PROVIDERS[public_key] = state
        elif state is None:
            state = {"sig": sig, "provider": TracerProvider()}
            _LANGFUSE_PROVIDERS[public_key] = state
    return Langfuse(
        public_key=public_key,
        secret_key=secret_key,
        host=host,
        tracing_enabled=True,
        tracer_provider=state["provider"],
    )


def _reinit_team_langfuse(team: str, replace: bool = False):
    cfg = store.data["teams"].get(team, {}).get("langfuse") or {}
    public_key = cfg.get("public_key") or ""
    before = _LANGFUSE_PROVIDERS.get(public_key, {}).get("provider")
    _LANGFUSE_CLIENTS[team] = _build_langfuse(cfg, replace=replace)
    after = _LANGFUSE_PROVIDERS.get(public_key, {}).get("provider")
    # If that rebuilt the exporter for this public key, any other team
    # sharing the key still holds a client wired to the old (now shut down)
    # one -- rebuild those too, against the new exporter.
    if before is not None and after is not before:
        for other, other_cfg in store.data["teams"].items():
            if other != team and (other_cfg.get("langfuse") or {}).get("public_key") == public_key:
                _LANGFUSE_CLIENTS[other] = _build_langfuse(other_cfg.get("langfuse"))


def _get_langfuse(team: str) -> Langfuse:
    client = _LANGFUSE_CLIENTS.get(team)
    if client is None:
        client = _build_langfuse(store.data["teams"].get(team, {}).get("langfuse"))
        _LANGFUSE_CLIENTS[team] = client
    return client


for _team_name in store.data["teams"]:
    _reinit_team_langfuse(_team_name)


def _span_team(span) -> str:
    # The root span carries metadata={"service": team, ...}; the SDK stores
    # metadata either flattened per key or as one JSON blob depending on version.
    attrs = span.attributes or {}
    team = attrs.get("langfuse.observation.metadata.service")
    if team is None and attrs.get("langfuse.observation.metadata"):
        try:
            team = json.loads(attrs["langfuse.observation.metadata"]).get("service")
        except Exception:
            team = None
    return str(team).strip('"') if team else ""


def _install_trace_export_log():
    """
    Logs one line per batch the Langfuse SDK uploads -- OK or FAILED, which
    team(s), how many traces/spans, how long it took, and the trace ids (paste
    one into Langfuse's search to jump straight to it). Wraps the uploader's
    public export() method rather than replacing the uploader, so the SDK's
    own auth, endpoint, compression and batching are untouched.
    Set LOG_TRACE_EXPORTS=false to turn it off.
    """
    if os.getenv("LOG_TRACE_EXPORTS", "true").lower() in ("false", "0", "no"):
        return
    try:
        from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
        from opentelemetry.sdk.trace.export import SpanExportResult
    except Exception as e:
        print(f"Langfuse: trace export logging unavailable: {e}", flush=True)
        return
    if getattr(OTLPSpanExporter.export, "_proxy_logged", False):
        return
    original_export = OTLPSpanExporter.export

    def export(self, spans, *args, **kwargs):
        started = time.time()
        try:
            result = original_export(self, spans, *args, **kwargs)
        except Exception as e:
            result, error = None, e
        else:
            error = None
        ms = int((time.time() - started) * 1000)
        try:
            spans = list(spans)
            trace_ids = sorted({format(s.context.trace_id, "032x") for s in spans})
            teams = sorted({t for t in (_span_team(s) for s in spans) if t}) or ["?"]
            ok = error is None and result == SpanExportResult.SUCCESS
            ts = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
            print(f"[{ts}] TRACE {'OK' if ok else 'FAILED'} | team={','.join(teams)} | "
                  f"{len(trace_ids)} trace(s), {len(spans)} span(s) | {ms}ms | "
                  f"-> {getattr(self, '_endpoint', '?')}"
                  + (f" | error={error}" if error else "")
                  + f" | trace_id={','.join(trace_ids[:5])}"
                  + (f",... (+{len(trace_ids) - 5} more)" if len(trace_ids) > 5 else ""), flush=True)
        except Exception:
            pass
        if error is not None:
            raise error
        return result

    export._proxy_logged = True
    OTLPSpanExporter.export = export


_install_trace_export_log()


def _log_background_flush_result(future: "asyncio.Future"):
    # run_in_executor's Future silently swallows an exception raised inside
    # the worker thread unless something actually inspects it -- without
    # this callback, a failed export (e.g. the Langfuse host being down)
    # would fail completely invisibly, with no log line at all, which
    # makes "the trace never showed up" impossible to diagnose.
    exc = future.exception()
    if exc is not None:
        print(f"Langfuse error: {exc}", flush=True)


def _flush_langfuse_background(langfuse: Langfuse):
    """
    langfuse.flush() is synchronous -- it blocks the calling thread until
    pending spans are exported, including Langfuse's own internal retry/
    backoff loop on a failed export. Called directly from the request path,
    that means an unreachable Langfuse host doesn't just fail to trace --
    it blocks this process's single event loop, stalling every other
    request being served concurrently, not just the one being flushed.
    Running it in a worker thread keeps a dead tracing backend from ever
    being able to affect real traffic, which synchronous flush() calls in
    the request path could not guarantee.
    """
    try:
        future = asyncio.get_running_loop().run_in_executor(None, langfuse.flush)
        future.add_done_callback(_log_background_flush_result)
    except Exception as e:
        print(f"Langfuse error: {e}", flush=True)

# ── httpx ─────────────────────────────────────────────────────────────────────
HTTPX_TIMEOUT = httpx.Timeout(connect=30.0, read=float(PROXY_TIMEOUT), write=60.0, pool=10.0)
http_client: httpx.AsyncClient = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global http_client
    http_client = httpx.AsyncClient(
        timeout=HTTPX_TIMEOUT,
        limits=httpx.Limits(max_connections=100, max_keepalive_connections=20, keepalive_expiry=30),
    )
    traced = [n for n, c in store.data["teams"].items() if c.get("langfuse", {}).get("enabled")]
    print(f"""
╔══════════════════════════════════════════════════╗
║        LLM Proxy  — port {PROXY_PORT}                       ║
╠══════════════════════════════════════════════════╣
║  Backends  : {", ".join(sorted(BACKENDS)) or "(none yet — add via /admin)":<36}║
║  Teams     : {", ".join(sorted(TEAMS)) or "(none yet — add via /admin)":<36}║
║  Tracing   : {(", ".join(sorted(traced)) + " (per-team Langfuse project)") if traced else "no team has tracing configured yet":<36}║
╚══════════════════════════════════════════════════╝
Admin UI:                http://<this-host>:{PROXY_PORT}/admin
Point developers at:     http://<this-host>:{PROXY_PORT}/<backend>/...
Per-team curl + models:  GET http://<this-host>:{PROXY_PORT}/whoami  (with their token)
Health check:            http://<this-host>:{PROXY_PORT}/health
""", flush=True)
    yield
    await http_client.aclose()
    for _client in _LANGFUSE_CLIENTS.values():
        _client.flush()


app = FastAPI(lifespan=lifespan)

# Only these paths are known to accept OpenAI's Chat-Completions-style
# stream_options.include_usage — used solely to decide whether it's safe to
# inject that field before streaming (see the "Streaming" section below).
# The Responses API always includes usage with no opt-in and rejects
# unrecognized top-level fields; Ollama's native routes don't know this
# field either and have their own way of reporting usage in the stream.
CHAT_COMPLETIONS_STREAM_PATHS = {
    "v1/chat/completions", "chat/completions",
    "v1/completions",      "completions",
}

RESERVED_BACKEND_NAMES = {"admin", "health", "whoami", "v1", "models"}


# ── Proxy-native endpoints (declared before the catch-all route below) ──────
@app.get("/health")
async def health():
    return {
        "status":          "ok",
        "public_base_url": PUBLIC_BASE_URL or None,
        "backends": {name: {"type": b.get("type", "openai"), "base_url": b["base_url"], "models": b["models"]}
                     for name, b in BACKENDS.items()},
        "teams":    {name: {"backends": cfg["backends"],
                             "tracing": store.data["teams"].get(name, {}).get("langfuse", {}).get("enabled", False)}
                     for name, cfg in TEAMS.items()},
    }


def _public_base_url(request: Request) -> str:
    """
    The base URL teams should use to reach this proxy. PUBLIC_BASE_URL, if
    set, wins outright — simplest and most predictable when you have a real
    domain, since it doesn't depend on a reverse proxy correctly forwarding
    X-Forwarded-Proto/Host. Otherwise this is derived from the request
    itself (trusting those headers, then falling back to the plain Host
    header) — never a hardcoded placeholder.
    """
    if PUBLIC_BASE_URL:
        return PUBLIC_BASE_URL
    proto = request.headers.get("x-forwarded-proto", request.url.scheme)
    host  = request.headers.get("x-forwarded-host") or request.headers.get("host") or request.url.netloc
    return f"{proto}://{host}"


@app.get("/whoami")
async def whoami(request: Request):
    """
    What a team needs to get started: their allowed backends, the models on
    each, and a ready-to-run curl example. Point a team at
    `GET /whoami -H "Authorization: Bearer <their token>"` instead of writing
    them onboarding docs by hand.
    """
    team, denied = _authenticate(request)
    if denied is not None:
        return denied

    token = _presented_token(request)
    base  = _public_base_url(request)
    out = []
    for name in TEAMS[team]["backends"]:
        b = BACKENDS.get(name)
        if not b:
            continue
        example_model = b["models"][0]["id"] if b["models"] else "<model>"
        out.append({
            "backend":  name,
            "base_url": f"{base}/{name}",
            "models":   b["models"],
            "example_curl": (
                f'curl {base}/{name}/v1/chat/completions '
                f'-H "Authorization: Bearer {token}" -H "Content-Type: application/json" '
                f'-d \'{{"model": "{example_model}", "messages": [{{"role": "user", "content": "hi"}}]}}\''
            ),
        })
    return {"team": team, "backends": out}


@app.get("/v1/models")
@app.get("/models")
async def list_models(request: Request):
    team, denied = _authenticate(request)
    if denied is not None:
        return denied

    data = []
    for name in TEAMS[team]["backends"]:
        b = BACKENDS.get(name)
        if not b:
            continue
        for m in b["models"]:
            data.append({
                "id": m["id"], "object": "model", "owned_by": name, "backend": name,
                "input_price_per_1m":  m.get("input_price", 0.0),
                "output_price_per_1m": m.get("output_price", 0.0),
            })
    return {"object": "list", "data": data}


# ── Admin UI + API ────────────────────────────────────────────────────────────
_ADMIN_HTML = (Path(__file__).parent / "admin.html").read_text()


@app.get("/admin", response_class=HTMLResponse)
async def admin_page():
    # The page itself carries no secrets — it prompts for the admin token and
    # calls /admin/api/* with it. Auth is enforced on those API calls, not here.
    return _ADMIN_HTML


def _mask(secret: str) -> str:
    if not secret:
        return ""
    return secret[:4] + "…" + secret[-4:] if len(secret) > 10 else "•" * len(secret)


class ModelPriceIn(BaseModel):
    id: str
    input_price: float = 0.0    # USD per 1,000,000 input tokens
    output_price: float = 0.0   # USD per 1,000,000 output tokens


class BackendIn(BaseModel):
    type: str = "openai"        # openai | vllm | ollama | anthropic | openai-compatible
    base_url: str
    api_key: str = ""           # PUT with "" on an existing backend keeps the current key
    models: list[ModelPriceIn] = []
    max_concurrency: int = DEFAULT_BACKEND_CONCURRENCY  # cap on requests in flight to this backend at once


class TeamIn(BaseModel):
    token: str = ""                    # PUT with "" on an existing team keeps its current token;
    backends: list[str] = []           # on create, "" means "generate one server-side"
    langfuse_host: str = ""            # this team's own Langfuse project — blank host/public_key
    langfuse_public_key: str = ""      # disables tracing for the team, same as leaving the whole
    langfuse_secret_key: str = ""      # section empty. PUT with secret_key="" keeps the current one.


@app.get("/admin/api/backends", dependencies=[Depends(require_admin)])
async def admin_list_backends():
    return {name: {**cfg, "api_key": _mask(cfg.get("api_key", ""))}
            for name, cfg in store.data["backends"].items()}


@app.put("/admin/api/backends/{name}", dependencies=[Depends(require_admin)])
async def admin_upsert_backend(name: str, body: BackendIn):
    name = name.strip()
    if not name or "/" in name or name in RESERVED_BACKEND_NAMES:
        raise HTTPException(status_code=400, detail=f"invalid or reserved backend name '{name}'")
    existing = store.data["backends"].get(name, {})
    store.data["backends"][name] = {
        "type":            body.type,
        "base_url":        body.base_url.rstrip("/"),
        "api_key":         body.api_key or existing.get("api_key", ""),
        "models":          [m.model_dump() for m in body.models],
        "max_concurrency": body.max_concurrency or DEFAULT_BACKEND_CONCURRENCY,
    }
    await store.save()
    _rebuild_indexes()
    _rebuild_backend_semaphore(name)
    return {"ok": True, "name": name}


@app.delete("/admin/api/backends/{name}", dependencies=[Depends(require_admin)])
async def admin_delete_backend(name: str):
    if name not in store.data["backends"]:
        raise HTTPException(status_code=404, detail="not found")
    in_use = [t for t, cfg in store.data["teams"].items() if name in cfg.get("backends", [])]
    if in_use:
        raise HTTPException(status_code=409, detail=f"backend in use by teams: {in_use}")
    del store.data["backends"][name]
    await store.save()
    _rebuild_indexes()
    _BACKEND_SEMAPHORES.pop(name, None)
    return {"ok": True}


@app.post("/admin/api/backends/{name}/test", dependencies=[Depends(require_admin)])
async def admin_test_backend(name: str, body: BackendIn):
    """
    Probes the backend with its own native model-listing route and returns
    what it finds, so the UI can offer "use these" instead of hand-typing a
    model list. openai/vllm/openai-compatible/anthropic all speak
    GET /v1/models (Anthropic's returns the same {"data": [...]} shape,
    just needs its own auth headers); ollama speaks GET /api/tags.
    """
    base_url = body.base_url.rstrip("/")
    existing = store.data["backends"].get(name, {})
    api_key  = body.api_key or existing.get("api_key", "")
    if body.type == "anthropic":
        headers = {"anthropic-version": ANTHROPIC_VERSION}
        if api_key:
            headers["x-api-key"] = api_key
    else:
        headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
    try:
        if body.type == "ollama":
            r = await http_client.get(f"{base_url}/api/tags", headers=headers, timeout=10)
            r.raise_for_status()
            models = sorted(m["name"] for m in r.json().get("models", []))
        else:
            r = await http_client.get(f"{base_url}/v1/models", headers=headers, timeout=10)
            r.raise_for_status()
            # Most servers wrap the list as {"data": [...]}; some OpenAI-
            # compatible ones (e.g. Together AI) return the bare list.
            data  = r.json()
            items = data if isinstance(data, list) else data.get("data", [])
            models = sorted(m["id"] for m in items)
        return {"ok": True, "models": models}
    except Exception as e:
        return {"ok": False, "error": str(e)}


@app.get("/admin/api/teams", dependencies=[Depends(require_admin)])
async def admin_list_teams():
    out = {}
    for name, cfg in store.data["teams"].items():
        lf = dict(cfg.get("langfuse", {}))
        lf["secret_key"] = _mask(lf.get("secret_key", ""))
        out[name] = {**cfg, "langfuse": lf}
    return out


@app.put("/admin/api/teams/{name}", dependencies=[Depends(require_admin)])
async def admin_upsert_team(name: str, body: TeamIn):
    name = name.strip()
    if not name:
        raise HTTPException(status_code=400, detail="invalid team name")
    unknown = [b for b in body.backends if b not in store.data["backends"]]
    if unknown:
        raise HTTPException(status_code=400, detail=f"unknown backend(s): {unknown}")
    existing    = store.data["teams"].get(name, {})
    existing_lf = existing.get("langfuse", {})
    token       = body.token or existing.get("token") or secrets.token_hex(24)
    # Blank public_key/secret_key on an update means "keep the current
    # value" (same convention as a blank token or backend api_key) — a
    # caller updating just one field (e.g. only the host, via a direct API
    # call rather than the admin UI, which always resends the full form)
    # must not silently wipe the others. host has no such fallback: an
    # explicitly blank host is the deliberate way to disable tracing
    # without discarding the keys, so it can be turned back on later by
    # setting the host again alone.
    public_key  = body.langfuse_public_key or existing_lf.get("public_key", "")
    secret_key  = body.langfuse_secret_key or existing_lf.get("secret_key", "")
    langfuse_cfg = {
        "enabled":    bool(body.langfuse_host and public_key and secret_key),
        "host":       body.langfuse_host,
        "public_key": public_key,
        "secret_key": secret_key,
    }
    store.data["teams"][name] = {"token": token, "backends": body.backends, "langfuse": langfuse_cfg}
    await store.save()
    _rebuild_indexes()
    _reinit_team_langfuse(name, replace=True)
    return {"ok": True, "name": name, "token": token, "tracing_enabled": langfuse_cfg["enabled"]}


@app.delete("/admin/api/teams/{name}", dependencies=[Depends(require_admin)])
async def admin_delete_team(name: str):
    if name not in store.data["teams"]:
        raise HTTPException(status_code=404, detail="not found")
    del store.data["teams"][name]
    await store.save()
    _rebuild_indexes()
    _LANGFUSE_CLIENTS.pop(name, None)
    return {"ok": True}


class SettingsIn(BaseModel):
    llm_endpoints: list[str]


@app.get("/admin/api/settings", dependencies=[Depends(require_admin)])
async def admin_get_settings():
    return {"llm_endpoints": sorted(store.data.get("llm_endpoints") or DEFAULT_LLM_ENDPOINTS)}


@app.put("/admin/api/settings", dependencies=[Depends(require_admin)])
async def admin_update_settings(body: SettingsIn):
    # Paths are matched with a leading "/" already stripped (see rest_key in
    # the proxy route below), so normalize here too — "/v1/rerank" and
    # "v1/rerank" typed into the same field should both work.
    cleaned = sorted({e.strip().lstrip("/") for e in body.llm_endpoints if e.strip()})
    if not cleaned:
        raise HTTPException(status_code=400, detail="at least one endpoint must remain allowed")
    store.data["llm_endpoints"] = cleaned
    await store.save()
    _rebuild_indexes()
    return {"ok": True, "llm_endpoints": cleaned}


# ── Helpers ────────────────────────────────────────────────────────────────────
def _fwd_headers(h: dict, api_key: str, backend_type: str = "openai") -> dict:
    """
    Strip hop-by-hop headers and the caller's proxy token, then attach the
    real backend key (if that backend needs one). Their token must never
    reach a real backend.

    Anthropic doesn't use OpenAI-style bearer auth at all -- it authenticates
    with an `x-api-key` header and requires an `anthropic-version` header on
    every request, so it gets its own branch here instead of the usual
    `Authorization: Bearer` every other backend type uses.
    """
    out = dict(h)
    for k in ("host", "content-length", "transfer-encoding", "connection",
              "authorization", "api-key", "x-api-key", "anthropic-version"):
        out.pop(k, None)
    if backend_type == "anthropic":
        if api_key:
            out["x-api-key"] = api_key
        out["anthropic-version"] = ANTHROPIC_VERSION
    elif api_key:
        out["authorization"] = f"Bearer {api_key}"
    return out


def _resp_headers(h) -> dict:
    """
    httpx already decompressed the body, so forwarding the upstream's
    content-encoding makes clients try to gunzip plain JSON and fail.
    """
    out = dict(h)
    for k in ("content-encoding", "content-length", "transfer-encoding", "connection"):
        out.pop(k, None)
    return out


def _find_usage_dict(parsed: dict):
    """
    Locates whatever holds the token counts, regardless of API shape.
    Chat Completions and the Responses API both nest a "usage" object;
    Ollama's native /api/chat and /api/generate don't nest anything — the
    counts (prompt_eval_count, eval_count) sit directly on the response body.
    """
    if not isinstance(parsed, dict):
        return None
    if isinstance(parsed.get("usage"), dict):
        return parsed["usage"]
    if "prompt_eval_count" in parsed or "eval_count" in parsed:
        return parsed
    return None


def _parse_usage(usage):
    """
    Every API names these token counts differently:
    - Chat Completions:    prompt_tokens / completion_tokens
    - Responses API:       input_tokens / output_tokens
    - Ollama (native):     prompt_eval_count / eval_count
    Accept any of them.

    Anthropic also reports prompt-cache tokens separately from input_tokens
    (cache_read_input_tokens / cache_creation_input_tokens); those are kept
    as their own usage keys so they're counted and priced, not dropped.
    """
    if not usage:
        return None

    def first(*keys):
        # First count that's actually present and non-null -- some servers
        # send e.g. "prompt_tokens": null, which used to become None and then
        # break the cost arithmetic.
        for k in keys:
            if usage.get(k) is not None:
                return usage[k]
        return 0

    input_tok  = first("prompt_tokens", "input_tokens", "prompt_eval_count")
    output_tok = first("completion_tokens", "output_tokens", "eval_count")
    cache_read  = usage.get("cache_read_input_tokens") or 0
    cache_write = usage.get("cache_creation_input_tokens") or 0
    total_tok  = usage.get("total_tokens") or (input_tok + output_tok + cache_read + cache_write)
    out = {"input": input_tok, "output": output_tok, "total": total_tok}
    if cache_read:
        out["cache_read_input_tokens"] = cache_read
    if cache_write:
        out["cache_creation_input_tokens"] = cache_write
    return out


def _compute_cost(backend_cfg: dict, model: str, usage: dict):
    """
    Langfuse only auto-prices models it recognizes by name — a self-hosted
    model, a custom fine-tune, or a typo'd model id all silently show $0.
    So pricing is looked up from what's configured on the backend itself
    (set per-model in /admin, in USD per 1,000,000 tokens) and forwarded as
    an explicit cost, bypassing Langfuse's own lookup entirely.

    Returns None — meaning "don't override, let Langfuse do its own
    lookup" — whenever the model isn't registered here or has no price set,
    so a genuinely free/self-hosted model or a model Langfuse already prices
    correctly (e.g. real gpt-4o with nothing entered) isn't clobbered with
    an explicit $0.
    """
    if not usage:
        return None
    models = backend_cfg.get("models", [])
    price = next((m for m in models if m.get("id") == model), None)
    if price is None and isinstance(model, str):
        # "GLM-5.3-Flash" vs a configured "glm-5.3-flash" is the same model;
        # an exact match above still wins when both exist.
        price = next((m for m in models if str(m.get("id", "")).lower() == model.lower()), None)
    if not price:
        return None
    input_price  = price.get("input_price")  or 0.0
    output_price = price.get("output_price") or 0.0
    if not input_price and not output_price:
        return None
    input_cost  = round((usage["input"]  / 1_000_000) * input_price,  8)
    output_cost = round((usage["output"] / 1_000_000) * output_price, 8)
    cost = {"input": input_cost, "output": output_cost}
    # Anthropic prompt caching, priced off the model's input price at
    # Anthropic's standard rates: cache reads 0.1x, (5-minute) cache writes 1.25x.
    if usage.get("cache_read_input_tokens"):
        cost["cache_read_input_tokens"] = round(
            (usage["cache_read_input_tokens"] / 1_000_000) * input_price * 0.1, 8)
    if usage.get("cache_creation_input_tokens"):
        cost["cache_creation_input_tokens"] = round(
            (usage["cache_creation_input_tokens"] / 1_000_000) * input_price * 1.25, 8)
    cost["total"] = round(sum(cost.values()), 8)
    return cost


def _extract_output(parsed):
    if not isinstance(parsed, dict):
        return parsed
    choices = parsed.get("choices")
    if choices:
        return choices[0].get("message") or choices[0].get("delta") or parsed
    # Responses API: no "choices" — "output" is a list of items (messages,
    # tool calls, reasoning blocks, ...); "output_text" is a plain-text
    # convenience field some clients add, not guaranteed to be present.
    if "output" in parsed:
        return parsed.get("output_text") or parsed["output"]
    # Ollama native: /api/chat replies with "message", /api/generate with
    # a plain "response" string.
    if "message" in parsed:
        return parsed["message"]
    if "response" in parsed:
        return parsed["response"]
    # Anthropic Messages API: reply text lives in a list of content blocks
    # (text/tool_use/...), not under "choices"/"output"/"message".
    if "content" in parsed:
        return parsed["content"]
    # Rerank (Cohere/vLLM/TEI-object-shaped): no generated text, just
    # ranked {index, relevance_score}/{index, score} items.
    if "results" in parsed:
        return parsed["results"]
    return parsed


def _log(method, path, status, ms, model, service):
    ts = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    print(f"[{ts}] {method} /{path} | {status} | {ms}ms | model={model} service={service}", flush=True)


def _parse_stream_buffer(full_bytes: bytes) -> tuple:
    """
    Parses a streamed response body regardless of which of these four
    shapes it's actually in — the proxy doesn't know in advance which one
    a given backend/endpoint uses:
    - Chat Completions (SSE, "data: {...}" lines): a bare
      {"usage": ..., "choices": [...]} chunk when
      stream_options.include_usage is set.
    - Responses API (SSE, "data: {...}" lines): a sequence of typed
      events; the final one nests the complete response — including
      usage — under a "response" key instead of at the top level, with
      no opt-in needed.
    - Ollama native (NDJSON, no "data:" prefix — one raw JSON object per
      line): the final line (done: true) carries prompt_eval_count /
      eval_count and the assembled message directly at the top level.
    - Anthropic Messages API (SSE, "data: {...}" lines): a sequence of
      typed events with no single chunk carrying the full picture —
      input_tokens arrives once on "message_start" (nested under
      message.usage), output_tokens is only ever a running total on each
      "message_delta" (so the last one wins), and the reply text itself
      arrives as incremental "content_block_delta" fragments that have to
      be concatenated back into one string, unlike every other shape here
      where some single chunk already holds the complete text.

    Returns (last_relevant_chunk_parsed, usage_dict_or_None).
    """
    parsed     = None
    usage_data = None
    anthropic_input_tokens  = None
    anthropic_output_tokens = None
    anthropic_cache_read    = None
    anthropic_cache_write   = None
    anthropic_text          = []
    for line in full_bytes.decode(errors="ignore").splitlines():
        line = line.strip()
        if not line or "[DONE]" in line:
            continue
        chunk_str = line[len("data:"):].strip() if line.startswith("data:") else line
        try:
            chunk_json = json.loads(chunk_str)
        except Exception:
            continue
        if chunk_json.get("usage"):
            usage_data = chunk_json["usage"]
        if chunk_json.get("choices"):
            parsed = chunk_json
        response_obj = chunk_json.get("response")
        if isinstance(response_obj, dict):
            if response_obj.get("usage"):
                usage_data = response_obj["usage"]
            if response_obj.get("output") is not None:
                parsed = response_obj
        if "prompt_eval_count" in chunk_json or "eval_count" in chunk_json:
            usage_data = chunk_json
        if "message" in chunk_json or ("response" in chunk_json and not isinstance(response_obj, dict)):
            parsed = chunk_json
        message_obj = chunk_json.get("message")
        if isinstance(message_obj, dict) and isinstance(message_obj.get("usage"), dict):
            anthropic_input_tokens = message_obj["usage"].get("input_tokens", anthropic_input_tokens)
            anthropic_output_tokens = message_obj["usage"].get("output_tokens", anthropic_output_tokens)
            anthropic_cache_read  = message_obj["usage"].get("cache_read_input_tokens", anthropic_cache_read)
            anthropic_cache_write = message_obj["usage"].get("cache_creation_input_tokens", anthropic_cache_write)
        if chunk_json.get("type") == "message_delta" and isinstance(chunk_json.get("usage"), dict):
            anthropic_output_tokens = chunk_json["usage"].get("output_tokens", anthropic_output_tokens)
        delta_obj = chunk_json.get("delta")
        if isinstance(delta_obj, dict) and delta_obj.get("type") == "text_delta":
            anthropic_text.append(delta_obj.get("text", ""))
    if anthropic_input_tokens is not None or anthropic_output_tokens is not None:
        usage_data = {"input_tokens": anthropic_input_tokens or 0, "output_tokens": anthropic_output_tokens or 0}
        if anthropic_cache_read:
            usage_data["cache_read_input_tokens"] = anthropic_cache_read
        if anthropic_cache_write:
            usage_data["cache_creation_input_tokens"] = anthropic_cache_write
    if anthropic_text:
        parsed = {"content": [{"type": "text", "text": "".join(anthropic_text)}]}
    return parsed, usage_data


# ── Proxy ──────────────────────────────────────────────────────────────────────
# Requests are routed by path: /<backend>/<rest...> forwards to that
# backend's base_url + /<rest...>, once the caller's token is confirmed
# authorised for <backend>. This keeps a client's `model` field untouched
# (many real model ids already contain a "/", e.g. "meta-llama/Llama-3-70b"),
# so routing never has to guess where a slash belongs.
@app.api_route("/{backend}/{rest:path}", methods=["GET","POST","PUT","DELETE","PATCH","HEAD","OPTIONS"])
async def proxy(backend: str, rest: str, request: Request, background_tasks: BackgroundTasks):
    if backend not in BACKENDS:
        return JSONResponse(status_code=404, content={"error": {
            "message": f"Unknown backend '{backend}'. See /health or /whoami for available backends.",
            "type":    "invalid_request_error",
        }})

    if request.method != "OPTIONS":
        team, denied = _authenticate(request)
        if denied is not None:
            _log(request.method, f"{backend}/{rest}", 401, 0, "-", "unauthorised")
            return denied
        if backend not in TEAMS[team]["backends"]:
            _log(request.method, f"{backend}/{rest}", 403, 0, "-", team)
            return JSONResponse(status_code=403, content={"error": {
                "message": f"Team '{team}' is not authorised for backend '{backend}'. "
                           f"Allowed: {TEAMS[team]['backends']}",
                "type":    "permission_error",
            }})
    else:
        team = "unknown"

    backend_cfg = BACKENDS[backend]
    url     = f"{backend_cfg['base_url']}/{rest}"
    headers = _fwd_headers(dict(request.headers), backend_cfg["api_key"], backend_cfg.get("type", "openai"))
    body_b  = await request.body()
    rest_key = rest.lstrip("/")
    is_llm   = rest_key in LLM_ENDPOINTS

    body_j = {}
    if body_b and is_llm:
        try:
            body_j = json.loads(body_b)
        except Exception:
            pass

    # ── Reject unsupported endpoints ─────────────────────────────────────────
    # Only the paths this proxy actually knows how to trace are allowed
    # through at all — e.g. /v1/completions and /v1/responses, not whatever
    # else a backend happens to expose. A path outside that list gets a
    # clean 400 instead of being silently forwarded untraced.
    if not is_llm:
        _log(request.method, f"{backend}/{rest}", 400, 0, "-", team)
        return JSONResponse(status_code=400, content={"error": {
            "message": f"'/{rest}' is not a supported endpoint on this proxy. "
                       f"Supported: {sorted(LLM_ENDPOINTS)}",
            "type":    "invalid_request_error",
            "code":    "unsupported_endpoint",
        }})

    # ── LLM path ──────────────────────────────────────────────────────────────
    model      = body_j.get("model", "unknown")
    # Chat Completions sends "messages"; the Responses API sends "input"
    # (a string or a list of role/content items) and no "messages" at all;
    # rerank sends neither — just "query" (+ "documents", left out of the
    # trace input to avoid bloating it with a full document list).
    messages   = body_j.get("messages") or body_j.get("input") or body_j.get("prompt") or body_j.get("query")
    trace_name = request.headers.get("x-trace-name") or model or "llm-request"
    user_id    = request.headers.get("x-user-id") or None
    # Derived from the validated token, so a client cannot spoof it.
    service    = team
    tags       = [t for t in [model, backend, service] if t]
    is_stream  = body_j.get("stream", False)
    start_ms   = time.time()
    # Each team traces into its own Langfuse project; a team with no keys
    # configured gets a no-op client, so nothing below needs to branch on it.
    langfuse   = _get_langfuse(team)

    # ── Streaming ─────────────────────────────────────────────────────────────
    if is_stream:
        if rest_key in CHAT_COMPLETIONS_STREAM_PATHS:
            # Ask the backend to include token usage in the final chunk.
            # Only done for paths known to accept this field — the
            # Responses API and Ollama's native routes either already
            # include usage with no opt-in, or might reject an
            # unrecognized top-level field outright.
            body_j["stream_options"] = {"include_usage": True}
            body_b = json.dumps(body_j).encode()

        collected = []
        err_box   = [None]

        # Deliberately NOT a `with` block. `stream_gen` below is an async
        # generator — its body doesn't run until Starlette iterates it while
        # writing the response, which happens after this function returns.
        # A `with langfuse.start_as_current_observation(...):` wrapped around
        # that `return` would exit (ending the span, locking in ~0 latency)
        # immediately on return, before any real streaming — or its eventual
        # .update() — ever happened. Managing the spans manually and ending
        # them from inside stream_gen()'s `finally` keeps them open for the
        # actual duration of the stream instead.
        with propagate_attributes(user_id=user_id, tags=tags):
            root_span = langfuse.start_observation(
                as_type  = "span",
                name     = trace_name,
                input    = messages,
                metadata = {"path": f"/{backend}/{rest}", "model": model, "backend": backend, "service": service},
            )
            gen_span = root_span.start_observation(
                as_type          = "generation",
                name             = "chat-completion",
                model            = model,
                input            = messages,
                model_parameters = {
                    "temperature": body_j.get("temperature"),
                    "max_tokens":  body_j.get("max_tokens"),
                    "top_p":       body_j.get("top_p"),
                },
            )

        async def stream_gen():
            try:
                # Held for the whole streamed call, not just the connect —
                # a slow backend keeps its slot occupied for as long as it's
                # actually generating, which is the point of the limit.
                async with _backend_semaphore(backend):
                    async with http_client.stream(
                        method=request.method, url=url, headers=headers, content=body_b
                    ) as r:
                        async for chunk in r.aiter_bytes():
                            collected.append(chunk)
                            yield chunk
            except Exception as e:
                err_box[0] = str(e)
                raise
            finally:
                ms     = int((time.time() - start_ms) * 1000)
                try:
                    # Parsed inside the try: a malformed stream must not be
                    # able to skip the .end() calls below and leave the
                    # trace open (and so never exported).
                    full   = b"".join(collected)
                    parsed, usage_data = _parse_stream_buffer(full)
                    usage  = _parse_usage(usage_data)
                    cost   = _compute_cost(backend_cfg, model, usage)
                    output = _extract_output(parsed)
                except Exception as e:
                    print(f"Langfuse error: could not parse stream for tracing: {e}", flush=True)
                    output, usage, cost = None, None, None
                try:
                    # Even when the connection errors out afterward, still
                    # attach whatever output/usage/cost was actually parsed
                    # from the bytes received so far. The backend may have
                    # already generated -- and been billed for -- the full
                    # completion before the connection hiccuped on the way
                    # out; a real cost sitting right here shouldn't be
                    # thrown away just because of what happened after it.
                    if err_box[0]:
                        gen_span.update(level="ERROR", status_message=err_box[0], output=output,
                                        usage_details=usage, cost_details=cost, metadata={"latency_ms": ms})
                        root_span.update(level="ERROR", output=output, metadata={"latency_ms": ms})
                    else:
                        gen_span.update(output=output, usage_details=usage, cost_details=cost,
                                        metadata={"latency_ms": ms})
                        root_span.update(output=output, metadata={"latency_ms": ms})
                    gen_span.end()
                    root_span.end()
                    _flush_langfuse_background(langfuse)
                except Exception as e:
                    print(f"Langfuse error: {e}", flush=True)
                _log(request.method, f"{backend}/{rest}", "stream", ms, model, service)

        return StreamingResponse(stream_gen(), media_type="text/event-stream")

    # ── Non-streaming ──────────────────────────────────────────────────────────
    with propagate_attributes(user_id=user_id, tags=tags):
        with langfuse.start_as_current_observation(
            as_type  = "span",
            name     = trace_name,
            input    = messages,
            metadata = {"path": f"/{backend}/{rest}", "model": model, "backend": backend, "service": service},
        ) as root_span:
            with langfuse.start_as_current_observation(
                as_type          = "generation",
                name             = "chat-completion",
                model            = model,
                input            = messages,
                model_parameters = {
                    "temperature": body_j.get("temperature"),
                    "max_tokens":  body_j.get("max_tokens"),
                    "top_p":       body_j.get("top_p"),
                },
            ) as gen_span:
                try:
                    async with _backend_semaphore(backend):
                        r = await http_client.request(
                            method=request.method, url=url, headers=headers, content=body_b
                        )
                    ms     = int((time.time() - start_ms) * 1000)
                    parsed = None
                    try:
                        parsed = r.json()
                    except Exception:
                        pass

                    level  = "ERROR" if r.status_code >= 400 else "DEFAULT"
                    output = _extract_output(parsed)
                    usage  = _parse_usage(_find_usage_dict(parsed))
                    cost   = _compute_cost(backend_cfg, model, usage)

                    gen_span.update(
                        output         = output,
                        level          = level,
                        status_message = r.text if r.status_code >= 400 else None,
                        usage_details  = usage,
                        cost_details   = cost,
                        metadata       = {"latency_ms": ms, "status_code": r.status_code},
                    )
                    root_span.update(output=output, level=level, metadata={"latency_ms": ms})

                    _log(request.method, f"{backend}/{rest}", r.status_code, ms, model, service)
                    result = Response(content=r.content, status_code=r.status_code,
                                      headers=_resp_headers(r.headers))

                except Exception as e:
                    ms = int((time.time() - start_ms) * 1000)
                    gen_span.update(level="ERROR", status_message=str(e), metadata={"latency_ms": ms})
                    root_span.update(level="ERROR", metadata={"latency_ms": ms})
                    print(f"Proxy error: {e}", flush=True)
                    result = JSONResponse(status_code=502, content={"error": "proxy_error", "message": str(e)})

    # Flushed only now that both `with` blocks have exited and ended the
    # spans -- flushing inside them exported everything except this
    # request's own trace, which then waited for the SDK's next timed export.
    _flush_langfuse_background(langfuse)
    return result


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=PROXY_PORT)