codex_api_proxy.py

#!/usr/bin/env python3
"""
Codex API Proxy - 用 codex exec 作为后端，暴露 OpenAI 兼容 API
解决 sharedchat.cc 封禁非 Codex 客户端 TLS 指纹的问题
"""

import subprocess
import json
import time
import uuid
import os
import sys
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse

# 配置
PROXY_PORT = int(os.environ.get("CODEX_PROXY_PORT", "19080"))
CODEX_BIN = os.environ.get("CODEX_BIN", "codex")
CODEX_TIMEOUT = int(os.environ.get("CODEX_TIMEOUT", "30"))

# codex 环境变量 (会继承当前环境，这里确保关键变量)
CODEX_ENV = os.environ.copy()


class CodexHandler(BaseHTTPRequestHandler):
    """处理 OpenAI 兼容的 API 请求"""

    def log_message(self, format, *args):
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        print(f"[{timestamp}] {self.address_string()} - {format % args}", flush=True)

    def do_POST(self):
        path = urlparse(self.path).path
        # LOG HEADERS ONLY (don't consume rfile!)
        auth_hdr = self.headers.get('Authorization', 'none')
        ct = self.headers.get('Content-Type', 'none')
        cl = self.headers.get('Content-Length', '0')
        print(f"=== INCOMING: POST path={path} ct={ct} auth={auth_hdr[:40]} cl={cl}", flush=True)

        if path == "/v1/responses":
            self._handle_responses()
        elif path == "/v1/chat/completions":
            self._handle_chat_completions()
        elif path == "/health":
            self._handle_health()
        else:
            self._send_json(404, {"error": f"Unknown path: {path}"})

    def do_GET(self):
        path = urlparse(self.path).path
        if path == "/health":
            self._handle_health()
        elif path == "/v1/models":
            self._handle_models()
        else:
            self._send_json(404, {"error": f"Unknown path: {path}"})

    def _handle_health(self):
        self._send_json(200, {"status": "ok", "backend": "codex_exec"})

    def _handle_models(self):
        self._send_json(200, {
            "object": "list",
            "data": [
                {"id": "gpt-5.5", "object": "model", "owned_by": "codex-proxy"},
                {"id": "o3", "object": "model", "owned_by": "codex-proxy"},
                {"id": "o4-mini", "object": "model", "owned_by": "codex-proxy"},
            ]
        })

    def _handle_responses(self):
        """处理 OpenAI Responses API 格式的请求"""
        try:
            body = self._read_body()
            # 提取文本输入
            prompt = self._extract_prompt_responses(body)
            if not prompt:
                self._send_json(400, {"error": "No prompt found in request"})
                return

            model = body.get("model", "gpt-5.5")
            stream = body.get("stream", False)

            self.log_message("RESponses request: model=%s stream=%s prompt=%.100s...", model, stream, prompt)

            # 调用 codex exec
            result = self._call_codex(prompt, model)

            if result["success"]:
                if stream:
                    self._send_sse_responses(result["output"], model)
                else:
                    response = self._format_responses_api(result["output"], model)
                    self._send_json(200, response)
            else:
                self._send_json(500, {
                    "error": {
                        "message": result["error"],
                        "type": "proxy_error",
                        "code": "codex_exec_failed"
                    }
                })

        except Exception as e:
            self.log_message("ERROR in _handle_responses: %s", str(e))
            self._send_json(500, {"error": {"message": str(e), "type": "proxy_error"}})

    def _handle_chat_completions(self):
        """处理 OpenAI Chat Completions API 格式的请求"""
        try:
            body = self._read_body()
            prompt = self._extract_prompt_chat(body)
            if not prompt:
                self._send_json(400, {"error": "No prompt found in request"})
                return

            model = body.get("model", "gpt-5.5")
            stream = body.get("stream", False)

            self.log_message("Chat request: model=%s prompt=%.100s...", model, prompt)

            result = self._call_codex(prompt, model)

            if result["success"]:
                response = self._format_chat_api(result["output"], model)
                self._send_json(200, response)
            else:
                self._send_json(500, {
                    "error": {"message": result["error"], "type": "proxy_error"}
                })

        except Exception as e:
            self.log_message("ERROR in _handle_chat_completions: %s", str(e))
            self._send_json(500, {"error": {"message": str(e)}})

    def _extract_prompt_responses(self, body):
        """从 Responses API 格式提取 prompt"""
        # 方式1: input 字段 (字符串)
        inp = body.get("input")
        if isinstance(inp, str):
            return inp
        if isinstance(inp, list):
            # input 是消息列表
            parts = []
            for item in inp:
                if isinstance(item, dict):
                    # Standard Responses API: type=message
                    if item.get("type") == "message":
                        content = item.get("content", "")
                        if isinstance(content, str):
                            parts.append(content)
                        elif isinstance(content, list):
                            for c in content:
                                if isinstance(c, dict) and c.get("type") == "output_text":
                                    parts.append(c.get("text", ""))
                                elif isinstance(c, str):
                                    parts.append(c)
                    elif item.get("type") == "output_text":
                        parts.append(item.get("text", ""))
                    # Sub2API format: {role: "user", content: "hi"}
                    elif item.get("role") == "user":
                        content = item.get("content", "")
                        if isinstance(content, str):
                            parts.append(content)
                        elif isinstance(content, list):
                            for c in content:
                                if isinstance(c, dict):
                                    parts.append(c.get("text", ""))
                                elif isinstance(c, str):
                                    parts.append(c)
                elif isinstance(item, str):
                    parts.append(item)
            if parts:
                return "\n".join(parts)

        # 方式2: messages 字段
        messages = body.get("messages")
        if messages:
            for msg in reversed(messages):
                if msg.get("role") == "user":
                    content = msg.get("content", "")
                    if isinstance(content, str):
                        return content
                    if isinstance(content, list):
                        texts = [c.get("text", "") for c in content if isinstance(c, dict)]
                        return " ".join(texts)

        return None

    def _extract_prompt_chat(self, body):
        """从 Chat Completions API 格式提取 prompt"""
        messages = body.get("messages", [])
        if messages:
            # 取最后一条 user message
            for msg in reversed(messages):
                if msg.get("role") == "user":
                    content = msg.get("content", "")
                    if isinstance(content, str):
                        return content
                    if isinstance(content, list):
                        texts = [c.get("text", "") for c in content if isinstance(c, dict)]
                        return " ".join(texts)
        return None

    def _call_codex(self, prompt, model="gpt-5.5", api_key=None):
        """调用 codex exec"""
        try:
            # 构建命令
            cmd = [CODEX_BIN, "exec", "--sandbox", "workspace-write", "--skip-git-repo-check", "-"]

            # 环境变量
            env = CODEX_ENV.copy()
            # 确保 codex 能找到正确的配置
            codex_home = os.path.expanduser("~/.codex")
            if os.path.exists(os.path.join(codex_home, "config.toml")):
                env["HOME"] = os.path.expanduser("~")
            
            # 不传Sub2API的API key给codex exec——让它使用keychain OAuth认证
            # sharedchat.cc对API key返回403，但keychain OAuth可以正常工作
            # 如果codex exec报auth错误，用户需要运行 `codex login` 刷新token
            self.log_message("Using keychain auth (no API key override)")

            self.log_message("Executing: %s (prompt length=%d)", " ".join(cmd), len(prompt))

            start_time = time.time()

            proc = subprocess.run(
                cmd,
                input=prompt,
                capture_output=True,
                text=True,
                timeout=CODEX_TIMEOUT,
                env=env,
                cwd=os.path.expanduser("~")
            )

            elapsed = time.time() - start_time
            self.log_message("codex exec completed in %.1fs, rc=%d", elapsed, proc.returncode)

            if proc.returncode == 0:
                output = proc.stdout.strip()
                self.log_message("codex output: %.200s", output)
                return {"success": True, "output": output}
            else:
                error = proc.stderr.strip() or proc.stdout.strip()
                self.log_message("codex ERROR: %.200s", error)
                return {"success": False, "error": error or f"codex exec failed with rc={proc.returncode}"}

        except subprocess.TimeoutExpired:
            return {"success": False, "error": f"codex exec timed out after {CODEX_TIMEOUT}s"}
        except FileNotFoundError:
            return {"success": False, "error": f"codex binary not found at {CODEX_BIN}"}
        except Exception as e:
            return {"success": False, "error": str(e)}

    def _format_responses_api(self, output, model):
        """格式化为 OpenAI Responses API 格式"""
        return {
            "id": f"resp_{uuid.uuid4().hex[:24]}",
            "object": "response",
            "created_at": int(time.time()),
            "status": "completed",
            "model": model,
            "output": [
                {
                    "type": "message",
                    "id": f"msg_{uuid.uuid4().hex[:24]}",
                    "role": "assistant",
                    "content": [
                        {
                            "type": "output_text",
                            "text": output
                        }
                    ]
                }
            ],
            "usage": {
                "input_tokens": 0,
                "output_tokens": 0,
                "total_tokens": 0
            }
        }

    def _format_chat_api(self, output, model):
        """格式化为 OpenAI Chat Completions API 格式"""
        return {
            "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": model,
            "choices": [
                {
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": output
                    },
                    "finish_reason": "stop"
                }
            ],
            "usage": {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0
            }
        }

    def _send_sse_responses(self, text, model):
        """发送 Responses API SSE 流式响应"""
        resp_id = f"resp_{int(time.time()*1000)}"
        now = int(time.time())
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.send_header("Cache-Control", "no-cache")
        self.send_header("Connection", "keep-alive")
        self.end_headers()

        def sse(event, data):
            line = f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
            self.wfile.write(line.encode("utf-8"))
            self.wfile.flush()

        # response.created
        sse("response.created", {"type": "response.created", "response": {
            "id": resp_id, "object": "response", "created_at": now,
            "model": model, "status": "in_progress",
            "output": [], "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
        }})
        # output_item.added
        sse("response.output_item.added", {"type": "response.output_item.added",
            "output_index": 0, "item": {"type": "message", "role": "assistant", "content": []}})
        # content_part.added
        sse("response.content_part.added", {"type": "response.content_part.added",
            "output_index": 0, "content_index": 0,
            "part": {"type": "output_text", "text": ""}})
        # output_text.delta (chunk the text)
        chunk_size = 20
        for i in range(0, len(text), chunk_size):
            sse("response.output_text.delta", {"type": "response.output_text.delta",
                "output_index": 0, "content_index": 0, "delta": text[i:i+chunk_size]})
        # output_text.done
        sse("response.output_text.done", {"type": "response.output_text.done",
            "output_index": 0, "content_index": 0,
            "text": text})
        # content_part.done
        sse("response.content_part.done", {"type": "response.content_part.done",
            "output_index": 0, "content_index": 0,
            "part": {"type": "output_text", "text": text}})
        # output_item.done
        sse("response.output_item.done", {"type": "response.output_item.done",
            "output_index": 0,
            "item": {"type": "message", "role": "assistant",
                     "content": [{"type": "output_text", "text": text}]}})
        # response.completed
        sse("response.completed", {"type": "response.completed", "response": {
            "id": resp_id, "object": "response", "created_at": now,
            "model": model, "status": "completed",
            "output": [{"type": "message", "role": "assistant",
                        "content": [{"type": "output_text", "text": text}]}],
            "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
        }})
        self.wfile.write(b"\n")
        self.wfile.flush()

    def _read_body(self):
        content_length = int(self.headers.get("Content-Length", 0))
        if content_length == 0:
            return {}
        raw = self.rfile.read(content_length)
        return json.loads(raw)

    def _send_json(self, status, data):
        body = json.dumps(data, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _send_sse(self, events):
        """发送 SSE 流式响应"""
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.send_header("Cache-Control", "no-cache")
        self.send_header("Connection", "keep-alive")
        self.end_headers()
        for event_type, data in events:
            line = f"event: {event_type}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
            self.wfile.write(line.encode("utf-8"))
            self.wfile.flush()
        self.wfile.write(b"data: [DONE]\n\n")
        self.wfile.flush()


def main():
    server = HTTPServer(("0.0.0.0", PROXY_PORT), CodexHandler)
    print(f"[*] Codex API Proxy started on port {PROXY_PORT}", flush=True)
    print(f"[*] Backend: {CODEX_BIN} exec", flush=True)
    print(f"[*] Timeout: {CODEX_TIMEOUT}s", flush=True)
    print(f"[*] Endpoints:", flush=True)
    print(f"    POST /v1/responses       (OpenAI Responses API)", flush=True)
    print(f"    POST /v1/chat/completions (OpenAI Chat API)", flush=True)
    print(f"    GET  /v1/models", flush=True)
    print(f"    GET  /health", flush=True)
    print(flush=True)

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n[*] Shutting down...", flush=True)
        server.server_close()


if __name__ == "__main__":
    main()