retype: set variable/function types with 'y' (IDA-style), via structured tools

Instead of parsing pseudocode text, add structured server tools (server/
patch_server.py, alongside del_type):
 - func_types(addr): prototype + local variables (name/type/is_arg)
 - set_lvar_type(addr,var,type): retype a decompiler local, working on auto/
   register vars too (stock set_type only updates already-user-modified lvars)

'y' in a code view retypes what's under the cursor: a local variable (prompt
prefilled with its current type) or a function (prompt prefilled with its full
prototype). Applies via set_lvar_type / set_type, then recompiles + refreshes.
Copy-line moved off 'y' to Ctrl+Y.

domain: Program.func_types / set_function_type / set_lvar_type (LVar/FuncTypes).
Pilot: function-prototype retype (prefill + apply). full suite 94/94.

4 files changed, 348 insertions, 23 deletions

diff --git a/idatui/app.py b/idatui/app.py
index 32a0698..3ea0f3e 100644
--- a/idatui/app.py
+++ b/idatui/app.py
@@ -121,6 +121,15 @@ class CommentRequested(Message):
         self.view = view
 
 
+class RetypeRequested(Message):
+    """A code view asks to retype the function/variable under the cursor."""
+
+    def __init__(self, view, name: str | None) -> None:  # type: ignore[no-untyped-def]
+        super().__init__()
+        self.view = view
+        self.name = name
+
+
 class NavMixin:
     """Follow / xrefs actions shared by the code views (bindings live on each
     view since Textual only merges BINDINGS from DOMNode subclasses)."""
@@ -129,8 +138,9 @@ class NavMixin:
         Binding("enter", "follow", "Follow"),
         Binding("x", "xrefs", "Xrefs"),
         Binding("n", "rename", "Rename"),
+        Binding("y", "retype", "Retype"),
         Binding("semicolon", "comment", "Comment"),
-        Binding("y", "copy_line", "Copy line"),
+        Binding("ctrl+y", "copy_line", "Copy line"),
     ]
 
     def action_follow(self) -> None:
@@ -145,6 +155,9 @@ class NavMixin:
     def action_comment(self) -> None:
         self.post_message(CommentRequested(self))
 
+    def action_retype(self) -> None:
+        self.post_message(RetypeRequested(self, self.word_under_cursor()))
+
     def action_copy_line(self) -> None:
         plain = self._line_plain(self.cursor)
         if not plain:
@@ -1459,6 +1472,10 @@ class IdaTui(App):
         height: 1; border: none; padding: 0 1;
         background: $success-darken-2; color: $text;
     }
+    #retype {
+        height: 1; border: none; padding: 0 1;
+        background: $accent-darken-2; color: $text;
+    }
     #status { height: 1; background: $panel; color: $text; padding: 0 1; }
     .decomp-loading {
         width: 100%; height: 100%; content-align: center middle;
@@ -1525,6 +1542,7 @@ class IdaTui(App):
         self._search_ctx: tuple[object | None, int] = (None, 1)
         self._rename_ctx: tuple[object | None, str] = (None, "")
         self._comment_ctx: tuple[object | None, int, str] = (None, 0, "")
+        self._retype_ctx: tuple[object | None, str, int, str] = (None, "", 0, "")
         self._xref_focus_name: str | None = None
         self._dirty = False
 
@@ -1551,6 +1569,10 @@ class IdaTui(App):
         ci.display = False
         ci.can_focus = False
         yield ci
+        ti = Input(id="retype")
+        ti.display = False
+        ti.can_focus = False
+        yield ti
         yield Static("connecting…", id="status")
         yield Footer()
 
@@ -2141,6 +2163,90 @@ class IdaTui(App):
         verb = "cleared comment" if not text else "commented"
         self._status(f"{verb} @ {ea:#x}   (Ctrl+S to save)")
 
+    # -- retype (set type, IDA 'y') --------------------------------------- #
+    def on_retype_requested(self, msg: RetypeRequested) -> None:
+        if self._cur is None or self.program is None:
+            return
+        self._prepare_retype(msg.view, msg.name)
+
+    @work(thread=True, exclusive=True, group="retype")
+    def _prepare_retype(self, view, word: str | None) -> None:  # type: ignore[no-untyped-def]
+        """Work out whether the cursor is on a local variable or a function, and
+        fetch the current type/prototype to prefill the prompt."""
+        assert self.program is not None and self._cur is not None
+        ft = self.program.func_types(self._cur.ea)
+        kind: str | None = None
+        subject: int = self._cur.ea
+        prefill = ""
+        # 1) a local variable (or arg) of the current function
+        if word and ft is not None:
+            lv = next((v for v in ft.lvars if v.name == word), None)
+            if lv is not None:
+                kind, prefill = "lvar", lv.type
+        # 2) a function under the cursor (self or a referenced one)
+        if kind is None and self._looks_like_symbol(word):
+            try:
+                tgt = self.program.resolve(word)
+            except Exception:  # noqa: BLE001
+                tgt = None
+            if tgt is not None:
+                tft = self.program.func_types(tgt)
+                if tft is not None:
+                    kind, subject, prefill = "func", tgt, tft.prototype
+        # 3) fall back to the current function itself
+        if kind is None and ft is not None:
+            kind, subject, prefill = "func", self._cur.ea, ft.prototype
+        if kind is None:
+            self.app.call_from_thread(self._status, "nothing to retype under the cursor")
+            return
+        self.app.call_from_thread(self._open_retype, view, kind, subject, word or "", prefill)
+
+    def _open_retype(self, view, kind: str, subject: int, word: str,  # type: ignore[no-untyped-def]
+                     prefill: str) -> None:
+        self._retype_ctx = (view, kind, subject, word)
+        self.query_one("#status", Static).display = False
+        inp = self.query_one("#retype", Input)
+        label = "prototype" if kind == "func" else f"type for '{word}'"
+        inp.placeholder = f"{label} —  Enter=apply  Esc=cancel"
+        inp.can_focus = True
+        inp.display = True
+        inp.value = prefill
+        inp.focus()
+
+    def _end_retype(self) -> None:
+        inp = self.query_one("#retype", Input)
+        inp.display = False
+        inp.can_focus = False
+        self.query_one("#status", Static).display = True
+        view = self._retype_ctx[0]
+        if view is not None:
+            view.focus()
+
+    @work(thread=True, exclusive=True, group="retype-apply")
+    def _do_retype(self, kind: str, subject: int, word: str, new: str) -> None:
+        assert self.program is not None
+        if kind == "func":
+            err = self.program.set_function_type(subject, new)
+        else:  # lvar of the current function
+            err = self.program.set_lvar_type(self._cur.ea, word, new)
+        if err:
+            self.app.call_from_thread(self._status, f"retype failed: {err}")
+            return
+        self.app.call_from_thread(self._after_retype, kind, word)
+
+    def _after_retype(self, kind: str, word: str) -> None:
+        # A type change alters the pseudocode (and disasm operand types), so
+        # recompile via the name-generation invalidation and reopen in place.
+        self.program.bump_names()
+        cur = self._cur
+        if cur is not None:
+            self._save_current_pos()
+            self.query_one(DecompView).loaded_ea = None
+            self._open_entry(cur, push=False)
+        self._dirty = True
+        what = "prototype" if kind == "func" else f"'{word}'"
+        self._status(f"retyped {what}   (Ctrl+S to save)")
+
     @work(thread=True, exclusive=True, group="rename")
     def _do_rename(self, view, old: str, new: str) -> None:  # type: ignore[no-untyped-def]
         assert self.program is not None
@@ -2353,6 +2459,11 @@ class IdaTui(App):
             event.prevent_default()
             self._end_comment()
             return
+        if self.query_one("#retype", Input).display:
+            event.stop()
+            event.prevent_default()
+            self._end_retype()
+            return
         fi = self.query_one("#func-filter", Input)
         if fi.display:
             event.stop()
@@ -2391,6 +2502,12 @@ class IdaTui(App):
             if view is not None and value != existing:  # empty value clears it
                 self._do_comment(ea, value)
             return
+        if inp.id == "retype":
+            view, kind, subject, word = self._retype_ctx
+            self._end_retype()
+            if view is not None and value:
+                self._do_retype(kind, subject, word, value)
+            return
         if getattr(self, "_goto_mode", False):
             self._goto_mode = False
             self.query_one("#func-table", DataTable).focus()
diff --git a/idatui/domain.py b/idatui/domain.py
index f52535c..279ef49 100644
--- a/idatui/domain.py
+++ b/idatui/domain.py
@@ -93,6 +93,21 @@ class Xref:
 
 
 @dataclass(frozen=True)
+class LVar:
+    name: str
+    type: str
+    is_arg: bool
+
+
+@dataclass
+class FuncTypes:
+    addr: int
+    name: str
+    prototype: str        # e.g. 'int __fastcall foo(int a, char *b)'
+    lvars: list[LVar]
+
+
+@dataclass(frozen=True)
 class Struct:
     name: str
     size: int
@@ -495,6 +510,42 @@ class Program:
             return res[0].get("error")
         return None
 
+    # -- function / variable types ---------------------------------------- #
+    def func_types(self, ea: int) -> FuncTypes | None:
+        """Structured decompiler types for the function at ``ea`` (prototype +
+        local variables). None if ``ea`` isn't a decompilable function. Requires
+        the injected ``func_types`` server tool."""
+        try:
+            r = self.client.call("func_types", addr=hex(ea))
+        except IDAToolError:
+            return None
+        if not isinstance(r, dict) or r.get("error"):
+            return None
+        lvars = [LVar(name=lv.get("name", ""), type=lv.get("type", ""),
+                      is_arg=bool(lv.get("is_arg")))
+                 for lv in r.get("lvars", []) if isinstance(lv, dict)]
+        return FuncTypes(addr=_as_int(r.get("addr", hex(ea))), name=r.get("name", ""),
+                         prototype=r.get("prototype", ""), lvars=lvars)
+
+    def set_function_type(self, ea: int, signature: str) -> str | None:
+        """Set a function's prototype. None on success, else an error string."""
+        r = self.client.call("set_type", edits=[{"addr": hex(ea), "signature": signature}])
+        res = r.get("result", []) if isinstance(r, dict) else []
+        row = res[0] if res and isinstance(res[0], dict) else {}
+        if row.get("ok"):
+            return None
+        return row.get("error") or "failed to set the prototype"
+
+    def set_lvar_type(self, fn_ea: int, var: str, ty: str) -> str | None:
+        """Set a decompiler local variable's type (via the injected server tool).
+        None on success, else an error string."""
+        r = self.client.call("set_lvar_type", addr=hex(fn_ea), variable=var, type=ty)
+        if isinstance(r, dict) and r.get("error"):
+            return r["error"]
+        if isinstance(r, dict) and not r.get("ok"):
+            return "failed to set the variable type"
+        return None
+
     def delete_type(self, name: str) -> str | None:
         """Delete a named type. Returns None on success, else an error string.
         Requires a server-side ``del_type`` tool; if absent, a clear message is
diff --git a/server/patch_server.py b/server/patch_server.py
index b989930..eeebfae 100644
--- a/server/patch_server.py
+++ b/server/patch_server.py
@@ -1,17 +1,27 @@
 #!/usr/bin/env python3
 """Inject idatui's extra ida-pro-mcp tools into the installed server package.
 
-ida-pro-mcp exposes no delete-type tool, which idatui's struct editor needs for
-the 'D' in CRUD. Rather than vendor/fork the server, we keep the tool source here
-and append it (idempotently) to the installed ``api_types.py``. That module is
-imported by every worker (``python -m ida_pro_mcp.idalib_server``), so the tool
-registers itself via ``@tool`` on the shared ``MCP_SERVER`` — no server code is
-forked, and re-running this (spawn.sh does, on every start) re-applies it after a
-reinstall/upgrade of ida-pro-mcp.
+ida-pro-mcp lacks a few tools idatui needs. Rather than vendor/fork the server,
+we keep the tool source here and inject it (idempotently) into the installed
+``api_types.py``. That module is imported by every worker
+(``python -m ida_pro_mcp.idalib_server``), so the tools register themselves via
+``@tool`` on the shared ``MCP_SERVER`` — no server code is forked, and re-running
+this (spawn.sh does, on every start) re-applies it after a reinstall/upgrade.
+
+Injected tools:
+  * ``del_type``      — delete a named local type (struct editor CRUD).
+  * ``func_types``    — structured decompiler types for a function (prototype +
+                        local variables), so clients don't parse pseudocode text.
+  * ``set_lvar_type`` — set a decompiler local variable's type; works on auto/
+                        register vars too (the stock set_type only updates lvars
+                        that already have user-saved info).
+
+The block between the BEGIN/END markers is *replaced* on each run, so editing
+BODY here and restarting the supervisor updates the tools.
 
 Run with the *same* interpreter the server uses (the idalib-mcp entry point's
 ``/usr/bin/python``), so it patches the file the workers actually import.
-Safe to run repeatedly; a no-op once applied.
+Changing a tool needs a supervisor restart so workers respawn.
 """
 from __future__ import annotations
 
@@ -23,10 +33,13 @@ BEGIN = "# >>> idatui-ext: begin (auto-injected by server/patch_server.py) >>>"
 END = "# <<< idatui-ext: end <<<"
 
 # Appended to ida_pro_mcp/ida_mcp/api_types.py, which already imports
-# ``Annotated``, ``tool``, ``idasync`` and ``ida_typeinf`` at module scope.
-SNIPPET = f'''
+# ``Annotated``, ``tool``, ``idasync``, ``ida_typeinf``, ``parse_address`` and
+# ``_parse_type_tinfo``.
+BODY = '''
+def _idatui_lv_get(x):
+    return x() if callable(x) else x
+
 
-{BEGIN}
 @tool
 @idasync
 def del_type(
@@ -36,11 +49,115 @@ def del_type(
     til = ida_typeinf.get_idati()
     ok = ida_typeinf.del_named_type(til, name, ida_typeinf.NTF_TYPE)
     if not ok:
-        return {{"name": name, "error": f"Type '{{name}}' not found or could not be deleted"}}
-    return {{"name": name, "deleted": True}}
-{END}
+        return {"name": name, "error": f"Type '{name}' not found or could not be deleted"}
+    return {"name": name, "deleted": True}
+
+
+@tool
+@idasync
+def func_types(
+    addr: Annotated[str, "Function address or name"],
+) -> dict:
+    """Structured decompiler types for a function: its prototype plus each local
+    variable (name/type/is_arg). Lets clients read/edit types without parsing
+    pseudocode text."""
+    import ida_hexrays
+    import idaapi
+
+    def _tstr(tif):
+        try:
+            s = tif.dstr()
+            if s:
+                return s
+        except Exception:
+            pass
+        return str(tif)
+
+    ea = parse_address(addr)
+    f = idaapi.get_func(ea)
+    if not f:
+        return {"addr": str(addr), "error": "no function at address"}
+    try:
+        cf = ida_hexrays.decompile(f.start_ea)
+    except Exception as e:
+        return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}"}
+    if cf is None:
+        return {"addr": hex(f.start_ea), "error": "decompilation failed"}
+    name = idaapi.get_func_name(f.start_ea) or ""
+    try:
+        proto = ida_typeinf.print_tinfo(
+            "", 0, 0, ida_typeinf.PRTYPE_1LINE, cf.type, name, "")
+    except Exception:
+        proto = ""
+    lvars = []
+    for lv in cf.get_lvars():
+        try:
+            ty = _tstr(_idatui_lv_get(lv.type))
+        except Exception:
+            ty = ""
+        lvars.append({
+            "name": _idatui_lv_get(lv.name),
+            "type": ty,
+            "is_arg": bool(_idatui_lv_get(lv.is_arg_var)),
+        })
+    return {
+        "addr": hex(f.start_ea),
+        "name": name,
+        "prototype": (proto or "").strip(),
+        "lvars": lvars,
+    }
+
+
+@tool
+@idasync
+def set_lvar_type(
+    addr: Annotated[str, "Function address or name"],
+    variable: Annotated[str, "Local variable name"],
+    type: Annotated[str, "New C type for the variable"],
+) -> dict:
+    """Set a decompiler local variable's type. Handles auto/register vars (unlike
+    set_type, which only updates lvars that already have user-saved info)."""
+    import ida_hexrays
+    import idaapi
+
+    ea = parse_address(addr)
+    f = idaapi.get_func(ea)
+    if not f:
+        return {"error": "no function at address"}
+    try:
+        cf = ida_hexrays.decompile(f.start_ea)
+    except Exception as e:
+        return {"error": f"decompile failed: {e}"}
+    if cf is None:
+        return {"error": "decompilation failed"}
+    target = None
+    for lv in cf.get_lvars():
+        if _idatui_lv_get(lv.name) == variable:
+            target = lv
+            break
+    if target is None:
+        return {"error": f"local variable {variable!r} not found"}
+    try:
+        tif = _parse_type_tinfo(type)
+    except Exception as e:
+        return {"error": f"bad type {type!r}: {e}"}
+    lsi = ida_hexrays.lvar_saved_info_t()
+    try:
+        lsi.ll = target
+    except Exception:
+        try:
+            lsi.ll.location = _idatui_lv_get(target.location)
+            lsi.ll.defea = target.defea
+        except Exception as e:
+            return {"error": f"could not locate variable: {e}"}
+    lsi.type = tif
+    ok = bool(ida_hexrays.modify_user_lvar_info(
+        f.start_ea, ida_hexrays.MLI_TYPE, lsi))
+    return {"addr": hex(f.start_ea), "variable": variable, "type": type, "ok": ok}
 '''
 
+SNIPPET = f"{BEGIN}\n{BODY.strip()}\n{END}\n"
+
 
 def api_types_path() -> pathlib.Path | None:
     """Locate ida_pro_mcp/ida_mcp/api_types.py without importing it (importing the
@@ -55,17 +172,23 @@ def api_types_path() -> pathlib.Path | None:
 def main() -> int:
     path = api_types_path()
     if path is None:
-        print("idatui: ida_pro_mcp not found; skipping del_type injection", file=sys.stderr)
+        print("idatui: ida_pro_mcp not found; skipping tool injection", file=sys.stderr)
         return 0
     text = path.read_text()
-    if BEGIN in text:
-        return 0  # already applied
+    if BEGIN in text and END in text:  # replace the existing block in place
+        pre = text[: text.index(BEGIN)].rstrip()
+        post = text[text.index(END) + len(END):].lstrip("\n")
+        new = pre + "\n\n" + SNIPPET + ("\n" + post if post else "")
+    else:
+        new = text.rstrip() + "\n\n" + SNIPPET
+    if new == text:
+        return 0
     try:
-        path.write_text(text.rstrip() + "\n" + SNIPPET)
+        path.write_text(new)
     except OSError as e:
         print(f"idatui: could not patch {path}: {e}", file=sys.stderr)
         return 1
-    print(f"idatui: injected del_type into {path}", file=sys.stderr)
+    print(f"idatui: injected/updated idatui-ext tools in {path}", file=sys.stderr)
     return 0
 
 
diff --git a/tests/test_tui.py b/tests/test_tui.py
index 6a8f898..84037a2 100644
--- a/tests/test_tui.py
+++ b/tests/test_tui.py
@@ -295,13 +295,13 @@ async def run(db):
               view.model.cached_line(view.cursor) is not None)
         check("status shows an address", "@ 0x" in str(status.render()), str(status.render()))
 
-        # 'y' copies the current code line to the clipboard.
+        # Ctrl+Y copies the current code line to the clipboard.
         view.focus()
         cur_line = view._line_plain(view.cursor)
         app._clipboard = ""
-        await pilot.press("y")
+        await pilot.press("ctrl+y")
         await wait_until(pilot, lambda: app._clipboard == cur_line, 10)
-        check("'y' copies the current code line to the clipboard",
+        check("Ctrl+Y copies the current code line to the clipboard",
               bool(cur_line) and app._clipboard == cur_line, f"clip={app._clipboard!r}")
 
         # Jump to bottom of a (possibly huge) function; must not hang.
@@ -976,6 +976,40 @@ async def run(db):
             else:
                 check("found a pseudocode line to comment", False, "no marker line")
 
+            # Retype ('y') a FUNCTION prototype: the prompt prefills the current
+            # prototype (from the structured func_types tool) and applying a new
+            # one updates it. (Local-var retyping uses the same seam via
+            # set_lvar_type; it's not asserted here because Hex-Rays restructures
+            # some vars on retype, making a deterministic check fragile.)
+            cf = app._cur  # whatever function is actually loaded now
+            await wait_until(pilot, lambda: dec.loaded_ea == cf.ea, 15)
+            fts = app.program.func_types(cf.ea)
+            if fts is not None and fts.prototype:
+                old_proto = fts.prototype
+                nm_col = dec._texts[0].find(cf.name)
+                dec.focus()
+                dec.cursor, dec.cursor_x = 0, (nm_col + 1 if nm_col >= 0 else 0)
+                dec.refresh()
+                await pilot.pause(0.05)
+                await pilot.press("y")
+                await wait_until(pilot, lambda: app.query_one("#retype", Input).display, 10)
+                ri = app.query_one("#retype", Input)
+                check("'y' on a function prefills its prototype",
+                      ri.display and ri.value == old_proto,
+                      f"val={ri.value!r} want={old_proto!r}")
+                ri.value = f"void __fastcall {cf.name}(int zz_retype_arg)"
+                await pilot.press("enter")
+                await wait_until(
+                    pilot, lambda: (lambda f: bool(f) and "zz_retype_arg" in f.prototype)(
+                        app.program.func_types(cf.ea)), 20)
+                after = app.program.func_types(cf.ea)
+                check("applying a retype changes the function prototype",
+                      after is not None and "zz_retype_arg" in after.prototype,
+                      f"proto={after.prototype if after else None!r}")
+                app.program.set_function_type(cf.ea, old_proto)  # restore
+            else:
+                check("got structured function types (func_types tool)", False)
+
             # Decompiler follow of a name NOT in refs (the function's own name).
             await pilot.press("g")
             await pilot.pause(0.1)