fix(job): surface review/spend-limited states and cost warning

.vscode/settings.json

@@ -6,8 +6,8 @@
         "-p",
         "test_*.py"
     ],
-    "python.testing.unittestEnabled": true,
-    "python.testing.pytestEnabled": false,
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true,
     "python.autoComplete.extraPaths": [
         "./src"
     ],
@@ -19,5 +19,8 @@
     "cursorpyright.analysis.extraPaths": [
         "./src"
     ],
-    "cursorpyright.analysis.typeCheckingMode": "basic"
+    "cursorpyright.analysis.typeCheckingMode": "basic",
+    "python.testing.pytestArgs": [
+        "tests"
+    ]
 }

docs/audiences.md

@@ -99,6 +99,19 @@ results = job.get_results()
 print(results)
 ```
 
+### Cost warnings and jobs under review
+
+`assign_job` never blocks on funds: the job is always created. If its estimated cost
+exceeds your account balance, `assign_job` logs a warning with the estimate, your
+balance, and the expected shortfall — the job still runs, but may pause partway until
+you top up.
+
+Some jobs don't go straight to running. A job can enter manual review
+(`ManualApproval`) or, once out of funds mid-run, become spend-limited
+(`SpendLimited`). Neither state completes on its own, so `get_results()` raises an
+informative error naming the state (and the review reason, when available) instead of
+blocking indefinitely — top up or wait for a reviewer, then call it again.
+
 ## Complete Example
 
 Here's the full workflow — creating a custom audience, adding qualification examples, and running a labeling job:

docs/quickstart.md

@@ -105,7 +105,7 @@ results = job.get_results() # (3)!
 
 1. Assigns the job definition to the audience and starts collecting responses.
 2. Opens your browser on the running job, where you can watch responses come in and monitor progress.
-3. Blocks until the job is complete and returns the results. You can also monitor progress on the [Rapidata Dashboard](https://app.rapidata.ai/dashboard).
+3. Blocks until the job is complete and returns the results. If the job needs manual review or runs out of funds, this raises an informative error instead of blocking — see [Cost warnings and jobs under review](audiences.md#cost-warnings-and-jobs-under-review). You can also monitor progress on the [Rapidata Dashboard](https://app.rapidata.ai/dashboard).
 
 To understand the results format, see the [Understanding the Results](understanding_the_results.md) guide.
 

pyproject.toml

@@ -34,6 +34,7 @@ dev = [
     "ipykernel>=7.1.0,<8",
     "jupyter>=1.1.1,<2",
     "pyright>=1.1.407",
+    "pytest>=8.0.0,<9",
 ]
 docs = [
     "mkdocs-material>=9.5.34,<10",
@@ -49,6 +50,9 @@ docs = [
     "mkdocs-llmstxt-md>=0.1.0",
 ]
 
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
 [tool.hatch.build.targets.sdist]
 include = ["src/rapidata"]
 

src/rapidata/rapidata_client/audience/_audience_base.py

@@ -10,6 +10,9 @@
         RapidataJobDefinition,
     )
     from rapidata.rapidata_client.job.rapidata_job import RapidataJob
+    from rapidata.api_client.models.create_job_endpoint_cost_warning_model import (
+        CreateJobEndpointCostWarningModel,
+    )
 
 
 class RapidataAudienceBase:
@@ -89,8 +92,31 @@ def assign_job(self, job_definition: RapidataJobDefinition) -> RapidataJob:
             managed_print(
                 f"Job '{job.name}' is now viewable under: {job.job_details_page}"
             )
+            self._warn_if_cost_exceeds_balance(job, response.cost_warning)
             return job
 
+    @staticmethod
+    def _warn_if_cost_exceeds_balance(
+        job: RapidataJob, cost_warning: CreateJobEndpointCostWarningModel | None
+    ) -> None:
+        """Surface the create response's optional cost warning via the SDK logger.
+
+        The job is created and runs regardless — this is an advisory estimate that it
+        may pause for funds before finishing, not an error.
+        """
+        if cost_warning is None:
+            return
+        logger.warning(
+            "Job '%s' has an estimated cost of %.2f, but the account balance is %.2f — "
+            "it will likely pause about %.2f short of finishing. The job was created and "
+            "runs as far as the balance allows; top up the account to let it complete. "
+            "This is an estimate.",
+            job.name,
+            cost_warning.estimated_cost,
+            cost_warning.available_balance,
+            cost_warning.shortfall,
+        )
+
     def find_jobs(
         self, name: str = "", amount: int = 10, page: int = 1
     ) -> list[RapidataJob]:

src/rapidata/rapidata_client/job/rapidata_job.py

@@ -9,6 +9,7 @@
 from colorama import Fore
 from tqdm.auto import tqdm
 
+from rapidata.api_client.models.audience_job_state import AudienceJobState
 from rapidata.service.openapi_service import OpenAPIService
 from rapidata.rapidata_client.config import (
     logger,
@@ -27,6 +28,9 @@
 )
 
 if TYPE_CHECKING:
+    from rapidata.api_client.models.get_job_by_id_endpoint_output import (
+        GetJobByIdEndpointOutput,
+    )
     from rapidata.rapidata_client.results.rapidata_results import RapidataResults
 
 T = TypeVar("T")
@@ -70,15 +74,47 @@ def __init__(
         self.job_details_page = f"https://app.{self._openapi_service.environment}/audiences/{self.audience_id}/job/{self.id}"
         logger.debug("RapidataJob initialized")
 
+    # States a job can settle into that will never progress to Completed/Failed on
+    # their own: ManualApproval needs a Rapidata reviewer to act, SpendLimited needs
+    # an account top-up. Waiting on either (e.g. from get_results) would hang the
+    # caller forever, so we surface them as informative errors instead.
+    _BLOCKING_STATUSES = (
+        AudienceJobState.MANUALAPPROVAL,
+        AudienceJobState.SPENDLIMITED,
+    )
+
+    def _fetch_job(self) -> GetJobByIdEndpointOutput:
+        """Fetches the job's GET output (state, failure message, review reason)."""
+        return self._openapi_service.order.job_api.job_job_id_get(self.id)
+
     def _get_job_failure_message(self) -> str | None:
         """Retrieves the failure message from the job if available."""
         try:
-            job = self._openapi_service.order.job_api.job_job_id_get(self.id)
-            return job.failure_message
+            return self._fetch_job().failure_message
         except Exception:
             logger.debug("Failed to get job failure message", self, exc_info=True)
             return None
 
+    def _raise_for_blocking_status(self, job: GetJobByIdEndpointOutput) -> None:
+        """Raises an informative error for a job state that can't reach completion
+        on its own, instead of letting the caller block on it indefinitely."""
+        if job.state == AudienceJobState.SPENDLIMITED:
+            raise Exception(
+                f"Job '{self}' is spend-limited: the account ran out of funds while "
+                f"running, so it stopped collecting responses. Partial results remain "
+                f"available; top up the account to resume and let the job finish."
+            )
+
+        # ManualApproval — reviewReason is optional; a job can legitimately be under
+        # review with no customer-facing reason recorded yet.
+        reason = job.review_reason
+        reason_detail = f" ({reason.value})" if reason else ""
+        raise Exception(
+            f"Job '{self}' is being reviewed{reason_detail}: a Rapidata reviewer "
+            f"needs to approve it before it runs. This can take a couple of hours; "
+            f"once approved, call this again."
+        )
+
     def _retry_operation(
         self,
         operation: Callable[[], T],
@@ -133,14 +169,23 @@ def _wait_for_status(
 
         Returns:
             The final status reached
+
+        Raises:
+            Exception: If the job enters a state it can't progress out of on its own
+                (``ManualApproval`` or ``SpendLimited``) while a different status is
+                being awaited — with the review reason when the API provides one.
         """
-        while (current_status := self.get_status()) not in target_statuses:
+        while True:
+            job = self._fetch_job()
+            current_status = job.state.value
+            if current_status in target_statuses:
+                return current_status
+            if current_status in self._BLOCKING_STATUSES:
+                self._raise_for_blocking_status(job)
             if status_message:
                 logger.debug(status_message, self, current_status)
             sleep(check_interval)
 
-        return current_status
-
     @property
     def completed_at(self) -> datetime | None:
         """Returns the completion date of the job, or None if not completed."""
@@ -190,9 +235,7 @@ def get_status(self) -> str:
             The current status of the job as a string.
         """
         with tracer.start_as_current_span("RapidataJob.get_status"):
-            return self._openapi_service.order.job_api.job_job_id_get(
-                self.id
-            ).state.value
+            return self._fetch_job().state.value
 
     def _regenerate_results(self) -> None:
         """Triggers regeneration of a job whose results have gone stale.
@@ -221,7 +264,9 @@ def get_results(self) -> RapidataResults:
             RapidataResults: The results of the job.
 
         Raises:
-            Exception: If failed to get job results.
+            Exception: If failed to get job results, or if the job is in manual
+                review (``ManualApproval``) or spend-limited (``SpendLimited``) and
+                therefore cannot complete without intervention.
         """
         with tracer.start_as_current_span("RapidataJob.get_results"):
             from rapidata.api_client.exceptions import ApiException
@@ -260,18 +305,22 @@ def display_progress_bar(self, refresh_rate: int = 5) -> None:
 
         Raises:
             ValueError: If refresh_rate is less than 1.
+            Exception: If the job has failed, or is in a state it can't progress out
+                of on its own (``ManualApproval`` or ``SpendLimited``).
         """
         if refresh_rate < 1:
             raise ValueError("refresh_rate must be at least 1")
 
-        current_status = self.get_status()
-        if current_status == "Completed":
+        job = self._fetch_job()
+        if job.state == AudienceJobState.COMPLETED:
             managed_print(f"Job '{self}' is already completed.")
             return
 
-        if current_status == "Failed":
-            failure_message = self._get_job_failure_message()
-            raise Exception(f"Job '{self}' has failed: {failure_message}")
+        if job.state == AudienceJobState.FAILED:
+            raise Exception(f"Job '{self}' has failed: {job.failure_message}")
+
+        if job.state in self._BLOCKING_STATUSES:
+            self._raise_for_blocking_status(job)
 
         # Get progress from pipeline if available
         with tqdm(
@@ -283,15 +332,17 @@ def display_progress_bar(self, refresh_rate: int = 5) -> None:
         ) as pbar:
             last_percentage = 0
             while True:
-                current_status = self.get_status()
+                job = self._fetch_job()
 
-                if current_status == "Completed":
+                if job.state == AudienceJobState.COMPLETED:
                     pbar.update(100 - last_percentage)
                     break
 
-                if current_status == "Failed":
-                    failure_message = self._get_job_failure_message()
-                    raise Exception(f"Job '{self}' has failed: {failure_message}")
+                if job.state == AudienceJobState.FAILED:
+                    raise Exception(f"Job '{self}' has failed: {job.failure_message}")
+
+                if job.state in self._BLOCKING_STATUSES:
+                    self._raise_for_blocking_status(job)
 
                 # Try to get progress from workflow
                 try:

tests/rapidata_client/audience/test_audience_base.py

@@ -0,0 +1,64 @@
+"""Tests for the create-time cost warning surfaced by assign_job.
+
+When job creation returns a costWarning (estimate exceeds balance), assign_job
+logs an advisory warning. The job is created regardless; behaviour is otherwise
+unchanged.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import rapidata.rapidata_client.audience._audience_base as base_module
+from rapidata.api_client.models.create_job_endpoint_cost_warning_model import (
+    CreateJobEndpointCostWarningModel,
+)
+from rapidata.rapidata_client.audience._audience_base import RapidataAudienceBase
+
+
+def _make_audience() -> tuple[RapidataAudienceBase, MagicMock]:
+    openapi_service = MagicMock()
+    openapi_service.environment = "rapidata.ai"
+    audience = RapidataAudienceBase(
+        id="aud-1", name="Aud", filters=[], openapi_service=openapi_service
+    )
+    return audience, openapi_service
+
+
+def test_assign_job_warns_on_cost_warning(monkeypatch):
+    audience, openapi_service = _make_audience()
+    response = MagicMock()
+    response.job_id = "job-1"
+    response.cost_warning = CreateJobEndpointCostWarningModel(
+        estimatedCost=120.0, availableBalance=50.0, shortfall=70.0
+    )
+    openapi_service.order.job_api.job_post.return_value = response
+
+    warn = MagicMock()
+    monkeypatch.setattr(base_module.logger, "warning", warn)
+
+    audience.assign_job(MagicMock(id="def-1", name="My Job"))
+
+    warn.assert_called_once()
+    args = warn.call_args.args
+    assert "estimate" in args[0].lower()
+    # Values are passed as lazy %-args, not pre-formatted into the message.
+    assert 120.0 in args
+    assert 50.0 in args
+    assert 70.0 in args
+
+
+def test_assign_job_no_warning_when_cost_within_balance(monkeypatch):
+    audience, openapi_service = _make_audience()
+    response = MagicMock()
+    response.job_id = "job-1"
+    response.cost_warning = None
+    openapi_service.order.job_api.job_post.return_value = response
+
+    warn = MagicMock()
+    monkeypatch.setattr(base_module.logger, "warning", warn)
+
+    job = audience.assign_job(MagicMock(id="def-1", name="My Job"))
+
+    warn.assert_not_called()
+    assert job.id == "job-1"