From c09ce32defd113777bbc3b558cb336dc630e3c9c Mon Sep 17 00:00:00 2001 From: DavdaJames Date: Sun, 10 Aug 2025 13:42:32 +0530 Subject: [PATCH 1/3] feature added for logging of mcp client --- agentic_security/mcp/client.py | 81 +++++++++++++++++++--------------- 1 file changed, 45 insertions(+), 36 deletions(-) diff --git a/agentic_security/mcp/client.py b/agentic_security/mcp/client.py index 38aa497..47f4770 100644 --- a/agentic_security/mcp/client.py +++ b/agentic_security/mcp/client.py @@ -1,4 +1,6 @@ + import asyncio +from agentic_security.logutils import logger from mcp import ClientSession, StdioServerParameters from mcp.client.stdio import stdio_client @@ -12,42 +14,49 @@ 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: + logger.info("Initializing client session...") + await session.initialize() + + 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}") + + 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__": From 40ff7f9dfb86b5648644261282ecdd0e6396f908 Mon Sep 17 00:00:00 2001 From: DavdaJames Date: Sun, 10 Aug 2025 13:49:08 +0530 Subject: [PATCH 2/3] added the comments back --- agentic_security/mcp/client.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/agentic_security/mcp/client.py b/agentic_security/mcp/client.py index 47f4770..6293b6a 100644 --- a/agentic_security/mcp/client.py +++ b/agentic_security/mcp/client.py @@ -18,9 +18,11 @@ async def run() -> None: 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}") @@ -33,6 +35,7 @@ async def run() -> None: 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!"} From a02aed2c2b1b875665936bedeb3d9db6e23d085a Mon Sep 17 00:00:00 2001 From: DavdaJames Date: Sun, 10 Aug 2025 14:33:25 +0530 Subject: [PATCH 3/3] changes done by pre-commit hooks --- agentic_security/mcp/client.py | 10 ++++++---- agentic_security/probe_data/data.py | 2 +- tests/test_registry.py | 2 +- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/agentic_security/mcp/client.py b/agentic_security/mcp/client.py index 6293b6a..e8a29cc 100644 --- a/agentic_security/mcp/client.py +++ b/agentic_security/mcp/client.py @@ -1,10 +1,10 @@ - import asyncio -from agentic_security.logutils import logger 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 @@ -15,13 +15,15 @@ async def run() -> None: try: - logger.info("Starting stdio client session with server parameters: %s", server_params) + 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() 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