Coverage for core / src / sensorkit / webapi / fastapi.py: 85%
378 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-02 00:03 +0000
1# SPDX-License-Identifier: Apache-2.0
2import asyncio
3import contextlib
4from collections.abc import AsyncGenerator
5from typing import Any, Callable, Iterable
7import uuid_utils.compat as uuid
8import uvicorn
9from fastapi import Depends, FastAPI, HTTPException, Request, Response
10from fastapi.middleware.cors import CORSMiddleware
11from fastapi.sse import EventSourceResponse, ServerSentEvent
12from loguru import logger
13from pydantic import BaseModel, Field, model_validator
15import sensorkit.api as sk
16from sensorkit.auto.agent import AgentConfigureRequest, AgentState, agent_configure_request
17from sensorkit.backend.base import (
18 BackendError,
19 KeyNotFound,
20 RemoteRequestError,
21 UnregisteredResponder,
22)
23from sensorkit.core.controller import ControllerState
24from sensorkit.core.device import DeviceState
25from sensorkit.core.entity import DeviceDetails, EntityInfo
26from sensorkit.core.program import ProgramState
27from sensorkit.webapi.forwarder import (
28 Forwarder,
29 KeyValueForwarder,
30 ProductForwarder,
31 RecordQueueSet,
32 SKRecord,
33 StreamForwarder,
34)
35from sensorkit.webapi.preview import PreviewJPEG
36from sensorkit.webapi.schema import add_sensorkit_schema
37from sensorkit.webapi.security import (
38 LOOPBACK_HOSTS,
39 AuthConfig,
40 AuthMiddleware,
41 CORSConfig,
42 NoAuthConfig,
43 SecurityHeadersMiddleware,
44 TLSConfig,
45)
46from sensorkit.webapi.serve import ProductInfo, ServeDataConfig, ServeHandler
48PRODUCT_LISTING_TIMEOUT = 10.0
51class WebAPIConfig(BaseModel):
52 """Configuration for the WebAPI service."""
54 port: int = 8000
55 host: str = "127.0.0.1"
56 agent: sk.EntityRef = sk.EntityRef("agent")
57 serve_data_products: list[ServeDataConfig] = []
58 tls: TLSConfig | None = None
59 auth: AuthConfig = NoAuthConfig()
60 cors: CORSConfig = CORSConfig()
61 expose_docs: bool = False
62 max_stream_clients: int = 32
63 stream_queue_size: int = 4096
66class AgentOverrideRequest(BaseModel):
67 """Request body for controller demand override."""
69 state: bool | None
72class EntityListing(EntityInfo):
73 """Entity listing response."""
75 name: str
76 online: bool
79class DeviceListing(EntityInfo):
80 """Device listing response."""
82 name: str
83 online: bool
84 archetype: str | None = None
85 traits: list[str] = Field(default_factory=list)
87 @model_validator(mode="after")
88 def _populate_from_details(self):
89 if isinstance(self.details, DeviceDetails):
90 self.archetype = self.details.archetype.name if self.details.archetype else None
91 self.traits = [trait.name for trait in self.details.traits]
93 return self
96def _status_ok(**kwargs) -> dict[str, Any]:
97 return {"status": "ok", **kwargs}
100@contextlib.contextmanager
101def _error_handler():
102 try:
103 yield
104 except KeyNotFound as err:
105 raise HTTPException(status_code=404, detail="Requested entity not found") from err
106 except RemoteRequestError as err:
107 raise HTTPException(status_code=500, detail="Remote exception occurred") from err
108 except UnregisteredResponder as err:
109 raise HTTPException(status_code=503, detail="Endpoint unavailable") from err
110 except BackendError as err:
111 raise HTTPException(status_code=503, detail="Backend not available") from err
112 except sk.CallError as err:
113 raise HTTPException(status_code=409, detail=str(err)) from err
116def _snapshot(forwarders: Iterable[Forwarder], entity_id: str | None = None) -> list[SKRecord]:
117 """Return the cached records of the given forwarders, optionally for a single entity."""
118 return [record for forwarder in forwarders for record in forwarder.snapshot(entity_id)]
121class WebAPI:
122 """
123 FastAPI-based web service for SensorKit.
125 Provides REST API endpoints for interacting with devices, controllers, programs,
126 and agents, as well as Server-Sent Events (SSE) for real-time state updates.
127 """
129 def __init__(self, kit: sk.SensorKit, config: WebAPIConfig):
130 self.kit = kit
131 self.config = config
132 self.client_queues: RecordQueueSet = set()
133 self.kv_forwarder = KeyValueForwarder(kit, targets=self.client_queues)
134 self.stream_forwarder = StreamForwarder(kit, targets=self.client_queues)
135 self._serve_handlers: list[ServeHandler] = []
136 self._product_forwarders: list[ProductForwarder] = []
137 self.authenticator = config.auth.create_authenticator()
138 self._stopped = asyncio.Event()
139 self._stopped.set()
141 ssl_options = config.tls.uvicorn_options() if config.tls else {}
143 self.app = self._create_fastapi_app()
144 self.server = uvicorn.Server(
145 uvicorn.Config(
146 self.app,
147 host=self.config.host,
148 port=self.config.port,
149 log_level="info",
150 **ssl_options,
151 )
152 )
154 # Resolve the configured agent entity reference.
155 self.config.agent.resolve(kit)
157 @property
158 def forwarders(self) -> tuple[Forwarder, ...]:
159 """Every active forwarder: key-value, stream, and one per data-product handler."""
160 return self.kv_forwarder, self.stream_forwarder, *self._product_forwarders
162 async def _check_stream_capacity(self):
163 """Reject a new subscriber once the concurrent stream limit is reached."""
164 if len(self.client_queues) >= self.config.max_stream_clients:
165 raise HTTPException(status_code=503, detail="Too many concurrent stream subscribers")
167 async def _stream_records(
168 self,
169 request: Request,
170 sources: Iterable[Forwarder],
171 match: Callable[[SKRecord], bool] = lambda _: True,
172 ) -> AsyncGenerator[ServerSentEvent]:
173 """Yield matching records as SSE: the snapshot of `sources`, then live updates.
175 Every forwarder feeds the same client queue, so live updates can only be selected
176 by predicate; `sources` narrows the snapshot alone and `match` must cover both.
177 """
178 request.state.is_sse = True
179 queue: asyncio.Queue[SKRecord | None] = asyncio.Queue(self.config.stream_queue_size)
180 self.client_queues.add(queue)
182 try:
183 for record in filter(match, _snapshot(sources)):
184 yield ServerSentEvent(raw_data=record.serialize())
186 while True:
187 record = await queue.get()
189 if record is None:
190 break
192 try:
193 if match(record):
194 yield ServerSentEvent(raw_data=record.serialize())
195 finally:
196 queue.task_done()
197 except asyncio.CancelledError:
198 host = request.client.host if request.client else "unknown"
199 logger.debug(f"SSE client disconnected: {host}")
200 finally:
201 self.client_queues.discard(queue)
203 async def _product_listing(self) -> list[ProductInfo]:
204 """Return every handler's product listing.
206 A handler's listing may not be ready immediately. Rather than report an empty
207 listing — which would falsely imply there are no products — we wait up to
208 PRODUCT_LISTING_TIMEOUT and raise 503 if it is still not available.
209 """
210 try:
211 async with asyncio.timeout(PRODUCT_LISTING_TIMEOUT):
212 listings = [await handler.get_listing() for handler in self._serve_handlers]
213 except TimeoutError as err:
214 raise HTTPException(
215 status_code=503,
216 detail="Data product listing not yet available; try again shortly",
217 ) from err
219 return [info for listing in listings for info in listing]
221 def _product_handler(self, controller_id: str, product_id: str) -> ServeHandler:
222 """Return the handler serving a product, or raise 404 if none does."""
223 for handler in self._serve_handlers:
224 if handler.has_product(controller_id, product_id):
225 return handler
227 raise HTTPException(status_code=404, detail="Product not found")
229 async def _product_data(self, controller_id: str, product_id: str) -> bytes:
230 """Return a product's raw bytes, or raise 404 if it is not being served.
232 A product can leave the handler's listing between the lookup and the read, so
233 the miss is reported the same way an unknown product is.
234 """
235 handler = self._product_handler(controller_id, product_id)
237 try:
238 return await handler.get_data(controller_id, product_id)
239 except KeyError as err:
240 raise HTTPException(status_code=404, detail="Product not found") from err
242 def _product_metadata(self, controller_id: str, product_id: str) -> dict:
243 """Return a product's cached metadata, or raise 404 if it is not being served."""
244 handler = self._product_handler(controller_id, product_id)
246 try:
247 return handler.get_metadata(controller_id, product_id)
248 except KeyError as err:
249 raise HTTPException(status_code=404, detail="Product not found") from err
251 def _create_fastapi_app(self):
252 hidden = {"docs_url": None, "redoc_url": None}
253 app = FastAPI(title="SensorKit Web API", **({} if self.config.expose_docs else hidden))
255 # CORS goes last so it answers preflight requests itself, which carry no credentials to
256 # authorize.
257 app.add_middleware(AuthMiddleware, authenticator=self.authenticator)
258 app.add_middleware(SecurityHeadersMiddleware, hsts=self.config.tls is not None)
259 app.add_middleware(CORSMiddleware, **self.config.cors.middleware_options())
261 self._create_global_endpoints(app)
262 self._create_device_endpoints(app)
263 self._create_controller_endpoints(app)
264 self._create_program_endpoints(app)
265 self._create_agent_endpoints(app)
266 self._finalize_openapi(app)
268 return app
270 def _finalize_openapi(self, app: FastAPI):
271 """Generate the OpenAPI document and fold in the SensorKit and security schemas."""
272 app.openapi()
273 schema = app.openapi_schema
274 add_sensorkit_schema(schema["components"]["schemas"])
276 if scheme := self.authenticator.openapi_scheme():
277 schema["components"]["securitySchemes"] = {"bearerAuth": scheme}
278 schema["security"] = [{"bearerAuth": []}]
280 def _create_global_endpoints(self, app: FastAPI):
281 @app.get(
282 "/data/subscribe",
283 tags=["Global"],
284 response_class=EventSourceResponse,
285 dependencies=[Depends(self._check_stream_capacity)],
286 )
287 async def firehose_subscription(request: Request):
288 """SSE endpoint for real-time updates."""
289 async with contextlib.aclosing(self._stream_records(request, self.forwarders)) as recs:
290 async for event in recs:
291 yield event
293 @app.get("/data/snapshot", tags=["Global"])
294 async def get_full_system_snapshot() -> list[SKRecord]:
295 """Return the all cached data for all entities."""
296 return _snapshot(self.forwarders)
298 @app.get("/data/snapshot/{entity_id}", tags=["Global"])
299 async def get_entity_snapshot(entity_id: str) -> list[SKRecord]:
300 """Return all cached data for a specific entity."""
301 return _snapshot(self.forwarders, entity_id)
303 @app.get("/entities", tags=["Global"])
304 async def get_entity_list() -> list[EntityListing | DeviceListing]:
305 """Return a list of all entity IDs."""
306 output = []
307 cache = self.kv_forwarder.cache
308 entities = tuple(cache.keys())
310 for entity in entities:
311 if record := cache[entity].get("EntityInfo"):
312 if record.payload is None:
313 continue
315 info = record.payload.copy()
316 info["name"] = entity
317 info["online"] = "EntityLease" in cache[entity]
319 output.append(
320 DeviceListing.model_validate(info)
321 if info["entity_type"] == "device"
322 else EntityListing.model_validate(info)
323 )
325 return output
327 def _create_device_endpoints(self, app: FastAPI):
328 @app.get("/device/{device_id}/state", tags=["Device"])
329 async def get_device_state(device_id: str) -> DeviceState:
330 """Return the current state of a device."""
331 with _error_handler():
332 return await self.kit.device(device_id).kv_get_model(DeviceState)
334 @app.post("/device/{device_id}/command", tags=["Device"])
335 async def run_device_command(device_id: str, command: sk.DeviceCommand):
336 """Execute a command on a device."""
337 with _error_handler():
338 await self.kit.device(device_id).command(command)
340 return _status_ok()
342 def _create_controller_endpoints(self, app: FastAPI):
343 @app.get("/controller/{controller_id}/state", tags=["Controller"])
344 async def get_controller_state(controller_id: str) -> ControllerState:
345 """Return the current state of a controller."""
346 with _error_handler():
347 return await self.kit.controller(controller_id).kv_get_model(ControllerState)
349 @app.post("/controller/{controller_id}/execute", tags=["Controller"])
350 async def run_controller_task(controller_id: str, body: sk.TaskSubmission | sk.Task):
351 """Execute a task on a controller.
353 The body may be a bare task or a `TaskSubmission` envelope. The envelope carries the
354 client-supplied execution parameters (`context`, `expiry_time`) that the controller
355 records on the minted `TaskExecution` — for example the file-naming keywords a caller
356 needs threaded onto the execution. A bare task is equivalent to an envelope with no
357 parameters.
358 """
359 submission = body if isinstance(body, sk.TaskSubmission) else sk.TaskSubmission(task=body)
361 with _error_handler():
362 # The controller mints the task_id; learn it from the returned execution.
363 execution = await self.kit.controller(controller_id).start_task(
364 submission.task,
365 context=submission.context,
366 expiry_time=submission.expiry_time,
367 )
368 await execution
370 return _status_ok(task_id=execution.task_id)
372 @app.post("/controller/{controller_id}/wait", tags=["Controller"])
373 async def wait_for_current_task(controller_id: str):
374 """Wait for the currently executing task to complete on a controller."""
375 with _error_handler():
376 await self.kit.controller(controller_id).wait_for_task()
378 return _status_ok()
380 @app.post("/controller/{controller_id}/wait/{task_id}", tags=["Controller"])
381 async def wait_for_task(controller_id: str, task_id: uuid.UUID):
382 """Wait for a task to complete on a controller."""
383 with _error_handler():
384 await self.kit.controller(controller_id).wait_for_task(task_id)
386 return _status_ok()
388 @app.post("/controller/{controller_id}/abort", tags=["Controller"])
389 async def abort_controller_task(controller_id: str, task_id: uuid.UUID | None = None):
390 """Abort a specific task or the current task on a controller."""
391 with _error_handler():
392 await self.kit.controller(controller_id).abort_task(task_id)
394 return _status_ok()
396 @app.get("/controller/{controller_id}/products", tags=["Controller"])
397 async def list_controller_products(controller_id: str) -> list[ProductInfo]:
398 """List data products associated with a controller."""
399 return [
400 info
401 for info in await self._product_listing()
402 if info.controller_id == controller_id
403 ]
405 @app.get(
406 "/controller/{controller_id}/products/subscribe",
407 tags=["Controller"],
408 response_class=EventSourceResponse,
409 dependencies=[Depends(self._check_stream_capacity)],
410 )
411 async def subscribe_controller_products(request: Request, controller_id: str):
412 """SSE stream of product arrivals for a controller, including initial listing."""
414 def is_controller_product(record: SKRecord) -> bool:
415 return record.kind == "product" and str(record.subject.entity()) == controller_id
417 stream = self._stream_records(request, self._product_forwarders, is_controller_product)
419 async with contextlib.aclosing(stream) as recs:
420 async for event in recs:
421 yield event
423 @app.get("/controller/{controller_id}/product/{product_id}/data", tags=["Controller"])
424 async def get_controller_product_data(controller_id: str, product_id: str):
425 """Return the raw FITS data for a product."""
426 raw_bytes = await self._product_data(controller_id, product_id)
428 return Response(content=raw_bytes, media_type="application/fits")
430 @app.get("/controller/{controller_id}/product/{product_id}/preview", tags=["Controller"])
431 async def get_controller_product_preview(controller_id: str, product_id: str):
432 """Return a JPEG preview image generated from the FITS product data."""
433 data = await self._product_data(controller_id, product_id)
435 try:
436 preview = await PreviewJPEG.from_fits(data)
437 except Exception as err:
438 logger.warning(f"Could not generate preview for {product_id}: {err}")
439 raise HTTPException(
440 status_code=422, detail="Could not generate preview"
441 ) from err
443 return Response(content=preview.jpeg_bytes, media_type="image/jpeg")
445 @app.get("/controller/{controller_id}/product/{product_id}/metadata", tags=["Controller"])
446 async def get_controller_product_metadata(controller_id: str, product_id: str) -> dict:
447 """Return the cached metadata for a product as a JSON object."""
448 metadata = self._product_metadata(controller_id, product_id)
450 if not metadata.get(ProductInfo):
451 logger.warning(f"Product {product_id} has no ProductInfo keyword set")
453 return metadata
455 def _create_program_endpoints(self, app: FastAPI):
456 @app.get("/program/{program_id}/state", tags=["Program"])
457 async def get_program_state(program_id: str) -> ProgramState:
458 """Return the current state of a program."""
459 with _error_handler():
460 return await self.kit.program(program_id).kv_get_model(ProgramState)
462 @app.post("/program/{program_id}/enable", tags=["Program"])
463 async def enable_program(program_id: str, controller_id: str | None = None):
464 """Enable a program."""
465 if controller_id is None:
466 # TODO: When the program enable() method supports a None controller, i.e. leave
467 # the associated controller unchanged, this part can be removed.
468 with _error_handler():
469 state = await self.kit.program(program_id).kv_get_model(ProgramState)
470 controller_id = state.enable_state.controller
472 if controller_id is None:
473 raise HTTPException(status_code=422, detail="Controller ID is required")
475 with _error_handler():
476 await self.kit.program(program_id).enable(controller_id)
478 return _status_ok()
480 @app.post("/program/{program_id}/disable", tags=["Program"])
481 async def disable_program(program_id: str):
482 """Disable a program."""
483 with _error_handler():
484 await self.kit.program(program_id).disable()
486 return _status_ok()
488 @app.post("/program/{program_id}/activate", tags=["Program"])
489 async def activate_program(program_id: str):
490 """Start tasking for a program."""
491 with _error_handler():
492 await self.kit.program(program_id).start_tasking()
494 return _status_ok()
496 @app.post("/program/{program_id}/deactivate", tags=["Program"])
497 async def deactivate_program(program_id: str):
498 """Stop tasking for a program."""
499 with _error_handler():
500 await self.kit.program(program_id).stop_tasking()
502 return _status_ok()
504 @app.post("/program/{program_id}/abort", tags=["Program"])
505 async def abort_program(program_id: str):
506 """Abort program tasking immediately."""
507 with _error_handler():
508 await self.kit.program(program_id).abort_tasking()
510 return _status_ok()
512 def _create_agent_endpoints(self, app: FastAPI):
513 @app.get("/agent/state", tags=["Agent"])
514 async def get_agent_state() -> AgentState:
515 """Get the current status of the agent."""
516 with _error_handler():
517 return await self.config.agent.require().kv_get_model(AgentState)
519 @app.post("/agent/enable", tags=["Agent"])
520 async def enable_agent():
521 """Enable all agent control."""
522 with _error_handler():
523 await self.config.agent.require().call(
524 agent_configure_request, AgentConfigureRequest(global_control_enabled=True)
525 )
527 return _status_ok()
529 @app.post("/agent/enable/{controller_id}", tags=["Agent"])
530 async def enable_controller(controller_id: str):
531 """Enable agent control for a specific controller."""
532 with _error_handler():
533 await self.config.agent.require().call(
534 agent_configure_request,
535 AgentConfigureRequest(controller_control_enabled={controller_id: True}),
536 )
538 return _status_ok()
540 @app.post("/agent/disable", tags=["Agent"])
541 async def disable_agent():
542 """Disable all agent control."""
543 with _error_handler():
544 await self.config.agent.require().call(
545 agent_configure_request, AgentConfigureRequest(global_control_enabled=False)
546 )
548 return _status_ok()
550 @app.post("/agent/disable/{controller_id}", tags=["Agent"])
551 async def disable_controller(controller_id: str):
552 """Disable agent control for a specific controller."""
553 with _error_handler():
554 await self.config.agent.require().call(
555 agent_configure_request,
556 AgentConfigureRequest(controller_control_enabled={controller_id: False}),
557 )
559 return _status_ok()
561 @app.post("/agent/override/{controller_id}", tags=["Agent"])
562 async def override_controller(controller_id: str, req: AgentOverrideRequest):
563 """Override controller demand."""
564 with _error_handler():
565 await self.config.agent.require().call(
566 agent_configure_request,
567 AgentConfigureRequest(controller_demand_override={controller_id: req.state}),
568 )
570 return _status_ok()
572 @app.post("/agent/scheduler/enable", tags=["Agent"])
573 async def enable_scheduler():
574 """Enable agent scheduling."""
575 with _error_handler():
576 await self.config.agent.require().call(
577 agent_configure_request,
578 AgentConfigureRequest(scheduling_enabled=True),
579 )
581 return _status_ok()
583 @app.post("/agent/scheduler/disable", tags=["Agent"])
584 async def disable_scheduler():
585 """Disable agent scheduling."""
586 with _error_handler():
587 await self.config.agent.require().call(
588 agent_configure_request,
589 AgentConfigureRequest(scheduling_enabled=False),
590 )
592 return _status_ok()
594 @app.post("/agent/scheduler/include/{program_id}", tags=["Agent"])
595 async def include_program(program_id: str):
596 """Include a program for agent consideration."""
597 with _error_handler():
598 await self.config.agent.require().call(
599 agent_configure_request,
600 AgentConfigureRequest(remove_program_exclusions={program_id}),
601 )
603 return _status_ok()
605 @app.post("/agent/scheduler/exclude/{program_id}", tags=["Agent"])
606 async def exclude_program(program_id: str):
607 """Exclude a program from agent consideration."""
608 with _error_handler():
609 await self.config.agent.require().call(
610 agent_configure_request,
611 AgentConfigureRequest(add_program_exclusions={program_id}),
612 )
614 return _status_ok()
616 async def _start_forwarders(self, *, task_group: asyncio.TaskGroup):
617 """Start the KV/stream forwarders and any configured data-product handlers.
619 Split out from `serve` only so tests can drive it without running the uvicorn
620 server. Not for use elsewhere: calling this and `serve` both would start two sets
621 of forwarder tasks.
622 """
623 await self.kv_forwarder.start(task_group=task_group)
624 await self.stream_forwarder.start(task_group=task_group)
626 self._serve_handlers = []
627 self._product_forwarders = []
629 for serve_config in self.config.serve_data_products:
630 handler = serve_config.create_handler()
631 forwarder = ProductForwarder(handler, targets=self.client_queues)
633 await forwarder.start(task_group=task_group)
634 handler.start_monitor(task_group=task_group)
636 self._serve_handlers.append(handler)
637 self._product_forwarders.append(forwarder)
639 def _warn_if_exposed(self):
640 """Warn when the service is reachable off-host without TLS or authentication."""
641 if self.config.host in LOOPBACK_HOSTS:
642 return
644 missing = [
645 name
646 for name, present in (
647 ("TLS", self.config.tls is not None),
648 ("authentication", self.authenticator.enabled),
649 )
650 if not present
651 ]
653 if missing:
654 logger.warning(
655 f"web API is bound to {self.config.host}:{self.config.port} with no "
656 f"{' and no '.join(missing)}; anyone who can reach it can control this system"
657 )
659 async def serve(self, *, task_group: asyncio.TaskGroup):
660 """Run the FastAPI server."""
661 # Already shut down or shutting down.
662 if self.server.should_exit:
663 return
665 self._stopped.clear()
667 try:
668 await self._start_forwarders(task_group=task_group)
670 if self.config.tls is not None:
671 # Loading the config is what builds the SSL context, and also what surfaces
672 # an unreadable or mismatched certificate and key.
673 self.server.config.load()
675 if (context := self.server.config.ssl) is not None:
676 self.config.tls.apply_minimum_version(context)
678 self._warn_if_exposed()
680 await self.server.serve()
681 finally:
682 self._stopped.set()
684 async def shutdown(self):
685 # Unblock any parked firehose generators so uvicorn's graceful shutdown
686 # isn't stuck waiting on a never-ending SSE response. Snapshot the set
687 # because each generator removes its own queue on exit.
688 for queue in tuple(self.client_queues):
689 if queue.full():
690 # Give up a record to make room for the sentinel.
691 with contextlib.suppress(asyncio.QueueEmpty):
692 queue.get_nowait()
694 queue.put_nowait(None)
696 # Signal a graceful shutdown to Uvicorn. Note that explicitly shutting down the server here
697 # does not cause `serve` to return!
698 self.server.should_exit = True
699 await self._stopped.wait()
701 for handler in self._serve_handlers:
702 await handler.stop_monitor()
704 for forwarder in self._product_forwarders:
705 await forwarder.stop()
707 await self.kv_forwarder.stop()
708 await self.stream_forwarder.stop()