@boundless-oss/atlas/src/modules/memory-management/code-memory.ts

Atlas - MCP Server for comprehensive startup project management

import { ConfigManager } from '../../config/config-manager.js';
import { CodeContext } from './types.js';
import { promises as fs } from 'fs';
import path from 'path';

export class CodeMemory {
  private configManager: ConfigManager;
  private codeContextPath: string;
  private codeContexts: Map<string, CodeContext>;

  constructor(configManager: ConfigManager) {
    this.configManager = configManager;
    this.codeContextPath = '';
    this.codeContexts = new Map();
  }

  async initialize(): Promise<void> {
    const storageManager = this.configManager.getStorageManager();
    const location = await storageManager.getStorageLocation();
    
    this.codeContextPath = path.join(location.data, 'memory', 'code-contexts.json');
    
    // Ensure directory exists
    await fs.mkdir(path.dirname(this.codeContextPath), { recursive: true });
    
    // Load existing code contexts
    await this.loadCodeContexts();
  }

  async analyzeContext(options: {
    filePath: string;
    functionName?: string;
    includeRelated?: boolean;
    depth?: number;
  }): Promise<CodeContext> {
    const fullPath = path.resolve(options.filePath);
    const contextKey = options.functionName ? `${fullPath}:${options.functionName}` : fullPath;

    // Check if we have cached context
    const cached = this.codeContexts.get(contextKey);
    if (cached) {
      return cached;
    }

    // Analyze the code file
    const context = await this.performCodeAnalysis(options);
    
    // Cache the context
    this.codeContexts.set(contextKey, context);
    await this.saveCodeContexts();

    return context;
  }

  async updateCodeContext(filePath: string, context: CodeContext): Promise<void> {
    this.codeContexts.set(filePath, context);
    await this.saveCodeContexts();
  }

  async findRelatedCode(filePath: string): Promise<CodeContext[]> {
    const targetContext = this.codeContexts.get(filePath);
    if (!targetContext) return [];

    const related: { context: CodeContext; score: number }[] = [];

    for (const [path, context] of this.codeContexts) {
      if (path === filePath) continue;

      let score = 0;

      // Check if files import each other
      const hasImportRelation = context.dependencies.some(dep => 
        dep.filePath === filePath
      ) || targetContext.dependencies.some(dep => 
        dep.filePath === path.replace(/:[^:]*$/, '') // Remove function name if present
      );

      if (hasImportRelation) {
        score += 0.5;
      }

      // Check for similar usage patterns
      const sharedPatterns = targetContext.usagePatterns.filter(pattern =>
        context.usagePatterns.includes(pattern)
      );
      score += sharedPatterns.length * 0.2;

      // Check for related code references
      const relatedCodeMatch = targetContext.relatedCode.some(related =>
        related.file === context.filePath
      );
      if (relatedCodeMatch) {
        score += 0.3;
      }

      if (score > 0.2) {
        related.push({ context, score });
      }
    }

    return related
      .sort((a, b) => b.score - a.score)
      .slice(0, 5)
      .map(item => item.context);
  }

  async getCodeStats(): Promise<{
    totalContexts: number;
    fileTypes: Record<string, number>;
    topPatterns: Array<{ pattern: string; count: number }>;
    dependencyGraph: Record<string, string[]>;
  }> {
    const contexts = Array.from(this.codeContexts.values());
    
    const fileTypes: Record<string, number> = {};
    const patternCounts: Record<string, number> = {};
    const dependencyGraph: Record<string, string[]> = {};

    contexts.forEach(context => {
      const ext = path.extname(context.filePath);
      fileTypes[ext] = (fileTypes[ext] || 0) + 1;

      context.usagePatterns.forEach(pattern => {
        patternCounts[pattern] = (patternCounts[pattern] || 0) + 1;
      });

      const deps = context.dependencies
        .filter(dep => dep.filePath)
        .map(dep => dep.filePath!);
      
      if (deps.length > 0) {
        dependencyGraph[context.filePath] = deps;
      }
    });

    const topPatterns = Object.entries(patternCounts)
      .sort(([, a], [, b]) => b - a)
      .slice(0, 10)
      .map(([pattern, count]) => ({ pattern, count }));

    return {
      totalContexts: contexts.length,
      fileTypes,
      topPatterns,
      dependencyGraph,
    };
  }

  private async performCodeAnalysis(options: {
    filePath: string;
    functionName?: string;
    includeRelated?: boolean;
    depth?: number;
  }): Promise<CodeContext> {
    try {
      const content = await fs.readFile(options.filePath, 'utf-8');
      const ext = path.extname(options.filePath);

      let context: CodeContext;

      switch (ext) {
        case '.ts':
        case '.tsx':
        case '.js':
        case '.jsx':
          context = await this.analyzeJavaScriptTypeScript(content, options);
          break;
        case '.py':
          context = await this.analyzePython(content, options);
          break;
        default:
          context = await this.analyzeGeneric(content, options);
      }

      // Add related code if requested
      if (options.includeRelated) {
        context.relatedCode = await this.findRelatedCodeReferences(context, options.depth || 2);
      }

      return context;
    } catch (error) {
      // Return minimal context if analysis fails
      return {
        filePath: options.filePath,
        functionName: options.functionName,
        summary: `Failed to analyze: ${(error as Error).message}`,
        dependencies: [],
        usagePatterns: [],
        relatedCode: [],
        recommendations: ['File could not be analyzed - check file permissions and format'],
      };
    }
  }

  private async analyzeJavaScriptTypeScript(
    content: string, 
    options: { filePath: string; functionName?: string }
  ): Promise<CodeContext> {
    const lines = content.split('\n');
    const dependencies: CodeContext['dependencies'] = [];
    const usagePatterns: string[] = [];
    const recommendations: string[] = [];

    // Extract imports
    const importRegex = /^import\s+(?:.*\s+from\s+)?['"]([^'"]+)['"];?/gm;
    let match;
    while ((match = importRegex.exec(content)) !== null) {
      dependencies.push({
        name: match[1],
        type: 'import',
        description: `Imported module: ${match[1]}`,
        filePath: this.resolveImportPath(match[1], options.filePath),
      });
    }

    // Extract function/class definitions
    const functionRegex = /^(?:export\s+)?(?:async\s+)?function\s+(\w+)/gm;
    while ((match = functionRegex.exec(content)) !== null) {
      dependencies.push({
        name: match[1],
        type: 'function',
        description: `Function: ${match[1]}`,
      });
    }

    const classRegex = /^(?:export\s+)?class\s+(\w+)/gm;
    while ((match = classRegex.exec(content)) !== null) {
      dependencies.push({
        name: match[1],
        type: 'class',
        description: `Class: ${match[1]}`,
      });
    }

    // Detect usage patterns
    if (content.includes('React')) usagePatterns.push('React component');
    if (content.includes('express')) usagePatterns.push('Express.js server');
    if (content.includes('async/await')) usagePatterns.push('Async/await pattern');
    if (content.includes('jest') || content.includes('describe(')) usagePatterns.push('Jest testing');
    if (content.includes('useState') || content.includes('useEffect')) usagePatterns.push('React hooks');
    if (content.includes('try/catch')) usagePatterns.push('Error handling');

    // Generate recommendations
    if (!content.includes('use strict') && content.includes('function')) {
      recommendations.push('Consider using strict mode');
    }
    if (content.includes('any') && options.filePath.endsWith('.ts')) {
      recommendations.push('Avoid using "any" type in TypeScript');
    }
    if (content.includes('console.log')) {
      recommendations.push('Remove console.log statements before production');
    }

    let summary = `${path.basename(options.filePath)} - `;
    if (usagePatterns.length > 0) {
      summary += usagePatterns[0];
    } else {
      summary += 'JavaScript/TypeScript file';
    }

    if (options.functionName) {
      const functionMatch = content.match(new RegExp(`function\\s+${options.functionName}\\s*\\(`));
      if (functionMatch) {
        summary += ` (analyzing function: ${options.functionName})`;
      }
    }

    return {
      filePath: options.filePath,
      functionName: options.functionName,
      summary,
      dependencies,
      usagePatterns,
      relatedCode: [],
      recommendations,
    };
  }

  private async analyzePython(
    content: string,
    options: { filePath: string; functionName?: string }
  ): Promise<CodeContext> {
    const dependencies: CodeContext['dependencies'] = [];
    const usagePatterns: string[] = [];
    const recommendations: string[] = [];

    // Extract imports
    const importRegex = /^(?:from\s+(\S+)\s+)?import\s+(.+)$/gm;
    let match;
    while ((match = importRegex.exec(content)) !== null) {
      const module = match[1] || match[2].split(',')[0].trim();
      dependencies.push({
        name: module,
        type: 'import',
        description: `Imported: ${module}`,
      });
    }

    // Extract function definitions
    const functionRegex = /^def\s+(\w+)\s*\(/gm;
    while ((match = functionRegex.exec(content)) !== null) {
      dependencies.push({
        name: match[1],
        type: 'function',
        description: `Function: ${match[1]}`,
      });
    }

    // Extract class definitions
    const classRegex = /^class\s+(\w+)/gm;
    while ((match = classRegex.exec(content)) !== null) {
      dependencies.push({
        name: match[1],
        type: 'class',
        description: `Class: ${match[1]}`,
      });
    }

    // Detect patterns
    if (content.includes('flask')) usagePatterns.push('Flask web application');
    if (content.includes('django')) usagePatterns.push('Django web application');
    if (content.includes('pandas')) usagePatterns.push('Data analysis with pandas');
    if (content.includes('numpy')) usagePatterns.push('Numerical computing');
    if (content.includes('async def')) usagePatterns.push('Async/await pattern');
    if (content.includes('unittest') || content.includes('pytest')) usagePatterns.push('Unit testing');

    return {
      filePath: options.filePath,
      functionName: options.functionName,
      summary: `${path.basename(options.filePath)} - Python ${usagePatterns[0] || 'script'}`,
      dependencies,
      usagePatterns,
      relatedCode: [],
      recommendations,
    };
  }

  private async analyzeGeneric(
    content: string,
    options: { filePath: string; functionName?: string }
  ): Promise<CodeContext> {
    const ext = path.extname(options.filePath);
    const dependencies: CodeContext['dependencies'] = [];
    const usagePatterns: string[] = [];

    // Basic analysis for any file type
    const lineCount = content.split('\n').length;
    const wordCount = content.split(/\s+/).length;

    if (ext === '.md') {
      usagePatterns.push('Markdown documentation');
    } else if (ext === '.json') {
      usagePatterns.push('JSON configuration');
    } else if (ext === '.yml' || ext === '.yaml') {
      usagePatterns.push('YAML configuration');
    }

    return {
      filePath: options.filePath,
      functionName: options.functionName,
      summary: `${path.basename(options.filePath)} - ${lineCount} lines, ${wordCount} words`,
      dependencies,
      usagePatterns,
      relatedCode: [],
      recommendations: [],
    };
  }

  private async findRelatedCodeReferences(
    context: CodeContext,
    depth: number
  ): Promise<CodeContext['relatedCode']> {
    const related: CodeContext['relatedCode'] = [];

    // Find files that import this file or are imported by this file
    for (const dep of context.dependencies) {
      if (dep.filePath && dep.type === 'import') {
        try {
          await fs.access(dep.filePath);
          related.push({
            file: dep.filePath,
            description: `Imported dependency: ${dep.name}`,
            relationship: 'imports',
          });
        } catch {
          // File doesn't exist or can't be accessed
        }
      }
    }

    // Look for files that might import this file
    if (depth > 1) {
      // This would require scanning other files in the project
      // For now, we'll keep it simple
    }

    return related.slice(0, 10); // Limit to 10 related files
  }

  private resolveImportPath(importPath: string, currentFile: string): string | undefined {
    if (importPath.startsWith('.')) {
      // Relative import
      const currentDir = path.dirname(currentFile);
      return path.resolve(currentDir, importPath);
    }
    
    // For node_modules or absolute imports, we'd need more sophisticated resolution
    return undefined;
  }

  private async loadCodeContexts(): Promise<void> {
    try {
      const data = await fs.readFile(this.codeContextPath, 'utf-8');
      const contextsArray: CodeContext[] = JSON.parse(data);
      
      this.codeContexts.clear();
      for (const context of contextsArray) {
        const key = context.functionName ? `${context.filePath}:${context.functionName}` : context.filePath;
        this.codeContexts.set(key, context);
      }
    } catch (error) {
      // File doesn't exist or is invalid, start with empty contexts
      this.codeContexts.clear();
    }
  }

  private async saveCodeContexts(): Promise<void> {
    const contextsArray = Array.from(this.codeContexts.values());
    await fs.writeFile(this.codeContextPath, JSON.stringify(contextsArray, null, 2));
  }
}