bamboo-coding 0.1.0__py3-none-any.whl

bamboo_coding/client/main.py

@@ -0,0 +1,425 @@
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import logging
+import os
+import signal
+import sys
+from dataclasses import asdict
+from pathlib import Path
+from typing import Awaitable, Callable
+
+from bamboo_coding.client.config import AgentConfig, load_config
+from bamboo_coding.client.controller_client import ControllerConnection
+from bamboo_coding.client.terminal_runtime import ClientTerminalRuntime
+from bamboo_coding.client.journal import RuntimeJournal
+from bamboo_coding.client.status import build_register_message, build_status_message
+from bamboo_coding.shared.protocol import (
+    TaskEventMessage,
+    TaskMessage,
+    TaskResultMessage,
+    TerminalCloseMessage,
+    TerminalInputMessage,
+    TerminalOpenMessage,
+    TerminalResizeMessage,
+)
+
+
+LOGGER = logging.getLogger(__name__)
+INITIAL_RETRY_DELAY = 1
+MAX_RETRY_DELAY = 8
+TaskExecutor = Callable[[TaskMessage], Awaitable[dict]]
+
+
+class ConfigError(RuntimeError):
+    pass
+
+
+class AgentRuntime:
+    def __init__(self, executor: TaskExecutor, config: AgentConfig | None = None):
+        self.config = config or load_config()
+        self.executor = executor
+        self.journal = RuntimeJournal(self.config.journal_path)
+        self.connection = ControllerConnection(self.config.controller_url, token=self.config.token)
+        self.terminals = ClientTerminalRuntime(
+            send_status=self.connection.send_terminal_status,
+            send_output=self.connection.send_terminal_output,
+            send_closed=self.connection.send_terminal_closed,
+            send_heartbeat=self.connection.send_terminal_heartbeat,
+        )
+        self.last_registration_message = ""
+
+    def capabilities(self) -> list[dict]:
+        return [
+            {"capability": "fs.tree", "actions": ["get"]},
+            {"capability": "fs.browse", "actions": ["get"]},
+            {"capability": "fs.file.read", "actions": ["get"]},
+            {"capability": "fs.file.read_raw", "actions": ["get"]},
+            {"capability": "fs.file.write", "actions": ["put"]},
+            {"capability": "fs.folder.create", "actions": ["post"]},
+            {"capability": "fs.path.delete", "actions": ["delete"]},
+            {"capability": "fs.path.rename", "actions": ["put"]},
+            {"capability": "git.diff", "actions": ["get"]},
+            {"capability": "git.status", "actions": ["get"]},
+            {"capability": "git.commits", "actions": ["list"]},
+            {"capability": "git.commit_details", "actions": ["get"]},
+            {"capability": "git.stage", "actions": ["post"]},
+            {"capability": "git.unstage", "actions": ["post"]},
+            {"capability": "git.branches.list", "actions": ["get"]},
+            {"capability": "git.branches.create", "actions": ["post"]},
+            {"capability": "git.branches.checkout", "actions": ["put"]},
+            {"capability": "git.branches.delete", "actions": ["delete"]},
+            {"capability": "git.branches.merge", "actions": ["post"]},
+            {"capability": "git.commit.create", "actions": ["post"]},
+            {"capability": "git.remotes.list", "actions": ["get"]},
+            {"capability": "git.remotes.add", "actions": ["post"]},
+            {"capability": "git.remotes.remove", "actions": ["delete"]},
+            {"capability": "git.remotes.push", "actions": ["post"]},
+            {"capability": "git.remotes.pull", "actions": ["post"]},
+            {"capability": "repo.tree", "actions": ["get"]},
+            {"capability": "repo.browse", "actions": ["get"]},
+            {"capability": "repo.diff", "actions": ["get"]},
+            {"capability": "repo.status", "actions": ["get"]},
+            {"capability": "repo.commits", "actions": ["list"]},
+            {"capability": "repo.commit_details", "actions": ["get"]},
+            {"capability": "repo.file.read", "actions": ["get"]},
+            {"capability": "repo.file.write", "actions": ["put"]},
+            {"capability": "repo.stage", "actions": ["post"]},
+            {"capability": "repo.unstage", "actions": ["post"]},
+            {"capability": "repo.branches.list", "actions": ["get"]},
+            {"capability": "repo.branches.create", "actions": ["post"]},
+            {"capability": "repo.branches.checkout", "actions": ["put"]},
+            {"capability": "repo.branches.delete", "actions": ["delete"]},
+            {"capability": "repo.branches.merge", "actions": ["post"]},
+            {"capability": "repo.commit.create", "actions": ["post"]},
+            {"capability": "terminal.open", "actions": ["post"]},
+            {"capability": "terminal.input", "actions": ["post"]},
+            {"capability": "terminal.resize", "actions": ["post"]},
+            {"capability": "terminal.close", "actions": ["post"]},
+        ]
+
+    def roots(self) -> list[dict]:
+        return [{"path": item.path, "name": item.name} for item in self.config.repositories]
+
+    def repositories(self) -> list[dict]:
+        return [asdict(item) for item in self.config.repositories]
+
+    def runtime_health(self) -> tuple[int, int, int]:
+        active = len(self.journal.list_active_tasks())
+        queued = 0
+        return active, queued, self.config.max_concurrent_tasks
+
+    def _persist_token(self, token: str) -> None:
+        self.config.token_path.parent.mkdir(parents=True, exist_ok=True)
+        self.config.token_path.write_text(json.dumps({"token": token}), encoding="utf-8")
+        self.connection.token = token
+
+    def validate_config(self) -> None:
+        errors = validate_runtime_config(self.config)
+        if errors:
+            raise ConfigError("\n".join(errors))
+
+    def prepare_storage(self) -> None:
+        self.config.token_path.parent.mkdir(parents=True, exist_ok=True)
+        self.config.journal_path.parent.mkdir(parents=True, exist_ok=True)
+
+    async def register_once(self) -> bool:
+        await self.connection.connect()
+        active, queued, max_concurrent = self.runtime_health()
+        register = build_register_message(
+            client_id=self.config.client_id,
+            hostname=self.config.hostname,
+            version=self.config.version,
+            platform_name="python",
+            capabilities=self.capabilities(),
+            roots=self.roots(),
+            active_tasks=active,
+            queued_tasks=queued,
+            max_concurrent_tasks=max_concurrent,
+            connected=True,
+        )
+        registration = await self.connection.send_register(register)
+        self.last_registration_message = str(registration.get("message") or "")
+        approved = bool(registration.get("approved", False))
+        token = registration.get("token")
+        if approved and isinstance(token, str) and token:
+            self._persist_token(token)
+        return approved
+
+    async def _sleep_before_retry(self, delay: int, reason: str, stop_event: asyncio.Event | None = None) -> bool:
+        LOGGER.warning("%s; retrying in %ss", reason, delay)
+        if stop_event is None:
+            await asyncio.sleep(delay)
+            return False
+        try:
+            await asyncio.wait_for(stop_event.wait(), timeout=delay)
+            return True
+        except asyncio.TimeoutError:
+            return False
+
+    async def run(self, stop_event: asyncio.Event | None = None) -> None:
+        self.validate_config()
+        self.prepare_storage()
+        retry_delay = INITIAL_RETRY_DELAY
+
+        while True:
+            if stop_event is not None and stop_event.is_set():
+                break
+            try:
+                approved = await self.register_once()
+                if not approved:
+                    await self.connection.close()
+                    reason = self.last_registration_message or "Client approval is still pending"
+                    if await self._sleep_before_retry(retry_delay, reason, stop_event):
+                        break
+                    retry_delay = min(retry_delay * 2, MAX_RETRY_DELAY)
+                    continue
+                retry_delay = INITIAL_RETRY_DELAY
+                await self.connection.listen(self.handle_task, self.handle_terminal_message)
+                await self.terminals.close_all(reason="controller_disconnected")
+                await self.connection.close()
+                if stop_event is not None and stop_event.is_set():
+                    break
+                if await self._sleep_before_retry(retry_delay, "Controller connection closed", stop_event):
+                    break
+                retry_delay = min(retry_delay * 2, MAX_RETRY_DELAY)
+            except asyncio.CancelledError:
+                await self.terminals.close_all(reason="cancelled")
+                await self.connection.close()
+                raise
+            except ConnectionRefusedError as exc:
+                await self.terminals.close_all(reason="connection_failed")
+                await self.connection.close()
+                if await self._sleep_before_retry(retry_delay, f"Controller connection failed: {exc}", stop_event):
+                    break
+                retry_delay = min(retry_delay * 2, MAX_RETRY_DELAY)
+            except Exception as exc:
+                await self.terminals.close_all(reason="runtime_error")
+                await self.connection.close()
+                if await self._sleep_before_retry(retry_delay, f"Client runtime error: {exc}", stop_event):
+                    break
+                retry_delay = min(retry_delay * 2, MAX_RETRY_DELAY)
+
+        await self.terminals.close_all(reason="stopped")
+        await self.connection.close()
+
+    async def handle_terminal_message(
+        self,
+        message: TerminalOpenMessage | TerminalInputMessage | TerminalResizeMessage | TerminalCloseMessage,
+    ) -> None:
+        if isinstance(message, TerminalOpenMessage):
+            await self.terminals.open_terminal(message)
+        elif isinstance(message, TerminalInputMessage):
+            await self.terminals.input_terminal(message)
+        elif isinstance(message, TerminalResizeMessage):
+            await self.terminals.resize_terminal(message.terminal_id, cols=message.cols, rows=message.rows)
+        elif isinstance(message, TerminalCloseMessage):
+            await self.terminals.close_terminal(message.terminal_id, reason=message.reason or "controller_closed")
+
+    async def handle_task(self, task: TaskMessage) -> None:
+        self.journal.record_task_received(
+            task_id=task.task_id,
+            repo_id=task.repo_id,
+            capability=task.capability,
+            action=task.action,
+            params=task.params,
+        )
+        self.journal.append_event(task_id=task.task_id, event_type="accepted", status="running", payload={"message": "started"})
+        await self.connection.send_event(TaskEventMessage(task_id=task.task_id, status="running", message="started"))
+        try:
+            result = await self.executor(task)
+            self.journal.record_terminal_result(task_id=task.task_id, status="succeeded", result=result)
+            await self.connection.send_result(TaskResultMessage(task_id=task.task_id, success=True, result=result))
+        except Exception as exc:
+            self.journal.record_terminal_result(task_id=task.task_id, status="failed", error=str(exc))
+            await self.connection.send_result(TaskResultMessage(task_id=task.task_id, success=False, error=str(exc)))
+        active, queued, max_concurrent = self.runtime_health()
+        await self.connection.send_status(
+            build_status_message(
+                client_id=self.config.client_id,
+                hostname=self.config.hostname,
+                version=self.config.version,
+                platform_name="python",
+                reason="task_update",
+                capabilities=self.capabilities(),
+                roots=self.roots(),
+                active_tasks=active,
+                queued_tasks=queued,
+                max_concurrent_tasks=max_concurrent,
+                connected=True,
+            )
+        )
+
+
+def validate_runtime_config(config: AgentConfig) -> list[str]:
+    errors: list[str] = []
+    if not config.config_path.exists():
+        errors.append(f"Config file not found at {config.config_path}. Run `bamboo-coding setup` first.")
+    if not config.allowed_roots:
+        errors.append(
+            f"No allowed roots configured in {config.config_path}. Configure [client].allowed_roots or run `bamboo-coding setup`."
+        )
+    return errors
+
+
+def _display_path(path: Path) -> str:
+    resolved = path.expanduser().resolve()
+    home = Path.home().expanduser().resolve()
+    try:
+        relative = resolved.relative_to(home)
+    except ValueError:
+        return str(resolved)
+    return f"~/{relative.as_posix()}"
+
+
+def _toml_string(value: str) -> str:
+    return json.dumps(value)
+
+
+def _render_config_text(
+    *,
+    controller_url: str,
+    client_id: str,
+    hostname: str,
+    allowed_roots: list[str],
+    max_concurrent_tasks: int,
+    token_path: Path,
+    journal_path: Path,
+) -> str:
+    roots_value = ", ".join(_toml_string(root) for root in allowed_roots)
+    return "\n".join(
+        [
+            "[client]",
+            f"controller_url = {_toml_string(controller_url)}",
+            f"client_id = {_toml_string(client_id)}",
+            f"hostname = {_toml_string(hostname)}",
+            f"allowed_roots = [{roots_value}]",
+            f"max_concurrent_tasks = {max_concurrent_tasks}",
+            "",
+            "[storage]",
+            f"token_path = {_toml_string(_display_path(token_path))}",
+            f"journal_path = {_toml_string(_display_path(journal_path))}",
+            "",
+        ]
+    )
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Bamboo Coding client")
+    parser.add_argument("--config", help="Path to bamboo-coding.toml")
+    subparsers = parser.add_subparsers(dest="command")
+
+    setup_parser = subparsers.add_parser("setup", help="Create the default Bamboo Coding client config")
+    setup_parser.add_argument("--controller-url", help="Controller websocket URL")
+    setup_parser.add_argument("--root", action="append", dest="roots", help="Allowed root path; pass more than once for multiple roots")
+    setup_parser.add_argument("--client-id", help="Client identifier advertised to the controller")
+    setup_parser.add_argument("--hostname", help="Hostname advertised to the controller")
+    setup_parser.add_argument("--force", action="store_true", help="Overwrite the config file if it already exists")
+
+    subparsers.add_parser("run", help="Run the Bamboo Coding client")
+    return parser
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    if args.command is None:
+        args.command = "run"
+    return args
+
+
+def setup_command(args: argparse.Namespace) -> int:
+    config = load_config()
+    config_path = config.config_path
+
+    if config_path.exists() and not args.force:
+        print(f"Config already exists at {config_path}. Use --force to overwrite it.", file=sys.stderr)
+        return 1
+
+    roots = [str(Path(root).expanduser().resolve()) for root in (args.roots or list(config.allowed_roots) or [str(Path.cwd().resolve())])]
+    client_id = args.client_id or config.client_id
+    hostname = args.hostname or config.hostname or client_id
+    controller_url = args.controller_url or config.controller_url
+    config_path.parent.mkdir(parents=True, exist_ok=True)
+    config_path.write_text(
+        _render_config_text(
+            controller_url=controller_url,
+            client_id=client_id,
+            hostname=hostname,
+            allowed_roots=roots,
+            max_concurrent_tasks=config.max_concurrent_tasks,
+            token_path=config.token_path,
+            journal_path=config.journal_path,
+        ),
+        encoding="utf-8",
+    )
+    print(f"Wrote config to {config_path}")
+    print("Next: bamboo-coding run")
+    return 0
+
+
+async def _run_runtime(runtime: AgentRuntime) -> None:
+    stop_event = asyncio.Event()
+    loop = asyncio.get_running_loop()
+    registered_signals: list[signal.Signals] = []
+
+    for signum in (signal.SIGINT, signal.SIGTERM):
+        try:
+            loop.add_signal_handler(signum, stop_event.set)
+            registered_signals.append(signum)
+        except (NotImplementedError, RuntimeError):
+            continue
+
+    try:
+        await runtime.run(stop_event=stop_event)
+    finally:
+        for signum in registered_signals:
+            try:
+                loop.remove_signal_handler(signum)
+            except RuntimeError:
+                continue
+
+
+def run_command(_args: argparse.Namespace) -> int:
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
+
+    from bamboo_coding.client.capabilities.repo import execute_task
+
+    config = load_config()
+    errors = validate_runtime_config(config)
+    if errors:
+        for error in errors:
+            print(error, file=sys.stderr)
+        return 1
+
+    repositories = {item.repo_id: item.path for item in config.repositories}
+
+    async def _execute(task: TaskMessage) -> dict:
+        return execute_task(task, repositories)
+
+    runtime = AgentRuntime(_execute, config=config)
+    try:
+        asyncio.run(_run_runtime(runtime))
+    except KeyboardInterrupt:
+        LOGGER.info("Client interrupted")
+    return 0
+
+
+def main(argv: list[str] | None = None) -> None:
+    args = parse_args(argv)
+    if args.config:
+        os.environ["BAMBOO_CODING_CONFIG_PATH"] = args.config
+
+    if args.command == "setup":
+        code = setup_command(args)
+    else:
+        code = run_command(args)
+
+    if code:
+        raise SystemExit(code)
+
+
+if __name__ == "__main__":
+    main()

bamboo_coding/client/status.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from bamboo_coding.shared.protocol import AgentRegisterMessage, ClientStatusMessage
+
+
+def _build_common_payload(*, client_id: str, hostname: str, version: str, platform_name: str, capabilities: list[dict], roots: list[dict] | None, repositories: list[dict] | None, active_tasks: int, queued_tasks: int, max_concurrent_tasks: int, connected: bool) -> dict:
+    root_payload = roots if roots is not None else (repositories or [])
+    return {
+        "protocol_version": 2,
+        "client_id": client_id,
+        "client_info": {
+            "hostname": hostname,
+            "version": version,
+            "platform": platform_name,
+        },
+        "capabilities": capabilities,
+        "roots": root_payload,
+        "connectivity": {"controller": {"connected": connected}},
+        "runtime_health": {
+            "active_tasks": active_tasks,
+            "queued_tasks": queued_tasks,
+            "max_concurrent_tasks": max_concurrent_tasks,
+        },
+    }
+
+
+def build_register_message(*, client_id: str, hostname: str, version: str, platform_name: str, capabilities: list[dict], roots: list[dict] | None = None, repositories: list[dict] | None = None, active_tasks: int, queued_tasks: int, max_concurrent_tasks: int, connected: bool = True) -> AgentRegisterMessage:
+    payload = _build_common_payload(
+        client_id=client_id,
+        hostname=hostname,
+        version=version,
+        platform_name=platform_name,
+        capabilities=capabilities,
+        roots=roots,
+        repositories=repositories,
+        active_tasks=active_tasks,
+        queued_tasks=queued_tasks,
+        max_concurrent_tasks=max_concurrent_tasks,
+        connected=connected,
+    )
+    return AgentRegisterMessage(**payload)
+
+
+def build_status_message(*, client_id: str, hostname: str, version: str, platform_name: str, reason: str, capabilities: list[dict], roots: list[dict] | None = None, repositories: list[dict] | None = None, active_tasks: int, queued_tasks: int, max_concurrent_tasks: int, connected: bool) -> ClientStatusMessage:
+    payload = _build_common_payload(
+        client_id=client_id,
+        hostname=hostname,
+        version=version,
+        platform_name=platform_name,
+        capabilities=capabilities,
+        roots=roots,
+        repositories=repositories,
+        active_tasks=active_tasks,
+        queued_tasks=queued_tasks,
+        max_concurrent_tasks=max_concurrent_tasks,
+        connected=connected,
+    )
+    payload["reason"] = reason
+    return ClientStatusMessage(**payload)