From 332d94ca6f7e12bd4350ebe5a21f3188c42f19f0 Mon Sep 17 00:00:00 2001
From: Ryo Ota <nwtgck@nwtgck.org>
Date: Sat, 14 Mar 2026 11:22:23 +0900
Subject: [PATCH] Add allowed-origins to the server (#987)

---
 mlx_lm/server.py     | 14 +++++++++++++-
 tests/test_server.py |  1 +
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/mlx_lm/server.py b/mlx_lm/server.py
index aff52af..c5f1847 100644
--- a/mlx_lm/server.py
+++ b/mlx_lm/server.py
@@ -1127,7 +1127,13 @@ class APIHandler(BaseHTTPRequestHandler):
         super().__init__(*args, **kwargs)
 
     def _set_cors_headers(self):
-        self.send_header("Access-Control-Allow-Origin", "*")
+        allowed_origins = self.response_generator.cli_args.allowed_origins
+        origin = self.headers.get("Origin")
+        if "*" in allowed_origins:
+            self.send_header("Access-Control-Allow-Origin", "*")
+        elif origin in allowed_origins:
+            self.send_header("Access-Control-Allow-Origin", origin)
+            self.send_header("Vary", "Origin")
         self.send_header("Access-Control-Allow-Methods", "*")
         self.send_header("Access-Control-Allow-Headers", "*")
 
@@ -1895,6 +1901,12 @@ def main():
         default=8080,
         help="Port for the HTTP server (default: 8080)",
     )
+    parser.add_argument(
+        "--allowed-origins",
+        type=lambda x: x.split(","),
+        default="*",
+        help="Allowed origins (default: *)",
+    )
     parser.add_argument(
         "--draft-model",
         type=str,
diff --git a/tests/test_server.py b/tests/test_server.py
index a63ae1a..c5a815e 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -47,6 +47,7 @@ class DummyModelProvider:
                 "prompt_cache_size": 10,
                 "prompt_cache_bytes": 1 << 63,
                 "prompt_cache_total_bytes": None,
+                "allowed_origins": ["*"],
             },
         )