vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Define browser manager API (#1497)

Browse Source 
Co-authored-by: Shuchang Zheng <shu@skyvern.com>
This commit is contained in:
Maksim Ivanov
2025-01-08 18:14:38 +01:00
committed by GitHub
parent 008b57d6f4
commit 7bfb1e9b21
8 changed files with 764 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
skyvern/forge/app.py
View File

@@ -16,6 +16,7 @@ from skyvern.forge.sdk.settings_manager import SettingsManager
from skyvern.forge.sdk.workflow.context_manager import WorkflowContextManager
from skyvern.forge.sdk.workflow.service import WorkflowService
from skyvern.webeye.browser_manager import BrowserManager
from skyvern.webeye.persistent_sessions_manager import PersistentSessionsManager
from skyvern.webeye.scraper.scraper import ScrapeExcludeFunc
SETTINGS_MANAGER = SettingsManager.get_settings()
@@ -37,6 +38,7 @@ SECONDARY_LLM_API_HANDLER = LLMAPIHandlerFactory.get_llm_api_handler(
WORKFLOW_CONTEXT_MANAGER = WorkflowContextManager()
WORKFLOW_SERVICE = WorkflowService()
AGENT_FUNCTION = AgentFunction()
PERSISTENT_SESSIONS_MANAGER = PersistentSessionsManager(database=DATABASE)
scrape_exclude: ScrapeExcludeFunc | None = None
authentication_function: Callable[[str], Awaitable[Organization]] | None = None
setup_api_app: Callable[[FastAPI], None] | None = None

174
skyvern/forge/sdk/db/client.py
View File

@@ -1,6 +1,6 @@
import json
from datetime import datetime, timedelta
from typing import Any, Sequence
from typing import Any, List, Optional, Sequence
import structlog
from sqlalchemy import and_, delete, func, select, update
@@ -24,6 +24,7 @@ from skyvern.forge.sdk.db.models import (
OrganizationAuthTokenModel,
OrganizationModel,
OutputParameterModel,
PersistentBrowserSessionModel,
StepModel,
TaskGenerationModel,
TaskModel,
@@ -62,6 +63,7 @@ from skyvern.forge.sdk.schemas.observers import (
ObserverThoughtType,
)
from skyvern.forge.sdk.schemas.organizations import Organization, OrganizationAuthToken
from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession
from skyvern.forge.sdk.schemas.task_generations import TaskGeneration
from skyvern.forge.sdk.schemas.tasks import OrderBy, ProxyLocation, SortDirection, Task, TaskStatus
from skyvern.forge.sdk.schemas.totp_codes import TOTPCode
@@ -2251,3 +2253,173 @@ class AgentDB:
convert_to_workflow_run_block(workflow_run_block, task=tasks_dict.get(workflow_run_block.task_id))
for workflow_run_block in workflow_run_blocks
]
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:
result = await session.execute(
select(PersistentBrowserSessionModel)
.filter_by(organization_id=organization_id)
.filter_by(deleted_at=None)
)
sessions = result.scalars().all()
return [PersistentBrowserSession.model_validate(session) for session in sessions]
except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True)
raise
except Exception:
LOG.error("UnexpectedError", exc_info=True)
raise
async def get_persistent_browser_session(
self, session_id: str, organization_id: str
) -> Optional[PersistentBrowserSessionModel]:
"""Get a specific persistent browser session."""
try:
async with self.Session() as session:
persistent_browser_session = (
await session.scalars(
select(PersistentBrowserSessionModel)
.filter_by(persistent_browser_session_id=session_id)
.filter_by(organization_id=organization_id)
.filter_by(deleted_at=None)
)
).first()
if persistent_browser_session:
return PersistentBrowserSession.model_validate(persistent_browser_session)
raise NotFoundError(f"PersistentBrowserSession {session_id} not found")
except NotFoundError:
LOG.error("NotFoundError", exc_info=True)
raise
except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True)
raise
except Exception:
LOG.error("UnexpectedError", exc_info=True)
raise
async def create_persistent_browser_session(
self,
organization_id: str,
runnable_type: str | None = None,
runnable_id: str | None = None,
) -> PersistentBrowserSessionModel:
"""Create a new persistent browser session."""
try:
async with self.Session() as session:
browser_session = PersistentBrowserSessionModel(
organization_id=organization_id,
runnable_type=runnable_type,
runnable_id=runnable_id,
)
session.add(browser_session)
await session.commit()
await session.refresh(browser_session)
return PersistentBrowserSession.model_validate(browser_session)
except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True)
raise
except Exception:
LOG.error("UnexpectedError", exc_info=True)
raise
async def mark_persistent_browser_session_deleted(self, session_id: str, organization_id: str) -> None:
"""Mark a persistent browser session as deleted."""
try:
async with self.Session() as session:
persistent_browser_session = (
await session.scalars(
select(PersistentBrowserSessionModel)
.filter_by(persistent_browser_session_id=session_id)
.filter_by(organization_id=organization_id)
)
).first()
if persistent_browser_session:
persistent_browser_session.deleted_at = datetime.utcnow()
await session.commit()
await session.refresh(persistent_browser_session)
else:
raise NotFoundError(f"PersistentBrowserSession {session_id} not found")
except NotFoundError:
LOG.error("NotFoundError", exc_info=True)
raise
except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True)
raise
except Exception:
LOG.error("UnexpectedError", exc_info=True)
raise
async def occupy_persistent_browser_session(
self, session_id: str, runnable_type: str, runnable_id: str, organization_id: str
) -> None:
"""Occupy a specific persistent browser session."""
try:
async with self.Session() as session:
persistent_browser_session = (
await session.scalars(
select(PersistentBrowserSessionModel)
.filter_by(persistent_browser_session_id=session_id)
.filter_by(organization_id=organization_id)
.filter_by(deleted_at=None)
)
).first()
if persistent_browser_session:
persistent_browser_session.runnable_type = runnable_type
persistent_browser_session.runnable_id = runnable_id
await session.commit()
await session.refresh(persistent_browser_session)
else:
raise NotFoundError(f"PersistentBrowserSession {session_id} not found")
except NotFoundError:
LOG.error("NotFoundError", exc_info=True)
raise
except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True)
raise
except Exception:
LOG.error("UnexpectedError", exc_info=True)
raise
async def release_persistent_browser_session(self, session_id: str, organization_id: str) -> None:
"""Release a specific persistent browser session."""
try:
async with self.Session() as session:
persistent_browser_session = (
await session.scalars(
select(PersistentBrowserSessionModel)
.filter_by(persistent_browser_session_id=session_id)
.filter_by(organization_id=organization_id)
.filter_by(deleted_at=None)
)
).first()
if persistent_browser_session:
persistent_browser_session.runnable_type = None
persistent_browser_session.runnable_id = None
await session.commit()
await session.refresh(persistent_browser_session)
else:
raise NotFoundError(f"PersistentBrowserSession {session_id} not found")
except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True)
raise
except NotFoundError:
LOG.error("NotFoundError", exc_info=True)
raise
except Exception:
LOG.error("UnexpectedError", exc_info=True)
raise
async def get_all_active_persistent_browser_sessions(self) -> List[PersistentBrowserSessionModel]:
"""Get all active persistent browser sessions across all organizations."""
try:
async with self.Session() as session:
result = await session.execute(select(PersistentBrowserSessionModel).filter_by(deleted_at=None))
return result.scalars().all()
except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True)
raise
except Exception:
LOG.error("UnexpectedError", exc_info=True)
raise

95
skyvern/forge/sdk/routes/agent_protocol.py
View File

@@ -69,6 +69,7 @@ from skyvern.forge.sdk.workflow.models.workflow import (
)
from skyvern.forge.sdk.workflow.models.yaml import WorkflowCreateYAMLRequest
from skyvern.webeye.actions.actions import Action
from skyvern.webeye.schemas import BrowserSessionResponse
base_router = APIRouter()
@@ -1123,3 +1124,97 @@ async def get_observer_cruise(
if not observer_cruise:
raise HTTPException(status_code=404, detail=f"Observer cruise {observer_cruise_id} not found")
return observer_cruise
@base_router.get(
"/browser_sessions/{browser_session_id}",
response_model=BrowserSessionResponse,
)
@base_router.get(
"/browser_sessions/{browser_session_id}/",
response_model=BrowserSessionResponse,
include_in_schema=False,
)
async def get_browser_session_by_id(
browser_session_id: str,
current_org: Organization = Depends(org_auth_service.get_current_org),
) -> BrowserSessionResponse:
analytics.capture("skyvern-oss-agent-workflow-run-get")
browser_session = await app.PERSISTENT_SESSIONS_MANAGER.get_session(
browser_session_id,
current_org.organization_id,
)
if not browser_session:
raise HTTPException(status_code=404, detail=f"Browser session {browser_session_id} not found")
return BrowserSessionResponse.from_browser_session(browser_session)
@base_router.get(
"/browser_sessions",
response_model=list[BrowserSessionResponse],
)
@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]:
"""Get all active browser sessions for the organization"""
analytics.capture("skyvern-oss-agent-browser-sessions-get")
browser_sessions = await app.PERSISTENT_SESSIONS_MANAGER.get_active_sessions(current_org.organization_id)
return [BrowserSessionResponse.from_browser_session(browser_session) for browser_session in browser_sessions]
@base_router.post(
"/browser_sessions",
response_model=BrowserSessionResponse,
)
@base_router.post(
"/browser_sessions/",
response_model=BrowserSessionResponse,
include_in_schema=False,
)
async def create_browser_session(
current_org: Organization = Depends(org_auth_service.get_current_org),
) -> BrowserSessionResponse:
browser_session, _ = await app.PERSISTENT_SESSIONS_MANAGER.create_session(current_org.organization_id)
return BrowserSessionResponse.from_browser_session(browser_session)
@base_router.post(
"/browser_sessions/close",
)
@base_router.post(
"/browser_sessions/close/",
include_in_schema=False,
)
async def close_browser_sessions(
current_org: Organization = Depends(org_auth_service.get_current_org),
) -> ORJSONResponse:
await app.PERSISTENT_SESSIONS_MANAGER.close_all_sessions(current_org.organization_id)
return ORJSONResponse(
content={"message": "All browser sessions closed"},
status_code=200,
media_type="application/json",
)
@base_router.post(
"/browser_sessions/{session_id}/close",
)
@base_router.post(
"/browser_sessions/{session_id}/close/",
include_in_schema=False,
)
async def close_browser_session(
session_id: str,
current_org: Organization = Depends(org_auth_service.get_current_org),
) -> ORJSONResponse:
await app.PERSISTENT_SESSIONS_MANAGER.close_session(current_org.organization_id, session_id)
return ORJSONResponse(
content={"message": "Browser session closed"},
status_code=200,
media_type="application/json",
)

15
skyvern/forge/sdk/schemas/persistent_browser_sessions.py Normal file
View File

@@ -0,0 +1,15 @@
from datetime import datetime
from pydantic import BaseModel, ConfigDict
class PersistentBrowserSession(BaseModel):
model_config = ConfigDict(from_attributes=True)
persistent_browser_session_id: str
organization_id: str
runnable_type: str | None = None
runnable_id: str | None = None
created_at: datetime
modified_at: datetime
deleted_at: datetime | None = None