betterbot/background_tasks.py

"""Background task manager — spawns long-running Copilot SDK sessions outside the request cycle."""

import asyncio
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Awaitable, Callable

from config import settings
from copilot_runtime import copilot, stream_session
from llm_costs import extract_usage_and_cost
from model_selection import ModelSelection, build_provider_config, resolve_selection
from ux import extract_final_text

logger = logging.getLogger(__name__)

BACKGROUND_TIMEOUT = 600  # 10 minutes


@dataclass
class BackgroundTask:
    task_id: str
    description: str
    thread_id: str
    started_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    result: str | None = None
    error: str | None = None
    usage: dict | None = None
    _asyncio_task: asyncio.Task[None] | None = field(default=None, repr=False)

    @property
    def done(self) -> bool:
        return self._asyncio_task is not None and self._asyncio_task.done()

    @property
    def elapsed_seconds(self) -> float:
        return (datetime.now(timezone.utc) - self.started_at).total_seconds()


class BackgroundTaskManager:
    """Tracks background agent tasks per thread. One active task per thread."""

    def __init__(self) -> None:
        self._tasks: dict[str, BackgroundTask] = {}
        self._thread_tasks: dict[str, list[str]] = {}
        self._counter = 0

    def _next_id(self) -> str:
        self._counter += 1
        return f"bg-{self._counter}"

    def start(
        self,
        *,
        thread_id: str,
        description: str,
        selection: ModelSelection,
        system_message: str,
        on_complete: Callable[["BackgroundTask"], Awaitable[None]],
    ) -> BackgroundTask:
        task_id = self._next_id()
        task = BackgroundTask(task_id=task_id, description=description, thread_id=thread_id)

        # Background agents use a dedicated model/provider (not the parent session's provider)
        bg_selection = resolve_selection(model=settings.BACKGROUND_MODEL)

        bg_system = (
            system_message.rstrip() + "\n\nYou are running as a background agent. "
            "Complete the task fully and return your findings. Be thorough."
        )

        async def _run() -> None:
            try:

                async def _collect() -> str:
                    events: list = []
                    async for ev in stream_session(
                        copilot,
                        model=bg_selection.model,
                        provider_config=build_provider_config(bg_selection),
                        system_message=bg_system,
                        prompt=description,
                    ):
                        events.append(ev)
                    # Extract usage before final text
                    task.usage = extract_usage_and_cost(bg_selection.model, bg_selection.provider, events)
                    return extract_final_text(events) or "Task completed but produced no output."

                task.result = await asyncio.wait_for(_collect(), timeout=BACKGROUND_TIMEOUT)
            except asyncio.TimeoutError:
                task.error = f"Timed out after {BACKGROUND_TIMEOUT}s"
            except Exception as exc:
                task.error = str(exc)
                logger.exception("Background task %s failed", task_id)
            finally:
                try:
                    await on_complete(task)
                except Exception:
                    logger.exception("Callback failed for background task %s", task_id)

        task._asyncio_task = asyncio.create_task(_run())
        self._tasks[task_id] = task
        self._thread_tasks.setdefault(thread_id, []).append(task_id)
        return task

    def get_active(self, thread_id: str) -> BackgroundTask | None:
        for task_id in reversed(self._thread_tasks.get(thread_id, [])):
            task = self._tasks.get(task_id)
            if task and not task.done:
                return task
        return None

    def get_latest(self, thread_id: str) -> BackgroundTask | None:
        ids = self._thread_tasks.get(thread_id, [])
        if not ids:
            return None
        return self._tasks.get(ids[-1])

    def context_summary(self, thread_id: str) -> str | None:
        """Return background-agent context to inject into the system message, or None."""
        active = self.get_active(thread_id)
        if active:
            return (
                f"A background agent is currently running ({active.elapsed_seconds:.0f}s elapsed).\n"
                f"Task: {active.description}\n"
                "Its results will be posted to the chat when done."
            )
        latest = self.get_latest(thread_id)
        if latest is None:
            return None
        if latest.error:
            return f"The last background agent failed: {latest.error}"
        if latest.result:
            snippet = latest.result[:2000] + ("..." if len(latest.result) > 2000 else "")
            return f"The last background agent completed.\nTask: {latest.description}\nResult:\n{snippet}"
        return None

    def format_status(self, thread_id: str) -> str:
        active = self.get_active(thread_id)
        if active:
            return f"A background agent is running ({active.elapsed_seconds:.0f}s elapsed).\nTask: {active.description}"
        latest = self.get_latest(thread_id)
        if latest is None:
            return "No background agent has run in this thread."
        if latest.error:
            return f"Last background agent failed: {latest.error}"
        return f"Last background agent completed.\nTask: {latest.description}"