support workflow_run_id in totp code (#2199)

alembic/versions/2025_04_21_0019-511d9da18f5d_add_workflow_run_id_to_totp_code_model.py

@@ -0,0 +1,33 @@
+"""add workflow_run_id to totp code model
+
+Revision ID: 511d9da18f5d
+Revises: 3aa168d1ffa5
+Create Date: 2025-04-21 00:19:48.989020+00:00
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "511d9da18f5d"
+down_revision: Union[str, None] = "3aa168d1ffa5"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("totp_codes", sa.Column("workflow_run_id", sa.String(), nullable=True))
+    op.create_foreign_key(None, "totp_codes", "workflow_runs", ["workflow_run_id"], ["workflow_run_id"])
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(None, "totp_codes", type_="foreignkey")
+    op.drop_column("totp_codes", "workflow_run_id")
+    # ### end Alembic commands ###

skyvern/forge/sdk/db/client.py

@@ -2089,6 +2089,7 @@ class AgentDB:
         code: str,
         task_id: str | None = None,
         workflow_id: str | None = None,
+        workflow_run_id: str | None = None,
         source: str | None = None,
         expired_at: datetime | None = None,
     ) -> TOTPCode:
@@ -2100,6 +2101,7 @@ class AgentDB:
                 code=code,
                 task_id=task_id,
                 workflow_id=workflow_id,
+                workflow_run_id=workflow_run_id,
                 source=source,
                 expired_at=expired_at,
             )

skyvern/forge/sdk/db/models.py

@@ -491,6 +491,7 @@ class TOTPCodeModel(Base):
     organization_id = Column(String, ForeignKey("organizations.organization_id"))
     task_id = Column(String, ForeignKey("tasks.task_id"))
     workflow_id = Column(String, ForeignKey("workflows.workflow_id"))
+    workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"))
     content = Column(String, nullable=False)
     code = Column(String, nullable=False)
     source = Column(String)

skyvern/forge/sdk/routes/totp.py

@@ -41,6 +41,7 @@ async def send_totp_code(
         code=code,
         task_id=data.task_id,
         workflow_id=data.workflow_id,
+        workflow_run_id=data.workflow_run_id,
         source=data.source,
         expired_at=data.expired_at,
     )

skyvern/forge/sdk/schemas/totp_codes.py

@@ -11,6 +11,7 @@ class TOTPCodeBase(BaseModel):
     totp_identifier: str | None = None
     task_id: str | None = None
     workflow_id: str | None = None
+    workflow_run_id: str | None = None
     source: str | None = None
     content: str | None = None
 

skyvern/webeye/actions/handler.py

@@ -3128,7 +3128,11 @@ async def poll_verification_code(
             )
         elif totp_identifier:
             verification_code = await _get_verification_code_from_db(
-                task_id, organization_id, totp_identifier, workflow_id=workflow_id
+                task_id,
+                organization_id,
+                totp_identifier,
+                workflow_id=workflow_id,
+                workflow_run_id=workflow_run_id,
             )
         if verification_code:
             LOG.info("Got verification code", verification_code=verification_code)
@@ -3169,9 +3173,12 @@ async def _get_verification_code_from_db(
     organization_id: str,
     totp_identifier: str,
     workflow_id: str | None = None,
+    workflow_run_id: str | None = None,
 ) -> str | None:
     totp_codes = await app.DATABASE.get_totp_codes(organization_id=organization_id, totp_identifier=totp_identifier)
     for totp_code in totp_codes:
+        if totp_code.workflow_run_id and workflow_run_id and totp_code.workflow_run_id != workflow_run_id:
+            continue
         if totp_code.workflow_id and workflow_id and totp_code.workflow_id != workflow_id:
             continue
         if totp_code.task_id and totp_code.task_id != task_id: