stack/apps/orchestrator/src/api/agents/agent-provider.registry.ts

import { Injectable, Logger, OnModuleInit } from "@nestjs/common";
import type { AgentSession, IAgentProvider } from "@mosaic/shared";
import { InternalAgentProvider } from "./internal-agent.provider";

@Injectable()
export class AgentProviderRegistry implements OnModuleInit {
  private readonly logger = new Logger(AgentProviderRegistry.name);
  private readonly providers = new Map<string, IAgentProvider>();

  constructor(private readonly internalProvider: InternalAgentProvider) {}

  onModuleInit(): void {
    this.registerProvider(this.internalProvider);
  }

  registerProvider(provider: IAgentProvider): void {
    const existingProvider = this.providers.get(provider.providerId);
    if (existingProvider !== undefined) {
      this.logger.warn(`Replacing existing provider registration for ${provider.providerId}`);
    }

    this.providers.set(provider.providerId, provider);
  }

  getProvider(providerId: string): IAgentProvider | null {
    return this.providers.get(providerId) ?? null;
  }

  async getProviderForSession(
    sessionId: string
  ): Promise<{ provider: IAgentProvider; session: AgentSession } | null> {
    for (const provider of this.providers.values()) {
      try {
        const session = await provider.getSession(sessionId);
        if (session !== null) {
          return {
            provider,
            session,
          };
        }
      } catch (error) {
        this.logger.warn(
          `Failed to get session ${sessionId} for provider ${provider.providerId}: ${this.toErrorMessage(error)}`
        );
      }
    }

    return null;
  }

  async listAllSessions(): Promise<AgentSession[]> {
    const providers = [...this.providers.values()];
    const sessionsByProvider = await Promise.all(
      providers.map(async (provider) => {
        try {
          const { sessions } = await provider.listSessions();
          return sessions;
        } catch (error) {
          this.logger.warn(
            `Failed to list sessions for provider ${provider.providerId}: ${this.toErrorMessage(error)}`
          );
          return [];
        }
      })
    );

    return sessionsByProvider
      .flat()
      .sort((left, right) => right.updatedAt.getTime() - left.updatedAt.getTime());
  }

  private toErrorMessage(error: unknown): string {
    if (error instanceof Error) {
      return error.message;
    }

    return String(error);
  }
}