diff --git a/tests/test_live_core.py b/tests/test_live_core.py index 4dae487..a1711e1 100644 --- a/tests/test_live_core.py +++ b/tests/test_live_core.py @@ -265,7 +265,7 @@ def test_from_a_plain_terminal_te_is_about_the_agent_that_last_worked_here(tmp_p Code state of a repository only Codex (or Pi) had worked in.""" from tokeneconomizer.plugin import cli from tokeneconomizer.plugin.paths import repo_identity, repo_state_dir - for k in ("TE_HOST", "CLAUDECODE", "COPILOT_CLI"): + for k in ("TE_HOST", "CLAUDECODE", "COPILOT_CLI", "OMPCODE"): monkeypatch.delenv(k, raising=False) monkeypatch.setenv("TOKENECONOMIZER_HOME", str(tmp_path / "te")) work = tmp_path / "work" diff --git a/tests/test_pi_history.py b/tests/test_pi_history.py new file mode 100644 index 0000000..29d4786 --- /dev/null +++ b/tests/test_pi_history.py @@ -0,0 +1,126 @@ +"""Issue #166: omp (a Pi fork) session logs -- history import and host detection. + +omp keeps Pi-style sessions under `~/.omp/agent/sessions//_.jsonl`, +with two differences from Pi's own logs: the first line is `{"type": "title", ...}`, and +extra non-conversation entry types (`custom`, `custom_message`, `developer`, +`credential_pin`, `model_change`, ...) appear between the messages. Live sessions fold under +the `pi` host (omp loads the Pi extension), and its shells set `OMPCODE=1` as well as +`CLAUDECODE=1`. +""" +import json +import os + +from tokeneconomizer.plugin import hosts +from tokeneconomizer.plugin.hosts import HOSTS, Pi, default_host +from tokeneconomizer import trace + + +OMP_LOG = [ + # omp writes this first; Pi writes the session line first + {"type": "title", "v": 1, "title": "Fix widget", "source": "auto"}, + {"type": "session", "version": 3, "id": "01AAA-1111-2222-3333", + "timestamp": "2026-11-09T08:00:00.000Z", "cwd": "/home/u/widget", "title": "Fix widget"}, + {"type": "model_change", "model": "some/model"}, # ignored + {"type": "message", "id": "m1", "message": {"role": "user", + "content": "list the files"}}, + {"type": "custom", "payload": {}}, # ignored + {"type": "message", "id": "m2", "parentId": "m1", + "message": {"role": "assistant", "content": [{"type": "text", "text": "ok"}]}}, +] + + +def _omp_log(tmp_path, name="2026-11-09T08-00-00-000Z_01AAA-1111-2222-3333.jsonl"): + p = tmp_path / "sessions" / "-home-u-widget" / name + p.parent.mkdir(parents=True) + with open(p, "w", encoding="utf-8") as fh: + for ev in OMP_LOG: + fh.write(json.dumps(ev) + "\n") + return p + + +def test_the_header_scan_skips_omps_title_line(tmp_path): + p = _omp_log(tmp_path) + head = trace.pi_session_header(str(p)) + assert head.get("id") == "01AAA-1111-2222-3333" + assert head.get("cwd") == "/home/u/widget" + + +def test_the_prompt_count_ignores_omps_extra_entry_types(tmp_path): + p = _omp_log(tmp_path) + assert trace.pi_prompt_count(str(p)) == 1 + + +def test_pi_history_lists_omp_sessions_by_their_id(tmp_path): + p = _omp_log(tmp_path) + pi = Pi() + logs = pi.history(env={"HOME": "/nonexistent"}, extra_dirs=(str(tmp_path),)) + assert len(logs) == 1 + # the session id comes after the last underscore, not the whole file stem + assert logs[0]["session_id"] == "01AAA-1111-2222-3333" + assert logs[0]["path"] == str(p) + + +def test_the_log_head_reads_the_session_line(tmp_path): + p = _omp_log(tmp_path) + head = Pi().log_head(str(p)) + assert head["session_id"] == "01AAA-1111-2222-3333" + assert head["cwd"] == "/home/u/widget" + assert head["started"] is not None # epoch seconds from the timestamp + + +def test_a_resumed_import_still_accepts_omps_file_name(tmp_path): + from tokeneconomizer.plugin.history import _still_eligible + + class _Ctx: + class store: # noqa: N801 - a stub the function reads + @staticmethod + def read(name): + return {} + + @staticmethod + def path(*_a): + return str(tmp_path / "none.json") + + item = {"id": "01AAA-1111-2222-3333", + "path": str(_omp_log(tmp_path))} + assert _still_eligible(_Ctx(), item) is None # not "unsafe_id" + assert _still_eligible(_Ctx(), {"id": "../evil", "path": "x"}) == "unsafe_id" + + +def test_the_converted_log_carries_the_timestamp_and_cwd(tmp_path): + p = _omp_log(tmp_path) + out = tmp_path / "converted.jsonl" + assert trace.pi_session_to_claude(str(p), str(out)) == 0 # no tool calls + lines = [json.loads(l) for l in open(out, encoding="utf-8")] + user = next(l for l in lines if l.get("type") == "user") + assert user.get("timestamp") == "2026-11-09T08:00:00.000Z" + assert user.get("cwd") == "/home/u/widget" + assert trace.claude_log_head(str(out))["started"] is not None + assert trace.claude_log_head(str(out))["cwd"] == "/home/u/widget" + + +def test_an_imported_omp_session_folds_with_its_prompts(tmp_path): + from tokeneconomizer.plugin import history + from tokeneconomizer.plugin.context import LiveContext + repo = tmp_path / "repo" + (repo / ".git").mkdir(parents=True) + env = {"TOKENECONOMIZER_HOME": str(tmp_path / "home"), "TE_FOLD_SYNC": "1", + "TE_MICROSKILLS": "envfacts", "GIT_CEILING_DIRECTORIES": str(tmp_path), + "PATH": os.environ.get("PATH", "")} + ctx = LiveContext.build(Pi(), str(repo), env) + p = _omp_log(tmp_path) + item = {"id": "01AAA-1111-2222-3333", "path": str(p), "started": 1.0, "ended": 2.0} + skipped = {} + history._import_one(ctx, {}, item, skipped, None) + assert skipped == {} + folded = ctx.store.read("folded.json") + assert folded["01AAA-1111-2222-3333"]["imported"] + rec = [r for r in ctx.store.read_lines("ledger.jsonl") if r["kind"] == "session"] + assert len(rec) == 1 and rec[0]["prompts"] == 1 and rec[0]["imported"] + + +def test_omps_shell_detects_as_the_pi_host(): + env = {"OMPCODE": "1", "CLAUDECODE": "1"} # omp sets both; OMPCODE wins + assert default_host(env) == "pi" + assert default_host({"CLAUDECODE": "1"}) == "claude" + assert default_host({}) in HOSTS # the plain-terminal fallback diff --git a/tokeneconomizer/plugin/history.py b/tokeneconomizer/plugin/history.py index 0b37827..3911e70 100644 --- a/tokeneconomizer/plugin/history.py +++ b/tokeneconomizer/plugin/history.py @@ -52,6 +52,7 @@ "over_limit", "attempt": {"id", "n"}, "started", "heartbeat", "finished", "noticed"} """ import os +import tempfile import time import traceback @@ -184,7 +185,7 @@ def plan(ctx, now=None): def skip(why): skipped[why] = skipped.get(why, 0) + 1 for log in ctx.host.history(ctx.env, cfg.history_dirs): - head = trace.claude_log_head(log["path"]) + head = ctx.host.log_head(log["path"]) if not is_this_repository(ctx, log["folder"], head.get("cwd"), cache): continue sid = trace.claude_safe_session_id(log["session_id"]) @@ -202,8 +203,11 @@ def skip(why): elif os.path.exists(ctx.store.path("sessions", f"{sid}.json")): skip("recorded_live") else: - prompts = trace.claude_prompt_count(log["path"]) - if prompts <= 0: + try: + prompts = ctx.host.prompt_count(log["path"]) + except (OSError, ValueError): + prompts = None + if not prompts or prompts <= 0: skip("no_prompt") elif _headless(head, prompts): skip("headless") @@ -219,7 +223,11 @@ def _still_eligible(ctx, item): """Why a planned session is no longer to be imported (a resumed plan), or None.""" sid = trace.claude_safe_session_id(item.get("id")) if isinstance(item, dict) else None path = item.get("path") if isinstance(item, dict) else None - if sid is None or not isinstance(path, str) or os.path.basename(path) != f"{sid}.jsonl": + # Pi's and omp's logs are named `_.jsonl`, Claude Code's + # `.jsonl`; either basename names the session. + base = os.path.basename(path) if isinstance(path, str) else None + named = base is not None and base.rsplit("_", 1)[-1] == f"{sid}.jsonl" + if sid is None or not isinstance(path, str) or base != f"{sid}.jsonl" and not named: return "unsafe_id" if sid in (ctx.store.read("folded.json") or {}): return "already_folded" @@ -250,6 +258,8 @@ def replay_envelope(ctx, transcript_path): level, overrides = cfg.loopguard_aggr, cfg.loopguard_thresholds thr = loopguard.base_thresholds(level, overrides) recorded = 0 + # The mechanism trace already converted a non-Claude log to the envelope this reader + # takes; an imported session folds from the SAME converted copy rather than the raw log. turns = trace.claude_prompt_turns(transcript_path) turn = next(turns, None) while turn is not None: @@ -422,17 +432,22 @@ def _import_one(ctx, doc, item, skipped, log): doc["attempt"] = {"id": sid, "n": tries + 1} _save(ctx, doc) try: - rec = fold.fold_locked(ctx, sid, item["path"], imported=True, - span=(item.get("started"), item.get("ended"))) - if rec is not None and not rec.get("error"): - doc["imported"] = int(_num(doc.get("imported"))) + 1 - try: - replay_envelope(ctx, item["path"]) - except Exception: - _record_error(ctx.root) - _say(log, f"read {sid}") - elif rec is not None: - skipped["unreadable"] = skipped.get("unreadable", 0) + 1 + # A non-Claude log (Pi's, omp's) is folded from the raw path, which the host converts + # itself; only the envelope replay reads the Claude-format copy, so it converts here. + with tempfile.TemporaryDirectory(prefix="te-import-", ignore_cleanup_errors=True) as workdir: + rec = fold.fold_locked(ctx, sid, item["path"], imported=True, + span=(item.get("started"), item.get("ended"))) + if rec is not None and not rec.get("error"): + doc["imported"] = int(_num(doc.get("imported"))) + 1 + try: + converted = (None if ctx.host.name == "claude" + else ctx.host.mechanism_trace(item["path"], workdir)) + replay_envelope(ctx, converted or item["path"]) + except Exception: + _record_error(ctx.root) + _say(log, f"read {sid}") + elif rec is not None: + skipped["unreadable"] = skipped.get("unreadable", 0) + 1 except Exception: _record_error(ctx.root) skipped["error"] = skipped.get("error", 0) + 1 diff --git a/tokeneconomizer/plugin/hosts.py b/tokeneconomizer/plugin/hosts.py index 556d40e..379f5af 100644 --- a/tokeneconomizer/plugin/hosts.py +++ b/tokeneconomizer/plugin/hosts.py @@ -195,6 +195,12 @@ def mechanism_trace(self, transcript_path, workdir): def session_usage(self, transcript_path): raise NotImplementedError + def log_head(self, path): + """What the head of a session log says about the session, for history import: + {"session_id", "cwd", "started", "entrypoint"} -- the Claude Code envelope's four; + a host whose logs are not in that format overrides this.""" + return trace.claude_log_head(path) + def prompt_count(self, transcript_path): """How many prompts the finished session ran; None where the log does not say.""" return None @@ -427,6 +433,58 @@ def mechanism_trace(self, transcript_path, workdir): n = trace.pi_session_to_claude(transcript_path, out) return out if n is not None else None + def log_head(self, path): + """Pi's log is not Claude Code's envelope: its `session` line says the same four.""" + head = trace.pi_session_header(path) + started = head.get("timestamp") + if isinstance(started, str): + started = trace.iso_epoch(started) + else: + started = None + return {"session_id": head.get("id"), "cwd": head.get("cwd"), + "started": started, "entrypoint": ""} + + keeps_history = True + + def history(self, env, extra_dirs=()): + """Every session log under `~/.pi/agent/sessions` and omp's `~/.omp/agent/sessions` + (omp is a Pi fork whose extension this plugin rides), then the `extra_dirs`: each a + `sessions` folder or a folder holding one. Files are named `_.jsonl`; + the id comes after the last `_`. A session filed in two of them is listed once, from + the first.""" + import glob + home = os.path.expanduser("~") + dirs = [os.path.join(home, ".pi", "agent", "sessions"), + os.path.join(home, ".omp", "agent", "sessions")] + for d in extra_dirs or (): + d = os.path.expanduser(str(d)) + dirs.append(os.path.join(d, "sessions") if os.path.isdir(os.path.join(d, "sessions")) + else d) + out, seen, seen_dirs = [], set(), set() + for d in dirs: + real = os.path.realpath(d) + if real in seen_dirs: + continue + seen_dirs.add(real) + for path in sorted(glob.glob(os.path.join(d, "*", "*.jsonl"))): + name = os.path.basename(path) + if not name.endswith(".jsonl"): + continue + sid = name[:-len(".jsonl")] + sid = sid.rsplit("_", 1)[-1] + if sid in seen: + continue + try: + st = os.stat(path) + except OSError: + continue + if not os.path.isfile(path): + continue + seen.add(sid) + out.append({"session_id": sid, "folder": os.path.basename(os.path.dirname(path)), + "path": path, "mtime": st.st_mtime, "size": st.st_size}) + return out + def prompt_count(self, transcript_path): return trace.pi_prompt_count(transcript_path) @@ -1204,6 +1262,10 @@ def default_host(env=None, recent=None): env = os.environ if env is None else env if env.get("TE_HOST") in HOSTS: return env["TE_HOST"] + if env.get("OMPCODE") == "1": + # omp (a Pi fork) marks its shells OMPCODE=1 as well as CLAUDECODE=1, and its + # sessions fold under the Pi host: check it before Claude Code's marker. + return "pi" if env.get("CLAUDECODE") == "1": return "claude" if env.get("COPILOT_CLI") == "1": diff --git a/tokeneconomizer/trace/__init__.py b/tokeneconomizer/trace/__init__.py index b5c380a..0933453 100644 --- a/tokeneconomizer/trace/__init__.py +++ b/tokeneconomizer/trace/__init__.py @@ -20,7 +20,7 @@ claude_log_head, claude_project_folder, claude_projects_dir, claude_prompt_count, claude_prompt_turns, claude_safe_session_id, claude_session_usage, claude_subagent_files, claude_tool_usage, - claude_turn_answered, claude_turn_interrupted) + claude_turn_answered, claude_turn_interrupted, iso_epoch) from .codex import (codex_capability_trace, codex_exec_calls, codex_prompt_count, codex_rollout_to_exec, codex_session_meta, codex_session_usage, rollout_item_to_exec) @@ -45,7 +45,7 @@ "CODEX_DEFAULT_TOOL_CATALOG", "catalog_sources", "parse_catalog", "parse_catalog_with_source", "parse_tool_catalog", "parse_tool_catalog_claude", "parse_tool_catalog_flat", "parse_tool_names_copilot", - "claude_calls", "claude_deferred_tools", "claude_history_logs", "claude_log_head", + "iso_epoch", "claude_calls", "claude_deferred_tools", "claude_history_logs", "claude_log_head", "claude_project_folder", "claude_projects_dir", "claude_prompt_count", "claude_prompt_turns", "claude_safe_session_id", "claude_session_usage", "claude_subagent_files", "claude_tool_usage", "claude_turn_answered", "claude_turn_interrupted", diff --git a/tokeneconomizer/trace/claude.py b/tokeneconomizer/trace/claude.py index d6c6dcd..e08f6d4 100644 --- a/tokeneconomizer/trace/claude.py +++ b/tokeneconomizer/trace/claude.py @@ -156,7 +156,7 @@ def claude_history_logs(projects_dir): return out -def _epoch(ts): +def iso_epoch(ts): if not isinstance(ts, str) or not ts: return None try: @@ -198,7 +198,7 @@ def claude_log_head(path): if out["cwd"] is None and isinstance(ev.get("cwd"), str) and ev["cwd"]: out["cwd"] = ev["cwd"] if out["started"] is None: - out["started"] = _epoch(ev.get("timestamp")) + out["started"] = iso_epoch(ev.get("timestamp")) if not out["entrypoint"] and isinstance(ev.get("entrypoint"), str): out["entrypoint"] = ev["entrypoint"].strip() if out["session_id"] and out["cwd"] and out["started"] and out["entrypoint"]: diff --git a/tokeneconomizer/trace/pi.py b/tokeneconomizer/trace/pi.py index 0eb182e..2a23d21 100644 --- a/tokeneconomizer/trace/pi.py +++ b/tokeneconomizer/trace/pi.py @@ -21,8 +21,11 @@ def _pi_entries(path): def pi_session_header(path): - """The `session` header line ({id, cwd, version, ...}), {} if absent.""" - for ev in _pi_entries(path)[:1]: + """The `session` header line ({id, cwd, version, ...}), {} if absent. + + omp (a Pi fork) starts its log with a `{"type":"title", ...}` line before the `session` + one, so the first few entries are scanned, not just the first.""" + for ev in _pi_entries(path)[:8]: if ev.get("type") == "session": return ev return {} @@ -78,8 +81,9 @@ def pi_session_to_claude(path, out_path): return None n = 0 entries = _pi_entries(path) - head = entries[0] if entries and entries[0].get("type") == "session" else {} + head = pi_session_header(path) cwd = head.get("cwd") if isinstance(head.get("cwd"), str) else None + started = head.get("timestamp") if isinstance(head.get("timestamp"), str) else None with open(out_path, "w", encoding="utf-8") as out: for e in pi_active_branch(entries): if e.get("type") != "message": @@ -87,8 +91,13 @@ def pi_session_to_claude(path, out_path): m = e.get("message") or {} role = m.get("role") if role == "user": - out.write(json.dumps({"type": "user", "message": { - "role": "user", "content": _pi_text(m.get("content"))}}) + "\n") + user = {"role": "user", "content": _pi_text(m.get("content"))} + line = {"type": "user", "message": user} + if started: + line["timestamp"] = started + if cwd: + line["cwd"] = cwd + out.write(json.dumps(line) + "\n") elif role == "assistant": blocks = [] for b in m.get("content") or []: