fix completing steps when they are not ordered by type

core/agents/code_monkey.py

@@ -208,7 +208,7 @@ class CodeMonkey(BaseAgent):
         await self.ui.generate_diff(file_path, old_content, new_content, n_new_lines, n_del_lines)
 
         await self.state_manager.save_file(file_path, new_content)
-        self.next_state.complete_step()
+        self.next_state.complete_step("save_file")
 
         input_required = self.state_manager.get_input_required(new_content)
         if input_required:

core/agents/executor.py

@@ -176,4 +176,4 @@ class Executor(BaseAgent):
         information we give it.
         """
         self.step = None
-        self.next_state.complete_step()
+        self.next_state.complete_step("command")

core/agents/human_input.py

@@ -21,7 +21,7 @@ class HumanInput(BaseAgent):
             default="continue",
             buttons_only=True,
         )
-        self.next_state.complete_step()
+        self.next_state.complete_step("human_intervention")
         return AgentResponse.done(self)
 
     async def input_required(self, files: list[dict]) -> AgentResponse:

core/agents/legacy_handler.py

@@ -8,7 +8,7 @@ class LegacyHandler(BaseAgent):
 
     async def run(self) -> AgentResponse:
         if self.data["type"] == "review_task":
-            self.next_state.complete_step()
+            self.next_state.complete_step("review_task")
             return AgentResponse.done(self)
 
         raise ValueError(f"Unknown reason for calling Legacy Handler with data: {self.data}")

core/db/models/project_state.py

@@ -257,14 +257,14 @@ class ProjectState(Base):
 
         return new_state
 
-    def complete_step(self):
+    def complete_step(self, step_type: str):
         if not self.unfinished_steps:
             raise ValueError("There are no unfinished steps to complete")
         if "next_state" in self.__dict__:
             raise ValueError("Current state is read-only (already has a next state).")
 
         log.debug(f"Completing step {self.unfinished_steps[0]['type']}")
-        self.unfinished_steps[0]["completed"] = True
+        self.get_steps_of_type(step_type)[0]["completed"] = True
         flag_modified(self, "steps")
 
     def complete_task(self):

tests/db/test_project_state.py

@@ -134,7 +134,7 @@ async def test_completing_unfinished_steps(testdb):
 
     assert state.unfinished_steps == state.steps
     assert state.current_step["id"] == "abc"
-    state.complete_step()
+    state.complete_step("create_readme")
     assert state.unfinished_steps == []
     assert state.current_step is None
     await testdb.commit()