codex_notifier.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import signal
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any


STATE_VERSION = 1
DEFAULT_POLL_INTERVAL = 1.0
DEFAULT_DISCOVER_INTERVAL = 5.0
APPROVAL_GRACE_SECONDS = 2.5
NORMAL_SOUND_ID = "message-new-instant"
CRITICAL_SOUND_ID = "dialog-warning"
DEFAULT_NORMAL_SOUND_FILE = Path("/usr/share/sounds/sound-icons/prompt.wav")
DEFAULT_CRITICAL_SOUND_FILE = Path("/usr/share/sounds/sound-icons/prompt.wav")


@dataclass(frozen=True)
class SoundProfile:
    normal_loops: int
    critical_loops: int
    normal_volume: float
    critical_volume: float
    normal_volume_db: float
    critical_volume_db: float
    normal_gap_seconds: float
    critical_gap_seconds: float


SOUND_PROFILES: dict[str, SoundProfile] = {
    "quiet": SoundProfile(
        normal_loops=1,
        critical_loops=1,
        normal_volume=0.25,
        critical_volume=0.4,
        normal_volume_db=-2.0,
        critical_volume_db=0.0,
        normal_gap_seconds=0.0,
        critical_gap_seconds=0.0,
    ),
    "balanced": SoundProfile(
        normal_loops=1,
        critical_loops=2,
        normal_volume=0.45,
        critical_volume=0.7,
        normal_volume_db=0.0,
        critical_volume_db=3.0,
        normal_gap_seconds=0.0,
        critical_gap_seconds=0.16,
    ),
    "loud": SoundProfile(
        normal_loops=2,
        critical_loops=3,
        normal_volume=0.7,
        critical_volume=1.0,
        normal_volume_db=2.0,
        critical_volume_db=5.0,
        normal_gap_seconds=0.1,
        critical_gap_seconds=0.16,
    ),
}


@dataclass
class SessionFileState:
    path: Path
    offset: int
    inode: int
    remainder: str = ""
    session_id: str | None = None
    cwd: str | None = None


@dataclass
class PendingApproval:
    call_id: str
    session_label: str
    justification: str
    command: str
    created_at: float


def expand_path(raw_path: str) -> Path:
    return Path(raw_path).expanduser().resolve()


def parse_json(text: str) -> Any | None:
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return None


def text_preview(text: str, limit: int = 260) -> str:
    lines = [line.strip() for line in text.replace("\r", "\n").splitlines() if line.strip()]
    preview = "\n".join(lines[:4]).strip()
    if len(preview) <= limit:
        return preview
    return preview[: limit - 3].rstrip() + "..."


def content_text(payload: dict[str, Any]) -> str:
    parts: list[str] = []
    for item in payload.get("content", []):
        text = item.get("text")
        if isinstance(text, str) and text.strip():
            parts.append(text.strip())
    return "\n".join(parts).strip()


def short_workspace(cwd: str | None, fallback: str) -> str:
    if not cwd:
        return fallback
    path = Path(cwd)
    parts = [part for part in path.parts if part not in ("/", "")]
    if not parts:
        return cwd
    return "/".join(parts[-2:]) if len(parts) >= 2 else parts[-1]


class CodexNotifier:
    def __init__(self, args: argparse.Namespace) -> None:
        self.watch_path = expand_path(args.watch)
        self.poll_interval = args.poll_interval
        self.discover_interval = args.discover_interval
        self.read_existing = args.read_existing
        self.notify_final_answers = not args.no_final_answer
        self.dry_run = args.dry_run
        self.enable_sound = not args.no_sound
        self.sound_profile_name = args.sound_profile
        self.sound_profile = SOUND_PROFILES[self.sound_profile_name]
        self.once = args.once
        self.state_file = expand_path(args.state_file)
        self.running = True
        self.last_discovery = 0.0
        self.last_state_save = 0.0
        self.pending_approvals: dict[str, PendingApproval] = {}
        self.files: dict[str, SessionFileState] = {}
        self.state_enabled = True
        self.saved_state = self.load_state()
        self.script_root = Path(__file__).resolve().parent
        default_icon = self.script_root / "assets" / "codex-notifier.svg"
        self.icon_path = expand_path(args.icon) if args.icon else default_icon
        self.pw_play_path = self.find_command("pw-play")
        self.aplay_path = self.find_command("aplay")
        self.canberra_path = self.find_command("canberra-gtk-play")
        self.notify_send_path = self.find_command("notify-send")
        self.normal_sound_file = self.pick_sound_file(DEFAULT_NORMAL_SOUND_FILE)
        self.critical_sound_file = self.pick_sound_file(DEFAULT_CRITICAL_SOUND_FILE)
        if not self.notify_send_path and not self.dry_run:
            raise RuntimeError("notify-send is required but was not found in PATH.")

    @staticmethod
    def find_command(name: str) -> str | None:
        result = subprocess.run(
            ["bash", "-lc", f"command -v {name}"],
            check=False,
            capture_output=True,
            text=True,
        )
        output = result.stdout.strip()
        return output or None

    def log(self, message: str) -> None:
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        print(f"[{timestamp}] {message}", flush=True)

    @staticmethod
    def pick_sound_file(preferred: Path) -> Path | None:
        if preferred.exists():
            return preferred
        return None

    def load_state(self) -> dict[str, Any]:
        try:
            if not self.state_file.exists():
                return {"version": STATE_VERSION, "files": {}}
            raw_text = self.state_file.read_text(encoding="utf-8")
        except OSError as exc:
            self.state_enabled = False
            self.log(f"state persistence disabled: {exc}")
            return {"version": STATE_VERSION, "files": {}}

        data = parse_json(raw_text)
        if not isinstance(data, dict):
            return {"version": STATE_VERSION, "files": {}}
        if data.get("version") != STATE_VERSION:
            return {"version": STATE_VERSION, "files": {}}
        files = data.get("files")
        if not isinstance(files, dict):
            return {"version": STATE_VERSION, "files": {}}
        return {"version": STATE_VERSION, "files": files}

    def save_state(self, force: bool = False) -> None:
        if not self.state_enabled:
            return
        now = time.monotonic()
        if not force and now - self.last_state_save < 5:
            return
        payload = {
            "version": STATE_VERSION,
            "files": {
                path: {
                    "offset": state.offset,
                    "inode": state.inode,
                    "session_id": state.session_id,
                    "cwd": state.cwd,
                }
                for path, state in self.files.items()
            },
        }
        try:
            self.state_file.parent.mkdir(parents=True, exist_ok=True)
            self.state_file.write_text(json.dumps(payload, indent=2), encoding="utf-8")
            self.last_state_save = now
        except OSError as exc:
            self.state_enabled = False
            self.log(f"state persistence disabled: {exc}")

    def discover_files(self, force: bool = False) -> None:
        now = time.monotonic()
        if not force and now - self.last_discovery < self.discover_interval:
            return
        self.last_discovery = now

        if self.watch_path.is_file():
            candidates = [self.watch_path]
        elif self.watch_path.is_dir():
            candidates = sorted(path for path in self.watch_path.rglob("*.jsonl") if path.is_file())
        else:
            self.log(f"watch path does not exist yet: {self.watch_path}")
            return

        known_paths = set(self.files)
        for path in candidates:
            path_key = str(path)
            if path_key in known_paths:
                continue
            stat = path.stat()
            saved = self.saved_state.get("files", {}).get(path_key, {})
            start_offset = stat.st_size
            session_id = None
            cwd = None
            if isinstance(saved, dict) and saved.get("inode") == stat.st_ino:
                saved_offset = saved.get("offset")
                if isinstance(saved_offset, int) and 0 <= saved_offset <= stat.st_size:
                    start_offset = saved_offset
                session_id = saved.get("session_id")
                cwd = saved.get("cwd")
            elif self.read_existing:
                start_offset = 0

            state = SessionFileState(
                path=path,
                offset=start_offset,
                inode=stat.st_ino,
                session_id=session_id if isinstance(session_id, str) else None,
                cwd=cwd if isinstance(cwd, str) else None,
            )
            self.files[path_key] = state
            if not state.cwd or not state.session_id:
                self.prime_session_metadata(state)

    def prime_session_metadata(self, state: SessionFileState) -> None:
        try:
            with state.path.open("r", encoding="utf-8", errors="replace") as handle:
                for _ in range(12):
                    line = handle.readline()
                    if not line:
                        break
                    event = parse_json(line)
                    if not isinstance(event, dict):
                        continue
                    self.update_metadata(event, state)
                    if state.cwd and state.session_id:
                        return
        except OSError:
            return

    def read_updates(self, state: SessionFileState) -> None:
        try:
            stat = state.path.stat()
        except FileNotFoundError:
            return

        if stat.st_ino != state.inode:
            state.inode = stat.st_ino
            state.offset = 0 if self.read_existing else stat.st_size
            state.remainder = ""
            self.prime_session_metadata(state)
            return

        if stat.st_size < state.offset:
            state.offset = 0 if self.read_existing else stat.st_size
            state.remainder = ""

        if stat.st_size == state.offset:
            return

        try:
            with state.path.open("r", encoding="utf-8", errors="replace") as handle:
                handle.seek(state.offset)
                chunk = handle.read()
                state.offset = handle.tell()
        except OSError:
            return

        if not chunk:
            return

        text = state.remainder + chunk
        lines = text.splitlines(keepends=True)
        if lines and not lines[-1].endswith("\n"):
            state.remainder = lines.pop()
        else:
            state.remainder = ""

        for line in lines:
            self.handle_line(line.rstrip("\n"), state)

    def handle_line(self, line: str, state: SessionFileState) -> None:
        event = parse_json(line)
        if not isinstance(event, dict):
            return
        self.update_metadata(event, state)
        if event.get("type") != "response_item":
            return

        payload = event.get("payload")
        if not isinstance(payload, dict):
            return

        payload_type = payload.get("type")
        if payload_type == "message":
            self.handle_message_event(payload, state)
            return
        if payload_type == "function_call":
            self.handle_function_call(payload, state)
            return
        if payload_type == "function_call_output":
            self.handle_function_output(payload)

    def update_metadata(self, event: dict[str, Any], state: SessionFileState) -> None:
        event_type = event.get("type")
        payload = event.get("payload")
        if not isinstance(payload, dict):
            return

        if event_type == "session_meta":
            session_id = payload.get("id")
            cwd = payload.get("cwd")
            if isinstance(session_id, str):
                state.session_id = session_id
            if isinstance(cwd, str):
                state.cwd = cwd
            return

        if event_type == "turn_context":
            session_id = payload.get("thread_id") or payload.get("turn_id")
            cwd = payload.get("cwd")
            if isinstance(session_id, str) and not state.session_id:
                state.session_id = session_id
            if isinstance(cwd, str):
                state.cwd = cwd

    def handle_message_event(self, payload: dict[str, Any], state: SessionFileState) -> None:
        if payload.get("role") != "assistant":
            return
        if payload.get("phase") != "final_answer":
            return
        if not self.notify_final_answers:
            return
        text = content_text(payload)
        if not text:
            return
        workspace = short_workspace(state.cwd, state.path.stem)
        summary = f"Codex is waiting: {workspace}"
        self.send_notification(summary, text_preview(text), urgency="normal")

    def handle_function_call(self, payload: dict[str, Any], state: SessionFileState) -> None:
        name = payload.get("name")
        call_id = payload.get("call_id")
        arguments = payload.get("arguments")
        if not isinstance(name, str) or not isinstance(call_id, str) or not isinstance(arguments, str):
            return
        args = parse_json(arguments)
        if not isinstance(args, dict):
            return

        workspace = short_workspace(state.cwd or args.get("workdir"), state.path.stem)
        if name == "request_user_input":
            self.notify_request_for_input(workspace, args)
            return

        if name not in {"exec_command", "shell_command"}:
            return

        needs_approval = args.get("sandbox_permissions") == "require_escalated"
        if not needs_approval and args.get("with_escalated_permissions") is True:
            needs_approval = True
        if not needs_approval:
            return

        justification = args.get("justification")
        command = args.get("cmd") or args.get("command")
        pending = PendingApproval(
            call_id=call_id,
            session_label=workspace,
            justification=justification if isinstance(justification, str) else "Codex asked for approval.",
            command=command if isinstance(command, str) else "",
            created_at=time.monotonic(),
        )
        self.pending_approvals[call_id] = pending

    def handle_function_output(self, payload: dict[str, Any]) -> None:
        call_id = payload.get("call_id")
        if isinstance(call_id, str):
            self.pending_approvals.pop(call_id, None)

    def notify_request_for_input(self, workspace: str, args: dict[str, Any]) -> None:
        questions = args.get("questions")
        if not isinstance(questions, list) or not questions:
            body = "Codex requested more input."
        else:
            prompts: list[str] = []
            for question in questions:
                if not isinstance(question, dict):
                    continue
                prompt = question.get("question")
                if isinstance(prompt, str) and prompt.strip():
                    prompts.append(prompt.strip())
            body = "\n".join(prompts[:3]) or "Codex requested more input."
        summary = f"Codex needs input: {workspace}"
        self.send_notification(summary, text_preview(body), urgency="critical")

    def flush_pending_approvals(self, force: bool = False) -> None:
        now = time.monotonic()
        ready_call_ids: list[str] = []
        for call_id, pending in self.pending_approvals.items():
            if not force and now - pending.created_at < APPROVAL_GRACE_SECONDS:
                continue
            ready_call_ids.append(call_id)

        for call_id in ready_call_ids:
            pending = self.pending_approvals.pop(call_id, None)
            if not pending:
                continue
            body = pending.justification
            if pending.command:
                body = f"{body}\n\nCommand: {text_preview(pending.command, limit=180)}"
            summary = f"Codex needs approval: {pending.session_label}"
            self.send_notification(summary, text_preview(body, limit=320), urgency="critical")

    def send_notification(self, summary: str, body: str, urgency: str) -> None:
        if self.dry_run:
            self.log(f"[dry-run] {urgency.upper()} {summary}")
            if body:
                self.log(f"[dry-run] {body}")
            return

        command = [
            self.notify_send_path,
            "-a",
            "Codex Notifier",
            "-u",
            urgency,
            "-i",
            str(self.icon_path),
            "-t",
            "0" if urgency == "critical" else "10000",
            summary,
        ]
        if body:
            command.append(body)
        subprocess.run(command, check=False)

        if self.enable_sound and (self.pw_play_path or self.aplay_path or self.canberra_path):
            self.play_sound(urgency)

    def play_sound(self, urgency: str) -> None:
        if urgency == "critical":
            sound_id = CRITICAL_SOUND_ID
            loop_count = self.sound_profile.critical_loops
            volume = self.sound_profile.critical_volume
            volume_db = self.sound_profile.critical_volume_db
            gap_seconds = self.sound_profile.critical_gap_seconds
            sound_file = self.critical_sound_file
        else:
            sound_id = NORMAL_SOUND_ID
            loop_count = self.sound_profile.normal_loops
            volume = self.sound_profile.normal_volume
            volume_db = self.sound_profile.normal_volume_db
            gap_seconds = self.sound_profile.normal_gap_seconds
            sound_file = self.normal_sound_file

        if sound_file:
            if self.pw_play_path:
                self.spawn_audio(
                    [
                        self.pw_play_path,
                        "--volume",
                        str(volume),
                        str(sound_file),
                    ]
                )
                if loop_count > 1:
                    for _ in range(loop_count - 1):
                        time.sleep(gap_seconds)
                        self.spawn_audio(
                            [
                                self.pw_play_path,
                                "--volume",
                                str(volume),
                                str(sound_file),
                            ]
                        )
                return

            if self.aplay_path:
                self.spawn_audio([self.aplay_path, "-q", str(sound_file)])
                if loop_count > 1:
                    for _ in range(loop_count - 1):
                        time.sleep(gap_seconds)
                        self.spawn_audio([self.aplay_path, "-q", str(sound_file)])
                return

        if self.canberra_path:
            self.spawn_audio(
                [
                    self.canberra_path,
                    "-i",
                    sound_id,
                    "-d",
                    "Codex Notifier",
                    "-l",
                    str(loop_count),
                    "-V",
                    str(volume_db),
                ]
            )

    def spawn_audio(self, command: list[str]) -> None:
        subprocess.Popen(
            command,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            start_new_session=True,
        )

    def send_test_notifications(self) -> None:
        self.send_notification(
            "Codex is waiting: sample/repo",
            "This is a sample completion notification from Codex Notifier.",
            urgency="normal",
        )
        self.send_notification(
            "Codex needs approval: sample/repo",
            "This is a sample approval notification from Codex Notifier.",
            urgency="critical",
        )

    def run(self) -> int:
        self.discover_files(force=True)
        if self.once:
            for state in list(self.files.values()):
                self.read_updates(state)
            self.flush_pending_approvals(force=True)
            self.save_state(force=True)
            return 0

        self.log(f"watching {self.watch_path}")
        while self.running:
            self.discover_files()
            for state in list(self.files.values()):
                self.read_updates(state)
            self.flush_pending_approvals()
            self.save_state()
            time.sleep(self.poll_interval)
        self.flush_pending_approvals(force=True)
        self.save_state(force=True)
        return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Watch Codex session files and raise Ubuntu desktop notifications."
    )
    parser.add_argument(
        "--watch",
        default="~/.codex/sessions",
        help="Codex session directory or a single session JSONL file to watch.",
    )
    parser.add_argument(
        "--state-file",
        default="~/.cache/codex-notifier/state.json",
        help="State file used to resume offsets across restarts.",
    )
    parser.add_argument(
        "--poll-interval",
        type=float,
        default=DEFAULT_POLL_INTERVAL,
        help="Seconds between reads of known files.",
    )
    parser.add_argument(
        "--discover-interval",
        type=float,
        default=DEFAULT_DISCOVER_INTERVAL,
        help="Seconds between scans for new session files.",
    )
    parser.add_argument(
        "--read-existing",
        action="store_true",
        help="Start new watches from byte 0 instead of the end of the file.",
    )
    parser.add_argument(
        "--once",
        action="store_true",
        help="Scan current files once and exit.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print notifications instead of sending them.",
    )
    parser.add_argument(
        "--no-sound",
        action="store_true",
        help="Disable notification sounds.",
    )
    parser.add_argument(
        "--sound-profile",
        choices=sorted(SOUND_PROFILES),
        default="balanced",
        help="Notification sound profile to use when sound is enabled.",
    )
    parser.add_argument(
        "--no-final-answer",
        action="store_true",
        help="Do not notify on assistant final answers; only approvals/input requests.",
    )
    parser.add_argument(
        "--icon",
        help="Optional icon path to use for desktop notifications.",
    )
    parser.add_argument(
        "--test-notification",
        action="store_true",
        help="Send sample notifications and exit.",
    )
    return parser


def install_signal_handlers(notifier: CodexNotifier) -> None:
    def stop(_signum: int, _frame: Any) -> None:
        notifier.running = False

    signal.signal(signal.SIGINT, stop)
    signal.signal(signal.SIGTERM, stop)


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    try:
        notifier = CodexNotifier(args)
    except RuntimeError as exc:
        print(str(exc), file=sys.stderr)
        return 1

    install_signal_handlers(notifier)
    if args.test_notification:
        notifier.send_test_notifications()
        return 0
    return notifier.run()


if __name__ == "__main__":
    raise SystemExit(main())