From 1546ac088d44893d7d89bd47ed11754e918ccd09 Mon Sep 17 00:00:00 2001 From: "Aaron K. Clark" Date: Thu, 16 Jul 2026 22:22:04 -0500 Subject: [PATCH 1/3] feat: add AI token accounting and expense profiles (v4.0.0) --- CHANGELOG.md | 11 ++ README.md | 28 +++ pyproject.toml | 2 +- src/omind/__init__.py | 2 +- src/omind/ai_usage.py | 336 ++++++++++++++++++++++++++++++++ src/omind/checkpoint.py | 65 +++--- src/omind/cli.py | 68 ++++++- src/omind/hooks.py | 114 +++++++++-- src/omind/verify.py | 62 ++++-- src/omind/web/app.py | 22 ++- src/omind/web/static/app.css | 52 +++++ src/omind/web/static/app.js | 108 +++++++++- src/omind/web/static/index.html | 1 + tests/test_ai_usage.py | 172 ++++++++++++++++ tests/test_checkpoint.py | 18 ++ tests/test_hooks.py | 44 ++++- tests/test_web.py | 17 +- uv.lock | 2 +- 18 files changed, 1054 insertions(+), 70 deletions(-) create mode 100644 src/omind/ai_usage.py create mode 100644 tests/test_ai_usage.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 700363f..d2e7364 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [4.0.0] - 2026-07-16 + +### Added +- **OMI-attributable AI token accounting and expense profiles.** A privacy-safe, + per-vault JSONL ledger records session-priming estimates and provider-reported + verifier/checkpoint usage without storing prompts or responses. `omind ai + profile low|medium|high` applies hard context budgets and progressively disables + optional model work; `omind ai usage` reports exact, estimated, cached, and + avoided tokens in text or JSON. The local web app adds the same profile control + and 24-hour/7-day/30-day/all-time usage dashboard through `/api/ai/*` endpoints. + ## [3.8.6] - 2026-07-06 ### Changed diff --git a/README.md b/README.md index 4957bbb..695eb29 100644 --- a/README.md +++ b/README.md @@ -62,6 +62,8 @@ reads and writes as long-term memory. `omind` does two things with it: - **`omind checkpoint`** — a scheduled job that records what the agent has been doing every N minutes into a daily worklog note, by mining the trails the hooks already capture (see [Activity checkpoints](#activity-checkpoints) below). +- **`omind ai`** — account for tokens attributable to OMI and select a manual + low/medium/high model-expense profile that bounds priming and optional model calls. The web UI works **fully offline** (fonts, styles, and the Markdown renderer are vendored — no CDN). It shows **backlinks** for the open note, refreshes the list @@ -360,6 +362,32 @@ text. Because it's a scheduled job — the same systemd-user-timer mechanism as `omind backup` and `omind mesh` — it doesn't depend on the agent's cooperation, which is what makes it a reliable *record* rather than a hopeful instruction. +## AI token usage and expense profiles + +omind keeps a machine-local, per-vault ledger for the AI tokens it causes: session +priming injected into an agent plus the verifier and optional checkpoint +`claude -p` calls. It does **not** parse or claim the rest of an agent session. +Provider-reported subprocess usage is recorded exactly; provider-neutral priming +is shown as an estimate (`ceil(characters / 4)`) and never mixed into the +provider-reported subtotal. The ledger contains counts and operational metadata +only — never prompts, responses, note contents, or credentials. + +```bash +omind ai profile # show saved/effective profile +omind ai profile medium # persist low, medium, or high per vault +omind ai usage --since 7d # 24h, 7d, 30d, or all +omind ai usage --since all --json # machine-readable report +``` + +Profiles describe the expense of the model already selected; they do not select a +model. `low` is the backward-compatible default (48,000-character priming cap and +all optional calls), `medium` halves the priming/prompt inputs, and `high` caps +priming at 8,000 characters and uses deterministic verifier/checkpoint behavior +without optional model calls. Set `OMI_AI_EXPENSE=low|medium|high` for a temporary +override; it wins over the saved profile. The web app's **AI Usage** view exposes +the same profile control, time windows, exact/estimated breakdown, per-operation +totals, and estimated avoided tokens. + ## Other agents: Hermes, OpenClaw, OpenCode, Codex, Gemini, Claude Desktop, Kiro, VS Code, Amazon Q [Claude Code](https://github.com/anthropics/claude-code) is the default, but the diff --git a/pyproject.toml b/pyproject.toml index effad3c..f3e3089 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "omind" -version = "3.8.6" +version = "4.0.0" description = "Reproduce the OMI/Obsidian memory integration for AI agents, plus a local web app to view, edit, and add memory entries." readme = "README.md" requires-python = ">=3.10" diff --git a/src/omind/__init__.py b/src/omind/__init__.py index 66f633c..c09ca6f 100644 --- a/src/omind/__init__.py +++ b/src/omind/__init__.py @@ -2,4 +2,4 @@ # Copyright 2026 Aaron K. Clark """omind — OMI/Obsidian memory tooling for AI agents.""" -__version__ = "3.8.6" +__version__ = "4.0.0" diff --git a/src/omind/ai_usage.py b/src/omind/ai_usage.py new file mode 100644 index 0000000..74cbd58 --- /dev/null +++ b/src/omind/ai_usage.py @@ -0,0 +1,336 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 Aaron K. Clark +"""OMI-attributable AI token accounting and model-expense profiles. + +The ledger deliberately stores counts and operational metadata only: never the +prompt, response, note body, or user text that produced those counts. Provider +usage is exact when ``claude -p --output-format json`` reports it; priming and +legacy/malformed responses use a clearly-labelled, provider-neutral estimate. +""" + +from __future__ import annotations + +import json +import math +import os +import shutil +import subprocess +from dataclasses import asdict, dataclass +from datetime import datetime, timedelta +from pathlib import Path +from typing import Any + +from omind import filelock, paths + +PROFILES = ("low", "medium", "high") +PROFILE_ENV = "OMI_AI_EXPENSE" +DEFAULT_PROFILE = "low" + + +@dataclass(frozen=True) +class ProfilePolicy: + context_chars: int + verifier_task_chars: int + verifier_material_chars: int + checkpoint_actions: int + checkpoint_guard_events: int + verifier_llm: bool + checkpoint_llm: bool + + +PROFILE_POLICIES: dict[str, ProfilePolicy] = { + "low": ProfilePolicy(48_000, 1_000, 2_000, 60, 30, True, True), + "medium": ProfilePolicy(24_000, 500, 1_000, 30, 15, True, True), + "high": ProfilePolicy(8_000, 1_000, 2_000, 0, 0, False, False), +} + + +def _vault_key(omi_dir: Path | str) -> str: + import hashlib + + resolved = str(Path(omi_dir).expanduser().resolve()) + return hashlib.sha256(resolved.encode()).hexdigest()[:12] + + +def profile_path(omi_dir: Path | str) -> Path: + return paths.state_dir() / f"ai-profile-{_vault_key(omi_dir)}.json" + + +def usage_path(omi_dir: Path | str) -> Path: + return paths.state_dir() / f"ai-usage-{_vault_key(omi_dir)}.jsonl" + + +def saved_profile(omi_dir: Path | str) -> str: + try: + data = json.loads(profile_path(omi_dir).read_text(encoding="utf-8")) + except (OSError, ValueError, TypeError): + return DEFAULT_PROFILE + value = data.get("profile") if isinstance(data, dict) else None + return str(value) if value in PROFILES else DEFAULT_PROFILE + + +def profile_info(omi_dir: Path | str) -> dict[str, str]: + saved = saved_profile(omi_dir) + override = os.environ.get(PROFILE_ENV, "").strip().lower() + if override in PROFILES: + return {"saved": saved, "effective": override, "source": "environment"} + return { + "saved": saved, + "effective": saved, + "source": "saved" if profile_path(omi_dir).exists() else "default", + } + + +def effective_profile(omi_dir: Path | str) -> str: + return profile_info(omi_dir)["effective"] + + +def policy(omi_dir: Path | str) -> ProfilePolicy: + return PROFILE_POLICIES[effective_profile(omi_dir)] + + +def set_profile(omi_dir: Path | str, profile: str) -> dict[str, str]: + value = profile.strip().lower() + if value not in PROFILES: + raise ValueError(f"profile must be one of: {', '.join(PROFILES)}") + path = profile_path(omi_dir) + paths.atomic_write_text(path, json.dumps({"profile": value}, indent=2) + "\n", mode=0o600) + return profile_info(omi_dir) + + +def estimate_tokens(text_or_chars: str | int) -> int: + """Provider-neutral estimate used only when exact tokenizer usage is absent.""" + chars = text_or_chars if isinstance(text_or_chars, int) else len(text_or_chars) + return math.ceil(max(0, chars) / 4) + + +def log_event( + omi_dir: Path | str, + operation: str, + *, + status: str = "executed", + measurement: str = "exact", + model: str = "", + input_tokens: int = 0, + output_tokens: int = 0, + cache_read_tokens: int = 0, + cache_write_tokens: int = 0, + characters: int = 0, + avoided_tokens: int = 0, + reason: str = "", + now: datetime | None = None, +) -> None: + """Append a privacy-safe usage record. Never raises into an agent hook.""" + record: dict[str, Any] = { + "ts": (now or datetime.now()).isoformat(timespec="seconds"), + "operation": operation, + "profile": effective_profile(omi_dir), + "status": status, + "measurement": measurement, + "input_tokens": max(0, int(input_tokens)), + "output_tokens": max(0, int(output_tokens)), + "cache_read_tokens": max(0, int(cache_read_tokens)), + "cache_write_tokens": max(0, int(cache_write_tokens)), + "characters": max(0, int(characters)), + "avoided_tokens": max(0, int(avoided_tokens)), + } + if model: + record["model"] = model[:120] + if reason: + record["reason"] = reason[:160] + try: + path = usage_path(omi_dir) + path.parent.mkdir(parents=True, exist_ok=True) + binary = getattr(os, "O_BINARY", 0) + fd = os.open(path, os.O_WRONLY | os.O_APPEND | os.O_CREAT | binary, 0o600) + try: + filelock.lock_fd(fd) + os.write(fd, (json.dumps(record, separators=(",", ":")) + "\n").encode()) + finally: + filelock.unlock_fd(fd) + os.close(fd) + except (OSError, ValueError, TypeError): + return + + +def record_priming(omi_dir: Path | str, characters: int, *, avoided_characters: int = 0) -> None: + log_event( + omi_dir, + "priming", + measurement="estimated", + characters=characters, + input_tokens=estimate_tokens(characters), + avoided_tokens=estimate_tokens(avoided_characters), + ) + + +def read_events(omi_dir: Path | str) -> list[dict[str, Any]]: + try: + lines = usage_path(omi_dir).read_text(encoding="utf-8", errors="replace").splitlines() + except OSError: + return [] + events: list[dict[str, Any]] = [] + for line in lines: + try: + value = json.loads(line) + except (ValueError, TypeError): + continue + if isinstance(value, dict): + events.append(value) + return events + + +def parse_window(value: str) -> timedelta | None: + clean = (value or "").strip().lower() + if clean == "all": + return None + import re + + match = re.fullmatch(r"(\d+)([hd])", clean) + if not match: + raise ValueError("--since must be 24h, 7d, 30d, or all") + number = int(match.group(1)) + return timedelta(hours=number) if match.group(2) == "h" else timedelta(days=number) + + +def usage_summary( + omi_dir: Path | str, *, since: str = "7d", now: datetime | None = None +) -> dict[str, Any]: + window = parse_window(since) + current = now or datetime.now() + cutoff = current - window if window is not None else None + events: list[dict[str, Any]] = [] + for event in read_events(omi_dir): + try: + stamp = datetime.fromisoformat(str(event.get("ts") or "")) + except ValueError: + continue + if stamp.tzinfo is not None: + stamp = stamp.astimezone().replace(tzinfo=None) + if cutoff is None or stamp >= cutoff: + events.append(event) + + numeric = ( + "input_tokens", + "output_tokens", + "cache_read_tokens", + "cache_write_tokens", + "avoided_tokens", + ) + + def totals(rows: list[dict[str, Any]]) -> dict[str, int]: + result: dict[str, int] = {} + for key in numeric: + value = 0 + for row in rows: + try: + value += max(0, int(row.get(key) or 0)) + except (ValueError, TypeError): + continue + result[key] = value + return result + + operations: dict[str, dict[str, int]] = {} + for operation in ("priming", "verifier", "checkpoint"): + operations[operation] = totals([e for e in events if e.get("operation") == operation]) + return { + "since": since, + "profile": profile_info(omi_dir), + "events": len(events), + "totals": totals(events), + "exact": totals([e for e in events if e.get("measurement") == "exact"]), + "estimated": totals([e for e in events if e.get("measurement") == "estimated"]), + "operations": operations, + } + + +def _usage_int(usage: dict[str, Any], *names: str) -> int: + for name in names: + value = usage.get(name) + if isinstance(value, (int, float)): + return max(0, int(value)) + return 0 + + +def run_claude( + omi_dir: Path | str, + operation: str, + prompt: str, + *, + timeout: int, + allowed: bool = True, +) -> str | None: + """Run a headless Claude call, account for it, and return response text. + + Any failure preserves the historic fail-open contract by returning ``None``. + """ + if not allowed: + log_event( + omi_dir, + operation, + status="skipped", + measurement="estimated", + characters=len(prompt), + avoided_tokens=estimate_tokens(prompt), + reason="disabled by expense profile", + ) + return None + claude = shutil.which("claude") + if not claude: + return None + try: + result = subprocess.run( + [claude, "-p", "--output-format", "json", prompt], + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + timeout=timeout, + ) + except (subprocess.TimeoutExpired, OSError): + return None + if result.returncode != 0: + return None + stdout = (result.stdout or "").strip() + if not stdout: + return None + try: + payload = json.loads(stdout) + except ValueError: + log_event( + omi_dir, + operation, + measurement="estimated", + characters=len(prompt) + len(stdout), + input_tokens=estimate_tokens(prompt), + output_tokens=estimate_tokens(stdout), + ) + return stdout + if not isinstance(payload, dict): + return None + text = payload.get("result") + if not isinstance(text, str) or not text.strip(): + return None + usage = payload.get("usage") + usage = usage if isinstance(usage, dict) else {} + has_usage = any(isinstance(v, (int, float)) for v in usage.values()) + log_event( + omi_dir, + operation, + measurement="exact" if has_usage else "estimated", + model=str(payload.get("model") or ""), + characters=len(prompt) + len(text), + input_tokens=(_usage_int(usage, "input_tokens") if has_usage else estimate_tokens(prompt)), + output_tokens=(_usage_int(usage, "output_tokens") if has_usage else estimate_tokens(text)), + cache_read_tokens=_usage_int(usage, "cache_read_input_tokens", "cache_read_tokens"), + cache_write_tokens=_usage_int( + usage, "cache_creation_input_tokens", "cache_write_input_tokens" + ), + ) + return text.strip() + + +def profile_payload(omi_dir: Path | str) -> dict[str, Any]: + info: dict[str, Any] = profile_info(omi_dir) + info["policies"] = {name: asdict(value) for name, value in PROFILE_POLICIES.items()} + return info diff --git a/src/omind/checkpoint.py b/src/omind/checkpoint.py index 46399c1..225e8de 100644 --- a/src/omind/checkpoint.py +++ b/src/omind/checkpoint.py @@ -144,47 +144,68 @@ def gather_activity(omi_dir: Path | str, cutoff: datetime, now: datetime) -> Act return Activity(actions=actions, guard_events=guard) -def _llm_narrative(activity: Activity, since: str) -> str | None: +def _llm_narrative( + activity: Activity, since: str, omi_dir: Path | str | None = None +) -> str | None: """A one-paragraph narrative from headless ``claude -p``; ``None`` on any unavailability/error/timeout (caller falls back to the deterministic summary).""" - claude = shutil.which("claude") - if not claude: - return None - lines = [f"{a['time']} {a['event']} {a['tool']} {a['detail']}" for a in activity.actions[-60:]] - guard = [f"{e.get('ts')} {e.get('outcome')} {e.get('tool')} {e.get('command')}" for e in - activity.guard_events[-30:]] + from omind import ai_usage + + limits = ai_usage.policy(omi_dir) if omi_dir is not None else None + action_limit = limits.checkpoint_actions if limits else 60 + guard_limit = limits.checkpoint_guard_events if limits else 30 + # A high-expense profile still builds the prompt shape so the skipped event + # can report an avoided-token estimate without exposing its contents. + lines = [ + f"{a['time']} {a['event']} {a['tool']} {a['detail']}" + for a in activity.actions[-max(action_limit, 60 if action_limit == 0 else action_limit) :] + ] + guard = [ + f"{e.get('ts')} {e.get('outcome')} {e.get('tool')} {e.get('command')}" + for e in activity.guard_events[ + -max(guard_limit, 30 if guard_limit == 0 else guard_limit) : + ] + ] prompt = ( "Summarize what this agent worked on in the last " f"{since}, in 1-3 sentences, factual and concise. No preamble.\n\n" "ACTIONS:\n" + "\n".join(lines) + "\n\nGUARD EVENTS:\n" + "\n".join(guard) + "\n" ) + if omi_dir is not None: + return ai_usage.run_claude( + omi_dir, + "checkpoint", + prompt, + timeout=_LLM_TIMEOUT, + allowed=limits.checkpoint_llm if limits else True, + ) + claude = shutil.which("claude") + if not claude: + return None try: result = subprocess.run( - [claude, "-p", prompt], - capture_output=True, - text=True, - encoding="utf-8", - errors="replace", - timeout=_LLM_TIMEOUT, + [claude, "-p", prompt], capture_output=True, text=True, timeout=_LLM_TIMEOUT ) except (subprocess.TimeoutExpired, OSError): return None - # A non-zero exit means claude printed a diagnostic, not a summary — don't - # adopt an error message as the day's narrative. - if result.returncode != 0: - return None - text = (result.stdout or "").strip() - return text or None + return (result.stdout or "").strip() or None if result.returncode == 0 else None -def render_section(activity: Activity, since: str, now: datetime, *, llm: bool = False) -> str: +def render_section( + activity: Activity, + since: str, + now: datetime, + *, + llm: bool = False, + omi_dir: Path | str | None = None, +) -> str: """One worklog section for this run (deterministic; ``llm`` adds a narrative).""" header = f"### {now.strftime('%H:%M')} — last {since}" if activity.is_empty(): return f"{header}\n- no recorded activity in this window\n" lines = [header] if llm: - narrative = _llm_narrative(activity, since) + narrative = _llm_narrative(activity, since, omi_dir) if narrative: lines.append(narrative) tools = Counter(a["tool"] for a in activity.actions if a["tool"]) @@ -234,7 +255,7 @@ def write_checkpoint( now = now or datetime.now() cutoff = now - parse_since(since) activity = gather_activity(omi_dir, cutoff, now) - section = render_section(activity, since, now, llm=llm) + section = render_section(activity, since, now, llm=llm, omi_dir=omi_dir) title = f"Worklog {now.strftime('%Y-%m-%d')}" details = _append_section(_existing_details(omi_dir, title), section) fields = NoteFields( diff --git a/src/omind/cli.py b/src/omind/cli.py index ceb6bfd..5110f73 100644 --- a/src/omind/cli.py +++ b/src/omind/cli.py @@ -64,7 +64,7 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--version", action="version", version=f"omind {__version__}") sub = parser.add_subparsers( dest="command", - metavar="{setup,quickstart,serve,doctor,self-update,backup,export,import,reindex,note,rollup,hook}", + metavar="{setup,quickstart,serve,doctor,self-update,backup,ai,export,import,reindex,note,rollup,hook}", ) setup = sub.add_parser( @@ -400,6 +400,20 @@ def build_parser() -> argparse.ArgumentParser: ) _add_vault_args(checkpoint) + ai = sub.add_parser( + "ai", help="inspect OMI-attributable token usage and select a model-expense profile" + ) + ai_sub = ai.add_subparsers(dest="ai_command", required=True) + ai_profile = ai_sub.add_parser("profile", help="show or set low/medium/high expense mode") + ai_profile.add_argument("profile", nargs="?", choices=("low", "medium", "high")) + _add_vault_args(ai_profile) + ai_usage = ai_sub.add_parser("usage", help="summarize OMI-attributable token usage") + ai_usage.add_argument( + "--since", choices=("24h", "7d", "30d", "all"), default="7d" + ) + ai_usage.add_argument("--json", action="store_true", help="emit machine-readable JSON") + _add_vault_args(ai_usage) + rollup = sub.add_parser( "rollup", help="compact weeks of daily session journals into one summary note each, " @@ -935,6 +949,56 @@ def _run_checkpoint(args: argparse.Namespace) -> int: return 0 +def _run_ai(args: argparse.Namespace) -> int: + import json + + from omind import ai_usage + + omi_dir = (args.vault / args.folder).expanduser() + if args.ai_command == "profile": + try: + info = ( + ai_usage.set_profile(omi_dir, args.profile) + if args.profile + else ai_usage.profile_info(omi_dir) + ) + except (OSError, ValueError) as exc: + print(f"error: could not set AI expense profile: {exc}", file=sys.stderr) + return 1 + print( + f"saved={info['saved']} effective={info['effective']} source={info['source']}" + ) + if info["source"] == "environment": + print(f"note: {ai_usage.PROFILE_ENV} overrides the saved profile") + return 0 + summary = ai_usage.usage_summary(omi_dir, since=args.since) + if args.json: + print(json.dumps(summary, indent=2)) + return 0 + totals = summary["totals"] + profile = summary["profile"] + print( + f"AI usage ({args.since}) — profile {profile['effective']} " + f"({profile['source']})" + ) + print( + f"input {totals['input_tokens']:,} | output {totals['output_tokens']:,} | " + f"cache read {totals['cache_read_tokens']:,} | " + f"cache write {totals['cache_write_tokens']:,}" + ) + print( + f"exact input {summary['exact']['input_tokens']:,} | " + f"estimated input {summary['estimated']['input_tokens']:,} | " + f"estimated avoided {totals['avoided_tokens']:,}" + ) + for name, values in summary["operations"].items(): + print( + f" {name}: input {values['input_tokens']:,}, " + f"output {values['output_tokens']:,}, avoided {values['avoided_tokens']:,}" + ) + return 0 + + def _split_csv(value: str) -> list[str]: """Split a comma-separated CLI flag into a clean list.""" return [item.strip() for item in value.split(",") if item.strip()] @@ -1107,6 +1171,8 @@ def main(argv: list[str] | None = None) -> int: return _run_graph(args) if args.command == "checkpoint": return _run_checkpoint(args) + if args.command == "ai": + return _run_ai(args) if args.command == "reindex": return _run_reindex(args) if args.command == "convert": diff --git a/src/omind/hooks.py b/src/omind/hooks.py index 1df068e..9673941 100644 --- a/src/omind/hooks.py +++ b/src/omind/hooks.py @@ -64,6 +64,8 @@ _SESSION_STATE_GLOB = "Session State *.md" _JOURNAL_GLOB = paths.JOURNAL_GLOB _JOURNAL_TAIL_BULLETS = 20 +# Backward-compatible name for the low-expense/default cap. Runtime selection +# lives in :mod:`omind.ai_usage`. _TOTAL_CONTEXT_CHAR_CAP = 48_000 _TRUNCATION_MARKER = "\n…[truncated]" @@ -341,7 +343,44 @@ def _update_nudge_line() -> str | None: return f"⚠️ {nudge}" if nudge else None -def build_session_start_context(omi_dir: Path | str) -> str: +def _allocate_sections( + sections: list[str], budget: int, *, minimum: int = 1_000 +) -> list[str]: + """Fit sections into ``budget`` while giving every section a useful floor. + + The floor prevents a long index from consuming the whole context before the + Playbook/workflow/persona sections get a voice. Unused space is then handed + out in priority order (the caller's order). + """ + if not sections or budget <= len(_TRUNCATION_MARKER): + return [] + allocations = [min(len(section), minimum) for section in sections] + total = sum(allocations) + if total > budget: + share = max(0, budget // len(sections)) + allocations = [min(len(section), share) for section in sections] + remaining = max(0, budget - sum(allocations)) + for index, section in enumerate(sections): + extra = min(remaining, len(section) - allocations[index]) + allocations[index] += extra + remaining -= extra + if remaining <= 0: + break + fitted: list[str] = [] + for section, size in zip(sections, allocations, strict=True): + if size <= len(_TRUNCATION_MARKER): + continue + fitted.append( + section + if size >= len(section) + else section[: size - len(_TRUNCATION_MARKER)] + _TRUNCATION_MARKER + ) + return fitted + + +def build_session_start_context( + omi_dir: Path | str, *, _context_cap: int | None = None +) -> str: """Build the SessionStart ``additionalContext`` payload. Injects the *content* of the OMI priming notes (:data:`PRIMING_FILES`) @@ -349,9 +388,10 @@ def build_session_start_context(omi_dir: Path | str) -> str: the agent issuing reads, then two dynamic sections: the newest ``Session State YYYY-MM-DD`` handoff note and the last :data:`_JOURNAL_TAIL_BULLETS` bullets of the newest auto-journal. The whole - payload is capped at :data:`_TOTAL_CONTEXT_CHAR_CAP` chars — static files - always win the budget; dynamic sections truncate (or drop) first. Falls - back to a read-the-vault reminder if nothing can be read. Never raises. + payload is capped by the active AI expense profile. Every static note gets a + useful floor and 20% of the remaining budget is reserved for current dynamic + state. Falls back to a read-the-vault reminder if nothing can be read. Never + raises. """ directory = Path(omi_dir) sections: list[str] = [] @@ -393,35 +433,64 @@ def build_session_start_context(omi_dir: Path | str) -> str: "already read. Read any [[wikilinked]] note you need before acting." ) if not sections and not dynamic: - return prefix + ( + fallback = prefix + ( header + " (Priming notes could not be read this session; read index.md, " "Memory Workflow.md, and CLAUDE CODE PERSONALITY.md from the vault " "directly.)" ) - - payload = "\n\n".join([header, *sections]) # static sections are never cut - for section in dynamic: - remaining = _TOTAL_CONTEXT_CHAR_CAP - len(payload) - len("\n\n") - if remaining <= len(_TRUNCATION_MARKER): - break # no useful room left for dynamic content - if len(section) > remaining: - section = section[: remaining - len(_TRUNCATION_MARKER)] + _TRUNCATION_MARKER - payload += "\n\n" + section - return prefix + payload + from omind import ai_usage + + return fallback[: (_context_cap or ai_usage.policy(omi_dir).context_chars)] + + from omind import ai_usage + + cap = _context_cap or ai_usage.policy(omi_dir).context_chars + base = prefix + header + remaining = max(0, cap - len(base) - 2) + # Keep 20% of the available section budget for current state. The newest + # handoff gets 60% of that pool and the journal tail gets 40%; when either is + # absent the other naturally receives the whole pool. + dynamic_budget = int(remaining * 0.20) if dynamic else 0 + static_budget = remaining - dynamic_budget + fitted_static = _allocate_sections(sections, static_budget) + fitted_dynamic: list[str] = [] + if len(dynamic) == 1: + fitted_dynamic = _allocate_sections(dynamic, dynamic_budget, minimum=0) + elif len(dynamic) >= 2: + first = _allocate_sections(dynamic[:1], int(dynamic_budget * 0.60), minimum=0) + second = _allocate_sections( + dynamic[1:2], dynamic_budget - int(dynamic_budget * 0.60), minimum=0 + ) + fitted_dynamic = [*first, *second] + payload = "\n\n".join([base, *fitted_static, *fitted_dynamic]) + if len(payload) > cap: + payload = payload[: cap - len(_TRUNCATION_MARKER)] + _TRUNCATION_MARKER + return payload def emit_session_start_context(omi_dir: Path | str, out: TextIO | None = None) -> None: """Emit OMI priming-note content as SessionStart ``additionalContext``. Never raises.""" sink = out if out is not None else sys.stdout + context = build_session_start_context(omi_dir) payload = { "hookSpecificOutput": { "hookEventName": "SessionStart", - "additionalContext": build_session_start_context(omi_dir), + "additionalContext": context, } } try: sink.write(json.dumps(payload) + "\n") + from omind import ai_usage + + baseline = ( + build_session_start_context(omi_dir, _context_cap=_TOTAL_CONTEXT_CHAR_CAP) + if ai_usage.policy(omi_dir).context_chars < _TOTAL_CONTEXT_CHAR_CAP + else context + ) + ai_usage.record_priming( + omi_dir, len(context), avoided_characters=max(0, len(baseline) - len(context)) + ) except Exception as exc: _record_failure(f"emit_session_start_context({omi_dir})", exc) @@ -482,8 +551,19 @@ def emit_pre_llm_call_context( guard.begin_turn(session, str(event.get("prompt") or "")) if _already_primed(session): return - payload = {"context": build_session_start_context(omi_dir)} + context = build_session_start_context(omi_dir) + payload = {"context": context} sink.write(json.dumps(payload) + "\n") + from omind import ai_usage + + baseline = ( + build_session_start_context(omi_dir, _context_cap=_TOTAL_CONTEXT_CHAR_CAP) + if ai_usage.policy(omi_dir).context_chars < _TOTAL_CONTEXT_CHAR_CAP + else context + ) + ai_usage.record_priming( + omi_dir, len(context), avoided_characters=max(0, len(baseline) - len(context)) + ) except Exception as exc: _record_failure(f"emit_pre_llm_call_context({omi_dir})", exc) diff --git a/src/omind/verify.py b/src/omind/verify.py index fee5d8a..cf04944 100644 --- a/src/omind/verify.py +++ b/src/omind/verify.py @@ -271,34 +271,48 @@ def _parse_verdict(text: str) -> bool | None: return None -def _ask_model(task: str, text: str) -> bool | None: +def _ask_model(task: str, text: str, omi_dir: Path | str | None = None) -> bool | None: """Ask headless ``claude -p`` whether the consult was relevant. ``None`` on any unavailability/error/timeout (the caller fails open).""" - claude = shutil.which("claude") - if not claude: - return None + from omind import ai_usage + + limits = ai_usage.policy(omi_dir) if omi_dir is not None else None + task_cap = limits.verifier_task_chars if limits else 1_000 + material_cap = limits.verifier_material_chars if limits else 2_000 prompt = ( "You are an OMI-compliance relevance checker. An agent was told to consult " "its memory (OMI) before acting on a task, and it consulted the material " "below. Answer with exactly one word — RELEVANT or IRRELEVANT — for whether " "that material is relevant to the task.\n\n" f"{_past_mistakes_context()}" - f"TASK:\n{task[:1000]}\n\n" - f"CONSULTED MATERIAL:\n{text[:2000]}\n" + f"TASK:\n{task[:task_cap]}\n\n" + f"CONSULTED MATERIAL:\n{text[:material_cap]}\n" ) try: timeout = int(os.environ.get(_TIMEOUT_ENV) or _DEFAULT_TIMEOUT) except ValueError: timeout = _DEFAULT_TIMEOUT - try: - result = subprocess.run( - [claude, "-p", prompt], capture_output=True, text=True, timeout=timeout - ) - except (subprocess.TimeoutExpired, OSError): - return None - if result.returncode != 0: - return None - return _parse_verdict(result.stdout or "") + if omi_dir is None: + # Compatibility path for the public pure ``judge`` helper. Real hook + # calls always supply the vault and therefore use the accounted wrapper. + claude = shutil.which("claude") + if not claude: + return None + try: + result = subprocess.run( + [claude, "-p", prompt], capture_output=True, text=True, timeout=timeout + ) + except (subprocess.TimeoutExpired, OSError): + return None + return _parse_verdict(result.stdout or "") if result.returncode == 0 else None + response = ai_usage.run_claude( + omi_dir, + "verifier", + prompt, + timeout=timeout, + allowed=limits.verifier_llm if limits else True, + ) + return _parse_verdict(response or "") def _signal_score(signal: str, text: str) -> float: @@ -316,7 +330,13 @@ def _signal_score(signal: str, text: str) -> float: return max(keyword, semantic) if semantic is not None else keyword -def _judge_scored(task: str, activity: str, text: str, pending: str = "") -> tuple[bool, float]: +def _judge_scored( + task: str, + activity: str, + text: str, + pending: str = "", + omi_dir: Path | str | None = None, +) -> tuple[bool, float]: """:func:`judge_with_activity` plus the deterministic score it decided on, so the off-topic log can carry WHAT the consult was judged against (#148) — without it, a sampling pass cannot adjudicate false positives after the fact.""" @@ -333,7 +353,11 @@ def _judge_scored(task: str, activity: str, text: str, pending: str = "") -> tup return True, score if score <= low: return False, score - verdict = _ask_model(task or activity or pending, text) + verdict = ( + _ask_model(task or activity or pending, text) + if omi_dir is None + else _ask_model(task or activity or pending, text, omi_dir) + ) return (True if verdict is None else verdict), score @@ -419,7 +443,9 @@ def verify_consult( activity = recent_activity(session, omi_dir, now=now) # #96/#97: the gate-blocked action (path noise stripped) — the agent's freshest intent. pending = retrieve.normalize_intent(guard.pending_intent(session)) - judged, score = _judge_scored(task, activity, _consult_text(kind, target, omi_dir), pending) + judged, score = _judge_scored( + task, activity, _consult_text(kind, target, omi_dir), pending, omi_dir + ) relevant = _always_relevant(target) or _guard_demanded(session, target) or judged guard.record_consult(session, kind=kind, target=target, relevant=relevant) if relevant: diff --git a/src/omind/web/app.py b/src/omind/web/app.py index ca93883..b674d2e 100644 --- a/src/omind/web/app.py +++ b/src/omind/web/app.py @@ -13,13 +13,14 @@ from collections.abc import Callable from dataclasses import asdict from pathlib import Path -from typing import TypeVar +from typing import Literal, TypeVar from fastapi import FastAPI, HTTPException from fastapi.staticfiles import StaticFiles from pydantic import BaseModel from starlette.middleware.trustedhost import TrustedHostMiddleware +from omind import ai_usage from omind import graph as graph_mod from omind.store import ( NoteConflictError, @@ -63,6 +64,10 @@ class RawUpdate(BaseModel): content: str +class AIProfileUpdate(BaseModel): + profile: Literal["low", "medium", "high"] + + #: Host headers accepted by default (a localhost bind). ``testserver`` is #: Starlette's TestClient host. DEFAULT_ALLOWED_HOSTS = ["localhost", "127.0.0.1", "[::1]", "testserver"] @@ -93,6 +98,21 @@ def get_meta() -> dict[str, object]: # mesh tells the UI whether DELETE archives (restorable) or removes. return {"mesh": store.mesh_mode()} + @app.get("/api/ai/profile") + async def get_ai_profile() -> dict[str, object]: + return ai_usage.profile_payload(store.omi_dir) + + @app.put("/api/ai/profile") + async def put_ai_profile(payload: AIProfileUpdate) -> dict[str, str]: + try: + return ai_usage.set_profile(store.omi_dir, payload.profile) + except OSError as exc: + raise HTTPException(status_code=500, detail="could not save AI profile") from exc + + @app.get("/api/ai/usage") + async def get_ai_usage(since: Literal["24h", "7d", "30d", "all"] = "7d") -> dict[str, object]: + return ai_usage.usage_summary(store.omi_dir, since=since) + @app.get("/api/notes/{name}") def get_note(name: str) -> dict[str, object]: raw = _guard(lambda: store.read_note(name)) diff --git a/src/omind/web/static/app.css b/src/omind/web/static/app.css index a1e9f43..1bdb819 100644 --- a/src/omind/web/static/app.css +++ b/src/omind/web/static/app.css @@ -177,6 +177,58 @@ body, color: var(--text); } +/* ---- AI usage ---------------------------------------------------------- */ + +.ai-sheet { max-width: 1040px; } +.ai-heading, +.ai-controls, +.ai-measurement { + display: flex; + align-items: center; + justify-content: space-between; + gap: 1rem; +} +.ai-controls label { + display: grid; + gap: 0.35rem; + color: var(--text-faint); + font: 600 10px var(--font-ui); + letter-spacing: 0.1em; + text-transform: uppercase; +} +.ai-stats { + display: grid; + grid-template-columns: repeat(4, minmax(0, 1fr)); + gap: 0.75rem; + margin: 2rem 0 1rem; +} +.ai-stat { + display: grid; + gap: 0.5rem; + padding: 1rem; + border: 1px solid var(--border); + border-radius: var(--radius); + background: var(--surface-2); +} +.ai-stat span, +.ai-measurement, +.ai-table th { + color: var(--text-faint); + font: 600 10px var(--font-ui); + letter-spacing: 0.08em; + text-transform: uppercase; +} +.ai-stat strong { font: 800 1.45rem var(--font-ui); } +.ai-measurement { justify-content: flex-start; margin: 1rem 0; } +.ai-table { width: 100%; border-collapse: collapse; } +.ai-table th, +.ai-table td { padding: 0.75rem; border-bottom: 1px solid var(--border); text-align: start; } +.ai-table td { font-family: "JetBrains Mono", monospace; font-size: 12px; } +@media (max-width: 900px) { + .ai-stats { grid-template-columns: repeat(2, minmax(0, 1fr)); } + .ai-heading { align-items: flex-start; flex-direction: column; } +} + /* ---- Theme swatches ---------------------------------------------------- */ .swatch { diff --git a/src/omind/web/static/app.js b/src/omind/web/static/app.js index c1908ed..5cfc6cf 100644 --- a/src/omind/web/static/app.js +++ b/src/omind/web/static/app.js @@ -9,7 +9,7 @@ const state = { activeTag: null, current: null, // filename currentVersion: "", // mtime+size token of the open note, for conflict detection - mode: "empty", // empty | view | edit | raw | new + mode: "empty", // empty | view | edit | raw | new | ai lang: "en", showArchived: false, // include soft-deleted (Disabled: true) notes in the list mesh: false, // server-reported: DELETE archives (restorable) instead of removing @@ -66,6 +66,12 @@ const I18N = { archivedToggle: "archived", archivedBadge: "archived", restore: "Restore", restoredToast: "Restored.", archivedToast: "Archived.", confirmArchive: 'Archive "{name}"? It stays restorable.', + aiUsage: "AI Usage", expenseProfile: "Model expense", window: "Window", + inputTokens: "Input tokens", outputTokens: "Output tokens", + cacheTokens: "Cache tokens", avoidedTokens: "Estimated avoided", + exactUsage: "Provider-reported", estimatedUsage: "Estimated", + attributionNote: "OMI overhead only · profiles do not select the model", + operation: "Operation", profileSaved: "Profile saved.", }, es: { tagline: "memoria", search: "buscar…", theme: "Tema", language: "Idioma", @@ -93,6 +99,12 @@ const I18N = { archivedToggle: "archivadas", archivedBadge: "archivada", restore: "Restaurar", restoredToast: "Restaurada.", archivedToast: "Archivada.", confirmArchive: "¿Archivar «{name}»? Podrás restaurarla.", + aiUsage: "Uso de IA", expenseProfile: "Coste del modelo", window: "Periodo", + inputTokens: "Tokens de entrada", outputTokens: "Tokens de salida", + cacheTokens: "Tokens en caché", avoidedTokens: "Ahorro estimado", + exactUsage: "Informado por el proveedor", estimatedUsage: "Estimado", + attributionNote: "Solo sobrecarga de OMI · los perfiles no eligen el modelo", + operation: "Operación", profileSaved: "Perfil guardado.", }, fr: { tagline: "mémoire", search: "rechercher…", theme: "Thème", language: "Langue", @@ -121,6 +133,12 @@ const I18N = { archivedToggle: "archivées", archivedBadge: "archivée", restore: "Restaurer", restoredToast: "Restaurée.", archivedToast: "Archivée.", confirmArchive: "Archiver « {name} » ? Restauration possible.", + aiUsage: "Utilisation IA", expenseProfile: "Coût du modèle", window: "Période", + inputTokens: "Jetons d’entrée", outputTokens: "Jetons de sortie", + cacheTokens: "Jetons en cache", avoidedTokens: "Économie estimée", + exactUsage: "Rapporté par le fournisseur", estimatedUsage: "Estimé", + attributionNote: "Surcharge OMI uniquement · le profil ne choisit pas le modèle", + operation: "Opération", profileSaved: "Profil enregistré.", }, ar: { tagline: "ذاكرة", search: "بحث…", theme: "السمة", language: "اللغة", @@ -147,6 +165,12 @@ const I18N = { archivedToggle: "المؤرشفة", archivedBadge: "مؤرشفة", restore: "استعادة", restoredToast: "تمت الاستعادة.", archivedToast: "تمت الأرشفة.", confirmArchive: "أرشفة «{name}»؟ يمكن استعادتها لاحقًا.", + aiUsage: "استخدام الذكاء الاصطناعي", expenseProfile: "تكلفة النموذج", window: "الفترة", + inputTokens: "رموز الإدخال", outputTokens: "رموز الإخراج", + cacheTokens: "رموز التخزين المؤقت", avoidedTokens: "التوفير التقديري", + exactUsage: "حسب تقرير المزوّد", estimatedUsage: "تقديري", + attributionNote: "تكلفة OMI فقط · الملفات لا تختار النموذج", + operation: "العملية", profileSaved: "تم حفظ الملف.", }, ru: { tagline: "память", search: "поиск…", theme: "Тема", language: "Язык", @@ -173,6 +197,12 @@ const I18N = { archivedToggle: "архив", archivedBadge: "в архиве", restore: "Восстановить", restoredToast: "Восстановлено.", archivedToast: "В архиве.", confirmArchive: "Архивировать «{name}»? Её можно будет восстановить.", + aiUsage: "Использование ИИ", expenseProfile: "Стоимость модели", window: "Период", + inputTokens: "Входные токены", outputTokens: "Выходные токены", + cacheTokens: "Токены кэша", avoidedTokens: "Оценка экономии", + exactUsage: "Отчёт провайдера", estimatedUsage: "Оценка", + attributionNote: "Только накладные расходы OMI · профиль не выбирает модель", + operation: "Операция", profileSaved: "Профиль сохранён.", }, zh: { tagline: "记忆", search: "搜索…", theme: "主题", language: "语言", @@ -197,6 +227,12 @@ const I18N = { archivedToggle: "已归档", archivedBadge: "已归档", restore: "恢复", restoredToast: "已恢复。", archivedToast: "已归档。", confirmArchive: "归档“{name}”？之后仍可恢复。", + aiUsage: "AI 用量", expenseProfile: "模型费用", window: "时间范围", + inputTokens: "输入令牌", outputTokens: "输出令牌", + cacheTokens: "缓存令牌", avoidedTokens: "预计节省", + exactUsage: "提供商报告", estimatedUsage: "估算", + attributionNote: "仅 OMI 开销 · 配置不会选择模型", + operation: "操作", profileSaved: "配置已保存。", }, }; @@ -371,6 +407,7 @@ function renderSidebar() { function renderEmpty() { teardownGraph(); + $("#ai-btn").classList.remove("active"); state.mode = "empty"; state.current = null; renderSidebar(); @@ -386,6 +423,7 @@ function renderEmpty() { async function openNote(name) { teardownGraph(); + $("#ai-btn").classList.remove("active"); try { const data = await api("GET", `/api/notes/${encodeURIComponent(name)}`); state.current = name; @@ -597,6 +635,7 @@ function openEdit(data) { function openNew() { teardownGraph(); + $("#ai-btn").classList.remove("active"); state.mode = "new"; state.current = null; renderSidebar(); @@ -719,12 +758,76 @@ searchEl.addEventListener("input", () => { }); $("#new-btn").addEventListener("click", openNew); +// ---- AI usage -------------------------------------------------------------- + +const tokenNumber = (value) => Number(value || 0).toLocaleString(); + +async function openAI(since = "7d") { + teardownGraph(); + state.current = null; + state.mode = "ai"; + $("#ai-btn").classList.add("active"); + $("#graph-btn").classList.remove("active"); + try { history.replaceState(null, "", "#ai"); } catch (_) {} + renderSidebar(); + try { + const usage = await api("GET", `/api/ai/usage?since=${encodeURIComponent(since)}`); + const totals = usage.totals; + const exact = usage.exact; + const estimated = usage.estimated; + const profile = usage.profile; + const rows = Object.entries(usage.operations).map(([name, values]) => ` + ${escapeHtml(name)}${tokenNumber(values.input_tokens)} + ${tokenNumber(values.output_tokens)}${tokenNumber(values.avoided_tokens)}`).join(""); + contentEl.innerHTML = ` +
+
+
${escapeHtml(t("attributionNote"))}
+

${escapeHtml(t("aiUsage"))}

+
+ + +
+
+
+
${escapeHtml(t("inputTokens"))}${tokenNumber(totals.input_tokens)}
+
${escapeHtml(t("outputTokens"))}${tokenNumber(totals.output_tokens)}
+
${escapeHtml(t("cacheTokens"))}${tokenNumber(totals.cache_read_tokens + totals.cache_write_tokens)}
+
${escapeHtml(t("avoidedTokens"))}${tokenNumber(totals.avoided_tokens)}
+
+
+ ${escapeHtml(t("exactUsage"))}: ${tokenNumber(exact.input_tokens + exact.output_tokens)} + ${escapeHtml(t("estimatedUsage"))}: ${tokenNumber(estimated.input_tokens + estimated.output_tokens)} +
+ + + ${rows}
${escapeHtml(t("operation"))}${escapeHtml(t("inputTokens"))}${escapeHtml(t("outputTokens"))}${escapeHtml(t("avoidedTokens"))}
+
`; + $("#ai-window").addEventListener("change", (event) => openAI(event.target.value)); + $("#ai-profile").addEventListener("change", async (event) => { + try { + await api("PUT", "/api/ai/profile", { profile: event.target.value }); + toast(t("profileSaved")); + await openAI(since); + } catch (error) { toast(error.message); } + }); + } catch (error) { + toast(error.message); + } +} +$("#ai-btn").addEventListener("click", () => openAI()); + // ---- Graph view ------------------------------------------------------------- async function openGraph() { if (!window.OmindGraph) return; teardownGraph(); // destroy any previous graph before rendering a new one state.current = null; + $("#ai-btn").classList.remove("active"); state.graphActive = true; try { history.replaceState(null, "", "#graph"); } catch (_) {} renderSidebar(); @@ -883,6 +986,7 @@ function applyStaticI18n() { if (tagline) tagline.textContent = t("tagline"); searchEl.placeholder = t("search"); $("#new-btn").textContent = t("new"); + $("#ai-btn").textContent = t("aiUsage"); const tp = $("#theme-picker"); if (tp) tp.title = t("theme"); const ls = $("#lang-select"); @@ -895,6 +999,7 @@ function applyStaticI18n() { function rerenderForLang() { if (state.mode === "view" && state.current) openNote(state.current); else if (state.mode === "empty") renderEmpty(); + else if (state.mode === "ai") openAI(); else renderSidebar(); } @@ -939,6 +1044,7 @@ initI18n(); .catch(() => {}); await refresh(); if (location.hash === "#graph") openGraph(); + else if (location.hash === "#ai") openAI(); else renderEmpty(); } catch (e) { contentEl.innerHTML = `
${escapeHtml( diff --git a/src/omind/web/static/index.html b/src/omind/web/static/index.html index c289772..df9b8ae 100644 --- a/src/omind/web/static/index.html +++ b/src/omind/web/static/index.html @@ -56,6 +56,7 @@

+
diff --git a/tests/test_ai_usage.py b/tests/test_ai_usage.py new file mode 100644 index 0000000..5f482b1 --- /dev/null +++ b/tests/test_ai_usage.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 Aaron K. Clark + +from __future__ import annotations + +import json +import subprocess +from datetime import datetime, timedelta +from pathlib import Path + +import pytest + +from omind import ai_usage +from omind.cli import main + + +def test_profile_default_saved_and_environment_override( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + omi = tmp_path / "OMI" + assert ai_usage.profile_info(omi) == { + "saved": "low", + "effective": "low", + "source": "default", + } + assert ai_usage.set_profile(omi, "medium")["effective"] == "medium" + monkeypatch.setenv(ai_usage.PROFILE_ENV, "high") + assert ai_usage.profile_info(omi) == { + "saved": "medium", + "effective": "high", + "source": "environment", + } + with pytest.raises(ValueError): + ai_usage.set_profile(omi, "pricey") + + +def test_ledger_is_per_vault_private_and_skips_torn_lines(tmp_path: Path) -> None: + first = tmp_path / "one" / "OMI" + second = tmp_path / "two" / "OMI" + ai_usage.record_priming(first, 9) + assert ai_usage.read_events(second) == [] + path = ai_usage.usage_path(first) + assert path.stat().st_mode & 0o777 == 0o600 + with path.open("ab") as stream: + stream.write(b"{torn\xff\n") + events = ai_usage.read_events(first) + assert len(events) == 1 + assert events[0]["input_tokens"] == 3 + serialized = json.dumps(events) + assert "prompt" not in serialized and "response" not in serialized + + +def test_usage_summary_separates_exact_estimated_and_avoided(tmp_path: Path) -> None: + omi = tmp_path / "OMI" + now = datetime(2026, 7, 16, 20, 0) + ai_usage.log_event( + omi, + "verifier", + input_tokens=10, + output_tokens=2, + now=now - timedelta(hours=1), + ) + ai_usage.log_event( + omi, + "priming", + measurement="estimated", + input_tokens=25, + now=now - timedelta(hours=2), + ) + ai_usage.log_event( + omi, + "checkpoint", + status="skipped", + measurement="estimated", + avoided_tokens=50, + now=now - timedelta(days=2), + ) + day = ai_usage.usage_summary(omi, since="24h", now=now) + assert day["totals"]["input_tokens"] == 35 + assert day["exact"]["input_tokens"] == 10 + assert day["estimated"]["input_tokens"] == 25 + assert day["totals"]["avoided_tokens"] == 0 + all_time = ai_usage.usage_summary(omi, since="all", now=now) + assert all_time["totals"]["avoided_tokens"] == 50 + with pytest.raises(ValueError): + ai_usage.usage_summary(omi, since="weekly") + + +def test_run_claude_records_provider_usage(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + omi = tmp_path / "OMI" + monkeypatch.setattr(ai_usage.shutil, "which", lambda _name: "/usr/bin/claude") + + def fake_run(*_args: object, **_kwargs: object) -> subprocess.CompletedProcess[str]: + payload = { + "result": "RELEVANT", + "model": "claude-test", + "usage": { + "input_tokens": 12, + "output_tokens": 3, + "cache_read_input_tokens": 4, + "cache_creation_input_tokens": 5, + }, + } + return subprocess.CompletedProcess([], 0, json.dumps(payload), "") + + monkeypatch.setattr(ai_usage.subprocess, "run", fake_run) + assert ai_usage.run_claude(omi, "verifier", "secret prompt", timeout=2) == "RELEVANT" + event = ai_usage.read_events(omi)[0] + assert event["measurement"] == "exact" + assert event["input_tokens"] == 12 + assert event["cache_read_tokens"] == 4 + assert "secret prompt" not in json.dumps(event) + + +def test_run_claude_estimates_malformed_json_and_records_profile_skip( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + omi = tmp_path / "OMI" + monkeypatch.setattr(ai_usage.shutil, "which", lambda _name: "/usr/bin/claude") + monkeypatch.setattr( + ai_usage.subprocess, + "run", + lambda *_a, **_k: subprocess.CompletedProcess([], 0, "RELEVANT", ""), + ) + assert ai_usage.run_claude(omi, "verifier", "abcd", timeout=2) == "RELEVANT" + assert ai_usage.read_events(omi)[0]["measurement"] == "estimated" + assert ai_usage.run_claude(omi, "checkpoint", "abcdefgh", timeout=2, allowed=False) is None + skipped = ai_usage.read_events(omi)[-1] + assert skipped["status"] == "skipped" + assert skipped["avoided_tokens"] == 2 + + +def test_cli_profile_and_json_usage(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + vault = tmp_path / "vault" + args = ["--vault", str(vault), "--folder", "OMI"] + assert main(["ai", "profile", "medium", *args]) == 0 + assert "effective=medium" in capsys.readouterr().out + assert main(["ai", "usage", "--since", "all", "--json", *args]) == 0 + payload = json.loads(capsys.readouterr().out) + assert payload["profile"]["effective"] == "medium" + assert payload["totals"]["input_tokens"] == 0 + + +def test_medium_verifier_caps_prompt_and_high_skips( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from omind import verify + + omi = tmp_path / "OMI" + captured: list[tuple[str, bool]] = [] + + def fake_run( + _omi: Path | str, + _operation: str, + prompt: str, + *, + timeout: int, + allowed: bool = True, + ) -> str: + del timeout + captured.append((prompt, allowed)) + return "RELEVANT" + + monkeypatch.setattr(ai_usage, "run_claude", fake_run) + ai_usage.set_profile(omi, "medium") + assert verify._ask_model("t" * 2_000, "m" * 4_000, omi) is True + assert captured[-1][1] is True + assert "t" * 501 not in captured[-1][0] + assert "m" * 1_001 not in captured[-1][0] + ai_usage.set_profile(omi, "high") + assert verify._ask_model("task", "material", omi) is True + assert captured[-1][1] is False diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py index 088e565..35ff9ac 100644 --- a/tests/test_checkpoint.py +++ b/tests/test_checkpoint.py @@ -149,3 +149,21 @@ def test_checkpoint_cli_install_timer(tmp_path: Path, monkeypatch: pytest.Monkey assert rc == 0 timer = (checkpoint.systemd_user_dir() / checkpoint.TIMER_UNIT_NAME).read_text(encoding="utf-8") assert "OnUnitActiveSec=1800s" in timer + + +def test_high_expense_checkpoint_skips_llm_and_records_avoided_tokens(tmp_path: Path) -> None: + from omind import ai_usage + + omi = _omi(tmp_path) + ai_usage.set_profile(omi, "high") + activity = checkpoint.Activity( + actions=[{"time": "12:00", "event": "PostToolUse", "tool": "Bash", "detail": "work"}] + ) + rendered = checkpoint.render_section( + activity, "15m", datetime(2026, 6, 20, 12, 30), llm=True, omi_dir=omi + ) + assert "1 action(s)" in rendered + event = ai_usage.read_events(omi)[-1] + assert event["operation"] == "checkpoint" + assert event["status"] == "skipped" + assert event["avoided_tokens"] > 0 diff --git a/tests/test_hooks.py b/tests/test_hooks.py index fddc466..0db65b2 100644 --- a/tests/test_hooks.py +++ b/tests/test_hooks.py @@ -369,17 +369,49 @@ def test_session_start_total_cap_truncates_dynamic_first(tmp_path: Path) -> None state.write_text("d" * hooks._PRIMING_FILE_CHAR_CAP, encoding="utf-8") ctx = hooks.build_session_start_context(tmp_path) assert len(ctx) <= hooks._TOTAL_CONTEXT_CHAR_CAP - assert ctx.count("STATIC-END") == len(hooks.PRIMING_FILES) # static never cut - assert "dddd" in ctx # session state partially injected … - assert ctx.endswith("…[truncated]") # … and truncated to fit the budget + for name in hooks.PRIMING_FILES: + assert f"===== OMI/{name} =====" in ctx # every static note gets a floor + assert "dddd" in ctx # 20% remains reserved for dynamic state + assert "…[truncated]" in ctx -def test_session_start_drops_dynamic_when_static_fills_cap(tmp_path: Path) -> None: +def test_session_start_reserves_dynamic_when_static_fills_cap(tmp_path: Path) -> None: _write_priming_files(tmp_path, body="s" * hooks._PRIMING_FILE_CHAR_CAP) # ~48k static (tmp_path / "Session State 2026-06-09.md").write_text("DYNAMIC-STATE", encoding="utf-8") ctx = hooks.build_session_start_context(tmp_path) - assert "DYNAMIC-STATE" not in ctx # no room: dynamic dropped, static intact - assert ctx.count("s" * hooks._PRIMING_FILE_CHAR_CAP) == len(hooks.PRIMING_FILES) + assert "DYNAMIC-STATE" in ctx + assert len(ctx) <= hooks._TOTAL_CONTEXT_CHAR_CAP + + +def test_session_start_expense_profiles_apply_hard_caps(tmp_path: Path) -> None: + from omind import ai_usage + + _write_priming_files(tmp_path, body="x" * hooks._PRIMING_FILE_CHAR_CAP) + for name, cap in (("medium", 24_000), ("high", 8_000)): + ai_usage.set_profile(tmp_path, name) + context = hooks.build_session_start_context(tmp_path) + assert len(context) <= cap + for priming in hooks.PRIMING_FILES: + assert f"===== OMI/{priming} =====" in context + hooks.emit_session_start_context(tmp_path, out=io.StringIO()) + event = ai_usage.read_events(tmp_path)[-1] + assert event["profile"] == "high" + assert event["avoided_tokens"] > 0 + + +def test_session_start_records_priming_once_per_emission( + tmp_path: Path, isolated_state: Path +) -> None: + from omind import ai_usage + + (tmp_path / "index.md").write_text("remember", encoding="utf-8") + event = '{"session_id":"token-session"}' + hooks.run_hook("pre_llm_call", tmp_path, stdin=io.StringIO(event), stdout=io.StringIO()) + hooks.run_hook("pre_llm_call", tmp_path, stdin=io.StringIO(event), stdout=io.StringIO()) + events = ai_usage.read_events(tmp_path) + assert len(events) == 1 + assert events[0]["operation"] == "priming" + assert events[0]["measurement"] == "estimated" def test_run_hook_stop_records_turn_line(tmp_path: Path) -> None: diff --git a/tests/test_web.py b/tests/test_web.py index 9d043c9..a1158a1 100644 --- a/tests/test_web.py +++ b/tests/test_web.py @@ -10,7 +10,7 @@ import pytest from fastapi.testclient import TestClient -from omind import paths +from omind import ai_usage, paths from omind.web.app import create_app @@ -31,6 +31,21 @@ def test_list_empty(client: TestClient) -> None: assert client.get("/api/notes").json() == [] +def test_ai_profile_and_usage_api(client: TestClient, omi_dir: Path) -> None: + profile = client.get("/api/ai/profile") + assert profile.status_code == 200 + assert profile.json()["effective"] == "low" + changed = client.put("/api/ai/profile", json={"profile": "high"}) + assert changed.status_code == 200 + assert changed.json()["effective"] == "high" + assert client.put("/api/ai/profile", json={"profile": "unknown"}).status_code == 422 + ai_usage.record_priming(omi_dir, 400) + usage = client.get("/api/ai/usage", params={"since": "all"}) + assert usage.status_code == 200 + assert usage.json()["totals"]["input_tokens"] == 100 + assert client.get("/api/ai/usage", params={"since": "forever"}).status_code == 422 + + def test_foreign_host_header_is_rejected(omi_dir: Path) -> None: """DNS-rebinding defence: a Host not on the allowlist gets 400 (#125).""" with TestClient(create_app(omi_dir)) as c: diff --git a/uv.lock b/uv.lock index 1089b52..c14fac1 100644 --- a/uv.lock +++ b/uv.lock @@ -2364,7 +2364,7 @@ wheels = [ [[package]] name = "omind" -version = "3.8.5" +version = "4.0.0" source = { editable = "." } dependencies = [ { name = "fastapi" }, From 21e941ced35b023976ebf9ac632f67020081e95a Mon Sep 17 00:00:00 2001 From: "Aaron K. Clark" Date: Thu, 16 Jul 2026 22:27:52 -0500 Subject: [PATCH 2/3] fix: bound AI usage window parsing --- CHANGELOG.md | 5 +++++ src/omind/ai_usage.py | 15 +++++++++------ tests/test_ai_usage.py | 2 ++ 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d2e7364..6f3b855 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 avoided tokens in text or JSON. The local web app adds the same profile control and 24-hour/7-day/30-day/all-time usage dashboard through `/api/ai/*` endpoints. +### Security +- Parse AI-usage reporting windows with a bounded direct grammar rather than a + regular expression over public API input, eliminating the polynomial-runtime + pattern identified by CodeQL during the 4.0.0 release review. + ## [3.8.6] - 2026-07-06 ### Changed diff --git a/src/omind/ai_usage.py b/src/omind/ai_usage.py index 74cbd58..f8c05ef 100644 --- a/src/omind/ai_usage.py +++ b/src/omind/ai_usage.py @@ -184,13 +184,16 @@ def parse_window(value: str) -> timedelta | None: clean = (value or "").strip().lower() if clean == "all": return None - import re - - match = re.fullmatch(r"(\d+)([hd])", clean) - if not match: + # This is a tiny public-input grammar; parse it directly instead of using a + # backtracking regex over attacker-controlled text (CodeQL py/redos). + if ( + not 2 <= len(clean) <= 10 + or clean[-1] not in {"h", "d"} + or any(char < "0" or char > "9" for char in clean[:-1]) + ): raise ValueError("--since must be 24h, 7d, 30d, or all") - number = int(match.group(1)) - return timedelta(hours=number) if match.group(2) == "h" else timedelta(days=number) + number = int(clean[:-1]) + return timedelta(hours=number) if clean[-1] == "h" else timedelta(days=number) def usage_summary( diff --git a/tests/test_ai_usage.py b/tests/test_ai_usage.py index 5f482b1..f4aa2ac 100644 --- a/tests/test_ai_usage.py +++ b/tests/test_ai_usage.py @@ -84,6 +84,8 @@ def test_usage_summary_separates_exact_estimated_and_avoided(tmp_path: Path) -> assert all_time["totals"]["avoided_tokens"] == 50 with pytest.raises(ValueError): ai_usage.usage_summary(omi, since="weekly") + with pytest.raises(ValueError): + ai_usage.parse_window("9" * 100_000 + "x") def test_run_claude_records_provider_usage(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: From bca08cbf10cc484e395ba185dcc7dbc3e971a148 Mon Sep 17 00:00:00 2001 From: "Aaron K. Clark" Date: Thu, 16 Jul 2026 22:33:36 -0500 Subject: [PATCH 3/3] fix: make release checks portable on Windows --- CHANGELOG.md | 4 ++++ src/omind/guard.py | 13 ++++++++++++- tests/test_ai_usage.py | 5 ++++- tests/test_provision.py | 11 ++++++----- 4 files changed, 26 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6f3b855..c35449a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 regular expression over public API input, eliminating the polynomial-runtime pattern identified by CodeQL during the 4.0.0 release review. +### Fixed +- Preserve backslashes when resolving `git -C` targets on Windows, and keep + POSIX permission-bit assertions scoped to platforms that implement them. + ## [3.8.6] - 2026-07-06 ### Changed diff --git a/src/omind/guard.py b/src/omind/guard.py index f3facc6..bd1373c 100644 --- a/src/omind/guard.py +++ b/src/omind/guard.py @@ -745,7 +745,18 @@ def _git_dash_c_path(command: str) -> Path | None: parts = _split_simple_commands(command) if not parts: return None - tokens = shlex.split(parts[0]) + # POSIX shlex treats every backslash as an escape and turns an unquoted + # Windows path such as ``C:\\repo`` into ``C:repo``. PowerShell/cmd do + # not use backslashes that way, so retain them on Windows. Non-POSIX + # shlex keeps surrounding quotes; remove only a matching outer pair. + tokens = shlex.split(parts[0], posix=os.name != "nt") + if os.name == "nt": + tokens = [ + token[1:-1] + if len(token) >= 2 and token[0] == token[-1] and token[0] in {'"', "'"} + else token + for token in tokens + ] except ValueError: return None if not tokens or tokens[0] != "git": diff --git a/tests/test_ai_usage.py b/tests/test_ai_usage.py index f4aa2ac..5e96a7c 100644 --- a/tests/test_ai_usage.py +++ b/tests/test_ai_usage.py @@ -4,6 +4,7 @@ from __future__ import annotations import json +import os import subprocess from datetime import datetime, timedelta from pathlib import Path @@ -40,7 +41,9 @@ def test_ledger_is_per_vault_private_and_skips_torn_lines(tmp_path: Path) -> Non ai_usage.record_priming(first, 9) assert ai_usage.read_events(second) == [] path = ai_usage.usage_path(first) - assert path.stat().st_mode & 0o777 == 0o600 + # Windows models file privacy with ACLs rather than POSIX mode bits. + if os.name != "nt": + assert path.stat().st_mode & 0o777 == 0o600 with path.open("ab") as stream: stream.write(b"{torn\xff\n") events = ai_usage.read_events(first) diff --git a/tests/test_provision.py b/tests/test_provision.py index da500c5..1e65ca1 100644 --- a/tests/test_provision.py +++ b/tests/test_provision.py @@ -154,11 +154,12 @@ def test_managed_hook_scripts_are_written_0755_atomically( ): f = hooks / name assert f.exists(), f"{name} was not provisioned" - perms = f.stat().st_mode & 0o777 - assert perms == 0o755, ( - f"{name} is {oct(perms)}; expected 0o755 — a hook must be o+r+x so a " - "chown-root can't render it unreadable to the agent user" - ) + if os.name != "nt": + perms = f.stat().st_mode & 0o777 + assert perms == 0o755, ( + f"{name} is {oct(perms)}; expected 0o755 — a hook must be o+r+x so a " + "chown-root can't render it unreadable to the agent user" + ) def test_default_vault_path_shape() -> None: