From 0a2e82f9527d463f6b24885f04032269641e42c5 Mon Sep 17 00:00:00 2001
From: LawyZheng <lawyzheng1106@gmail.com>
Date: Thu, 17 Jul 2025 14:50:05 +0800
Subject: [PATCH] dns for browser session (#2972)

---
 ...76a3aa_add_more_info_to_browser_session.py | 33 +++++++++++++++++++
 skyvern/forge/sdk/db/client.py                |  4 +++
 skyvern/forge/sdk/db/models.py                |  2 ++
 3 files changed, 39 insertions(+)
 create mode 100644 alembic/versions/2025_07_17_0612-2e997076a3aa_add_more_info_to_browser_session.py

diff --git a/alembic/versions/2025_07_17_0612-2e997076a3aa_add_more_info_to_browser_session.py b/alembic/versions/2025_07_17_0612-2e997076a3aa_add_more_info_to_browser_session.py
new file mode 100644
index 00000000..545d9a0a
--- /dev/null
+++ b/alembic/versions/2025_07_17_0612-2e997076a3aa_add_more_info_to_browser_session.py
@@ -0,0 +1,33 @@
+"""add_more_info_to_browser_session
+
+Revision ID: 2e997076a3aa
+Revises: f63892bfc429
+Create Date: 2025-07-17 06:12:04.974095+00:00
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "2e997076a3aa"
+down_revision: Union[str, None] = "f63892bfc429"
+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("persistent_browser_sessions", sa.Column("ip_address", sa.String(), nullable=True))
+    op.add_column("persistent_browser_sessions", sa.Column("ecs_task_arn", sa.String(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("persistent_browser_sessions", "ecs_task_arn")
+    op.drop_column("persistent_browser_sessions", "ip_address")
+    # ### end Alembic commands ###
diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py
index a849e357..084eb984 100644
--- a/skyvern/forge/sdk/db/client.py
+++ b/skyvern/forge/sdk/db/client.py
@@ -3011,6 +3011,8 @@ class AgentDB:
         self,
         browser_session_id: str,
         browser_address: str,
+        ip_address: str,
+        ecs_task_arn: str | None,
         organization_id: str | None = None,
     ) -> None:
         """Set the browser address for a persistent browser session."""
@@ -3026,6 +3028,8 @@ class AgentDB:
                 ).first()
                 if persistent_browser_session:
                     persistent_browser_session.browser_address = browser_address
+                    persistent_browser_session.ip_address = ip_address
+                    persistent_browser_session.ecs_task_arn = ecs_task_arn
                     # once the address is set, the session is started
                     persistent_browser_session.started_at = datetime.utcnow()
                     await session.commit()
diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py
index b47f4207..fa51ba45 100644
--- a/skyvern/forge/sdk/db/models.py
+++ b/skyvern/forge/sdk/db/models.py
@@ -696,6 +696,8 @@ class PersistentBrowserSessionModel(Base):
     browser_address = Column(String, nullable=True)
     status = Column(String, nullable=True, default="created")
     timeout_minutes = Column(Integer, nullable=True)
+    ip_address = Column(String, nullable=True)
+    ecs_task_arn = Column(String, nullable=True)
     started_at = Column(DateTime, nullable=True)
     completed_at = Column(DateTime, nullable=True)
     created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False, index=True)