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.
7.0 kB · 177 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178from __future__ import annotations
import jsonimport osimport selectorsimport signalimport subprocessimport timefrom collections import dequefrom contextlib import suppressfrom pathlib import Path
CELL_REPEATS = 50SHELL_REPEATS = 5MIXED_CELLS = 35OUTPUT_BYTES = 32 * 1024FRAME_ROWS = 10_000FRAME_COLUMNS = 8
BASH_EMPTY = "_r = await bash(':')\nassert _r.exit_code == 0 and _r.output == ''"GIT_STATUS = "_r = await bash('git status --porcelain')\nassert _r.exit_code == 0 and _r.output == ''"BASH_OUTPUT = ( f"_r = await bash(\"printf '%{OUTPUT_BYTES}s' ''\")\n" f"assert _r.exit_code == 0 and len(_r.output) == {OUTPUT_BYTES}\nprint(_r.output, end='')")
class Kernel: def __init__(self, command: list[str], cwd: Path, env: dict[str, str], transcript: Path): self.transcript = transcript self.stderr = transcript.with_suffix(".stderr").open("wb") self.selector = selectors.DefaultSelector() self.raw: list[bytes] = [] self.pending = b"" self.lines: deque[bytes] = deque() self.bytes = 0 self.sequence = 0 self.started = time.perf_counter() try: self.process = subprocess.Popen( command, cwd=cwd, env=env, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=self.stderr, bufsize=0, start_new_session=True, ) except Exception: self.stderr.close() self.selector.close() raise assert self.process.stdout is not None self.selector.register(self.process.stdout, selectors.EVENT_READ)
def event(self, deadline: float) -> dict: while not self.lines: remaining = deadline - time.perf_counter() if remaining <= 0 or not self.selector.select(remaining): raise TimeoutError("Timed out waiting for the kernel protocol") assert self.process.stdout is not None chunk = os.read(self.process.stdout.fileno(), 65536) if not chunk: raise RuntimeError("Kernel exited before completing the request") self.bytes += len(chunk) if self.bytes > 8_000_000: raise RuntimeError("Kernel transcript exceeds 8 MB") self.raw.append(chunk) chunks = (self.pending + chunk).split(b"\n") self.pending = chunks.pop() self.lines.extend(chunks) event = json.loads(self.lines.popleft()) if not isinstance(event, dict) or not isinstance(event.get("event"), str): raise ValueError("Invalid kernel protocol event") return event
def ready(self) -> tuple[float, str]: event = self.event(time.perf_counter() + 30) if event.get("event") != "ready" or event.get("protocol") != 3: raise ValueError("Expected kernel protocol 3 ready event") return time.perf_counter() - self.started, str(event.get("python", "unknown"))
def send(self, request: dict) -> None: assert self.process.stdin is not None data = memoryview((json.dumps(request) + "\n").encode()) while data: written = self.process.stdin.write(data) if not written: raise RuntimeError("Kernel protocol write failed") data = data[written:]
def done(self, request_id: str, *, interrupted: bool = False) -> list[dict]: deadline = time.perf_counter() + 30 events = [] while True: event = self.event(deadline) events.append(event) if event.get("event") == "done" and event.get("id") == request_id: if interrupted: if event.get("status") != "error" or not any( e.get("ename") == "KeyboardInterrupt" and e.get("id") == request_id for e in events ): raise RuntimeError("Kernel did not acknowledge the interrupt") elif event.get("status") != "ok": raise RuntimeError(f"Kernel request failed: {str(events)[-400:]}") return events
def request(self, kind: str, **fields: object) -> list[dict]: self.sequence += 1 request_id = str(self.sequence) self.send({"type": kind, "id": request_id, **fields}) return self.done(request_id)
def execute(self, code: str) -> list[dict]: return self.request("execute", code=code)
def batch(self, code: str, repeats: int, *, output_bytes: int | None = None) -> float: started = time.perf_counter() for _ in range(repeats): events = self.execute(code) if output_bytes is not None: output = "".join(e.get("text", "") for e in events if e["event"] == "stdout") if len(output.encode()) != output_bytes: raise RuntimeError("Kernel output was truncated or missing") return (time.perf_counter() - started) / repeats
def mixed(self) -> float: started = time.perf_counter() for index in range(MIXED_CELLS): self.execute(GIT_STATUS if index % 4 == 0 else "pass") return time.perf_counter() - started
def interrupt(self) -> float: self.sequence += 1 request_id = str(self.sequence) self.send( { "type": "execute", "id": request_id, "code": "import asyncio\nprint('interrupt-ready', flush=True)\nawait asyncio.sleep(60)", } ) deadline = time.perf_counter() + 30 while True: event = self.event(deadline) if ( event.get("id") == request_id and event["event"] == "stdout" and "interrupt-ready" in event.get("text", "") ): break if event["event"] in ("error", "done"): raise RuntimeError("Interrupt fixture did not reach the running state") started = time.perf_counter() self.send({"type": "interrupt", "id": request_id}) self.done(request_id, interrupted=True) elapsed = time.perf_counter() - started self.execute("assert 1 + 1 == 2") return elapsed
def close(self) -> None: try: if self.process.poll() is None: self.send({"type": "shutdown"}) self.process.wait(timeout=5) except (OSError, subprocess.TimeoutExpired): with suppress(ProcessLookupError): os.killpg(self.process.pid, signal.SIGKILL) self.process.wait(timeout=5) finally: for stream in (self.process.stdin, self.process.stdout): if stream is not None: stream.close() self.selector.close() self.stderr.close() self.transcript.with_suffix(".jsonl").write_bytes(b"".join(self.raw))