add timeout to browser session request (#2338)

skyvern/forge/sdk/db/client.py

@@ -1,6 +1,6 @@
 import json
 from datetime import datetime, timedelta
-from typing import Any, List, Optional, Sequence
+from typing import Any, List, Sequence
 
 import structlog
 from sqlalchemy import and_, delete, distinct, func, pool, select, tuple_, update
@@ -2656,7 +2656,7 @@ class AgentDB:
                 for workflow_run_block in workflow_run_blocks
             ]
 
-    async def get_active_persistent_browser_sessions(self, organization_id: str) -> List[PersistentBrowserSession]:
+    async def get_active_persistent_browser_sessions(self, organization_id: str) -> list[PersistentBrowserSession]:
         """Get all active persistent browser sessions for an organization."""
         try:
             async with self.Session() as session:
@@ -2676,7 +2676,7 @@ class AgentDB:
 
     async def get_persistent_browser_session_by_id(
         self, session_id: str, organization_id: str | None = None
-    ) -> Optional[PersistentBrowserSession]:
+    ) -> PersistentBrowserSession | None:
         """Get a specific persistent browser session."""
         try:
             async with self.Session() as session:
@@ -2703,7 +2703,7 @@ class AgentDB:
 
     async def get_persistent_browser_session(
         self, session_id: str, organization_id: str
-    ) -> Optional[PersistentBrowserSessionModel]:
+    ) -> PersistentBrowserSession | None:
         """Get a specific persistent browser session."""
         try:
             async with self.Session() as session:
@@ -2733,6 +2733,7 @@ class AgentDB:
         organization_id: str,
         runnable_type: str | None = None,
         runnable_id: str | None = None,
+        timeout_minutes: int | None = None,
     ) -> PersistentBrowserSessionModel:
         """Create a new persistent browser session."""
         try:
@@ -2741,6 +2742,7 @@ class AgentDB:
                     organization_id=organization_id,
                     runnable_type=runnable_type,
                     runnable_id=runnable_id,
+                    timeout_minutes=timeout_minutes,
                 )
                 session.add(browser_session)
                 await session.commit()

skyvern/forge/sdk/routes/browser_sessions.py

@@ -3,9 +3,10 @@ from fastapi.responses import ORJSONResponse
 
 from skyvern import analytics
 from skyvern.forge import app
-from skyvern.forge.sdk.routes.routers import base_router, legacy_base_router
+from skyvern.forge.sdk.routes.routers import base_router
 from skyvern.forge.sdk.schemas.organizations import Organization
 from skyvern.forge.sdk.services import org_auth_service
+from skyvern.schemas.browser_sessions import CreateBrowserSessionRequest
 from skyvern.webeye.schemas import BrowserSessionResponse
 
 
@@ -25,20 +26,6 @@ from skyvern.webeye.schemas import BrowserSessionResponse
         401: {"description": "Unauthorized - Invalid or missing authentication"},
     },
 )
-@legacy_base_router.get(
-    "/browser_sessions/{browser_session_id}",
-    response_model=BrowserSessionResponse,
-    tags=["session"],
-    openapi_extra={
-        "x-fern-sdk-group-name": "session",
-        "x-fern-sdk-method-name": "get_browser_session",
-    },
-)
-@legacy_base_router.get(
-    "/browser_sessions/{browser_session_id}/",
-    response_model=BrowserSessionResponse,
-    include_in_schema=False,
-)
 async def get_browser_session(
     browser_session_id: str,
     current_org: Organization = Depends(org_auth_service.get_current_org),
@@ -68,20 +55,6 @@ async def get_browser_session(
         401: {"description": "Unauthorized - Invalid or missing authentication"},
     },
 )
-@legacy_base_router.get(
-    "/browser_sessions",
-    response_model=list[BrowserSessionResponse],
-    tags=["session"],
-    openapi_extra={
-        "x-fern-sdk-group-name": "browser_sessions",
-        "x-fern-sdk-method-name": "get_browser_sessions",
-    },
-)
-@legacy_base_router.get(
-    "/browser_sessions/",
-    response_model=list[BrowserSessionResponse],
-    include_in_schema=False,
-)
 async def get_browser_sessions(
     current_org: Organization = Depends(org_auth_service.get_current_org),
 ) -> list[BrowserSessionResponse]:
@@ -106,24 +79,14 @@ async def get_browser_sessions(
         401: {"description": "Unauthorized - Invalid or missing authentication"},
     },
 )
-@legacy_base_router.post(
-    "/browser_sessions",
-    response_model=BrowserSessionResponse,
-    tags=["Browser Sessions"],
-    openapi_extra={
-        "x-fern-sdk-group-name": "session",
-        "x-fern-sdk-method-name": "create_browser_session",
-    },
-)
-@legacy_base_router.post(
-    "/browser_sessions/",
-    response_model=BrowserSessionResponse,
-    include_in_schema=False,
-)
 async def create_browser_session(
+    browser_session_request: CreateBrowserSessionRequest,
     current_org: Organization = Depends(org_auth_service.get_current_org),
 ) -> BrowserSessionResponse:
-    browser_session = await app.PERSISTENT_SESSIONS_MANAGER.create_session(current_org.organization_id)
+    browser_session = await app.PERSISTENT_SESSIONS_MANAGER.create_session(
+        organization_id=current_org.organization_id,
+        timeout_minutes=browser_session_request.timeout,
+    )
     return BrowserSessionResponse.from_browser_session(browser_session)
 
 
@@ -141,18 +104,6 @@ async def create_browser_session(
         401: {"description": "Unauthorized - Invalid or missing authentication"},
     },
 )
-@legacy_base_router.post(
-    "/browser_sessions/{browser_session_id}/close",
-    tags=["Browser Sessions"],
-    openapi_extra={
-        "x-fern-sdk-group-name": "browser_session",
-        "x-fern-sdk-method-name": "close_browser_session",
-    },
-)
-@legacy_base_router.post(
-    "/browser_sessions/{browser_session_id}/close/",
-    include_in_schema=False,
-)
 async def close_browser_session(
     browser_session_id: str,
     current_org: Organization = Depends(org_auth_service.get_current_org),

skyvern/forge/sdk/schemas/persistent_browser_sessions.py

@@ -12,6 +12,9 @@ class PersistentBrowserSession(BaseModel):
     runnable_id: str | None = None
     browser_address: str | None = None
     status: str | None = None
+    timeout_minutes: int | None = None
+    started_at: datetime | None = None
+    completed_at: datetime | None = None
     created_at: datetime
     modified_at: datetime
     deleted_at: datetime | None = None

skyvern/schemas/browser_sessions.py

@@ -0,0 +1,14 @@
+from pydantic import BaseModel, Field
+
+MIN_TIMEOUT = 5
+MAX_TIMEOUT = 10080
+DEFAULT_TIMEOUT = 60
+
+
+class CreateBrowserSessionRequest(BaseModel):
+    timeout: int = Field(
+        default=DEFAULT_TIMEOUT,
+        description=f"Timeout in minutes for the session. Timeout is applied after the session is started. Must be between {MIN_TIMEOUT} and {MAX_TIMEOUT}. Defaults to {DEFAULT_TIMEOUT}.",
+        ge=MIN_TIMEOUT,
+        le=MAX_TIMEOUT,
+    )

skyvern/webeye/persistent_sessions_manager.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, Optional, Tuple
 
 import structlog
 from playwright._impl._errors import TargetClosedError
@@ -31,7 +31,7 @@ class PersistentSessionsManager:
         cls.instance.database = database
         return cls.instance
 
-    async def get_active_sessions(self, organization_id: str) -> List[PersistentBrowserSession]:
+    async def get_active_sessions(self, organization_id: str) -> list[PersistentBrowserSession]:
         """Get all active sessions for an organization."""
         return await self.database.get_active_persistent_browser_sessions(organization_id)
 
@@ -40,7 +40,7 @@ class PersistentSessionsManager:
         browser_session = self._browser_sessions.get(session_id)
         return browser_session.browser_state if browser_session else None
 
-    async def get_session(self, session_id: str, organization_id: str) -> Optional[PersistentBrowserSession]:
+    async def get_session(self, session_id: str, organization_id: str) -> PersistentBrowserSession | None:
         """Get a specific browser session by session ID."""
         return await self.database.get_persistent_browser_session(session_id, organization_id)
 
@@ -49,6 +49,7 @@ class PersistentSessionsManager:
         organization_id: str,
         runnable_id: str | None = None,
         runnable_type: str | None = None,
+        timeout_minutes: int | None = None,
     ) -> PersistentBrowserSession:
         """Create a new browser session for an organization and return its ID with the browser state."""
 
@@ -61,6 +62,7 @@ class PersistentSessionsManager:
             organization_id=organization_id,
             runnable_type=runnable_type,
             runnable_id=runnable_id,
+            timeout_minutes=timeout_minutes,
         )
 
         return browser_session_db

skyvern/webeye/schemas.py

@@ -16,7 +16,15 @@ class BrowserSessionResponse(BaseModel):
         None, description="Type of runnable associated with this session (workflow, task etc)"
     )
     runnable_id: str | None = Field(None, description="ID of the associated runnable")
-    created_at: datetime = Field(description="Timestamp when the session was created")
+    timeout: int | None = Field(
+        None,
+        description="Timeout in minutes for the session. Timeout is applied after the session is started.",
+    )
+    started_at: datetime | None = Field(None, description="Timestamp when the session was started")
+    completed_at: datetime | None = Field(None, description="Timestamp when the session was completed")
+    created_at: datetime = Field(
+        description="Timestamp when the session was created (the timestamp for the initial request)"
+    )
     modified_at: datetime = Field(description="Timestamp when the session was last modified")
     deleted_at: datetime | None = Field(None, description="Timestamp when the session was deleted, if applicable")
 
@@ -36,6 +44,9 @@ class BrowserSessionResponse(BaseModel):
             organization_id=browser_session.organization_id,
             runnable_type=browser_session.runnable_type,
             runnable_id=browser_session.runnable_id,
+            timeout=browser_session.timeout_minutes,
+            started_at=browser_session.started_at,
+            completed_at=browser_session.completed_at,
             created_at=browser_session.created_at,
             modified_at=browser_session.modified_at,
             deleted_at=browser_session.deleted_at,