Add main.py

main.py

@@ -0,0 +1,144 @@
+#!/usr/bin/env python3
+"""
+CORS proxy for Unsloth Studio API.
+
+Usage:
+    python cors_proxy.py                    # proxies to http://127.0.0.1:8888 on port 8080
+    python cors_proxy.py --target 8000      # proxies to http://127.0.0.1:8000
+    python cors_proxy.py --listen 9090      # listens on port 9090
+    python cors_proxy.py --target 10.0.0.5:8000 --listen 9000
+"""
+
+import argparse
+import json
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from urllib.request import Request, urlopen, URLError
+
+CORS_HEADERS = {
+    "Access-Control-Allow-Origin": "*",
+    "Access-Control-Allow-Methods": "POST, OPTIONS, GET",
+    "Access-Control-Allow-Headers": "Content-Type, Authorization",
+    "Access-Control-Max-Age": "86400",
+}
+
+
+class ProxyHandler(BaseHTTPRequestHandler):
+    target_host = "127.0.0.1"
+    target_port = 8888
+
+    def _target_url(self):
+        return f"http://{self.target_host}:{self.target_port}{self.path}"
+
+    def _set_cors(self, status=200):
+        self.send_response(status)
+        for k, v in CORS_HEADERS.items():
+            self.send_header(k, v)
+
+    def do_OPTIONS(self):
+        self._set_cors(204)
+        self.end_headers()
+
+    def do_POST(self):
+        content_length = int(self.headers.get("Content-Length", 0))
+        body = self.rfile.read(content_length)
+
+        target = self._target_url()
+        req = Request(target, data=body, method="POST")
+
+        for h in ("Content-Type", "Authorization"):
+            val = self.headers.get(h)
+            if val:
+                req.add_header(h, val)
+
+        try:
+            resp = urlopen(req, timeout=60)
+            resp_body = resp.read()
+            status = resp.status
+        except URLError as e:
+            detail = f"Upstream server at {self.target_host}:{self.target_port} is unreachable."
+            if hasattr(e, "reason"):
+                detail += f" Reason: {e.reason}"
+            self._set_cors(502)
+            self.send_header("Content-Type", "application/json")
+            self.end_headers()
+            self.wfile.write(json.dumps({
+                "error": {"message": detail, "type": "proxy_error"}
+            }).encode())
+            return
+        except Exception as e:
+            self._set_cors(502)
+            self.send_header("Content-Type", "application/json")
+            self.end_headers()
+            self.wfile.write(json.dumps({
+                "error": {"message": f"Proxy error: {e}", "type": "proxy_error"}
+            }).encode())
+            return
+
+        self._set_cors(status)
+        ctype = resp.headers.get("Content-Type", "application/json")
+        self.send_header("Content-Type", ctype)
+        self.end_headers()
+        self.wfile.write(resp_body)
+
+    def log_message(self, fmt, *args):
+        method = self.command
+        path = self.path
+        print(f"  [{method}] {path}  ->  {self._target_url()}")
+
+    def do_GET(self):
+        if self.path == "/health":
+            self._set_cors(200)
+            self.send_header("Content-Type", "application/json")
+            self.end_headers()
+            self.wfile.write(json.dumps({
+                "status": "ok",
+                "proxy_to": f"{self.target_host}:{self.target_port}",
+            }).encode())
+            return
+        self._set_cors(405)
+        self.end_headers()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="CORS proxy for Unsloth Studio")
+    parser.add_argument(
+        "--target", default="127.0.0.1:8888",
+        help="Unsloth Studio address (default: 127.0.0.1:8888)",
+    )
+    parser.add_argument(
+        "--listen", default=8080, type=int,
+        help="Port to listen on (default: 8080)",
+    )
+    args = parser.parse_args()
+
+    host, port_str = args.target, "8888"
+    if ":" in args.target:
+        host, port_str = args.target.rsplit(":", 1)
+    try:
+        port = int(port_str)
+    except ValueError:
+        print(f"Invalid target port: {port_str}")
+        sys.exit(1)
+
+    ProxyHandler.target_host = host
+    ProxyHandler.target_port = port
+
+    server = HTTPServer(("0.0.0.0", args.listen), ProxyHandler)
+
+    print(f"  Unsloth Studio CORS Proxy")
+    print(f"  ─────────────────────────")
+    print(f"  Listening on:   http://127.0.0.1:{args.listen}")
+    print(f"  Forwarding to:  http://{host}:{port}")
+    print(f"  Plugin API URL: http://127.0.0.1:{args.listen}")
+    print(f"  Health check:   http://127.0.0.1:{args.listen}/health")
+    print()
+
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        print("\nShutting down.")
+        server.server_close()
+
+
+if __name__ == "__main__":
+    main()