SKY-9488: extract SDK run schema types (#5837)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

alembic/versions/2026_05_06_1930-2704b02dede9_add_max_steps_per_workflow_run_to_organizations.py

@@ -0,0 +1,30 @@
+"""add max_steps_per_workflow_run to organizations
+
+Revision ID: 2704b02dede9
+Revises: 53a306b96ed7
+Create Date: 2026-05-06T19:30:35.812486+00:00
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "2704b02dede9"
+down_revision: Union[str, None] = "53a306b96ed7"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "organizations",
+        sa.Column("max_steps_per_workflow_run", sa.Integer(), nullable=True),
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("organizations", "max_steps_per_workflow_run")

conftest.py

@@ -0,0 +1,8 @@
+from skyvern.forge.sdk.forge_log import setup_logger
+
+# After SKY-7947 (lightweight Skyvern wheel), `import skyvern` no longer
+# configures structlog as a side effect. Without this call the default
+# rich-based exception renderer is active during tests, which raises
+# `TypeError: 'Mock' object is not iterable` when it tries to introspect
+# Mock locals in tracebacks (e.g. utils_test.py::TestParseApiResponse).
+setup_logger()

skyvern-frontend/src/api/types.ts

@@ -208,6 +208,7 @@ export type OrganizationApiResponse = {
   // Optional because the field is added in a backend image rollout; until
   // every API host has the new schema this property may be absent from the
   // response payload.
+  max_steps_per_workflow_run?: number | null;
   artifact_url_expiry_seconds?: number | null;
   organization_id: string;
   organization_name: string;

skyvern-frontend/src/routes/workflows/workflowRun/TaskBlockParameters.tsx

@@ -281,7 +281,7 @@ function TaskBlockParameters({ block, definitionBlock }: Props) {
           {typeof maxStepsPerRun === "number" ? (
             <div className="flex gap-16">
               <div className="w-80">
-                <h1 className="text-lg">Max Steps Per Run</h1>
+                <h1 className="text-lg">Max Steps Per Block</h1>
               </div>
               <Input value={maxStepsPerRun.toString()} readOnly />
             </div>

skyvern-frontend/src/routes/workflows/workflowRun/blockInfo/FileDownloadBlockParameters.tsx

@@ -66,7 +66,7 @@ function FileDownloadBlockParameters({
       {typeof maxStepsPerRun === "number" ? (
         <div className="flex gap-16">
           <div className="w-80">
-            <h1 className="text-lg">Max Steps Per Run</h1>
+            <h1 className="text-lg">Max Steps Per Block</h1>
           </div>
           <Input value={maxStepsPerRun.toString()} readOnly />
         </div>

skyvern/__init__.py

@@ -5,51 +5,20 @@ from skyvern._version import __version__
 from skyvern.utils import setup_windows_event_loop_policy
 
 if typing.TYPE_CHECKING:
-    from skyvern.library import Skyvern  # noqa: E402
-
-
-def _setup_package_logging() -> None:
-    from skyvern.forge.sdk.forge_log import setup_logger  # noqa: PLC0415
-
-    setup_logger()
+    from skyvern.library import Skyvern  # noqa: E402,F401
 
 
 setup_windows_event_loop_policy()
-_setup_package_logging()
+# Server entrypoints configure package logging; base SDK imports stay side-effect-light.
 
-# noinspection PyUnresolvedReferences
-__all__ = [
-    "__version__",
-    "Skyvern",
-    "SkyvernPage",
-    "RunContext",
-    "action",
-    "cached",
-    "conditional",
-    "download",
-    "extract",
-    "http_request",
-    "goto",
-    "login",
-    "loop",
-    "parse_file",
-    "parse_pdf",
-    "prompt",
-    "render_list",
-    "render_template",
-    "run_code",
-    "run_script",
-    "run_task",
-    "send_email",
-    "setup",
-    "upload_file",
-    "validate",
-    "wait",
-    "workflow",
-]
-
-_lazy_imports = {
+_base_lazy_imports = {
     "Skyvern": "skyvern.library",
+    "SkyvernEnvironment": "skyvern.client",
+    "RunEngine": "skyvern.schemas.run_enums",
+    "RunStatus": "skyvern.schemas.run_enums",
+}
+
+_server_lazy_imports = {
     "SkyvernPage": "skyvern.core.script_generations.skyvern_page",
     "RunContext": "skyvern.core.script_generations.skyvern_page",
     "setup": "skyvern.core.script_generations.run_initializer",
@@ -76,14 +45,25 @@ _lazy_imports = {
     "validate": "skyvern.services.script_service",
     "wait": "skyvern.services.script_service",
 }
+# Keep server-only names in __all__ for backwards-compatible wildcard exports;
+# direct access remains lazily gated by the skyvern[server] extra.
+_all_lazy_imports = {**_base_lazy_imports, **_server_lazy_imports}
+__all__ = ["__version__", *_base_lazy_imports, *_server_lazy_imports]
 
 
 def __getattr__(name: str) -> Any:
-    if name in _lazy_imports:
-        module_path = _lazy_imports[name]
+    if name in _all_lazy_imports:
+        module_path = _all_lazy_imports[name]
         from importlib import import_module  # noqa: PLC0415
 
-        module = import_module(module_path)
+        try:
+            module = import_module(module_path)
+        except ImportError as exc:
+            if name in _server_lazy_imports:
+                from skyvern.exceptions import raise_server_extra_required  # noqa: PLC0415
+
+                raise_server_extra_required(f"skyvern.{name}", exc)
+            raise
 
         # For attributes that need to be extracted from the module
         if hasattr(module, name):

skyvern/analytics.py

@@ -1,26 +1,44 @@
+from __future__ import annotations
+
 import atexit
 import functools
 import importlib.metadata
 import platform
 import threading
 import traceback
-from typing import Any, Dict, Optional
+from typing import TYPE_CHECKING, Any
 
 import structlog
 import typer
-from posthog import Posthog
 
 from skyvern._version import __version__ as _build_version
 from skyvern.config import settings
 
+if TYPE_CHECKING:
+    from posthog import Posthog
+
 LOG = structlog.get_logger(__name__)
 
 _FLUSH_TIMEOUT_SECONDS = 2
 
 
-def _build_posthog_client(api_key: str, host: str) -> Posthog | None:
+def _load_posthog_class() -> type[Posthog] | None:
     try:
-        client = Posthog(api_key, host=host, disable_geoip=False, timeout=2, max_retries=1)
+        from posthog import Posthog  # noqa: PLC0415
+    except ModuleNotFoundError as exc:
+        if exc.name == "posthog":
+            return None
+        raise
+    return Posthog
+
+
+def _build_posthog_client(api_key: str, host: str) -> Posthog | None:
+    posthog_cls = _load_posthog_class()
+    if posthog_cls is None:
+        return None
+
+    try:
+        client = posthog_cls(api_key, host=host, disable_geoip=False, timeout=2, max_retries=1)
         atexit.unregister(client.join)
         return client
     except Exception:
@@ -48,7 +66,7 @@ def get_oss_version() -> str:
 
 
 @functools.lru_cache(maxsize=1)
-def analytics_metadata() -> Dict[str, Any]:
+def analytics_metadata() -> dict[str, Any]:
     # Cached: all fields are process-lifetime constants. Do not add dynamic fields here.
     return {
         "os": platform.system().lower(),
@@ -60,7 +78,7 @@ def analytics_metadata() -> Dict[str, Any]:
     }
 
 
-def dynamic_analytics_metadata() -> Dict[str, Any]:
+def dynamic_analytics_metadata() -> dict[str, Any]:
     metadata: dict[str, Any] = {}
     if settings.ANALYTICS_TEST_ID:
         metadata["analytics_test_id"] = settings.ANALYTICS_TEST_ID
@@ -132,9 +150,9 @@ def capture(
 def capture_setup_event(
     event_name: str,
     success: bool = True,
-    error_type: Optional[str] = None,
-    error_message: Optional[str] = None,
-    extra_data: Optional[dict[str, Any]] = None,
+    error_type: str | None = None,
+    error_message: str | None = None,
+    extra_data: dict[str, Any] | None = None,
 ) -> None:
     """Capture a setup-related analytics event.
 
@@ -162,8 +180,8 @@ def capture_setup_event(
 def capture_setup_error(
     event_name: str,
     error: Exception,
-    error_type: Optional[str] = None,
-    extra_data: Optional[dict[str, Any]] = None,
+    error_type: str | None = None,
+    extra_data: dict[str, Any] | None = None,
 ) -> None:
     """Capture a setup error with exception details.
 

skyvern/cli/mcp_tools/org.py

@@ -47,7 +47,9 @@ async def skyvern_org_update(
         Field(
             description=(
                 "Partial settings dict. Allowed keys: "
-                "max_steps_per_run (int >= 1), "
+                "max_steps_per_run (int >= 1, per-block cap), "
+                "max_steps_per_workflow_run (int >= 1, per-workflow-run total cap), "
+                "clear_max_steps_per_workflow_run (bool — resets the run-level cap to NULL), "
                 "max_retries_per_step (int >= 0), "
                 "webhook_callback_url (string), "
                 "artifact_url_expiry_seconds (int 3600-604800), "

skyvern/exceptions.py

@@ -1,6 +1,32 @@
 import re
+from http import HTTPStatus
+from importlib.util import find_spec
+from typing import NoReturn
 
-from fastapi import status
+# Representative modules that indicate the server extra is installed enough for
+# server/local/browser import graphs. Keep this list intentionally small, but
+# include the heavy modules users commonly have partially installed.
+_SERVER_EXTRA_SENTINELS = (
+    "fastapi",
+    "jinja2",
+    "libcst",
+    "litellm",
+    "playwright",
+    "sqlalchemy",
+    "starlette",
+    "starlette_context",
+)
+
+
+def _missing_server_extra_dependency(module_name: str) -> bool:
+    root_module = module_name.split(".", maxsplit=1)[0]
+    if root_module in _SERVER_EXTRA_SENTINELS:
+        return find_spec(root_module) is None
+    if root_module == "skyvern":
+        return False
+    # Unknown missing modules may be genuine dependency bugs, so only known
+    # server-extra sentinels are rewritten to the install hint.
+    return False
 
 
 class SkyvernException(Exception):
@@ -9,6 +35,28 @@ class SkyvernException(Exception):
         super().__init__(message)
 
 
+class SkyvernExtraNotInstalled(ImportError):
+    def __init__(self, feature: str, extra: str = "server"):
+        self.feature = feature
+        self.extra = extra
+        super().__init__(f"{feature} requires server support. Install it with `pip install skyvern[{extra}]`.")
+
+
+def raise_server_extra_required(feature: str, exc: ImportError) -> NoReturn:
+    if isinstance(exc, ModuleNotFoundError) and exc.name is not None and _missing_server_extra_dependency(exc.name):
+        raise SkyvernExtraNotInstalled(feature) from exc
+    raise exc
+
+
+def require_server_extra_modules(feature: str, module_names: tuple[str, ...] = ()) -> None:
+    # Server/local/browser APIs require the full server extra, not a partial Playwright-only install.
+    required_modules = dict.fromkeys((*_SERVER_EXTRA_SENTINELS, *module_names))
+    for module_name in required_modules:
+        if find_spec(module_name) is None:
+            missing = ModuleNotFoundError(f"No module named '{module_name}'", name=module_name)
+            raise SkyvernExtraNotInstalled(feature) from missing
+
+
 class SkyvernClientException(SkyvernException):
     def __init__(self, message: str | None = None, status_code: int | None = None):
         self.status_code = status_code
@@ -16,8 +64,8 @@ class SkyvernClientException(SkyvernException):
 
 
 class SkyvernHTTPException(SkyvernException):
-    def __init__(self, message: str | None = None, status_code: int = status.HTTP_400_BAD_REQUEST):
-        self.status_code = status_code
+    def __init__(self, message: str | None = None, status_code: int | HTTPStatus = HTTPStatus.BAD_REQUEST):
+        self.status_code = int(status_code)
         super().__init__(message)
 
 
@@ -55,7 +103,7 @@ def get_user_facing_exception_message(exception: Exception) -> str:
 
 class DisabledBlockExecutionError(SkyvernHTTPException):
     def __init__(self, message: str | None = None):
-        super().__init__(message, status_code=status.HTTP_400_BAD_REQUEST)
+        super().__init__(message, status_code=HTTPStatus.BAD_REQUEST)
 
 
 class RateLimitExceeded(SkyvernHTTPException):
@@ -64,7 +112,7 @@ class RateLimitExceeded(SkyvernHTTPException):
             f"Rate limit exceeded for organization {organization_id}. "
             f"Maximum {max_requests} requests per {window_seconds} seconds allowed."
         )
-        super().__init__(message, status_code=status.HTTP_429_TOO_MANY_REQUESTS)
+        super().__init__(message, status_code=HTTPStatus.TOO_MANY_REQUESTS)
 
 
 class InvalidOpenAIResponseFormat(SkyvernException):
@@ -97,7 +145,7 @@ class WebhookReplayError(SkyvernHTTPException):
         self,
         message: str | None = None,
         *,
-        status_code: int = status.HTTP_400_BAD_REQUEST,
+        status_code: int | HTTPStatus = HTTPStatus.BAD_REQUEST,
     ):
         super().__init__(message=message or "Webhook replay failed.", status_code=status_code)
 
@@ -114,7 +162,7 @@ class MissingApiKey(WebhookReplayError):
 
 class TaskNotFound(SkyvernHTTPException):
     def __init__(self, task_id: str | None = None):
-        super().__init__(f"Task {task_id} not found", status_code=status.HTTP_404_NOT_FOUND)
+        super().__init__(f"Task {task_id} not found", status_code=HTTPStatus.NOT_FOUND)
 
 
 class MissingElement(SkyvernException):
@@ -212,7 +260,7 @@ class WorkflowNotFound(SkyvernHTTPException):
 
         super().__init__(
             f"Workflow not found. {workflow_repr}",
-            status_code=status.HTTP_404_NOT_FOUND,
+            status_code=HTTPStatus.NOT_FOUND,
         )
 
 
@@ -223,20 +271,20 @@ class WorkflowNotFoundForWorkflowRun(SkyvernHTTPException):
     ) -> None:
         super().__init__(
             f"Workflow not found for workflow run {workflow_run_id}",
-            status_code=status.HTTP_404_NOT_FOUND,
+            status_code=HTTPStatus.NOT_FOUND,
         )
 
 
 class WorkflowRunNotFound(SkyvernHTTPException):
     def __init__(self, workflow_run_id: str) -> None:
-        super().__init__(f"WorkflowRun {workflow_run_id} not found", status_code=status.HTTP_404_NOT_FOUND)
+        super().__init__(f"WorkflowRun {workflow_run_id} not found", status_code=HTTPStatus.NOT_FOUND)
 
 
 class MissingValueForParameter(SkyvernHTTPException):
     def __init__(self, parameter_key: str, workflow_id: str, workflow_run_id: str) -> None:
         super().__init__(
             f"Missing value for parameter {parameter_key} in workflow run {workflow_run_id} of workflow {workflow_id}",
-            status_code=status.HTTP_400_BAD_REQUEST,
+            status_code=HTTPStatus.BAD_REQUEST,
         )
 
 
@@ -276,7 +324,7 @@ class InvalidCredentialId(SkyvernHTTPException):
         super().__init__(
             f"Invalid credential ID: {sanitize_credential_for_error(credential_id)}."
             " Failed to resolve to a valid credential.",
-            status_code=status.HTTP_400_BAD_REQUEST,
+            status_code=HTTPStatus.BAD_REQUEST,
         )
 
 
@@ -284,7 +332,7 @@ class WorkflowParameterNotFound(SkyvernHTTPException):
     def __init__(self, workflow_parameter_id: str) -> None:
         super().__init__(
             f"Workflow parameter {workflow_parameter_id} not found",
-            status_code=status.HTTP_404_NOT_FOUND,
+            status_code=HTTPStatus.NOT_FOUND,
         )
 
 
@@ -411,7 +459,7 @@ class OrganizationNotFound(SkyvernHTTPException):
     def __init__(self, organization_id: str) -> None:
         super().__init__(
             f"Organization {organization_id} not found",
-            status_code=status.HTTP_404_NOT_FOUND,
+            status_code=HTTPStatus.NOT_FOUND,
         )
 
 
@@ -419,7 +467,7 @@ class StepNotFound(SkyvernHTTPException):
     def __init__(self, organization_id: str, task_id: str, step_id: str | None = None) -> None:
         super().__init__(
             f"Step {step_id or 'latest'} not found. organization_id={organization_id} task_id={task_id}",
-            status_code=status.HTTP_404_NOT_FOUND,
+            status_code=HTTPStatus.NOT_FOUND,
         )
 
 
@@ -856,7 +904,7 @@ class BlockedHost(SkyvernHTTPException):
     def __init__(self, host: str) -> None:
         super().__init__(
             f"The host in your url is blocked: {host}",
-            status_code=status.HTTP_400_BAD_REQUEST,
+            status_code=HTTPStatus.BAD_REQUEST,
         )
 
 
@@ -867,7 +915,7 @@ class InvalidWorkflowParameter(SkyvernHTTPException):
             message += f" Workflow permanent id: {workflow_permanent_id}"
         super().__init__(
             message,
-            status_code=status.HTTP_400_BAD_REQUEST,
+            status_code=HTTPStatus.BAD_REQUEST,
         )
 
 
@@ -1007,7 +1055,7 @@ class BrowserSessionClosed(SkyvernHTTPException):
     def __init__(self, browser_session_id: str) -> None:
         super().__init__(
             f"Browser session {browser_session_id} is closed.",
-            status_code=status.HTTP_410_GONE,
+            status_code=HTTPStatus.GONE,
         )
 
 
@@ -1015,7 +1063,7 @@ class BrowserSessionNotFound(SkyvernHTTPException):
     def __init__(self, browser_session_id: str) -> None:
         super().__init__(
             f"Browser session {browser_session_id} does not exist or is not live.",
-            status_code=status.HTTP_404_NOT_FOUND,
+            status_code=HTTPStatus.NOT_FOUND,
         )
 
 
@@ -1023,7 +1071,7 @@ class BrowserSessionStartupTimeout(SkyvernHTTPException):
     def __init__(self, browser_session_id: str) -> None:
         super().__init__(
             f"Browser session {browser_session_id} failed to start within the timeout period.",
-            status_code=status.HTTP_504_GATEWAY_TIMEOUT,
+            status_code=HTTPStatus.GATEWAY_TIMEOUT,
         )
 
 
@@ -1032,7 +1080,7 @@ class BrowserProfileNotFound(SkyvernHTTPException):
         message = f"Browser profile {profile_id} not found"
         if organization_id:
             message += f" for organization {organization_id}"
-        super().__init__(message, status_code=status.HTTP_404_NOT_FOUND)
+        super().__init__(message, status_code=HTTPStatus.NOT_FOUND)
 
 
 class APIKeyNotFound(SkyvernHTTPException):
@@ -1064,7 +1112,7 @@ class OutputParameterNotFound(SkyvernHTTPException):
     def __init__(self, block_label: str, workflow_permanent_id: str) -> None:
         super().__init__(
             f"Output parameter for {block_label} not found in workflow {workflow_permanent_id}",
-            status_code=status.HTTP_400_BAD_REQUEST,
+            status_code=HTTPStatus.BAD_REQUEST,
         )
 
 
@@ -1073,7 +1121,7 @@ class TemporalSubmissionFailed(SkyvernHTTPException):
         workflow_run_str = f" for workflow_run_id={workflow_run_id}" if workflow_run_id else ""
         super().__init__(
             f"Failed to submit {workflow_type} to Temporal{workflow_run_str}",
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
         )
 
 

skyvern/forge/__main__.py

@@ -3,13 +3,11 @@ import sys
 from pathlib import Path
 
 import structlog
-import uvicorn
 from dotenv import load_dotenv
-from uvicorn.config import Config as _UvicornConfig
-from uvicorn.supervisors.watchfilesreload import FileFilter as _FileFilter
 
 from skyvern import analytics
 from skyvern.config import settings
+from skyvern.exceptions import require_server_extra_modules
 
 LOG = structlog.stdlib.get_logger()
 
@@ -28,6 +26,9 @@ def _build_reload_excludes() -> list[str]:
 
 
 def _verify_reload_excludes_cover_artifacts(reload_excludes: list[str]) -> None:
+    from uvicorn.config import Config as _UvicornConfig  # noqa: PLC0415
+    from uvicorn.supervisors.watchfilesreload import FileFilter as _FileFilter  # noqa: PLC0415
+
     # A miss lets watchfiles restart on artifact writes during long browser cleanups,
     # deadlocking the supervisor on Process.join.
     if not settings.ARTIFACT_STORAGE_PATH:
@@ -55,6 +56,13 @@ def _verify_reload_excludes_cover_artifacts(reload_excludes: list[str]) -> None:
 
 
 if __name__ == "__main__":
+    require_server_extra_modules("skyvern.forge", ("uvicorn",))
+
+    import uvicorn
+
+    from skyvern.forge.forge_app_initializer import _ensure_server_logging_configured
+
+    _ensure_server_logging_configured()
     analytics.capture("skyvern-oss-run-server")
     port = settings.PORT
     LOG.info("Agent server starting.", host="0.0.0.0", port=port)

skyvern/forge/agent.py

@@ -4564,6 +4564,18 @@ class ForgeAgent:
             or settings.MAX_STEPS_PER_RUN
         )
 
+        workflow_run_budget = await self._check_workflow_run_step_budget(organization, task)
+        if workflow_run_budget is not None and workflow_run_budget[0] >= workflow_run_budget[1]:
+            last_step = await self._terminate_for_workflow_run_step_budget(
+                organization=organization,
+                task=task,
+                step=step,
+                total_workflow_run_steps=workflow_run_budget[0],
+                max_steps_per_workflow_run=workflow_run_budget[1],
+                speculative_next_step=next_step,
+            )
+            return False, last_step, None
+
         if step.order + 1 >= max_steps_per_run:
             LOG.info(
                 "Step completed but max steps reached, marking task as failed",
@@ -5007,6 +5019,94 @@ class ForgeAgent:
                 errors=[],
             )
 
+    async def _check_workflow_run_step_budget(
+        self,
+        organization: Organization,
+        task: Task,
+    ) -> tuple[int, int] | None:
+        """Returns (total_steps_so_far, max_steps_per_workflow_run) when the org has a
+        run-level cap and this task belongs to a workflow run; otherwise None.
+
+        ``max_steps_per_workflow_run`` is a per-org cap on total step count summed across
+        all task blocks within a single workflow run. Distinct from the per-block
+        ``max_steps_per_run`` ceiling.
+        """
+        if not task.workflow_run_id or not organization.max_steps_per_workflow_run:
+            return None
+        workflow_run_tasks = await app.DATABASE.tasks.get_tasks_by_workflow_run_id(task.workflow_run_id)
+        task_ids = [t.task_id for t in workflow_run_tasks]
+        if not task_ids:
+            return 0, organization.max_steps_per_workflow_run
+        total_steps = await app.DATABASE.tasks.get_total_unique_step_order_count_by_task_ids(
+            task_ids=task_ids,
+            organization_id=organization.organization_id,
+        )
+        return total_steps or 0, organization.max_steps_per_workflow_run
+
+    async def _terminate_for_workflow_run_step_budget(
+        self,
+        *,
+        organization: Organization,
+        task: Task,
+        step: Step,
+        total_workflow_run_steps: int,
+        max_steps_per_workflow_run: int,
+        speculative_next_step: Step | None = None,
+    ) -> Step:
+        """Mark ``step`` as the last step of ``task`` and fail the task because the
+        workflow run has hit its total-step budget. Returns the updated last step.
+
+        ``speculative_next_step`` is the parallel-verification path's pre-created next
+        step that must be cancelled when present.
+        """
+        LOG.info(
+            "Step completed but workflow-run max steps reached, marking task as failed",
+            step_order=step.order,
+            step_retry=step.retry_index,
+            workflow_run_id=task.workflow_run_id,
+            total_workflow_run_steps=total_workflow_run_steps,
+            max_steps_per_workflow_run=max_steps_per_workflow_run,
+        )
+        if speculative_next_step is not None:
+            final_status = step.speculative_original_status or StepStatus.completed
+            step.speculative_original_status = None
+            step.status = final_status
+            last_step = await self.update_step(
+                step,
+                status=final_status,
+                output=step.output,
+                is_last=True,
+            )
+        else:
+            last_step = await self.update_step(step, is_last=True)
+
+        failure_reason = (
+            f"Workflow run reached the maximum total steps ({max_steps_per_workflow_run}) set on the organization."
+        )
+        errors = [ReachMaxStepsError().model_dump()]
+        failure_category = classify_from_failure_reason(failure_reason, fallback_to_unknown=True)
+        LOG.info(
+            "Task failure classified",
+            task_id=task.task_id,
+            workflow_run_id=task.workflow_run_id,
+            organization_id=task.organization_id,
+            task_status="failed",
+            failure_category=failure_category,
+            primary_failure_category=failure_category[0].get("category") if failure_category else None,
+            failure_category_source="code_level",
+            failure_category_path="max_steps_per_workflow_run",
+        )
+        if speculative_next_step is not None:
+            await self._cancel_speculative_step(speculative_next_step)
+        await self.update_task(
+            task,
+            status=TaskStatus.failed,
+            failure_reason=failure_reason,
+            errors=errors,
+            failure_category=failure_category,
+        )
+        return last_step
+
     async def handle_completed_step(
         self,
         organization: Organization,
@@ -5124,6 +5224,17 @@ class ForgeAgent:
             )
             return True, last_step, None
 
+        workflow_run_budget = await self._check_workflow_run_step_budget(organization, task)
+        if workflow_run_budget is not None and workflow_run_budget[0] >= workflow_run_budget[1]:
+            last_step = await self._terminate_for_workflow_run_step_budget(
+                organization=organization,
+                task=task,
+                step=step,
+                total_workflow_run_steps=workflow_run_budget[0],
+                max_steps_per_workflow_run=workflow_run_budget[1],
+            )
+            return False, last_step, None
+
         if step.order + 1 >= max_steps_per_run:
             LOG.info(
                 "Step completed but max steps reached, marking task as failed",

skyvern/forge/api_app.py

@@ -1,3 +1,8 @@
+"""FastAPI application factory for server installs."""
+
+# The server-extra guard must run before FastAPI/Starlette imports.
+# ruff: noqa: E402
+
 import os
 import uuid
 from contextlib import asynccontextmanager
@@ -5,11 +10,16 @@ from datetime import datetime
 from typing import Any, AsyncGenerator, Awaitable, Callable
 
 import structlog
+from pydantic import ValidationError
+
+from skyvern.exceptions import require_server_extra_modules
+
+require_server_extra_modules("skyvern.forge.api_app", ("fastapi", "starlette", "starlette_context"))
+
 from fastapi import FastAPI, Response, status
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.openapi.utils import get_openapi
 from fastapi.responses import JSONResponse
-from pydantic import ValidationError
 from starlette.requests import HTTPConnection, Request
 from starlette_context.middleware import RawContextMiddleware
 from starlette_context.plugins.base import Plugin

skyvern/forge/forge_app_initializer.py

@@ -1,3 +1,5 @@
+from threading import Lock
+
 import structlog
 
 from skyvern.config import settings
@@ -5,9 +7,25 @@ from skyvern.forge import set_force_app_instance
 from skyvern.forge.forge_app import ForgeApp, create_forge_app
 
 LOG = structlog.get_logger()
+_SERVER_LOGGING_CONFIGURED = False
+_SERVER_LOGGING_LOCK = Lock()
+
+
+def _ensure_server_logging_configured() -> None:
+    global _SERVER_LOGGING_CONFIGURED
+    with _SERVER_LOGGING_LOCK:
+        if _SERVER_LOGGING_CONFIGURED:
+            return
+
+        from skyvern.forge.sdk.forge_log import setup_logger  # noqa: PLC0415
+
+        setup_logger()
+        _SERVER_LOGGING_CONFIGURED = True
 
 
 def start_forge_app() -> ForgeApp:
+    _ensure_server_logging_configured()
+
     force_app_instance = create_forge_app()
     set_force_app_instance(force_app_instance)
 

skyvern/forge/sdk/api/llm/__init__.py

@@ -1,4 +1,5 @@
-import litellm
+"""LLM package namespace.
 
-# litellm's aiohttp_transport drops per-request timeout; httpx default honors it.
-litellm.disable_aiohttp_transport = True
+Keep this package import side-effect free so legacy type shims can be imported
+from a base SDK install without requiring the server extra.
+"""

skyvern/forge/sdk/api/llm/api_handler_factory.py

@@ -30,11 +30,7 @@ from skyvern.forge.sdk.api.llm.exceptions import (
     LLMProviderError,
     LLMProviderErrorRetryableTask,
 )
-from skyvern.forge.sdk.api.llm.models import (
-    LLMAllowedFailsPolicy,
-    LLMConfig,
-    LLMRouterConfig,
-)
+from skyvern.forge.sdk.api.llm.litellm_transport import configure_litellm_transport
 from skyvern.forge.sdk.api.llm.ui_tars_response import UITarsResponse
 from skyvern.forge.sdk.api.llm.utils import (
     is_image_message,
@@ -50,8 +46,17 @@ from skyvern.forge.sdk.models import SpeculativeLLMMetadata, Step
 from skyvern.forge.sdk.schemas.ai_suggestions import AISuggestion
 from skyvern.forge.sdk.schemas.task_v2 import TaskV2, Thought
 from skyvern.forge.sdk.trace import apply_context_attrs, traced
+from skyvern.schemas.llm import (
+    LLMAllowedFailsPolicy,
+    LLMConfig,
+    LLMRouterConfig,
+)
 from skyvern.utils.image_resizer import Resolution, get_resize_target_dimension, resize_screenshots
 
+# Keep this server-only side effect out of the package __init__ so the legacy
+# models shim can import without litellm. Legacy LLM calls enter this module.
+configure_litellm_transport()
+
 LOG = structlog.get_logger()
 
 # Canonical span name for all LLM chokepoints. Milestone 1 of the agent

skyvern/forge/sdk/api/llm/config_registry.py

@@ -6,7 +6,7 @@ from skyvern.forge.sdk.api.llm.exceptions import (
     InvalidLLMConfigError,
     MissingLLMProviderEnvVarsError,
 )
-from skyvern.forge.sdk.api.llm.models import LiteLLMParams, LLMConfig, LLMRouterConfig
+from skyvern.schemas.llm import LiteLLMParams, LLMConfig, LLMRouterConfig
 
 LOG = structlog.get_logger()
 

skyvern/forge/sdk/api/llm/litellm_transport.py

@@ -0,0 +1,5 @@
+def configure_litellm_transport() -> None:
+    import litellm
+
+    # litellm's aiohttp transport drops per-request timeout; httpx honors it.
+    litellm.disable_aiohttp_transport = True

skyvern/forge/sdk/api/llm/models.py

@@ -1,93 +1,37 @@
-from dataclasses import dataclass, field
-from typing import Any, Literal, Optional, TypedDict
+from __future__ import annotations
 
-from skyvern.forge.sdk.settings_manager import SettingsManager
+import warnings
+from typing import TYPE_CHECKING, Any
+
+from skyvern.schemas import llm as _llm
+
+__all__ = [
+    "LiteLLMParams",
+    "LLMAllowedFailsPolicy",
+    "LLMConfig",
+    "LLMConfigBase",
+    "LLMRouterConfig",
+    "LLMRouterModelConfig",
+]
+
+_DEPRECATED_EXPORTS = frozenset(__all__)
+_DEPRECATION_MESSAGE = (
+    "skyvern.forge.sdk.api.llm.models is deprecated; import LLM types from skyvern.schemas.llm instead."
+)
+
+if TYPE_CHECKING:
+    from skyvern.schemas.llm import (  # noqa: F401
+        LiteLLMParams,
+        LLMAllowedFailsPolicy,
+        LLMConfig,
+        LLMConfigBase,
+        LLMRouterConfig,
+        LLMRouterModelConfig,
+    )
 
 
-class LiteLLMParams(TypedDict, total=False):
-    api_key: str | None
-    api_version: str | None
-    api_base: str | None
-    model_info: dict[str, Any] | None
-    vertex_credentials: str | None
-    vertex_location: str | None
-    thinking: dict[str, Any] | None
-    thinking_level: str | None
-    service_tier: str | None
-    extra_headers: dict[str, str] | None
-    timeout: float | None
-
-
-@dataclass(frozen=True)
-class LLMConfigBase:
-    model_name: str
-    required_env_vars: list[str]
-    supports_vision: bool
-    add_assistant_prefix: bool
-
-    def get_missing_env_vars(self) -> list[str]:
-        missing_env_vars = []
-        for env_var in self.required_env_vars:
-            env_var_value = getattr(SettingsManager.get_settings(), env_var, None)
-            if not env_var_value:
-                missing_env_vars.append(env_var)
-
-        return missing_env_vars
-
-
-@dataclass(frozen=True)
-class LLMConfig(LLMConfigBase):
-    litellm_params: Optional[LiteLLMParams] = field(default=None)
-    max_tokens: int | None = SettingsManager.get_settings().LLM_CONFIG_MAX_TOKENS
-    max_completion_tokens: int | None = None
-    temperature: float | None = SettingsManager.get_settings().LLM_CONFIG_TEMPERATURE
-    reasoning_effort: str | None = None
-
-
-@dataclass(frozen=True)
-class LLMAllowedFailsPolicy:
-    bad_request_error_allowed_fails: int | None = None
-    authentication_error_allowed_fails: int | None = None
-    timeout_error_allowed_fails: int | None = None
-    rate_limit_error_allowed_fails: int | None = None
-    content_policy_violation_error_allowed_fails: int | None = None
-    internal_server_error_allowed_fails: int | None = None
-
-
-@dataclass(frozen=True)
-class LLMRouterModelConfig:
-    model_name: str
-    # https://litellm.vercel.app/docs/routing
-    litellm_params: dict[str, Any]
-    model_info: dict[str, Any] = field(default_factory=dict)
-    tpm: int | None = None
-    rpm: int | None = None
-
-
-@dataclass(frozen=True)
-class LLMRouterConfig(LLMConfigBase):
-    model_list: list[LLMRouterModelConfig]
-    # All three redis parameters are required. Even if there isn't a password, it should be an empty string.
-    main_model_group: str
-    redis_host: str | None = None
-    redis_port: int | None = None
-    redis_password: str | None = None
-    fallback_model_group: str | None = None
-    routing_strategy: Literal[
-        "simple-shuffle",
-        "least-busy",
-        "usage-based-routing",
-        "usage-based-routing-v2",
-        "latency-based-routing",
-    ] = "usage-based-routing"
-    num_retries: int = 1
-    retry_delay_seconds: int = 15
-    set_verbose: bool = False
-    disable_cooldowns: bool | None = None
-    allowed_fails: int | None = None
-    allowed_fails_policy: LLMAllowedFailsPolicy | None = None
-    cooldown_time: float | None = None
-    max_tokens: int | None = SettingsManager.get_settings().LLM_CONFIG_MAX_TOKENS
-    max_completion_tokens: int | None = None
-    reasoning_effort: str | None = None
-    temperature: float | None = SettingsManager.get_settings().LLM_CONFIG_TEMPERATURE
+def __getattr__(name: str) -> Any:
+    if name in _DEPRECATED_EXPORTS:
+        warnings.warn(_DEPRECATION_MESSAGE, DeprecationWarning, stacklevel=2)
+        return getattr(_llm, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

skyvern/forge/sdk/copilot/model_resolver.py

@@ -26,12 +26,13 @@ from agents.run_config import RunConfig
 from skyvern.config import settings
 from skyvern.forge.sdk.api.llm.config_registry import LLMConfigRegistry
 from skyvern.forge.sdk.api.llm.exceptions import InvalidLLMConfigError
-from skyvern.forge.sdk.api.llm.models import LLMConfig, LLMRouterConfig
+from skyvern.forge.sdk.api.llm.litellm_transport import configure_litellm_transport
 from skyvern.forge.sdk.copilot.session_factory import (
     copilot_call_model_input_filter,
     copilot_session_input_callback,
 )
 from skyvern.forge.sdk.copilot.tracing_setup import is_tracing_enabled
+from skyvern.schemas.llm import LLMConfig, LLMRouterConfig
 
 LOG = structlog.get_logger()
 
@@ -125,6 +126,8 @@ def resolve_model_config(llm_api_handler: Any) -> tuple[str, RunConfig, str, boo
 
     Returns (model_name, run_config, llm_key, supports_vision).
     """
+    configure_litellm_transport()
+
     llm_key = getattr(llm_api_handler, "llm_key", None) or settings.LLM_KEY
     config = LLMConfigRegistry.get_config(llm_key)
 

skyvern/forge/sdk/db/models.py

@@ -170,6 +170,7 @@ class OrganizationModel(Base):
     organization_name = Column(String, nullable=False)
     webhook_callback_url = Column(UnicodeText)
     max_steps_per_run = Column(Integer, nullable=True)
+    max_steps_per_workflow_run = Column(Integer, nullable=True)
     max_retries_per_step = Column(Integer, nullable=True)
     domain = Column(String, nullable=True, index=True)
     bw_organization_id = Column(String, nullable=True, default=None)

skyvern/forge/sdk/db/repositories/organizations.py

@@ -153,6 +153,8 @@ class OrganizationsRepository(BaseRepository):
         organization_name: str | None = None,
         webhook_callback_url: str | None = None,
         max_steps_per_run: int | None = None,
+        max_steps_per_workflow_run: int | None = None,
+        clear_max_steps_per_workflow_run: bool = False,
         max_retries_per_step: int | None = None,
         artifact_url_expiry_seconds: int | None = None,
         clear_artifact_url_expiry_seconds: bool = False,
@@ -170,6 +172,10 @@ class OrganizationsRepository(BaseRepository):
                 organization.webhook_callback_url = webhook_callback_url
             if max_steps_per_run is not None:
                 organization.max_steps_per_run = max_steps_per_run
+            if clear_max_steps_per_workflow_run:
+                organization.max_steps_per_workflow_run = None
+            elif max_steps_per_workflow_run is not None:
+                organization.max_steps_per_workflow_run = max_steps_per_workflow_run
             if max_retries_per_step is not None:
                 organization.max_retries_per_step = max_retries_per_step
             # ``clear_*`` decouples "don't update" (None) from "explicitly clear":

skyvern/forge/sdk/db/utils.py

@@ -322,6 +322,7 @@ def convert_to_organization(org_model: OrganizationModel) -> Organization:
         organization_name=org_model.organization_name,
         webhook_callback_url=org_model.webhook_callback_url,
         max_steps_per_run=org_model.max_steps_per_run,
+        max_steps_per_workflow_run=org_model.max_steps_per_workflow_run,
         max_retries_per_step=org_model.max_retries_per_step,
         domain=org_model.domain,
         bw_organization_id=org_model.bw_organization_id,

skyvern/forge/sdk/routes/agent_protocol.py

@@ -3277,9 +3277,19 @@ async def update_organization(
                 "artifact_url_expiry_seconds — pick one"
             ),
         )
+    if org_update.clear_max_steps_per_workflow_run and org_update.max_steps_per_workflow_run is not None:
+        raise HTTPException(
+            status_code=http_status.HTTP_400_BAD_REQUEST,
+            detail=(
+                "clear_max_steps_per_workflow_run cannot be combined with a non-null "
+                "max_steps_per_workflow_run — pick one"
+            ),
+        )
     updated = await app.DATABASE.organizations.update_organization(
         current_org.organization_id,
         max_steps_per_run=org_update.max_steps_per_run,
+        max_steps_per_workflow_run=org_update.max_steps_per_workflow_run,
+        clear_max_steps_per_workflow_run=org_update.clear_max_steps_per_workflow_run,
         max_retries_per_step=org_update.max_retries_per_step,
         webhook_callback_url=org_update.webhook_callback_url,
         artifact_url_expiry_seconds=org_update.artifact_url_expiry_seconds,

skyvern/forge/sdk/schemas/organizations.py

@@ -12,6 +12,7 @@ class Organization(BaseModel):
     organization_name: str
     webhook_callback_url: str | None = None
     max_steps_per_run: int | None = None
+    max_steps_per_workflow_run: int | None = None
     max_retries_per_step: int | None = None
     domain: str | None = None
     bw_organization_id: str | None = None
@@ -177,6 +178,16 @@ class GetOrganizationAPIKeysResponse(BaseModel):
 
 class OrganizationUpdate(BaseModel):
     max_steps_per_run: int | None = Field(default=None, ge=1)
+    max_steps_per_workflow_run: int | None = Field(default=None, ge=1)
+    clear_max_steps_per_workflow_run: bool = Field(
+        False,
+        description=(
+            "When true, resets ``max_steps_per_workflow_run`` to NULL — there will be no "
+            "run-level cap, only the per-block ``max_steps_per_run`` ceiling. Mutually "
+            "exclusive with a non-null value in ``max_steps_per_workflow_run`` (the clear "
+            "flag wins)."
+        ),
+    )
     # 0 is a valid "disable retries" value — see ForgeAgent.execute_step.
     max_retries_per_step: int | None = Field(default=None, ge=0)
     webhook_callback_url: str | None = None

skyvern/forge/sdk/settings_manager.py

@@ -1,14 +1,3 @@
-from skyvern.config import Settings
-from skyvern.config import settings as base_settings
+from skyvern.settings_manager import SettingsManager
 
-
-class SettingsManager:
-    __instance: Settings = base_settings
-
-    @staticmethod
-    def get_settings() -> Settings:
-        return SettingsManager.__instance
-
-    @staticmethod
-    def set_settings(settings: Settings) -> None:
-        SettingsManager.__instance = settings
+__all__ = ["SettingsManager"]

skyvern/library/ai_locator.py

@@ -1,7 +1,12 @@
 from __future__ import annotations
 
+# ruff: noqa: E402
 from typing import Any, Callable
 
+from skyvern.exceptions import require_server_extra_modules
+
+require_server_extra_modules("skyvern.library.ai_locator")
+
 from playwright.async_api import Locator, Page
 
 from skyvern.core.script_generations.skyvern_page_ai import SkyvernPageAi

skyvern/library/embedded_server_factory.py

@@ -11,7 +11,7 @@ from httpx import ASGITransport
 from skyvern.config import settings
 from skyvern.forge.api_app import create_api_app
 from skyvern.forge.sdk.api.llm.config_registry import LLMConfigRegistry
-from skyvern.forge.sdk.api.llm.models import LLMConfig, LLMRouterConfig
+from skyvern.schemas.llm import LLMConfig, LLMRouterConfig
 
 _EMBEDDED_ACTIVE: bool = False
 

skyvern/library/skyvern.py

@@ -1,26 +1,32 @@
+from __future__ import annotations
+
 import asyncio
 import os
 import pathlib
 import tempfile
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import httpx
 import structlog
 from dotenv import load_dotenv
-from playwright.async_api import Playwright, async_playwright
 
 from skyvern.client import AsyncSkyvern, BrowserSessionResponse, SkyvernEnvironment
 from skyvern.client.core import RequestOptions
 from skyvern.client.types.task_run_response import TaskRunResponse
 from skyvern.client.types.workflow_run_response import WorkflowRunResponse
-from skyvern.forge.sdk.api.llm.models import LLMConfig, LLMRouterConfig
 from skyvern.library.constants import DEFAULT_AGENT_HEARTBEAT_INTERVAL, DEFAULT_AGENT_TIMEOUT, DEFAULT_CDP_PORT
-from skyvern.library.skyvern_browser import SkyvernBrowser
-from skyvern.schemas.run_blocks import CredentialType
-from skyvern.schemas.runs import ProxyLocationInput, RunEngine, RunStatus, proxy_location_to_request
+from skyvern.schemas.credential_type import CredentialType
+from skyvern.schemas.proxy_location import ProxyLocationInput, proxy_location_to_request
+from skyvern.schemas.run_enums import RunEngine, RunStatus
 
 LOG = structlog.get_logger()
 
+if TYPE_CHECKING:
+    from playwright.async_api import Playwright
+
+    from skyvern.library.skyvern_browser import SkyvernBrowser
+    from skyvern.schemas.llm import LLMConfig, LLMRouterConfig
+
 
 def _get_browser_session_url(browser_session_id: str) -> str:
     return f"https://app.skyvern.com/browser-session/{browser_session_id}"
@@ -130,7 +136,7 @@ class Skyvern(AsyncSkyvern):
         llm_config: LLMRouterConfig | LLMConfig | None = None,
         settings: dict[str, Any] | None = None,
         use_in_memory_db: bool | None = None,
-    ) -> "Skyvern":
+    ) -> Skyvern:
         """Local/embedded mode: Run Skyvern locally in-process.
 
         Args:
@@ -171,7 +177,7 @@ class Skyvern(AsyncSkyvern):
                 Example 4 - Custom LLM with environment variables:
                     ```python
                     from skyvern import Skyvern
-                    from skyvern.forge.sdk.api.llm.models import LLMConfig
+                    from skyvern.schemas.llm import LLMConfig
 
                     # Assumes OPENAI_API_KEY is set in your environment
                     skyvern = Skyvern.local(
@@ -215,7 +221,13 @@ class Skyvern(AsyncSkyvern):
         """
         from dotenv import dotenv_values  # noqa: PLC0415
 
-        from skyvern.library.embedded_server_factory import create_embedded_server  # noqa: PLC0415
+        try:
+            from skyvern.library.embedded_server_factory import create_embedded_server  # noqa: PLC0415
+        except ImportError as exc:
+            from skyvern.exceptions import raise_server_extra_required  # noqa: PLC0415
+
+            raise_server_extra_required("Skyvern.local()", exc)
+
         from skyvern.utils.env_paths import resolve_backend_env_path  # noqa: PLC0415
 
         # Auto-detect mode when use_in_memory_db is not explicitly set.
@@ -463,6 +475,8 @@ class Skyvern(AsyncSkyvern):
             SkyvernBrowser: A browser instance with Skyvern capabilities.
         """
 
+        from skyvern.library.skyvern_browser import SkyvernBrowser  # noqa: PLC0415
+
         playwright = await self._get_playwright()
 
         if user_data_dir:
@@ -496,6 +510,8 @@ class Skyvern(AsyncSkyvern):
         Returns:
             SkyvernBrowser: A browser instance connected to the existing browser.
         """
+        from skyvern.library.skyvern_browser import SkyvernBrowser  # noqa: PLC0415
+
         playwright = await self._get_playwright()
         browser = await playwright.chromium.connect_over_cdp(cdp_url)
         browser_context = browser.contexts[0] if browser.contexts else await browser.new_context()
@@ -619,6 +635,8 @@ class Skyvern(AsyncSkyvern):
         if browser_session.browser_address is None:
             raise ValueError(f"Browser address is missing for session {browser_session.browser_session_id}")
 
+        from skyvern.library.skyvern_browser import SkyvernBrowser  # noqa: PLC0415
+
         playwright = await self._get_playwright()
         browser = await playwright.chromium.connect_over_cdp(
             browser_session.browser_address, headers={"x-api-key": self._api_key}
@@ -628,6 +646,15 @@ class Skyvern(AsyncSkyvern):
 
     async def _get_playwright(self) -> Playwright:
         if self._playwright is None:
+            from skyvern.exceptions import raise_server_extra_required  # noqa: PLC0415
+
+            # Cloud-API-only SDK users can instantiate Skyvern without server extras,
+            # so browser startup keeps its own extra hint at the Playwright boundary.
+            try:
+                from playwright.async_api import async_playwright  # noqa: PLC0415
+            except ImportError as exc:
+                raise_server_extra_required("Browser APIs", exc)
+
             self._playwright = await async_playwright().start()
         return self._playwright
 

skyvern/library/skyvern_browser.py

@@ -1,5 +1,10 @@
+# ruff: noqa: E402
 from typing import TYPE_CHECKING, Any
 
+from skyvern.exceptions import require_server_extra_modules
+
+require_server_extra_modules("skyvern.library.skyvern_browser")
+
 from playwright.async_api import BrowserContext, Page
 
 from skyvern.library.skyvern_browser_page import SkyvernBrowserPage

skyvern/library/skyvern_browser_page.py

@@ -1,5 +1,10 @@
+# ruff: noqa: E402
 from typing import TYPE_CHECKING, Any
 
+from skyvern.exceptions import require_server_extra_modules
+
+require_server_extra_modules("skyvern.library.skyvern_browser_page")
+
 from playwright.async_api import Frame, Page
 
 from skyvern.core.script_generations.skyvern_page import SkyvernPage

skyvern/library/skyvern_browser_page_agent.py

@@ -3,13 +3,19 @@ import typing
 from typing import Any, Literal, overload
 
 import structlog
-from playwright.async_api import Page
 
-from skyvern.client import GetRunResponse, SkyvernEnvironment, WorkflowRunResponse
-from skyvern.client.core import RequestOptions
-from skyvern.library.constants import DEFAULT_AGENT_HEARTBEAT_INTERVAL, DEFAULT_AGENT_TIMEOUT
-from skyvern.schemas.run_blocks import CredentialType
-from skyvern.schemas.runs import RunEngine, RunStatus, TaskRunResponse
+from skyvern.exceptions import require_server_extra_modules
+
+require_server_extra_modules("skyvern.library.skyvern_browser_page_agent")
+
+from playwright.async_api import Page  # noqa: E402
+
+from skyvern.client import GetRunResponse, SkyvernEnvironment, WorkflowRunResponse  # noqa: E402
+from skyvern.client.core import RequestOptions  # noqa: E402
+from skyvern.library.constants import DEFAULT_AGENT_HEARTBEAT_INTERVAL, DEFAULT_AGENT_TIMEOUT  # noqa: E402
+from skyvern.schemas.credential_type import CredentialType  # noqa: E402
+from skyvern.schemas.run_enums import RunEngine, RunStatus  # noqa: E402
+from skyvern.schemas.runs import TaskRunResponse  # noqa: E402
 
 if typing.TYPE_CHECKING:
     from skyvern.library.skyvern_browser import SkyvernBrowser
@@ -102,7 +108,7 @@ class SkyvernBrowserPageAgent:
     async def login(
         self,
         *,
-        credential_type: Literal[CredentialType.skyvern] = CredentialType.skyvern,
+        credential_type: Literal[CredentialType.skyvern, "skyvern"] = CredentialType.skyvern,  # type: ignore[valid-type]
         credential_id: str,
         url: str | None = None,
         prompt: str | None = None,
@@ -117,7 +123,7 @@ class SkyvernBrowserPageAgent:
     async def login(
         self,
         *,
-        credential_type: Literal[CredentialType.bitwarden],
+        credential_type: Literal[CredentialType.bitwarden, "bitwarden"],  # type: ignore[valid-type]
         bitwarden_item_id: str,
         bitwarden_collection_id: str | None = None,
         url: str | None = None,
@@ -133,7 +139,7 @@ class SkyvernBrowserPageAgent:
     async def login(
         self,
         *,
-        credential_type: Literal[CredentialType.onepassword],
+        credential_type: Literal[CredentialType.onepassword, "1password"],  # type: ignore[valid-type]
         onepassword_vault_id: str,
         onepassword_item_id: str,
         url: str | None = None,
@@ -149,7 +155,7 @@ class SkyvernBrowserPageAgent:
     async def login(
         self,
         *,
-        credential_type: Literal[CredentialType.azure_vault],
+        credential_type: Literal[CredentialType.azure_vault, "azure_vault"],  # type: ignore[valid-type]
         azure_vault_name: str,
         azure_vault_username_key: str,
         azure_vault_password_key: str,
@@ -166,7 +172,7 @@ class SkyvernBrowserPageAgent:
     async def login(
         self,
         *,
-        credential_type: CredentialType = CredentialType.skyvern,
+        credential_type: CredentialType | str = CredentialType.skyvern,
         url: str | None = None,
         credential_id: str | None = None,
         bitwarden_collection_id: str | None = None,
@@ -244,6 +250,8 @@ class SkyvernBrowserPageAgent:
             WorkflowRunResponse containing the login workflow execution results.
         """
 
+        credential_type = CredentialType(credential_type)
+
         LOG.info("Starting AI login workflow", credential_type=credential_type)
 
         workflow_run = await self._browser.skyvern.login(

skyvern/library/skyvern_browser_page_ai.py

@@ -1,6 +1,12 @@
+# ruff: noqa: E402
 from typing import TYPE_CHECKING, Any
 
 import structlog
+
+from skyvern.exceptions import require_server_extra_modules
+
+require_server_extra_modules("skyvern.library.skyvern_browser_page_ai")
+
 from playwright.async_api import Page
 
 from skyvern.client import (

skyvern/library/skyvern_locator.py

@@ -1,5 +1,10 @@
+# ruff: noqa: E402
 from typing import Any, Pattern
 
+from skyvern.exceptions import require_server_extra_modules
+
+require_server_extra_modules("skyvern.library.skyvern_locator")
+
 from playwright.async_api import Locator
 
 

skyvern/schemas/credential_type.py

@@ -0,0 +1,8 @@
+from enum import StrEnum
+
+
+class CredentialType(StrEnum):
+    skyvern = "skyvern"
+    bitwarden = "bitwarden"
+    onepassword = "1password"
+    azure_vault = "azure_vault"

skyvern/schemas/llm.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal, TypedDict, cast
+
+__all__ = [
+    "LiteLLMParams",
+    "LLMAllowedFailsPolicy",
+    "LLMConfig",
+    "LLMConfigBase",
+    "LLMRouterConfig",
+    "LLMRouterModelConfig",
+]
+
+_SETTINGS_DEFAULT = object()
+# Sentinel replaced in __post_init__; distinct from callers explicitly passing None.
+_DEFAULT_MAX_TOKENS = cast("int | None", _SETTINGS_DEFAULT)
+_DEFAULT_TEMPERATURE = cast("float | None", _SETTINGS_DEFAULT)
+
+
+def _assert_settings_defaults_resolved(*values: object) -> None:
+    if any(value is _SETTINGS_DEFAULT for value in values):
+        raise RuntimeError("settings default sentinel was not resolved")
+
+
+def _resolve_generation_defaults(config: object, max_tokens: object, temperature: object) -> None:
+    if max_tokens is _SETTINGS_DEFAULT or temperature is _SETTINGS_DEFAULT:
+        settings = _settings()
+        if max_tokens is _SETTINGS_DEFAULT:
+            object.__setattr__(config, "max_tokens", settings.LLM_CONFIG_MAX_TOKENS)
+        if temperature is _SETTINGS_DEFAULT:
+            object.__setattr__(config, "temperature", settings.LLM_CONFIG_TEMPERATURE)
+
+    _assert_settings_defaults_resolved(
+        getattr(config, "max_tokens"),
+        getattr(config, "temperature"),
+    )
+
+
+class LiteLLMParams(TypedDict, total=False):
+    api_key: str | None
+    api_version: str | None
+    api_base: str | None
+    model_info: dict[str, Any] | None
+    vertex_credentials: str | None
+    vertex_location: str | None
+    thinking: dict[str, Any] | None
+    thinking_level: str | None
+    service_tier: str | None
+    extra_headers: dict[str, str] | None
+    timeout: float | None
+
+
+def _settings() -> Any:
+    # Keep settings resolution lazy so importing skyvern.schemas.llm only defines the
+    # public dataclasses; config defaults are read when a config is constructed.
+    from skyvern.settings_manager import SettingsManager  # noqa: PLC0415
+
+    return SettingsManager.get_settings()
+
+
+@dataclass(frozen=True)
+class LLMConfigBase:
+    model_name: str
+    required_env_vars: list[str]
+    supports_vision: bool
+    add_assistant_prefix: bool
+
+    def get_missing_env_vars(self) -> list[str]:
+        settings = _settings()
+        missing_env_vars = []
+        for env_var in self.required_env_vars:
+            env_var_value = getattr(settings, env_var, None)
+            if not env_var_value:
+                missing_env_vars.append(env_var)
+
+        return missing_env_vars
+
+
+@dataclass(frozen=True)
+class LLMConfig(LLMConfigBase):
+    """Base-safe LLM config shared by SDK and server import paths.
+
+    Default max-token and temperature fields use a private sentinel so callers can
+    still pass None explicitly. __post_init__ must replace the sentinel before the
+    frozen dataclass instance is observable.
+    """
+
+    litellm_params: LiteLLMParams | None = field(default=None)
+    max_tokens: int | None = _DEFAULT_MAX_TOKENS
+    max_completion_tokens: int | None = None
+    temperature: float | None = _DEFAULT_TEMPERATURE
+    reasoning_effort: str | None = None
+
+    def __post_init__(self) -> None:
+        _resolve_generation_defaults(self, self.max_tokens, self.temperature)
+
+
+@dataclass(frozen=True)
+class LLMAllowedFailsPolicy:
+    bad_request_error_allowed_fails: int | None = None
+    authentication_error_allowed_fails: int | None = None
+    timeout_error_allowed_fails: int | None = None
+    rate_limit_error_allowed_fails: int | None = None
+    content_policy_violation_error_allowed_fails: int | None = None
+    internal_server_error_allowed_fails: int | None = None
+
+
+@dataclass(frozen=True)
+class LLMRouterModelConfig:
+    model_name: str
+    # https://litellm.vercel.app/docs/routing
+    litellm_params: dict[str, Any]
+    model_info: dict[str, Any] = field(default_factory=dict)
+    tpm: int | None = None
+    rpm: int | None = None
+
+
+@dataclass(frozen=True)
+class LLMRouterConfig(LLMConfigBase):
+    """Base-safe router config with the same settings-default sentinel invariant as LLMConfig."""
+
+    model_list: list[LLMRouterModelConfig]
+    # All three redis parameters are required. Even if there isn't a password, it should be an empty string.
+    main_model_group: str
+    redis_host: str | None = None
+    redis_port: int | None = None
+    redis_password: str | None = None
+    fallback_model_group: str | None = None
+    routing_strategy: Literal[
+        "simple-shuffle",
+        "least-busy",
+        "usage-based-routing",
+        "usage-based-routing-v2",
+        "latency-based-routing",
+    ] = "usage-based-routing"
+    num_retries: int = 1
+    retry_delay_seconds: int = 15
+    set_verbose: bool = False
+    disable_cooldowns: bool | None = None
+    allowed_fails: int | None = None
+    allowed_fails_policy: LLMAllowedFailsPolicy | None = None
+    cooldown_time: float | None = None
+    max_tokens: int | None = _DEFAULT_MAX_TOKENS
+    max_completion_tokens: int | None = None
+    reasoning_effort: str | None = None
+    temperature: float | None = _DEFAULT_TEMPERATURE
+
+    def __post_init__(self) -> None:
+        _resolve_generation_defaults(self, self.max_tokens, self.temperature)

skyvern/schemas/proxy_location.py

@@ -0,0 +1,304 @@
+from __future__ import annotations
+
+from enum import StrEnum
+from typing import Any
+from zoneinfo import ZoneInfo
+
+from pydantic import BaseModel, Field, field_validator
+
+
+class ProxyLocation(StrEnum):
+    RESIDENTIAL = "RESIDENTIAL"
+    US_CA = "US-CA"
+    US_NY = "US-NY"
+    US_TX = "US-TX"
+    US_FL = "US-FL"
+    US_WA = "US-WA"
+    RESIDENTIAL_ES = "RESIDENTIAL_ES"
+    RESIDENTIAL_IE = "RESIDENTIAL_IE"
+    RESIDENTIAL_GB = "RESIDENTIAL_GB"
+    RESIDENTIAL_IN = "RESIDENTIAL_IN"
+    RESIDENTIAL_JP = "RESIDENTIAL_JP"
+    RESIDENTIAL_FR = "RESIDENTIAL_FR"
+    RESIDENTIAL_DE = "RESIDENTIAL_DE"
+    RESIDENTIAL_NZ = "RESIDENTIAL_NZ"
+    RESIDENTIAL_ZA = "RESIDENTIAL_ZA"
+    RESIDENTIAL_AR = "RESIDENTIAL_AR"
+    RESIDENTIAL_AU = "RESIDENTIAL_AU"
+    RESIDENTIAL_BR = "RESIDENTIAL_BR"
+    RESIDENTIAL_TR = "RESIDENTIAL_TR"
+    RESIDENTIAL_CA = "RESIDENTIAL_CA"
+    RESIDENTIAL_MX = "RESIDENTIAL_MX"
+    RESIDENTIAL_IT = "RESIDENTIAL_IT"
+    RESIDENTIAL_NL = "RESIDENTIAL_NL"
+    RESIDENTIAL_PH = "RESIDENTIAL_PH"
+    RESIDENTIAL_KR = "RESIDENTIAL_KR"
+    RESIDENTIAL_SA = "RESIDENTIAL_SA"
+    RESIDENTIAL_ISP = "RESIDENTIAL_ISP"
+    NONE = "NONE"
+
+    @staticmethod
+    def get_zone(proxy_location: ProxyLocation) -> str:
+        zone_mapping = {
+            ProxyLocation.US_CA: "california",
+            ProxyLocation.US_NY: "newyork",
+            ProxyLocation.US_TX: "texas",
+            ProxyLocation.US_FL: "florida",
+            ProxyLocation.US_WA: "washington",
+            ProxyLocation.RESIDENTIAL: "residential_long-country-us",
+        }
+        if proxy_location in zone_mapping:
+            return zone_mapping[proxy_location]
+        raise ValueError(f"No zone mapping for proxy location: {proxy_location}")
+
+    @classmethod
+    def residential_country_locations(cls) -> set[ProxyLocation]:
+        return {
+            cls.RESIDENTIAL,
+            cls.RESIDENTIAL_ES,
+            cls.RESIDENTIAL_IE,
+            cls.RESIDENTIAL_GB,
+            cls.RESIDENTIAL_IN,
+            cls.RESIDENTIAL_JP,
+            cls.RESIDENTIAL_FR,
+            cls.RESIDENTIAL_DE,
+            cls.RESIDENTIAL_NZ,
+            cls.RESIDENTIAL_ZA,
+            cls.RESIDENTIAL_AR,
+            cls.RESIDENTIAL_AU,
+            cls.RESIDENTIAL_BR,
+            cls.RESIDENTIAL_TR,
+            cls.RESIDENTIAL_CA,
+            cls.RESIDENTIAL_MX,
+            cls.RESIDENTIAL_IT,
+            cls.RESIDENTIAL_NL,
+            cls.RESIDENTIAL_PH,
+            cls.RESIDENTIAL_KR,
+            cls.RESIDENTIAL_SA,
+        }
+
+    @staticmethod
+    def get_proxy_count(proxy_location: ProxyLocation) -> int:
+        counts = {
+            ProxyLocation.RESIDENTIAL: 10000,
+            ProxyLocation.RESIDENTIAL_ES: 2000,
+            ProxyLocation.RESIDENTIAL_IE: 2000,
+            ProxyLocation.RESIDENTIAL_GB: 2000,
+            ProxyLocation.RESIDENTIAL_IN: 2000,
+            ProxyLocation.RESIDENTIAL_JP: 2000,
+            ProxyLocation.RESIDENTIAL_FR: 2000,
+            ProxyLocation.RESIDENTIAL_DE: 2000,
+            ProxyLocation.RESIDENTIAL_NZ: 2000,
+            ProxyLocation.RESIDENTIAL_ZA: 2000,
+            ProxyLocation.RESIDENTIAL_AR: 2000,
+            ProxyLocation.RESIDENTIAL_AU: 2000,
+            ProxyLocation.RESIDENTIAL_BR: 2000,
+            ProxyLocation.RESIDENTIAL_TR: 2000,
+            ProxyLocation.RESIDENTIAL_CA: 2000,
+            ProxyLocation.RESIDENTIAL_MX: 2000,
+            ProxyLocation.RESIDENTIAL_IT: 2000,
+            ProxyLocation.RESIDENTIAL_NL: 2000,
+            ProxyLocation.RESIDENTIAL_PH: 2000,
+            ProxyLocation.RESIDENTIAL_KR: 2000,
+            ProxyLocation.RESIDENTIAL_SA: 2000,
+        }
+        return counts.get(proxy_location, 10000)
+
+    @staticmethod
+    def get_country_code(proxy_location: ProxyLocation) -> str:
+        mapping = {
+            ProxyLocation.RESIDENTIAL: "US",
+            ProxyLocation.RESIDENTIAL_ES: "ES",
+            ProxyLocation.RESIDENTIAL_IE: "IE",
+            ProxyLocation.RESIDENTIAL_GB: "GB",
+            ProxyLocation.RESIDENTIAL_IN: "IN",
+            ProxyLocation.RESIDENTIAL_JP: "JP",
+            ProxyLocation.RESIDENTIAL_FR: "FR",
+            ProxyLocation.RESIDENTIAL_DE: "DE",
+            ProxyLocation.RESIDENTIAL_NZ: "NZ",
+            ProxyLocation.RESIDENTIAL_ZA: "ZA",
+            ProxyLocation.RESIDENTIAL_AR: "AR",
+            ProxyLocation.RESIDENTIAL_AU: "AU",
+            ProxyLocation.RESIDENTIAL_BR: "BR",
+            ProxyLocation.RESIDENTIAL_TR: "TR",
+            ProxyLocation.RESIDENTIAL_CA: "CA",
+            ProxyLocation.RESIDENTIAL_MX: "MX",
+            ProxyLocation.RESIDENTIAL_IT: "IT",
+            ProxyLocation.RESIDENTIAL_NL: "NL",
+            ProxyLocation.RESIDENTIAL_PH: "PH",
+            ProxyLocation.RESIDENTIAL_KR: "KR",
+            ProxyLocation.RESIDENTIAL_SA: "SA",
+        }
+        return mapping.get(proxy_location, "US")
+
+
+# Supported countries for granular geo-targeting.
+SUPPORTED_GEO_COUNTRIES = frozenset(
+    {
+        "US",
+        "AR",
+        "AU",
+        "BR",
+        "CA",
+        "DE",
+        "ES",
+        "FR",
+        "GB",
+        "IE",
+        "IN",
+        "IT",
+        "JP",
+        "MX",
+        "NL",
+        "NZ",
+        "PH",
+        "KR",
+        "SA",
+        "TR",
+        "ZA",
+    }
+)
+
+
+class GeoTarget(BaseModel):
+    """Granular proxy geo-targeting request with country, optional subdivision, and optional city."""
+
+    country: str = Field(
+        description="ISO 3166-1 alpha-2 country code (e.g., 'US', 'GB', 'DE')",
+        examples=["US", "GB", "DE", "FR"],
+        min_length=2,
+        max_length=2,
+    )
+    subdivision: str | None = Field(
+        default=None,
+        description="ISO 3166-2 subdivision code without country prefix (e.g., 'CA' for California, 'NY' for New York)",
+        examples=["CA", "NY", "TX", "ENG"],
+        max_length=10,
+    )
+    city: str | None = Field(
+        default=None,
+        description="City name in English from GeoNames (e.g., 'New York', 'Los Angeles', 'London')",
+        examples=["New York", "Los Angeles", "London", "Berlin"],
+        max_length=100,
+    )
+
+    @field_validator("country")
+    @classmethod
+    def validate_country(cls, v: str) -> str:
+        """Validate country is in supported list and normalize to uppercase."""
+        v = v.upper()
+        if v not in SUPPORTED_GEO_COUNTRIES:
+            raise ValueError(
+                f"Country '{v}' is not supported for geo targeting. "
+                f"Supported countries: {sorted(SUPPORTED_GEO_COUNTRIES)}"
+            )
+        return v
+
+    @field_validator("subdivision")
+    @classmethod
+    def validate_subdivision(cls, v: str | None) -> str | None:
+        """Normalize subdivision code to uppercase and strip country prefix if present."""
+        if v is None:
+            return v
+        v = v.upper()
+        # Strip country prefix if accidentally included (e.g., "US-CA" -> "CA")
+        if "-" in v:
+            v = v.split("-", 1)[1]
+        return v
+
+
+ProxyLocationInput = ProxyLocation | GeoTarget | dict[str, Any] | None
+
+
+def proxy_location_to_request(proxy_location: ProxyLocationInput) -> ProxyLocation | dict[str, Any] | None:
+    if isinstance(proxy_location, GeoTarget):
+        return proxy_location.model_dump()
+    return proxy_location
+
+
+def get_tzinfo_from_proxy(proxy_location: ProxyLocation) -> ZoneInfo | None:
+    if proxy_location == ProxyLocation.NONE:
+        return None
+
+    if proxy_location == ProxyLocation.US_CA:
+        return ZoneInfo("America/Los_Angeles")
+
+    if proxy_location == ProxyLocation.US_NY:
+        return ZoneInfo("America/New_York")
+
+    if proxy_location == ProxyLocation.US_TX:
+        return ZoneInfo("America/Chicago")
+
+    if proxy_location == ProxyLocation.US_FL:
+        return ZoneInfo("America/New_York")
+
+    if proxy_location == ProxyLocation.US_WA:
+        return ZoneInfo("America/Los_Angeles")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL:
+        return ZoneInfo("America/New_York")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_ES:
+        return ZoneInfo("Europe/Madrid")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_IE:
+        return ZoneInfo("Europe/Dublin")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_GB:
+        return ZoneInfo("Europe/London")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_IN:
+        return ZoneInfo("Asia/Kolkata")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_JP:
+        return ZoneInfo("Asia/Tokyo")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_FR:
+        return ZoneInfo("Europe/Paris")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_DE:
+        return ZoneInfo("Europe/Berlin")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_NZ:
+        return ZoneInfo("Pacific/Auckland")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_ZA:
+        return ZoneInfo("Africa/Johannesburg")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_AR:
+        return ZoneInfo("America/Argentina/Buenos_Aires")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_AU:
+        return ZoneInfo("Australia/Sydney")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_BR:
+        return ZoneInfo("America/Sao_Paulo")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_TR:
+        return ZoneInfo("Europe/Istanbul")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_CA:
+        return ZoneInfo("America/Toronto")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_MX:
+        return ZoneInfo("America/Mexico_City")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_IT:
+        return ZoneInfo("Europe/Rome")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_NL:
+        return ZoneInfo("Europe/Amsterdam")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_PH:
+        return ZoneInfo("Asia/Manila")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_KR:
+        return ZoneInfo("Asia/Seoul")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_SA:
+        return ZoneInfo("Asia/Riyadh")
+
+    if proxy_location == ProxyLocation.RESIDENTIAL_ISP:
+        return ZoneInfo("America/New_York")
+
+    return None

skyvern/schemas/run_blocks.py

@@ -1,15 +1,7 @@
-from enum import StrEnum
-
 from pydantic import BaseModel, Field
 
-from skyvern.schemas.runs import ProxyLocation
-
-
-class CredentialType(StrEnum):
-    skyvern = "skyvern"
-    bitwarden = "bitwarden"
-    onepassword = "1password"
-    azure_vault = "azure_vault"
+from skyvern.schemas.credential_type import CredentialType
+from skyvern.schemas.proxy_location import ProxyLocation
 
 
 class BaseRunBlockRequest(BaseModel):

skyvern/schemas/run_enums.py

@@ -0,0 +1,40 @@
+from enum import StrEnum
+
+
+class RunType(StrEnum):
+    task_v1 = "task_v1"
+    task_v2 = "task_v2"
+    workflow_run = "workflow_run"
+    openai_cua = "openai_cua"
+    anthropic_cua = "anthropic_cua"
+    ui_tars = "ui_tars"
+
+
+class RunEngine(StrEnum):
+    skyvern_v1 = "skyvern-1.0"
+    skyvern_v2 = "skyvern-2.0"
+    openai_cua = "openai-cua"
+    anthropic_cua = "anthropic-cua"
+    ui_tars = "ui-tars"
+
+
+CUA_ENGINES = (RunEngine.openai_cua, RunEngine.anthropic_cua, RunEngine.ui_tars)
+CUA_RUN_TYPES = (RunType.openai_cua, RunType.anthropic_cua, RunType.ui_tars)
+
+# Statuses that are final; once a row reaches one of these, it never changes.
+# Single source of truth used by sync cron, partial indexes, and run response helpers.
+TERMINAL_STATUSES = ("completed", "failed", "terminated", "canceled", "timed_out")
+
+
+class RunStatus(StrEnum):
+    created = "created"
+    queued = "queued"
+    running = "running"
+    timed_out = "timed_out"
+    failed = "failed"
+    terminated = "terminated"
+    completed = "completed"
+    canceled = "canceled"
+
+    def is_final(self) -> bool:
+        return self.value in TERMINAL_STATUSES

skyvern/schemas/runs.py

@@ -1,9 +1,7 @@
 from __future__ import annotations
 
 from datetime import datetime
-from enum import StrEnum
-from typing import Annotated, Any, Literal, Union
-from zoneinfo import ZoneInfo
+from typing import TYPE_CHECKING, Annotated, Any, Literal, TypeAlias, Union
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
@@ -31,358 +29,38 @@ from skyvern.schemas.docs.doc_strings import (
     TOTP_URL_DOC_STRING,
     WEBHOOK_URL_DOC_STRING,
 )
+from skyvern.schemas.proxy_location import (  # noqa: F401
+    SUPPORTED_GEO_COUNTRIES,
+    GeoTarget,
+    ProxyLocation,
+    ProxyLocationInput,
+    get_tzinfo_from_proxy,
+    proxy_location_to_request,
+)
+from skyvern.schemas.run_enums import (  # noqa: F401
+    CUA_ENGINES,
+    CUA_RUN_TYPES,
+    TERMINAL_STATUSES,
+    RunEngine,
+    RunStatus,
+    RunType,
+)
 from skyvern.utils.url_validators import validate_url
 
-
-class ProxyLocation(StrEnum):
-    RESIDENTIAL = "RESIDENTIAL"
-    US_CA = "US-CA"
-    US_NY = "US-NY"
-    US_TX = "US-TX"
-    US_FL = "US-FL"
-    US_WA = "US-WA"
-    RESIDENTIAL_ES = "RESIDENTIAL_ES"
-    RESIDENTIAL_IE = "RESIDENTIAL_IE"
-    RESIDENTIAL_GB = "RESIDENTIAL_GB"
-    RESIDENTIAL_IN = "RESIDENTIAL_IN"
-    RESIDENTIAL_JP = "RESIDENTIAL_JP"
-    RESIDENTIAL_FR = "RESIDENTIAL_FR"
-    RESIDENTIAL_DE = "RESIDENTIAL_DE"
-    RESIDENTIAL_NZ = "RESIDENTIAL_NZ"
-    RESIDENTIAL_ZA = "RESIDENTIAL_ZA"
-    RESIDENTIAL_AR = "RESIDENTIAL_AR"
-    RESIDENTIAL_AU = "RESIDENTIAL_AU"
-    RESIDENTIAL_BR = "RESIDENTIAL_BR"
-    RESIDENTIAL_TR = "RESIDENTIAL_TR"
-    RESIDENTIAL_CA = "RESIDENTIAL_CA"
-    RESIDENTIAL_MX = "RESIDENTIAL_MX"
-    RESIDENTIAL_IT = "RESIDENTIAL_IT"
-    RESIDENTIAL_NL = "RESIDENTIAL_NL"
-    RESIDENTIAL_PH = "RESIDENTIAL_PH"
-    RESIDENTIAL_KR = "RESIDENTIAL_KR"
-    RESIDENTIAL_SA = "RESIDENTIAL_SA"
-    RESIDENTIAL_ISP = "RESIDENTIAL_ISP"
-    NONE = "NONE"
-
-    @staticmethod
-    def get_zone(proxy_location: ProxyLocation) -> str:
-        zone_mapping = {
-            ProxyLocation.US_CA: "california",
-            ProxyLocation.US_NY: "newyork",
-            ProxyLocation.US_TX: "texas",
-            ProxyLocation.US_FL: "florida",
-            ProxyLocation.US_WA: "washington",
-            ProxyLocation.RESIDENTIAL: "residential_long-country-us",
-        }
-        if proxy_location in zone_mapping:
-            return zone_mapping[proxy_location]
-        raise ValueError(f"No zone mapping for proxy location: {proxy_location}")
-
-    @classmethod
-    def residential_country_locations(cls) -> set[ProxyLocation]:
-        return {
-            cls.RESIDENTIAL,
-            cls.RESIDENTIAL_ES,
-            cls.RESIDENTIAL_IE,
-            cls.RESIDENTIAL_GB,
-            cls.RESIDENTIAL_IN,
-            cls.RESIDENTIAL_JP,
-            cls.RESIDENTIAL_FR,
-            cls.RESIDENTIAL_DE,
-            cls.RESIDENTIAL_NZ,
-            cls.RESIDENTIAL_ZA,
-            cls.RESIDENTIAL_AR,
-            cls.RESIDENTIAL_AU,
-            cls.RESIDENTIAL_BR,
-            cls.RESIDENTIAL_TR,
-            cls.RESIDENTIAL_CA,
-            cls.RESIDENTIAL_MX,
-            cls.RESIDENTIAL_IT,
-            cls.RESIDENTIAL_NL,
-            cls.RESIDENTIAL_PH,
-            cls.RESIDENTIAL_KR,
-            cls.RESIDENTIAL_SA,
-        }
-
-    @staticmethod
-    def get_proxy_count(proxy_location: ProxyLocation) -> int:
-        counts = {
-            ProxyLocation.RESIDENTIAL: 10000,
-            ProxyLocation.RESIDENTIAL_ES: 2000,
-            ProxyLocation.RESIDENTIAL_IE: 2000,
-            ProxyLocation.RESIDENTIAL_GB: 2000,
-            ProxyLocation.RESIDENTIAL_IN: 2000,
-            ProxyLocation.RESIDENTIAL_JP: 2000,
-            ProxyLocation.RESIDENTIAL_FR: 2000,
-            ProxyLocation.RESIDENTIAL_DE: 2000,
-            ProxyLocation.RESIDENTIAL_NZ: 2000,
-            ProxyLocation.RESIDENTIAL_ZA: 2000,
-            ProxyLocation.RESIDENTIAL_AR: 2000,
-            ProxyLocation.RESIDENTIAL_AU: 2000,
-            ProxyLocation.RESIDENTIAL_BR: 2000,
-            ProxyLocation.RESIDENTIAL_TR: 2000,
-            ProxyLocation.RESIDENTIAL_CA: 2000,
-            ProxyLocation.RESIDENTIAL_MX: 2000,
-            ProxyLocation.RESIDENTIAL_IT: 2000,
-            ProxyLocation.RESIDENTIAL_NL: 2000,
-            ProxyLocation.RESIDENTIAL_PH: 2000,
-            ProxyLocation.RESIDENTIAL_KR: 2000,
-            ProxyLocation.RESIDENTIAL_SA: 2000,
-        }
-        return counts.get(proxy_location, 10000)
-
-    @staticmethod
-    def get_country_code(proxy_location: ProxyLocation) -> str:
-        mapping = {
-            ProxyLocation.RESIDENTIAL: "US",
-            ProxyLocation.RESIDENTIAL_ES: "ES",
-            ProxyLocation.RESIDENTIAL_IE: "IE",
-            ProxyLocation.RESIDENTIAL_GB: "GB",
-            ProxyLocation.RESIDENTIAL_IN: "IN",
-            ProxyLocation.RESIDENTIAL_JP: "JP",
-            ProxyLocation.RESIDENTIAL_FR: "FR",
-            ProxyLocation.RESIDENTIAL_DE: "DE",
-            ProxyLocation.RESIDENTIAL_NZ: "NZ",
-            ProxyLocation.RESIDENTIAL_ZA: "ZA",
-            ProxyLocation.RESIDENTIAL_AR: "AR",
-            ProxyLocation.RESIDENTIAL_AU: "AU",
-            ProxyLocation.RESIDENTIAL_BR: "BR",
-            ProxyLocation.RESIDENTIAL_TR: "TR",
-            ProxyLocation.RESIDENTIAL_CA: "CA",
-            ProxyLocation.RESIDENTIAL_MX: "MX",
-            ProxyLocation.RESIDENTIAL_IT: "IT",
-            ProxyLocation.RESIDENTIAL_NL: "NL",
-            ProxyLocation.RESIDENTIAL_PH: "PH",
-            ProxyLocation.RESIDENTIAL_KR: "KR",
-            ProxyLocation.RESIDENTIAL_SA: "SA",
-        }
-        return mapping.get(proxy_location, "US")
-
-
-# Supported countries for GeoTarget - must match Massive's coverage
-SUPPORTED_GEO_COUNTRIES = frozenset(
-    {
-        "US",
-        "AR",
-        "AU",
-        "BR",
-        "CA",
-        "DE",
-        "ES",
-        "FR",
-        "GB",
-        "IE",
-        "IN",
-        "IT",
-        "JP",
-        "MX",
-        "NL",
-        "NZ",
-        "PH",
-        "KR",
-        "SA",
-        "TR",
-        "ZA",
-    }
-)
-
-
-class GeoTarget(BaseModel):
-    """
-    Granular geographic targeting for proxy selection.
-
-    Supports country, subdivision (state/region), and city level targeting.
-    Uses ISO 3166-1 alpha-2 for countries, ISO 3166-2 for subdivisions,
-    and GeoNames English names for cities.
-
-    Examples:
-        - {"country": "US"} - United States (same as RESIDENTIAL)
-        - {"country": "US", "subdivision": "CA"} - California, US
-        - {"country": "US", "subdivision": "NY", "city": "New York"} - New York City
-        - {"country": "GB", "city": "London"} - London, UK
-    """
-
-    country: str = Field(
-        description="ISO 3166-1 alpha-2 country code (e.g., 'US', 'GB', 'DE')",
-        examples=["US", "GB", "DE", "FR"],
-        min_length=2,
-        max_length=2,
-    )
-    subdivision: str | None = Field(
-        default=None,
-        description="ISO 3166-2 subdivision code without country prefix (e.g., 'CA' for California, 'NY' for New York)",
-        examples=["CA", "NY", "TX", "ENG"],
-        max_length=10,
-    )
-    city: str | None = Field(
-        default=None,
-        description="City name in English from GeoNames (e.g., 'New York', 'Los Angeles', 'London')",
-        examples=["New York", "Los Angeles", "London", "Berlin"],
-        max_length=100,
-    )
-
-    @field_validator("country")
-    @classmethod
-    def validate_country(cls, v: str) -> str:
-        """Validate country is in supported list and normalize to uppercase."""
-        v = v.upper()
-        if v not in SUPPORTED_GEO_COUNTRIES:
-            raise ValueError(
-                f"Country '{v}' is not supported for geo targeting. "
-                f"Supported countries: {sorted(SUPPORTED_GEO_COUNTRIES)}"
-            )
-        return v
-
-    @field_validator("subdivision")
-    @classmethod
-    def validate_subdivision(cls, v: str | None) -> str | None:
-        """Normalize subdivision code to uppercase and strip country prefix if present."""
-        if v is None:
-            return v
-        v = v.upper()
-        # Strip country prefix if accidentally included (e.g., "US-CA" -> "CA")
-        if "-" in v:
-            v = v.split("-", 1)[1]
-        return v
-
-
-# Type alias for proxy location that accepts either legacy enum or new GeoTarget
-ProxyLocationInput = ProxyLocation | GeoTarget | dict | None
-
-
-def proxy_location_to_request(proxy_location: ProxyLocationInput) -> ProxyLocation | dict | None:
-    if isinstance(proxy_location, GeoTarget):
-        return proxy_location.model_dump()
-    return proxy_location
-
-
-def get_tzinfo_from_proxy(proxy_location: ProxyLocation) -> ZoneInfo | None:
-    if proxy_location == ProxyLocation.NONE:
-        return None
-
-    if proxy_location == ProxyLocation.US_CA:
-        return ZoneInfo("America/Los_Angeles")
-
-    if proxy_location == ProxyLocation.US_NY:
-        return ZoneInfo("America/New_York")
-
-    if proxy_location == ProxyLocation.US_TX:
-        return ZoneInfo("America/Chicago")
-
-    if proxy_location == ProxyLocation.US_FL:
-        return ZoneInfo("America/New_York")
-
-    if proxy_location == ProxyLocation.US_WA:
-        return ZoneInfo("America/Los_Angeles")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL:
-        return ZoneInfo("America/New_York")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_ES:
-        return ZoneInfo("Europe/Madrid")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_IE:
-        return ZoneInfo("Europe/Dublin")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_GB:
-        return ZoneInfo("Europe/London")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_IN:
-        return ZoneInfo("Asia/Kolkata")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_JP:
-        return ZoneInfo("Asia/Tokyo")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_FR:
-        return ZoneInfo("Europe/Paris")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_DE:
-        return ZoneInfo("Europe/Berlin")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_NZ:
-        return ZoneInfo("Pacific/Auckland")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_ZA:
-        return ZoneInfo("Africa/Johannesburg")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_AR:
-        return ZoneInfo("America/Argentina/Buenos_Aires")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_AU:
-        return ZoneInfo("Australia/Sydney")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_BR:
-        return ZoneInfo("America/Sao_Paulo")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_TR:
-        return ZoneInfo("Europe/Istanbul")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_CA:
-        return ZoneInfo("America/Toronto")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_MX:
-        return ZoneInfo("America/Mexico_City")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_IT:
-        return ZoneInfo("Europe/Rome")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_NL:
-        return ZoneInfo("Europe/Amsterdam")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_PH:
-        return ZoneInfo("Asia/Manila")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_KR:
-        return ZoneInfo("Asia/Seoul")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_SA:
-        return ZoneInfo("Asia/Riyadh")
-
-    if proxy_location == ProxyLocation.RESIDENTIAL_ISP:
-        return ZoneInfo("America/New_York")
-
-    return None
-
-
-class RunType(StrEnum):
-    task_v1 = "task_v1"
-    task_v2 = "task_v2"
-    workflow_run = "workflow_run"
-    openai_cua = "openai_cua"
-    anthropic_cua = "anthropic_cua"
-    ui_tars = "ui_tars"
-
-
-class RunEngine(StrEnum):
-    skyvern_v1 = "skyvern-1.0"
-    skyvern_v2 = "skyvern-2.0"
-    openai_cua = "openai-cua"
-    anthropic_cua = "anthropic-cua"
-    ui_tars = "ui-tars"
-
-
-CUA_ENGINES = [RunEngine.openai_cua, RunEngine.anthropic_cua, RunEngine.ui_tars]
-CUA_RUN_TYPES = [RunType.openai_cua, RunType.anthropic_cua, RunType.ui_tars]
-
-
-class RunStatus(StrEnum):
-    created = "created"
-    queued = "queued"
-    running = "running"
-    timed_out = "timed_out"
-    failed = "failed"
-    terminated = "terminated"
-    completed = "completed"
-    canceled = "canceled"
-
-    def is_final(self) -> bool:
-        return self.value in TERMINAL_STATUSES
-
-
-# Statuses that are final — once a row reaches one of these, it never changes.
-# Single source of truth: used by the sync cron, the partial index, and any
-# code that needs to know whether a run is "done".
-TERMINAL_STATUSES = ("completed", "failed", "terminated", "canceled", "timed_out")
+# Type checkers need string Literal values, while pydantic's discriminated
+# union preserves enum instances when runtime Literals use the enum members.
+if TYPE_CHECKING:
+    TaskRunTypeField: TypeAlias = Literal["task_v1", "task_v2", "openai_cua", "anthropic_cua", "ui_tars"]
+    WorkflowRunTypeField: TypeAlias = Literal["workflow_run"]
+else:
+    TaskRunTypeField = Literal[
+        RunType.task_v1,
+        RunType.task_v2,
+        RunType.openai_cua,
+        RunType.anthropic_cua,
+        RunType.ui_tars,
+    ]
+    WorkflowRunTypeField = Literal[RunType.workflow_run]
 
 
 class TaskRunRequest(BaseModel):
@@ -402,7 +80,7 @@ class TaskRunRequest(BaseModel):
     title: str | None = Field(
         default=None, description="The title for the task", examples=["The title of my first skyvern task"]
     )
-    proxy_location: ProxyLocation | GeoTarget | dict | None = Field(
+    proxy_location: ProxyLocationInput = Field(
         default=ProxyLocation.RESIDENTIAL,
         description=PROXY_LOCATION_DOC_STRING + " Can also be a GeoTarget object for granular city/state targeting: "
         '{"country": "US", "subdivision": "CA", "city": "San Francisco"}',
@@ -509,7 +187,7 @@ class WorkflowRunRequest(BaseModel):
     )
     parameters: dict[str, Any] | None = Field(default=None, description="Parameters to pass to the workflow")
     title: str | None = Field(default=None, description="The title for this workflow run")
-    proxy_location: ProxyLocation | GeoTarget | dict | None = Field(
+    proxy_location: ProxyLocationInput = Field(
         default=ProxyLocation.RESIDENTIAL,
         description=PROXY_LOCATION_DOC_STRING + " Can also be a GeoTarget object for granular city/state targeting: "
         '{"country": "US", "subdivision": "CA", "city": "San Francisco"}',
@@ -689,8 +367,8 @@ class BaseRunResponse(BaseModel):
 
 
 class TaskRunResponse(BaseRunResponse):
-    run_type: Literal[RunType.task_v1, RunType.task_v2, RunType.openai_cua, RunType.anthropic_cua, RunType.ui_tars] = (
-        Field(description="Types of a task run - task_v1, task_v2, openai_cua, anthropic_cua, ui_tars")
+    run_type: TaskRunTypeField = Field(
+        description="Types of a task run - task_v1, task_v2, openai_cua, anthropic_cua, ui_tars"
     )
     run_request: TaskRunRequest | None = Field(
         default=None, description="The original request parameters used to start this task run"
@@ -698,7 +376,7 @@ class TaskRunResponse(BaseRunResponse):
 
 
 class WorkflowRunResponse(BaseRunResponse):
-    run_type: Literal[RunType.workflow_run] = Field(description="Type of run - always workflow_run for workflow runs")
+    run_type: WorkflowRunTypeField = Field(description="Type of run - always workflow_run for workflow runs")
     run_with: str = Field(
         default="agent",
         description="Whether the workflow run was executed with agent or code",

skyvern/settings_manager.py

@@ -0,0 +1,14 @@
+from skyvern.config import Settings
+from skyvern.config import settings as base_settings
+
+
+class SettingsManager:
+    __instance: Settings = base_settings
+
+    @staticmethod
+    def get_settings() -> Settings:
+        return SettingsManager.__instance
+
+    @staticmethod
+    def set_settings(settings: Settings) -> None:
+        SettingsManager.__instance = settings

skyvern/utils/url_validators.py

@@ -1,7 +1,7 @@
 import ipaddress
+from http import HTTPStatus
 from urllib.parse import quote, urlparse, urlsplit, urlunsplit
 
-from fastapi import status
 from pydantic import HttpUrl, ValidationError
 
 from skyvern.config import settings
@@ -78,7 +78,7 @@ def validate_url(url: str) -> str | None:
         url = prepend_scheme_and_validate_url(url=url)
         v = HttpUrl(url=url)
     except Exception as e:
-        raise SkyvernHTTPException(message=str(e), status_code=status.HTTP_400_BAD_REQUEST)
+        raise SkyvernHTTPException(message=str(e), status_code=HTTPStatus.BAD_REQUEST)
 
     if not v.host:
         return None

tests/unit/test_analytics.py

@@ -6,6 +6,17 @@ from unittest.mock import Mock, patch
 from skyvern import analytics
 
 
+class FakePosthog:
+    def __init__(self, api_key: str, **_: object) -> None:
+        self.api_key = api_key
+
+    def join(self) -> None:
+        pass
+
+    def flush(self) -> None:
+        pass
+
+
 def test_capture_includes_dynamic_test_id() -> None:
     original_test_id = analytics.settings.ANALYTICS_TEST_ID
     analytics.settings.ANALYTICS_TEST_ID = "smoke-test-123"
@@ -33,8 +44,10 @@ def test_reconfigure_posthog_client_uses_project_settings() -> None:
     analytics.settings.POSTHOG_PROJECT_HOST = "https://app.posthog.com"
 
     try:
-        analytics.reconfigure_posthog_client()
-        assert analytics.posthog.api_key == "phc_test_project_key"
+        with patch.object(analytics, "_load_posthog_class", return_value=FakePosthog):
+            analytics.reconfigure_posthog_client()
+            assert analytics.posthog is not None
+            assert analytics.posthog.api_key == "phc_test_project_key"
     finally:
         analytics.settings.POSTHOG_PROJECT_API_KEY = original_api_key
         analytics.settings.POSTHOG_PROJECT_HOST = original_host

tests/unit/test_browser_page_agent_login.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Literal
+from unittest.mock import AsyncMock
+
+import pytest
+
+from skyvern.client import SkyvernEnvironment, WorkflowRunResponse
+from skyvern.library.skyvern_browser_page_agent import SkyvernBrowserPageAgent
+from skyvern.schemas.credential_type import CredentialType
+
+
+def test_login_overload_literal_values_match_credential_type_values() -> None:
+    assert CredentialType.skyvern.value == "skyvern"
+    assert CredentialType.bitwarden.value == "bitwarden"
+    assert CredentialType.onepassword.value == "1password"
+    assert CredentialType.azure_vault.value == "azure_vault"
+
+
+def _workflow_run_response() -> WorkflowRunResponse:
+    return WorkflowRunResponse.model_validate(
+        {
+            "run_id": "wr_123",
+            "run_type": "workflow_run",
+            "status": "completed",
+            "created_at": "2026-05-05T00:00:00Z",
+            "modified_at": "2026-05-05T00:00:00Z",
+        }
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("credential_type", [CredentialType.skyvern, "skyvern"])
+async def test_browser_page_agent_login_normalizes_credential_type(
+    credential_type: CredentialType | Literal["skyvern"],
+) -> None:
+    workflow_run_response = _workflow_run_response()
+    skyvern = SimpleNamespace(
+        environment=SkyvernEnvironment.LOCAL,
+        login=AsyncMock(return_value=workflow_run_response),
+    )
+    browser = SimpleNamespace(
+        skyvern=skyvern,
+        browser_session_id="pbs_123",
+        browser_address="ws://127.0.0.1:9222",
+    )
+    agent = SkyvernBrowserPageAgent(browser, page=object())  # type: ignore[arg-type]
+    agent._wait_for_run_completion = AsyncMock(return_value=workflow_run_response)  # type: ignore[method-assign]
+
+    response = await agent.login(
+        credential_type=credential_type,
+        credential_id="cred_123",
+        url="https://example.com/login",
+    )
+
+    assert response.run_id == "wr_123"
+    assert skyvern.login.await_args.kwargs["credential_type"] is CredentialType.skyvern

tests/unit/test_copilot_model_resolver.py

@@ -6,12 +6,13 @@ from typing import Any
 from unittest.mock import MagicMock
 
 import pytest
+from structlog.testing import capture_logs
 
 
 class TestModelResolver:
     def test_router_config_empty_model_list_raises(self, monkeypatch: pytest.MonkeyPatch) -> None:
         from skyvern.forge.sdk.api.llm.exceptions import InvalidLLMConfigError
-        from skyvern.forge.sdk.api.llm.models import LLMRouterConfig
+        from skyvern.schemas.llm import LLMRouterConfig
 
         router_config = LLMRouterConfig(
             model_name="test",
@@ -38,7 +39,7 @@ class TestModelResolver:
         """Shim for SKY-9257: a router key resolves to its main_model_group entry as a direct
         LLMConfig so copilot-v2 can run until SKY-9256 lands the real bridge.
         """
-        from skyvern.forge.sdk.api.llm.models import LLMRouterConfig, LLMRouterModelConfig
+        from skyvern.schemas.llm import LLMRouterConfig, LLMRouterModelConfig
 
         main = LLMRouterModelConfig(
             model_name="vertex-gemini-2.5-flash",  # router group alias
@@ -84,12 +85,8 @@ class TestModelResolver:
         assert run_config.model_settings.extra_args is not None
         assert run_config.model_settings.extra_args["timeout"] == 900.0
 
-    def test_router_config_no_main_group_match_falls_back_to_first_entry(
-        self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
-    ) -> None:
-        import logging
-
-        from skyvern.forge.sdk.api.llm.models import LLMRouterConfig, LLMRouterModelConfig
+    def test_router_config_no_main_group_match_falls_back_to_first_entry(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        from skyvern.schemas.llm import LLMRouterConfig, LLMRouterModelConfig
 
         entry = LLMRouterModelConfig(
             model_name="some-group",
@@ -113,15 +110,15 @@ class TestModelResolver:
         handler = MagicMock()
         handler.llm_key = "MISCONFIGURED_ROUTER"
 
-        with caplog.at_level(logging.WARNING, logger="skyvern.forge.sdk.copilot.model_resolver"):
+        with capture_logs() as logs:
             model_name, _, _, _ = resolve_model_config(handler)
 
         assert model_name == "vertex_ai/gemini-2.5-flash"
-        joined = " ".join(record.getMessage() for record in caplog.records)
+        joined = " ".join(str(record.get("event", "")) for record in logs)
         assert "main_model_group has no matching" in joined
 
     def test_maps_basic_config(self, monkeypatch: pytest.MonkeyPatch) -> None:
-        from skyvern.forge.sdk.api.llm.models import LLMConfig
+        from skyvern.schemas.llm import LLMConfig
 
         monkeypatch.delenv("COPILOT_TRACING_ENABLED", raising=False)
         config = LLMConfig(
@@ -153,7 +150,7 @@ class TestModelResolver:
         assert run_config.model_settings.max_tokens == 4096
 
     def test_maps_basic_config_with_tracing_enabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
-        from skyvern.forge.sdk.api.llm.models import LLMConfig
+        from skyvern.schemas.llm import LLMConfig
 
         monkeypatch.setenv("COPILOT_TRACING_ENABLED", "1")
         config = LLMConfig(
@@ -179,7 +176,7 @@ class TestModelResolver:
         assert run_config.tracing_disabled is False
 
     def test_returns_supports_vision_false(self, monkeypatch: pytest.MonkeyPatch) -> None:
-        from skyvern.forge.sdk.api.llm.models import LLMConfig
+        from skyvern.schemas.llm import LLMConfig
 
         config = LLMConfig(
             model_name="openai/gpt-4-turbo",
@@ -200,9 +197,9 @@ class TestModelResolver:
         _, _, _, supports_vision = resolve_model_config(handler)
         assert supports_vision is False
 
-    def test_routes_all_litellm_params(self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture) -> None:
-        from skyvern.forge.sdk.api.llm.models import LiteLLMParams, LLMConfig
+    def test_routes_all_litellm_params(self, monkeypatch: pytest.MonkeyPatch) -> None:
         from skyvern.forge.sdk.copilot import model_resolver as model_resolver_module
+        from skyvern.schemas.llm import LiteLLMParams, LLMConfig
 
         # Reset the per-process warn-once gate so the caplog assertion is
         # deterministic regardless of test ordering.
@@ -238,7 +235,7 @@ class TestModelResolver:
         handler = MagicMock()
         handler.llm_key = "VERTEX_KEY"
 
-        with caplog.at_level("WARNING"):
+        with capture_logs() as logs:
             _, run_config, _, _ = resolve_model_config(handler)
         ms = run_config.model_settings
         assert ms is not None
@@ -255,10 +252,7 @@ class TestModelResolver:
         assert "thinking_level" not in ms.extra_args
         if ms.extra_body is not None:
             assert "thinking_level" not in ms.extra_body
-        assert any(
-            isinstance(record.msg, dict) and record.msg.get("dropped_key") == "thinking_level"
-            for record in caplog.records
-        )
+        assert any(record.get("dropped_key") == "thinking_level" for record in logs)
 
         for field in ("api_version", "model_info", "vertex_credentials", "vertex_location", "timeout"):
             assert ms.extra_args[field] == lp[field]
@@ -266,7 +260,7 @@ class TestModelResolver:
     def test_default_timeout_fallback(self, monkeypatch: pytest.MonkeyPatch) -> None:
         """When litellm_params has no timeout, inject settings.LLM_CONFIG_TIMEOUT."""
         from skyvern.config import settings
-        from skyvern.forge.sdk.api.llm.models import LLMConfig
+        from skyvern.schemas.llm import LLMConfig
 
         config = LLMConfig(
             model_name="openai/gpt-4",
@@ -289,8 +283,34 @@ class TestModelResolver:
         assert run_config.model_settings.extra_args is not None
         assert run_config.model_settings.extra_args["timeout"] == settings.LLM_CONFIG_TIMEOUT
 
+    def test_disables_litellm_aiohttp_transport(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        import litellm
+
+        from skyvern.schemas.llm import LLMConfig
+
+        monkeypatch.setattr(litellm, "disable_aiohttp_transport", False)
+        config = LLMConfig(
+            model_name="openai/gpt-4",
+            required_env_vars=[],
+            supports_vision=True,
+            add_assistant_prefix=False,
+        )
+        monkeypatch.setattr(
+            "skyvern.forge.sdk.copilot.model_resolver.LLMConfigRegistry.get_config",
+            lambda key: config,
+        )
+
+        from skyvern.forge.sdk.copilot.model_resolver import resolve_model_config
+
+        handler = MagicMock()
+        handler.llm_key = "BASIC_KEY"
+
+        resolve_model_config(handler)
+
+        assert litellm.disable_aiohttp_transport is True
+
     def test_explicit_timeout_wins_over_fallback(self, monkeypatch: pytest.MonkeyPatch) -> None:
-        from skyvern.forge.sdk.api.llm.models import LiteLLMParams, LLMConfig
+        from skyvern.schemas.llm import LiteLLMParams, LLMConfig
 
         lp: LiteLLMParams = {"timeout": 123.0}
         config = LLMConfig(
@@ -315,16 +335,12 @@ class TestModelResolver:
         assert run_config.model_settings.extra_args is not None
         assert run_config.model_settings.extra_args["timeout"] == 123.0
 
-    def test_warns_on_unrouted_litellm_params_keys(
-        self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
-    ) -> None:
+    def test_warns_on_unrouted_litellm_params_keys(self, monkeypatch: pytest.MonkeyPatch) -> None:
         """Keys in litellm_params that aren't explicitly routed should produce
         a LOG.warning listing the dropped keys — covers typos, dynamically
         injected values, and future additions to LiteLLMParams that we
         haven't updated the routing for."""
-        import logging
-
-        from skyvern.forge.sdk.api.llm.models import LLMConfig
+        from skyvern.schemas.llm import LLMConfig
 
         # Build a dict that bypasses TypedDict type-checking for the unknown key.
         lp: dict[str, Any] = {
@@ -337,7 +353,7 @@ class TestModelResolver:
             required_env_vars=[],
             supports_vision=True,
             add_assistant_prefix=False,
-            litellm_params=lp,  # type: ignore[typeddict-item]
+            litellm_params=lp,  # type: ignore[arg-type]
         )
         monkeypatch.setattr(
             "skyvern.forge.sdk.copilot.model_resolver.LLMConfigRegistry.get_config",
@@ -349,19 +365,15 @@ class TestModelResolver:
         handler = MagicMock()
         handler.llm_key = "WITH_TYPO_KEY"
 
-        with caplog.at_level(logging.WARNING, logger="skyvern.forge.sdk.copilot.model_resolver"):
+        with capture_logs() as logs:
             resolve_model_config(handler)
 
-        joined = " ".join(record.getMessage() for record in caplog.records)
+        joined = " ".join(str(record.get("unrouted_keys", "")) for record in logs)
         assert "future_litellm_addition" in joined
         assert "typo_feild_name" in joined
 
-    def test_no_warning_when_all_litellm_params_are_routed(
-        self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
-    ) -> None:
-        import logging
-
-        from skyvern.forge.sdk.api.llm.models import LiteLLMParams, LLMConfig
+    def test_no_warning_when_all_litellm_params_are_routed(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        from skyvern.schemas.llm import LiteLLMParams, LLMConfig
 
         lp: LiteLLMParams = {"api_base": "https://example.com", "timeout": 60.0}
         config = LLMConfig(
@@ -381,8 +393,8 @@ class TestModelResolver:
         handler = MagicMock()
         handler.llm_key = "CLEAN_KEY"
 
-        with caplog.at_level(logging.WARNING, logger="skyvern.forge.sdk.copilot.model_resolver"):
+        with capture_logs() as logs:
             resolve_model_config(handler)
 
-        joined = " ".join(record.getMessage() for record in caplog.records)
+        joined = " ".join(str(record.get("event", "")) for record in logs)
         assert "unrouted" not in joined

tests/unit/test_exception_messages.py

@@ -1,7 +1,15 @@
+from http import HTTPStatus
+
+import pytest
+
 from skyvern.exceptions import (
     SkyvernException,
+    SkyvernExtraNotInstalled,
+    SkyvernHTTPException,
     UnknownErrorWhileCreatingBrowserContext,
     get_user_facing_exception_message,
+    raise_server_extra_required,
+    require_server_extra_modules,
 )
 
 
@@ -35,6 +43,109 @@ def test_get_user_facing_exception_message_for_generic_exception() -> None:
     assert message == "Unexpected error: raw error"
 
 
+def test_skyvern_http_exception_normalizes_status_code_to_plain_int() -> None:
+    error = SkyvernHTTPException("bad request", status_code=HTTPStatus.BAD_REQUEST)
+
+    assert error.status_code == 400
+    assert type(error.status_code) is int
+
+
+def test_raise_server_extra_required_translates_when_server_extra_missing(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        "skyvern.exceptions.find_spec",
+        lambda module_name: None,
+    )
+    missing = ModuleNotFoundError("No module named 'starlette_context'", name="starlette_context")
+
+    with pytest.raises(SkyvernExtraNotInstalled, match=r"pip install skyvern\[server\]"):
+        raise_server_extra_required("skyvern.library.skyvern_browser", missing)
+
+
+def test_raise_server_extra_required_translates_missing_server_marker(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        "skyvern.exceptions.find_spec",
+        lambda module_name: None if module_name == "playwright" else object(),
+    )
+    missing = ModuleNotFoundError("No module named 'playwright'", name="playwright")
+
+    with pytest.raises(SkyvernExtraNotInstalled, match=r"pip install skyvern\[server\]"):
+        raise_server_extra_required("skyvern.library.skyvern_browser", missing)
+
+
+def test_raise_server_extra_required_preserves_installed_marker_submodule_failure(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr("skyvern.exceptions.find_spec", lambda module_name: object())
+    missing = ModuleNotFoundError("No module named 'playwright._impl._broken'", name="playwright._impl._broken")
+
+    with pytest.raises(ModuleNotFoundError) as exc_info:
+        raise_server_extra_required("skyvern.library.skyvern_browser", missing)
+
+    assert exc_info.value is missing
+
+
+def test_raise_server_extra_required_preserves_unknown_missing_dependency_when_server_extra_incomplete(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(
+        "skyvern.exceptions.find_spec",
+        lambda module_name: None if module_name == "playwright" else object(),
+    )
+    missing = ModuleNotFoundError("No module named 'bogus_internal_dep'", name="bogus_internal_dep")
+
+    with pytest.raises(ModuleNotFoundError) as exc_info:
+        raise_server_extra_required("skyvern.services.script_service", missing)
+
+    assert exc_info.value is missing
+
+
+def test_raise_server_extra_required_preserves_missing_dependency_when_server_markers_present(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr("skyvern.exceptions.find_spec", lambda module_name: object())
+    missing = ModuleNotFoundError("No module named 'bogus_internal_dep'", name="bogus_internal_dep")
+
+    with pytest.raises(ModuleNotFoundError) as exc_info:
+        raise_server_extra_required("skyvern.services.script_service", missing)
+
+    assert exc_info.value is missing
+
+
+def test_raise_server_extra_required_preserves_internal_skyvern_import_failure_when_server_extra_incomplete(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(
+        "skyvern.exceptions.find_spec",
+        lambda module_name: None if module_name == "playwright" else object(),
+    )
+    missing = ModuleNotFoundError("No module named 'skyvern.typo'", name="skyvern.typo")
+
+    with pytest.raises(ModuleNotFoundError) as exc_info:
+        raise_server_extra_required("skyvern.services.script_service", missing)
+
+    assert exc_info.value is missing
+
+
+def test_require_server_extra_modules_requires_server_sentinels(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        "skyvern.exceptions.find_spec",
+        lambda module_name: None if module_name == "sqlalchemy" else object(),
+    )
+
+    with pytest.raises(SkyvernExtraNotInstalled, match=r"pip install skyvern\[server\]"):
+        require_server_extra_modules("skyvern.library.skyvern_browser_page")
+
+
+def test_require_server_extra_modules_catches_partial_server_graph(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        "skyvern.exceptions.find_spec",
+        lambda module_name: None if module_name == "jinja2" else object(),
+    )
+
+    with pytest.raises(SkyvernExtraNotInstalled, match=r"pip install skyvern\[server\]"):
+        require_server_extra_modules("skyvern.library.skyvern_browser_page")
+
+
 def test_browser_connection_error_connect_over_cdp_websocket() -> None:
     """The exact error from SKY-8578: connect_over_cdp fails with 502 Bad Gateway."""
     raw_error = (

tests/unit/test_forge_app_initializer.py

@@ -0,0 +1,55 @@
+import runpy
+import sys
+from threading import Lock
+from types import SimpleNamespace
+
+import pytest
+
+from skyvern.forge import forge_app_initializer
+
+
+def test_start_forge_app_configures_server_logging_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    logging_calls: list[str] = []
+    forced_app_instances: list[object] = []
+    fake_app = object()
+
+    monkeypatch.setattr(forge_app_initializer, "_SERVER_LOGGING_CONFIGURED", False)
+    monkeypatch.setattr(forge_app_initializer, "_SERVER_LOGGING_LOCK", Lock())
+    monkeypatch.setattr(forge_app_initializer, "create_forge_app", lambda: fake_app)
+    monkeypatch.setattr(forge_app_initializer, "set_force_app_instance", forced_app_instances.append)
+    monkeypatch.setattr(forge_app_initializer.settings, "ADDITIONAL_MODULES", [])
+
+    import skyvern.forge.sdk.forge_log as forge_log
+
+    monkeypatch.setattr(forge_log, "setup_logger", lambda: logging_calls.append("setup"))
+
+    assert forge_app_initializer.start_forge_app() is fake_app
+    assert forge_app_initializer.start_forge_app() is fake_app
+
+    assert logging_calls == ["setup"]
+    assert forced_app_instances == [fake_app, fake_app]
+
+
+def test_forge_main_uses_shared_logging_initializer(monkeypatch: pytest.MonkeyPatch) -> None:
+    logging_calls: list[str] = []
+    uvicorn_calls: list[dict[str, object]] = []
+
+    monkeypatch.setattr(forge_app_initializer, "_SERVER_LOGGING_CONFIGURED", False)
+    monkeypatch.setattr(forge_app_initializer, "_SERVER_LOGGING_LOCK", Lock())
+
+    import skyvern.forge.sdk.forge_log as forge_log
+
+    monkeypatch.setattr(forge_log, "setup_logger", lambda: logging_calls.append("setup"))
+    monkeypatch.setattr("skyvern.exceptions.require_server_extra_modules", lambda *_args, **_kwargs: None)
+    monkeypatch.setattr("skyvern.analytics.capture", lambda *_args, **_kwargs: None)
+    monkeypatch.setitem(
+        sys.modules,
+        "uvicorn",
+        SimpleNamespace(run=lambda *args, **kwargs: uvicorn_calls.append({"args": args, "kwargs": kwargs})),
+    )
+
+    runpy.run_module("skyvern.forge.__main__", run_name="__main__")
+    forge_app_initializer._ensure_server_logging_configured()
+
+    assert logging_calls == ["setup"]
+    assert len(uvicorn_calls) == 1

tests/unit/test_library_skyvern.py

@@ -0,0 +1,34 @@
+from skyvern.schemas import runs
+from skyvern.schemas.proxy_location import (
+    SUPPORTED_GEO_COUNTRIES,
+    GeoTarget,
+    ProxyLocation,
+    proxy_location_to_request,
+)
+
+
+def test_runs_reexports_proxy_location_types() -> None:
+    assert runs.GeoTarget is GeoTarget
+    assert runs.ProxyLocation is ProxyLocation
+    assert runs.SUPPORTED_GEO_COUNTRIES is SUPPORTED_GEO_COUNTRIES
+    assert runs.proxy_location_to_request is proxy_location_to_request
+
+
+def test_proxy_location_to_request_dumps_geotarget_values() -> None:
+    assert proxy_location_to_request(GeoTarget(country="us")) == {
+        "country": "US",
+        "subdivision": None,
+        "city": None,
+    }
+
+
+def test_geotarget_schema_keeps_public_description() -> None:
+    assert GeoTarget.model_json_schema()["description"]
+
+
+def test_proxy_location_to_request_preserves_non_geotarget_values() -> None:
+    location = {"country": "GB"}
+
+    assert proxy_location_to_request(None) is None
+    assert proxy_location_to_request(ProxyLocation.RESIDENTIAL) is ProxyLocation.RESIDENTIAL
+    assert proxy_location_to_request(location) is location

tests/unit/test_llm_types_public_shim.py

@@ -0,0 +1,109 @@
+import copy
+import importlib
+import pickle
+
+import pytest
+
+from skyvern.schemas.llm import LLMConfig as CanonicalLLMConfig
+from skyvern.schemas.llm import LLMRouterConfig as CanonicalLLMRouterConfig
+
+
+def test_forge_llm_models_shim_warns_and_preserves_class_identity() -> None:
+    models = importlib.import_module("skyvern.forge.sdk.api.llm.models")
+
+    with pytest.warns(DeprecationWarning, match="skyvern.schemas.llm"):
+        legacy_llm_config = models.LLMConfig
+
+    with pytest.warns(DeprecationWarning, match="skyvern.schemas.llm"):
+        legacy_router_config = models.LLMRouterConfig
+
+    assert legacy_llm_config is CanonicalLLMConfig
+    assert legacy_router_config is CanonicalLLMRouterConfig
+
+
+def test_llm_config_defaults_resolve_at_construction() -> None:
+    config = CanonicalLLMConfig(
+        model_name="gpt-test",
+        required_env_vars=[],
+        supports_vision=True,
+        add_assistant_prefix=False,
+    )
+
+    assert config.max_tokens is not None
+    assert config.temperature is not None
+
+
+def test_llm_config_preserves_explicit_none_generation_params() -> None:
+    config = CanonicalLLMConfig(
+        model_name="o-series-test",
+        required_env_vars=[],
+        supports_vision=True,
+        add_assistant_prefix=False,
+        max_tokens=None,
+        temperature=None,
+    )
+
+    assert config.max_tokens is None
+    assert config.temperature is None
+
+
+def test_llm_router_config_preserves_explicit_none_generation_params() -> None:
+    config = CanonicalLLMRouterConfig(
+        model_name="router-test",
+        required_env_vars=[],
+        supports_vision=True,
+        add_assistant_prefix=False,
+        model_list=[],
+        main_model_group="router-test",
+        max_tokens=None,
+        temperature=None,
+    )
+
+    assert config.max_tokens is None
+    assert config.temperature is None
+
+
+def test_llm_config_generation_defaults_survive_copy_and_pickle() -> None:
+    configs: list[CanonicalLLMConfig | CanonicalLLMRouterConfig] = [
+        CanonicalLLMConfig(
+            model_name="gpt-test",
+            required_env_vars=[],
+            supports_vision=True,
+            add_assistant_prefix=False,
+        ),
+        CanonicalLLMRouterConfig(
+            model_name="router-test",
+            required_env_vars=[],
+            supports_vision=True,
+            add_assistant_prefix=False,
+            model_list=[],
+            main_model_group="router-test",
+        ),
+    ]
+
+    for config in configs:
+        for clone in (copy.deepcopy(config), pickle.loads(pickle.dumps(config))):
+            assert clone.max_tokens == config.max_tokens
+            assert clone.temperature == config.temperature
+
+
+def test_cloud_router_config_instances_keep_public_type_identity() -> None:
+    from skyvern.forge.sdk.api.llm.config_registry import LLMConfigRegistry
+
+    llm_key = "TEST_PUBLIC_TYPE_IDENTITY"
+    LLMConfigRegistry.deregister_config(llm_key)
+    LLMConfigRegistry.register_config(
+        llm_key,
+        CanonicalLLMConfig(
+            model_name="gpt-test",
+            required_env_vars=[],
+            supports_vision=True,
+            add_assistant_prefix=False,
+        ),
+    )
+
+    try:
+        config = LLMConfigRegistry.get_config(llm_key)
+        assert isinstance(config, CanonicalLLMConfig)
+    finally:
+        LLMConfigRegistry.deregister_config(llm_key)

tests/unit/test_organization_update_schema.py

@@ -32,6 +32,25 @@ class TestOrganizationUpdateSchema:
     def test_clear_artifact_flag_defaults_false(self) -> None:
         assert OrganizationUpdate().clear_artifact_url_expiry_seconds is False
 
+    def test_clear_max_steps_per_workflow_run_defaults_false(self) -> None:
+        assert OrganizationUpdate().clear_max_steps_per_workflow_run is False
+
+    def test_accepts_max_steps_per_workflow_run(self) -> None:
+        update = OrganizationUpdate(max_steps_per_workflow_run=42)
+        assert update.model_dump(exclude_unset=True) == {"max_steps_per_workflow_run": 42}
+
+    def test_rejects_zero_max_steps_per_workflow_run(self) -> None:
+        with pytest.raises(ValueError):
+            OrganizationUpdate(max_steps_per_workflow_run=0)
+
+    def test_rejects_negative_max_steps_per_workflow_run(self) -> None:
+        with pytest.raises(ValueError):
+            OrganizationUpdate(max_steps_per_workflow_run=-5)
+
+    def test_clear_max_steps_per_workflow_run_can_be_set(self) -> None:
+        update = OrganizationUpdate(clear_max_steps_per_workflow_run=True)
+        assert update.clear_max_steps_per_workflow_run is True
+
     def test_rejects_non_int_max_steps(self) -> None:
         with pytest.raises(ValueError):
             OrganizationUpdate(max_steps_per_run="not a number")  # type: ignore[arg-type]

tests/unit/test_run_response_types.py

@@ -0,0 +1,51 @@
+from datetime import UTC, datetime
+
+from pydantic import TypeAdapter
+
+from skyvern.schemas.runs import RunResponse, RunStatus, RunType, TaskRunResponse, WorkflowRunResponse
+
+
+def test_task_run_response_preserves_run_type_enum() -> None:
+    response = TaskRunResponse.model_validate(
+        {
+            "run_id": "tr_123",
+            "run_type": "task_v2",
+            "status": "completed",
+            "created_at": datetime.now(UTC).isoformat(),
+            "modified_at": datetime.now(UTC).isoformat(),
+        }
+    )
+
+    assert response.run_type is RunType.task_v2
+    assert response.run_type.value == "task_v2"
+    assert response.status is RunStatus.completed
+
+
+def test_workflow_run_response_preserves_run_type_enum() -> None:
+    response = WorkflowRunResponse.model_validate(
+        {
+            "run_id": "wr_123",
+            "run_type": "workflow_run",
+            "status": "completed",
+            "created_at": datetime.now(UTC).isoformat(),
+            "modified_at": datetime.now(UTC).isoformat(),
+        }
+    )
+
+    assert response.run_type is RunType.workflow_run
+    assert response.run_type.value == "workflow_run"
+
+
+def test_run_response_discriminator_preserves_run_type_enum() -> None:
+    response = TypeAdapter(RunResponse).validate_python(
+        {
+            "run_id": "tr_123",
+            "run_type": "task_v2",
+            "status": "completed",
+            "created_at": datetime.now(UTC).isoformat(),
+            "modified_at": datetime.now(UTC).isoformat(),
+        }
+    )
+
+    assert isinstance(response, TaskRunResponse)
+    assert response.run_type is RunType.task_v2

tests/unit/test_workflow_run_step_budget.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from skyvern.forge.agent import ForgeAgent
+
+
+def _make_organization(*, organization_id: str = "org-1", max_steps_per_workflow_run: int | None = None) -> MagicMock:
+    org = MagicMock()
+    org.organization_id = organization_id
+    org.max_steps_per_workflow_run = max_steps_per_workflow_run
+    return org
+
+
+def _make_task(
+    *, task_id: str = "task-1", organization_id: str = "org-1", workflow_run_id: str | None = "wr-1"
+) -> MagicMock:
+    task = MagicMock()
+    task.task_id = task_id
+    task.organization_id = organization_id
+    task.workflow_run_id = workflow_run_id
+    return task
+
+
+@pytest.mark.asyncio
+async def test_returns_none_when_task_has_no_workflow_run_id() -> None:
+    agent = ForgeAgent()
+    org = _make_organization(max_steps_per_workflow_run=20)
+    task = _make_task(workflow_run_id=None)
+
+    with patch("skyvern.forge.agent.app") as mock_app:
+        mock_app.DATABASE.tasks.get_tasks_by_workflow_run_id = AsyncMock()
+        mock_app.DATABASE.tasks.get_total_unique_step_order_count_by_task_ids = AsyncMock()
+
+        result = await agent._check_workflow_run_step_budget(org, task)
+
+    assert result is None
+    mock_app.DATABASE.tasks.get_tasks_by_workflow_run_id.assert_not_called()
+    mock_app.DATABASE.tasks.get_total_unique_step_order_count_by_task_ids.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_returns_none_when_org_has_no_cap() -> None:
+    agent = ForgeAgent()
+    org = _make_organization(max_steps_per_workflow_run=None)
+    task = _make_task(workflow_run_id="wr-1")
+
+    with patch("skyvern.forge.agent.app") as mock_app:
+        mock_app.DATABASE.tasks.get_tasks_by_workflow_run_id = AsyncMock()
+        mock_app.DATABASE.tasks.get_total_unique_step_order_count_by_task_ids = AsyncMock()
+
+        result = await agent._check_workflow_run_step_budget(org, task)
+
+    assert result is None
+    mock_app.DATABASE.tasks.get_tasks_by_workflow_run_id.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_returns_zero_count_when_workflow_run_has_no_tasks() -> None:
+    agent = ForgeAgent()
+    org = _make_organization(max_steps_per_workflow_run=20)
+    task = _make_task(workflow_run_id="wr-1")
+
+    with patch("skyvern.forge.agent.app") as mock_app:
+        mock_app.DATABASE.tasks.get_tasks_by_workflow_run_id = AsyncMock(return_value=[])
+        mock_app.DATABASE.tasks.get_total_unique_step_order_count_by_task_ids = AsyncMock()
+
+        result = await agent._check_workflow_run_step_budget(org, task)
+
+    assert result == (0, 20)
+    # Skip the count query when there are no task ids to filter on.
+    mock_app.DATABASE.tasks.get_total_unique_step_order_count_by_task_ids.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_returns_count_and_cap_when_tasks_present() -> None:
+    agent = ForgeAgent()
+    org = _make_organization(max_steps_per_workflow_run=20)
+    task = _make_task(workflow_run_id="wr-1")
+    sibling_tasks = [
+        MagicMock(task_id="task-1"),
+        MagicMock(task_id="task-2"),
+    ]
+
+    with patch("skyvern.forge.agent.app") as mock_app:
+        mock_app.DATABASE.tasks.get_tasks_by_workflow_run_id = AsyncMock(return_value=sibling_tasks)
+        mock_app.DATABASE.tasks.get_total_unique_step_order_count_by_task_ids = AsyncMock(return_value=18)
+
+        result = await agent._check_workflow_run_step_budget(org, task)
+
+    assert result == (18, 20)
+    mock_app.DATABASE.tasks.get_total_unique_step_order_count_by_task_ids.assert_awaited_once_with(
+        task_ids=["task-1", "task-2"],
+        organization_id="org-1",
+    )
+
+
+@pytest.mark.asyncio
+async def test_treats_none_count_as_zero() -> None:
+    # ``get_total_unique_step_order_count_by_task_ids`` can return None for a no-row count.
+    agent = ForgeAgent()
+    org = _make_organization(max_steps_per_workflow_run=20)
+    task = _make_task(workflow_run_id="wr-1")
+
+    with patch("skyvern.forge.agent.app") as mock_app:
+        mock_app.DATABASE.tasks.get_tasks_by_workflow_run_id = AsyncMock(return_value=[MagicMock(task_id="task-1")])
+        mock_app.DATABASE.tasks.get_total_unique_step_order_count_by_task_ids = AsyncMock(return_value=None)
+
+        result = await agent._check_workflow_run_step_budget(org, task)
+
+    assert result == (0, 20)