omega-code 0.4.0__py3-none-any.whl

omega/trace.py

@@ -0,0 +1,156 @@
+"""Per-session JSONL event trace: every `events.Event` a session emits is
+appended as one line to `~/.omega/sessions/<id>/trace.jsonl`, independent of
+what either UI chooses to render -- so `omega trace <id>` can replay a
+session's full timeline (including events a UI drops on the floor) after the
+fact, and price it against `omega.eval.prices`."""
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import asdict, is_dataclass
+from pathlib import Path
+from typing import Any
+
+from . import events, session
+
+_TOOL_TYPES = {"ToolStart", "ToolEnd"}
+
+_GLYPH = {
+    "ToolStart": "●", "ToolEnd": "└", "SubagentSpawned": "●", "SubagentDone": "✓",
+    "Compacted": "⏺", "MemoryWrite": "◆", "MemoryConsolidated": "◆", "Error": "✗",
+    "Done": "●", "Fallback": "⇄", "ModelUsed": "◇", "Usage": "◇", "Phase": "·",
+    "Checkpoint": "⎘", "Verified": "✓", "JobStarted": "⟳", "JobFinished": "✓",
+}
+
+
+def trace_path(session_id: str) -> Path:
+    return session.DIR / session_id / "trace.jsonl"
+
+
+def append(session_id: str, ev: events.Event, turn: int) -> None:
+    """Fire-and-forget: observability must never be able to break a turn, so
+    any failure here (disk full, permissions, an event with no dataclass
+    fields) is swallowed rather than raised."""
+    try:
+        if not is_dataclass(ev):
+            return
+        path = trace_path(session_id)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        line = {"t": time.time(), "turn": turn, "type": type(ev).__name__, **asdict(ev)}
+        with path.open("a") as f:
+            f.write(json.dumps(line) + "\n")
+    except Exception:
+        pass
+
+
+def _read(session_id: str) -> list[dict[str, Any]]:
+    """Tolerates a truncated trailing line the same way `session._read_jsonl`
+    does -- appends are not fsynced, so a crash mid-write leaves one."""
+    path = trace_path(session_id)
+    if not path.exists():
+        return []
+    rows: list[dict[str, Any]] = []
+    for line in path.read_text().splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            rows.append(json.loads(line))
+        except json.JSONDecodeError:
+            break
+    return rows
+
+
+def _summary(row: dict[str, Any]) -> str:
+    t = row.get("type", "")
+    if t == "ToolStart":
+        return f"{row.get('name', '')}  {row.get('args_preview', '')}"
+    if t == "ToolEnd":
+        dur = float(row.get("duration_s") or 0.0)
+        return f"{row.get('name', '')}  {row.get('outcome', '')}  ({dur:.2f}s)".strip()
+    if t == "Done":
+        return " ".join(str(row.get("text", "")).split())[:80]
+    if t == "Error":
+        message = str(row.get("message", ""))
+        return message.splitlines()[0] if message else ""
+    if t == "SubagentSpawned":
+        return f"subagent({row.get('tier', '')})  {row.get('task_preview', '')}"
+    if t == "SubagentDone":
+        return f"subagent {row.get('subagent_id', '')} done"
+    if t == "Compacted":
+        return str(row.get("note", ""))
+    if t == "MemoryWrite":
+        return f"{row.get('type', '')} '{row.get('title', '')}' ({row.get('scope', '')})"
+    if t == "MemoryConsolidated":
+        return str(row.get("summary", ""))
+    if t == "ModelUsed":
+        alias = row.get("alias") or ""
+        return f"{alias} {row.get('model', '')}".strip()
+    if t == "Usage":
+        return f"{row.get('prompt_tokens', 0)} in / {row.get('completion_tokens', 0)} out"
+    if t == "Fallback":
+        return f"{row.get('from_model', '')} -> {row.get('to_model', '')}: {row.get('reason', '')}"
+    if t == "Phase":
+        return str(row.get("state", ""))
+    if t == "Checkpoint":
+        return f"checkpoint {row.get('id', '')}"
+    if t == "Verified":
+        return str(row.get("results_summary", ""))
+    if t == "JobStarted":
+        return f"{row.get('id', '')}  {row.get('command', '')}"
+    if t == "JobFinished":
+        return f"{row.get('id', '')}  exit {row.get('exit_code', '')}"
+    return ""
+
+
+def _turn_totals(rows: list[dict[str, Any]]) -> dict[int, tuple[int, int, float | None]]:
+    """`turn -> (tokens_in, tokens_out, cost_usd)`, summed from `Usage` rows
+    and priced against whichever alias the last `ModelUsed` row announced --
+    `cost_usd` is None once that model is absent from `eval.prices.PRICES`."""
+    from .eval import prices
+
+    totals: dict[int, tuple[int, int, float | None]] = {}
+    alias: str | None = None
+    for row in rows:
+        if row.get("type") == "ModelUsed":
+            alias = row.get("alias")
+        elif row.get("type") == "Usage":
+            turn = int(row.get("turn", 0))
+            prev_in, prev_out, _ = totals.get(turn, (0, 0, None))
+            tin = prev_in + int(row.get("prompt_tokens", 0))
+            tout = prev_out + int(row.get("completion_tokens", 0))
+            cost = prices.estimate_cost(alias, tin, tout) if alias else None
+            totals[turn] = (tin, tout, cost)
+    return totals
+
+
+def render_timeline(session_id: str, *, tools_only: bool = False, raw_json: bool = False) -> str:
+    rows = _read(session_id)
+    if raw_json:
+        return "\n".join(json.dumps(r) for r in rows)
+    if not rows:
+        return f"no trace for session {session_id}"
+
+    totals = _turn_totals(rows)
+    view = [r for r in rows if r.get("type") in _TOOL_TYPES] if tools_only else rows
+    t0 = float(rows[0].get("t") or 0.0)
+
+    lines = []
+    for row in view:
+        t = str(row.get("type", ""))
+        if t == "TextDelta":
+            continue
+        offset = float(row.get("t") or t0) - t0
+        turn = row.get("turn", 0)
+        glyph = _GLYPH.get(t, "·")
+        lines.append(f"+{offset:7.2f}s  turn {turn:<3} {glyph} {t:<18}{_summary(row)}")
+
+    if totals:
+        lines.append("")
+        lines.append("per-turn totals:")
+        for turn in sorted(totals):
+            tin, tout, cost = totals[turn]
+            cost_text = f"${cost:.4f}" if cost is not None else "unknown"
+            lines.append(f"  turn {turn}: {tin} in / {tout} out tokens · {cost_text}")
+
+    return "\n".join(lines)

omega/trajectory.py

@@ -0,0 +1,146 @@
+"""The trajectory ledger: a deterministic, non-LLM record of every tool call
+made this session, derived by walking `history` rather than kept as separate
+state -- that way it survives session resume for free and can never drift
+from what actually happened.
+
+Two consumers:
+  - `render()` -- the live "## Recent actions" block `loop.run_turn` injects
+    into the volatile half of the system prompt every round.
+  - `compaction_lines()` -- the full-fidelity, uncapped-per-field lines
+    `compact.py` folds into its deterministic ledger for a dropped range.
+"""
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+from . import tools
+from .session import Message
+from .ui import format
+
+MAX_ENTRIES = 50
+ACTION_LINE_CAP = 25
+BLOCK_MAX_CHARS = 2000
+ARGS_DIGEST_CHARS = 64
+SUMMARY_CHARS = 120
+
+_DEDUPE_NOTE = ("If an identical call already appears above, do not re-run it "
+                "-- call fetch_result on its recorded artifact instead.")
+
+
+@dataclass(frozen=True)
+class Entry:
+    tool: str
+    args_digest: str
+    summary: str
+    artifact_id: str | None
+    call_id: str
+
+
+@dataclass(frozen=True)
+class _Raw:
+    tool: str
+    args: dict[str, Any]
+    result: str
+    artifact_id: str | None
+    call_id: str
+
+
+def _flatten(text: str) -> str:
+    return " ".join(text.split())
+
+
+def _cap(text: str, limit: int) -> str:
+    text = _flatten(text)
+    return text if len(text) <= limit else text[:limit - 1] + "…"
+
+
+def _describe_args(name: str, args: dict[str, Any]) -> str:
+    try:
+        described = format.describe_call(name, args)
+    except Exception:
+        described = ""
+    if described.startswith(name):
+        described = described[len(name):].strip()
+    if not described:
+        described = json.dumps(args, default=str)
+    return _flatten(described)
+
+
+def _tool_calls_by_id(history: list[Message]) -> dict[str, tuple[str, str]]:
+    out: dict[str, tuple[str, str]] = {}
+    for m in history:
+        if m.get("role") != "assistant":
+            continue
+        for tc in m.get("tool_calls") or []:
+            fn = tc.get("function") or {}
+            out[tc["id"]] = (fn.get("name", ""), fn.get("arguments", "{}"))
+    return out
+
+
+def _raw_entries(history: list[Message]) -> list[_Raw]:
+    calls = _tool_calls_by_id(history)
+    out: list[_Raw] = []
+    for m in history:
+        if m.get("role") != "tool":
+            continue
+        call_id = str(m.get("tool_call_id") or "")
+        name, args_json = calls.get(call_id, ("", "{}"))
+        if not name:
+            continue
+        try:
+            args = dict(json.loads(args_json or "{}"))
+        except json.JSONDecodeError:
+            args = {}
+        text = str(m.get("content") or "")
+        _offloaded, artifact_id = tools.offload_info(text)
+        out.append(_Raw(tool=name, args=args, result=text,
+                        artifact_id=artifact_id, call_id=call_id))
+    return out
+
+
+def build(history: list[Message]) -> list[Entry]:
+    """The last `MAX_ENTRIES` tool calls this session, oldest dropped."""
+    raws = _raw_entries(history)[-MAX_ENTRIES:]
+    return [Entry(tool=r.tool,
+                  args_digest=_cap(_describe_args(r.tool, r.args), ARGS_DIGEST_CHARS),
+                  summary=_cap(r.result, SUMMARY_CHARS),
+                  artifact_id=r.artifact_id, call_id=r.call_id)
+            for r in raws]
+
+
+def _line(tool: str, args_text: str, summary_text: str, artifact_id: str | None) -> str:
+    suffix = f" [artifact {artifact_id}]" if artifact_id else ""
+    return f"{tool}({args_text}) → {summary_text}{suffix}"
+
+
+def render(history: list[Message]) -> str:
+    """`## Recent actions` block for the volatile half of the system prompt.
+    Empty string when there is nothing to show yet."""
+    entries = build(history)[-ACTION_LINE_CAP:]
+    if not entries:
+        return ""
+    header = "## Recent actions\n"
+    footer = "\n\n" + _DEDUPE_NOTE
+    lines = [f"- {_line(e.tool, e.args_digest, e.summary, e.artifact_id)}" for e in entries]
+    # Oldest-first drop keeps the most recent (most relevant) actions when the
+    # block would otherwise blow the char budget.
+    while lines and len(header) + len("\n".join(lines)) + len(footer) > BLOCK_MAX_CHARS:
+        lines.pop(0)
+    if not lines:
+        return ""
+    return header + "\n".join(lines) + footer
+
+
+def compaction_lines(history: list[Message], entry_char_cap: int) -> list[str]:
+    """Every tool call in `history` (typically the dropped range of a
+    compaction), one deterministic line each -- no LLM involved, so it cannot
+    hallucinate. Unlike `render()`, nothing here is capped to the live
+    ledger's tight per-field limits; only the whole line is capped."""
+    lines = []
+    for r in _raw_entries(history):
+        args_text = _describe_args(r.tool, r.args)
+        line = _line(r.tool, args_text, _flatten(r.result), r.artifact_id)
+        lines.append(_cap(line, entry_char_cap) if len(line) > entry_char_cap else line)
+    return lines

omega/ui/composer.py

@@ -0,0 +1,140 @@
+"""Shared composer helpers for the TUI (and any other Python front-end).
+
+Mirrors the Swift `ComposerPrompt` helpers: `@mention` parsing, workspace
+file ranking, and flattening dropped paths into the text-only prompt the
+daemon accepts.
+"""
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+_TOKEN = re.compile(r"[A-Za-z0-9._\-++/\\]")
+_SKIP_DIRS = {
+    ".git", ".hg", ".svn",
+    "node_modules", ".build", "DerivedData", ".swiftpm",
+    "dist", "build", "Pods", "Carthage",
+    ".venv", "venv", "__pycache__", ".mypy_cache", ".pytest_cache",
+    ".next", ".nuxt", "target", ".omega", ".idea", ".vscode",
+    "xcuserdata",
+}
+_MAX_FILES = 8_000
+_MAX_DEPTH = 10
+_IMAGE_EXT = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".tif", ".tiff", ".heic", ".bmp"}
+
+
+@dataclass(frozen=True)
+class MentionQuery:
+    query: str
+    start: int  # character offset of '@'
+    end: int    # character offset of the cursor (exclusive of the token)
+
+
+def mention_query(text: str, cursor: int) -> MentionQuery | None:
+    """If `cursor` sits inside an `@token` at a word boundary, return it."""
+    if not text:
+        return None
+    cursor = max(0, min(cursor, len(text)))
+    start = cursor
+    while start > 0 and _TOKEN.match(text[start - 1]):
+        start -= 1
+    if start == 0 or text[start - 1] != "@":
+        return None
+    at = start - 1
+    if at > 0 and not text[at - 1].isspace():
+        return None
+    return MentionQuery(query=text[start:cursor], start=at, end=cursor)
+
+
+def rank_files(files: list[str], query: str, limit: int = 12) -> list[str]:
+    q = query.strip().lower()
+    if not q:
+        return sorted(files, key=lambda p: (p.count("/"), p.lower()))[:limit]
+
+    hits: list[tuple[int, int, str, str]] = []
+    for path in files:
+        filename = path.rsplit("/", 1)[-1]
+        path_l, file_l = path.lower(), filename.lower()
+        if file_l == q:
+            score = 0
+        elif file_l.startswith(q):
+            score = 1
+        elif path_l.startswith(q):
+            score = 2
+        elif q in file_l:
+            score = 3
+        elif q in path_l:
+            score = 4
+        else:
+            continue
+        hits.append((score, len(filename), file_l, path))
+    hits.sort()
+    return [h[3] for h in hits[:limit]]
+
+
+def insert_mention(text: str, cursor: int, path: str) -> tuple[str, int]:
+    """Replace the active `@token` (or append) with `` `path` ``. Returns (text, cursor)."""
+    q = mention_query(text, cursor)
+    replacement = f"`{path}` "
+    if q is None:
+        new = text + replacement
+        return new, len(new)
+    new = text[:q.start] + replacement + text[q.end:]
+    return new, q.start + len(replacement)
+
+
+def build_prompt(text: str, attachment_paths: list[str]) -> str:
+    trimmed = text.strip()
+    paths = [p.strip() for p in attachment_paths if p.strip()]
+    if not paths:
+        return trimmed
+    lines: list[str] = []
+    if trimmed:
+        lines.append(trimmed)
+        lines.append("")
+    lines.append("Attached:")
+    lines.extend(f"- {p}" for p in paths)
+    return "\n".join(lines)
+
+
+def display_path(path: str, workspace_root: str | None) -> str:
+    if not workspace_root:
+        return path
+    root = os.path.realpath(workspace_root)
+    file = os.path.realpath(path)
+    if file == root:
+        return "."
+    prefix = root if root.endswith(os.sep) else root + os.sep
+    if file.startswith(prefix):
+        return file[len(prefix):]
+    return path
+
+
+def is_image_path(path: str) -> bool:
+    return Path(path).suffix.lower() in _IMAGE_EXT
+
+
+def list_workspace_files(root: str) -> list[str]:
+    root_path = Path(root).expanduser().resolve()
+    if not root_path.is_dir():
+        return []
+    files: list[str] = []
+    root_s = str(root_path)
+    prefix = root_s if root_s.endswith(os.sep) else root_s + os.sep
+    for dirpath, dirnames, filenames in os.walk(root_path):
+        rel_dir = dirpath[len(prefix):] if dirpath.startswith(prefix) else ""
+        depth = 0 if not rel_dir else rel_dir.count(os.sep) + 1
+        if depth > _MAX_DEPTH:
+            dirnames[:] = []
+            continue
+        dirnames[:] = [d for d in dirnames if d not in _SKIP_DIRS and not d.startswith(".")]
+        for name in filenames:
+            if name.startswith("."):
+                continue
+            rel = name if not rel_dir else f"{rel_dir}/{name}".replace("\\", "/")
+            files.append(rel)
+            if len(files) >= _MAX_FILES:
+                return files
+    return files