codemode: close the performance gap with the old worker (heads 35x -> 2.2x)

Two changes, both about work that was never ours to do, found by profiling the
A/B benchmark rather than guessing.

1. Serialise inside the database process.

   Code Mode runs to_jsonable() over whatever a snippet returns, walking the
   entire structure to make it JSON-safe. Our answers are already JSON-safe and
   they are large: a 200-row listing page is ~10k small objects, and walking
   them cost 66ms of the page's 92ms -- 114x what json.dumps of the very same
   data costs (0.58ms). Snippets now return one pre-serialised string, so that
   walk is O(1) and the client parses a payload it was going to parse anyway.
   heads(200): 92ms -> 24.7ms.

2. Detach the runtime's trace hook while our snippet runs.

   ida_codemode.runtime wraps every execute_python in sys.settrace(timeout_trace)
   to enforce deadlines, and timeout_trace RETURNS ITSELF -- which switches on
   LINE tracing in every frame it sees. Every line of every function we call
   pays a Python-level callback. Measured here:

     ida_bytes.get_flags   0.106us untraced   5.49us traced      52x
     (plain idalib, no Code Mode: 0.119us -- i.e. untraced == native)
     heads(200 rows)       2.0ms   untraced   20.2ms traced      10x

   That one hook was the entire residual gap against the old unix-socket worker.
   The snippet now detaches it and restores it in a finally.

   What that gives up, stated plainly: the deadline is no longer enforced for a
   pure-Python loop inside our snippet. The runtime's other cancellation path --
   a threading.Timer calling ida_kernwin.set_cancelled() -- does not go through
   the trace and still fires, so a long IDA operation remains interruptible, and
   every operation here is bounded by its own count/limit argument. Set
   IDATUI_CODEMODE_TRACE=1 to keep the stock behaviour.

Against the worker backend, same box, targets/echo (worker -> codemode):

  heads_200          2.65ms ->  5.85ms   2.2x   (was 35x)
  heads_500          6.21ms -> 10.67ms   1.7x
  heads expect-hit   2.16ms ->  4.61ms   2.1x
  disasm_200         9.03ms ->  5.25ms   0.6x   faster
  decompile_cold   162.62ms -> 30.66ms   0.2x   faster
  decompile_warm    30.02ms -> 25.39ms   0.8x   faster
  decomp_map        45.82ms -> 47.85ms   1.0x   parity
  pc_nums           19.90ms -> 22.56ms   1.1x   parity
  rename_func      254.08ms -> 255.54ms  1.0x   parity
  connect          550.0ms  -> 410.0ms   0.7x   faster

What is left is the transport floor: an empty execute_python round trip is
2.0ms, so trivial calls (data_type 0.07ms -> 2.63ms, force_recompile, a single
xref query) look like 40x while being 2.5ms of wall clock. Reducing those needs
fewer calls, not faster ones -- the digest/expect path already does that for the
listing, which is where call volume actually is.

Full suite: 788 passed, 0 failed, 115.3s (was 146.3s; the pilot alone went
80.9s -> 62.2s).

1 files changed, 66 insertions, 5 deletions

diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py
index 9d91335..88d7b31 100644
--- a/idatui/codemode_client.py
+++ b/idatui/codemode_client.py
@@ -22,7 +22,7 @@ import shlex
 import threading
 import time
 from pathlib import Path
-from textwrap import dedent
+from textwrap import dedent, indent
 from typing import Any
 
 from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session
@@ -165,10 +165,63 @@ def _parse_load_args(value: str) -> tuple[str | None, int | None, str | None]:
     return processor, loading_address, file_type
 
 
+#: Key of the pre-serialised payload envelope. See _script().
+_PACKED = "__idatui_json__"
+
+#: Serialise the answer INSIDE the database process and hand back one string.
+#:
+#: Code Mode runs to_jsonable() over whatever a snippet returns, walking the
+#: whole structure to make it JSON-safe. Our answers are already JSON-safe, and
+#: they are big: a 200-row listing page is ~10k small objects, which costs 66ms
+#: to walk -- 72% of the page's total cost, and 114x what json.dumps of the very
+#: same data costs (0.58ms). Returning a STRING makes that walk O(1); the client
+#: parses it, which it was going to do at the transport layer anyway.
+_PACK_EPILOGUE = (
+    '\n{"' + _PACKED + '": json.dumps(result, separators=(",", ":"), default=str)}\n'
+)
+
+
+#: Keep Code Mode's per-line trace hook installed while our snippet runs.
+#: Set IDATUI_CODEMODE_TRACE=1 to restore the stock behaviour.
+_KEEP_TRACE = os.environ.get("IDATUI_CODEMODE_TRACE", "") not in ("", "0")
+
+
 def _script(args: dict[str, Any], body: str) -> str:
-    """Bind JSON arguments without interpolating user text into Python code."""
+    """Bind JSON arguments without interpolating user text into Python code.
+
+    Also runs the body with Code Mode's trace hook detached, which is worth an
+    order of magnitude. The runtime wraps every execute_python in
+    sys.settrace(timeout_trace), and that trace function RETURNS ITSELF, which
+    turns on line tracing in every frame it sees -- so every line of every
+    function we call pays a Python-level callback. Measured on this box:
+    ida_bytes.get_flags is 0.106us untraced (0.119us in a plain idalib process)
+    and 5.49us traced, 52x; a 200-row listing page is 2.0ms untraced and 20.2ms
+    traced. That single hook was the whole residual gap against the old worker.
+
+    What this gives up: the deadline is no longer enforced for a pure-Python
+    loop inside our snippet. The runtime's OTHER cancellation path -- a
+    threading.Timer that calls ida_kernwin.set_cancelled() -- is independent of
+    the trace and still fires, so a long IDA operation is still interruptible;
+    and every operation here is bounded by its own count/limit argument. The
+    trace is restored in a finally, so a raising snippet cannot leak the change.
+    """
     encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":"))
-    return f"import json\na = json.loads({encoded!r})\n{dedent(body).strip()}\n"
+    head = f"import json\na = json.loads({encoded!r})\n"
+    if _KEEP_TRACE:
+        return f"{head}{dedent(body).strip()}\n{_PACK_EPILOGUE}"
+    return (
+        f"{head}"
+        "import sys\n"
+        "_idatui_trace = sys.gettrace()\n"
+        "sys.settrace(None)\n"
+        "try:\n"
+        f"{indent(dedent(body).strip(), '    ')}\n"
+        '    _idatui_packed = {"' + _PACKED + '": json.dumps('
+        'result, separators=(",", ":"), default=str)}\n'
+        "finally:\n"
+        "    sys.settrace(_idatui_trace)\n"
+        "_idatui_packed\n"
+    )
 
 
 _OPERATIONS: dict[str, str] = {
@@ -1179,6 +1232,13 @@ class CodeModeClient:
             raise IDAToolError("execute_python", "Code Mode returned an invalid execution result")
         return response["result"]
 
+    @staticmethod
+    def _unpack(answer: Any) -> Any:
+        """Undo _PACK_EPILOGUE. Anything else passes through untouched."""
+        if isinstance(answer, dict) and _PACKED in answer:
+            return json.loads(answer[_PACKED])
+        return answer
+
     def invoke(self, operation: str, *, timeout: float | None = None, **args) -> Any:
         """Execute one TUI domain operation through Code Mode."""
         if operation in ("idb_save", "save"):
@@ -1189,12 +1249,13 @@ class CodeModeClient:
         if body is None:
             raise IDAToolError(operation, f"unknown ida-tui Code Mode operation: {operation}")
         try:
-            answer = self.execute_python(_script(args, body), timeout=timeout)
+            answer = self._unpack(self.execute_python(_script(args, body), timeout=timeout))
             if isinstance(answer, dict) and answer.get(_NEED_LIB):
                 # First call against this database process (or a restarted one).
                 self.execute_python(_script({"source": _REMOTE_LIB}, _INSTALL_LIB),
                                     timeout=timeout)
-                answer = self.execute_python(_script(args, body), timeout=timeout)
+                answer = self._unpack(
+                    self.execute_python(_script(args, body), timeout=timeout))
             return answer
         except IDAToolError as exc:
             if exc.tool == "execute_python":