diff --git a/cognee-mcp/src/tools/clear_registry.py b/cognee-mcp/src/tools/clear_registry.py
index 2bdb967fb..6199773ec 100644
--- a/cognee-mcp/src/tools/clear_registry.py
+++ b/cognee-mcp/src/tools/clear_registry.py
@@ -1,6 +1,6 @@
 from contextlib import redirect_stdout
 import sys
-from src.utils.context import cognee_client
+from src.utils import context
 from cognee.shared.logging_utils import get_logger
 import mcp.types as types
 
@@ -21,8 +21,8 @@ async def clear_registry() -> list:
 
     with redirect_stdout(sys.stderr):
         try:
-            await cognee_client.prune_data()
-            await cognee_client.prune_system(metadata=True)
+            await context.cognee_client.prune_data()
+            await context.cognee_client.prune_system(metadata=True)
             logger.info("MCP server registry cleared")
             return [
                 types.TextContent(
diff --git a/cognee-mcp/src/tools/find_mcp_server.py b/cognee-mcp/src/tools/find_mcp_server.py
index 5ce49bc10..60bc36ed0 100644
--- a/cognee-mcp/src/tools/find_mcp_server.py
+++ b/cognee-mcp/src/tools/find_mcp_server.py
@@ -4,7 +4,7 @@ import sys
 from cognee.shared.logging_utils import get_logger
 import mcp.types as types
 
-from src.utils.context import cognee_client
+from src.utils import context
 
 logger = get_logger()
 
@@ -51,13 +51,13 @@ async def find_mcp_server(requirements: str, max_results: int = 5) -> list:
             logger.info(f"Searching for MCP servers matching: {requirements}")
 
             # Search using GRAPH_COMPLETION for intelligent matching
-            search_results = await cognee_client.search(
+            search_results = await context.cognee_client.search(
                 query_text=f"Find MCP servers that can: {requirements}. Include their capabilities, installation instructions, and documentation.",
                 query_type="GRAPH_COMPLETION",
             )
 
             # Format the results
-            if cognee_client.use_api:
+            if context.cognee_client.use_api:
                 if isinstance(search_results, str):
                     result_text = search_results
                 elif isinstance(search_results, list) and len(search_results) > 0:
diff --git a/cognee-mcp/src/tools/list_mcp_servers.py b/cognee-mcp/src/tools/list_mcp_servers.py
index 827d8adc6..8eb07146c 100644
--- a/cognee-mcp/src/tools/list_mcp_servers.py
+++ b/cognee-mcp/src/tools/list_mcp_servers.py
@@ -4,7 +4,7 @@ import sys
 from cognee.shared.logging_utils import get_logger
 import mcp.types as types
 
-from src.utils.context import cognee_client
+from src.utils import context
 
 logger = get_logger()
 
@@ -28,13 +28,13 @@ async def list_mcp_servers() -> list:
             logger.info("Listing all MCP servers")
 
             # Search for all MCP servers with connection details
-            search_results = await cognee_client.search(
+            search_results = await context.cognee_client.search(
                 query_text="List all MCP servers with their names, descriptions, capabilities, connection information (URL, command, args), installation instructions, and documentation links",
                 query_type="GRAPH_COMPLETION",
             )
 
             # Format the results
-            if cognee_client.use_api:
+            if context.cognee_client.use_api:
                 if isinstance(search_results, str):
                     result_text = search_results
                 elif isinstance(search_results, list) and len(search_results) > 0:
diff --git a/cognee-mcp/src/tools/remember_mcp_server.py b/cognee-mcp/src/tools/remember_mcp_server.py
index eeb951dbd..42b0971b7 100644
--- a/cognee-mcp/src/tools/remember_mcp_server.py
+++ b/cognee-mcp/src/tools/remember_mcp_server.py
@@ -4,7 +4,7 @@ import sys
 from cognee.shared.logging_utils import get_logger, get_log_file_location
 import mcp.types as types
 
-from src.utils.context import cognee_client
+from src.utils import context
 
 logger = get_logger()
 
@@ -115,10 +115,12 @@ Args: {args or "Not provided"}
 
             try:
                 # Add to knowledge graph with special node set
-                await cognee_client.add(server_content, node_set=["mcp_servers", server_name])
+                await context.cognee_client.add(
+                    server_content, node_set=["mcp_servers", server_name]
+                )
 
                 # Process into knowledge graph
-                await cognee_client.cognify()
+                await context.cognee_client.cognify()
 
                 logger.info(f"Successfully stored MCP server: {server_name}")
             except Exception as e: