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.
9.2 kB · 235 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236from __future__ import annotations
import osimport secretsimport termiosimport timefrom collections.abc import Callablefrom pathlib import Path
import pexpectimport pyte
# Bound input-handler detection delay without adding a fixed wait to fast startup.PROBE_INTERVAL = 0.02MAX_PENDING_INPUT = 65_536
QUERIES = { "\x1b[6n": "\x1b[1;1R", "\x1b[?6n": "\x1b[?1;1R", "\x1b[c": "\x1b[?1;2c", "\x1b[0c": "\x1b[?1;2c", "\x1b[>c": "\x1b[>0;276;0c", "\x1b[>0c": "\x1b[>0;276;0c", "\x1b[?u": "\x1b[?0u", "\x1b[?2026$p": "\x1b[?2026;2$y", "\x1b[14t": "\x1b[4;640;960t", "\x1b[16t": "\x1b[6;16;8t", "\x1b[18t": "\x1b[8;40;120t",}for terminator in ("\x07", "\x1b\\"): QUERIES[f"\x1b]11;?{terminator}"] = "\x1b]11;rgb:0000/0000/0000\x1b\\" QUERIES[f"\x1b]10;?{terminator}"] = "\x1b]10;rgb:eeee/eeee/eeee\x1b\\"
class Display: def __init__(self, reply: Callable[[str], object]): self.screen = pyte.Screen(120, 40) self.stream = pyte.Stream(self.screen) self.reply = reply self.pending = ""
def feed(self, chunk: str) -> None: self.pending += chunk while True: matches = [(self.pending.find(query), query) for query in QUERIES if query in self.pending] if not matches: break position, query = min(matches) self.stream.feed(self.pending[:position]) self.reply(QUERIES[query]) self.pending = self.pending[position + len(query) :] keep = max( ( length for query in QUERIES for length in range(1, len(query)) if self.pending.endswith(query[:length]) ), default=0, ) if keep: self.stream.feed(self.pending[:-keep]) self.pending = self.pending[-keep:] else: self.stream.feed(self.pending) self.pending = ""
def text(self) -> str: return "\n".join(row.rstrip() for row in self.screen.display)
class Terminal: def __init__(self, command: list[str], cwd: Path, env: dict[str, str], transcript: Path): self.started = time.perf_counter() self.child = pexpect.spawn( command[0], command[1:], cwd=str(cwd), env=env, encoding="utf-8", codec_errors="replace", dimensions=(40, 120), echo=False, ) os.set_blocking(self.child.child_fd, False) self._pending_input = b"" self.display = Display(self._send) self.transcript = transcript self.raw: list[str] = [] self.bytes = 0
def _send(self, text: str) -> None: data = text.encode("utf-8") if len(self._pending_input) + len(data) > MAX_PENDING_INPUT: raise RuntimeError("Terminal input queue exceeds 64 KiB") self._pending_input += data self._flush_input()
def _flush_input(self) -> None: # Never wait for a raw-mode child to read. The next pump retries the exact tail. if not self._pending_input: return try: written = os.write(self.child.child_fd, self._pending_input) except (BlockingIOError, InterruptedError): return self._pending_input = self._pending_input[written:]
def pump(self) -> bool: self._flush_input() try: chunk = self.child.read_nonblocking(65536, timeout=0.005) except pexpect.TIMEOUT: return False except pexpect.EOF as error: raise RuntimeError("Prime Agent exited before the measurement completed") from error self.bytes += len(chunk) if self.bytes > 8_000_000: raise RuntimeError("Terminal transcript exceeds 8 MB") self.raw.append(chunk) self.display.feed(chunk) return True
def settle(self, seconds: float) -> None: deadline = time.perf_counter() + seconds while time.perf_counter() < deadline: self.pump()
def until(self, predicate: Callable[[Display], bool], seconds: float) -> None: deadline = time.perf_counter() + seconds while (2026 << 5) in self.display.screen.mode or not predicate(self.display): if time.perf_counter() >= deadline: raise TimeoutError( "Timed out waiting for terminal display to reach the expected terminal state" ) self.pump()
def until_output(self, predicate: Callable[[str], bool], seconds: float) -> None: index = len(self.raw) output = "" deadline = time.perf_counter() + seconds while True: output += "".join(self.raw[index:]) index = len(self.raw) if predicate(output): return if time.perf_counter() >= deadline: raise TimeoutError("Timed out waiting for new terminal output") self.pump()
def ready(self, seconds: float = 30) -> float: deadline = time.perf_counter() + seconds probes: list[str] = [] next_probe = 0.0 echoed: float | None = None origin: tuple[int, int] | None = None needs_clear = False stage = "noncanonical, no-echo terminal input" # Keep the ten-character input workload identical across benchmark revisions. erase = "\x7f" * 10 try: while time.perf_counter() < deadline: updated = self.pump() # pyte shifts DEC private modes by five bits. Do not inspect a partial # synchronized TUI frame, especially while the editor row is cleared. updated = updated and (2026 << 5) not in self.display.screen.mode flags = termios.tcgetattr(self.child.child_fd)[3] if flags & (termios.ICANON | termios.ECHO): continue now = time.perf_counter() if origin is not None: cursor = self.display.screen.cursor # A shortened marker is not a cleared editor. Verify the cursor # returns to the insertion point, allowing any empty-editor placeholder. if ( updated and not self._pending_input and (cursor.y, cursor.x) == origin and not any(probe in self.display.text() for probe in probes) ): assert echoed is not None return echoed - self.started continue if updated and probes: text = self.display.text() if echoed is None and any(probe in text for probe in probes): echoed = now stage = "latest editor probe rendering" if echoed is not None: for row, line in enumerate(self.display.screen.display): column = line.find(probes[-1]) if column >= 0: origin = (row, column) self._send(erase) needs_clear = False stage = "editor probe cleanup" break if echoed is None and not self._pending_input and now >= next_probe: # Raw mode can precede the editor's input handler. Balanced retries # replace dropped/partial probes without accumulating editor text. if needs_clear: self._send(erase) probe = "b" + secrets.token_hex(4) + "r" self._send(probe) probes.append(probe) needs_clear = True next_probe = now + PROBE_INTERVAL stage = "editor input rendering" # Once any probe renders, stop retrying and acknowledge the newest one # before clearing: older queued retries must not outlive cleanup. raise TimeoutError(f"Timed out waiting for {stage}") finally: if needs_clear: try: # Drop unsent retries; cleanup must not extend the shared deadline. self._pending_input = b"" self._send(erase) except (OSError, pexpect.EOF): pass
def close(self) -> None: raw, snapshot = "".join(self.raw), self.display.text() self.transcript.with_suffix(".raw").write_text(raw) self.transcript.with_suffix(".txt").write_text(snapshot) try: if self.child.isalive(): self._pending_input = b"" self._send("\x03") time.sleep(0.15) if self.child.isalive(): self._send("\x03") except (OSError, pexpect.EOF): pass finally: self.child.close(force=True)