checkbot/app/runtime.py

import asyncio
import time
from contextlib import asynccontextmanager

from fastapi import FastAPI

from . import audit
from .alerts import (AlertActions, AlertAnalyzer, AlertPipeline, AlertRepository, SendLimiter,
                     SlidingWindowLimiter, SystemNotifier, base_url_provider, content_mode_provider,
                     send_limit_provider)
from .chat import AdapterRegistry
from .chatops import ChatOps
from .constants import (ALERTS_BUNDLE_SECONDS, ALERTS_RECOVERY_ENABLED, ALERTS_REMINDER_HOURS, DEFAULT_BUNDLE_SECONDS,
                        DEFAULT_REMINDER_HOURS, FLAG_OFF, FLAG_ON, MAX_REMINDER_HOURS, SELFTEST_ENABLED,
                        SELFTEST_NOTIFY)
from .db import Database
from .agent_runs import configured_secrets
from .mutes import MuteService
from .ollama import OllamaClient
from .poller import BaselineGuard, Poller, PollerStatus, StateStore
from .prompts import (PROMPT_AGENT, PROMPT_ALERT, PromptRepository, PromptService, PromptTester, UserRateLimiter,
                      default_registry)
from .retention import RetentionRepository, RetentionRunner, RetentionService, policy_from
from .reminders import ReminderWorker
from .routing import RouteService, seed_default_rules
from .selftest import SelftestRepository, SelftestRunner, SelftestService


class Runtime:
    def __init__(self, db: Database, store):
        self._db = db
        self.status = PollerStatus()
        self.alerts = AlertRepository(db)
        self.state = StateStore(db)
        self.registry = AdapterRegistry(store)
        self.mutes = MuteService(db)
        self.routes = RouteService(db)
        self.prompts = PromptService(PromptRepository(db), default_registry(), audit.recorder(db), store.get)
        self.retention = RetentionService(RetentionRepository(db), lambda: policy_from(store.get_int),
                                          audit.recorder(db))
        self.retention_runner = RetentionRunner(self.retention)
        analyzer = AlertAnalyzer(store, lambda: OllamaClient.from_settings(store),
                                 self.state.recent_changes, SlidingWindowLimiter(time.time),
                                 lambda: self.prompts.effective_text(PROMPT_ALERT))
        actions, limiter, send_limit = AlertActions(store), SendLimiter(time.time), send_limit_provider(store)
        self.pipeline = AlertPipeline(
            self.alerts, self.registry.current, analyzer,
            lambda: store.get_int(ALERTS_BUNDLE_SECONDS, DEFAULT_BUNDLE_SECONDS),
            routes=self.routes, mutes=self.mutes, actions=actions,
            recovery_enabled=lambda: store.get(ALERTS_RECOVERY_ENABLED, FLAG_ON) != FLAG_OFF,
            limiter=limiter, send_limit=send_limit, content_mode=content_mode_provider(store),
            base_url=base_url_provider(store),
        )
        self.baseline = BaselineGuard(db)
        self.notifier = SystemNotifier(self.alerts, self.registry.current, limiter, send_limit, actions,
                                       on_baseline=self.pipeline.reset_reported)
        self.poller = Poller(store, self.state, self.pipeline.handle, self.status,
                             rules=self.routes.enabled_routes, notifier=self.notifier, baseline=self.baseline)
        self.selftest_repo = SelftestRepository(db)
        self.selftest = SelftestService(store, self.selftest_repo)
        self.selftest_runner = SelftestRunner(
            db, self.selftest, self.notifier,
            enabled=lambda: store.get(SELFTEST_ENABLED, FLAG_ON) != FLAG_OFF,
            notify_enabled=lambda: store.get(SELFTEST_NOTIFY, FLAG_ON) != FLAG_OFF,
        )
        self.reminders = ReminderWorker(
            self.state.load, self.routes.enabled_routes, self.mutes.active,
            lambda: min(store.get_int(ALERTS_REMINDER_HOURS, DEFAULT_REMINDER_HOURS), MAX_REMINDER_HOURS),
            self.pipeline.handle, self.state.mark_reminded,
        )
        self.chatops = ChatOps(db, store, self.state, self.alerts, self.mutes,
                               lambda: self.prompts.effective_text(PROMPT_AGENT))
        self.prompt_tester = PromptTester(
            store, lambda: OllamaClient.from_settings(store), self.state.recent_changes, self.chatops.agent,
            lambda: configured_secrets(store), UserRateLimiter())

    @asynccontextmanager
    async def lifespan(self, app: FastAPI):
        await asyncio.to_thread(seed_default_rules, self._db)
        poller_task = asyncio.create_task(self.poller.run())
        reminder_task = asyncio.create_task(self.reminders.run())
        pipeline_task = asyncio.create_task(self.pipeline.run())
        listener_tasks = [asyncio.create_task(run()) for run in self.chatops.listener_runners()]
        retention_task = asyncio.create_task(self.retention_runner.run())
        selftest_task = asyncio.create_task(self.selftest_runner.run())
        try:
            yield
        finally:
            for task in listener_tasks:
                task.cancel()
            await asyncio.gather(*listener_tasks, return_exceptions=True)
            retention_task.cancel()
            await asyncio.gather(retention_task, return_exceptions=True)
            selftest_task.cancel()
            await asyncio.gather(selftest_task, return_exceptions=True)
            poller_task.cancel()
            reminder_task.cancel()
            await asyncio.gather(poller_task, reminder_task, return_exceptions=True)
            await self.pipeline.finish()
            pipeline_task.cancel()
            await asyncio.gather(pipeline_task, return_exceptions=True)
            await self.pipeline.aclose()
            await self.registry.aclose()