feat: add in-app LM Remote configuration
This commit is contained in:
@@ -0,0 +1,617 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def modules():
|
||||
root = Path(__file__).resolve().parents[1]
|
||||
package_name = "lm_remote_proxy_test_package"
|
||||
package = types.ModuleType(package_name)
|
||||
package.__path__ = [str(root)]
|
||||
sys.modules[package_name] = package
|
||||
|
||||
loaded = {}
|
||||
for name in ("config", "remote_client", "proxy"):
|
||||
full_name = f"{package_name}.{name}"
|
||||
spec = importlib.util.spec_from_file_location(full_name, root / f"{name}.py")
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[full_name] = module
|
||||
spec.loader.exec_module(module)
|
||||
loaded[name] = module
|
||||
|
||||
yield types.SimpleNamespace(**loaded)
|
||||
for name in ("proxy", "remote_client", "config"):
|
||||
sys.modules.pop(f"{package_name}.{name}", None)
|
||||
sys.modules.pop(package_name, None)
|
||||
|
||||
|
||||
class DummyContent:
|
||||
def __init__(self, body: bytes):
|
||||
self.body = body
|
||||
self.offset = 0
|
||||
|
||||
async def read(self, size: int) -> bytes:
|
||||
chunk = self.body[self.offset : self.offset + size]
|
||||
self.offset += len(chunk)
|
||||
return chunk
|
||||
|
||||
|
||||
class FragmentedContent:
|
||||
def __init__(self, chunks: list[bytes]):
|
||||
self.chunks = list(chunks)
|
||||
|
||||
async def read(self, size: int) -> bytes:
|
||||
if not self.chunks:
|
||||
return b""
|
||||
chunk = self.chunks.pop(0)
|
||||
if len(chunk) <= size:
|
||||
return chunk
|
||||
self.chunks.insert(0, chunk[size:])
|
||||
return chunk[:size]
|
||||
|
||||
|
||||
class DummyRequest:
|
||||
def __init__(
|
||||
self,
|
||||
method="GET",
|
||||
path="/",
|
||||
payload=None,
|
||||
headers=None,
|
||||
*,
|
||||
raw_body: bytes | None = None,
|
||||
content_length: int | None | object = ...,
|
||||
):
|
||||
self.method = method
|
||||
self.path = path
|
||||
self._payload = payload
|
||||
self.headers = headers or {}
|
||||
self.query_string = ""
|
||||
self.can_read_body = payload is not None
|
||||
self.content_type = (
|
||||
"application/json" if payload is not None else "application/octet-stream"
|
||||
)
|
||||
self._body = (
|
||||
raw_body
|
||||
if raw_body is not None
|
||||
else json.dumps(payload).encode("utf-8")
|
||||
if payload is not None
|
||||
else b""
|
||||
)
|
||||
self.content = DummyContent(self._body)
|
||||
if content_length is ...:
|
||||
self.content_length = len(self._body) if self._body else None
|
||||
else:
|
||||
self.content_length = content_length
|
||||
if raw_body is not None:
|
||||
self.content_type = "application/json"
|
||||
|
||||
async def json(self):
|
||||
return self._payload
|
||||
|
||||
async def read(self):
|
||||
return self._body
|
||||
|
||||
|
||||
def response_json(response: web.Response) -> dict:
|
||||
return json.loads(response.body.decode("utf-8"))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_proxy(modules, tmp_path, monkeypatch):
|
||||
config = modules.config.RemoteConfig(
|
||||
tmp_path / "user" / "config.json",
|
||||
tmp_path / "missing-legacy.json",
|
||||
environ={},
|
||||
)
|
||||
monkeypatch.setattr(modules.proxy, "remote_config", config)
|
||||
monkeypatch.setattr(modules.remote_client, "remote_config", config)
|
||||
modules.proxy.RemoteLoraClient._instance = None
|
||||
modules.proxy._proxy_sessions = {}
|
||||
modules.proxy._proxy_session_lock = None
|
||||
modules.proxy._active_proxy_websockets = set()
|
||||
modules.proxy._active_proxy_websockets_lock = None
|
||||
return modules.proxy, config
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_endpoint_saves_and_hot_enables(isolated_proxy):
|
||||
proxy, config = isolated_proxy
|
||||
get_response = await proxy._handle_config(DummyRequest())
|
||||
initial = response_json(get_response)
|
||||
assert initial["configured"]["remote_url"] == ""
|
||||
|
||||
put_response = await proxy._handle_config(
|
||||
DummyRequest(
|
||||
"PUT",
|
||||
proxy._CONFIG_ROUTE,
|
||||
{
|
||||
"revision": initial["revision"],
|
||||
"config": {
|
||||
"remote_url": "http://manager.local:8188/",
|
||||
"timeout": 40,
|
||||
"path_mappings": {"/remote": "/local"},
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
payload = response_json(put_response)
|
||||
assert put_response.status == 200
|
||||
assert payload["effective"]["remote_url"] == "http://manager.local:8188"
|
||||
assert payload["restart_required"] is False
|
||||
assert config.is_configured
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_middleware_falls_through_when_disabled_then_uses_new_url(
|
||||
isolated_proxy, monkeypatch
|
||||
):
|
||||
proxy, config = isolated_proxy
|
||||
|
||||
async def local_handler(request):
|
||||
return web.Response(text="local")
|
||||
|
||||
disabled_response = await proxy.lm_remote_proxy_middleware(
|
||||
DummyRequest(path="/loras"), local_handler
|
||||
)
|
||||
assert disabled_response.text == "local"
|
||||
|
||||
state = config.as_dict()
|
||||
config.save(
|
||||
{"remote_url": "http://manager.local", "timeout": 30, "path_mappings": {}},
|
||||
expected_revision=state["revision"],
|
||||
)
|
||||
captured = {}
|
||||
|
||||
async def fake_proxy_http(request, snapshot):
|
||||
captured["url"] = snapshot.remote_url
|
||||
return web.Response(text="remote")
|
||||
|
||||
monkeypatch.setattr(proxy, "_proxy_http", fake_proxy_http)
|
||||
enabled_response = await proxy.lm_remote_proxy_middleware(
|
||||
DummyRequest(path="/loras"), local_handler
|
||||
)
|
||||
assert enabled_response.text == "remote"
|
||||
assert captured["url"] == "http://manager.local"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_endpoint_rejects_changes_to_environment_managed_field(
|
||||
modules, tmp_path, monkeypatch
|
||||
):
|
||||
config = modules.config.RemoteConfig(
|
||||
tmp_path / "user.json",
|
||||
tmp_path / "legacy.json",
|
||||
environ={"LM_REMOTE_URL": "http://managed.local"},
|
||||
)
|
||||
monkeypatch.setattr(modules.proxy, "remote_config", config)
|
||||
initial = config.as_dict()
|
||||
response = await modules.proxy._handle_config(
|
||||
DummyRequest(
|
||||
"PUT",
|
||||
modules.proxy._CONFIG_ROUTE,
|
||||
{
|
||||
"revision": initial["revision"],
|
||||
"config": {
|
||||
"remote_url": "http://changed.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {},
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
payload = response_json(response)
|
||||
assert response.status == 409
|
||||
assert payload["field"] == "remote_url"
|
||||
assert "LM_REMOTE_URL" in payload["error"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conflict_reload_rotates_runtime_generation(isolated_proxy, monkeypatch):
|
||||
proxy, config = isolated_proxy
|
||||
initial = config.as_dict()
|
||||
config._config_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
config._config_file.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"remote_url": "http://external.local",
|
||||
"timeout": 20,
|
||||
"path_mappings": {},
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
retired = []
|
||||
rotated = []
|
||||
|
||||
async def fake_retire(generation):
|
||||
retired.append(generation)
|
||||
|
||||
async def fake_rotate(generation):
|
||||
rotated.append(generation)
|
||||
|
||||
monkeypatch.setattr(proxy, "_retire_proxy_sessions", fake_retire)
|
||||
monkeypatch.setattr(proxy, "_rotate_active_websockets", fake_rotate)
|
||||
response = await proxy._handle_config(
|
||||
DummyRequest(
|
||||
"PUT",
|
||||
proxy._CONFIG_ROUTE,
|
||||
{
|
||||
"revision": initial["revision"],
|
||||
"config": {
|
||||
"remote_url": "http://browser.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {},
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
payload = response_json(response)
|
||||
|
||||
assert response.status == 409
|
||||
assert payload["latest"]["effective"]["remote_url"] == "http://external.local"
|
||||
assert retired == [config.generation]
|
||||
assert rotated == [config.generation]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_test_uses_unsaved_draft(isolated_proxy, monkeypatch):
|
||||
proxy, _ = isolated_proxy
|
||||
captured = {}
|
||||
|
||||
async def fake_test(remote_url, timeout):
|
||||
captured.update(remote_url=remote_url, timeout=timeout)
|
||||
return 17
|
||||
|
||||
monkeypatch.setattr(proxy, "_perform_connection_test", fake_test)
|
||||
response = await proxy._handle_test_connection(
|
||||
DummyRequest(
|
||||
"POST",
|
||||
proxy._TEST_CONNECTION_ROUTE,
|
||||
{"remote_url": "http://draft.local:8188/", "timeout": 8},
|
||||
)
|
||||
)
|
||||
payload = response_json(response)
|
||||
assert response.status == 200
|
||||
assert payload["latency_ms"] == 17
|
||||
assert captured == {"remote_url": "http://draft.local:8188", "timeout": 8}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_test_rejects_empty_url(isolated_proxy):
|
||||
proxy, _ = isolated_proxy
|
||||
response = await proxy._handle_test_connection(
|
||||
DummyRequest(
|
||||
"POST",
|
||||
proxy._TEST_CONNECTION_ROUTE,
|
||||
{"remote_url": "", "timeout": 30},
|
||||
)
|
||||
)
|
||||
assert response.status == 400
|
||||
assert response_json(response)["field"] == "remote_url"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_response_reader_accepts_fragmented_json(isolated_proxy):
|
||||
proxy, _ = isolated_proxy
|
||||
response = types.SimpleNamespace(
|
||||
content=FragmentedContent([b'{"sta', b'tus":"', b'ok"}'])
|
||||
)
|
||||
|
||||
assert await proxy._read_small_json(response) == {"status": "ok"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_response_reader_enforces_hard_limit(isolated_proxy):
|
||||
proxy, _ = isolated_proxy
|
||||
response = types.SimpleNamespace(
|
||||
content=FragmentedContent([b" " * proxy._MAX_TEST_RESPONSE, b" "])
|
||||
)
|
||||
|
||||
with pytest.raises(proxy._ConnectionTestError, match="unexpectedly large"):
|
||||
await proxy._read_small_json(response)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_loop_header_is_rejected(isolated_proxy):
|
||||
proxy, config = isolated_proxy
|
||||
state = config.as_dict()
|
||||
config.save(
|
||||
{"remote_url": "http://manager.local", "timeout": 30, "path_mappings": {}},
|
||||
expected_revision=state["revision"],
|
||||
)
|
||||
|
||||
async def local_handler(request):
|
||||
return web.Response(text="local")
|
||||
|
||||
response = await proxy.lm_remote_proxy_middleware(
|
||||
DummyRequest(
|
||||
path="/api/lm/health-check",
|
||||
headers={proxy._PROXY_HOP_HEADER: "1"},
|
||||
),
|
||||
local_handler,
|
||||
)
|
||||
assert response.status == 508
|
||||
|
||||
|
||||
def test_register_proxy_is_available_while_unconfigured(isolated_proxy):
|
||||
proxy, config = isolated_proxy
|
||||
assert not config.is_configured
|
||||
app = web.Application()
|
||||
proxy.register_proxy(app)
|
||||
assert proxy.lm_remote_proxy_middleware in app.middlewares
|
||||
assert proxy._cleanup_proxy_session in app.on_shutdown
|
||||
|
||||
|
||||
def test_register_proxy_keeps_existing_security_middleware_first(isolated_proxy):
|
||||
proxy, _ = isolated_proxy
|
||||
|
||||
@web.middleware
|
||||
async def security_guard(request, handler):
|
||||
return await handler(request)
|
||||
|
||||
app = web.Application(middlewares=[security_guard])
|
||||
proxy.register_proxy(app)
|
||||
|
||||
assert list(app.middlewares) == [security_guard, proxy.lm_remote_proxy_middleware]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunked_config_body_is_hard_limited(isolated_proxy):
|
||||
proxy, _ = isolated_proxy
|
||||
response = await proxy._handle_config(
|
||||
DummyRequest(
|
||||
"PUT",
|
||||
proxy._CONFIG_ROUTE,
|
||||
raw_body=b" " * (proxy._MAX_CONFIG_BODY + 1),
|
||||
content_length=None,
|
||||
)
|
||||
)
|
||||
|
||||
assert response.status == 413
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remote_client_fetches_all_pages_beyond_server_cap(
|
||||
isolated_proxy, modules, monkeypatch
|
||||
):
|
||||
_, config = isolated_proxy
|
||||
client = modules.remote_client.RemoteLoraClient()
|
||||
calls = []
|
||||
|
||||
async def fake_get_json(path, params=None, *, snapshot=None):
|
||||
page = int(params["page"])
|
||||
calls.append((page, int(params["page_size"]), snapshot.generation))
|
||||
start = (page - 1) * 100
|
||||
count = 100 if page == 1 else 55
|
||||
return {
|
||||
"items": [
|
||||
{"file_name": f"model-{index}"} for index in range(start, start + count)
|
||||
],
|
||||
"total_pages": 2,
|
||||
}
|
||||
|
||||
monkeypatch.setattr(client, "_get_json", fake_get_json)
|
||||
items = await client._get_lora_list_cached(snapshot=config.snapshot)
|
||||
|
||||
assert len(items) == 155
|
||||
assert items[-1]["file_name"] == "model-154"
|
||||
assert [call[:2] for call in calls] == [(1, 100), (2, 100)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("method_name", "expected_path"),
|
||||
[
|
||||
("_get_lora_list_cached", "/api/lm/loras/list"),
|
||||
("_get_checkpoint_list_cached", "/api/lm/checkpoints/list"),
|
||||
],
|
||||
)
|
||||
async def test_successful_empty_listing_is_cached(
|
||||
isolated_proxy, modules, monkeypatch, method_name, expected_path
|
||||
):
|
||||
_, config = isolated_proxy
|
||||
client = modules.remote_client.RemoteLoraClient()
|
||||
calls = []
|
||||
|
||||
async def fake_get_all_pages(path, *, snapshot):
|
||||
calls.append((path, snapshot.generation))
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(client, "_get_all_pages", fake_get_all_pages)
|
||||
cached_method = getattr(client, method_name)
|
||||
|
||||
assert await cached_method(snapshot=config.snapshot) == []
|
||||
assert await cached_method(snapshot=config.snapshot) == []
|
||||
assert calls == [(expected_path, config.generation)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_lookup_never_returns_another_generations_cache(
|
||||
isolated_proxy, modules, monkeypatch
|
||||
):
|
||||
_, config = isolated_proxy
|
||||
client = modules.remote_client.RemoteLoraClient()
|
||||
client._lora_cache = [{"file_name": "new-generation"}]
|
||||
client._lora_cache_generation = config.generation
|
||||
old_snapshot = modules.config.ConfigSnapshot(
|
||||
config.generation - 1,
|
||||
"http://old.local",
|
||||
30,
|
||||
(("/remote", "/old-local"),),
|
||||
)
|
||||
|
||||
async def fail_fetch(path, *, snapshot):
|
||||
raise OSError("old server unavailable")
|
||||
|
||||
monkeypatch.setattr(client, "_get_all_pages", fail_fetch)
|
||||
|
||||
assert await client._get_lora_list_cached(snapshot=old_snapshot) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lora_info_maps_with_the_fetch_generation(
|
||||
isolated_proxy, modules, monkeypatch
|
||||
):
|
||||
_, config = isolated_proxy
|
||||
initial = config.as_dict()
|
||||
config.save(
|
||||
{
|
||||
"remote_url": "http://one.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {"/remote": "/local-one"},
|
||||
},
|
||||
expected_revision=initial["revision"],
|
||||
)
|
||||
client = modules.remote_client.RemoteLoraClient()
|
||||
|
||||
async def fake_list(*, snapshot=None):
|
||||
latest = config.as_dict()
|
||||
config.save(
|
||||
{
|
||||
"remote_url": "http://two.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {"/remote": "/local-two"},
|
||||
},
|
||||
expected_revision=latest["revision"],
|
||||
)
|
||||
return [
|
||||
{
|
||||
"file_name": "portrait",
|
||||
"file_path": "/remote/portrait.safetensors",
|
||||
"folder": "",
|
||||
"civitai": {},
|
||||
}
|
||||
]
|
||||
|
||||
monkeypatch.setattr(client, "_get_lora_list_cached", fake_list)
|
||||
monkeypatch.setattr(client, "_relative_lora_path", lambda path, folder: path)
|
||||
|
||||
relative, _ = await client.get_lora_info("portrait")
|
||||
|
||||
assert relative == "/local-one/portrait.safetensors"
|
||||
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_session_does_not_retain_remote_cookies(
|
||||
isolated_proxy, monkeypatch
|
||||
):
|
||||
proxy, config = isolated_proxy
|
||||
captured = {}
|
||||
|
||||
def create_session(*args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return FakeSession()
|
||||
|
||||
monkeypatch.setattr(proxy.aiohttp, "ClientSession", create_session)
|
||||
lease = proxy._proxy_session_lease(config.snapshot)
|
||||
await lease.__aenter__()
|
||||
await lease.__aexit__(None, None, None)
|
||||
|
||||
assert isinstance(captured["cookie_jar"], proxy.aiohttp.DummyCookieJar)
|
||||
await proxy._close_all_proxy_sessions()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_session_rotation_waits_for_inflight_request(
|
||||
isolated_proxy, monkeypatch
|
||||
):
|
||||
proxy, config = isolated_proxy
|
||||
created = []
|
||||
|
||||
def create_session(*args, **kwargs):
|
||||
session = FakeSession()
|
||||
created.append(session)
|
||||
return session
|
||||
|
||||
monkeypatch.setattr(proxy.aiohttp, "ClientSession", create_session)
|
||||
old_snapshot = config.snapshot
|
||||
old_lease = proxy._proxy_session_lease(old_snapshot)
|
||||
old_session = await old_lease.__aenter__()
|
||||
|
||||
state = config.as_dict()
|
||||
new_snapshot = config.save(
|
||||
{"remote_url": "http://new.local", "timeout": 30, "path_mappings": {}},
|
||||
expected_revision=state["revision"],
|
||||
)
|
||||
await proxy._retire_proxy_sessions(new_snapshot.generation)
|
||||
assert old_session.closed is False
|
||||
|
||||
new_lease = proxy._proxy_session_lease(new_snapshot)
|
||||
new_session = await new_lease.__aenter__()
|
||||
assert new_session is not old_session
|
||||
await new_lease.__aexit__(None, None, None)
|
||||
assert new_session.closed is False
|
||||
|
||||
await old_lease.__aexit__(None, None, None)
|
||||
assert old_session.closed is True
|
||||
assert new_session.closed is False
|
||||
|
||||
await proxy._close_all_proxy_sessions()
|
||||
assert new_session.closed is True
|
||||
|
||||
|
||||
class FakeWebSocket:
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
self.close_code = None
|
||||
|
||||
async def close(self, *, code=None, message=None):
|
||||
self.closed = True
|
||||
self.close_code = code
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_rotation_closes_only_retired_generation(isolated_proxy):
|
||||
proxy, config = isolated_proxy
|
||||
current_generation = config.generation
|
||||
old_bridge = proxy._ActiveWebSocket(
|
||||
current_generation - 1, FakeWebSocket(), FakeSession()
|
||||
)
|
||||
current_bridge = proxy._ActiveWebSocket(
|
||||
current_generation, FakeWebSocket(), FakeSession()
|
||||
)
|
||||
proxy._active_proxy_websockets.update({old_bridge, current_bridge})
|
||||
|
||||
await proxy._rotate_active_websockets(current_generation)
|
||||
|
||||
assert old_bridge.local_ws.closed is True
|
||||
assert old_bridge.local_ws.close_code == 1012
|
||||
assert old_bridge.session.closed is True
|
||||
assert current_bridge.local_ws.closed is False
|
||||
assert proxy._active_proxy_websockets == {current_bridge}
|
||||
|
||||
await proxy._rotate_active_websockets(None)
|
||||
assert current_bridge.local_ws.closed is True
|
||||
assert current_bridge.session.closed is True
|
||||
|
||||
|
||||
def test_mapped_local_path_is_resolved_against_comfy_lora_roots(
|
||||
modules, tmp_path, monkeypatch
|
||||
):
|
||||
root = tmp_path / "models" / "loras"
|
||||
fake_folder_paths = types.SimpleNamespace(
|
||||
get_folder_paths=lambda model_type: [str(root)]
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", fake_folder_paths)
|
||||
relative = modules.remote_client.RemoteLoraClient._relative_lora_path(
|
||||
str(root / "styles" / "portrait.safetensors"),
|
||||
"wrong-remote-folder",
|
||||
)
|
||||
assert relative == "styles/portrait.safetensors"
|
||||
Reference in New Issue
Block a user