diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index 1127dc78b..7e6483371 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -23,9 +23,7 @@ from metagpt.utils.common import create_func_config
 
 
 class BaseWriteAnalysisCode(Action):
-    async def run(
-        self, context: List[Message], plan: Plan = None, task_guide: str = ""
-    ) -> str:
+    async def run(self, context: List[Message], plan: Plan = None) -> str:
         """Run of a code writing action, used in data analysis or modeling
 
         Args:
@@ -85,7 +83,6 @@ class WriteCodeByGenerate(BaseWriteAnalysisCode):
         self,
         context: [List[Message]],
         plan: Plan = None,
-        code_steps: str = "",
         system_msg: str = None,
         **kwargs,
     ) -> str:
@@ -155,11 +152,11 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         self,
         context: List[Message],
         plan: Plan = None,
-        code_steps: str = "",
         data_desc: str = "",
     ) -> str:
         task_type = plan.current_task.task_type
         task = plan.current_task.instruction
+        code_steps = plan.current_task.code_steps
         available_tools = registry.get_all_schema_by_module(task_type)
         available_tools = [
             {k: tool[k] for k in ["name", "description"] if k in tool}
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index de649e857..3260dd43f 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -59,7 +59,7 @@ class MLEngineer(Role):
             logger.info(f"ready to take on task {task}")
 
             # take on current task
-            code, result, success, code_steps = await self._write_and_exec_code()
+            code, result, success = await self._write_and_exec_code()
 
             # ask for acceptance, users can other refuse and change tasks in the plan
             review, task_result_confirmed = await self._ask_review(trigger=ReviewConst.TASK_REVIEW_TRIGGER)
@@ -73,7 +73,6 @@ class MLEngineer(Role):
                 # tick off this task and record progress
                 task.code = code
                 task.result = result
-                task.code_steps = code_steps
                 self.plan.finish_current_task()
                 self.working_memory.clear()
 
@@ -102,7 +101,7 @@ class MLEngineer(Role):
         return rsp
 
     async def _write_and_exec_code(self, max_retry: int = 3):
-        code_steps = (
+        self.plan.current_task.code_steps = (
             await WriteCodeSteps().run(self.plan)
             if self.use_code_steps
             else ""
@@ -121,12 +120,12 @@ class MLEngineer(Role):
             if not self.use_tools or self.plan.current_task.task_type == "other":
                 # code = "print('abc')"
                 code = await WriteCodeByGenerate().run(
-                    context=context, plan=self.plan, code_steps=code_steps, temperature=0.0
+                    context=context, plan=self.plan, temperature=0.0
                 )
                 cause_by = WriteCodeByGenerate
             else:
                 code = await WriteCodeWithTools().run(
-                    context=context, plan=self.plan, code_steps=code_steps, data_desc=""
+                    context=context, plan=self.plan, data_desc=""
                 )
                 cause_by = WriteCodeWithTools
 
@@ -151,7 +150,7 @@ class MLEngineer(Role):
                 if ReviewConst.CHANGE_WORD[0] in review:
                     counter = 0  # redo the task again with help of human suggestions
 
-        return code, result, success, code_steps
+        return code, result, success
 
     async def _ask_review(self, auto_run: bool = None, trigger: str = ReviewConst.TASK_REVIEW_TRIGGER):
         auto_run = auto_run or self.auto_run
diff --git a/metagpt/schema.py b/metagpt/schema.py
index f91922535..8eb7e31ca 100644
--- a/metagpt/schema.py
+++ b/metagpt/schema.py
@@ -78,10 +78,10 @@ class Task(BaseModel):
     dependent_task_ids: list[str] = [] # Tasks prerequisite to this Task
     instruction: str = ""
     task_type: str = ""
+    code_steps: str = ""
     code: str = ""
     result: str = ""
     is_finished: bool = False
-    code_steps: str = ""
 
 
 class Plan(BaseModel):