fix: refresh the AI decompilation when a predicted name is accepted

reai_toolkit/app/coordinators/ai_decomp_coordinator.py

@@ -17,6 +17,7 @@
     RENAME_NOT_READY,
     RenderModel,
     display_rows_for_source_lines,
+    own_function_placeholder,
     render_progress,
     render_stream,
     render_view_with_map,
@@ -339,19 +340,36 @@ def apply_predicted_name(self, name: str) -> None:
             return
 
         if self.ai_decomp_service.update_function_name(ea, name):
-            self.ai_decomp_service.tag_function_as_renamed(name)
-            self._hide_predicted_name()
-            self.refresh_disassembly_view()
+            self._after_function_renamed(ea, name)
             return
 
         final = self.ai_decomp_service.apply_deduped_name(ea, name)
         if final is None:
             self.show_info_dialog(msg=f"Could not rename this function to '{name}'.")
             return
-        self.ai_decomp_service.tag_function_as_renamed(final)
+        self._after_function_renamed(ea, final)
+
+    def _after_function_renamed(self, ea: int, name: str) -> None:
+        self.ai_decomp_service.tag_function_as_renamed(name)
         self._hide_predicted_name()
         self.refresh_disassembly_view()
 
+        placeholder = (
+            own_function_placeholder(self._current_tokenised)
+            if self._current_tokenised is not None
+            else None
+        )
+        if placeholder is None:
+            self.refresh_current()
+            return
+
+        self.ai_decomp_service.apply_overrides(
+            ea=ea,
+            overrides={placeholder: name},
+            on_decomp=lambda response: self._on_decomp_complete(ea, response),
+            on_tokenised=lambda response: self._on_tokenised_complete(ea, response),
+        )
+
     def request_rename(self, display_line: int, word: str) -> None:
         ea = self._current_func_vaddr
         if ea is None:

reai_toolkit/app/coordinators/ai_decomp_render.py

@@ -20,6 +20,8 @@
 
 _IDENT_RE = re.compile(r"[A-Za-z_]\w*")
 
+OWN_FUNCTION_KIND = "own_function"
+
 RENAME_NOT_READY = "The AI decompilation is still loading — try again once it finishes."
 RENAME_NOT_DECOMP_LINE = "This line is not part of the AI decompilation."
 RENAME_NOT_CODE_LINE = (
@@ -251,6 +253,14 @@ def effective_values(tokens: "GetTokensResponse") -> dict[str, str]:
     return values
 
 
+def own_function_placeholder(tokens: "GetTokensResponse") -> Optional[str]:
+    rendered = tokens.placeholder_to_rendered_token or {}
+    for placeholder, token in rendered.items():
+        if token.kind == OWN_FUNCTION_KIND:
+            return placeholder
+    return None
+
+
 def names_token(ident: str, rendered_value: str) -> bool:
     return ident == rendered_value or ident in _IDENT_RE.findall(rendered_value)
 

tests/unit/ai_decomp/test_render.py

@@ -18,6 +18,7 @@
     RENAME_NOT_ON_LINE,
     RENAME_UNRESOLVED,
     display_rows_for_source_lines,
+    own_function_placeholder,
     effective_values,
     find_token,
     index_of_identifier,
@@ -420,3 +421,25 @@ def test_render_stream_of_a_finished_run_keeps_the_source():
 def test_render_progress_without_steps_falls_back_to_status():
     text = render_progress(_wp(step="", step_index=0, steps_total=0, status="PENDING"))
     assert "// PENDING" in text.split("\n")
+
+
+def test_own_function_placeholder_finds_the_token_naming_the_function_itself():
+    tokens = _tokens(
+        rendered={
+            "@@V_v5@@": _rt("v5"),
+            "@@F_self@@": _rt("FUN_140004780", kind="own_function"),
+            "@@F_other@@": _rt("FUN_140001000", kind="function", function_id=9),
+        }
+    )
+
+    assert own_function_placeholder(tokens) == "@@F_self@@"
+
+
+def test_own_function_placeholder_is_none_when_no_own_function_token_exists():
+    tokens = _tokens(rendered={"@@V_v5@@": _rt("v5"), "@@P_a1@@": _rt("a1", kind="param")})
+
+    assert own_function_placeholder(tokens) is None
+
+
+def test_own_function_placeholder_is_none_when_there_are_no_tokens_yet():
+    assert own_function_placeholder(_tokens(rendered={})) is None