87 lines
3.0 KiB
Python
87 lines
3.0 KiB
Python
"""The only subprocess code: run one `claude -p` call in the foreground and tee its stream."""
|
|
import json
|
|
import signal
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
|
|
POLL_S = 0.2
|
|
READER_JOIN_S = 10.0 # an orphaned grandchild can hold stdout open after claude exits
|
|
STOP_SIGNALS = tuple(s for s in (signal.SIGTERM, signal.SIGINT, getattr(signal, "SIGBREAK", None))
|
|
if s is not None)
|
|
|
|
|
|
def _feed(stdin, data: bytes) -> None:
|
|
"""Write the prompt and close stdin; a child that exits early is not an error here."""
|
|
try:
|
|
stdin.write(data)
|
|
except (OSError, ValueError):
|
|
pass
|
|
try:
|
|
stdin.close()
|
|
except (OSError, ValueError):
|
|
pass
|
|
|
|
|
|
def _read(stdout, raw, ss) -> None:
|
|
"""Tee every stream line to the raw file and feed each JSON event to `ss`."""
|
|
try:
|
|
for line in iter(stdout.readline, b""):
|
|
raw.write(line)
|
|
raw.flush()
|
|
try:
|
|
ev = json.loads(line.decode("utf-8", "replace"))
|
|
except ValueError:
|
|
continue
|
|
ss.feed(ev, time.time())
|
|
except (OSError, ValueError):
|
|
pass # the raw file or pipe was closed after READER_JOIN_S
|
|
|
|
|
|
def run_claude(argv, prompt: str, env, cwd, raw_path, on_start, ss=None):
|
|
"""Run one call to its end. Returns (ss, exit_code, stopped)."""
|
|
if ss is None:
|
|
from grm.stream import StreamState
|
|
ss = StreamState()
|
|
raw_path = Path(raw_path)
|
|
with open(raw_path, "wb") as raw, open(raw_path.with_suffix(".err"), "wb") as err:
|
|
child = subprocess.Popen(argv, cwd=cwd, env=env, stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE, stderr=err)
|
|
stopped = False
|
|
|
|
def _stop(signum, frame):
|
|
nonlocal stopped
|
|
try:
|
|
if stopped: # a second signal: stop waiting for a graceful exit
|
|
child.kill()
|
|
else:
|
|
child.terminate()
|
|
except OSError:
|
|
pass
|
|
stopped = True
|
|
|
|
old = {}
|
|
if threading.current_thread() is threading.main_thread():
|
|
old = {s: signal.signal(s, _stop) for s in STOP_SIGNALS}
|
|
try:
|
|
reader = threading.Thread(target=_read, args=(child.stdout, raw, ss), daemon=True)
|
|
reader.start()
|
|
threading.Thread(target=_feed, args=(child.stdin, prompt.encode("utf-8")),
|
|
daemon=True).start()
|
|
on_start(child.pid)
|
|
while child.poll() is None: # a bare wait() would block signal handlers on Windows
|
|
time.sleep(POLL_S)
|
|
except BaseException:
|
|
child.kill()
|
|
child.wait()
|
|
raise
|
|
finally:
|
|
for s, h in old.items():
|
|
if h is not None:
|
|
signal.signal(s, h)
|
|
reader.join(READER_JOIN_S)
|
|
if not reader.is_alive():
|
|
child.stdout.close()
|
|
return ss, child.returncode, stopped
|