Add allowed-origins to the server (#987)

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,

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": ["*"],
             },
         )