From 79676dbb5252ed818ba49e76f9a1f364e02d1e2c Mon Sep 17 00:00:00 2001
From: LeonOstrez <leon.ostrez@gmail.com>
Date: Tue, 8 Jul 2025 22:16:54 +0200
Subject: [PATCH] optimize get_fe_states

---
 core/cli/main.py                |  2 +-
 core/db/models/project_state.py | 26 +++++++++++++++++++-------
 core/state/state_manager.py     |  4 ++--
 3 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/core/cli/main.py b/core/cli/main.py
index 79629971..006eff35 100644
--- a/core/cli/main.py
+++ b/core/cli/main.py
@@ -209,7 +209,7 @@ async def run_pythagora_session(sm: StateManager, ui: UIBase, args: Namespace):
             return False
 
         # SPECIFICATION
-        fe_states = await sm.get_fe_states()
+        fe_states = await sm.get_fe_states(limit=10)
         be_back_logs, last_task_in_db = await sm.get_be_back_logs()
 
         if sm.current_state.specification and sm.current_state.specification.original_description:
diff --git a/core/db/models/project_state.py b/core/db/models/project_state.py
index fb6eeac2..90ae3b7c 100644
--- a/core/db/models/project_state.py
+++ b/core/db/models/project_state.py
@@ -903,7 +903,9 @@ class ProjectState(Base):
             return results
 
     @staticmethod
-    async def get_fe_states(session: "AsyncSession", branch_id: UUID) -> Optional["ProjectState"]:
+    async def get_fe_states(
+        session: "AsyncSession", branch_id: UUID, limit: Optional[int] = None
+    ) -> Optional["ProjectState"]:
         query = select(ProjectState).where(
             and_(
                 ProjectState.branch_id == branch_id,
@@ -931,16 +933,26 @@ class ProjectState(Base):
         result = await session.execute(query)
         fe_end = result.scalars().one_or_none()
 
-        query = select(ProjectState).where(
-            and_(
-                ProjectState.branch_id == branch_id,
-                ProjectState.step_index >= fe_start.step_index,
-                ProjectState.step_index <= fe_end.step_index,
+        query = (
+            select(ProjectState)
+            .where(
+                and_(
+                    ProjectState.branch_id == branch_id,
+                    ProjectState.step_index >= fe_start.step_index,
+                    ProjectState.step_index <= fe_end.step_index,
+                )
             )
+            .order_by(ProjectState.step_index.desc())
         )
 
+        if limit:
+            query = query.limit(limit)
+
         results = await session.execute(query)
-        return results.scalars().all()
+        states = results.scalars().all()
+
+        # Since we ordered by step_index desc and limited, we need to reverse to get chronological order
+        return list(reversed(states))
 
     @staticmethod
     def get_epic_task_number(state, current_task) -> (int, int):
diff --git a/core/state/state_manager.py b/core/state/state_manager.py
index dab3d952..8e98301e 100644
--- a/core/state/state_manager.py
+++ b/core/state/state_manager.py
@@ -159,8 +159,8 @@ class StateManager:
             self.current_session, self.current_state.branch_id, start_state_id, end_state_id, limit
         )
 
-    async def get_fe_states(self) -> Optional[ProjectState]:
-        return await ProjectState.get_fe_states(self.current_session, self.current_state.branch_id)
+    async def get_fe_states(self, limit: Optional[int] = None) -> Optional[ProjectState]:
+        return await ProjectState.get_fe_states(self.current_session, self.current_state.branch_id, limit)
 
     async def get_be_back_logs(self):
         """