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

> Explore the Go backend purpose in AI-Infra-Guard. It's the core engine managing scans, data, and agent communication via REST API, WebSocket, and UI.

- Repository: [Tencent/AI-Infra-Guard](https://github.com/tencent/AI-Infra-Guard)
- Tags: architecture
- Published: 2026-08-22

---

**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`](https://github.com/Tencent/AI-Infra-Guard/blob/main/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`](https://github.com/Tencent/AI-Infra-Guard/blob/main/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`](https://github.com/Tencent/AI-Infra-Guard/blob/main/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`](https://github.com/Tencent/AI-Infra-Guard/blob/main/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`](https://github.com/Tencent/AI-Infra-Guard/blob/main/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`](https://github.com/Tencent/AI-Infra-Guard/blob/main/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`](https://github.com/Tencent/AI-Infra-Guard/blob/main/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`](https://github.com/Tencent/AI-Infra-Guard/blob/main/cmd/cli/cmd/webserver.go)) launches the full platform:

```bash
./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`](https://github.com/Tencent/AI-Infra-Guard/blob/main/cmd/cli/cmd/scan.go)) enables direct execution without the web interface:

```bash
./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:

```go
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:

```go
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:

```http
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:

```go
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`](https://github.com/Tencent/AI-Infra-Guard/blob/main/cmd/cli/main.go) through the full server stack.
- It manages **configuration** through [`internal/options/options.go`](https://github.com/Tencent/AI-Infra-Guard/blob/main/internal/options/options.go), centralizing timeouts, rate limits, and template paths.
- **Data persistence** relies on [`pkg/database/database.go`](https://github.com/Tencent/AI-Infra-Guard/blob/main/pkg/database/database.go) with SQLite support for tasks and model storage.
- **Real-time communication** occurs via WebSocket endpoints in [`common/websocket/agent.go`](https://github.com/Tencent/AI-Infra-Guard/blob/main/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`](https://github.com/Tencent/AI-Infra-Guard/blob/main/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`](https://github.com/Tencent/AI-Infra-Guard/blob/main/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`](https://github.com/Tencent/AI-Infra-Guard/blob/main/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.