Stop ida-pro-mcp installing a sys.setprofile hook around every tool call. Its deadline mechanism profiles every python call/return so a pure-python tool loop can be interrupted; our tools are call-heavy, so it taxed the whole backend 3.3x. Worker now sets IDA_MCP_TOOL_TIMEOUT_SEC=0 and arms the deadline itself with one polling watchdog thread + ida_kernwin.set_cancelled() (the half that actually frees the IDA main thread). Also rewrote _idatui_spans to jump between colour tags instead of walking characters (byte-identical over 258k real lines).

Result: {"status":"keep","total_ms":26923.9,"lg_boot_ms":762.2,"lg_decomp_ms":2631.7,"lg_graph_ms":941.8,"lg_hex_ms":1052,"lg_index_ms":67.2,"lg_listing_cold_ms":440.6,"lg_listing_warm_ms":530.5,"lg_nav_ms":10598.3,"lg_palette_ms":4.6,"lg_render_ms":227.8,"lg_search_ms":5265.5,"pure_graph_ms":237.9,"sm_boot_ms":535.3,"sm_decomp_ms":631.8,"sm_graph_ms":702.1,"sm_hex_ms":841.1,"sm_index_ms":0,"sm_listing_cold_ms":268.2,"sm_listing_warm_ms":269.4,"sm_nav_ms":443.6,"sm_palette_ms":0.3,"sm_render_ms":277.8,"sm_search_ms":194.4,"fails":0}

2 files changed, 132 insertions, 53 deletions

diff --git a/idatui/worker.py b/idatui/worker.py
index a4e3509..556e69a 100644
--- a/idatui/worker.py
+++ b/idatui/worker.py
@@ -26,8 +26,68 @@ import pickle
 import socket
 import struct
 import sys
+import threading
+import time
 import uuid
 
+#: Seconds a single tool call may run before it is cancelled. 0 disables the
+#: deadline entirely.
+TOOL_TIMEOUT_SEC = float(os.environ.get("IDATUI_TOOL_TIMEOUT_SEC") or 60)
+
+# ida-pro-mcp enforces its own tool deadline by installing a `sys.setprofile`
+# hook for the duration of every call, so that a pure-python loop inside a tool
+# body can be interrupted. That hook runs a python function on EVERY python call
+# and return -- and our tools are exactly the call-heavy kind: `heads` renders
+# hundreds of items per request and measured 92us/row with the hook against
+# 28us/row without it. A 3.3x tax on the whole backend to bound loops that are
+# already bounded by their `count` argument.
+#
+# So: turn the upstream mechanism off and re-arm the half that does the real
+# work ourselves (see _Deadline). ida_kernwin.set_cancelled() is what actually
+# frees the IDA main thread -- decompile, auto_wait, find_bytes and friends poll
+# user_cancelled() and bail within a poll cycle -- and it costs nothing until it
+# fires.
+os.environ["IDA_MCP_TOOL_TIMEOUT_SEC"] = "0"
+
+
+class _Deadline:
+    """A single watchdog thread that cancels a tool call which overruns.
+
+    Arming is two attribute writes, because it is on the path of every call the
+    TUI makes (a scroll is dozens of them). The watchdog polls instead of being
+    signalled for the same reason: waking a thread per call costs more than the
+    0.25s of granularity it buys on a 60s deadline.
+    """
+
+    TICK = 0.25
+
+    def __init__(self, seconds: float) -> None:
+        import ida_kernwin
+        self._kernwin = ida_kernwin
+        self.seconds = seconds
+        self._until: float | None = None
+        t = threading.Thread(target=self._run, name="idatui-deadline",
+                             daemon=True)
+        t.start()
+
+    def _run(self) -> None:
+        while True:
+            time.sleep(self.TICK)
+            until = self._until
+            if until is not None and time.monotonic() >= until:
+                self._until = None
+                # THREAD_SAFE in the IDA SDK; upstream fires it off a Timer too.
+                self._kernwin.set_cancelled()
+
+    def arm(self) -> None:
+        # Clear unconditionally: the flag is sticky, and one left set would make
+        # every later user_cancelled() true forever.
+        self._kernwin.clr_cancelled()
+        self._until = time.monotonic() + self.seconds
+
+    def disarm(self) -> None:
+        self._until = None
+
 
 # --------------------------------------------------------------------------- #
 # framing
@@ -147,6 +207,7 @@ def _open_and_register(binpath: str, load_args: str = ""):
 def serve(sockpath: str, binpath: str, load_args: str = "") -> None:
     tools, module, save = _open_and_register(binpath, load_args)
     sid = uuid.uuid4().hex[:8]
+    deadline = _Deadline(TOOL_TIMEOUT_SEC) if TOOL_TIMEOUT_SEC > 0 else None
 
     def dispatch(name: str, args: dict):
         args = dict(args)
@@ -167,7 +228,14 @@ def serve(sockpath: str, binpath: str, load_args: str = "") -> None:
         fn = tools.get(name)
         if fn is None:
             raise KeyError(f"unknown tool: {name!r}")
-        result = fn(**args)
+        if deadline is None:
+            result = fn(**args)
+        else:
+            deadline.arm()
+            try:
+                result = fn(**args)
+            finally:
+                deadline.disarm()
         # Match the MCP server's structuredContent: a dict passes through, any
         # other return (list/scalar) is wrapped as {"result": ...}. domain.py
         # parses that exact shape (e.g. lookup_funcs -> payload["result"]).
diff --git a/server/patch_server.py b/server/patch_server.py
index 73147d0..b806be2 100644
--- a/server/patch_server.py
+++ b/server/patch_server.py
@@ -364,6 +364,8 @@ def _idatui_tag_map():
 
 _IDATUI_TAGS = None
 _IDATUI_OPND_TAGS = None
+_IDATUI_CTL = None   # re: the three control characters a tagged line can hold
+_IDATUI_WS = None    # re: a run of whitespace, exactly what str.isspace() calls one
 
 
 def _idatui_opnd_tag_map():
@@ -393,75 +395,84 @@ def _idatui_spans(line):
     Unknown tags become 'text' rather than being dropped: a processor module can
     emit a colour we don't classify, and losing the characters would corrupt the
     line."""
-    global _IDATUI_TAGS, _IDATUI_OPND_TAGS
+    global _IDATUI_TAGS, _IDATUI_OPND_TAGS, _IDATUI_CTL, _IDATUI_WS
     import ida_lines
     if _IDATUI_TAGS is None:
         _IDATUI_TAGS = _idatui_tag_map()
     if _IDATUI_OPND_TAGS is None:
         _IDATUI_OPND_TAGS = _idatui_opnd_tag_map()
+    if _IDATUI_CTL is None:
+        import re as _re
+        _IDATUI_CTL = _re.compile("[\\x01\\x02\\x03]")
+        # str.isspace() is true for \\x1c-\\x1f and \\x85 as well as the \\s
+        # class, so spell those out: this substitution has to agree with the
+        # plain-text collapse character for character (checked over every
+        # codepoint) or the row silently loses its highlighting.
+        _IDATUI_WS = _re.compile("[\\\\s\\x1c\\x1d\\x1e\\x1f\\x85]+")
+    tags, opnds, ctl = _IDATUI_TAGS, _IDATUI_OPND_TAGS, _IDATUI_CTL
     on, off, esc = "\x01", "\x02", "\x03"
     addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28))
     addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16))
+    # Jump between control characters and take the text in between as one slice.
+    # A per-character loop here was 68% of the whole `heads` tool: a disasm line
+    # is ~50 characters but only ~15 tags, and everything between two tags is
+    # already exactly one span's worth of text.
     spans, stack, buf = [], [], []   # stack entries: (kind, operand index|None)
+    kind, opnd = "text", None        # state the current run of text belongs to
     i, n = 0, len(line)
-
-    def _opnd():
-        for _k, o in reversed(stack):
-            if o is not None:
-                return o
-        return None
-
-    def flush():
-        if buf:
-            spans.append([stack[-1][0] if stack else "text", "".join(buf),
-                          _opnd()])
-            del buf[:]
-
-    while i < n:
-        ch = line[i]
-        if ch == on and i + 1 < n:
-            tag = line[i + 1]
-            if tag == addr_tag:
-                # An embedded target address, not display text: 16 hex digits
-                # that must not reach the screen.
-                i += 2 + addr_len
-                continue
-            flush()
-            stack.append((_IDATUI_TAGS.get(tag, "text"),
-                          _IDATUI_OPND_TAGS.get(tag)))
-            i += 2
+    for m in ctl.finditer(line):
+        j = m.start()
+        if j < i:                    # inside an address payload / after an esc
             continue
-        if ch == off and i + 1 < n:
-            flush()
-            if stack:
-                stack.pop()
-            i += 2
+        if j + 1 >= n:               # a trailing control char is literal text
+            break
+        ch = line[j]
+        if ch == esc:                # escaped literal: keep the char it guards
+            buf.append(line[i:j])
+            buf.append(line[j + 1])
+            i = j + 2
             continue
-        if ch == esc and i + 1 < n:      # escaped literal
-            buf.append(line[i + 1])
-            i += 2
+        tag = line[j + 1]
+        if ch == on and tag == addr_tag:
+            # An embedded target address, not display text: 16 hex digits that
+            # must not reach the screen. Deliberately NOT a span boundary.
+            buf.append(line[i:j])
+            i = j + 2 + addr_len
             continue
-        buf.append(ch)
-        i += 1
-    flush()
+        buf.append(line[i:j])
+        i = j + 2
+        txt = "".join(buf)
+        if txt:
+            spans.append([kind, txt, opnd])
+        del buf[:]
+        if ch == on:
+            stack.append((kind, opnd))
+            kind = tags.get(tag, "text")
+            o = opnds.get(tag)
+            if o is not None:
+                opnd = o     # operands nest: an inner colour keeps the operand
+        elif stack:
+            kind, opnd = stack.pop()
+        else:
+            kind, opnd = "text", None
+    if i < n:
+        buf.append(line[i:])
+    txt = "".join(buf)
+    if txt:
+        spans.append([kind, txt, opnd])
     # Collapse IDA's column padding EXACTLY as the plain text does. A run of
-    # spaces can straddle two spans, so this walks characters rather than
-    # collapsing each span on its own — otherwise the spans and `text` disagree
-    # about the line and the row silently loses its highlighting.
+    # spaces can straddle two spans, so the leading space of a span is dropped
+    # when the previous one ended in space — otherwise the spans and `text`
+    # disagree about the line and the row silently loses its highlighting.
     out, prev_space = [], False
+    ws = _IDATUI_WS
     for kind, txt, opnd in spans:
-        acc = []
-        for ch in txt:
-            if ch.isspace():
-                if prev_space:
-                    continue
-                acc.append(" ")
-                prev_space = True
-            else:
-                acc.append(ch)
-                prev_space = False
+        acc = ws.sub(" ", txt)
+        if prev_space and acc[:1] == " ":
+            acc = acc[1:]
         if acc:
-            out.append([kind, "".join(acc), opnd])
+            prev_space = acc[-1] == " "
+            out.append([kind, acc, opnd])
     while out and out[0][1] == " ":
         out.pop(0)
     while out and out[-1][1] == " ":