From 0b22c81e6c9c24d006a5b24b2fc5a2370e6c0baa Mon Sep 17 00:00:00 2001 From: kstyagi Date: Fri, 25 Sep 2026 23:03:25 +0530 Subject: [PATCH] feat: add tool catalog and installation endpoints with models for tool management --- uvm/api_models.py | 88 ++++++++++++++++++++++++++++++++++++++++++++ uvm/routers/tools.py | 46 +++++++++++++++++++++++ uvm/routers/vms.py | 8 ++++ uvm/server.py | 16 ++++++++ 4 files changed, 158 insertions(+) create mode 100644 uvm/routers/tools.py diff --git a/uvm/api_models.py b/uvm/api_models.py index 686801e..9517baf 100644 --- a/uvm/api_models.py +++ b/uvm/api_models.py @@ -11,6 +11,7 @@ from .domain import VmRecord from .errors import ValidationError from .guest_agent import DEFAULT_TIMEOUT_SECONDS, MAX_COMMAND_BYTES, MAX_TIMEOUT_SECONDS from .guest_exec import validate_exec +from .tool_catalog import MAX_SELECTIONS, ToolSelection class HealthResponse(BaseModel): @@ -42,6 +43,19 @@ class PortRuleModel(BaseModel): protocol: Literal["TCP", "HTTP", "HTTPS"] = "TCP" +class ToolSelectionModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + id: str = Field(strict=True, min_length=1, max_length=32) + version: str | None = Field(default=None, strict=True, max_length=16, + description="Required for Python: 3.11, 3.12, or 3.13. Omit for other tools.") + + @model_validator(mode="after") + def validate_selection(self): + ToolSelection(self.id, self.version) + return self + + class VmCreateRequest(BaseModel): model_config = ConfigDict(extra="forbid") @@ -54,6 +68,8 @@ class VmCreateRequest(BaseModel): name: str | None = None ports: list[PortRuleModel] | None = None no_default_ports: bool = Field(default=False, strict=True) + tools: list[ToolSelectionModel] = Field(default_factory=list, max_length=MAX_SELECTIONS, + description="Optional guest tools to install asynchronously after boot") class VmResizeRequest(BaseModel): @@ -149,6 +165,78 @@ class VmAgentInstallResponse(BaseModel): restart_required: bool = True +class ToolCatalogResponse(BaseModel): + id: str + name: str + description: str + version_label: str + versions: list[str] + source: str + + +class ToolInstallRequest(BaseModel): + model_config = ConfigDict(extra="forbid") + + tools: list[ToolSelectionModel] = Field(min_length=1, max_length=MAX_SELECTIONS) + start_if_stopped: bool = Field(default=False, strict=True) + restart: bool = Field(default=False, strict=True, description="Restart first to refresh the guest agent") + request_id: str | None = Field(default=None, strict=True, max_length=80, pattern=r"^[A-Za-z0-9_-]{1,80}$", + description="Optional idempotency key, retained with the most recent 10 jobs") + + +class ToolInventoryResponse(BaseModel): + key: str + id: str + version: str | None + installed: bool + detected_version: str | None + error: str | None + usage: str + + +class ToolJobItemResponse(BaseModel): + key: str + id: str + version: str | None + status: Literal["queued", "installing", "installed", "already_installed", "failed", "interrupted", "skipped"] + detected_version: str | None + error: str | None + output: str + truncated: bool + + +class ToolJobSummaryResponse(BaseModel): + id: str + status: Literal["queued", "starting", "waiting", "installing", "succeeded", "failed", "interrupted"] + created_at: float + completed_at: float | None + + +class ToolJobResponse(ToolJobSummaryResponse): + vm_id: str + source: Literal["create", "user"] + updated_at: float + start_if_stopped: bool + restart: bool + request_id: str | None + error: str | None + outcome_unknown: bool + items: list[ToolJobItemResponse] + + +class VmToolsResponse(BaseModel): + vm_id: str + inventory: list[ToolInventoryResponse] + checked_at: float | None + agent_ready: bool + guest_busy: bool + supported: bool | None + error: str | None + stale: bool + latest_job: ToolJobResponse | None + jobs: list[ToolJobSummaryResponse] + + class DiagnosticResponse(BaseModel): name: str ok: bool diff --git a/uvm/routers/tools.py b/uvm/routers/tools.py new file mode 100644 index 0000000..a463bc3 --- /dev/null +++ b/uvm/routers/tools.py @@ -0,0 +1,46 @@ +"""Authenticated tool catalog, installation jobs, and guest inventory.""" + +from __future__ import annotations + +from fastapi import APIRouter, Depends, status + +from ..api_models import ToolCatalogResponse, ToolInstallRequest, ToolJobResponse, VmToolsResponse +from ..app import Application +from ..errors import UvmError +from ..tool_catalog import CATALOG, normalize_tools +from .dependencies import get_authorized_application + + +router = APIRouter(tags=["tools"]) + + +def tool_service(application): + service = getattr(application, "tools", None) + if service is None: + raise UvmError("guest tool installation is unavailable in this application") + return service + + +@router.get("/tools", response_model=list[ToolCatalogResponse]) +def catalog(application: Application = Depends(get_authorized_application)): + return list(CATALOG) + + +@router.get("/vms/{vm_id}/tools", response_model=VmToolsResponse) +def vm_tools(vm_id: str, refresh: bool = False, application: Application = Depends(get_authorized_application)): + """Return cached observations; refresh=true probes the running guest without installing.""" + return tool_service(application).snapshot(vm_id, refresh=refresh) + + +@router.post("/vms/{vm_id}/tools/install", response_model=ToolJobResponse, status_code=status.HTTP_202_ACCEPTED) +def install_tools(vm_id: str, request: ToolInstallRequest, + application: Application = Depends(get_authorized_application)): + return tool_service(application).submit( + vm_id, normalize_tools([item.model_dump() for item in request.tools]), + start_if_stopped=request.start_if_stopped, restart=request.restart, request_id=request.request_id, + ) + + +@router.get("/vms/{vm_id}/tools/jobs/{job_id}", response_model=ToolJobResponse) +def tool_job(vm_id: str, job_id: str, application: Application = Depends(get_authorized_application)): + return tool_service(application).store.get_job(vm_id, job_id) diff --git a/uvm/routers/vms.py b/uvm/routers/vms.py index 7475ad4..a9b63b3 100644 --- a/uvm/routers/vms.py +++ b/uvm/routers/vms.py @@ -22,6 +22,7 @@ from ..api_models import ( from ..app import Application from ..domain import PortRule, VmRecord, VmSpec from ..errors import ValidationError +from ..tool_catalog import normalize_tools from ..validation import ( parse_cpu, parse_name, @@ -33,6 +34,7 @@ from ..validation import ( parse_username, ) from .dependencies import get_authorized_application +from .tools import tool_service router = APIRouter(prefix="/vms", tags=["vms"]) @@ -73,6 +75,9 @@ def create_vm( ports = parse_ports([f"{rule.port}:{rule.protocol}" for rule in request.ports or ()]) if not request.no_default_ports and not any(rule.port == 22 for rule in ports): ports = tuple(sorted((*ports, PortRule(22)), key=lambda rule: rule.port)) + tools = normalize_tools([item.model_dump() for item in request.tools], allow_empty=True) + if tools: + tool_service(application).start() vm = application.lifecycle.create( VmSpec( cpu=parse_cpu(str(request.cpu)), @@ -83,8 +88,11 @@ def create_vm( password=parse_password(request.password.get_secret_value()), name=parse_name(request.name) if request.name is not None else None, ports=ports, + tools=tools, ) ) + if tools: + tool_service(application).notify() return vm_response(vm, policy_status=_mutation_policy_status(vm)) diff --git a/uvm/server.py b/uvm/server.py index b87c218..d807498 100644 --- a/uvm/server.py +++ b/uvm/server.py @@ -3,6 +3,7 @@ from __future__ import annotations from collections.abc import Callable +from contextlib import asynccontextmanager from ipaddress import ip_address from pathlib import Path import ssl @@ -32,6 +33,7 @@ def create_api(application: Application | None = None, *, host: str | None = Non try: from fastapi import FastAPI from fastapi.responses import JSONResponse + from starlette.concurrency import run_in_threadpool except ModuleNotFoundError as error: raise UvmError( "FastAPI server support is not installed. Install this project with its dependencies." @@ -42,11 +44,24 @@ def create_api(application: Application | None = None, *, host: str | None = Non from .routers.installation import router as installation_router from .routers.vms import router as vms_router from .routers.diagnostics import router as diagnostics_router + from .routers.tools import router as tools_router + + @asynccontextmanager + async def lifespan(_api): + tools = getattr(resolved_application, "tools", None) + if tools is not None: + await run_in_threadpool(tools.start) + try: + yield + finally: + if tools is not None: + await run_in_threadpool(tools.shutdown) api = FastAPI( title="uvm", version="0.1.0", description="Local Firecracker microVM management API.", + lifespan=lifespan, ) api.state.uvm_application = resolved_application @@ -65,6 +80,7 @@ def create_api(application: Application | None = None, *, host: str | None = Non api.include_router(installation_router) api.include_router(vms_router) api.include_router(diagnostics_router) + api.include_router(tools_router) return api