cloud_tasks.py

"""Cloud Tasks dispatch + OIDC verification for the v2 sync pipeline.

The /v2/sync-local-files fast path enqueues one named task per sync job;
Cloud Tasks POSTs it back to /v2/sync-jobs/run on the backend-sync service
with an OIDC token minted for SYNC_TASKS_INVOKER_SA.

All functions fail closed when the SYNC_TASKS_* env vars are unset: enqueue
raises and verification returns 403 — the handler ships in the shared image
to services that must never accept task traffic. A caller that has already
staged audio must not start an inline worker after an enqueue exception: a
lost create-task acknowledgement can mean the deterministic named task exists.
"""

import json
import logging
import hashlib
import os
import uuid
from typing import Any, Dict, NamedTuple, Optional

from fastapi import HTTPException, Request
from google.api_core.exceptions import AlreadyExists, NotFound
from google.auth.transport import requests as google_auth_requests
from google.cloud import tasks_v2
from google.oauth2 import id_token
from google.protobuf import duration_pb2

from utils.log_sanitizer import sanitize

logger = logging.getLogger(__name__)

# Must match the queue's dispatchDeadline and the handler's request timeout
# (HTTP_SYNC_JOBS_RUN_TIMEOUT); see the run-lock TTL invariant in sync_jobs.py.
DISPATCH_DEADLINE_SECONDS = 1500

# Shared by the production admission boundary and the hermetic recorder. A
# recorder-local allowlist previously rejected new durable fields only after
# admission had already returned 202.
SYNC_JOB_TASK_PAYLOAD_KEYS = frozenset(
    {
        'schema_version',
        'job_id',
        'uid',
        'raw_blob_paths',
        'source',
        'should_lock',
        'conversation_id',
        'geolocation',
        'client_device_id',
        'client_platform',
        'enqueued_at',
        'lane',
        'capture_time_trust',
        'recording_age_seconds',
        'content_id',
        'ledger_fence_mode',
    }
)

_tasks_client: Optional[tasks_v2.CloudTasksClient] = None
_google_auth_request: Optional[google_auth_requests.Request] = None


class AccountDeletionTaskAuthentication(NamedTuple):
    """Verified Cloud Tasks identity plus its narrowly scoped audience lane."""

    retry_count: int


def _get_tasks_client() -> tasks_v2.CloudTasksClient:
    global _tasks_client
    if _tasks_client is None:
        _tasks_client = tasks_v2.CloudTasksClient()
    return _tasks_client


def _get_auth_request() -> google_auth_requests.Request:
    global _google_auth_request
    if _google_auth_request is None:
        _google_auth_request = google_auth_requests.Request()
    return _google_auth_request


def _handler_url() -> str:
    return os.getenv('SYNC_TASKS_HANDLER_URL', '')


def _oidc_audience() -> str:
    return os.getenv('SYNC_TASKS_OIDC_AUDIENCE') or _handler_url()


def _account_deletion_oidc_audience() -> str:
    return os.getenv('ACCOUNT_DELETION_TASKS_OIDC_AUDIENCE') or os.getenv('ACCOUNT_DELETION_HANDLER_URL', '')


def _invoker_sa() -> str:
    return os.getenv('SYNC_TASKS_INVOKER_SA', '')


def get_sync_tasks_max_attempts() -> int:
    # Must mirror the queue's maxAttempts (documented invariant).
    return int(os.getenv('SYNC_TASKS_MAX_ATTEMPTS', '5'))


def is_cloud_tasks_dispatch_enabled() -> bool:
    return os.getenv('SYNC_DISPATCH_MODE', 'inline') == 'cloud_tasks'


def is_sync_backfill_routing_enabled() -> bool:
    return os.getenv('SYNC_BACKFILL_ROUTING_ENABLED', 'false').lower() == 'true'


def is_audio_merge_dispatch_enabled() -> bool:
    return os.getenv('AUDIO_MERGE_DISPATCH_MODE', 'inline') == 'cloud_tasks'


# The production customer data plane, per INV-DATA-1
# (product/invariants/data-plane-continuity.md).
PRODUCTION_DATA_PROJECTS = frozenset({'based-hardware'})


def is_account_deletion_dispatch_enabled() -> bool:
    return os.getenv('ACCOUNT_DELETION_DISPATCH_MODE', 'inline') == 'cloud_tasks'


def assert_inline_account_deletion_permitted() -> None:
    """Refuse to execute a wipe in-process against production data.

    ``OMI_ENV_STAGE`` is unset on a developer machine, so the production guard
    below returns early there while ``.env`` still points at the production
    project. That combination made a local backend run a wipe executor for real
    accounts. The project a process is pointed at is the honest test, and it is
    one no local run can forget to set.
    """
    project = (os.getenv('GOOGLE_CLOUD_PROJECT') or os.getenv('SYNC_TASKS_PROJECT') or '').strip()
    if project and project in PRODUCTION_DATA_PROJECTS:
        raise RuntimeError(
            f'refusing inline account-deletion execution against production project {project!r}; '
            'set ACCOUNT_DELETION_DISPATCH_MODE=cloud_tasks so the OIDC handler owns the wipe'
        )


def validate_account_deletion_dispatch_configuration() -> None:
    """Reject a production process that could execute deletion wipes inline.

    Account deletion is intentionally different from sync's staged rollout: an
    accepted deletion request must have one durable, OIDC-protected execution
    owner. Keeping this check at process startup prevents a missing deploy
    binding from silently falling back to the in-process dispatcher.
    """
    if os.getenv('OMI_ENV_STAGE', '').strip().lower() != 'prod':
        return

    if not is_account_deletion_dispatch_enabled():
        raise RuntimeError('production requires ACCOUNT_DELETION_DISPATCH_MODE=cloud_tasks')

    required_env = (
        'SYNC_TASKS_PROJECT',
        'SYNC_TASKS_LOCATION',
        'SYNC_TASKS_INVOKER_SA',
        'SYNC_TASKS_HANDLER_URL',
        'ACCOUNT_DELETION_TASKS_QUEUE',
        'ACCOUNT_DELETION_HANDLER_URL',
    )
    missing = [name for name in required_env if not os.getenv(name, '').strip()]
    if missing:
        raise RuntimeError(f'production account-deletion Cloud Tasks config is incomplete: {", ".join(missing)}')

    assert_account_deletion_queue_exists()


def assert_account_deletion_queue_exists(client: Any = None) -> None:
    """Prove the configured queue resolves, not merely that its name is set.

    Reading env vars said "configured" for a month while the queue did not
    exist, so every dispatch 404'd behind an accepted deletion request. Only a
    definitive NotFound fails startup; an unreachable Cloud Tasks API is an
    unanswered question, not a proven absence.
    """
    project = os.getenv('SYNC_TASKS_PROJECT', '').strip()
    location = os.getenv('SYNC_TASKS_LOCATION', '').strip()
    queue = os.getenv('ACCOUNT_DELETION_TASKS_QUEUE', '').strip()
    if not all([project, location, queue]):
        return
    resolved = client or _get_tasks_client()
    try:
        resolved.get_queue(name=resolved.queue_path(project, location, queue))
    except NotFound as exc:
        raise RuntimeError(
            f'account-deletion Cloud Tasks queue {queue!r} does not exist in {project}/{location}; '
            'an accepted deletion request would have no executor'
        ) from exc
    except Exception as exc:  # noqa: BLE001 - availability is not absence
        logger.warning('account-deletion queue existence probe inconclusive: %s', sanitize(str(exc)))


def is_listen_finalization_dispatch_enabled() -> bool:
    """Whether platform-key listen finalization uses its durable worker."""
    return os.getenv('LISTEN_FINALIZATION_DISPATCH_MODE', 'inline') == 'cloud_tasks'


def is_listen_finalization_dispatch_configured() -> bool:
    """Whether the durable finalizer can be admitted without an inline fallback."""
    return is_listen_finalization_dispatch_enabled() and all(
        (
            os.getenv('SYNC_TASKS_PROJECT', ''),
            os.getenv('SYNC_TASKS_LOCATION', ''),
            os.getenv('LISTEN_FINALIZATION_TASKS_QUEUE', ''),
            _listen_finalization_handler_url(),
            _listen_finalization_invoker_sa(),
        )
    )


def get_account_deletion_tasks_max_attempts() -> int:
    return int(os.getenv('ACCOUNT_DELETION_TASKS_MAX_ATTEMPTS', get_sync_tasks_max_attempts()))


def get_listen_finalization_tasks_max_attempts() -> int:
    """Must mirror the dedicated finalization queue's maxAttempts setting."""
    return int(os.getenv('LISTEN_FINALIZATION_TASKS_MAX_ATTEMPTS', get_sync_tasks_max_attempts()))


def _enqueue_named_task(
    queue: str,
    url: str,
    task_id: str,
    payload: Dict[str, Any],
    *,
    audience: Optional[str] = None,
    invoker_sa: Optional[str] = None,
) -> None:
    """Enqueue one named HTTP task. Duplicate names are treated as success —
    Cloud Tasks deduplicates named tasks. Any other failure raises."""
    project = os.getenv('SYNC_TASKS_PROJECT', '')
    location = os.getenv('SYNC_TASKS_LOCATION', '')
    selected_invoker_sa = invoker_sa or _invoker_sa()
    if not all([project, location, queue, url, selected_invoker_sa]):
        raise RuntimeError('Cloud Tasks dispatch enabled but task env vars are incomplete')

    client = _get_tasks_client()
    parent = client.queue_path(project, location, queue)
    task = tasks_v2.Task(
        name=client.task_path(project, location, queue, task_id),
        http_request=tasks_v2.HttpRequest(
            http_method=tasks_v2.HttpMethod.POST,
            url=url,
            headers={'Content-Type': 'application/json'},
            body=json.dumps(payload).encode(),
            oidc_token=tasks_v2.OidcToken(
                service_account_email=selected_invoker_sa,
                audience=audience or _oidc_audience(),
            ),
        ),
        dispatch_deadline=duration_pb2.Duration(seconds=DISPATCH_DEADLINE_SECONDS),
    )
    try:
        client.create_task(parent=parent, task=task)  # type: ignore[reportUnknownMemberType]  # google.cloud.tasks_v2 partially untyped
    except AlreadyExists:
        logger.info('task %s already enqueued, skipping duplicate', task_id)


def enqueue_sync_job(payload: Dict[str, Any]) -> None:
    """Enqueue one named HTTP task (task id = job_id) for a sync job.

    Duplicate names are success. Callers retry the same name a bounded number
    of times, then retain staged retry material if acknowledgement remains
    uncertain; they never fall back inline after submitting this task.

    Queue selection stays on the main queue unless SYNC_BACKFILL_ROUTING_ENABLED
    is true, the payload lane is backfill, and both SYNC_BACKFILL_TASKS_QUEUE
    and SYNC_BACKFILL_TASKS_HANDLER_URL are set. Missing backfill env falls
    back to the main queue so a job is never dropped.

    The two-lane split was collapsed in #10400 after a customer incident: every
    offline upload classifies as backfill (no server capture proof), and the
    backfill worker then admitted only a few jobs at once, so Cloud Tasks
    retried the surplus with exponential backoff until recordings sat
    unprocessed for many hours. Restoring the split is gated default-off
    because the backfill worker is now maxScale 30 / concurrency 1
    (request-based) rather than the ~4-dispatch lane that caused the incident.
    The lane label is always carried on the payload for metering and reporting.
    """
    if frozenset(payload) != SYNC_JOB_TASK_PAYLOAD_KEYS:
        raise ValueError('sync job payload does not match the durable worker schema')
    if payload.get('lane') == 'backfill' and is_sync_backfill_routing_enabled():
        queue = os.getenv('SYNC_BACKFILL_TASKS_QUEUE', '').strip()
        handler_url = os.getenv('SYNC_BACKFILL_TASKS_HANDLER_URL', '').strip()
        if queue and handler_url:
            _enqueue_named_task(
                queue,
                handler_url,
                str(payload['job_id']),
                payload,
                audience=os.getenv('SYNC_BACKFILL_TASKS_OIDC_AUDIENCE') or handler_url,
            )
            return
    _enqueue_named_task(os.getenv('SYNC_TASKS_QUEUE', ''), _handler_url(), str(payload['job_id']), payload)


def enqueue_audio_merge_job(payload: Dict[str, Any]) -> None:
    """Enqueue one named merge task per (conversation, audio_file).

    Task name am-{conversation_id}-{audio_file_id} dedupes concurrent enqueues
    from /urls polling; the handler's artifact-exists check covers the rest.
    Tokens are minted with the same audience as sync tasks so a single
    verify_cloud_tasks_oidc dependency covers both handlers.

    schema_version 2 = conversation-level artifact build: the name embeds the
    audio_files fingerprint so a rebuild after late chunks gets a fresh name
    and isn't swallowed by the named-task tombstone. 'amc-' cannot collide with
    per-part names (audio_file ids are UUIDv4).
    """
    if payload.get('schema_version') == 2:
        task_id = f"amc-{payload['conversation_id']}-{payload['fingerprint']}"
    else:
        task_id = f"am-{payload['conversation_id']}-{payload['audio_file_id']}"
    _enqueue_named_task(
        os.getenv('AUDIO_MERGE_TASKS_QUEUE', ''),
        os.getenv('AUDIO_MERGE_HANDLER_URL', ''),
        task_id,
        payload,
    )


def enqueue_account_deletion_wipe(wipe_job_id: str) -> None:
    """Wake one durable deletion job without exposing a user identifier.

    The Firestore job is canonical. Cloud Tasks diagnostics must not contain a
    Firebase uid; the OIDC handler resolves the uid only after looking up this
    opaque job identifier.
    """
    if not wipe_job_id:
        raise ValueError('wipe_job_id must be non-empty')
    job_hash = hashlib.sha256(wipe_job_id.encode('utf-8')).hexdigest()[:32]
    task_id = f"account-delete-{job_hash}-{uuid.uuid4().hex}"
    _enqueue_named_task(
        os.getenv('ACCOUNT_DELETION_TASKS_QUEUE', ''),
        os.getenv('ACCOUNT_DELETION_HANDLER_URL', ''),
        task_id,
        {'job_id': wipe_job_id},
        audience=_account_deletion_oidc_audience(),
    )


def _listen_finalization_handler_url() -> str:
    return os.getenv('LISTEN_FINALIZATION_TASKS_HANDLER_URL', '')


def _listen_finalization_audience() -> str:
    return os.getenv('LISTEN_FINALIZATION_TASKS_OIDC_AUDIENCE') or _listen_finalization_handler_url()


def _listen_finalization_invoker_sa() -> str:
    return os.getenv('LISTEN_FINALIZATION_TASKS_INVOKER_SA') or _invoker_sa()


def enqueue_listen_finalization_job(job_id: str, dispatch_generation: int) -> None:
    """Wake the finalizer with opaque routing data only.

    The Firestore job is canonical.  The task intentionally contains neither a
    uid nor any conversation/BYOK material so Cloud Tasks diagnostics cannot
    expose user content or credentials.
    """
    _enqueue_named_task(
        os.getenv('LISTEN_FINALIZATION_TASKS_QUEUE', ''),
        _listen_finalization_handler_url(),
        f'listen-finalization-{job_id}-{dispatch_generation}',
        {'job_id': job_id, 'dispatch_generation': dispatch_generation},
        audience=_listen_finalization_audience(),
        invoker_sa=_listen_finalization_invoker_sa(),
    )


def _verify_cloud_tasks_oidc(request: Request, *, audience: str, invoker_sa: str, log_failure: bool = True) -> int:
    """Verify a configured task audience and issuer; returns task retry count.

    Sync function on purpose — verify_oauth2_token fetches Google certs over
    HTTP, and FastAPI runs sync dependencies in the threadpool.
    """
    if not audience or not invoker_sa:
        # Env unset: this service is not a task target (e.g. main backend
        # running the shared image) — never accept task traffic.
        raise HTTPException(status_code=403, detail='Task dispatch not configured on this service')

    auth_header = request.headers.get('authorization', '')
    if not auth_header.startswith('Bearer '):
        raise HTTPException(status_code=403, detail='Missing bearer token')

    try:
        claims: Any = id_token.verify_oauth2_token(auth_header[len('Bearer ') :], _get_auth_request(), audience=audience)  # type: ignore[reportUnknownMemberType]  # google.oauth2.id_token partially untyped
    except Exception as e:
        # Distinguishes bad tokens from transient JWKS-fetch failures in logs
        if log_failure:
            logger.warning('OIDC token verification failed: %s', e)
        raise HTTPException(status_code=403, detail='Invalid OIDC token')

    if claims.get('email') != invoker_sa or not claims.get('email_verified'):
        raise HTTPException(status_code=403, detail='Unexpected token identity')

    try:
        return int(request.headers.get('x-cloudtasks-taskretrycount', '0'))
    except ValueError:
        return 0


def verify_cloud_tasks_oidc(request: Request) -> int:
    """FastAPI dependency for sync and merge task routes."""
    return _verify_cloud_tasks_oidc(request, audience=_oidc_audience(), invoker_sa=_invoker_sa())


def verify_account_deletion_cloud_tasks_oidc(request: Request) -> AccountDeletionTaskAuthentication:
    """Verify deletion tasks."""
    deletion_audience = _account_deletion_oidc_audience()
    retry_count = _verify_cloud_tasks_oidc(
        request,
        audience=deletion_audience,
        invoker_sa=_invoker_sa(),
        log_failure=False,
    )
    return AccountDeletionTaskAuthentication(retry_count=retry_count)


def verify_listen_finalization_cloud_tasks_oidc(request: Request) -> int:
    """FastAPI dependency for the isolated listen finalization task route."""
    return _verify_cloud_tasks_oidc(
        request,
        audience=_listen_finalization_audience(),
        invoker_sa=_listen_finalization_invoker_sa(),
    )