main.py

import os
import random
import time
import logging

import httpx
from fastapi import FastAPI, Request, Response
from fastapi.responses import StreamingResponse

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("oc-proxy")

UPSTREAM_URL = os.environ.get("UPSTREAM_URL", "https://opencode.ai/zen/go/v1")
MAX_RETRIES = int(os.environ.get("MAX_RETRIES", "3"))
CROCKFORD_ALPHABET = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"

app = FastAPI()


def encode_base32(value: int, length: int = 10) -> str:
    result = ""
    for _ in range(length):
        result = CROCKFORD_ALPHABET[value & 31] + result
        value >>= 5
    return result


def random_base32(length: int = 17) -> str:
    return "".join(CROCKFORD_ALPHABET[random.randint(0, 31)] for _ in range(length))


def generate_session_id() -> str:
    ts = int(time.time() * 1000)
    encoded = encode_base32(0xFFFFFFFFFFFF - ts)
    return f"ses_{encoded}{random_base32()}"


def generate_request_id() -> str:
    ts = int(time.time() * 1000)
    encoded = encode_base32(ts)
    return f"msg_{encoded}{random_base32()}"


session_id = generate_session_id()
logger.info(f"Initial session ID: {session_id}")


def regenerate_credentials():
    global session_id
    session_id = generate_session_id()
    logger.info(f"Regenerated session ID: {session_id}")


@app.post("/v1/chat/completions")
async def proxy_chat(request: Request):
    return await _do_proxy(request, "chat/completions")


@app.get("/v1/models")
async def proxy_models(request: Request):
    response = await _do_proxy(request, "models")
    if response.status_code == 200:
        import json
        data = json.loads(response.body)
        data["data"] = [m for m in data.get("data", []) if "free" in m.get("id", "").lower()]
        return Response(
            content=json.dumps(data),
            status_code=200,
            headers=dict(response.headers),
            media_type="application/json",
        )
    return response


async def _do_proxy(request: Request, path: str):
    global session_id

    body = await request.body()
    headers = dict(request.headers)
    headers.pop("host", None)
    headers.pop("content-length", None)
    headers.pop("accept-encoding", None)

    target_url = f"{UPSTREAM_URL}/{path}"
    if request.url.query:
        target_url += f"?{request.url.query}"

    is_stream = b'"stream":true' in body or b'"stream": true' in body

    for attempt in range(MAX_RETRIES):
        request_id = generate_request_id()
        req_headers = {
            **headers,
            "x-opencode-session": session_id,
            "x-opencode-request": request_id,
        }

        logger.info(f"Attempt {attempt + 1}/{MAX_RETRIES} -> {request.method} {target_url}")

        if is_stream:
            client = httpx.AsyncClient(timeout=300.0)
            req = client.build_request(
                method=request.method,
                url=target_url,
                headers=req_headers,
                content=body,
            )
            response = await client.send(req, stream=True)

            if response.status_code == 429 or response.status_code >= 500:
                await response.aclose()
                await client.aclose()
                logger.warning(f"Got {response.status_code}, regenerating credentials...")
                regenerate_credentials()
                continue

            resp_headers = dict(response.headers)
            resp_headers.pop("content-length", None)
            resp_headers.pop("transfer-encoding", None)
            resp_headers.pop("content-encoding", None)

            async def stream_response():
                try:
                    async for chunk in response.aiter_bytes():
                        yield chunk
                finally:
                    await response.aclose()
                    await client.aclose()

            return StreamingResponse(
                stream_response(),
                status_code=response.status_code,
                headers=resp_headers,
            )
        else:
            async with httpx.AsyncClient(timeout=300.0) as client:
                response = await client.request(
                    method=request.method,
                    url=target_url,
                    headers=req_headers,
                    content=body,
                )

            if response.status_code == 429 or response.status_code >= 500:
                logger.warning(f"Got {response.status_code}, regenerating credentials...")
                regenerate_credentials()
                continue

            resp_headers = dict(response.headers)
            resp_headers.pop("content-length", None)
            resp_headers.pop("transfer-encoding", None)
            resp_headers.pop("content-encoding", None)

            return Response(
                content=response.content,
                status_code=response.status_code,
                headers=resp_headers,
            )

    logger.error("All retries exhausted")
    return Response(content=b'{"error": "all retries exhausted"}', status_code=502)


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


def run():
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8080)