diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py
index 1c797b5a2..199544f58 100644
--- a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py
+++ b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py
@@ -482,13 +482,7 @@ def display_output(self, output, display_format, md):
             lang_indicator = r"^```[a-zA-Z0-9]*\n"
             output = re.sub(lang_indicator, "", output)
             output = re.sub(r"\n```$", "", output)
-            new_cell_payload = dict(
-                source="set_next_input",
-                text=output,
-                replace=False,
-            )
-            ip = self.shell
-            ip.payload_manager.write_payload(new_cell_payload)
+            self.shell.set_next_input(output, replace=False)
             return HTML(
                 "AI generated code inserted below ⬇️", metadata=md
             )