vrex-flow-engine 0.2.4__tar.gz → 0.3.0__tar.gz

{vrex_flow_engine-0.2.4 → vrex_flow_engine-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vrex-flow-engine
-Version: 0.2.4
+Version: 0.3.0
 Summary: Launcher + supervisor and OpenAI-compatible image/video generation service (Vrex Flow Engine).
 Author: Vrex
 License: Proprietary

{vrex_flow_engine-0.2.4 → vrex_flow_engine-0.3.0}/flow_engine/__init__.py

@@ -1,4 +1,4 @@
 """flow_engine — host a Flowboard browser-extension instance and expose
 OpenAI-compatible image/video generation endpoints over Google Flow."""
 
-__version__ = "0.2.4"
+__version__ = "0.2.5"

vrex_flow_engine-0.3.0/flow_engine/bridge/flow_client.py

@@ -0,0 +1,315 @@
+"""Bridge to the Chrome MV3 extension over WebSocket.
+
+Control flow:
+1. Extension opens WS to :9223 and sends ``register_instance``.
+2. Agent replies ``{type:"callback_secret", secret}``.
+3. To call Flow, the agent sends ``{id, method:"rpc_request", params:{rpcId,
+   payload, captchaAction?, captchaPaths?, sourcePath?}}``. The extension
+   solves reCAPTCHA in a flow.google.com tab when asked, POSTs the
+   ``batchexecute`` request with the user's cookies + XSRF (``at``) token,
+   parses the response envelope, and POSTs ``{id, status, data, rpcError}``
+   to ``/api/ext/callback`` with ``X-Callback-Secret``. Video bytes take a
+   separate ``upload_request`` command (resumable upload, not batchexecute).
+4. That HTTP handler resolves the pending future by id.
+5. Inbound WS messages (``extension_ready``, ``session_captured``,
+   ``session_lost``, ``pong``) update per-instance state. A fresh session
+   triggers an account refresh (credits + paygate tier) through the bridge.
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import secrets
+import time
+import uuid
+from typing import Any, Optional
+
+from flow_engine.bridge import flow_rpc
+
+logger = logging.getLogger(__name__)
+
+# Minimum gap between account refreshes for the same session — the extension
+# re-announces its session on every reconnect and on periodic re-bootstraps.
+_ACCOUNT_REFRESH_MIN_INTERVAL_S = 60.0
+
+
+class FlowClient:
+    """Bridge client. One instance per pool entry (one per Chrome profile)."""
+
+    DEFAULT_TIMEOUT = 180.0  # seconds
+
+    def __init__(self, semaphore: Optional[asyncio.Semaphore] = None) -> None:
+        self._ws: Optional[Any] = None
+        self._pending: dict[str, asyncio.Future] = {}
+        self._callback_secret: str = secrets.token_urlsafe(32)
+        from flow_engine.config import POOL_MAX_CONCURRENCY
+        self._semaphore: asyncio.Semaphore = semaphore or asyncio.Semaphore(
+            POOL_MAX_CONCURRENCY
+        )
+        # Set by PoolInstance so _send can update in_flight without importing pool.
+        self._on_send_start: Any = lambda: None
+        self._on_send_end: Any = lambda: None
+
+        self._session_present: bool = False
+        self._session_captured_at: Optional[float] = None
+        self._last_account_refresh_at: Optional[float] = None
+        # Identity pushed by the extension from the page bootstrap data.
+        self._user_info: Optional[dict] = None
+        # Account facts from the credits RPC. paygate_tier is the integer tier
+        # id Flow's catalog keys its credit mapping by (1 / 2 / 3).
+        self._paygate_tier: Optional[int] = None
+        self._credits: Optional[int] = None
+        self._plan_code: Optional[int] = None
+        self._request_count = 0
+        self._success_count = 0
+        self._failed_count = 0
+        self._last_error: Optional[str] = None
+
+    # ── connection ─────────────────────────────────────────────────────────
+    @property
+    def connected(self) -> bool:
+        return self._ws is not None
+
+    @property
+    def callback_secret(self) -> str:
+        return self._callback_secret
+
+    def set_extension(self, ws: Any) -> None:
+        self._ws = ws
+
+    def clear_extension(self) -> None:
+        self._ws = None
+        self._clear_session()
+        for fut in self._pending.values():
+            if not fut.done():
+                fut.set_exception(ConnectionError("extension_disconnected"))
+        self._pending.clear()
+
+    def _clear_session(self) -> None:
+        self._session_present = False
+        self._user_info = None
+        self._paygate_tier = None
+        self._credits = None
+        self._plan_code = None
+
+    @property
+    def user_info(self) -> Optional[dict]:
+        return self._user_info
+
+    @property
+    def paygate_tier(self) -> Optional[int]:
+        return self._paygate_tier
+
+    @property
+    def credits(self) -> Optional[int]:
+        return self._credits
+
+    @property
+    def session_present(self) -> bool:
+        return self._session_present
+
+    async def refresh_account(self) -> bool:
+        """Resolve credits + paygate tier via the credits RPC. Returns True when
+        the tier is cached. Runs through the extension, so it needs a live
+        signed-in flow.google.com session."""
+        resp = await self.rpc(flow_rpc.RPC_ACCOUNT_CREDITS, flow_rpc.credits_payload(), timeout=30.0)
+        if resp.get("error"):
+            logger.warning("refresh_account failed: %s", resp["error"])
+            return False
+        account = flow_rpc.parse_credits(resp.get("data"))
+        tier = account.get("paygate_tier")
+        if tier is None:
+            logger.warning("refresh_account: no paygate tier in response %r", resp.get("data"))
+            return False
+        self._paygate_tier = tier
+        self._credits = account.get("credits")
+        self._plan_code = account.get("plan_code")
+        logger.info("account resolved tier=%s credits=%s", tier, self._credits)
+        return True
+
+    # ── inbound handling ───────────────────────────────────────────────────
+    async def handle_message(self, data: dict) -> None:
+        t = data.get("type")
+        if t == "extension_ready":
+            self._session_present = bool(data.get("sessionPresent"))
+            logger.info("extension_ready sessionPresent=%s", self._session_present)
+            return
+        if t == "session_captured":
+            self._session_present = True
+            self._session_captured_at = time.time()
+            email = data.get("email")
+            if isinstance(email, str) and email:
+                self._user_info = {"email": email}
+            now = time.time()
+            last = self._last_account_refresh_at or 0.0
+            if self._paygate_tier is None or (now - last) > _ACCOUNT_REFRESH_MIN_INTERVAL_S:
+                self._last_account_refresh_at = now
+                logger.info("session_captured for %s", email or "<no email>")
+                # Don't await: the WS handler must stay responsive.
+                asyncio.create_task(self.refresh_account())
+            return
+        if t == "session_lost":
+            logger.warning("session_lost: %s", data.get("reason"))
+            self._clear_session()
+            return
+        if t == "pong":
+            return
+        # Inbound response (legacy path; production flow uses HTTP callback)
+        req_id = data.get("id")
+        if req_id and req_id in self._pending:
+            self._resolve(req_id, data)
+
+    def resolve_callback(self, data: dict) -> bool:
+        """Called by the HTTP callback endpoint after validating the secret.
+        Returns True if a pending future matched."""
+        req_id = data.get("id")
+        if not req_id or req_id not in self._pending:
+            return False
+        self._resolve(req_id, data)
+        return True
+
+    def _resolve(self, req_id: str, data: dict) -> None:
+        fut = self._pending.pop(req_id, None)
+        if not fut or fut.done():
+            return
+        status = data.get("status")
+        http_error = isinstance(status, int) and status >= 400
+        if http_error or data.get("error") or data.get("rpcError"):
+            self._failed_count += 1
+            msg = data.get("error") or data.get("rpcError") or f"API_{status}"
+            self._last_error = str(msg)[:200]
+        else:
+            self._success_count += 1
+        fut.set_result(data)
+
+    # ── outbound ──────────────────────────────────────────────────────────
+    async def notify(self, message: dict) -> bool:
+        """Fire-and-forget WS push to the extension."""
+        if not self.connected or self._ws is None:
+            return False
+        try:
+            await self._ws.send(json.dumps(message))
+            return True
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("notify failed: %s", exc)
+            return False
+
+    async def _send(self, method: str, params: dict, timeout: Optional[float] = None) -> dict:
+        if not self.connected:
+            return {"error": "extension_disconnected"}
+        # Per-instance concurrency semaphore — excess callers queue, no 503.
+        async with self._semaphore:
+            self._on_send_start()
+            try:
+                return await self._send_inner(method, params, timeout)
+            finally:
+                self._on_send_end()
+
+    async def _send_inner(
+        self, method: str, params: dict, timeout: Optional[float] = None
+    ) -> dict:
+        req_id = str(uuid.uuid4())
+        fut: asyncio.Future = asyncio.get_running_loop().create_future()
+        self._pending[req_id] = fut
+        self._request_count += 1
+        payload = {"id": req_id, "method": method, "params": params}
+        try:
+            await self._ws.send(json.dumps(payload))
+            return await asyncio.wait_for(fut, timeout=timeout or self.DEFAULT_TIMEOUT)
+        except asyncio.TimeoutError:
+            self._pending.pop(req_id, None)
+            self._failed_count += 1
+            self._last_error = "timeout"
+            return {"error": "timeout"}
+        except Exception as exc:  # noqa: BLE001
+            self._pending.pop(req_id, None)
+            self._failed_count += 1
+            self._last_error = str(exc)
+            return {"error": str(exc)}
+
+    async def rpc(
+        self,
+        rpc_id: str,
+        payload: Any,
+        *,
+        captcha_action: Optional[str] = None,
+        captcha_path: Optional[list[int]] = None,
+        captcha_paths: Optional[list[list[int]]] = None,
+        source_path: Optional[str] = None,
+        timeout: Optional[float] = None,
+    ) -> dict:
+        """Run one batchexecute RPC through the extension.
+
+        Returns ``{"data": <parsed rpc result>, "raw": <envelope>}`` or
+        ``{"error": <str>, "raw": <envelope>}``. ``captcha_action`` makes the
+        extension solve reCAPTCHA first and write the token at every path in
+        ``captcha_paths`` (lists of array indices into ``payload``;
+        ``captcha_path`` is the single-path shorthand).
+        """
+        params: dict[str, Any] = {"rpcId": rpc_id, "payload": payload}
+        if captcha_action:
+            params["captchaAction"] = captcha_action
+            params["captchaPaths"] = captcha_paths or ([captcha_path] if captcha_path else [])
+        if source_path:
+            params["sourcePath"] = source_path
+        resp = await self._send("rpc_request", params, timeout=timeout)
+        if not isinstance(resp, dict):
+            return {"error": "invalid_bridge_response", "raw": resp}
+        if resp.get("error"):
+            return {"error": str(resp["error"]), "raw": resp}
+        if resp.get("rpcError"):
+            return {"error": f"rpc_{rpc_id}_failed: {resp['rpcError']}", "raw": resp}
+        status = resp.get("status")
+        if isinstance(status, int) and status >= 400:
+            return {"error": f"rpc_{rpc_id}_http_{status}", "raw": resp}
+        return {"data": resp.get("data"), "raw": resp}
+
+    async def upload_video(
+        self,
+        project_id: str,
+        file_name: str,
+        content_type: str,
+        data_base64: str,
+        timeout: Optional[float] = None,
+    ) -> dict:
+        """Run Flow's resumable video upload through the extension (start +
+        upload/finalize with the user's cookies). Returns ``{"data": <finalize
+        JSON>, "raw"}`` or ``{"error", "raw"}``. Bytes travel as base64 since
+        the bridge is JSON."""
+        resp = await self._send(
+            "upload_request",
+            {
+                "projectId": project_id,
+                "fileName": file_name,
+                "contentType": content_type,
+                "dataBase64": data_base64,
+            },
+            timeout=timeout or self.DEFAULT_TIMEOUT,
+        )
+        if not isinstance(resp, dict):
+            return {"error": "invalid_bridge_response", "raw": resp}
+        if resp.get("error"):
+            return {"error": str(resp["error"]), "raw": resp}
+        status = resp.get("status")
+        if isinstance(status, int) and status >= 400:
+            return {"error": f"upload_http_{status}", "raw": resp}
+        return {"data": resp.get("data"), "raw": resp}
+
+    # ── observability ─────────────────────────────────────────────────────
+    @property
+    def ws_stats(self) -> dict:
+        age = (
+            int(time.time() - self._session_captured_at)
+            if self._session_captured_at is not None else None
+        )
+        return {
+            "connected": self.connected,
+            "session_present": self._session_present,
+            "session_age_s": age,
+            "pending": len(self._pending),
+            "request_count": self._request_count,
+            "success_count": self._success_count,
+            "failed_count": self._failed_count,
+            "last_error": self._last_error,
+        }