fleet-framework 0.1.0__py3-none-any.whl

fleet/core/store.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+from fleet.core.backend import Backend
+
+
+@dataclass
+class WorkerRec:
+    worker_id: str
+    automation_type: str
+    hardware: dict[str, Any]
+    config: dict[str, Any]
+    config_gen: int
+    state: str
+    last_seen: float
+    last_error: str | None
+    stats: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "worker_id": self.worker_id,
+            "automation_type": self.automation_type,
+            "hardware": self.hardware,
+            "config": self.config,
+            "config_gen": self.config_gen,
+            "state": self.state,
+            "last_seen": self.last_seen,
+            "last_error": self.last_error,
+            "current_stats": self.stats,
+        }
+
+
+class Store:
+    """Master-side facade over Backend that produces WorkerRec objects."""
+
+    def __init__(self, backend: Backend) -> None:
+        self._backend = backend
+
+    async def upsert_register(
+        self, automation_type: str, worker_id: str, hardware: dict[str, Any]
+    ) -> None:
+        await self._backend.worker_register(automation_type, worker_id, hardware)
+
+    async def get_config(self, automation_type: str, worker_id: str) -> tuple[dict[str, Any], int]:
+        return await self._backend.worker_get_config(automation_type, worker_id)
+
+    async def patch_config(
+        self, automation_type: str, worker_id: str, partial: dict[str, Any]
+    ) -> tuple[dict[str, Any], int]:
+        current, _ = await self._backend.worker_get_config(automation_type, worker_id)
+        merged = {**current, **partial}
+        return await self._backend.worker_set_config(automation_type, worker_id, merged)
+
+    async def replace_config(
+        self, automation_type: str, worker_id: str, full: dict[str, Any]
+    ) -> tuple[dict[str, Any], int]:
+        return await self._backend.worker_set_config(automation_type, worker_id, full)
+
+    async def report_state(
+        self,
+        automation_type: str,
+        worker_id: str,
+        state: str,
+        last_error: str | None,
+    ) -> None:
+        await self._backend.worker_set_state(
+            automation_type, worker_id, state, last_error, last_seen=time.time()
+        )
+
+    async def report_stats(
+        self, automation_type: str, worker_id: str, stats: dict[str, Any]
+    ) -> None:
+        await self._backend.worker_set_stats(automation_type, worker_id, stats)
+
+    async def get_worker(self, automation_type: str, worker_id: str) -> WorkerRec | None:
+        raw = await self._backend.worker_get(automation_type, worker_id)
+        if raw is None:
+            return None
+        return WorkerRec(
+            worker_id=worker_id,
+            automation_type=automation_type,
+            hardware=raw["hardware"],
+            config=raw["config"],
+            config_gen=raw["config_gen"],
+            state=raw["state"],
+            last_seen=raw["last_seen"],
+            last_error=raw["last_error"],
+            stats=raw["stats"],
+        )
+
+    async def list_workers(self, automation_type: str | None = None) -> list[WorkerRec]:
+        pairs = await self._backend.worker_list(automation_type)
+        out: list[WorkerRec] = []
+        for at, wid in pairs:
+            rec = await self.get_worker(at, wid)
+            if rec is not None:
+                out.append(rec)
+        return out
+
+    async def remove_worker(self, automation_type: str, worker_id: str) -> bool:
+        return await self._backend.worker_remove(automation_type, worker_id)

fleet/master/__init__.py

@@ -0,0 +1,3 @@
+from fleet.master.app import create_app
+
+__all__ = ["create_app"]

fleet/master/api.py

@@ -0,0 +1,324 @@
+from __future__ import annotations
+
+import json
+import logging
+import uuid
+from typing import Any
+
+from fastapi import APIRouter, Depends, HTTPException
+from pydantic import BaseModel, ValidationError
+
+from fleet.core.automation import BaseAutomation, catalog_doc, get_registry
+from fleet.core.backend import Backend
+from fleet.core.otel import span
+from fleet.core.protocol import ConfigChanged, Drain
+from fleet.core.store import Store
+from fleet.master.auth import require_admin, require_scope, require_worker
+from fleet.master.broadcaster import Broadcaster
+from fleet.master.ratelimit import RateLimiter, make_token_rate_limit
+
+logger = logging.getLogger(__name__)
+
+
+class PopRequest(BaseModel):
+    n: int = 1
+
+
+class SubmitRequest(BaseModel):
+    payload: dict[str, Any]
+    task_id: str | None = None
+    max_attempts: int | None = None
+    priority: int = 0
+    """Higher = more urgent. >0 routes to the high-priority bucket;
+    <0 to the low-priority bucket; 0 (default) to the normal bucket."""
+    ttl_seconds: int | None = None
+    """If set, the task is silently dropped on reserve when older than
+    this many seconds. Useful for "act on now or never" events."""
+    idempotency_key: str | None = None
+    """If set, a second submission with the same (automation_type, key) within
+    the idempotency window returns the original task_id without enqueueing."""
+
+
+def _decode_envelopes(raw_list: list[bytes]) -> list[dict[str, Any]]:
+    out: list[dict[str, Any]] = []
+    for raw in raw_list:
+        try:
+            out.append(json.loads(raw.decode("utf-8") if isinstance(raw, bytes) else raw))
+        except (json.JSONDecodeError, UnicodeDecodeError):
+            continue
+    return out
+
+
+def make_api_router(
+    store: Store, broadcaster: Broadcaster, backend: Backend
+) -> APIRouter:
+    router = APIRouter()
+    registry: dict[str, type[BaseAutomation]] = get_registry()
+    import os as _os
+    submit_rps = float(_os.environ.get("FLEET_SUBMIT_RPS_PER_TOKEN", "0"))
+    submit_limiter = RateLimiter(submit_rps)
+    submit_rate_dep = make_token_rate_limit(submit_limiter)
+    idempotency_ttl = int(_os.environ.get("FLEET_IDEMPOTENCY_TTL_SECONDS", "86400"))
+
+    @router.get("/api/v1/automations", dependencies=[Depends(require_admin)])
+    async def list_automations() -> list[dict[str, Any]]:
+        return [
+            {
+                "type": name,
+                "class": cls.__name__,
+                "config_schema": cls.Config.model_json_schema(),
+                "kind": cls.__mro__[1].__name__,
+            }
+            for name, cls in registry.items()
+        ]
+
+    @router.get(
+        "/api/v1/automations/{automation_type}/workers",
+        dependencies=[Depends(require_scope)],
+    )
+    async def list_workers(automation_type: str) -> list[dict[str, Any]]:
+        _require_type(automation_type, registry)
+        recs = await store.list_workers(automation_type)
+        return [r.to_dict() for r in recs]
+
+    @router.get(
+        "/api/v1/automations/{automation_type}/workers/{worker_id}",
+        dependencies=[Depends(require_scope)],
+    )
+    async def get_worker(automation_type: str, worker_id: str) -> dict[str, Any]:
+        _require_type(automation_type, registry)
+        rec = await store.get_worker(automation_type, worker_id)
+        if rec is None:
+            raise HTTPException(404, "worker not found")
+        return rec.to_dict()
+
+    @router.patch(
+        "/api/v1/automations/{automation_type}/workers/{worker_id}/config",
+        dependencies=[Depends(require_scope)],
+    )
+    async def patch_config(
+        automation_type: str, worker_id: str, partial: dict[str, Any]
+    ) -> dict[str, Any]:
+        cls = _require_type(automation_type, registry)
+        current, _ = await store.get_config(automation_type, worker_id)
+        merged = {**current, **partial}
+        try:
+            cls.Config.model_validate(merged)
+        except ValidationError as e:
+            raise HTTPException(422, f"invalid config: {e}") from e
+        cfg, gen = await store.replace_config(automation_type, worker_id, merged)
+        await broadcaster.publish(
+            automation_type,
+            worker_id,
+            ConfigChanged(config_gen=gen).model_dump(),
+        )
+        return {"config": cfg, "config_gen": gen}
+
+    @router.get(
+        "/api/v1/automations/{automation_type}/output/length",
+        dependencies=[Depends(require_scope)],
+    )
+    async def stream_len(automation_type: str) -> dict[str, int]:
+        _require_type(automation_type, registry)
+        return {"length": await backend.stream_length(automation_type)}
+
+    @router.get(
+        "/api/v1/automations/{automation_type}/queue",
+        dependencies=[Depends(require_scope)],
+    )
+    async def queue_inspect(automation_type: str, n: int = 20) -> dict[str, Any]:
+        _require_type(automation_type, registry)
+        items = _decode_envelopes(await backend.queue_peek(automation_type, n))
+        length = await backend.queue_size(automation_type)
+        return {"length": length, "items": items}
+
+    @router.get(
+        "/api/v1/pools/{pool_name}",
+        dependencies=[Depends(require_admin)],
+    )
+    async def pool_inspect(pool_name: str, n: int = 20) -> dict[str, Any]:
+        raws = await backend.pool_list(pool_name)
+        size = await backend.pool_size(pool_name)
+        items: list[dict[str, Any]] = []
+        for raw in raws[:n]:
+            try:
+                payload = json.loads(raw.payload.decode("utf-8"))
+            except (json.JSONDecodeError, UnicodeDecodeError):
+                payload = None
+            items.append({
+                "id": raw.id,
+                "tags": raw.tags,
+                "payload": payload,
+                "expires_at": raw.expires_at,
+                "in_use_until": raw.in_use_until,
+            })
+        return {"size": size, "items": items}
+
+    @router.get(
+        "/api/v1/automations/{automation_type}/output/peek",
+        dependencies=[Depends(require_scope)],
+    )
+    async def stream_pk(automation_type: str, n: int = 10) -> list[dict[str, Any]]:
+        _require_type(automation_type, registry)
+        return _decode_envelopes(await backend.stream_peek(automation_type, n))
+
+    @router.post(
+        "/api/v1/automations/{automation_type}/output/pop",
+        dependencies=[Depends(require_scope)],
+    )
+    async def stream_pp(automation_type: str, req: PopRequest) -> list[dict[str, Any]]:
+        _require_type(automation_type, registry)
+        return _decode_envelopes(await backend.stream_pop(automation_type, req.n))
+
+    @router.get("/api/v1/catalog", dependencies=[Depends(require_admin)])
+    async def get_catalog() -> dict[str, Any]:
+        return {name: catalog_doc(cls) for name, cls in registry.items()}
+
+    @router.get(
+        "/api/v1/catalog/{automation_type}",
+        dependencies=[Depends(require_scope)],
+    )
+    async def get_catalog_one(automation_type: str) -> dict[str, Any]:
+        cls = _require_type(automation_type, registry)
+        return catalog_doc(cls)
+
+    @router.post(
+        "/api/v1/automations/{automation_type}/tasks",
+        dependencies=[Depends(require_scope), Depends(submit_rate_dep)],
+    )
+    async def submit_task(automation_type: str, req: SubmitRequest) -> dict[str, Any]:
+        cls = _require_type(automation_type, registry)
+        if cls.TaskPayload is None:
+            raise HTTPException(
+                400,
+                f"automation '{automation_type}' did not declare TaskPayload; "
+                "only BatchAutomations with a declared task schema accept submissions",
+            )
+        try:
+            cls.TaskPayload.model_validate(req.payload)
+        except ValidationError as e:
+            raise HTTPException(422, f"invalid task payload: {e}") from e
+        tid = req.task_id or uuid.uuid4().hex
+        if req.idempotency_key:
+            ns = f"idem:{automation_type}"
+            claimed = await backend.kv_setnx(
+                ns, req.idempotency_key, tid.encode("utf-8"), idempotency_ttl,
+            )
+            if not claimed:
+                existing = await backend.kv_get(ns, req.idempotency_key)
+                if existing is not None:
+                    return {
+                        "task_id": existing.decode("utf-8"),
+                        "queued": False,
+                        "duplicate": True,
+                    }
+        import time as _time
+        envelope: dict[str, Any] = {
+            "id": tid,
+            "payload": req.payload,
+            "attempt": 0,
+            "submitted_at": _time.time(),
+        }
+        if req.max_attempts is not None:
+            envelope["max_attempts"] = req.max_attempts
+        if req.ttl_seconds is not None:
+            envelope["ttl_seconds"] = req.ttl_seconds
+        body = json.dumps(envelope).encode("utf-8")
+        with span(
+            "api.submit_task",
+            automation_type=automation_type,
+            task_id=tid,
+            priority=req.priority,
+        ):
+            await backend.queue_push(automation_type, body, priority=req.priority)
+        return {"task_id": tid, "queued": True, "priority": req.priority}
+
+    @router.get(
+        "/api/v1/automations/{automation_type}/dead",
+        dependencies=[Depends(require_scope)],
+    )
+    async def dlq_peek(automation_type: str, n: int = 50) -> dict[str, Any]:
+        _require_type(automation_type, registry)
+        items = _decode_envelopes(await backend.dlq_peek(automation_type, n))
+        length = await backend.dlq_length(automation_type)
+        return {"length": length, "items": items}
+
+    @router.post(
+        "/api/v1/automations/{automation_type}/dead/replay",
+        dependencies=[Depends(require_scope)],
+    )
+    async def dlq_replay(automation_type: str, req: PopRequest) -> dict[str, int]:
+        _require_type(automation_type, registry)
+        popped = await backend.dlq_pop(automation_type, req.n)
+        for body in popped:
+            try:
+                env = json.loads(body.decode("utf-8") if isinstance(body, bytes) else body)
+                env["attempt"] = 0  # reset attempts on replay
+                body = json.dumps(env).encode("utf-8")
+            except Exception:
+                pass
+            await backend.queue_push(automation_type, body)
+        return {"replayed": len(popped)}
+
+    @router.delete(
+        "/api/v1/automations/{automation_type}/dead",
+        dependencies=[Depends(require_scope)],
+    )
+    async def dlq_drain(automation_type: str) -> dict[str, int]:
+        _require_type(automation_type, registry)
+        return {"cleared": await backend.dlq_clear(automation_type)}
+
+    @router.post(
+        "/api/v1/automations/{automation_type}/workers/{worker_id}/drain",
+        dependencies=[Depends(require_scope)],
+    )
+    async def drain_worker(automation_type: str, worker_id: str) -> dict[str, Any]:
+        _require_type(automation_type, registry)
+        rec = await store.get_worker(automation_type, worker_id)
+        if rec is None:
+            raise HTTPException(404, "worker not found")
+        await broadcaster.publish(
+            automation_type, worker_id, Drain().model_dump()
+        )
+        return {"worker_id": worker_id, "drain": "requested"}
+
+    @router.delete(
+        "/api/v1/automations/{automation_type}/workers/{worker_id}",
+        dependencies=[Depends(require_scope)],
+    )
+    async def delete_worker(automation_type: str, worker_id: str) -> dict[str, Any]:
+        _require_type(automation_type, registry)
+        removed = await store.remove_worker(automation_type, worker_id)
+        if not removed:
+            raise HTTPException(404, "worker not found")
+        return {"worker_id": worker_id, "removed": True}
+
+    @router.get(
+        "/api/v1/automations/{automation_type}/workers/{worker_id}/stats",
+        dependencies=[Depends(require_scope)],
+    )
+    async def worker_stats_history(
+        automation_type: str, worker_id: str, since: float = 0, limit: int = 240
+    ) -> list[dict[str, Any]]:
+        _require_type(automation_type, registry)
+        return await backend.worker_stats_history(
+            automation_type, worker_id, since=since, limit=limit,
+        )
+
+    @router.get(
+        "/api/v1/automations/{automation_type}/workers/{worker_id}/config",
+        dependencies=[Depends(require_worker)],
+    )
+    async def worker_poll(automation_type: str, worker_id: str) -> dict[str, Any]:
+        _require_type(automation_type, registry)
+        cfg, gen = await store.get_config(automation_type, worker_id)
+        return {"config": cfg, "config_gen": gen}
+
+    return router
+
+
+def _require_type(automation_type: str, registry: dict[str, type[BaseAutomation]]) -> type[BaseAutomation]:
+    cls = registry.get(automation_type)
+    if cls is None:
+        raise HTTPException(404, f"unknown automation type: {automation_type}")
+    return cls

fleet/master/app.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+
+from fastapi import FastAPI
+
+from fleet.core.automation import catalog_doc, get_registry, load_entry_points
+from fleet.core.backend import backend_from_url
+from fleet.core.proxy import load_provider_entry_points
+from fleet.core.store import Store
+from fleet.master.api import make_api_router
+from fleet.master.broadcaster import Broadcaster
+from fleet.master.dashboard import make_dashboard_router
+from fleet.master.metrics_route import make_metrics_router
+from fleet.master.ws_router import make_ws_router
+
+logger = logging.getLogger(__name__)
+
+
+def create_app(backend_url: str | None = None) -> FastAPI:
+    # Backwards-compatible: callers passing redis://... still work because
+    # backend_from_url handles redis:// natively. New callers can pass
+    # memory:// or any registered scheme.
+    backend_url = (
+        backend_url
+        or os.environ.get("FLEET_BACKEND_URL")
+        or os.environ.get("REDIS_URL", "redis://localhost:6379/0")
+    )
+    registry = load_entry_points()
+    providers = load_provider_entry_points()
+    logger.info("loaded automations: %s", list(registry.keys()))
+    logger.info("loaded proxy providers: %s", list(providers.keys()))
+
+    backend = backend_from_url(backend_url)
+    store = Store(backend)
+    broadcaster = Broadcaster()
+
+    app = FastAPI(title="fleet")
+    app.state.store = store
+    app.state.backend = backend
+    app.state.broadcaster = broadcaster
+
+    # publish the catalog: schemas for every installed automation. dashboards
+    # and operators read this; consumers don't need it for typing because they
+    # import the same contract refs as the producer.
+    @app.on_event("startup")
+    async def _publish_catalog() -> None:
+        for name, cls in get_registry().items():
+            try:
+                doc = json.dumps(catalog_doc(cls)).encode("utf-8")
+                await backend.catalog_set(name, doc)
+            except Exception:
+                logger.exception("failed to publish catalog entry for %s", name)
+
+    sweeper_task: "asyncio.Task | None" = None
+    worker_ttl_seconds = int(os.environ.get("FLEET_WORKER_IDLE_TTL_SECONDS", "0"))
+
+    @app.on_event("startup")
+    async def _start_sweeper() -> None:
+        nonlocal sweeper_task
+        import time as _time
+
+        async def _loop() -> None:
+            while True:
+                try:
+                    for name in get_registry():
+                        moved = await backend.queue_sweep_expired(name)
+                        if moved:
+                            logger.info("swept %d expired reservation(s) from %s", moved, name)
+                    if worker_ttl_seconds > 0:
+                        cutoff = _time.time() - worker_ttl_seconds
+                        for at in get_registry():
+                            for rec in await store.list_workers(at):
+                                if rec.last_seen and rec.last_seen < cutoff:
+                                    if await store.remove_worker(at, rec.worker_id):
+                                        logger.info(
+                                            "evicted idle worker %s/%s (last_seen %.0fs ago)",
+                                            at, rec.worker_id, _time.time() - rec.last_seen,
+                                        )
+                except asyncio.CancelledError:
+                    raise
+                except Exception:
+                    logger.exception("queue sweep failed")
+                await asyncio.sleep(10)
+
+        sweeper_task = asyncio.create_task(_loop(), name="queue-sweeper")
+
+    @app.on_event("shutdown")
+    async def _stop_sweeper() -> None:
+        if sweeper_task is not None and not sweeper_task.done():
+            sweeper_task.cancel()
+
+    app.include_router(make_api_router(store, broadcaster, backend))
+    app.include_router(make_ws_router(store, broadcaster, backend))
+    app.include_router(make_metrics_router(store, backend))
+    app.include_router(make_dashboard_router())
+
+    @app.get("/healthz")
+    async def healthz() -> dict[str, str]:
+        return {"status": "ok"}
+
+    return app

fleet/master/auth.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import json
+import logging
+import os
+import secrets
+
+from fastapi import Header, HTTPException
+
+logger = logging.getLogger(__name__)
+
+
+def _expected(env_name: str, fallback: str | None) -> str:
+    val = os.environ.get(env_name) or fallback
+    if not val:
+        val = secrets.token_hex(16)
+        os.environ[env_name] = val
+        print(f"[fleet] generated {env_name}={val}")
+    return val
+
+
+def admin_token() -> str:
+    return _expected("FLEET_ADMIN_TOKEN", None)
+
+
+def worker_token() -> str:
+    """Primary worker token. Kept for the WS-token check path that needs
+    a single canonical value (it always falls back to the first accepted
+    one). For request-time auth, use `worker_tokens()` which also accepts
+    rotation peers from FLEET_WORKER_TOKENS."""
+    return _expected("FLEET_WORKER_TOKEN", None)
+
+
+def worker_tokens() -> tuple[str, ...]:
+    """All currently-valid worker tokens. The primary token from
+    FLEET_WORKER_TOKEN is always included; additional rotation tokens
+    can be provided via FLEET_WORKER_TOKENS as a comma-separated string
+    or a JSON list. During a rotation, set the new token first, roll
+    workers across to it, then drop the old one."""
+    primary = worker_token()
+    extras: list[str] = []
+    raw = os.environ.get("FLEET_WORKER_TOKENS")
+    if raw:
+        raw = raw.strip()
+        if raw.startswith("["):
+            try:
+                parsed = json.loads(raw)
+                if isinstance(parsed, list):
+                    extras = [t for t in parsed if isinstance(t, str) and t]
+            except json.JSONDecodeError:
+                logger.warning("FLEET_WORKER_TOKENS is not valid JSON; ignoring")
+        else:
+            extras = [t.strip() for t in raw.split(",") if t.strip()]
+    seen: set[str] = set()
+    out: list[str] = []
+    for t in (primary, *extras):
+        if t not in seen:
+            seen.add(t)
+            out.append(t)
+    return tuple(out)
+
+
+def _matches_any(token: str, candidates: tuple[str, ...]) -> bool:
+    matched = False
+    for c in candidates:
+        if secrets.compare_digest(token, c):
+            matched = True
+    return matched
+
+
+def _token_scopes() -> dict[str, set[str]]:
+    """Parse FLEET_TOKEN_SCOPES as a JSON map of token -> list of automation types.
+
+    The admin token implicitly has the wildcard scope `*`. Bad JSON is logged
+    and ignored (the admin token still works).
+    """
+    raw = os.environ.get("FLEET_TOKEN_SCOPES")
+    out: dict[str, set[str]] = {admin_token(): {"*"}}
+    if not raw:
+        return out
+    try:
+        parsed = json.loads(raw)
+    except json.JSONDecodeError:
+        logger.warning("FLEET_TOKEN_SCOPES is not valid JSON; ignoring")
+        return out
+    if not isinstance(parsed, dict):
+        logger.warning("FLEET_TOKEN_SCOPES must be a JSON object {token: [types]}")
+        return out
+    for tok, scopes in parsed.items():
+        if not isinstance(tok, str) or not isinstance(scopes, list):
+            continue
+        out[tok] = {s for s in scopes if isinstance(s, str)}
+    return out
+
+
+def _bearer(authorization: str | None) -> str:
+    if not authorization:
+        raise HTTPException(401, "missing Authorization header")
+    if not authorization.lower().startswith("bearer "):
+        raise HTTPException(401, "Authorization must be Bearer <token>")
+    return authorization.split(" ", 1)[1].strip()
+
+
+async def require_admin(authorization: str | None = Header(default=None)) -> None:
+    token = _bearer(authorization)
+    if not secrets.compare_digest(token, admin_token()):
+        raise HTTPException(403, "invalid token")
+
+
+async def require_worker(authorization: str | None = Header(default=None)) -> None:
+    token = _bearer(authorization)
+    if not _matches_any(token, worker_tokens()):
+        raise HTTPException(403, "invalid token")
+
+
+async def require_scope(
+    automation_type: str,
+    authorization: str | None = Header(default=None),
+) -> None:
+    token = _bearer(authorization)
+    scopes_by_token = _token_scopes()
+    # find the token by constant-time compare against each known one
+    granted: set[str] | None = None
+    for known, scopes in scopes_by_token.items():
+        if secrets.compare_digest(token, known):
+            granted = scopes
+            break
+    if granted is None:
+        raise HTTPException(403, "invalid token")
+    if "*" in granted or automation_type in granted:
+        return
+    raise HTTPException(403, f"token has no scope for '{automation_type}'")