diff --git a/run_streaming.py b/run_streaming.py index 6048295a5..1d645de3a 100644 --- a/run_streaming.py +++ b/run_streaming.py @@ -6,7 +6,7 @@ import structlog import typer from skyvern.forge import app -from skyvern.forge.forge_app_initializer import start_forge_app +from skyvern.forge.forge_app_initializer import start_streaming_worker_app from skyvern.forge.sdk.api.files import get_skyvern_temp_dir from skyvern.forge.sdk.workflow.models.workflow import WorkflowRunStatus from skyvern.utils.files import get_json_from_file, get_skyvern_state_file_path, initialize_skyvern_state_file @@ -16,62 +16,73 @@ LOG = structlog.get_logger() async def run() -> None: - start_forge_app() + start_streaming_worker_app() await initialize_skyvern_state_file(task_id=None, workflow_run_id=None, organization_id=None) - while True: - await asyncio.sleep(INTERVAL) - try: - current_json = get_json_from_file(get_skyvern_state_file_path()) - except Exception: - continue - - task_id = current_json.get("task_id") - workflow_run_id = current_json.get("workflow_run_id") - organization_id = current_json.get("organization_id") - if not organization_id or (not task_id and not workflow_run_id): - continue - - try: - if workflow_run_id: - workflow_run = await app.DATABASE.workflow_runs.get_workflow_run(workflow_run_id=workflow_run_id) - if not workflow_run or workflow_run.status in [ - WorkflowRunStatus.completed, - WorkflowRunStatus.failed, - WorkflowRunStatus.terminated, - ]: - continue - file_name = f"{workflow_run_id}.png" - - elif task_id: - task = await app.DATABASE.tasks.get_task(task_id=task_id, organization_id=organization_id) - if not task or task.status.is_final(): - continue - file_name = f"{task_id}.png" - else: + try: + while True: + await asyncio.sleep(INTERVAL) + try: + current_json = get_json_from_file(get_skyvern_state_file_path()) + except Exception: continue - except Exception: - LOG.exception( - "Failed to get task or workflow run while taking streaming screenshot in worker", - task_id=task_id, - workflow_run_id=workflow_run_id, - organization_id=organization_id, + + task_id = current_json.get("task_id") + workflow_run_id = current_json.get("workflow_run_id") + organization_id = current_json.get("organization_id") + if not organization_id or (not task_id and not workflow_run_id): + continue + + try: + if workflow_run_id: + workflow_run = await app.DATABASE.workflow_runs.get_workflow_run(workflow_run_id=workflow_run_id) + if not workflow_run or workflow_run.status in [ + WorkflowRunStatus.completed, + WorkflowRunStatus.failed, + WorkflowRunStatus.terminated, + ]: + continue + file_name = f"{workflow_run_id}.png" + + elif task_id: + task = await app.DATABASE.tasks.get_task(task_id=task_id, organization_id=organization_id) + if not task or task.status.is_final(): + continue + file_name = f"{task_id}.png" + else: + continue + except Exception: + LOG.exception( + "Failed to get task or workflow run while taking streaming screenshot in worker", + task_id=task_id, + workflow_run_id=workflow_run_id, + organization_id=organization_id, + ) + continue + + # create f"{get_skyvern_temp_dir()}/{organization_id}" directory if it does not exists + os.makedirs(f"{get_skyvern_temp_dir()}/{organization_id}", exist_ok=True) + png_file_path = f"{get_skyvern_temp_dir()}/{organization_id}/{file_name}" + + # run subprocess to take screenshot + subprocess.run( + f"xwd -root | xwdtopnm 2>/dev/null | pnmtopng > {png_file_path}", shell=True, env={"DISPLAY": ":99"} ) - continue - - # create f"{get_skyvern_temp_dir()}/{organization_id}" directory if it does not exists - os.makedirs(f"{get_skyvern_temp_dir()}/{organization_id}", exist_ok=True) - png_file_path = f"{get_skyvern_temp_dir()}/{organization_id}/{file_name}" - - # run subprocess to take screenshot - subprocess.run( - f"xwd -root | xwdtopnm 2>/dev/null | pnmtopng > {png_file_path}", shell=True, env={"DISPLAY": ":99"} - ) + try: + await app.STORAGE.save_streaming_file(organization_id, file_name) + except Exception: + LOG.debug("Failed to upload screenshot", organization_id=organization_id, file_name=file_name) + finally: + # Dispose the DB engine/pool this worker owns. The minimal streaming-worker + # runtime makes this worker the sole owner of app.DATABASE, so releasing it + # when the loop is cancelled or exits (task cancellation / SIGINT / an + # unhandled error) frees the connection pool deterministically. A SIGKILL + # container stop bypasses this, but the OS reclaims the pool with the process. try: - await app.STORAGE.save_streaming_file(organization_id, file_name) + await app.DATABASE.engine.dispose() except Exception: - LOG.debug("Failed to upload screenshot", organization_id=organization_id, file_name=file_name) + LOG.debug("Failed to dispose streaming worker database engine") def main() -> None: diff --git a/skyvern/forge/forge_app_initializer.py b/skyvern/forge/forge_app_initializer.py index f2bc1f86f..133863c18 100644 --- a/skyvern/forge/forge_app_initializer.py +++ b/skyvern/forge/forge_app_initializer.py @@ -5,6 +5,11 @@ import structlog from skyvern.config import settings from skyvern.forge import set_force_app_instance from skyvern.forge.forge_app import ForgeApp, create_forge_app +from skyvern.forge.sdk.artifact.storage.azure import AzureStorage +from skyvern.forge.sdk.artifact.storage.factory import StorageFactory +from skyvern.forge.sdk.artifact.storage.gcs import GcsStorage +from skyvern.forge.sdk.artifact.storage.s3 import S3Storage +from skyvern.forge.sdk.db.agent_db import AgentDB LOG = structlog.get_logger() _SERVER_LOGGING_CONFIGURED = False @@ -44,3 +49,39 @@ def start_forge_app() -> ForgeApp: ) return force_app_instance + + +def start_streaming_worker_app() -> ForgeApp: + """Initialize the minimal app graph the screenshot streaming worker needs. + + The all-in-one container runs ``run_streaming.py`` as a second process + alongside the main server. That worker only reads run/task rows and writes + screenshot files, so it needs ONLY ``app.DATABASE`` and ``app.STORAGE`` -- + not the full ``create_forge_app()`` object graph (browser manager, LLM + clients/handlers, persistent-sessions manager, credential vaults, workflow + service, agent, replica DB, cache, ...). Building only the proven-minimal + bundle avoids duplicating that heavyweight fixed process state in a second + process. + + A startup failure here (e.g. an unreachable database) propagates and fails + the process closed, rather than being swallowed into a silent screenshot + loop; ``set_force_app_instance`` runs only once the minimal bundle is built. + """ + _ensure_server_logging_configured() + + force_app_instance = ForgeApp() + force_app_instance.SETTINGS_MANAGER = settings + force_app_instance.DATABASE = AgentDB(settings.DATABASE_STRING, debug_enabled=settings.DEBUG_MODE) + + # Storage backend selection mirrors create_forge_app(); keep the two in sync + # so a new backend is never silently downgraded to LocalStorage in the worker. + if settings.SKYVERN_STORAGE_TYPE == "s3": + StorageFactory.set_storage(S3Storage()) + elif settings.SKYVERN_STORAGE_TYPE == "azureblob": + StorageFactory.set_storage(AzureStorage()) + elif settings.SKYVERN_STORAGE_TYPE == "gcs": + StorageFactory.set_storage(GcsStorage()) + force_app_instance.STORAGE = StorageFactory.get_storage() + + set_force_app_instance(force_app_instance) + return force_app_instance diff --git a/tests/unit/test_run_streaming.py b/tests/unit/test_run_streaming.py new file mode 100644 index 000000000..ab23277d1 --- /dev/null +++ b/tests/unit/test_run_streaming.py @@ -0,0 +1,162 @@ +"""Unit tests for the screenshot streaming worker (``run_streaming.py``). + +These lock in two guarantees for the minimal-runtime rewrite: + 1. the screenshot loop still retrieves the same run data and emits the same + ``save_streaming_file`` payload as before, and + 2. the database engine the worker owns is deterministically disposed when the + loop exits (deterministic shutdown of the only heavyweight resource it now + initializes). +""" + +import asyncio +import importlib.util +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from skyvern.forge.sdk.workflow.models.workflow import WorkflowRunStatus + +_RUN_STREAMING_PATH = Path(__file__).resolve().parents[2] / "run_streaming.py" +_spec = importlib.util.spec_from_file_location("run_streaming", _RUN_STREAMING_PATH) +assert _spec and _spec.loader +run_streaming = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(run_streaming) + +_UNSET = object() + + +def _install_worker_fakes( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + state_payloads: list[object], + *, + workflow_run: object = _UNSET, + task: object = _UNSET, +) -> SimpleNamespace: + """Wire ``run_streaming`` with in-memory fakes and a scripted state-file reader. + + ``state_payloads`` is consumed one entry per loop iteration: a dict is + returned from ``get_json_from_file``; an exception instance is raised (use + ``asyncio.CancelledError`` to break the otherwise-infinite loop). + + ``workflow_run`` / ``task`` override the row returned by the DB reads (default: + an active run and a non-final task); pass ``None`` to simulate a missing row. + """ + dispose = AsyncMock() + wr_result = SimpleNamespace(status=WorkflowRunStatus.running) if workflow_run is _UNSET else workflow_run + task_result = SimpleNamespace(status=SimpleNamespace(is_final=lambda: False)) if task is _UNSET else task + get_workflow_run = AsyncMock(return_value=wr_result) + get_task = AsyncMock(return_value=task_result) + save_streaming_file = AsyncMock() + + fake_app = SimpleNamespace( + DATABASE=SimpleNamespace( + workflow_runs=SimpleNamespace(get_workflow_run=get_workflow_run), + tasks=SimpleNamespace(get_task=get_task), + engine=SimpleNamespace(dispose=dispose), + ), + STORAGE=SimpleNamespace(save_streaming_file=save_streaming_file), + ) + + monkeypatch.setattr(run_streaming, "start_streaming_worker_app", lambda: None) + monkeypatch.setattr(run_streaming, "app", fake_app) + monkeypatch.setattr(run_streaming, "initialize_skyvern_state_file", AsyncMock()) + monkeypatch.setattr(run_streaming, "INTERVAL", 0) + monkeypatch.setattr(run_streaming, "get_skyvern_state_file_path", lambda: str(tmp_path / "state.json")) + monkeypatch.setattr(run_streaming, "get_skyvern_temp_dir", lambda: str(tmp_path)) + monkeypatch.setattr(run_streaming, "subprocess", MagicMock()) + monkeypatch.setattr(run_streaming, "os", MagicMock()) + + payloads = iter(state_payloads) + + def _next_state(_path: str) -> object: + value = next(payloads) + if isinstance(value, BaseException): + raise value + return value + + monkeypatch.setattr(run_streaming, "get_json_from_file", _next_state) + + return SimpleNamespace( + app=fake_app, + dispose=dispose, + get_workflow_run=get_workflow_run, + get_task=get_task, + save_streaming_file=save_streaming_file, + ) + + +@pytest.mark.asyncio +async def test_run_streams_screenshot_for_active_workflow_run(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + fakes = _install_worker_fakes( + monkeypatch, + tmp_path, + state_payloads=[ + {"task_id": None, "workflow_run_id": "wr_1", "organization_id": "o_1"}, + asyncio.CancelledError(), + ], + ) + + with pytest.raises(asyncio.CancelledError): + await run_streaming.run() + + fakes.get_workflow_run.assert_awaited_once_with(workflow_run_id="wr_1") + fakes.save_streaming_file.assert_awaited_once_with("o_1", "wr_1.png") + + +@pytest.mark.asyncio +async def test_run_streams_screenshot_for_active_task(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + # The task path uses a different DB call (get_task, scoped by organization_id) + # and a different file name ({task_id}.png) than the workflow-run path. + fakes = _install_worker_fakes( + monkeypatch, + tmp_path, + state_payloads=[ + {"task_id": "t_1", "workflow_run_id": None, "organization_id": "o_1"}, + asyncio.CancelledError(), + ], + ) + + with pytest.raises(asyncio.CancelledError): + await run_streaming.run() + + fakes.get_task.assert_awaited_once_with(task_id="t_1", organization_id="o_1") + fakes.save_streaming_file.assert_awaited_once_with("o_1", "t_1.png") + + +@pytest.mark.asyncio +async def test_run_skips_screenshot_when_workflow_run_finalized( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + # A finalized run must short-circuit (continue) before any screenshot is saved. + fakes = _install_worker_fakes( + monkeypatch, + tmp_path, + state_payloads=[ + {"task_id": None, "workflow_run_id": "wr_1", "organization_id": "o_1"}, + asyncio.CancelledError(), + ], + workflow_run=SimpleNamespace(status=WorkflowRunStatus.completed), + ) + + with pytest.raises(asyncio.CancelledError): + await run_streaming.run() + + fakes.get_workflow_run.assert_awaited_once_with(workflow_run_id="wr_1") + fakes.save_streaming_file.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_disposes_database_engine_on_shutdown(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + fakes = _install_worker_fakes( + monkeypatch, + tmp_path, + state_payloads=[asyncio.CancelledError()], + ) + + with pytest.raises(asyncio.CancelledError): + await run_streaming.run() + + fakes.dispose.assert_awaited_once() diff --git a/tests/unit/test_streaming_worker_app.py b/tests/unit/test_streaming_worker_app.py new file mode 100644 index 000000000..2994f1a91 --- /dev/null +++ b/tests/unit/test_streaming_worker_app.py @@ -0,0 +1,110 @@ +"""Unit tests for the minimal streaming-worker runtime. + +The all-in-one container runs ``run_streaming.py`` alongside the main server. +That worker only reads run/task rows and writes screenshot files, so it must be +bootstrapped with a *minimal* app graph (``app.DATABASE`` + ``app.STORAGE``) +rather than the full ``create_forge_app()`` object graph (browser manager, LLM +clients/handlers, persistent-sessions manager, credential vaults, agent, ...). +""" + +from threading import Lock +from types import SimpleNamespace + +import pytest + +from skyvern.forge import forge_app_initializer + +# Heavyweight components that ``create_forge_app()`` constructs but the screenshot +# streaming worker never touches. If any of these show up on the streaming-worker +# app instance, the minimal-runtime guarantee has regressed. +HEAVY_COMPONENTS = [ + "REPLICA_DATABASE", + "CACHE", + "ARTIFACT_MANAGER", + "BROWSER_MANAGER", + "LLM_API_HANDLER", + "OPENAI_CLIENT", + "ANTHROPIC_CLIENT", + "SECONDARY_LLM_API_HANDLER", + "WORKFLOW_CONTEXT_MANAGER", + "WORKFLOW_SERVICE", + "AGENT_FUNCTION", + "PERSISTENT_SESSIONS_MANAGER", + "BROWSER_SESSION_RECORDING_SERVICE", + "BITWARDEN_CREDENTIAL_VAULT_SERVICE", + "agent", +] + + +def _stub_logging(monkeypatch: pytest.MonkeyPatch) -> None: + # Keep the initializer from reconfiguring process-wide logging during tests. + monkeypatch.setattr(forge_app_initializer, "_SERVER_LOGGING_CONFIGURED", True) + monkeypatch.setattr(forge_app_initializer, "_SERVER_LOGGING_LOCK", Lock()) + + +def _spy_agent_db(monkeypatch: pytest.MonkeyPatch) -> tuple[object, list[tuple[str, bool]]]: + calls: list[tuple[str, bool]] = [] + fake_db = SimpleNamespace(engine=SimpleNamespace()) + + def _fake_agent_db(database_string: str, debug_enabled: bool = False) -> object: + calls.append((database_string, debug_enabled)) + return fake_db + + monkeypatch.setattr(forge_app_initializer, "AgentDB", _fake_agent_db) + return fake_db, calls + + +def test_start_streaming_worker_app_does_not_build_full_forge_app(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_logging(monkeypatch) + fake_db, db_calls = _spy_agent_db(monkeypatch) + + create_forge_app_calls: list[int] = [] + monkeypatch.setattr(forge_app_initializer, "create_forge_app", lambda: create_forge_app_calls.append(1)) + + installed: list[object] = [] + monkeypatch.setattr(forge_app_initializer, "set_force_app_instance", installed.append) + + app_instance = forge_app_initializer.start_streaming_worker_app() + + # The whole point: never build the full app graph. + assert create_forge_app_calls == [] + # It constructs exactly one DB, wired from settings, and installs the app once. + assert db_calls == [(forge_app_initializer.settings.DATABASE_STRING, forge_app_initializer.settings.DEBUG_MODE)] + assert app_instance.DATABASE is fake_db + assert installed == [app_instance] + + +def test_start_streaming_worker_app_only_sets_minimal_attributes(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_logging(monkeypatch) + _spy_agent_db(monkeypatch) + monkeypatch.setattr(forge_app_initializer, "set_force_app_instance", lambda inst: None) + + app_instance = forge_app_initializer.start_streaming_worker_app() + + # Minimal dependency surface the worker actually reads. + assert hasattr(app_instance, "DATABASE") + assert hasattr(app_instance, "STORAGE") + assert app_instance.STORAGE is forge_app_initializer.StorageFactory.get_storage() + + # Everything heavy stays unconstructed (ForgeApp declares these as annotations + # only, so an un-set attribute raises AttributeError -> hasattr is False). + for component in HEAVY_COMPONENTS: + assert not hasattr(app_instance, component), f"{component} must not be initialized in the streaming worker" + + +def test_start_streaming_worker_app_startup_failure_is_fail_closed(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_logging(monkeypatch) + + def _boom(*_args: object, **_kwargs: object) -> object: + raise RuntimeError("database unreachable") + + monkeypatch.setattr(forge_app_initializer, "AgentDB", _boom) + + installed: list[object] = [] + monkeypatch.setattr(forge_app_initializer, "set_force_app_instance", installed.append) + + # A startup failure must surface (process fails to start), never be swallowed + # into a silent screenshot loop, and must never install a half-built app. + with pytest.raises(RuntimeError, match="database unreachable"): + forge_app_initializer.start_streaming_worker_app() + assert installed == []