diff --git a/.github/workflows/license-check-python.yml b/.github/workflows/license-check-python.yml index ea84e2d8..4b6ccd3a 100644 --- a/.github/workflows/license-check-python.yml +++ b/.github/workflows/license-check-python.yml @@ -8,6 +8,7 @@ on: - 'osprey_rpc/pyproject.toml' - 'osprey_worker/pyproject.toml' - 'example_plugins/pyproject.toml' + - 'example_atproto_plugins/pyproject.toml' - '.github/allowed-licenses.txt' - '.github/pip-licenses-classifiers.txt' permissions: @@ -44,7 +45,7 @@ jobs: # google-crc32c: Apache-2.0; no metadata in 1.8.0 (googleapis/google-cloud-python#16515) # ddtrace: Apache-2.0 OR BSD-3-Clause; sets License field to a filename ("LICENSE.BSD3") # First-party packages (no license metadata needed): - # osprey-rpc, osprey-worker, example_plugins + # osprey-rpc, osprey-worker, example_plugins, example_atproto_plugins # Packages pending license policy review: # unidecode: GPL v2+ (#354) # tld: GPL v2 / LGPL / MPL 1.1 (#355) @@ -55,5 +56,5 @@ jobs: --ignore-packages \ google-crc32c \ ddtrace \ - osprey-rpc osprey-worker example_plugins \ + osprey-rpc osprey-worker example_plugins example_atproto_plugins \ unidecode tld psycopg2-binary simplejson text-unidecode diff --git a/AGENTS.md b/AGENTS.md index e7fae7b5..8fabf029 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -12,14 +12,16 @@ Top-level modules: - `osprey_coordinator/` — Rust gRPC coordinator (tokio, tonic, etcd, rdkafka). Rust code belongs here. - `proto/osprey/rpc/` — protobuf source of truth for `osprey_rpc` and `osprey_coordinator` types. - `example_plugins/` — reference plugins (UDFs, output sinks, labels service) using the pluggy-based plugin system. Do not add production code here. +- `example_atproto_plugins/` — reference plugin demonstrating a custom input stream that consumes the Bluesky firehose. Stack `docker-compose.atproto.yaml` on top of the main compose file (or use `./run-atproto.sh`) to run Osprey against live ATProto traffic. Do not add production code here. - `example_rules/` — sample SML rules and YAML config. +- `example_atproto_rules/` — sample SML rules paired with `example_atproto_plugins/`. Reference files: `docs/DEVELOPMENT.md` (setup), `example_plugins/src/register_plugins.py` (plugin patterns), `example_plugins/src/services/labels_service.py` (labels service example). ## Design - API: gRPC between `osprey_coordinator` and workers; HTTP/Flask for `osprey-ui-api` (port 5004); protobuf definitions under `proto/osprey/rpc/` are authoritative. -- Rules: SML (Osprey's rule language) with user-defined functions registered via pluggy hooks (`@hookimpl_osprey`): `register_udfs`, `register_output_sinks`, `register_labels_service_or_provider`. +- Rules: SML (Osprey's rule language) with user-defined functions registered via pluggy hooks (`@hookimpl_osprey`): `register_udfs`, `register_output_sinks`, `register_labels_service_or_provider`, `register_input_stream` (custom event source; see `example_atproto_plugins/`). - Data model conventions: Pydantic for models, SQLAlchemy for persistence (versions pinned in `pyproject.toml`). ## Build and run diff --git a/CHANGELOG.md b/CHANGELOG.md index 3dc19cf5..a26bd293 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ For more information about each release including git tags and artifacts, see [R - Per-action health metrics in the executor ([#191](https://github.com/roostorg/osprey/pull/191) by [@cmttt](https://github.com/cmttt)) - Option to suppress cached errors to reduce metric bloat ([#180](https://github.com/roostorg/osprey/pull/180) by [@lithium-powered](https://github.com/lithium-powered)) - Experimental asyncio-native worker with metrics and engine/coordinator improvements ([#341](https://github.com/roostorg/osprey/pull/341) by [@cmttt](https://github.com/cmttt)) +- ATProto JetStream example plugins and rules ([#236](https://github.com/roostorg/osprey/pull/236) by [@haileyok](https://github.com/haileyok)) - `osprey-stress` CLI: closed-loop stress harness that produces synthetic events at a configurable rate, observes their `ExecutionResult`s on the output topic, and reports drop rate and p50/p95/p99 latency, exiting non-zero on threshold breach so it can gate CI on pipeline health ([#367](https://github.com/roostorg/osprey/pull/367) by [@julietshen](https://github.com/julietshen), closes [#324](https://github.com/roostorg/osprey/issues/324)) ### Changed diff --git a/docker-compose.atproto.yaml b/docker-compose.atproto.yaml new file mode 100644 index 00000000..9e03c4f5 --- /dev/null +++ b/docker-compose.atproto.yaml @@ -0,0 +1,21 @@ +# Override that swaps the synthetic Kafka producer for the live Bluesky JetStream +# firehose, via the example_atproto_plugins package. Stack on top of the main +# compose file: +# +# docker compose -f docker-compose.yaml -f docker-compose.atproto.yaml up +# +# Or use the convenience wrapper: ./run-atproto.sh +services: + osprey-worker: + environment: + OSPREY_INPUT_STREAM_SOURCE: plugin + OSPREY_RULES_PATH: /osprey/example_atproto_rules + volumes: + - ./example_atproto_rules:/osprey/example_atproto_rules + - ./example_atproto_plugins:/osprey/example_atproto_plugins + + osprey-ui-api: + environment: + OSPREY_RULES_PATH: /osprey/example_atproto_rules + volumes: + - ./example_atproto_rules:/osprey/example_atproto_rules diff --git a/example_atproto_plugins/README.md b/example_atproto_plugins/README.md new file mode 100644 index 00000000..6d1fee1b --- /dev/null +++ b/example_atproto_plugins/README.md @@ -0,0 +1,123 @@ +# example_atproto_plugins + +A sample Osprey plugin that consumes ATProto's [JetStream](https://docs.bsky.app/blog/jetstream) as the input event source. It gives you: + +- a `register_input_stream` hook implementation that subscribes to JetStream over WebSocket and yields Osprey `Action`s with the JetStream JSON event passed through as-is, +- realistic per-second event volume from the live Bluesky network, which is useful for load and soak testing changes that the synthetic 1-event/second producer doesn't exercise, +- a companion `example_atproto_rules/` tree showing how to organize rules against ATProto event shapes, with file structure modeled on [haileyok/atproto-ruleset](https://github.com/haileyok/atproto-ruleset). + +This package registers the **input stream** plus two optional enrichment UDFs (see below). The sample rules also use a UDF (`TextContains`), a labels service, and an output sink that are provided by the sibling `example_plugins/` package, so the two run together: the worker image installs both, and Osprey loads every registered plugin, so `example_plugins` supplies those pieces automatically in the docker stack. If you lift this sample into a setup without `example_plugins`, provide those yourself (a labels provider and output sink) or restrict the rules to stdlib UDFs. + +## Running + +From the repo root: + +```sh +./run-atproto.sh +``` + +This brings up the full Osprey local stack (Druid, Postgres, MinIO, Kafka) along with a JetStream websocket override, and swaps the worker's input source from Kafka to the JetStream plugin, pointing it at `example_atproto_rules` instead of `example_rules`. First-run startup takes a few minutes. + +## Configuration + +| Env var | Default | Description | +| --- | --- | --- | +| `OSPREY_INPUT_STREAM_SOURCE` | (must be) `plugin` | Selects the plugin-provided stream. | +| `OSPREY_JETSTREAM_ENDPOINT` | `wss://jetstream2.us-west.bsky.network/subscribe` | JetStream WebSocket URL. | +| `OSPREY_JETSTREAM_WANTED_COLLECTIONS` | `app.bsky.feed.post,app.bsky.feed.like,app.bsky.feed.repost,app.bsky.graph.follow,app.bsky.actor.profile` | Comma-separated collections to subscribe to (server-side filter). | + +## Action shape + +The JetStream JSON event is passed through unchanged as the Action's `data` dict, so rules read JetStream-native paths directly. `action_name` is `_` for commit events (`create_post`, `delete_like`, `update_profile`, …) using the short names defined in `COLLECTION_NAMES`, or `identity` for identity events. + +### Commit events (e.g. `create_post`, `delete_like`) + +``` +{ + "did": "did:plc:...", + "time_us": 1714500000000000, + "kind": "commit", + "commit": { + "rev": "...", + "operation": "create" | "update" | "delete", + "collection": "app.bsky.feed.post", + "rkey": "...", + "cid": "...", + "record": { ... raw ATProto record ... } + } +} +``` + +### Identity events (`action_name='identity'`) + +``` +{ + "did": "did:plc:...", + "time_us": ..., + "kind": "identity", + "identity": {"did": "...", "seq": ..., "time": "..."} +} +``` + +JetStream identity events carry only `did` / `seq` / `time` — not the handle. Resolve the handle from the DID via the opt-in enrichment below. + +Account events, commits for collections not in `COLLECTION_NAMES`, and commits with operations other than `create` / `update` / `delete` are skipped. + +### Profile enrichment (opt-in) + +JetStream events identify the actor only by DID, which isn't searchable the way a handle or display name is. The plugin ships two UDFs, `AtprotoHandle` and `AtprotoDisplayName`, that resolve a DID to those fields via Bluesky's public, unauthenticated AppView (`app.bsky.actor.getProfile`). Results are cached per DID and lookups fail soft (the feature is simply absent) when the API errors or rate-limits. The whole profile is fetched once per DID: because async UDFs run concurrently, a rule that reads both fields would otherwise fire two `getProfile` calls at once, so a fetch already in progress for a DID is shared rather than duplicated. Cached entries expire after an hour since handles and display names change; the fuller approach is to bust a DID's entry when an identity or profile-update event comes through JetStream, left out here to keep the example focused. + +**It is off by default.** Each unique DID costs an external API call, which is fine for a demo but is exactly the kind of dependency you don't want in a load test — so the default rules run against the raw firehose with no outbound calls. To turn enrichment on: + +1. Import `models/enrichment.sml` in `example_atproto_rules/main.sml`. Imports must stay lexicographically sorted, so the list becomes: + + ``` + Import( + rules=[ + 'models/base.sml', + 'models/enrichment.sml', + 'models/record/base.sml', + 'models/record/post.sml', + ], + ) + ``` + +2. Add `Handle` and `DisplayName` to the `['*']` feature list in `example_atproto_rules/config/ui_config.yaml` so they show in the event stream. + +For a smoother demo once enabled, narrow `OSPREY_JETSTREAM_WANTED_COLLECTIONS` to lower the unique-DID (and thus request) rate. + +### Extending the enrichment + +`getProfile` returns the whole profile, and `enrichment_udfs.py` already caches it per DID, so more trust & safety signals are cheap to add — a new UDF just reads another field off the same cached fetch. For example, an account-age signal: + +```python +from datetime import datetime, timezone + + +class AtprotoAccountAgeDays(UDFBase[DidArguments, int]): + """Whole days since the account's profile was created.""" + + category = _ATPROTO_CATEGORY + execute_async = True + + def execute(self, execution_context: ExecutionContext, arguments: DidArguments) -> int: + created_at = _profile_or_skip(arguments.did).get('createdAt') + if not isinstance(created_at, str): + raise ExpectedUdfException() + created = datetime.fromisoformat(created_at.replace('Z', '+00:00')) + return max(0, (datetime.now(timezone.utc) - created).days) +``` + +Register it in `register_plugins.py`'s `register_udfs`, then reference it from `enrichment.sml`. The same pattern exposes `followersCount` / `followsCount` / `postsCount` (bot/spam heuristics), `description` (a scannable bio), or `labels` (moderation labels already applied to the account). + +### UI default features + +`example_atproto_rules/config/ui_config.yaml` declares the per-action default features the Osprey UI surfaces in the event stream — e.g. `UserId` for every action, `PostText` for `create_post`, `Subject` for like / repost / follow events. Add new entries there to expose more fields without touching rule code. + +`action_id` is minted from `snowflake-id-worker` in batches of 250. The plugin therefore needs `SNOWFLAKE_API_ENDPOINT` to be set (the local docker-compose stack provides it). + +## Caveats + +- **Not production-ready.** No durable cursor on process restart, no zstd compression, no DID-level filtering. Good for sample / load-testing purposes; not a drop-in for a real ATProto deployment. +- **Enrichment is off by default and best-effort.** JetStream carries no handle/profile/account-age data. The opt-in `Handle` / `DisplayName` UDFs resolve a DID against the public AppView on demand (cached, fail-soft), which is enough for demos but will rate-limit at full firehose volume — so it stays off unless you enable it, keeping load tests dependency-free. Rulesets that need reliable, complete enrichment (such as much of [atproto-ruleset](https://github.com/haileyok/atproto-ruleset)) still want a dedicated enrichment pipeline in front of this one rather than per-event API lookups. +- **Connection health.** WebSocket-level PING/PONG keepalive runs every 20s with a 10s pong timeout (`websocket-client`'s `WebSocketApp.run_forever(ping_interval, ping_timeout)`). A stalled or dead connection is detected within ~30s and triggers a reconnect from the last seen `time_us` cursor. diff --git a/example_atproto_plugins/__init__.py b/example_atproto_plugins/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/example_atproto_plugins/pyproject.toml b/example_atproto_plugins/pyproject.toml new file mode 100644 index 00000000..246b1a0f --- /dev/null +++ b/example_atproto_plugins/pyproject.toml @@ -0,0 +1,18 @@ +[project] +name = "example_atproto_plugins" +version = "0.1.0" +description = "Example Osprey plugin that consumes Bluesky's ATProto JetStream firehose" +requires-python = ">=3.11" +dependencies = [ + "pluggy==1.5.0", + "websocket-client==1.8.0", +] + +[tool.setuptools] +package-dir = {"" = "src"} + +[tool.setuptools.packages.find] +where = ["src"] + +[project.entry-points.osprey_plugin] +atproto_plugins = "atproto_plugin.register_plugins" diff --git a/example_atproto_plugins/src/atproto_plugin/__init__.py b/example_atproto_plugins/src/atproto_plugin/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/example_atproto_plugins/src/atproto_plugin/enrichment_udfs.py b/example_atproto_plugins/src/atproto_plugin/enrichment_udfs.py new file mode 100644 index 00000000..fffc2f5a --- /dev/null +++ b/example_atproto_plugins/src/atproto_plugin/enrichment_udfs.py @@ -0,0 +1,164 @@ +"""Optional enrichment UDFs that resolve an ATProto DID to profile fields. + +JetStream events identify the actor only by DID, which isn't searchable the way a +handle or display name is. These UDFs resolve a DID to those fields via Bluesky's +public, unauthenticated AppView (`app.bsky.actor.getProfile`). + +They are registered by the plugin but wired into rules only via the opt-in +`models/enrichment.sml`, because each unique DID costs an external API call -- +great for demos, but a dependency you don't want in a load test. The whole +profile is fetched once per DID and cached, and lookups fail soft (the feature is +simply absent) when the API errors or rate-limits. + +Async UDFs run concurrently in a gevent pool, so a rule that reads both the handle +and the display name would fire `AtprotoHandle` and `AtprotoDisplayName` at the +same time. Both would miss a cold cache and each make its own `getProfile` call. +To avoid that, a fetch in progress for a DID is shared: the second greenlet waits +on the first one's result instead of making a duplicate request. Cached entries +expire after `_CACHE_TTL_SECONDS`, since handles and display names change; the +fuller approach is to bust a DID's entry when an identity or profile-update event +comes through JetStream, which is left out here to keep the example focused. + +See the README's "Extending the enrichment" section for how to expose more of the +profile (account age, follower counts, existing labels) from the same cached fetch. +""" + +import time +from collections import OrderedDict +from threading import Event, Lock +from typing import Any + +import requests +from osprey.engine.executor.execution_context import ExecutionContext, ExpectedUdfException +from osprey.engine.udf.arguments import ArgumentsBase +from osprey.engine.udf.base import UDFBase + +_ATPROTO_CATEGORY = 'ATProto' +_GET_PROFILE_URL = 'https://public.api.bsky.app/xrpc/app.bsky.actor.getProfile' +_REQUEST_TIMEOUT_SECONDS = 5 +_CACHE_MAX_SIZE = 100_000 +_CACHE_TTL_SECONDS = 60 * 60 + +_session = requests.Session() +# did -> (profile dict, monotonic time at which the entry expires). Ordered so the +# least-recently-used entry is evicted first once the cache is full. +_profile_cache: 'OrderedDict[str, tuple[dict[str, Any], float]]' = OrderedDict() +# did -> a fetch currently in progress, so concurrent misses for the same DID +# (e.g. AtprotoHandle and AtprotoDisplayName on one event) share one API call. +_inflight: dict[str, '_InflightFetch'] = {} +_cache_lock = Lock() + + +class _InflightFetch: + """A single `getProfile` call in progress, shared by every greenlet awaiting it. + + The greenlet that created it does the fetch and populates `profile` or `error` + before setting `done`; waiters block on `done`, then read the result. Under + gevent's cooperative scheduling this needs no memory barrier -- the waiter only + runs after `done.set()` yields back to it. + """ + + __slots__ = ('done', 'profile', 'error') + + def __init__(self) -> None: + self.done = Event() + self.profile: dict[str, Any] | None = None + self.error: Exception | None = None + + +def _get_profile_from_api(did: str) -> dict[str, Any]: + """Hit the public getProfile endpoint. Raises on any transport/HTTP/parse error.""" + response = _session.get(_GET_PROFILE_URL, params={'actor': did}, timeout=_REQUEST_TIMEOUT_SECONDS) + response.raise_for_status() + profile = response.json() + if not isinstance(profile, dict): + raise ValueError('getProfile did not return an object') + return profile + + +def _fetch_profile(did: str) -> dict[str, Any]: + """Return the getProfile response for a DID, coalescing concurrent cache misses. + + Raises on any transport/HTTP/parse error so callers can fail soft. + """ + with _cache_lock: + cached = _profile_cache.get(did) + if cached is not None: + profile, expires_at = cached + if time.monotonic() < expires_at: + _profile_cache.move_to_end(did) + return profile + del _profile_cache[did] + + inflight = _inflight.get(did) + is_leader = inflight is None + if inflight is None: + inflight = _InflightFetch() + _inflight[did] = inflight + + if not is_leader: + # Someone else is already fetching this DID; wait for their result. + inflight.done.wait() + if inflight.error is not None: + raise inflight.error + assert inflight.profile is not None + return inflight.profile + + # Leader: do the request outside the lock so a slow call doesn't block other DIDs. + try: + profile = _get_profile_from_api(did) + except Exception as exc: + inflight.error = exc + with _cache_lock: + _inflight.pop(did, None) + inflight.done.set() + raise + + with _cache_lock: + _profile_cache[did] = (profile, time.monotonic() + _CACHE_TTL_SECONDS) + _profile_cache.move_to_end(did) + while len(_profile_cache) > _CACHE_MAX_SIZE: + _profile_cache.popitem(last=False) + _inflight.pop(did, None) + inflight.profile = profile + inflight.done.set() + return profile + + +def _profile_or_skip(did: str) -> dict[str, Any]: + """Fetch the cached profile, converting any lookup failure into a soft skip.""" + try: + return _fetch_profile(did) + except (requests.RequestException, ValueError): + raise ExpectedUdfException() + + +class DidArguments(ArgumentsBase): + did: str + """The ATProto DID to resolve (e.g. the actor's `$.did`).""" + + +class AtprotoHandle(UDFBase[DidArguments, str]): + """Resolves an ATProto DID to its current handle.""" + + category = _ATPROTO_CATEGORY + execute_async = True + + def execute(self, execution_context: ExecutionContext, arguments: DidArguments) -> str: + handle = _profile_or_skip(arguments.did).get('handle') + if not handle: + raise ExpectedUdfException() + return handle + + +class AtprotoDisplayName(UDFBase[DidArguments, str]): + """Resolves an ATProto DID to its display name.""" + + category = _ATPROTO_CATEGORY + execute_async = True + + def execute(self, execution_context: ExecutionContext, arguments: DidArguments) -> str: + display_name = _profile_or_skip(arguments.did).get('displayName') + if not display_name: + raise ExpectedUdfException() + return display_name diff --git a/example_atproto_plugins/src/atproto_plugin/jetstream_input_stream.py b/example_atproto_plugins/src/atproto_plugin/jetstream_input_stream.py new file mode 100644 index 00000000..77066f03 --- /dev/null +++ b/example_atproto_plugins/src/atproto_plugin/jetstream_input_stream.py @@ -0,0 +1,190 @@ +import json +import time +from datetime import datetime, timezone +from typing import Any, Dict, Iterator, List, Optional +from urllib.parse import urlencode + +import gevent +import websocket +from gevent.queue import Queue +from osprey.engine.executor.execution_context import Action +from osprey.worker.lib.backoff import Backoff +from osprey.worker.lib.instruments import metrics +from osprey.worker.lib.osprey_shared.logging import get_logger +from osprey.worker.lib.snowflake import generate_snowflake_batch +from osprey.worker.sinks.sink.input_stream import BaseInputStream +from osprey.worker.sinks.utils.acking_contexts import BaseAckingContext, NoopAckingContext + +logger = get_logger() + + +DEFAULT_ENDPOINT = 'wss://jetstream2.us-west.bsky.network/subscribe' +COLLECTION_NAMES = { + 'app.bsky.feed.post': 'post', + 'app.bsky.feed.like': 'like', + 'app.bsky.feed.repost': 'repost', + 'app.bsky.graph.follow': 'follow', + 'app.bsky.actor.profile': 'profile', +} +DEFAULT_COLLECTIONS = tuple(COLLECTION_NAMES.keys()) +SNOWFLAKE_BATCH_SIZE = 250 +PING_INTERVAL_SECONDS = 20 +PING_TIMEOUT_SECONDS = 10 + + +class JetStreamInputStream(BaseInputStream[BaseAckingContext[Action]]): + """An Osprey event input stream that subscribes to the ATProto JetStream websocket and yields + Osprey actions. + + The JetStream JSON event is passed through unchanged as the Action's data dict, so rules may + target the JetStream-native paths directly, like $.did, $.kind, or $.commit.operation. + """ + + def __init__( + self, + endpoint: Optional[str] = None, + wanted_collections: Optional[List[str]] = None, + reconnect_seconds: float = 2.0, + max_reconnect_seconds: float = 60.0, + ): + super().__init__() + self._endpoint = endpoint or DEFAULT_ENDPOINT + self._wanted_collections = list(wanted_collections) if wanted_collections else list(DEFAULT_COLLECTIONS) + self._backoff = Backoff(min_delay=reconnect_seconds, max_delay=max_reconnect_seconds) + self._last_time_us: Optional[int] = None + self._snowflake_buffer: List[int] = [] + + def _next_action_id(self) -> int: + if not self._snowflake_buffer: + batch = generate_snowflake_batch(count=SNOWFLAKE_BATCH_SIZE, retries=3) + self._snowflake_buffer = [s.to_int() for s in batch] + return self._snowflake_buffer.pop() + + def _build_url(self) -> str: + params = [('wantedCollections', c) for c in self._wanted_collections] + if self._last_time_us is not None: + params.append(('cursor', str(self._last_time_us))) + return f'{self._endpoint}?{urlencode(params)}' + + def _gen(self) -> Iterator[BaseAckingContext[Action]]: + while True: + had_event = False + try: + url = self._build_url() + for ctx in self._stream_one_connection(url): + had_event = True + yield ctx + except Exception as e: + logger.exception(f'JetStream stream error: {e}') + + if had_event: + self._backoff.succeed() + delay = self._backoff.current + else: + delay = self._backoff.fail() + logger.info(f'Reconnecting in {delay:.1f}s') + time.sleep(delay) + + def _stream_one_connection(self, url: str) -> Iterator[BaseAckingContext[Action]]: + # WebSocketApp drives PING/PONG keepalive on its own greenlet; we bridge its + # callback API into this generator via a gevent.queue.Queue. The 'done' sentinel + # is pushed by on_close/on_error, and also as a safety net by greenlet.link in + # case run_forever exits without firing on_close (e.g., uncaught exception). + queue: 'Queue[tuple[str, Optional[bytes]]]' = Queue() + + def on_open(ws: Any) -> None: + logger.info('JetStream connection established') + + def on_message(ws: Any, raw: Any) -> None: + queue.put(('message', raw)) + + def on_close(ws: Any, status: Any, msg: Any) -> None: + logger.info(f'JetStream connection closed (status={status}); will reconnect') + queue.put(('done', None)) + + def on_error(ws: Any, err: Any) -> None: + logger.warning(f'JetStream connection error: {err}; will reconnect') + queue.put(('done', None)) + + logger.info(f'Connecting to JetStream at {url}') + app = websocket.WebSocketApp( + url, + on_open=on_open, + on_message=on_message, + on_close=on_close, + on_error=on_error, + ) + runner = gevent.spawn(app.run_forever, ping_interval=PING_INTERVAL_SECONDS, ping_timeout=PING_TIMEOUT_SECONDS) + runner.link(lambda _g: queue.put(('done', None))) + + try: + while True: + kind, raw = queue.get() + if kind == 'done': + return + if not raw: + continue + try: + event = json.loads(raw) + except json.JSONDecodeError as e: + raw_bytes = raw if isinstance(raw, bytes) else str(raw).encode('utf-8', errors='replace') + logger.warning(f'JetStream payload was not valid JSON ({e}); first 200 bytes: {raw_bytes[:200]!r}') + continue + if not isinstance(event, dict): + logger.warning( + f'JetStream payload parsed to non-object JSON (got {type(event).__name__}); skipping' + ) + continue + try: + action_id = self._next_action_id() + except Exception: + # Drop the connection rather than re-minting per message: _gen()'s reconnect + # backoff then throttles retries during a snowflake-id-worker outage. + logger.exception( + 'failed to mint action_id from snowflake-id-worker; dropping connection to back off' + ) + return + action = _event_to_action(event, action_id=action_id) + if action is None: + continue + time_us = event.get('time_us') + if time_us and isinstance(time_us, int) and time_us > 0: + self._last_time_us = time_us + metrics.increment('jetstream_input_stream.events', tags=[f'action_name:{action.action_name}']) + yield NoopAckingContext(action) + finally: + try: + app.close() + except Exception: + logger.info('ignored error while closing JetStream WebSocketApp', exc_info=True) + runner.join(timeout=5) + + +def _event_to_action(event: Dict[str, Any], action_id: int) -> Optional[Action]: + """Wraps a JetStream event as an Osprey action, or returns None if it should be skipped.""" + kind = event.get('kind') + if kind not in ('commit', 'identity'): + return None + time_us = event.get('time_us') + if not isinstance(time_us, int) or time_us <= 0: + return None + try: + timestamp = datetime.fromtimestamp(time_us / 1_000_000, tz=timezone.utc) + except (OverflowError, OSError, ValueError): + # A single out-of-range time_us should skip that event, not tear down the connection. + return None + if kind == 'commit': + commit = event.get('commit') or {} + operation = commit.get('operation') + short = COLLECTION_NAMES.get(commit.get('collection', '')) + if short is None or operation not in ('create', 'update', 'delete'): + return None + action_name = f'{operation}_{short}' + else: + action_name = 'identity' + return Action( + action_id=action_id, + action_name=action_name, + data=event, + timestamp=timestamp, + ) diff --git a/example_atproto_plugins/src/atproto_plugin/py.typed b/example_atproto_plugins/src/atproto_plugin/py.typed new file mode 100644 index 00000000..e69de29b diff --git a/example_atproto_plugins/src/atproto_plugin/register_plugins.py b/example_atproto_plugins/src/atproto_plugin/register_plugins.py new file mode 100644 index 00000000..bcb1dd9f --- /dev/null +++ b/example_atproto_plugins/src/atproto_plugin/register_plugins.py @@ -0,0 +1,25 @@ +from collections.abc import Sequence +from typing import Any, Type + +from osprey.engine.executor.execution_context import Action +from osprey.engine.udf.base import UDFBase +from osprey.worker.adaptor.plugin_manager import hookimpl_osprey +from osprey.worker.lib.config import Config +from osprey.worker.sinks.sink.input_stream import BaseInputStream +from osprey.worker.sinks.utils.acking_contexts import BaseAckingContext + +from atproto_plugin.enrichment_udfs import AtprotoDisplayName, AtprotoHandle +from atproto_plugin.jetstream_input_stream import JetStreamInputStream + + +@hookimpl_osprey +def register_input_stream(config: Config) -> BaseInputStream[BaseAckingContext[Action]]: + endpoint = config.get_optional_str('OSPREY_JETSTREAM_ENDPOINT') + raw_collections = config.get_optional_str('OSPREY_JETSTREAM_WANTED_COLLECTIONS') + wanted = [c.strip() for c in raw_collections.split(',') if c.strip()] if raw_collections else None + return JetStreamInputStream(endpoint=endpoint, wanted_collections=wanted) + + +@hookimpl_osprey +def register_udfs() -> Sequence[Type[UDFBase[Any, Any]]]: + return [AtprotoHandle, AtprotoDisplayName] diff --git a/example_atproto_plugins/tests/__init__.py b/example_atproto_plugins/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/example_atproto_plugins/tests/test_enrichment_udfs.py b/example_atproto_plugins/tests/test_enrichment_udfs.py new file mode 100644 index 00000000..f147aa64 --- /dev/null +++ b/example_atproto_plugins/tests/test_enrichment_udfs.py @@ -0,0 +1,116 @@ +import time +from types import SimpleNamespace +from typing import Iterator +from unittest.mock import MagicMock, patch + +import pytest +import requests +from atproto_plugin import enrichment_udfs +from atproto_plugin.enrichment_udfs import AtprotoDisplayName, AtprotoHandle +from osprey.engine.executor.execution_context import ExpectedUdfException + +PLACEHOLDER_DID = 'did:plc:aaaaaaaaaaaaaaaaaaaaaaaa' +SAMPLE_PROFILE = {'did': PLACEHOLDER_DID, 'handle': 'alice.bsky.social', 'displayName': 'Alice'} + + +@pytest.fixture(autouse=True) +def clear_cache() -> Iterator[None]: + enrichment_udfs._profile_cache.clear() + enrichment_udfs._inflight.clear() + yield + enrichment_udfs._profile_cache.clear() + enrichment_udfs._inflight.clear() + + +def _mock_response(payload: object) -> MagicMock: + response = MagicMock() + response.json.return_value = payload + return response + + +def test_fetch_profile_fetches_and_caches() -> None: + with patch.object(enrichment_udfs._session, 'get', return_value=_mock_response(SAMPLE_PROFILE)) as get: + first = enrichment_udfs._fetch_profile(PLACEHOLDER_DID) + second = enrichment_udfs._fetch_profile(PLACEHOLDER_DID) + assert first == SAMPLE_PROFILE + assert second == first + # Cached: the second lookup does not hit the network. + get.assert_called_once() + + +def test_fetch_profile_propagates_http_error() -> None: + response = MagicMock() + response.raise_for_status.side_effect = requests.HTTPError('400') + with patch.object(enrichment_udfs._session, 'get', return_value=response): + with pytest.raises(requests.HTTPError): + enrichment_udfs._fetch_profile(PLACEHOLDER_DID) + # A failed fetch leaves nothing behind, so the next lookup retries. + assert PLACEHOLDER_DID not in enrichment_udfs._inflight + + +def test_fetch_rides_on_in_progress_request() -> None: + # Simulate another greenlet already fetching this DID: the entry is present in + # _inflight with its result set. A second caller must ride on it, not re-request. + inflight = enrichment_udfs._InflightFetch() + inflight.profile = SAMPLE_PROFILE + inflight.done.set() + enrichment_udfs._inflight[PLACEHOLDER_DID] = inflight + + with patch.object(enrichment_udfs._session, 'get') as get: + result = enrichment_udfs._fetch_profile(PLACEHOLDER_DID) + + assert result == SAMPLE_PROFILE + get.assert_not_called() + + +def test_fetch_reraises_in_progress_error() -> None: + inflight = enrichment_udfs._InflightFetch() + inflight.error = requests.ConnectionError() + inflight.done.set() + enrichment_udfs._inflight[PLACEHOLDER_DID] = inflight + + with patch.object(enrichment_udfs._session, 'get') as get: + with pytest.raises(requests.ConnectionError): + enrichment_udfs._fetch_profile(PLACEHOLDER_DID) + get.assert_not_called() + + +def test_expired_entry_triggers_refetch() -> None: + with patch.object(enrichment_udfs._session, 'get', return_value=_mock_response(SAMPLE_PROFILE)) as get: + enrichment_udfs._fetch_profile(PLACEHOLDER_DID) + assert get.call_count == 1 + + # Jump past the TTL so the cached entry is considered stale and refetched. + stale = time.monotonic() + enrichment_udfs._CACHE_TTL_SECONDS + 1 + with patch.object(enrichment_udfs.time, 'monotonic', return_value=stale): + enrichment_udfs._fetch_profile(PLACEHOLDER_DID) + assert get.call_count == 2 + + +def test_atproto_handle_returns_handle() -> None: + udf = AtprotoHandle.__new__(AtprotoHandle) + with patch.object(enrichment_udfs, '_fetch_profile', return_value=SAMPLE_PROFILE): + assert udf.execute(None, SimpleNamespace(did=PLACEHOLDER_DID)) == 'alice.bsky.social' + + +def test_atproto_display_name_returns_display_name() -> None: + udf = AtprotoDisplayName.__new__(AtprotoDisplayName) + with patch.object(enrichment_udfs, '_fetch_profile', return_value=SAMPLE_PROFILE): + assert udf.execute(None, SimpleNamespace(did=PLACEHOLDER_DID)) == 'Alice' + + +def test_missing_field_raises_expected_udf_exception() -> None: + handle_udf = AtprotoHandle.__new__(AtprotoHandle) + name_udf = AtprotoDisplayName.__new__(AtprotoDisplayName) + with patch.object(enrichment_udfs, '_fetch_profile', return_value={'did': PLACEHOLDER_DID}): + with pytest.raises(ExpectedUdfException): + handle_udf.execute(None, SimpleNamespace(did=PLACEHOLDER_DID)) + with pytest.raises(ExpectedUdfException): + name_udf.execute(None, SimpleNamespace(did=PLACEHOLDER_DID)) + + +def test_transport_error_raises_expected_udf_exception() -> None: + udf = AtprotoHandle.__new__(AtprotoHandle) + with patch.object(enrichment_udfs, '_fetch_profile', side_effect=requests.ConnectionError()): + with pytest.raises(ExpectedUdfException): + udf.execute(None, SimpleNamespace(did=PLACEHOLDER_DID)) diff --git a/example_atproto_plugins/tests/test_jetstream_input_stream.py b/example_atproto_plugins/tests/test_jetstream_input_stream.py new file mode 100644 index 00000000..943afe3b --- /dev/null +++ b/example_atproto_plugins/tests/test_jetstream_input_stream.py @@ -0,0 +1,229 @@ +from unittest import mock + +import pytest +from atproto_plugin.jetstream_input_stream import JetStreamInputStream, _event_to_action + +SAMPLE_POST_COMMIT = { + 'did': 'did:plc:abc123', + 'time_us': 1714500000000000, + 'kind': 'commit', + 'commit': { + 'rev': '3kf...', + 'operation': 'create', + 'collection': 'app.bsky.feed.post', + 'rkey': 'abcdefg', + 'cid': 'bafyrei...', + 'record': { + '$type': 'app.bsky.feed.post', + 'text': 'this is a test post', + 'createdAt': '2024-04-30T12:00:00Z', + }, + }, +} + + +def test_post_commit_passes_through_jetstream_payload(): + action = _event_to_action(SAMPLE_POST_COMMIT, action_id=42) + + assert action is not None + assert action.action_id == 42 + assert action.action_name == 'create_post' + # Action.data is the raw JetStream event — rules read $.did, $.commit.*, etc. + assert action.data is SAMPLE_POST_COMMIT + assert action.data['did'] == 'did:plc:abc123' + assert action.data['commit']['operation'] == 'create' + assert action.data['commit']['collection'] == 'app.bsky.feed.post' + assert action.data['commit']['record']['text'] == 'this is a test post' + + +def test_like_commit_passes_through_jetstream_payload(): + like = { + **SAMPLE_POST_COMMIT, + 'commit': { + **SAMPLE_POST_COMMIT['commit'], + 'collection': 'app.bsky.feed.like', + 'record': {'subject': {'uri': 'at://...', 'cid': 'bafy...'}}, + }, + } + action = _event_to_action(like, action_id=1) + + assert action is not None + assert action.action_name == 'create_like' + assert action.data['commit']['collection'] == 'app.bsky.feed.like' + assert 'text' not in action.data['commit']['record'] + + +def test_delete_commit_passes_through_without_record(): + delete = { + **SAMPLE_POST_COMMIT, + 'commit': { + 'operation': 'delete', + 'collection': 'app.bsky.feed.post', + 'rkey': 'abcdefg', + 'rev': '3kf...', + }, + } + action = _event_to_action(delete, action_id=2) + + assert action is not None + assert action.action_name == 'delete_post' + assert action.data['commit']['operation'] == 'delete' + assert 'record' not in action.data['commit'] + + +def test_identity_event_passes_through(): + identity = { + 'did': 'did:plc:xyz', + 'time_us': 1714500000000001, + 'kind': 'identity', + 'identity': {'did': 'did:plc:xyz', 'handle': 'someone.bsky.social', 'seq': 100}, + } + action = _event_to_action(identity, action_id=3) + + assert action is not None + assert action.action_name == 'identity' + assert action.data['identity']['handle'] == 'someone.bsky.social' + + +def test_account_event_is_skipped(): + account = { + 'did': 'did:plc:xyz', + 'time_us': 1714500000000002, + 'kind': 'account', + 'account': {'active': True, 'did': 'did:plc:xyz', 'seq': 100}, + } + assert _event_to_action(account, action_id=4) is None + + +def test_event_without_time_us_is_skipped(): + bad = {k: v for k, v in SAMPLE_POST_COMMIT.items() if k != 'time_us'} + assert _event_to_action(bad, action_id=5) is None + + +def test_event_with_zero_time_us_is_skipped(): + zero_time = {**SAMPLE_POST_COMMIT, 'time_us': 0} + assert _event_to_action(zero_time, action_id=5) is None + + +def test_event_with_negative_time_us_is_skipped(): + neg_time = {**SAMPLE_POST_COMMIT, 'time_us': -1} + assert _event_to_action(neg_time, action_id=5) is None + + +def test_event_with_out_of_range_time_us_is_skipped(): + out_of_range = {**SAMPLE_POST_COMMIT, 'time_us': 10**30} + assert _event_to_action(out_of_range, action_id=5) is None + + +def test_unknown_kind_is_skipped(): + weird = {'did': 'did:plc:xyz', 'time_us': 1, 'kind': 'something-new'} + assert _event_to_action(weird, action_id=6) is None + + +@pytest.mark.parametrize( + 'collection,operation,expected_name', + [ + ('app.bsky.feed.post', 'create', 'create_post'), + ('app.bsky.feed.post', 'update', 'update_post'), + ('app.bsky.feed.post', 'delete', 'delete_post'), + ('app.bsky.feed.like', 'create', 'create_like'), + ('app.bsky.feed.like', 'delete', 'delete_like'), + ('app.bsky.feed.repost', 'create', 'create_repost'), + ('app.bsky.feed.repost', 'delete', 'delete_repost'), + ('app.bsky.graph.follow', 'create', 'create_follow'), + ('app.bsky.graph.follow', 'delete', 'delete_follow'), + ('app.bsky.actor.profile', 'update', 'update_profile'), + ], +) +def test_commit_action_name_combines_operation_and_collection(collection, operation, expected_name): + event = { + **SAMPLE_POST_COMMIT, + 'commit': {**SAMPLE_POST_COMMIT['commit'], 'collection': collection, 'operation': operation}, + } + action = _event_to_action(event, action_id=1) + assert action is not None + assert action.action_name == expected_name + + +def test_commit_for_unmapped_collection_is_skipped(): + event = { + **SAMPLE_POST_COMMIT, + 'commit': {**SAMPLE_POST_COMMIT['commit'], 'collection': 'app.bsky.graph.starterpack'}, + } + assert _event_to_action(event, action_id=1) is None + + +def test_commit_with_unexpected_operation_is_skipped(): + event = { + **SAMPLE_POST_COMMIT, + 'commit': {**SAMPLE_POST_COMMIT['commit'], 'operation': 'wat'}, + } + assert _event_to_action(event, action_id=1) is None + + +def test_build_url_includes_wanted_collections(): + stream = JetStreamInputStream( + endpoint='wss://example.com/sub', + wanted_collections=['app.bsky.feed.post', 'app.bsky.feed.like'], + ) + url = stream._build_url() + assert 'wss://example.com/sub' in url + assert 'wantedCollections=app.bsky.feed.post' in url + assert 'wantedCollections=app.bsky.feed.like' in url + assert 'cursor=' not in url + + +def test_build_url_includes_cursor_when_last_time_us_set(): + stream = JetStreamInputStream( + endpoint='wss://example.com/sub', + wanted_collections=['app.bsky.feed.post'], + ) + stream._last_time_us = 1714500000000000 + url = stream._build_url() + assert 'wss://example.com/sub' in url + assert 'wantedCollections=app.bsky.feed.post' in url + assert 'cursor=1714500000000000' in url + + +def test_stream_one_connection_handles_malformed_json_and_gracefully_closes(): + class FakeWebSocketApp: + def __init__(self, url, on_open=None, on_message=None, on_close=None, on_error=None): + self.url = url + self.on_open = on_open + self.on_message = on_message + self.on_close = on_close + self.closed = False + + def run_forever(self, **kwargs): + self.on_open(self) + self.on_message(self, b'not-json') + self.on_message( + self, + b'{"did": "did:plc:x", "time_us": 1714500000000000, "kind": "commit",' + b' "commit": {"operation": "create", "collection": "app.bsky.feed.post",' + b' "rkey": "abc", "rev": "3kf...", "cid": "bafy...",' + b' "record": {"$type": "app.bsky.feed.post", "text": "hello"}}}', + ) + self.on_close(self, 1000, 'normal') + + def close(self): + self.closed = True + + captured = {} + + def factory(*args, **kwargs): + captured['app'] = FakeWebSocketApp(*args, **kwargs) + return captured['app'] + + stream = JetStreamInputStream(reconnect_seconds=0.1) + # Pre-fill so _next_action_id() doesn't hit the snowflake-id-worker. + stream._snowflake_buffer = [12345] + + with mock.patch('websocket.WebSocketApp', factory): + actions = list(stream._stream_one_connection('wss://example.com/sub')) + + assert len(actions) == 1 + assert actions[0]._item.action_id == 12345 + assert actions[0]._item.action_name == 'create_post' + assert actions[0]._item.data['did'] == 'did:plc:x' + assert captured['app'].closed diff --git a/example_atproto_rules/config/labels.yaml b/example_atproto_rules/config/labels.yaml new file mode 100644 index 00000000..799d5495 --- /dev/null +++ b/example_atproto_rules/config/labels.yaml @@ -0,0 +1,5 @@ +labels: + test-poster: + valid_for: [UserId] + connotation: neutral + description: ATProto user who posted a record containing the word 'test' diff --git a/example_atproto_rules/config/ui_config.yaml b/example_atproto_rules/config/ui_config.yaml new file mode 100644 index 00000000..0b0868cb --- /dev/null +++ b/example_atproto_rules/config/ui_config.yaml @@ -0,0 +1,10 @@ +ui_config: + default_summary_features: + - actions: ['*'] + features: [UserId] + - actions: ['create_post', 'update_post'] + features: [PostText, Rkey] + - actions: ['delete_post'] + features: [Rkey] + - actions: ['create_like', 'delete_like', 'create_repost', 'delete_repost', 'create_follow', 'delete_follow'] + features: [Subject, Rkey] diff --git a/example_atproto_rules/main.sml b/example_atproto_rules/main.sml new file mode 100644 index 00000000..49d08c01 --- /dev/null +++ b/example_atproto_rules/main.sml @@ -0,0 +1,9 @@ +Import( + rules=[ + 'models/base.sml', + 'models/record/base.sml', + 'models/record/post.sml', + ], +) + +Require(rule='rules/index.sml') diff --git a/example_atproto_rules/models/base.sml b/example_atproto_rules/models/base.sml new file mode 100644 index 00000000..cafc322b --- /dev/null +++ b/example_atproto_rules/models/base.sml @@ -0,0 +1,20 @@ +ActionName = GetActionName() + +UserId: Entity[str] = EntityJson( + type='UserId', + path='$.did', + required=False, +) + +OperationKind: Optional[str] = JsonData( + path='$.commit.operation', + required=False, +) + +IsOperation = OperationKind != None + +Second: int = 1 +Minute: int = Second * 60 +Hour: int = Minute * 60 +Day: int = Hour * 24 +Week: int = Day * 7 diff --git a/example_atproto_rules/models/enrichment.sml b/example_atproto_rules/models/enrichment.sml new file mode 100644 index 00000000..c7c6b26b --- /dev/null +++ b/example_atproto_rules/models/enrichment.sml @@ -0,0 +1,20 @@ +# Opt-in profile enrichment. +# +# JetStream carries only the actor DID, so these features resolve it to a handle +# and display name via a per-event call to the Bluesky public API (cached per DID, +# absent when the account can't be resolved or the API rate-limits). +# +# Because that is an external dependency you don't want in a load test, this file +# is NOT imported by main.sml by default. To enable it for demos, add +# 'models/enrichment.sml' to main.sml's imports and add Handle / DisplayName to +# config/ui_config.yaml. See the plugin README for how to extend it with more +# profile fields (account age, follower counts, existing labels). + +Did: str = JsonData( + path='$.did', + required=False, +) + +Handle: str = AtprotoHandle(did=Did) + +DisplayName: str = AtprotoDisplayName(did=Did) diff --git a/example_atproto_rules/models/record/base.sml b/example_atproto_rules/models/record/base.sml new file mode 100644 index 00000000..f768b31a --- /dev/null +++ b/example_atproto_rules/models/record/base.sml @@ -0,0 +1,37 @@ +Import(rules=['models/base.sml']) + +IsCreate = OperationKind == 'create' +IsUpdate = OperationKind == 'update' +IsDelete = OperationKind == 'delete' + +Collection: str = JsonData( + path='$.commit.collection', + required=False, +) + +Rkey: str = JsonData( + path='$.commit.rkey', + required=False, +) + +Cid: str = JsonData( + path='$.commit.cid', + required=False, +) + +# Some collections (like, repost) have `record.subject = {uri, cid}`; others (follow) +# have `record.subject = "did:plc:..."`. Prefer the URI when present, else fall back to +# the raw subject value (coerced to str — for follows that's the DID; for the dict shape +# the URI path resolves first so the coerced repr is never used). +SubjectUri: Optional[str] = JsonData( + path='$.commit.record.subject.uri', + required=False, +) + +SubjectRaw: Optional[str] = JsonData( + path='$.commit.record.subject', + required=False, + coerce_type=True, +) + +Subject: str = ResolveOptional(optional_value=SubjectUri, default_value=SubjectRaw) diff --git a/example_atproto_rules/models/record/post.sml b/example_atproto_rules/models/record/post.sml new file mode 100644 index 00000000..405eda0c --- /dev/null +++ b/example_atproto_rules/models/record/post.sml @@ -0,0 +1,6 @@ +Import(rules=['models/base.sml']) + +PostText: str = JsonData( + path='$.commit.record.text', + required=False, +) diff --git a/example_atproto_rules/rules/index.sml b/example_atproto_rules/rules/index.sml new file mode 100644 index 00000000..113cb722 --- /dev/null +++ b/example_atproto_rules/rules/index.sml @@ -0,0 +1,6 @@ +Import(rules=['models/base.sml']) + +Require( + rule='rules/record/index.sml', + require_if=IsOperation, +) diff --git a/example_atproto_rules/rules/record/index.sml b/example_atproto_rules/rules/record/index.sml new file mode 100644 index 00000000..0fabb9d5 --- /dev/null +++ b/example_atproto_rules/rules/record/index.sml @@ -0,0 +1,11 @@ +Import( + rules=[ + 'models/base.sml', + 'models/record/base.sml', + ], +) + +Require( + rule='rules/record/post/index.sml', + require_if=(IsCreate or IsUpdate) and Collection == 'app.bsky.feed.post', +) diff --git a/example_atproto_rules/rules/record/post/index.sml b/example_atproto_rules/rules/record/post/index.sml new file mode 100644 index 00000000..1d84f15e --- /dev/null +++ b/example_atproto_rules/rules/record/post/index.sml @@ -0,0 +1,9 @@ +Import( + rules=[ + 'models/base.sml', + 'models/record/base.sml', + 'models/record/post.sml', + ], +) + +Require(rule='rules/record/post/post_contains_test.sml') diff --git a/example_atproto_rules/rules/record/post/post_contains_test.sml b/example_atproto_rules/rules/record/post/post_contains_test.sml new file mode 100644 index 00000000..886f068c --- /dev/null +++ b/example_atproto_rules/rules/record/post/post_contains_test.sml @@ -0,0 +1,21 @@ +Import( + rules=[ + 'models/base.sml', + 'models/record/base.sml', + 'models/record/post.sml', + ], +) + +PostContainsTestRule = Rule( + when_all=[ + TextContains(text=PostText, phrase='test'), + ], + description='ATProto post contains the word "test"', +) + +WhenRules( + rules_any=[PostContainsTestRule], + then=[ + LabelAdd(entity=UserId, label='test-poster'), + ], +) diff --git a/osprey_worker/Dockerfile b/osprey_worker/Dockerfile index afaa68d7..76878a3b 100644 --- a/osprey_worker/Dockerfile +++ b/osprey_worker/Dockerfile @@ -40,10 +40,11 @@ ADD osprey_rpc/pyproject.toml /osprey/osprey_rpc/pyproject.toml ADD osprey_worker/pyproject.toml /osprey/osprey_worker/pyproject.toml ADD osprey_async_worker/pyproject.toml /osprey/osprey_async_worker/pyproject.toml ADD example_plugins/pyproject.toml /osprey/example_plugins/pyproject.toml +ADD example_atproto_plugins/pyproject.toml /osprey/example_atproto_plugins/pyproject.toml # Create minimal package structure required by uv -RUN mkdir -p /osprey/osprey_worker /osprey/osprey_rpc /osprey/example_plugins/src && \ - touch /osprey/osprey_worker/__init__.py /osprey/osprey_rpc/__init__.py /osprey/example_plugins/src/__init__.py +RUN mkdir -p /osprey/osprey_worker /osprey/osprey_rpc /osprey/example_plugins/src /osprey/example_atproto_plugins/src/atproto_plugin && \ + touch /osprey/osprey_worker/__init__.py /osprey/osprey_rpc/__init__.py /osprey/example_plugins/src/__init__.py /osprey/example_atproto_plugins/src/atproto_plugin/__init__.py # Install dependencies first (this layer will be cached when only source code changes). # Exclude osprey_async_worker: the asyncio worker is experimental and lives in its own @@ -59,9 +60,11 @@ RUN . .venv/bin/activate && update-tld-names # Add source code after dependencies are installed (osprey_async_worker is # intentionally omitted — it is not installed in this image). ADD example_rules /osprey/example_rules +ADD example_atproto_rules /osprey/example_atproto_rules ADD osprey_worker /osprey/osprey_worker ADD osprey_rpc /osprey/osprey_rpc ADD example_plugins /osprey/example_plugins +ADD example_atproto_plugins /osprey/example_atproto_plugins COPY entrypoint.sh /osprey/entrypoint.sh diff --git a/pyproject.toml b/pyproject.toml index cf10a8b9..36061921 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -125,6 +125,7 @@ osprey_rpc = { workspace = true } osprey_worker = { workspace = true } osprey_async_worker = { workspace = true } example_plugins = { workspace = true } +example_atproto_plugins = { workspace = true } [tool.uv.workspace] members = [ @@ -132,6 +133,7 @@ members = [ "osprey_worker", "osprey_async_worker", "example_plugins", + "example_atproto_plugins", ] @@ -164,10 +166,10 @@ ignore = [ ] [tool.ruff.lint.isort] -known-first-party = ["osprey_worker", "osprey_async_worker", "osprey_rpc", "example_plugins"] +known-first-party = ["osprey_worker", "osprey_async_worker", "osprey_rpc", "example_plugins", "example_atproto_plugins"] [tool.fawltydeps] -code = ["osprey_worker/src", "osprey_async_worker/src", "osprey_rpc/src", "example_plugins/src"] +code = ["osprey_worker/src", "osprey_async_worker/src", "osprey_rpc/src", "example_plugins/src", "example_atproto_plugins/src"] deps = ["pyproject.toml"] ignore_unused = [ # Type stubs: used by mypy, never imported directly @@ -222,7 +224,7 @@ ignore_unused = [ ] [tool.pytest.ini_options] -testpaths = ["osprey_worker", "example_plugins"] +testpaths = ["osprey_worker", "example_plugins", "example_atproto_plugins"] asyncio_mode = "auto" [tool.mypy] @@ -237,6 +239,7 @@ mypy_path = [ "osprey_worker/src", "osprey_async_worker/src", "example_plugins/src", + "example_atproto_plugins/src", ] # Strict mode includes the following flags. When these are all True, they can be replaced with strict mode. diff --git a/run-atproto.sh b/run-atproto.sh new file mode 100755 index 00000000..ba299324 --- /dev/null +++ b/run-atproto.sh @@ -0,0 +1,5 @@ +#!/bin/bash +# Run Osprey against the live Bluesky JetStream firehose by stacking +# docker-compose.atproto.yaml on top of the main compose file. +set -e +exec docker compose -f docker-compose.yaml -f docker-compose.atproto.yaml "${@:-up}" diff --git a/uv.lock b/uv.lock index 94209888..047fba96 100644 --- a/uv.lock +++ b/uv.lock @@ -16,6 +16,7 @@ resolution-markers = [ [manifest] members = [ + "example-atproto-plugins", "example-plugins", "osprey-async-worker", "osprey-rpc", @@ -528,6 +529,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/39/74/99f5d33b4ec2eed9d19d975ab61e3eeddc71c8780a3e46ec9e0f17095451/envier-0.5.2-py3-none-any.whl", hash = "sha256:65099cf3aa9b3b3b4b92db2f7d29e2910672e085b76f7e587d2167561a834add", size = 10050, upload-time = "2024-07-05T10:34:34.247Z" }, ] +[[package]] +name = "example-atproto-plugins" +version = "0.1.0" +source = { editable = "example_atproto_plugins" } +dependencies = [ + { name = "pluggy" }, + { name = "websocket-client" }, +] + +[package.metadata] +requires-dist = [ + { name = "pluggy", specifier = "==1.5.0" }, + { name = "websocket-client", specifier = "==1.8.0" }, +] + [[package]] name = "example-plugins" version = "0.1.0" @@ -2722,6 +2738,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/41/52/e465037f5375f43533d1a80b6923955201596a99142ed524d77b571a1418/wcwidth-0.7.0-py3-none-any.whl", hash = "sha256:5d69154c429a82910e241c738cd0e2976fac8a2dd47a1a805f4afed1c0f136f2", size = 110825, upload-time = "2026-05-02T16:04:11.033Z" }, ] +[[package]] +name = "websocket-client" +version = "1.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e6/30/fba0d96b4b5fbf5948ed3f4681f7da2f9f64512e1d303f94b4cc174c24a5/websocket_client-1.8.0.tar.gz", hash = "sha256:3239df9f44da632f96012472805d40a23281a991027ce11d2f45a6f24ac4c3da", size = 54648, upload-time = "2024-04-23T22:16:16.976Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/84/44687a29792a70e111c5c477230a72c4b957d88d16141199bf9acb7537a3/websocket_client-1.8.0-py3-none-any.whl", hash = "sha256:17b44cc997f5c498e809b22cdf2d9c7a9e71c02c8cc2b6c56e7c2d1239bfa526", size = 58826, upload-time = "2024-04-23T22:16:14.422Z" }, +] + [[package]] name = "werkzeug" version = "1.0.1"