major bug fixes and new integration tests

Author: Lekssays

CONTRIBUTING.md

@@ -4,7 +4,7 @@ services:
   codebadger-toolkit:
     build:
       context: .
-      dockerfile: Dockerfile.joern
+      dockerfile: Dockerfile
     image: codebadger-toolkit:latest
     container_name: codebadger-toolkit-server
     ports:

DOCS.md

@@ -0,0 +1,209 @@
+#!/usr/bin/env python3
+"""
+Sample Client for CodeBadger Toolkit Server
+
+This client demonstrates basic usage of the CodeBadger Toolkit MCP server:
+1. Generate a CPG for a local codebase
+2. List methods in the codebase
+3. Run a simple CPGQL query
+"""
+
+import asyncio
+import logging
+import sys
+import os
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+try:
+    from fastmcp import Client
+except ImportError:
+    logger.error("FastMCP not found. Install with: pip install fastmcp")
+    sys.exit(1)
+
+
+def extract_tool_result(result):
+    """Extract dictionary data from CallToolResult"""
+    if hasattr(result, 'content') and result.content:
+        content_text = result.content[0].text
+        try:
+            import json
+            parsed = json.loads(content_text)
+
+            # Handle complex results that return Scala output with embedded JSON
+            if isinstance(parsed, dict) and 'value' in parsed:
+                value = parsed['value']
+                if isinstance(value, str):
+                    # Look for embedded JSON in the Scala output
+                    import re
+                    # Match the escaped JSON string between quotes
+                    json_match = re.search(r'val res\d+: String = ("\{.*\}")', value)
+                    if json_match:
+                        try:
+                            # Extract the escaped JSON string and unescape it
+                            escaped_json = json_match.group(1)
+                            # Remove the surrounding quotes and unescape
+                            json_str = escaped_json[1:-1]  # Remove quotes
+                            json_str = json_str.replace('\\"', '"')  # Unescape quotes
+                            json_str = json_str.replace('\\\\', '\\')  # Unescape backslashes
+                            return json.loads(json_str)
+                        except json.JSONDecodeError:
+                            pass
+                    # If no embedded JSON found, return the original parsed result
+                    return parsed
+                else:
+                    return parsed
+            else:
+                return parsed
+        except json.JSONDecodeError:
+            return {"error": content_text}
+    return {}
+
+
+async def main():
+    """Main client function"""
+    # Server URL - adjust if running on different host/port
+    server_url = "http://localhost:4242/mcp"
+
+    # Path to the codebase - use container path since server runs in Docker
+    # Host path: playground/codebases/core -> Container path: /app/playground/codebases/core
+    codebase_path = "/app/playground/codebases/core"
+
+    logger.info("="*60)
+    logger.info("CODEBADGER TOOLKIT SAMPLE CLIENT")
+    logger.info("="*60)
+    logger.info(f"Server URL: {server_url}")
+    logger.info(f"Codebase: {codebase_path}")
+
+    try:
+        async with Client(server_url) as client:
+            logger.info("\n[1] Testing server connectivity...")
+            await client.ping()
+            logger.info("✓ Server is responding")
+
+            # ===== GENERATE CPG =====
+            logger.info("\n[2] Generating CPG for codebase...")
+            cpg_result = await client.call_tool("generate_cpg", {
+                "source_type": "local",
+                "source_path": codebase_path,
+                "language": "c"
+            })
+
+            cpg_dict = extract_tool_result(cpg_result)
+            logger.info(f"CPG generation result: {cpg_dict}")
+
+            if "codebase_hash" not in cpg_dict:
+                logger.error("❌ No codebase_hash returned")
+                return
+
+            codebase_hash = cpg_dict["codebase_hash"]
+            logger.info(f"✓ CPG generation initiated. Hash: {codebase_hash}")
+
+            # ===== WAIT FOR CPG TO BE READY =====
+            logger.info("\n[3] Waiting for CPG to be ready...")
+            cpg_ready = False
+            max_attempts = 30
+
+            for attempt in range(max_attempts):
+                await asyncio.sleep(2)  # Wait 2 seconds between checks
+
+                status_result = await client.call_tool("get_cpg_status", {
+                    "codebase_hash": codebase_hash
+                })
+
+                status_dict = extract_tool_result(status_result)
+                status = status_dict.get("status")
+                exists = status_dict.get("exists", False)
+
+                logger.info(f"  Attempt {attempt + 1}/{max_attempts}: status={status}, exists={exists}")
+
+                if status in ["ready", "cached"] and exists:
+                    cpg_ready = True
+                    logger.info("✓ CPG is ready!")
+                    break
+
+            if not cpg_ready:
+                logger.error("❌ CPG not ready after waiting")
+                return
+
+            # ===== LIST METHODS =====
+            logger.info("\n[4] Listing methods in the codebase...")
+            methods_result = await client.call_tool("list_methods", {
+                "codebase_hash": codebase_hash,
+                "limit": 20
+            })
+
+            methods_dict = extract_tool_result(methods_result)
+
+            if methods_dict.get("success"):
+                methods = methods_dict.get("methods", [])
+                total = methods_dict.get("total", 0)
+                logger.info(f"✓ Found {total} methods total, showing up to 20:")
+
+                for method in methods[:10]:  # Show first 10
+                    logger.info(f"  - {method.get('name', 'unknown')} in {method.get('filename', 'unknown')}")
+
+                if len(methods) > 10:
+                    logger.info(f"  ... and {len(methods) - 10} more methods")
+            else:
+                logger.error(f"❌ Failed to list methods: {methods_dict}")
+
+                        # ===== RUN SIMPLE QUERY - GET CODEBASE SUMMARY =====
+            logger.info("\n[5] Getting codebase summary...")
+
+            summary_result = await client.call_tool("get_codebase_summary", {
+                "codebase_hash": codebase_hash
+            })
+
+            summary_dict = extract_tool_result(summary_result)
+
+            if summary_dict.get("success"):
+                summary = summary_dict.get("summary", {})
+                logger.info("✓ Codebase summary retrieved:")
+                logger.info(f"  Language: {summary.get('language')}")
+                logger.info(f"  Files: {summary.get('total_files')}")
+                logger.info(f"  Methods: {summary.get('total_methods')}")
+                logger.info(f"  Calls: {summary.get('total_calls')}")
+            else:
+                logger.error(f"❌ Failed to get summary: {summary_dict}")
+
+            # ===== RUN ANOTHER QUERY - LIST CALLS =====
+            logger.info("\n[6] Listing function calls...")
+
+            calls_result = await client.call_tool("list_calls", {
+                "codebase_hash": codebase_hash,
+                "limit": 10
+            })
+
+            calls_dict = extract_tool_result(calls_result)
+
+            if calls_dict.get("success"):
+                calls = calls_dict.get("calls", [])
+                total = calls_dict.get("total", 0)
+                logger.info(f"✓ Found {total} calls total, showing up to 10:")
+
+                for call in calls[:5]:  # Show first 5
+                    logger.info(f"  - {call.get('caller', 'unknown')} -> {call.get('callee', 'unknown')}")
+
+                if len(calls) > 5:
+                    logger.info(f"  ... and {len(calls) - 5} more calls")
+            else:
+                logger.error(f"❌ Failed to list calls: {calls_dict}")
+
+            logger.info("\n" + "="*60)
+            logger.info("SAMPLE CLIENT COMPLETED SUCCESSFULLY!")
+            logger.info("="*60)
+
+    except Exception as e:
+        logger.error(f"❌ Client error: {e}", exc_info=True)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        logger.info("\n🛑 Client interrupted by user")
+        sys.exit(1)

README.md

@@ -18,9 +18,12 @@
     CodebaseTracker,
     GitManager,
     CPGGenerator,
+    JoernServerClient,
+    JoernServerManager,
+    PortManager,
     QueryExecutor
 )
-from src.utils import RedisClient, setup_logging
+from src.utils import RedisClient, SyncRedisClient, setup_logging
 from src.tools import register_tools
 
 # Version information - bump this when releasing new versions
@@ -47,30 +50,40 @@ async def lifespan(mcp: FastMCP):
     logger.info("Created required directories")
 
     try:
-        # Initialize Redis
+        # Initialize Redis clients
         redis_client = RedisClient(config.redis)
         await redis_client.connect()
-        logger.info("Redis client connected")
+        
+        sync_redis_client = SyncRedisClient(config.redis)
+        sync_redis_client.connect()
+        
+        logger.info("Redis clients connected")
 
         # Initialize services
         services['config'] = config
         services['redis'] = redis_client
-        services['codebase_tracker'] = CodebaseTracker(redis_client)
+        services['sync_redis'] = sync_redis_client
+        services['codebase_tracker'] = CodebaseTracker(sync_redis_client)
         services['git_manager'] = GitManager(config.storage.workspace_root)
 
-        # Initialize CPG generator (runs Joern CLI directly in container)
-        services['cpg_generator'] = CPGGenerator(config=config)
-        # Skip initialize() - no Docker needed
+        # Initialize port manager for Joern servers
+        services['port_manager'] = PortManager()
 
-        # Initialize query executor (runs Joern servers as local subprocesses)
-        services['query_executor'] = QueryExecutor(
-            config.query,
-            config.joern,
-            redis_client,
-            docker_orchestrator=None  # Will start Joern servers directly
+        # Initialize Joern server manager
+        services['joern_server_manager'] = JoernServerManager(
+            joern_binary_path=config.joern.binary_path
         )
+        
+        # Initialize CPG generator (runs Joern CLI directly in container)
+        services['cpg_generator'] = CPGGenerator(config=config, joern_server_manager=services['joern_server_manager'])
         # Skip initialize() - no Docker needed
 
+        # Initialize query executor with Joern server manager
+        services['query_executor'] = QueryExecutor(services['joern_server_manager'], config=config.query)
+        
+        # Register MCP tools now that services are initialized
+        register_tools(mcp, services)
+        
         logger.info("All services initialized")
         logger.info("CodeBadger Toolkit Server is ready")
 
@@ -79,12 +92,9 @@ async def lifespan(mcp: FastMCP):
         # Shutdown
         logger.info("Shutting down CodeBadger Toolkit Server")
 
-        # Cleanup query executor (stops any running Joern server subprocesses)
-        if 'query_executor' in services:
-            await services['query_executor'].cleanup()
-        
         # Close connections
         await redis_client.close()
+        sync_redis_client.close()
 
         logger.info("CodeBadger Toolkit Server shutdown complete")
 
@@ -99,8 +109,8 @@ async def lifespan(mcp: FastMCP):
     lifespan=lifespan
 )
 
-# Register MCP tools
-register_tools(mcp, services)
+# Note: Tools are registered inside the lifespan function
+# register_tools(mcp, services)
 
 
 # Health check endpoint