proxy.py

"""
opencode2claude - Anthropic Messages API → OpenCode (Session API) 代理服务

OpenCode 是本地运行的 AI 编程助手，默认监听 localhost:4096。
本代理将 Claude Code 的 Anthropic 格式请求转发到 OpenCode 的 Session API。

启动方式：
    1. 先启动 OpenCode: opencode
    2. 再启动代理: uvicorn proxy:app --host 0.0.0.0 --port 8082

配置 Claude Code：
    $env:ANTHROPIC_BASE_URL = "http://localhost:8082"
    $env:ANTHROPIC_API_KEY = "placeholder"
    claude
"""

from __future__ import annotations

import json
import os
import uuid
import logging
import asyncio
from contextlib import asynccontextmanager
from typing import AsyncIterator

import httpx
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from dotenv import load_dotenv

from converter.stream import estimate_input_tokens

# ── 加载环境变量 ──
load_dotenv()

# ── 配置 ──
OPENCODE_BASE_URL = os.getenv("OPENCODE_BASE_URL", "http://localhost:4096")
TARGET_MODEL = os.getenv(
    "TARGET_MODEL", "minimax/MiniMax-M1"
)  # OpenCode 模型格式: provider/model
FAKE_MODEL = os.getenv("FAKE_MODEL", "claude-opus-4-5-20250213")
PROXY_PORT = int(os.getenv("PROXY_PORT", "8082"))
DEBUG = os.getenv("DEBUG", "false").lower() == "true"

# ── 日志 ──
logging.basicConfig(
    level=logging.DEBUG if DEBUG else logging.INFO,
    format="%(asctime)s | %(levelname)-7s | %(message)s",
)
logger = logging.getLogger("opencode2claude")


# ── HTTP 客户端 ──
@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理。"""
    app.state.http_client = httpx.AsyncClient(
        timeout=httpx.Timeout(connect=10.0, read=300.0, write=30.0, pool=10.0),
        limits=httpx.Limits(max_keepalive_connections=20, max_connections=100),
    )
    logger.info(f"代理服务启动 | OpenCode: {OPENCODE_BASE_URL} | 模型: {TARGET_MODEL}")

    # 检查 OpenCode 是否运行
    try:
        resp = await app.state.http_client.get(f"{OPENCODE_BASE_URL}/session")
        if resp.status_code == 200:
            logger.info(f"OpenCode 连接成功")
        else:
            logger.warning(f"OpenCode 返回: {resp.status_code}")
    except Exception as e:
        logger.warning(f"无法连接 OpenCode ({OPENCODE_BASE_URL}): {e}")
        logger.warning("请确保 OpenCode 已启动: opencode")

    yield
    await app.state.http_client.aclose()
    logger.info("代理服务已关闭")


app = FastAPI(
    title="opencode2claude",
    description="将 Claude Code 请求转发到本地 OpenCode",
    version="2.0.0",
    lifespan=lifespan,
)


# ────────────────────────────────────────────────────────────────────────────────
# OpenCode Session 管理
# ────────────────────────────────────────────────────────────────────────────────


async def get_or_create_session(client: httpx.AsyncClient) -> str:
    """每次请求都创建新的 OpenCode session，避免上下文累积导致超时。"""
    # 每次都创建新 session，不复用旧的
    # 因为 OpenCode session 会累积上下文，导致后续请求越来越慢
    resp = await client.post(
        f"{OPENCODE_BASE_URL}/session",
        json={"title": f"Claude Code Proxy - {uuid.uuid4().hex[:8]}"},
    )
    if resp.status_code not in (200, 201):
        raise HTTPException(
            status_code=502, detail=f"Failed to create OpenCode session: {resp.text}"
        )

    data = resp.json()
    session_id = data.get("id") or data.get("data", {}).get("id")
    if not session_id:
        raise HTTPException(
            status_code=502, detail=f"OpenCode session creation returned no ID: {data}"
        )

    logger.info(f"创建新 OpenCode session: {session_id}")
    return session_id


# ────────────────────────────────────────────────────────────────────────────────
# 格式转换
# ────────────────────────────────────────────────────────────────────────────────


def anthropic_to_opencode_parts(body: dict) -> list[dict]:
    """将 Anthropic messages 转换为 OpenCode parts 格式。"""
    parts = []

    # 添加 system prompt
    system = body.get("system")
    if system:
        if isinstance(system, str):
            parts.append({"type": "text", "text": f"[System]\n{system}\n\n"})
        elif isinstance(system, list):
            text = "\n".join(
                b.get("text", "") for b in system if b.get("type") == "text"
            )
            if text:
                parts.append({"type": "text", "text": f"[System]\n{text}\n\n"})

    # 转换 messages
    for msg in body.get("messages", []):
        role = msg.get("role", "user")
        content = msg.get("content", "")

        if isinstance(content, str):
            text = content
        elif isinstance(content, list):
            # 提取文本内容
            text_parts = []
            for block in content:
                if block.get("type") == "text":
                    text_parts.append(block.get("text", ""))
                elif block.get("type") == "tool_result":
                    # 工具结果
                    result = block.get("content", "")
                    if isinstance(result, list):
                        result = "\n".join(
                            b.get("text", "") for b in result if b.get("type") == "text"
                        )
                    text_parts.append(
                        f"[Tool Result: {block.get('tool_use_id', '')}]\n{result}"
                    )
                elif block.get("type") == "tool_use":
                    # 工具调用
                    text_parts.append(
                        f"[Tool Call: {block.get('name', '')}]\n{json.dumps(block.get('input', {}), ensure_ascii=False)}"
                    )
            text = "\n".join(text_parts)
        else:
            text = str(content)

        if text:
            # OpenCode 不区分 user/assistant，我们加个前缀
            prefix = "[User]" if role == "user" else "[Assistant]"
            parts.append({"type": "text", "text": f"{prefix}\n{text}\n\n"})

    return parts


def opencode_to_anthropic_response(opencode_resp: dict, message_id: str) -> dict:
    """将 OpenCode 响应转换为 Anthropic Messages 格式。"""
    # OpenCode 响应格式: { info: {...}, parts: [...] }
    info = opencode_resp.get("info", {}) or opencode_resp.get("data", {}).get(
        "info", {}
    )
    parts = opencode_resp.get("parts", []) or opencode_resp.get("data", {}).get(
        "parts", []
    )

    # 提取文本内容
    text_content = []
    for part in parts:
        if part.get("type") == "text":
            text_content.append(part.get("text", ""))

    full_text = "\n".join(text_content) if text_content else ""

    # 估算 tokens
    output_tokens = max(1, len(full_text) // 4)

    return {
        "id": message_id,
        "type": "message",
        "role": "assistant",
        "content": [{"type": "text", "text": full_text}] if full_text else [],
        "model": FAKE_MODEL,
        "stop_reason": "end_turn",
        "stop_sequence": None,
        "usage": {
            "input_tokens": 100,  # 估算
            "output_tokens": output_tokens,
        },
    }


# ────────────────────────────────────────────────────────────────────────────────
# 核心路由
# ────────────────────────────────────────────────────────────────────────────────


@app.post("/v1/messages")
async def messages(request: Request):
    """处理 Anthropic Messages API 请求，转发到 OpenCode。"""
    client: httpx.AsyncClient = request.app.state.http_client

    # 解析请求
    try:
        body = await request.json()
    except json.JSONDecodeError as e:
        raise HTTPException(status_code=400, detail="Invalid JSON body")

    stream = body.get("stream", False)
    logger.info(f"收到请求 | stream={stream} | model={body.get('model')}")

    if DEBUG:
        logger.debug(f"原始请求体:\n{json.dumps(body, ensure_ascii=False, indent=2)}")

    # 获取 session
    try:
        session_id = await get_or_create_session(client)
    except Exception as e:
        logger.error(f"获取 OpenCode session 失败: {e}")
        raise HTTPException(
            status_code=502, detail=f"Failed to get OpenCode session: {e}"
        )

    # 转换为 OpenCode 格式 (不传 model，使用 OpenCode 默认模型)
    parts = anthropic_to_opencode_parts(body)

    opencode_body = {
        "parts": parts,
    }

    if DEBUG:
        logger.debug(
            f"OpenCode 请求体:\n{json.dumps(opencode_body, ensure_ascii=False, indent=2)}"
        )

    message_id = f"msg_{uuid.uuid4().hex[:24]}"

    # 流式处理
    if stream:
        return StreamingResponse(
            stream_opencode_response(
                client, session_id, opencode_body, message_id, body
            ),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
            },
        )

    # 非流式处理
    try:
        resp = await client.post(
            f"{OPENCODE_BASE_URL}/session/{session_id}/message",
            json=opencode_body,
        )
        resp.raise_for_status()
    except httpx.HTTPStatusError as e:
        logger.error(f"OpenCode 返回错误: {e.response.status_code} - {e.response.text}")
        raise HTTPException(status_code=502, detail=e.response.text)
    except httpx.RequestError as e:
        logger.error(f"请求 OpenCode 失败: {e}")
        raise HTTPException(status_code=502, detail=str(e))

    # 转换响应
    try:
        opencode_resp = resp.json()
        if DEBUG:
            logger.debug(
                f"OpenCode 响应:\n{json.dumps(opencode_resp, ensure_ascii=False, indent=2)}"
            )
        anthropic_resp = opencode_to_anthropic_response(opencode_resp, message_id)
    except Exception as e:
        logger.error(f"响应转换失败: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Response conversion failed: {e}")

    return JSONResponse(content=anthropic_resp)


async def stream_opencode_response(
    client: httpx.AsyncClient,
    session_id: str,
    opencode_body: dict,
    message_id: str,
    original_body: dict,
) -> AsyncIterator[str]:
    """流式代理 OpenCode 响应。"""

    # 估算 input tokens
    input_tokens = estimate_input_tokens(
        original_body.get("system"),
        original_body.get("messages", []),
    )

    # 发送 message_start
    yield format_sse(
        "message_start",
        {
            "type": "message_start",
            "message": {
                "id": message_id,
                "type": "message",
                "role": "assistant",
                "content": [],
                "model": FAKE_MODEL,
                "stop_reason": None,
                "stop_sequence": None,
                "usage": {"input_tokens": input_tokens, "output_tokens": 0},
            },
        },
    )

    # 发送 content_block_start
    yield format_sse(
        "content_block_start",
        {
            "type": "content_block_start",
            "index": 0,
            "content_block": {"type": "text", "text": ""},
        },
    )

    output_tokens = 0

    try:
        # OpenCode 的 message API 是同步的，我们模拟流式输出
        resp = await client.post(
            f"{OPENCODE_BASE_URL}/session/{session_id}/message",
            json=opencode_body,
        )
        resp.raise_for_status()

        opencode_resp = resp.json()
        parts = opencode_resp.get("parts", []) or opencode_resp.get("data", {}).get(
            "parts", []
        )

        # 提取文本并分块输出
        for part in parts:
            if part.get("type") == "text":
                text = part.get("text", "")
                # 分成小块模拟流式
                chunk_size = 20
                for i in range(0, len(text), chunk_size):
                    chunk = text[i : i + chunk_size]
                    output_tokens += max(1, len(chunk) // 4)
                    yield format_sse(
                        "content_block_delta",
                        {
                            "type": "content_block_delta",
                            "index": 0,
                            "delta": {"type": "text_delta", "text": chunk},
                        },
                    )
                    await asyncio.sleep(0.01)  # 小延迟模拟打字效果

    except Exception as e:
        logger.error(f"流式请求失败: {e}")
        yield format_sse(
            "content_block_delta",
            {
                "type": "content_block_delta",
                "index": 0,
                "delta": {"type": "text_delta", "text": f"[Error: {e}]"},
            },
        )

    # 发送 content_block_stop
    yield format_sse(
        "content_block_stop",
        {
            "type": "content_block_stop",
            "index": 0,
        },
    )

    # 发送 message_delta
    yield format_sse(
        "message_delta",
        {
            "type": "message_delta",
            "delta": {"stop_reason": "end_turn", "stop_sequence": None},
            "usage": {"output_tokens": output_tokens},
        },
    )

    # 发送 message_stop
    yield format_sse("message_stop", {"type": "message_stop"})


def format_sse(event: str, data: dict) -> str:
    """格式化为 SSE 格式。"""
    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"


# ────────────────────────────────────────────────────────────────────────────────
# 辅助路由
# ────────────────────────────────────────────────────────────────────────────────


@app.get("/")
async def root():
    return {"status": "ok", "service": "opencode2claude", "target": TARGET_MODEL}


@app.get("/health")
async def health():
    return {"status": "healthy"}


@app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [
            {
                "id": FAKE_MODEL,
                "object": "model",
                "created": 1700000000,
                "owned_by": "anthropic",
            }
        ],
    }


@app.post("/v1/messages/count_tokens")
async def count_tokens(request: Request):
    """Token 计数端点。"""
    try:
        body = await request.json()
    except json.JSONDecodeError:
        raise HTTPException(status_code=400, detail="Invalid JSON body")

    input_tokens = estimate_input_tokens(
        body.get("system"),
        body.get("messages", []),
    )
    return {"input_tokens": input_tokens}


# ────────────────────────────────────────────────────────────────────────────────
# 入口
# ────────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "proxy:app",
        host="0.0.0.0",
        port=PROXY_PORT,
        reload=DEBUG,
        log_level="debug" if DEBUG else "info",
    )