"""Async-by-default shell execution: bash() spawns immediately and returns a live handle."""
from __future__ import annotations
import asyncio
import atexit
import json
import os
import re
import secrets
import selectors
import shutil
import signal
import socket
import struct
import subprocess
import sys
import threading
import time
import uuid
from collections import deque
from collections.abc import Callable, Generator
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, cast
from . import _winjob
_IS_POSIX = os.name == "posix"
if _IS_POSIX:
import fcntl
import termios
_HEAD_CAP = 512 * 1024
_TAIL_CAP = 3 * 512 * 1024
_READ_CHUNK = 65536
# Fixed child-side fd for the status channel; POSIX shells (notably dash) only
# guarantee single-digit fds in redirection syntax.
_STATUS_FD = 9
_OUTPUT_FD = 8
_COMPLETION_PREFIX = b"\x1eprime-agent-complete:"
_COMPLETION_SUFFIX = b"\x1f"
# Cancelled one-shot awaits: TERM grace before the group KILL, then the bounded
# wait for a confirmed group exit before CancelledError propagates.
_CANCEL_TERM_GRACE = 0.5
_CANCEL_KILL_WAIT = 2.0
# A `cd` prefix in a bash() command is always a mistake: each call is its own
# process, so cd there cannot move the persistent kernel cwd. Catch the leading
# form and the chained form (`x && cd y`) and point at os.chdir / repo.cd.
_CD_PREFIX_RE = re.compile(r"^\s*cd(?:\s|$)")
_CD_CHAIN_RE = re.compile(r"(?:&&|\|\||;|\n)\s*cd(?:\s|$)")
_live_handles: set["BashHandle"] = set()
_live_lock = threading.Lock()
_hook_installed = False
_hook_lock = threading.Lock()
BASH_HANDLE_EVENT_MIME = "application/vnd.prime-agent.bash-handle+json"
_completion_notifier: Callable[[dict[str, Any]], None] | None = None
def _set_completion_notifier(notifier: Callable[[dict[str, Any]], None] | None) -> None:
global _completion_notifier
_completion_notifier = notifier
def _notify_completion(event: dict[str, Any]) -> None:
notifier = _completion_notifier
if notifier is None:
return
try:
notifier(event)
except Exception:
pass # A broken observer must not prevent BashResult callbacks from settling.
@dataclass(frozen=True)
class BashResult:
exit_code: int
output: str
duration: float
class _BoundedBuffer:
"""First _HEAD_CAP bytes plus a rolling _TAIL_CAP-byte tail; the middle is dropped."""
def __init__(self) -> None:
self._head = bytearray()
self._tail: deque[bytes] = deque()
self._tail_size = 0
self._dropped = 0
self._lock = threading.Lock()
def write(self, chunk: bytes) -> None:
with self._lock:
if len(self._head) < _HEAD_CAP:
take = _HEAD_CAP - len(self._head)
self._head.extend(chunk[:take])
chunk = chunk[take:]
if not chunk:
return
self._tail.append(chunk)
self._tail_size += len(chunk)
# Trim the oldest chunk instead of dropping it whole so exactly _TAIL_CAP bytes stay.
while self._tail_size > _TAIL_CAP:
excess = self._tail_size - _TAIL_CAP
oldest = self._tail[0]
if len(oldest) <= excess:
self._tail.popleft()
self._tail_size -= len(oldest)
self._dropped += len(oldest)
else:
self._tail[0] = oldest[excess:]
self._tail_size -= excess
self._dropped += excess
def size(self) -> int:
with self._lock:
return len(self._head) + self._tail_size
def text(self) -> str:
with self._lock:
head = bytes(self._head)
tail = b"".join(self._tail)
dropped = self._dropped
if not dropped:
return (head + tail).decode("utf-8", errors="replace")
marker = f"\n... [{dropped} bytes dropped] ...\n"
return head.decode("utf-8", errors="replace") + marker + tail.decode("utf-8", errors="replace")
class BashHandle:
"""Live handle to a shell command; await it for the BashResult.
A handle awaited before any other API use (the `await bash(cmd)` one-shot
form, including `h = bash(cmd)` awaited immediately) owns the command:
cancelling that await kills the process group. Touching .pid/.running/
.output()/.tail()/.poll()/.kill() first marks the handle as a background
handle; later awaits only wait and cancelling them leaves it running.
"""
def __init__(self, command: str) -> None:
self.command = command
self._buffer = _BoundedBuffer()
self._done = threading.Event()
self._eof = threading.Event()
self._completion_terminal = threading.Event()
self._completion_output: str | None = None
self._completion_lock = threading.Lock()
self._completion_pending = b""
self._status: int | None = None
self._status_known = threading.Event()
self._reaped = False
self._result: BashResult | None = None
self._callbacks: list[Callable[[], None]] = []
self._callback_lock = threading.Lock()
self._handle_id = uuid.uuid4().hex
self._completion_observed = False
self._completion_notified = False
# Serializes kill/reap so a pid fallback can never outlive the process handle.
self._kill_lock = threading.Lock()
self._started = time.monotonic()
# POSIX: own process group so kill() signals the whole pipeline; Windows
# contains the tree in a kill-on-close job object.
self._status_read = -1
self._wake_read = -1
self._wake_write = -1
# True only while the pump moves a chunk from the pipe into the buffer.
self._pump_transfer = False
self._job: int | None = None
self._completion_marker: bytes | None = None
status_write = -1
if _IS_POSIX:
# Full-duplex status channel: the child end rides in as stdin (fd 0)
# and the script remaps it to _STATUS_FD before swapping in /dev/null
# (dash rejects multi-digit fds in redirections at parse time). The
# parent end doubles as the gate: the child blocks on it until the
# pid is journaled, so a kernel kill in that window cannot leak an
# unjournaled command (parent death closes the socket -> child exits).
parent_sock, child_sock = socket.socketpair()
self._status_read = parent_sock.detach()
status_write = child_sock.detach()
try:
self._wake_read, self._wake_write = os.pipe()
except BaseException:
os.close(self._status_read)
os.close(status_write)
raise
completion_token = secrets.token_hex(32)
# Halves stop passive echoes; a deliberate forgery freezes only this call while later bytes stay live.
token_midpoint = len(completion_token) // 2
self._completion_marker = (
_COMPLETION_PREFIX + completion_token.encode("ascii") + _COMPLETION_SUFFIX
)
script = _status_script(
_with_prefix(command),
completion_token[:token_midpoint],
completion_token[token_midpoint:],
)
else:
# Windows lacks a foreground-status channel, so its exit drain stays best-effort.
script = _with_prefix(command)
self._job = _winjob.create_job()
if self._job is None:
# Nothing spawned yet, so nothing can leak: refuse to start.
raise RuntimeError("bash(): Windows job containment could not be established")
try:
self._proc: subprocess.Popen[bytes] | _winjob.JobProcess
if _IS_POSIX:
self._proc = subprocess.Popen(
[_shell(), "-c", script],
cwd=os.getcwd(),
env=_child_env(),
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
start_new_session=True,
stdin=status_write,
)
else:
self._proc = _winjob.spawn_in_job(
self._job, [_shell(), "-c", script], cwd=os.getcwd(), env=_child_env()
)
except BaseException:
for fd in (self._status_read, self._wake_read, self._wake_write):
if fd >= 0:
os.close(fd)
if self._job is not None:
job, self._job = self._job, None
_winjob.close(job)
raise
finally:
if status_write >= 0:
os.close(status_write)
self._pid: int = self._proc.pid
self._released = False
with _live_lock:
_live_handles.add(self)
enrolled = _record_journal(self._pid, active=True)
if not enrolled:
# Fail closed: a configured journal that cannot enroll the pid must
# not let the command run (the host reaper would never see it).
self._abort_spawn()
raise RuntimeError(
"bash(): orphan-journal enrollment failed (journal configured but the "
"pid could not be recorded); the spawned process was killed"
)
if _IS_POSIX:
# Journal first, then open the gate: the child does not run the user
# command until this byte arrives. A failed write means the child
# already died; the status/EOF paths report that normally.
try:
os.write(self._status_read, b"\n")
except OSError:
pass
else:
# The child is already job-contained and journaled; resume is the
# last step. A failed resume would strand a permanently suspended
# child: fail closed via the assigned job.
if not cast("_winjob.JobProcess", self._proc).resume():
self._abort_spawn()
raise RuntimeError("bash(): Windows job containment could not be established")
threading.Thread(target=self._pump, daemon=True).start()
threading.Thread(target=self._report, daemon=True).start()
threading.Thread(target=self._watch, daemon=True).start()
def _mark_completion_observed(self) -> None:
with self._callback_lock:
if self._completion_observed:
return
self._completion_observed = True
notified = self._completion_notified
if notified:
_notify_completion({"kind": "observed", "handleId": self._handle_id})
@property
def pid(self) -> int:
self._released = True
return self._pid
@property
def running(self) -> bool:
# Group liveness, matching kill()'s guard and the journal; poll()/await
# keep foreground result semantics after `cmd &` returns early.
self._released = True
running = not self._reaped
if not running and self._done.is_set():
self._mark_completion_observed()
return running
def output(self) -> str:
self._released = True
text = self._buffer.text()
if self._done.is_set():
self._mark_completion_observed()
return text
def tail(self, n: int = 50) -> str:
self._released = True
text = "\n".join(self._buffer.text().splitlines()[-n:])
if self._done.is_set():
self._mark_completion_observed()
return text
def poll(self) -> BashResult | None:
self._released = True
if not self._done.is_set():
return None
self._mark_completion_observed()
return self._result
def kill(self, sig: int = signal.SIGTERM, grace: float = 5.0) -> None:
# Guard on group death, not _done: kill() must still reach a lingering
# background group after the foreground result was already delivered.
self._released = True
self._mark_completion_observed()
if self._reaped:
return
if not _IS_POSIX:
with self._kill_lock:
if self._reaped: # re-check: _watch may have reaped while we waited
return
if self._job is not None and _winjob.terminate(self._job):
return
# TerminateJobObject failed or reap raced: taskkill fallback.
if not _taskkill_tree(self._pid):
try:
self._proc.kill()
except OSError:
pass
return
_signal_group(self._pid, sig)
if sig == signal.SIGTERM:
timer = threading.Timer(grace, self._force_kill)
timer.daemon = True
timer.start()
def _force_kill(self) -> None:
if not self._reaped:
_signal_group(self._pid, signal.SIGKILL)
def _pump(self) -> None:
stdout = self._proc.stdout
assert stdout is not None
if not _IS_POSIX:
try:
while chunk := stdout.read1(_READ_CHUNK):
self._buffer.write(chunk)
except (OSError, ValueError):
pass
stdout.close()
self._eof.set()
return
fd = stdout.fileno()
try:
with selectors.DefaultSelector() as sel:
sel.register(fd, selectors.EVENT_READ)
while True:
sel.select()
self._pump_transfer = True
try:
chunk = os.read(fd, _READ_CHUNK)
if not chunk:
break
self._consume_output(chunk)
finally:
self._pump_transfer = False
except (OSError, ValueError):
pass
self._abandon_completion()
try:
stdout.close()
except OSError:
pass
self._eof.set()
def _consume_output(self, chunk: bytes) -> None:
marker = self._completion_marker
assert marker is not None
with self._completion_lock:
if self._completion_terminal.is_set():
self._buffer.write(chunk)
return
data = self._completion_pending + chunk
marker_at = data.find(marker)
if marker_at >= 0:
self._buffer.write(data[:marker_at])
self._completion_pending = b""
self._completion_output = self._buffer.text()
self._completion_terminal.set()
self._buffer.write(data[marker_at + len(marker) :])
return
retained = 0
for size in range(min(len(data), len(marker) - 1), 0, -1):
if data.endswith(marker[:size]):
retained = size
break
self._buffer.write(data[:-retained] if retained else data)
self._completion_pending = data[-retained:] if retained else b""
def _abandon_completion(self) -> None:
with self._completion_lock:
if self._completion_terminal.is_set():
return
self._buffer.write(self._completion_pending)
self._completion_pending = b""
self._completion_terminal.set()
def _wait_for_completion(self) -> str | None:
self._completion_terminal.wait()
return self._completion_output
def _report(self) -> None:
# Finalize at foreground completion (status channel), not EOF, so
# `cmd &` does not hang the await; the shell then `wait`s for its
# background jobs, keeping the journaled group identity alive.
status: int | None = None
try:
status = self._read_status()
# Reserve the delivered status before draining so a shell death during
# the drain window cannot override it with wait()'s signal exit code.
with self._callback_lock:
self._status = status
finally:
# _watch blocks on this event without a timeout, so every exit path
# (parsed status, EOF, garbage, exception) must set it.
self._status_known.set()
if status is not None:
output = self._wait_for_completion()
if output is None:
self._drain_grace()
self._finalize(status, output)
def _watch(self) -> None:
# Observe shell death independently of the status socket: an early
# `exit`/`exec`/`set -e`/fatal signal skips `printf`, and background
# children can hold the socket open past the shell's lifetime.
exit_code = self._proc.wait()
if self._wake_write >= 0:
# Unblock _read_status: background children can hold the status socket
# open past the shell's lifetime via bash's saved-fd duplicate.
try:
os.write(self._wake_write, b"x")
except OSError:
pass
os.close(self._wake_write)
# _report always sets _status_known (try/finally), so wait indefinitely:
# a slow reporter can never lose a delivered status to wait()'s code.
self._status_known.wait()
with self._callback_lock:
delivered = self._status
if delivered is None and not self._done.is_set():
self._abandon_completion()
self._drain_grace()
self._finalize(exit_code)
with self._kill_lock:
delivered = self._reap_group()
self._reaped = True
if not _IS_POSIX:
# Reaped: pid fallbacks are gone, so the handle may finally close.
cast("_winjob.JobProcess", self._proc).close()
if delivered:
_record_journal(self._pid, active=False)
with _live_lock:
_live_handles.discard(self)
def _reap_group(self) -> bool:
# Group liveness, not leader death, gates the inactive record: members
# that outlive the leader would leak behind a stale journal anchor.
if not _IS_POSIX:
# Terminate then close the last handle: kill-on-close reaps
# stragglers. An unproven terminate falls back to taskkill; if
# that also fails the record stays active for the host reaper.
delivered = False
if self._job is not None:
delivered = _winjob.terminate(self._job)
job, self._job = self._job, None
_winjob.close(job)
return delivered or _taskkill_tree(self._pid)
try:
os.killpg(self._pid, 0)
except ProcessLookupError:
return True # group already gone
except PermissionError:
pass
return _signal_group(self._pid, signal.SIGKILL)
def _read_status(self) -> int | None:
if self._status_read < 0:
return None
try:
# DefaultSelector (kqueue/epoll) instead of select(): select() rejects
# fds >= FD_SETSIZE (1024) even when the process fd limit is higher.
with selectors.DefaultSelector() as sel:
sel.register(self._status_read, selectors.EVENT_READ)
sel.register(self._wake_read, selectors.EVENT_READ)
line = b""
while b"\n" not in line:
ready = {key.fd for key, _ in sel.select()}
# Prefer status bytes: any status write happens before shell exit,
# so it is already readable whenever the wake fd fires.
if self._status_read not in ready:
break # shell died without writing a status
chunk = os.read(self._status_read, 64)
if not chunk:
break # EOF without a full status line
line += chunk
return int(line)
except (OSError, ValueError):
return None
finally:
os.close(self._status_read)
os.close(self._wake_read)
def _drain_grace(self) -> None:
# Best-effort fallback when process exit/EOF arrives without a sentinel.
deadline = time.monotonic() + 0.5
size = self._buffer.size()
while time.monotonic() < deadline:
if self._eof.wait(0.05):
return
# A chunk between pipe read and buffer commit (transfer flag) is
# invisible to both FIONREAD and the buffer size; wait it out.
if self._pipe_pending() or self._pump_transfer:
size = self._buffer.size()
continue
current = self._buffer.size()
if current == size:
return
size = current
def _pipe_pending(self) -> bool:
# POSIX only: FIONREAD on the capture pipe; Windows keeps the
# quiescence heuristic (best-effort parity).
if not _IS_POSIX or self._eof.is_set():
return False
stdout = self._proc.stdout
if stdout is None:
return False
try:
pending = struct.unpack("i", fcntl.ioctl(stdout.fileno(), termios.FIONREAD, struct.pack("i", 0)))[0]
except (OSError, ValueError):
return False
return pending > 0
def _finalize(self, exit_code: int, output: str | None = None) -> None:
with self._callback_lock:
if self._done.is_set():
return
result = BashResult(
exit_code=exit_code,
output=self._buffer.text() if output is None else output,
duration=time.monotonic() - self._started,
)
self._result = result
self._done.set()
notify = not self._completion_observed
if notify:
self._completion_notified = True
callbacks = self._callbacks
self._callbacks = []
if notify:
_notify_completion(
{
"kind": "completed",
"handleId": self._handle_id,
"pid": self._pid,
"exitCode": result.exit_code,
"duration": result.duration,
}
)
for callback in callbacks:
callback()
def _add_done_callback(self, callback: Callable[[], None]) -> None:
with self._callback_lock:
if not self._done.is_set():
self._callbacks.append(callback)
return
callback()
async def _wait(self) -> BashResult:
# Asyncio-native wakeup: no executor thread is parked for the command's
# duration, so many concurrent awaits cannot exhaust the default pool.
loop = asyncio.get_running_loop()
fut: asyncio.Future[None] = loop.create_future()
def _wake() -> None:
try:
loop.call_soon_threadsafe(lambda: fut.done() or fut.set_result(None))
except RuntimeError:
pass # awaiting loop already closed
self._add_done_callback(_wake)
await fut
assert self._result is not None
return self._result
async def _wait_owned(self) -> BashResult:
# One-shot `await bash(cmd)` owns the process: a cancelled await (e.g.
# a kernel interrupt) must not leave the command running. TERM, bounded
# grace, group KILL, then a bounded confirmed-exit wait before the
# CancelledError propagates, so no side effect can land after it.
try:
return await self._wait()
except asyncio.CancelledError:
# Signal synchronously first: even if the cleanup awaits below are
# re-cancelled, TERM is already delivered and the escalation timer
# armed. The confirm wait runs as a shielded task so repeated
# cancels of this task cannot skip it (they re-raise into awaits
# inside this except block); the loop re-awaits until it finishes
# (the confirm coroutine itself is bounded).
self.kill(grace=_CANCEL_TERM_GRACE)
confirm = asyncio.ensure_future(self._confirm_group_exit())
while not confirm.done():
try:
await asyncio.shield(confirm)
except asyncio.CancelledError:
continue
raise
async def _confirm_group_exit(self) -> None:
if not await self._await_group_death(_CANCEL_TERM_GRACE):
if _IS_POSIX:
_signal_group(self._pid, signal.SIGKILL)
else:
# kill() holds the escalation lock; to_thread keeps the loop free.
await asyncio.to_thread(self.kill)
await self._await_group_death(_CANCEL_KILL_WAIT)
def _group_alive(self) -> bool:
if not _IS_POSIX:
job = self._job # snapshot: _watch may clear it concurrently
if job is not None:
# Job accounting sees detached descendants a dead leader hides.
empty = _winjob.is_empty(job)
if empty is not None:
return not empty
return self._proc.poll() is None
try:
os.killpg(self._pid, 0)
except ProcessLookupError:
return False
except PermissionError:
pass
return True
async def _await_group_death(self, timeout: float) -> bool:
deadline = time.monotonic() + timeout
while self._group_alive():
if time.monotonic() >= deadline:
return False
await asyncio.sleep(0.02)
return True
def _abort_spawn(self) -> None:
# Enrollment or containment failed before the gate opened (POSIX) or
# while the child is still suspended, before resume (Windows): kill
# the child and unwind the handle before threads start.
if _IS_POSIX:
for fd in (self._status_read, self._wake_read, self._wake_write):
if fd >= 0:
try:
os.close(fd)
except OSError:
pass
self._status_read = self._wake_read = self._wake_write = -1
delivered = _signal_group(self._pid, signal.SIGKILL)
else:
with self._kill_lock:
delivered = False
if self._job is not None:
delivered = _winjob.terminate(self._job)
job, self._job = self._job, None
_winjob.close(job)
if not delivered:
# Pre-resume abort: the never-run leader has no descendants, so a
# delivered kill retires the journal record.
try:
self._proc.kill()
delivered = True
except OSError:
pass
if self._proc.stdout is not None:
self._proc.stdout.close()
# The blocking wait stays outside the lock: hProcess is still open, so a
# concurrent raw-pid fallback stays pinned to the right process.
try:
self._proc.wait(timeout=5)
except (OSError, subprocess.SubprocessError):
pass
with self._kill_lock:
self._reaped = True
if not _IS_POSIX:
# Reaped commits before close: later lock holders skip raw-pid fallbacks.
cast("_winjob.JobProcess", self._proc).close()
with _live_lock:
_live_handles.discard(self)
if delivered:
_record_journal(self._pid, active=False)
def __await__(self) -> Generator[Any, None, BashResult]:
# A handle awaited before any other API use is a one-shot command tied
# to the await (kill-on-cancel); touching the handle API first marks it
# as a deliberate background handle whose awaits only wait.
owned = not self._released
self._released = True
self._mark_completion_observed()
return (self._wait_owned() if owned else self._wait()).__await__()
def __repr__(self) -> str:
if self._result is not None:
self._mark_completion_observed()
state = f"exit_code={self._result.exit_code}" if self._result else "running"
return f"<BashHandle pid={self._pid} {state} command={self.command!r}>"
def bash(command: str) -> BashHandle:
"""Start a shell command immediately; await the handle for the result.
`await bash(cmd)` is a one-shot: cancelling the await (e.g. an interrupt)
kills the command's process group. `h = bash(cmd)` used as a background
handle (any .pid/.running/.output()/.tail()/.poll()/.kill() access before
the first await) survives cancellation; awaiting it only waits. Leak
containment is per-platform: process groups plus the orphan journal on
POSIX; a kill-on-close job object on Windows entered while the child is
still suspended, so no descendant can escape it and kill()/crash cleanup
are unconditional -- bash() raises if containment cannot be established.
Output written after the completion fence (e.g. by an EXIT trap or a
background job) is not in BashResult.output but stays visible via
handle.output()/tail().
"""
if not isinstance(command, str) or not command:
raise TypeError("command must be a non-empty str")
if _CD_PREFIX_RE.match(command) or _CD_CHAIN_RE.search(command):
raise ValueError(
"bash() rejects `cd` prefixes: each bash() call runs in its own process, "
"so `cd` here cannot move the persistent kernel working directory. Move it "
"once with os.chdir(...) or repo.cd(...), then call bash() without the prefix. "
"For git, use the git skill (git.<cmd>), which resolves the repo root itself."
)
_install_shutdown_hook()
return BashHandle(command)
def _shell() -> str:
# Read per call so env changes made in the REPL apply to later commands.
override = os.environ.get("PRIME_AGENT_BASH_SHELL")
if override:
if not os.path.isabs(override):
raise ValueError("PRIME_AGENT_BASH_SHELL must be an absolute path")
return override
if not _IS_POSIX:
# Never consult PATH on Windows: a repo-controlled PATH could supply
# the shell. The host injects PRIME_AGENT_BASH_SHELL when one exists.
raise RuntimeError(
"bash() needs PRIME_AGENT_BASH_SHELL set to the absolute path of a "
"POSIX shell on Windows (e.g. install Git Bash in its default "
"location so the host injects it)"
)
# PATH fallback only serves bare/standalone POSIX runtime use: the host
# always injects PRIME_AGENT_BASH_SHELL (an absolute path) when a shell exists.
shell = shutil.which("bash")
return shell or "/bin/sh"
def _with_prefix(command: str) -> str:
prefix = os.environ.get("PRIME_AGENT_BASH_COMMAND_PREFIX")
return f"{prefix}\n{command}" if prefix else command
def _fence_printf() -> str:
# `\command -p printf` defeats alias expansion but not a user-defined shell
# function named `command`, which would swallow both fence frames and leave
# the await hanging until the shell dies (wedged behind background jobs). A
# slash-qualified command name bypasses function and alias lookup for
# ordinary command names, so resolve printf on the system default utility PATH.
path = shutil.which("printf", path=os.confstr("CS_PATH") or os.defpath)
if path and "'" not in path:
return f"'{path}'"
return "\\command -p printf"
def _status_script(command: str, completion_a: str, completion_b: str) -> str:
# Closed control fds preserve background behavior; supported shells atomically write the frame.
emit = _fence_printf()
return (
f"exec {_STATUS_FD}>&0 {_OUTPUT_FD}>&1 0</dev/null\n"
f"read -r _prime_agent_gate <&{_STATUS_FD} || exit 127\n"
"{\n"
f"{command}\n"
f"}} {_OUTPUT_FD}>&- {_STATUS_FD}>&-\n"
"__prime_status=$?\n"
"\\set +x\n"
f"{emit} '\\036prime-agent-complete:%s%s\\037' "
f"'{completion_a}' '{completion_b}' >&{_OUTPUT_FD} || exit \"$__prime_status\"\n"
f"{emit} '%s\\n' \"$__prime_status\" >&{_STATUS_FD}\n"
f"exec {_OUTPUT_FD}>&- {_STATUS_FD}>&-\n"
"wait\n"
'exit "$__prime_status"\n'
)
def _child_env() -> dict[str, str]:
return {**os.environ, "NO_COLOR": "1", "TERM": "dumb", "CLICOLOR": "0", "FORCE_COLOR": "0"}
def _signal_group(pid: int, sig: int) -> bool:
"""True when the signal was delivered or the group is already gone."""
try:
os.killpg(pid, sig)
except ProcessLookupError:
return True # already dead: safe to mark the journal record inactive
except OSError:
return False # not delivered: the record must stay active for the host reaper
return True
def _system32(*parts: str) -> str:
# Absolute paths for Windows helper binaries: PATH (and CWD on Windows
# CPython) lookup could resolve a planted taskkill.exe/powershell.exe.
root = os.environ.get("SystemRoot", r"C:\Windows")
return os.path.join(root, "System32", *parts)
def _helper_env() -> dict[str, str]:
return {**os.environ, "NoDefaultCurrentDirectoryInExePath": "1"}
def _taskkill_tree(pid: int) -> bool:
# Windows has no process groups to signal; taskkill /T kills the whole tree.
try:
return (
subprocess.run(
[_system32("taskkill.exe"), "/PID", str(pid), "/T", "/F"],
capture_output=True,
timeout=10,
env=_helper_env(),
).returncode
== 0
)
except (OSError, subprocess.SubprocessError):
return False
def _process_start_id(pid: int) -> str | None:
if os.name == "nt":
# Mirrors getWindowsProcessStartId in session-lease.ts byte-for-byte so
# the host's identity comparison matches the journaled string.
try:
out = subprocess.run(
[
_system32("WindowsPowerShell", "v1.0", "powershell.exe"),
"-NoLogo",
"-NoProfile",
"-NonInteractive",
"-Command",
f"([System.Diagnostics.Process]::GetProcessById({pid})).StartTime.ToUniversalTime().Ticks",
],
capture_output=True,
text=True,
timeout=5,
env=_helper_env(),
).stdout.strip()
return f"win:{out}" if out.isdigit() else None
except (OSError, subprocess.SubprocessError):
return None
try:
with open(f"/proc/{pid}/stat", "r") as f:
stat = f.read()
fields = stat[stat.rindex(")") + 2 :].split(" ")
if len(fields) > 19 and fields[19]:
return f"proc:{fields[19]}"
except (OSError, ValueError):
pass
try:
# macOS has no /proc; /bin/ps is always present there, so use the
# absolute path (bare `ps` stays only as the exotic-POSIX last resort).
ps = "/bin/ps" if sys.platform == "darwin" else "ps"
out = subprocess.run(
[ps, "-p", str(pid), "-o", "lstart="], capture_output=True, text=True, timeout=5
).stdout.strip()
return f"ps:{out}" if out else None
except (OSError, subprocess.SubprocessError):
return None
def _record_journal(pid: int, active: bool) -> bool:
# Returns False only when the journal is configured but enrollment failed;
# active-record callers must then fail closed. Active records always carry
# a processStartId so host reaping stays identity-verified.
path = os.environ.get("PRIME_AGENT_INTERNAL_ORPHAN_PROCESS_JOURNAL")
owner = os.environ.get("PRIME_AGENT_KERNEL_OWNER_PID")
if not path or not owner:
return True
try:
owner_pid = int(owner)
except ValueError:
return False
start_id = _process_start_id(pid) if active else None
if active and start_id is None:
return False
record: dict[str, Any] = {
"version": 1,
"pid": pid,
"ownerPid": owner_pid,
# The host reaps bash children per kernel pid when it kills or loses this kernel.
"kernelPid": os.getpid(),
**({"processStartId": start_id} if start_id else {}),
"active": active,
"recordedAt": datetime.now(timezone.utc).isoformat(),
}
data = (json.dumps(record) + "\n").encode()
try:
fd = os.open(path, os.O_WRONLY | os.O_APPEND | os.O_CREAT, 0o600)
try:
# Complete-write loop: a short write would leave a truncated JSON
# line that the host discards, which must count as failure.
view = memoryview(data)
while view:
written = os.write(fd, view)
if written <= 0:
return False
view = view[written:]
os.fsync(fd)
finally:
os.close(fd)
except OSError:
return False
return True
def _kill_live_handles() -> None:
with _live_lock:
handles = list(_live_handles)
for handle in handles:
if _IS_POSIX:
delivered = _signal_group(handle._pid, signal.SIGKILL)
else:
with handle._kill_lock:
if handle._reaped:
continue
delivered = handle._job is not None and _winjob.terminate(handle._job)
if not delivered:
delivered = _taskkill_tree(handle._pid)
if not delivered:
# Leader-only fallback cannot prove the tree died: never
# justifies an inactive record.
try:
handle._proc.kill()
except OSError:
pass
if delivered:
_record_journal(handle._pid, active=False)
def _install_shutdown_hook() -> None:
global _hook_installed
with _hook_lock:
if _hook_installed:
return
_hook_installed = True
atexit.register(_kill_live_handles)