diff --git a/agentic_security/mcp/client.py b/agentic_security/mcp/client.py
index 38aa497..e8a29cc 100644
--- a/agentic_security/mcp/client.py
+++ b/agentic_security/mcp/client.py
@@ -3,6 +3,8 @@
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
 
+from agentic_security.logutils import logger
+
 # Create server parameters for stdio connection
 server_params = StdioServerParameters(
     command="python",  # Executable
@@ -12,42 +14,54 @@
 
 
 async def run() -> None:
-    async with stdio_client(server_params) as (read, write):
-        async with ClientSession(read, write) as session:
-            # Initialize the connection --> connection does not work
-            await session.initialize()
-
-            # List available prompts, resources, and tools --> no avalialbe tools
-            prompts = await session.list_prompts()
-            print(f"Available prompts: {prompts}")
-
-            resources = await session.list_resources()
-            print(f"Available resources: {resources}")
-
-            tools = await session.list_tools()
-            print(f"Available tools: {tools}")
-
-            # Call the echo tool --> echo tool iisue
-            echo_result = await session.call_tool(
-                "echo_tool", arguments={"message": "Hello from client!"}
-            )
-            print(f"Tool result: {echo_result}")
-
-            # # Read the echo resource
-            # echo_content, mime_type = await session.read_resource(
-            #     "echo://Hello_resource"
-            # )
-            # print(f"Resource content: {echo_content}")
-            # print(f"Resource MIME type: {mime_type}")
-
-            # # Get and use the echo prompt
-            # prompt_result = await session.get_prompt(
-            #     "echo_prompt", arguments={"message": "Hello prompt!"}
-            # )
-            # print(f"Prompt result: {prompt_result}")
-
-            # You can perform additional operations here as needed
-            return prompts, resources, tools
+    try:
+        logger.info(
+            "Starting stdio client session with server parameters: %s", server_params
+        )
+        async with stdio_client(server_params) as (read, write):
+            async with ClientSession(read, write) as session:
+                # Initialize the connection --> connection does not work
+                logger.info("Initializing client session...")
+                await session.initialize()
+
+                # List available prompts, resources, and tools --> no avalialbe tools
+                logger.info("Listing available prompts...")
+                prompts = await session.list_prompts()
+                logger.info(f"Available prompts: {prompts}")
+
+                logger.info("Listing available resources...")
+                resources = await session.list_resources()
+                logger.info(f"Available resources: {resources}")
+
+                logger.info("Listing available tools...")
+                tools = await session.list_tools()
+                logger.info(f"Available tools: {tools}")
+
+                # Call the echo tool --> echo tool issue
+                logger.info("Calling echo_tool with message...")
+                echo_result = await session.call_tool(
+                    "echo_tool", arguments={"message": "Hello from client!"}
+                )
+                logger.info(f"Tool result: {echo_result}")
+
+                # # Read the echo resource
+                # echo_content, mime_type = await session.read_resource(
+                #     "echo://Hello_resource"
+                # )
+                # logger.info(f"Resource content: {echo_content}")
+                # logger.info(f"Resource MIME type: {mime_type}")
+
+                # # Get and use the echo prompt
+                # prompt_result = await session.get_prompt(
+                #     "echo_prompt", arguments={"message": "Hello prompt!"}
+                # )
+                # logger.info(f"Prompt result: {prompt_result}")
+
+                logger.info("Client operations completed successfully.")
+                return prompts, resources, tools
+    except Exception as e:
+        logger.error(f"An error occurred during client operations: {e}", exc_info=True)
+        raise
 
 
 if __name__ == "__main__":
diff --git a/agentic_security/probe_data/data.py b/agentic_security/probe_data/data.py
index 66ba25e..5d7351e 100644
--- a/agentic_security/probe_data/data.py
+++ b/agentic_security/probe_data/data.py
@@ -8,7 +8,6 @@
 import httpx
 import pandas as pd
 from cache_to_disk import cache_to_disk
-from datasets import load_dataset
 
 from agentic_security.logutils import logger
 from agentic_security.probe_data import stenography_fn
@@ -20,6 +19,7 @@
     inspect_ai_tool,
     rl_model,
 )
+from datasets import load_dataset
 
 # Type aliases for clarity
 T = TypeVar("T")
diff --git a/tests/test_registry.py b/tests/test_registry.py
index 2cf0cbc..6338e2b 100644
--- a/tests/test_registry.py
+++ b/tests/test_registry.py
@@ -1,7 +1,7 @@
 import pytest
-from datasets import load_dataset
 
 from agentic_security.probe_data import REGISTRY
+from datasets import load_dataset
 
 
 @pytest.mark.slow