feat: sniff handshake and publish Kafka events in livepeer proxy

Intercept the WebSocket handshake in _proxy_ws to extract manifest_id
and user_id. Publish websocket_connected (after handshake) and
websocket_disconnected (on close) events with connection_info,
matching fal_app.py event parity.

Signed-off-by: emranemran <emran@livepeer.org>
Signed-off-by: emranemran <emran.mah@gmail.com>

Author: emranemran

src/scope/cloud/livepeer_fal_app.py

@@ -271,17 +271,58 @@ def _build_runner_command() -> list[str]:
     ]
 
 
-async def _proxy_ws(client_ws: WebSocket) -> None:
+async def _proxy_ws(client_ws: WebSocket, metadata: dict | None = None) -> None:
     """Connect to the local runner and proxy traffic bidirectionally.
 
+    If *metadata* is provided, the initial handshake messages are sniffed
+    (best-effort) to extract ``manifest_id`` and ``user_id`` before the
+    bidirectional proxy loop begins.
+
     Raises WebSocketDisconnect if the client disconnects.
     Returns normally if the runner connection drops.
     """
+    import json
+
     import websockets
     from websockets.exceptions import ConnectionClosed
 
     async with websockets.connect(RUNNER_LOCAL_WS_URL) as runner_ws:
-
+        # --- Sniff the two-message handshake (best-effort) ---
+        # 1. Runner → Client: {"type": "ready", "connection_id": "..."}
+        ready_msg = await runner_ws.recv()
+        if isinstance(ready_msg, bytes):
+            await client_ws.send_bytes(ready_msg)
+        else:
+            if metadata is not None:
+                try:
+                    parsed = json.loads(ready_msg)
+                    metadata["connection_id"] = parsed.get("connection_id")
+                except (json.JSONDecodeError, AttributeError):
+                    pass
+            await client_ws.send_text(ready_msg)
+
+        # 2. Client → Runner: job_info with manifest_id, params.daydream_user_id
+        job_msg = await client_ws.receive()
+        job_msg_type = job_msg.get("type")
+        if job_msg_type == "websocket.receive":
+            text_data = job_msg.get("text")
+            bytes_data = job_msg.get("bytes")
+            if text_data is not None:
+                if metadata is not None:
+                    try:
+                        parsed = json.loads(text_data)
+                        metadata["manifest_id"] = parsed.get("manifest_id")
+                        params = parsed.get("params") or {}
+                        metadata["user_id"] = params.get("daydream_user_id")
+                    except (json.JSONDecodeError, AttributeError):
+                        pass
+                await runner_ws.send(text_data)
+            elif bytes_data is not None:
+                await runner_ws.send(bytes_data)
+        elif job_msg_type == "websocket.disconnect":
+            raise WebSocketDisconnect()
+
+        # --- Bidirectional proxy loop ---
         async def client_to_runner() -> None:
             while True:
                 message = await client_ws.receive()
@@ -437,6 +478,33 @@ async def websocket_handler(self, client_ws: WebSocket) -> None:
 
         await client_ws.accept()
 
+        # Initialize Kafka publisher (lazy, once per process)
+        global kafka_publisher
+        if kafka_publisher is None:
+            kafka_publisher = KafkaPublisher()
+            await kafka_publisher.start()
+
+        connection_start_time = time.time()
+        metadata: dict = {}
+        connected_event_published = False
+
+        import json
+
+        fal_log_labels_raw = os.getenv("FAL_LOG_LABELS", "unknown")
+        try:
+            fal_log_labels = json.loads(fal_log_labels_raw)
+        except (json.JSONDecodeError, TypeError):
+            fal_log_labels = fal_log_labels_raw
+
+        connection_info = {
+            "gpu_type": LivepeerScopeApp.machine_type,
+            "fal_region": os.getenv("NOMAD_DC", "unknown"),
+            "fal_runner_id": os.getenv(
+                "FAL_JOB_ID", os.getenv("FAL_RUNNER_ID", "unknown")
+            ),
+            "fal_log_labels": fal_log_labels,
+        }
+
         # Ensure any previous session data is cleaned up
         event = _get_cleanup_event()
         await event.wait()
@@ -448,7 +516,7 @@ async def websocket_handler(self, client_ws: WebSocket) -> None:
             while True:
                 print(f"Connecting proxy to runner websocket at {RUNNER_LOCAL_WS_URL}")
                 try:
-                    await _proxy_ws(client_ws)
+                    await _proxy_ws(client_ws, metadata=metadata)
                 except (
                     ConnectionClosed,
                     InvalidStatus,
@@ -457,6 +525,23 @@ async def websocket_handler(self, client_ws: WebSocket) -> None:
                 ) as exc:
                     print(f"Livepeer fal ws runner connection failed: {exc}")
 
+                # Publish websocket_connected once after handshake metadata is available
+                if (
+                    not connected_event_published
+                    and metadata.get("manifest_id")
+                    and kafka_publisher
+                    and kafka_publisher.is_running
+                ):
+                    await kafka_publisher.publish(
+                        "websocket_connected",
+                        {
+                            "user_id": metadata.get("user_id"),
+                            "connection_id": metadata.get("manifest_id"),
+                            "connection_info": connection_info,
+                        },
+                    )
+                    connected_event_published = True
+
                 now = time.monotonic()
                 cutoff = now - RUNNER_FAILURE_WINDOW_SECONDS
                 failure_timestamps.append(now)
@@ -478,6 +563,22 @@ async def websocket_handler(self, client_ws: WebSocket) -> None:
         except Exception as exc:
             print(f"Livepeer fal ws proxy error: {type(exc).__name__}: {exc}")
         finally:
+            # Publish websocket_disconnected event
+            if kafka_publisher and kafka_publisher.is_running:
+                end_time = time.time()
+                elapsed_ms = int((end_time - connection_start_time) * 1000)
+                await kafka_publisher.publish(
+                    "websocket_disconnected",
+                    {
+                        "user_id": metadata.get("user_id"),
+                        "connection_id": metadata.get("manifest_id"),
+                        "connection_info": connection_info,
+                        "duration_ms": elapsed_ms,
+                        "session_start_time_ms": int(connection_start_time * 1000),
+                        "session_end_time_ms": int(end_time * 1000),
+                    },
+                )
+
             await run_cleanup()
             with suppress(Exception):
                 await client_ws.close()