"""The generation job runner, shared by the request path and the reaper. Every media generation is a **durable job**, not an HTTP call that happens to take a while. The row exists in ``queued`` before the vendor is contacted, so the states are: queued ──claim──> running ──> succeeded | failed ▲ │ └──── reaper ◄─────┘ (lease expired: previous runner died) Two callers, one code path: * the request handler spawns :func:`run_job` right after the row commits, so the employee sees progress immediately; * ``ai_labs.reap_generations`` calls the same function for any job whose lease has expired, which is what makes a backend restart mid-video a delay rather than a lost job. **Cooperative leasing** is what keeps those two from colliding. Claiming is a single conditional ``UPDATE`` on ``lease_expires_at``; only the writer whose ``rowcount`` is 1 proceeds. The lease is short (180s) and heartbeat- extended while polling, so a dead runner's job is picked up in minutes rather than after the job's whole deadline. **Artifacts are re-homed before the job is called done.** DashScope's result links expire ~24h out (measured), so a row that stored one would be a time-bomb: correct in review, broken for the employee next week. The bytes are pulled into Helix's own bucket and only the resulting object key is persisted. """ from __future__ import annotations import asyncio import logging import uuid from collections.abc import Awaitable, Callable from datetime import UTC, datetime, timedelta from typing import Any from sqlalchemy import select, update from sqlalchemy.ext.asyncio import AsyncSession from app.core.exceptions import AppException from app.core.storage import put_object_bytes from app.domains.ai_labs import attachments as attach from app.domains.ai_labs import catalog, usage from app.domains.ai_labs.models import ( GEN_FAILED, GEN_QUEUED, GEN_RUNNING, GEN_SUCCEEDED, GEN_TERMINAL_STATUSES, AiLabGeneration, ) from app.domains.ai_labs.settings import AiLabsConfig, resolve_config from app.domains.bailian.media_client import ( TASK_FAILED, TASK_RUNNING, TASK_SUCCEEDED, BailianMediaClient, ) logger = logging.getLogger(__name__) #: How long a claim is good for. Short on purpose: a runner that dies must #: not park its job for the job's whole deadline. Extended on every poll. LEASE_SECONDS = 180 #: Gap between lease heartbeats. #: #: The heartbeat is UNCONDITIONAL — it extends on every tick rather than #: only when the remaining lease drops below a threshold. (A #: ``LEASE_REFRESH_WHEN_BELOW = 60`` constant used to sit here describing #: the threshold policy; nothing ever read it, so it documented a #: mechanism this module does not have. Removed rather than implemented: #: unconditional is the safer of the two, and a knob nobody reads is a #: standing invitation to "fix" the loop to honour it.) #: #: A sixth of the lease, not a third. A swallowed heartbeat error costs one #: extension, so how many consecutive failures the lease can absorb IS the #: safety margin — and for a sync job the cost of the lease lapsing is the #: reaper re-submitting a paid render. Six chances instead of three, for #: one extra cheap UPDATE per minute. #: def heartbeat_interval(lease_seconds: float = 0.0) -> float: """Seconds between lease heartbeats for a lease of this length. Takes the lease length as an argument so a test that shortens ``LEASE_SECONDS`` to compress a timeline gets a proportionally shorter interval, and so the loop has no arithmetic of its own to drift from. """ lease = lease_seconds or LEASE_SECONDS return max(1.0, lease / 6) #: Consecutive heartbeat errors before the keeper stops heartbeating. #: With the interval at ``LEASE_SECONDS / 6`` this is still comfortably #: inside the lease, so the runner learns it is losing the lease *before* #: the row becomes claimable rather than after. MAX_HEARTBEAT_FAILURES = 3 #: Gap between vendor task polls. The vendor's own state changes on the #: order of tens of seconds (measured: 29s for a 480P/5s clip, 95s for #: 1080P/10s), so polling faster only spends request budget. POLL_INTERVAL_SECONDS = 6.0 #: Attempts a job gets before it is failed for good. A transport blip on #: submit is worth one retry; a vendor rejection is not, and is failed #: immediately regardless of this count. MAX_ATTEMPTS = 3 #: Consecutive transport failures tolerated *inside* one polling loop before #: the job is handed back to the queue. Polling a video for ten minutes #: across a rolling deploy of the vendor's gateway will drop a request or #: two; giving up on the first one turns a clip the vendor is about to #: finish into a failure the employee has to pay for again. The deadline #: stays authoritative, so this cannot extend a job past its own limit. MAX_CONSECUTIVE_POLL_ERRORS = 5 #: ``error_code`` written while a job waits to be retried. Not a terminal #: code: the reaper replaces it when the attempts run out, so it must be #: recognisable rather than merely truthy. RETRY_SENTINEL = "RETRYING" #: Transport-shaped vendor errors: worth another attempt, unlike a #: rejection of the request itself. _RETRYABLE_CODES = frozenset({ "BAILIAN_MEDIA_NETWORK", "BAILIAN_MEDIA_RESULT_UNAVAILABLE", }) #: Content type → file extension for the stored artifact. _EXTENSIONS = { "image/png": "png", "image/jpeg": "jpg", "image/webp": "webp", "video/mp4": "mp4", "audio/wav": "wav", "audio/mpeg": "mp3", } #: Fallback extension per capability output kind, for a vendor response #: whose ``Content-Type`` is missing or generic. _FALLBACK_EXTENSION = { catalog.OUTPUT_IMAGE: "png", catalog.OUTPUT_VIDEO: "mp4", catalog.OUTPUT_AUDIO: "wav", } _FALLBACK_CONTENT_TYPE = { catalog.OUTPUT_IMAGE: "image/png", catalog.OUTPUT_VIDEO: "video/mp4", catalog.OUTPUT_AUDIO: "audio/wav", } def _now() -> datetime: return datetime.now(UTC) def _as_utc(value: datetime | None) -> datetime | None: """Attach UTC to a naive timestamp read back from MySQL. ``DateTime(timezone=True)`` does not survive a MySQL round trip with its tzinfo, so an arithmetic comparison against an aware ``now()`` would raise. Every read of a stored timestamp in this module goes through here. """ if value is None: return None return value.replace(tzinfo=UTC) if value.tzinfo is None else value # ── leasing ──────────────────────────────────────────────────── async def _claim( db: AsyncSession, *, generation_id: uuid.UUID, token: str, ) -> bool: """Take the lease. ``True`` iff this caller now owns the job. The ``WHERE`` clause is the whole mechanism: a job is claimable only when it is not finished and either unleased or expired. Two runners racing produce exactly one ``rowcount == 1``. """ now = _now() stmt = ( update(AiLabGeneration) .where( AiLabGeneration.id == generation_id, AiLabGeneration.status.in_((GEN_QUEUED, GEN_RUNNING)), ( AiLabGeneration.lease_expires_at.is_(None) | (AiLabGeneration.lease_expires_at < now) ), ) .values( status=GEN_RUNNING, lease_owner=token, lease_expires_at=now + timedelta(seconds=LEASE_SECONDS), attempts=AiLabGeneration.attempts + 1, updated_at=now, ) ) result = await db.execute( stmt.execution_options(synchronize_session=False), ) await db.commit() # ``rowcount`` lives on ``CursorResult``, which is what an UPDATE # actually returns; the declared ``Result`` union does not carry it # (repo idiom — see ``security_awareness.campaign_service``). return (getattr(result, "rowcount", 0) or 0) == 1 async def _heartbeat( db: AsyncSession, *, generation_id: uuid.UUID, token: str, ) -> bool: """Extend the lease. ``False`` means someone else owns the job now. A runner that loses its lease must stop rather than keep writing: the new owner is already polling the same vendor task, and two writers would fight over the terminal state. """ now = _now() stmt = ( update(AiLabGeneration) .where( AiLabGeneration.id == generation_id, AiLabGeneration.lease_owner == token, ) .values( lease_expires_at=now + timedelta(seconds=LEASE_SECONDS), last_polled_at=now, updated_at=now, ) ) result = await db.execute( stmt.execution_options(synchronize_session=False), ) await db.commit() # ``rowcount`` lives on ``CursorResult``, which is what an UPDATE # actually returns; the declared ``Result`` union does not carry it # (repo idiom — see ``security_awareness.campaign_service``). return (getattr(result, "rowcount", 0) or 0) == 1 class _LeaseKeeper: """Hold a job's lease while a long ``await`` is in flight. The polling loop of an async job heartbeats between polls, but two stretches of work do not get to do that: a synchronous generation is one single vendor call (measured up to ~70s for a 4K image, and the artifact download follows it), and the success path of an async job downloads a clip that can be tens of megabytes. Both can outlast ``LEASE_SECONDS``, and an expired lease is claimable — the reaper would take a job whose vendor call is still running and **submit it a second time**, which bills the company twice and hands the employee whichever of the two results lands last. Runs on **its own session**. Sharing the caller's ``AsyncSession`` across two concurrent tasks corrupts it, and the whole point of this class is to write while the caller is blocked on something else. ``lost`` turns true **only** when a heartbeat matched no row — meaning another runner owns the job, or the employee cancelled it. A heartbeat that could not reach the database does NOT set it: "I cannot see the row" is a different claim from "the row is not mine", and only the latter justifies discarding a finished render. The keeper stops heartbeating after :data:`MAX_HEARTBEAT_FAILURES` consecutive errors and says so loudly, but the caller still attempts its terminal write, which is lease-gated and therefore safe either way. The caller checks the flag rather than being interrupted: aborting mid-``await`` would leave the vendor's work unaccounted for. """ def __init__(self, *, generation_id: uuid.UUID, token: str) -> None: self.generation_id = generation_id self.token = token self.lost = False self._task: asyncio.Task[None] | None = None async def _loop(self) -> None: from app.core.database import async_session_factory interval = heartbeat_interval(LEASE_SECONDS) failures = 0 while True: await asyncio.sleep(interval) try: async with async_session_factory() as own: if not await _heartbeat( own, generation_id=self.generation_id, token=self.token, ): self.lost = True return failures = 0 except asyncio.CancelledError: raise except Exception: # noqa: BLE001 — a keeper must not kill the job failures += 1 logger.warning( "ai_labs: lease heartbeat %s/%s failed for generation %s", failures, MAX_HEARTBEAT_FAILURES, self.generation_id, ) if failures >= MAX_HEARTBEAT_FAILURES: # Loud, but ``lost`` is deliberately NOT set. # # A failed heartbeat means "we could not reach the # database", which is not the same claim as "somebody # else owns this job" — and only the second one # justifies throwing work away. Setting ``lost`` here # made the caller skip its terminal write and DELETE # the finished artifact, so a 90-second DB blip during # a render that then succeeded turned into: artifact # discarded, row left ``running``, and a full-price # re-render on the next reap. # # There is nothing to protect by giving up early. # ``_finish`` is itself gated on ``lease_owner == # token``, so attempting it costs one no-op UPDATE if # the lease really did move, and saves the job if it # did not. Stop heartbeating (the connection is not # answering), keep the outcome. logger.error( "ai_labs: %s consecutive heartbeat failures on " "generation %s — the lease may lapse and the reaper " "may reclaim it. Still attempting the terminal " "write, which is lease-gated and therefore safe.", failures, self.generation_id, ) return async def __aenter__(self) -> _LeaseKeeper: self._task = asyncio.create_task( self._loop(), name=f"ai_labs.lease.{self.generation_id}", ) return self async def __aexit__(self, *_exc: object) -> None: if self._task is not None: self._task.cancel() try: await self._task except (asyncio.CancelledError, Exception): # noqa: BLE001 pass self._task = None async def _finish( db: AsyncSession, *, generation_id: uuid.UUID, token: str, status: str, artifacts: list[dict[str, Any]] | None = None, output_text: str | None = None, error_code: str | None = None, error_message: str | None = None, ) -> bool: """Write the terminal state, but only while still holding the lease.""" now = _now() values: dict[str, Any] = { "status": status, "finished_at": now, "lease_owner": None, "lease_expires_at": None, "updated_at": now, } if artifacts is not None: values["artifacts"] = artifacts if output_text is not None: values["output_text"] = output_text if error_code is not None: values["error_code"] = error_code[:64] if error_message is not None: values["error_message"] = error_message[:1000] stmt = ( update(AiLabGeneration) .where( AiLabGeneration.id == generation_id, AiLabGeneration.lease_owner == token, ) .values(**values) ) result = await db.execute( stmt.execution_options(synchronize_session=False), ) await db.commit() # ``rowcount`` lives on ``CursorResult``, which is what an UPDATE # actually returns; the declared ``Result`` union does not carry it # (repo idiom — see ``security_awareness.campaign_service``). return (getattr(result, "rowcount", 0) or 0) == 1 async def _release_for_retry( db: AsyncSession, *, generation_id: uuid.UUID, token: str, reason: str, ) -> None: """Hand the job back to the queue for another attempt.""" now = _now() stmt = ( update(AiLabGeneration) .where( AiLabGeneration.id == generation_id, AiLabGeneration.lease_owner == token, ) .values( status=GEN_QUEUED, lease_owner=None, lease_expires_at=None, error_code=RETRY_SENTINEL, error_message=reason[:1000], updated_at=now, ) ) await db.execute(stmt.execution_options(synchronize_session=False)) await db.commit() async def _reset_session( db: AsyncSession, *, generation_id: uuid.UUID, ) -> None: """Make the session writable again before recording a failure. A failed ``execute``/``flush`` leaves the transaction aborted: every subsequent statement on it raises, including the one that would record the reason. Rolling back first is the only way the terminal write can land — and a job whose failure was never written is worse than a job that failed, because the employee sees a spinner instead of an error. Never raises: if even the rollback fails, the connection is gone and the reaper will pick the row up when its lease expires. """ try: await db.rollback() except Exception as exc: # noqa: BLE001 — savepoint-exempt: last resort logger.warning( "ai_labs: rollback failed while failing generation %s: %s", generation_id, exc, ) async def _record_task_id( db: AsyncSession, *, generation_id: uuid.UUID, token: str, task_id: str, ) -> None: """Persist the vendor task id **immediately** after submitting. This is the difference between a crash costing a poll and a crash costing a generation: without the id on the row, a resumed job cannot find the work the vendor is already doing and would submit a second one. """ now = _now() stmt = ( update(AiLabGeneration) .where( AiLabGeneration.id == generation_id, AiLabGeneration.lease_owner == token, ) .values( vendor_task_id=task_id[:128], started_at=AiLabGeneration.started_at, updated_at=now, ) ) await db.execute(stmt.execution_options(synchronize_session=False)) await db.commit() # ── vendor request construction ──────────────────────────────── async def _build_vendor_input( db: AsyncSession, *, generation: AiLabGeneration, capability: catalog.GenerationCapability, ) -> tuple[dict[str, Any], dict[str, Any]]: """``(input, parameters)`` for one job. The employee's ``prompt`` is placed verbatim. ``prompt_extend`` is pinned off in every capability's ``fixed_params`` — that flag is the vendor rewriting the prompt, and this product does not do that. """ inputs = await attach.load( db, object_keys=list(generation.input_object_keys or []), owner_id=generation.person_id, allowed_kinds=( {attach.KIND_IMAGE} if capability.input_kind == catalog.INPUT_IMAGE else {attach.KIND_AUDIO} if capability.input_kind == catalog.INPUT_AUDIO else set() ), ) stored_params = dict(generation.params or {}) parameters: dict[str, Any] = dict(capability.fixed_params) for param in capability.params: if param.name in stored_params: parameters[param.wire_name] = stored_params[param.name] key = capability.key if key == catalog.CAP_IMAGE: return ( {"messages": [{ "role": "user", "content": [{"text": generation.prompt}], }]}, parameters, ) if key == catalog.CAP_IMAGE_EDIT: content: list[dict[str, Any]] = [ {"image": item.data_uri} for item in inputs ] content.append({"text": generation.prompt}) return ({"messages": [{"role": "user", "content": content}]}, parameters) if key == catalog.CAP_VIDEO: vendor_input: dict[str, Any] = {"prompt": generation.prompt} if inputs: # One reference image turns text-to-video into image-to-video on # the same model and endpoint (verified 2026-09-03 with an # inline data URI). vendor_input["img_url"] = inputs[0].data_uri return (vendor_input, parameters) if key == catalog.CAP_SPEECH: vendor_input = {"text": generation.prompt} voice = stored_params.get("voice") if voice: vendor_input["voice"] = voice return (vendor_input, {}) if key == catalog.CAP_TRANSCRIBE: return ( {"messages": [{ "role": "user", "content": [{"audio": inputs[0].data_uri}], }]}, parameters, ) raise AppException( code="AI_LABS_CAPABILITY_UNSUPPORTED", message=f"No vendor request shape is defined for '{key}'.", status_code=500, ) # ── artifact persistence ─────────────────────────────────────── async def _store_artifacts( db: AsyncSession, *, client: BailianMediaClient, generation: AiLabGeneration, capability: catalog.GenerationCapability, urls: list[str], ) -> list[dict[str, Any]]: """Download each vendor artifact into Helix's bucket. Returns the persisted descriptors. Raises if nothing could be stored — a job with no readable output is a failure, not a success with an empty gallery. """ stored: list[dict[str, Any]] = [] for index, url in enumerate(urls[: max(1, capability.max_outputs)]): raw, content_type = await client.fetch_result(url) if not content_type or content_type == "application/octet-stream": content_type = _FALLBACK_CONTENT_TYPE.get( capability.output_kind, "application/octet-stream", ) extension = _EXTENSIONS.get(content_type) or _FALLBACK_EXTENSION.get( capability.output_kind, "bin", ) created = _as_utc(generation.created_at) or _now() relative_key = ( f"ai_labs_output/{created:%Y/%m/%d}/{generation.person_id}/" f"{generation.id}-{index}.{extension}" ) object_key = await put_object_bytes( db, relative_key=relative_key, content=raw, content_type=content_type, ) stored.append({ "object_key": object_key, "content_type": content_type, "size_bytes": len(raw), }) if not stored: raise AppException( code="AI_LABS_NO_ARTIFACT", message="The vendor reported success but produced no file.", status_code=502, ) return stored # ── the runner ───────────────────────────────────────────────── async def run_job(generation_id: uuid.UUID) -> None: """Drive one generation to a terminal state. Never raises. Opens its own session: the request that created the row has already returned by the time this runs, and the reaper hands over only an id. Every exception is converted into a persisted failure, because a job that dies without writing its reason is indistinguishable to the employee from one that is still working. """ token = uuid.uuid4().hex # Looked up at call time so the test harness's patched, schema-isolated # factory applies — see the note in ``chat._persist_answer``. from app.core.database import async_session_factory try: async with async_session_factory() as db: config = await resolve_config(db) if not config.available: # Nothing to do and nothing to blame the employee for: the # feature was switched off (or the key removed) between the # click and the run. await _fail_unclaimed( db, generation_id, code="AI_LABS_UNAVAILABLE", message="AI Labs is not currently configured.", ) return if not await _claim(db, generation_id=generation_id, token=token): # Someone else owns it, or it already finished. Both fine. return generation = await db.get(AiLabGeneration, generation_id) if generation is None: return await _run_claimed( db, generation=generation, token=token, config=config, ) except Exception: # noqa: BLE001 — a runner must not take the loop down logger.exception( "ai_labs: generation %s crashed outside its own error handling", generation_id, ) await _fail_detached( generation_id, code="AI_LABS_INTERNAL", message="The job failed unexpectedly.", # This runner may still hold the lease it claimed before # crashing, so pass the token — otherwise the guarded write # matches nothing and the reason is never recorded. token=token, ) async def _run_claimed( db: AsyncSession, *, generation: AiLabGeneration, token: str, config: AiLabsConfig, ) -> None: capability = catalog.capability(generation.capability) if capability is None: # Copied out before the write: ``_finish`` commits, which expires # every loaded instance, and reading them afterwards would raise # ``MissingGreenlet``. removed_person = generation.person_id removed_capability = generation.capability removed_model = generation.model await _finish( db, generation_id=generation.id, token=token, status=GEN_FAILED, error_code="AI_LABS_CAPABILITY_REMOVED", error_message=( "This capability is no longer offered, so the job cannot be " "run." ), ) # ``requests=0``: the vendor was never called, so nothing was # billed — but the job did end badly, and this exit used to record # nothing at all. A catalogue entry pulled out from under a queued # job is precisely the kind of failure an operator has to be able # to see, and it was the one class of failure invisible in every # count on the admin console. await usage.record( person_id=removed_person, capability_key=removed_capability, model=removed_model, timezone_name=config.reporting_timezone, requests=0, failures=1, error_code="AI_LABS_CAPABILITY_REMOVED", ) return person_id = generation.person_id generation_id = generation.id model = generation.model attempts = generation.attempts existing_task = (generation.vendor_task_id or "").strip() if generation.started_at is None: generation.started_at = _now() await db.commit() client = BailianMediaClient(config.api_key, region=config.region) # ── how spend gets counted ──────────────────────────────── # # **The charge is recorded when it is incurred, not when the job # reaches a terminal state.** Those are different moments, and every # previous version of this conflated them: # # * gating on ``_finish``'s bool meant a cancel detected by # ``_finish`` recorded nothing, while the same event detected by # the lease keeper did — and a sync render shorter than one # heartbeat is *always* detected by ``_finish``; # * replacing that with a "did this run bill?" predicate evaluated at # ENTRY was worse: it is right about who paid, but the paths that # leave before the terminal write never reach the recording at all. # An async runner that submitted (charged) and then lost its lease # returns through ``payload is None``, and the reaper that takes # over sees an existing task id and declines to record — so **every # async generation that changed hands, which is the documented # recovery path after any restart mid-render, was delivered to the # employee with no usage row, no cost attribution and no quota # consumed at all.** # # So: ``_note_charge`` fires the moment the vendor accepts work, from # inside the call that accepts it. Whoever pays, records — immediately, # before anything can go wrong afterwards. The terminal writes then add # only the outcome (artifacts, media seconds, failure), never another # ``requests``, so the counters accumulate to exactly one charge. # # ``requests`` therefore means "the vendor accepted billable work" and # ``failures`` means "the job ended badly"; a job that fails before the # vendor accepts anything records the second without the first. charged = False async def _note_charge() -> None: """The vendor has accepted work. Count it now, once.""" nonlocal charged if charged: return charged = True await usage.record( person_id=person_id, capability_key=capability.key, model=model, timezone_name=config.reporting_timezone, requests=1, ) async def _note_outcome( outcome: dict[str, Any] | None, *, failed: bool = False, error_code: str | None = None, ) -> None: """Add what the run produced. Never counts a second request. Books against the local day it is called on, which for a job that spans midnight is a different row from its charge — a ten-minute video submitted at 23:59 counts its request on one day and its seconds on the next. Accepted: both numbers are right, the attribution of one of them is a few minutes early, and pinning the outcome to the charge's day would mean carrying that date through every exit for a boundary case nobody reads per-day. """ data = outcome or {} artifacts = len(data.get("artifacts") or []) media_seconds = int(data.get("media_seconds") or 0) if not failed and not artifacts and not media_seconds: return await usage.record( person_id=person_id, capability_key=capability.key, model=model, timezone_name=config.reporting_timezone, requests=0, failures=1 if failed else 0, artifacts=artifacts, media_seconds=media_seconds, # The reason dimension for ``ai_labs_failures_daily``. It is # the code that was just written to the row's ``error_code``, # so the counters and the employee's own view of the job agree # on why it ended. Ignored when ``failed`` is false. error_code=error_code if failed else None, ) # ``best_effort`` is the wrong tool for this shape: it protects an # *outer* transaction from a fallible inner step, whereas the failure # being handled here is the outer work itself, and what must survive is # the write that records why. Both handlers therefore begin with # ``_reset_session`` — an explicit ROLLBACK — because an exception that # originated in a failed DB write leaves the transaction aborted, and # ``_finish`` would then fail too, leaving the job ``running`` behind a # held lease until it expired. Every value the handlers need was copied # into a local **before** this try, since a rollback expires every # loaded ORM instance and reading ``generation.model`` afterwards would # raise ``MissingGreenlet``. # savepoint-exempt: recovery is rollback + fresh Core UPDATE (see above) try: # Explicitly optional. ``_run_async`` returns ``None`` for "still # running, but not mine any more", and inferring the type from the # sync branch alone is what let a ``payload.get(...)`` past review. payload: dict[str, Any] | None async with _LeaseKeeper( generation_id=generation_id, token=token, ) as keeper: if capability.transport == catalog.TRANSPORT_SYNC: payload = await _run_sync( db, client=client, generation=generation, capability=capability, on_charge=_note_charge, ) else: payload = await _run_async( db, client=client, generation=generation, capability=capability, token=token, existing_task_id=existing_task, on_charge=_note_charge, ) if keeper.lost: # The lease went out from under us while the vendor was # working — cancelled by the employee, or taken over. The # terminal write below would match no row anyway. logger.info( "ai_labs: generation %s lost its lease during the vendor " "call — leaving the terminal state to its owner", generation_id, ) # **The spend still happened.** The vendor rendered and # billed whether or not anyone is still waiting, so the # counter records it — ``cancel_generation`` cannot, because # it cannot know whether the vendor had already been called. # # Through ``_note_outcome``, which tolerates a ``None`` # payload. # ``_run_async`` returns ``None`` for "still running, but # not mine any more", and BOTH that and ``keeper.lost`` are # set by the same condition — a heartbeat matching no row — # detected by two watchers with no ordering between them. # Reaching into ``payload`` directly here raised # ``AttributeError`` whenever the keeper won that race on an # async job: swallowed by the outer handler, so the usage row # this branch exists to write was never written and the paid # artifact leaked. mypy had flagged the type; mypy is in no # CI workflow. await _note_outcome(payload) # The vendor may already have produced a file and # ``_store_artifacts`` may already have re-homed it. Nothing # will reference those objects now. await _discard_artifacts(payload, generation_id=generation_id) return except AppException as exc: await _reset_session(db, generation_id=generation_id) retryable = exc.code in _RETRYABLE_CODES if retryable and attempts < MAX_ATTEMPTS: await _release_for_retry( db, generation_id=generation_id, token=token, # Safe to persist: every ``AppException`` that reaches here # carries one of our own sentences plus at most a vendor # error *code* — the prose is logged, never relayed. reason=exc.message, ) return # The FAILURE is recorded only by the runner whose terminal write # landed — otherwise a job another runner is about to finish # successfully would also be counted as failed. The spend itself is # not recorded here at all: it was already counted at charge time, # by whichever call the vendor accepted work from. if await _finish( db, generation_id=generation_id, token=token, status=GEN_FAILED, error_code=exc.code, error_message=exc.message, ): await _note_outcome(None, failed=True, error_code=exc.code) return except Exception: # noqa: BLE001 — persist, never vanish logger.exception("ai_labs: generation %s failed", generation_id) await _reset_session(db, generation_id=generation_id) # ``str(exc)`` is not relayed: an unexpected exception's text can # carry a vendor body, a crafted redirect URL, or a fragment of the # employee's own input, and this string is persisted and returned. # The full detail is in the ``logger.exception`` above. if await _finish( db, generation_id=generation_id, token=token, status=GEN_FAILED, error_code="AI_LABS_INTERNAL", error_message=( "Something went wrong on our side. Please try again." ), ): await _note_outcome( None, failed=True, error_code="AI_LABS_INTERNAL", ) return if payload is None: # Async job still running and the lease was lost — the new owner is # driving it. Leave the row alone. return # The artifacts are already in Helix's bucket and the vendor has already # been paid, so the only remaining failure mode worth guarding is the # terminal write itself. # # **Whether a retry is free depends on the transport, and only the # async half is.** An async job re-polls its stored ``vendor_task_id``, # which costs nothing. A sync job has no task id: the reaper re-runs # ``_run_sync``, which calls the vendor again — a second full-price # render for one click. The earlier version released BOTH for retry and # justified it with the async argument alone, which is the same # asymmetry ``_finish_sync_output`` was split out to fix, left open on # its twin path. # # The handler's first act is an explicit ROLLBACK (``_reset_session``), # and the recovery write that follows is a fresh Core UPDATE on the # reset session — the same shape, and the same reason, as the guarded # block further up this function. # savepoint-exempt: recovery is rollback + fresh Core UPDATE (see above) try: finished = await _finish( db, generation_id=generation_id, token=token, status=GEN_SUCCEEDED, artifacts=payload.get("artifacts"), output_text=payload.get("output_text"), ) except Exception: # noqa: BLE001 — savepoint-exempt: rollback + requeue logger.exception( "ai_labs: could not store the result of generation %s", generation_id, ) await _reset_session(db, generation_id=generation_id) if capability.transport == catalog.TRANSPORT_SYNC: # Fail it. The employee resubmits deliberately if they still # want it, rather than the company paying twice for a render # nobody asked to repeat. await _finish( db, generation_id=generation_id, token=token, status=GEN_FAILED, error_code="AI_LABS_OUTPUT_UNAVAILABLE", error_message=( "The model produced a result but it could not be saved. " "Please submit it again." ), ) # The row is terminally failed, so nothing will ever reference # the bytes already in the bucket and — ``ai_labs_output/`` # having no grant ledger row — no sweep can find them. And the # render was paid for, so it belongs in the counters like every # other terminal failure here. await _note_outcome( None, failed=True, error_code="AI_LABS_OUTPUT_UNAVAILABLE", ) await _discard_artifacts(payload, generation_id=generation_id) else: await _release_for_retry( db, generation_id=generation_id, token=token, reason="The result could not be saved. Retrying.", ) return # Only the outcome: the charge was counted the moment the vendor took # the work, by whichever runner paid for it. await _note_outcome(payload) if not finished: # The row moved on, so the objects this runner stored are # unreferenced and uncollectable — same as ``keeper.lost``. await _discard_artifacts(payload, generation_id=generation_id) async def _run_sync( db: AsyncSession, *, client: BailianMediaClient, generation: AiLabGeneration, capability: catalog.GenerationCapability, on_charge: Callable[[], Awaitable[None]], ) -> dict[str, Any]: """One synchronous multimodal call, plus artifact storage. ``on_charge`` is awaited the moment the vendor returns a result — the moment the money is spent — so the spend is recorded even if everything after this point fails or the lease moves on. """ vendor_input, parameters = await _build_vendor_input( db, generation=generation, capability=capability, ) # Loading the attachments opened a read transaction, and the vendor call # below takes tens of seconds. Committing first returns the connection # to the pool instead of parking it idle across an external call — # MySQL/TiDB drops such connections and the *next* statement fails with # a 2013 that looks nothing like its cause. # # ``generation`` IS still read after this commit (here and inside # ``_store_artifacts``), and that is safe only because the session # factory sets ``expire_on_commit=False`` — an earlier version of this # comment claimed the values were all in locals, which they are not. # If that factory setting ever changes, this is one of the places that # breaks, with a ``MissingGreenlet`` from a lazy refresh. await db.commit() envelope = await client.multimodal_generate( model=generation.model, input_=vendor_input, parameters=parameters or None, ) # Paid. Count it here rather than at a terminal state: a cancel, a lost # lease or a failed download after this line must not make the charge # disappear from the counters. await on_charge() # **Everything below this line happens after the vendor was paid.** # # A transport failure while DOWNLOADING the result must not be retried # the way a transport failure while SUBMITTING is. For a sync # capability there is no task id to resume, so ``_release_for_retry`` # sends the job back to the reaper and ``_run_sync`` calls # ``multimodal_generate`` again — a second and third paid generation # for one click, when the first one demonstrably succeeded and only # its download failed. The async family is immune because it re-polls # a stored ``vendor_task_id``, which costs nothing. # # So the retryable codes are re-labelled here. The employee gets a # clear failure on a job the company paid for exactly once, which is # the better of the two bad outcomes. try: return await _finish_sync_output( db, client=client, generation=generation, capability=capability, envelope=envelope, ) except AppException as exc: if exc.code not in _RETRYABLE_CODES: raise logger.warning( "ai_labs: generation %s produced output that could not be " "stored (%s); not retrying — the vendor has already been paid " "for this render", generation.id, exc.code, ) raise AppException( code="AI_LABS_OUTPUT_UNAVAILABLE", message=( "The model produced a result but it could not be saved. " "Please submit it again." ), status_code=502, ) from exc async def _finish_sync_output( db: AsyncSession, *, client: BailianMediaClient, generation: AiLabGeneration, capability: catalog.GenerationCapability, envelope: dict[str, Any], ) -> dict[str, Any]: """Extract and store the output of one synchronous call. Split out so :func:`_run_sync` can re-label a retryable failure that happens on this side of the vendor call — see the note at its call site. """ if capability.output_kind == catalog.OUTPUT_TEXT: text = client.extract_text(envelope) if not text: raise AppException( code="AI_LABS_EMPTY_TRANSCRIPT", message=( "No speech was recognised in that recording." ), status_code=502, ) return {"output_text": text, "artifacts": []} if capability.output_kind == catalog.OUTPUT_AUDIO: url = client.extract_audio_url(envelope) urls = [url] if url else [] else: urls = client.extract_media_urls(envelope) if not urls: raise AppException( code="AI_LABS_NO_ARTIFACT", message="The vendor returned no file for this request.", status_code=502, ) artifacts = await _store_artifacts( db, client=client, generation=generation, capability=capability, urls=urls, ) return {"artifacts": artifacts} async def _run_async( db: AsyncSession, *, client: BailianMediaClient, generation: AiLabGeneration, capability: catalog.GenerationCapability, token: str, existing_task_id: str, on_charge: Callable[[], Awaitable[None]], ) -> dict[str, Any] | None: """Submit (if needed) and poll a vendor task to its terminal state. ``None`` means "still running, but this runner no longer holds the lease" — the caller must leave the row for whoever does. ``on_charge`` is awaited immediately after a SUBMIT, and only then: a run that resumes an existing task id re-polls work already paid for and must not count a second charge. Recording at submit rather than at the terminal state is what makes the charge survive this function returning ``None``, which is exactly what happens when the job changes hands mid-render — the documented recovery path after any restart. """ task_id = existing_task_id if not task_id: vendor_input, parameters = await _build_vendor_input( db, generation=generation, capability=capability, ) # Release the read transaction the attachment load opened before # going out to the vendor — see the note in ``_run_sync``. Safe to # keep using ``generation`` afterwards: the session factory sets # ``expire_on_commit=False``. await db.commit() task_id = await client.submit_task( path=capability.task_path or "", model=generation.model, input_=vendor_input, parameters=parameters or None, ) # Charge first, task id second. Both writes matter, but they fail # differently: if ``_record_task_id`` goes down and the charge has # already been counted, the re-submit counts its own second charge # and the books stay right. In the other order a DB blip between # the two loses a charge the vendor did make and nothing ever # counts it. await on_charge() await _record_task_id( db, generation_id=generation.id, token=token, task_id=task_id, ) started = _as_utc(generation.started_at) or _now() deadline = started + timedelta(seconds=capability.deadline_seconds) poll_errors = 0 while True: try: verdict = await client.get_task(task_id, model=generation.model) except AppException as exc: # A dropped poll is not a failed generation. The vendor is # still rendering the clip we already paid for, and its task id # is on the row, so the only thing at stake is this loop. Give # up only after several consecutive failures — or immediately # if the vendor rejected the request rather than dropping it. if exc.code not in _RETRYABLE_CODES: raise poll_errors += 1 if poll_errors >= MAX_CONSECUTIVE_POLL_ERRORS or _now() >= deadline: raise logger.warning( "ai_labs: poll %s/%s failed for generation %s (task %s): %s", poll_errors, MAX_CONSECUTIVE_POLL_ERRORS, generation.id, task_id, exc.code, ) await asyncio.sleep(POLL_INTERVAL_SECONDS) continue poll_errors = 0 status = verdict["status"] if status == TASK_SUCCEEDED: artifacts = await _store_artifacts( db, client=client, generation=generation, capability=capability, urls=verdict["media_urls"], ) vendor_usage = verdict.get("usage") or {} return { "artifacts": artifacts, "media_seconds": _seconds_from_usage(vendor_usage), } if status == TASK_FAILED: vendor_code = str(verdict.get("error_code") or "").strip() # The vendor's own prose is NOT persisted. This is the async # family — video and the async image models — which is exactly # where a content-policy rejection happens, and DashScope # quotes the offending input back: # ``"Input data may contain inappropriate content: ''"``. ``_run_claimed`` writes ``exc.message`` into # ``ai_labs_generations.error_message``, which the API returns, # so relaying it puts a piece of the employee's prompt in a # durable column on the one surface that promises nobody reads # what they typed. The sync path was sanitised in the first # round; this branch bypassed it because ``get_task`` returns # the verdict instead of raising. logger.warning( "ai_labs: vendor failed generation %s (task %s) code=%s: %s", generation.id, task_id, vendor_code or "?", str(verdict.get("error_message") or "")[:400], ) raise AppException( code=vendor_code or "AI_LABS_VENDOR_FAILED", message=( "The model provider could not complete this generation" + (f" (code {vendor_code[:64]})." if vendor_code else ".") ), status_code=502, ) if status != TASK_RUNNING: # pragma: no cover — defensive raise AppException( code="AI_LABS_VENDOR_STATE", message=f"Unexpected vendor task state: {status!r}", status_code=502, ) if _now() >= deadline: raise AppException( code="AI_LABS_TIMEOUT", message=( "The vendor did not finish this generation within " f"{capability.deadline_seconds} seconds." ), status_code=504, ) await asyncio.sleep(POLL_INTERVAL_SECONDS) if not await _heartbeat( db, generation_id=generation.id, token=token, ): logger.info( "ai_labs: generation %s lease lost — another runner has it", generation.id, ) return None def _seconds_from_usage(vendor_usage: dict[str, Any]) -> int: """Produced media seconds, from whichever key the vendor used.""" for key in ("output_video_duration", "duration", "video_duration"): value = vendor_usage.get(key) try: if value is not None: return max(0, int(float(value))) except (TypeError, ValueError): continue return 0 # ── failure paths that do not hold a lease ───────────────────── async def _fail_unclaimed( db: AsyncSession, generation_id: uuid.UUID, *, code: str, message: str, token: str | None = None, ) -> None: """Fail a job we never claimed (config gone, capability withdrawn). The lease condition is the missing half of this module's leasing discipline. Every other write keys on ``lease_owner == token``; this one keyed only on "not terminal", so a runner reaching the last-resort failure path *after* its lease had moved on would stamp ``failed`` over the new owner's ``running`` row — and the new owner's ``_finish`` would then match nothing, recording a generation that actually succeeded as a failure. The weakest gate is the real boundary, so this one now matches the others. It writes a row that is **unleased, expired, or leased to ``token``**. Passing the token matters on the crash path: the runner did claim the job, so without it the failure would not be recorded at all and the employee would watch a spinner until the lease expired. """ now = _now() conditions = [ AiLabGeneration.id == generation_id, AiLabGeneration.status.notin_(tuple(GEN_TERMINAL_STATUSES)), ] lease_is_free = AiLabGeneration.lease_owner.is_(None) | ( AiLabGeneration.lease_expires_at < now ) conditions.append( lease_is_free if token is None else (lease_is_free | (AiLabGeneration.lease_owner == token)) ) await db.execute( update(AiLabGeneration) .where(*conditions) .values( status=GEN_FAILED, error_code=code[:64], error_message=message[:1000], finished_at=now, lease_owner=None, lease_expires_at=None, updated_at=now, ) # See the note on the other lease writes: without this, SQLAlchemy # evaluates the WHERE clause in Python against any loaded instance # and the naive/aware timestamp comparison raises. .execution_options(synchronize_session=False), ) await db.commit() async def _fail_detached( generation_id: uuid.UUID, *, code: str, message: str, token: str | None = None, ) -> None: """Last-resort failure write on its own session. Reached when the outer handler caught something that may have left the original session unusable — a poisoned transaction cannot be the thing that records why. """ from app.core.database import async_session_factory try: async with async_session_factory() as own: await _fail_unclaimed( own, generation_id, code=code, message=message, token=token, ) except Exception: # noqa: BLE001 logger.exception( "ai_labs: could not record failure for generation %s", generation_id, ) async def _discard_artifacts( payload: dict[str, Any] | None, *, generation_id: uuid.UUID, ) -> None: """Delete artifacts this runner stored but did not get to reference. Reached when the work succeeded at the vendor and the bytes were re-homed into Helix's bucket, but the terminal write did not land: the employee cancelled, or another runner took the lease. Either way **these** objects are referenced by nothing, and nothing else will collect them — ``storage.orphan_upload_cleanup`` walks the ``storage_upload_grants`` ledger, and ``ai_labs_output/`` keys were never granted because the executor writes them itself. **A key is only deleted if the row does not name it.** Output keys are deterministic — ``ai_labs_output/{date}/{person}/{generation}-{i}.{ext}`` — so a runner that took over writes the *same* key. Deleting blindly would remove the object the winning runner had just stored and leave a ``succeeded`` row pointing at nothing, which is worse than the leak this function exists to prevent. Re-reading the row is also what makes the deterministic-key case a no-op rather than a hazard: the winner's key is in ``artifacts``, so it is skipped. Best effort and never raises: the job's state is already correct, and a bucket hiccup here is a storage-cost problem, not a reason to turn a settled job into an error. """ from app.core.database import async_session_factory from app.core.storage import delete_object keys = [ key for item in ((payload or {}).get("artifacts") or []) if isinstance(key := (item or {}).get("object_key"), str) and key ] if not keys: return try: async with async_session_factory() as own: row = await own.get(AiLabGeneration, generation_id) referenced = { key for item in ((row.artifacts if row else None) or []) if isinstance(key := (item or {}).get("object_key"), str) } except Exception: # noqa: BLE001 — cannot read the row, so delete none logger.warning( "ai_labs: could not check which artifacts generation %s " "references; leaving all of them in place", generation_id, ) return for key in keys: if key in referenced: # The row names this object — whoever won the race stored it # (at the same deterministic key), and it is live. continue try: async with async_session_factory() as own: await delete_object(own, object_key=key) except Exception: # noqa: BLE001 — cleanup must not raise logger.warning( "ai_labs: could not delete unreferenced artifact %s", key, ) # ── reaper support ───────────────────────────────────────────── async def find_reapable( db: AsyncSession, *, limit: int = 50, ) -> list[uuid.UUID]: """Ids of unfinished jobs whose lease is free — oldest first. Ordering by ``created_at`` matters: a burst that outlives one deploy should drain in the order employees submitted it, not newest-first. """ now = _now() stmt = ( select(AiLabGeneration.id) .where( AiLabGeneration.status.in_((GEN_QUEUED, GEN_RUNNING)), ( AiLabGeneration.lease_expires_at.is_(None) | (AiLabGeneration.lease_expires_at < now) ), ) .order_by(AiLabGeneration.created_at) .limit(limit) ) return list((await db.execute(stmt)).scalars().all()) __all__ = [ "LEASE_SECONDS", "MAX_ATTEMPTS", "POLL_INTERVAL_SECONDS", "find_reapable", "run_job", ]