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.
120 kB · 2916 lines
Python
at main
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917"""Async-by-default shell execution: bash() spawns immediately and returns a live handle."""
from __future__ import annotations
import asyncioimport atexitimport functoolsimport inspectimport jsonimport mathimport 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 _winjobfrom .host_safety import ( HostProfile, ProcessLease, adopt_process, check_power_control, check_vram_command, clear_owned_processes, get_host_profile, get_owned_process, is_pid_owned, kill_process, list_owned_pids, list_owned_processes, mark_owned_process_reaped, preflight_vram, register_owned_process, reset_host_profile, set_host_profile, signal_process,)
_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"}
# A kill is safe exactly when every target sits inside the command's own process# group: the command's shell, a child of that shell, or that group itself._KILL_SEGMENT_TOKENS = frozenset({"kill", "pkill", "killall"})_KILL_SEGMENT_RE = re.compile(r"^\s*(?:kill|pkill|killall)(?:\s|$)", re.IGNORECASE)_SELF_KILL_TARGET_RE = re.compile(r"^-?(?:\$\$|\$!|%\d+|%[A-Za-z][A-Za-z0-9_-]*)$")_KILL_GROUP_PID_RE = re.compile(r"^-\d{2,}$")# pkill matches by name or pattern, so only an explicit parent scope proves the# match stays under this command's own shell._KILL_PARENT_FLAGS = ("-P", "--parent")_XARGS_VALUE_FLAGS = frozenset({ "-a", "-d", "-E", "-I", "-L", "-n", "-P", "-s", "--arg-file", "--delimiter", "--eof", "--replace", "--max-lines", "--max-args", "--max-procs", "--max-chars",})# ssh options that take a separate value, so the destination stays findable._SSH_VALUE_FLAGS = frozenset({ "-B", "-b", "-c", "-D", "-E", "-e", "-F", "-I", "-i", "-J", "-L", "-l", "-m", "-O", "-o", "-p", "-Q", "-R", "-S", "-W", "-w",})# ssh and mosh carry a remote command; scp and sftp carry paths._REMOTE_COMMAND_TOKENS = frozenset({"ssh", "mosh"})# Pipeline stages that only read their input and write stdout. A stage outside# this set is refused: the guard cannot see what it writes or what it runs._READ_ONLY_FILTER_TOKENS = frozenset({ "cat", "column", "cut", "grep", "egrep", "fgrep", "rg", "ripgrep", "ag", "ack", "head", "tail", "jq", "nl", "od", "rev", "sort", "tac", "tr", "uniq", "wc", "xxd",})
# Network perimeter. A dev or admin server bound to a public interface answers the# whole network (01a09bf6 routed a private admin endpoint through a public worker),# so that bind is refused with the loopback form and one explicit opt-in. A hardcoded# dev port only warns: a single-run server on 5173 is legitimate, but 01a09351 lost# two parallel agents to EADDRINUSE on the same port._PUBLIC_BIND_HOSTS = frozenset({"0.0.0.0", "::", "*"})_SHORT_BIND_FLAGS = frozenset({"-H", "-b", "-p"})_BIND_FLAG_RE = re.compile( r"^--?[a-z0-9.-]*(?:host|address|addr|bind|listen|interface|ip|port|publish)[a-z0-9.-]*$", re.IGNORECASE,)# PORT and *_PORT (never EXPORT); HOST/HOSTNAME/ADDR/... and *_HOST/... ._PORT_ENV_RE = re.compile(r"^(?:PORT|[A-Z0-9]+_PORT)$", re.IGNORECASE)_HOST_ENV_RE = re.compile( r"^(?:HOST|HOSTNAME|ADDRESS|ADDR|BIND|LISTEN|IP|[A-Z0-9]+_(?:HOST|HOSTNAME|ADDRESS|ADDR|BIND|LISTEN|IP))$", re.IGNORECASE,)# A bare `host:port`: no scheme and no path, so a URL the command merely reads is not# mistaken for an endpoint it binds._ENDPOINT_TOKEN_RE = re.compile(r"^[A-Za-z0-9_.\-\[\]:]+:\d{1,5}$")_PUBLIC_LISTEN_OPT_IN_RE = re.compile( r"(?:^|[\s;&|(])RLM_ALLOW_PUBLIC_LISTEN=(?:1|true|yes|on)\b", re.IGNORECASE)# Dev servers two agents collide on in one worktree: vite (5173-5178 and the 24678# hmr port), next/express (3000), astro (4321), storybook (6006), http.server (8000),# expo (19000, 19006), parcel (1234), tauri (1420), plus the hand-rolled 3001-3003,# 4000, 4200, 5000-5001, 8080-8081, 8888 and 9000._DEV_SERVER_PORTS = frozenset({ 1234, 1420, 3000, 3001, 3002, 3003, 4000, 4200, 4321, 5000, 5001, 5173, 5174, 5175, 5176, 5177, 5178, 6006, 8000, 8080, 8081, 8888, 9000, 19000, 19006, 24678,})
_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 _start_background_completion_task( handle: BashHandle, loop: asyncio.AbstractEventLoop, notice: Any, activity: dict[str, Any],) -> None: """Create a bash completion notice task on the loop's own thread.
A handle may be started from an executor thread, which has no loop to create a task on; this runs there instead, and mirrors the inline construction path's rollback when the task cannot be created. """ from . import repl
try: task = loop.create_task(notice) except BaseException: handle.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)
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)
def _live_cell_owner() -> asyncio.Task[Any] | None: """Body task of the cell executing right now, ignoring detached context copies.""" try: from . import repl
if repl.is_active(): return repl.active_cell_task() except (ImportError, RuntimeError): pass return None
@dataclass(frozen=True)class BashResult: exit_code: int _output: str duration: float stdout: str = "" stderr: str = "" output_file: str | None = None timed_out: bool = False stdin_delivered: int | None = None stdin_expected: int | None = None
def text(self) -> str: """Return the command output text.""" return self._output
@property def stdin_bytes(self) -> int | None: """Alias for stdin_delivered.""" return self.stdin_delivered
@property def stdin_unconsumed(self) -> int | None: """Bytes the command never read, when delivery was accounted for.""" if self.stdin_expected is None or self.stdin_delivered is None: return None return max(0, self.stdin_expected - self.stdin_delivered)
@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_text(self) -> str: """Return the complete unclipped text, 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 full_output(self) -> str: """Alias to full_text().""" return self.full_text()
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 "" stdin_info = f" stdin={self.stdin_delivered}/{self.stdin_expected}" if self.stdin_expected is not None else "" if byte_count <= 200 and line_count <= 4: return f"<BashResult exit_code={self.exit_code} duration={dur_str}{file_info}{stdin_info} text={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}{stdin_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 total_bytes(self) -> int: with self._lock: return len(self._head) + self._tail_size + self._dropped
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/ .text()/.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, timeout: float | None = None, ) -> None: self._input_path: str | None = None self._fifo_dir: str | None = None self._input_expected: int | None = None self._input_delivered: int | None = None self._input_error: str | None = None self._is_fifo: bool = False self._source_handle: BashHandle | None = None self._source_channel: 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
self._is_fifo = True self._source_handle = source_handle self._source_channel = 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) self._input_expected = len(data) 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) self._input_expected = len(data) self._is_fifo = False 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._result_consumed = False self._consumed_notice: 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() # Deadline, in seconds: the whole command is killed like a cancelled await, # because the cell deadline cannot tell an honest long build from a wedge. self._timeout = timeout self._timed_out = False self._deadline_timer: threading.Timer | None = None # 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, is_fifo=self._is_fifo, ) 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 register_owned_process(self._pid, self.command, handle=self) 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() if self._timeout is not None: # Armed last, so a spawn failure can never leave a live timer behind. self._deadline_timer = threading.Timer(self._timeout, self._expire) self._deadline_timer.daemon = True self._deadline_timer.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 _get_total_bytes(self, channel: str) -> int: if channel == "stdout": return self._stdout_buffer.total_bytes() elif channel == "stderr": return self._stderr_buffer.total_bytes() else: return self._combined_buffer.total_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 | None] = queue.Queue()
def listener(ch: str, chunk: bytes) -> None: if channel == "combined" or channel == ch: if chunk: q.put(chunk) else: q.put(None)
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)
written = 0 write_err: str | None = None try: with open(fifo_path, "wb") as writer: if initial_bytes: writer.write(initial_bytes) writer.flush() written += len(initial_bytes) if not already_done: while True: chunk = q.get() if chunk is None: break writer.write(chunk) writer.flush() written += len(chunk) except BrokenPipeError: write_err = f"broken pipe after {written} bytes" except OSError as e: write_err = f"write error after {written} bytes: {e}" finally: self._input_delivered = written if write_err: self._input_error = write_err 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 text(self) -> str: """Single uniform text accessor for BashHandle, BashStream, and BashResult.""" self._released = True self._note_result_consumed() return self._buffer.text()
def __str__(self) -> str: return self.text()
def tail(self, n: int = 50) -> str: self._released = True self._note_result_consumed() return "\n".join(self._buffer.text().splitlines()[-n:])
def poll(self) -> BashResult | None: self._released = True self._note_result_consumed() 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 _expire(self) -> None: # The deadline path: same escalation as a cancelled `await bash(cmd)` (TERM, # bounded grace, group KILL), but it also has to reach a handle nobody is # awaiting, which is exactly the detached long-work path. if self._done.is_set(): return self._timed_out = True self.kill(grace=_CANCEL_TERM_GRACE)
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 unconsumed: int = 0 try: status, unconsumed = self._read_status() if self._input_expected is not None and not self._is_fifo: self._input_delivered = max(0, self._input_expected - unconsumed) # 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(): if self._input_delivered is None and self._input_expected is not None: self._input_delivered = 0 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: mark_owned_process_reaped(self._pid) _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) -> tuple[int | None, int]: if self._status_read < 0: return None, 0 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 parts = line.strip().split() if not parts: return None, 0 code = int(parts[0]) unconsumed = int(parts[1]) if len(parts) > 1 else 0 return code, unconsumed except (OSError, ValueError): return None, 0 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()
# Stdin delivery accounting: delivered = self._input_delivered expected = self._input_expected if self._is_fifo and self._source_handle is not None and self._source_channel is not None: expected = self._source_handle._get_total_bytes(self._source_channel) self._input_expected = expected
if expected is not None: if delivered is None: delivered = expected if not _IS_POSIX else 0 self._input_delivered = delivered # A consumer that reads part of its input and exits is legitimate (`head -1`, # `grep -q`), so the failure is delivery, not consumption: a write error, or # nothing consumed at all while bytes were expected. Partial consumption is # reported through stdin_delivered/stdin_expected instead of failing the call. if self._input_error is not None or (delivered == 0 and expected > 0): err_suffix = f": {self._input_error}" if self._input_error else "" msg = ( f"stdin delivery incomplete: delivered {delivered} of {expected} bytes{err_suffix}\n" ) if exit_code == 0: exit_code = 1 if msg not in stderr_text: stderr_text = f"{msg}{stderr_text}" if stderr_text else msg if msg not in combined_text: combined_text = f"{msg}{combined_text}" if combined_text else msg
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, timed_out=self._timed_out, stdin_delivered=delivered, stdin_expected=expected, ) self._result = result self._done.set() if self._deadline_timer is not None: self._deadline_timer.cancel() 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 _note_result_consumed(self, awaiter: asyncio.Task[Any] | None = None) -> None: """Record a result read that reaches the model: only reads during a live cell count (a detached reader between turns must keep the notice — it is the idle session's only wake-up), and an awaiting reader must be one the live cell waits for.""" if not self._done.is_set(): return owner = _live_cell_owner() if owner is None: return if awaiter is not None and not _creating_cell_waits_for(owner, awaiter): return with self._callback_lock: if self._result_consumed: return self._result_consumed = True notice, self._consumed_notice = self._consumed_notice, None if notice is not None: notice()
def _schedule_background_completion_notice(self) -> None: cell_finished = self._creating_cell_finished if cell_finished is None: return from . import repl
try: running = asyncio.get_running_loop() except RuntimeError: running = None # A creating thread without a loop of its own (an executor worker) still owns # this notice: the task belongs on the serving loop, which owns the cell. serving = getattr(repl, "serving_loop", None) loop = running if running is not None else (serving() if serving is not None else None) if loop is None: return 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) if running is None: try: # Only the loop's own thread may create a task on it. loop.call_soon_threadsafe(_start_background_completion_task, self, loop, notice, activity) except RuntimeError: # Loop already closed: the notice is undeliverable, but the activity # event that was already published must not look permanently active. notice.close() repl.emit({"application/vnd.prime-agent.bash-activity+json": {**activity, "active": False}}) return 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 self._result_consumed 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 isinstance(reply, dict) and reply.get("status") == "ok": # Notice accepted by the host; later reads must ask it to withdraw. self._arm_consumed_notice(command) else: 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}})
def _arm_consumed_notice(self, command: str) -> None: # Armed only post-acceptance: the withdrawal can never overtake its notice. loop = asyncio.get_running_loop()
def dispatch() -> None: def start() -> None: task = loop.create_task(self._notify_result_consumed(command)) task.add_done_callback(_consume_notice_task)
try: loop.call_soon_threadsafe(start) except RuntimeError: pass # notifying loop already closed
with self._callback_lock: if not self._result_consumed: self._consumed_notice = dispatch return dispatch()
async def _notify_result_consumed(self, command: str) -> None: from . import repl
if not repl.is_active(): return try: await repl.host_request( {"type": "bash.consumed", "pid": self._pid, "command": command} ) except (OSError, RuntimeError): return # bridge closed at teardown; old hosts error-reply — both fine
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: mark_owned_process_reaped(self._pid) _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 if completed: self._note_result_consumed(current_task)
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, timeout: float | 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/.text()/.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 text but stays visible via handle.text()/tail().
`timeout` (seconds) arms a deadline at spawn: the command's process group is killed the same way a cancelled await kills it, and BashResult.timed_out says why the result is a kill rather than an exit. It bounds one command, not the cell, and it is the only bound that reaches a detached `h = bash(cmd)` handle, which no cell deadline can interrupt. """ 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") if timeout is not None: if isinstance(timeout, bool) or not isinstance(timeout, (int, float)): raise TypeError("timeout must be a number of seconds or None") if not math.isfinite(timeout) or timeout <= 0: raise ValueError("timeout must be positive and finite")
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, timeout=timeout) return BashHandle(command, timeout=timeout)
_EXEMPT_SKILL_MODULES = { "git", "jj", "search", "nix", "ssh", "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 _tokenize_command(command: str) -> list[str] | None: """Shell tokens for guard inspection, or None when the quoting is unbalanced.""" try: lexer = shlex.shlex(command, posix=True, punctuation_chars=True) lexer.whitespace_split = True return list(lexer) except ValueError: try: return shlex.split(command) except ValueError: return None
def _segment_tokens(tokens: list[str]) -> list[list[str]]: """Top-level segments, split on the operators that start a new command.""" 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) return segments
def _subshell_commands(command: str) -> list[str]: return [match.group(1) or match.group(2) for match in re.finditer(r"\$\((.*?)\)|`([^`]+)`", command)]
def _pipe_stages(tokens: list[str]) -> list[list[str]]: stages: list[list[str]] = [] current: list[str] = [] for tok in tokens: if tok in ("|", "|&"): stages.append(current) current = [] else: current.append(tok) stages.append(current) return stages
def _strip_wrappers(tokens: list[str]) -> list[str]: """Drop the prefixes that do not change which command runs.""" while tokens and tokens[0].lower() in ("sudo", "time", "nohup", "exec", "command", "builtin"): tokens = tokens[1:] while tokens and "=" in tokens[0] and not tokens[0].startswith(("-", "=")): tokens = tokens[1:] # `VAR=value cmd` sets the environment for cmd if tokens and tokens[0].lower() == "env": tokens = tokens[1:] while tokens: tok = tokens[0] if tok.startswith("-") or ("=" in tok and not tok.startswith("=")): tokens = tokens[1:] else: break return tokens
def _kill_targets(words: list[str]) -> list[str]: """What a `kill` names: everything that is not a signal spec or a flag.""" targets: list[str] = [] index = 1 while index < len(words): tok = words[index] if tok == "--": index += 1 continue if tok in ("-s", "--signal"): index += 2 continue if tok.startswith("-") and not (_KILL_GROUP_PID_RE.match(tok) or _SELF_KILL_TARGET_RE.match(tok)): index += 1 # a signal spec (`-9`, `-KILL`) or another flag continue targets.append(tok) index += 1 return targets
def _kill_is_parent_scoped_to_self(words: list[str]) -> bool: """True for `pkill -P $$`, whose matches are children of this command's shell.""" self_targets = ("$$", "$!") for index, tok in enumerate(words): if tok in _KILL_PARENT_FLAGS: value = words[index + 1] if index + 1 < len(words) else "" elif tok.startswith("--parent="): value = tok.split("=", 1)[1] elif tok.startswith("-P"): value = tok[2:] else: continue return value in self_targets return False
def _unsafe_kill_error(command: str, reason: str, *, remote: bool) -> str: transport = ( " It runs on a remote host over the ssh transport, where a match that reaches the ssh " "wrapper or its ControlMaster socket severs this kernel's transport." if remote else "" ) return ( f"bash() rejects `{command}`: {reason}.{transport} Kill only what this command owns: " "`h = bash(cmd); h.kill()`, scoped `kill_process(pid)`, or adopt via `rlm.adopt_process(pid, reason=...)` " "for a command you started, `kill -TERM $!` or `pkill -P $$ name` for a child of this command, or `kill -- -$$` for this command's group." )
def _reject_unsafe_kill(words: list[str], *, remote: bool) -> None: """Raise unless every target provably sits inside the command's own process group.""" command = " ".join(words) name = Path(words[0]).name.lower() if any(tok in ("-l", "-L", "--list", "--table") for tok in words[1:]): return # listing signal names signals nothing if name == "pkill": if _kill_is_parent_scoped_to_self(words): return raise ValueError( _unsafe_kill_error( command, "`pkill` matches by name or pattern, so it can signal anything on the machine, " "not just this command's own process group", remote=remote, ) ) if name == "killall": raise ValueError( _unsafe_kill_error( command, "`killall` matches by name, so it can signal anything on the machine, " "not just this command's own process group", remote=remote, ) ) targets = _kill_targets(words) if not targets: raise ValueError( _unsafe_kill_error( command, "`kill` names no target inside this command's own process group", remote=remote ) ) outside = next((tok for tok in targets if not _SELF_KILL_TARGET_RE.match(tok)), None) if outside is not None: raise ValueError( _unsafe_kill_error(command, f"target `{outside}` is not this command's own process group", remote=remote) )
def _strip_timeout(tokens: list[str]) -> list[str]: """The command `timeout` runs, with its own flags and duration removed.""" index = 1 while index < len(tokens): tok = tokens[index] if tok in ("-s", "--signal", "-k", "--kill-after"): index += 2 continue if tok.startswith("-"): index += 1 continue index += 1 # the duration break return tokens[index:]
def _strip_xargs(tokens: list[str]) -> list[str]: """The command `xargs` runs, with xargs' own flags removed.""" index = 1 while index < len(tokens): tok = tokens[index] if not tok.startswith("-"): break index += 1 if tok in _XARGS_VALUE_FLAGS: index += 1 return tokens[index:]
def _remote_command_payload(words: list[str]) -> str | None: """The command an ssh/mosh invocation runs on the remote host, if it has one.""" index = 1 while index < len(words): tok = words[index] if tok in _SSH_VALUE_FLAGS: index += 2 continue if tok.startswith("-"): index += 1 continue break payload = words[index + 1:] return " ".join(payload) if payload else None
def _scan_unsafe_kills(command: str, *, remote: bool = False) -> None: """Refuse a kill that can reach outside the command's own process group.
Every bash() command runs in its own process group (start_new_session), so the only targets that are provably safe are that group's members: `$$` (its leader), `$!` and `%N` (children of its shell). In 01a09d7f a broad `pkill` during cleanup matched the agent's own ssh wrapper and ControlMaster sockets, severed the transport, and wedged the kernel for 92 minutes. """ tokens = _tokenize_command(command) if tokens is None: # Unbalanced quotes: the token pass cannot run, so refuse from the raw # segments, like the other fallbacks in _check_command_guards. for segment in _COMMAND_SPLIT_RE.split(command): if _KILL_SEGMENT_RE.match(segment): raise ValueError( _unsafe_kill_error( segment.strip(), "its quoting is unbalanced, so its target cannot be checked", remote=remote ) ) return for seg_tokens in _segment_tokens(tokens): words = _strip_wrappers(seg_tokens) if not words: continue name = Path(words[0]).name.lower() if name in _KILL_SEGMENT_TOKENS: _reject_unsafe_kill(words, remote=remote) continue if name == "timeout": remainder = _strip_timeout(words) if remainder: _scan_unsafe_kills(" ".join(remainder), remote=remote) continue if name in _REMOTE_COMMAND_TOKENS: payload = _remote_command_payload(words) if payload: _scan_unsafe_kills(payload, remote=True) continue for stage in _pipe_stages(words): stage_words = _strip_wrappers(stage) if not stage_words: continue stage_name = Path(stage_words[0]).name.lower() if stage_name in _KILL_SEGMENT_TOKENS: _reject_unsafe_kill(stage_words, remote=remote) continue if stage_name == "xargs": xargs_command = _strip_xargs(stage_words) if xargs_command and Path(xargs_command[0]).name.lower() in _KILL_SEGMENT_TOKENS: _reject_unsafe_kill(xargs_command, remote=remote) # Subshells last: a kill inside a remote payload then reports the transport # it would sever, not only the local rule. for sub_cmd in _subshell_commands(command): _scan_unsafe_kills(sub_cmd, remote=remote)
def _stage_writes_a_file(words: list[str]) -> str | None: """The token by which a pipeline stage writes a file, if it does.""" if Path(words[0]).name.lower() == "sort": if any(tok == "-o" or tok.startswith(("-o", "--output")) for tok in words[1:]): return "-o" for index, tok in enumerate(words): if tok in ("&>", ">>", ">|", ">"): return tok if tok == ">&" and not (index + 1 < len(words) and words[index + 1].isdigit()): return tok return None
def _pipeline_error(command: str, stage: int, reason: str) -> str: return ( f"bash() rejects `{command}`: stage {stage} {reason}. A pipeline may only end in read-only " "stream filters (cat, head, tail, grep, rg, wc, jq, sort, uniq, ...). Read the output in " "Python instead: `h = bash(cmd)`, then `await h` and work on `h.text()`, or feed one " "stage with `input=h` / `h.stdout.pipe('...')`." )
def _check_pipeline(seg_tokens: list[str], perimeter: _Perimeter) -> None: """Permit a pipeline whose every stage only reads its input.
The head stage keeps every single-command rule; later stages must be read-only stream filters; no stage may write a file. A pipeline cannot be inspected partially, so one bad stage refuses the whole command. """ stages = _pipe_stages(seg_tokens) head = _strip_wrappers(stages[0]) if head: head_name = Path(head[0]).name.lower() if head_name in ("ls", "find", "dir"): for stage in stages[1:]: stage_words = _strip_wrappers(stage) filter_name = Path(stage_words[0]).name.lower() if stage_words else "" if filter_name in ("grep", "egrep", "fgrep", "rg", "ripgrep"): raise ValueError( f"bash() rejects `{head_name} | {filter_name}`: use search.files('pattern', path='...') " "or search.ls('path', pattern='...') instead." ) command = " ".join(seg_tokens) for index, stage in enumerate(stages, 1): words = _strip_wrappers(stage) if not words: raise ValueError(_pipeline_error(command, index, "is empty")) write_token = _stage_writes_a_file(words) if write_token is not None: raise ValueError(_pipeline_error(command, index, f"writes a file (`{write_token}`)")) if index == 1: _check_segment(words, perimeter) continue name = Path(words[0]).name.lower() if name in _KILL_SEGMENT_TOKENS: _reject_unsafe_kill(words, remote=False) if name in _READ_ONLY_FILTER_TOKENS: continue _check_segment(words, perimeter) raise ValueError(_pipeline_error(command, index, f"runs `{name}`, which is not a read-only stream filter"))
@dataclass(frozen=True)class _Endpoint: """One endpoint a command names, as written: `host`, `port`, or `host:port`."""
text: str host: str | None port: int | None
class _Perimeter: """Per-command network state: the explicit public opt-in and the warned ports."""
def __init__(self, command: str) -> None: self.public_allowed = bool(_PUBLIC_LISTEN_OPT_IN_RE.search(command)) self.warned: set[int] = set()
def _host_key(host: str) -> str: """A bind host as written, without quoting or IPv6 brackets.""" return host.strip().strip("'\"").removeprefix("[").removesuffix("]").lower()
def _port_number(value: str) -> int | None: return int(value) if value.isdigit() and len(value) <= 5 and int(value) < 65536 else None
def _endpoint_host(value: str) -> str | None: """The host part of `host:port`, `[host]:port`, `host:port:port`, or a bare host.""" if value.startswith("["): closing = value.find("]") return value[: closing + 1] if closing > 0 else None if ":" not in value: return value or None if value.startswith("::"): return value # a bare IPv6 host: `::` and `::1` are different perimeters return value.split(":", 1)[0] or None
def _endpoint(value: str) -> _Endpoint | None: text = value.strip().strip("'\"") if not text: return None if text.isdigit(): return _Endpoint(text=text, host=None, port=_port_number(text)) if ":" not in text: return _Endpoint(text=text, host=text, port=None) return _Endpoint( text=text, host=_endpoint_host(text), port=_port_number(text.rpartition(":")[2]) )
def _is_bind_flag(token: str) -> bool: """True for a flag whose value names a listening endpoint.""" return token in _SHORT_BIND_FLAGS or bool(_BIND_FLAG_RE.match(token))
def _segment_endpoints(tokens: list[str]) -> list[_Endpoint]: """Endpoints this segment names: bind-flag values, HOST/PORT assignments, bare host:port.""" endpoints: list[_Endpoint] = [] for index, token in enumerate(tokens): flag, separator, inline = token.partition("=") if _is_bind_flag(flag): value = inline if separator else (tokens[index + 1] if index + 1 < len(tokens) else "") endpoint = _endpoint(value) if endpoint is not None: endpoints.append(endpoint) continue if not separator: if _ENDPOINT_TOKEN_RE.match(token): endpoint = _endpoint(token) if endpoint is not None: endpoints.append(endpoint) continue if token.startswith("-"): continue if _PORT_ENV_RE.match(flag): endpoints.append(_Endpoint(text=token, host=None, port=_port_number(inline))) elif _HOST_ENV_RE.match(flag): endpoints.append(_Endpoint(text=token, host=inline or None, port=None)) return endpoints
def _check_network_perimeter(command: str, tokens: list[str], perimeter: _Perimeter) -> None: """Refuse a public bind; warn about a hardcoded dev port.
Only the public bind is refused: a bind the user wants public says so with the opt-in prefix, and the dev port is allowed because one server on 5173 is ordinary. """ for endpoint in _segment_endpoints(tokens): if endpoint.host is not None and _host_key(endpoint.host) in _PUBLIC_BIND_HOSTS: if perimeter.public_allowed: continue raise ValueError(_public_bind_error(command, endpoint.text)) port = endpoint.port if port is not None and port in _DEV_SERVER_PORTS and port not in perimeter.warned: perimeter.warned.add(port) _warn_advisory(_dev_port_warning(command, port))
def _public_bind_error(command: str, endpoint: str) -> str: return ( f"bash() rejects `{command}`: `{endpoint}` listens on a public interface. Bind loopback " "instead (`--host 127.0.0.1`, `HOST=127.0.0.1 cmd`, or `-p 127.0.0.1:3000:3000`); a dev or " "admin server on 0.0.0.0 answers the whole network, so a private route reaches the " "internet. If the public bind is deliberate, prefix the command with " "`RLM_ALLOW_PUBLIC_LISTEN=1`." )
def _dev_port_warning(command: str, port: int) -> str: return ( f"bash() warns `{command}`: port {port} is a well-known dev port, so a second server in " "this worktree collides with it (`EADDRINUSE`). Take an ephemeral port instead " "(`PORT=0`, or `--port 0`) and read the port the server prints." )
def _warn_advisory(message: str) -> None: """Advisory output; the cell's own stderr reaches the model in the same tool result.""" try: print(message, file=sys.stderr) except Exception: # noqa: BLE001 - a warning must never break the command pass
def _check_command_guards(command: str, perimeter: _Perimeter | None = None) -> None: # Kill-family commands are checked separately: they must be provably # confined to this command's own process group. _scan_unsafe_kills(command)
# The network perimeter is one state per command: the opt-in and the warned # ports are read from the whole command text and shared with nested segments. if perimeter is None: perimeter = _Perimeter(command)
# 1. Check for subshells $(...) and `...` for sub_cmd in _subshell_commands(command): _check_command_guards(sub_cmd, perimeter)
tokens = _tokenize_command(command) if tokens is None: # 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 = []
for seg_tokens in _segment_tokens(tokens): _check_segment(seg_tokens, perimeter)
def _check_segment(seg_tokens: list[str], perimeter: _Perimeter) -> None: raw_seg_str = " ".join(seg_tokens) _check_network_perimeter(raw_seg_str, seg_tokens, perimeter) seg_tokens = _strip_wrappers(seg_tokens) if not seg_tokens: return
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: return
# Check power control and desktop VRAM safety check_power_control(raw_seg_str, seg_tokens) check_vram_command(raw_seg_str, seg_tokens)
# 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:]), perimeter) return
# 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), perimeter) 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), perimeter) 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), perimeter) break return
# 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, perimeter) break return
# Pipelines: every stage must only read its input (see _check_pipeline) if any(tok in ("|", "|&") for tok in seg_tokens): _check_pipeline(seg_tokens, perimeter) return
# 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. A child that needs its own checkout gets one with " "rlm.spawn(..., workspace='isolated')." ) 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=_default_utility_path()) 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, is_fifo: bool = False,) -> 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" if input_file and not is_fifo: unconsumed_check = "__prime_unconsumed=$(cat | wc -c)\n" else: unconsumed_check = "__prime_unconsumed=0\n"
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"{unconsumed_check}" f"{emit} '\\036prime-agent-complete:%s%s\\037' " f"'{completion_a}' '{completion_b}' >&{_OUTPUT_FD} || exit \"$__prime_status\"\n" f"{emit} '%s %s\\n' \"$__prime_status\" \"$__prime_unconsumed\" >&{_STATUS_FD}\n" f"exec {_OUTPUT_FD}>&- {_STATUS_FD}>&-\n" "wait\n" 'exit "$__prime_status"\n' )
def _default_utility_path() -> str: """The POSIX default utility path: portable, never a platform-shaped guess.""" confstr = getattr(os, "confstr", None) if confstr is not None: try: return confstr("CS_PATH") or os.defpath except (ValueError, OSError): pass return os.defpath
# Co-located caches and ephemeral workspaces. Parallel sessions of one# repository share a build cache and, when a caller asks for isolation, an# ephemeral checkout; both live inside the repository, not in the system temp# directory. A checkout that loses its repository (a /tmp tree, a plain copy)# makes Nix copy a store path per build, and a per-session build cache makes# every session recompile from cold._WORKTREES_DIRNAME = ".worktrees"_SHARED_CACHE_RELATIVE = Path("target") / "subagents-shared"_SHARED_CACHE_ENV = "CARGO_TARGET_DIR"_SHARED_CACHE_SWITCH = "PRIME_AGENT_SHARED_CACHE"_SHARED_CACHE_DIR_ENV = "PRIME_AGENT_SHARED_CACHE_DIR"_EPHEMERAL_IGNORE_MARKER = "# prime-agent ephemeral workspaces"_VCS_COMMAND_TIMEOUT = 120.0
def _env_flag(name: str, default: bool) -> bool: """Read a boolean switch; off for the usual falsy spellings and for unset.""" raw = os.environ.get(name) if raw is None or not raw.strip(): return default return raw.strip().lower() not in {"0", "false", "no", "off"}
def _repo_root(cwd: str | os.PathLike[str] | None = None) -> Path | None: """The nearest directory holding a `.jj` or `.git` entry, or None.
Resolution walks up from the current directory only: no glob, no walk of a tree that could be a multi-terabyte home. """ try: start = (Path(cwd) if cwd is not None else Path(os.getcwd())).resolve() except OSError: return None for directory in (start, *start.parents): if (directory / ".jj").exists() or (directory / ".git").exists(): return directory return None
def _shared_repo_root(root: Path) -> Path: """The main working copy `root` belongs to, through a linked checkout.
`git worktree add` writes a `.git` file naming the linked git dir, and `jj workspace add` writes a `.jj/repo` file naming the main store. Both resolve back to the main working copy, so every checkout of one repository shares one cache. Anything else is its own root. """ jj_repo = root / ".jj" / "repo" if jj_repo.is_file(): try: pointer = (jj_repo.parent / jj_repo.read_text(encoding="utf-8").strip()).resolve() except OSError: return root return pointer.parent.parent if pointer.name == "repo" else root git_entry = root / ".git" if git_entry.is_file(): try: _, _, gitdir = git_entry.read_text(encoding="utf-8").strip().partition(":") except OSError: return root if not gitdir.strip(): return root linked = Path(gitdir.strip()) git_dir = (linked if linked.is_absolute() else root / linked).resolve().parent.parent return git_dir.parent if git_dir.name == ".git" else root return root
def _shared_cache_dir() -> Path | None: """The repository's shared build cache, or None when there is nothing to share.
`PRIME_AGENT_SHARED_CACHE=0` turns sharing off for this session, `PRIME_AGENT_SHARED_CACHE_DIR` puts the cache somewhere else entirely, and a repository that does not build with cargo keeps its own layout. """ if not _env_flag(_SHARED_CACHE_SWITCH, True): return None override = os.environ.get(_SHARED_CACHE_DIR_ENV, "").strip() if override: return Path(override).expanduser() root = _repo_root() if root is None: return None main_root = _shared_repo_root(root) if not (main_root / "Cargo.toml").is_file(): return None return main_root / _SHARED_CACHE_RELATIVE
def _prepare_shared_cache(path: Path) -> None: """Create the shared cache directory with a self-ignoring `.gitignore`.
The cache sits inside the repository, so a working copy that does not already ignore the directory would snapshot a multi-gigabyte build tree into its next commit; a `.gitignore` holding `*` keeps the tree out of git and jj alike. Failures are ignored: cargo creates the directory itself, and a cache that cannot be prepared must not stop a command. """ try: marker = path / ".gitignore" if marker.exists(): return path.mkdir(parents=True, exist_ok=True) marker.write_text("*\n", encoding="utf-8") except OSError: pass
def _run_vcs(command: list[str], cwd: Path) -> subprocess.CompletedProcess[str]: try: return subprocess.run( command, cwd=str(cwd), capture_output=True, text=True, timeout=_VCS_COMMAND_TIMEOUT, ) except (OSError, subprocess.SubprocessError) as exc: raise RuntimeError(f"workspace command failed: {' '.join(command)}: {exc}") from None
def _vcs_message(result: subprocess.CompletedProcess[str]) -> str: return (result.stderr or result.stdout or "").strip() or f"exit {result.returncode}"
def _vcs_kind(root: Path) -> Literal["git", "jj"] | None: """Which VCS owns `root`; jj wins a colocated repository.""" if (root / ".jj").exists(): return "jj" if (root / ".git").exists(): return "git" return None
def _exclude_ephemeral_paths(root: Path, patterns: tuple[str, ...]) -> None: """Add local ignore patterns so helper directories never reach VCS status.
`.git/info/exclude` is the untracked, per-clone place for this: editing a tracked `.gitignore` would be a change to the user's repository. jj reads the same file in a colocated repository and skips a nested working copy on its own, so a jj-only repository needs nothing here. """ exclude = _shared_repo_root(root) / ".git" / "info" / "exclude" try: if not exclude.parent.is_dir(): return current = exclude.read_text(encoding="utf-8") if exclude.is_file() else "" known = {line.strip() for line in current.splitlines()} missing = [pattern for pattern in patterns if pattern not in known] if not missing: return text = "" if not current or current.endswith("\n") else "\n" if _EPHEMERAL_IGNORE_MARKER not in known: text += f"{_EPHEMERAL_IGNORE_MARKER}\n" text += "".join(f"{pattern}\n" for pattern in missing) with exclude.open("a", encoding="utf-8") as handle: handle.write(text) except OSError: pass
@dataclass(frozen=True)class Workspace: """One ephemeral co-located checkout, and the way to remove it again."""
name: str path: Path vcs: Literal["git", "jj"] root: Path
def prune(self) -> str | None: """Forget and delete the checkout, and never raise.
Removal is forced: a caller prunes a checkout it has already decided to throw away, and both VCSes otherwise refuse to discard an edited tree. jj keeps the forgotten working-copy commit in the repository, so `jj log` still reaches work that was never committed; a git checkout goes outright. Returns a note when removal only partly succeeded, None when the checkout is gone. """ note: str | None = None try: if self.vcs == "jj": result = _run_vcs(["jj", "workspace", "forget", self.name], self.root) if result.returncode != 0: note = f"jj workspace forget {self.name} failed: {_vcs_message(result)}" else: result = _run_vcs(["git", "worktree", "remove", "--force", str(self.path)], self.root) if result.returncode != 0: _run_vcs(["git", "worktree", "prune"], self.root) note = f"git worktree remove {self.path} failed: {_vcs_message(result)}" except RuntimeError as exc: note = str(exc) shutil.rmtree(self.path, ignore_errors=True) if self.path.exists(): return note or f"workspace directory {self.path} could not be removed" return note
def create_workspace(name: str, *, repo_root: Path | None = None) -> Workspace: """Create an ephemeral checkout under `<repo_root>/.worktrees/<name>`.
The checkout is a linked checkout of the repository that contains it — `jj workspace add` for a jj repository, `git worktree add --detach` for a git one — so it shares the repository's object store and its `.git`/`.jj` pointers resolve. It is never created under the system temp directory, because a checkout that loses its repository makes a Nix build copy a store path for every build, and never outside its repository.
Raises RuntimeError when the directory is not a repository, when the path would leave the repository, or when the VCS command fails. The caller owns the result and removes it with `Workspace.prune()`. """ if not isinstance(name, str) or not name.strip(): raise ValueError("workspace name must be a non-empty str") name = name.strip() root = Path(repo_root) if repo_root is not None else _repo_root() if root is None: raise RuntimeError("workspace='isolated' needs a git or jj repository to add a workspace to") root = root.resolve() vcs = _vcs_kind(root) if vcs is None: raise RuntimeError(f"{root} is neither a git nor a jj repository") # resolve() collapses a name that walks out with `..`, so containment is a # fact about the final path rather than about how it is spelled. path = (root / _WORKTREES_DIRNAME / name).resolve() if root not in path.parents: raise RuntimeError(f"workspace path {path} must stay inside {root}") _exclude_ephemeral_paths(root, (f"{_WORKTREES_DIRNAME}/",)) path.parent.mkdir(parents=True, exist_ok=True) command = ( ["jj", "workspace", "add", "--name", name, str(path)] if vcs == "jj" else ["git", "worktree", "add", "--detach", str(path), "HEAD"] ) result = _run_vcs(command, root) if result.returncode != 0: raise RuntimeError(f"workspace add failed: {' '.join(command)}: {_vcs_message(result)}") return Workspace(name=name, path=path, vcs=vcs, root=root)
def _child_env() -> dict[str, str]: """The launcher's environment, verbatim, plus the non-interactive guard.
Children inherit the environment this runtime was started with: the launcher's or daemon's environment, PATH and every toolchain variable included, with no login-shell probe (no rc file is read, no subprocess runs, no network is touched). PATH keeps its inherited left-to-right order, so a toolchain the launcher put first stays first; only a runtime started with no PATH at all falls back to the POSIX default utility path, so a missing tool fails as "not found" instead of every tool failing that way.
Same non-interactive guard as the coding-agent shell tool (packages/coding-agent/src/utils/shell.ts): agent shell commands have no usable stdin, so interactive prompts (git commit without -m opening $EDITOR, credential asks, pagers) can only hang. Fail fast or no-op instead. Deliberately overrides inherited terminal settings; a per-command inline assignment (`GIT_EDITOR=vim git commit`) still wins because it replaces the exported value for that command.
A repository that builds with cargo also gets one shared warm cache: every session and every child of the same repository builds into `<repo_root>/target/subagents-shared`, so a fresh child does not recompile a cold tree and parallel children do not partition the cache per session. A value the launcher already set for that variable wins, and `PRIME_AGENT_SHARED_CACHE=0` turns sharing off for a caller who wants isolation. """ env = dict(os.environ) if not env.get("PATH"): env["PATH"] = _default_utility_path() env.update( { "NO_COLOR": "1", "TERM": "dumb", "CLICOLOR": "0", "FORCE_COLOR": "0", "GIT_EDITOR": "true", "GIT_SEQUENCE_EDITOR": "true", "GIT_TERMINAL_PROMPTS": "0", "GIT_ASKPASS": "true", "SSH_ASKPASS_REQUIRE": "never", "EDITOR": "true", "VISUAL": "true", "PAGER": "cat", "GIT_PAGER": "cat", "DEBIAN_FRONTEND": "noninteractive", } ) env["NO_COLOR"] = "1" env["TERM"] = "dumb" env["CLICOLOR"] = "0" env["FORCE_COLOR"] = "0" shared_cache = _shared_cache_dir() if shared_cache is not None and not env.get(_SHARED_CACHE_ENV): _prepare_shared_cache(shared_cache) env[_SHARED_CACHE_ENV] = str(shared_cache) return env
def _pid_owns_group(pid: int) -> bool: """False only when the pid provably belongs to another process group.
Every bash() child leads its own group (start_new_session=True), so pgid == pid. Any other pid is somebody else's group member, and signalling it as a group would reach the kernel, the daemon, or the ssh wrapper and ControlMaster socket that carry this kernel's transport. """ try: return os.getpgid(pid) == pid except ProcessLookupError: # A reaped leader can still have live group members, and pgid == pid is # the only shape bash() creates, so the group kill stays valid. return True except OSError: return True
def _signal_group(pid: int, sig: int) -> bool: """True when the signal was delivered or the group is already gone.
A pid that does not lead its own process group is refused: a group kill of any other pid lands on the group that contains it. A refusal reports False, which keeps the journal record active and leaves the host reaper to decide. """ if not _pid_owns_group(pid): return False from .host_safety import _is_group_fully_owned, get_owned_process
lease = get_owned_process(pid) if lease is not None and not lease.created_by_session: if not _is_group_fully_owned(pid): return False 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)