diff --git a/.gitignore b/.gitignore index 535ad65..c0c061a 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,6 @@ graphify-out/ .claude/ .graphifyignore .mypy_cache/ + +# local leftover: old provider build dir (image moved to Ceki-me/docker-browser) +docker/ diff --git a/README.md b/README.md index 65e8d78..4d3a3aa 100644 --- a/README.md +++ b/README.md @@ -354,6 +354,41 @@ Successful commands write a single JSON line to stdout. Errors go to stderr as ` Full reference (with EN+RU): https://browser.ceki.me/docs#cli +### `ceki provider` — rent out your browser + +Turn a machine you control into a **provider**: it runs a real Chromium with the +Ceki extension, injects your browser token and brings the browser online so it +can be rented out as a public browser. + +The provider itself is **not** reimplemented in this SDK. It lives in the +public repo [Ceki-me/docker-browser](https://github.com/Ceki-me/docker-browser) +and ships as the Docker Hub image `ceki/provider`. `ceki provider run` pulls +and runs that image — the launcher stays a single source of truth. + +```bash +export CEKI_PROVIDER_TOKEN= + +ceki provider run # pull + run, stays online until stopped +ceki provider run --timeout 600 # stop after 10 minutes +ceki provider run --build ~/docker-browser # build from a local checkout instead of pulling +``` + +The token is issued for one specific browser and cannot be reused for another. + +#### Provider environment variables + +| Variable | Required | Purpose | +|---|---|---| +| `CEKI_PROVIDER_TOKEN` | yes | Extension token issued for this browser | +| `CEKI_PROVIDER_IMAGE` | no | Image tag (default `ceki/provider:latest`) | +| `CEKI_PROVIDER_VIEWPORT` | no | Browser viewport / resolution WxH (default `1920x1080`) | +| `CEKI_PROVIDER_LOG_LEVEL` | no | Container log verbosity (`DEBUG`/`INFO`/`WARNING`/`ERROR`) | +| `TZ` | no | Browser timezone (kept consistent with your location) | +| `DISPLAY` | no | X display (the container starts its own virtual screen if unset) | + +`docker stop` (or Ctrl-C) sends a clean shutdown signal: the rented browser is +closed and your browser goes **offline**. + ### `ceki contract` — participate in contracts via `/mcp/agent` For AI agents executing tasks inside a contract: list contracts/jobs, post diff --git a/ceki_sdk/__init__.py b/ceki_sdk/__init__.py index a5893ff..00ce8da 100644 --- a/ceki_sdk/__init__.py +++ b/ceki_sdk/__init__.py @@ -19,6 +19,7 @@ ) from ._models import BrowserOption, ChatMessage, Match, ReadReceipt, SessionInfo, Snapshot from ._profile import BrowserProfile +from ._provider import ProviderError, run_provider from .humanize import HumanProfile __version__ = "2.36.2" @@ -38,6 +39,8 @@ "AuthFailed", "ConnectionLost", "ProviderDisconnected", + "run_provider", + "ProviderError", "SessionNotFound", "SessionExpired", "NotOwner", diff --git a/ceki_sdk/_client.py b/ceki_sdk/_client.py index 12757e4..325a99b 100644 --- a/ceki_sdk/_client.py +++ b/ceki_sdk/_client.py @@ -5,6 +5,7 @@ import logging import os import time +from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Any import httpx @@ -67,6 +68,11 @@ def __init__( self._closed = False self._stashed_first_frame: str | None = None + # Optional hook invoked on ``session.ended``/``session_end``. The + # daemon uses it to drop the session from its registry and close the + # shared WebSocket once the last session for a client is gone. + self._on_session_ended: Callable[[str], Awaitable[None]] | None = None + # P2P WebRTC transport (primary, WS = fallback) self._p2p: WebRTCTransport | None = None self._p2p_init_lock = asyncio.Lock() @@ -210,6 +216,12 @@ async def rent( self._pending_rent_queue.remove(fut) except ValueError: pass + # ``rent_pending`` already moved *fut* into ``_pending_rents`` keyed + # by the relay's event_id — drop it there too, or a long-running + # client (the daemon) accumulates a dead future per timed-out rent. + for eid, pfut in list(self._pending_rents.items()): + if pfut is fut: + del self._pending_rents[eid] raise TimeoutError("rent timed out waiting for match") # Wait for P2P WebRTC transport to initialize before returning Browser. @@ -486,8 +498,14 @@ async def _dispatch(self, msg: dict[str, Any]) -> None: if mtype == "cdp_response": session_id = msg.get("session_id", "") browser = self._active_browsers.get(session_id) - log.debug("WS cdp_response: sid=%s browser=%s active=%s msg_id=%s ok=%s", - session_id, bool(browser), list(self._active_browsers.keys()), msg.get("id"), msg.get("ok")) + log.debug( + "WS cdp_response: sid=%s browser=%s active=%s msg_id=%s ok=%s", + session_id, + bool(browser), + list(self._active_browsers.keys()), + msg.get("id"), + msg.get("ok"), + ) if browser: await browser._on_cdp_response(msg) return @@ -503,11 +521,25 @@ async def _dispatch(self, msg: dict[str, Any]) -> None: if browser: await browser._on_tab_opened(msg) return - if mtype in ("session.ended", "session_end"): - session_id = msg.get("session_id", "") + if mtype in ("session.ended", "session_end", "session_ended"): + # The relay's session-end message is ``session_ended`` with the id in + # ``event_id`` (older aliases used ``session_id``). Accept every form + # so relay-initiated ends (provider death, admin stop, backend reaper) + # are never dropped — otherwise the daemon would keep the session and + # its shared WS alive forever. + sid = msg.get("session_id") or msg.get("event_id") + session_id = str(sid) if sid else "" browser = self._active_browsers.get(session_id) if browser: await browser._on_session_ended(msg) + # Notify the daemon so it can drop the session from its registry and + # close the shared WS once the last session for this client is gone. + hook = self._on_session_ended + if hook is not None: + try: + await hook(session_id) + except Exception as exc: + log.error("session.ended hook failed: %s", exc) return if mtype == "session.provider_disconnected": session_id = msg.get("session_id", "") @@ -552,9 +584,23 @@ async def _dispatch(self, msg: dict[str, Any]) -> None: asyncio.create_task(browser.chat._on_send_error(msg)) return if mtype == "error": - session_id = msg.get("session_id") - if session_id and session_id in self._active_browsers: - await self._active_browsers[session_id]._on_error(msg) + sid = msg.get("session_id") or msg.get("event_id") + session_id = str(sid) if sid else "" + browser = self._active_browsers.get(session_id) if session_id else None + if browser is not None and msg.get("code", 0) in (-1011, -1018): + # Relay reports a session end as ``error -1011/-1018`` (provider + # death, grace expiry, admin kill). Clean up exactly like + # ``session_ended`` so the daemon never keeps a dead session. + await browser._on_session_ended(msg) + hook = self._on_session_ended + if hook is not None: + try: + await hook(session_id) + except Exception as exc: + log.error("session.ended hook failed: %s", exc) + return + if browser is not None: + await browser._on_error(msg) else: self._handle_error(msg) return diff --git a/ceki_sdk/_provider.py b/ceki_sdk/_provider.py new file mode 100644 index 0000000..b866a03 --- /dev/null +++ b/ceki_sdk/_provider.py @@ -0,0 +1,188 @@ +"""Provider mode: rent out this machine's browser through Ceki. + +A thin wrapper around the official provider image. The provider itself — real +Chromium + the Ceki extension + token handshake + online poll + liveness — +lives in the public repo ``Ceki-me/docker-browser``. The SDK does NOT +reimplement the provider; it pulls an image and runs it, so the launcher stays a +single source of truth in docker-browser. The image is pulled from Docker Hub +(``ceki/provider:latest``) first and, if that pull fails, from the GHCR build +(``ghcr.io/ceki-me/docker-browser:latest``) that docker-browser CI publishes on +every GitHub Release. + +CLI entry: + ceki provider run [--token TOKEN] [--image IMAGE] [--build DIR] + [--viewport WxH] [--timeout SECONDS] [--verbose] + +Environment variables: + CEKI_PROVIDER_TOKEN extension token for this browser (required) + CEKI_PROVIDER_IMAGE image tag (default ``ceki/provider:latest``) + CEKI_PROVIDER_VIEWPORT browser viewport WxH (default 1920x1080) + CEKI_PROVIDER_LOG_LEVEL container log verbosity (default INFO) + TZ timezone passed into the container + DISPLAY X display (the container starts Xvfb if unset) + +Only the public provider envs are passed through — internal envs +(CEKI_WS_URL / CEKI_API_URL / update knobs) are not part of the SDK contract. +""" + +from __future__ import annotations + +import os +import shutil +import subprocess + +DEFAULT_IMAGE = "ceki/provider:latest" + +# GHCR copy of the same launcher, published by docker-browser CI on every +# GitHub Release. Used when the Docker Hub tag is not (yet) available. +FALLBACK_IMAGE = "ghcr.io/ceki-me/docker-browser:latest" + +# Public docker-browser envs forwarded from the caller's environment. +_PUBLIC_ENVS = ("CEKI_PROVIDER_VIEWPORT", "CEKI_PROVIDER_LOG_LEVEL", "TZ", "DISPLAY") + +# Default command the image runs. ``docker run`` args override the image CMD, +# so ``--timeout`` is passed by appending this command + the flag. +_APP_CMD = ("python", "-m", "ceki_browser_provider.app") + + +class ProviderError(Exception): + """Raised when the provider cannot be deployed or brought online.""" + + +def resolve_token(token: str | None = None) -> str: + """Resolve the provider token from arg or environment.""" + value = (token or os.environ.get("CEKI_PROVIDER_TOKEN") or "").strip() + if not value: + raise ProviderError( + "Provider token is required: set CEKI_PROVIDER_TOKEN or pass --token" + ) + return value + + +def resolve_image(explicit: str | None = None) -> str: + """Resolve the image tag from arg, env or the default.""" + return (explicit or os.environ.get("CEKI_PROVIDER_IMAGE") or DEFAULT_IMAGE).strip() + + +def _docker() -> str: + binary = shutil.which("docker") + if binary is None: + raise ProviderError( + "Docker is required to run a provider. Install Docker, or build the " + "provider manually from https://github.com/Ceki-me/docker-browser" + ) + return binary + + +def _env_map(token: str, viewport: str | None = None, verbose: bool = False) -> dict[str, str]: + """Build the container env: token + public env pass-through + explicit args.""" + env = {"CEKI_PROVIDER_TOKEN": token} + for name in _PUBLIC_ENVS: + if os.environ.get(name): + env[name] = os.environ[name] + if viewport: + env["CEKI_PROVIDER_VIEWPORT"] = viewport + if verbose: + env["CEKI_PROVIDER_LOG_LEVEL"] = "DEBUG" + return env + + +def _run_cmd( + docker: str, + image: str, + env: dict[str, str], + timeout: int | None = None, +) -> list[str]: + """Build the ``docker run`` command line (token never logged).""" + cmd = [docker, "run", "--rm"] + for name, value in env.items(): + cmd += ["-e", f"{name}={value}"] + cmd.append(image) + if timeout: + # docker run args replace the image CMD — keep the image's default + # entry command and append --timeout so the container self-stops. + cmd.extend([*_APP_CMD, f"--timeout={timeout}"]) + return cmd + + +def _build_image(build_dir: str) -> None: + """Build ``ceki/provider:latest`` from a local docker-browser checkout.""" + build_sh = os.path.join(build_dir, "build.sh") + if not os.path.isfile(build_sh): + raise ProviderError( + f"{build_sh} not found — pass the docker-browser repo directory " + "(https://github.com/Ceki-me/docker-browser)" + ) + print(f"[ceki-provider] building image from {build_dir} (./build.sh) ...") + if subprocess.call([build_sh], cwd=build_dir) != 0: + raise ProviderError("docker-browser build.sh failed") + + +def _pull_image(docker: str, image: str) -> bool: + """Pull ``image`` if not present locally; True if it is runnable now.""" + if ( + subprocess.run([docker, "image", "inspect", image], capture_output=True) + .returncode + == 0 + ): + return True + print(f"[ceki-provider] pulling {image} ...") + return subprocess.run([docker, "pull", image]).returncode == 0 + + +def _pull_or_fallback(docker: str, image: str) -> str: + """Return the first image tag that is runnable, trying the GHCR fallback. + + An explicit tag (``--image`` / ``$CEKI_PROVIDER_IMAGE``) is used as-is and + never swapped; only the default Docker Hub tag gets the GHCR fallback, + since the Hub image may not be published yet. + """ + candidates = [image] + if image == DEFAULT_IMAGE: + candidates.append(FALLBACK_IMAGE) + for candidate in candidates: + if _pull_image(docker, candidate): + return candidate + raise ProviderError( + "failed to pull " + + " and ".join(f"'{candidate}'" for candidate in candidates) + + " — check the image name and network access" + ) + + +def run_provider( + *, + token: str | None = None, + image: str | None = None, + build: str | None = None, + viewport: str | None = None, + timeout: int | None = None, + verbose: bool = False, +) -> int: + """Pull and run the docker-browser provider image until stopped. + + Returns a process exit code (0 on clean shutdown, 130 on Ctrl-C). + """ + token_value = resolve_token(token) + image_value = resolve_image(image) + docker = _docker() + + if build: + _build_image(build) + + # Resolve a runnable image: the Docker Hub tag first, the GHCR copy as a + # fallback when the Hub tag is not published (or not reachable) yet. + run_image = _pull_or_fallback(docker, image_value) + + env = _env_map(token_value, viewport=viewport, verbose=verbose) + cmd = _run_cmd(docker, run_image, env, timeout=timeout) + + print( + f"[ceki-provider] starting {run_image} — browser online until stopped " + "(Ctrl-C / docker stop)" + ) + try: + code = subprocess.call(cmd) + except KeyboardInterrupt: + return 130 + return 0 if code in (0, 130) else code diff --git a/ceki_sdk/cli.py b/ceki_sdk/cli.py index 4bb9d12..24d57db 100644 --- a/ceki_sdk/cli.py +++ b/ceki_sdk/cli.py @@ -112,19 +112,26 @@ async def _daemon_request( path: str, params: dict[str, Any] | None = None, timeout: float = 120.0, -) -> Any: +) -> tuple[bool, Any]: """Send an IPC request to a running daemon. - Returns ``None`` when the daemon is not running (clean fallback for the - caller). Raises ``CekiError`` when the daemon *was* expected to be + Returns ``(False, None)`` when the daemon is not running (clean fallback + for the caller). Raises ``CekiError`` when the daemon *was* expected to be reachable but isn't — the caller shows the error to the user instead of falling back to one-shot mode. + Returns ``(True, result)`` on success. ``result`` may be ``None`` for + void commands (navigate/click/type/scroll/switch_tab/configure/stop/ + profile_import) — a null result is a SUCCESS, not a failure signal. The + caller MUST branch on the first tuple element, never on ``result`` alone, + otherwise a successful void command falls through to the one-shot resume + fallback which disconnects and kills the rented session. + The function checks ``PID_FILE`` first as a fast-path; if absent there is no running daemon. If present but unreachable we clean the stale file. """ if not PID_FILE.exists(): - return None # daemon not running → clean fallback + return False, None # daemon not running → clean fallback port = daemon_port() url = f"http://{DAEMON_HOST}:{port}{path}" @@ -139,10 +146,10 @@ async def _daemon_request( body = resp.json() if not body.get("ok"): raise CekiError(body.get("error", "daemon error")) - return body.get("result") + return True, body.get("result") except httpx.ConnectError: PID_FILE.unlink(missing_ok=True) - return None # stale PID → clean fallback + return False, None # stale PID → clean fallback except httpx.TimeoutException: raise CekiError("daemon not responding (timeout), start daemon first") except httpx.HTTPError as e: @@ -232,6 +239,33 @@ def _cmd_daemon(args: argparse.Namespace) -> int: return 1 +def _cmd_provider(args: argparse.Namespace) -> int: + """Run a browser provider via the official docker-browser image. + + Long-running: pulls the provider image (Docker Hub ``ceki/provider`` with a + GHCR fallback) and runs the container — Chromium + Ceki extension + token + handshake — keeping the browser online until SIGTERM/SIGINT. The provider + launcher itself lives in the public repo Ceki-me/docker-browser; the SDK + only orchestrates docker. + """ + from ._provider import ProviderError, run_provider + + try: + return run_provider( + token=args.token, + image=args.image, + build=args.build, + viewport=args.viewport, + timeout=args.timeout, + verbose=args.verbose, + ) + except ProviderError as e: + _err(str(e), "provider") + return 1 + except KeyboardInterrupt: + return 130 + + async def _cmd_rent(args: argparse.Namespace) -> None: # Auto-start daemon on rent — subsequent commands use the persistent WS if not _ensure_daemon(): @@ -241,12 +275,12 @@ async def _cmd_rent(args: argparse.Namespace) -> None: # Try daemon IPC fp_from = str(Path(args.fingerprint_from).resolve()) if args.fingerprint_from else None try: - result = await _daemon_request("/rent", { + ok, result = await _daemon_request("/rent", { "schedule": args.schedule, "mode": args.mode, "fingerprint_from": fp_from, }) - if result is not None: + if ok: sid = result["session_id"] save_session(sid, { "session_id": sid, @@ -295,8 +329,8 @@ async def _resume_browser(api_key: str, session_id: str): async def _cmd_snapshot(args: argparse.Namespace) -> None: # Try daemon IPC try: - result = await _daemon_request("/snapshot", {"session_id": args.session_id}) - if result is not None: + ok, result = await _daemon_request("/snapshot", {"session_id": args.session_id}) + if ok: png_bytes = base64.b64decode(result["screenshot"]) if result.get("screenshot") else b"" out_path = args.output with open(out_path, "wb") as f: @@ -340,12 +374,12 @@ def _human_flag(args: argparse.Namespace) -> bool | None: async def _cmd_navigate(args: argparse.Namespace) -> None: # Try daemon IPC try: - result = await _daemon_request("/navigate", { + ok, result = await _daemon_request("/navigate", { "session_id": args.session_id, "url": args.url, "human": _human_flag(args), }) - if result is not None: + if ok: _out({"ok": True}) return except CekiError as e: @@ -366,13 +400,13 @@ async def _cmd_navigate(args: argparse.Namespace) -> None: async def _cmd_click(args: argparse.Namespace) -> None: # Try daemon IPC try: - result = await _daemon_request("/click", { + ok, result = await _daemon_request("/click", { "session_id": args.session_id, "x": args.x, "y": args.y, "human": _human_flag(args), }) - if result is not None: + if ok: _out({"ok": True, "pointer": [args.x, args.y]}) return except CekiError as e: @@ -397,13 +431,13 @@ async def _cmd_type(args: argparse.Namespace) -> None: # --natural is a no-op alias kept for backwards compatibility. # Try daemon IPC try: - result = await _daemon_request("/type", { + ok, result = await _daemon_request("/type", { "session_id": args.session_id, "text": args.text, "selector": args.selector, "human": _human_flag(args), }) - if result is not None: + if ok: _out({"ok": True}) return except CekiError as e: @@ -424,14 +458,14 @@ async def _cmd_type(args: argparse.Namespace) -> None: async def _cmd_scroll(args: argparse.Namespace) -> None: # Try daemon IPC try: - result = await _daemon_request("/scroll", { + ok, result = await _daemon_request("/scroll", { "session_id": args.session_id, "x": args.x, "y": args.y, "dy": args.dy, "human": _human_flag(args), }) - if result is not None: + if ok: _out({"ok": True}) return except CekiError as e: @@ -453,21 +487,21 @@ async def _cmd_chat(args: argparse.Namespace) -> None: # Try daemon IPC try: if args.chat_action == "send": - result = await _daemon_request("/chat/send", { + ok, result = await _daemon_request("/chat/send", { "session_id": args.session_id, "text": args.text, }) - if result is not None: + if ok: _out({"ok": True, "message_id": result.get("message_id")}) return elif args.chat_action == "next": last_seen = get_last_seen_ts(args.session_id) - result = await _daemon_request("/chat/next", { + ok, result = await _daemon_request("/chat/next", { "session_id": args.session_id, "timeout": args.timeout, "since": last_seen, }) - if result is not None: + if ok: if result: # has message update_last_seen_ts(args.session_id, result["ts"]) _out(result) # None → no message @@ -481,12 +515,12 @@ async def _cmd_chat(args: argparse.Namespace) -> None: since = datetime.fromtimestamp(ts_val, tz=timezone.utc).isoformat() except ValueError: since = args.since - result = await _daemon_request("/chat/history", { + ok, result = await _daemon_request("/chat/history", { "session_id": args.session_id, "since": since, "limit": args.limit, }) - if result is not None: + if ok: _out(result) return except CekiError as e: @@ -547,8 +581,8 @@ async def on_msg(msg): async def _cmd_stop(args: argparse.Namespace) -> None: # Try daemon IPC try: - result = await _daemon_request("/stop", {"session_id": args.session_id}) - if result is not None: + ok, result = await _daemon_request("/stop", {"session_id": args.session_id}) + if ok: delete_session(args.session_id) _out({"ok": True}) return @@ -573,12 +607,12 @@ async def _cmd_profile(args: argparse.Namespace) -> None: try: if args.profile_action == "export": domains = ",".join(args.domains) if args.domains else None - result = await _daemon_request("/profile/export", { + ok, result = await _daemon_request("/profile/export", { "session_id": args.session_id, "domains": domains, "no_session_storage": args.no_session_storage, }) - if result is not None: + if ok: with open(args.output, "w") as f: json.dump(result, f) _out({"ok": True, "path": args.output}) @@ -586,11 +620,11 @@ async def _cmd_profile(args: argparse.Namespace) -> None: elif args.profile_action == "import": with open(args.input, "r") as f: profile_dict = json.load(f) - result = await _daemon_request("/profile/import", { + ok, result = await _daemon_request("/profile/import", { "session_id": args.session_id, "profile": profile_dict, }) - if result is not None: + if ok: _out({"ok": True}) return except CekiError as e: @@ -695,11 +729,11 @@ async def _cmd_wait(args: argparse.Namespace) -> None: async def _cmd_screenshot(args: argparse.Namespace) -> None: # Try daemon IPC try: - result = await _daemon_request("/screenshot", { + ok, result = await _daemon_request("/screenshot", { "session_id": args.session_id, "full": args.full, }) - if result is not None: + if ok: data = base64.b64decode(_unwrap_screenshot_data(result)) with open(args.output, "wb") as f: f.write(data) @@ -725,8 +759,8 @@ async def _cmd_screenshot(args: argparse.Namespace) -> None: async def _cmd_switch_tab(args: argparse.Namespace) -> None: # Try daemon IPC try: - result = await _daemon_request("/switch-tab", {"session_id": args.session_id}) - if result is not None: + ok, result = await _daemon_request("/switch-tab", {"session_id": args.session_id}) + if ok: _out({"ok": True}) return except CekiError as e: @@ -752,8 +786,8 @@ async def _cmd_configure(args: argparse.Namespace) -> None: params["masking_mode"] = args.masking_mode if args.fingerprint is not None: params["fingerprint"] = args.fingerprint - result = await _daemon_request("/configure", params) - if result is not None: + ok, result = await _daemon_request("/configure", params) + if ok: _out({"ok": True}) return except CekiError as e: @@ -1154,12 +1188,12 @@ async def _cmd_cdp(args: argparse.Namespace) -> None: # Try daemon IPC params = json.loads(args.params) if args.params else {} try: - result = await _daemon_request("/cdp", { + ok, result = await _daemon_request("/cdp", { "session_id": args.session_id, "method": args.method, "params": params, }) - if result is not None: + if ok: _out(result) return except CekiError as e: @@ -1605,6 +1639,47 @@ def build_parser() -> argparse.ArgumentParser: dsub.add_parser("stop", help="Stop daemon (SIGTERM)") dsub.add_parser("status", help="Check daemon status") + # ── provider subcommand ─────────────────────────────────────────── + p_provider = sub.add_parser( + "provider", + help="Run a browser provider — rent out this machine's browser", + ) + psub = p_provider.add_subparsers(dest="provider_action", required=True) + p_run = psub.add_parser( + "run", + help=( + "Pull and run the official provider image (ceki/provider on " + "Docker Hub, falling back to ghcr.io/ceki-me/docker-browser) and " + "keep the browser online, auto-accepting rentals" + ), + ) + p_run.add_argument( + "--token", + help="Provider extension token (default: $CEKI_PROVIDER_TOKEN)", + ) + p_run.add_argument( + "--image", + help=( + "Provider image tag (default: $CEKI_PROVIDER_IMAGE or " + "ceki/provider:latest, falling back to the GHCR build if the " + "Docker Hub pull fails)" + ), + ) + p_run.add_argument( + "--build", + help="Build the image from a local docker-browser checkout instead of pulling", + ) + p_run.add_argument( + "--viewport", + help="Browser viewport WxH (default 1920x1080)", + ) + p_run.add_argument( + "--timeout", + type=int, + help="Run for N seconds then stop (default: run until stopped)", + ) + p_run.add_argument("--verbose", action="store_true", help="Verbose provider logging (DEBUG)") + return parser @@ -1646,6 +1721,9 @@ def main() -> None: if args.command == "daemon": sys.exit(_cmd_daemon(args)) + if args.command == "provider": + sys.exit(_cmd_provider(args)) + handler = handlers.get(args.command) if not handler: _err(f"Unknown command: {args.command}") diff --git a/ceki_sdk/contract.py b/ceki_sdk/contract.py index f76d60f..b19648f 100644 --- a/ceki_sdk/contract.py +++ b/ceki_sdk/contract.py @@ -250,6 +250,10 @@ def tools(self) -> Any: return body def raw(self, tool: str, args: dict[str, Any] | None = None) -> Any: + # prompts/list and prompts/get are MCP protocol methods (JSON-RPC), + # not tools. Call _rpc directly instead of tools/call. + if tool.startswith("prompts/"): + return self._rpc(tool, args or {}) return self.call(tool, args) # ── domain helpers ──────────────────────────────────────────── diff --git a/ceki_sdk/daemon.py b/ceki_sdk/daemon.py index 69e97d8..51a1797 100644 --- a/ceki_sdk/daemon.py +++ b/ceki_sdk/daemon.py @@ -20,7 +20,7 @@ from . import ConnectOptions, connect from ._browser import _unwrap_screenshot_data -from ._exceptions import SessionNotFound +from ._exceptions import ConnectionLost, SessionNotFound log = logging.getLogger(__name__) @@ -159,9 +159,52 @@ async def _handle_rent(self, params: dict) -> dict: profile = json.load(f) fp_data = profile.get("fingerprint") or True - client = await connect(api_key, _connect_options()) - browser = await client.rent(schedule, mode=mode, fingerprint=fp_data) - self.server.daemon_server._sessions[browser.session_id] = (client, browser) + # Reuse ONE shared Client per api_key — all sessions multiplex over a + # single WebSocket. Old per-rent clients were never closed, leaking a + # live WS per rent and confusing relay cdp_response routing. + daemon = self.server.daemon_server + + async def _shared_client(): + client = daemon._clients.get(api_key) + # A cached client may be closed already (WS torn down after the + # last session ended, or by a failed rent) — never reuse it. + if client is None or client._closed or client._ws is None: + client = await connect(api_key, _connect_options()) + client._on_session_ended = daemon._on_session_ended + daemon._clients[api_key] = client + return client + + client = await _shared_client() + try: + browser = await client.rent(schedule, mode=mode, fingerprint=fp_data) + except (TimeoutError, ConnectionLost) as exc: + # The shared WS is half-dead: the relay stopped routing rent/match + # without a close frame, so the TCP socket stays ESTABLISHED, + # recv() never raises and pongs keep coming — neither the reader + # nor the heartbeat notices. Every rent through it hangs 90s and + # 504s forever until the daemon restarts. Drop the poisoned client + # and retry ONCE on a fresh connection. + log.warning( + "rent failed for %s (%s) — recreating shared client", + api_key, type(exc).__name__, + ) + await daemon._drop_client(api_key, client) + client = await _shared_client() + try: + browser = await client.rent(schedule, mode=mode, fingerprint=fp_data) + except Exception: + await daemon._drop_client(api_key, client) + raise + except Exception: + if not daemon._client_has_sessions(client): + if daemon._clients.get(api_key) is client: + daemon._clients.pop(api_key, None) + try: + await client.disconnect() + except Exception: + pass + raise + daemon._sessions[browser.session_id] = browser return { "session_id": browser.session_id, "chat_topic_id": browser.chat_topic_id, @@ -233,17 +276,14 @@ async def _handle_snapshot(self, params: dict) -> dict: async def _handle_stop(self, params: dict) -> None: session_id = params.get("session_id", "") - entry = self.server.daemon_server._sessions.pop(session_id, None) - if entry is None: + browser = self.server.daemon_server._sessions.pop(session_id, None) + if browser is None: raise ValueError(f"session not found: {session_id}") - client, browser = entry try: await browser.close() finally: - try: - await client.disconnect() - except Exception: - pass + # Shared client is disconnected only when the LAST session ends. + await self.server.daemon_server._maybe_disconnect_clients() async def _handle_chat_send(self, params: dict) -> dict: browser = await self._resolve_browser(params) @@ -342,14 +382,14 @@ async def _handle_request_captcha(self, params: dict) -> dict: # ── session resolution ───────────────────────────────────────────── async def _resolve_browser(self, params: dict): - """Look up a stored (Client, Browser) pair by session_id.""" + """Look up a stored Browser by session_id.""" session_id = params.get("session_id", "") if not session_id: raise ValueError("session_id required") - entry = self.server.daemon_server._sessions.get(session_id) - if entry is None: + browser = self.server.daemon_server._sessions.get(session_id) + if browser is None: raise SessionNotFound(f"session not found: {session_id}") - return entry[1] + return browser _ENDPOINTS: dict[str, str] = { @@ -389,9 +429,13 @@ class DaemonServer: - A daemon thread runs a ``ThreadingHTTPServer`` that accepts IPC requests. - The HTTP handler calls :meth:`run_async` to schedule a coroutine on the event-loop and waits for its result — bridging sync → async boundaries. - - Sessions are stored in memory as ``{session_id: (Client, Browser)}``. + - Sessions are stored in memory as ``{session_id: Browser}``. + - Clients are shared per ``api_key`` in ``{api_key: Client}`` — all sessions + of one key multiplex over a single WebSocket, so there is never more than + one live connection per key. - ``SIGTERM`` / ``SIGINT`` triggers a graceful shutdown: all sessions are - closed, the PID file is removed, and the event-loop stops. + closed, shared clients are disconnected, the PID file is removed, and the + event-loop stops. """ def __init__(self, host: str = DAEMON_HOST, port: int | None = None) -> None: @@ -400,7 +444,8 @@ def __init__(self, host: str = DAEMON_HOST, port: int | None = None) -> None: self._httpd: ThreadingHTTPServer | None = None self._thread: threading.Thread | None = None self._loop: asyncio.AbstractEventLoop | None = None - self._sessions: dict[str, tuple[Any, Any]] = {} + self._sessions: dict[str, Any] = {} + self._clients: dict[str, Any] = {} # ── public API ───────────────────────────────────────────────────── @@ -451,21 +496,87 @@ def start(self) -> None: async def _shutdown(self) -> None: log.info("shutting down (closing %d session(s))", len(self._sessions)) # Close all sessions - for session_id, (client, browser) in list(self._sessions.items()): + for session_id, browser in list(self._sessions.items()): try: await browser.close() except Exception as exc: log.debug("close session %s: %s", session_id, exc, exc_info=True) + self._sessions.clear() + # Disconnect all shared clients + for api_key, client in list(self._clients.items()): try: await client.disconnect() except Exception: pass - self._sessions.clear() + self._clients.clear() # Stop HTTP server (blocking call offloaded to thread pool) if self._httpd: await asyncio.to_thread(self._httpd.shutdown) self._loop.stop() + def _client_has_sessions(self, client: Any) -> bool: + """True if any live session belongs to *client* (via its Browser).""" + return any(browser._client is client for browser in self._sessions.values()) + + async def _disconnect_client(self, client: Any) -> None: + try: + await client.disconnect() + except Exception as exc: + log.debug("disconnect shared client: %s", exc, exc_info=True) + + async def _drop_client(self, api_key: str, client: Any) -> None: + """Remove a poisoned shared client and drop its sessions. + + Called when a rent times out on a client whose WebSocket went half-dead + (relay stopped routing without a close frame). That WS can no longer + carry rent/match messages, so the client is dropped from the cache and + any sessions it owned are discarded — they are unreachable anyway. + """ + if self._clients.get(api_key) is client: + self._clients.pop(api_key, None) + for sid in [ + sid for sid, browser in self._sessions.items() + if browser._client is client + ]: + self._sessions.pop(sid, None) + await self._disconnect_client(client) + + async def _maybe_disconnect_clients(self) -> None: + """Disconnect shared clients once the last session for them is gone. + + Called from the HTTP handler (not the client's own reader task), so + awaiting ``client.disconnect()`` here is safe. + """ + if self._sessions: + return + clients = list(self._clients.items()) + self._clients.clear() + for _, client in clients: + await self._disconnect_client(client) + + async def _on_session_ended(self, session_id: str) -> None: + """Daemon-side cleanup when a rented session ends on the relay. + + Invoked by the shared :class:`Client` on ``session.ended``/``session_end`` + (see ``_client.py``). Removes the session from ``_sessions`` and, when + the last session is gone, closes the shared client's WebSocket so the + relay never accumulates orphan connections. + + This runs inside the client's own reader task, so disconnecting the + client must be deferred to a separate task — ``disconnect()`` cancels + the reader task, which would otherwise cancel this very coroutine and + skip the actual WS/P2P teardown. + """ + self._sessions.pop(session_id, None) + if self._sessions: + return + clients = list(self._clients.items()) + self._clients.clear() + for _, client in clients: + # Stop the reader loop synchronously; teardown happens in a task. + client._closed = True + asyncio.create_task(self._disconnect_client(client)) + def _cleanup(self) -> None: PID_FILE.unlink(missing_ok=True) log.info("daemon stopped") diff --git a/tests/test_daemon.py b/tests/test_daemon.py new file mode 100644 index 0000000..cc16904 --- /dev/null +++ b/tests/test_daemon.py @@ -0,0 +1,470 @@ +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, Mock, patch + +import pytest + +from ceki_sdk._client import Client +from ceki_sdk.daemon import DaemonHTTPHandler, DaemonServer + + +def _make_handler(daemon: DaemonServer) -> DaemonHTTPHandler: + """Construct an HTTP handler wired to *daemon* without a real socket.""" + handler = DaemonHTTPHandler.__new__(DaemonHTTPHandler) + server = Mock() + server.daemon_server = daemon + handler.server = server + return handler + + +def _make_browser(session_id: str, client: Mock) -> Mock: + browser = Mock() + browser.session_id = session_id + browser.chat_topic_id = f"topic-{session_id}" + browser.schedule_id = 1 + browser.close = AsyncMock() + browser._client = client + return browser + + +@pytest.mark.asyncio +async def test_daemon_reuses_shared_client_per_api_key(): + """Renting two sessions with the same key must use ONE shared Client.""" + daemon = DaemonServer() + handler = _make_handler(daemon) + + shared = AsyncMock() + shared._closed = False + shared.rent = AsyncMock(side_effect=[ + _make_browser("s1", shared), + _make_browser("s2", shared), + ]) + + with patch("ceki_sdk.daemon.connect", AsyncMock(return_value=shared)) as connect_mock: + r1 = await handler._handle_rent({"api_key": "key", "schedule": 5}) + r2 = await handler._handle_rent({"api_key": "key", "schedule": 6}) + + assert r1["session_id"] == "s1" + assert r2["session_id"] == "s2" + + # One shared client registered, reused across both rents — connect() once. + assert list(daemon._clients) == ["key"] + assert daemon._clients["key"] is shared + connect_mock.assert_awaited_once() + + # Both sessions registered as plain Browsers on the same client. + assert set(daemon._sessions) == {"s1", "s2"} + assert daemon._sessions["s1"]._client is shared + assert daemon._sessions["s2"]._client is shared + + # The session.ended hook is wired so the daemon learns about relay ends. + assert shared._on_session_ended.__func__ is DaemonServer._on_session_ended + assert shared._on_session_ended.__self__ is daemon + + +@pytest.mark.asyncio +async def test_daemon_rent_failure_does_not_leak_new_client(): + """A failed rent on a fresh client must clean up that client.""" + daemon = DaemonServer() + handler = _make_handler(daemon) + + shared = AsyncMock() + shared.rent = AsyncMock(side_effect=RuntimeError("no providers")) + + with patch("ceki_sdk.daemon.connect", AsyncMock(return_value=shared)): + with pytest.raises(RuntimeError): + await handler._handle_rent({"api_key": "key", "schedule": 5}) + + assert daemon._clients == {} + assert daemon._sessions == {} + shared.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_daemon_rent_timeout_recreates_poisoned_shared_client(): + """A rent that times out on a half-dead WS must drop the client and retry + on a fresh connection instead of 504ing forever (ev report: POST /rent 504 + every time once the shared WS stopped routing rent/match).""" + daemon = DaemonServer() + handler = _make_handler(daemon) + + poisoned = AsyncMock() + poisoned._closed = False + poisoned.rent = AsyncMock(side_effect=TimeoutError("rent timed out")) + + fresh = AsyncMock() + fresh._closed = False + fresh.rent = AsyncMock(side_effect=[_make_browser("s1", fresh)]) + + connect_mock = AsyncMock(side_effect=[poisoned, fresh]) + with patch("ceki_sdk.daemon.connect", connect_mock): + r = await handler._handle_rent({"api_key": "key", "schedule": 5}) + + assert r["session_id"] == "s1" + # Poisoned client dropped from cache + disconnected, fresh one registered. + assert daemon._clients["key"] is fresh + assert set(daemon._sessions) == {"s1"} + assert connect_mock.await_count == 2 + poisoned.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_daemon_rent_timeout_retry_failure_also_drops_fresh_client(): + """If the retry on the fresh client times out too, that client is dropped + as well — the daemon is left with no cached client so the NEXT request + reconnects instead of hitting a half-dead WS again.""" + daemon = DaemonServer() + handler = _make_handler(daemon) + + poisoned = AsyncMock() + poisoned._closed = False + poisoned.rent = AsyncMock(side_effect=TimeoutError("rent timed out")) + + fresh = AsyncMock() + fresh._closed = False + fresh.rent = AsyncMock(side_effect=TimeoutError("rent timed out")) + + connect_mock = AsyncMock(side_effect=[poisoned, fresh]) + with patch("ceki_sdk.daemon.connect", connect_mock): + with pytest.raises(TimeoutError): + await handler._handle_rent({"api_key": "key", "schedule": 5}) + + assert daemon._clients == {} + assert daemon._sessions == {} + poisoned.disconnect.assert_awaited_once() + fresh.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_daemon_session_ended_disconnects_only_when_empty(): + """session.ended removes the sid; shared client is closed only after the + last session is gone.""" + daemon = DaemonServer() + client_a = AsyncMock() + client_b = AsyncMock() + daemon._clients["key-a"] = client_a + daemon._clients["key-b"] = client_b + b1 = _make_browser("s1", client_a) + b2 = _make_browser("s2", client_a) + daemon._sessions = {"s1": b1, "s2": b2} + + # First session ends → still one live session → no disconnect. + await daemon._on_session_ended("s1") + assert daemon._sessions == {"s2": b2} + assert daemon._clients == {"key-a": client_a, "key-b": client_b} + client_a.disconnect.assert_not_awaited() + client_b.disconnect.assert_not_awaited() + + # Last session ends → all shared clients torn down (as a scheduled task). + await daemon._on_session_ended("s2") + assert daemon._sessions == {} + assert daemon._clients == {} + assert client_a._closed is True + assert client_b._closed is True + await asyncio.sleep(0.05) # let the scheduled disconnect tasks run + client_a.disconnect.assert_awaited_once() + client_b.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_daemon_stop_disconnects_only_when_last_session(): + """/stop pops the sid and disconnects the shared client only on the last.""" + daemon = DaemonServer() + handler = _make_handler(daemon) + shared = AsyncMock() + daemon._clients["key"] = shared + b1 = _make_browser("s1", shared) + b2 = _make_browser("s2", shared) + daemon._sessions = {"s1": b1, "s2": b2} + + await handler._handle_stop({"session_id": "s1"}) + assert daemon._sessions == {"s2": b2} + b1.close.assert_awaited_once() + shared.disconnect.assert_not_awaited() + + await handler._handle_stop({"session_id": "s2"}) + assert daemon._sessions == {} + b2.close.assert_awaited_once() + shared.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_daemon_shutdown_closes_sessions_and_disconnects_clients(): + """Shutdown closes every browser and disconnects every shared client.""" + daemon = DaemonServer() + shared = AsyncMock() + daemon._clients["key"] = shared + b1 = _make_browser("s1", shared) + daemon._sessions = {"s1": b1} + daemon._httpd = None + daemon._loop = Mock() + + await daemon._shutdown() + + b1.close.assert_awaited_once() + shared.disconnect.assert_awaited_once() + assert daemon._sessions == {} + assert daemon._clients == {} + + +@pytest.mark.asyncio +async def test_client_dispatch_invokes_session_ended_hook(): + """The shared client fires _on_session_ended when the relay ends a session.""" + client = Client( + api_key="k", + relay_url="wss://relay/ws/agent", + api_url="https://api", + chat_url="https://chat", + reconnect=False, + ) + client._ws = AsyncMock() + client._ws.send = AsyncMock() + + ended: list[str] = [] + + async def hook(session_id: str) -> None: + ended.append(session_id) + + client._on_session_ended = hook + await client._dispatch({"type": "session.ended", "session_id": "s9", "reason": "completed"}) + assert ended == ["s9"] + + +@pytest.mark.asyncio +async def test_client_dispatch_session_end_alias_invokes_hook(): + """session_end (alias) also reaches the daemon hook.""" + client = Client( + api_key="k", + relay_url="wss://relay/ws/agent", + api_url="https://api", + chat_url="https://chat", + reconnect=False, + ) + client._ws = AsyncMock() + client._ws.send = AsyncMock() + + ended: list[str] = [] + + async def hook(session_id: str) -> None: + ended.append(session_id) + + client._on_session_ended = hook + await client._dispatch({"type": "session_end", "session_id": "s10"}) + assert ended == ["s10"] + + +@pytest.mark.asyncio +async def test_client_dispatch_relay_session_ended_event_id_invokes_hook(): + """The relay's real ``session_ended`` (underscore, id in ``event_id``) must + reach the daemon hook — otherwise relay-initiated ends (provider death, + admin stop, backend reaper) leak the session and its shared WS.""" + client = Client( + api_key="k", + relay_url="wss://relay/ws/agent", + api_url="https://api", + chat_url="https://chat", + reconnect=False, + ) + client._ws = AsyncMock() + client._ws.send = AsyncMock() + + browser = Mock() + browser.session_id = "s11" + browser._on_session_ended = AsyncMock() + client._active_browsers["s11"] = browser + + ended: list[str] = [] + + async def hook(session_id: str) -> None: + ended.append(session_id) + + client._on_session_ended = hook + await client._dispatch({ + "type": "session_ended", + "event_id": "s11", + "reason": "provider_disconnected", + }) + assert ended == ["s11"] + browser._on_session_ended.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_client_dispatch_error_1011_with_event_id_cleans_session(): + """error -1011 (provider death, grace expiry) is a session end: it must + clean the browser AND notify the daemon hook, not just log an unhandled + relay error (which previously left the session/WS alive forever).""" + client = Client( + api_key="k", + relay_url="wss://relay/ws/agent", + api_url="https://api", + chat_url="https://chat", + reconnect=False, + ) + client._ws = AsyncMock() + client._ws.send = AsyncMock() + + browser = Mock() + browser.session_id = "s12" + browser._on_session_ended = AsyncMock() + client._active_browsers["s12"] = browser + + ended: list[str] = [] + + async def hook(session_id: str) -> None: + ended.append(session_id) + + client._on_session_ended = hook + await client._dispatch({ + "type": "error", + "code": -1011, + "event_id": "s12", + "reason": "provider_disconnected", + }) + assert ended == ["s12"] + browser._on_session_ended.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_shared_client_keeps_exactly_one_ws(mock_relay): + """After N rents+closes on a shared client the relay sees EXACTLY 1 WS (not N), + and it drops to 0 once the last session ends. End-to-end over a real local + WebSocket: rent() multiplexes over one connection, session.ended feeds the + daemon hook, and the daemon disconnects the shared client when idle.""" + from ceki_sdk import ConnectOptions, connect + + url = f"ws://127.0.0.1:{mock_relay.port}/ws/agent" + daemon = DaemonServer() + + with patch.dict( + "os.environ", + {"CEKI_FORCE_WS": "1", "CEKI_HUMAN_DISABLE": "1"}, + ): + client = await connect("test-key", ConnectOptions(relay_url=url)) + client._on_session_ended = daemon._on_session_ended + + async def ack_rent(session_id: str) -> None: + await asyncio.sleep(0.05) + ev_id = f"ev-{session_id}" + await mock_relay.send_to_all( + {"type": "rent_pending", "event_id": ev_id, "schedule_id": 1} + ) + await asyncio.sleep(0.02) + await mock_relay.send_to_all({ + "type": "match", + "event_id": ev_id, + "session_id": session_id, + "schedule_id": 1, + "chat_topic_id": None, + "browser_info": {}, + }) + + t1 = asyncio.create_task(ack_rent("sess-A")) + b1 = await client.rent(1) + await t1 + t2 = asyncio.create_task(ack_rent("sess-B")) + b2 = await client.rent(2) + await t2 + + # Both sessions multiplex over the SAME shared WebSocket. + assert len(mock_relay.connections) == 1 + + # Mirror what _handle_rent does: register the shared client + sessions. + daemon._clients = {"test-key": client} + daemon._sessions = {"sess-A": b1, "sess-B": b2} + + # First session ends → still one live session → WS stays up. + await mock_relay.send_to_all( + {"type": "session.ended", "session_id": "sess-A", "reason": "completed"} + ) + await asyncio.sleep(0.15) + assert len(mock_relay.connections) == 1 + assert daemon._sessions == {"sess-B": b2} + + # Last session ends via the relay's REAL format (``session_ended`` + + # ``event_id``, as sent by finishSession on provider death / admin stop) + # → shared client disconnects → no orphan WS. + await mock_relay.send_to_all( + {"type": "session_ended", "event_id": "sess-B", "reason": "completed"} + ) + for _ in range(50): + if len(mock_relay.connections) == 0: + break + await asyncio.sleep(0.05) + assert len(mock_relay.connections) == 0 + assert daemon._sessions == {} + + +@pytest.mark.asyncio +async def test_provider_death_cleans_session_and_ws(mock_relay): + """QA repro: on provider death the relay sends ``session.provider_disconnected``, + then after grace ``error -1011`` + ``session_ended`` (id in ``event_id``). + The daemon must drop the dead session and close the shared WS — previously + ``session_ended`` was silently ignored, leaving a live WS and a zombie entry.""" + from ceki_sdk import ConnectOptions, connect + + url = f"ws://127.0.0.1:{mock_relay.port}/ws/agent" + daemon = DaemonServer() + + with patch.dict( + "os.environ", + {"CEKI_FORCE_WS": "1", "CEKI_HUMAN_DISABLE": "1"}, + ): + client = await connect("test-key", ConnectOptions(relay_url=url)) + client._on_session_ended = daemon._on_session_ended + + async def ack_rent(session_id: str) -> None: + await asyncio.sleep(0.05) + ev_id = f"ev-{session_id}" + await mock_relay.send_to_all( + {"type": "rent_pending", "event_id": ev_id, "schedule_id": 1} + ) + await asyncio.sleep(0.02) + await mock_relay.send_to_all({ + "type": "match", + "event_id": ev_id, + "session_id": session_id, + "schedule_id": 1, + "chat_topic_id": None, + "browser_info": {}, + }) + + t = asyncio.create_task(ack_rent("sess-C")) + browser = await client.rent(1) + await t + + assert len(mock_relay.connections) == 1 + daemon._clients = {"test-key": client} + daemon._sessions = {"sess-C": browser} + + # 1) Provider goes down → grace starts. Session stays tracked (provider + # may rejoin), the WS stays up. + await mock_relay.send_to_all({ + "type": "session.provider_disconnected", + "session_id": "sess-C", + "retry_within_ms": 60000, + }) + await asyncio.sleep(0.1) + assert daemon._sessions == {"sess-C": browser} + assert len(mock_relay.connections) == 1 + + # 2) Grace expires → relay reports the end (exact finishSession payloads). + await mock_relay.send_to_all({ + "type": "error", + "code": -1011, + "event_id": "sess-C", + "reason": "provider_disconnected", + }) + await mock_relay.send_to_all({ + "type": "session_ended", + "event_id": "sess-C", + "reason": "provider_disconnected", + }) + + # Session dropped + shared client disconnected → no orphan WS. + for _ in range(50): + if not daemon._sessions and len(mock_relay.connections) == 0: + break + await asyncio.sleep(0.05) + assert daemon._sessions == {} + assert len(mock_relay.connections) == 0 diff --git a/tests/test_provider.py b/tests/test_provider.py new file mode 100644 index 0000000..4de5910 --- /dev/null +++ b/tests/test_provider.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +from unittest.mock import Mock, patch + +import pytest + +from ceki_sdk._provider import ( + DEFAULT_IMAGE, + ProviderError, + _env_map, + _run_cmd, + resolve_image, + resolve_token, + run_provider, +) + + +def _clean_provider_env(monkeypatch): + for key in ( + "CEKI_PROVIDER_TOKEN", + "PROVIDER_TOKEN", + "CEKI_PROVIDER_IMAGE", + "CEKI_PROVIDER_VIEWPORT", + "CEKI_PROVIDER_LOG_LEVEL", + "TZ", + "DISPLAY", + "CEKI_API_URL", + "CEKI_WS_URL", + ): + monkeypatch.delenv(key, raising=False) + + +# ── token resolution ──────────────────────────────────────────────────────── + + +def test_resolve_token_from_arg(monkeypatch): + _clean_provider_env(monkeypatch) + assert resolve_token("tok-1") == "tok-1" + + +def test_resolve_token_from_env(monkeypatch): + _clean_provider_env(monkeypatch) + monkeypatch.setenv("CEKI_PROVIDER_TOKEN", "env-tok") + assert resolve_token(None) == "env-tok" + assert resolve_token("") == "env-tok" + + +def test_resolve_token_required(monkeypatch): + _clean_provider_env(monkeypatch) + with pytest.raises(ProviderError): + resolve_token(None) + with pytest.raises(ProviderError): + resolve_token(" ") + + +# ── image resolution ──────────────────────────────────────────────────────── + + +def test_resolve_image_default(monkeypatch): + _clean_provider_env(monkeypatch) + assert resolve_image(None) == DEFAULT_IMAGE + + +def test_resolve_image_from_env(monkeypatch): + _clean_provider_env(monkeypatch) + monkeypatch.setenv("CEKI_PROVIDER_IMAGE", "ceki/provider:test") + assert resolve_image(None) == "ceki/provider:test" + + +def test_resolve_image_from_arg(monkeypatch): + _clean_provider_env(monkeypatch) + monkeypatch.setenv("CEKI_PROVIDER_IMAGE", "ceki/provider:env") + assert resolve_image("ceki/provider:arg") == "ceki/provider:arg" + + +# ── container env map ─────────────────────────────────────────────────────── + + +def test_env_map_token(monkeypatch): + _clean_provider_env(monkeypatch) + assert _env_map("tok") == {"CEKI_PROVIDER_TOKEN": "tok"} + + +def test_env_map_public_pass_through(monkeypatch): + _clean_provider_env(monkeypatch) + monkeypatch.setenv("TZ", "Europe/Moscow") + monkeypatch.setenv("DISPLAY", ":1") + env = _env_map("tok") + assert env["TZ"] == "Europe/Moscow" + assert env["DISPLAY"] == ":1" + + +def test_env_map_internal_not_forwarded(monkeypatch): + """Internal docker-browser envs must NOT leak into the public SDK contract.""" + _clean_provider_env(monkeypatch) + monkeypatch.setenv("CEKI_API_URL", "https://api.internal.example") + monkeypatch.setenv("CEKI_WS_URL", "wss://ws.internal.example") + env = _env_map("tok") + assert "CEKI_API_URL" not in env + assert "CEKI_WS_URL" not in env + + +def test_env_map_viewport(monkeypatch): + _clean_provider_env(monkeypatch) + env = _env_map("tok", viewport="1280x720") + assert env["CEKI_PROVIDER_VIEWPORT"] == "1280x720" + + +def test_env_map_verbose(monkeypatch): + _clean_provider_env(monkeypatch) + env = _env_map("tok", verbose=True) + assert env["CEKI_PROVIDER_LOG_LEVEL"] == "DEBUG" + + +# ── docker run command ────────────────────────────────────────────────────── + + +def test_run_cmd_basic(): + cmd = _run_cmd("/usr/bin/docker", DEFAULT_IMAGE, {"CEKI_PROVIDER_TOKEN": "tok"}) + assert cmd == [ + "/usr/bin/docker", "run", "--rm", + "-e", "CEKI_PROVIDER_TOKEN=tok", + DEFAULT_IMAGE, + ] + + +def test_run_cmd_timeout_appends_app_command(): + cmd = _run_cmd( + "docker", + DEFAULT_IMAGE, + {"CEKI_PROVIDER_TOKEN": "tok"}, + timeout=600, + ) + assert cmd[-4:] == [ + "python", "-m", "ceki_browser_provider.app", "--timeout=600", + ] + + +# ── run_provider orchestration ───────────────────────────────────────────── + + +def test_docker_missing_raises(monkeypatch): + _clean_provider_env(monkeypatch) + monkeypatch.setenv("CEKI_PROVIDER_TOKEN", "tok") + with patch("ceki_sdk._provider.shutil.which", return_value=None): + with pytest.raises(ProviderError): + run_provider() + + +def test_run_provider_token_required(monkeypatch): + _clean_provider_env(monkeypatch) + with pytest.raises(ProviderError): + run_provider() + + +def test_run_provider_inspects_pulls_and_runs(monkeypatch): + _clean_provider_env(monkeypatch) + monkeypatch.setenv("CEKI_PROVIDER_TOKEN", "tok") + docker_bin = "/usr/bin/docker" + runs: list[list[str]] = [] + + def fake_run(cmd, **kwargs): + if cmd[:3] == [docker_bin, "image", "inspect"]: + return Mock(returncode=1) # not present → pull next + if cmd[:2] == [docker_bin, "pull"]: + return Mock(returncode=0) + raise AssertionError(f"unexpected run: {cmd}") + + def fake_call(cmd, **kwargs): + runs.append(cmd) + return 0 + + with patch("ceki_sdk._provider.shutil.which", return_value=docker_bin), \ + patch("ceki_sdk._provider.subprocess.run", side_effect=fake_run), \ + patch("ceki_sdk._provider.subprocess.call", side_effect=fake_call): + code = run_provider() + + assert code == 0 + assert runs == [ + [ + docker_bin, "run", "--rm", + "-e", "CEKI_PROVIDER_TOKEN=tok", + DEFAULT_IMAGE, + ] + ] + + +def test_run_provider_pull_failure(monkeypatch): + _clean_provider_env(monkeypatch) + monkeypatch.setenv("CEKI_PROVIDER_TOKEN", "tok") + docker_bin = "/usr/bin/docker" + + def fake_run(cmd, **kwargs): + if cmd[:3] == [docker_bin, "image", "inspect"]: + return Mock(returncode=1) + if cmd[:2] == [docker_bin, "pull"]: + return Mock(returncode=1) + raise AssertionError(f"unexpected run: {cmd}") + + with patch("ceki_sdk._provider.shutil.which", return_value=docker_bin), \ + patch("ceki_sdk._provider.subprocess.run", side_effect=fake_run): + with pytest.raises(ProviderError): + run_provider() + + +def test_run_provider_keyboard_interrupt(monkeypatch): + _clean_provider_env(monkeypatch) + monkeypatch.setenv("CEKI_PROVIDER_TOKEN", "tok") + + def fake_call(cmd, **kwargs): + raise KeyboardInterrupt() + + with patch("ceki_sdk._provider.shutil.which", return_value="/usr/bin/docker"), \ + patch("ceki_sdk._provider.subprocess.run", return_value=Mock(returncode=0)), \ + patch("ceki_sdk._provider.subprocess.call", side_effect=fake_call): + code = run_provider() + assert code == 130