Files

461 lines
14 KiB
Python
Executable File

#!/usr/bin/env python3
"""Present one review artifact through Orca or a reusable Herdr mcat pane."""
from __future__ import annotations
import argparse
import base64
import fcntl
import hashlib
import json
import os
import re
import secrets
import shutil
import subprocess
import tempfile
import time
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Mapping, Sequence
PANE_LABEL = "artifact presenter"
VIEWER_PROCESS_NAME = "mcat"
TEXT_SAMPLE_BYTES = 8192
ORCA_SESSION_ENVIRONMENT_KEYS = (
"ORCA_WORKSPACE_ID",
"ORCA_WORKTREE_ID",
"ORCA_PANE_KEY",
"ORCA_TERMINAL_HANDLE",
)
PANE_READY_ATTEMPTS = 40
PANE_READY_DELAY_SECONDS = 0.05
PANE_SHELL_STABLE_CHECKS = 2
VIEWER_START_ATTEMPTS = 40
SAFE_COMMAND_TOKEN = re.compile(r"^[A-Za-z0-9_./=-]+$")
VIEWER_LAUNCH_SCRIPT = """#!/usr/bin/env python3
import base64
import os
import sys
from pathlib import Path
def decode(value: str) -> str:
return base64.urlsafe_b64decode(value.encode("ascii")).decode("utf-8")
acknowledgment, token, viewer, artifact = map(decode, sys.argv[1:])
Path(acknowledgment).write_text(token, encoding="utf-8")
os.execv(viewer, [viewer, artifact])
"""
RunCommand = Callable[..., subprocess.CompletedProcess[str]]
FindExecutable = Callable[[str], str | None]
@dataclass(frozen=True)
class PresentationResult:
status: str
artifact: Path
pane_id: str | None = None
reason: str | None = None
def message(self) -> str:
if self.status == "artifact-only":
return f"artifact-only: {self.reason}; {self.artifact}"
if self.status == "orca-opened":
return f"orca-opened: {self.artifact}"
return f"{self.status}: pane {self.pane_id}; {self.artifact}"
def is_orca_managed_session(env: Mapping[str, str]) -> bool:
"""Return whether this process belongs to an Orca-managed terminal session."""
if env.get("TERM_PROGRAM", "").casefold() == "orca":
return True
return any(env.get(key) for key in ORCA_SESSION_ENVIRONMENT_KEYS)
def is_editor_friendly_text(artifact: Path) -> bool:
"""Recognize UTF-8 text without maintaining a source-extension allowlist."""
with artifact.open("rb") as handle:
sample = handle.read(TEXT_SAMPLE_BYTES)
if b"\0" in sample:
return False
try:
sample.decode("utf-8")
except UnicodeDecodeError:
return False
return True
def orca_executable(
env: Mapping[str, str],
which: FindExecutable,
) -> str | None:
configured = env.get("ORCA_CLI_COMMAND")
if configured:
return which(configured)
if env.get("ORCA_DEV_REPO_ROOT"):
return which("orca-dev")
return which("orca")
def open_in_orca(
orca: str,
artifact: Path,
run: RunCommand,
) -> PresentationResult:
completed = run(
[orca, "file", "open", str(artifact), "--json"],
check=True,
capture_output=True,
text=True,
)
response = json.loads(completed.stdout)
result = response.get("result") if isinstance(response, dict) else None
if (
not isinstance(response, dict)
or response.get("ok") is not True
or not isinstance(result, dict)
or result.get("opened") is not True
):
raise ValueError("Orca did not confirm that it opened the file")
return PresentationResult("orca-opened", artifact)
def run_json(
executable: str,
args: Sequence[str],
run: RunCommand,
) -> dict[str, object]:
completed = run(
[executable, *args],
check=True,
capture_output=True,
text=True,
)
value = json.loads(completed.stdout)
if not isinstance(value, dict):
raise ValueError("Herdr returned a non-object response")
return value
def panes_with_label(response: Mapping[str, object]) -> list[str]:
result = response.get("result")
if not isinstance(result, dict):
return []
panes = result.get("panes")
if not isinstance(panes, list):
return []
pane_ids: list[str] = []
for pane in panes:
if not isinstance(pane, dict) or pane.get("label") != PANE_LABEL:
continue
pane_id = pane.get("pane_id")
if isinstance(pane_id, str) and pane_id:
pane_ids.append(pane_id)
return pane_ids
def created_pane_id(response: Mapping[str, object]) -> str:
result = response.get("result")
pane = result.get("pane") if isinstance(result, dict) else None
pane_id = pane.get("pane_id") if isinstance(pane, dict) else None
if not isinstance(pane_id, str) or not pane_id:
raise ValueError("Herdr split response is missing result.pane.pane_id")
return pane_id
@dataclass(frozen=True)
class PaneProcessState:
shell_pid: int
foreground_pids: frozenset[int]
foreground_names: frozenset[str]
@property
def shell_is_ready(self) -> bool:
return self.shell_pid in self.foreground_pids
@property
def viewer_is_active(self) -> bool:
return VIEWER_PROCESS_NAME in self.foreground_names
def pane_process_state(response: Mapping[str, object]) -> PaneProcessState:
result = response.get("result")
process_info = result.get("process_info") if isinstance(result, dict) else None
if not isinstance(process_info, dict):
raise ValueError("Herdr process response is missing process info")
shell_pid = process_info.get("shell_pid")
processes = process_info.get("foreground_processes")
if not isinstance(shell_pid, int) or not isinstance(processes, list):
raise ValueError("Herdr process response is missing shell state")
pids: set[int] = set()
names: set[str] = set()
for process in processes:
if not isinstance(process, dict):
continue
pid = process.get("pid")
name = process.get("name")
if isinstance(pid, int):
pids.add(pid)
if isinstance(name, str) and name:
names.add(name)
return PaneProcessState(shell_pid, frozenset(pids), frozenset(names))
def prepare_pane_for_command(
herdr: str,
pane_id: str,
run: RunCommand,
*,
allow_startup_transition: bool,
) -> None:
viewer_closed = False
stable_shell_checks = 0
for attempt in range(PANE_READY_ATTEMPTS):
response = run_json(herdr, ["pane", "process-info", "--pane", pane_id], run)
state = pane_process_state(response)
if state.shell_is_ready:
stable_shell_checks += 1
if stable_shell_checks >= PANE_SHELL_STABLE_CHECKS:
return
else:
stable_shell_checks = 0
if state.viewer_is_active:
if not viewer_closed:
run(
[herdr, "pane", "send-keys", pane_id, "q"],
check=True,
capture_output=True,
text=True,
)
viewer_closed = True
elif (
state.foreground_pids
and not state.shell_is_ready
and not viewer_closed
and not allow_startup_transition
):
processes = ", ".join(sorted(state.foreground_names)) or "unknown"
raise ValueError(f"pane is running an unrelated foreground process: {processes}")
if attempt + 1 < PANE_READY_ATTEMPTS:
time.sleep(PANE_READY_DELAY_SECONDS)
raise ValueError("pane shell did not become ready for the artifact command")
def encode_argument(value: str) -> str:
return base64.urlsafe_b64encode(value.encode("utf-8")).decode("ascii")
@contextmanager
def lock_pane(pane_id: str):
runtime_directory = Path(os.environ.get("XDG_RUNTIME_DIR", tempfile.gettempdir()))
lock_directory = runtime_directory / "pi-artifact-presenter-locks"
lock_directory.mkdir(mode=0o700, exist_ok=True)
lock_name = hashlib.sha256(pane_id.encode("utf-8")).hexdigest() + ".lock"
descriptor = os.open(
lock_directory / lock_name,
os.O_CREAT | os.O_RDWR | os.O_CLOEXEC,
0o600,
)
with os.fdopen(descriptor, "w") as handle:
fcntl.flock(handle, fcntl.LOCK_EX)
try:
yield
finally:
fcntl.flock(handle, fcntl.LOCK_UN)
def run_mcat(
herdr: str,
mcat: str,
pane_id: str,
artifact: Path,
run: RunCommand,
*,
reused: bool = False,
) -> None:
with lock_pane(pane_id):
prepare_pane_for_command(
herdr,
pane_id,
run,
allow_startup_transition=not reused,
)
token = secrets.token_hex(12)
with tempfile.TemporaryDirectory(
prefix="pi-artifact-presenter-",
dir="/tmp",
) as acknowledgment_directory:
temporary = Path(acknowledgment_directory)
acknowledgment = temporary / "viewer-started"
launcher = temporary / "launch-viewer.py"
launcher.write_text(VIEWER_LAUNCH_SCRIPT, encoding="utf-8")
command_parts = [
"python3",
str(launcher),
encode_argument(str(acknowledgment)),
encode_argument(token),
encode_argument(mcat),
encode_argument(str(artifact)),
]
if not all(SAFE_COMMAND_TOKEN.fullmatch(part) for part in command_parts):
raise ValueError("viewer launcher produced an unsafe terminal token")
command = " ".join(command_parts)
run(
[herdr, "pane", "send-text", pane_id, command + "\r"],
check=True,
capture_output=True,
text=True,
)
for attempt in range(VIEWER_START_ATTEMPTS):
try:
if acknowledgment.read_text(encoding="utf-8") == token:
return
except FileNotFoundError:
pass
if attempt + 1 < VIEWER_START_ATTEMPTS:
time.sleep(PANE_READY_DELAY_SECONDS)
raise ValueError("artifact command was delivered but did not start")
def present_in_workspace(
artifact: Path,
herdr: str,
mcat: str,
workspace_id: str,
source_pane_id: str,
run: RunCommand,
) -> PresentationResult:
try:
pane_list = run_json(herdr, ["pane", "list", "--workspace", workspace_id], run)
existing_panes = panes_with_label(pane_list)
except (json.JSONDecodeError, OSError, subprocess.SubprocessError, ValueError) as error:
return PresentationResult("artifact-only", artifact, reason=f"could not list Herdr panes: {error}")
for existing_pane in existing_panes:
try:
run_mcat(herdr, mcat, existing_pane, artifact, run, reused=True)
return PresentationResult("reused", artifact, pane_id=existing_pane)
except (OSError, subprocess.SubprocessError, ValueError):
try:
run(
[herdr, "pane", "rename", existing_pane, f"{PANE_LABEL} stale"],
check=False,
capture_output=True,
text=True,
)
except OSError:
pass
try:
split = run_json(
herdr,
[
"pane",
"split",
"--pane",
source_pane_id,
"--direction",
"right",
"--cwd",
str(artifact.parent),
"--no-focus",
],
run,
)
pane_id = created_pane_id(split)
except (json.JSONDecodeError, OSError, subprocess.SubprocessError, ValueError) as error:
return PresentationResult("artifact-only", artifact, reason=f"could not create Herdr viewer: {error}")
try:
run(
[herdr, "pane", "rename", pane_id, PANE_LABEL],
check=True,
capture_output=True,
text=True,
)
run_mcat(herdr, mcat, pane_id, artifact, run)
return PresentationResult("opened", artifact, pane_id=pane_id)
except (OSError, subprocess.SubprocessError, ValueError) as error:
try:
run(
[herdr, "pane", "close", pane_id],
check=False,
capture_output=True,
text=True,
)
except OSError:
pass
return PresentationResult("artifact-only", artifact, reason=f"could not initialize Herdr viewer: {error}")
def present_document(
document: Path,
*,
env: Mapping[str, str] = os.environ,
which: FindExecutable = shutil.which,
run: RunCommand = subprocess.run,
) -> PresentationResult:
artifact = document.expanduser().resolve()
if not artifact.is_file():
raise ValueError(f"artifact is not a file: {artifact}")
if is_orca_managed_session(env) and is_editor_friendly_text(artifact):
orca = orca_executable(env, which)
if orca:
try:
return open_in_orca(orca, artifact, run)
except (json.JSONDecodeError, OSError, subprocess.SubprocessError, ValueError):
pass
herdr = which("herdr")
mcat = which("mcat")
workspace_id = env.get("HERDR_WORKSPACE_ID")
source_pane_id = env.get("HERDR_PANE_ID")
if not herdr:
return PresentationResult("artifact-only", artifact, reason="herdr is unavailable")
if not mcat:
return PresentationResult("artifact-only", artifact, reason="mcat is unavailable")
if not workspace_id or not source_pane_id:
return PresentationResult("artifact-only", artifact, reason="not running in a Herdr pane")
with lock_pane(f"workspace:{workspace_id}"):
return present_in_workspace(
artifact,
herdr,
mcat,
workspace_id,
source_pane_id,
run,
)
def create_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("artifact", type=Path, help="Local file to present with mcat")
return parser
def main() -> int:
args = create_parser().parse_args()
try:
result = present_document(args.artifact)
except ValueError as error:
print(error, file=os.sys.stderr)
return 2
print(result.message())
return 0
if __name__ == "__main__":
raise SystemExit(main())