support workflow_run_id in totp code (#2199)

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