From 95ae86f8204b9d71b3fa5e8a9728bf65ed9823c1 Mon Sep 17 00:00:00 2001
From: Daksha1611 <mehtadaksha1611@gmail.com>
Date: Fri, 6 Nov 2026 21:55:12 +0530
Subject: [PATCH] runtime: accept a config object in the download endpoints

Signed-off-by: Daksha1611 <mehtadaksha1611@gmail.com>
---
 python/kthena/downloader/app.py     | 13 +++--
 python/kthena/tests/test_runtime.py | 75 +++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+), 4 deletions(-)

diff --git a/python/kthena/downloader/app.py b/python/kthena/downloader/app.py
index 69730dcf94..a1bb5f24aa 100644
--- a/python/kthena/downloader/app.py
+++ b/python/kthena/downloader/app.py
@@ -23,15 +23,20 @@
 logger = setup_logger()
 
 
-def load_config(config_str: str = None) -> dict:
+def load_config(raw_config: str | dict | None = None) -> dict:
+    """Build the downloader config from a JSON string (CLI) or a dict (runtime API)."""
     config = {}
-    if config_str:
+    if isinstance(raw_config, dict):
+        config = dict(raw_config)
+    elif raw_config:
         try:
-            config = json.loads(config_str)
+            config = json.loads(raw_config)
             logger.info("Loaded configuration from JSON string.")
-        except json.JSONDecodeError as e:
+        except (json.JSONDecodeError, TypeError) as e:
             logger.error(f"Error parsing configuration JSON: {e}")
             raise ValueError("Invalid configuration JSON format.") from e
+        if not isinstance(config, dict):
+            raise ValueError("Configuration must be a JSON object.")
 
     env_config = {
         "hf_token": os.getenv("HF_AUTH_TOKEN"),
diff --git a/python/kthena/tests/test_runtime.py b/python/kthena/tests/test_runtime.py
index 2e1262f3d8..eadef4218d 100644
--- a/python/kthena/tests/test_runtime.py
+++ b/python/kthena/tests/test_runtime.py
@@ -23,6 +23,7 @@
 import uvicorn
 from fastapi.testclient import TestClient
 
+from kthena.downloader.app import load_config
 from kthena.runtime import app as runtime_app
 from kthena.runtime import memory_kv_manager as mkm
 from kthena.runtime.app import (
@@ -170,6 +171,80 @@ def test_load_lora_adapter_async_download_success(monkeypatch):
         assert resp.status_code == 202
 
 
+@pytest.mark.parametrize("path, extra_body", [
+    ("/v1/download_model", {}),
+    ("/v1/load_lora_adapter", {"lora_name": "lora-sql"}),
+])
+def test_download_endpoints_pass_object_config_to_downloader(monkeypatch, path, extra_body):
+    # Arrange: record the config the endpoint hands to the downloader
+    monkeypatch.delenv("HF_AUTH_TOKEN", raising=False)
+    app = create_application(_make_args())
+    state = get_app_state(app)
+    monkeypatch.setattr("httpx.AsyncClient", _FakeAsyncClient)
+    state.client = _FakeAsyncClient()
+    monkeypatch.setattr("httpx.Client", _FakeSyncClient)
+
+    received = []
+    monkeypatch.setattr(
+        "kthena.runtime.app.download_model",
+        lambda source, output_dir, config, max_workers: received.append(config),
+    )
+
+    with TestClient(app) as client:
+        body = {
+            "source": "hf://org/repo",
+            "output_dir": "/tmp/model",
+            "config": {"hf_token": "from-request"},
+            **extra_body,
+        }
+
+        # Act
+        resp = client.post(path, json=body)
+
+    # Assert
+    assert resp.status_code == 200
+    assert received == [{"hf_token": "from-request"}]
+
+
+@pytest.mark.parametrize("config, expected", [
+    (None, {}),
+    ("", {}),
+    ({}, {}),
+    ({"hf_token": "from-config"}, {"hf_token": "from-config"}),
+    ('{"hf_token": "from-config"}', {"hf_token": "from-config"}),
+])
+def test_load_config_accepts_supported_forms(monkeypatch, config, expected):
+    monkeypatch.delenv("HF_AUTH_TOKEN", raising=False)
+
+    assert load_config(config) == expected
+
+
+@pytest.mark.parametrize("config", [
+    {"hf_token": "from-config"},
+    '{"hf_token": "from-config"}',
+])
+def test_load_config_prefers_explicit_values_over_environment(monkeypatch, config):
+    monkeypatch.setenv("HF_AUTH_TOKEN", "from-env")
+
+    assert load_config(config) == {"hf_token": "from-config"}
+
+
+def test_load_config_fills_missing_values_without_mutating_input(monkeypatch):
+    monkeypatch.setenv("HF_AUTH_TOKEN", "from-env")
+    config = {"endpoint": "http://storage.local"}
+
+    result = load_config(config)
+
+    assert result == {"endpoint": "http://storage.local", "hf_token": "from-env"}
+    assert config == {"endpoint": "http://storage.local"}
+
+
+@pytest.mark.parametrize("config", [123, ["hf_token"], "[1, 2]", "not json"])
+def test_load_config_rejects_values_that_are_not_objects(config):
+    with pytest.raises(ValueError):
+        load_config(config)
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("mode, body, expected_status", [
     ("redis", {"router_id": "router-a"}, 409),