From 556ee9a440a8a0d693d1c405c4fb8c72c9437c1d Mon Sep 17 00:00:00 2001
From: Rohan Agarwal <47861399+roaga@users.noreply.github.com>
Date: Mon, 10 Feb 2025 13:15:11 -0500
Subject: [PATCH] fix(autofix): Fix coding step rethinking (#1896)

We were using the wrong memory key when rethinking in the coding step,
causing it to be missing the memory and error out.
---
 .../autofix/components/coding/component.py       | 16 +++++++++++++---
 src/seer/automation/autofix/tasks.py             |  4 ++--
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/src/seer/automation/autofix/components/coding/component.py b/src/seer/automation/autofix/components/coding/component.py
index a22b71ff2..89a43e398 100644
--- a/src/seer/automation/autofix/components/coding/component.py
+++ b/src/seer/automation/autofix/components/coding/component.py
@@ -177,12 +177,22 @@ def invoke(self, request: CodingRequest) -> CodingOutput | None:
 
             custom_solution = request.solution if isinstance(request.solution, str) else None
 
+            if not request.initial_memory:
+                agent.memory.insert(
+                    0,
+                    Message(
+                        role="user",
+                        content=CodingPrompts.format_fix_msg(
+                            has_tools=not is_obvious,
+                            custom_solution=custom_solution,
+                            mode=request.mode,
+                        ),
+                    ),
+                )
+
             response = agent.run(
                 RunConfig(
                     system_prompt=CodingPrompts.format_system_msg(has_tools=not is_obvious),
-                    prompt=CodingPrompts.format_fix_msg(
-                        has_tools=not is_obvious, custom_solution=custom_solution, mode=request.mode
-                    ),
                     model=AnthropicProvider.model("claude-3-5-sonnet-v2@20241022"),
                     memory_storage_key="code",
                     run_name="Code",
diff --git a/src/seer/automation/autofix/tasks.py b/src/seer/automation/autofix/tasks.py
index d66269cac..a48603fe4 100644
--- a/src/seer/automation/autofix/tasks.py
+++ b/src/seer/automation/autofix/tasks.py
@@ -337,7 +337,7 @@ def receive_user_message(request: AutofixUpdateRequest):
         is_solution_step = step_to_restart.key == "solution_processing"
         is_root_cause_step = step_to_restart.key == "root_cause_analysis_processing"
         memory = (
-            context.get_memory("plan_and_code")
+            context.get_memory("code")
             if is_coding_step
             else (
                 context.get_memory("root_cause_analysis")
@@ -461,7 +461,7 @@ def restart_from_point_with_feedback(
     is_solution_step = step_to_restart.key == "solution_processing"
     is_root_cause_step = step_to_restart.key == "root_cause_analysis_processing"
     memory = (
-        context.get_memory("plan_and_code")
+        context.get_memory("code")
         if is_coding_step
         else (
             context.get_memory("root_cause_analysis")