packages/core/src/context/memoryContextManager.ts

/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import {
  loadJitSubdirectoryMemory,
  concatenateInstructions,
  getGlobalMemoryPaths,
  getUserProjectMemoryPaths,
  getExtensionMemoryPaths,
  getEnvironmentMemoryPaths,
  readGeminiMdFiles,
  categorizeAndConcatenate,
  type GeminiFileContent,
  deduplicatePathsByFileIdentity,
} from '../utils/memoryDiscovery.js';
import type { Config } from '../config/config.js';
import { coreEvents, CoreEvent } from '../utils/events.js';

export class MemoryContextManager {
  private readonly loadedPaths: Set<string> = new Set();
  private readonly loadedFileIdentities: Set<string> = new Set();
  private readonly config: Config;
  private globalMemory: string = '';
  private extensionMemory: string = '';
  private projectMemory: string = '';
  private userProjectMemoryContent: string = '';

  constructor(config: Config) {
    this.config = config;
  }

  /**
   * Refreshes the memory by reloading global, extension, and project memory.
   */
  async refresh(): Promise<void> {
    this.loadedPaths.clear();
    this.loadedFileIdentities.clear();

    const paths = await this.discoverMemoryPaths();
    const contentsMap = await this.loadMemoryContents(paths);

    this.categorizeMemoryContents(paths, contentsMap);
    this.emitMemoryChanged();
  }

  private async discoverMemoryPaths() {
    const [global, extension, project, userProjectMemory] = await Promise.all([
      getGlobalMemoryPaths(),
      Promise.resolve(
        getExtensionMemoryPaths(this.config.getExtensionLoader()),
      ),
      this.config.isTrustedFolder()
        ? getEnvironmentMemoryPaths(
            [...this.config.getWorkspaceContext().getDirectories()],
            this.config.getMemoryBoundaryMarkers(),
          )
        : Promise.resolve([]),
      getUserProjectMemoryPaths(this.config.storage.getProjectMemoryDir()),
    ]);

    return { global, extension, project, userProjectMemory };
  }

  private async loadMemoryContents(paths: {
    global: string[];
    extension: string[];
    project: string[];
    userProjectMemory: string[];
  }) {
    const allPathsStringDeduped = Array.from(
      new Set([
        ...paths.global,
        ...paths.extension,
        ...paths.project,
        ...paths.userProjectMemory,
      ]),
    );

    // deduplicate by file identity to handle case-insensitive filesystems
    const { paths: allPaths, identityMap: pathIdentityMap } =
      await deduplicatePathsByFileIdentity(allPathsStringDeduped);

    const allContents = await readGeminiMdFiles(
      allPaths,
      this.config.getImportFormat(),
      this.config.getMemoryBoundaryMarkers(),
    );

    const loadedFilePaths = allContents
      .filter((c) => c.content !== null)
      .map((c) => c.filePath);
    this.markAsLoaded(loadedFilePaths);

    // Cache file identities for performance optimization
    for (const filePath of loadedFilePaths) {
      const identity = pathIdentityMap.get(filePath);
      if (identity) {
        this.loadedFileIdentities.add(identity);
      }
    }

    return new Map(allContents.map((c) => [c.filePath, c]));
  }

  private categorizeMemoryContents(
    paths: {
      global: string[];
      extension: string[];
      project: string[];
      userProjectMemory: string[];
    },
    contentsMap: Map<string, GeminiFileContent>,
  ) {
    const hierarchicalMemory = categorizeAndConcatenate(paths, contentsMap);

    this.globalMemory = hierarchicalMemory.global || '';
    this.extensionMemory = hierarchicalMemory.extension || '';
    this.userProjectMemoryContent = hierarchicalMemory.userProjectMemory || '';

    const mcpInstructions =
      this.config.getMcpClientManager()?.getMcpInstructions() || '';
    const projectMemoryWithMcp = [
      hierarchicalMemory.project,
      mcpInstructions.trimStart(),
    ]
      .filter(Boolean)
      .join('\n\n');

    this.projectMemory = this.config.isTrustedFolder()
      ? projectMemoryWithMcp
      : '';
  }

  /**
   * Discovers and loads context for a specific accessed path (Tier 3 - JIT).
   * Traverses upwards from the accessed path to the project root.
   */
  async discoverContext(
    accessedPath: string,
    trustedRoots: string[],
  ): Promise<string> {
    if (!this.config.isTrustedFolder()) {
      return '';
    }
    const result = await loadJitSubdirectoryMemory(
      accessedPath,
      trustedRoots,
      this.loadedPaths,
      this.loadedFileIdentities,
      this.config.getMemoryBoundaryMarkers(),
    );

    if (result.files.length === 0) {
      return '';
    }

    const newFilePaths = result.files.map((f) => f.path);
    this.markAsLoaded(newFilePaths);

    // Cache identities for newly loaded files
    if (result.fileIdentities) {
      for (const identity of result.fileIdentities) {
        this.loadedFileIdentities.add(identity);
      }
    }
    return concatenateInstructions(
      result.files.map((f) => ({ filePath: f.path, content: f.content })),
    );
  }

  private emitMemoryChanged(): void {
    coreEvents.emit(CoreEvent.MemoryChanged, {
      fileCount: this.loadedPaths.size,
    });
  }

  getGlobalMemory(): string {
    return this.globalMemory;
  }

  getExtensionMemory(): string {
    return this.extensionMemory;
  }

  getEnvironmentMemory(): string {
    return this.projectMemory;
  }

  getUserProjectMemory(): string {
    return this.userProjectMemoryContent;
  }

  private markAsLoaded(paths: string[]): void {
    paths.forEach((p) => this.loadedPaths.add(p));
  }

  getLoadedPaths(): ReadonlySet<string> {
    return this.loadedPaths;
  }
}