From 9038e6e047e4578a08fc979548b687f07696f76b Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Mon, 21 Apr 2025 08:24:29 +0800 Subject: [PATCH] support workflow_run_id in totp code (#2199) --- ..._add_workflow_run_id_to_totp_code_model.py | 33 +++++++++++++++++++ skyvern/forge/sdk/db/client.py | 2 ++ skyvern/forge/sdk/db/models.py | 1 + skyvern/forge/sdk/routes/totp.py | 1 + skyvern/forge/sdk/schemas/totp_codes.py | 1 + skyvern/webeye/actions/handler.py | 9 ++++- 6 files changed, 46 insertions(+), 1 deletion(-) create mode 100644 alembic/versions/2025_04_21_0019-511d9da18f5d_add_workflow_run_id_to_totp_code_model.py diff --git a/alembic/versions/2025_04_21_0019-511d9da18f5d_add_workflow_run_id_to_totp_code_model.py b/alembic/versions/2025_04_21_0019-511d9da18f5d_add_workflow_run_id_to_totp_code_model.py new file mode 100644 index 00000000..818a9f11 --- /dev/null +++ b/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 ### diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 91c3123a..2071d7a8 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/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, ) diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index 695d1d85..595146e0 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/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) diff --git a/skyvern/forge/sdk/routes/totp.py b/skyvern/forge/sdk/routes/totp.py index 58d722ec..c7c3fe83 100644 --- a/skyvern/forge/sdk/routes/totp.py +++ b/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, ) diff --git a/skyvern/forge/sdk/schemas/totp_codes.py b/skyvern/forge/sdk/schemas/totp_codes.py index 8acdddbd..7b058d65 100644 --- a/skyvern/forge/sdk/schemas/totp_codes.py +++ b/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 diff --git a/skyvern/webeye/actions/handler.py b/skyvern/webeye/actions/handler.py index 36703775..c4866cf5 100644 --- a/skyvern/webeye/actions/handler.py +++ b/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: