diff --git a/vision_agent/agent/vision_agent_coder.py b/vision_agent/agent/vision_agent_coder.py
index 52e0fafa..4ef6b07e 100644
--- a/vision_agent/agent/vision_agent_coder.py
+++ b/vision_agent/agent/vision_agent_coder.py
@@ -186,6 +186,7 @@ def pick_plan(
                 if tool_output.success
                 else "Code execution failed"
             ),
+            "code": DefaultImports.prepend_imports(code),
             # "payload": tool_output.to_json(),
             "status": "completed" if tool_output.success else "failed",
         }