Purpose of the Go Backend in AI-Infra-Guard: Architecture and Core Functions

The Go backend serves as the standalone core engine of AI-Infra-Guard, coordinating all scanning operations, data management, and agent communication through a REST API, WebSocket interface, and modern Web UI.

The Go backend powers the Tencent AI-Infra-Guard platform as its central orchestration layer for AI infrastructure security scanning. This component provides the essential bridge between user interfaces, automated agents, and the scanning engine, exposing functionality through both command-line tools and programmable APIs according to the source code.

Core Architecture Overview

The Go backend operates as a standalone service that initializes the entire platform stack. In cmd/cli/main.go, the binary entry point builds the CLI application and dispatches sub-commands through cmd.Execute(), supporting both interactive webserver mode and direct scan execution.

The server initialization occurs in common/websocket/server.go, where the RunWebServer function orchestrates multiple critical subsystems:

• Initializes the trpc-go framework and Gin router
• Loads the SQLite database via database.InitDB
• Registers REST API endpoints and serves static UI assets
• Starts the WebSocket server for real-time agent communication

Configuration Management

Scanning behavior and runtime parameters are centralized in the options.Options struct defined in internal/options/options.go. This configuration holder manages:

• Target specifications and timeout values
• Rate limiting parameters
• Template directories for fingerprints and vulnerability signatures
• WebSocket addresses and external API-checker URLs

The options struct propagates through both the web server and the standalone scanner, ensuring consistent behavior across CLI and API-driven workflows.

Data Persistence Layer

The backend implements persistent storage through pkg/database/database.go, which initializes SQLite (or alternative database) tables for tasks, models, and file uploads. Key components include:

• database.InitDB: Establishes database connections and schema
• NewTaskStore: Manages scanning job records and results
• NewModelStore: Handles AI model definitions with auto-population of known entries

This layer ensures scan results, task metadata, and model configurations survive process restarts and remain available for historical analysis.

Task and Agent Orchestration

TaskManager Coordination

Located in common/websocket/task_manager.go, the TaskManager coordinates scanning jobs across the platform. It handles task creation, stores intermediate results, and pushes progress updates to clients via Server-Sent Events (SSE) through endpoints like GET /api/v1/app/tasks/sse/{sessionId}.

Real-Time Agent Communication

The AgentManager in common/websocket/agent.go exposes the /agents/ws WebSocket endpoint, enabling AI agents written in any language to:

• Receive task assignments in JSON format
• Report scanning results in real-time
• Maintain persistent connections for bidirectional communication

This architecture decouples the scanning logic from the backend while maintaining centralized control and data collection.

API Infrastructure and Documentation

REST API Endpoints

The backend registers comprehensive REST endpoints in common/websocket/api.go, including POST /api/v1/app/tasks for creating scanning jobs. The server also optionally proxies model-resolution requests to external services via apichecker.NewWithModelStore as implemented in common/websocket/server.go#L81-L88.

Swagger UI Integration

Automatic API documentation is served at /docs/* through Swagger UI integration, generated from Go annotations in common/websocket/server.go#L31-L33. This provides interactive documentation for the entire REST interface without manual maintenance.

Command-Line Interface

The backend exposes two primary CLI commands through cmd/cli/cmd/:

Webserver Command

The webserver sub-command (defined in cmd/cli/cmd/webserver.go) launches the full platform:

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

This initializes the Gin-based server, opens the configured port, and serves both the API and static frontend assets.

Scan Command

The scan sub-command (in cmd/cli/cmd/scan.go) enables direct execution without the web interface:

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

This command constructs an options.Options instance and invokes the Go runner directly for standalone scanning operations.

Practical Implementation Examples

Starting the Web UI Programmatically

To launch the platform from within Go code:

package main

import (
	"github.com/Tencent/AI-Infra-Guard/cmd/cli/cmd"
)

func main() {
	// Launches the full web UI on 127.0.0.1:8088
	cmd.Execute()
}

Configuring Scan Options

When implementing custom scanning logic:

scanOptions := &options.Options{
	Target:        []string{"http://127.0.0.1:8000"},
	FPTemplates:   "data/fingerprints",
	AdvTemplates:  "data/vuln",
	TimeOut:       10,
	RateLimit:     200,
}
r, _ := runner.New(scanOptions)
r.RunEnumeration()

Creating Tasks via REST API

External systems can trigger scans through the API:

POST http://localhost:8088/api/v1/app/tasks
Content-Type: application/json

{
  "targets": ["http://127.0.0.1:8000"],
  "fpTemplates": "data/fingerprints",
  "advTemplates": "data/vuln"
}

Connecting AI Agents

Agents establish WebSocket connections to receive tasks:

wsURL := "ws://127.0.0.1:8088/agents/ws"
conn, _ := websocket.Dial(wsURL, "", "http://localhost/")

Summary

• The Go backend functions as a standalone service orchestrating all AI-Infra-Guard operations, from cmd/cli/main.go through the full server stack.
• It manages configuration through internal/options/options.go, centralizing timeouts, rate limits, and template paths.
• Data persistence relies on pkg/database/database.go with SQLite support for tasks and model storage.
• Real-time communication occurs via WebSocket endpoints in common/websocket/agent.go and SSE streams managed by the TaskManager.
• The backend exposes functionality through both a CLI interface (scan and webserver commands) and a comprehensive REST API with Swagger documentation.
• All Python MCP-scan, Agent-scan, and Prompt-Security components depend on this Go backend for task scheduling and result persistence.

Frequently Asked Questions

What protocols does the AI-Infra-Guard Go backend use for communication?

The Go backend utilizes HTTP/HTTPS for REST API endpoints and WebSocket connections for real-time agent communication. It also implements Server-Sent Events (SSE) for streaming task progress updates to web clients, as defined in common/websocket/task_manager.go.

How does the Go backend handle database operations?

The backend initializes SQLite (or configured alternatives) through database.InitDB in pkg/database/database.go. It uses specialized stores including NewTaskStore for scanning jobs and NewModelStore for AI model definitions, providing persistent storage for all scan results and metadata.

Can the Go backend run independently of the Web UI?

Yes, the backend supports standalone operation through the scan CLI command implemented in cmd/cli/cmd/scan.go. This allows direct execution of vulnerability scans without launching the webserver, making it suitable for CI/CD pipelines and automated security testing.

Where is the API documentation generated in the Go backend?

Automatic API documentation is generated from Go annotations and served via Swagger UI at the /docs/* endpoint, as configured in common/websocket/server.go#L31-L33. This provides interactive documentation for all REST endpoints without requiring separate documentation maintenance.

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