helm-server 0.1.0__py3-none-any.whl

helm_server/acp.py

@@ -0,0 +1,722 @@
+"""ACP (Agent Client Protocol) runner for long-lived cursor-agent sessions.
+
+`agent acp` speaks JSON-RPC on stdio. That is the only path that can surface
+permission prompts, AskQuestion, plan approval, and todos. The print runner
+stays the default (`AGENT_SESSION_RUNNER=print`); set `acp` to use this.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import os
+import queue
+import subprocess
+import threading
+import time
+from typing import Callable, Optional
+
+HOME = os.path.expanduser("~")
+LineCallback = Callable[[str], None]
+
+
+def _agent_bin() -> str:
+    from helm_server import runners
+
+    return runners.cursor_acp_cli()
+
+
+class AcpError(RuntimeError):
+    pass
+
+
+EXITED = "cursor-agent acp exited; run `cursor-agent update` and `cursor-agent status`"
+
+
+_INSTRUCTIONS_MARK = "Repo instructions (loaded for this session):"
+_INSTRUCTIONS_LIMIT = 60_000
+
+
+def _split_frontmatter(text: str) -> tuple[dict, str]:
+    if not text.startswith("---\n"):
+        return {}, text
+    end = text.find("\n---\n", 4)
+    if end < 0:
+        return {}, text
+    header: dict[str, str] = {}
+    for line in text[4:end].splitlines():
+        if ":" not in line:
+            continue
+        key, value = line.split(":", 1)
+        header[key.strip().lower()] = value.strip()
+    return header, text[end + 5 :]
+
+
+def always_apply_rules(cwd: str) -> list[tuple[str, str]]:
+    """alwaysApply rule bodies from the repo, in filename order."""
+    rules_dir = os.path.join(cwd, ".cursor", "rules")
+    if not os.path.isdir(rules_dir):
+        return []
+    found: list[tuple[str, str]] = []
+    for name in sorted(os.listdir(rules_dir)):
+        if not name.endswith((".mdc", ".md")):
+            continue
+        path = os.path.join(rules_dir, name)
+        try:
+            with open(path, encoding="utf-8") as handle:
+                raw = handle.read()
+        except OSError:
+            continue
+        header, body = _split_frontmatter(raw)
+        if header.get("alwaysapply", "").lower() != "true":
+            continue
+        body = body.strip()
+        if body:
+            found.append((name, body))
+    return found
+
+
+def project_instructions(cwd: str, limit: int = _INSTRUCTIONS_LIMIT) -> str:
+    """AGENTS.md plus always-on rules. `cursor-agent -p` loads these; ACP does not."""
+    chunks: list[str] = []
+    for name, body in always_apply_rules(cwd):
+        chunks.append("## .cursor/rules/%s\n%s" % (name, body))
+    agents = os.path.join(cwd, "AGENTS.md")
+    if os.path.isfile(agents):
+        try:
+            with open(agents, encoding="utf-8") as handle:
+                text = handle.read().strip()
+        except OSError:
+            text = ""
+        if text:
+            chunks.append("## AGENTS.md\n" + text)
+    blob = "\n\n".join(chunks).strip()
+    if len(blob) > limit:
+        blob = blob[:limit].rstrip() + "\n\n[truncated]"
+    return blob
+
+
+def with_project_instructions(cwd: str, prompt: str, *, first_turn: bool = True) -> str:
+    """Attach repo rules on the first ACP turn only.
+
+    Re-wrapping every follow-up as "loaded for this session" / "User request:"
+    made short answers like "B" look like a brand-new chat with no prior A/B.
+    """
+    prompt = prompt or ""
+    if not first_turn:
+        return prompt
+    if prompt.startswith(_INSTRUCTIONS_MARK):
+        return prompt
+    body = project_instructions(cwd)
+    if not body:
+        return prompt
+    return _INSTRUCTIONS_MARK + "\n\n" + body + "\n\nUser request:\n" + prompt
+
+
+_IMAGE_MIME = {
+    ".jpg": "image/jpeg",
+    ".jpeg": "image/jpeg",
+    ".png": "image/png",
+    ".webp": "image/webp",
+    ".gif": "image/gif",
+}
+
+
+def image_mime(path: str) -> str:
+    ext = os.path.splitext(path or "")[1].lower()
+    return _IMAGE_MIME.get(ext, "image/png")
+
+
+def image_prompt_block(path: str) -> Optional[dict]:
+    """ACP `session/prompt` image block. file:// resource URIs are rejected."""
+    if not path or not os.path.isfile(path):
+        return None
+    with open(path, "rb") as handle:
+        raw = handle.read()
+    if not raw:
+        return None
+    return {
+        "type": "image",
+        "mimeType": image_mime(path),
+        "data": base64.b64encode(raw).decode("ascii"),
+        "uri": "file://%s" % os.path.abspath(path),
+    }
+
+
+def prompt_blocks(prompt: str, images: Optional[list] = None) -> list:
+    blocks = [{"type": "text", "text": prompt}]
+    for path in images or []:
+        block = image_prompt_block(path)
+        if block:
+            blocks.append(block)
+    return blocks
+
+
+def _model_choice_id(item) -> str:
+    if isinstance(item, str):
+        return item.strip()
+    if isinstance(item, dict):
+        return str(item.get("modelId") or item.get("id") or "").strip()
+    return ""
+
+
+def _model_choice_name(item) -> str:
+    if isinstance(item, str):
+        return item.strip()
+    if isinstance(item, dict):
+        return str(item.get("name") or item.get("label") or "").strip()
+    return ""
+
+
+def _compact_model_token(value: str) -> str:
+    return "".join(ch for ch in (value or "").lower() if ch.isalnum())
+
+
+def resolve_model_id(requested: str, available) -> str:
+    """Map a CLI `--list-models` id onto an ACP `modelId` when the labels differ."""
+    needle = (requested or "").strip()
+    if not needle:
+        return ""
+    pairs = []
+    for item in available or []:
+        mid = _model_choice_id(item)
+        name = _model_choice_name(item)
+        if mid or name:
+            pairs.append((mid, name))
+    for mid, name in pairs:
+        if mid == needle or name == needle:
+            return mid or needle
+    lower = needle.lower()
+    exact = [mid or name for mid, name in pairs if mid.lower() == lower or name.lower() == lower]
+    if len(exact) == 1:
+        return exact[0]
+    compact = _compact_model_token(needle)
+    fuzzy = [
+        mid or name
+        for mid, name in pairs
+        if compact
+        and compact
+        in (_compact_model_token(mid) + _compact_model_token(name))
+    ]
+    if len(fuzzy) == 1:
+        return fuzzy[0]
+    return needle
+
+
+class AcpClient:
+    """One `agent acp` process. Tests inject a transport instead of spawning."""
+
+    def __init__(self, cwd: str, env: Optional[dict] = None, transport=None):
+        self.cwd = cwd
+        self.env = env or os.environ.copy()
+        self._id = 0
+        self._pending: dict[int, queue.Queue] = {}
+        self._lock = threading.Lock()
+        self._inbox: queue.Queue = queue.Queue()
+        self._closed = threading.Event()
+        self.session_id = ""
+        self.models: list[dict] = []
+        self.modes: list[dict] = []
+        self.current_model = ""
+        self.current_mode = "agent"
+        # Rules go out once per live ACP process; follow-ups keep the bare prompt.
+        self.instructions_sent = False
+        if transport is not None:
+            self.proc = transport
+            self._stdin = transport.stdin
+            self._stdout = transport.stdout
+        else:
+            self.proc = subprocess.Popen(
+                [_agent_bin(), "acp"],
+                cwd=cwd,
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.DEVNULL,
+                text=True,
+                bufsize=1,
+                env=self.env,
+            )
+            self._stdin = self.proc.stdin
+            self._stdout = self.proc.stdout
+        self._reader = threading.Thread(target=self._read_loop, daemon=True, name="acp-read")
+        self._reader.start()
+
+    def close(self) -> None:
+        self._closed.set()
+        try:
+            if self._stdin:
+                self._stdin.close()
+        except Exception:
+            pass
+        try:
+            if getattr(self.proc, "poll", None) is None or self.proc.poll() is None:
+                self.proc.terminate()
+        except Exception:
+            pass
+
+    def _read_loop(self) -> None:
+        stream = self._stdout
+        if stream is None:
+            return
+        try:
+            while not self._closed.is_set():
+                line = stream.readline()
+                if not line:
+                    break
+                try:
+                    msg = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+                if not isinstance(msg, dict):
+                    continue
+                mid = msg.get("id")
+                if mid is not None and ("result" in msg or "error" in msg):
+                    waiter = self._pending.get(int(mid))
+                    if waiter is not None:
+                        waiter.put(msg)
+                    continue
+                self._inbox.put(msg)
+        finally:
+            self._closed.set()
+            with self._lock:
+                waiters = list(self._pending.values())
+            for waiter in waiters:
+                waiter.put({"error": {"message": EXITED}})
+
+    def request(self, method: str, params=None, timeout: float = 60) -> dict:
+        with self._lock:
+            if self._closed.is_set():
+                raise AcpError(EXITED)
+            self._id += 1
+            rid = self._id
+            waiter: queue.Queue = queue.Queue()
+            self._pending[rid] = waiter
+        payload = {"jsonrpc": "2.0", "id": rid, "method": method}
+        if params is not None:
+            payload["params"] = params
+        try:
+            if self._stdin is None:
+                raise AcpError("ACP stdin closed")
+            self._stdin.write(json.dumps(payload) + "\n")
+            self._stdin.flush()
+        except (OSError, ValueError) as exc:
+            self._pending.pop(rid, None)
+            raise AcpError(EXITED) from exc
+        try:
+            msg = waiter.get(timeout=timeout)
+        except queue.Empty:
+            raise AcpError("ACP timed out on %s" % method)
+        finally:
+            self._pending.pop(rid, None)
+        if msg.get("error"):
+            err = msg["error"]
+            raise AcpError(str(err.get("message") or err))
+        result = msg.get("result")
+        return result if isinstance(result, dict) else {}
+
+    def respond(self, req_id, result) -> None:
+        if self._stdin is None:
+            return
+        self._stdin.write(json.dumps({"jsonrpc": "2.0", "id": req_id, "result": result}) + "\n")
+        self._stdin.flush()
+
+    def poll_notice(self, timeout: float = 0.05) -> Optional[dict]:
+        try:
+            return self._inbox.get(timeout=max(0.0, timeout))
+        except queue.Empty:
+            return None
+
+    def initialize(self) -> dict:
+        return self.request(
+            "initialize",
+            {
+                "protocolVersion": 1,
+                "clientCapabilities": {
+                    "fs": {"readTextFile": False, "writeTextFile": False},
+                    "terminal": False,
+                },
+                "clientInfo": {"name": "agent-session", "version": "1"},
+            },
+            timeout=30,
+        )
+
+    def authenticate(self) -> dict:
+        return self.request("authenticate", {"methodId": "cursor_login"}, timeout=30)
+
+    def session_new(self, cwd: str, mcp_servers=None) -> dict:
+        result = self.request(
+            "session/new",
+            {"cwd": cwd, "mcpServers": mcp_servers or []},
+            timeout=30,
+        )
+        self._remember_session(result)
+        return result
+
+    def session_load(self, session_id: str, cwd: str = "", mcp_servers=None) -> dict:
+        # cursor-agent answers "Internal error" when cwd or mcpServers is missing.
+        result = self.request(
+            "session/load",
+            {"sessionId": session_id, "cwd": cwd or self.cwd, "mcpServers": mcp_servers or []},
+            timeout=30,
+        )
+        self._remember_session(result, fallback=session_id)
+        return result
+
+    def session_prompt(self, prompt: str, images: Optional[list] = None, timeout: float = 1800) -> dict:
+        return self.request(
+            "session/prompt",
+            {"sessionId": self.session_id, "prompt": prompt_blocks(prompt, images)},
+            timeout=timeout,
+        )
+
+    def session_cancel(self) -> None:
+        if not self.session_id or self._stdin is None:
+            return
+        self._stdin.write(
+            json.dumps(
+                {
+                    "jsonrpc": "2.0",
+                    "method": "session/cancel",
+                    "params": {"sessionId": self.session_id},
+                }
+            )
+            + "\n"
+        )
+        self._stdin.flush()
+
+    def set_mode(self, mode: str) -> None:
+        if mode not in ("agent", "plan", "ask") or not self.session_id:
+            return
+        try:
+            self.request("session/set_mode", {"sessionId": self.session_id, "modeId": mode}, timeout=15)
+            self.current_mode = mode
+        except AcpError:
+            pass
+
+    def set_model(self, model: str) -> None:
+        if not model or not self.session_id:
+            return
+        candidates = []
+        resolved = resolve_model_id(model, self.models)
+        for candidate in (resolved, model):
+            if candidate and candidate not in candidates:
+                candidates.append(candidate)
+        for candidate in candidates:
+            try:
+                self.request(
+                    "session/set_model",
+                    {"sessionId": self.session_id, "modelId": candidate},
+                    timeout=15,
+                )
+                self.current_model = candidate
+                return
+            except AcpError:
+                continue
+
+    def _remember_session(self, result: dict, fallback: str = "") -> None:
+        self.session_id = str(result.get("sessionId") or fallback)
+        modes = result.get("modes") if isinstance(result.get("modes"), dict) else {}
+        self.current_mode = str(modes.get("currentModeId") or self.current_mode or "agent")
+        self.modes = list(modes.get("availableModes") or [])
+        models = result.get("models") if isinstance(result.get("models"), dict) else {}
+        self.current_model = str(models.get("currentModelId") or self.current_model)
+        self.models = list(models.get("availableModels") or [])
+
+
+def _emit(on_line: LineCallback, payload: dict) -> None:
+    on_line(json.dumps(payload))
+
+
+def acp_notice_to_lines(msg: dict) -> list[dict]:
+    """Map an ACP notification to stream-json-shaped objects."""
+    method = msg.get("method")
+    params = msg.get("params") if isinstance(msg.get("params"), dict) else {}
+    if method == "session/update":
+        update = params.get("update") if isinstance(params.get("update"), dict) else params
+        kind = update.get("sessionUpdate") or update.get("session_update")
+        content = update.get("content") if isinstance(update.get("content"), dict) else {}
+        text = content.get("text") if isinstance(content.get("text"), str) else update.get("text")
+        if kind in ("agent_message_chunk", "agent_message") and text:
+            return [
+                {
+                    "type": "assistant",
+                    "timestamp_ms": int(time.time() * 1000),
+                    "message": {"role": "assistant", "content": [{"type": "text", "text": text}]},
+                }
+            ]
+        if kind in ("agent_thought_chunk", "agent_thought") and text:
+            return [{"type": "thinking", "subtype": "delta", "text": text}]
+        if kind in ("tool_call", "tool_call_update"):
+            status = "completed" if kind == "tool_call_update" and update.get("status") in (
+                "completed",
+                "failed",
+            ) else "started"
+            title = str(update.get("title") or update.get("kind") or "tool")
+            call_id = str(update.get("toolCallId") or update.get("tool_call_id") or "")
+            return [
+                {
+                    "type": "tool_call",
+                    "subtype": status,
+                    "call_id": call_id,
+                    "tool_call": {"function": {"name": title, "arguments": json.dumps(update.get("rawInput") or {})}},
+                }
+            ]
+        if kind == "plan":
+            return [{"type": "thinking", "subtype": "delta", "text": str(update.get("text") or update.get("plan") or "")}]
+        return []
+    if method == "cursor/update_todos":
+        return [{"type": "thinking", "subtype": "delta", "text": ""}]
+    return []
+
+
+def pending_from_notice(msg: dict) -> Optional[dict]:
+    method = msg.get("method")
+    params = msg.get("params") if isinstance(msg.get("params"), dict) else {}
+    req_id = msg.get("id")
+    if method == "session/request_permission":
+        options = params.get("options") or [
+            {"optionId": "allow-once", "name": "Allow once"},
+            {"optionId": "allow-always", "name": "Always allow"},
+            {"optionId": "reject-once", "name": "Reject"},
+        ]
+        return {
+            "kind": "permission",
+            "rpcId": req_id,
+            "title": str(params.get("title") or "Permission"),
+            "detail": str(params.get("toolCall") or params.get("description") or ""),
+            "options": options,
+        }
+    if method == "cursor/ask_question":
+        return {
+            "kind": "question",
+            "rpcId": req_id,
+            "title": str(params.get("title") or "Question"),
+            "questions": params.get("questions") or [],
+        }
+    if method == "cursor/create_plan":
+        return {
+            "kind": "plan",
+            "rpcId": req_id,
+            "title": str(params.get("name") or "Plan"),
+            "overview": str(params.get("overview") or ""),
+            "plan": str(params.get("plan") or ""),
+            "todos": params.get("todos") or [],
+        }
+    if method == "cursor/update_todos":
+        return {
+            "kind": "todos",
+            "rpcId": None,
+            "todos": params.get("todos") or [],
+            "merge": bool(params.get("merge")),
+        }
+    if method == "cursor/task":
+        return {
+            "kind": "task",
+            "rpcId": None,
+            "description": str(params.get("description") or ""),
+            "subagentType": params.get("subagentType"),
+        }
+    return None
+
+
+def auto_permission_result(policy: str) -> Optional[dict]:
+    if policy in ("yolo", "auto"):
+        return {"outcome": {"outcome": "selected", "optionId": "allow-always"}}
+    if policy == "auto-review":
+        return {"outcome": {"outcome": "selected", "optionId": "allow-once"}}
+    return None
+
+
+class AcpRunner:
+    """CursorRunner-compatible wrapper around one ACP process per chat."""
+
+    def __init__(self):
+        self._clients: dict[str, AcpClient] = {}
+        self._lock = threading.Lock()
+        self.on_pending: Optional[Callable[[dict], None]] = None
+        self.answer: Optional[Callable[[str], Optional[dict]]] = None
+        self.policy = "yolo"
+
+    def create_chat(self, cwd: str) -> Optional[str]:
+        client = AcpClient(cwd)
+        try:
+            client.initialize()
+            client.authenticate()
+            client.session_new(cwd)
+        except Exception:
+            client.close()
+            return None
+        if not client.session_id:
+            client.close()
+            return None
+        with self._lock:
+            self._clients[client.session_id] = client
+        return client.session_id
+
+    def persist_list(self) -> str:
+        try:
+            proc = subprocess.run(
+                [_agent_bin(), "persist", "list"],
+                capture_output=True,
+                text=True,
+                timeout=30,
+            )
+        except Exception as exc:
+            return str(exc)
+        return ((proc.stdout or "") + (proc.stderr or "")).strip()
+
+    def list_models(self) -> str:
+        try:
+            proc = subprocess.run(
+                [_agent_bin(), "--list-models"],
+                capture_output=True,
+                text=True,
+                timeout=30,
+            )
+        except Exception as exc:
+            return str(exc)
+        return ((proc.stdout or "") + (proc.stderr or "")).strip()
+
+    def _client(self, cwd: str, resume_id: Optional[str], env: Optional[dict]) -> AcpClient:
+        if resume_id:
+            with self._lock:
+                existing = self._clients.get(resume_id)
+            if existing and not existing._closed.is_set():
+                return existing
+        client = AcpClient(cwd, env=env)
+        client.initialize()
+        client.authenticate()
+        if resume_id:
+            try:
+                client.session_load(resume_id, cwd)
+            except AcpError:
+                client.session_new(cwd)
+        else:
+            client.session_new(cwd)
+        with self._lock:
+            if client.session_id:
+                self._clients[client.session_id] = client
+            if resume_id and resume_id != client.session_id:
+                self._clients[resume_id] = client
+        return client
+
+    def run(
+        self,
+        cwd: str,
+        prompt: str,
+        resume_id: Optional[str],
+        mode: Optional[str],
+        on_line: LineCallback,
+        env: Optional[dict] = None,
+        timeout: int = 1800,
+        model: Optional[str] = None,
+        cancel: Optional[threading.Event] = None,
+        images: Optional[list] = None,
+    ) -> int:
+        try:
+            client = self._client(cwd, resume_id, env)
+        except Exception as exc:
+            _emit(on_line, {"type": "text", "text": "\n[acp start failed: %s]\n" % exc})
+            return 1
+        if mode:
+            client.set_mode(mode)
+        if model:
+            client.set_model(model)
+        if client.session_id:
+            _emit(
+                on_line,
+                {
+                    "type": "system",
+                    "subtype": "init",
+                    "session_id": client.session_id,
+                    "model": client.current_model or model or "",
+                    "cwd": cwd,
+                    "permissionMode": self.policy,
+                },
+            )
+
+        done = threading.Event()
+        outcome: dict = {"rc": 1, "final": []}
+
+        def prompt_thread():
+            try:
+                first_turn = not client.instructions_sent
+                instructed = with_project_instructions(cwd, prompt, first_turn=first_turn)
+                client.instructions_sent = True
+                result = client.session_prompt(instructed, images=images, timeout=timeout)
+                outcome["rc"] = 0 if result.get("stopReason") != "cancelled" else 130
+                text = result.get("result") or result.get("text")
+                if isinstance(text, str) and text:
+                    outcome["final"].append(
+                        {
+                            "type": "result",
+                            "subtype": "success",
+                            "is_error": False,
+                            "result": text,
+                            "session_id": client.session_id,
+                        }
+                    )
+            except AcpError as exc:
+                outcome["final"].append({"type": "text", "text": "\n[%s]\n" % exc})
+                outcome["rc"] = 1
+            finally:
+                done.set()
+
+        thread = threading.Thread(target=prompt_thread, daemon=True, name="acp-prompt")
+        thread.start()
+        deadline = time.time() + timeout
+        while not done.is_set():
+            if cancel is not None and cancel.is_set():
+                client.session_cancel()
+                done.wait(5)
+                return 130
+            if time.time() >= deadline:
+                client.session_cancel()
+                _emit(on_line, {"type": "text", "text": "\n[timed out]\n"})
+                return 124
+            notice = client.poll_notice(0.1)
+            if not notice:
+                continue
+            pending = pending_from_notice(notice)
+            if pending and pending.get("kind") == "permission":
+                auto = auto_permission_result(self.policy)
+                if auto and pending.get("rpcId") is not None:
+                    client.respond(pending["rpcId"], auto)
+                    pending = None
+            if pending and pending.get("kind") in ("permission", "question", "plan"):
+                if self.on_pending:
+                    self.on_pending(pending)
+                answered = None
+                wait_until = time.time() + 300
+                while time.time() < wait_until and not done.is_set():
+                    if cancel is not None and cancel.is_set():
+                        if pending.get("rpcId") is not None:
+                            client.respond(pending["rpcId"], {"outcome": {"outcome": "cancelled"}})
+                        client.session_cancel()
+                        return 130
+                    if self.answer:
+                        answered = self.answer(str(pending.get("rpcId")))
+                    if answered is not None:
+                        break
+                    time.sleep(0.1)
+                if pending.get("rpcId") is not None:
+                    client.respond(pending["rpcId"], answered or {"outcome": {"outcome": "cancelled"}})
+                continue
+            self._emit_notice(notice, pending, on_line)
+        # The CLI writes its last updates just before the prompt result, so
+        # they can still be queued when the result ends the loop above.
+        while True:
+            notice = client.poll_notice(0)
+            if not notice:
+                break
+            self._emit_notice(notice, pending_from_notice(notice), on_line)
+        for payload in outcome["final"]:
+            _emit(on_line, payload)
+        thread.join(2)
+        return int(outcome["rc"])
+
+    def _emit_notice(self, notice: dict, pending: Optional[dict], on_line: LineCallback) -> None:
+        if pending and pending.get("kind") in ("todos", "task") and self.on_pending:
+            self.on_pending(pending)
+        for payload in acp_notice_to_lines(notice):
+            _emit(on_line, payload)