litellm/tests/code_coverage_tests/memory_test.py

"""
Memory Violation Detection Test

Detects bad memory patterns in the LiteLLM codebase that can lead to memory leaks or OOMs.

The detector uses a modular pattern-based system. To add detection for new memory patterns:

1. Create a Pattern subclass implementing get_pattern_name(), visit_assign(), and check_cleanup()
   - You can extend the Pattern class with additional methods as needed for your detection logic
2. Add the pattern to MemoryViolationDetector.DEFAULT_PATTERNS

Currently detects:
- queue.get() / queue.get_nowait() operations where variables aren't set to None
- Class-level data structures that have add operations during runtime without size limits:
  * Built-in: list, dict, set
  * Collections: deque, defaultdict, Counter, OrderedDict, ChainMap
  * Queues: queue.Queue, asyncio.Queue (if unbounded, i.e., no maxsize parameter)
  * Heap operations: heapq.heappush(), heapq.heapreplace(), heapq.heappushpop() on class-level lists
"""

import ast
import os
from abc import ABC, abstractmethod
from typing import List, Dict, Any, Optional, Sequence


class Pattern(ABC):
    """Base class for memory violation detection patterns"""
    
    @abstractmethod
    def get_pattern_name(self) -> str:
        """Return unique identifier for this violation type"""
        pass
    
    @abstractmethod
    def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Detect memory-sensitive operations in assignment. Returns list of {line, var_name, call} dicts."""
        pass
    
    @abstractmethod
    def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt], 
                     context: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Verify variables are set to None. Returns list of violation dicts."""
        pass


class QueueGetPattern(Pattern):
    """Detects queue.get()/get_nowait() operations that aren't cleared"""
    
    def get_pattern_name(self) -> str:
        return "queue_reference_not_cleared"
    
    def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Detect queue.get() or queue.get_nowait() calls where object name contains 'queue'"""
        operations = []
        
        if isinstance(node.value, ast.Call):
            func = node.value.func
            if isinstance(func, ast.Attribute) and func.attr in ("get", "get_nowait"):
                obj_name = context["get_attr_string"](func.value)
                if "queue" in obj_name.lower() and node.targets and isinstance(node.targets[0], ast.Name):
                    operations.append({
                        "line": node.lineno,
                        "var_name": node.targets[0].id,
                        "call": context["get_call_string"](node.value),
                    })
        
        return operations
    
    def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt],
                     context: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Flag queue variables that aren't set to None"""
        violations = []
        is_var_set_to_none = context["is_var_set_to_none"]
        current_function = context["current_function"]
        file_path = context["file_path"]
        
        queue_vars = {op["var_name"]: op["line"] for op in operations}
        
        for var_name, line_num in queue_vars.items():
            if not is_var_set_to_none(var_name, function_body):
                violations.append({
                    "line": line_num,
                    "type": self.get_pattern_name(),
                    "var_name": var_name,
                    "function": current_function,
                    "file_path": file_path,
                    "message": (
                        f"Queue variable '{var_name}' in function "
                        f"'{current_function}' is not set to None after use. "
                        f"If the runtime is overwhelmed, this can cause OOM (Out of Memory) errors."
                    ),
                })
        
        return violations


class UnboundedDataStructurePattern(Pattern):
    """Detects class-level data structures (lists, dicts, sets) that can grow unbounded"""
    
    def get_pattern_name(self) -> str:
        return "unbounded_data_structure"
    
    def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Detect list/dict/set creations that are at class level"""
        operations = []
        
        # Check if this is a data structure creation
        is_data_structure = False
        structure_type = None
        
        if isinstance(node.value, (ast.List, ast.Dict, ast.Set)):
            is_data_structure = True
            if isinstance(node.value, ast.List):
                structure_type = "list"
            elif isinstance(node.value, ast.Dict):
                structure_type = "dict"
            elif isinstance(node.value, ast.Set):
                structure_type = "set"
        elif isinstance(node.value, ast.Call):
            # Check for list(), dict(), set() calls
            func = node.value.func
            if isinstance(func, ast.Name):
                if func.id in ("list", "dict", "set"):
                    is_data_structure = True
                    structure_type = func.id
            elif isinstance(func, ast.Attribute):
                # Handle cases like collections.defaultdict(list), collections.deque(), etc.
                obj_name = context["get_attr_string"](func.value)
                attr_name = func.attr
                
                # Check for collections module data structures
                if "collections" in obj_name.lower() or "collections" in str(func.value):
                    if attr_name in ("deque", "defaultdict", "Counter", "OrderedDict", "ChainMap"):
                        # For deque, we track it and let size checks determine if it's bounded
                        # (deque with maxlen parameter is bounded, but we detect that via size checks)
                        is_data_structure = True
                        structure_type = attr_name
                    elif attr_name in ("list", "dict", "set"):
                        # collections.defaultdict(list) pattern
                        is_data_structure = True
                        structure_type = "defaultdict" if "defaultdict" in obj_name.lower() else attr_name
                # Check for queue.Queue, asyncio.Queue (if unbounded)
                elif "queue" in obj_name.lower() or "asyncio" in obj_name.lower():
                    if attr_name == "Queue":
                        # Check if maxsize is set (bounded queue)
                        has_maxsize = False
                        for keyword in node.value.keywords:
                            if keyword.arg == "maxsize":
                                has_maxsize = True
                                break
                        if not has_maxsize:
                            is_data_structure = True
                            structure_type = "queue"
                # Direct attribute access like deque(), Counter(), etc.
                elif attr_name in ("deque", "defaultdict", "Counter", "OrderedDict", "ChainMap"):
                    is_data_structure = True
                    structure_type = attr_name
        
        if is_data_structure and node.targets and isinstance(node.targets[0], ast.Name):
            scope = context.get("current_scope", "function")
            # Only track if it's at class level (not module level)
            if scope == "class":
                operations.append({
                    "line": node.lineno,
                    "var_name": node.targets[0].id,
                    "structure_type": structure_type,
                    "scope": scope,
                    "call": context["get_call_string"](node.value) if isinstance(node.value, ast.Call) else f"{structure_type}()",
                })
        
        return operations
    
    def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt],
                     context: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Flag persistent data structures that have add operations without size limits"""
        violations = []
        current_function = context["current_function"]
        current_scope = context.get("current_scope", "function")
        file_path = context["file_path"]
        get_attr_string = context["get_attr_string"]
        
        # Skip if this is initialization code (module-level, class-level, or __init__ methods)
        # Only flag operations in regular methods/functions that can be called during runtime
        is_initialization = (
            current_scope in ("module", "class") or
            current_function in ("__init__", "__new__", "__class_init__") or
            current_function is None  # Module-level code
        )
        
        if is_initialization:
            return violations  # Don't flag initialization code
        
        # Track which variables have add operations and size checks
        var_add_operations = {}  # var_name -> list of lines with add operations
        var_size_checks = {}  # var_name -> has size limit check
        
        # Build a set of variable names to check
        tracked_vars = {op["var_name"]: op for op in operations}
        
        # Scan body for operations on these variables
        for stmt in function_body:
            for node in ast.walk(stmt):
                # Check for method calls that add items
                if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
                    attr_name = node.func.attr
                    obj_name = get_attr_string(node.func.value)
                    
                    # Check if this is an add operation on one of our tracked variables
                    for var_name, op in tracked_vars.items():
                        structure_type = op["structure_type"]
                        
                        # Match variable name (exact or as attribute)
                        if obj_name == var_name or obj_name.endswith(f".{var_name}") or obj_name.endswith(f"['{var_name}']"):
                            # Check for add operations
                            add_ops = {
                                "list": ["append", "extend", "insert"],
                                "dict": ["update", "setdefault"],
                                "set": ["add", "update"],
                                "deque": ["append", "appendleft", "extend", "extendleft", "insert"],
                                "defaultdict": ["update", "setdefault"],
                                "Counter": ["update"],
                                "OrderedDict": ["update", "setdefault"],
                                "ChainMap": ["new_child"],
                                "queue": ["put", "put_nowait"],
                            }
                            
                            if attr_name in add_ops.get(structure_type, []):
                                if var_name not in var_add_operations:
                                    var_add_operations[var_name] = []
                                var_add_operations[var_name].append(node.lineno)
                            
                            # Check for size limit checks (len() calls, maxsize/maxlen attributes)
                            if (attr_name in ("__len__",) or 
                                "maxsize" in attr_name.lower() or 
                                "max_size" in attr_name.lower() or
                                attr_name == "maxlen"):  # For deque
                                var_size_checks[var_name] = True
                
                # Check for heapq operations on tracked lists (heapq.heappush, heapq.heappop)
                if isinstance(node, ast.Call):
                    func = node.func
                    # Check for heapq.heappush(list_var, item) or heapq.heappop(list_var)
                    if isinstance(func, ast.Attribute):
                        func_obj = get_attr_string(func.value)
                        func_name = func.attr
                        # Check if it's a heapq operation
                        if func_obj == "heapq" and func_name in ("heappush", "heapreplace", "heappushpop"):
                            # First argument should be our tracked variable
                            if len(node.args) > 0:
                                arg_name = get_attr_string(node.args[0])
                                for var_name, op in tracked_vars.items():
                                    if op["structure_type"] == "list" and (
                                        arg_name == var_name or arg_name.endswith(f".{var_name}")
                                    ):
                                        if var_name not in var_add_operations:
                                            var_add_operations[var_name] = []
                                        var_add_operations[var_name].append(node.lineno)
                
                # Check for dict item assignment: dict[key] = value
                if isinstance(node, ast.Assign):
                    for target in node.targets:
                        if isinstance(target, ast.Subscript):
                            target_name = get_attr_string(target.value)
                            for var_name in tracked_vars:
                                if target_name == var_name or target_name.endswith(f".{var_name}"):
                                    if var_name not in var_add_operations:
                                        var_add_operations[var_name] = []
                                    var_add_operations[var_name].append(node.lineno)
                
                # Check for augmented assignment: list += [...]
                if isinstance(node, ast.AugAssign):
                    target_name = get_attr_string(node.target)
                    for var_name in tracked_vars:
                        if target_name == var_name or target_name.endswith(f".{var_name}"):
                            if var_name not in var_add_operations:
                                var_add_operations[var_name] = []
                            var_add_operations[var_name].append(node.lineno)
                
                # Check for size comparisons in conditionals
                if isinstance(node, (ast.If, ast.While, ast.Assert)):
                    test = getattr(node, "test", None)
                    if test:
                        for comp_node in ast.walk(test):
                            if isinstance(comp_node, ast.Compare):
                                left_str = get_attr_string(comp_node.left) if hasattr(comp_node, "left") else ""
                                # Check for len() calls
                                if isinstance(comp_node.left, ast.Call):
                                    call_func = comp_node.left.func
                                    if isinstance(call_func, ast.Name) and call_func.id == "len":
                                        if len(comp_node.left.args) > 0:
                                            arg_name = get_attr_string(comp_node.left.args[0])
                                            for var_name in tracked_vars:
                                                if arg_name == var_name or arg_name.endswith(f".{var_name}"):
                                                    # Check if comparing to a limit
                                                    for comparator in comp_node.comparators:
                                                        if isinstance(comparator, ast.Constant):
                                                            var_size_checks[var_name] = True
                                                        elif isinstance(comparator, ast.Name):
                                                            # Could be a constant like MAX_SIZE
                                                            if "max" in comparator.id.lower() or "limit" in comparator.id.lower():
                                                                var_size_checks[var_name] = True
                                                        # Handle deprecated ast.Num for Python < 3.8
                                                        try:
                                                            Num = getattr(ast, "Num", None)
                                                            if Num and isinstance(comparator, Num):
                                                                var_size_checks[var_name] = True
                                                        except (AttributeError, TypeError):
                                                            pass
                                # Check for direct variable comparisons
                                for var_name in tracked_vars:
                                    if var_name in left_str:
                                        for comparator in comp_node.comparators:
                                            if isinstance(comparator, ast.Constant):
                                                var_size_checks[var_name] = True
                                            # Handle deprecated ast.Num for Python < 3.8
                                            try:
                                                Num = getattr(ast, "Num", None)
                                                if Num and isinstance(comparator, Num):
                                                    var_size_checks[var_name] = True
                                            except (AttributeError, TypeError):
                                                pass
        
        # Flag violations: persistent structures with add operations but no size checks
        for op in operations:
            var_name = op["var_name"]
            structure_type = op["structure_type"]
            
            if var_name in var_add_operations and var_name not in var_size_checks:
                violations.append({
                    "line": op["line"],
                    "type": self.get_pattern_name(),
                    "var_name": var_name,
                    "function": current_function or "class-level",
                    "file_path": file_path,
                    "message": (
                        f"Class-level {structure_type} '{var_name}' "
                        f"has add operations (lines {var_add_operations[var_name]}) but no size limit checks. "
                        f"This can lead to unbounded memory growth and OOM errors during runtime."
                    ),
                })
        
        return violations


class MemoryViolationDetector(ast.NodeVisitor):
    """AST visitor that detects memory violations using registered patterns"""
    
    DEFAULT_PATTERNS: List[Pattern] = [QueueGetPattern(), UnboundedDataStructurePattern()]

    def __init__(self, file_path: str, patterns: Optional[Sequence[Pattern]] = None):
        self.file_path = file_path
        self.violations: List[Dict[str, Any]] = []
        self.current_function: Optional[str] = None
        self.current_scope: str = "module"  # Track current scope: module, class, function
        self.patterns = self.DEFAULT_PATTERNS if patterns is None else patterns
        self.ast_tree: Optional[ast.Module] = None  # Store full AST for module-level checks
        
        self.pattern_operations: Dict[str, List[Dict[str, Any]]] = {
            pattern.get_pattern_name(): [] for pattern in self.patterns
        }
        
        # Track class-level operations separately (for checking in functions)
        self.class_level_operations: Dict[str, List[Dict[str, Any]]] = {
            pattern.get_pattern_name(): [] for pattern in self.patterns
        }
        
        self._context = {
            "get_call_string": self._get_call_string,
            "get_attr_string": self._get_attr_string,
            "is_var_set_to_none": self._is_var_set_to_none,
            "current_function": None,
            "current_scope": "module",
            "file_path": file_path,
        }

    def visit_ClassDef(self, node):
        """Track class scope"""
        old_scope = self.current_scope
        self.current_scope = "class"
        self._context["current_scope"] = "class"
        
        self.generic_visit(node)
        
        self.current_scope = old_scope
        self._context["current_scope"] = old_scope

    def visit_FunctionDef(self, node):
        """Track function scope and check cleanup after visiting"""
        old_function = self.current_function
        old_scope = self.current_scope
        self.current_function = node.name
        self.current_scope = "function"
        self._context["current_function"] = node.name
        self._context["current_scope"] = "function"
        
        for pattern_name in self.pattern_operations:
            self.pattern_operations[pattern_name] = []
        
        self.generic_visit(node)
        self._check_function_cleanup(node)
        
        self.current_function = old_function
        self.current_scope = old_scope
        self._context["current_function"] = old_function
        self._context["current_scope"] = old_scope

    def visit_AsyncFunctionDef(self, node):
        """Track async function scope and check cleanup after visiting"""
        old_function = self.current_function
        old_scope = self.current_scope
        self.current_function = node.name
        self.current_scope = "function"
        self._context["current_function"] = node.name
        self._context["current_scope"] = "function"
        
        for pattern_name in self.pattern_operations:
            self.pattern_operations[pattern_name] = []
        
        self.generic_visit(node)
        self._check_function_cleanup(node)
        
        self.current_function = old_function
        self.current_scope = old_scope
        self._context["current_function"] = old_function
        self._context["current_scope"] = old_scope

    def visit_Assign(self, node):
        """Detect memory-sensitive operations in assignments"""
        for pattern in self.patterns:
            operations = pattern.visit_assign(node, self._context)
            # Track function-level operations
            self.pattern_operations[pattern.get_pattern_name()].extend(operations)
            # Track class-level operations separately (for checking in functions)
            for op in operations:
                if op.get("scope") == "class":
                    self.class_level_operations[pattern.get_pattern_name()].append(op)
        
        self.generic_visit(node)

    def _check_function_cleanup(self, node):
        """Check cleanup for all detected operations"""
        for pattern in self.patterns:
            operations = self.pattern_operations[pattern.get_pattern_name()]
            if operations:
                violations = pattern.check_cleanup(operations, node.body, self._context)
                self.violations.extend(violations)
            
            # For UnboundedDataStructurePattern, also check if this function modifies class-level structures
            if isinstance(pattern, UnboundedDataStructurePattern):
                class_ops = self.class_level_operations[pattern.get_pattern_name()]
                if class_ops and self.current_function not in ("__init__", "__new__", "__class_init__", None):
                    # Check if this regular function modifies class-level structures
                    violations = pattern.check_cleanup(class_ops, node.body, self._context)
                    self.violations.extend(violations)
    
    def _check_module_level_cleanup(self):
        """Check cleanup for module/class level operations"""
        # Module-level operations are now checked when visiting functions
        # This method is kept for potential future use but doesn't need to do anything
        # since we only want to flag runtime modifications in functions, not initialization code
        pass

    def _is_var_set_to_none(self, var_name: str, body: List[ast.stmt]) -> bool:
        """Check if variable is set to None after its initial assignment"""
        assignment_line = None
        for stmt in body:
            for node in ast.walk(stmt):
                if isinstance(node, ast.Assign):
                    for target in node.targets:
                        if isinstance(target, ast.Name) and target.id == var_name:
                            assignment_line = node.lineno
                            break
                    if assignment_line:
                        break
            if assignment_line:
                break
        
        if not assignment_line:
            return False
        
        for stmt in body:
            for node in ast.walk(stmt):
                if isinstance(node, ast.Assign):
                    for target in node.targets:
                        if isinstance(target, ast.Name) and target.id == var_name and node.lineno > assignment_line:
                            if isinstance(node.value, ast.Constant) and node.value.value is None:
                                return True
                            try:
                                NameConstant = getattr(ast, "NameConstant", None)
                                if NameConstant and isinstance(node.value, NameConstant):
                                    if getattr(node.value, "value", None) is None:
                                        return True
                            except (AttributeError, TypeError):
                                pass
        return False

    def _get_call_string(self, node: ast.Call) -> str:
        """Get string representation of function call"""
        try:
            if hasattr(ast, "unparse"):
                return ast.unparse(node)
            elif isinstance(node.func, ast.Attribute):
                return f"{self._get_attr_string(node.func.value)}.{node.func.attr}()"
            return str(node)
        except Exception:
            return str(node)

    def _get_attr_string(self, node: ast.AST) -> str:
        """Get string representation of attribute access"""
        if isinstance(node, ast.Name):
            return node.id
        elif isinstance(node, ast.Attribute):
            return f"{self._get_attr_string(node.value)}.{node.attr}"
        return str(node)


def check_file_for_memory_violations(file_path: str, patterns: Optional[Sequence[Pattern]] = None) -> List[Dict[str, Any]]:
    """Check a single file for memory violations"""
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read()
        
        if "test" in file_path.lower() or "__pycache__" in file_path:
            return []
        
        tree = ast.parse(content, filename=file_path)
        detector = MemoryViolationDetector(file_path, patterns)
        detector.ast_tree = tree  # Store AST for potential future use
        detector.visit(tree)
        # Class-level operations are checked when visiting functions
        return detector.violations
    except Exception as e:
        print(f"Error parsing {file_path}: {e}")
        return []


def check_directory_for_memory_violations(directory_path: str, ignore_patterns: Optional[List[str]] = None,
                                          patterns: Optional[Sequence[Pattern]] = None) -> List[Dict[str, Any]]:
    """Recursively scan directory for memory violations"""
    if ignore_patterns is None:
        ignore_patterns = ["__pycache__", ".pyc", "site-packages", "venv", ".venv", "env", ".env", "node_modules", "tests"]
    
    all_violations = []
    for root, _dirs, files in os.walk(directory_path):
        if any(pattern in root for pattern in ignore_patterns):
            continue
        for file in files:
            if file.endswith(".py"):
                violations = check_file_for_memory_violations(os.path.join(root, file), patterns)
                all_violations.extend(violations)
    return all_violations


def main():
    """Run memory violation detection on codebase"""
    codebase_path = "./litellm"
    
    print("=" * 80)
    print("MEMORY VIOLATION DETECTION TEST")
    print("=" * 80)
    print(f"Scanning: {codebase_path}")
    print(f"Active patterns: {', '.join(p.get_pattern_name() for p in MemoryViolationDetector.DEFAULT_PATTERNS)}")
    print()
    
    violations = check_directory_for_memory_violations(codebase_path)
    
    if violations:
        by_type = {}
        for v in violations:
            vtype = v["type"]
            if vtype not in by_type:
                by_type[vtype] = []
            by_type[vtype].append(v)
        
        print("MEMORY VIOLATIONS FOUND:")
        print("=" * 80)
        
        total = len(violations)
        for vtype, vlist in by_type.items():
            print(f"\n{vtype.upper().replace('_', ' ')}: {len(vlist)} violation(s)")
            print("-" * 80)
            for v in vlist[:10]:
                print(f"  [VIOLATION] {v['file_path'] if 'file_path' in v else 'unknown'}:{v['line']}")
                print(f"     Function: {v['function']}")
                print(f"     Variable: {v['var_name']}")
                print(f"     {v['message']}")
                print()
            if len(vlist) > 10:
                print(f"     ... and {len(vlist) - 10} more violations of this type")
        
        print("=" * 80)
        print(f"TOTAL VIOLATIONS: {total}")
        print()
        print("RECOMMENDATIONS:")
        print("   1. Set queue variables to None after use: obj = queue.get(); ...; obj = None")
        print("   2. Use bounded queues to prevent unbounded accumulation")
        print("   3. Process items faster than they're added, or drain queues periodically")
        print("   4. For class-level data structures (lists, dicts, sets) that are modified at runtime:")
        print("      - Add size limit checks: if len(data) >= MAX_SIZE: ...")
        print("      - Implement periodic cleanup or use bounded collections")
        print("      - Consider using collections.deque with maxlen for lists")
        print("=" * 80)
        
        first_v = violations[0]
        raise Exception(
            f"Found {total} memory violations! "
            f"First violation: {first_v.get('file_path', 'unknown')}:{first_v['line']} - "
            f"{first_v['message']}"
        )
    else:
        print("OK No memory violations found!")


if __name__ == "__main__":
    main()