133 lines
4.6 KiB
Python
133 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
||
"""repl_drive.py - 用 pty 驱动 ligent 的 REPL，跑多轮对话并打印全部输出。
|
||
|
||
交互式行为（多轮上下文、/命令、行编辑）没法用管道测 —— 管道模式会把
|
||
整段 stdin 当成一个任务。所以这里开一个伪终端，逐条发送提示词。
|
||
|
||
用法：
|
||
repl_drive.py --bin ./ligent --base-url URL --prompt "第一轮" --prompt "第二轮"
|
||
"""
|
||
import argparse
|
||
import os
|
||
import pty
|
||
import re
|
||
import select
|
||
import sys
|
||
import time
|
||
|
||
|
||
def drive(binary, args, prompts, wait_per_prompt=2.0, quiet_after=1.5,
|
||
raw_keys=None):
|
||
pid, fd = pty.fork()
|
||
if pid == 0:
|
||
os.execv(binary, [binary] + args)
|
||
os._exit(127)
|
||
|
||
out = bytearray()
|
||
# 行编辑器每次回到提示符都会打开括号粘贴，所以这个序列就是"轮到用户输入了"
|
||
# 的信号。用它做同步点，比"等输出安静下来"可靠也快得多。
|
||
PROMPT_READY = b"\x1b[?2004h"
|
||
ready_seen = 0
|
||
tail = b""
|
||
|
||
def note(chunk):
|
||
# 标记可能被 read 切断，所以要连着上一块的尾巴一起找
|
||
nonlocal tail, ready_seen
|
||
buf = tail + chunk
|
||
ready_seen += buf.count(PROMPT_READY)
|
||
tail = buf[-(len(PROMPT_READY) - 1):]
|
||
|
||
def drain(timeout, idle=0.5, until_ready=False):
|
||
"""等到"回到提示符"（until_ready），或者输出彻底安静下来。
|
||
|
||
必须按标记等，不能只按安静等：等待动画每 250ms 画一帧，pty 永远
|
||
不会安静超过 idle，于是每次都要白等满 timeout —— 而 mock 其实几
|
||
毫秒就答复了。这是整个 e2e 从 126 秒掉下来的关键。
|
||
"""
|
||
want = ready_seen + 1 if until_ready else None
|
||
deadline = time.time() + timeout
|
||
last = time.time()
|
||
while time.time() < deadline:
|
||
r, _, _ = select.select([fd], [], [], 0.02)
|
||
if r:
|
||
try:
|
||
chunk = os.read(fd, 65536)
|
||
except OSError:
|
||
return False
|
||
if not chunk:
|
||
return False
|
||
out.extend(chunk)
|
||
note(chunk)
|
||
last = time.time()
|
||
if want is not None and ready_seen >= want:
|
||
return True
|
||
elif time.time() - last > idle:
|
||
return True
|
||
return True
|
||
|
||
drain(quiet_after, until_ready=True) # 等横幅画完、第一次回到提示符
|
||
if raw_keys is not None:
|
||
# 逐字节发送原始按键序列（用于测方向键等编辑行为）
|
||
for chunk, pause in raw_keys:
|
||
os.write(fd, chunk if isinstance(chunk, bytes) else chunk.encode("utf-8"))
|
||
drain(pause)
|
||
os.write(fd, b"/exit\n")
|
||
drain(0.8)
|
||
try:
|
||
os.close(fd)
|
||
except OSError:
|
||
pass
|
||
try:
|
||
os.waitpid(pid, 0)
|
||
except ChildProcessError:
|
||
pass
|
||
return re.sub(r"\x1b\[[0-9;]*m", "", out.decode("utf-8", "replace"))
|
||
|
||
for p in prompts:
|
||
os.write(fd, (p + "\n").encode("utf-8"))
|
||
if not drain(wait_per_prompt, until_ready=True):
|
||
break
|
||
os.write(fd, b"/exit\n")
|
||
drain(0.8)
|
||
|
||
try:
|
||
os.close(fd)
|
||
except OSError:
|
||
pass
|
||
try:
|
||
os.waitpid(pid, 0)
|
||
except ChildProcessError:
|
||
pass
|
||
|
||
text = out.decode("utf-8", "replace")
|
||
return re.sub(r"\x1b\[[0-9;]*m", "", text) # 去掉 ANSI 样式
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--bin", default="./ligent")
|
||
ap.add_argument("--base-url", required=True)
|
||
ap.add_argument("--api-key", default="test")
|
||
ap.add_argument("--prompt", action="append", default=[])
|
||
ap.add_argument("--extra", action="append", default=[])
|
||
ap.add_argument("--wait", type=float, default=2.0)
|
||
ap.add_argument("--keys", action="append", default=[],
|
||
help="原始按键序列（支持 \\x1b 转义），按顺序发送")
|
||
a = ap.parse_args()
|
||
|
||
# --no-history：被驱动的实例不该往用户真实的 ligent.history 里写测试垃圾。
|
||
# 需要历史的调用方用 --extra=--history PATH（排在后面，会覆盖这一条）。
|
||
args = ["--api-key", a.api_key, "--base-url", a.base_url, "--no-color",
|
||
"--no-confirm", "--no-auto-save", "--no-history"] + a.extra
|
||
if a.keys:
|
||
raw = [(k.encode("utf-8").decode("unicode_escape").encode("latin-1"), 0.25)
|
||
for k in a.keys]
|
||
sys.stdout.write(drive(a.bin, args, [], a.wait, raw_keys=raw))
|
||
else:
|
||
sys.stdout.write(drive(a.bin, args, a.prompt, a.wait))
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|