omega-code 0.4.0__py3-none-any.whl

omega/artifacts.py

@@ -0,0 +1,151 @@
+import json
+import secrets
+import time
+from collections import OrderedDict
+from pathlib import Path
+from typing import Any
+
+from .tools import truncate
+
+DIR = Path.home() / ".omega" / "sessions"
+
+Meta = dict[str, Any]
+
+# Threshold above which a tool result is offloaded to an artifact instead of
+# returned inline.
+OFFLOAD_THRESHOLD = 4000
+
+# A single stored artifact is capped here even though the tool result that
+# produced it may be larger -- a clipped result on disk still beats no
+# artifact at all, and beats forcing the whole thing through the preview.
+RESULT_MAX_CHARS = 100_000
+
+# Default page size for fetch_result -- comfortably under tools.MAX_INLINE_CHARS
+# so a default-sized page never itself needs the inline re-entry cap.
+PAGE_CHARS = 18_000
+
+# Per-process, per-session content cache checked before disk. Bounded by
+# entry count rather than bytes: RESULT_MAX_CHARS already caps a single
+# entry, so MAX_CACHE_ENTRIES * RESULT_MAX_CHARS is a safe worst case.
+MAX_CACHE_ENTRIES = 50
+_CACHE: OrderedDict[str, str] = OrderedDict()
+
+
+def _cache_key(session_id: str, artifact_id: str) -> str:
+    return f"{session_id}/{artifact_id}"
+
+
+def _cache_put(session_id: str, artifact_id: str, content: str) -> None:
+    key = _cache_key(session_id, artifact_id)
+    _CACHE[key] = content
+    _CACHE.move_to_end(key)
+    while len(_CACHE) > MAX_CACHE_ENTRIES:
+        _CACHE.popitem(last=False)
+
+
+def _cache_get(session_id: str, artifact_id: str) -> str | None:
+    key = _cache_key(session_id, artifact_id)
+    if key not in _CACHE:
+        return None
+    _CACHE.move_to_end(key)
+    return _CACHE[key]
+
+
+def _session_dir(session_id: str) -> Path:
+    return DIR / session_id / "artifacts"
+
+
+def _paths(session_id: str, artifact_id: str) -> tuple[Path, Path]:
+    d = _session_dir(session_id)
+    return d / f"{artifact_id}.txt", d / f"{artifact_id}.meta.json"
+
+
+def _write_meta(meta_path: Path, meta: Meta) -> None:
+    tmp = meta_path.with_suffix(".tmp")
+    tmp.write_text(json.dumps(meta, indent=1))
+    tmp.replace(meta_path)
+
+
+def save(session_id: str, content: str, title: str | None = None, kind: str = "offload") -> str:
+    d = _session_dir(session_id)
+    d.mkdir(parents=True, exist_ok=True)
+    artifact_id = secrets.token_hex(4)
+    text_path, meta_path = _paths(session_id, artifact_id)
+    stored = content if len(content) <= RESULT_MAX_CHARS else truncate(content, RESULT_MAX_CHARS)
+    text_path.write_text(stored)
+    now = time.time()
+    _write_meta(meta_path, {
+        "title": title or "",
+        "kind": kind,
+        "created": now,
+        "updated": now,
+        "char_count": len(stored),
+    })
+    _cache_put(session_id, artifact_id, stored)
+    return artifact_id
+
+
+def fetch(session_id: str, artifact_id: str, offset: int = 0, limit: int = PAGE_CHARS) -> str:
+    """Returns `content[offset:offset+limit]` plus a trailer line: either
+    `[end]` when the page reaches the artifact's end, or
+    `[chars <start>-<end> of <total>; next_offset=<n>]` so the model knows how
+    to page for more without re-fetching from the start."""
+    content = _cache_get(session_id, artifact_id)
+    if content is None:
+        text_path, _meta_path = _paths(session_id, artifact_id)
+        if not text_path.exists():
+            return f"error: no artifact {artifact_id!r} in session {session_id!r}"
+        content = text_path.read_text()
+        _cache_put(session_id, artifact_id, content)
+
+    total = len(content)
+    end = min(offset + limit, total)
+    chunk = content[offset:end]
+    trailer = "[end]" if end >= total else f"[chars {offset}-{end} of {total}; next_offset={end}]"
+    return f"{chunk}\n{trailer}"
+
+
+def update(session_id: str, artifact_id: str, content: str) -> str:
+    text_path, meta_path = _paths(session_id, artifact_id)
+    if not meta_path.exists():
+        return f"error: no artifact {artifact_id!r} in session {session_id!r}"
+    stored = content if len(content) <= RESULT_MAX_CHARS else truncate(content, RESULT_MAX_CHARS)
+    text_path.write_text(stored)
+    meta = json.loads(meta_path.read_text())
+    meta["char_count"] = len(stored)
+    meta["updated"] = time.time()
+    _write_meta(meta_path, meta)
+    _cache_put(session_id, artifact_id, stored)
+    return f"updated artifact {artifact_id}"
+
+
+def list_artifacts(session_id: str) -> list[Meta]:
+    d = _session_dir(session_id)
+    if not d.exists():
+        return []
+    out: list[Meta] = []
+    for meta_path in sorted(d.glob("*.meta.json")):
+        artifact_id = meta_path.name.removesuffix(".meta.json")
+        try:
+            meta = json.loads(meta_path.read_text())
+        except Exception:
+            continue
+        out.append({
+            "id": artifact_id,
+            "title": meta.get("title", ""),
+            "kind": meta.get("kind", "offload"),
+            "size": meta.get("char_count", 0),
+            "created": meta.get("created", 0),
+        })
+    return out
+
+
+def offload_if_large(text: str, session_id: str, threshold: int = OFFLOAD_THRESHOLD,
+                     preview_limit: int = 1200) -> str:
+    if len(text) <= threshold:
+        return text
+    artifact_id = save(session_id, text, kind="offload")
+    clipped_note = f" (clipped to {RESULT_MAX_CHARS} chars on save)" if len(text) > RESULT_MAX_CHARS else ""
+    return (f"{truncate(text, preview_limit)}\n"
+            f"[full output: {len(text)} chars{clipped_note}, saved as artifact {artifact_id} "
+            f"— fetch_result({artifact_id}) to read more]")

omega/checkpoint.py

@@ -0,0 +1,246 @@
+"""Working-tree checkpoints for BUILD-mode turns.
+
+Each checkpoint snapshots the working tree into a git tree object via a
+throwaway index (GIT_INDEX_FILE pointed at a temp file) so the user's real
+index/staging area and HEAD are never touched. `.omega/` is excluded from
+every snapshot so undo/diff never reads or writes omega's own state.
+"""
+import json
+import os
+import secrets
+import shutil
+import subprocess
+import tempfile
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+DIR = Path.home() / ".omega" / "sessions"
+
+NOT_A_REPO = "not a git repository"
+
+# Pathspec passed to every `git add -A` snapshot: everything except omega's
+# own project-local state directory.
+_PATHSPEC = ("--", ".", ":!.omega")
+
+
+@dataclass(frozen=True)
+class Checkpoint:
+    id: str
+    turn: int
+    created: float
+    tree_sha: str
+
+
+def _checkpoints_path(session_id: str) -> Path:
+    return DIR / session_id / "checkpoints.json"
+
+
+def _load(session_id: str) -> list[Checkpoint]:
+    path = _checkpoints_path(session_id)
+    if not path.exists():
+        return []
+    try:
+        raw = json.loads(path.read_text())
+    except json.JSONDecodeError:
+        return []
+    return [Checkpoint(**c) for c in raw]
+
+
+def _save(session_id: str, checkpoints: list[Checkpoint]) -> None:
+    path = _checkpoints_path(session_id)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(".tmp")
+    tmp.write_text(json.dumps([asdict(c) for c in checkpoints], indent=1))
+    tmp.replace(path)
+
+
+def _find(checkpoints: list[Checkpoint], since_turn: int | None) -> Checkpoint | None:
+    if since_turn is None:
+        return checkpoints[0] if checkpoints else None
+    for cp in checkpoints:
+        if cp.turn == since_turn:
+            return cp
+    return None
+
+
+def _repo_root(cwd: str) -> Path | None:
+    if shutil.which("git") is None:
+        return None
+    try:
+        result = subprocess.run(["git", "rev-parse", "--show-toplevel"], cwd=cwd,
+                                capture_output=True, text=True, timeout=10)
+    except (OSError, subprocess.SubprocessError):
+        return None
+    if result.returncode != 0 or not result.stdout.strip():
+        return None
+    return Path(result.stdout.strip())
+
+
+def _temp_index() -> str:
+    """A path git can treat as a fresh index -- created then removed so
+    `git add`/`read-tree` populate it from nothing, not a stale prior run."""
+    fd, path = tempfile.mkstemp(prefix="omega-checkpoint-idx-")
+    os.close(fd)
+    os.remove(path)
+    return path
+
+
+def _write_tree(repo_root: Path) -> str | None:
+    idx_path = _temp_index()
+    env = {**os.environ, "GIT_INDEX_FILE": idx_path}
+    try:
+        added = subprocess.run(["git", "add", "-A", *_PATHSPEC], cwd=repo_root, env=env,
+                               capture_output=True, text=True, timeout=60)
+        if added.returncode != 0:
+            return None
+        written = subprocess.run(["git", "write-tree"], cwd=repo_root, env=env,
+                                 capture_output=True, text=True, timeout=30)
+        if written.returncode != 0:
+            return None
+        return written.stdout.strip()
+    except (OSError, subprocess.SubprocessError):
+        return None
+    finally:
+        Path(idx_path).unlink(missing_ok=True)
+
+
+def create(session_id: str, turn: int, cwd: str | None = None) -> Checkpoint | None:
+    """Snapshots the working tree without touching the index or HEAD. Returns
+    None (never raises) when `cwd` isn't inside a git repo."""
+    root = _repo_root(cwd or os.getcwd())
+    if root is None:
+        return None
+    tree_sha = _write_tree(root)
+    if tree_sha is None:
+        return None
+    cp = Checkpoint(id=secrets.token_hex(4), turn=turn, created=time.time(), tree_sha=tree_sha)
+    checkpoints = _load(session_id)
+    checkpoints.append(cp)
+    _save(session_id, checkpoints)
+    return cp
+
+
+def undo(session_id: str, steps: int = 1, cwd: str | None = None) -> str:
+    """Restores the working tree to the checkpoint `steps` turns back. Never
+    touches .git, untracked ignored files, or .omega/."""
+    root = _repo_root(cwd or os.getcwd())
+    if root is None:
+        return NOT_A_REPO
+    checkpoints = _load(session_id)
+    if steps < 1 or steps > len(checkpoints):
+        return f"no checkpoint {steps} turn(s) back (this session has {len(checkpoints)})"
+    target = checkpoints[-steps]
+
+    current_tree = _write_tree(root)
+    if current_tree is None:
+        return "error: could not snapshot the current working tree"
+
+    diffed = subprocess.run(
+        ["git", "diff-tree", "-r", "--no-renames", "--name-status", current_tree, target.tree_sha],
+        cwd=root, capture_output=True, text=True, timeout=30)
+    to_delete = [line[2:] for line in diffed.stdout.splitlines() if line.startswith("D\t")]
+    touched = len(diffed.stdout.splitlines())
+
+    idx_path = _temp_index()
+    env = {**os.environ, "GIT_INDEX_FILE": idx_path}
+    try:
+        read = subprocess.run(["git", "read-tree", target.tree_sha], cwd=root, env=env,
+                              capture_output=True, text=True, timeout=30)
+        if read.returncode != 0:
+            return f"error: git read-tree failed: {read.stderr.strip()}"
+        checked_out = subprocess.run(["git", "checkout-index", "-a", "-f"], cwd=root, env=env,
+                                     capture_output=True, text=True, timeout=60)
+        if checked_out.returncode != 0:
+            return f"error: git checkout-index failed: {checked_out.stderr.strip()}"
+    except (OSError, subprocess.SubprocessError) as e:
+        return f"error: {type(e).__name__}: {e}"
+    finally:
+        Path(idx_path).unlink(missing_ok=True)
+
+    removed = 0
+    for rel in to_delete:
+        if rel.startswith(".omega/") or rel.startswith(".git/"):
+            continue
+        path = root / rel
+        try:
+            if path.exists() or path.is_symlink():
+                path.unlink()
+                removed += 1
+        except OSError:
+            continue
+        _prune_empty_parents(path.parent, root)
+
+    return (f"reverted {steps} turn(s) to the checkpoint from turn {target.turn} "
+           f"({touched} file(s) touched, {removed} removed)")
+
+
+def _prune_empty_parents(start: Path, root: Path) -> None:
+    d = start
+    while d != root and root in d.parents and d.exists():
+        try:
+            next(d.iterdir())
+            return  # not empty
+        except StopIteration:
+            pass
+        try:
+            d.rmdir()
+        except OSError:
+            return
+        d = d.parent
+
+
+def diff(session_id: str, since_turn: int | None = None, cwd: str | None = None) -> str:
+    """Unified diff between the checkpoint before `since_turn` (default: the
+    session's first checkpoint) and the current working tree, untracked files
+    included."""
+    root = _repo_root(cwd or os.getcwd())
+    if root is None:
+        return NOT_A_REPO
+    checkpoints = _load(session_id)
+    target = _find(checkpoints, since_turn)
+    if target is None:
+        return "no checkpoint recorded yet this session" if since_turn is None \
+            else f"no checkpoint recorded for turn {since_turn}"
+
+    idx_path = _temp_index()
+    env = {**os.environ, "GIT_INDEX_FILE": idx_path}
+    try:
+        added = subprocess.run(["git", "add", "-A", *_PATHSPEC], cwd=root, env=env,
+                               capture_output=True, text=True, timeout=60)
+        if added.returncode != 0:
+            return f"error: git add failed: {added.stderr.strip()}"
+        diffed = subprocess.run(["git", "diff", "--cached", target.tree_sha], cwd=root, env=env,
+                                capture_output=True, text=True, timeout=60)
+        return diffed.stdout or "(no changes since that checkpoint)"
+    except (OSError, subprocess.SubprocessError) as e:
+        return f"error: {type(e).__name__}: {e}"
+    finally:
+        Path(idx_path).unlink(missing_ok=True)
+
+
+def changed_files(session_id: str, since_turn: int | None = None, cwd: str | None = None) -> list[str]:
+    """Empty list for a non-git cwd or a session with no checkpoints -- unlike
+    undo()/diff(), the return type has no room for a sentinel message."""
+    root = _repo_root(cwd or os.getcwd())
+    if root is None:
+        return []
+    checkpoints = _load(session_id)
+    target = _find(checkpoints, since_turn)
+    if target is None:
+        return []
+
+    idx_path = _temp_index()
+    env = {**os.environ, "GIT_INDEX_FILE": idx_path}
+    try:
+        added = subprocess.run(["git", "add", "-A", *_PATHSPEC], cwd=root, env=env,
+                               capture_output=True, text=True, timeout=60)
+        if added.returncode != 0:
+            return []
+        listed = subprocess.run(["git", "diff", "--cached", "--name-only", target.tree_sha],
+                                cwd=root, env=env, capture_output=True, text=True, timeout=60)
+        return [line for line in listed.stdout.splitlines() if line]
+    except (OSError, subprocess.SubprocessError):
+        return []
+    finally:
+        Path(idx_path).unlink(missing_ok=True)

omega/compact.py

@@ -0,0 +1,106 @@
+import json
+from typing import cast
+
+from . import llm, trajectory
+from .config import Config
+from .llm import Turn
+from .session import Message
+
+FRACTION = 0.75
+KEEP_LAST = 6
+
+# Cap on what the summariser LLM ever sees: the transcript already truncates
+# each message to 1500 chars, but a long-running turn can still pile up far
+# more messages than that keeps sane for one summarisation call.
+_COMPACT_SUMMARY_INPUT_CHARS = 40_000
+
+# Per-line cap for the deterministic ledger appended alongside the LLM
+# summary -- generous compared to the live trajectory block's 64/120 char
+# fields, since this is the durable record of a range that is about to be
+# dropped from history for good.
+_COMPACT_ENTRY_CHARS = 1000
+
+SYSTEM = """Summarise this conversation excerpt for an agent that will keep working.
+
+Preserve: what was asked, what was done, file paths touched, decisions made,
+errors hit and how they were resolved, and anything still outstanding.
+Drop: raw file contents, long command output, and restated code.
+
+Write dense prose. No preamble."""
+
+
+def estimate_tokens(messages: list[Message]) -> int:
+    return sum(len(json.dumps(m)) for m in messages) // 4
+
+
+def safe_split(history: list[Message], keep_last: int) -> int:
+    """Return an index to cut at that never separates an assistant message
+    carrying tool_calls from the tool results that answer it.
+
+    Only the position immediately before a user message is safe, because tool
+    results always follow their assistant message inside a turn.
+    """
+    if len(history) <= keep_last:
+        return 0
+    for i in range(max(1, len(history) - keep_last), 0, -1):
+        m = history[i]
+        # Safe boundaries: a user message, or an assistant message with no
+        # tool_calls (its group is complete). Cutting anywhere else orphans a
+        # tool_call_id. Assistant boundaries matter because an agentic turn has
+        # exactly ONE user message -- at index 0, which can never be a cut.
+        if m.get("role") == "user":
+            return i
+        if m.get("role") == "assistant" and not m.get("tool_calls"):
+            return i
+    return 0
+
+
+async def maybe_compact(cfg: Config, history: list[Message], used: int, limit: int,
+                        fraction: float = FRACTION, keep_last: int = KEEP_LAST) -> str | None:
+    """Shrink history in place. Returns a note if it compacted, else None."""
+    if limit <= 0 or used < limit * fraction:
+        return None
+
+    cut = safe_split(history, keep_last)
+    if cut == 0:
+        return None
+
+    older, recent = history[:cut], history[cut:]
+    transcript = "\n\n".join(
+        f"[{m.get('role')}] {str(m.get('content'))[:1500]}" for m in older
+        if m.get("content"))[:_COMPACT_SUMMARY_INPUT_CHARS]
+
+    role = cfg.role("compact") if "compact" in cfg.roles else cfg.role("main")
+    summary = ""
+    async for kind, payload in llm.stream(
+            role, [{"role": "system", "content": SYSTEM},
+                   {"role": "user", "content": transcript}]):
+        if kind == "done":
+            summary = cast(Turn, payload).text
+
+    if not summary.strip():
+        return None
+
+    # Anthropic's "preserved thinking" check binds a Claude Fable 5.1 thinking
+    # block's signature to the exact conversation prefix that produced it;
+    # summarising `older` into one message is itself a history edit, which
+    # invalidates every thinking block still carried by `recent`. Dropping the
+    # blocks (unbilled, and explicitly sanctioned as the no-beta recovery path)
+    # is simpler and safer than keeping them alive across a rewritten prefix --
+    # the model just answers the next turn without that carried-over reasoning.
+    for m in recent:
+        m.pop("thinking", None)
+
+    ledger_lines = trajectory.compaction_lines(older, _COMPACT_ENTRY_CHARS)
+    ledger_body = "\n".join(ledger_lines) if ledger_lines else "(no tool calls in this range)"
+
+    # The ledger is its own message, never folded into the summary message
+    # above: mutating an already-emitted message's content would invalidate
+    # the Anthropic prompt cache for everything after it, while appending a
+    # brand-new message only ever extends the cached prefix.
+    history[:] = [
+        {"role": "user", "content": f"[Summary of earlier conversation]\n{summary}"},
+        {"role": "user", "content": f"[Action ledger for dropped range]\n{ledger_body}"},
+        *recent,
+    ]
+    return f"compacted {len(older)} messages → summary"