diff --git a/example_atproto_plugins/README.md b/example_atproto_plugins/README.md index 93561d18..116e8ff8 100644 --- a/example_atproto_plugins/README.md +++ b/example_atproto_plugins/README.md @@ -6,7 +6,7 @@ A sample Osprey plugin that consumes ATProto's [JetStream](https://docs.bsky.app - 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 **only the input stream**. 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. +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 @@ -55,20 +55,73 @@ The JetStream JSON event is passed through unchanged as the Action's `data` dict "did": "did:plc:...", "time_us": ..., "kind": "identity", - "identity": {"did": "...", "handle": "...", "seq": ..., "time": "..."} + "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. The cache holds up to 100,000 DIDs and evicts least-recently-used entries past that, so at very high DID cardinality an evicted DID is re-fetched the next time it appears. + +**It is off by default.** Each cache miss (a DID not seen within the last hour, or evicted once the cache is full) 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. + +Being "off" here means not invoked, not unregistered. The plugin registers `AtprotoHandle` / `AtprotoDisplayName` whenever it is installed (the same as every other example UDF), so the SML compiler can resolve them, but registration is inert and makes no API calls. A `getProfile` lookup happens only when a rule references a UDF, which happens only through `models/enrichment.sml`. So the import list in `main.sml` is the only switch, and by default it is off. Registration can't be gated on the import instead, since the compiler has to know the UDF exists before it can resolve the reference. + +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. `PostText` for `create_post`, `IdentityHandle` for `identity`, `Subject` for like / repost / follow events. Add new entries there to expose more fields without touching rule code. +`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. -- **No event enrichment.** JetStream only carries what's in the commit itself; rulesets that depend on handle / profile / account age (such as much of [atproto-ruleset](https://github.com/haileyok/atproto-ruleset)) are fed by a separate enrichment pipeline, not JetStream directly. This plugin emits JetStream-native paths ($.did, $.commit.collection, etc.); enrichment-fed rulesets would need an enrichment service in front of this one or a different plugin. +- **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/src/atproto_plugin/enrichment_udfs.py b/example_atproto_plugins/src/atproto_plugin/enrichment_udfs.py new file mode 100644 index 00000000..4e8c4e14 --- /dev/null +++ b/example_atproto_plugins/src/atproto_plugin/enrichment_udfs.py @@ -0,0 +1,167 @@ +"""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') + # getProfile is untrusted JSON, so guard the type: a truthy non-string + # would otherwise cross the str return boundary. Treat it like a missing + # field and skip. + if not isinstance(handle, str) or 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 isinstance(display_name, str) or not display_name: + raise ExpectedUdfException() + return display_name diff --git a/example_atproto_plugins/src/atproto_plugin/register_plugins.py b/example_atproto_plugins/src/atproto_plugin/register_plugins.py index dc6b0736..bcb1dd9f 100644 --- a/example_atproto_plugins/src/atproto_plugin/register_plugins.py +++ b/example_atproto_plugins/src/atproto_plugin/register_plugins.py @@ -1,9 +1,14 @@ +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 @@ -13,3 +18,8 @@ def register_input_stream(config: Config) -> BaseInputStream[BaseAckingContext[A 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/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_rules/config/ui_config.yaml b/example_atproto_rules/config/ui_config.yaml index 8b1d6036..0b0868cb 100644 --- a/example_atproto_rules/config/ui_config.yaml +++ b/example_atproto_rules/config/ui_config.yaml @@ -8,5 +8,3 @@ ui_config: features: [Rkey] - actions: ['create_like', 'delete_like', 'create_repost', 'delete_repost', 'create_follow', 'delete_follow'] features: [Subject, Rkey] - - actions: ['identity'] - features: [IdentityHandle] diff --git a/example_atproto_rules/main.sml b/example_atproto_rules/main.sml index 63227315..49d08c01 100644 --- a/example_atproto_rules/main.sml +++ b/example_atproto_rules/main.sml @@ -1,7 +1,6 @@ Import( rules=[ 'models/base.sml', - 'models/identity.sml', 'models/record/base.sml', 'models/record/post.sml', ], diff --git a/example_atproto_rules/models/enrichment.sml b/example_atproto_rules/models/enrichment.sml new file mode 100644 index 00000000..dc7d5ea5 --- /dev/null +++ b/example_atproto_rules/models/enrichment.sml @@ -0,0 +1,27 @@ +# 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). +# +# Note the difference between a UDF being registered and being invoked. The plugin +# registers AtprotoHandle / AtprotoDisplayName whenever it is installed, the same +# as every other example UDF, so the compiler can resolve them. That registration +# is inert: it makes no API calls. The getProfile lookups happen only when a rule +# references the UDFs, which happens only through this file. So leaving it out of +# main.sml's imports is what keeps enrichment off; there is no separate switch. + +Did: str = JsonData( + path='$.did', + required=False, +) + +Handle: str = AtprotoHandle(did=Did) + +DisplayName: str = AtprotoDisplayName(did=Did) diff --git a/example_atproto_rules/models/identity.sml b/example_atproto_rules/models/identity.sml deleted file mode 100644 index 73be9cce..00000000 --- a/example_atproto_rules/models/identity.sml +++ /dev/null @@ -1,6 +0,0 @@ -Import(rules=['models/base.sml']) - -IdentityHandle: str = JsonData( - path='$.identity.handle', - required=False, -)