diff --git a/core/database.py b/core/database.py index 1621492c7..1c9fd2d98 100644 --- a/core/database.py +++ b/core/database.py @@ -568,6 +568,7 @@ class ModelEndpoint(TimestampMixin, Base): hidden_models = Column(Text, nullable=True) # JSON list of model IDs that failed probing cached_models = Column(Text, nullable=True) # JSON list of last-known model IDs (avoids probe on list) pinned_models = Column(Text, nullable=True) # JSON list of admin-pinned model IDs (manual, may not appear in /v1/models) + cookbook_serve_session_id = Column(String, nullable=True) # Current Cookbook serve identity model_type = Column(String, nullable=True, default="llm") # "llm" or "image" # auto = classify by URL; local = self-hosted server; api/proxy = external # OpenAI-compatible API even when reachable through a private/tailnet IP. @@ -1378,6 +1379,26 @@ def _migrate_add_pinned_models_column(): except Exception: pass + +def _migrate_add_cookbook_serve_session_id_column(): + """Bind Cookbook cleanup to the serve session that owns the endpoint.""" + from sqlalchemy import inspect + + with engine.begin() as conn: + columns = { + column["name"] + for column in inspect(conn).get_columns("model_endpoints") + } + + if "cookbook_serve_session_id" not in columns: + conn.execute( + text( + "ALTER TABLE model_endpoints " + "ADD COLUMN cookbook_serve_session_id VARCHAR" + ) + ) + + def _migrate_add_notes_sort_order(): """Add sort_order, image_url, repeat columns to notes if they don't exist.""" import sqlite3 @@ -2390,6 +2411,7 @@ def init_db(): _migrate_add_hidden_models_column() _migrate_add_cached_models_column() _migrate_add_pinned_models_column() + _migrate_add_cookbook_serve_session_id_column() _migrate_add_notes_sort_order() _migrate_add_model_type_column() _migrate_add_model_endpoint_refresh_columns() diff --git a/routes/cookbook_routes.py b/routes/cookbook_routes.py index 89f892dd7..408693cb3 100644 --- a/routes/cookbook_routes.py +++ b/routes/cookbook_routes.py @@ -1408,7 +1408,11 @@ def setup_cookbook_routes() -> APIRouter: return {"models": models, "host": host or "local"} - def _auto_register_image_endpoint(req: ServeRequest, remote: str | None) -> str | None: + def _auto_register_image_endpoint( + req: ServeRequest, + remote: str | None, + serve_session_id: str | None = None, + ) -> str | None: """Register a diffusion model as an image endpoint so it appears in the model selector.""" import re from core.database import SessionLocal, ModelEndpoint @@ -1442,6 +1446,7 @@ def setup_cookbook_routes() -> APIRouter: existing.name = display_name existing.endpoint_kind = "local" existing.model_refresh_mode = "manual" + existing.cookbook_serve_session_id = serve_session_id if pinned_models: existing.cached_models = json.dumps(pinned_models) existing.pinned_models = json.dumps(pinned_models) @@ -1465,6 +1470,7 @@ def setup_cookbook_routes() -> APIRouter: model_refresh_mode="manual", cached_models=json.dumps(pinned_models) if pinned_models else None, pinned_models=json.dumps(pinned_models) if pinned_models else None, + cookbook_serve_session_id=serve_session_id, ) db.add(ep) db.commit() @@ -1610,43 +1616,71 @@ def setup_cookbook_routes() -> APIRouter: # let the probe layer mark it offline if nothing's listening. logger.info(f"crash-watchdog: serve {session_id} exited cleanly (0); leaving endpoint {endpoint_id}") return - # Non-zero exit — drop the endpoint. + # A nonzero exit may delete only the endpoint owned by this serve. try: from core.database import SessionLocal as _SL, ModelEndpoint as _ME + from core.constants import internal_api_base as _internal_api_base + from core.middleware import INTERNAL_TOOL_HEADER, INTERNAL_TOOL_TOKEN + from urllib.parse import quote + import httpx + db = _SL() try: ep = db.query(_ME).filter(_ME.id == endpoint_id).first() - if ep: - # A scheduled serve can leave old non-zero exit markers - # in tmux scrollback while the current OpenAI endpoint is - # actually alive. Verify reachability before deleting the - # endpoint row; otherwise chats fall back even though the - # served model is ready. - try: - probe_url = ep.base_url.rstrip("/") + "/models" - with urllib.request.urlopen(probe_url, timeout=3) as resp: - if 200 <= getattr(resp, "status", 0) < 300: - logger.info( - f"crash-watchdog: serve {session_id} has exit marker {exit_code} " - f"but endpoint {ep.id} is reachable; leaving it registered" - ) - return - except Exception: - pass + if not ep or ep.cookbook_serve_session_id != session_id: logger.info( - f"crash-watchdog: dropping endpoint {endpoint_id} " - f"({ep.name} @ {ep.base_url}) — serve exited {exit_code}" + "crash-watchdog: endpoint %s no longer owned by %s", + endpoint_id, session_id, ) - db.delete(ep) - db.commit() + return + probe_url = ep.base_url.rstrip("/") + "/models" finally: db.close() + + # Preserve an endpoint that is already healthy, including + # after an old runner emits a stale failure marker. + try: + with urllib.request.urlopen(probe_url, timeout=3) as resp: + if 200 <= getattr(resp, "status", 0) < 300: + logger.info( + "crash-watchdog: endpoint %s is reachable; keeping it", + endpoint_id, + ) + return + except Exception: + pass + + async with httpx.AsyncClient(timeout=8) as client: + response = await client.delete( + f"{_internal_api_base()}/api/model-endpoints/{quote(endpoint_id, safe='')}", + params={"expected_cookbook_serve_session_id": session_id}, + headers={INTERNAL_TOOL_HEADER: INTERNAL_TOOL_TOKEN}, + ) + if response.status_code == 200: + logger.info( + "crash-watchdog: deleted endpoint %s after serve %s exited %s", + endpoint_id, session_id, exit_code, + ) + elif response.status_code in (404, 409): + logger.info( + "crash-watchdog: endpoint %s absent or reassigned; keeping it", + endpoint_id, + ) + else: + logger.warning( + "crash-watchdog: endpoint %s deletion returned HTTP %s", + endpoint_id, response.status_code, + ) except Exception as e: - logger.warning(f"crash-watchdog: endpoint cleanup failed: {e!r}") + logger.warning("crash-watchdog: endpoint cleanup failed: %r", e) return logger.debug(f"crash-watchdog: no exit marker for {session_id} within window; leaving endpoint {endpoint_id}") - def _auto_register_llm_endpoint(req: ServeRequest, remote: str | None) -> str | None: + def _auto_register_llm_endpoint( + req: ServeRequest, + remote: str | None, + serve_session_id: str | None = None, + ) -> str | None: """Register a freshly-served LLM as a model endpoint so it appears in the model picker without a manual /setup step — the text-model sibling of _auto_register_image_endpoint. @@ -1663,6 +1697,7 @@ def setup_cookbook_routes() -> APIRouter: ) import re from core.database import SessionLocal, ModelEndpoint + from routes.model_routes import _normalize_model_ids # Port: ordered fallbacks so we match whatever the user actually # asked for, not a hardcoded default: @@ -1724,8 +1759,12 @@ def setup_cookbook_routes() -> APIRouter: # deployment alias from /v1/models, and a stale server can answer on the # same port while the new launch failed. Keeping the requested model id # pinned makes the picker reflect the actual launch intent. - pinned_models = [mlx_shim_model_id] if mlx_shim_model_id else ([req.repo_id] if req.repo_id else []) + launched_model = mlx_shim_model_id or req.repo_id + pinned_models = [launched_model] if launched_model else [] + # Launch registration must not delete other endpoints. + # Matching names, model pins, and failed probes do not prove that + # another serve has stopped. Lifecycle cleanup owns termination. db = SessionLocal() try: # Reuse an endpoint already pointed at this URL instead of duplicating. @@ -1736,16 +1775,16 @@ def setup_cookbook_routes() -> APIRouter: existing.name = display_name existing.endpoint_kind = "local" existing.model_refresh_mode = "auto" - if pinned_models: - try: - existing_pinned = json.loads(existing.pinned_models or "[]") - except Exception: - existing_pinned = [] - merged_pinned = [] - for mid in [*existing_pinned, *pinned_models]: - if mid and mid not in merged_pinned: - merged_pinned.append(mid) - existing.pinned_models = json.dumps(merged_pinned) if merged_pinned else None + # A reused endpoint belongs to the newest serve session. + existing.cookbook_serve_session_id = serve_session_id + if launched_model: + # Pins represent model selection, not launch history. + existing_pins = _normalize_model_ids( + existing.pinned_models + ) + if launched_model not in existing_pins: + existing_pins.append(launched_model) + existing.pinned_models = json.dumps(existing_pins) if is_ollama_endpoint: existing.endpoint_kind = "ollama" if pinned_models: @@ -1771,22 +1810,6 @@ def setup_cookbook_routes() -> APIRouter: db.commit() except Exception as _pe: logger.warning(f"Re-probe failed for {base_url}: {_pe!r}") - # Sweep stale dupes: other endpoints with the same display name - # at DIFFERENT URLs (likely failed earlier-attempt ports) get - # deleted so the picker doesn't show an offline ghost next to - # the working one. Only sweeps endpoints whose id starts with - # `local-` so we never touch a user's hand-added DeepSeek/OpenAI/ - # etc. entry with a coincidentally matching name. - stale = (db.query(ModelEndpoint) - .filter(ModelEndpoint.name == display_name) - .filter(ModelEndpoint.base_url != base_url) - .filter(ModelEndpoint.id.like("local-%")) - .all()) - for s in stale: - logger.info(f"Sweeping stale local endpoint {s.id} ({s.base_url})") - db.delete(s) - if stale: - db.commit() return existing.id ep_id = f"local-{uuid.uuid4().hex[:8]}" @@ -1802,22 +1825,11 @@ def setup_cookbook_routes() -> APIRouter: cached_models=json.dumps(pinned_models) if pinned_models else None, pinned_models=json.dumps(pinned_models) if pinned_models else None, supports_tools=supports_tools, + cookbook_serve_session_id=serve_session_id, ) db.add(ep) db.commit() logger.info(f"Auto-registered local model endpoint: {display_name} @ {base_url}") - # Same sweep on first-register path: drop any pre-existing local-* - # endpoints with this display name pointed elsewhere. - stale = (db.query(ModelEndpoint) - .filter(ModelEndpoint.name == display_name) - .filter(ModelEndpoint.id != ep_id) - .filter(ModelEndpoint.id.like("local-%")) - .all()) - for s in stale: - logger.info(f"Sweeping stale local endpoint {s.id} ({s.base_url})") - db.delete(s) - if stale: - db.commit() # Probe /v1/models NOW and write cached_models so the chat # picker actually shows the model on the next /api/models # call. Without this immediate probe, the endpoint has empty @@ -2639,9 +2651,17 @@ def setup_cookbook_routes() -> APIRouter: endpoint_id = None is_image_endpoint = "diffusion_server.py" in req.cmd or "mlx_image_server.py" in req.cmd if is_image_endpoint: - endpoint_id = _auto_register_image_endpoint(req, remote) + endpoint_id = _auto_register_image_endpoint( + req, + remote, + serve_session_id=session_id, + ) elif not is_pip_install: - endpoint_id = _auto_register_llm_endpoint(req, remote) + endpoint_id = _auto_register_llm_endpoint( + req, + remote, + serve_session_id=session_id, + ) # Crash watchdog: the auto-register above writes the endpoint row # IMMEDIATELY (before the server has even bound its port) so the diff --git a/routes/model_routes.py b/routes/model_routes.py index c79049b36..bacb575ff 100644 --- a/routes/model_routes.py +++ b/routes/model_routes.py @@ -3044,20 +3044,44 @@ def setup_model_routes(model_discovery): return {"dependents": _settings_using_endpoint(ep_id)} @router.delete("/model-endpoints/{ep_id}") - def delete_model_endpoint(ep_id: str, request: Request): + def delete_model_endpoint( + ep_id: str, + request: Request, + expected_cookbook_serve_session_id: Optional[str] = None, + ): require_admin(request) db = SessionLocal() try: ep = db.query(ModelEndpoint).filter(ModelEndpoint.id == ep_id).first() if not ep or not _chatgpt_endpoint_visible(ep, request): raise HTTPException(404, "Endpoint not found") + if expected_cookbook_serve_session_id is not None: + # The DELETE itself checks ownership. An old serve must never + # mutate the settings of an endpoint reused by another serve. + if ( + not ep_id.startswith(("local-", "img-")) + or not expected_cookbook_serve_session_id + ): + raise HTTPException(409, "Cookbook endpoint ownership mismatch") + deleted = ( + db.query(ModelEndpoint) + .filter( + ModelEndpoint.id == ep_id, + ModelEndpoint.cookbook_serve_session_id + == expected_cookbook_serve_session_id, + ) + .delete(synchronize_session=False) + ) + if not deleted: + raise HTTPException(409, "Cookbook endpoint ownership changed") # Clean up any settings that reference this endpoint cleared = _clear_settings_for_endpoint(ep_id) cleared_user_preferences = _clear_user_prefs_for_endpoint(ep_id) cleared_sessions = _clear_sessions_for_endpoint(db, ep.base_url, ep.id, ep.owner) cleared_loaded_sessions = _clear_loaded_sessions_for_endpoint(ep.base_url, ep.id, ep.owner) auth_id = getattr(ep, "provider_auth_id", None) - db.delete(ep) + if expected_cookbook_serve_session_id is None: + db.delete(ep) cleared_provider_auth = _delete_orphaned_provider_auth(db, auth_id, exclude_ep_id=ep_id) db.commit() _invalidate_models_cache() diff --git a/src/cookbook_serve_lifecycle.py b/src/cookbook_serve_lifecycle.py index a65893cab..e76ae956b 100644 --- a/src/cookbook_serve_lifecycle.py +++ b/src/cookbook_serve_lifecycle.py @@ -31,63 +31,43 @@ def _internal_headers() -> dict: async def _delete_endpoint_for_task(task: dict) -> None: - """Drop the auto-registered model endpoint for a scheduled-stop serve. + """Remove a stopped serve's endpoint using the normal deletion contract.""" + import re + from urllib.parse import quote - Without this, killing the tmux session leaves the endpoint sitting in - the picker (probe goes offline; chats still try to route there) and - the user has to delete it by hand in Settings -> Endpoints. - """ - endpoint_id = (task.get("_endpointId") or task.get("endpointId") or "").strip() - if not endpoint_id: - logger.info( - "cookbook_serve_lifecycle: task %s has no endpoint id; skipping endpoint deletion", - task.get("sessionId") or task.get("id") or "", - ) + endpoint_id = str(task.get("_endpointId") or task.get("endpointId") or "").strip() + serve_session_id = str(task.get("sessionId") or task.get("id") or "").strip() + if ( + not re.fullmatch(r"(?:local|img)-[0-9a-f]{8}", endpoint_id) + or not serve_session_id + ): + logger.info("cookbook_serve_lifecycle: missing endpoint/session identity") return - import re as _re - payload = task.get("payload") or {} - cmd = str(payload.get("_cmd") or "") - remote = task.get("remoteHost") or "" - # Build host the same way _auto_register_llm_endpoint does so URL match wins. - if remote: - host = remote.split("@")[-1] if "@" in remote else remote - else: - host = "host.docker.internal" - port_match = _re.search(r"--port\s+(\d+)", cmd) - ollama_host_match = _re.search(r"OLLAMA_HOST=[^\s]*?:(\d+)", cmd) - if port_match: - port = int(port_match.group(1)) - elif ollama_host_match: - port = int(ollama_host_match.group(1)) - elif "ollama" in cmd: - port = 11434 - else: - port = 8080 - base_url = f"http://{host}:{port}/v1" + try: async with httpx.AsyncClient(timeout=8) as client: - r = await client.get( - f"{internal_api_base()}/api/model-endpoints", + response = await client.delete( + f"{internal_api_base()}/api/model-endpoints/{quote(endpoint_id, safe='')}", + params={"expected_cookbook_serve_session_id": serve_session_id}, headers=_internal_headers(), ) - if r.status_code >= 400: - return - eps = r.json() if r.content else [] - # Delete only the endpoint created by this scheduled serve. URL - # matching is unsafe because a later scheduled serve can reuse the - # same host:port after an older task has gone stale. - ep = next((e for e in eps if e.get("id") == endpoint_id), None) - if ep: - await client.delete( - f"{internal_api_base()}/api/model-endpoints/{ep['id']}", - headers=_internal_headers(), - ) - logger.info( - f"cookbook_serve_lifecycle: deleted endpoint {ep.get('id')} " - f"({ep.get('base_url')}) after scheduled stop" - ) - except Exception as e: - logger.warning(f"cookbook_serve_lifecycle: endpoint delete failed: {e}") + if response.status_code == 200: + logger.info( + "cookbook_serve_lifecycle: deleted endpoint %s for serve %s", + endpoint_id, serve_session_id, + ) + elif response.status_code in (404, 409): + logger.info( + "cookbook_serve_lifecycle: endpoint %s absent or reassigned; keeping it", + endpoint_id, + ) + else: + logger.warning( + "cookbook_serve_lifecycle: endpoint %s deletion returned HTTP %s", + endpoint_id, response.status_code, + ) + except Exception as exc: + logger.warning("cookbook_serve_lifecycle: cleanup failed: %s", exc) async def _stop_serve(session_id: str, remote_host: str = "", ssh_port: str = "") -> bool: diff --git a/tests/test_cookbook_endpoint_registration.py b/tests/test_cookbook_endpoint_registration.py index 8e3a9b994..f413eff25 100644 --- a/tests/test_cookbook_endpoint_registration.py +++ b/tests/test_cookbook_endpoint_registration.py @@ -1,5 +1,9 @@ from pathlib import Path +import pytest + +from tests.helpers.database import disposable_database + ROOT = Path(__file__).resolve().parents[1] COOKBOOK_RUNNING = ROOT / "static" / "js" / "cookbookRunning.js" @@ -28,3 +32,416 @@ def test_cookbook_advertised_bind_urls_keep_connectable_host(): assert "function _endpointFromAdvertisedUrl" in src assert "_isAnyBindHost(u.hostname) ? currentHost" in src assert "host = u.hostname || host;" not in src + + +@pytest.fixture +def endpoint_db(tmp_path, monkeypatch): + import core.database as database + import routes.model_routes as model_routes + + with disposable_database(tmp_path) as factory: + monkeypatch.setattr(database, "SessionLocal", factory) + monkeypatch.setattr(model_routes, "_probe_endpoint", lambda *args, **kwargs: None) + yield factory + + +def _launch_llm( + repo_id: str, + port: int, + serve_session_id: str | None = None, +) -> str: + import routes.cookbook_routes as cookbook_routes + from routes.cookbook_helpers import ServeRequest + + router = cookbook_routes.setup_cookbook_routes() + serve = next( + route.endpoint for route in router.routes + if route.path == "/api/model/serve" and "POST" in route.methods + ) + cells = dict(zip(serve.__code__.co_freevars, serve.__closure__ or ())) + register = cells["_auto_register_llm_endpoint"].cell_contents + endpoint_id = register( + ServeRequest(repo_id=repo_id, cmd=f"vllm serve {repo_id} --port {port}"), + None, + serve_session_id=serve_session_id, + ) + assert endpoint_id + return endpoint_id + + +def _endpoint(factory, endpoint_id: str): + import core.database as database + + with factory() as db: + return db.get(database.ModelEndpoint, endpoint_id) + + +@pytest.mark.parametrize("reuse", [False, True]) +def test_launch_keeps_endpoint_of_other_model_with_same_short_name(endpoint_db, reuse): + if reuse: + mirror = _launch_llm("unsloth/Llama-3.1-8B-Instruct", 8002) + original = _launch_llm("meta-llama/Llama-3.1-8B-Instruct", 8001) + + launched = _launch_llm("unsloth/Llama-3.1-8B-Instruct", 8002) + + if reuse: + assert launched == mirror + assert _endpoint(endpoint_db, original) is not None + assert _endpoint(endpoint_db, launched) is not None + + +@pytest.mark.parametrize("reuse", [False, True]) +def test_relaunch_on_another_port_preserves_existing_endpoint(endpoint_db, reuse): + if reuse: + current = _launch_llm("meta-llama/Llama-3.1-8B-Instruct", 8002) + stale = _launch_llm("Qwen/Qwen3-8B", 8001) + + launched = _launch_llm("Qwen/Qwen3-8B", 8002) + + if reuse: + assert launched == current + assert _endpoint(endpoint_db, stale) is not None + assert _endpoint(endpoint_db, launched) is not None + + +def test_sweep_ignores_endpoint_now_serving_a_different_model(endpoint_db): + endpoint = _launch_llm("Qwen/Qwen3-8B", 8001) + _launch_llm("meta-llama/Llama-3.1-8B", 8001) + _launch_llm("Qwen/Qwen3-8B", 8001) + + _launch_llm("meta-llama/Llama-3.1-8B", 8002) + + assert _endpoint(endpoint_db, endpoint) is not None + + +def test_sweep_ignores_endpoints_not_created_by_cookbook(endpoint_db): + import core.database as database + + with endpoint_db() as db: + db.add(database.ModelEndpoint( + id="manual", name="Qwen3-8B", base_url="http://localhost:9000/v1", + pinned_models='["Qwen/Qwen3-8B"]', + )) + db.commit() + + _launch_llm("Qwen/Qwen3-8B", 8002) + + assert _endpoint(endpoint_db, "manual") is not None + + +def test_manual_pin_order_does_not_change_on_relaunch(endpoint_db): + import json + from core.database import ModelEndpoint + + endpoint_id = _launch_llm("organization-a/Model-A", 8001) + + with endpoint_db() as db: + endpoint = db.get(ModelEndpoint, endpoint_id) + endpoint.pinned_models = json.dumps([ + "organization-a/Model-A", + "organization-b/Model-B", + ]) + db.commit() + + _launch_llm("organization-a/Model-A", 8001) + + with endpoint_db() as db: + endpoint = db.get(ModelEndpoint, endpoint_id) + assert json.loads(endpoint.pinned_models) == [ + "organization-a/Model-A", + "organization-b/Model-B", + ] + + +def test_manual_pins_cannot_trigger_endpoint_deletion(endpoint_db): + import json + from core.database import ModelEndpoint + + original = _launch_llm("organization-a/Model-A", 8001) + + with endpoint_db() as db: + endpoint = db.get(ModelEndpoint, original) + endpoint.pinned_models = json.dumps([ + "organization-a/Model-A", + "organization-b/Model-B", + ]) + db.commit() + + _launch_llm("organization-b/Model-B", 8002) + + assert _endpoint(endpoint_db, original) is not None + + +def test_unavailable_model_probe_never_deletes_another_endpoint( + endpoint_db, monkeypatch +): + import routes.model_routes as model_routes + + original = _launch_llm("organization-a/Model-A", 8001) + + monkeypatch.setattr( + model_routes, + "_probe_endpoint", + lambda *args, **kwargs: [], + ) + + _launch_llm("organization-a/Model-A", 8002) + + assert _endpoint(endpoint_db, original) is not None + + +@pytest.mark.asyncio +async def test_old_watchdog_cannot_delete_reused_endpoint( + endpoint_db, + monkeypatch, + endpoint_http_api, +): + import asyncio + import urllib.request + import routes.cookbook_routes as cookbook_routes + + original = _launch_llm( + "organization-a/Model-A", + 8001, + serve_session_id="serve-old", + ) + + reused = _launch_llm( + "organization-b/Model-B", + 8001, + serve_session_id="serve-new", + ) + + assert original == reused + + with endpoint_db() as db: + from core.database import ModelEndpoint + + endpoint = db.get(ModelEndpoint, reused) + assert endpoint.cookbook_serve_session_id == "serve-new" + + router = cookbook_routes.setup_cookbook_routes() + serve = next( + route.endpoint + for route in router.routes + if route.path == "/api/model/serve" + and "POST" in route.methods + ) + + cells = dict( + zip( + serve.__code__.co_freevars, + serve.__closure__ or (), + ) + ) + + watchdog = cells["_serve_crash_watchdog"].cell_contents + + class FakeProcess: + async def communicate(self): + return ( + b"=== Process exited with code 1 ===\n", + b"", + ) + + async def fake_subprocess(*args, **kwargs): + return FakeProcess() + + async def fake_sleep(seconds): + return None + + def unavailable(*args, **kwargs): + raise OSError("Server still warming up") + + monkeypatch.setattr( + asyncio, + "create_subprocess_exec", + fake_subprocess, + ) + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + monkeypatch.setattr( + urllib.request, + "urlopen", + unavailable, + ) + + await watchdog( + endpoint_id=original, + session_id="serve-old", + remote=None, + ssh_port=None, + is_windows=False, + ) + + assert _endpoint(endpoint_db, reused) is not None + + # The current serve's own crash watchdog must still be able + # to remove its endpoint when it actually fails. + await watchdog( + endpoint_id=reused, + session_id="serve-new", + remote=None, + ssh_port=None, + is_windows=False, + ) + + assert _endpoint(endpoint_db, reused) is None + + +def test_cookbook_serve_session_migration_is_idempotent( + tmp_path, + monkeypatch, +): + from sqlalchemy import create_engine, inspect, text + import core.database as database + + engine = create_engine( + f"sqlite:///{tmp_path / 'legacy-endpoints.db'}" + ) + + try: + with engine.begin() as conn: + conn.execute( + text( + "CREATE TABLE model_endpoints " + "(id VARCHAR PRIMARY KEY)" + ) + ) + + monkeypatch.setattr(database, "engine", engine) + + database._migrate_add_cookbook_serve_session_id_column() + database._migrate_add_cookbook_serve_session_id_column() + + columns = { + column["name"] + for column in inspect(engine).get_columns( + "model_endpoints" + ) + } + + assert "cookbook_serve_session_id" in columns + + finally: + engine.dispose() + + +@pytest.fixture +def endpoint_http_api(endpoint_db, monkeypatch): + """Route deletion through a disposable in-process FastAPI app.""" + import httpx + from fastapi import FastAPI + import routes.model_routes as model_routes + import routes.prefs_routes as prefs_routes + import core.constants as constants + import src.cookbook_serve_lifecycle as lifecycle + + monkeypatch.setattr(model_routes, "SessionLocal", endpoint_db) + monkeypatch.setattr(model_routes, "require_admin", lambda request: None) + + settings = {} + preferences = {} + monkeypatch.setattr(model_routes, "_load_settings", lambda: settings) + monkeypatch.setattr(model_routes, "_save_settings", lambda value: settings.update(value)) + monkeypatch.setattr(prefs_routes, "_load", lambda: preferences) + monkeypatch.setattr(prefs_routes, "_save", lambda value: preferences.update(value)) + + # Endpoint deletion must not touch the user's running sessions. + import src.ai_interaction as ai_interaction + monkeypatch.setattr(ai_interaction, "get_session_manager", lambda: None) + + app = FastAPI() + router = model_routes.setup_model_routes(None) + app.include_router(router) + + original_client = httpx.AsyncClient + transport = httpx.ASGITransport(app=app) + + def client_factory(*args, **kwargs): + kwargs["transport"] = transport + return original_client(*args, **kwargs) + + monkeypatch.setattr(httpx, "AsyncClient", client_factory) + monkeypatch.setattr(constants, "internal_api_base", lambda: "http://testserver") + monkeypatch.setattr(lifecycle, "internal_api_base", lambda: "http://testserver") + + return {"settings": settings, "preferences": preferences, "router": router} + + +@pytest.mark.asyncio +async def test_ownership_condition_preserves_settings_on_mismatch( + endpoint_db, endpoint_http_api, +): + import httpx + original = _launch_llm("organization-a/Model-A", 8001, "serve-old") + reused = _launch_llm("organization-b/Model-B", 8001, "serve-new") + assert reused == original + settings = endpoint_http_api["settings"] + settings.update({"default_endpoint_id": reused, "default_model": "organization-b/Model-B"}) + async with httpx.AsyncClient() as client: + response = await client.delete( + f"http://testserver/api/model-endpoints/{reused}", + params={"expected_cookbook_serve_session_id": "serve-old"}, + ) + assert response.status_code == 409 + assert _endpoint(endpoint_db, reused) is not None + assert settings["default_endpoint_id"] == reused + + +@pytest.mark.asyncio +async def test_owned_deletion_clears_settings_and_model_cache( + endpoint_db, endpoint_http_api, +): + import httpx + from core.database import Session as DbSession + + endpoint_id = _launch_llm("organization-a/Model-A", 8001, "serve-current") + settings = endpoint_http_api["settings"] + settings.update({ + "default_endpoint_id": endpoint_id, + "default_model": "organization-a/Model-A", + "foreground_model_fallbacks": [{"endpoint_id": endpoint_id, "model": "x"}], + }) + with endpoint_db() as db: + db.add(DbSession( + id="owned-delete-session", name="Test", owner=None, + endpoint_url="http://localhost:8001/v1", + endpoint_id=endpoint_id, model="organization-a/Model-A", + headers={"Authorization": "Bearer test-only"}, + )) + db.commit() + + # Seed the actual model listing cache captured by the route closure. + router = endpoint_http_api["router"] + route = next(r for r in router.routes if r.path == "/api/model-endpoints/{ep_id}" and "DELETE" in r.methods) + closure = dict(zip(route.endpoint.__code__.co_freevars, route.endpoint.__closure__ or ())) + invalidate = closure["_invalidate_models_cache"].cell_contents + cache_closure = dict(zip(invalidate.__code__.co_freevars, invalidate.__closure__ or ())) + models_cache = cache_closure["_models_cache"].cell_contents + models_cache["test"] = {"data": [], "time": 0} + + async with httpx.AsyncClient() as client: + response = await client.delete( + f"http://testserver/api/model-endpoints/{endpoint_id}", + params={"expected_cookbook_serve_session_id": "serve-current"}, + ) + assert response.status_code == 200, response.text + assert _endpoint(endpoint_db, endpoint_id) is None + assert settings["default_endpoint_id"] == "" + assert settings["foreground_model_fallbacks"] == [] + assert not models_cache + with endpoint_db() as db: + sess = db.get(DbSession, "owned-delete-session") + assert sess.headers == {} + + +@pytest.mark.asyncio +async def test_ordinary_admin_delete_does_not_require_serve_owner( + endpoint_db, endpoint_http_api, +): + import httpx + endpoint_id = _launch_llm("organization-a/Model-A", 8001, "serve-new") + async with httpx.AsyncClient() as client: + response = await client.delete(f"http://testserver/api/model-endpoints/{endpoint_id}") + assert response.status_code == 200, response.text + assert _endpoint(endpoint_db, endpoint_id) is None diff --git a/tests/test_cookbook_serve_lifecycle.py b/tests/test_cookbook_serve_lifecycle.py index cd32d4314..f8fd25720 100644 --- a/tests/test_cookbook_serve_lifecycle.py +++ b/tests/test_cookbook_serve_lifecycle.py @@ -4,6 +4,13 @@ import pytest from src import cookbook_serve_lifecycle as lifecycle +from tests.test_cookbook_endpoint_registration import ( + endpoint_db, + endpoint_http_api, + _endpoint, + _launch_llm, +) + @pytest.mark.asyncio async def test_tick_persists_only_successfully_stopped_serves(tmp_path, monkeypatch): @@ -48,3 +55,219 @@ async def test_tick_persists_only_successfully_stopped_serves(tmp_path, monkeypa assert tasks["stop-succeeds"]["_scheduledStopAtMs"] is None assert tasks["stop-fails"]["status"] == "running" assert tasks["stop-fails"]["_scheduledStopAtMs"] == 0 + + + +@pytest.mark.asyncio +async def test_scheduled_stop_preserves_endpoint_reused_by_newer_serve( + endpoint_db, + endpoint_http_api, +): + original = _launch_llm( + "organization-a/Model-A", + 8001, + serve_session_id="serve-old", + ) + + reused = _launch_llm( + "organization-b/Model-B", + 8001, + serve_session_id="serve-new", + ) + + assert original == reused + + await lifecycle._delete_endpoint_for_task({ + "_endpointId": original, + "sessionId": "serve-old", + "type": "serve", + }) + + assert _endpoint(endpoint_db, reused) is not None + + await lifecycle._delete_endpoint_for_task({ + "_endpointId": reused, + "sessionId": "serve-new", + "type": "serve", + }) + + assert _endpoint(endpoint_db, reused) is None + + +@pytest.mark.asyncio +async def test_scheduled_stop_preserves_legacy_endpoint_without_owner( + endpoint_db, + endpoint_http_api, +): + endpoint_id = _launch_llm( + "organization-a/Model-A", + 8002, + ) + + await lifecycle._delete_endpoint_for_task({ + "_endpointId": endpoint_id, + "sessionId": "serve-old", + "type": "serve", + }) + + assert _endpoint(endpoint_db, endpoint_id) is not None + + +@pytest.mark.asyncio +async def test_scheduled_stop_cannot_delete_manual_endpoint( + endpoint_db, + endpoint_http_api, +): + from core.database import ModelEndpoint + + with endpoint_db() as db: + db.add(ModelEndpoint( + id="manual-endpoint", + name="Manual", + base_url="http://localhost:9000/v1", + cookbook_serve_session_id="serve-old", + )) + db.commit() + + await lifecycle._delete_endpoint_for_task({ + "_endpointId": "manual-endpoint", + "sessionId": "serve-old", + "type": "serve", + }) + + assert _endpoint(endpoint_db, "manual-endpoint") is not None + + +@pytest.mark.asyncio +async def test_scheduled_stop_requires_session_identity(endpoint_db, endpoint_http_api): + endpoint_id = _launch_llm( + "organization-a/Model-A", + 8003, + serve_session_id="serve-current", + ) + + await lifecycle._delete_endpoint_for_task({ + "_endpointId": endpoint_id, + "type": "serve", + }) + + assert _endpoint(endpoint_db, endpoint_id) is not None + + + +def _launch_image(repo_id, port, serve_session_id): + import routes.cookbook_routes as cookbook_routes + from routes.cookbook_helpers import ServeRequest + + router = cookbook_routes.setup_cookbook_routes() + + serve = next( + route.endpoint + for route in router.routes + if route.path == "/api/model/serve" + and "POST" in route.methods + ) + + cells = dict( + zip( + serve.__code__.co_freevars, + serve.__closure__ or (), + ) + ) + + register = cells["_auto_register_image_endpoint"].cell_contents + + endpoint_id = register( + ServeRequest( + repo_id=repo_id, + cmd=f"python scripts/diffusion_server.py --port {port}", + ), + None, + serve_session_id=serve_session_id, + ) + + assert endpoint_id + return endpoint_id + + +@pytest.mark.asyncio +async def test_scheduled_stop_deletes_owned_image_endpoint( + endpoint_db, + monkeypatch, + endpoint_http_api, +): + import src.settings as settings + from core.database import ModelEndpoint + + monkeypatch.setattr(settings, "load_settings", lambda: {}) + monkeypatch.setattr( + settings, + "save_settings", + lambda value: None, + ) + + endpoint_id = _launch_image( + "organization-a/Image-A", + 8101, + "serve-image-a", + ) + + assert endpoint_id.startswith("img-") + + with endpoint_db() as db: + endpoint = db.get(ModelEndpoint, endpoint_id) + assert endpoint.cookbook_serve_session_id == "serve-image-a" + + await lifecycle._delete_endpoint_for_task({ + "_endpointId": endpoint_id, + "sessionId": "serve-image-a", + "type": "serve", + }) + + assert _endpoint(endpoint_db, endpoint_id) is None + + +@pytest.mark.asyncio +async def test_old_image_serve_cannot_delete_reused_endpoint( + endpoint_db, + monkeypatch, + endpoint_http_api, +): + import src.settings as settings + + monkeypatch.setattr(settings, "load_settings", lambda: {}) + monkeypatch.setattr( + settings, + "save_settings", + lambda value: None, + ) + + original = _launch_image( + "organization-a/Image-A", + 8102, + "serve-old", + ) + + reused = _launch_image( + "organization-b/Image-B", + 8102, + "serve-new", + ) + + assert original == reused + + await lifecycle._delete_endpoint_for_task({ + "_endpointId": original, + "sessionId": "serve-old", + "type": "serve", + }) + + assert _endpoint(endpoint_db, reused) is not None + + await lifecycle._delete_endpoint_for_task({ + "_endpointId": reused, + "sessionId": "serve-new", + "type": "serve", + }) + + assert _endpoint(endpoint_db, reused) is None