baserun-cli 0.1.0__py3-none-any.whl

baserun_cli/channel.py

@@ -0,0 +1,350 @@
+"""Channel client: connects to nchan over WebSocket with HMAC auth + auto-reconnect.
+
+Models on the fieldshortcut openclaw channel plugin (client.ts/gateway.ts):
+  - signed message = "{app_id}:{timestamp}", HMAC-SHA256 hex
+  - 10-minute timestamp window enforced server-side
+  - reconnect: 5s delay, infinite attempts
+  - receives `type: "task"` frames; dispatches to a handler callback
+
+The client subscribes to its agent channel `agent:{app_id}`. To publish run
+events back, it uses the same WS connection (nchan echoes publishes to
+subscribers) OR a separate HTTP publisher to `run:{run_id}`.
+"""
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import hmac
+import json
+import logging
+import time
+from typing import Any, Awaitable, Callable
+from urllib.parse import urlparse
+
+import httpx
+import websockets
+from websockets.asyncio.client import connect
+
+log = logging.getLogger(__name__)
+
+RECONNECT_DELAY = 5.0
+MAX_RECONNECT_ATTEMPTS = 0  # 0 = infinite
+
+
+def sign(app_secret: str, app_id: str, timestamp: int) -> str:
+    """HMAC-SHA256 hex of '{app_id}:{timestamp}'."""
+    msg = f"{app_id}:{timestamp}".encode()
+    return hmac.new(app_secret.encode(), msg, hashlib.sha256).hexdigest()
+
+
+TaskHandler = Callable[[dict[str, Any]], Awaitable[None]]
+
+
+class ChannelClient:
+    """Connects to nchan as a subscriber on `agent:{app_id}`, dispatches tasks."""
+
+    def __init__(
+        self,
+        nchan_url: str,            # e.g. ws://localhost:9390  (or https→wss)
+        app_id: str,
+        app_secret: str,
+        on_task: TaskHandler | None = None,
+        claim_interval: float = 3.0,  # HTTP poll interval for missed tasks
+    ) -> None:
+        self.nchan_url = nchan_url
+        self.app_id = app_id
+        self.app_secret = app_secret
+        self.on_task = on_task
+        self._stopped = False
+        self._claim_interval = claim_interval
+        self._ws_connected = False  # tracks WS state for claim loop
+        # WS publishers for run:{run_id} events (one connection per run, ordered)
+        pub_base = nchan_url.replace("ws://", "http://").replace("wss://", "https://")
+        self._publish_base = pub_base
+        self._http = httpx.AsyncClient(timeout=10.0)
+        self._run_pubs: dict[str, Any] = {}
+
+    def _ws_url(self) -> str:
+        # nchan channel path: /agent/{app_id}
+        base = self.nchan_url
+        if base.startswith("https://"):
+            base = "wss://" + base[len("https://"):]
+        elif base.startswith("http://"):
+            base = "ws://" + base[len("http://"):]
+        return f"{base.rstrip('/')}/agent/{self.app_id}"
+
+    def _auth_headers(self) -> dict[str, str]:
+        ts = int(time.time())
+        return {
+            "X-Agent-App-Id": self.app_id,
+            "X-Agent-Timestamp": str(ts),
+            "X-Agent-Signature": sign(self.app_secret, self.app_id, ts),
+        }
+
+    async def ensure_run_pub(self, run_id: str) -> None:
+        """Pre-connect the WS publisher for a run (fire-and-forget).
+
+        Called at run start so the WS handshake completes before the first
+        event arrives. If this fails, publish_event will retry on demand.
+        """
+        if run_id in self._run_pubs:
+            return
+        try:
+            ws = await self._open_run_pub(run_id)
+            self._run_pubs[run_id] = ws
+        except Exception as e:
+            log.debug("pre-connect ws pub for run %s failed (will retry on first event): %s", run_id, e)
+
+    async def publish_event(self, run_id: str, payload: dict[str, Any]) -> bool:
+        """Publish a run event via WebSocket (guarantees ordering within a run).
+
+        All events (including terminal) go through the same WS connection to
+        preserve event order in the nchan channel buffer. Mixing WS and HTTP
+        can cause ordering issues because HTTP requests may be processed by
+        nchan at different times relative to buffered WS messages.
+
+        Returns True if published successfully.
+        """
+        data = json.dumps(payload, ensure_ascii=False)
+        is_terminal = (payload.get("data", {}).get("finished") is True)
+        max_retries = 4 if is_terminal else 2
+
+        for attempt in range(max_retries):
+            ws = self._run_pubs.get(run_id)
+            if ws is None:
+                try:
+                    ws = await self._open_run_pub(run_id)
+                    self._run_pubs[run_id] = ws
+                except Exception as e:
+                    log.warning("ws pub connect for run %s failed (attempt %d): %s", run_id, attempt + 1, e)
+                    if attempt < max_retries - 1:
+                        await asyncio.sleep(1.0)
+                    continue
+            try:
+                await ws.send(data)
+                return True
+            except Exception as e:
+                log.warning("ws publish to run %s failed (attempt %d): %s", run_id, attempt + 1, e)
+                self._run_pubs.pop(run_id, None)
+                if attempt < max_retries - 1:
+                    await asyncio.sleep(1.0)
+
+        log.error("publish failed for run %s seq=%s after %d attempts",
+                   run_id, payload.get("data", {}).get("seq"), max_retries)
+        return False
+
+    async def _open_run_pub(self, run_id: str) -> Any:
+        """Open a WS publisher connection for run:{run_id}."""
+        base = self.nchan_url
+        if base.startswith("https://"):
+            base = "wss://" + base[len("https://"):]
+        elif base.startswith("http://"):
+            base = "ws://" + base[len("http://"):]
+        url = f"{base.rstrip('/')}/internal/run/{run_id}/publish"
+        log.info("opening ws publisher for run %s", run_id)
+        return await connect(url, ping_interval=30, open_timeout=30)
+
+    async def _http_publish(self, run_id: str, payload: dict[str, Any]) -> None:
+        """HTTP fallback for publishing events (legacy, not used in normal flow)."""
+        url = f"{self._publish_base.rstrip('/')}/internal/run/{run_id}/publish"
+        resp = await self._http.post(url, content=json.dumps(payload, ensure_ascii=False))
+        resp.raise_for_status()
+
+    async def close_run_pub(self, run_id: str) -> None:
+        """Close the WS publisher for a run."""
+        ws = self._run_pubs.pop(run_id, None)
+        if ws is not None:
+            try:
+                await ws.close()
+            except Exception:
+                pass
+
+    async def verify_and_replay(self, run_id: str, local_events: list[dict]) -> None:
+        """对比 nchan channel 实际状态与本地 events，补发缺失的事件。
+
+        1. 订阅 nchan channel，读取已有消息的 seq 集合
+        2. 对比本地 JSONL 的 seq 集合
+        3. 缺失的事件按 seq 顺序通过 WS 重新发送
+
+        这比盲目重发 failed_events 更可靠——能捕获 ws.send() 返回成功
+        但数据实际未到达 nchan 的情况（TCP buffer 问题）。
+        """
+        if not local_events:
+            return
+
+        # 读取 nchan channel 当前状态
+        url = f"{self._publish_base.rstrip('/')}/internal/run/{run_id}"
+        nchan_seqs: set[int] = set()
+        has_terminal = False
+        try:
+            async with httpx.AsyncClient(timeout=httpx.Timeout(10.0, connect=5.0)) as c:
+                async with c.stream("GET", url, headers={"Accept": "text/event-stream"}) as resp:
+                    async for line in resp.aiter_lines():
+                        if not line.startswith("data:"):
+                            continue
+                        raw = line[5:].strip()
+                        if not raw:
+                            continue
+                        try:
+                            msg = json.loads(raw)
+                        except json.JSONDecodeError:
+                            continue
+                        data = msg.get("data") if isinstance(msg, dict) else msg
+                        if not isinstance(data, dict):
+                            continue
+                        seq = data.get("seq", 0)
+                        if seq:
+                            nchan_seqs.add(seq)
+                        if data.get("finished") and data.get("kind") in ("final", "error"):
+                            has_terminal = True
+                            break  # terminal found = all events present
+        except Exception as e:
+            log.warning("verify: failed to read nchan for run %s: %s", run_id, e)
+            # nchan 读取失败，保守重发所有事件
+            nchan_seqs = set()
+
+        # terminal 已在 nchan 中 = 全部到齐
+        if has_terminal:
+            log.info("verify: run %s terminal present in nchan, all good", run_id)
+            return
+
+        # 找出缺失的事件
+        local_seqs = {ev.get("data", {}).get("seq", 0) for ev in local_events}
+        missing_seqs = local_seqs - nchan_seqs
+
+        if not missing_seqs and has_terminal:
+            return
+
+        if not missing_seqs:
+            # nchan 有所有非 terminal 事件但缺 terminal → 重发 terminal
+            missing_seqs = {ev.get("data", {}).get("seq", 0) for ev in local_events
+                           if ev.get("data", {}).get("finished") is True} - nchan_seqs
+
+        if not missing_seqs:
+            log.info("verify: run %s no missing events", run_id)
+            return
+
+        # 按 seq 排序重发
+        missing_events = sorted(
+            [ev for ev in local_events if ev.get("data", {}).get("seq", 0) in missing_seqs],
+            key=lambda ev: ev.get("data", {}).get("seq", 0),
+        )
+        log.info(
+            "verify: run %s replaying %d missing events (seqs: %s, nchan has %d/%d)",
+            run_id, len(missing_events),
+            [ev["data"]["seq"] for ev in missing_events],
+            len(nchan_seqs), len(local_seqs),
+        )
+
+        for ev in missing_events:
+            ok = await self.publish_event(run_id, ev)
+            if not ok:
+                log.error("verify: replay still failed for run %s seq=%s", run_id, ev["data"]["seq"])
+
+    async def run(self) -> None:
+        """Main loop: WS subscriber + HTTP claim-task poller, running in parallel.
+
+        - WS subscriber: instant task delivery when connected (push path)
+        - HTTP poller: claims missed tasks every N seconds (pull path)
+        Both dispatch to the same handler; dedup is handled by TaskRunner.
+        """
+        await asyncio.gather(
+            self._ws_loop(),
+            self._claim_loop(),
+        )
+
+    async def _ws_loop(self) -> None:
+        """WS connect-dispatch-reconnect loop (push path)."""
+        attempts = 0
+        while not self._stopped:
+            try:
+                await self._connect_and_serve()
+                attempts = 0
+            except Exception as e:
+                log.error("channel connection error: %s", e)
+            if self._stopped:
+                break
+            if MAX_RECONNECT_ATTEMPTS > 0 and attempts >= MAX_RECONNECT_ATTEMPTS:
+                break
+            attempts += 1
+            log.info("reconnecting in %.0fs (attempt %d)...", RECONNECT_DELAY, attempts)
+            await asyncio.sleep(RECONNECT_DELAY)
+
+    async def _claim_once(self) -> None:
+        """Single claim sweep — used after WS reconnect to catch missed tasks."""
+        try:
+            url = f"{self._publish_base.rstrip('/')}/api/agent/{self.app_id}/claim-task"
+            resp = await self._http.post(url)
+            if resp.status_code == 200:
+                msg = resp.json()
+                if msg.get("type") == "task":
+                    log.info("claimed missed task after reconnect: run %s", msg.get("data", {}).get("run_id"))
+                    if self.on_task:
+                        await self._safe_dispatch(msg)
+                    # recurse to drain the queue
+                    await self._claim_once()
+        except Exception as e:
+            log.debug("claim-once failed: %s", e)
+
+    async def _claim_loop(self) -> None:
+        """HTTP poll for missed tasks (pull path).
+
+        Only polls when WS is disconnected. On reconnect, does one immediate
+        claim sweep (in case tasks were queued during disconnection), then
+        goes idle while WS handles real-time delivery.
+        """
+        await asyncio.sleep(1.0)  # initial delay (let WS connect first)
+        while not self._stopped:
+            if self._ws_connected:
+                # WS is connected — push path handles delivery, no need to poll
+                await asyncio.sleep(self._claim_interval)
+                continue
+
+            # WS disconnected — poll for queued tasks
+            try:
+                url = f"{self._publish_base.rstrip('/')}/api/agent/{self.app_id}/claim-task"
+                resp = await self._http.post(url)
+                if resp.status_code == 200:
+                    msg = resp.json()
+                    if msg.get("type") == "task":
+                        log.info("claimed task via HTTP: run %s", msg.get("data", {}).get("run_id"))
+                        if self.on_task:
+                            await self._safe_dispatch(msg)
+                        continue  # immediately try again (might be more queued)
+            except Exception as e:
+                log.debug("claim-task poll failed: %s", e)
+            await asyncio.sleep(self._claim_interval)
+
+    async def _connect_and_serve(self) -> None:
+        url = self._ws_url()
+        headers = self._auth_headers()
+        log.info("connecting to %s as %s", url, self.app_id)
+        async with connect(url, additional_headers=headers, ping_interval=30) as ws:
+            self._ws_connected = True
+            log.info("connected; waiting for tasks")
+            # sweep for tasks queued during disconnection
+            asyncio.create_task(self._claim_once())
+            async for raw in ws:
+                try:
+                    msg = json.loads(raw)
+                except (json.JSONDecodeError, TypeError):
+                    continue
+                if msg.get("type") != "task":
+                    continue
+                if self.on_task is None:
+                    log.warning("received task but no handler bound; dropping")
+                    continue
+                # dispatch without blocking the recv loop
+                asyncio.create_task(self._safe_dispatch(msg))
+            # WS dropped — reset flag so claim loop starts polling
+            self._ws_connected = False
+
+    async def _safe_dispatch(self, msg: dict[str, Any]) -> None:
+        try:
+            assert self.on_task is not None
+            await self.on_task(msg)
+        except Exception:
+            log.exception("task handler error for %s", msg.get("data", {}).get("run_id"))
+
+    def stop(self) -> None:
+        self._stopped = True

baserun_cli/main.py

@@ -0,0 +1,98 @@
+"""Agent client entry point.
+
+Connects to the nchan bus as the agent's channel subscriber, receives task
+envelopes, spawns the CLI agent (claude/codex/...), parses streaming output,
+and publishes run events back to run:{run_id}.
+
+Config via env:
+  NCHAN_URL        e.g. ws://localhost:9390
+  AGENT_APP_ID     channel credential (issued by the server)
+  AGENT_APP_SECRET HMAC secret
+  CONNECTOR_TYPE   claude_code | codex | <custom>
+  CLI_CONFIG_*     JSON override for the CLISpec (bin, workdir, env, ...)
+  CONCURRENCY      how many runs to run in parallel (default 1)
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import sys
+
+from .channel import ChannelClient
+from .runner import TaskRunner
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(name)s %(levelname)s %(message)s",
+)
+log = logging.getLogger("baserun-cli")
+
+
+def _load_cli_config() -> dict:
+    """Read CLISpec overrides from CLI_CONFIG env (JSON) if present."""
+    raw = os.environ.get("CLI_CONFIG")
+    if not raw:
+        return {}
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError as e:
+        log.warning("CLI_CONFIG invalid JSON, ignoring: %s", e)
+        return {}
+
+
+async def main() -> None:
+    nchan_url = os.environ.get("NCHAN_URL", "ws://localhost:9390")
+    app_id = os.environ.get("AGENT_APP_ID", "")
+    app_secret = os.environ.get("AGENT_APP_SECRET", "")
+    connector_type = os.environ.get("CONNECTOR_TYPE", "claude_code")
+    concurrency = int(os.environ.get("CONCURRENCY", "1"))
+
+    if not app_id or not app_secret:
+        log.error("AGENT_APP_ID / AGENT_APP_SECRET must be set")
+        sys.exit(1)
+
+    cli_config = _load_cli_config()
+
+    # channel created first (handler bound after runner exists), breaking the
+    # chicken-and-egg (channel needs on_task; on_task needs runner; runner needs channel).
+    channel = ChannelClient(nchan_url, app_id, app_secret)
+    runner = TaskRunner(channel, concurrency=concurrency, app_id=app_id)
+
+    async def on_task(envelope: dict) -> None:
+        # enrich the data with our connector_type + cli_config so the runner
+        # builds the right CLISpec
+        data = envelope.get("data") or envelope
+        data.setdefault("connector_type", connector_type)
+        # merge cli_config under config.cli_spec so resolve_spec picks it up
+        cfg = data.get("config") or {}
+        if cli_config:
+            cfg.setdefault("cli_spec", {}).update(cli_config.get("cli_spec", cli_config))
+            if workdir := cli_config.get("workdir"):
+                cfg.setdefault("workdir", workdir)
+            if bin_ := cli_config.get("bin"):
+                cfg.setdefault("bin", bin_)
+        data["config"] = cfg
+        envelope["data"] = data
+        await runner.handle(envelope)
+
+    channel.on_task = on_task
+
+    log.info(
+        "agent client starting: app_id=%s connector=%s concurrency=%d nchan=%s",
+        app_id, connector_type, concurrency, nchan_url,
+    )
+    await channel.run()
+
+
+def main_sync() -> None:
+    """Synchronous entry point for the console script."""
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        log.info("shutting down")
+
+
+if __name__ == "__main__":
+    main_sync()