Implement pass-through route registration and removal to prevent duplicates and memory leaks

Author: jugaldb

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -57,6 +57,9 @@
 
 pass_through_endpoint_logging = PassThroughEndpointLogging()
 
+# Global registry to track registered pass-through routes and prevent memory leaks
+_registered_pass_through_routes: Dict[str, Dict[str, str]] = {}
+
 
 def get_response_body(response: httpx.Response) -> Optional[dict]:
     try:
@@ -952,8 +955,19 @@ def add_exact_path_route(
         merge_query_params: Optional[bool],
         dependencies: Optional[List],
         cost_per_request: Optional[float],
+        endpoint_id: str,
     ):
         """Add exact path route for pass-through endpoint"""
+        route_key = f"{endpoint_id}:exact:{path}"
+        
+        # Check if this exact route is already registered
+        if route_key in _registered_pass_through_routes:
+            verbose_proxy_logger.debug(
+                "Skipping duplicate exact pass through endpoint: %s (already registered)",
+                path,
+            )
+            return
+            
         verbose_proxy_logger.debug(
             "adding exact pass through endpoint: %s, dependencies: %s",
             path,
@@ -974,6 +988,13 @@ def add_exact_path_route(
             methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
             dependencies=dependencies,
         )
+        
+        # Register the route to prevent duplicates
+        _registered_pass_through_routes[route_key] = {
+            "endpoint_id": endpoint_id,
+            "path": path,
+            "type": "exact"
+        }
 
     @staticmethod
     def add_subpath_route(
@@ -985,9 +1006,20 @@ def add_subpath_route(
         merge_query_params: Optional[bool],
         dependencies: Optional[List],
         cost_per_request: Optional[float],
+        endpoint_id: str,
     ):
         """Add wildcard route for sub-paths"""
         wildcard_path = f"{path}/{{subpath:path}}"
+        route_key = f"{endpoint_id}:subpath:{path}"
+        
+        # Check if this subpath route is already registered
+        if route_key in _registered_pass_through_routes:
+            verbose_proxy_logger.debug(
+                "Skipping duplicate wildcard pass through endpoint: %s (already registered)",
+                wildcard_path,
+            )
+            return
+            
         verbose_proxy_logger.debug(
             "adding wildcard pass through endpoint: %s, dependencies: %s",
             wildcard_path,
@@ -1009,6 +1041,26 @@ def add_subpath_route(
             methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
             dependencies=dependencies,
         )
+        
+        # Register the route to prevent duplicates
+        _registered_pass_through_routes[route_key] = {
+            "endpoint_id": endpoint_id,
+            "path": path,
+            "type": "subpath"
+        }
+
+    @staticmethod
+    def remove_endpoint_routes(endpoint_id: str):
+        """Remove all routes for a specific endpoint ID from the registry"""
+        keys_to_remove = [
+            key for key, value in _registered_pass_through_routes.items()
+            if value["endpoint_id"] == endpoint_id
+        ]
+        for key in keys_to_remove:
+            del _registered_pass_through_routes[key]
+            verbose_proxy_logger.debug(
+                "Removed pass-through route from registry: %s", key
+            )
 
 
 async def initialize_pass_through_endpoints(
@@ -1072,6 +1124,7 @@ async def initialize_pass_through_endpoints(
             merge_query_params=_merge_query_params,
             dependencies=_dependencies,
             cost_per_request=endpoint.get("cost_per_request", None),
+            endpoint_id=endpoint.get("id"),
         )
 
         # Add wildcard route for sub-paths
@@ -1085,6 +1138,7 @@ async def initialize_pass_through_endpoints(
                 merge_query_params=_merge_query_params,
                 dependencies=_dependencies,
                 cost_per_request=endpoint.get("cost_per_request", None),
+                endpoint_id=endpoint.get("id"),
             )
 
         verbose_proxy_logger.debug("Added new pass through endpoint: %s (ID: %s)", _path, endpoint.get("id"))
@@ -1229,6 +1283,9 @@ async def update_pass_through_endpoints(
     # Update the list
     pass_through_endpoint_data[endpoint_index] = endpoint_dict
 
+    # Remove old routes from registry before they get re-registered
+    InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id)
+
     ## Update db
     updated_data = ConfigFieldUpdate(
         field_name="pass_through_endpoints",
@@ -1369,6 +1426,9 @@ async def delete_pass_through_endpoints(
     pass_through_endpoint_data.pop(endpoint_index)
     response_obj = found_endpoint
 
+    # Remove routes from registry  
+    InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id)
+
     ## Update db
     updated_data = ConfigFieldUpdate(
         field_name="pass_through_endpoints",