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

[RapidPoseidon]

What & why

Workstream E (Wave 2) of the RFC Job submission: cost warning + transparent review outcomes. Two customer-facing SDK gaps:

1. Silent hang. RapidataJob._wait_for_status waited only for Completed/Failed, so a job that entered ManualApproval or SpendLimited blocked the caller forever (most visibly via get_results()). It now raises an informative error instead:

   • ManualApproval → names the review reason when the API returns one (reviewReason), and falls back to a generic "is being reviewed" message when it's null (a job can legitimately be under review with no reason yet).
   • SpendLimited → says partial results remain available and the account needs a top-up to resume.
   • display_progress_bar had the identical loop and gets the same guard, so the two entry points behave consistently.

2. Cost warning. When job creation returns the new costWarning { estimatedCost, availableBalance, shortfall }, assign_job now emits a Python warnings.warn (estimate, balance, shortfall), phrased as an estimate. The job is still created and behavior is otherwise unchanged — absent costWarning → no warning.

Follow-up commit folds in code-review findings: the stale "blocks until complete" promise in quickstart.md now points at the new review/cost-warning docs; the polling loops fetch the job GET output once per iteration instead of re-fetching it for the review reason / failure message when raising; and the blocking-state comparisons use the generated AudienceJobState enum instead of raw strings (matching how RapidataOrder uses OrderState).

Tests

Introduces pytest (dev dep) and tests/test_job_review_and_cost_warning.py:

• hang fix on ManualApproval with and without a reviewReason
• hang fix on SpendLimited (+ display_progress_bar)
• costWarning emission and the no-warning path
• unchanged happy path (Completed → results returned)

pyright, black, pytest (7/7), and mkdocs build all clean; docs notes in audiences.md and quickstart.md.

Depends on the generated models from Update OpenAPI Schemas #670 — merged, and this branch is now rebased on top of it, so the diff is only the hand-written changes.

Note: no CI job runs pytest today (CI is pyright-only); wiring the new suite into CI is a sensible follow-up.

🔗 Session: https://node-07688719.poseidon.rapidata.internal/
🔗 Review session: https://cli-karl-57b3b766.poseidon.rapidata.internal/

[claude]

Review

I focused on the hand-written diff (rapidata_job.py, _audience_base.py, tests, docs) since the api_client/openapi changes belong to #670 per the PR description.

Bugs / correctness

1. Paused job state still hangs forever — the exact bug this PR fixes, just for a different state.
AudienceJobState has a PAUSED value (audience_job_state.py:34), and job_api.py exposes job_job_id_pause_post/job_job_id_resume_post, so a job can enter Paused and, like ManualApproval/SpendLimited, will never leave it on its own. But RapidataJob._BLOCKING_STATUSES (rapidata_job.py:77) only lists ("ManualApproval", "SpendLimited"). If a job is Paused, both _wait_for_status (used by get_results()) and display_progress_bar's polling loop will spin indefinitely — reintroducing the silent-hang problem this PR is meant to eliminate.

Notably, the sibling RapidataOrder class already treats Paused as a state that must not be waited on: its _wait_for_state for get_results() includes OrderState.PAUSED as a target/stop state (rapidata_order.py:421-429), and display_progress_bar has an explicit comment: "Without this check a Failed or Paused order would pin the caller's thread forever" (rapidata_order.py:334-341). It looks like the job-side fix missed porting this case over from the order-side pattern it's otherwise mirroring.

2. The SpendLimited message promises "partial results remain available," but nothing in the SDK can deliver them.
_raise_for_blocking_status unconditionally raises for SpendLimited before get_results() ever attempts job_job_id_download_results_get (rapidata_job.py:101-106, 284-287). There's also no job-side equivalent of RapidataOrder's preliminary_results=True / _get_preliminary_results() mechanism (rapidata_order.py:388-411), and no "preliminary results" endpoint exists at all in the generated job_api.py. So a user who hits this exception has no way, via this SDK, to actually retrieve the "partial results" the message references. Either the message should be softened (e.g., point to the dashboard/support instead of implying an SDK-retrievable result), or a preliminary-results path should be added for jobs to back up the claim.

Minor / nitpicks

• _get_job_failure_message and _get_review_reason are near-duplicates that each make their own job_job_id_get call, and _raise_for_blocking_status triggers yet another job_job_id_get right after _wait_for_status's own get_status() call already fetched the same job. Not a bug, but each blocking-status raise does 2 network round trips where 1 would do — could consolidate into a single job fetch.
• Docs and error-message wording are otherwise clear and match the actual field names (reviewReason, costWarning) coming from the generated models.

Tests

Good coverage for the paths that are implemented: ManualApproval with/without a reviewReason, SpendLimited (both get_results and display_progress_bar), the unchanged Completed happy path, and both branches of the cost-warning warning. The mocks correctly stand in for the generated pydantic models and the stacklevel=3 in _warn_if_cost_exceeds_balance correctly attributes the warning to the caller of assign_job.

Given finding #1, a test for the Paused state (mirroring the existing SpendLimited/ManualApproval tests) would be worth adding once that's fixed — right now there's no coverage for it because the production code doesn't handle it either.

As called out in the PR description, no CI job runs pytest yet, so this suite currently only runs when someone remembers to invoke it locally — worth wiring in as a fast follow.

Positive notes

• The hang fix for ManualApproval/SpendLimited and the cost-warning surfacing are both good, targeted fixes for real footguns (silent indefinite blocking on get_results()).
• Fallback to a generic "is being reviewed" message when reviewReason is null is handled correctly and is tested.
• warnings.warn usage, stacklevel, and the "job is still created" behavior are all correct and consistent with the PR's stated intent of not changing behavior otherwise.

[claude]

Review

Solid fix for a real footgun (jobs hanging forever in ManualApproval/SpendLimited), and the cost-warning surfacing is a nice, low-risk addition. A few notes below, nothing blocking.

What the PR does

• RapidataJob._wait_for_status (used by get_results) and display_progress_bar now recognize ManualApproval and SpendLimited as terminal-for-polling-purposes states and raise an informative Exception instead of looping forever.
• assign_job surfaces the create endpoint's optional costWarning via warnings.warn, without changing job-creation behavior.
• Polling now fetches the job GET payload once per iteration (_fetch_job) and reuses it for state/failure-message/review-reason instead of issuing a second request, which is a nice incidental efficiency win.
• Blocking-state comparisons use the generated AudienceJobState enum rather than raw strings, matching the RapidataOrder/OrderState convention.
• New tests/test_job_review_and_cost_warning.py covers the happy path, both blocking states, the reason/no-reason branch, and the warn/no-warn branch — good coverage of the new behavior.

Code quality / correctness

• _BLOCKING_STATUSES is a tuple of AudienceJobState enum members, while _wait_for_status compares against job.state.value (a plain str). This works today because AudienceJobState subclasses str, so "ManualApproval" in (AudienceJobState.MANUALAPPROVAL, ...) is True — but it's a bit fragile/implicit. Comparing job.state (the enum) directly, like display_progress_bar already does, would be more consistent within the same file (_wait_for_status mixes .value-string comparisons for target_statuses with enum-based _BLOCKING_STATUSES, whereas display_progress_bar compares the enum throughout).
• _raise_for_blocking_status raises a bare Exception in both branches, consistent with the rest of this file (e.g. the "has failed" errors), so this matches existing conventions rather than being a new issue — but a dedicated exception type (or types) for these two cases would let callers programmatically distinguish "needs review" from "needs funds" instead of parsing the message string. Worth considering as a future follow-up rather than blocking this PR.
• _wait_for_status's target_statuses: list[str] combined with an internal blocking check against enum members is a little inconsistent with RapidataOrder, which threads OrderState enums through its equivalent wait helper. Not a bug, just a minor inconsistency between the two "wait for status" implementations that a future refactor could reconcile.

Docs

• docs/audiences.md and docs/quickstart.md updates read well, and the anchor link (audiences.md#cost-warnings-and-jobs-under-review) matches the new heading.

Tests

• Good use of MagicMock to stand in for the OpenAPI client without touching the network; the catch_warnings/simplefilter("error") trick in test_assign_job_no_warning_when_cost_within_balance is a solid way to assert the negative case.
• As called out in the PR description, pytest isn't wired into CI yet (CI is pyright-only) — worth a fast follow so this suite doesn't silently bit-rot.
• No test exercises display_progress_bar's in-loop transition into SpendLimited/ManualApproval (only the initial pre-loop check is covered by test_display_progress_bar_raises_on_spend_limited, which returns the blocking state on the very first _fetch_job() call). A test where the first job_job_id_get call returns Running and a subsequent call returns SpendLimited would cover the loop-body branch that isn't currently exercised.

Security / performance

• No security concerns — no new external input parsing, and the warning message only echoes values already present in the trusted API response.
• The per-iteration single-fetch change is a small performance improvement over the previous double-fetch (get_status() + _get_job_failure_message()).

Overall: clean, well-tested, and appropriately scoped fix for a genuine hang bug. The enum/string mixing in _wait_for_status and the missing in-loop-transition test are the only things I'd consider addressing, and neither is blocking.

[LinoGiger]

why does this not just use a logger?

[LinoGiger]

since when do we have test files on the sdk? like i don't mind it in general. but please create a better folder structure than this ^^

[RapidPoseidon]

Thanks @LinoGiger — both addressed in 7ffa270 (on top of your 535d8d6 polling refactor, which I kept):

1. Logger, not warnings.warn. assign_job now surfaces the cost warning via logger.warning(...) (lazy %-args), matching the SDK's existing user-facing warning convention (e.g. rapidata_client.py). Docs wording updated accordingly.

2. Test folder structure. Split the flat file into a package layout mirroring src/:

   • tests/rapidata_client/job/test_rapidata_job.py — hang fixes + display_progress_bar + happy path
   • tests/rapidata_client/audience/test_audience_base.py — cost-warning emit / no-emit

   Kept your AudienceJobState(state) test-helper change so the enum-based blocking checks are exercised correctly.

pyright, black, pytest (7 passing) and mkdocs build all green locally.

[claude]

Review

Went through the diff in detail (rapidata_job.py, _audience_base.py, docs, and the new tests). Overall this is a solid, well-scoped fix — the hang-on-ManualApproval/SpendLimited bug is real and the fix is correct, the refactor to fetch the job once per poll iteration is a nice cleanup, and using the generated AudienceJobState enum instead of raw strings is the right call for consistency with RapidataOrder. I traced the enum/model wiring (AudienceJobState, ReviewReasonModel, CreateJobEndpointCostWarningModel, GetJobByIdEndpointOutput.review_reason) and everything lines up correctly with the generated models — no attribute-name or type mismatches.

Correctness

• _wait_for_status/display_progress_bar correctly stop polling and raise instead of looping forever on ManualApproval/SpendLimited. Comparing job.state.value (a plain str) against _BLOCKING_STATUSES (a tuple of AudienceJobState members) works because AudienceJobState is a str, Enum subclass, so ==/in fall back to string equality — worth double-checking that was intentional rather than incidental, since it's a bit implicit.
• _raise_for_blocking_status assumes any non-SpendLimited blocking status is ManualApproval (no elif check). That's true today given _BLOCKING_STATUSES only has two members, but it's a latent trap if a third blocking state is ever added — the code would silently mislabel it as "being reviewed." A short comment or explicit branch would make the invariant clearer.
• Good catch on avoiding the double-fetch (status check + separate fetch for the failure/review reason) — the single _fetch_job() per iteration is both simpler and cuts an API call per poll tick.

Docs

• docs/audiences.md and quickstart.md updates are accurate and the anchor link (#cost-warnings-and-jobs-under-review) matches the new heading. Good that the stale "blocks until complete" claim was corrected rather than left dangling.

Cost warning

• _warn_if_cost_exceeds_balance correctly no-ops when cost_warning is None, and logging via logger.warning (rather than warnings.warn) matches the SDK's existing user-facing-warning convention — good change based on review feedback.
• Minor: this new method sits between assign_job and find_jobs without a blank-line/visual break issue — no, actually formatting looks fine on inspection.

Tests

• Good coverage: both ManualApproval variants (with/without reviewReason), SpendLimited (including display_progress_bar), the cost-warning emission/no-emission paths, and the unchanged happy path. Mocking job_job_id_get and asserting on the raised message contents is a reasonable approach given there's no real backend to hit here.
• Since CI doesn't run pytest yet (per the PR description), consider following up soon so this suite doesn't silently bit-rot — a hang-fix regression here is exactly the kind of thing you want caught automatically.

Nit

• The hardcoded "This can take a couple of hours" in the ManualApproval message is a guess dressed up as a fact; if that number isn't backed by an SLA it might be worth softening (e.g. "can take some time") so it doesn't read as a commitment.

Nothing blocking — nice, focused fix with good test coverage and matching docs updates.

[claude]

Review

Solid fix for a real production issue (jobs stuck in ManualApproval/SpendLimited hanging get_results()/display_progress_bar() forever), with good test coverage and docs updates per the repo's CLAUDE.md conventions (generated AudienceJobState enum used instead of raw strings, docs kept in sync, tests added).

Correctness — looks right

• AudienceJobState/ReviewReasonModel are str, Enum subclasses, so the current_status in self._BLOCKING_STATUSES string/enum mix in _wait_for_status (rapidata_job.py:183) and the job.state in self._BLOCKING_STATUSES checks elsewhere work correctly — confirmed by reading the generated models.
• cost_warning.estimated_cost/available_balance/shortfall and job.review_reason field names match the generated CreateJobEndpointCostWarningModel / GetJobByIdEndpointOutput models, and both are Optional with sensible defaults, so older backends without these fields degrade gracefully to "no warning" / "no reason."
• Fetching the job once per poll iteration (_fetch_job()) and reusing it for state + failure message + review reason removes the double-GET the old code had in display_progress_bar.

One thing worth double-checking: PR description vs. implementation

The PR description says cost warnings are surfaced via a Python warnings.warn, but _warn_if_cost_exceeds_balance (_audience_base.py:109) actually calls logger.warning(...) — the SDK's own logger, not the stdlib warnings module. Functionally this is fine (default log level is WARNING, so it prints by default), but it means:

• pytest.warns() / warnings.catch_warnings() won't catch it — a user can't -W error on it or filter it the way they would a real UserWarning.
• Visibility depends on the SDK's logging config rather than Python's warnings machinery.

If a real warnings.warn(..., UserWarning) was the intent (matching the PR description), consider switching; if logger.warning is intentional (arguably more consistent with the rest of this class, which uses logger/managed_print everywhere), it'd be worth correcting the PR description so future readers aren't misled.

Minor / non-blocking

• _raise_for_blocking_status (rapidata_job.py:98) branches on SPENDLIMITED and then unconditionally treats everything else as ManualApproval. Safe today since it's only called when state in _BLOCKING_STATUSES (just these two values), but it's an implicit assumption — if a third blocking status is ever added, it'll silently get the "being reviewed" message. An explicit elif job.state == AudienceJobState.MANUALAPPROVAL (or an assert/else: raise AssertionError for unhandled states) would make that safer against future changes.
• Could annotate _raise_for_blocking_status as -> NoReturn (it never returns) for self-documentation.
• Test coverage: display_progress_bar's ManualApproval path isn't exercised (only SpendLimited is, at test_rapidata_job.py), and no test drives _wait_for_status's loop across a state transition (e.g. Running → Completed) since the mocks return a fixed status each call. Not a big deal given the logic is shared with the well-tested get_results path, but would close the gap.

Nits

• pyproject.toml/uv.lock additions (pytest as a dev dep) and the .vscode/settings.json switch from unittest to pytest runner look consistent with the new tests/ suite; no concerns there.
• Docs changes (audiences.md, quickstart.md) are clear and appropriately scoped — good adherence to the CLAUDE.md guidance about keeping docs focused and not repeating information.

Nothing here blocks the change — the core hang-fix and cost-warning logic are correct and well-tested. The warnings.warn vs logger.warning discrepancy is worth a quick confirmation of intent, and the other points are optional hardening for future maintainers.