mcp: refactor and extend agent loop (google#14135)

- Use a stdio server instead of http, simplifying how we run the server
- Enable multiple runs for fixing builds if first chat run didn't work
- Remove misleading comments

Signed-off-by: David Korczynski <[email protected]>

Author: DavidKorczynski

infra/experimental/mcp/client.py

@@ -23,22 +23,20 @@
 import subprocess
 import argparse
 import time
-
+import sys
 import httpx
-from pydantic_ai import Agent
-from pydantic_ai.mcp import MCPServerSSE
+import pathlib
 
-import logfire
+from pydantic_ai import Agent
+from pydantic_ai.mcp import MCPServerStdio
 
 import config as oss_fuzz_mcp_config
 
-logfire.configure(send_to_logfire='if-token-present')
-logfire.instrument_pydantic_ai()
-
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    format="[CLIENT] %(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    stream=sys.stderr)
 logger = logging.getLogger("mcp-server")
 
 MCP_SERVER_URL = "http://localhost:8000/sse"
@@ -178,9 +176,9 @@
 """
 
 
-async def chat_with_agent(prompt: str) -> list:
+async def run_agent_loop(prompt: str) -> list:
   """
-    Send a message to the LLM with access to the MCP tools.
+    Performs a run with the LLM.
     
     Args:
         prompt: The user's message
@@ -190,23 +188,24 @@ async def chat_with_agent(prompt: str) -> list:
     """
   nodes = []
   try:
-    server = MCPServerSSE(url=MCP_SERVER_URL,
-                          timeout=5200.0,
-                          read_timeout=5000.0)
+    server = MCPServerStdio(
+        'python3',
+        [str(pathlib.Path(__file__).parent.resolve()) + '/oss_fuzz_server.py'],
+        timeout=5200.0)
 
-    agent = Agent(model="openai:gpt-4o", toolsets=[server], retries=30)
+    agent = Agent(model="openai:gpt-4", toolsets=[server], retries=30)
 
     # Run the agent with the MCP server context
     logger.info('Starting agent run')
     async with agent.iter(prompt) as agent_run:
       logger.info('Agent run started')
-
       async for node in agent_run:
-        logger.info('Running node %s', node.__class__.__name__)
+        logger.info('Running node [%d] %s', len(nodes), node.__class__.__name__)
         time.sleep(3)
         nodes.append(node)
   except Exception as e:
     logger.info('Error during agent run: %s', e)
+    sys.exit(1)
 
   return nodes
 
@@ -365,7 +364,7 @@ async def does_project_build(project: str) -> bool:
   return True
 
 
-async def fix_project_build(project: str):
+async def fix_project_build(project: str, max_tries: int = 3):
   """Runs an agent to fix the build of an OSS-Fuzz project."""
 
   project_language = _detect_language(project)
@@ -382,8 +381,11 @@ async def fix_project_build(project: str):
   if oss_fuzz_filetree:
     extra_project_text += f'The files in the OSS-Fuzz project for {project} are:\n{oss_fuzz_filetree}\n'
 
-  nodes = await chat_with_agent(
-      f"""Fix the OSS-Fuzz project {project} that currently has a broken build.
+  nodes = []
+  for _attempt in range(max_tries):
+    logger.info('Attempt %d to fix project %s', _attempt + 1, project)
+    nodes += await run_agent_loop(
+        f"""Fix the OSS-Fuzz project {project} that currently has a broken build.
 Use the build logs from OSS-Fuzz's project {project} and determine why it fails, then 
 proceed to adjust Dockerfile and build.sh scripts until the project builds.
 
@@ -410,7 +412,10 @@ async def fix_project_build(project: str):
 - Continue adjusting the files in {oss_fuzz_mcp_config.BASE_OSS_FUZZ_DIR}/projects/{project}/ until "fuzzer-check" passes.
 """)
 
-  fix_success = await does_project_build(project)
+    fix_success = await does_project_build(project)
+    if fix_success:
+      logger.info('Project %s build fixed successfully.', project)
+      break
   return nodes, fix_success
 
 
@@ -501,7 +506,7 @@ async def add_run_tests_command(project_name: str):
 
   os.chmod(run_tests_path, 0o755)
 
-  await chat_with_agent(f"""
+  await run_agent_loop(f"""
 You are an expert software security engineer that is specialized in OSS-Fuzz.
 You are tasked with adding a run_tests.sh script to an OSS-Fuzz project.
 This script should run the tests of the project, and ensure that the project is working correctly.
@@ -530,7 +535,7 @@ async def expand_existing_project(project_name: str):
     logger.info('Failed to prepare %s. Exiting.', project_name)
     return
 
-  nodes = await chat_with_agent(
+  nodes = await run_agent_loop(
       f"""You are a security engineer that is an expert in fuzzing development, and your goal is to expand on the
 fuzzing harnesses of OSS-Fuzz project {project_name}.
 Use the tools to understand the fuzzing harnesses of the {project_name}'s OSS-Fuzz integration.
@@ -571,7 +576,8 @@ def _log_nodes(logfile, nodes, header_text=''):
 
 async def fix_oss_fuzz_projects(projects_to_fix=None,
                                 max_projects_to_fix=4,
-                                language=''):
+                                language='',
+                                max_tries=3):
   """Fixes the build of a list of OSS-Fuzz projects."""
 
   if projects_to_fix is None:
@@ -598,7 +604,7 @@ async def fix_oss_fuzz_projects(projects_to_fix=None,
         continue
     except:
       continue
-    nodes, fix_success = await fix_project_build(project)
+    nodes, fix_success = await fix_project_build(project, max_tries)
     responses.append({'project': project, 'fix_success': fix_success})
     if nodes:
       _log_nodes(f'responses-fix-build-{project}.json',
@@ -707,7 +713,7 @@ async def initiate_project_creation(project: str, project_repo: str,
 available to extract code coverage of the project when you're creating the harness, and either
 add more fuzzing harnesses to the project or extend the harness to cover more functions."""
 
-  nodes = await chat_with_agent(
+  nodes = await run_agent_loop(
       f"""You are an expert software security engineer and you are tasked with creating an OSS-Fuzz project.
 I have set up an initial project structure at {oss_fuzz_mcp_config.BASE_OSS_FUZZ_DIR}/projects/{project}/. This structure
 includes a Dockerfile, build.sh, and project.yaml file. The Dockerfile clones the target
@@ -882,6 +888,11 @@ def parse_arguments():
   fix_builds = subparsers.add_parser(
       'fix-builds',
       help='Fix the builds of OSS-Fuzz projects that are currently broken.')
+  fix_builds.add_argument(
+      '--max_attempts',
+      type=int,
+      default=3,
+      help='Maximum number of attempts to fix each project (default: 3)')
 
   fix_builds.add_argument('--max-projects',
                           type=int,
@@ -939,7 +950,8 @@ async def main():
 
   initialize_oss_fuzz()
   if args.command == 'fix-builds':
-    await fix_oss_fuzz_projects(args.projects, args.max_projects, args.language)
+    await fix_oss_fuzz_projects(args.projects, args.max_projects, args.language,
+                                args.max_attempts)
   elif args.command == 'create-project':
     logger.info('Creating OSS-Fuzz project for URL: %s', args.project_url)
     await create_oss_fuzz_integration_for_project(args.project_url,

infra/experimental/mcp/oss_fuzz_server.py

@@ -18,28 +18,24 @@
 import logging
 import asyncio
 import os
+import sys
 import json
 import time
 import subprocess
-
-import httpx
 from mcp.server.fastmcp import FastMCP
 
 import config as oss_fuzz_mcp_config
 
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    format=
+    "[SERVER] %(asctime)s - %(name)s - '%(module)s - %(funcName)s - %(levelname)s - %(message)s",
+    stream=sys.stderr)
 logger = logging.getLogger("mcp-server")
 
-# Create a shared HTTP client for API requests
-http_client = httpx.AsyncClient(timeout=10.0)
-
 # Create an MCP server with a name
-mcp = FastMCP("OSS-Fuzz tools with relevant file system utilities.",
-              host="0.0.0.0",
-              port=8000)
+mcp = FastMCP("OSS-Fuzz tools with relevant file system utilities.")
 
 FILE_ACCESS_ERROR = f"""Error: Cannot access directories outside of the base directory.
 Remember, all paths accessible by you must be prefixed with {oss_fuzz_mcp_config.BASE_DIR}.
@@ -96,14 +92,16 @@ async def check_if_oss_fuzz_project_builds(project_name: str) -> bool:
                           project_name,
                           cwd=oss_fuzz_mcp_config.BASE_OSS_FUZZ_DIR,
                           shell=True,
+                          stdout=subprocess.DEVNULL,
+                          stderr=subprocess.STDOUT,
                           timeout=60 * 20)
     return True
   except subprocess.CalledProcessError as e:
     logger.info("Build failed for project '%s': {%s}", project_name, str(e))
     return False
   except subprocess.TimeoutExpired:
     logger.info(f"Building project {project_name} timed out.")
-    return False
+  return False
 
 
 def shorten_logs_if_needed(log_string: str) -> str:
@@ -227,7 +225,6 @@ async def check_run_tests(
     Returns:
         The logs from building the project with custom artifacts.
     """
-  logger.info('Running test 1')
   clone_oss_fuzz_if_it_does_not_exist()
   logger.info(
       "Checking if OSS-Fuzz project '%s' builds with custom artifacts...",
@@ -236,12 +233,10 @@ async def check_run_tests(
   os.makedirs(oss_fuzz_mcp_config.BASE_TMP_LOGS, exist_ok=True)
   target_logs = os.path.join(oss_fuzz_mcp_config.BASE_TMP_LOGS,
                              'check-fuzz-run-tests-log.txt')
-  logger.info('Running test 2')
   if os.path.isfile(target_logs):
     os.remove(target_logs)
   log_stdout = open(target_logs, 'w', encoding='utf-8')
   try:
-    logger.info('Running test 3')
     subprocess.check_call(
         f'infra/experimental/chronos/check_tests.sh {project_name} c++',
         cwd=oss_fuzz_mcp_config.BASE_OSS_FUZZ_DIR,
@@ -251,7 +246,6 @@ async def check_run_tests(
         timeout=60 * 20)
 
   except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
-    logger.info('Running test 4')
     logger.info("Build failed for project '%s': {%s}", project_name, str(e))
     log_stdout.write("\n\nrun-tests.sh failed!!\n")
     with open(target_logs, 'r', encoding='utf-8') as f:
@@ -260,14 +254,12 @@ async def check_run_tests(
     logger.info("run-tests.sh logs for project '%s': {%s}", project_name,
                 logs_to_return)
     return logs_to_return
-  logger.info('Running test 5.01')
+
   with open(target_logs, 'r', encoding='utf-8') as f:
     logs = f.read()
   logs_to_return = shorten_logs_if_needed(logs)
-  logger.info('Running test 5')
   logger.info("run-tests.sh for project '%s': {%s}", project_name,
               logs_to_return)
-  logger.info('Running test 6')
   return logs_to_return
 
 
@@ -353,7 +345,6 @@ async def list_files(path: str = "") -> str:
   _internal_delay()
   target_dir = os.path.normpath(path)
   if not target_dir.startswith(oss_fuzz_mcp_config.BASE_DIR):
-    # Security check to prevent directory traversal
     return FILE_ACCESS_ERROR
 
   logger.info("Listing files in directory: %s", target_dir)
@@ -388,7 +379,6 @@ async def get_file_size(file_path) -> str:
   _internal_delay()
   target_file = os.path.normpath(file_path)
   if not target_file.startswith(oss_fuzz_mcp_config.BASE_DIR):
-    # Security check to prevent directory traversal
     return FILE_ACCESS_ERROR
 
   logger.info("Getting file size: %s", target_file)
@@ -420,7 +410,6 @@ async def read_file(file_path: str, start_idx: int, end_idx: int) -> str:
   _internal_delay()
   target_file = os.path.normpath(file_path)
   if not target_file.startswith(oss_fuzz_mcp_config.BASE_DIR):
-    # Security check to prevent directory traversal
     return FILE_ACCESS_ERROR
 
   logger.info("Reading file: %s", target_file)
@@ -462,7 +451,6 @@ async def write_file(file_path: str, content: str) -> str:
   logger.info("Writing to file: %s", file_path)
   target_file = os.path.normpath(file_path)
 
-  # Security check to prevent directory traversal
   if not target_file.startswith(oss_fuzz_mcp_config.BASE_DIR):
     return FILE_ACCESS_ERROR
 
@@ -489,7 +477,6 @@ async def delete_file(file_path: str) -> str:
   logger.info("Deleting file: %s", file_path)
   target_file = os.path.normpath(file_path)
 
-  # Security check to prevent directory traversal
   if not target_file.startswith(oss_fuzz_mcp_config.BASE_DIR):
     return FILE_ACCESS_ERROR
 
@@ -664,19 +651,17 @@ async def get_coverage_of_oss_fuzz_project(project_name):
   logger.info('Refined coverage dict: %s', json.dumps(refined_cov_dict,
                                                       indent=2))
 
-  # Split up the coverage
   return json.dumps(refined_cov_dict, indent=2)
 
 
 def start_mcp_server():
   """Starts the MCP server."""
   try:
-    logger.info("Starting MCP server on port 8000...")
-    # Close the HTTP client when the server shuts down
-    mcp.run(transport="sse")
+    logger.info("Starting MCP server.")
+    mcp.run(transport="stdio")
   except KeyboardInterrupt:
-    logger.info("Server shutting down...")
-    asyncio.run(http_client.aclose())
+    logger.info("Caught KeyboardInterrupt.")
+  logger.info('Server shut down.')
 
 
 if __name__ == "__main__":