vikarti.anatra/Skyvern
1
0
Fork 0
mirror of https://github.com/Skyvern-AI/skyvern.git synced 2026-10-02 19:57:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 6

refactor(SKY-12367): boot the screenshot streaming worker with a minimal app runtime (#7378)

Browse source 
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
LawyZheng 2026-07-17 00:02:35 +08:00 • committed by GitHub
parent d6c3f6caa3
commit 9baa005141
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 374 additions and 50 deletions
Download patch file Download diff file Expand all files Collapse all files

111
run_streaming.py
View file

@ -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:

41
skyvern/forge/forge_app_initializer.py
View file

@ -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

162
tests/unit/test_run_streaming.py Normal file
View file

@ -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()

110
tests/unit/test_streaming_worker_app.py Normal file
View file

@ -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 == []