llm-to-agent/src/kernel/kernel.ts

import { EventBus, type EventHandler, type Unsubscribe } from "./events";
import type {
  Extension,
  ExtensionConfig,
  ExtensionFactory,
  ExtensionRuntimeContext,
  ExtensionSetupContext,
} from "./extension";
import { ExtensionRegistry } from "./registry";

export interface KernelOptions {
  extensionConfigs?: Record<string, ExtensionConfig>;
}

interface InstalledExtension {
  id: string;
  extension: Extension;
}

export class Kernel {
  #extensions: InstalledExtension[] = [];
  #activeExtensions: Extension[] = [];
  #registry = new ExtensionRegistry();
  #events = new EventBus();
  #setupContext: ExtensionSetupContext;
  #runtimeContext: ExtensionRuntimeContext;
  #extensionConfigs: Record<string, ExtensionConfig>;
  #state = "created";

  constructor(options: KernelOptions = {}) {
    const registry = this.#registry;
    const events = this.#events;
    this.#extensionConfigs = options.extensionConfigs ?? {};

    this.#setupContext = {
      add(name, value) {
        registry.add(name, value);
      },
      on(event, handler) {
        return events.on(event, handler);
      },
    };

    this.#runtimeContext = {
      get(name) {
        return registry.get(name);
      },
      all(name) {
        return registry.all(name);
      },
      on(event, handler) {
        return events.on(event, handler);
      },
      emit(event, payload) {
        return events.emit(event, payload);
      },
    };
  }

  get state(): string {
    return this.#state;
  }

  get installedExtensionIds(): string[] {
    return this.#extensions.map(({ id }) => id);
  }

  use(...factories: ExtensionFactory[]): this {
    if (this.#state !== "created") {
      throw new Error(`Cannot install extensions while kernel is ${this.#state}.`);
    }

    const ids = new Set(this.#extensions.map(({ id }) => id));

    for (const factory of factories) {
      if (ids.has(factory.id)) {
        throw new Error(`Extension "${factory.id}" is already installed.`);
      }
      ids.add(factory.id);
    }

    const extensions = factories.map((factory) => ({
      id: factory.id,
      extension: factory({ ...this.#extensionConfigs[factory.id] }),
    }));
    this.#extensions.push(...extensions);
    return this;
  }

  get<T>(name: string): T {
    return this.#registry.get(name);
  }

  all<T>(name: string): T[] {
    return this.#registry.all(name);
  }

  on<T>(event: string, handler: EventHandler<T>): Unsubscribe {
    return this.#events.on(event, handler);
  }

  emit<T>(event: string, payload: T): Promise<unknown[]> {
    return this.#events.emit(event, payload);
  }

  async start(): Promise<void> {
    if (this.#state !== "created") {
      throw new Error(`Cannot start kernel while it is ${this.#state}.`);
    }

    this.#state = "starting";

    try {
      for (const { extension } of this.#extensions) {
        await extension.setup(this.#setupContext);
      }

      for (const { extension } of this.#extensions) {
        this.#activeExtensions.push(extension);
        await extension.start?.(this.#runtimeContext);
      }

      this.#state = "running";
    } catch (error) {
      const cleanupErrors = await this.#stopActiveExtensions();
      this.#events.clear();
      this.#registry.clear();
      this.#state = "stopped";

      if (cleanupErrors.length > 0) {
        throw new AggregateError(
          [error, ...cleanupErrors],
          "Kernel failed to start and one or more extensions failed to clean up.",
        );
      }

      throw error;
    }
  }

  async stop(): Promise<void> {
    if (this.#state === "stopped") return;

    if (this.#state !== "running") {
      throw new Error(`Cannot stop kernel while it is ${this.#state}.`);
    }

    this.#state = "stopping";
    const errors = await this.#stopActiveExtensions();
    this.#events.clear();
    this.#registry.clear();
    this.#state = "stopped";

    if (errors.length > 0) {
      throw new AggregateError(errors, "One or more extensions failed to stop.");
    }
  }

  async #stopActiveExtensions(): Promise<unknown[]> {
    const errors: unknown[] = [];

    for (const extension of [...this.#activeExtensions].reverse()) {
      try {
        await extension.stop?.(this.#runtimeContext);
      } catch (error) {
        errors.push(error);
      }
    }

    this.#activeExtensions = [];
    return errors;
  }
}