rockycode 0.1.0__py3-none-any.whl

rockycode/engine/server.py

@@ -0,0 +1,431 @@
+"""JSON-RPC 2.0 server over stdin/stdout — wraps the Engine for editor clients.
+
+One line per JSON message (NDJSON, same transport as LSP). The server manages
+multiple sessions; each session has its own Engine instance and its own
+run_turn task. Events stream out as notifications; permission requests become
+blocking request/response pairs.
+
+This is the backend the VS Code extension connects through.
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import sys
+import traceback
+import uuid
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import AsyncIterator, Optional
+
+from rockycode.engine import Engine
+from rockycode.engine.events import (
+    Compacted,
+    EngineError,
+    StateChanged,
+    TextDelta,
+    ThinkingDelta,
+    ToolFinished,
+    ToolStarted,
+    TurnFinished,
+)
+
+# ─────────────────────────────────────────────────────────────────────────────
+# protocol helpers
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def _notify(method: str, params: dict) -> dict:
+    return {"jsonrpc": "2.0", "method": method, "params": params}
+
+
+def _response(msg_id, result):
+    return {"jsonrpc": "2.0", "id": msg_id, "result": result}
+
+
+def _error(msg_id, code: int, message: str):
+    return {"jsonrpc": "2.0", "id": msg_id, "error": {"code": code, "message": message}}
+
+
+_DEBUG = os.getenv("ROCKYCODE_SERVE_DEBUG") == "1"
+
+
+def _client_error() -> str:
+    """A message safe to send to the client. The full traceback goes to stderr
+    for the operator; the client sees a generic string unless
+    ROCKYCODE_SERVE_DEBUG=1 — so a hostile client can't harvest local paths /
+    internals from an error over the wire. Call only from inside an `except`."""
+    tb = traceback.format_exc()
+    print(tb, file=sys.stderr, flush=True)
+    return tb if _DEBUG else "internal error"
+
+
+def _scrub_message(msg: str) -> str:
+    """Client-safe error text for streamed session/error events: redact secrets
+    and the user's home path so a hostile editor client can't harvest local
+    paths / internals (parity with _client_error). Full text under
+    ROCKYCODE_SERVE_DEBUG=1."""
+    if _DEBUG:
+        return msg
+    try:
+        from rockycode.engine.redact import redact
+        msg = redact(msg)
+    except Exception:  # noqa: BLE001 — best-effort
+        pass
+    try:
+        home = str(Path.home())
+        if home and home != "/":
+            msg = msg.replace(home, "~")
+    except Exception:  # noqa: BLE001
+        pass
+    return msg
+
+
+def _event_to_notification(session_id: str, ev) -> Optional[dict]:
+    """Convert an engine Event to a JSON-RPC notification dict, or None to skip."""
+    if isinstance(ev, StateChanged):
+        return _notify("session/state_changed", {
+            "session_id": session_id, "state": ev.state.value,
+        })
+    if isinstance(ev, ThinkingDelta):
+        return _notify("session/thinking_delta", {
+            "session_id": session_id, "text": ev.text,
+        })
+    if isinstance(ev, TextDelta):
+        return _notify("session/text_delta", {
+            "session_id": session_id, "text": ev.text,
+        })
+    if isinstance(ev, ToolStarted):
+        return _notify("session/tool_started", {
+            "session_id": session_id, "call_id": ev.call_id, "tool": ev.tool,
+            "args": ev.args,
+        })
+    if isinstance(ev, ToolFinished):
+        return _notify("session/tool_finished", {
+            "session_id": session_id, "call_id": ev.call_id, "tool": ev.tool,
+            "output": ev.output, "ok": ev.ok, "duration_s": ev.duration_s,
+        })
+    if isinstance(ev, Compacted):
+        return _notify("session/compacted", {
+            "session_id": session_id, "strategy": ev.strategy,
+            "tokens_before": ev.tokens_before, "tokens_after": ev.tokens_after,
+        })
+    if isinstance(ev, TurnFinished):
+        return _notify("session/turn_finished", {
+            "session_id": session_id, "steps": ev.steps, "usage": ev.usage,
+        })
+    if isinstance(ev, EngineError):
+        return _notify("session/error", {
+            "session_id": session_id, "message": _scrub_message(ev.message),
+        })
+    return None
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# output — thread-safe single-line JSON to stdout
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _write_line(msg: dict) -> None:
+    """Write one JSON-RPC message line to stdout. NOT async — called from both
+    the main server loop and session tasks. The write+flush has no await between
+    them, so within one event loop two coroutines can't interleave bytes. stdout
+    carries ONLY protocol JSON; logs and errors go to stderr."""
+    line = json.dumps(msg, ensure_ascii=False, default=str) + "\n"
+    sys.stdout.write(line)
+    sys.stdout.flush()
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# session
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+class Session:
+    def __init__(self, session_id: str, engine: Engine):
+        self.session_id = session_id
+        self.engine = engine
+        self._task: Optional[asyncio.Task] = None
+        self.created_at = datetime.now(timezone.utc).isoformat()
+
+    def cancel(self) -> None:
+        if self._task and not self._task.done():
+            self._task.cancel()
+
+    @property
+    def running(self) -> bool:
+        return self._task is not None and not self._task.done()
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# session manager
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+class SessionManager:
+    def __init__(self, model: str, workdir: Path, **engine_kwargs):
+        self._model = model
+        self._workdir = workdir
+        self._engine_kwargs = engine_kwargs
+        self._system_prompt = engine_kwargs.pop("system_prompt", "") or ""
+        self._sessions: dict[str, Session] = {}
+        # Permission bridge: engine approver → JSON-RPC notification → wait
+        # for client's session/permission_response.
+        self._perm_events: dict[str, asyncio.Event] = {}
+        self._perm_results: dict[str, bool] = {}
+        self._perm_sessions: dict[str, str] = {}  # event_id → session_id
+
+    def _new_session(self) -> Session:
+        sid = uuid.uuid4().hex[:12]
+        kwargs = dict(self._engine_kwargs)
+        if self._system_prompt:
+            kwargs["system_prompt"] = self._system_prompt
+        engine = Engine(model=self._model, workdir=self._workdir, **kwargs)
+
+        # Wire up artifact support: create_artifact tool + live server
+        try:
+            from rockycode.engine.artifact import ArtifactServer, build_artifact_tools
+            tools = build_artifact_tools(workdir=self._workdir, engine=engine)
+            engine.registry.update(tools)
+            server = ArtifactServer(self._workdir)
+            async def _artifact_target():
+                await server.start()
+                return server
+            engine.artifact_target = _artifact_target
+        except Exception:
+            pass  # artifacts are optional
+
+        sess = Session(sid, engine)
+        self._sessions[sid] = sess
+        return sess
+
+    def _make_approver(self, session_id: str):
+        """Return an async approve(name, args) → bool that bridges to the
+        JSON-RPC client via a notification + wait."""
+
+        async def approve(name: str, args: dict) -> bool:
+            eid = uuid.uuid4().hex[:8]
+            ev = asyncio.Event()
+            self._perm_events[eid] = ev
+            self._perm_sessions[eid] = session_id
+            risk = getattr(
+                self._sessions[session_id].engine.registry.get(name), "risk", "risky"
+            )
+            _write_line(_notify("session/request_permission", {
+                "session_id": session_id, "event_id": eid,
+                "tool": name, "args": args, "risk": risk,
+            }))
+            try:
+                await asyncio.wait_for(ev.wait(), timeout=120.0)
+            except asyncio.TimeoutError:
+                self._cleanup_perm(eid)
+                return False
+            result = self._perm_results.pop(eid, False)
+            self._cleanup_perm(eid)
+            return result
+
+        return approve
+
+    def resolve_permission(self, event_id: str, allowed: bool) -> bool:
+        if event_id in self._perm_events:
+            self._perm_results[event_id] = allowed
+            self._perm_events[event_id].set()
+            return True
+        return False
+
+    def _cleanup_perm(self, eid: str) -> None:
+        self._perm_events.pop(eid, None)
+        self._perm_sessions.pop(eid, None)
+
+    def get(self, session_id: str) -> Optional[Session]:
+        return self._sessions.get(session_id)
+
+    async def handle_chat(self, session_id: str, message: str) -> None:
+        sess = self.get(session_id)
+        if sess is None:
+            return
+        sess.engine.approver = self._make_approver(session_id)
+        sess._task = asyncio.current_task()
+        try:
+            async for ev in sess.engine.run_turn(message):
+                n = _event_to_notification(session_id, ev)
+                if n is not None:
+                    _write_line(n)
+        except asyncio.CancelledError:
+            pass  # _repair_history runs in run_turn's finally
+        except Exception:
+            _write_line(_notify("session/error", {
+                "session_id": session_id,
+                "message": _client_error(),
+            }))
+        finally:
+            sess._task = None
+
+    def list_sessions(self) -> list[dict]:
+        return [
+            {
+                "session_id": s.session_id,
+                "created_at": s.created_at,
+                "n_messages": len(s.engine.history),
+                "running": s.running,
+            }
+            for s in self._sessions.values()
+        ]
+
+    def cancel_session(self, session_id: str) -> bool:
+        sess = self.get(session_id)
+        if sess is None:
+            return False
+        sess.cancel()
+        return True
+
+    async def shutdown(self) -> None:
+        tasks = []
+        for sess in self._sessions.values():
+            sess.cancel()
+            if sess._task is not None:
+                tasks.append(sess._task)
+        # Await the cancellations rather than sleeping — otherwise the tasks are
+        # still pending at interpreter exit ("Task was destroyed but it is
+        # pending" noise on stderr).
+        for t in tasks:
+            try:
+                await t
+            except asyncio.CancelledError:
+                pass
+            except Exception:  # a dying session must not block shutdown
+                pass
+        # Reward lines for the dream (self-evolve): flush each session's
+        # heuristic outcome now that its turns are done. Idempotent, and a
+        # no-turn session writes nothing.
+        for sess in self._sessions.values():
+            try:
+                sess.engine.finalize_outcome()
+            except Exception:  # noqa: BLE001 — best-effort at teardown
+                pass
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# server loop
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+async def _read_messages() -> AsyncIterator[dict]:
+    """Read NDJSON messages from stdin using a thread so we don't block the
+    event loop on synchronous stdin reads."""
+    loop = asyncio.get_running_loop()
+
+    def _read() -> str:
+        return sys.stdin.readline()
+
+    while True:
+        line = await loop.run_in_executor(None, _read)
+        if not line:
+            break
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            msg = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        # A hostile/buggy client can send valid JSON that isn't an object
+        # (42, "x", null, []). Drop it here so the dispatch loop only ever sees
+        # dicts — one bad line must not crash the server on msg.get(...).
+        if isinstance(msg, dict):
+            yield msg
+
+
+async def run_server(
+    model: str,
+    workdir: Path,
+    thinking: bool = True,
+    reasoning_effort: str = "max",
+    max_tokens: int = 16384,
+    context_window: int = 131072,
+    max_steps: int = 0,
+    system_prompt: str = "",
+) -> None:
+    """Blocking: reads JSON-RPC from stdin, writes to stdout, runs until EOF."""
+
+    mgr = SessionManager(
+        model=model, workdir=workdir,
+        thinking=thinking, reasoning_effort=reasoning_effort,
+        max_tokens=max_tokens, context_window=context_window,
+        max_steps=max_steps, system_prompt=system_prompt,
+    )
+
+    chat_tasks: dict[str, asyncio.Task] = {}
+
+    async for msg in _read_messages():
+        method = msg.get("method", "")
+        msg_id = msg.get("id")
+        params = msg.get("params") or {}   # tolerate a missing or null "params"
+
+        try:
+            if method == "initialize":
+                sess = mgr._new_session()
+                # Report whether a usable key is present (rocky's credential
+                # chain: env, .env files serve loaded, keychain) so the client's
+                # "setup required" card reflects reality instead of guessing
+                # from the editor's own environment.
+                from rockycode.onboarding import is_configured
+                _write_line(_response(msg_id, {
+                    "version": "0.1.0",
+                    "session_id": sess.session_id,
+                    "model": model,
+                    "configured": is_configured(),
+                }))
+
+            elif method == "chat":
+                sid = params.get("session_id", "")
+                if not sid or sid not in mgr._sessions:
+                    sess = mgr._new_session()
+                    sid = sess.session_id
+                # Cancel any running turn on this session before starting a new one
+                prev = chat_tasks.get(sid)
+                if prev and not prev.done():
+                    mgr.cancel_session(sid)
+                    prev.cancel()
+                    try:
+                        await prev
+                    except asyncio.CancelledError:
+                        pass
+                task = asyncio.create_task(mgr.handle_chat(sid, params.get("message", "")))
+                chat_tasks[sid] = task
+                _write_line(_response(msg_id, {"session_id": sid}))
+
+            elif method == "cancel":
+                sid = params.get("session_id", "")
+                ok = mgr.cancel_session(sid)
+                _write_line(_response(msg_id, {"cancelled": ok}))
+
+            elif method == "list_sessions":
+                _write_line(_response(msg_id, {"sessions": mgr.list_sessions()}))
+
+            elif method == "get_status":
+                sid = params.get("session_id", "")
+                sess = mgr.get(sid)
+                _write_line(_response(msg_id, {
+                    "session_id": sid,
+                    "state": "busy" if (sess and sess.running) else "idle",
+                }))
+
+            elif method == "session/permission_response":
+                eid = params.get("event_id", "")
+                allowed = params.get("allowed", False)
+                mgr.resolve_permission(eid, allowed)
+                _write_line(_response(msg_id, {"ack": True}))
+
+            elif method == "shutdown":
+                _write_line(_response(msg_id, {"ok": True}))
+                break
+
+            else:
+                _write_line(_error(msg_id, -32601, f"unknown method: {method}"))
+
+        except Exception:
+            if msg_id is not None:
+                _write_line(_error(msg_id, -32603, _client_error()))
+
+    await mgr.shutdown()   # cancel any in-flight turns on EOF / shutdown

rockycode/engine/skills.py

@@ -0,0 +1,178 @@
+"""Skills: reuse what users already have, zero migration.
+
+Discovery, in priority order (first definition of a name wins):
+
+  1. <project>/.claude/skills/<name>/SKILL.md   — Claude Code project scope
+  2. <project>/.rockycode/skills/<name>/SKILL.md — rockycode's own
+  3. ~/.rockycode/skills/<name>/SKILL.md        — rocky global ($ROCKYCODE_HOME;
+                                                   dream-approved skills install here)
+  4. ~/.claude/skills/<name>/SKILL.md           — Claude Code user scope
+  5. ~/.codex/prompts/<name>.md                 — Codex custom prompts
+  6. rockycode/skills/<name>/SKILL.md           — packaged built-ins, last on
+     purpose: any project or user copy of the same name shadows a built-in
+
+Progressive disclosure, same pattern as Claude Code: only `name — description`
+lines go into the system prompt; the model calls the `skill` tool to load a
+skill's full instructions on demand, so fifty installed skills cost fifty
+lines of context, not fifty files.
+
+SKILL.md frontmatter is parsed leniently (name:/description: lines between
+--- markers) — no YAML dependency. Codex prompts have no frontmatter; the
+filename is the name and the first body line is the description.
+
+Chat-only, like MCP: bench never loads skills, so scores measure the harness.
+"""
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+from rockycode.engine.tools import Tool, _truncate
+
+MAX_SKILLS_IN_PROMPT = 100
+MAX_DESCRIPTION_CHARS = 200
+
+# Skills that ship inside the package (rockycode/skills/<name>/SKILL.md).
+BUILTIN_DIR = Path(__file__).resolve().parents[1] / "skills"
+
+_FRONTMATTER = re.compile(r"\A---\s*\n(.*?)\n---\s*\n", re.DOTALL)
+
+
+@dataclass
+class Skill:
+    name: str
+    description: str
+    path: Path  # the markdown file
+    source: str
+
+
+def _parse_frontmatter(text: str) -> tuple[dict[str, str], str]:
+    """Lenient `key: value` frontmatter parse. Returns (fields, body)."""
+    m = _FRONTMATTER.match(text)
+    if not m:
+        return {}, text
+    fields: dict[str, str] = {}
+    for line in m.group(1).splitlines():
+        if ":" in line and not line.startswith((" ", "\t", "#")):
+            key, _, value = line.partition(":")
+            fields[key.strip().lower()] = value.strip().strip("'\"")
+    return fields, text[m.end():]
+
+
+def _first_line(text: str) -> str:
+    for line in text.splitlines():
+        line = line.strip().lstrip("#").strip()
+        if line:
+            return line
+    return "(no description)"
+
+
+def _load_skill_md(path: Path, fallback_name: str, source: str) -> Optional[Skill]:
+    try:
+        text = path.read_text(errors="replace")
+    except OSError:
+        return None
+    fields, body = _parse_frontmatter(text)
+    return Skill(
+        name=fields.get("name") or fallback_name,
+        description=(fields.get("description") or _first_line(body))[:MAX_DESCRIPTION_CHARS],
+        path=path,
+        source=source,
+    )
+
+
+def discover_skills(workdir: Path, *, home: Optional[Path] = None,
+                    builtin: bool = True) -> list[Skill]:
+    """Collect skills from project + user locations, then packaged built-ins.
+    `home=None` disables user-level lookup (tests); pass Path.home() for the
+    real thing. `builtin=False` disables the packaged set (tests)."""
+    found: dict[str, Skill] = {}
+
+    def add(skill: Optional[Skill]) -> None:
+        if skill is not None and skill.name not in found:
+            found[skill.name] = skill
+
+    def scan_skill_dirs(root: Path, source: str) -> None:
+        if not root.is_dir():
+            return
+        for d in sorted(root.iterdir()):
+            md = d / "SKILL.md"
+            if d.is_dir() and md.exists():
+                add(_load_skill_md(md, d.name, source))
+
+    scan_skill_dirs(workdir / ".claude" / "skills", "project .claude/skills")
+    scan_skill_dirs(workdir / ".rockycode" / "skills", "project .rockycode/skills")
+    if home is not None:
+        # Rocky's own global store ($ROCKYCODE_HOME-aware, like trajectories) —
+        # where dream-approved proposals install. After the project dirs on
+        # purpose: a project skill always wins over a dream-drafted one.
+        base = os.environ.get("ROCKYCODE_HOME")
+        rocky_global = (Path(base).expanduser() if base else home / ".rockycode") / "skills"
+        scan_skill_dirs(rocky_global, "~/.rockycode/skills")
+        scan_skill_dirs(home / ".claude" / "skills", "~/.claude/skills")
+        codex_prompts = home / ".codex" / "prompts"
+        if codex_prompts.is_dir():
+            for f in sorted(codex_prompts.glob("*.md")):
+                add(_load_skill_md(f, f.stem, "~/.codex/prompts"))
+    if builtin:
+        scan_skill_dirs(BUILTIN_DIR, "built-in")
+
+    return list(found.values())[:MAX_SKILLS_IN_PROMPT]
+
+
+def skills_prompt_section(skills: list[Skill]) -> str:
+    lines = [f"- {s.name} — {s.description}" for s in skills]
+    return (
+        "\n\n# Skills available\n\n"
+        "These are installed playbooks — rocky's built-ins plus the user's own. "
+        "When a task matches one, call the "
+        "`skill` tool with its name FIRST and follow its instructions.\n\n"
+        + "\n".join(lines)
+    )
+
+
+def build_skill_tool(skills: list[Skill]) -> Tool:
+    by_name = {s.name: s for s in skills}
+    schema = {
+        "type": "function",
+        "function": {
+            "name": "skill",
+            "description": (
+                "Load the full instructions of an installed skill. "
+                "Call this before attempting a task that matches a skill's description."
+            ),
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "name": {
+                        "type": "string",
+                        "description": "Skill name from the 'Skills available' list.",
+                        "enum": sorted(by_name),
+                    },
+                },
+                "required": ["name"],
+            },
+        },
+    }
+
+    async def fn(name: str) -> str:
+        skill = by_name.get(name)
+        if skill is None:
+            return f"[error] no skill named '{name}'. available: {', '.join(sorted(by_name))}"
+        try:
+            text = skill.path.read_text(errors="replace")
+        except OSError as e:
+            return f"[error] could not read skill: {e}"
+        _, body = _parse_frontmatter(text)
+        # The directory matters: skills often ship scripts referenced
+        # relative to their own folder.
+        return _truncate(
+            f"# skill: {skill.name}\n# directory: {skill.path.parent}\n\n{body.strip()}"
+        )
+
+    # Loading a skill only returns SKILL.md text — it runs nothing. The real risk
+    # is whatever bash/web_fetch the model then attempts, which is gated on its own.
+    return Tool(name="skill", schema=schema, fn=fn, risk="safe")

rockycode/engine/titler.py

@@ -0,0 +1,46 @@
+"""Session titles: one tiny flash call, fired-and-forgotten after the first
+exchange, so the resume picker and exit card show "credential namespace fix"
+instead of the first message echoed back.
+
+Best-effort by design: offline, no key, a slow provider, or a weird reply all
+mean "no title" — the first-message summary remains the fallback everywhere,
+and nothing here may ever surface an error into the chat.
+"""
+from __future__ import annotations
+
+import os
+
+_MAX_CHARS = 60
+
+
+def _title_model() -> str:
+    return os.getenv("ROCKYCODE_TITLE_MODEL", "deepseek-v4-flash")
+
+
+def _clean(raw: str) -> str | None:
+    """First line, quotes/periods stripped, hard cap — a title, not a reply."""
+    t = raw.strip().splitlines()[0].strip().strip('"\'“”「」').rstrip(".。").strip()
+    if not t:
+        return None
+    return t[:_MAX_CHARS]
+
+
+async def generate_title(client, first_user: str, first_reply: str) -> str | None:
+    """3-8 word session title from the opening exchange, or None."""
+    try:
+        r = await client.chat.completions.create(
+            model=_title_model(),
+            max_tokens=24,
+            messages=[
+                {"role": "system", "content": (
+                    "Write a 3-8 word title for this coding-chat session, in the "
+                    "user's language. Output ONLY the title — no quotes, no period."
+                )},
+                {"role": "user", "content": (
+                    f"[user]\n{first_user[:1000]}\n\n[assistant]\n{first_reply[:1000]}"
+                )},
+            ],
+        )
+        return _clean(r.choices[0].message.content or "")
+    except Exception:  # noqa: BLE001 — titles are decoration; never break a session
+        return None