From eb6119c7971fa1a38cad102e6bab77c8c23a3ecc Mon Sep 17 00:00:00 2001 From: kstyagi Date: Wed, 23 Sep 2026 06:30:58 +0530 Subject: [PATCH] feat: add graceful shutdown endpoint for UVM API server and update documentation --- DEVELOPER.md | 7 +++++++ README.md | 13 ++++++++++++ USAGE.md | 15 ++++++++++++++ tests/test_server.py | 47 ++++++++++++++++++++++++++++++++++++-------- uvm/api_models.py | 4 ++++ uvm/routers/exit.py | 35 +++++++++++++++++++++++++++++++++ uvm/server.py | 7 ++++++- 7 files changed, 119 insertions(+), 9 deletions(-) create mode 100644 uvm/routers/exit.py diff --git a/DEVELOPER.md b/DEVELOPER.md index 990b8a1..8074151 100644 --- a/DEVELOPER.md +++ b/DEVELOPER.md @@ -808,6 +808,7 @@ or secret manager when that matters. | `GET` | `/` | `server.py` | Redirect to `/dashboard` | | `GET` | `/dashboard` | `dashboard.py` | Public static dashboard shell | | `GET` | `/health` | `health.py` | Unauthenticated liveness response | +| `POST` | `/exit` | `exit.py` | Request graceful Uvicorn shutdown | | `GET` | `/diagnostics` | `diagnostics.py` | Host readiness checks | | `POST` | `/install` | `installation.py` | `{ "force": false }` body | | `GET` | `/vms` | `vms.py` | List persisted VM records and observed status | @@ -829,6 +830,12 @@ All management routes use `get_authorized_application()` and require: X-UVM-Token: ``` +`POST /exit` is a management route. `run_server()` creates the Uvicorn server, +stores it on the FastAPI application state, and runs it in-process. The route +returns `202` before a background task sets `should_exit`, allowing Uvicorn to +deliver the response and perform its normal graceful shutdown. It does not +stop or delete any managed VM. + Request schemas reject unexpected fields. `UvmError` instances are rendered as JSON errors by `server.py`; Pydantic request validation uses FastAPI's normal 422 response format. Create accepts `name: str | null`, `ports: list | null`, diff --git a/README.md b/README.md index e1e228f..6b8070e 100644 --- a/README.md +++ b/README.md @@ -303,6 +303,7 @@ The server exposes these routes: | `GET` | `/` | Redirect to the local dashboard | | `GET` | `/dashboard` | Serve the local management dashboard | | `GET` | `/health` | Unauthenticated liveness response | +| `POST` | `/exit` | Gracefully stop the UVM API server | | `GET` | `/diagnostics` | Run host readiness checks for the dashboard/system view | | `POST` | `/install` | Install or refresh host assets; body: `{ "force": false }` | | `GET` | `/vms` | List VM records and observed status | @@ -364,6 +365,18 @@ curl \ http://127.0.0.1:8000/vms ``` +Stop the API server without stopping its managed VMs: + +```sh +curl -X POST \ + -H "X-UVM-Token: $UVM_API_TOKEN" \ + http://127.0.0.1:8000/exit +``` + +The endpoint returns `202 {"status":"stopping"}` and then gracefully terminates +the UVM API process. It is a protected management endpoint; it does not stop or +delete managed VMs. + Do not expose this server directly to an untrusted network. For remote access, keep TLS enabled and consider placing it behind an additional authenticated reverse proxy. diff --git a/USAGE.md b/USAGE.md index c608c84..ab95ce8 100644 --- a/USAGE.md +++ b/USAGE.md @@ -141,6 +141,21 @@ go to stderr. VM JSON omits passwords. Lifecycle and sync results are arrays of `changes`). `ports list --json` returns an array even for one VM. JSON does not disable confirmations, and errors may still be plain text on stderr. +### Stop the Management API + +The API can be stopped remotely through its authenticated shutdown endpoint. +This stops only the UVM API process; it does not stop or delete any VM: + +```sh +curl -X POST \ + -H "X-UVM-Token: $UVM_API_TOKEN" \ + http://127.0.0.1:8000/exit +``` + +The server first returns `202 {"status":"stopping"}`, then performs Uvicorn's +graceful shutdown. Use the actual API host and port when they differ from the +defaults. `GET /exit` is intentionally not supported. + ## 4. Install UVM Assets UVM installs host packages, Firecracker, Jailer, a guest kernel, and a guest diff --git a/tests/test_server.py b/tests/test_server.py index fc49cf1..d028ea8 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -7,7 +7,7 @@ import tempfile import unittest from pathlib import Path from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch from uvm.cli import build_parser, main from uvm.config import Settings @@ -99,19 +99,31 @@ class ServerCliTests(unittest.TestCase): application = fake_application(api_token="test-token") application.settings.api_tls_cert = certificate application.settings.api_tls_key = key - fake_uvicorn = SimpleNamespace(run=lambda *_args, **_kwargs: None) + api = SimpleNamespace(state=SimpleNamespace()) + server = SimpleNamespace(run=Mock()) + config = object() + fake_uvicorn = SimpleNamespace( + Config=Mock(return_value=config), + Server=Mock(return_value=server), + ) with ( patch("uvm.server._validate_server_settings"), - patch("uvm.server.create_api", return_value=object()), + patch("uvm.server.create_api", return_value=api), patch.dict(sys.modules, {"uvicorn": fake_uvicorn}), - patch.object(fake_uvicorn, "run") as run_mock, ): run_server(application, host="0.0.0.0", port=8443) - run_mock.assert_called_once() - self.assertEqual(run_mock.call_args.kwargs["ssl_certfile"], str(certificate)) - self.assertEqual(run_mock.call_args.kwargs["ssl_keyfile"], str(key)) + fake_uvicorn.Config.assert_called_once_with( + api, + host="0.0.0.0", + port=8443, + ssl_certfile=str(certificate), + ssl_keyfile=str(key), + ) + fake_uvicorn.Server.assert_called_once_with(config) + self.assertIs(api.state.uvicorn_server, server) + server.run.assert_called_once_with() def test_api_token_rejects_non_ascii_or_whitespace(self) -> None: for token in ("s\u00e9cret", "contains space", ""): @@ -300,7 +312,8 @@ class ApiRouterTests(unittest.TestCase): ) application.state_store = SimpleNamespace(initialize=lambda: None) self.headers = {"X-UVM-Token": "test-token"} - self.client = TestClient(create_api(application, host="127.0.0.1")) + self.api = create_api(application, host="127.0.0.1") + self.client = TestClient(self.api) self.addCleanup(self.client.close) def test_health_is_available_without_credentials(self) -> None: @@ -339,6 +352,7 @@ class ApiRouterTests(unittest.TestCase): ("POST", "/vms/test-vm/ports", {}), ("POST", "/vms/test-vm/ports/sync", {}), ("GET", "/diagnostics", {}), + ("POST", "/exit", {}), ("POST", "/install", {}), ): for headers in ({}, {"X-UVM-Token": "wrong-token"}): @@ -353,6 +367,22 @@ class ApiRouterTests(unittest.TestCase): self.assertEqual(authorized.json()[0]["id"], "vm-test") self.assertNotIn("password", authorized.json()[0]) + def test_exit_requests_a_graceful_authenticated_server_shutdown(self) -> None: + server = SimpleNamespace(should_exit=False) + self.api.state.uvicorn_server = server + + response = self.client.post("/exit", headers=self.headers) + + self.assertEqual(response.status_code, 202) + self.assertEqual(response.json(), {"status": "stopping"}) + self.assertTrue(server.should_exit) + + def test_exit_requires_a_running_uvicorn_server(self) -> None: + response = self.client.post("/exit", headers=self.headers) + + self.assertEqual(response.status_code, 503) + self.assertEqual(response.json(), {"detail": "server shutdown is unavailable outside uvm --serve"}) + def test_vm_lifecycle_routes_and_extra_field_validation(self) -> None: created = self.client.post( "/vms", @@ -752,6 +782,7 @@ class ApiRouterTests(unittest.TestCase): ("/", "get"), ("/dashboard", "get"), ("/health", "get"), + ("/exit", "post"), ("/install", "post"), ("/vms", "get"), ("/vms", "post"), diff --git a/uvm/api_models.py b/uvm/api_models.py index aee7751..2c7317a 100644 --- a/uvm/api_models.py +++ b/uvm/api_models.py @@ -14,6 +14,10 @@ class HealthResponse(BaseModel): status: str = "ok" +class ServerExitResponse(BaseModel): + status: str = "stopping" + + class InstallRequest(BaseModel): model_config = ConfigDict(extra="forbid") diff --git a/uvm/routers/exit.py b/uvm/routers/exit.py new file mode 100644 index 0000000..8226021 --- /dev/null +++ b/uvm/routers/exit.py @@ -0,0 +1,35 @@ +"""Authenticated graceful shutdown endpoint for the UVM API server.""" + +from __future__ import annotations + +from typing import Any + +from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Request, status + +from ..api_models import ServerExitResponse +from ..app import Application +from .dependencies import get_authorized_application + + +router = APIRouter(tags=["server"]) + + +@router.post("/exit", response_model=ServerExitResponse, status_code=status.HTTP_202_ACCEPTED) +def exit_server( + background_tasks: BackgroundTasks, + request: Request, + _application: Application = Depends(get_authorized_application), +) -> ServerExitResponse: + """Stop the Uvicorn process after this response has been delivered.""" + server = getattr(request.app.state, "uvicorn_server", None) + if server is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="server shutdown is unavailable outside uvm --serve", + ) + background_tasks.add_task(_request_server_exit, server) + return ServerExitResponse() + + +def _request_server_exit(server: Any) -> None: + server.should_exit = True diff --git a/uvm/server.py b/uvm/server.py index 95b0145..51e9ecd 100644 --- a/uvm/server.py +++ b/uvm/server.py @@ -35,6 +35,7 @@ def create_api(application: Application | None = None, *, host: str | None = Non ) from error from .routers.health import router as health_router + from .routers.exit import router as exit_router from .routers.installation import router as installation_router from .routers.vms import router as vms_router from .routers.diagnostics import router as diagnostics_router @@ -55,6 +56,7 @@ def create_api(application: Application | None = None, *, host: str | None = Non ) api.include_router(health_router) + api.include_router(exit_router) api.include_router(installation_router) api.include_router(vms_router) api.include_router(diagnostics_router) @@ -91,7 +93,10 @@ def run_server(application: Application, *, host: str, port: int) -> None: if settings.api_tls_cert is not None: options["ssl_certfile"] = str(settings.api_tls_cert) options["ssl_keyfile"] = str(settings.api_tls_key) - uvicorn.run(create_api(application, host=host), **options) + api = create_api(application, host=host) + server = uvicorn.Server(uvicorn.Config(api, **options)) + api.state.uvicorn_server = server + server.run() def _is_loopback_host(host: str) -> bool: