main.py

import asyncio
import json
import logging
import os

from utils.env_loader import firebase_admin_options, load_backend_env
from utils.firebase_admin_runtime import (
    firebase_verify_only_credential,
    install_firebase_auth_mutation_guard,
    install_google_adc_guard,
)
from config.chat_first_e2e_fixture import is_chat_first_e2e_harness_runtime

load_backend_env()  # No-op if no env files exist (production); stage + local overrides otherwise
install_google_adc_guard()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

import firebase_admin
from fastapi import FastAPI
from starlette.middleware.cors import CORSMiddleware
from starlette.staticfiles import StaticFiles

from database.google_credentials import prepare_google_credentials

prepare_google_credentials()
install_firebase_auth_mutation_guard()

from routers import (
    chat,
    firmware,
    static_map,
    transcribe,
    omni_relay,
    auto_model,
    notifications,
    speech_profile,
    agents,
    users,
    trends,
    sync,
    apps,
    payment,
    integration,
    conversations,
    conversation_mutations,
    memories,
    api_key_management,
    mcp,
    mcp_sse,
    oauth,
    auth,
    action_items,
    account_cutover,
    candidates,
    chat_first,
    chat_first_e2e,
    daily_summary_e2e,
    task_integrations,
    integrations,
    x_connector,
    other,
    developer,
    updates,
    calendar_meetings,
    google_calendar,
    calendar_onboarding,
    imports,
    knowledge_graph,
    wrapped,
    folders,
    goals,
    workstreams,
    announcements,
    phone_calls,
    agent_tools,
    tools,
    metrics,
    fair_use_admin,
    feedback_admin,
    staged_tasks,
    focus_sessions,
    advice,
    chat_sessions,
    chat_generation,
    desktop_agent_vm,
    desktop_chat,
    desktop_core,
    desktop_prompts,
    desktop_proxy,
    desktop_realtime,
    desktop_screen_crisp,
    frame_requests,
    referrals,
    desktop_tts_updates,
    scores,
    stt,
    tts,
    memory_admin,
    memory_product,
    task_recommendations,
    conversation_finalization,
    public_shared_conversation_chat,
    screen_frames,
    jit_ledger_snapshot,
    csat,
    jit_rollout,
    email_preferences,
)
from routers.listen.registry import proactive_message_dispatcher

from utils.other.timeout import TimeoutMiddleware
from utils.observability import log_langsmith_status
from utils.subscription import validate_stripe_price_ids
from utils.http_client import close_all_clients
from utils.jit_rollout import close_posthog_control_plane
from utils.metrics import start_metrics_sidecar_server, stop_metrics_sidecar_server
from utils.executors import (
    drain_background_tasks,
    log_executor_health,
    run_blocking,
    db_executor,
)
from utils.executors import start_background_task
from utils.cloud_tasks import validate_account_deletion_dispatch_configuration
from utils.llm.managed_spend_ledger import shutdown_managed_spend_ledger
from services.conversation_finalization import reconcile_abandoned_byok_finalization_jobs
from services.conversation_finalization import reconcile_listen_finalization_jobs
from services.conversation_finalization import reconcile_meeting_receipts
from services.conversation_finalization import reconcile_stale_processing_conversations
from database.durable_queue_age import publish_all_queue_oldest_ready_ages
from services.users.account_deletion import reconcile_pending_deletion_wipes
from utils.other.local_storage import local_storage_root_from_env

# Log LangSmith tracing status at startup
log_langsmith_status()

# Validate Stripe price IDs so misconfigured plans fail loud
validate_stripe_price_ids()

_auth_emulator_host = os.environ.get("FIREBASE_AUTH_EMULATOR_HOST", "").strip()
_firebase_admin_options = firebase_admin_options()
_verify_only_credential = firebase_verify_only_credential()
if _verify_only_credential is not None:
    firebase_admin.initialize_app(_verify_only_credential, options=_firebase_admin_options)
elif _auth_emulator_host:
    for _adc_key in ("GOOGLE_APPLICATION_CREDENTIALS", "SERVICE_ACCOUNT_JSON", "FIREBASE_AUTH_CREDENTIALS_PATH"):
        os.environ.pop(_adc_key, None)
    _firebase_project_id = (
        os.environ.get("FIREBASE_AUTH_PROJECT_ID") or os.environ.get("FIREBASE_PROJECT_ID") or "demo-omi-local"
    )
    firebase_admin.initialize_app(options={"projectId": _firebase_project_id})  # type: ignore[reportUnknownMemberType]  # firebase_admin untyped
elif os.environ.get("SERVICE_ACCOUNT_JSON"):
    service_account_info = json.loads(os.environ["SERVICE_ACCOUNT_JSON"])
    credentials = firebase_admin.credentials.Certificate(service_account_info)
    firebase_admin.initialize_app(credentials, options=_firebase_admin_options)  # type: ignore[reportUnknownMemberType]  # firebase_admin untyped
else:
    firebase_admin.initialize_app(options=_firebase_admin_options)  # type: ignore[reportUnknownMemberType]  # firebase_admin untyped

app = FastAPI()

_local_storage_root = local_storage_root_from_env()
if _local_storage_root is not None:
    _local_storage_root.mkdir(parents=True, exist_ok=True)
    app.mount('/_local/storage', StaticFiles(directory=_local_storage_root), name='local-storage')

# Explicit, default-deny CORS: this API is Bearer-token authenticated (mobile/
# desktop apps, not ambient browser cookies), so no cross-origin browser
# caller needs to be allowed by default. CORS_ALLOWED_ORIGINS lets an operator
# opt a specific web frontend in (comma-separated exact origins — never "*",
# and never combined with allow_credentials, which would let any site read
# authenticated responses for a signed-in visitor).
_cors_allowed_origins = [o.strip() for o in os.getenv('CORS_ALLOWED_ORIGINS', '').split(',') if o.strip()]
if '*' in _cors_allowed_origins:
    raise RuntimeError('CORS_ALLOWED_ORIGINS must not contain "*" — list explicit origins instead')
app.add_middleware(
    CORSMiddleware,
    allow_origins=_cors_allowed_origins,
    allow_credentials=False,
    allow_methods=['*'],
    allow_headers=['*'],
)

app.include_router(transcribe.router)
app.include_router(static_map.router)
app.include_router(omni_relay.router)
app.include_router(auto_model.router)
app.include_router(conversations.router)
app.include_router(conversation_mutations.router)
app.include_router(public_shared_conversation_chat.router)
app.include_router(action_items.router)
app.include_router(account_cutover.router)
app.include_router(candidates.router)
app.include_router(chat_first.router)
if is_chat_first_e2e_harness_runtime():
    # The fixture router has its own runtime check as defense in depth.  It is
    # intentionally absent from dev/prod route tables, not merely disabled.
    app.include_router(chat_first_e2e.router)
    # Same stage boundary, same defense in depth: the desktop memory-review flow
    # needs a daily summary carrying `memories_learned`, which only the nightly
    # job produces in a deployable environment.
    app.include_router(daily_summary_e2e.router)
app.include_router(task_integrations.router)
app.include_router(integrations.router)
app.include_router(x_connector.router)
app.include_router(memories.router)
app.include_router(chat.router)
app.include_router(speech_profile.router)
app.include_router(notifications.router)
app.include_router(integration.router)
app.include_router(agents.router)
app.include_router(users.router)
app.include_router(referrals.router)
app.include_router(csat.router)
app.include_router(feedback_admin.router)
app.include_router(email_preferences.router)
app.include_router(desktop_prompts.router)
app.include_router(conversation_finalization.router)
app.include_router(trends.router)

app.include_router(other.router)

app.include_router(firmware.router)
app.include_router(updates.router)
app.include_router(sync.router)

app.include_router(apps.router)
app.include_router(calendar_meetings.router)
app.include_router(google_calendar.router)
app.include_router(calendar_onboarding.router)
app.include_router(oauth.router)  # Added oauth router (for Omi Apps)
app.include_router(auth.router)  # Added auth router (for the main Omi App, this is the core auth router)


app.include_router(payment.router)
app.include_router(api_key_management.mcp_router)
app.include_router(mcp.router)
app.include_router(mcp_sse.router)
app.include_router(api_key_management.developer_router)
app.include_router(developer.router)
app.include_router(imports.router)
app.include_router(wrapped.router)
app.include_router(folders.router)
app.include_router(knowledge_graph.router)
app.include_router(goals.router)
app.include_router(workstreams.router)
app.include_router(announcements.router)
app.include_router(phone_calls.router)
app.include_router(agent_tools.router)
app.include_router(tools.router)
app.include_router(metrics.router)
app.include_router(fair_use_admin.router)
app.include_router(staged_tasks.router)
app.include_router(focus_sessions.router)
app.include_router(advice.router)
app.include_router(chat_sessions.router)
app.include_router(chat_generation.router)
app.include_router(scores.router)
app.include_router(stt.router)
app.include_router(tts.router)
app.include_router(memory_admin.router)
app.include_router(memory_product.router)
app.include_router(task_recommendations.router)
app.include_router(jit_ledger_snapshot.router)
app.include_router(jit_rollout.router)
app.include_router(desktop_core.router)
app.include_router(desktop_agent_vm.router)
app.include_router(desktop_chat.router)
app.include_router(desktop_proxy.router)
app.include_router(desktop_realtime.router)
app.include_router(desktop_screen_crisp.router)
app.include_router(frame_requests.router)
app.include_router(desktop_tts_updates.router)
app.include_router(screen_frames.router)
jit_rollout.validate_jit_rollout_contract(app)


methods_timeout = {
    "GET": os.environ.get('HTTP_GET_TIMEOUT'),
    "POST": os.environ.get('HTTP_POST_TIMEOUT'),
    "PUT": os.environ.get('HTTP_PUT_TIMEOUT'),
    "PATCH": os.environ.get('HTTP_PATCH_TIMEOUT'),
    "DELETE": os.environ.get('HTTP_DELETE_TIMEOUT'),
}

# The Cloud Tasks sync-job handler runs the whole pipeline inside the request,
# so it needs a much higher cap than the default. Must stay below the job run
# lock TTL (1800s) so a lock can never expire under a live run.
paths_timeout = {
    "/v2/sync-jobs/run": os.environ.get('HTTP_SYNC_JOBS_RUN_TIMEOUT', 1500),
    "/v2/audio-merge-jobs/run": os.environ.get('HTTP_AUDIO_MERGE_RUN_TIMEOUT', 600),
    "/v1/users/account-deletion-wipes/run": os.environ.get('HTTP_ACCOUNT_DELETION_WIPE_RUN_TIMEOUT', 1500),
    "/v1/conversation-finalization-jobs/run": os.environ.get('HTTP_LISTEN_FINALIZATION_RUN_TIMEOUT', 1500),
    # STT proxy: 30s slot wait + 300s parakeet client budget (get_stt_proxy_client)
    # + headroom for auth and the multipart spool read; the default POST timeout
    # would cut long files off mid-transcription.
    "/v1/stt/transcribe": os.environ.get('HTTP_STT_TRANSCRIBE_TIMEOUT', 350),
}

app.add_middleware(TimeoutMiddleware, methods_timeout=methods_timeout, paths_timeout=paths_timeout)

from utils.byok import BYOKMiddleware

app.add_middleware(BYOKMiddleware)


@app.on_event("startup")  # type: ignore[reportDeprecated]  # FastAPI on_event still functional; lifespan migration would change app wiring
async def startup_event():
    start_metrics_sidecar_server()
    validate_account_deletion_dispatch_configuration()
    asyncio.create_task(log_executor_health())
    # Drain account-deletion wipes orphaned by a previous deploy/restart. Offloaded
    # to db_executor so the blocking Firestore queries don't stall event-loop startup.
    start_background_task(
        run_blocking(db_executor, _drain_pending_deletion_wipes),
        name='startup_deletion_wipe_reconcile',
    )
    # Periodic reconciliation ensures stale retrying claims (worker crashed) and
    # new pending/failed wipes are retried without requiring a restart.
    start_background_task(_periodic_deletion_wipe_reconcile(), name='periodic_deletion_wipe_reconcile')
    start_background_task(
        run_blocking(db_executor, _drain_listen_finalization_jobs),
        name='startup_listen_finalization_reconcile',
    )
    start_background_task(
        run_blocking(db_executor, _drain_stale_processing_conversations),
        name='startup_stale_processing_reconcile',
    )
    start_background_task(
        run_blocking(db_executor, _drain_abandoned_byok_finalization_jobs),
        name='startup_byok_abandonment_reconcile',
    )
    start_background_task(
        run_blocking(db_executor, _drain_meeting_receipts),
        name='startup_meeting_receipt_reconcile',
    )
    start_background_task(_periodic_listen_finalization_reconcile(), name='periodic_listen_finalization_reconcile')
    start_background_task(
        proactive_message_dispatcher(),
        name='proactive_message_dispatcher',
    )


def _drain_pending_deletion_wipes():
    """Best-effort reconciliation of pending/failed account-deletion wipes on startup."""
    try:
        result = reconcile_pending_deletion_wipes()
        if result.get('requeued'):
            logger.info(f"Startup deletion-wipe reconciliation: {result}")
    except Exception as e:
        logger.error(f"Startup deletion-wipe reconciliation failed: {e}")


async def _periodic_deletion_wipe_reconcile(interval_seconds: int = 300):
    """Periodically reconcile orphaned or failed account-deletion wipes.

    Runs every 5 minutes (default) so stale retrying claims and new
    pending/failed wipes are retried without requiring a restart.
    """
    while True:
        await asyncio.sleep(interval_seconds)
        try:
            result = await run_blocking(db_executor, reconcile_pending_deletion_wipes)
            if result.get('requeued'):
                logger.info(f"Periodic deletion-wipe reconciliation: {result}")
        except Exception as e:
            logger.error(f"Periodic deletion-wipe reconciliation failed: {e}")


def _drain_listen_finalization_jobs():
    """Best-effort durable finalization recovery after a restart/deploy."""
    try:
        result = reconcile_listen_finalization_jobs()
        if result.get('requeued'):
            logger.info(f"Startup listen-finalization reconciliation: {result}")
    except Exception as e:
        logger.error(f"Startup listen-finalization reconciliation failed: {e}")


def _drain_stale_processing_conversations():
    """Best-effort recovery of bare-`processing` conversations orphaned by a sync-route crash."""
    try:
        result = reconcile_stale_processing_conversations()
        if result.get('completed') or result.get('migrated'):
            logger.info(f"Startup stale-processing reconciliation: {result}")
    except Exception as e:
        logger.error(f"Startup stale-processing reconciliation failed: {e}")


def _drain_abandoned_byok_finalization_jobs():
    """Best-effort disposition of BYOK finalization jobs no live session can claim."""
    try:
        result = reconcile_abandoned_byok_finalization_jobs()
        if result.get('abandoned'):
            logger.info(f"Startup byok-abandonment reconciliation: {result}")
    except Exception as e:
        logger.error(f"Startup byok-abandonment reconciliation failed: {e}")


def _drain_meeting_receipts():
    """Best-effort repair of missing meeting receipt intents and historical receipts."""
    try:
        result = reconcile_meeting_receipts()
        if result.get('repaired') or result.get('backfilled'):
            logger.info(f"Startup meeting-receipt reconciliation: {result}")
    except Exception as e:
        logger.error(f"Startup meeting-receipt reconciliation failed: {e}")


def _listen_finalization_reconcile_interval_seconds() -> int:
    """Periodic reconcile cadence; overridable for hermetic behavioral tests."""
    try:
        seconds = int(os.getenv('LISTEN_FINALIZATION_RECONCILE_INTERVAL_SECONDS', '300'))
    except ValueError:
        seconds = 300
    return max(1, seconds)


async def _periodic_listen_finalization_reconcile(interval_seconds: int | None = None):
    """Replay stale finalization leases and publish durable backlog metrics."""
    if interval_seconds is None:
        interval_seconds = _listen_finalization_reconcile_interval_seconds()
    while True:
        await asyncio.sleep(interval_seconds)
        try:
            result = await run_blocking(db_executor, reconcile_listen_finalization_jobs)
            if result.get('requeued'):
                logger.info(f"Periodic listen-finalization reconciliation: {result}")
        except Exception as e:
            logger.error(f"Periodic listen-finalization reconciliation failed: {e}")
        try:
            stale_result = await run_blocking(db_executor, reconcile_stale_processing_conversations)
            if stale_result.get('completed') or stale_result.get('migrated'):
                logger.info(f"Periodic stale-processing reconciliation: {stale_result}")
        except Exception as e:
            logger.error(f"Periodic stale-processing reconciliation failed: {e}")
        try:
            byok_result = await run_blocking(db_executor, reconcile_abandoned_byok_finalization_jobs)
            if byok_result.get('abandoned'):
                logger.info(f"Periodic byok-abandonment reconciliation: {byok_result}")
        except Exception as e:
            logger.error(f"Periodic byok-abandonment reconciliation failed: {e}")
        try:
            receipt_result = await run_blocking(db_executor, reconcile_meeting_receipts)
            if receipt_result.get('repaired') or receipt_result.get('backfilled'):
                logger.info(f"Periodic meeting-receipt reconciliation: {receipt_result}")
        except Exception as e:
            logger.error(f"Periodic meeting-receipt reconciliation failed: {e}")
        try:
            await run_blocking(db_executor, publish_all_queue_oldest_ready_ages)
        except Exception as e:
            logger.error(f"Periodic durable-queue age publish failed: {e}")


@app.on_event("shutdown")  # type: ignore[reportDeprecated]  # FastAPI on_event still functional; lifespan migration would change app wiring
async def shutdown_event():
    await drain_background_tasks(timeout=10.0)
    await shutdown_managed_spend_ledger()
    await close_all_clients()
    close_posthog_control_plane()
    stop_metrics_sidecar_server()


paths = ['_temp', '_samples', '_segments', '_speech_profiles']
for path in paths:
    if not os.path.exists(path):
        os.makedirs(path)