stack-mos-dt-0/apps/gateway/src/agent/agent.service.ts

import { Inject, Injectable, Logger, type OnModuleDestroy } from '@nestjs/common';
import {
  createAgentSession,
  SessionManager,
  type AgentSession as PiAgentSession,
  type AgentSessionEvent,
  type ToolDefinition,
} from '@mariozechner/pi-coding-agent';
import type { Brain } from '@mosaic/brain';
import { BRAIN } from '../brain/brain.tokens.js';
import { ProviderService } from './provider.service.js';
import { createBrainTools } from './tools/brain-tools.js';

export interface AgentSessionOptions {
  provider?: string;
  modelId?: string;
}

export interface AgentSession {
  id: string;
  provider: string;
  modelId: string;
  piSession: PiAgentSession;
  listeners: Set<(event: AgentSessionEvent) => void>;
  unsubscribe: () => void;
}

@Injectable()
export class AgentService implements OnModuleDestroy {
  private readonly logger = new Logger(AgentService.name);
  private readonly sessions = new Map<string, AgentSession>();
  private readonly creating = new Map<string, Promise<AgentSession>>();

  private readonly customTools: ToolDefinition[];

  constructor(
    private readonly providerService: ProviderService,
    @Inject(BRAIN) private readonly brain: Brain,
  ) {
    this.customTools = createBrainTools(brain);
    this.logger.log(`Registered ${this.customTools.length} custom tools`);
  }

  async createSession(sessionId: string, options?: AgentSessionOptions): Promise<AgentSession> {
    const existing = this.sessions.get(sessionId);
    if (existing) return existing;

    const inflight = this.creating.get(sessionId);
    if (inflight) return inflight;

    const promise = this.doCreateSession(sessionId, options).finally(() => {
      this.creating.delete(sessionId);
    });
    this.creating.set(sessionId, promise);
    return promise;
  }

  private async doCreateSession(
    sessionId: string,
    options?: AgentSessionOptions,
  ): Promise<AgentSession> {
    const model = this.resolveModel(options);
    const providerName = model?.provider ?? 'default';
    const modelId = model?.id ?? 'default';

    this.logger.log(
      `Creating agent session: ${sessionId} (provider=${providerName}, model=${modelId})`,
    );

    let piSession: PiAgentSession;
    try {
      const result = await createAgentSession({
        sessionManager: SessionManager.inMemory(),
        modelRegistry: this.providerService.getRegistry(),
        model: model ?? undefined,
        tools: [],
        customTools: this.customTools,
      });
      piSession = result.session;
    } catch (err) {
      this.logger.error(
        `Failed to create agent session for ${sessionId}`,
        err instanceof Error ? err.stack : String(err),
      );
      throw new Error(`Agent session creation failed for ${sessionId}: ${String(err)}`);
    }

    const listeners = new Set<(event: AgentSessionEvent) => void>();

    const unsubscribe = piSession.subscribe((event) => {
      for (const listener of listeners) {
        try {
          listener(event);
        } catch (err) {
          this.logger.error(`Event listener error in session ${sessionId}`, err);
        }
      }
    });

    const session: AgentSession = {
      id: sessionId,
      provider: providerName,
      modelId,
      piSession,
      listeners,
      unsubscribe,
    };

    this.sessions.set(sessionId, session);
    this.logger.log(`Agent session ${sessionId} ready (${providerName}/${modelId})`);

    return session;
  }

  private resolveModel(options?: AgentSessionOptions) {
    if (!options?.provider && !options?.modelId) {
      return this.providerService.getDefaultModel() ?? null;
    }

    if (options.provider && options.modelId) {
      const model = this.providerService.findModel(options.provider, options.modelId);
      if (!model) {
        throw new Error(`Model not found: ${options.provider}/${options.modelId}`);
      }
      return model;
    }

    if (options.modelId) {
      const available = this.providerService.listAvailableModels();
      const match = available.find((m) => m.id === options.modelId);
      if (match) {
        return this.providerService.findModel(match.provider, match.id) ?? null;
      }
    }

    return this.providerService.getDefaultModel() ?? null;
  }

  getSession(sessionId: string): AgentSession | undefined {
    return this.sessions.get(sessionId);
  }

  async prompt(sessionId: string, message: string): Promise<void> {
    const session = this.sessions.get(sessionId);
    if (!session) {
      throw new Error(`No agent session found: ${sessionId}`);
    }
    try {
      await session.piSession.prompt(message);
    } catch (err) {
      this.logger.error(
        `Prompt failed for session=${sessionId}, messageLength=${message.length}`,
        err instanceof Error ? err.stack : String(err),
      );
      throw err;
    }
  }

  onEvent(sessionId: string, listener: (event: AgentSessionEvent) => void): () => void {
    const session = this.sessions.get(sessionId);
    if (!session) {
      throw new Error(`No agent session found: ${sessionId}`);
    }
    session.listeners.add(listener);
    return () => session.listeners.delete(listener);
  }

  async destroySession(sessionId: string): Promise<void> {
    const session = this.sessions.get(sessionId);
    if (!session) return;
    this.logger.log(`Destroying agent session ${sessionId}`);
    try {
      session.unsubscribe();
    } catch (err) {
      this.logger.error(`Failed to unsubscribe session ${sessionId}`, String(err));
    }
    session.listeners.clear();
    this.sessions.delete(sessionId);
  }

  async onModuleDestroy(): Promise<void> {
    this.logger.log('Shutting down all agent sessions');
    const stops = Array.from(this.sessions.keys()).map((id) => this.destroySession(id));
    const results = await Promise.allSettled(stops);
    for (const result of results) {
      if (result.status === 'rejected') {
        this.logger.error('Session shutdown failure', String(result.reason));
      }
    }
  }
}