main.py

"""
FastAPI proxy service for routing Claude Code requests to Z.AI API
"""
import os
import asyncio
import logging
import random
from typing import Any, Dict
from fastapi import FastAPI, Request, Response, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
import httpx

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

app = FastAPI(title="Claude Code → Z.AI Proxy")

# Z.AI API configuration
ZAI_API_KEY = os.getenv("ZAI_API_KEY")
ZAI_BASE_URL = "https://api.z.ai/api/anthropic/v1/messages"

# Model mapping: Claude models → Z.AI GLM models
MODEL_MAP = {
    "claude-haiku-4-5-20251001": "glm-4.7-flash",
    "claude-sonnet-4-5-20250929": "glm-4.6",
    "claude-opus-4-5-20251101": "glm-4.7",
}

if not ZAI_API_KEY:
    logger.error("❌ ZAI_API_KEY environment variable is not set!")
    raise ValueError("ZAI_API_KEY environment variable is required")

logger.info(f"✓ Z.AI API Key loaded: {ZAI_API_KEY[:20]}...")


@app.get("/")
async def health_check():
    """Health check endpoint."""
    return {"status": "ok", "service": "Claude Code → Z.AI Proxy"}


async def call_z_ai_with_retry(
    client: httpx.AsyncClient,
    url: str,
    json_data: dict,
    headers: dict,
    max_retries: int = 5,
    base_delay: float = 2.0
) -> httpx.Response:
    """
    Calls Z.AI API with exponential backoff on 429 errors.
    
    Args:
        client: httpx AsyncClient
        url: Z.AI API endpoint URL
        json_data: Request body
        headers: Request headers
        max_retries: Maximum number of retry attempts (default: 5)
        base_delay: Base delay in seconds (default: 2.0, will double: 2s, 4s, 8s, 16s, 32s)
    
    Returns:
        httpx.Response object
    """
    for attempt in range(max_retries):
        try:
            response = await client.post(url, json=json_data, headers=headers)
            
            if response.status_code == 429:
                if attempt < max_retries - 1:
                    delay = base_delay * (2 ** attempt)
                    # Add random jitter: ±20%
                    jitter = delay * random.uniform(0.8, 1.2)
                    logger.warning(
                        f"⚠️  429 Too Many Requests (attempt {attempt + 1}/{max_retries}). "
                        f"Retrying in {jitter:.1f}s..."
                    )
                    await asyncio.sleep(jitter)
                    continue
                else:
                    logger.error("❌ Max retries exceeded for 429 error")
                    return response
            
            # Success or other non-429 error
            return response
            
        except Exception as e:
            logger.error(f"Request exception (attempt {attempt + 1}): {e}")
            if attempt < max_retries - 1:
                delay = base_delay * (2 ** attempt)
                jitter = delay * random.uniform(0.8, 1.2)
                await asyncio.sleep(jitter)
                continue
            raise
    
    return response


@app.post("/v1/messages")
async def proxy_messages(request: Request):
    """
    Main proxy endpoint for Claude Code messages.
    Transforms Claude models to Z.AI GLM models and forwards requests.
    """
    try:
        body = await request.json()
        model = body.get("model", "")
        
        # Log request details
        logger.info(f"REQUEST: model={model}, stream={body.get('stream', False)}")
        
        # Map Claude model to Z.AI model
        z_ai_model = MODEL_MAP.get(model, "glm-4.7-flash")
        logger.info(f"PROXY CALL: from={model} -> to={z_ai_model}, temp={body.get('temperature', 1.0)}")
        
        # Update model in request body
        body["model"] = z_ai_model
        
        # Prepare headers
        headers = {
            "Authorization": f"Bearer {ZAI_API_KEY}",
            "Content-Type": "application/json",
        }
        
        # Log stream mode
        logger.info(f"Stream: {body.get('stream', False)}")
        
        # Call Z.AI API with retry logic
        async with httpx.AsyncClient(timeout=120.0) as client:
            response = await call_z_ai_with_retry(
                client=client,
                url=ZAI_BASE_URL,
                json_data=body,
                headers=headers
            )
        
        # Handle response
        if response.status_code == 200:
            logger.info("✓ Z.AI response: status=200, success=True")
            
            # Stream or non-stream response
            if body.get("stream", False):
                return StreamingResponse(
                    response.aiter_bytes(),
                    media_type="application/x-ndjson"
                )
            else:
                return JSONResponse(response.json())
        
        else:
            error_body = response.text
            try:
                error_body = response.json()
            except:
                pass
            
            logger.error(
                f"❌ Z.AI API error: status={response.status_code}, body={error_body}"
            )
            
            return JSONResponse(
                status_code=response.status_code,
                content={
                    "error": {
                        "type": "api_error",
                        "message": f"Z.AI API error: {response.status_code}",
                        "details": error_body
                    }
                }
            )
    
    except Exception as e:
        logger.error(f"❌ Proxy error: {e}", exc_info=True)
        return JSONResponse(
            status_code=500,
            content={
                "error": {
                    "type": "proxy_error",
                    "message": str(e)
                }
            }
        )


@app.post("/api/event_logging/batch")
async def event_logging_batch(request: Request):
    """
    Dummy event logging endpoint.
    Claude Code CLI sends telemetry/analytics here.
    We accept it but don't process it.
    """
    try:
        body = await request.json()
        logger.debug(f"📊 Event logging: {len(body.get('events', []))} events received")
        return {"success": True}
    except Exception as e:
        logger.debug(f"Event logging error (ignored): {e}")
        return {"success": True}


@app.api_route("/{path:path}", methods=["POST", "GET", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"])
async def catch_all(path: str, request: Request):
    """
    Catch-all route for unhandled endpoints.
    Returns 200 OK to prevent errors from unknown Claude Code requests.
    """
    method = request.method
    logger.debug(f"🔄 Catch-all: {method} /{path}")
    
    # For most requests, just return success
    if method in ["GET", "HEAD"]:
        return {"status": "ok"}
    else:
        # For POST/PUT/DELETE, try to parse body for debugging
        try:
            body = await request.json()
            logger.debug(f"   Body: {body}")
        except:
            pass
        
        return {"status": "ok"}


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