zut_client.py RawHistoryFolder

1"""Minimal Python client for the `zut rpc` JSON protocol.
2
3Usage:
4    python zut_client.py "fix the failing test"
5
6Spawns `zut rpc`, sends one prompt, prints assistant text as it streams,
7and exits when the turn finishes. No external dependencies — stdlib
8only. Implements just enough of the protocol to be useful as a
9starting point; see docs/rpc.md for the full schema.
10"""
11from __future__ import annotations
12
13import json
14import os
15import subprocess
16import sys
17import threading
18import uuid
19
20
21class ZutClient:
22    def __init__(self, *flags: str) -> None:
23        argv = ["zut", "rpc", *flags]
24        env = os.environ.copy()
25        self.proc = subprocess.Popen(
26            argv,
27            stdin=subprocess.PIPE,
28            stdout=subprocess.PIPE,
29            stderr=sys.stderr,
30            env=env,
31            text=True,
32            bufsize=1,  # line-buffered
33        )
34        self._lock = threading.Lock()
35
36    def send(self, **command) -> str:
37        """Send one command, return its `id`."""
38        if "id" not in command:
39            command["id"] = uuid.uuid4().hex[:8]
40        line = json.dumps(command)
41        with self._lock:
42            assert self.proc.stdin is not None
43            self.proc.stdin.write(line + "\n")
44            self.proc.stdin.flush()
45        return command["id"]
46
47    def events(self):
48        """Yield every JSON object the server emits, until stdout closes."""
49        assert self.proc.stdout is not None
50        for line in self.proc.stdout:
51            line = line.strip()
52            if not line:
53                continue
54            try:
55                yield json.loads(line)
56            except json.JSONDecodeError:
57                continue
58
59    def close(self) -> None:
60        if self.proc.stdin is not None:
61            try:
62                self.proc.stdin.close()
63            except Exception:
64                pass
65        self.proc.wait(timeout=5)
66
67
68def main() -> int:
69    if len(sys.argv) < 2:
70        print("usage: zut_client.py <prompt>", file=sys.stderr)
71        return 2
72    prompt = " ".join(sys.argv[1:])
73
74    client = ZutClient()
75    try:
76        token = os.environ.get("ZUTCORE_RPC_TOKEN")
77        if token:
78            client.send(type="hello", token=token)
79
80        client.send(type="prompt", message=prompt)
81
82        for ev in client.events():
83            t = ev.get("type")
84            if t == "text_delta":
85                sys.stdout.write(ev.get("delta", ""))
86                sys.stdout.flush()
87            elif t == "tool_call":
88                print(
89                    f"\n[tool] {ev.get('name')}({json.dumps(ev.get('args', {}))})",
90                    file=sys.stderr,
91                )
92            elif t == "tool_result":
93                if ev.get("is_error"):
94                    print("[tool error]", file=sys.stderr)
95            elif t == "usage":
96                cum = ev.get("cumulative", {})
97                print(
98                    f"\n[usage] cum input={cum.get('input')} "
99                    f"output={cum.get('output')} cost=${cum.get('cost_usd', 0):.4f}",
100                    file=sys.stderr,
101                )
102            elif t == "done":
103                break
104            elif t == "error":
105                print(f"\n[error] {ev.get('message')}", file=sys.stderr)
106                return 1
107        print()
108        return 0
109    finally:
110        client.close()
111
112
113if __name__ == "__main__":
114    sys.exit(main())