api/background_tasks.py

"""Supervision utilities for long-lived background asyncio tasks."""
from __future__ import annotations

import asyncio
import logging
from collections.abc import Awaitable
from typing import Any

_logger = logging.getLogger("agente_ai.background_tasks")
_tasks: dict[str, asyncio.Task[Any]] = {}


def spawn_background_task(coro: Awaitable[Any], *, name: str) -> asyncio.Task[Any]:
    """Start one named background task and retain it for lifecycle shutdown.

    A live task with the same name is reused. The passed coroutine is closed in
    that case so duplicate startup calls do not leak an un-awaited coroutine.
    """
    current = _tasks.get(name)
    if current is not None and not current.done():
        close = getattr(coro, "close", None)
        if close is not None:
            close()
        return current

    task = asyncio.create_task(coro, name=name)
    _tasks[name] = task

    def _report(task_result: asyncio.Task[Any]) -> None:
        if task_result.cancelled():
            return
        try:
            error = task_result.exception()
        except asyncio.CancelledError:
            return
        if error is not None:
            _logger.error("background task %s failed: %s", name, error, exc_info=error)

    task.add_done_callback(_report)
    return task


async def shutdown_background_tasks() -> None:
    """Cancel and await all supervised background tasks."""
    tasks = [task for task in _tasks.values() if not task.done()]
    for task in tasks:
        task.cancel()
    if tasks:
        await asyncio.gather(*tasks, return_exceptions=True)
    _tasks.clear()


__all__ = ["spawn_background_task", "shutdown_background_tasks"]