@samitouri / QOSami_Agents_Zero-1 / commits / 2565497d

tty session prototype

frdel committed Aug 1, 2025 at 12:28 UTC 2565497da030a3e80601fac1cfa82577c8cdd4a3
4 files changed +544 -3
python/helpers/tty_session.py new
+273
@@ -0,0 +1,273 @@
1 +import asyncio, os, sys, platform, errno
2 +
3 +_IS_WIN = platform.system() == "Windows"
4 +if _IS_WIN:
5 + import winpty # pip install pywinpty # type: ignore
6 + import msvcrt
7 +
8 +
9 +# Make stdin / stdout tolerant to broken UTF-8 so input() never aborts
10 +sys.stdin.reconfigure(errors="replace") # type: ignore
11 +sys.stdout.reconfigure(errors="replace") # type: ignore
12 +
13 +
14 +# ──────────────────────────── PUBLIC CLASS ────────────────────────────
15 +
16 +
17 +class TTYSession:
18 + def __init__(
19 + self, cmd, *, cwd=None, env=None, encoding="utf-8", echo=False
20 + ): # ← NEW kw-arg `echo`
21 + self.cmd = cmd if isinstance(cmd, str) else " ".join(cmd)
22 + self.cwd = cwd
23 + self.env = env or os.environ.copy()
24 + self.encoding = encoding
25 + self.echo = echo # ← store preference
26 + self._proc = None
27 + self._buf = asyncio.Queue()
28 +
29 + # ── user-facing coroutines ────────────────────────────────────────
30 + async def start(self):
31 + if _IS_WIN:
32 + self._proc = await _spawn_winpty(
33 + self.cmd, self.cwd, self.env, self.echo
34 + ) # ← pass echo
35 + else:
36 + self._proc = await _spawn_posix_pty(
37 + self.cmd, self.cwd, self.env, self.echo
38 + ) # ← pass echo
39 + asyncio.create_task(self._pump_stdout())
40 +
41 + async def send(self, data: str | bytes):
42 + if self._proc is None:
43 + raise RuntimeError("TTYSpawn is not started")
44 + if isinstance(data, str):
45 + data = data.encode(self.encoding)
46 + self._proc.stdin.write(data) # type: ignore
47 + await self._proc.stdin.drain() # type: ignore
48 +
49 + async def sendline(self, line: str):
50 + await self.send(line + "\n")
51 +
52 + async def wait(self):
53 + if self._proc is None:
54 + raise RuntimeError("TTYSpawn is not started")
55 + return await self._proc.wait()
56 +
57 + def kill(self):
58 + if self._proc is None:
59 + raise RuntimeError("TTYSpawn is not started")
60 + self._proc.kill()
61 +
62 + async def read(self, timeout=None):
63 + # Return any decoded text the child produced, or None on timeout
64 + try:
65 + return await asyncio.wait_for(self._buf.get(), timeout)
66 + except asyncio.TimeoutError:
67 + return None
68 +
69 + # backward-compat alias:
70 + readline = read
71 +
72 + async def read_full_until_idle(self, idle_timeout, total_timeout):
73 + # Collect child output using iter_until_idle to avoid duplicate logic
74 + return "".join(
75 + [
76 + chunk
77 + async for chunk in self.read_chunks_until_idle(
78 + idle_timeout, total_timeout
79 + )
80 + ]
81 + )
82 +
83 + async def read_chunks_until_idle(self, idle_timeout, total_timeout):
84 + # Yield each chunk as soon as it arrives until idle or total timeout
85 + import time
86 +
87 + start = time.monotonic()
88 + while True:
89 + if time.monotonic() - start > total_timeout:
90 + break
91 + chunk = await self.read(timeout=idle_timeout)
92 + if chunk is None:
93 + break
94 + yield chunk
95 +
96 + # ── internal: stream raw output into the queue ────────────────────
97 + async def _pump_stdout(self):
98 + if self._proc is None:
99 + raise RuntimeError("TTYSpawn is not started")
100 + reader = self._proc.stdout
101 + while True:
102 + chunk = await reader.read(4096) # grab whatever is ready # type: ignore
103 + if not chunk:
104 + break
105 + self._buf.put_nowait(chunk.decode(self.encoding, "replace"))
106 +
107 +
108 +# ──────────────────────────── POSIX IMPLEMENTATION ────────────────────
109 +
110 +
111 +async def _spawn_posix_pty(cmd, cwd, env, echo):
112 + import pty, asyncio, os, termios
113 +
114 + master, slave = pty.openpty()
115 +
116 + # ── Disable ECHO on the slave side if requested ──
117 + if not echo:
118 + attrs = termios.tcgetattr(slave)
119 + attrs[3] &= ~termios.ECHO # lflag
120 + termios.tcsetattr(slave, termios.TCSANOW, attrs)
121 +
122 + proc = await asyncio.create_subprocess_shell(
123 + cmd,
124 + stdin=slave,
125 + stdout=slave,
126 + stderr=slave,
127 + cwd=cwd,
128 + env=env,
129 + close_fds=True,
130 + )
131 + os.close(slave)
132 +
133 + loop = asyncio.get_running_loop()
134 + reader = asyncio.StreamReader()
135 +
136 + def _on_data():
137 + try:
138 + data = os.read(master, 1 << 16)
139 + except OSError as e:
140 + if e.errno != errno.EIO: # EIO == EOF on some systems
141 + raise
142 + data = b""
143 + if data:
144 + reader.feed_data(data)
145 + else:
146 + reader.feed_eof()
147 + loop.remove_reader(master)
148 +
149 + loop.add_reader(master, _on_data)
150 +
151 + class _Stdin:
152 + def write(self, d):
153 + os.write(master, d)
154 +
155 + async def drain(self):
156 + await asyncio.sleep(0)
157 +
158 + proc.stdin = _Stdin() # type: ignore
159 + proc.stdout = reader
160 + return proc
161 +
162 +
163 +# ──────────────────────────── WINDOWS IMPLEMENTATION ──────────────────
164 +
165 +
166 +async def _spawn_winpty(cmd, cwd, env, echo):
167 + # A quick way to silence command echo in cmd.exe is /Q (quiet)
168 + if not echo and cmd.strip().lower().startswith("cmd") and "/q" not in cmd.lower():
169 + cmd = cmd.replace("cmd.exe", "cmd.exe /Q")
170 +
171 + cols, rows = 80, 25
172 + pty = winpty.PTY(cols, rows) # type: ignore
173 + child = pty.spawn(cmd, cwd=cwd or os.getcwd(), env=env)
174 +
175 + master_r_fd = msvcrt.open_osfhandle(child.conout_pipe, os.O_RDONLY) # type: ignore
176 + master_w_fd = msvcrt.open_osfhandle(child.conin_pipe, 0) # type: ignore
177 +
178 + loop = asyncio.get_running_loop()
179 + reader = asyncio.StreamReader()
180 +
181 + def _on_data():
182 + try:
183 + data = os.read(master_r_fd, 1 << 16)
184 + except OSError:
185 + data = b""
186 + if data:
187 + reader.feed_data(data)
188 + else:
189 + reader.feed_eof()
190 + loop.remove_reader(master_r_fd)
191 +
192 + loop.add_reader(master_r_fd, _on_data)
193 +
194 + class _Stdin:
195 + def write(self, d):
196 + os.write(master_w_fd, d)
197 +
198 + async def drain(self):
199 + await asyncio.sleep(0)
200 +
201 + class _Proc(asyncio.subprocess.Process):
202 + def __init__(self):
203 + self.stdin = _Stdin() # type: ignore
204 + self.stdout = reader
205 + self.pid = child.pid
206 +
207 + async def wait(self):
208 + while child.isalive():
209 + await asyncio.sleep(0.2)
210 + return 0
211 +
212 + def kill(self):
213 + child.kill()
214 +
215 + return _Proc()
216 +
217 +
218 +# ───────────────────────── INTERACTIVE DRIVER ─────────────────────────
219 +if __name__ == "__main__":
220 +
221 + async def interactive_shell():
222 + shell_cmd, prompt_hint = ("cmd.exe", "$") if _IS_WIN else ("/bin/bash", "$")
223 +
224 + # echo=False → suppress the shell’s own echo of commands
225 + term = TTYSession(shell_cmd)
226 + await term.start()
227 +
228 + timeout = 1.0
229 +
230 + print(f"Connected to {shell_cmd}.")
231 + print("Type commands for the shell.")
232 + print("• /t=<seconds> → change idle timeout")
233 + print("• /exit → quit helper\n")
234 +
235 + await term.sendline(" ")
236 + print(await term.read_full_until_idle(timeout, timeout), end="", flush=True)
237 +
238 + while True:
239 + try:
240 + user = input(f"(timeout={timeout}) {prompt_hint} ")
241 + except (EOFError, KeyboardInterrupt):
242 + print("\nLeaving…")
243 + break
244 +
245 + if user.lower() == "/exit":
246 + break
247 + if user.startswith("/t="):
248 + try:
249 + timeout = float(user.split("=", 1)[1])
250 + print(f"[helper] idle timeout set to {timeout}s")
251 + except ValueError:
252 + print("[helper] invalid number")
253 + continue
254 +
255 + idle_timeout = timeout
256 + total_timeout = 10 * idle_timeout
257 + if user == "":
258 + # Just read output, do not send empty line
259 + async for chunk in term.read_chunks_until_idle(
260 + idle_timeout, total_timeout
261 + ):
262 + print(chunk, end="", flush=True)
263 + else:
264 + await term.sendline(user)
265 + async for chunk in term.read_chunks_until_idle(
266 + idle_timeout, total_timeout
267 + ):
268 + print(chunk, end="", flush=True)
269 +
270 + await term.sendline("exit")
271 + await term.wait()
272 +
273 + asyncio.run(interactive_shell())
tests/rate_limiter_test.py
+1 -2
@@ -1,6 +1,5 @@
1
2 -import sys
3 -import os
2 +import sys, os
3 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
4 import models
5
tests/test_fasta2a_client.py renamed
+4 -1
@@ -3,8 +3,11 @@
3 Test script to verify FastA2A agent card routing and authentication.
4 """
5
6 +import sys, os
7 +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
8 +
9 +
10 import asyncio
7 -import sys
11 from python.helpers import settings
12
13
tests/tty_test.py new
+266
@@ -0,0 +1,266 @@
1 +"""
2 +tty_spawn.py – drive any console program through a real TTY
3 +cross-platform (Windows via pywinpty, POSIX via pty).
4 +
5 +API (async):
6 + term = TTYSpawn(cmd, *, cwd=None, env=None, encoding="utf-8")
7 + await term.start() # launch child
8 + chunk = await term.read(timeout=1.0) # None on timeout
9 + await term.send("data") # raw bytes
10 + await term.sendline("text") # adds '\n'
11 + await term.wait() # exit code
12 + term.kill() # abort
13 +"""
14 +# ────────────────── NO ORIGINAL LINES REMOVED – ONLY APPENDED CODE ──────────────────
15 +
16 +import asyncio, os, sys, platform, errno
17 +
18 +_IS_WIN = platform.system() == "Windows"
19 +if _IS_WIN:
20 + import winpty # pip install pywinpty
21 + import msvcrt
22 +
23 +
24 +# Make stdin / stdout tolerant to broken UTF-8 so input() never aborts
25 +try:
26 + # Python 3.7+ has reconfigure()
27 + sys.stdin .reconfigure(errors="replace")
28 + sys.stdout.reconfigure(errors="replace")
29 +except AttributeError:
30 + # Older Pythons: wrap them manually
31 + import io
32 + sys.stdin = io.TextIOWrapper(sys.stdin.buffer,
33 + encoding=sys.stdin.encoding or "utf-8",
34 + errors="replace",
35 + line_buffering=True)
36 + sys.stdout = io.TextIOWrapper(sys.stdout.buffer,
37 + encoding=sys.stdout.encoding or "utf-8",
38 + errors="replace",
39 + line_buffering=True)
40 +# ─────────────────────────────────────────────────────────────────────
41 +
42 +
43 +# ──────────────────────────── PUBLIC CLASS ────────────────────────────
44 +
45 +class TTYSpawn:
46 + def __init__(self, cmd, *, cwd=None, env=None,
47 + encoding="utf-8", echo=True): # ← NEW kw-arg `echo`
48 + self.cmd = cmd if isinstance(cmd, str) else " ".join(cmd)
49 + self.cwd = cwd
50 + self.env = env or os.environ.copy()
51 + self.encoding = encoding
52 + self.echo = echo # ← store preference
53 + self._proc = None
54 + self._buf = asyncio.Queue()
55 +
56 + # ── user-facing coroutines ────────────────────────────────────────
57 + async def start(self):
58 + if _IS_WIN:
59 + self._proc = await _spawn_winpty(
60 + self.cmd, self.cwd, self.env, self.echo) # ← pass echo
61 + else:
62 + self._proc = await _spawn_posix_pty(
63 + self.cmd, self.cwd, self.env, self.echo) # ← pass echo
64 + asyncio.create_task(self._pump_stdout())
65 +
66 + async def read(self, timeout=None):
67 + # Return any decoded text the child produced, or None on timeout
68 + try:
69 + return await asyncio.wait_for(self._buf.get(), timeout)
70 + except asyncio.TimeoutError:
71 + return None
72 +
73 + # backward-compat alias:
74 + readline = read
75 +
76 + async def send(self, data: str | bytes):
77 + if isinstance(data, str):
78 + data = data.encode(self.encoding)
79 + self._proc.stdin.write(data)
80 + await self._proc.stdin.drain()
81 +
82 + async def sendline(self, line: str):
83 + await self.send(line + "\n")
84 +
85 + async def wait(self):
86 + return await self._proc.wait()
87 +
88 + def kill(self):
89 + self._proc.kill()
90 +
91 + async def read_until_idle(self, idle_timeout, total_timeout):
92 + # Collect child output using iter_until_idle to avoid duplicate logic
93 + return "".join([chunk async for chunk in self.iter_until_idle(idle_timeout, total_timeout)])
94 +
95 + async def iter_until_idle(self, idle_timeout, total_timeout):
96 + # Yield each chunk as soon as it arrives until idle or total timeout
97 + import time
98 + start = time.monotonic()
99 + while True:
100 + if time.monotonic() - start > total_timeout:
101 + break
102 + chunk = await self.read(timeout=idle_timeout)
103 + if chunk is None:
104 + break
105 + yield chunk
106 +
107 + # ── internal: stream raw output into the queue ────────────────────
108 + async def _pump_stdout(self):
109 + reader = self._proc.stdout
110 + while True:
111 + chunk = await reader.read(4096) # grab whatever is ready
112 + if not chunk:
113 + break
114 + self._buf.put_nowait(chunk.decode(self.encoding, "replace"))
115 +
116 +# ──────────────────────────── POSIX IMPLEMENTATION ────────────────────
117 +
118 +async def _spawn_posix_pty(cmd, cwd, env, echo):
119 + import pty, asyncio, os, termios
120 + master, slave = pty.openpty()
121 +
122 + # ── Disable ECHO on the slave side if requested ──
123 + if not echo:
124 + attrs = termios.tcgetattr(slave)
125 + attrs[3] &= ~termios.ECHO # lflag
126 + termios.tcsetattr(slave, termios.TCSANOW, attrs)
127 +
128 + proc = await asyncio.create_subprocess_shell(
129 + cmd,
130 + stdin=slave,
131 + stdout=slave,
132 + stderr=slave,
133 + cwd=cwd,
134 + env=env,
135 + close_fds=True,
136 + )
137 + os.close(slave)
138 +
139 + loop = asyncio.get_running_loop()
140 + reader = asyncio.StreamReader()
141 +
142 + def _on_data():
143 + try:
144 + data = os.read(master, 1 << 16)
145 + except OSError as e:
146 + if e.errno != errno.EIO: # EIO == EOF on some systems
147 + raise
148 + data = b""
149 + if data:
150 + reader.feed_data(data)
151 + else:
152 + reader.feed_eof()
153 + loop.remove_reader(master)
154 +
155 + loop.add_reader(master, _on_data)
156 +
157 + class _Stdin:
158 + def write(self, d): os.write(master, d)
159 + async def drain(self): await asyncio.sleep(0)
160 +
161 + proc.stdin = _Stdin()
162 + proc.stdout = reader
163 + return proc
164 +
165 +# ──────────────────────────── WINDOWS IMPLEMENTATION ──────────────────
166 +
167 +async def _spawn_winpty(cmd, cwd, env, echo):
168 + # A quick way to silence command echo in cmd.exe is /Q (quiet)
169 + if not echo and cmd.strip().lower().startswith("cmd") and "/q" not in cmd.lower():
170 + cmd = cmd.replace("cmd.exe", "cmd.exe /Q")
171 +
172 + cols, rows = 80, 25
173 + pty = winpty.PTY(cols, rows)
174 + child = pty.spawn(cmd, cwd=cwd or os.getcwd(), env=env)
175 +
176 + master_r_fd = msvcrt.open_osfhandle(child.conout_pipe, os.O_RDONLY)
177 + master_w_fd = msvcrt.open_osfhandle(child.conin_pipe, 0)
178 +
179 + loop = asyncio.get_running_loop()
180 + reader = asyncio.StreamReader()
181 +
182 + def _on_data():
183 + try:
184 + data = os.read(master_r_fd, 1 << 16)
185 + except OSError:
186 + data = b""
187 + if data:
188 + reader.feed_data(data)
189 + else:
190 + reader.feed_eof()
191 + loop.remove_reader(master_r_fd)
192 +
193 + loop.add_reader(master_r_fd, _on_data)
194 +
195 + class _Stdin:
196 + def write(self, d): os.write(master_w_fd, d)
197 + async def drain(self): await asyncio.sleep(0)
198 +
199 + class _Proc(asyncio.subprocess.Process):
200 + def __init__(self):
201 + self.stdin = _Stdin()
202 + self.stdout = reader
203 + self.pid = child.pid
204 + async def wait(self):
205 + while child.isalive():
206 + await asyncio.sleep(0.2)
207 + return 0
208 + def kill(self):
209 + child.kill()
210 +
211 + return _Proc()
212 +
213 +# ───────────────────────── INTERACTIVE DRIVER ─────────────────────────
214 +
215 +async def interactive_shell():
216 + shell_cmd, prompt_hint = (
217 + ("cmd.exe", "$") if _IS_WIN else ("/bin/bash", "$")
218 + )
219 +
220 + # echo=False → suppress the shell’s own echo of commands
221 + term = TTYSpawn(shell_cmd, echo=False)
222 + await term.start()
223 +
224 + timeout = 1.0
225 +
226 + print(f"Connected to {shell_cmd}.")
227 + print("Type commands for the shell.")
228 + print("• t=<seconds> → change idle timeout")
229 + print("• exit → quit helper\n")
230 +
231 + await term.sendline(" ")
232 + print(await term.read_until_idle(timeout, timeout), end="", flush=True)
233 +
234 + while True:
235 + try:
236 + user = input(f"(timeout={timeout}) {prompt_hint} ")
237 + except (EOFError, KeyboardInterrupt):
238 + print("\nLeaving…")
239 + break
240 +
241 + if user.lower() == "/exit":
242 + break
243 + if user.startswith("/t="):
244 + try:
245 + timeout = float(user.split("=", 1)[1])
246 + print(f"[helper] idle timeout set to {timeout}s")
247 + except ValueError:
248 + print("[helper] invalid number")
249 + continue
250 +
251 + idle_timeout = timeout
252 + total_timeout = 10 * idle_timeout
253 + if user == "":
254 + # Just read output, do not send empty line
255 + async for chunk in term.iter_until_idle(idle_timeout, total_timeout):
256 + print(chunk, end="", flush=True)
257 + else:
258 + await term.sendline(user)
259 + async for chunk in term.iter_until_idle(idle_timeout, total_timeout):
260 + print(chunk, end="", flush=True)
261 +
262 + await term.sendline("exit")
263 + await term.wait()
264 +
265 +if __name__ == "__main__":
266 + asyncio.run(interactive_shell())