From fbc69e6a0c0b23ef462b29b13a900e151feb8227 Mon Sep 17 00:00:00 2001
From: Andy Li <55300002+cliandy@users.noreply.github.com>
Date: Thu, 24 Jul 2025 16:51:31 -0700
Subject: [PATCH] fix: error logging for stop reasons

---
 letta/agents/letta_agent.py | 6 ++++++
 tests/test_utils.py         | 2 ++
 2 files changed, 8 insertions(+)

diff --git a/letta/agents/letta_agent.py b/letta/agents/letta_agent.py
index 3c7960bc..3f7faab7 100644
--- a/letta/agents/letta_agent.py
+++ b/letta/agents/letta_agent.py
@@ -359,6 +359,8 @@ class LettaAgent(BaseAgent):
             # Update step if it needs to be updated
             finally:
                 if settings.track_stop_reason:
+                    if step_progression == StepProgression.FINISHED and should_continue:
+                        continue
                     self.logger.info("Running final update. Step Progression: %s", step_progression)
                     try:
                         if step_progression < StepProgression.STEP_LOGGED:
@@ -586,6 +588,8 @@ class LettaAgent(BaseAgent):
                 # Update step if it needs to be updated
             finally:
                 if settings.track_stop_reason:
+                    if step_progression == StepProgression.FINISHED and should_continue:
+                        continue
                     self.logger.info("Running final update. Step Progression: %s", step_progression)
                     try:
                         if step_progression < StepProgression.STEP_LOGGED:
@@ -897,6 +901,8 @@ class LettaAgent(BaseAgent):
             # Update step if it needs to be updated
             finally:
                 if settings.track_stop_reason:
+                    if step_progression == StepProgression.FINISHED and should_continue:
+                        continue
                     self.logger.info("Running final update. Step Progression: %s", step_progression)
                     try:
                         if step_progression < StepProgression.STEP_LOGGED:
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 84622882..87cfc127 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1,3 +1,5 @@
+import asyncio
+
 import pytest
 
 from letta.constants import MAX_FILENAME_LENGTH