feat(atproto): resolve DID to handle and display name via enrichment UDFs

.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

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

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

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

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 `<operation>_<short>` 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.
+
+**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.

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"

example_atproto_plugins/src/atproto_plugin/enrichment_udfs.py

@@ -0,0 +1,99 @@
+"""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.
+
+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.
+"""
+
+from collections import OrderedDict
+from threading import 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 = 10_000
+
+_session = requests.Session()
+# did -> profile dict (the raw getProfile response).
+_profile_cache: 'OrderedDict[str, dict[str, Any]]' = OrderedDict()
+_cache_lock = Lock()
+
+
+def _fetch_profile(did: str) -> dict[str, Any]:
+    """Return the getProfile response for a DID, hitting the public API on cache miss.
+
+    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_cache.move_to_end(did)
+            return cached
+
+    # Fetch outside the lock so a slow request does not block other greenlets.
+    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')
+
+    with _cache_lock:
+        _profile_cache[did] = profile
+        _profile_cache.move_to_end(did)
+        while len(_profile_cache) > _CACHE_MAX_SIZE:
+            _profile_cache.popitem(last=False)
+    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