From 46fdd2b91362532ef3cb818e42c25c1b856b9753 Mon Sep 17 00:00:00 2001
From: Marc Kelechava <marc@skyvern.com>
Date: Mon, 24 Nov 2025 10:04:14 -0800
Subject: [PATCH] Fix TOTP Code Sort (#4082)

---
 skyvern/forge/sdk/db/client.py          | 10 ++++------
 skyvern/forge/sdk/routes/credentials.py | 24 ++++++++----------------
 2 files changed, 12 insertions(+), 22 deletions(-)

diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py
index 271e32ed..75751ece 100644
--- a/skyvern/forge/sdk/db/client.py
+++ b/skyvern/forge/sdk/db/client.py
@@ -3343,16 +3343,12 @@ class AgentDB:
         valid_lifespan_minutes: int | None = None,
         otp_type: OTPType | None = None,
         workflow_run_id: str | None = None,
+        totp_identifier: str | None = None,
     ) -> list[TOTPCode]:
         """
         Return recent otp codes for an organization ordered by newest first with optional
         workflow_run_id filtering.
         """
-        all_null = and_(
-            TOTPCodeModel.task_id.is_(None),
-            TOTPCodeModel.workflow_id.is_(None),
-            TOTPCodeModel.workflow_run_id.is_(None),
-        )
         async with self.Session() as session:
             query = select(TOTPCodeModel).filter_by(organization_id=organization_id)
 
@@ -3365,7 +3361,9 @@ class AgentDB:
                 query = query.filter(TOTPCodeModel.otp_type == otp_type)
             if workflow_run_id is not None:
                 query = query.filter(TOTPCodeModel.workflow_run_id == workflow_run_id)
-            query = query.order_by(asc(all_null), TOTPCodeModel.created_at.desc()).limit(limit)
+            if totp_identifier:
+                query = query.filter(TOTPCodeModel.totp_identifier == totp_identifier)
+            query = query.order_by(TOTPCodeModel.created_at.desc()).limit(limit)
             totp_codes = (await session.scalars(query)).all()
             return [TOTPCode.model_validate(totp_code) for totp_code in totp_codes]
 
diff --git a/skyvern/forge/sdk/routes/credentials.py b/skyvern/forge/sdk/routes/credentials.py
index 2c097d28..d328c3de 100644
--- a/skyvern/forge/sdk/routes/credentials.py
+++ b/skyvern/forge/sdk/routes/credentials.py
@@ -178,22 +178,14 @@ async def get_totp_codes(
         description="Maximum number of codes to return.",
     ),
 ) -> list[TOTPCode]:
-    if totp_identifier:
-        codes = await app.DATABASE.get_otp_codes(
-            organization_id=curr_org.organization_id,
-            totp_identifier=totp_identifier,
-            otp_type=otp_type,
-            workflow_run_id=workflow_run_id,
-            limit=limit,
-        )
-    else:
-        codes = await app.DATABASE.get_recent_otp_codes(
-            organization_id=curr_org.organization_id,
-            limit=limit,
-            valid_lifespan_minutes=None,
-            otp_type=otp_type,
-            workflow_run_id=workflow_run_id,
-        )
+    codes = await app.DATABASE.get_recent_otp_codes(
+        organization_id=curr_org.organization_id,
+        limit=limit,
+        valid_lifespan_minutes=None,
+        otp_type=otp_type,
+        workflow_run_id=workflow_run_id,
+        totp_identifier=totp_identifier,
+    )
 
     return codes