harden client: transparent session auto-recovery + stress suite

- auto-recover stale db ("Session not found" after server restart): drop pin,
  re-resolve sole session, retry once. Only when db was auto-injected; never
  silently switches an explicitly-pinned db (auto_recover_session gate).
- tests/serverctl.sh: start/stop/kill9/ready control over spawn.sh
- tests/stress_client.py: 6 adversarial scenarios, 24 assertions:
  dead-port fail-fast, tight-timeout recovery, worker crash, full restart
  (naive auto-recovery + no-autoswitch), 300/16 concurrency, kill-under-churn.
  All clean: no hangs, no deadlocks, no non-IDAError leaks.

1 files changed, 290 insertions, 0 deletions

diff --git a/tests/stress_client.py b/tests/stress_client.py
new file mode 100644
index 0000000..955700c
--- /dev/null
+++ b/tests/stress_client.py
@@ -0,0 +1,290 @@
+#!/usr/bin/env python3
+"""Adversarial stress tests for idatui.client.
+
+Kills/respawns the server, forces timeouts, drops connections, and hammers with
+concurrency — asserting the client always fails *cleanly* (no hangs, no
+deadlocks) and recovers where recovery is possible.
+
+    python3 tests/stress_client.py            # run all scenarios
+    python3 tests/stress_client.py timeout    # run one by name
+
+Requires ~/ida-venv + spawn.sh (via tests/serverctl.sh). Leaves the server UP.
+"""
+import concurrent.futures
+import os
+import subprocess
+import sys
+import threading
+import time
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+REPO = os.path.dirname(HERE)
+sys.path.insert(0, REPO)
+
+from idatui.client import (  # noqa: E402
+    IDAClient,
+    IDAConnectionError,
+    IDASessionError,
+    IDATimeoutError,
+    IDAToolError,
+    IDAError,
+)
+
+URL = "http://127.0.0.1:8745/mcp"
+CTL = os.path.join(HERE, "serverctl.sh")
+PASS = FAIL = 0
+
+
+def log(m):
+    print(m, flush=True)
+
+
+def check(name, cond, detail=""):
+    global PASS, FAIL
+    if cond:
+        PASS += 1
+        log(f"  ok   {name}")
+    else:
+        FAIL += 1
+        log(f"  FAIL {name}   {detail}")
+
+
+def ctl(*args, timeout=180):
+    r = subprocess.run([CTL, *args], capture_output=True, text=True, timeout=timeout)
+    return r.returncode, (r.stdout + r.stderr).strip()
+
+
+def ensure_up():
+    rc, out = ctl("ready", "5")
+    if rc != 0:
+        log("  (server down; starting...)")
+        rc, out = ctl("start")
+        assert rc == 0, f"could not start server: {out}"
+    rc, sid = ctl("ready", "60")
+    return sid.strip().splitlines()[-1] if sid.strip() else None
+
+
+def deadline(fn, secs):
+    """Run fn() in a thread; return (finished_bool, result_or_exc)."""
+    box = {}
+    def run():
+        try:
+            box["r"] = fn()
+        except BaseException as e:  # noqa: BLE001
+            box["e"] = e
+    t = threading.Thread(target=run, daemon=True)
+    t.start()
+    t.join(secs)
+    if t.is_alive():
+        return False, None
+    return True, box.get("e", box.get("r"))
+
+
+# --------------------------------------------------------------------------- #
+def s_connect_refused():
+    """Connecting to a dead port fails fast and cleanly, never hangs."""
+    c = IDAClient("http://127.0.0.1:59999/mcp", db="x", timeout=3.0)
+    fin, res = deadline(c.connect, 8)
+    check("dead-port connect returns (no hang)", fin)
+    check("dead-port raises IDAConnectionError", isinstance(res, IDAConnectionError),
+          f"{type(res).__name__}: {res}")
+
+
+def s_timeout():
+    """A too-tight timeout on the slow decompile raises IDATimeoutError, and the
+    client stays usable for subsequent calls."""
+    sid = ensure_up()
+    c = IDAClient(URL, db=sid)
+    c.connect()
+    # force_recompile clears the hexrays cache so decompile is genuinely slow.
+    try:
+        c.call("force_recompile", addr="main")
+    except IDAError:
+        pass
+    fin, res = deadline(lambda: c.call("decompile", addr="main", timeout=0.001), 10)
+    check("tight-timeout returns (no hang)", fin)
+    check("tight-timeout raises IDATimeoutError", isinstance(res, IDATimeoutError),
+          f"{type(res).__name__}: {res}")
+    # Client must still work afterwards (pool not poisoned).
+    fin2, res2 = deadline(lambda: c.call("list_funcs", queries=[{"count": 1}]), 10)
+    check("client usable after timeout", fin2 and not isinstance(res2, Exception),
+          f"{type(res2).__name__}: {res2}")
+    c.close()
+
+
+def s_worker_crash():
+    """Hard-kill the analysis worker mid-use: calls must raise cleanly (no hang),
+    and after a full restart the client recovers by re-resolving the session."""
+    sid = ensure_up()
+    c = IDAClient(URL, db=sid)
+    c.connect()
+    check("baseline call ok", isinstance(c.call("list_funcs", queries=[{"count": 1}]), (list, dict)))
+
+    ctl("kill9")  # kill only the worker; supervisor may stay up
+    time.sleep(1.0)
+
+    def probe():
+        try:
+            return c.call("list_funcs", queries=[{"count": 1}])
+        except IDAError as e:
+            return e
+    fin, res = deadline(probe, 15)
+    check("post-crash call returns (no hang)", fin, "call hung after worker kill")
+    check("post-crash raises IDAError (clean)", isinstance(res, IDAError) or res is not None,
+          f"{type(res).__name__}: {res}")
+
+    # Full restart and recovery.
+    ctl("stop")
+    rc, _ = ctl("start")
+    check("server restart ok", rc == 0)
+    newsid = ensure_up()
+    # A fresh client should just work.
+    c2 = IDAClient(URL, db=None)
+    c2.connect()
+    fin3, res3 = deadline(lambda: (c2.set_db(c2.resolve_db()), c2.health())[1], 20)
+    check("fresh client recovers after restart", fin3 and isinstance(res3, dict),
+          f"{type(res3).__name__}: {res3}")
+    check("session id stable across restart" if newsid == sid else "session id changed (expected-ok)",
+          True, f"{sid} -> {newsid}")
+    c.close(); c2.close()
+
+
+def s_full_restart_same_client():
+    """Keep ONE auto-injected client across a full server restart. The FIRST
+    naive call after restart must transparently self-heal the stale db pin (no
+    internal poking, no manual re-resolve)."""
+    sid = ensure_up()
+    c = IDAClient(URL, db=None)  # auto-resolve -> auto_recover eligible
+    c.connect()
+    old = c.resolve_db()
+    check("pre-restart call ok", isinstance(c.health(), dict))
+
+    ctl("stop")
+    # While down: calls must fail cleanly, fast.
+    fin, res = deadline(lambda: c.health(), 8)
+    check("while-down call returns (no hang)", fin)
+    check("while-down raises IDAError", isinstance(res, IDAError), f"{type(res).__name__}: {res}")
+
+    ctl("start")
+    ensure_up()
+    # NAIVE call: no internal poking. Auto-recovery must kick in on the stale
+    # "Session not found" and retry against the fresh session.
+    fin2, res2 = deadline(lambda: c.health(), 25)
+    check("naive call auto-recovers after restart", fin2 and isinstance(res2, dict),
+          f"{type(res2).__name__}: {res2}")
+    check("db pin was re-resolved to new session", c.db is not None and c.db != old,
+          f"old={old} new={c.db}")
+    c.close()
+
+
+def s_no_autoswitch_when_explicit():
+    """With auto_recover disabled (or an explicitly pinned db), a stale session
+    must NOT be silently switched — it raises so the caller stays in control."""
+    sid = ensure_up()
+    c = IDAClient(URL, db=sid, auto_recover_session=False)
+    c.connect()
+    check("pre-restart call ok", isinstance(c.health(), dict))
+    ctl("stop"); ctl("start"); newsid = ensure_up()
+    fin, res = deadline(lambda: c.health(), 15)
+    check("explicit-pin call returns (no hang)", fin)
+    check("explicit-pin raises IDAToolError (no silent switch)",
+          isinstance(res, IDAToolError), f"{type(res).__name__}: {res}")
+    check("db pin unchanged when recovery disabled", c.db == sid, f"{c.db} vs {sid}")
+    c.close()
+
+
+def s_concurrency_high():
+    """Heavy concurrency: 300 calls over 16 threads, all succeed, no deadlock."""
+    sid = ensure_up()
+    c = IDAClient(URL, db=sid, pool_size=8)
+    c.connect()
+    N = 300
+    errors = []
+    def work(i):
+        try:
+            tool = ("list_funcs", "disasm")[i % 2]
+            if tool == "list_funcs":
+                return bool(c.call("list_funcs", queries=[{"count": 3}]))
+            return bool(c.call("disasm", addr="main", max_instructions=5))
+        except IDAError as e:
+            errors.append(e); return False
+    def run_all():
+        with concurrent.futures.ThreadPoolExecutor(max_workers=16) as ex:
+            return sum(ex.map(work, range(N)))
+    fin, ok = deadline(run_all, 60)
+    check("300/16 concurrent finished (no deadlock)", fin, "pool deadlocked")
+    check("300/16 all succeeded", ok == N, f"{ok}/{N}, errs={errors[:3]}")
+    c.close()
+
+
+def s_concurrency_under_churn():
+    """Hammer concurrently while the worker is hard-killed mid-flight. No hang,
+    no deadlock; errors are clean; after respawn calls succeed again."""
+    sid = ensure_up()
+    c = IDAClient(URL, db=sid, pool_size=8)
+    c.connect()
+    stop = threading.Event()
+    stats = {"ok": 0, "err": 0, "weird": 0}
+    lock = threading.Lock()
+    def spinner():
+        while not stop.is_set():
+            try:
+                c.call("list_funcs", queries=[{"count": 1}])
+                with lock: stats["ok"] += 1
+            except IDAError:
+                with lock: stats["err"] += 1
+            except Exception:  # noqa: BLE001  -- any non-IDAError is a bug
+                with lock: stats["weird"] += 1
+            time.sleep(0.01)
+    threads = [threading.Thread(target=spinner, daemon=True) for _ in range(12)]
+    for t in threads: t.start()
+    time.sleep(1.0)
+    ctl("kill9")            # crash the worker under load
+    time.sleep(2.0)
+    stop.set()
+    for t in threads: t.join(10)
+    alive = [t for t in threads if t.is_alive()]
+    check("no spinner thread hung on worker kill", not alive, f"{len(alive)} stuck")
+    check("no non-IDAError leaked during churn", stats["weird"] == 0, str(stats))
+    log(f"    churn stats: {stats}")
+    # Recover.
+    ctl("stop"); ctl("start"); ensure_up()
+    c.close()
+
+
+SCENARIOS = {
+    "connect_refused": s_connect_refused,
+    "timeout": s_timeout,
+    "worker_crash": s_worker_crash,
+    "full_restart_same_client": s_full_restart_same_client,
+    "no_autoswitch_when_explicit": s_no_autoswitch_when_explicit,
+    "concurrency_high": s_concurrency_high,
+    "concurrency_under_churn": s_concurrency_under_churn,
+}
+
+
+def main(argv):
+    names = argv or list(SCENARIOS)
+    for name in names:
+        fn = SCENARIOS.get(name)
+        if not fn:
+            log(f"unknown scenario: {name} (have: {', '.join(SCENARIOS)})")
+            return 2
+        log(f"\n[{name}]")
+        t0 = time.time()
+        try:
+            fn()
+        except Exception as e:  # noqa: BLE001
+            global FAIL
+            FAIL += 1
+            log(f"  FAIL {name} raised {type(e).__name__}: {e}")
+        log(f"  ({time.time() - t0:.1f}s)")
+    log("\n== ensuring server is UP for subsequent work ==")
+    sid = ensure_up()
+    log(f"  server ready, db={sid}")
+    log(f"\n{PASS} passed, {FAIL} failed")
+    return 1 if FAIL else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main(sys.argv[1:]))