OpenMAIC/instrumentation.ts

/**
 * Process-scoped startup work.
 *
 * Next calls `register` once per server instance, before it serves a request.
 * That makes it the only place in this app where a background schedule can
 * live: a route module has no such guarantee — it can be instantiated more than
 * once and gets no shutdown hook — so anything periodic started from one is
 * really started per instantiation.
 *
 * `register` must return before the server is ready, so nothing here may block
 * on I/O. Starting a timer does not.
 */
export async function register(): Promise<void> {
  // Also invoked for the Edge runtime, which has neither `pg` nor timers we
  // want; the persistence stack is Node-only.
  if (process.env.NEXT_RUNTIME !== 'nodejs') return;

  const { warnIfDefaultDatabasePasswordIsPublished } =
    await import('@/lib/server/database-password-warning');
  warnIfDefaultDatabasePasswordIsPublished();

  // A boot that fails here must stop the process. Next.js logs a throw from
  // `register` as "Failed to prepare server" but keeps listening and answers
  // every request with a 500, so the throw alone leaves a deployment that
  // looks alive and serves nothing. exitOnBootFailure prints the failure and
  // exits non-zero (see lib/server/boot-failure.ts): a refused configuration as
  // one line with the original message, anything else with its stack. The
  // rethrow is reached only where the exit is stubbed (tests).
  try {
    await validateBootConfiguration();
  } catch (error) {
    const { exitOnBootFailure } = await import('@/lib/server/boot-failure');
    await exitOnBootFailure(error);
    throw error;
  }

  // Warn-only checks on the configuration that passed validation: single-user
  // mode without ACCESS_CODE serves one library to whoever can reach the
  // server. Its warning names ACCESS_CODE itself, so the generic unset-code
  // warning is skipped then: one warning, not two.
  const { warnAboutOwnerIdentityConfiguration } = await import('@/lib/server/identity/registry');
  if (!warnAboutOwnerIdentityConfiguration()) {
    const { warnIfAccessCodeIsUnset } = await import('@/lib/server/access-code-warning');
    warnIfAccessCodeIsUnset(process.env.ACCESS_CODE);
  }

  // Every store's recorded schema versions, read in the background: a database
  // this release must not run against (upgraded by a newer release) stops the
  // process here rather than failing every request that touches the store
  // (lib/persistence/schema-boot-check.ts).
  const { startSchemaBootCheck } = await import('@/lib/persistence/schema-boot-check');
  void startSchemaBootCheck(process.env.DATABASE_URL ?? '');

  // Warn-only: the instance secret that seals keys saved in the model settings
  // (lib/server/instance-secret-check.ts), checked against the keys already
  // stored with one query in the background. It never stops the server.
  const { warnAboutInstanceSecret } = await import('@/lib/server/instance-secret-check');
  void warnAboutInstanceSecret();

  // The one-time import of classrooms earlier versions stored as files
  // (lib/server/legacy-classroom-import.ts). It reads the disk and the
  // database, so it runs in the background, retrying with backoff until it
  // completes: `register` must not wait on it.
  const { startLegacyClassroomImport } = await import('@/lib/server/legacy-classroom-import');
  const legacyClassroomImport = startLegacyClassroomImport();

  // Imported dynamically so the Edge bundle never pulls in `pg`.
  const { startAssetCollectorSchedule } =
    await import('@/lib/persistence/asset-collector-schedule');
  const assetSchedule = startAssetCollectorSchedule();

  // Warn-first boot-time validation of model routing config (MODEL_ROUTES,
  // DEFAULT_MODEL, <PREFIX>_MODELS). Cheap and non-throwing: broken config
  // surfaces here as [config] warnings instead of failing at request time.
  // Imported dynamically so the Edge bundle never pulls in the fs/js-yaml
  // backed provider config it reads.
  const { validateServerConfig } = await import('@/lib/server/config-validation');
  validateServerConfig();

  let runner: import('@/lib/server/agent-runtime/runner').AgentRunnerHandle | undefined;
  let extractionRunner:
    | import('@/lib/server/material-extraction/runner').MaterialExtractionRunnerHandle
    | undefined;
  let ownerExtractionRunner:
    | import('@/lib/server/material-extraction/owner-extraction').OwnerExtractionRunnerHandle
    | undefined;
  let stopAgentEventNotifyBus: (() => Promise<void>) | null = null;
  let generationRunner:
    | import('@/lib/server/generation/run/runner').GenerationRunnerHandle
    | undefined;
  let materialExtractor:
    | import('@/lib/server/materials/extractor-wake').OwnerMaterialExtractorHandle
    | undefined;
  try {
    // Course generation runs on the server in every deployment (classic
    // generation is the default product), so its worker does not depend on
    // the agent runtime being enabled. It only installs a timer; its tables
    // are provisioned by the first scan.
    const { startGenerationRunner } = await import('@/lib/server/generation/run/runner');
    generationRunner = startGenerationRunner();
  } catch (error) {
    console.error('[instrumentation] Generation runner startup failed', error);
  }
  try {
    // Uploaded materials are extracted in the background from their upload
    // on (lib/server/materials/extraction.ts), wherever uploads are served.
    const { isServerPersistenceConfigured } = await import('@/lib/config/feature-flags');
    if (isServerPersistenceConfigured()) {
      const { startOwnerMaterialExtractor } = await import('@/lib/server/materials/extraction');
      materialExtractor = startOwnerMaterialExtractor();
    }
  } catch (error) {
    console.error('[instrumentation] Material extractor startup failed', error);
  }
  try {
    const { isAgentRuntimeConfigured } = await import('@/lib/config/feature-flags');
    if (isAgentRuntimeConfigured()) {
      // One dedicated LISTEN connection per application instance. The HTTP
      // SSE routes and the runner share its in-process fanout registry; it is
      // not a pool client and never scales with the number of streams.
      const { startAgentEventNotifyBus } =
        await import('@/lib/server/agent-runtime/event-notify-bus');
      const eventNotifyBus = startAgentEventNotifyBus();
      stopAgentEventNotifyBus = () => eventNotifyBus.stop();
      // startAgentRunner only installs a timer. Store/schema initialization is
      // retained behind the store's lazy promise and never blocks register().
      const runtime = await import('@/lib/server/agent-runtime/runner');
      runner = runtime.startAgentRunner();
      const extraction = await import('@/lib/server/material-extraction/runner');
      extractionRunner = extraction.startMaterialExtractionRunner();
      // Owner-level extraction of library sources (RFC #1716): what a
      // conversation links is extracted once per owner, beside the session
      // chain, which keeps serving copies made before links.
      const ownerExtraction = await import('@/lib/server/material-extraction/owner-extraction');
      ownerExtractionRunner = ownerExtraction.startOwnerExtractionRunner();
      // Two bounded passes over old material objects, one per start each;
      // neither blocks register() (see lib/server/materials/migrate-to-pool.ts).
      const materialPasses = import('@/lib/server/materials/migrate-to-pool');
      // Moving pre-pool material uploads into the asset pool is opt-in: it
      // deletes their old objects, so an operator turns it on only once every
      // instance runs this release.
      if (process.env.MATERIALS_POOL_BACKFILL === '1') {
        void materialPasses
          .then(({ migrateOwnerMaterialsToPool }) => migrateOwnerMaterialsToPool())
          .then((report) => console.info('[material-backfill] pass finished', report))
          .catch((error) => console.error('[material-backfill] pass failed', error));
      }
      // Old originals of deleted sources whose cleanup after deletion failed.
      // No new read reaches a deleted row, and this only retries a cleanup
      // the deletion already authorized, so it needs no flag.
      void materialPasses
        .then(({ removeDeletedOriginals }) => removeDeletedOriginals())
        .then((report) => console.info('[material-delete] cleanup pass finished', report))
        .catch((error) => console.error('[material-delete] cleanup pass failed', error));
    }
  } catch (error) {
    console.error('[instrumentation] Agent runtime startup failed', error);
  }

  let shutdownPromise: Promise<void> | undefined;
  const shutdown = (): Promise<void> => {
    shutdownPromise ??= (async () => {
      // Park sessions before any pool they use is closed. This preserves the
      // last durable entry-tree checkpoint for immediate takeover.
      try {
        await extractionRunner?.stop();
      } catch (error) {
        console.error('[instrumentation] Material extraction runner drain failed', error);
      }
      try {
        // Bounded: a provider call cannot be cancelled. A run still going
        // when the pool closes cannot commit anything (every write of it
        // checks its claim in a transaction of its own); its lease expires
        // and another instance claims the source again.
        const drained = await ownerExtractionRunner?.stop();
        if (drained && !drained.drained) {
          console.warn(
            `[instrumentation] Owner extraction not drained: ${drained.running} run(s) still going`,
          );
        }
      } catch (error) {
        console.error('[instrumentation] Owner extraction runner drain failed', error);
      }
      try {
        await materialExtractor?.stop();
      } catch (error) {
        console.error('[instrumentation] Material extractor drain failed', error);
      }
      try {
        await runner?.stop();
      } catch (error) {
        console.error('[instrumentation] Agent runner drain failed', error);
      }
      try {
        await generationRunner?.stop();
      } catch (error) {
        console.error('[instrumentation] Generation runner drain failed', error);
      }
      try {
        // Started at boot with the agent runtime, or on demand by the first
        // stream that subscribes to it (a generation run's events).
        await (
          stopAgentEventNotifyBus ??
          (await import('@/lib/server/agent-runtime/event-notify-bus')).stopAgentEventNotifyBus
        )();
      } catch (error) {
        console.error('[instrumentation] Agent event notify bus drain failed', error);
      }
      try {
        await legacyClassroomImport.stop();
      } catch (error) {
        console.error('[instrumentation] Legacy classroom import drain failed', error);
      }
      try {
        await assetSchedule?.stop();
      } catch (error) {
        console.error('[instrumentation] Asset collector drain failed', error);
      }
      const connectionString = process.env.DATABASE_URL?.trim();
      if (connectionString) {
        try {
          const { getServerPersistenceProvider } =
            await import('@/lib/persistence/server-provider');
          const { pool } = await getServerPersistenceProvider(connectionString);
          await pool.end();
        } catch (error) {
          console.error('[instrumentation] Persistence pool shutdown failed', error);
        }
      }
    })();
    return shutdownPromise;
  };

  process.once('SIGTERM', () => void shutdown());
  process.once('SIGINT', () => void shutdown());
}

/**
 * The fatal boot validations: each throws on a configuration the server must
 * not start with, and the throw stops the process (see `register`). Warnings (the unset ACCESS_CODE warning, the model-routing
 * checks in `validateServerConfig`) are not here and never stop the process.
 */
async function validateBootConfiguration(): Promise<void> {
  // Each check runs through runConfigurationCheck, so only what the check
  // itself throws is reported as a refused configuration; a module that fails
  // to load is reported as a startup failure, with its stack.
  const { runConfigurationCheck } = await import('@/lib/server/boot-configuration-error');

  // The database, before anything else: every course, chat and asset lives in
  // it and there is no browser-storage fallback, so a server without one would
  // boot, pass its health check and then fail every persistence request. The
  // refusal names the fix (`pnpm db:up` locally, DATABASE_URL or
  // `docker compose up` for a deployment).
  const { requireDatabaseUrl } = await import('@/lib/server/database-requirement');
  runConfigurationCheck(() => requireDatabaseUrl());

  // The deployment's model configuration: openmaic.yml (or the file named by
  // OPENMAIC_CONFIG), else the legacy provider variables translated. A
  // malformed file, an unset `${VAR}` or a slot pointing at an undeclared
  // provider is refused here with every problem listed, and so is a
  // MODEL_ROUTES left without openmaic.yml, rather than discovered by the
  // first generation. Loaded once; the notices are printed here.
  const { deploymentConfig } = await import('@/lib/server/model-config/runtime');
  runConfigurationCheck(deploymentConfig);

  // The asset quota, read here rather than at the first persistence request.
  // The provider that consumes it is lazy and memoised, so a malformed ceiling
  // would otherwise let the process boot, pass its health check, and then fail
  // every persistence request -- documents and runtime included -- until it was
  // fixed and the process restarted. `register` runs before the server is
  // ready, so throwing here is what makes a misconfigured deployment fail to
  // start instead of failing to work. First, so the throw cannot skip the
  // teardown registration for something `register` has already started.
  const { resolveAssetQuotaBytes } = await import('@/lib/persistence/asset-quota');
  runConfigurationCheck(resolveAssetQuotaBytes);

  // The pending-allocation window, for the same reason and at the same moment.
  // Too short is worse than malformed: it silently expires allocations whose
  // document write was still coming, so it must fail the process rather than
  // the request that discovers it.
  const { resolveAssetPendingTtlMs } = await import('@/lib/persistence/asset-pending-ttl');
  runConfigurationCheck(resolveAssetPendingTtlMs);

  // Owner identity, for the same reason and at the same moment. A malformed
  // PERSISTENCE_SHARED_OWNER_ID or single-user setting, the two together, or a
  // setting a host registration would ignore,
  // would otherwise boot, pass its health check, and then fail (or silently
  // mis-identify) every owner-scoped request — and an operator has no way to
  // tell from the outside that their setting was not accepted. An empty value
  // is treated as unset, so this only rejects values that are present and
  // unusable.
  //
  // A host that brings its own identity registers its owner auth methods
  // here, in the order they are asked, before validation and before the
  // server serves a request, and with them any host extension hooks (course
  // creation, library listing, upload admission, the asset byte store, and
  // the admission, execution context and notifications of generation runs):
  //
  //   const { configureOwnerAuthentication } = await import('@/lib/server/identity');
  //   configureOwnerAuthentication({ methods: [myOwnerAuthMethod] });
  //   const { configurePersistenceHooks, configureAssetByteStore } =
  //     await import('@/lib/server/persistence-hooks');
  //   configurePersistenceHooks(myPersistenceHooks);
  //   configureAssetByteStore(myAssetByteStore);
  //   const { configureGenerationRunHooks } = await import('@/lib/server/generation-run-hooks');
  //   configureGenerationRunHooks(myGenerationRunHooks);
  //
  // A registration that throws stops the process too, reported as a startup
  // failure with its stack (it is host code, not a setting).
  const { validateOwnerIdentityConfiguration } = await import('@/lib/server/identity/registry');
  runConfigurationCheck(validateOwnerIdentityConfiguration);

  // The registered hooks, for the same reason and at the same moment: a host
  // byte store that cannot sign under ASSET_BYTE_EGRESS=redirect would
  // otherwise be discovered by the first asset read.
  const { validatePersistenceHooksConfiguration } =
    await import('@/lib/server/persistence-hooks/registry');
  runConfigurationCheck(validatePersistenceHooksConfiguration);
}