pane: reap leaked idalib workers; drive: friendlier name resolution (fixes spawn-hang + bare KeyError)

3 files changed, 143 insertions, 6 deletions

diff --git a/idatui/domain.py b/idatui/domain.py
index 3dd21ea..c374b71 100644
--- a/idatui/domain.py
+++ b/idatui/domain.py
@@ -904,9 +904,33 @@ class Program:
         res = payload.get("result", []) if isinstance(payload, dict) else []
         fn = res[0].get("fn") if res and isinstance(res[0], dict) else None
         if not fn:
-            raise KeyError(f"cannot resolve {target!r}")
+            raise KeyError(f"cannot resolve {target!r}{self._name_suggestion(s)}")
         return _as_int(fn["addr"])
 
+    def _name_suggestion(self, s: str) -> str:
+        """Best-effort ' — did you mean …?' hint for a failed name resolve.
+
+        ``lookup_funcs`` matches exact function names only, so a demangled or
+        partial name (``QuaziLies`` for ``_Z9QuaziLiesPcS_ii``) or a data symbol
+        (``checkKey``) resolves to nothing. Surface the substring matches from
+        the function index so the caller can retype the exact name instead of
+        getting a bare ``cannot resolve``. Never raises — suggestions are a
+        nicety, not a contract.
+        """
+        try:
+            idx = self.functions(filter=s)
+            idx.ensure(6)
+            cands = idx.window(0, 6)
+        except Exception:  # noqa: BLE001 -- suggestions are strictly optional
+            return ""
+        if not cands:
+            return (" (no function name contains it; it may be a data symbol or "
+                    "not a function — pass an address like 0x1234)")
+        shown = cands[:5]
+        names = ", ".join(f"{c.name} @ {c.addr:#x}" for c in shown)
+        more = " …" if len(cands) > len(shown) else ""
+        return f" — did you mean: {names}{more}?"
+
     # -- comments ---------------------------------------------------------- #
     def set_comment(self, ea: int, text: str):
         """Set (empty text clears) the comment at ``ea``; affects both the disasm
diff --git a/idatui/pane.py b/idatui/pane.py
index 6ee45ca..af81982 100644
--- a/idatui/pane.py
+++ b/idatui/pane.py
@@ -28,6 +28,7 @@ import argparse
 import json
 import os
 import secrets
+import signal
 import socket
 import subprocess
 import sys
@@ -78,6 +79,65 @@ def _tmux(*args: str) -> str:
 
 
 # --------------------------------------------------------------------------- #
+# idalib worker reaping
+#
+# ``pane stop`` kills the TUI pane, but the ida-pro-mcp supervisor does not
+# reliably reap the ``ida_pro_mcp.idalib_server`` worker it forked for that
+# session. Leaked workers accumulate against IDA_MCP_MAX_WORKERS until the next
+# ``spawn`` blocks forever waiting for a free slot (the TUI comes up but never
+# becomes ready). A worker is only *safe* to reap when no idatui pane is live
+# (then every worker is orphaned) — that mirrors the hand workaround
+# `pkill -f ida_pro_mcp.idalib_server` and avoids killing an in-use analyser.
+# --------------------------------------------------------------------------- #
+_WORKER_PATTERN = r"ida_pro_mcp\.idalib_server"
+
+
+def _worker_pids() -> list[int]:
+    """PIDs of the supervisor's per-binary idalib worker processes.
+
+    Matches the worker module invocation only (not the ``idalib-mcp``
+    supervisor, whose command line does not contain the module path), and
+    never our own PID.
+    """
+    try:
+        out = subprocess.run(["pgrep", "-f", _WORKER_PATTERN],
+                             capture_output=True, text=True)
+    except OSError:
+        return []
+    me = os.getpid()
+    pids: list[int] = []
+    for tok in out.stdout.split():
+        try:
+            pid = int(tok)
+        except ValueError:
+            continue
+        if pid != me:
+            pids.append(pid)
+    return pids
+
+
+def _count_live_panes() -> int:
+    return sum(1 for r in _load_registry() if _pane_alive(r.get("pane", "")))
+
+
+def _reap_orphan_workers(force: bool = False) -> int:
+    """Kill leaked idalib workers when it is safe (no live pane) or ``force``.
+
+    Returns the number of workers signalled. Best-effort; never raises.
+    """
+    if not force and _count_live_panes() > 0:
+        return 0
+    reaped = 0
+    for pid in _worker_pids():
+        try:
+            os.kill(pid, signal.SIGKILL)
+            reaped += 1
+        except OSError:
+            pass
+    return reaped
+
+
+# --------------------------------------------------------------------------- #
 # supervisor (ida-pro-mcp server) — auto-start if down
 # --------------------------------------------------------------------------- #
 def _server_addr(url: str) -> tuple[str, int]:
@@ -148,6 +208,14 @@ def spawn(args) -> int:
         print(f"error: no such binary: {target}", file=sys.stderr)
         return 2
 
+    # Reap workers leaked by previously-stopped/crashed panes so we don't spawn
+    # into a full IDA_MCP_MAX_WORKERS (which makes the new TUI hang forever,
+    # never reaching ready). No-op while any pane is live.
+    reaped = _reap_orphan_workers()
+    if reaped:
+        print(f"reaped {reaped} orphaned idalib worker(s) before spawn",
+              file=sys.stderr)
+
     # make sure the ida-pro-mcp supervisor is up (auto-start it if not)
     srv: dict[str, Any] = {"server_started": False, "server_up": True}
     if not args.no_ensure_server:
@@ -202,9 +270,17 @@ def _q(s: str) -> str:
     return shlex.quote(s)
 
 
-def _wait_ready(sock: str, timeout: float, pane: str) -> dict[str, Any]:
-    """Poll the socket + ping until the TUI reports ready (or timeout)."""
-    deadline = time.time() + timeout
+def _wait_ready(sock: str, timeout: float, pane: str,
+                stuck_after: float = 45.0) -> dict[str, Any]:
+    """Poll the socket + ping until the TUI reports ready (or timeout).
+
+    Emits a one-time hint to stderr if it's still not ready after ``stuck_after``
+    seconds, so a wedged idalib worker / full worker pool surfaces a diagnostic
+    instead of an unexplained silent hang.
+    """
+    start = time.time()
+    deadline = start + timeout
+    warned = False
     last: dict[str, Any] = {"ready": False}
     while time.time() < deadline:
         if not _pane_alive(pane):
@@ -217,6 +293,13 @@ def _wait_ready(sock: str, timeout: float, pane: str) -> dict[str, Any]:
                     return last
             except (OSError, RpcError, ConnectionError):
                 pass
+        if not warned and (time.time() - start) > stuck_after:
+            warned = True
+            why = ("RPC socket not created yet" if not os.path.exists(sock)
+                   else "TUI up but analysis not ready")
+            print(f"still waiting ({int(time.time() - start)}s): {why}. If this "
+                  f"hangs, the idalib worker may be stuck or IDA_MCP_MAX_WORKERS "
+                  f"is full — try `python -m idatui.pane reap`.", file=sys.stderr)
         time.sleep(0.4)
     last = dict(last)
     last["ready"] = False
@@ -255,7 +338,12 @@ def stop(args) -> int:
             except OSError:
                 pass
     _save_registry([r for r in reg if r not in rows])
-    print(json.dumps({"stopped": [r.get("sock") or r.get("pane") for r in rows]}))
+    # Reap the workers those panes leaked (safe: only fires once no pane is live).
+    reaped = _reap_orphan_workers()
+    out = {"stopped": [r.get("sock") or r.get("pane") for r in rows]}
+    if reaped:
+        out["reaped_workers"] = reaped
+    print(json.dumps(out))
     return 0
 
 
@@ -276,10 +364,24 @@ def list_panes(args) -> int:
         alive.append(r)
     if args.prune:
         _save_registry(alive)
+        reaped = _reap_orphan_workers()
+        if reaped:
+            print(f"reaped {reaped} orphaned idalib worker(s)", file=sys.stderr)
     print(json.dumps(alive, indent=2))
     return 0
 
 
+def reap(args) -> int:
+    """Kill leaked idalib workers (safe when no pane is live; --force overrides)."""
+    live = _count_live_panes()
+    n = _reap_orphan_workers(force=args.force)
+    print(json.dumps({"reaped_workers": n, "live_panes": live, "forced": args.force}))
+    if n == 0 and not args.force and live > 0:
+        print(f"note: {live} live pane(s) — not reaping in-use workers; pass "
+              f"--force to reap anyway", file=sys.stderr)
+    return 0
+
+
 def main(argv: list[str]) -> int:
     p = argparse.ArgumentParser(prog="idatui.pane",
                                 description="spawn/manage idatui TUI panes in tmux")
@@ -311,6 +413,11 @@ def main(argv: list[str]) -> int:
     ls.add_argument("--prune", action="store_true", help="drop dead panes (and their sockets)")
     ls.set_defaults(fn=list_panes)
 
+    rp = sub.add_parser("reap", help="kill leaked idalib workers (frees worker slots)")
+    rp.add_argument("--force", action="store_true",
+                    help="reap even while panes are live (may kill an in-use analyser)")
+    rp.set_defaults(fn=reap)
+
     args = p.parse_args(argv)
     return args.fn(args)
 
diff --git a/idatui/rpc.py b/idatui/rpc.py
index f41c734..e363ba3 100644
--- a/idatui/rpc.py
+++ b/idatui/rpc.py
@@ -438,7 +438,13 @@ class RpcServer:
             result = await self._dispatch(method, params)
             return {"id": rid, "result": result}
         except Exception as e:  # noqa: BLE001 — report, never kill the connection
-            return {"id": rid, "error": {"message": f"{type(e).__name__}: {e}"}}
+            # ``str(KeyError("msg"))`` returns ``repr("msg")`` (adds quotes), which
+            # mangles our friendly resolve messages; unwrap the single arg instead.
+            if isinstance(e, KeyError) and len(e.args) == 1 and isinstance(e.args[0], str):
+                msg = e.args[0]
+            else:
+                msg = str(e)
+            return {"id": rid, "error": {"message": f"{type(e).__name__}: {msg}"}}
 
     # -- composed helpers (semantic verbs) -------------------------------- #
     async def _press(self, keys, pred=None, timeout=20.0):