Readd retry_message

moatless/loop.py

@@ -267,7 +267,7 @@ def _set_state_loop(self, state: AgenticState):
     def retries(self) -> int:
         retries = 0
         for action in reversed(self._current_transition.actions):
-            if action.retry_message:
+            if action.trigger == "retry":
                 retries += 1
             else:
                 return retries
@@ -281,7 +281,7 @@ def retry_messages(self, state: AgenticState) -> list[Message]:
             return messages
 
         for action in self._current_transition.actions:
-            if action.retry_message:
+            if action.trigger == "retry":
                 if isinstance(action.action, Content):
                     messages.append(
                         AssistantMessage(
@@ -504,6 +504,7 @@ def _run(self):
             TrajectoryAction(
                 action=action,
                 trigger=response.trigger,
+                retry_message=response.retry_message,
                 completion_cost=cost,
                 input_tokens=input_tokens,
                 output_tokens=output_tokens,

moatless/state.py

@@ -100,7 +100,7 @@ def required_fields(cls) -> set[str]:
     def retries(self) -> int:
         retries = 0
         for action in reversed(self.loop._current_transition.actions):
-            if action.retry_message:
+            if action.trigger == "retry":
                 retries += 1
             else:
                 return retries

moatless/trajectory.py

@@ -17,6 +17,7 @@
 class TrajectoryAction(BaseModel):
     action: ActionRequest
     trigger: Optional[str]
+    retry_message: Optional[str]
     completion_cost: Optional[float] = None
     input_tokens: Optional[int] = None
     output_tokens: Optional[int] = None