omega-code 0.4.0__py3-none-any.whl

omega/subagent.py

@@ -0,0 +1,98 @@
+import secrets
+from collections.abc import Callable
+
+from . import events, tools
+from .config import Config
+from .session import Message
+
+CFG: Config | None = None
+# Set by run_turn at the start of every turn, mirroring tools.set_tainted(False).
+EMIT: Callable[[events.Event], None] | None = None
+
+TIERS = {"fast": "subagent_fast", "mid": "subagent_mid", "verify": "subagent_mid"}
+
+SYSTEM = """You are a research subagent. You have read-only tools.
+
+Investigate and answer the task concretely -- real paths, real line numbers, real
+symbol names. Do not speculate; if you could not determine something, say so.
+
+Your reply is the ONLY thing the orchestrator sees, so it must stand alone. Be
+dense and short: no preamble, no restating the task, no offers to help further."""
+
+REVIEW_SYSTEM = """You are a verification subagent reviewing a just-completed
+BUILD-mode turn. You have read-only tools.
+
+You are given the user's original request and a diff of the changes made to
+address it. Check the diff for: correctness against the ask, cases the diff
+misses, and leftover debug code (stray prints, commented-out code, TODO
+markers left by mistake). Read surrounding files with your tools when the
+diff alone doesn't tell you enough.
+
+Reply with exactly "OK" if the change looks right. Otherwise reply with a
+short numbered list of concrete issues -- no preamble, no restating the diff."""
+
+
+@tools.tool(
+    "subagent",
+    "Delegate a bounded read-only research task to a smaller, faster model. "
+    "Returns a dense summary instead of raw tool output, so use it for wide "
+    "searches, multi-file reading, and 'where is X' questions. "
+    "tier: 'fast' for lookups and greps, 'mid' for reasoning across files.",
+    {"task": {"type": "string"},
+     "tier": {"type": "string", "enum": ["fast", "mid"]}},
+    ["task"],
+)
+async def _subagent(task: str, tier: str = "fast") -> str:
+    from .loop import run_agent
+    if CFG is None:
+        return "error: subagent not wired"
+    role = TIERS.get(tier, "subagent_fast")
+    subagent_id = secrets.token_hex(3)
+    emit = EMIT or (lambda _e: None)
+
+    def forward(ev: events.Event) -> None:
+        # The subagent's prose is its return value, not transcript text --
+        # only its tool activity is visible upward. ModelUsed is also
+        # swallowed: it names the subagent's own tier model and would
+        # otherwise stomp the top-level turn's status-bar model on emit.
+        # Phase is swallowed too -- the subagent's own waiting/thinking/
+        # streaming churn is not the top-level turn's phase; the parent
+        # stays in "tools" for the whole duration of the subagent call.
+        if not isinstance(ev, events.TextDelta | events.Done | events.ModelUsed | events.Phase):
+            emit(ev)
+
+    emit(events.SubagentSpawned(subagent_id=subagent_id, tier=tier,
+                                task_preview=" ".join(task.split())[:80]))
+    history: list[Message] = [{"role": "user", "content": task}]
+    result = await run_agent(CFG, role, SYSTEM, history,
+                             tools.READ_ONLY - {"subagent"}, emit=forward,
+                             max_rounds=12, subagent_id=subagent_id, tier=tier)
+    emit(events.SubagentDone(subagent_id=subagent_id,
+                             summary_preview=" ".join(result.split())[:120]))
+    return result
+
+
+async def review(cfg: Config, request: str, diff_text: str,
+                 emit: Callable[[events.Event], None] | None = None) -> str:
+    """Invoked automatically by loop.run_agent at the end of a BUILD turn
+    that made a non-trivial change (see loop.py) -- not reachable from the
+    `subagent` tool's tier enum, which only offers "fast"/"mid" to the model.
+    Returns "OK" or a numbered list of issues; never raises."""
+    from .loop import run_agent
+    emit = emit or (lambda _e: None)
+    subagent_id = secrets.token_hex(3)
+
+    def forward(ev: events.Event) -> None:
+        if not isinstance(ev, events.TextDelta | events.Done | events.ModelUsed | events.Phase):
+            emit(ev)
+
+    emit(events.SubagentSpawned(subagent_id=subagent_id, tier="verify",
+                                task_preview=" ".join(request.split())[:80]))
+    task = f"Original request:\n{request}\n\nDiff of changes made this turn:\n{diff_text}"
+    history: list[Message] = [{"role": "user", "content": task}]
+    result = await run_agent(cfg, TIERS["verify"], REVIEW_SYSTEM, history,
+                             tools.READ_ONLY - {"subagent"}, emit=forward,
+                             max_rounds=6, subagent_id=subagent_id, tier="verify")
+    emit(events.SubagentDone(subagent_id=subagent_id,
+                             summary_preview=" ".join(result.split())[:120]))
+    return result

omega/tasks.py

@@ -0,0 +1,195 @@
+"""A Task pairs a session with a repo (optionally its own git worktree) and
+the run state (status/phase/tokens/cost) the D1 daemon updates as a worker
+process streams events -- see server/manager.py. Task id is always the
+underlying session id, so `session.load(task.id)` recovers full history.
+"""
+from __future__ import annotations
+
+import json
+import re
+import shutil
+import subprocess
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any, Literal
+
+from . import gitlog, session
+
+DIR = Path.home() / ".omega" / "tasks"
+WORKTREES_DIR = Path.home() / ".omega" / "worktrees"
+
+Status = Literal["idle", "running", "waiting_input", "done", "failed"]
+
+# `gh pr view` is a network round trip; a WS overview push or a sidebar repaint
+# must not pay for it every time -- see pr_info().
+_PR_CACHE_TTL = 60.0
+_pr_cache: dict[str, tuple[float, dict[str, Any] | None]] = {}
+
+
+@dataclass
+class Task:
+    id: str
+    title: str
+    repo: str
+    cwd: str
+    worktree: bool
+    branch: str | None = None
+    pr: dict[str, Any] | None = None
+    model: str | None = None
+    mode: str = "build"
+    status: Status = "idle"
+    phase: str = ""
+    created: float = 0.0
+    updated: float = 0.0
+    tokens_in: int = 0
+    tokens_out: int = 0
+    cost_usd: float | None = None
+    elapsed_s: float = 0.0
+
+    @property
+    def path(self) -> Path:
+        return DIR / f"{self.id}.json"
+
+    def save(self) -> None:
+        DIR.mkdir(parents=True, exist_ok=True)
+        self.updated = time.time()
+        tmp = self.path.with_suffix(".tmp")
+        tmp.write_text(json.dumps(asdict(self), indent=1))
+        tmp.replace(self.path)
+
+
+def _slugify(text: str, limit: int = 40) -> str:
+    slug = re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")
+    return slug[:limit].strip("-") or "task"
+
+
+def title_from_prompt(prompt: str, width: int = 80) -> str:
+    text = " ".join(prompt.split())
+    return text[: width - 1] + "…" if len(text) > width else text
+
+
+def _run_git(args: list[str], cwd: Path) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(["git", *args], cwd=cwd, capture_output=True, text=True, timeout=30)
+
+
+def _repo_root(repo: str) -> Path | None:
+    path = Path(repo).expanduser().resolve()
+    result = _run_git(["rev-parse", "--show-toplevel"], path)
+    if result.returncode != 0:
+        return None
+    return Path(result.stdout.strip())
+
+
+def create(repo: str, prompt: str | None = None, worktree: bool = True,
+          model: str | None = None, mode: str = "build") -> Task:
+    """Creates the underlying session first (its id becomes the task id), then
+    -- when `worktree` -- a git worktree branched from the repo's current HEAD
+    at `~/.omega/worktrees/<repo-basename>/<task-id>` and points the session's
+    cwd at it, so every tool the agent runs (bash's cwd included) stays inside
+    the isolated checkout."""
+    root = _repo_root(repo)
+    if worktree and root is None:
+        raise ValueError(f"{repo!r} is not inside a git repository")
+    root = root or Path(repo).expanduser().resolve()
+
+    sess = session.Session.new(cwd=str(root), mode=mode)
+    branch: str | None = None
+    cwd = str(root)
+
+    if worktree:
+        slug = _slugify(prompt) if prompt else sess.id
+        branch = f"omega/{slug}"
+        wt_path = WORKTREES_DIR / root.name / sess.id
+        wt_path.parent.mkdir(parents=True, exist_ok=True)
+        added = _run_git(["worktree", "add", str(wt_path), "-b", branch], root)
+        if added.returncode != 0:
+            raise RuntimeError(f"git worktree add failed: {added.stderr.strip()}")
+        cwd = str(wt_path)
+        sess.cwd = cwd
+    else:
+        branch = gitlog._branch(root) if (root / ".git").exists() else None
+        if branch == "?":
+            branch = None
+
+    sess.save()
+    task = Task(
+        id=sess.id, title=title_from_prompt(prompt) if prompt else "(no prompt yet)",
+        repo=str(root), cwd=cwd, worktree=worktree, branch=branch,
+        model=model, mode=mode, created=sess.created, updated=sess.created,
+    )
+    task.save()
+    return task
+
+
+def _load(path: Path) -> Task | None:
+    try:
+        raw = json.loads(path.read_text())
+    except (OSError, json.JSONDecodeError):
+        return None
+    known = set(Task.__dataclass_fields__)
+    try:
+        return Task(**{k: v for k, v in raw.items() if k in known})
+    except TypeError:
+        return None
+
+
+def list_tasks() -> list[Task]:
+    if not DIR.exists():
+        return []
+    found = [t for p in DIR.glob("*.json") if (t := _load(p)) is not None]
+    return sorted(found, key=lambda t: t.updated, reverse=True)
+
+
+def get(task_id: str) -> Task | None:
+    path = DIR / f"{task_id}.json"
+    if path.exists():
+        return _load(path)
+    matches = sorted(DIR.glob(f"{task_id}*.json")) if DIR.exists() else []
+    return _load(matches[-1]) if matches else None
+
+
+def update(task_id: str, **fields: Any) -> Task | None:
+    task = get(task_id)
+    if task is None:
+        return None
+    for key, value in fields.items():
+        if key in Task.__dataclass_fields__:
+            setattr(task, key, value)
+    task.save()
+    return task
+
+
+def remove(task_id: str, delete_worktree: bool = False) -> bool:
+    task = get(task_id)
+    if task is None:
+        return False
+    if delete_worktree and task.worktree:
+        _run_git(["worktree", "remove", "--force", task.cwd], Path(task.repo))
+    task.path.unlink(missing_ok=True)
+    return True
+
+
+def pr_info(task: Task) -> dict[str, Any] | None:
+    """`gh pr view` for this task's branch, cached for `_PR_CACHE_TTL` seconds.
+    Never raises: no branch, no `gh` on PATH, no PR yet, or a network hiccup
+    all just mean "no PR to show" rather than an error the caller must handle."""
+    if not task.branch or shutil.which("gh") is None:
+        return None
+    now = time.time()
+    cached = _pr_cache.get(task.id)
+    if cached is not None and now - cached[0] < _PR_CACHE_TTL:
+        return cached[1]
+
+    info: dict[str, Any] | None
+    try:
+        result = subprocess.run(
+            ["gh", "pr", "view", "--json", "number,url,state,statusCheckRollup,title",
+             "--head", task.branch],
+            cwd=task.repo, capture_output=True, text=True, timeout=15,
+        )
+        info = json.loads(result.stdout) if result.returncode == 0 and result.stdout.strip() else None
+    except (OSError, subprocess.SubprocessError, json.JSONDecodeError):
+        info = None
+    _pr_cache[task.id] = (now, info)
+    return info