|
| 1 | +from workers import WorkerEntrypoint, Response, DurableObject |
| 2 | +import js |
| 3 | +import json |
| 4 | +import time |
| 5 | +from pyodide.ffi import create_proxy |
| 6 | +from urllib.parse import urlparse |
| 7 | + |
| 8 | + |
| 9 | +class BlueskyFirehoseConsumer(DurableObject): |
| 10 | + """Durable Object that maintains a persistent WebSocket connection to Bluesky Jetstream.""" |
| 11 | + |
| 12 | + def __init__(self, state, env): |
| 13 | + super().__init__(state, env) |
| 14 | + self.websocket = None |
| 15 | + self.connected = False |
| 16 | + self.last_print_time = 0 # Track last time we printed a post |
| 17 | + |
| 18 | + async def fetch(self, request): |
| 19 | + """Handle incoming requests to the Durable Object.""" |
| 20 | + # If we're not connected then make sure we start a connection. |
| 21 | + if not self.connected: |
| 22 | + await self._schedule_next_alarm() |
| 23 | + await self._connect_to_jetstream() |
| 24 | + |
| 25 | + url = urlparse(request.url) |
| 26 | + path = url.path |
| 27 | + |
| 28 | + if path == "/status": |
| 29 | + status = "connected" if self.connected else "disconnected" |
| 30 | + return Response(f"Firehose status: {status}") |
| 31 | + else: |
| 32 | + return Response("Available endpoints: /status") |
| 33 | + |
| 34 | + async def alarm(self): |
| 35 | + """Handle alarm events - used to ensure that the DO stays alive and connected""" |
| 36 | + print("Alarm triggered - making sure we are connected to jetstream...") |
| 37 | + if not self.connected: |
| 38 | + await self._connect_to_jetstream() |
| 39 | + else: |
| 40 | + print("Already connected, skipping reconnection") |
| 41 | + |
| 42 | + # Schedule the next alarm to keep the DO alive |
| 43 | + await self._schedule_next_alarm() |
| 44 | + |
| 45 | + async def _schedule_next_alarm(self): |
| 46 | + """Schedule the next alarm to run in 1 minute to keep the DO alive.""" |
| 47 | + # Schedule alarm for 1 minute from now, overwriting any existing alarms |
| 48 | + next_alarm_time = int(time.time() * 1000) + 60000 |
| 49 | + return await self.ctx.storage.setAlarm(next_alarm_time) |
| 50 | + |
| 51 | + async def _on_open(self, event): |
| 52 | + """Handle WebSocket open event.""" |
| 53 | + self.connected = True |
| 54 | + print("Connected to Bluesky Jetstream firehose!") |
| 55 | + print("Filtering for: app.bsky.feed.post (post events, rate limited to 1/sec)") |
| 56 | + # Ensure alarm is set when we connect |
| 57 | + await self._schedule_next_alarm() |
| 58 | + |
| 59 | + def _on_message(self, event): |
| 60 | + """Handle incoming WebSocket messages.""" |
| 61 | + try: |
| 62 | + # Parse the JSON message |
| 63 | + data = json.loads(event.data) |
| 64 | + |
| 65 | + # Store the timestamp for resumption on reconnect |
| 66 | + if time_us := data.get("time_us"): |
| 67 | + # Store the timestamp asynchronously |
| 68 | + self.ctx.storage.kv.put("last_event_timestamp", time_us) |
| 69 | + |
| 70 | + # Jetstream sends different event types |
| 71 | + # We're interested in 'commit' events which contain posts |
| 72 | + if data.get("kind") != "commit": |
| 73 | + return |
| 74 | + |
| 75 | + commit = data.get("commit", {}) |
| 76 | + collection = commit.get("collection") |
| 77 | + |
| 78 | + # Filter for post events |
| 79 | + if collection != "app.bsky.feed.post": |
| 80 | + return |
| 81 | + |
| 82 | + # Rate limiting: only print at most 1 per second |
| 83 | + current_time = time.time() |
| 84 | + if current_time - self.last_print_time >= 1.0: |
| 85 | + record = commit.get("record", {}) |
| 86 | + print("Post record", record) |
| 87 | + |
| 88 | + # Update last print time |
| 89 | + self.last_print_time = current_time |
| 90 | + |
| 91 | + except Exception as e: |
| 92 | + print(f"Error processing message: {e}") |
| 93 | + |
| 94 | + def _on_error(self, event): |
| 95 | + """Handle WebSocket error event.""" |
| 96 | + print(f"WebSocket error: {event}") |
| 97 | + self.connected = False |
| 98 | + self.ctx.abort("WebSocket error occurred") |
| 99 | + |
| 100 | + async def _on_close(self, event): |
| 101 | + """Handle WebSocket close event.""" |
| 102 | + print(f"WebSocket closed: code={event.code}, reason={event.reason}") |
| 103 | + self.connected = False |
| 104 | + self.ctx.abort("WebSocket closed") |
| 105 | + |
| 106 | + async def _connect_to_jetstream(self): |
| 107 | + """Connect to the Bluesky Jetstream WebSocket and start consuming events.""" |
| 108 | + # Get the last event timestamp from storage to resume from the right position |
| 109 | + last_timestamp = self.ctx.storage.kv.get("last_event_timestamp") |
| 110 | + |
| 111 | + # Jetstream endpoint - we'll filter for posts |
| 112 | + # Using wantedCollections parameter to only get post events |
| 113 | + jetstream_url = "wss://jetstream2.us-east.bsky.network/subscribe?wantedCollections=app.bsky.feed.post" |
| 114 | + |
| 115 | + # If we have a last timestamp, add it to resume from that point |
| 116 | + if last_timestamp: |
| 117 | + jetstream_url += f"&cursor={last_timestamp}" |
| 118 | + print( |
| 119 | + f"Connecting to Bluesky Jetstream at {jetstream_url} (resuming from timestamp: {last_timestamp})" |
| 120 | + ) |
| 121 | + else: |
| 122 | + print( |
| 123 | + f"Connecting to Bluesky Jetstream at {jetstream_url} (starting fresh)" |
| 124 | + ) |
| 125 | + |
| 126 | + # Create WebSocket using JS FFI |
| 127 | + ws = js.WebSocket.new(jetstream_url) |
| 128 | + self.websocket = ws |
| 129 | + |
| 130 | + # Attach event handlers |
| 131 | + # |
| 132 | + # Note that ordinarily proxies need to be destroyed once they are no longer used. |
| 133 | + # However, in this Durable Object context, the WebSocket and its event listeners |
| 134 | + # persist for the lifetime of the Durable Object, so we don't explicitly destroy |
| 135 | + # the proxies here. When the websocket connection closes, the Durable Object |
| 136 | + # is restarted which destroys these proxies. |
| 137 | + # |
| 138 | + # In the future, we plan to provide support for native Python websocket APIs which |
| 139 | + # should eliminate the need for proxy wrappers. |
| 140 | + ws.addEventListener("open", create_proxy(self._on_open)) |
| 141 | + ws.addEventListener("message", create_proxy(self._on_message)) |
| 142 | + ws.addEventListener("error", create_proxy(self._on_error)) |
| 143 | + ws.addEventListener("close", create_proxy(self._on_close)) |
| 144 | + |
| 145 | + |
| 146 | +class Default(WorkerEntrypoint): |
| 147 | + """Main worker entry point that routes requests to the Durable Object.""" |
| 148 | + |
| 149 | + async def fetch(self, request): |
| 150 | + # Get the Durable Object namespace from the environment |
| 151 | + namespace = self.env.BLUESKY_FIREHOSE |
| 152 | + |
| 153 | + # Use a fixed ID so we always connect to the same Durable Object instance |
| 154 | + # This ensures we maintain a single persistent connection |
| 155 | + id = namespace.idFromName("bluesky-consumer") |
| 156 | + stub = namespace.get(id) |
| 157 | + |
| 158 | + # Forward the request to the Durable Object |
| 159 | + return await stub.fetch(request) |
0 commit comments