oh-aicoding-tool 0.1.0

package/codex_langfuse_notify.py

@@ -0,0 +1,591 @@
+#!/usr/bin/env python3
+"""
+Codex -> Langfuse notify hook.
+
+Codex calls the configured notify command near the end of a turn. This script
+uses that signal to incrementally read the matching Codex session JSONL file and
+emit the new assistant/user/tool events to Langfuse.
+"""
+
+import json
+import os
+import sys
+import time
+import hashlib
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+try:
+    from langfuse import Langfuse, propagate_attributes
+except Exception:
+    sys.exit(0)
+
+
+CODEX_DIR = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex"))
+STATE_DIR = CODEX_DIR / "langfuse"
+CONFIG_FILE = STATE_DIR / "config.json"
+STATE_FILE = STATE_DIR / "state.json"
+LOCK_FILE = STATE_DIR / "state.lock"
+LOG_FILE = STATE_DIR / "codex_langfuse_notify.log"
+
+DEBUG = os.environ.get("CODEX_LANGFUSE_DEBUG", "").lower() == "true"
+MAX_CHARS = int(os.environ.get("CODEX_LANGFUSE_MAX_CHARS", "20000"))
+
+
+def log(level: str, message: str) -> None:
+    try:
+        STATE_DIR.mkdir(parents=True, exist_ok=True)
+        ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        with open(LOG_FILE, "a", encoding="utf-8") as f:
+            f.write(f"{ts} [{level}] {message}\n")
+    except Exception:
+        pass
+
+
+def debug(message: str) -> None:
+    if DEBUG:
+        log("DEBUG", message)
+
+
+class FileLock:
+    def __init__(self, path: Path, timeout_s: float = 2.0):
+        self.path = path
+        self.timeout_s = timeout_s
+        self._fh = None
+
+    def __enter__(self):
+        STATE_DIR.mkdir(parents=True, exist_ok=True)
+        self._fh = open(self.path, "a+", encoding="utf-8")
+        try:
+            import fcntl
+
+            deadline = time.time() + self.timeout_s
+            while True:
+                try:
+                    fcntl.flock(self._fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+                    break
+                except BlockingIOError:
+                    if time.time() > deadline:
+                        break
+                    time.sleep(0.05)
+        except Exception:
+            pass
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        try:
+            import fcntl
+
+            fcntl.flock(self._fh.fileno(), fcntl.LOCK_UN)
+        except Exception:
+            pass
+        try:
+            self._fh.close()
+        except Exception:
+            pass
+
+
+@dataclass
+class SessionState:
+    offset: int = 0
+    buffer: str = ""
+    turn_count: int = 0
+
+
+def read_json_if_exists(path: Path) -> Dict[str, Any]:
+    try:
+        if not path.exists():
+            return {}
+        text = path.read_text(encoding="utf-8-sig")
+        if not text.strip():
+            return {}
+        value = json.loads(text)
+        return value if isinstance(value, dict) else {}
+    except Exception as e:
+        debug(f"read_json_if_exists failed for {path}: {e}")
+        return {}
+
+
+def write_json_atomic(path: Path, obj: Dict[str, Any]) -> None:
+    try:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        tmp = path.with_suffix(".tmp")
+        tmp.write_text(json.dumps(obj, indent=2, sort_keys=True), encoding="utf-8")
+        os.replace(tmp, path)
+    except Exception as e:
+        debug(f"write_json_atomic failed for {path}: {e}")
+
+
+def read_stdin_payload() -> Dict[str, Any]:
+    try:
+        text = sys.stdin.read()
+        if not text.strip():
+            return {}
+        data = json.loads(text)
+        return data if isinstance(data, dict) else {}
+    except Exception as e:
+        debug(f"stdin payload parse failed: {e}")
+        return {}
+
+
+def first_string(*values: Any) -> Optional[str]:
+    for value in values:
+        if isinstance(value, str) and value.strip():
+            return value.strip()
+    return None
+
+
+def find_value(obj: Any, keys: Tuple[str, ...]) -> Optional[str]:
+    if isinstance(obj, dict):
+        for key, value in obj.items():
+            if key in keys and isinstance(value, str) and value.strip():
+                return value.strip()
+            found = find_value(value, keys)
+            if found:
+                return found
+    elif isinstance(obj, list):
+        for item in obj:
+            found = find_value(item, keys)
+            if found:
+                return found
+    return None
+
+
+def normalize_session_path(raw: Optional[str]) -> Optional[Path]:
+    if not raw:
+        return None
+    try:
+        path = Path(raw).expanduser()
+        if path.exists() and path.is_file():
+            return path.resolve()
+    except Exception:
+        return None
+    return None
+
+
+def latest_session_file() -> Optional[Path]:
+    sessions_dir = CODEX_DIR / "sessions"
+    if not sessions_dir.exists():
+        return None
+    newest: Optional[Path] = None
+    newest_mtime = -1.0
+    try:
+        for path in sessions_dir.rglob("*.jsonl"):
+            try:
+                mtime = path.stat().st_mtime
+            except Exception:
+                continue
+            if mtime > newest_mtime:
+                newest = path
+                newest_mtime = mtime
+    except Exception as e:
+        debug(f"latest_session_file failed: {e}")
+    return newest
+
+
+def resolve_session_file(payload: Dict[str, Any]) -> Optional[Path]:
+    raw_path = first_string(
+        find_value(payload, ("session_path", "sessionPath", "transcript_path", "transcriptPath", "rollout_path", "rolloutPath")),
+        os.environ.get("CODEX_SESSION_PATH"),
+    )
+    session_path = normalize_session_path(raw_path)
+    if session_path:
+        return session_path
+
+    session_id = first_string(
+        find_value(payload, ("session_id", "sessionId", "conversation_id", "conversationId", "id")),
+        os.environ.get("CODEX_SESSION_ID"),
+    )
+    if session_id:
+        sessions_dir = CODEX_DIR / "sessions"
+        try:
+            matches = list(sessions_dir.rglob(f"*{session_id}*.jsonl"))
+            if matches:
+                return max(matches, key=lambda p: p.stat().st_mtime).resolve()
+        except Exception as e:
+            debug(f"session_id lookup failed: {e}")
+
+    return latest_session_file()
+
+
+def session_key(path: Path) -> str:
+    return hashlib.sha256(str(path).encode("utf-8")).hexdigest()
+
+
+def load_session_state(state: Dict[str, Any], key: str) -> SessionState:
+    raw = state.get(key, {})
+    if not isinstance(raw, dict):
+        raw = {}
+    return SessionState(
+        offset=int(raw.get("offset", 0) or 0),
+        buffer=str(raw.get("buffer", "") or ""),
+        turn_count=int(raw.get("turn_count", 0) or 0),
+    )
+
+
+def save_session_state(state: Dict[str, Any], key: str, ss: SessionState) -> None:
+    state[key] = {
+        "offset": ss.offset,
+        "buffer": ss.buffer,
+        "turn_count": ss.turn_count,
+        "updated": datetime.now(timezone.utc).isoformat(),
+    }
+
+
+def read_new_jsonl(path: Path, ss: SessionState) -> Tuple[List[Dict[str, Any]], SessionState]:
+    out: List[Dict[str, Any]] = []
+    try:
+        size = path.stat().st_size
+        if ss.offset > size:
+            ss = SessionState()
+        with open(path, "rb") as f:
+            f.seek(ss.offset)
+            chunk = f.read().decode("utf-8", errors="replace")
+            ss.offset = f.tell()
+    except Exception as e:
+        debug(f"read_new_jsonl failed: {e}")
+        return out, ss
+
+    text = ss.buffer + chunk
+    if not text:
+        return out, ss
+
+    lines = text.splitlines(keepends=True)
+    ss.buffer = ""
+    if lines and not lines[-1].endswith(("\n", "\r")):
+        ss.buffer = lines.pop()
+
+    for line in lines:
+        raw = line.strip()
+        if not raw:
+            continue
+        try:
+            item = json.loads(raw)
+            if isinstance(item, dict):
+                out.append(item)
+        except Exception as e:
+            debug(f"jsonl row parse failed: {e}")
+
+    return out, ss
+
+
+def extract_text(content: Any) -> str:
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts: List[str] = []
+        for item in content:
+            if isinstance(item, str):
+                parts.append(item)
+            elif isinstance(item, dict):
+                text = item.get("text") or item.get("output_text") or item.get("input_text")
+                if isinstance(text, str):
+                    parts.append(text)
+        return "\n".join(parts)
+    if isinstance(content, dict):
+        text = content.get("text") or content.get("message")
+        return text if isinstance(text, str) else ""
+    return ""
+
+
+def truncate(value: Any, max_chars: int = MAX_CHARS) -> Tuple[Any, Dict[str, Any]]:
+    if not isinstance(value, str):
+        try:
+            text = json.dumps(value, ensure_ascii=False)
+        except Exception:
+            text = str(value)
+    else:
+        text = value
+
+    orig_len = len(text)
+    if orig_len <= max_chars:
+        return value if isinstance(value, str) else value, {"truncated": False, "orig_len": orig_len}
+    kept = text[:max_chars]
+    return kept, {
+        "truncated": True,
+        "orig_len": orig_len,
+        "kept_len": len(kept),
+        "sha256": hashlib.sha256(text.encode("utf-8")).hexdigest(),
+    }
+
+
+def get_payload(row: Dict[str, Any]) -> Dict[str, Any]:
+    payload = row.get("payload")
+    return payload if isinstance(payload, dict) else {}
+
+
+def get_session_meta(rows: List[Dict[str, Any]], session_path: Path) -> Dict[str, Any]:
+    meta: Dict[str, Any] = {"session_path": str(session_path)}
+    for row in rows:
+        if row.get("type") == "session_meta":
+            payload = get_payload(row)
+            if payload:
+                meta.update(payload)
+    return meta
+
+
+def extract_usage(rows: List[Dict[str, Any]]) -> Dict[str, Any]:
+    usage: Dict[str, Any] = {}
+    for row in rows:
+        payload = get_payload(row)
+        if row.get("type") == "event_msg" and payload.get("type") == "token_count":
+            info = payload.get("info")
+            if isinstance(info, dict):
+                last = info.get("last_token_usage")
+                total = info.get("total_token_usage")
+                if isinstance(last, dict):
+                    usage["last_token_usage"] = last
+                if isinstance(total, dict):
+                    usage["total_token_usage"] = total
+    return usage
+
+
+def usage_details_from_codex(usage: Dict[str, Any]) -> Dict[str, int]:
+    raw = usage.get("last_token_usage")
+    if not isinstance(raw, dict):
+        return {}
+    out: Dict[str, int] = {}
+    mapping = {
+        "input_tokens": "input",
+        "output_tokens": "output",
+        "cached_input_tokens": "cache_read_input_tokens",
+        "reasoning_output_tokens": "reasoning_output_tokens",
+    }
+    for src, dst in mapping.items():
+        value = raw.get(src)
+        if isinstance(value, int) and value >= 0:
+            out[dst] = value
+    return out
+
+
+def collect_turn_material(rows: List[Dict[str, Any]]) -> Dict[str, Any]:
+    user_texts: List[str] = []
+    assistant_texts: List[str] = []
+    tool_calls: List[Dict[str, Any]] = []
+    tool_results: List[Dict[str, Any]] = []
+
+    for row in rows:
+        row_type = row.get("type")
+        payload = get_payload(row)
+
+        if row_type == "response_item":
+            item_type = payload.get("type")
+            if item_type == "message":
+                role = payload.get("role")
+                text = extract_text(payload.get("content"))
+                if text:
+                    if role == "user":
+                        user_texts.append(text)
+                    elif role == "assistant":
+                        assistant_texts.append(text)
+            elif item_type == "function_call":
+                tool_calls.append(
+                    {
+                        "id": payload.get("call_id") or payload.get("id") or "",
+                        "name": payload.get("name") or "tool",
+                        "input": payload.get("arguments") or payload.get("input") or {},
+                    }
+                )
+
+        if row_type == "event_msg":
+            event_type = payload.get("type")
+            if event_type == "agent_message" and isinstance(payload.get("message"), str):
+                assistant_texts.append(payload["message"])
+            elif isinstance(event_type, str) and event_type.endswith("_end"):
+                tool_results.append(
+                    {
+                        "id": payload.get("call_id") or "",
+                        "name": event_type,
+                        "output": {
+                            "stdout": payload.get("stdout"),
+                            "stderr": payload.get("stderr"),
+                            "success": payload.get("success"),
+                            "status": payload.get("status"),
+                            "aggregated_output": payload.get("aggregated_output"),
+                        },
+                    }
+                )
+
+    return {
+        "user_text": "\n\n".join(user_texts[-3:]),
+        "assistant_text": "\n\n".join(assistant_texts),
+        "tool_calls": tool_calls,
+        "tool_results": tool_results,
+    }
+
+
+def emit_codex_turn(
+    langfuse: Langfuse,
+    session_id: str,
+    user_id: Optional[str],
+    turn_num: int,
+    session_path: Path,
+    meta: Dict[str, Any],
+    material: Dict[str, Any],
+    usage: Dict[str, Any],
+) -> None:
+    user_text, user_meta = truncate(material.get("user_text") or "")
+    assistant_text, assistant_meta = truncate(material.get("assistant_text") or "")
+    usage_details = usage_details_from_codex(usage)
+    model = first_string(meta.get("model"), meta.get("model_provider")) or "codex"
+
+    with propagate_attributes(
+        user_id=user_id,
+        session_id=session_id,
+        trace_name=f"Codex - Turn {turn_num}",
+        tags=["codex"],
+    ):
+        with langfuse.start_as_current_observation(
+            name=f"Codex - Turn {turn_num}",
+            input={"role": "user", "content": user_text},
+            metadata={
+                "source": "codex",
+                "session_id": session_id,
+                "turn_number": turn_num,
+                "session_path": str(session_path),
+                "cwd": meta.get("cwd"),
+                "originator": meta.get("originator"),
+                "cli_version": meta.get("cli_version"),
+                "user_text": user_meta,
+                "usage": usage,
+            },
+        ) as trace_span:
+            with langfuse.start_as_current_observation(
+                name="Codex Response",
+                as_type="generation",
+                model=model,
+                input={"role": "user", "content": user_text},
+                output={"role": "assistant", "content": assistant_text},
+                usage_details=usage_details or None,
+                metadata={"assistant_text": assistant_meta},
+            ):
+                pass
+
+            for call in material.get("tool_calls") or []:
+                tool_input, input_meta = truncate(call.get("input"))
+                with langfuse.start_as_current_observation(
+                    name=f"Tool: {call.get('name') or 'tool'}",
+                    as_type="tool",
+                    input=tool_input,
+                    metadata={
+                        "tool_id": call.get("id"),
+                        "tool_name": call.get("name"),
+                        "input_meta": input_meta,
+                    },
+                ):
+                    pass
+
+            for result in material.get("tool_results") or []:
+                output, output_meta = truncate(result.get("output"))
+                with langfuse.start_as_current_observation(
+                    name=f"Tool Result: {result.get('name') or 'tool'}",
+                    as_type="tool",
+                    metadata={
+                        "tool_id": result.get("id"),
+                        "tool_name": result.get("name"),
+                        "output_meta": output_meta,
+                    },
+                ) as tool_obs:
+                    tool_obs.update(output=output)
+
+            trace_span.update(output={"role": "assistant", "content": assistant_text})
+
+
+def main() -> int:
+    if os.environ.get("TRACE_TO_LANGFUSE", "").lower() == "false":
+        return 0
+
+    payload = read_stdin_payload()
+    config = read_json_if_exists(CONFIG_FILE)
+
+    public_key = (
+        os.environ.get("CODEX_LANGFUSE_PUBLIC_KEY")
+        or os.environ.get("LANGFUSE_PUBLIC_KEY")
+        or config.get("publicKey")
+    )
+    secret_key = (
+        os.environ.get("CODEX_LANGFUSE_SECRET_KEY")
+        or os.environ.get("LANGFUSE_SECRET_KEY")
+        or config.get("secretKey")
+    )
+    host = (
+        os.environ.get("CODEX_LANGFUSE_BASE_URL")
+        or os.environ.get("LANGFUSE_BASEURL")
+        or os.environ.get("LANGFUSE_HOST")
+        or config.get("baseUrl")
+        or "https://cloud.langfuse.com"
+    )
+    user_id = (
+        find_value(payload, ("user_id", "userId", "username", "userName"))
+        or os.environ.get("CODEX_LANGFUSE_USER_ID")
+        or os.environ.get("LANGFUSE_USER_ID")
+        or config.get("userId")
+        or os.environ.get("USERNAME")
+    )
+
+    if not public_key or not secret_key:
+        debug("missing Langfuse credentials")
+        return 0
+
+    session_path = resolve_session_file(payload)
+    if not session_path or not session_path.exists():
+        debug("missing Codex session file")
+        return 0
+
+    try:
+        langfuse = Langfuse(public_key=public_key, secret_key=secret_key, host=host)
+    except Exception as e:
+        debug(f"Langfuse init failed: {e}")
+        return 0
+
+    try:
+        with FileLock(LOCK_FILE):
+            state = read_json_if_exists(STATE_FILE)
+            key = session_key(session_path)
+            ss = load_session_state(state, key)
+            rows, ss = read_new_jsonl(session_path, ss)
+            if not rows:
+                save_session_state(state, key, ss)
+                write_json_atomic(STATE_FILE, state)
+                return 0
+
+            material = collect_turn_material(rows)
+            if not material.get("assistant_text") and not material.get("tool_calls") and not material.get("tool_results"):
+                save_session_state(state, key, ss)
+                write_json_atomic(STATE_FILE, state)
+                return 0
+
+            meta = get_session_meta(rows, session_path)
+            session_id = first_string(str(meta.get("id")) if meta.get("id") else "", session_path.stem) or session_path.stem
+            usage = extract_usage(rows)
+            turn_num = ss.turn_count + 1
+
+            try:
+                emit_codex_turn(langfuse, session_id, user_id, turn_num, session_path, meta, material, usage)
+                ss.turn_count = turn_num
+            except Exception as e:
+                debug(f"emit_codex_turn failed: {e}")
+
+            save_session_state(state, key, ss)
+            write_json_atomic(STATE_FILE, state)
+
+        try:
+            langfuse.flush()
+        except Exception:
+            pass
+        log("INFO", f"Processed Codex turn {ss.turn_count} for {session_path}")
+        return 0
+    except Exception as e:
+        debug(f"unexpected failure: {e}")
+        return 0
+    finally:
+        try:
+            langfuse.shutdown()
+        except Exception:
+            pass
+
+
+if __name__ == "__main__":
+    sys.exit(main())