Components of the AI-Infra-Guard System Architecture: Tencent's Modular AI Security Platform

The AI-Infra-Guard system architecture comprises eight core components: a high-performance Go engine for web services and CLI operations, specialized Python scanners for MCP servers and AI agents, a YAML-based data rule system, WebSocket API layer, React frontend, Docker orchestration, utility libraries, and an extensible plugin framework.

The AI-Infra-Guard (AIG) platform developed by Tencent delivers comprehensive security scanning for AI infrastructure, model serving components, and autonomous agent workflows. Understanding the components of the AI-Infra-Guard system architecture reveals how this open-source project combines a high-performance Go core with flexible Python modules to enable six critical security capabilities: ClawScan, Agent Scan, MCP Server evaluation, AI-Infra Vulnerability Scanning, Jailbreak Evaluation, and Model API Relay Checking. The modular design allows security teams to deploy individual scanners or operate the full integrated platform via Docker Compose.

Core Architectural Components

Go Core Engine (Web and CLI)

The Go core provides the foundational infrastructure for the platform, handling task orchestration, rule engine operations, and both web service and CLI interfaces. In cmd/cli/main.go, the CLI entry point initializes the web server and scanning workflows using the urfave/cli/v2 framework, while cmd/agent/main.go manages agent processes that register with the central controller. The core implements a WebSocket gateway in common/websocket/websocket.go that enables real-time bidirectional communication between the server and distributed scanning agents.

Python Scanner Modules

Three specialized Python scanners extend the platform's detection capabilities for specific AI threat surfaces. The mcp-scan module (located in mcp-scan/main.py) performs static and dynamic analysis of Model Context Protocol (MCP) servers and agent skills, identifying risky permission sets and capability leaks. The agent-scan module (agent-scan/main.py) conducts automated red-team evaluations of AI agent platforms including Dify and Coze. AIG-PromptSecurity provides comprehensive prompt jailbreak assessment and adversarial safety testing against large language models.

Data-Driven Rule System

AI-Infra-Guard utilizes a hierarchical YAML-based rule structure stored in the data/ directory that drives all detection logic without code changes. The system organizes security intelligence into three categories: data/fingerprints/ for service identification (detecting Ollama, vLLM, ComfyUI, etc.), data/vuln/ for CVE vulnerability definitions, and data/mcp/ for MCP-specific security rules. The cmd/yamlcheck tool validates rule syntax and schema compliance before deployment, ensuring the knowledgebase remains consistent across distributed deployments.

WebSocket API and Task Management

The architecture implements a centralized task management system through common/websocket/task_manager.go, which maintains job queues and streams progress updates via WebSocket connections. This layer exposes JSON-RPC APIs under /api/v1/ endpoints, handling task submission, version reporting, and plugin registration. The WebSocket implementation enables long-running security assessments to report real-time status to the frontend while storing intermediate results in the database layer.

Frontend Visualization Layer

A React-based single-page application located in frontend/ provides the dashboard interface for security operators. The frontend renders real-time scan progress, vulnerability reports with severity classifications, and result visualizations by consuming the WebSocket API. This component runs as a standalone service within the Docker Compose stack, communicating exclusively through the Go backend's API endpoints.

Extensible Plugin Framework

The internal/mcp/plugins.go file implements a registration system allowing community contributors to extend scanning capabilities through a standardized interface. Plugins integrate with the existing task orchestration pipeline in common/websocket/task_manager.go, enabling custom scanners to utilize the WebSocket API, database storage, and result formatting infrastructure without modifying core source code.

Deployment and Infrastructure Components

Docker Orchestration

The platform ships with docker-compose.yml and docker-compose.images.yml files that containerize the Go web service, Python scanner workers, and React frontend into a unified deployment stack. This configuration enables single-command deployment (docker-compose up) of the complete AI-Infra-Guard architecture, automatically networking the WebSocket gateway, database volumes, and scanner modules.

Database and Utility Libraries

The pkg/database/ directory contains lightweight persistence implementations using SQLite and Bolt databases for storing scan results, task queues, and configuration state. HTTP client utilities in pkg/httpx/ provide standardized networking capabilities with custom transport configurations used across both the core engine and scanner modules for reliable asset discovery.

Implementation Examples

Running a full AI-Infra scan from the CLI:

# Build the Go binary from cmd/cli/main.go

go build -o ai-infra-guard ./cmd/cli/main.go

# Start the web server on port 8088

./ai-infra-guard webserver --server 127.0.0.1:8088 &

# Execute an AI-infrastructure scan against a vLLM instance

./ai-infra-guard scan -t http://127.0.0.1:8000

Invoking the Python MCP scanner directly:

pip install -r mcp-scan/requirements.txt
python mcp-scan/main.py --repo /path/to/mcp-server

Submitting a skill scan via the pip package:

pip install aig-skill-scan
export LLM_API_KEY="your-api-key"
aig-skill-scan --repo ./my-skill -m deepseek-v4-flash -o result.json

Validating custom fingerprint rules before deployment:

go build -o yamlcheck ./cmd/yamlcheck
./yamlcheck data/fingerprints

Summary

• Go Core Engine: Implements CLI (cmd/cli/main.go), agent processes (cmd/agent/main.go), and WebSocket gateway (common/websocket/websocket.go) for task orchestration
• Python Scanners: Specialized modules (mcp-scan/main.py, agent-scan/main.py) for MCP servers, AI agents, and prompt security testing
• YAML Rule System: Version-controlled detection logic in data/fingerprints/, data/vuln/, and data/mcp/ with validation via cmd/yamlcheck
• WebSocket API: Real-time communication and centralized task management through common/websocket/task_manager.go
• React Frontend: Visualization dashboard residing in frontend/ directory
• Plugin Framework: Extensible architecture via internal/mcp/plugins.go supporting custom scanner registration
• Docker Deployment: Containerized orchestration using docker-compose.yml for full-stack deployment
• Utility Layer: Database persistence (pkg/database/) and HTTP utilities (pkg/httpx/) supporting cross-platform operations

Frequently Asked Questions

What programming languages does AI-Infra-Guard use?

The platform implements performance-critical services in Go (including the web server, CLI, and task management in cmd/cli/main.go and common/websocket/task_manager.go), while utilizing Python for specialized security scanners targeting MCP servers and AI agents. This hybrid architecture balances the execution speed required for network scanning with Python's rich ecosystem of AI/ML libraries.

How does the plugin framework extend scanning capabilities?

The internal/mcp/plugins.go file implements a registration system that allows developers to add custom detection modules without modifying the core codebase. Plugins integrate with the central task manager in common/websocket/task_manager.go, enabling new scanners to utilize the existing WebSocket API, database storage, and result streaming infrastructure.

Where are detection rules stored and how are they validated?

Security rules reside as YAML files in data/fingerprints/, data/vuln/, and data/mcp/. The cmd/yamlcheck utility validates rule syntax and schema compliance before deployment, ensuring consistency across the knowledgebase. This data-driven approach allows rapid updates to detection logic for emerging AI vulnerabilities without recompiling the Go core.

Can AI-Infra-Guard run in containerized environments?

Yes, the repository includes docker-compose.yml configurations that orchestrate the Go web service, Python worker containers, and React frontend into a unified deployment. This architecture supports both single-node deployments and distributed scanning scenarios where agents connect to a central server via the WebSocket gateway implemented in common/websocket/websocket.go.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}