fleet-framework 0.1.0__py3-none-any.whl

fleet/core/config.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class BaseConfig(BaseModel):
+    # generic fields every automation has. plugin subclasses add domain knobs.
+    model_config = ConfigDict(extra="forbid")
+
+    enabled: bool = False
+    slots: int = Field(default=0, ge=0, le=1024)
+    recycle_seconds: int = Field(default=0, ge=0, le=86400)
+
+    max_attempts: int = Field(default=3, ge=1, le=100)
+    """Batch automations only: how many times a task is retried before the DLQ."""
+
+    visibility_seconds: int = Field(default=300, ge=10, le=86400)
+    """Batch automations only: reservation TTL on dequeued tasks. If the slot
+    crashes mid-run_one, the master sweeper requeues after this window."""
+
+    # simple-case proxy. wired into the built-in `static` provider when no
+    # provider is named. mutually compatible with proxy_provider/_config.
+    proxy_url: str | None = None
+
+    # provider-based proxy. set both fields together. when set, takes
+    # precedence over proxy_url.
+    proxy_provider: str | None = None
+    proxy_provider_config: dict[str, Any] | None = None
+
+    @classmethod
+    def restart_fields(cls) -> frozenset[str]:
+        return frozenset({"proxy_url", "proxy_provider", "proxy_provider_config"})
+
+    @classmethod
+    def reboot_fields(cls) -> frozenset[str]:
+        return frozenset()

fleet/core/context.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar
+
+from pydantic import BaseModel
+
+if TYPE_CHECKING:
+    from fleet.core.backend import Backend
+    from fleet.core.config import BaseConfig
+    from fleet.core.contract import Pool, Queue, Stream
+    from fleet.core.events import EventBus
+    from fleet.core.metrics import SlotMetrics
+    from fleet.core.primitives import KV, Counter, Lock, PoolHandle, StreamReader
+    from fleet.core.proxy import ProxyHandle
+
+C = TypeVar("C", bound="BaseConfig")
+P = TypeVar("P", bound=BaseModel)
+T = TypeVar("T", bound=BaseModel)
+
+
+@dataclass
+class Context(Generic[C]):
+    # passed into every slot/task callback. owns the lifecycle of one slot.
+    automation_type: str
+    worker_id: str
+    slot_id: int
+    config: C
+    shutdown: asyncio.Event
+    logger: logging.Logger
+    _emit_raw: Any  # async (payload_dict) -> None — pushes onto own stream via master
+    _backend: "Backend"
+    _events: "EventBus"
+    _metrics: "SlotMetrics"
+    _automation_cls: Any = None  # type[BaseAutomation], set by SlotRunner
+    _proxy: Optional["ProxyHandle"] = None
+    proxy_url: Optional[str] = None
+    extra: dict[str, Any] = field(default_factory=dict)
+
+    # ---- proxy ----
+    @property
+    def proxy(self) -> "ProxyHandle":
+        from fleet.core.proxy import ProxyHandle
+        if self._proxy is None:
+            self._proxy = ProxyHandle(provider=None, slot_id=self.slot_id)
+        return self._proxy
+
+    # ---- own output stream ----
+    async def emit(self, payload: BaseModel) -> None:
+        # the automation declared `Output = SomePydantic`. validate at the boundary.
+        cls = self._automation_cls
+        if cls is not None and cls.Output is not None and not isinstance(payload, cls.Output):
+            raise TypeError(
+                f"{self.automation_type}: emit() expected {cls.Output.__name__}, "
+                f"got {type(payload).__name__}"
+            )
+        if hasattr(payload, "model_dump"):
+            await self._emit_raw(payload.model_dump(mode="json"))
+        else:
+            raise TypeError("emit() requires a Pydantic model")
+
+    # ---- read another automation's stream ----
+    def stream(self, spec: "Stream[P]") -> "StreamReader[P]":
+        from fleet.core.primitives import StreamReader
+        return StreamReader(spec, self._backend)
+
+    # ---- pools ----
+    def pool(self, spec: "Pool[P, T]") -> "PoolHandle[P, T]":
+        from fleet.core.primitives import PoolHandle
+        return PoolHandle(spec, self._backend)
+
+    # ---- submit task to another automation ----
+    async def submit(self, spec: "Queue[P]", payload: P, *, task_id: Optional[str] = None) -> str:
+        from fleet.core.primitives import QueueHandle
+        return await QueueHandle(spec, self._backend).push(payload, task_id=task_id)
+
+    # ---- kv ----
+    def kv(self, namespace: str) -> "KV":
+        from fleet.core.primitives import KV
+        return KV(self._backend, namespace)
+
+    # ---- lock ----
+    def lock(self, name: str, *, hold_seconds: int = 30, wait_seconds: float = 5.0) -> "Lock":
+        from fleet.core.primitives import Lock
+        return Lock(self._backend, name, hold_seconds, wait_seconds)
+
+    # ---- counter ----
+    def counter(self, name: str) -> "Counter":
+        from fleet.core.primitives import Counter
+        return Counter(self._backend, name)
+
+    # ---- event bus ----
+    @property
+    def events(self) -> "EventBus":
+        return self._events
+
+    # ---- metrics ----
+    @property
+    def metrics(self) -> "SlotMetrics":
+        return self._metrics

fleet/core/contract.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar
+
+from pydantic import BaseModel, TypeAdapter
+
+P = TypeVar("P", bound=BaseModel)
+T = TypeVar("T", bound=BaseModel)
+
+
+@dataclass(frozen=True)
+class Pool(Generic[P, T]):
+    name: str
+    payload: type[P]
+    tags: type[T]
+
+    def __init__(self, name: str, *, payload: type[P], tags: type[T]) -> None:
+        object.__setattr__(self, "name", name)
+        object.__setattr__(self, "payload", payload)
+        object.__setattr__(self, "tags", tags)
+
+    def json_schema(self) -> dict[str, Any]:
+        return {
+            "kind": "pool",
+            "name": self.name,
+            "payload": self.payload.model_json_schema(),
+            "tags": self.tags.model_json_schema(),
+        }
+
+
+@dataclass(frozen=True)
+class Queue(Generic[P]):
+    """Task queue. `name` is the automation_type that consumes from it."""
+    name: str
+    payload: type[P]
+
+    def __init__(self, name: str, *, payload: type[P]) -> None:
+        object.__setattr__(self, "name", name)
+        object.__setattr__(self, "payload", payload)
+
+    def json_schema(self) -> dict[str, Any]:
+        return {
+            "kind": "queue",
+            "name": self.name,
+            "payload": self.payload.model_json_schema(),
+        }
+
+
+@dataclass(frozen=True)
+class Stream(Generic[P]):
+    """Output stream. `name` is the automation_type that emits to it."""
+    name: str
+    payload: type[P]
+
+    def __init__(self, name: str, *, payload: type[P]) -> None:
+        object.__setattr__(self, "name", name)
+        object.__setattr__(self, "payload", payload)
+
+    def json_schema(self) -> dict[str, Any]:
+        return {
+            "kind": "stream",
+            "name": self.name,
+            "payload": self.payload.model_json_schema(),
+        }
+
+
+def validate_tag_filter(tags_type: type[BaseModel], **kwargs: Any) -> dict[str, Any]:
+    """Validate kwargs against a Tags model. TypeError on unknown field,
+    ValueError on type mismatch. Returns coerced values."""
+    fields = tags_type.model_fields
+    out: dict[str, Any] = {}
+    for k, v in kwargs.items():
+        if k not in fields:
+            known = ", ".join(sorted(fields.keys())) or "(none)"
+            raise TypeError(
+                f"unknown tag '{k}' for {tags_type.__name__}; known tags: {known}"
+            )
+        annotation = fields[k].annotation
+        try:
+            out[k] = TypeAdapter(annotation).validate_python(v)
+        except Exception as e:
+            raise ValueError(f"tag '{k}' value {v!r} invalid: {e}") from e
+    return out
+
+
+__all__ = ["Pool", "Queue", "Stream", "validate_tag_filter", "P", "T"]

fleet/core/country_presets.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+# ISO 3166-1 alpha-2. used by proxy providers to resolve string presets to
+# concrete country lists. providers stringify the result however they need.
+
+EU_27 = [
+    "AT", "BE", "BG", "HR", "CY", "CZ", "DK", "EE", "FI", "FR",
+    "DE", "GR", "HU", "IE", "IT", "LV", "LT", "LU", "MT", "NL",
+    "PL", "PT", "RO", "SK", "SI", "ES", "SE",
+]
+EU_PLUS_GB = EU_27 + ["GB"]
+NA = ["US", "CA"]
+EU_US = EU_27 + ["US"]
+LATAM = ["AR", "BR", "CL", "CO", "MX", "PE", "UY"]
+APAC = ["AU", "JP", "KR", "SG", "HK", "TW", "NZ", "MY", "TH", "ID", "PH", "VN"]
+
+_PRESETS: dict[str, list[str]] = {
+    "EU": EU_27,
+    "EU+UK": EU_PLUS_GB,
+    "EU+GB": EU_PLUS_GB,
+    "NA": NA,
+    "EU+US": EU_US,
+    "EU+NA": EU_27 + NA,
+    "LATAM": LATAM,
+    "APAC": APAC,
+    "WORLDWIDE": [],
+    "GLOBAL": [],
+    "ANY": [],
+}
+
+
+def resolve_countries(spec: str | list[str] | None) -> list[str]:
+    # returns uppercase ISO codes. empty list means "no restriction" (worldwide).
+    if spec is None:
+        return []
+    if isinstance(spec, list):
+        return [c.strip().upper() for c in spec if c and c.strip()]
+    s = spec.strip()
+    if not s:
+        return []
+    key = s.upper()
+    if key in _PRESETS:
+        return list(_PRESETS[key])
+    if "," in s:
+        return [p.strip().upper() for p in s.split(",") if p.strip()]
+    return [key]
+
+
+def known_presets() -> list[str]:
+    return sorted(_PRESETS.keys())

fleet/core/events.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+import redis.asyncio as aioredis
+
+# lightweight pub/sub on redis. ephemeral — no persistence. for signals,
+# not for work hand-off (use stream.py for that).
+
+
+class EventBus:
+    def __init__(self, r: aioredis.Redis) -> None:
+        self._r = r
+
+    async def publish(self, topic: str, payload: dict[str, Any]) -> None:
+        await self._r.publish(f"event:{topic}", json.dumps(payload))
+
+    async def subscribe(self, topic: str) -> AsyncIterator[dict[str, Any]]:
+        # async generator that yields payloads forever. cancel the consumer
+        # task to stop. consumer is responsible for handling backpressure.
+        pubsub = self._r.pubsub()
+        await pubsub.subscribe(f"event:{topic}")
+        try:
+            async for msg in pubsub.listen():
+                if msg is None or msg.get("type") != "message":
+                    continue
+                data = msg.get("data")
+                if data is None:
+                    continue
+                try:
+                    yield json.loads(data)
+                except (json.JSONDecodeError, TypeError):
+                    continue
+        finally:
+            try:
+                await pubsub.unsubscribe(f"event:{topic}")
+                await pubsub.aclose()
+            except Exception:
+                pass
+
+
+async def first(topic: str, bus: EventBus, *, timeout: float = 5.0) -> dict[str, Any] | None:
+    # helper for tests / one-shot waits. returns None on timeout.
+    async def _consume() -> dict[str, Any]:
+        async for ev in bus.subscribe(topic):
+            return ev
+        raise RuntimeError("unreachable")
+
+    try:
+        return await asyncio.wait_for(_consume(), timeout=timeout)
+    except asyncio.TimeoutError:
+        return None

fleet/core/logging.py

@@ -0,0 +1,97 @@
+"""Logging setup: text or JSON, with contextvars threaded into every record.
+
+Set `FLEET_LOG_FORMAT=json` to emit structured logs. Set `LOG_LEVEL` for the
+threshold. Plugins call `bind_log_context(...)` to attach worker_id /
+slot_id / automation_type / task_id to records emitted from their slot.
+"""
+
+from __future__ import annotations
+
+import contextvars
+import json
+import logging
+import os
+from typing import Any
+
+_context: contextvars.ContextVar[dict[str, Any]] = contextvars.ContextVar(
+    "_fleet_log_context", default={}
+)
+
+
+def bind_log_context(**kwargs: Any) -> contextvars.Token:
+    """Add fields to the current async context's log records.
+
+    Returns a Token; pass it to reset_log_context to roll back.
+    """
+    merged = {**_context.get(), **kwargs}
+    return _context.set(merged)
+
+
+def reset_log_context(token: contextvars.Token) -> None:
+    _context.reset(token)
+
+
+class _ContextFilter(logging.Filter):
+    def filter(self, record: logging.LogRecord) -> bool:
+        ctx = _context.get()
+        for k, v in ctx.items():
+            if not hasattr(record, k):
+                setattr(record, k, v)
+        return True
+
+
+_STANDARD_LOGRECORD_FIELDS = frozenset({
+    "args", "asctime", "created", "exc_info", "exc_text", "filename",
+    "funcName", "levelname", "levelno", "lineno", "message", "module",
+    "msecs", "msg", "name", "pathname", "process", "processName",
+    "relativeCreated", "stack_info", "thread", "threadName",
+    "taskName",
+})
+
+
+class _JsonFormatter(logging.Formatter):
+    def format(self, record: logging.LogRecord) -> str:
+        # record.created is seconds-since-epoch float; iso-format with ms precision.
+        from datetime import datetime, timezone
+        ts = datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(timespec="milliseconds")
+        out: dict[str, Any] = {
+            "ts": ts,
+            "level": record.levelname,
+            "logger": record.name,
+            "msg": record.getMessage(),
+        }
+        if record.exc_info:
+            out["exc"] = self.formatException(record.exc_info)
+        for k, v in record.__dict__.items():
+            if k in _STANDARD_LOGRECORD_FIELDS or k.startswith("_"):
+                continue
+            out[k] = v
+        return json.dumps(out, default=str)
+
+
+def configure_logging(
+    level: str | None = None,
+    fmt: str | None = None,
+) -> None:
+    """Idempotent log setup. Reads LOG_LEVEL and FLEET_LOG_FORMAT from env
+    when args are None."""
+    lvl = (level or os.environ.get("LOG_LEVEL", "INFO")).upper()
+    chosen_fmt = (fmt or os.environ.get("FLEET_LOG_FORMAT", "text")).lower()
+
+    root = logging.getLogger()
+    for h in list(root.handlers):
+        root.removeHandler(h)
+
+    handler = logging.StreamHandler()
+    if chosen_fmt == "json":
+        handler.setFormatter(_JsonFormatter())
+    else:
+        handler.setFormatter(logging.Formatter(
+            "%(asctime)s %(levelname)s %(name)s: %(message)s"
+        ))
+    handler.addFilter(_ContextFilter())
+    root.addHandler(handler)
+    root.setLevel(lvl)
+
+
+__all__ = ["bind_log_context", "configure_logging", "reset_log_context"]