Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
75 kB · 1822 lines
Python
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823"""Async-by-default shell execution: bash() spawns immediately and returns a live handle."""
from __future__ import annotations
import asyncioimport atexitimport functoolsimport inspectimport jsonimport osimport queueimport reimport secretsimport selectorsimport shleximport shutilimport signalimport socketimport structimport subprocessimport sysimport tempfileimport threadingimport timefrom collections import dequefrom collections.abc import Callable, Generatorfrom dataclasses import dataclassfrom datetime import datetime, timezonefrom pathlib import Pathfrom typing import Any, Literal, 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_DEFAULT_SPILL_BYTES = 64 * 1024_DEFAULT_SPILL_LINES = 400_REMOTE_EXEC_TOKENS = {"ssh", "scp", "sftp", "mosh"}_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_COMPLETION_NOTICE_COMMAND_CAP = 1000_ASYNCIO_WRAPPER_CALLBACKS = { ("asyncio.tasks", "gather.<locals>._done_callback"), ("asyncio.tasks", "shield.<locals>._inner_done_callback"), ("asyncio.tasks", "_wait.<locals>._on_completion"), ("asyncio.tasks", "as_completed.<locals>._on_completion"), ("asyncio.tasks", "_release_waiter"),}
_CD_SEGMENT_RE = re.compile(r"^\s*cd(?:\s|$)", re.IGNORECASE)_VCS_SEGMENT_RE = re.compile(r"^\s*(?:git|jj)(?:\s|$)", re.IGNORECASE)_SEARCH_SEGMENT_RE = re.compile(r"^\s*(?:grep|egrep|fgrep|rg|ripgrep|ag|ack)(?:\s|$)", re.IGNORECASE)_FIND_SEGMENT_RE = re.compile(r"^\s*(?:find|fd)(?:\s|$)", re.IGNORECASE)_LIST_SEGMENT_RE = re.compile(r"^\s*(?:ls|dir)(?:\s|$)", re.IGNORECASE)_COMMAND_SPLIT_RE = re.compile(r"&&|\|\||;|\n")_SHELL_WRAPPER_TOKEN_RE = re.compile(r"^(?:bash|sh|zsh|dash|ksh)$", re.IGNORECASE)_WRAPPER_PREFIX_TOKENS = {"eval", "exec", "sudo", "nohup", "time"}
_live_handles: set["BashHandle"] = set()_live_lock = threading.Lock()_hook_installed = False_hook_lock = threading.Lock()
def _current_cell_completion_context() -> tuple[asyncio.Event, asyncio.Task[Any] | None] | None: """Get the creating REPL cell's lifecycle without coupling standalone use to repl.""" try: from . import repl
if repl.is_active(): return repl.current_cell_completion_context() except (ImportError, RuntimeError): pass return None
def _consume_notice_task(task: asyncio.Task[None]) -> None: """Retrieve detached notifier failures so they never become loop warnings.""" if not task.cancelled(): task.exception()
def _completion_reaches( start: asyncio.Future[Any], targets: tuple[asyncio.Future[Any], ...]) -> bool: """Follow asyncio's wrapper and TaskGroup ownership callbacks.""" pending = [start] seen_futures: set[int] = set() seen_values: set[int] = set()
def collect(value: Any, depth: int = 0) -> None: if isinstance(value, asyncio.Future): pending.append(value) return identity = id(value) if depth >= 4 or identity in seen_values: return seen_values.add(identity)
nested: list[Any] = [] if isinstance(value, asyncio.Queue): pending.extend(value._getters) elif isinstance(value, functools.partial): nested.extend((value.func, value.args, value.keywords)) elif isinstance(value, dict): nested.extend(value.keys()) nested.extend(value.values()) elif isinstance(value, (tuple, list, set, frozenset)): nested.extend(value) else: closure = getattr(value, "__closure__", None) or () for cell in closure: try: nested.append(cell.cell_contents) except ValueError: pass bound_self = getattr(value, "__self__", None) if bound_self is not None: nested.append(bound_self) for item in nested: collect(item, depth + 1)
while pending: future = pending.pop() if any(future is target for target in targets): return True if id(future) in seen_futures: continue seen_futures.add(id(future)) for entry in getattr(future, "_callbacks", None) or (): callback = entry[0] if isinstance(entry, tuple) else entry base = callback.func if isinstance(callback, functools.partial) else callback identity = (getattr(base, "__module__", None), getattr(base, "__qualname__", None)) if identity in _ASYNCIO_WRAPPER_CALLBACKS: collect(callback) elif identity == ("asyncio.tasks", "_AsCompletedIterator._handle_completion"): collect(base.__self__._done) elif identity == (None, "Task.task_wakeup"): task = getattr(callback, "__self__", None) if isinstance(task, asyncio.Task): pending.append(task) elif identity == ("asyncio.taskgroups", "TaskGroup._on_task_done"): parent = getattr(getattr(callback, "__self__", None), "_parent_task", None) if isinstance(parent, asyncio.Future): pending.append(parent) return False
def _creating_cell_waits_for( owner: asyncio.Task[Any] | None, awaiter: asyncio.Task[Any] | None) -> bool: """Return whether the cell owner directly or transitively waits for awaiter.""" if owner is None or awaiter is None: return False if owner is awaiter: return True waiter = getattr(owner, "_fut_waiter", None) targets: tuple[asyncio.Future[Any], ...] = (owner,) if isinstance(waiter, asyncio.Future): targets += (waiter,) return _completion_reaches(awaiter, targets)
@dataclass(frozen=True)class BashResult: exit_code: int output: str duration: float stdout: str = "" stderr: str = "" output_file: str | None = None
@property def file_path(self) -> str | None: """Alias to output_file if output was dumped to disk.""" return self.output_file
@property def path(self) -> Path | None: """Path object pointing to output_file if spilled, else None.""" return Path(self.output_file) if self.output_file else None
def full_output(self) -> str: """Return the complete unclipped output, reading from output_file if spilled.""" if self.output_file and os.path.exists(self.output_file): try: return Path(self.output_file).read_text(encoding="utf-8", errors="replace") except OSError: pass return self.output
def __str__(self) -> str: return self.output
def __repr__(self) -> str: lines = self.output.splitlines() line_count = len(lines) byte_count = len(self.output.encode("utf-8", errors="replace")) dur_str = f"{self.duration:.2f}s" file_info = f" output_file={self.output_file!r}" if self.output_file else "" if byte_count <= 200 and line_count <= 4: return f"<BashResult exit_code={self.exit_code} duration={dur_str}{file_info} output={self.output!r}>" preview = "\n".join(lines[:3]) if line_count > 3: preview += "\n..." return f"<BashResult exit_code={self.exit_code} duration={dur_str}{file_info} ({line_count} lines, {byte_count} bytes) preview:\n{preview}>"
@dataclass(frozen=True)class BashStream: """Stream descriptor for stdout or stderr from a BashHandle.""" handle: BashHandle channel: Literal["stdout", "stderr"]
def text(self) -> str: return self.handle._get_buffered_bytes(self.channel).decode("utf-8", errors="replace")
def pipe(self, command: str) -> BashHandle: return bash(command, input=self)
def __str__(self) -> str: return self.text()
def __repr__(self) -> str: return f"<BashStream channel={self.channel!r} handle={self.handle!r}>"
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")
def raw_bytes(self) -> bytes: with self._lock: return bytes(self._head) + b"".join(self._tail)
def _write_bash_spill_file(pid: int, text: str) -> str: session_dir = os.environ.get("RLM_SESSION_DIR") if session_dir and session_dir.strip(): spill_dir = Path(session_dir.strip()) / "artifacts" / "bash" else: spill_dir = Path(tempfile.gettempdir()) / f"prime-bash-{os.getuid()}" spill_dir.mkdir(parents=True, exist_ok=True) spill_path = spill_dir / f"bash-{pid}-{int(time.time() * 1000)}-{secrets.token_hex(4)}.log" spill_path.write_text(text, encoding="utf-8", errors="replace") return str(spill_path)
def _truncate_spill_preview(text: str, total_bytes: int, total_lines: int, file_path: str) -> str: lines = text.splitlines(keepends=True) if len(lines) <= 80: head = text[:2048] tail = text[-2048:] return f"{head}\n\n... [output truncated: {total_bytes} bytes dropped from inline view. Full output saved to {file_path}] ...\n\n{tail}" head_lines = lines[:50] tail_lines = lines[-30:] marker = f"\n... [output truncated: {total_lines} lines, {total_bytes} bytes dropped from inline view. Full output saved to {file_path}] ...\n" return "".join(head_lines) + marker + "".join(tail_lines)
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, input: str | bytes | BashHandle | BashStream | None = None, display_command: str | None = None, ) -> None: self._input_path: str | None = None self._fifo_dir: str | None = None self.command = display_command if display_command is not None else command completion_context = _current_cell_completion_context() self._creating_cell_finished = completion_context[0] if completion_context else None self._creating_cell_task = completion_context[1] if completion_context else None self._awaited_by_creating_cell = False if isinstance(input, (BashHandle, BashStream)): if isinstance(input, BashHandle): source_handle = input channel = "combined" else: source_handle = input.handle channel = input.channel
if _IS_POSIX: fifo_dir = tempfile.mkdtemp(prefix="prime-bash-fifo-") fifo_path = os.path.join(fifo_dir, "stream.fifo") os.mkfifo(fifo_path) self._input_path = fifo_path self._fifo_dir = fifo_dir
threading.Thread( target=self._feed_from_handle, args=(source_handle, channel, fifo_path), daemon=True, ).start() else: # Windows fallback: drain source buffer to a NamedTemporaryFile data = source_handle._get_buffered_bytes(channel) tf = tempfile.NamedTemporaryFile(delete=False, prefix="prime-bash-input-") try: tf.write(data) tf.flush() self._input_path = tf.name finally: tf.close() elif input is not None: data = input.encode("utf-8") if isinstance(input, str) else bytes(input) tf = tempfile.NamedTemporaryFile(delete=False, prefix="prime-bash-input-") try: tf.write(data) tf.flush() self._input_path = tf.name finally: tf.close() self._stdout_buffer = _BoundedBuffer() self._stderr_buffer = _BoundedBuffer() self._combined_buffer = _BoundedBuffer() self._buffer = self._combined_buffer self._listeners: list[Callable[[str, bytes], None]] = [] self._listener_lock = threading.Lock() 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._reap_callback: Callable[[], None] | None = None self._callback_lock = threading.Lock() # 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:], input_file=self._input_path, ) else: # Windows lacks a foreground-status channel, so its exit drain stays best-effort. script = _with_prefix(command) if self._input_path: script = f"{script} < {shlex.quote(self._input_path)}" 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.PIPE, 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: self._cleanup_input() 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() self._schedule_background_completion_notice()
@property def stdout(self) -> BashStream: return BashStream(self, "stdout")
@property def stderr(self) -> BashStream: return BashStream(self, "stderr")
def pipe(self, command: str) -> BashHandle: return bash(command, input=self)
def _get_buffered_bytes(self, channel: str) -> bytes: if channel == "stdout": return self._stdout_buffer.raw_bytes() elif channel == "stderr": return self._stderr_buffer.raw_bytes() else: return self._combined_buffer.raw_bytes()
def _write_stdout(self, data: bytes) -> None: if not data: return with self._listener_lock: self._stdout_buffer.write(data) self._combined_buffer.write(data) listeners = list(self._listeners) for l in listeners: try: l("stdout", data) except Exception: pass
def _write_stderr(self, data: bytes) -> None: if not data: return with self._listener_lock: self._stderr_buffer.write(data) self._combined_buffer.write(data) listeners = list(self._listeners) for l in listeners: try: l("stderr", data) except Exception: pass
def _emit_eof(self) -> None: with self._listener_lock: listeners = list(self._listeners) for l in listeners: try: l("stdout", b"") l("stderr", b"") except Exception: pass
def _feed_from_handle(self, source_handle: BashHandle, channel: str, fifo_path: str) -> None: q: queue.Queue[bytes] = queue.Queue()
def listener(ch: str, chunk: bytes) -> None: if channel == "combined" or channel == ch: q.put(chunk)
with source_handle._listener_lock: already_done = source_handle._eof.is_set() initial_bytes = source_handle._get_buffered_bytes(channel) if not already_done: source_handle._listeners.append(listener)
try: with open(fifo_path, "wb") as writer: if initial_bytes: writer.write(initial_bytes) writer.flush() if not already_done: while True: chunk = q.get() if not chunk: break writer.write(chunk) writer.flush() except (OSError, BrokenPipeError): pass finally: with source_handle._listener_lock: if listener in source_handle._listeners: source_handle._listeners.remove(listener)
@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 return not self._reaped
def output(self) -> str: self._released = True return self._buffer.text()
def text(self) -> str: """Unified text accessor so BashHandle and BashStream share the same protocol.""" return self.output()
def tail(self, n: int = 50) -> str: self._released = True return "\n".join(self._buffer.text().splitlines()[-n:])
def poll(self) -> BashResult | None: self._released = True return self._result if self._done.is_set() else None
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 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 stderr = getattr(self._proc, "stderr", None) assert stdout is not None if not _IS_POSIX or stderr is None: try: while chunk := stdout.read1(_READ_CHUNK): self._consume_output(chunk) except (OSError, ValueError): pass stdout.close() if stderr is not None: try: while chunk := stderr.read1(_READ_CHUNK): self._write_stderr(chunk) except (OSError, ValueError): pass stderr.close() self._emit_eof() self._eof.set() return
stdout_fd = stdout.fileno() stderr_fd = stderr.fileno() try: os.set_blocking(stdout_fd, False) os.set_blocking(stderr_fd, False) except OSError: pass
open_fds = {stdout_fd, stderr_fd} try: with selectors.DefaultSelector() as sel: sel.register(stdout_fd, selectors.EVENT_READ, data="stdout") sel.register(stderr_fd, selectors.EVENT_READ, data="stderr") while open_fds: for key, _ in sel.select(): fd = key.fd channel = key.data self._pump_transfer = True try: try: chunk = os.read(fd, _READ_CHUNK) except (BlockingIOError, InterruptedError): continue if not chunk: sel.unregister(fd) open_fds.discard(fd) if channel == "stdout": self._abandon_completion() continue if channel == "stdout": self._consume_output(chunk) else: self._write_stderr(chunk) finally: self._pump_transfer = False except (OSError, ValueError): pass self._abandon_completion() try: stdout.close() except OSError: pass try: stderr.close() except OSError: pass self._emit_eof() self._eof.set()
def _consume_output(self, chunk: bytes) -> None: marker = self._completion_marker if marker is None: self._write_stdout(chunk) return with self._completion_lock: if self._completion_terminal.is_set(): self._write_stdout(chunk) return data = self._completion_pending + chunk marker_at = data.find(marker) if marker_at >= 0: self._write_stdout(data[:marker_at]) self._completion_pending = b"" self._completion_output = self._combined_buffer.text() self._completion_terminal.set() post = data[marker_at + len(marker) :] if post: self._write_stdout(post) return retained = 0 for size in range(min(len(data), len(marker) - 1), 0, -1): if data.endswith(marker[:size]): retained = size break self._write_stdout(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._write_stdout(self._completion_pending) self._completion_pending = b"" self._completion_output = self._combined_buffer.text() 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() with self._callback_lock: callback, self._reap_callback = self._reap_callback, None if callback is not None: callback() 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 _cleanup_input(self) -> None: if self._input_path: path, self._input_path = self._input_path, None try: os.unlink(path) except OSError: pass if getattr(self, "_fifo_dir", None): fdir, self._fifo_dir = self._fifo_dir, None try: os.rmdir(fdir) except OSError: pass
def __del__(self) -> None: self._cleanup_input()
def _finalize(self, exit_code: int, output: str | None = None) -> None: self._cleanup_input() with self._callback_lock: if self._done.is_set(): return combined_text = self._combined_buffer.text() if output is None else output stdout_text = self._stdout_buffer.text() stderr_text = self._stderr_buffer.text()
spill_threshold_bytes = int(os.environ.get("PRIME_AGENT_BASH_SPILL_BYTES", _DEFAULT_SPILL_BYTES)) spill_threshold_lines = int(os.environ.get("PRIME_AGENT_BASH_SPILL_LINES", _DEFAULT_SPILL_LINES)) total_bytes = len(combined_text.encode("utf-8", errors="replace")) line_count = combined_text.count("\n")
output_file: str | None = None final_output = combined_text final_stdout = stdout_text final_stderr = stderr_text if total_bytes > spill_threshold_bytes or line_count > spill_threshold_lines: output_file = _write_bash_spill_file(self._pid, combined_text) final_output = _truncate_spill_preview(combined_text, total_bytes, line_count, output_file) stdout_bytes = len(stdout_text.encode("utf-8", errors="replace")) if stdout_bytes > spill_threshold_bytes or stdout_text.count("\n") > spill_threshold_lines: final_stdout = _truncate_spill_preview(stdout_text, stdout_bytes, stdout_text.count("\n"), output_file) stderr_bytes = len(stderr_text.encode("utf-8", errors="replace")) if stderr_bytes > spill_threshold_bytes or stderr_text.count("\n") > spill_threshold_lines: final_stderr = _truncate_spill_preview(stderr_text, stderr_bytes, stderr_text.count("\n"), output_file)
result = BashResult( exit_code=exit_code, output=final_output, duration=time.monotonic() - self._started, stdout=final_stdout, stderr=final_stderr, output_file=output_file, ) self._result = result self._done.set() callbacks = self._callbacks self._callbacks = [] 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()
def _schedule_background_completion_notice(self) -> None: cell_finished = self._creating_cell_finished if cell_finished is None: return try: loop = asyncio.get_running_loop() except RuntimeError: return from . import repl
activity = {"id": secrets.token_hex(16), "pid": self._pid, "active": True} # Publish synchronously before bash() returns and the creating cell can end. repl.emit({"application/vnd.prime-agent.bash-activity+json": activity}) notice = self._notify_background_completion(cell_finished, activity) try: task = loop.create_task(notice) except BaseException: self.kill(signal.SIGKILL if _IS_POSIX else signal.SIGTERM) notice.close() repl.emit({"application/vnd.prime-agent.bash-activity+json": {**activity, "active": False}}) raise task.add_done_callback(_consume_notice_task)
async def _notify_background_completion( self, cell_finished: asyncio.Event, activity: dict[str, Any] ) -> None: from . import repl
try: result = await self._wait() await self._wait_reaped() # The cell may do other work before awaiting this handle. Do not classify # it as detached until that whole cell has crossed its completion barrier. await cell_finished.wait() if self._awaited_by_creating_cell or not repl.is_active(): return command = self.command if len(command) > _COMPLETION_NOTICE_COMMAND_CAP: command = command[:_COMPLETION_NOTICE_COMMAND_CAP] + "\n... [command truncated]" reply = await repl.host_request( { "type": "bash.completed", "pid": self._pid, "command": command, "exitCode": result.exit_code, } ) if not isinstance(reply, dict) or reply.get("status") != "ok": sys.stderr.write( f"Background bash completion follow-up for pid {self._pid} was not accepted. " "Inspect the saved handle with poll(), output(), or tail().\n" ) except (OSError, RuntimeError): # Standalone runtimes have no host handler, and teardown can close # the bridge while a process is finishing. Shell results stay usable. return finally: # Reap and deliver (or report rejection) before releasing kernel residency. repl.emit({"application/vnd.prime-agent.bash-activity+json": {**activity, "active": False}})
async def _wait_reaped(self) -> None: loop = asyncio.get_running_loop() future: asyncio.Future[None] = loop.create_future()
def wake() -> None: try: loop.call_soon_threadsafe(lambda: future.done() or future.set_result(None)) except RuntimeError: pass
with self._callback_lock: if self._reaped: return self._reap_callback = wake try: await future finally: with self._callback_lock: if self._reap_callback is wake: self._reap_callback = None
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: self._cleanup_input() # 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. try: current_task = asyncio.current_task() except RuntimeError: current_task = None creating_cell_waited = _creating_cell_waits_for(self._creating_cell_task, current_task) owned = not self._released wait = self._wait_owned() if owned else self._wait() self._released = True completed = False try: result = yield from wait.__await__() completed = True return result finally: if (completed or owned) and ( creating_cell_waited or _creating_cell_waits_for(self._creating_cell_task, current_task) ): self._awaited_by_creating_cell = True
def __repr__(self) -> str: 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, input: str | bytes | BashHandle | BashStream | None = None, *, display_command: str | None = None,) -> 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 input is not None and not isinstance(input, (str, bytes, bytearray, BashHandle, BashStream)): raise TypeError("input must be a str, bytes, BashHandle, BashStream, or None")
frame = inspect.currentframe() caller = frame.f_back if frame else None if not caller or not _is_exempt_skill_caller(caller): _check_command_guards(command)
_install_shutdown_hook() if input is not None or display_command is not None: return BashHandle(command, input=input, display_command=display_command) return BashHandle(command)
_EXEMPT_SKILL_MODULES = { "git", "jj", "search", "nix", "ssh_control", "repo",}
def _is_exempt_skill_caller(frame: Any) -> bool: f = frame while f: mod = f.f_globals.get("__name__", "") filename = f.f_code.co_filename or ""
# User cells, interactive input, or evaluated strings are never exempt if mod == "__main__" or filename.startswith("<cell-") or filename in ("<stdin>", "<string>"): return False
# If a private helper (like _run) was invoked directly by an interactive cell, it is not exempt if f.f_code.co_name.startswith("_"): caller = f.f_back if caller: c_fn = caller.f_code.co_filename or "" c_mod = caller.f_globals.get("__name__", "") if c_mod == "__main__" or c_fn.startswith("<cell-") or c_fn in ("<stdin>", "<string>"): return False
if mod in _EXEMPT_SKILL_MODULES or (mod and any(mod.startswith(m + ".") for m in _EXEMPT_SKILL_MODULES)): return True if "/skills/" in filename or "/prime-agent/skills/" in filename or "/dist/skills/" in filename: return True f = f.f_back return False
def _get_inner_cmd(tokens: list[str], idx: int) -> str: if len(tokens) == idx + 2: return tokens[idx + 1] return " ".join(tokens[idx + 1:])
def _check_command_guards(command: str) -> None: # 1. Check for subshells $(...) and `...` for sub in re.findall(r"\$\((.*?)\)|`([^`]+)`", command): sub_cmd = sub[0] or sub[1] _check_command_guards(sub_cmd)
try: s = shlex.shlex(command, posix=True, punctuation_chars=True) s.whitespace_split = True tokens = list(s) except ValueError: try: tokens = shlex.split(command) except ValueError: # Unbalanced quotes: inspect raw command segments so malformed syntax cannot bypass guards for seg in _COMMAND_SPLIT_RE.split(command): seg_s = seg.strip() if not seg_s: continue if _CD_SEGMENT_RE.match(seg_s): raise ValueError( "bash() rejects `cd` prefixes: move the working directory once " "with os.chdir(...) or repo.cd(...), then call bash() without the prefix." ) vcs_m = _VCS_SEGMENT_RE.match(seg_s) if vcs_m: raise ValueError(_format_vcs_error(seg_s, vcs_m.group(0).strip().lower())) search_m = _SEARCH_SEGMENT_RE.match(seg_s) if search_m: raise ValueError( f"bash() rejects `{search_m.group(0).strip().lower()}`: use search('pattern', path='...') " f"or search.files('*.ext') instead." ) find_m = _FIND_SEGMENT_RE.match(seg_s) if find_m: raise ValueError( f"bash() rejects `{find_m.group(0).strip().lower()}`: use search.files('pattern', path='...') " f"or Path.glob() instead." ) list_m = _LIST_SEGMENT_RE.match(seg_s) if list_m: raise ValueError( f"bash() rejects `{list_m.group(0).strip().lower()}`: use search.ls('path', pattern='...') " f"or os.listdir(...) instead." ) tokens = []
# Split top-level segments respecting quotes segments: list[list[str]] = [] current: list[str] = [] for tok in tokens: if tok in ("&&", "||", ";", "\n", "\r\n", "&"): if current: segments.append(current) current = [] else: current.append(tok) if current: segments.append(current)
for seg_tokens in segments: if not seg_tokens: continue
# Unpeel prefix wrappers (sudo, time, nohup, exec, command, builtin) while seg_tokens and seg_tokens[0].lower() in ("sudo", "time", "nohup", "exec", "command", "builtin"): seg_tokens = seg_tokens[1:]
# Unpeel env prefixes (e.g. env -i FOO=bar ...) if seg_tokens and seg_tokens[0].lower() == "env": seg_tokens = seg_tokens[1:] while seg_tokens: t = seg_tokens[0] if t.startswith("-") or ("=" in t and not t.startswith("=")): seg_tokens = seg_tokens[1:] else: break
if not seg_tokens: continue
raw_first = seg_tokens[0] cmd_name = Path(raw_first).name.lower()
# If it's a remote command (ssh, scp, sftp, mosh), skip local worktree guards! if cmd_name in _REMOTE_EXEC_TOKENS: continue
# If it's eval, recurse on the evaluated arguments if cmd_name == "eval": if len(seg_tokens) > 1: _check_command_guards(" ".join(seg_tokens[1:])) continue
# Check nix-shell if cmd_name == "nix-shell": for i, tok in enumerate(seg_tokens[1:], 1): if tok in ("--run", "--command") and i + 1 < len(seg_tokens): _check_command_guards(_get_inner_cmd(seg_tokens, i)) break raise ValueError( "bash() rejects `nix-shell`: use `await nix.develop(...)` or `await nix.dev_env(...)` instead." )
# Check nix if cmd_name == "nix": subcmd = None subcmd_idx = 1 for i, tok in enumerate(seg_tokens[1:], 1): if not tok.startswith("-"): subcmd = tok.lower() subcmd_idx = i break if subcmd == "develop": for i, tok in enumerate(seg_tokens[subcmd_idx + 1:], subcmd_idx + 1): if tok in ("--command", "-c") and i + 1 < len(seg_tokens): _check_command_guards(_get_inner_cmd(seg_tokens, i)) break raise ValueError( "bash() rejects `nix develop`: use `shell = await nix.develop(...)` and `shell.run(...)` or `await nix.dev_env(...)` instead." ) if subcmd == "build": raise ValueError("bash() rejects `nix build`: use `nix.build(...)` instead.") if subcmd == "run": raise ValueError("bash() rejects `nix run`: use `nix.run(...)` instead.") if subcmd == "fmt": raise ValueError("bash() rejects `nix fmt`: use `nix.fmt(...)` instead.") if subcmd == "flake": flake_sub = None for tok in seg_tokens[subcmd_idx + 1:]: if not tok.startswith("-"): flake_sub = tok.lower() break if flake_sub == "check": raise ValueError("bash() rejects `nix flake check`: use `nix.check(...)` instead.") if flake_sub == "show": raise ValueError("bash() rejects `nix flake show`: use `nix.show(...)` instead.") if subcmd == "eval": raise ValueError("bash() rejects `nix eval`: use `nix.eval(...)` instead.") if subcmd: raise ValueError( f"bash() rejects `nix {subcmd}`: use the Python `nix` skill (`nix.show`, `nix.build`, `nix.develop`, `nix.run`, `nix.fmt`) instead." ) raise ValueError("bash() rejects `nix`: use the Python `nix` skill instead.")
# If it's nushell (nu -c '...'), check the inner command if cmd_name == "nu": for i, tok in enumerate(seg_tokens[1:], 1): if tok in ("-c", "--commands") and i + 1 < len(seg_tokens): _check_command_guards(_get_inner_cmd(seg_tokens, i)) break continue
# If it's a shell wrapper (bash -c '...'), check the inner command if _SHELL_WRAPPER_TOKEN_RE.match(cmd_name): for i, tok in enumerate(seg_tokens[1:], 1): if tok in ("-c", "--command") or (tok.startswith("-") and "c" in tok and not tok.startswith("--")): if i + 1 < len(seg_tokens): inner_cmd = _get_inner_cmd(seg_tokens, i) _check_command_guards(inner_cmd) break continue
# Check for unquoted pipe operators if any(tok in ("|", "|&") for tok in seg_tokens): pipe_stages: list[list[str]] = [] cur_stage: list[str] = [] for tok in seg_tokens: if tok in ("|", "|&"): if cur_stage: pipe_stages.append(cur_stage) cur_stage = [] else: cur_stage.append(tok) if cur_stage: pipe_stages.append(cur_stage)
first_cmd = Path(pipe_stages[0][0]).name.lower() if pipe_stages and pipe_stages[0] else "" if first_cmd in ("ls", "find", "dir"): for stage in pipe_stages[1:]: stage_cmd = Path(stage[0]).name.lower() if stage else "" if stage_cmd in ("grep", "egrep", "fgrep", "rg", "ripgrep"): raise ValueError( f"bash() rejects `{first_cmd} | {stage_cmd}`: use search.files('pattern', path='...') " f"or ls('path', 'pattern') instead." ) raise ValueError( "bash() rejects pipes: pass input=... or use .pipe('...') in Python instead." )
# Check single segment seg_str = shlex.join(seg_tokens) if _CD_SEGMENT_RE.match(seg_str): raise ValueError( "bash() rejects `cd` prefixes: move the working directory once " "with os.chdir(...) or repo.cd(...), then call bash() without the prefix." ) if cmd_name in ("git", "jj"): raise ValueError(_format_vcs_error(seg_str, cmd_name)) if cmd_name in ("grep", "egrep", "fgrep", "rg", "ripgrep", "ag", "ack"): raise ValueError( f"bash() rejects `{cmd_name}`: use search('pattern', path='...') " f"or search.files('*.ext') instead." ) if cmd_name in ("find", "fd"): raise ValueError( f"bash() rejects `{cmd_name}`: use search.files('pattern', path='...') " f"or Path.glob() instead." ) if cmd_name in ("ls", "dir"): raise ValueError( f"bash() rejects `{cmd_name}`: use search.ls('path', pattern='...') " f"or os.listdir(...) instead." ) vcs_match = _VCS_SEGMENT_RE.match(seg_str) if vcs_match: cmd_matched = vcs_match.group(0).strip().lower() raise ValueError(_format_vcs_error(seg_str, cmd_matched)) search_match = _SEARCH_SEGMENT_RE.match(seg_str) if search_match: cmd_matched = search_match.group(0).strip().lower() raise ValueError( f"bash() rejects `{cmd_matched}`: use search('pattern', path='...') " f"or search.files('*.ext') instead." ) find_match = _FIND_SEGMENT_RE.match(seg_str) if find_match: cmd_matched = find_match.group(0).strip().lower() raise ValueError( f"bash() rejects `{cmd_matched}`: use search.files('pattern', path='...') " f"or Path.glob() instead." ) list_match = _LIST_SEGMENT_RE.match(seg_str) if list_match: cmd_matched = list_match.group(0).strip().lower() raise ValueError( f"bash() rejects `{cmd_matched}`: use search.ls('path', pattern='...') " f"or os.listdir(...) instead." )
_GIT_NATIVE_SUBCOMMANDS = { "status", "st", "log", "diff", "show", "grep", "commit", "add", "checkout", "switch", "branch", "branches", "fetch", "pull", "push", "merge", "reset", "restore", "stash", "tag", "blame", "root", "version",}
_JJ_NATIVE_SUBCOMMANDS = { "status", "st", "log", "diff", "show", "rebase", "squash", "absorb", "describe", "edit", "new", "abandon", "restore", "bookmark", "bookmarks", "workspace", "workspaces", "op", "file", "grep", "root", "version",}
def _format_vcs_error(sub_seg: str, vcs_name: str) -> str: tokens = sub_seg.strip().split() subcmd = None for tok in tokens[1:]: if not tok.startswith("-"): subcmd = tok.lower() break
if vcs_name == "git": if subcmd: if subcmd in ("status", "st"): return "bash() rejects `git status`: use git.status(...) instead." if subcmd == "log": return "bash() rejects `git log`: use git.log(...) instead." if subcmd == "diff": return "bash() rejects `git diff`: use git.diff(...) instead." if subcmd == "show": return "bash() rejects `git show`: use git.show(...) instead." if subcmd == "grep": return "bash() rejects `git grep`: use git.grep(...) instead." if subcmd == "commit": return "bash() rejects `git commit`: use git.commit(...) instead." if subcmd == "add": return "bash() rejects `git add`: use git.add(...) instead." if subcmd in ("checkout", "switch"): return f"bash() rejects `git {subcmd}`: use git.{subcmd}(...) instead." if subcmd in ("branch", "branches"): return f"bash() rejects `git {subcmd}`: use git.branch(...) or git.branches(...) instead." if subcmd in ("fetch", "pull", "push", "merge", "reset", "restore", "stash", "tag"): return f"bash() rejects `git {subcmd}`: use git.{subcmd}(...) instead." if subcmd == "blame": return "bash() rejects `git blame`: use git.blame(path, rev=...) instead." if subcmd == "root": return "bash() rejects `git root`: use git.root() instead." if subcmd == "version": return "bash() rejects `git version`: use git.version() instead." return f"bash() rejects `git {subcmd}`: use git.run('{subcmd}', ...) instead." return "bash() rejects `git`: use git.status(...) instead."
if vcs_name == "jj": if subcmd: if subcmd in ("status", "st"): return f"bash() rejects `jj {subcmd}`: use jj.status(...) instead." if subcmd == "log": return "bash() rejects `jj log`: use jj.log(...) or jj.graph() instead." if subcmd == "diff": return "bash() rejects `jj diff`: use jj.diff(...) instead." if subcmd == "show": return "bash() rejects `jj show`: use jj.show(...) instead." if subcmd in ("rebase", "squash", "absorb", "describe", "edit", "new", "abandon", "restore"): return f"bash() rejects `jj {subcmd}`: use jj.{subcmd}(...) instead." if subcmd in ("bookmark", "bookmarks"): return f"bash() rejects `jj {subcmd}`: use jj.bookmark_set(...) or jj.bookmarks(...) instead." if subcmd in ("workspace", "workspaces"): return f"bash() rejects `jj {subcmd}`: use jj.workspaces(), jj.workspace_add(...), or jj.workspace_forget(...) instead." if subcmd == "op": return "bash() rejects `jj op`: use jj.op_log(...) instead." if subcmd == "file": return "bash() rejects `jj file`: use jj.show(rev, path) or jj.files(rev) instead." if subcmd == "grep": return "bash() rejects `jj grep`: use jj.grep(...) instead." if subcmd == "root": return "bash() rejects `jj root`: use jj.root() instead." if subcmd == "version": return "bash() rejects `jj version`: use jj.version() instead." return f"bash() rejects `jj {subcmd}`: use jj.run('{subcmd}', ...) instead." return "bash() rejects `jj`: use jj.status(...) instead."
return f"bash() rejects `{vcs_name}`: use {vcs_name}.status() instead."
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, input_file: str | None = None) -> str: # Closed control fds preserve background behavior; supported shells atomically write the frame. emit = _fence_printf() stdin_redir = f"0<{shlex.quote(input_file)}" if input_file else "0</dev/null" return ( f"exec {_STATUS_FD}>&0 {_OUTPUT_FD}>&1 {stdin_redir}\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)