stack/apps/gateway/src/agent/agent.service.ts

import { Inject, Injectable, Logger, type OnModuleDestroy } from '@nestjs/common';
import {
  createAgentSession,
  DefaultResourceLoader,
  SessionManager,
  type AgentSession as PiAgentSession,
  type AgentSessionEvent,
  type ToolDefinition,
} from '@mariozechner/pi-coding-agent';
import type { Brain } from '@mosaic/brain';
import type { Memory } from '@mosaic/memory';
import { BRAIN } from '../brain/brain.tokens.js';
import { MEMORY } from '../memory/memory.tokens.js';
import { EmbeddingService } from '../memory/embedding.service.js';
import { CoordService } from '../coord/coord.service.js';
import { ProviderService } from './provider.service.js';
import { McpClientService } from '../mcp-client/mcp-client.service.js';
import { SkillLoaderService } from './skill-loader.service.js';
import { createBrainTools } from './tools/brain-tools.js';
import { createCoordTools } from './tools/coord-tools.js';
import { createMemoryTools } from './tools/memory-tools.js';
import { createFileTools } from './tools/file-tools.js';
import { createGitTools } from './tools/git-tools.js';
import { createShellTools } from './tools/shell-tools.js';
import { createWebTools } from './tools/web-tools.js';
import type { SessionInfoDto } from './session.dto.js';
import { SessionGCService } from '../gc/session-gc.service.js';

export interface AgentSessionOptions {
  provider?: string;
  modelId?: string;
  /**
   * Sandbox working directory for the session.
   * File, git, and shell tools will be restricted to this directory.
   * Falls back to AGENT_FILE_SANDBOX_DIR env var or process.cwd().
   */
  sandboxDir?: string;
  /**
   * Platform-level system prompt for this session.
   * Merged with skill prompt additions (platform prompt first, then skills).
   * Falls back to AGENT_SYSTEM_PROMPT env var when omitted.
   */
  systemPrompt?: string;
  /**
   * Explicit allowlist of tool names available in this session.
   * When set, only listed tools are registered with the agent.
   * When omitted for non-admin users, falls back to AGENT_USER_TOOLS env var.
   * Admins (isAdmin=true) always receive the full tool set unless explicitly restricted.
   */
  allowedTools?: string[];
  /** Whether the requesting user has admin privileges. Controls default tool access. */
  isAdmin?: boolean;
  /**
   * DB agent config ID. When provided, loads agent config from DB and merges
   * provider, model, systemPrompt, and allowedTools. Explicit call-site options
   * take precedence over config values.
   */
  agentConfigId?: string;
}

export interface AgentSession {
  id: string;
  provider: string;
  modelId: string;
  piSession: PiAgentSession;
  listeners: Set<(event: AgentSessionEvent) => void>;
  unsubscribe: () => void;
  createdAt: number;
  promptCount: number;
  channels: Set<string>;
  /** System prompt additions injected from enabled prompt-type skills. */
  skillPromptAdditions: string[];
  /** Resolved sandbox directory for this session. */
  sandboxDir: string;
  /** Tool names available in this session, or null when all tools are available. */
  allowedTools: string[] | null;
}

@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>>();

  constructor(
    @Inject(ProviderService) private readonly providerService: ProviderService,
    @Inject(BRAIN) private readonly brain: Brain,
    @Inject(MEMORY) private readonly memory: Memory,
    @Inject(EmbeddingService) private readonly embeddingService: EmbeddingService,
    @Inject(CoordService) private readonly coordService: CoordService,
    @Inject(McpClientService) private readonly mcpClientService: McpClientService,
    @Inject(SkillLoaderService) private readonly skillLoaderService: SkillLoaderService,
    @Inject(SessionGCService) private readonly gc: SessionGCService,
  ) {}

  /**
   * Build the full set of custom tools scoped to the given sandbox directory.
   * Brain/coord/memory/web tools are stateless with respect to cwd; file/git/shell
   * tools receive the resolved sandboxDir so they operate within the sandbox.
   */
  private buildToolsForSandbox(sandboxDir: string): ToolDefinition[] {
    return [
      ...createBrainTools(this.brain),
      ...createCoordTools(this.coordService),
      ...createMemoryTools(
        this.memory,
        this.embeddingService.available ? this.embeddingService : null,
      ),
      ...createFileTools(sandboxDir),
      ...createGitTools(sandboxDir),
      ...createShellTools(sandboxDir),
      ...createWebTools(),
    ];
  }

  /**
   * Resolve the tool allowlist for a session.
   * - Admin users: all tools unless an explicit allowedTools list is passed.
   * - Regular users: use allowedTools if provided, otherwise parse AGENT_USER_TOOLS env var.
   * Returns null when all tools should be available.
   */
  private resolveAllowedTools(isAdmin: boolean, allowedTools?: string[]): string[] | null {
    if (allowedTools !== undefined) {
      return allowedTools.length === 0 ? [] : allowedTools;
    }
    if (isAdmin) {
      return null; // admins get everything
    }
    const envTools = process.env['AGENT_USER_TOOLS'];
    if (!envTools) {
      return null; // no restriction configured
    }
    return envTools
      .split(',')
      .map((t) => t.trim())
      .filter((t) => t.length > 0);
  }

  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> {
    // Merge DB agent config when agentConfigId is provided
    let mergedOptions = options;
    if (options?.agentConfigId) {
      const agentConfig = await this.brain.agents.findById(options.agentConfigId);
      if (agentConfig) {
        mergedOptions = {
          provider: options.provider ?? agentConfig.provider,
          modelId: options.modelId ?? agentConfig.model,
          systemPrompt: options.systemPrompt ?? agentConfig.systemPrompt ?? undefined,
          allowedTools: options.allowedTools ?? agentConfig.allowedTools ?? undefined,
          sandboxDir: options.sandboxDir,
          isAdmin: options.isAdmin,
          agentConfigId: options.agentConfigId,
        };
        this.logger.log(
          `Merged agent config "${agentConfig.name}" (${agentConfig.id}) into session ${sessionId}`,
        );
      }
    }

    const model = this.resolveModel(mergedOptions);
    const providerName = model?.provider ?? 'default';
    const modelId = model?.id ?? 'default';

    // Resolve sandbox directory: option > env var > process.cwd()
    const sandboxDir =
      mergedOptions?.sandboxDir ?? process.env['AGENT_FILE_SANDBOX_DIR'] ?? process.cwd();

    // Resolve allowed tool set
    const allowedTools = this.resolveAllowedTools(
      mergedOptions?.isAdmin ?? false,
      mergedOptions?.allowedTools,
    );

    this.logger.log(
      `Creating agent session: ${sessionId} (provider=${providerName}, model=${modelId}, sandbox=${sandboxDir}, tools=${allowedTools === null ? 'all' : allowedTools.join(',') || 'none'})`,
    );

    // Load skill tools from the catalog
    const { metaTools: skillMetaTools, promptAdditions } =
      await this.skillLoaderService.loadForSession();
    if (skillMetaTools.length > 0) {
      this.logger.log(`Attaching ${skillMetaTools.length} skill tool(s) to session ${sessionId}`);
    }
    if (promptAdditions.length > 0) {
      this.logger.log(
        `Injecting ${promptAdditions.length} skill prompt addition(s) into session ${sessionId}`,
      );
    }

    // Build per-session tools scoped to the sandbox directory
    const sandboxTools = this.buildToolsForSandbox(sandboxDir);

    // Combine static tools with dynamically discovered MCP client tools and skill tools
    const mcpTools = this.mcpClientService.getToolDefinitions();
    let allCustomTools = [...sandboxTools, ...skillMetaTools, ...mcpTools];
    if (mcpTools.length > 0) {
      this.logger.log(`Attaching ${mcpTools.length} MCP client tool(s) to session ${sessionId}`);
    }

    // Filter tools by allowlist when a restriction is in effect
    if (allowedTools !== null) {
      const allowedSet = new Set(allowedTools);
      const before = allCustomTools.length;
      allCustomTools = allCustomTools.filter((t) => allowedSet.has(t.name));
      this.logger.log(
        `Tool restriction applied: ${allCustomTools.length}/${before} tools allowed for session ${sessionId}`,
      );
    }

    // Build system prompt: platform prompt + skill additions appended
    const platformPrompt =
      mergedOptions?.systemPrompt ?? process.env['AGENT_SYSTEM_PROMPT'] ?? undefined;
    const appendSystemPrompt =
      promptAdditions.length > 0 ? promptAdditions.join('\n\n') : undefined;

    // Construct a resource loader that injects the configured system prompt
    const resourceLoader = new DefaultResourceLoader({
      cwd: sandboxDir,
      noExtensions: true,
      noSkills: true,
      noPromptTemplates: true,
      noThemes: true,
      systemPrompt: platformPrompt,
      appendSystemPrompt: appendSystemPrompt,
    });
    await resourceLoader.reload();

    let piSession: PiAgentSession;
    try {
      const result = await createAgentSession({
        sessionManager: SessionManager.inMemory(),
        modelRegistry: this.providerService.getRegistry(),
        model: model ?? undefined,
        cwd: sandboxDir,
        tools: [],
        customTools: allCustomTools,
        resourceLoader,
      });
      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,
      createdAt: Date.now(),
      promptCount: 0,
      channels: new Set(),
      skillPromptAdditions: promptAdditions,
      sandboxDir,
      allowedTools,
    };

    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);
  }

  listSessions(): SessionInfoDto[] {
    const now = Date.now();
    return Array.from(this.sessions.values()).map((s) => ({
      id: s.id,
      provider: s.provider,
      modelId: s.modelId,
      createdAt: new Date(s.createdAt).toISOString(),
      promptCount: s.promptCount,
      channels: Array.from(s.channels),
      durationMs: now - s.createdAt,
    }));
  }

  getSessionInfo(sessionId: string): SessionInfoDto | undefined {
    const s = this.sessions.get(sessionId);
    if (!s) return undefined;
    return {
      id: s.id,
      provider: s.provider,
      modelId: s.modelId,
      createdAt: new Date(s.createdAt).toISOString(),
      promptCount: s.promptCount,
      channels: Array.from(s.channels),
      durationMs: Date.now() - s.createdAt,
    };
  }

  addChannel(sessionId: string, channel: string): void {
    const session = this.sessions.get(sessionId);
    if (session) {
      session.channels.add(channel);
    }
  }

  removeChannel(sessionId: string, channel: string): void {
    const session = this.sessions.get(sessionId);
    if (session) {
      session.channels.delete(channel);
    }
  }

  async prompt(sessionId: string, message: string): Promise<void> {
    const session = this.sessions.get(sessionId);
    if (!session) {
      throw new Error(`No agent session found: ${sessionId}`);
    }
    session.promptCount += 1;
    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));
    }
    try {
      session.piSession.dispose();
    } catch (err) {
      this.logger.error(`Failed to dispose piSession for ${sessionId}`, String(err));
    }
    session.listeners.clear();
    session.channels.clear();
    this.sessions.delete(sessionId);

    // Run GC cleanup for this session (fire and forget, errors are logged)
    this.gc.collect(sessionId).catch((err: unknown) => {
      this.logger.error(
        `GC collect failed for session ${sessionId}`,
        err instanceof Error ? err.stack : String(err),
      );
    });
  }

  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));
      }
    }
  }
}