389 lines
12 KiB
Python
Executable File
389 lines
12 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Open one review artifact in a reusable right-side 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"
|
|
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}"
|
|
return f"{self.status}: pane {self.pane_id}; {self.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}")
|
|
|
|
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())
|