rockycode 0.1.0__py3-none-any.whl

rockycode/engine/goal_review.py

@@ -0,0 +1,161 @@
+"""Safe chat tools for reviewing and merging a goal branch.
+
+A goal run leaves its work COMMITTED on a `goal/<slug>` branch in a separate git
+worktree — never in the user's current files. Back in chat the model may be asked
+to "review it" or "merge it". Rather than trust it to hand-roll git (where a bad
+merge could leave the repo half-conflicted), these tools bake in the guards:
+
+  list_goal_branches   — discover what's there (read-only)
+  review_goal_branch   — the branch's diff vs where it forked (read-only)
+  merge_goal_branch    — merge into the current branch, but ONLY if the tree is
+                         clean, and ABORT (restoring the repo exactly) on conflict
+
+Every tool refuses any ref that isn't a `goal/...` branch, so the model can't be
+talked into merging an arbitrary branch. merge is risk="risky" → it still goes
+through the normal approval gate in ask/careful mode.
+"""
+from __future__ import annotations
+
+import asyncio
+import re
+from pathlib import Path
+
+from rockycode.engine.tools import Tool, _fn_schema
+
+_GOAL_BRANCH_RX = re.compile(r"^goal/[\w./-]+$")
+_MAX_DIFF_CHARS = 40_000
+
+
+async def _git(repo: str, *args: str) -> tuple[str, int]:
+    """Run `git -C repo <args>` off the event loop → (combined output, rc)."""
+    proc = await asyncio.create_subprocess_exec(
+        "git", "-C", repo, *args,
+        stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT,
+    )
+    out, _ = await proc.communicate()
+    return out.decode("utf-8", "replace"), proc.returncode
+
+
+def _is_goal_branch(name: str) -> bool:
+    return bool(_GOAL_BRANCH_RX.match(name or ""))
+
+
+async def _exists(repo: str, branch: str) -> bool:
+    _out, rc = await _git(repo, "rev-parse", "--verify", "--quiet", f"refs/heads/{branch}")
+    return rc == 0
+
+
+def build_goal_tools(*, workdir: Path, reviewer=None) -> dict[str, Tool]:
+    """Return {list_goal_branches, review_goal_branch, merge_goal_branch} bound to
+    the host repo at *workdir*.
+
+    *reviewer* (optional): an async callable (branch) -> report. When provided,
+    review_goal_branch delegates to it — a read-only explore child reads the
+    branch and returns cited, harness-verified findings — so the chat context
+    receives a REVIEW instead of a 40k-char raw diff. When None (tests, callers
+    that don't carry an engine), the plain truncated diff dump remains."""
+    repo = str(Path(workdir).resolve())
+
+    async def list_goal_branches() -> str:
+        """List the goal/* branches in this repo, newest first, with how many
+        commits each is ahead of the current branch."""
+        out, rc = await _git(
+            repo, "for-each-ref", "--sort=-committerdate",
+            "--format=%(refname:short)\t%(contents:subject)", "refs/heads/goal")
+        if rc != 0:
+            return f"[error] {out.strip() or 'could not list branches'}"
+        rows = [ln for ln in out.splitlines() if ln.strip()]
+        if not rows:
+            return "[ok] no goal branches — run /goal to create one."
+        lines = ["goal branches (newest first):"]
+        for row in rows:
+            branch, _, subject = row.partition("\t")
+            ahead, _ = await _git(repo, "rev-list", "--count", f"HEAD..{branch}")
+            lines.append(f"  {branch}  (+{ahead.strip() or '?'} commits)  {subject.strip()}")
+        return "\n".join(lines)
+
+    def _diff_dump(branch: str, out: str) -> str:
+        if len(out) > _MAX_DIFF_CHARS:
+            head = out[:_MAX_DIFF_CHARS]
+            return (f"[ok] diff of {branch} (truncated to {_MAX_DIFF_CHARS:,} chars — "
+                    f"full: git -C {repo} diff HEAD...{branch}):\n{head}\n…[truncated]")
+        return f"[ok] diff of {branch} vs the current branch:\n{out}"
+
+    async def review_goal_branch(branch: str, raw: bool = False) -> str:
+        """Review a goal branch. With a reviewer wired: delegate to a read-only
+        explore child (full diff + surrounding code, cited findings) and return
+        only its report. raw=True — or no reviewer — returns the plain
+        (truncated) diff instead. Use before merge_goal_branch."""
+        if not _is_goal_branch(branch):
+            return f"[refused] '{branch}' is not a goal branch (must look like goal/…)."
+        if not await _exists(repo, branch):
+            return f"[not found] no branch '{branch}' — call list_goal_branches to see what exists."
+        out, rc = await _git(repo, "diff", f"HEAD...{branch}")
+        if rc != 0:
+            return f"[error] {out.strip() or 'diff failed'}"
+        if not out.strip():
+            return f"[ok] '{branch}' has no changes vs the current branch."
+        if reviewer is None or raw:
+            return _diff_dump(branch, out)
+        try:
+            report = await reviewer(branch)
+        except Exception as e:  # noqa: BLE001 — a failed review degrades to the dump
+            note = f"[note] delegated review failed ({type(e).__name__}) — raw diff instead.\n"
+            return note + _diff_dump(branch, out)
+        return (f"[ok] delegated review of {branch} (the full diff stayed out of "
+                f"this context; raw=true for the diff itself):\n{report}")
+
+    async def merge_goal_branch(branch: str) -> str:
+        """Merge a goal branch into the CURRENT branch. Guarded: refuses if the
+        working tree has uncommitted changes, and aborts (restoring the repo
+        exactly) if the merge conflicts — so it can never leave a half-merged mess.
+        Only operates on goal/… branches."""
+        if not _is_goal_branch(branch):
+            return f"[refused] '{branch}' is not a goal branch (must look like goal/…). Nothing merged."
+        if not await _exists(repo, branch):
+            return f"[not found] no branch '{branch}' — call list_goal_branches to see what exists."
+
+        status, _ = await _git(repo, "status", "--porcelain")
+        if status.strip():
+            return ("[refused] your working tree has uncommitted changes — commit or stash them "
+                    "first, then merge again. Nothing was merged.")
+        target, _ = await _git(repo, "rev-parse", "--abbrev-ref", "HEAD")
+        target = target.strip() or "the current branch"
+
+        out, rc = await _git(repo, "merge", "--no-ff", branch, "-m", f"Merge goal branch {branch}")
+        if rc == 0:
+            return (f"[merged] {branch} → {target}. Your files now have the changes. "
+                    f"Tidy the leftover goal worktrees any time with: rockycode goal --clean "
+                    f"(keeps the branches).")
+        await _git(repo, "merge", "--abort")   # restore the repo byte-for-byte
+        return (f"[conflict] merging {branch} into {target} hit conflicts — I aborted it, so your "
+                f"repo is unchanged. Resolve manually if you want: git -C {repo} merge {branch}\n"
+                f"conflicting files:\n{out.strip()[:1500]}")
+
+    return {
+        "list_goal_branches": Tool(
+            name="list_goal_branches",
+            schema=_fn_schema("list_goal_branches",
+                              "List the goal/* branches produced by past /goal runs, with how many "
+                              "commits each is ahead. Read-only.", {}, []),
+            fn=list_goal_branches, risk="safe"),
+        "review_goal_branch": Tool(
+            name="review_goal_branch",
+            schema=_fn_schema("review_goal_branch",
+                              "Review a goal branch before merging. Read-only. Returns an independent "
+                              "reviewer's report with verified citations (the diff itself stays out of "
+                              "this context); pass raw=true if you need the plain truncated diff.",
+                              {"branch": {"type": "string", "description": "The goal branch, e.g. goal/20260706-2113."},
+                               "raw": {"type": "boolean", "description": "Return the raw (truncated) diff instead of the delegated review."}},
+                              ["branch"]),
+            fn=review_goal_branch, risk="safe"),
+        "merge_goal_branch": Tool(
+            name="merge_goal_branch",
+            schema=_fn_schema("merge_goal_branch",
+                              "Safely merge a goal branch into the current branch. Refuses a dirty tree "
+                              "and aborts cleanly on conflict — never leaves a half-merged repo. Only "
+                              "goal/… branches.",
+                              {"branch": {"type": "string", "description": "The goal branch to merge, e.g. goal/20260706-2113."}},
+                              ["branch"]),
+            fn=merge_goal_branch, risk="risky"),
+    }

rockycode/engine/goal_session.py

@@ -0,0 +1,259 @@
+"""Frontend-agnostic orchestration of ONE goal run.
+
+The sequence — isolated workspace → plan → derive permits from the plan →
+provision the sandbox → run the milestone loop → summarize — is identical whether
+it's driven from a terminal or from the in-app GoalScreen. This module owns that
+sequence behind a small seam (the methods GoalScreen calls), so the UI only has
+to render and collect y/e/n + edit text; it never touches Docker, git, or the
+models directly. That also makes the screen testable with a fake backend.
+
+Planning happens BEFORE the sandbox exists (it's LLM-only), so network/permits are
+decided from the REAL plan, not a guess at the user's wording — same rule as the
+CLI. `LiveGoalBackend` is the real implementation; tests substitute their own.
+"""
+from __future__ import annotations
+
+import time as _time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable, List, Optional
+
+from rockycode.engine.budget import GoalBudget
+from rockycode.engine.safety import Verdict, network_intent, pre_scan
+from rockycode.engine.worktree import GoalWorkspace
+from rockycode.pricing import UsageLedger
+
+
+@dataclass
+class Permits:
+    """What a plan will need, derived from the planner's REQUIRES line + a scan
+    backstop over the milestones. `blocked` set → the plan must not run."""
+    use_network: bool
+    net_reason: str
+    asks: List[Verdict]
+    approved: set
+    blocked: Optional[str] = None
+
+    @property
+    def needs_notice(self) -> bool:
+        return bool(self.use_network or self.asks or self.net_reason)
+
+
+@dataclass
+class GoalSummary:
+    status: str            # "done" | "budget" | "stalled" | "aborted" | "cancelled" | "error"
+    reason: str
+    milestones_done: int = 0
+    milestones_total: int = 0
+    branch: str = ""
+    origin: str = ""
+    workspace: str = ""
+    log: str = ""
+    currency: str = "usd"
+    spend: float = 0.0
+    base: str = ""
+
+    @property
+    def review_cmd(self) -> str:
+        if not self.branch:
+            return ""
+        return f"git -C {self.origin} diff {self.base or 'HEAD'}..{self.branch}"
+
+    @property
+    def tidy_cmd(self) -> str:
+        if not self.branch:
+            return ""
+        return f"git -C {self.origin} worktree remove {self.workspace}"
+
+
+class LiveGoalBackend:
+    """The real backend: git worktree isolation + Docker sandbox + EngineDriver.
+
+    Lifecycle the screen drives:  setup() → plan() → [discuss()…] → run() → done.
+    cleanup() is idempotent and safe to call on any exit path."""
+
+    def __init__(
+        self,
+        objective: str,
+        context: str = "",
+        *,
+        model: str,
+        reviewer_model: str,
+        budget: GoalBudget,
+        workdir: Path,
+        currency: str = "usd",
+        network: Optional[bool] = None,
+        review_every: int = 3,
+        plan_file: Optional[Path] = None,
+    ) -> None:
+        self.objective = objective
+        self.context = context
+        # When set, plan() reads THIS already-approved plan file (from chat /plan)
+        # instead of calling the LLM planner — the handoff skips re-planning.
+        self.plan_file = Path(plan_file) if plan_file else None
+        self.model = model
+        self.reviewer_model = reviewer_model
+        self.budget = budget
+        self.workdir = Path(workdir).resolve()
+        self.currency = currency
+        self.network = network            # None → decide from the plan; True/False → forced
+        self.review_every = review_every
+
+        self.ledger = UsageLedger()
+        self.ws: Optional[GoalWorkspace] = None
+        self.driver = None
+        self._client = None
+        self._log_path: Optional[Path] = None
+        self._sandbox = None
+
+    # ---- lifecycle ----------------------------------------------------------
+
+    async def setup(self) -> str:
+        """Create the isolated workspace + the driver (no sandbox yet). Returns a
+        one-line description of where the work will happen."""
+        from openai import AsyncOpenAI
+
+        from rockycode.engine.goal import EngineDriver
+        from rockycode.onboarding import require_base_url, require_key
+
+        slug = _time.strftime("%Y%m%d-%H%M%S")
+        self.ws = GoalWorkspace.create(self.workdir, slug)
+        log_dir = Path.home() / ".rockycode" / "goal-logs"
+        try:
+            log_dir.mkdir(parents=True, exist_ok=True)
+            self._log_path = log_dir / f"{slug}.log"
+        except OSError:
+            self._log_path = None
+        self.log(f"# goal: {self.objective}\n# {slug}  workspace={self.ws.path}  branch={self.ws.branch}")
+
+        # Explicit key AND endpoint, never the SDK's ambient env fallbacks.
+        self._client = AsyncOpenAI(api_key=require_key(), base_url=require_base_url(),
+                                   max_retries=5, timeout=300.0)
+        from rockycode.engine.explore import make_goal_verifier
+        self.driver = EngineDriver(
+            client=self._client, model=self.model, reviewer_model=self.reviewer_model,
+            workspace=self.ws, ledger=self.ledger, currency=self.currency,
+            verifier=make_goal_verifier(client=self._client, model=self.model,
+                                        workdir=self.ws.path, ledger=self.ledger),
+        )
+        where = str(self.ws.path) + (f"  ·  branch {self.ws.branch}" if self.ws.branch else "  (copy)")
+        return where
+
+    async def plan(self) -> tuple[list[str], str]:
+        # Handoff from chat /plan: the plan is already drafted + approved. Read it
+        # (shape-agnostic) and turn each phase into a milestone — no LLM re-plan.
+        if self.plan_file is not None:
+            from rockycode.engine.planmode import parse_plan_file
+            try:
+                text = self.plan_file.read_text(errors="replace")
+            except OSError:
+                text = ""
+            phases = parse_plan_file(text)
+            milestones = [
+                (title + ("\n" + "\n".join(f"- {s}" for s in steps) if steps else ""))
+                for title, steps in phases
+            ]
+            # Pass the FULL plan text as the requires-scan supplement so permits()
+            # catches an install/network mention wherever it sits (a REQUIRES line,
+            # or inside a step like "pip install requests"), not just in milestones.
+            return milestones[:12], text
+        plan_input = self.objective if not self.context else (
+            f"{self.objective}\n\n[Context from the chat that led here — use it to inform "
+            f"the plan; the objective above is the goal]:\n{self.context}")
+        plan, requires = await self.driver.plan(plan_input)
+        return plan, requires
+
+    def permits(self, plan: list[str], requires: str) -> Permits:
+        scan_text = "\n".join(plan) + (("\n" + requires) if requires else "")
+        flags = pre_scan(scan_text)
+        blocked = next((v.reason for v in flags if v.action == "block"), None)
+        asks = [v for v in flags if v.action == "ask"]
+        net_reason = network_intent(requires) or network_intent(scan_text)
+        use_network = self.network if self.network is not None else bool(net_reason)
+        return Permits(
+            use_network=bool(use_network), net_reason=net_reason or "",
+            asks=asks, approved={v.pattern for v in asks}, blocked=blocked,
+        )
+
+    async def discuss(self, plan: list[str], requires: str, msg: str) -> tuple[str, list[str], str]:
+        """Answer a question about the plan and return (reply, plan, requires) —
+        the plan may be revised. Errors become a short reply, never a crash."""
+        try:
+            return await self.driver.discuss(self.objective, plan, requires, msg)
+        except Exception as e:  # noqa: BLE001
+            return (f"(couldn't reason about that — {type(e).__name__})", plan, requires)
+
+    async def run(
+        self, plan: list[str], permits: Permits, on_event: Callable[[str], None]
+    ) -> GoalSummary:
+        """Provision the sandbox with the permit decision, attach the engine, and
+        run the milestone loop. Always stops the sandbox; never raises for a normal
+        run failure (returns an 'error' summary instead)."""
+        from rockycode.engine.goal import GoalRunner, safe_bash_tool
+        from rockycode.engine.loop import Engine
+        from rockycode.engine.sandbox import ChatSandbox, build_sandbox_registry
+
+        self.log("plan:\n" + "\n".join(f"  {i}. {m}" for i, m in enumerate(plan, 1)))
+        self.budget.start()
+
+        def emit(m: str) -> None:
+            on_event(m)
+            self.log(m)
+
+        try:
+            self._sandbox = await ChatSandbox.start(self.ws.path, network=permits.use_network)
+        except Exception as e:  # noqa: BLE001
+            return self._summary("error", f"sandbox (Docker) failed to start — {e}")
+
+        try:
+            reg = build_sandbox_registry(self._sandbox)
+            reg["bash"] = safe_bash_tool(self._sandbox, permits.approved)
+            engine = Engine(model=self.model, client=self._client, workdir=self.ws.path,
+                            registry=reg, trajectory_meta={"goal": self.objective, "runner": "goal"})
+            self.driver.attach(engine, network=permits.use_network)
+            runner = GoalRunner(self.objective, self.driver, self.budget, self.ws, self.ledger,
+                                review_every=self.review_every, on_event=emit, preplanned=plan)
+            result = await runner.run()
+            summary = self._summary(result.status, result.reason,
+                                    result.milestones_done, result.milestones_total)
+            self.log(f"result: {result.status} — {result.reason}  "
+                     f"({result.milestones_done}/{result.milestones_total} milestones)")
+            return summary
+        except Exception as e:  # noqa: BLE001
+            return self._summary("error", f"{type(e).__name__}: {e}")
+        finally:
+            try:
+                await self._sandbox.stop()
+            except Exception:  # noqa: BLE001
+                pass
+            self._sandbox = None
+
+    async def cleanup(self, keep: bool) -> None:
+        """Tidy the workspace. keep=True leaves the worktree + branch to review;
+        keep=False removes an unstarted run (cancel / plan error)."""
+        if self.ws is not None:
+            try:
+                self.ws.cleanup(keep=keep)
+            except Exception:  # noqa: BLE001
+                pass
+
+    # ---- helpers ------------------------------------------------------------
+
+    def _summary(self, status: str, reason: str, done: int = 0, total: int = 0) -> GoalSummary:
+        ws = self.ws
+        return GoalSummary(
+            status=status, reason=reason, milestones_done=done, milestones_total=total,
+            branch=(ws.branch if ws else ""), origin=(str(ws.origin) if ws else ""),
+            workspace=(str(ws.path) if ws else ""), log=(str(self._log_path) if self._log_path else ""),
+            currency=self.currency, spend=self.ledger.cost(self.currency),
+            base=(getattr(ws, "base", "") or "" if ws else ""),
+        )
+
+    def log(self, msg: str) -> None:
+        if self._log_path is None:
+            return
+        try:
+            with self._log_path.open("a", encoding="utf-8") as f:
+                f.write(msg + "\n")
+        except OSError:
+            pass