ai-gateway-go/cmd/agent-node/main.go

// Command agent-node 是智能体节点 Agent:向 AI Gateway 注册的心跳 + 任务执行器。
//
// 职责:
//   - 周期性心跳(上报版本/能力/最近错误),刷新在线状态;
//   - 轮询认领节点池任务(POST /v1/agent/nodes/{code}/tasks/claim);
//   - 按任务类型执行:prompt(本地/网关 LLM)、http(受控出站)、mcp_invoke /
//     skill_run / digital_employee(经网关运行时转发)、custom(回显);
//   - 凭认领令牌上报结果或错误,失败交给网关按退避重试。
//
// 安全边界:
//   - 不执行任意 shell 命令;http 任务只允许 http(s) 绝对地址且拒绝重定向;
//   - 任务载荷是网关下发的可信 JSON,节点只做类型化执行;
//   - 节点令牌经环境变量/配置文件注入,不经任务载荷传递。
package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

type config struct {
	GatewayURL   string
	Code         string
	Token        string
	Version      string
	Name         string
	Capabilities []string
	// LLM 执行:优先本地 OpenAI 兼容端点,缺省回退网关。
	LLMBaseURL string
	LLMAPIKey  string
	// 转发类任务(mcp/skill/数字员工)与回退 LLM 用的网关 API Key。
	GatewayKey     string
	HeartbeatEvery time.Duration
	ClaimEvery     time.Duration
	HTTPTimeout    time.Duration
}

func loadConfig() (config, error) {
	cfg := config{
		GatewayURL:     strings.TrimRight(env("AGENT_GATEWAY_URL", "http://127.0.0.1:8080"), "/"),
		Code:           env("AGENT_CODE", ""),
		Token:          env("AGENT_TOKEN", ""),
		Version:        env("AGENT_VERSION", "0.11.4-agent"),
		Name:           env("AGENT_NAME", ""),
		LLMBaseURL:     strings.TrimRight(env("AGENT_LLM_BASE_URL", ""), "/"),
		LLMAPIKey:      env("AGENT_LLM_API_KEY", ""),
		GatewayKey:     env("AGENT_GATEWAY_KEY", ""),
		HeartbeatEvery: durEnv("AGENT_HEARTBEAT_SECONDS", 30*time.Second),
		ClaimEvery:     durEnv("AGENT_CLAIM_SECONDS", 5*time.Second),
		HTTPTimeout:    durEnv("AGENT_HTTP_TIMEOUT_SECONDS", 60*time.Second),
	}
	for _, capability := range strings.Split(env("AGENT_CAPABILITIES", "tool_exec,llm,http"), ",") {
		if value := strings.TrimSpace(capability); value != "" {
			cfg.Capabilities = append(cfg.Capabilities, value)
		}
	}
	if cfg.Code == "" || cfg.Token == "" {
		return cfg, errors.New("AGENT_CODE 与 AGENT_TOKEN 必填(管理端登记节点后交付)")
	}
	if cfg.HeartbeatEvery < time.Second || cfg.ClaimEvery < time.Second || cfg.HTTPTimeout < time.Second {
		return cfg, errors.New("心跳/认领/超时秒数必须 >= 1")
	}
	return cfg, nil
}

func env(key, fallback string) string {
	if value := strings.TrimSpace(os.Getenv(key)); value != "" {
		return value
	}
	return fallback
}

func durEnv(key string, fallback time.Duration) time.Duration {
	if value := os.Getenv(key); value != "" {
		if parsed, err := time.ParseDuration(value + "s"); err == nil {
			return parsed
		}
	}
	return fallback
}

type agent struct {
	cfg    config
	client *http.Client
	logger *slog.Logger
	// lastError 是最近一次心跳/执行错误,随下一次心跳上报。
	lastError string
}

func main() {
	flag.Parse()
	logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
	cfg, err := loadConfig()
	if err != nil {
		logger.Error("invalid configuration", "error", err)
		os.Exit(1)
	}
	a := &agent{
		cfg:    cfg,
		client: &http.Client{Timeout: cfg.HTTPTimeout},
		logger: logger,
	}
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()
	logger.Info("agent node starting", "gateway", cfg.GatewayURL, "code", cfg.Code, "capabilities", cfg.Capabilities)
	run(ctx, a)
	logger.Info("agent node stopped")
}

func run(ctx context.Context, a *agent) {
	heartbeatTicker := time.NewTicker(a.cfg.HeartbeatEvery)
	claimTicker := time.NewTicker(a.cfg.ClaimEvery)
	defer heartbeatTicker.Stop()
	defer claimTicker.Stop()
	// 启动立即做一次心跳与认领,便于部署后快速验证。
	if err := a.heartbeat(ctx); err != nil {
		a.logger.Warn("initial heartbeat failed", "error", err)
		a.lastError = err.Error()
	} else {
		a.lastError = ""
	}
	for {
		select {
		case <-ctx.Done():
			return
		case <-heartbeatTicker.C:
			if err := a.heartbeat(ctx); err != nil {
				a.logger.Warn("heartbeat failed", "error", err)
				a.lastError = err.Error()
			} else {
				a.lastError = ""
			}
		case <-claimTicker.C:
			a.processOne(ctx)
		}
	}
}

// heartbeat 上报版本/能力/最近错误。
func (a *agent) heartbeat(ctx context.Context) error {
	payload, _ := json.Marshal(map[string]any{
		"version":      a.cfg.Version,
		"capabilities": capabilityMap(a.cfg.Capabilities),
		"metadata":     map[string]string{"name": a.cfg.Name, "agent": "agent-node"},
		"error":        a.lastError,
	})
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, a.cfg.GatewayURL+"/v1/agent/nodes/"+a.cfg.Code+"/heartbeat", bytes.NewReader(payload))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("X-Agent-Token", a.cfg.Token)
	response, err := a.client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode/100 != 2 {
		return fmt.Errorf("heartbeat HTTP %d", response.StatusCode)
	}
	return nil
}

// processOne 认领并执行一个任务(串行;一次一个)。
func (a *agent) processOne(ctx context.Context) {
	claimCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
	defer cancel()
	task, err := a.claim(claimCtx)
	if err != nil {
		a.logger.Warn("claim failed", "error", err)
		a.lastError = "claim: " + err.Error()
		return
	}
	if task == nil {
		return
	}
	a.logger.Info("task claimed", "task_id", task.ID, "task_type", task.TaskType)
	result, runErr := a.execute(ctx, *task)
	if runErr != nil {
		a.logger.Warn("task execution failed", "task_id", task.ID, "error", runErr)
		a.lastError = "task " + task.ID[:8] + ": " + runErr.Error()
	}
	if err := a.complete(ctx, task.ID, task.ClaimToken, result, runErr); err != nil {
		a.logger.Warn("task report failed", "task_id", task.ID, "error", err)
		a.lastError = "report: " + err.Error()
		return
	}
	a.logger.Info("task reported", "task_id", task.ID, "success", runErr == nil)
}

type claimedTask struct {
	ID         string          `json:"id"`
	TaskType   string          `json:"task_type"`
	Payload    json.RawMessage `json:"payload"`
	ClaimToken string          `json:"claim_token"`
}

func (a *agent) claim(ctx context.Context) (*claimedTask, error) {
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, a.cfg.GatewayURL+"/v1/agent/nodes/"+a.cfg.Code+"/tasks/claim", nil)
	if err != nil {
		return nil, err
	}
	request.Header.Set("X-Agent-Token", a.cfg.Token)
	response, err := a.client.Do(request)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(response.Body, 1<<20))
	if response.StatusCode/100 != 2 {
		return nil, fmt.Errorf("claim HTTP %d: %s", response.StatusCode, strings.TrimSpace(string(raw)))
	}
	// 网关统一 apiresponse 信封:{code,msg,data:{task}}。
	var envelope struct {
		Code int `json:"code"`
		Data struct {
			Task *claimedTask `json:"task"`
		} `json:"data"`
	}
	if err := json.Unmarshal(raw, &envelope); err != nil {
		return nil, err
	}
	if envelope.Code != 200 {
		return nil, fmt.Errorf("claim 业务错误 code=%d", envelope.Code)
	}
	return envelope.Data.Task, nil
}

func (a *agent) complete(ctx context.Context, taskID, claimToken string, result map[string]any, runErr error) error {
	var errorText string
	if runErr != nil {
		errorText = runErr.Error()
		if len(errorText) > 4000 {
			errorText = errorText[:4000]
		}
	}
	payload, _ := json.Marshal(map[string]any{"claim_token": claimToken, "result": result, "error": errorText})
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, a.cfg.GatewayURL+"/v1/agent/nodes/"+a.cfg.Code+"/tasks/"+taskID+"/complete", bytes.NewReader(payload))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("X-Agent-Token", a.cfg.Token)
	response, err := a.client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode/100 != 2 {
		return fmt.Errorf("complete HTTP %d", response.StatusCode)
	}
	return nil
}

// execute 按任务类型执行。载荷来自受管网关,字段类型不符返回错误。
func (a *agent) execute(ctx context.Context, task claimedTask) (map[string]any, error) {
	var payload map[string]any
	if err := json.Unmarshal(task.Payload, &payload); err != nil {
		return nil, fmt.Errorf("task payload 不是 JSON 对象: %w", err)
	}
	switch task.TaskType {
	case "prompt":
		return a.executePrompt(ctx, payload)
	case "http":
		return a.executeHTTP(ctx, payload)
	case "mcp_invoke":
		return a.forwardGateway(ctx, "/v1/mcp-servers/"+stringValue(payload, "code")+"/tools/"+stringValue(payload, "tool")+"/invoke", map[string]any{"input": payload["args"]})
	case "skill_run":
		return a.forwardGateway(ctx, "/v1/skills/"+stringValue(payload, "code")+"/render", map[string]any{"variables": payload["variables"]})
	case "digital_employee":
		return a.forwardGateway(ctx, "/v1/digital-employees/"+stringValue(payload, "code")+"/chat/completions", map[string]any{"messages": payload["messages"]})
	case "custom":
		return map[string]any{"echo": payload}, nil
	}
	return nil, fmt.Errorf("不支持的节点任务类型 %s", task.TaskType)
}

// executePrompt 调用本地 LLM(OpenAI 兼容);未配置时回退网关 /v1/chat/completions。
func (a *agent) executePrompt(ctx context.Context, payload map[string]any) (map[string]any, error) {
	model := stringValue(payload, "model")
	messages, _ := payload["messages"].([]any)
	if model == "" || len(messages) == 0 {
		return nil, errors.New("prompt 任务需要 model 与 messages")
	}
	if a.cfg.LLMBaseURL != "" {
		return a.callOpenAICompatible(ctx, a.cfg.LLMBaseURL+"/chat/completions", a.cfg.LLMAPIKey, model, messages)
	}
	if a.cfg.GatewayKey == "" {
		return nil, errors.New("未配置 AGENT_LLM_BASE_URL 或 AGENT_GATEWAY_KEY,无法执行 prompt 任务")
	}
	body, _ := json.Marshal(map[string]any{"model": model, "messages": messages, "stream": false})
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, a.cfg.GatewayURL+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer "+a.cfg.GatewayKey)
	response, err := a.client.Do(request)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(response.Body, 2<<20))
	if response.StatusCode/100 != 2 {
		return nil, fmt.Errorf("网关 LLM HTTP %d: %s", response.StatusCode, truncate(string(raw), 300))
	}
	var decoded map[string]any
	if err := json.Unmarshal(raw, &decoded); err != nil {
		return nil, err
	}
	answer, _ := firstChoiceContent(decoded)
	return map[string]any{"answer": answer, "source": "gateway"}, nil
}

// callOpenAICompatible 调用本地 OpenAI 兼容端点(节点侧 LLM)。
func (a *agent) callOpenAICompatible(ctx context.Context, endpoint, apiKey, model string, messages []any) (map[string]any, error) {
	body, _ := json.Marshal(map[string]any{"model": model, "messages": messages, "stream": false})
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	request.Header.Set("Content-Type", "application/json")
	if apiKey != "" {
		request.Header.Set("Authorization", "Bearer "+apiKey)
	}
	response, err := a.client.Do(request)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(response.Body, 2<<20))
	if response.StatusCode/100 != 2 {
		return nil, fmt.Errorf("本地 LLM HTTP %d: %s", response.StatusCode, truncate(string(raw), 300))
	}
	var decoded map[string]any
	if err := json.Unmarshal(raw, &decoded); err != nil {
		return nil, err
	}
	answer, _ := firstChoiceContent(decoded)
	return map[string]any{"answer": answer, "source": "local"}, nil
}

// executeHTTP 执行受控 HTTP 任务:只允许 http(s),拒绝重定向,响应限 1 MiB。
func (a *agent) executeHTTP(ctx context.Context, payload map[string]any) (map[string]any, error) {
	target := strings.TrimSpace(stringValue(payload, "url"))
	method := strings.ToUpper(strings.TrimSpace(stringValue(payload, "method")))
	if method == "" {
		method = "GET"
	}
	parsed, err := url.Parse(target)
	if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Hostname() == "" {
		return nil, errors.New("http 任务 url 必须是 http(s) 绝对地址")
	}
	var body io.Reader
	if value, ok := payload["body"].(string); ok && strings.TrimSpace(value) != "" {
		body = strings.NewReader(value)
	}
	request, err := http.NewRequestWithContext(ctx, method, parsed.String(), body)
	if err != nil {
		return nil, err
	}
	if headers, ok := payload["headers"].(map[string]any); ok {
		for key, value := range headers {
			request.Header.Set(key, fmt.Sprint(value))
		}
	}
	client := &http.Client{
		Timeout: a.cfg.HTTPTimeout,
		CheckRedirect: func(*http.Request, []*http.Request) error {
			return errors.New("http 任务不允许重定向")
		},
	}
	response, err := client.Do(request)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(response.Body, (1<<20)+1))
	if len(raw) > 1<<20 {
		return nil, errors.New("http 任务响应超过 1 MiB")
	}
	var decoded any = string(raw)
	if json.Unmarshal(raw, &decoded) != nil {
		decoded = string(raw)
	}
	return map[string]any{"status_code": response.StatusCode, "body": decoded}, nil
}

// forwardGateway 经网关运行时执行 mcp/skill/数字员工(需要 AGENT_GATEWAY_KEY)。
func (a *agent) forwardGateway(ctx context.Context, path string, bodyValue map[string]any) (map[string]any, error) {
	if a.cfg.GatewayKey == "" {
		return nil, errors.New("未配置 AGENT_GATEWAY_KEY,无法执行网关转发类任务")
	}
	body, _ := json.Marshal(bodyValue)
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, a.cfg.GatewayURL+path, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer "+a.cfg.GatewayKey)
	response, err := a.client.Do(request)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(response.Body, 2<<20))
	if response.StatusCode/100 != 2 {
		return nil, fmt.Errorf("网关转发 HTTP %d: %s", response.StatusCode, truncate(string(raw), 300))
	}
	var decoded map[string]any
	if err := json.Unmarshal(raw, &decoded); err != nil {
		return nil, err
	}
	return decoded, nil
}

func capabilityMap(capabilities []string) map[string]bool {
	result := map[string]bool{}
	for _, capability := range capabilities {
		result[capability] = true
	}
	return result
}

func stringValue(payload map[string]any, key string) string {
	value, _ := payload[key].(string)
	return strings.TrimSpace(value)
}

func firstChoiceContent(decoded map[string]any) (string, bool) {
	choices, _ := decoded["choices"].([]any)
	if len(choices) == 0 {
		return "", false
	}
	choice, _ := choices[0].(map[string]any)
	message, _ := choice["message"].(map[string]any)
	content, _ := message["content"].(string)
	return content, strings.TrimSpace(content) != ""
}

func truncate(value string, max int) string {
	runes := []rune(value)
	if len(runes) > max {
		return string(runes[:max]) + "…"
	}
	return value
}