ran black on python files

Author: geobio

services/agent-environment/dev_scripts/debug_and_concurrency_tests/get_running_servers.py

@@ -3,7 +3,7 @@
 Useful as a quick check to see if some mcp servers errored during startup.
 Does not check that they work tho (e.g. api key could be invalid and /call_tool will fail, but mcp server is running).
 
-Usage (requires server to be running on port 1984, you can do "make run" or "make run-docker"): 
+Usage (requires server to be running on port 1984, you can do "make run" or "make run-docker"):
 uv run python get_running_servers.py
 """
 
@@ -12,92 +12,100 @@
 import os
 import sys
 
+
 def load_configured_servers():
     """Load the configured servers from mcp_client.py"""
     try:
         # Add the src directory to Python path so we can import the module
-        src_path = os.path.join(os.path.dirname(__file__), '../../src')
+        src_path = os.path.join(os.path.dirname(__file__), "../../src")
         if src_path not in sys.path:
             sys.path.insert(0, src_path)
-        
+
         # Import the mcp_client module directly
         from agent_environment.mcp_client import config
-        
+
         # Extract server names from the config
         configured_servers = set(config["mcpServers"].keys())
         return configured_servers
-        
+
     except ImportError as e:
-        raise Exception(f"Could not import mcp_client module: {e}. Are you using \"uv run ...\"?")
+        raise Exception(
+            f'Could not import mcp_client module: {e}. Are you using "uv run ..."?'
+        )
     except KeyError as e:
         raise Exception(f"Config structure error: {e}")
     except Exception as e:
         raise Exception(f"Error loading configured servers: {e}")
 
+
 def main():
-    load_configured_servers() # to trigger the import warning at the top of this script
+    load_configured_servers()  # to trigger the import warning at the top of this script
 
     try:
         # Step 1: Run the curl script and extract tool names
         print("Calling ./curl_scripts/mcp__list_tools.sh to get running mcp servers...")
-        
+
         # Run the command: ./curl_scripts/mcp__list_tools.sh | jq | grep "^    \"name"
         curl_script_path = "./curl_scripts/mcp__list_tools.sh"
-        
+
         if not os.path.exists(curl_script_path):
             print(f"Error: {curl_script_path} not found")
             sys.exit(1)
-        
+
         # Run the command pipeline
-        result = subprocess.run([
-            "bash", "-c", 
-            f"{curl_script_path} | jq | grep '^    \"name'"
-        ], capture_output=True, text=True, check=True)
-        
+        result = subprocess.run(
+            ["bash", "-c", f"{curl_script_path} | jq | grep '^    \"name'"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+
         tool_names_content = result.stdout
-        
+
         # Write to temporary file
         with open("tool-names.txt", "w") as f:
             f.write(tool_names_content)
-        
+
         # Step 2: Process with regex to extract server names
-        pattern = r'.*\"([a-zA-Z0-9\-]+)_.*'
+        pattern = r".*\"([a-zA-Z0-9\-]+)_.*"
         server_names = set()
-        
+
         for line in tool_names_content.splitlines():
             line = line.strip()
             if line:
                 match = re.search(pattern, line)
                 if match:
                     server_name = match.group(1)
                     server_names.add(server_name)
-        
+
         # Step 3: Remove duplicates (already done with set) and save
         print(f"\nFound {len(server_names)} unique running servers\n")
-        
+
         # Sort for consistent output
         sorted_servers = sorted(server_names)
-        
+
         print(f"Running servers: {sorted_servers}\n")
-        
+
         # Step 3.5: Compare with configured servers
-        
+
         configured_servers = load_configured_servers()
         not_running = configured_servers - server_names
-        
+
         if not_running:
             print(f"⚠️  CONFIGURED BUT NOT RUNNING ({len(not_running)} servers):")
             for server in sorted(not_running):
                 print(f"  ❌ {server}")
         else:
-            print(f"✅ All {len(configured_servers)} mcp servers in mcp_client.py are running!")
-        
+            print(
+                f"✅ All {len(configured_servers)} mcp servers in mcp_client.py are running!"
+            )
+
         # Step 4: Clean up temporary file
         if os.path.exists("tool-names.txt"):
             os.remove("tool-names.txt")
-        
+
         print("\nDone!")
-        
+
     except subprocess.CalledProcessError as e:
         print(f"Error running command: {e}")
         print(f"stdout: {e.stdout}")
@@ -107,5 +115,6 @@ def main():
         print(f"Error: {e}")
         sys.exit(1)
 
+
 if __name__ == "__main__":
-    main() 
+    main()

services/agent-environment/src/agent_environment/main.py

@@ -19,9 +19,11 @@
 
 # Create cache with appropriate settings for the use case
 tool_cache = Cache(
-    maxsize=10000,          # Max 10000 unique requests (fits about 2000 tasks)
-    ttl=CACHE_TTL_HOURS * 60 * 60,       # 48 hours TTL by default (but each item will have some slight variation)
-    enable_stats=True       # Track cache performance
+    maxsize=10000,  # Max 10000 unique requests (fits about 2000 tasks)
+    ttl=CACHE_TTL_HOURS
+    * 60
+    * 60,  # 48 hours TTL by default (but each item will have some slight variation)
+    enable_stats=True,  # Track cache performance
 )
 
 # Tool name mappings - maps old invalid names to correct names
@@ -126,40 +128,45 @@ def should_cache_tool(tool_name: str) -> bool:
     server_name = tool_name.split("_", 1)[0]
     return server_name in CACHEABLE_SERVERS
 
+
 def generate_cache_key(tool_name: str, tool_args: dict) -> str:
     """Generate consistent cache key from tool call parameters."""
-    cache_data = {
-        "tool_name": tool_name,
-        "tool_args": tool_args
-    }
+    cache_data = {"tool_name": tool_name, "tool_args": tool_args}
     cache_str = json.dumps(cache_data, sort_keys=True)
     return hashlib.md5(cache_str.encode()).hexdigest()
 
+
 @app.post("/call-tool")
 async def call_tool(
     request: CallToolRequest,
 ) -> list[mcp.types.ContentBlock]:
     """Call a specific tool with the provided arguments."""
-    
+
     mapped_tool_name = TOOL_NAME_MAPPINGS.get(request.tool_name, request.tool_name)
-    
+
     # Generate cache key
     cache_key = generate_cache_key(mapped_tool_name, request.tool_args)
-    
+
     # Check cache first
     cached_result = tool_cache.get(cache_key)
-    if cached_result is not None and request.use_cache and should_cache_tool(mapped_tool_name):
+    if (
+        cached_result is not None
+        and request.use_cache
+        and should_cache_tool(mapped_tool_name)
+    ):
         logger.info(f"Returning cached result for tool '{request.tool_name}'")
         return cached_result
-    
+
     async with client:
         try:
             result = await client.call_tool(mapped_tool_name, request.tool_args)
 
             # Check for errors first (FastMCP best practice)
             if result.is_error:
                 error_msg = "Unknown error"
-                if result.content and isinstance(result.content[0], mcp.types.TextContent):
+                if result.content and isinstance(
+                    result.content[0], mcp.types.TextContent
+                ):
                     error_msg = result.content[0].text
                 raise HTTPException(
                     status_code=500,
@@ -172,24 +179,26 @@ async def call_tool(
                 # TTL is 70-100% of default TTL, to avoid all items expiring at the same time
                 random_ttl = int(CACHE_TTL_HOURS * 60 * 60 * random.uniform(0.7, 1.0))
                 tool_cache.set(cache_key, content_blocks, ttl=random_ttl)
-            
+
             return content_blocks
-            
+
         except Exception as e:
             raise HTTPException(
                 status_code=500,
                 detail=f"Failed to call tool '{request.tool_name}': {str(e)}",
             )
 
+
 @app.get("/cache-stats")
 async def get_cache_stats():
     """Get cache statistics for monitoring."""
     return {
         "cache_size": len(tool_cache),
         "max_size": tool_cache.maxsize,
-        "ttl_seconds": tool_cache.ttl
+        "ttl_seconds": tool_cache.ttl,
     }
 
+
 @app.post("/cache-clear")
 async def clear_cache():
     """Clear the entire cache."""
@@ -201,7 +210,7 @@ async def clear_cache():
 async def get_enabled_servers() -> dict[str, Any]:
     """Get list of configured MCP servers with their status (OK or ERROR_NOT_ONLINE)."""
     configured = set(config.get("mcpServers", {}).keys())
-    
+
     async with client:
         try:
             tools = await client.list_tools()
@@ -211,18 +220,18 @@ async def get_enabled_servers() -> dict[str, Any]:
                 if "_" in tool.name:
                     server_name = tool.name.split("_", 1)[0]
                     live_servers.add(server_name)
-            
+
             # Build status list for each configured server
             servers = [
                 (name, "OK" if name in live_servers else "ERROR_NOT_ONLINE")
                 for name in sorted(configured)
             ]
-            
+
             return {
                 "servers": servers,
                 "total": len(configured),
                 "online": len(live_servers),
-                "offline": len(configured - live_servers)
+                "offline": len(configured - live_servers),
             }
         except Exception as e:
             raise HTTPException(
@@ -234,6 +243,7 @@ async def get_enabled_servers() -> dict[str, Any]:
 async def health() -> dict[str, Any]:
     """Simple health check that verifies client is also ok. Timeout is 5 seconds."""
     try:
+
         async def _health_check_with_client():
             async with client:
                 return {