feat(serverless): address an environment on invoke and get_task (RUNSERV-1238)

README.md

@@ -597,6 +597,22 @@ task = await client.get_task("image-tools", accepted["id"])
 
 A task that runs and fails comes back with `status` set to `"failed"` and its `error` set. Only a failure to start raises.
 
+An app id names one copy of the app per environment. Calls go to `production` unless you name another environment, for the whole client or for one call. Waiting polls the task in the environment it was invoked in; when you call `get_task` yourself, pass that same environment. An environment paused at its monthly budget refuses an invocation with a `quota` error, which is not retried: it holds until the budget is raised or the UTC month turns.
+
+```python
+client = Runware(serverless_environment="staging")
+
+task = await client.invoke("image-tools", "generate", {"prompt": "a red bicycle"})
+
+task = await client.invoke(
+    "image-tools", "generate", {"prompt": "a red bicycle"}, environment="production",
+)
+accepted = await client.invoke(
+    "image-tools", "train", {"epochs": 10}, environment="dev", options=InvokeOptions(wait=False),
+)
+task = await client.get_task("image-tools", accepted["id"], environment="dev")
+```
+
 ## Content metadata
 
 `client.content.*` exposes Runware's curated model catalog as read-only metadata — names, AIRs, headlines, capabilities, pricing, examples. Public information, no extra cost.

runware/client.py

@@ -425,6 +425,7 @@ async def invoke(
         payload: dict[str, Any] | None = None,  # pyright: ignore[reportExplicitAny]
         *,
         task_id: str | None = None,
+        environment: str | None = None,
         options: InvokeOptions | None = None,
     ) -> ServerlessTask:
         """
@@ -444,28 +445,44 @@ async def invoke(
         Sending the same id again returns the task it already names instead of
         starting a second, so a call whose response was lost is safe to repeat.
 
+        ``environment`` selects which copy of the app to call, falling back to
+        ``config.serverless_environment`` and then to ``production``. Waiting
+        polls the task in that same environment.
+
         Raises a ``RunwareError`` when the task fails to start. A task that
         runs and fails comes back with ``status`` ``failed`` and its ``error``
         set, because that is an outcome rather than a broken call.
         """
         return await self._serverless.invoke(
-            app_id, endpoint_path, payload, task_id=task_id, options=options,
+            app_id,
+            endpoint_path,
+            payload,
+            task_id=task_id,
+            environment=environment,
+            options=options,
         )
 
     async def get_task(
         self,
         app_id: str,
         task_id: str,
         options: GetTaskOptions | None = None,
+        *,
+        environment: str | None = None,
     ) -> ServerlessTask:
         """
         Read one serverless task by id.
 
         ``invoke`` already waits for you, so reach for this when you want to
         poll yourself: after ``invoke`` with ``wait=False``, or from a
         different process than the one that submitted the task.
+
+        Pass the ``environment`` the task was invoked in. It falls back to
+        ``config.serverless_environment`` and then to ``production``.
         """
-        return await self._serverless.get_task(app_id, task_id, options)
+        return await self._serverless.get_task(
+            app_id, task_id, options, environment=environment,
+        )
 
     # ----------------------------------------------------------- internal
 

runware/errors.py

@@ -158,6 +158,7 @@ def derive_code(raw: str) -> ErrorCode:
         or "Balance" in raw
         or raw == "quotaExceeded"
         or raw == "paymentRequired"
+        or raw == "environmentPaused"
     ):
         return "quota"
     if "RateLimit" in raw or raw == "rateLimitExceeded":

runware/serverless.py

@@ -19,6 +19,7 @@
 import uuid
 from collections.abc import Awaitable, Callable
 from typing import Any, cast
+from urllib.parse import urlencode
 
 import aiohttp
 
@@ -36,6 +37,7 @@
 
 _APP_ID_PATTERN = re.compile(r"^[a-z][a-z0-9-]{4,28}[a-z0-9]$")
 _ENDPOINT_PATH_PATTERN = re.compile(r"^[a-z]([a-z0-9-]{0,62}[a-z0-9])?$")
+_ENVIRONMENT_PATTERN = re.compile(r"^[a-z0-9][a-z0-9._-]{0,62}$")
 _TASK_ID_PATTERN = re.compile(
     r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$",
 )
@@ -55,6 +57,10 @@
 
 _TERMINAL = frozenset({"completed", "failed"})
 
+# The 429 an environment paused at its monthly budget answers invoke with. It
+# holds until the budget is raised or the UTC month turns, so it is not retried.
+_ENVIRONMENT_PAUSED_TYPE = "https://docs.runware.ai/serverless/errors#environment-paused"
+
 
 def is_terminal(status: str) -> bool:
     """Whether a task status is final."""
@@ -121,6 +127,25 @@ def validate_endpoint_path(endpoint_path: object) -> str:
     return endpoint_path
 
 
+def validate_environment(
+    environment: object, parameter: str = "environment",
+) -> str | None:
+    """Check an environment name against the platform's own pattern."""
+    if environment is None:
+        return None
+    if not isinstance(environment, str) or not _ENVIRONMENT_PATTERN.fullmatch(
+        environment,
+    ):
+        raise create_runware_error(
+            "invalidParameter",
+            f'{parameter} "{environment}" is invalid: use 1 to 63 lowercase '
+            + "letters, digits, hyphens, underscores or periods, starting with "
+            + "a letter or digit",
+            parameter=parameter,
+        )
+    return environment
+
+
 def validate_task_id(task_id: object) -> str:
     """Check a task id is a canonical lowercase UUID."""
     if not isinstance(task_id, str) or not _TASK_ID_PATTERN.match(task_id):
@@ -187,10 +212,12 @@ def _problem_to_error(
     if lines:
         message = message + "\n" + "\n".join(lines)
 
-    error = create_runware_error(
-        _raw_code_for_status(status), message, status_code=status,
-    )
     problem_type = problem.get("type")
+    raw_code = (
+        "environmentPaused" if problem_type == _ENVIRONMENT_PAUSED_TYPE
+        else _raw_code_for_status(status)
+    )
+    error = create_runware_error(raw_code, message, status_code=status)
     if isinstance(problem_type, str) and problem_type != "about:blank":
         error.problem_type = problem_type
     request_id = problem.get("requestId")
@@ -223,20 +250,37 @@ class ServerlessApi:
     def __init__(
         self, session_provider: SessionProvider, config: SDKConfig,
     ) -> None:
+        _ = validate_environment(
+            config.serverless_environment, "serverless_environment",
+        )
         self._session_provider: SessionProvider = session_provider
         self._config: SDKConfig = config
 
     @property
     def _origin(self) -> str:
         return self._config.serverless_base_url.rstrip("/")
 
+    def _environment(self, environment: str | None) -> str | None:
+        """The per-call environment, else the client's, else none."""
+        return validate_environment(
+            environment if environment is not None
+            else self._config.serverless_environment,
+        )
+
+    def _app_url(self, app: str, path: str, environment: str | None) -> str:
+        url = f"{self._origin}/v1/apps/{app}/{path}"
+        if environment is None:
+            return url
+        return f"{url}?{urlencode({'environment': environment})}"
+
     async def invoke(
         self,
         app_id: str,
         endpoint_path: str,
         payload: dict[str, Any] | None = None,  # pyright: ignore[reportExplicitAny]
         *,
         task_id: str | None = None,
+        environment: str | None = None,
         options: InvokeOptions | None = None,
     ) -> ServerlessTask:
         """
@@ -248,6 +292,10 @@ async def invoke(
             )
             print(task["output"])
 
+        ``environment`` selects which copy of the app to call, falling back to
+        ``config.serverless_environment`` and then to ``production``. Waiting
+        polls the task in that same environment.
+
         ``delivery_method`` decides how the wait happens, not whether you get
         a result: ``async`` (the default) takes an acknowledgement and polls,
         ``sync`` holds one request open and is the fast path for work that
@@ -266,6 +314,7 @@ async def invoke(
         app = validate_app_id(app_id)
         endpoint = validate_endpoint_path(endpoint_path)
         resolved_task_id = _resolve_task_id(task_id)
+        resolved_environment = self._environment(environment)
         delivery_method = options.delivery_method if options.wait else "async"
         route = "invoke-sync" if delivery_method == "sync" else "invoke-async"
 
@@ -285,7 +334,7 @@ async def invoke(
         started_at = time.monotonic()
         _status, response_body = await self._request(
             "POST",
-            f"{self._origin}/v1/apps/{app}/{route}/{endpoint}",
+            self._app_url(app, f"{route}/{endpoint}", resolved_environment),
             body,
             timeout_ms,
             options.cancel_event,
@@ -296,27 +345,35 @@ async def invoke(
         if not options.wait or is_terminal(task["status"]):
             return task
 
-        return await self._wait_for_task(task, options, started_at)
+        return await self._wait_for_task(
+            task, resolved_environment, options, started_at,
+        )
 
     async def get_task(
         self,
         app_id: str,
         task_id: str,
         options: GetTaskOptions | None = None,
+        *,
+        environment: str | None = None,
     ) -> ServerlessTask:
         """
         Read one serverless task by id.
 
         ``invoke`` already waits for you, so reach for this when you want to
         poll yourself: after ``invoke`` with ``wait=False``, or from a
         different process than the one that submitted the task.
+
+        Pass the ``environment`` the task was invoked in. It falls back to
+        ``config.serverless_environment`` and then to ``production``.
         """
         options = options or GetTaskOptions()
         app = validate_app_id(app_id)
         resolved = validate_task_id(task_id)
+        resolved_environment = self._environment(environment)
         _status, body = await self._request(
             "GET",
-            f"{self._origin}/v1/apps/{app}/tasks/{resolved}",
+            self._app_url(app, f"tasks/{resolved}", resolved_environment),
             None,
             options.timeout or self._config.timeout,
             options.cancel_event,
@@ -326,12 +383,16 @@ async def get_task(
     async def _wait_for_task(
         self,
         task: ServerlessTask,
+        environment: str | None,
         options: InvokeOptions,
         started_at: float,
     ) -> ServerlessTask:
         """
         Poll until the task reaches a terminal state. Never resubmits: the task
         is already running, and a second invocation would be a second charge.
+
+        Polls in the environment the task was invoked in. A task carries no
+        environment of its own, and a poll without one would read production.
         """
         interval = (
             options.poll_interval
@@ -344,13 +405,17 @@ async def _wait_for_task(
 
         while not is_terminal(current["status"]):
             if time.monotonic() - started_at > budget_s:
+                environment_arg = (
+                    "" if environment is None
+                    else f', environment="{environment}"'
+                )
                 raise create_runware_error(
                     "timeout",
                     f'Task {current["id"]} is still running after '
                     + f"{options.poll_timeout or self._config.poll_timeout}ms. "
                     + "It was not cancelled: poll "
-                    + f'get_task("{current["appId"]}", "{current["id"]}") '
-                    + "to pick it up.",
+                    + f'get_task("{current["appId"]}", "{current["id"]}"'
+                    + f"{environment_arg}) to pick it up.",
                     task_uuid=current["id"],
                 )
 
@@ -364,6 +429,7 @@ async def _wait_for_task(
                         timeout=options.timeout,
                         cancel_event=options.cancel_event,
                     ),
+                    environment=environment,
                 )
                 not_found_since = None
             except RunwareError as error:
@@ -428,7 +494,7 @@ async def _request(
     @staticmethod
     def _should_retry(error: BaseException) -> bool:
         if is_runware_error(error):
-            if error.code == "aborted":
+            if error.code in ("aborted", "quota"):
                 return False
             if error.code == "timeout":
                 return True

runware/types/sdk.py

@@ -57,6 +57,10 @@ class SDKConfig:
     http_base_url: str = "https://api.runware.ai/v1"
     # Serverless API origin, without a version path. Used by invoke and get_task.
     serverless_base_url: str = "https://api.serverless.runware.ai"
+    # Serverless environment invoke and get_task address, such as "staging".
+    # An app id names one copy per environment. None sends no environment,
+    # which the platform reads as "production".
+    serverless_environment: str | None = None
 
     timeout: int = 1_200_000
     poll_timeout: int = 1_200_000

tests/integration/test_smoke.py

@@ -32,6 +32,8 @@
 SERVERLESS_BASE = os.environ.get("RUNWARE_SERVERLESS_BASE_URL")
 SERVERLESS_APP = os.environ.get("RUNWARE_SERVERLESS_TEST_APP")
 SERVERLESS_ENDPOINT = os.environ.get("RUNWARE_SERVERLESS_TEST_ENDPOINT")
+# Optional: the environment of the app copy to call. Unset means production.
+SERVERLESS_ENVIRONMENT = os.environ.get("RUNWARE_SERVERLESS_TEST_ENVIRONMENT") or None
 SERVERLESS_PAYLOAD: dict[str, object] = json.loads(
     os.environ.get("RUNWARE_SERVERLESS_TEST_PAYLOAD") or "{}",
 )
@@ -42,7 +44,11 @@
 
 
 def serverless_client() -> Runware:
-    return Runware(serverless_base_url=SERVERLESS_BASE or "", transport="rest")
+    return Runware(
+        serverless_base_url=SERVERLESS_BASE or "",
+        serverless_environment=SERVERLESS_ENVIRONMENT,
+        transport="rest",
+    )
 
 IMAGE_MODEL = "runware:400@2"  # Flux 2 Klein 9b — cheap and fast
 TEXT_MODEL = "google:gemma@4-31b"  # cheap and fast