virtuai-cli 0.8.1__tar.gz → 0.8.3__tar.gz

{virtuai_cli-0.8.1 → virtuai_cli-0.8.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: virtuai-cli
-Version: 0.8.1
+Version: 0.8.3
 Summary: Run VirtuAI deep agents on your local machine
 Author-email: uCloudStore <lmoreno@ucloudstore.com>
 License: Proprietary

{virtuai_cli-0.8.1 → virtuai_cli-0.8.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "virtuai-cli"
-version = "0.8.1"
+version = "0.8.3"
 description = "Run VirtuAI deep agents on your local machine"
 readme = "README.md"
 requires-python = ">=3.11"

{virtuai_cli-0.8.1 → virtuai_cli-0.8.3}/src/virtuai_cli/__init__.py

@@ -1,2 +1,2 @@
 """VirtuAI local CLI."""
-__version__ = "0.8.1"
+__version__ = "0.8.3"

{virtuai_cli-0.8.1 → virtuai_cli-0.8.3}/src/virtuai_cli/chat/tui.py

@@ -131,6 +131,9 @@ class ChatApp(App):
         # Pending attachments to include with the next message — flipped to the
         # multimodal endpoint when non-empty, cleared after a successful send.
         self._pending_attachments: list[Path] = []
+        # Buffer for execute_output chunks that arrive before the matching
+        # tool_start SSE event has mounted the tool card.
+        self._streaming_output: dict[str, list[str]] = {}
 
     # ── Layout ────────────────────────────────────────────────────────────
     def compose(self) -> ComposeResult:
@@ -228,8 +231,22 @@ class ChatApp(App):
                 self._status_message = msg
             self._refresh_status_bar()
 
+        def on_output(frame_id: str, chunk: str) -> None:
+            # Buffer the chunk — the SSE tool_start may not have mounted the
+            # card yet if the subprocess produces output almost immediately.
+            buf = self._streaming_output.setdefault(frame_id, [])
+            buf.append(chunk)
+            # Route directly to the tool card if it already exists.
+            if self._current_turn is not None:
+                card = self._current_turn._tools.get(frame_id)
+                if card is not None and hasattr(card, "append_output"):
+                    card.append_output(chunk)
+
         try:
-            await ws_runner.run_forever(self.server_url, self.token, self.workdir, on_status=on_status)
+            await ws_runner.run_forever(
+                self.server_url, self.token, self.workdir,
+                on_status=on_status, on_output=on_output,
+            )
         except asyncio.CancelledError:
             self._connection_state = "offline"
             self._refresh_status_bar()
@@ -738,14 +755,38 @@ class ChatApp(App):
                 model_id=self.current_model_id,
             )
 
+        # Pump the SSE stream from a background task into a queue.
+        # When the user cancels (Esc), we abandon the pump task without
+        # awaiting it — httpx teardown happens in the background instead of
+        # blocking the event loop while the user is trying to type.
+        _DONE: object = object()
+        queue: asyncio.Queue = asyncio.Queue()
+
+        async def _pump() -> None:
+            try:
+                async for event in stream_iter:
+                    queue.put_nowait(event)
+            except asyncio.CancelledError:
+                pass
+            except Exception as exc:
+                queue.put_nowait({"type": "error", "error": str(exc)})
+            finally:
+                queue.put_nowait(_DONE)
+
+        pump_task = asyncio.create_task(_pump())
+
         try:
-            async for event in stream_iter:
-                await self._handle_event(event, turn)
+            while True:
+                item = await queue.get()
+                if item is _DONE:
+                    break
+                await self._handle_event(item, turn)
                 self._follow_bottom()
         except asyncio.CancelledError:
-            # User hit Esc (or app is shutting down). Skip the expensive
-            # mark_final flush — see widgets.AssistantTurn.mark_final.
+            # User hit Esc. Cancel the pump but don't await it — let httpx
+            # close the connection in the background.
             cancelled = True
+            pump_task.cancel()
             raise
         except Exception as exc:
             try:
@@ -792,11 +833,12 @@ class ChatApp(App):
             # story. Empty == "", "{}", or anything that strips to nothing.
             input_val = event.get("input", "")
             input_str = input_val if isinstance(input_val, str) else json.dumps(input_val, default=str)
+            tool_id = event.get("id", "")
             if input_str.strip() in ("", "{}", "[]"):
-                turn._tools[event.get("id", "")] = None  # type: ignore[assignment]
+                turn._tools[tool_id] = None  # type: ignore[assignment]
             else:
                 await turn.start_tool(
-                    event.get("id", ""),
+                    tool_id,
                     event.get("name", "tool"),
                     input_val,
                     display=event.get("display", ""),
@@ -805,6 +847,14 @@ class ChatApp(App):
                     read_offset=event.get("read_offset"),
                     read_limit=event.get("read_limit"),
                 )
+                # Flush any execute_output chunks that raced ahead of this
+                # tool_start event (can happen when a subprocess starts
+                # printing immediately, e.g. a fast SSH connection).
+                card = turn._tools.get(tool_id)
+                if card is not None:
+                    for buffered_chunk in self._streaming_output.pop(tool_id, []):
+                        if hasattr(card, "append_output"):
+                            card.append_output(buffered_chunk)
         elif etype == "tool_end":
             # If we skipped the matching tool_start (empty input), don't try
             # to finalize a card that doesn't exist.

{virtuai_cli-0.8.1 → virtuai_cli-0.8.3}/src/virtuai_cli/chat/widgets.py

@@ -221,6 +221,7 @@ class TextSegment(Static):
         self._last_render_t: float = 0.0
         self._render_pending: bool = False
         self._render_timer: Any = None
+        self._cancelled: bool = False
 
     def append(self, chunk: str) -> None:
         self._buf += chunk
@@ -244,6 +245,7 @@ class TextSegment(Static):
     def cancel_pending(self) -> None:
         """Drop any scheduled render — used on stream cancel so the cancel
         doesn't trigger another expensive Markdown reparse a moment later."""
+        self._cancelled = True
         if self._render_timer is not None:
             try:
                 self._render_timer.stop()
@@ -253,6 +255,11 @@ class TextSegment(Static):
         self._render_pending = False
 
     def _flush(self) -> None:
+        # Guard: the timer may have already fired and queued this callback
+        # before cancel_pending() ran. Bail out so we don't do an expensive
+        # Markdown parse right as the user is trying to type.
+        if self._cancelled:
+            return
         self._render_pending = False
         self._render_timer = None
         self._last_render_t = time.monotonic()
@@ -439,6 +446,8 @@ class ToolCallCard(Static):
         self.tool_name = tool_name
         self.input_str = input_str or ""
         self.output_preview: Optional[str] = None
+        self._stream_buf: str = ""
+        self._last_stream_t: float = 0.0
         self.add_class("-running")
         self._redraw()
 
@@ -452,8 +461,24 @@ class ToolCallCard(Static):
         if self.output_preview:
             for line in self.output_preview.strip().splitlines()[:8]:
                 body_lines.append(Text(f"  {line[:160]}", style="dim"))
+        elif self._stream_buf:
+            # Show the last 8 lines of live streaming output while the tool runs.
+            for line in self._stream_buf.splitlines()[-8:]:
+                body_lines.append(Text(f"  {line[:160]}", style="dim cyan"))
         self.update(Text("\n").join(body_lines))
 
+    def append_output(self, chunk: str) -> None:
+        """Append a live streaming chunk from the WS execute_output frame."""
+        if "-done" in self.classes or "-error" in self.classes:
+            return
+        self._stream_buf += chunk
+        # Throttle redraws to ~10 Hz — fast enough to feel live, cheap enough
+        # not to stall the event loop on a command that floods stdout.
+        now = time.monotonic()
+        if now - self._last_stream_t >= 0.1:
+            self._last_stream_t = now
+            self._redraw()
+
     def finish(self, output_preview: Optional[str], errored: bool = False) -> None:
         self.output_preview = output_preview
         self.remove_class("-running")

{virtuai_cli-0.8.1 → virtuai_cli-0.8.3}/src/virtuai_cli/runner.py

@@ -9,7 +9,7 @@ import platform
 import sys
 import time
 from pathlib import Path
-from typing import Awaitable, Callable, Optional
+from typing import Callable, Optional
 
 _AUDIT_DIR = Path.home() / ".virtuai"
 
@@ -23,6 +23,8 @@ from virtuai_cli.security import resolve_workdir
 # Callback the TUI uses to capture connection events (level, message) instead
 # of printing to stdout (which would scramble the Textual screen).
 StatusCallback = Callable[[str, str], None]
+# Callback for streaming execute output: (frame_id, chunk)
+OutputCallback = Callable[[str, str], None]
 
 logger = logging.getLogger(__name__)
 console = Console()
@@ -40,6 +42,7 @@ def _audit(workdir: Path, command: str, exit_code: int, elapsed: float) -> None:
         pass
 
 _RECONNECT_DELAYS = [1, 2, 4, 8, 16, 30]  # backoff steps
+_RECV_TIMEOUT = 60.0  # seconds of silence before declaring the connection stale
 
 
 def _default_on_status(level: str, msg: str) -> None:
@@ -47,7 +50,13 @@ def _default_on_status(level: str, msg: str) -> None:
     console.print(f"[{style}]{msg}[/{style}]" if style else msg)
 
 
-async def _handle_frame(frame: dict, workdir: Path, ws, on_status: StatusCallback) -> None:
+async def _handle_frame(
+    frame: dict,
+    workdir: Path,
+    ws,
+    on_status: StatusCallback,
+    on_output: Optional[OutputCallback] = None,
+) -> None:
     """Dispatch a server frame and send the result back."""
     ftype = frame.get("type")
     fid = frame.get("id")
@@ -76,19 +85,97 @@ async def _handle_frame(frame: dict, workdir: Path, ws, on_status: StatusCallbac
         timeout = int(frame.get("timeout_secs", 300))
         preview = command[:80] + ("..." if len(command) > 80 else "")
         on_status("exec", preview)
+
+        # Security checks — same policy as the sync executor.
+        from virtuai_cli.security import check_command, scrub_env
+        denial = check_command(command)
+        if denial:
+            await ws.send(json.dumps({
+                "v": 1, "id": fid, "type": "execute_result",
+                "output": f"blocked by local denylist: {denial}",
+                "exit_code": 126,
+            }))
+            return
+
+        env = scrub_env(workdir)
+        wrapped = exec_._jail_wrap(command, workdir)
         t0 = time.monotonic()
-        loop = asyncio.get_event_loop()
-        result = await loop.run_in_executor(
-            None, lambda: exec_.execute(command, workdir, timeout=timeout)
-        )
+
+        try:
+            proc = await asyncio.create_subprocess_shell(
+                wrapped,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.STDOUT,
+                cwd=str(workdir),
+                env=env,
+            )
+        except Exception as exc:
+            await ws.send(json.dumps({
+                "v": 1, "id": fid, "type": "execute_result",
+                "output": f"executor error: {exc}",
+                "exit_code": 1,
+            }))
+            return
+
+        chunks: list[str] = []
+        total_bytes = 0
+        timed_out = False
+
+        try:
+            async with asyncio.timeout(timeout):
+                async for raw_line in proc.stdout:
+                    chunk = raw_line.decode("utf-8", errors="replace")
+                    chunks.append(chunk)
+                    total_bytes += len(chunk)
+                    if on_output and fid:
+                        on_output(fid, chunk)
+                    # Forward chunk to server — server may show it in the agent
+                    # trace or buffer it; either way the agent gets the full
+                    # execute_result below.
+                    try:
+                        await ws.send(json.dumps({
+                            "v": 1, "id": fid, "type": "execute_output",
+                            "output": chunk,
+                        }))
+                    except Exception:
+                        pass
+                    if total_bytes > exec_._MAX_OUTPUT_BYTES:
+                        proc.kill()
+                        break
+        except TimeoutError:
+            timed_out = True
+            try:
+                proc.kill()
+            except Exception:
+                pass
+
+        try:
+            await asyncio.wait_for(proc.wait(), timeout=5.0)
+        except Exception:
+            pass
+
+        exit_code = proc.returncode if proc.returncode is not None else (124 if timed_out else 1)
         elapsed = time.monotonic() - t0
-        _audit(workdir, command, result["exit_code"], elapsed)
+
+        full_output = "".join(chunks)
+        if timed_out:
+            full_output += f"\ncommand timed out after {timeout}s"
+            exit_code = 124
+        elif total_bytes > exec_._MAX_OUTPUT_BYTES:
+            keep = exec_._MAX_OUTPUT_BYTES // 2
+            full_output = (
+                full_output[:keep]
+                + f"\n\n... [output truncated: {total_bytes:,} bytes total] ...\n\n"
+                + full_output[-keep:]
+            )
+
+        _audit(workdir, command, exit_code, elapsed)
         await ws.send(json.dumps({
             "v": 1,
             "id": fid,
             "type": "execute_result",
-            "output": result["output"],
-            "exit_code": result["exit_code"],
+            "output": full_output,
+            "exit_code": exit_code,
         }))
         return
 
@@ -129,6 +216,7 @@ async def run_session(
     token: str,
     workdir: Path,
     on_status: Optional[StatusCallback] = None,
+    on_output: Optional[OutputCallback] = None,
 ) -> None:
     """Run a single connected session until the WebSocket closes."""
     on_status = on_status or _default_on_status
@@ -139,10 +227,17 @@ async def run_session(
     # JSON pings every 20 s to keep the connection alive (cli_api/router.py).
     ssl_ctx = _ssl_context() if ws_url.startswith("wss://") else None
     async with websockets.connect(url, ping_interval=None, ssl=ssl_ctx) as ws:
-        async for raw in ws:
+        while True:
+            try:
+                raw = await asyncio.wait_for(ws.recv(), timeout=_RECV_TIMEOUT)
+            except asyncio.TimeoutError:
+                # No frame received within the timeout window — the server
+                # sends a JSON ping every ~20 s, so 60 s of silence means the
+                # TCP connection is stale without a clean close.
+                raise ConnectionError("No data received for 60 s — stale connection")
             try:
                 frame = json.loads(raw)
-                await _handle_frame(frame, workdir, ws, on_status)
+                await _handle_frame(frame, workdir, ws, on_status, on_output)
             except SystemExit:
                 raise
             except Exception as exc:
@@ -154,11 +249,13 @@ async def run_forever(
     token: str,
     workdir: Path,
     on_status: Optional[StatusCallback] = None,
+    on_output: Optional[OutputCallback] = None,
 ) -> None:
     """Connect with exponential backoff, reconnecting on drop.
 
     Pass `on_status` to redirect connection messages to a TUI/log sink
     instead of printing to stdout (which would scramble a Textual screen).
+    Pass `on_output` to receive streaming execute_output chunks (frame_id, chunk).
     """
     on_status = on_status or _default_on_status
     ws_url = server_url.rstrip("/").replace("https://", "wss://").replace("http://", "ws://")
@@ -170,7 +267,7 @@ async def run_forever(
     attempt = 0
     while True:
         try:
-            await run_session(ws_url, token, workdir, on_status)
+            await run_session(ws_url, token, workdir, on_status, on_output)
             attempt = 0  # successful session resets backoff
         except SystemExit:
             return

{virtuai_cli-0.8.1 → virtuai_cli-0.8.3}/src/virtuai_cli.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: virtuai-cli
-Version: 0.8.1
+Version: 0.8.3
 Summary: Run VirtuAI deep agents on your local machine
 Author-email: uCloudStore <lmoreno@ucloudstore.com>
 License: Proprietary