Pi Bridge: fix Windows generations never finalizing (#127)

* pi_bridge: add TALEMATE_PI_BRIDGE_TRACE event-stream instrumentation for #126

* fix(pi_bridge): shut pi down via stdin EOF with tree-kill fallback so Windows generations finalize (#126)

* pi_bridge: address PR #127 review - bound all cleanup waits, survive repeat cancellation, cancellation tests, trace docs

* pi_bridge: address PR #127 round-2 review - incremental stderr buffer, unconditional drain disposal, shielded tree kill, annotations, test handshake

* pi_bridge: address PR #127 round-3 review - catch taskkill spawn failure, partial-stderr regression test, comment placement

CHANGELOG.yaml

@@ -28,6 +28,7 @@
   fixes:
   - "Node Editor: The Build Prompt node's `memory_prompt` input socket and the Generate Response node's `action_type` input socket were silently ignored — only the node property ever took effect, so a graph computing either value dynamically and wiring it in had no effect. Non-empty wired values now win over the property (the property still applies when the socket is unconnected or resolves empty). Note: graphs that already had something connected to these sockets will start honoring the connection."
   - "Scene Library: Deleting the save file of the currently loaded scene is now refused with an error - previously the file (and its version history) was deleted but silently recreated by the next save or autosave, making the delete appear to not work. Load a different scene first."
+  - "Pi Bridge Client: On Windows every generation hung at the very end — the response streamed in fully (visible as tokens-per-second slowly decaying to zero) but was never delivered, leaving the client busy until cancelled. pi is spawned through npm's pi.cmd shim there, and terminating that shim after the response orphaned the underlying node process, which kept the stdio pipes open and blocked the cleanup forever. The client now asks pi to shut down by closing its stdin (pi exits cleanly on stdin EOF on all platforms), force-kills the whole process tree if it doesn't exit within a grace period — so cancelling a generation no longer leaks a node process that keeps the provider request running — and no longer lets cleanup block response delivery."
   - "Pi Bridge Client: On Windows the client could not launch pi at all — the startup model fetch logged an unhelpful '[WinError 2] The system cannot find the file specified' and generations failed the same way, because pi installed via npm is a pi.cmd shim that Windows only executes by its full path. The pi executable is now resolved to its full path before spawning, and if a launch still fails the error names the executable that could not be run. A missing pi install is also no longer reported at startup — it only surfaces on a configured Pi Bridge client."
   - "Client Status: A single failed status check (e.g. a brief connection timeout to a busy local KoboldCpp) no longer flaps a connected client to 'Could not connect' for a few seconds — the client is only marked disconnected after consecutive status failures."
   - "Agent Settings Dialog: Closing the dialog without changing anything no longer re-saves its settings — previously the untouched dialog pushed the copy it was opened with back to the backend, silently reverting any setting changed elsewhere while it was open (such as by the help agent). The help agent likewise refuses to change settings of an agent whose dialog is currently open, asking to close it first."

docs/getting-started/advanced/environment-variables.md

@@ -31,6 +31,7 @@ This variable is read by Vite at build time and re-read at container start in th
 |----------|---------|---------|
 | `TALEMATE_DEBUG` | _unset_ | Set to `1` to enable `DEBUG`-level logging and write errors to a separate error log file. See [Debug logging](debug-logging.md). |
 | `TALEMATE_LOG_PROMPTS` | _unset_ | Set to any non-empty value to write full prompt + response data to `logs/prompt_log.jsonl`. See [Prompt logging](prompt-logging.md). |
+| `TALEMATE_PI_BRIDGE_TRACE` | _unset_ | Set to `1` to log the [Pi Bridge client's](../../user-guide/clients/types/pi-bridge.md) pi event stream as it is consumed (per-event sizes and types plus a heartbeat warning while the stream is silent). Diagnostic aid for generations that stream but never finish. |
 
 `start-backend.sh` and `start-backend.bat` set `TALEMATE_DEBUG=1` automatically; the production `start.sh` / `start.bat` do not.
 

docs/user-guide/clients/types/pi-bridge.md

@@ -89,3 +89,7 @@ Talemate's Docker image ships with pi preinstalled, so the client works without
 - `./pi/auth.json` — pi's stored credentials
 
 The directory persists across container recreations. Inside the container it is exposed via the `PI_CODING_AGENT_DIR` environment variable (`/app/pi`); when building the image manually, the pi version can be overridden with the `PI_VERSION` build argument.
+
+## Troubleshooting
+
+If generations misbehave at the stream level (tokens arrive but the response never finishes, or nothing arrives at all), set `TALEMATE_PI_BRIDGE_TRACE=1` before starting the backend to log the pi event stream as it is consumed — see [Environment variables](../../../getting-started/advanced/environment-variables.md#logging-debugging).

src/talemate/client/pi_bridge.py

@@ -3,6 +3,7 @@ import json
 import os
 import re
 import shutil
+import sys
 import tempfile
 import time
 from typing import Literal
@@ -37,6 +38,11 @@ log = structlog.get_logger("talemate.client.pi_bridge")
 
 PI_BINARY = "pi"
 
+# Set TALEMATE_PI_BRIDGE_TRACE=1 to log the pi event stream as it is consumed
+# (per-event sizes/types plus a silence heartbeat). Diagnostic aid for stream
+# stalls that only reproduce inside talemate (issue #126).
+PI_BRIDGE_TRACE = os.environ.get("TALEMATE_PI_BRIDGE_TRACE", "") == "1"
+
 # pi event lines carry entire messages (message_end includes the full
 # accumulated thinking block, agent_end the whole conversation), which
 # easily exceeds asyncio's 64KB default StreamReader line limit and would
@@ -45,6 +51,14 @@ PI_STDOUT_LIMIT = 2**26
 
 DEFAULT_PROVIDER = "openrouter"
 
+# how long pi gets to exit on its own after stdin closes before the process
+# tree is force-killed
+PI_SHUTDOWN_GRACE = 3.0
+
+# cap on waiting for the stderr drain during cleanup, in case something
+# still holds the pipe open
+PI_STDERR_DRAIN_TIMEOUT = 5.0
+
 THINKING_LEVELS = ["minimal", "low", "medium", "high", "xhigh", "max"]
 
 # provider -> [model ids], populated from `pi --list-models` (which includes
@@ -62,6 +76,69 @@ FETCH_RETRY_COOLDOWN = 30.0
 _models_last_attempt: float | None = None
 
 
+async def kill_pi_process_tree(proc: asyncio.subprocess.Process) -> None:
+    """Force-kill pi and any children it spawned.
+
+    On Windows the spawned process is npm's pi.cmd shim (cmd.exe); killing it
+    directly orphans the node process underneath, which keeps the stdio pipes
+    open forever (issue #126). taskkill /T takes down the whole tree.
+    """
+    try:
+        if sys.platform == "win32":
+            try:
+                killer = await asyncio.create_subprocess_exec(
+                    "taskkill",
+                    "/F",
+                    "/T",
+                    "/PID",
+                    str(proc.pid),
+                    stdout=asyncio.subprocess.DEVNULL,
+                    stderr=asyncio.subprocess.DEVNULL,
+                )
+                await killer.wait()
+            except OSError as e:
+                # degrade loudly - the proc.kill() below only kills the
+                # shim, so a failed taskkill means the node tree survives
+                log.warning("taskkill failed", pid=proc.pid, error=str(e))
+    finally:
+        # runs even if cancellation lands on the taskkill awaits above
+        try:
+            proc.kill()
+        except ProcessLookupError:
+            pass
+        except OSError as e:
+            log.warning("failed to kill pi process", pid=proc.pid, error=str(e))
+
+
+async def terminate_pi_process(proc: asyncio.subprocess.Process) -> None:
+    """End the pi subprocess without orphaning grandchildren.
+
+    Closing stdin asks pi's RPC mode to shut down (it exits on stdin EOF),
+    which ends the actual node process even when spawned through the pi.cmd
+    shim on Windows; the tree kill is the fallback for a process that will
+    not exit on its own.
+    """
+    if proc.returncode is None:
+        proc.stdin.close()
+        # the kill fallbacks in the except branches below are shielded so a
+        # further cancellation cannot interrupt them mid-way - the detached
+        # kill still runs to completion (at the cost of the trailing reap
+        # warning never firing on that path)
+        try:
+            await asyncio.wait_for(proc.wait(), timeout=PI_SHUTDOWN_GRACE)
+        except asyncio.TimeoutError:
+            await asyncio.shield(kill_pi_process_tree(proc))
+        except asyncio.CancelledError:
+            # a repeat cancellation while parked in the grace wait must not
+            # skip the kill, or the pi/node tree leaks with a live request
+            await asyncio.shield(kill_pi_process_tree(proc))
+            raise
+    try:
+        await asyncio.wait_for(proc.wait(), timeout=PI_SHUTDOWN_GRACE)
+    except asyncio.TimeoutError:
+        log.warning("pi process did not exit after kill", pid=proc.pid)
+
+
 def resolve_pi_binary() -> str | None:
     """Absolute path to the pi executable, or None when pi is not installed.
 
@@ -202,6 +279,63 @@ handlers["talemate_started"].connect(on_talemate_started)
 async_signals.get("config.saved").connect(on_config_saved)
 
 
+class PiBridgeTrace:
+    """Per-generation stream statistics logger for TALEMATE_PI_BRIDGE_TRACE."""
+
+    HEARTBEAT_INTERVAL = 5.0
+    # log every Nth delta event; every one would flood the log at
+    # generation speed while the interesting ones are the non-deltas
+    DELTA_LOG_EVERY = 25
+
+    def __init__(self, logger):
+        self.log = logger
+        self.started = time.monotonic()
+        self.lines = 0
+        self.bytes = 0
+        self.deltas = 0
+        self.last_type = "(none)"
+        self.last_line_at = self.started
+
+    def elapsed(self) -> float:
+        return round(time.monotonic() - self.started, 2)
+
+    def line(self, raw_line: bytes, event_type: str):
+        self.lines += 1
+        self.bytes += len(raw_line)
+        self.last_line_at = time.monotonic()
+        is_delta = event_type.startswith("message_update")
+        if is_delta:
+            self.deltas += 1
+        if not is_delta or self.deltas % self.DELTA_LOG_EVERY == 1:
+            self.log.info(
+                "pi_bridge trace: event",
+                elapsed=self.elapsed(),
+                line=self.lines,
+                size=len(raw_line),
+                type=event_type,
+                total_bytes=self.bytes,
+                deltas=self.deltas,
+            )
+        self.last_type = event_type
+
+    async def heartbeat(self, proc):
+        while True:
+            await asyncio.sleep(self.HEARTBEAT_INTERVAL)
+            quiet = time.monotonic() - self.last_line_at
+            if quiet < self.HEARTBEAT_INTERVAL:
+                continue
+            self.log.warning(
+                "pi_bridge trace: no output",
+                elapsed=self.elapsed(),
+                quiet_seconds=round(quiet, 1),
+                lines=self.lines,
+                total_bytes=self.bytes,
+                deltas=self.deltas,
+                last_event=self.last_type,
+                pi_alive=proc.returncode is None,
+            )
+
+
 REASONING_FIELD_GROUP = FieldGroup(
     name="reasoning",
     label="Reasoning",
@@ -460,15 +594,43 @@ class PiBridgeClient(ConcurrentInferenceMixin, ClientBase):
             # drain stderr continuously so a chatty pi process can never fill
             # the OS pipe buffer and deadlock against our stdout readline loop;
             # the collected output feeds the unexpected-exit error message.
-            stderr_task = asyncio.create_task(proc.stderr.read())
+            # drained incrementally into a buffer so whatever pi wrote is
+            # available even when the drain has to be abandoned mid-read.
+            stderr_buf = bytearray()
+
+            async def drain_stderr():
+                while True:
+                    chunk = await proc.stderr.read(4096)
+                    if not chunk:
+                        return
+                    stderr_buf.extend(chunk)
+
+            stderr_task = asyncio.create_task(drain_stderr())
+
+            trace = None
+            heartbeat_task = None
+            if PI_BRIDGE_TRACE:
+                trace = PiBridgeTrace(self.log)
+                heartbeat_task = asyncio.create_task(trace.heartbeat(proc))
+                self.log.info(
+                    "pi_bridge trace: spawned",
+                    argv=argv,
+                    prompt_bytes=len(prompt),
+                    pid=proc.pid,
+                )
 
             try:
                 rpc_command = {"id": "generate", "type": "prompt", "message": prompt}
                 proc.stdin.write((json.dumps(rpc_command) + "\n").encode("utf-8"))
                 await proc.stdin.drain()
+                if trace:
+                    self.log.info(
+                        "pi_bridge trace: prompt command sent",
+                        elapsed=trace.elapsed(),
+                    )
 
                 response_text, reasoning_text = await self._consume_events(
-                    proc, stderr_task
+                    proc, stderr_task, stderr_buf, trace=trace
                 )
 
                 self._reasoning_response = reasoning_text or None
@@ -483,16 +645,35 @@ class PiBridgeClient(ConcurrentInferenceMixin, ClientBase):
 
                 return response_text
             finally:
-                if proc.returncode is None:
-                    try:
-                        proc.kill()
-                    except ProcessLookupError:
-                        pass
-                await proc.wait()
-                # the kill/exit above closes the pipe, so the drain finishes
-                await stderr_task
+                if heartbeat_task:
+                    heartbeat_task.cancel()
+                try:
+                    await terminate_pi_process(proc)
+                    # the exit above closes the pipe, so the drain finishes;
+                    # bounded in case something still holds stderr open.
+                    # skipped when the unexpected-exit path already consumed
+                    # it - if that drain timed out the task is CANCELLED, and
+                    # re-awaiting it would raise CancelledError over the
+                    # in-flight PiBridgeError
+                    if not stderr_task.done():
+                        await asyncio.wait_for(
+                            stderr_task, timeout=PI_STDERR_DRAIN_TIMEOUT
+                        )
+                except asyncio.TimeoutError:
+                    pass
+                finally:
+                    # cancellation out of terminate_pi_process must not
+                    # abandon a pending drain (task-destroyed warnings)
+                    if not stderr_task.done():
+                        stderr_task.cancel()
 
-    async def _consume_events(self, proc, stderr_task) -> tuple[str, str]:
+    async def _consume_events(
+        self,
+        proc: asyncio.subprocess.Process,
+        stderr_task: asyncio.Task,
+        stderr_buf: bytearray,
+        trace: PiBridgeTrace | None = None,
+    ) -> tuple[str, str]:
         """Read the pi RPC event stream until the agent settles.
 
         Returns the accumulated assistant text and thinking content.
@@ -501,14 +682,20 @@ class PiBridgeClient(ConcurrentInferenceMixin, ClientBase):
         reasoning_text = ""
 
         while True:
-            line = await proc.stdout.readline()
-            if not line:
-                stderr = (await stderr_task).decode(errors="replace")
+            raw_line = await proc.stdout.readline()
+            if not raw_line:
+                # bounded: an orphaned grandchild can hold stderr open; the
+                # buffer still carries whatever pi wrote before dying
+                try:
+                    await asyncio.wait_for(stderr_task, timeout=PI_STDERR_DRAIN_TIMEOUT)
+                except asyncio.TimeoutError:
+                    pass
+                stderr = bytes(stderr_buf).decode(errors="replace")
                 raise PiBridgeError(
                     f"pi process exited unexpectedly: {stderr[:500].strip() or 'no error output'}"
                 )
 
-            line = line.decode("utf-8").strip()
+            line = raw_line.decode("utf-8").strip()
             if not line:
                 continue
 
@@ -520,6 +707,13 @@ class PiBridgeClient(ConcurrentInferenceMixin, ClientBase):
 
             event_type = event.get("type")
 
+            if trace:
+                sub = event.get("assistantMessageEvent", {}).get("type", "")
+                trace.line(
+                    raw_line,
+                    f"{event_type}/{sub}" if sub else str(event_type),
+                )
+
             if event_type == "response":
                 if not event.get("success", True):
                     raise PiBridgeError(

tests/test_client_pi_bridge.py

@@ -34,6 +34,8 @@ def _event_line(event: dict) -> bytes:
 class FakeStdin:
     def __init__(self):
         self.written = b""
+        self.closed = False
+        self.closed_event = asyncio.Event()
 
     def write(self, data: bytes):
         self.written += data
@@ -41,20 +43,35 @@ class FakeStdin:
     async def drain(self):
         pass
 
+    def close(self):
+        self.closed = True
+        self.closed_event.set()
+
 
 class FakeProcess:
     """Stands in for the pi RPC subprocess. The stdout reader uses the same
-    line limit the client requests from create_subprocess_exec."""
+    line limit the client requests from create_subprocess_exec. Models a pi
+    that exits as soon as it is waited on (i.e. honors the stdin-EOF
+    shutdown request)."""
 
-    def __init__(self, events: list[dict], stderr: bytes = b""):
+    def __init__(
+        self,
+        events: list[dict],
+        stderr: bytes = b"",
+        stderr_eof=True,
+        stdout_eof=True,
+    ):
+        self.pid = 4242
         self.stdin = FakeStdin()
         self.stdout = asyncio.StreamReader(limit=pi_bridge.PI_STDOUT_LIMIT)
         for event in events:
             self.stdout.feed_data(_event_line(event))
-        self.stdout.feed_eof()
+        if stdout_eof:
+            self.stdout.feed_eof()
         self.stderr = asyncio.StreamReader()
         self.stderr.feed_data(stderr)
-        self.stderr.feed_eof()
+        if stderr_eof:
+            self.stderr.feed_eof()
         self.returncode = None
         self.killed = False
 
@@ -68,6 +85,24 @@ class FakeProcess:
         return self.returncode
 
 
+class HangingProcess(FakeProcess):
+    """A pi that ignores the stdin-EOF shutdown request and only exits when
+    killed (models the orphaned-node hang behind the pi.cmd shim)."""
+
+    def __init__(self, events: list[dict], **kwargs):
+        super().__init__(events, **kwargs)
+        self._exited = asyncio.Event()
+
+    def kill(self):
+        super().kill()
+        self._exited.set()
+
+    async def wait(self):
+        if self.returncode is None:
+            await self._exited.wait()
+        return self.returncode
+
+
 @pytest.fixture
 def spawner(monkeypatch):
     """Patch subprocess creation in the pi_bridge module; tests queue fake
@@ -174,12 +209,128 @@ async def test_generate_sends_prompt_command_and_cleans_up(client, spawner):
     # the resolved binary path is spawned, not the bare name (Windows
     # CreateProcess cannot find npm's pi.cmd shim by bare name)
     assert spawner.calls[0]["args"][0] == "/usr/bin/pi"
-    # the subprocess is terminated after the response
-    assert proc.killed
+    # shutdown is requested via stdin EOF (killing the shim process would
+    # orphan node on Windows) and the process exits on its own
+    assert proc.stdin.closed
+    assert proc.returncode == 0
+    assert not proc.killed
     # the raised stream limit is requested from the subprocess reader
     assert spawner.calls[0]["kwargs"]["limit"] == pi_bridge.PI_STDOUT_LIMIT
 
 
+@pytest.mark.asyncio
+async def test_shutdown_falls_back_to_kill_when_pi_ignores_stdin_eof(
+    client, spawner, monkeypatch
+):
+    """A pi that does not exit on stdin EOF is force-killed after the grace
+    period instead of hanging generate() forever (issue #126)."""
+    monkeypatch.setattr(pi_bridge, "PI_SHUTDOWN_GRACE", 0.01)
+    message = _assistant_message([{"type": "text", "text": "ok"}])
+    proc = HangingProcess(_generation_events(message))
+    spawner.queue(proc)
+
+    assert await client.generate("hi", {}, "conversation") == "ok"
+    assert proc.stdin.closed
+    assert proc.killed
+
+
+@pytest.mark.asyncio
+async def test_shutdown_kills_process_tree_on_windows(client, spawner, monkeypatch):
+    """On Windows the fallback kill must take down the whole tree via
+    taskkill - killing only the pi.cmd shim orphans the node process, which
+    holds the stdio pipes open and hangs the generation (issue #126)."""
+    monkeypatch.setattr(pi_bridge, "PI_SHUTDOWN_GRACE", 0.01)
+    monkeypatch.setattr(pi_bridge.sys, "platform", "win32")
+    message = _assistant_message([{"type": "text", "text": "ok"}])
+    proc = HangingProcess(_generation_events(message))
+    spawner.queue(proc)
+    spawner.queue(FakeProcess([]))  # stands in for the taskkill process
+
+    assert await client.generate("hi", {}, "conversation") == "ok"
+    assert proc.killed
+    assert spawner.calls[1]["args"] == ("taskkill", "/F", "/T", "/PID", "4242")
+
+
+@pytest.mark.asyncio
+async def test_cancelled_generation_shuts_down_pi(client, spawner, monkeypatch):
+    """Cancelling mid-stream must still shut pi down (stdin EOF, then the
+    kill fallback) - on Windows a leaked node process keeps the provider
+    request running (issue #126)."""
+    monkeypatch.setattr(pi_bridge, "PI_SHUTDOWN_GRACE", 0.01)
+    # stream never settles: only the prompt ack arrives, stdout stays open
+    proc = HangingProcess(
+        [{"id": "generate", "type": "response", "command": "prompt", "success": True}],
+        stdout_eof=False,
+    )
+    spawner.queue(proc)
+
+    task = asyncio.ensure_future(client.generate("hi", {}, "conversation"))
+    await asyncio.sleep(0.05)
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert proc.stdin.closed
+    assert proc.killed
+
+
+@pytest.mark.asyncio
+async def test_repeated_cancel_during_grace_wait_still_kills(
+    client, spawner, monkeypatch
+):
+    """A second cancellation delivered while cleanup is parked in the
+    shutdown grace wait must not skip the kill fallback."""
+    monkeypatch.setattr(pi_bridge, "PI_SHUTDOWN_GRACE", 30.0)
+    proc = HangingProcess(
+        [{"id": "generate", "type": "response", "command": "prompt", "success": True}],
+        stdout_eof=False,
+    )
+    spawner.queue(proc)
+
+    task = asyncio.ensure_future(client.generate("hi", {}, "conversation"))
+    await asyncio.sleep(0.05)
+    task.cancel()  # lands in _consume_events, unwinds into cleanup
+    # deterministic handshake: stdin closing means cleanup reached
+    # terminate_pi_process; one extra tick parks it in the grace wait
+    await asyncio.wait_for(proc.stdin.closed_event.wait(), timeout=1)
+    await asyncio.sleep(0)
+    task.cancel()  # impatient repeat cancel
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert proc.killed
+
+
+@pytest.mark.asyncio
+async def test_unexpected_exit_reports_partial_stderr_on_drain_timeout(
+    client, spawner, monkeypatch
+):
+    """When stdout EOFs but stderr never closes, the bounded drain must
+    surface whatever pi wrote before dying (as a PiBridgeError, not a
+    CancelledError from re-awaiting the timed-out drain task)."""
+    monkeypatch.setattr(pi_bridge, "PI_STDERR_DRAIN_TIMEOUT", 0.05)
+    proc = FakeProcess([], stderr=b"Error: partial diagnosis", stderr_eof=False)
+    spawner.queue(proc)
+
+    with pytest.raises(pi_bridge.PiBridgeError, match="partial diagnosis"):
+        await client.generate("hi", {}, "conversation")
+
+
+@pytest.mark.asyncio
+async def test_generate_returns_even_if_stderr_never_closes(
+    client, spawner, monkeypatch
+):
+    """Cleanup must not block response delivery on the stderr drain - an
+    orphaned grandchild holding the pipe open was exactly the Windows hang:
+    the response was fully received but never returned (issue #126)."""
+    monkeypatch.setattr(pi_bridge, "PI_STDERR_DRAIN_TIMEOUT", 0.05)
+    message = _assistant_message([{"type": "text", "text": "ok"}])
+    proc = FakeProcess(_generation_events(message), stderr_eof=False)
+    spawner.queue(proc)
+
+    assert await client.generate("hi", {}, "conversation") == "ok"
+
+
 @pytest.mark.asyncio
 async def test_generate_handles_oversized_event_lines(client, spawner):
     """A long thinking session produces a message_end line far beyond