ai-gateway-go/internal/agentnode/store.go

package agentnode

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/url"
	"regexp"
	"sort"
	"strings"
	"time"

	platformid "aigateway.local/core/internal/platform/id"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
)

var (
	ErrNotFound     = errors.New("agent node not found")
	ErrConflict     = errors.New("agent node already exists")
	ErrInvalidToken = errors.New("agent node token invalid")
	ErrInvalidInput = errors.New("agent node input invalid")
	ErrStore        = errors.New("agent node store unavailable")
)

var nodeCodePattern = regexp.MustCompile(`^[a-z0-9][a-z0-9._-]{0,127}$`)

type Store struct{ pool *pgxpool.Pool }

func NewStore(pool *pgxpool.Pool) *Store { return &Store{pool: pool} }

type Node struct {
	ID              string          `json:"id"`
	Code            string          `json:"code"`
	Name            string          `json:"name"`
	Description     string          `json:"description"`
	Endpoint        string          `json:"endpoint"`
	NodeType        string          `json:"node_type"`
	PoolType        string          `json:"pool_type"`
	PoolCode        string          `json:"pool_code"`
	Enabled         bool            `json:"enabled"`
	Status          string          `json:"status"`
	TokenPrefix     string          `json:"token_prefix"`
	Version         string          `json:"version"`
	Capabilities    json.RawMessage `json:"capabilities"`
	Metadata        json.RawMessage `json:"metadata"`
	LastHeartbeatAt *time.Time      `json:"last_heartbeat_at,omitempty"`
	LastHeartbeatIP string          `json:"last_heartbeat_ip,omitempty"`
	LastError       string          `json:"last_error"`
	CreatedAt       time.Time       `json:"created_at"`
	UpdatedAt       time.Time       `json:"updated_at"`
}

type CreateInput struct {
	Code        string
	Name        string
	Description string
	Endpoint    string
	NodeType    string
	PoolType    string
	PoolCode    string
	Enabled     bool
}

type UpdateInput struct {
	ID          string
	Name        string
	Description string
	Endpoint    string
	NodeType    string
	PoolType    string
	PoolCode    string
	Enabled     bool
}

type HeartbeatInput struct {
	Version      string         `json:"version"`
	Capabilities map[string]any `json:"capabilities"`
	Metadata     map[string]any `json:"metadata"`
	Error        string         `json:"error"`
}

// RoutePreviewInput describes the node-pool constraints used by the read-only
// routing preview. It deliberately contains no task or request payload: the
// preview only validates candidate selection before a remote executor exists.
type RoutePreviewInput struct {
	PoolType             string   `json:"pool_type"`
	PoolCode             string   `json:"pool_code"`
	RequiredCapabilities []string `json:"required_capabilities"`
	RequestKey           string   `json:"request_key"`
}

type RoutePreview struct {
	PoolType             string   `json:"pool_type"`
	PoolCode             string   `json:"pool_code"`
	RequiredCapabilities []string `json:"required_capabilities"`
	RequestKey           string   `json:"request_key"`
	SelectionPolicy      string   `json:"selection_policy"`
	Reason               string   `json:"reason"`
	Selected             *Node    `json:"selected"`
	Candidates           []Node   `json:"candidates"`
}

const nodeSelect = `SELECT n.id::text,n.code,n.name,n.description,n.endpoint,n.node_type,n.pool_type,n.pool_code,n.enabled,
CASE WHEN NOT n.enabled THEN 'disabled' WHEN n.last_heartbeat_at IS NULL THEN 'pending' WHEN n.last_heartbeat_at < clock_timestamp()-interval '90 seconds' THEN 'offline' ELSE 'online' END,
n.token_prefix,n.version,n.capabilities,n.metadata,n.last_heartbeat_at,coalesce(host(n.last_heartbeat_ip),''),n.last_error,n.created_at,n.updated_at
FROM gateway.agent_nodes n`

func normalizeJSON(raw []byte) json.RawMessage {
	if len(raw) == 0 || !json.Valid(raw) {
		return json.RawMessage(`{}`)
	}
	return raw
}

func objectJSON(value map[string]any) ([]byte, error) {
	if value == nil {
		return nil, nil
	}
	raw, err := json.Marshal(value)
	if err != nil {
		return nil, fmt.Errorf("%w: metadata cannot be encoded", ErrInvalidInput)
	}
	return raw, nil
}

func validateCommon(code, name, description, endpoint, nodeType, poolType, poolCode string) error {
	if !nodeCodePattern.MatchString(code) || strings.ToLower(code) != code {
		return fmt.Errorf("%w: code must use lowercase letters, numbers, dot, underscore or hyphen", ErrInvalidInput)
	}
	if strings.TrimSpace(name) == "" || len(name) > 128 || len(description) > 4000 || len(endpoint) > 512 {
		return fmt.Errorf("%w: node fields exceed their limits", ErrInvalidInput)
	}
	if nodeType != "worker" && nodeType != "gateway" && nodeType != "executor" {
		return fmt.Errorf("%w: node type is invalid", ErrInvalidInput)
	}
	if poolType != "public" && poolType != "private" {
		return fmt.Errorf("%w: pool type is invalid", ErrInvalidInput)
	}
	if strings.TrimSpace(poolCode) == "" || len(poolCode) > 64 {
		return fmt.Errorf("%w: pool code is invalid", ErrInvalidInput)
	}
	// Endpoint 将来可能被节点池路由直接拨号,必须保证是干净的 http(s)
	// 绝对地址(无 userinfo/query/fragment)。不做 DNS 解析:节点本身常部署
	// 在内网,不能按公网规则校验。
	if strings.TrimSpace(endpoint) != "" {
		parsed, parseErr := url.Parse(strings.TrimSpace(endpoint))
		if parseErr != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Hostname() == "" ||
			parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" {
			return fmt.Errorf("%w: endpoint must be an absolute http(s) URL without user info, query or fragment", ErrInvalidInput)
		}
	}
	return nil
}

func generateToken() (string, string, []byte, error) {
	raw := make([]byte, 32)
	if _, err := rand.Read(raw); err != nil {
		return "", "", nil, err
	}
	secret := "agn_" + base64.RawURLEncoding.EncodeToString(raw)
	prefix := secret[:12]
	digest := sha256.Sum256([]byte(secret))
	return secret, prefix, digest[:], nil
}

func scanNode(row pgx.Row) (Node, error) {
	var item Node
	err := row.Scan(&item.ID, &item.Code, &item.Name, &item.Description, &item.Endpoint, &item.NodeType, &item.PoolType, &item.PoolCode, &item.Enabled, &item.Status, &item.TokenPrefix, &item.Version, &item.Capabilities, &item.Metadata, &item.LastHeartbeatAt, &item.LastHeartbeatIP, &item.LastError, &item.CreatedAt, &item.UpdatedAt)
	if errors.Is(err, pgx.ErrNoRows) {
		return Node{}, ErrNotFound
	}
	item.Capabilities = normalizeJSON(item.Capabilities)
	item.Metadata = normalizeJSON(item.Metadata)
	return item, err
}

func (s *Store) Create(ctx context.Context, input CreateInput, actorID string) (Node, string, error) {
	if s == nil || s.pool == nil {
		return Node{}, "", ErrStore
	}
	input.Code = strings.ToLower(strings.TrimSpace(input.Code))
	input.Name = strings.TrimSpace(input.Name)
	input.Description = strings.TrimSpace(input.Description)
	input.Endpoint = strings.TrimSpace(input.Endpoint)
	input.NodeType = strings.TrimSpace(input.NodeType)
	input.PoolType = strings.TrimSpace(input.PoolType)
	input.PoolCode = strings.TrimSpace(input.PoolCode)
	if input.NodeType == "" {
		input.NodeType = "worker"
	}
	if input.PoolType == "" {
		input.PoolType = "private"
	}
	if input.PoolCode == "" {
		input.PoolCode = "default"
	}
	if err := validateCommon(input.Code, input.Name, input.Description, input.Endpoint, input.NodeType, input.PoolType, input.PoolCode); err != nil {
		return Node{}, "", err
	}
	id, err := platformid.NewUUID()
	if err != nil {
		return Node{}, "", err
	}
	secret, prefix, digest, err := generateToken()
	if err != nil {
		return Node{}, "", err
	}
	_, err = s.pool.Exec(ctx, `INSERT INTO gateway.agent_nodes(id,code,name,description,endpoint,node_type,pool_type,pool_code,enabled,token_prefix,token_hash,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,nullif($12,'')::uuid)`, id, input.Code, input.Name, input.Description, input.Endpoint, input.NodeType, input.PoolType, input.PoolCode, input.Enabled, prefix, digest, actorID)
	if err != nil {
		var pgErr *pgconn.PgError
		if errors.As(err, &pgErr) && pgErr.Code == "23505" {
			return Node{}, "", ErrConflict
		}
		return Node{}, "", fmt.Errorf("%w: %v", ErrStore, err)
	}
	item, err := s.Get(ctx, id)
	return item, secret, err
}

func (s *Store) Get(ctx context.Context, id string) (Node, error) {
	if s == nil || s.pool == nil {
		return Node{}, ErrStore
	}
	return scanNode(s.pool.QueryRow(ctx, nodeSelect+` WHERE n.id=$1`, id))
}

func (s *Store) List(ctx context.Context) ([]Node, error) {
	if s == nil || s.pool == nil {
		return nil, ErrStore
	}
	rows, err := s.pool.Query(ctx, nodeSelect+` ORDER BY n.updated_at DESC,n.code`)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrStore, err)
	}
	defer rows.Close()
	items := make([]Node, 0)
	for rows.Next() {
		item, scanErr := scanNode(rows)
		if scanErr != nil {
			return nil, fmt.Errorf("%w: %v", ErrStore, scanErr)
		}
		items = append(items, item)
	}
	return items, rows.Err()
}

func normalizeRoutePreviewInput(input RoutePreviewInput) (RoutePreviewInput, error) {
	input.PoolType = strings.TrimSpace(strings.ToLower(input.PoolType))
	input.PoolCode = strings.TrimSpace(input.PoolCode)
	input.RequestKey = strings.TrimSpace(input.RequestKey)
	if input.PoolType != "public" && input.PoolType != "private" {
		return RoutePreviewInput{}, fmt.Errorf("%w: pool type is invalid", ErrInvalidInput)
	}
	if input.PoolCode == "" || len(input.PoolCode) > 64 {
		return RoutePreviewInput{}, fmt.Errorf("%w: pool code is invalid", ErrInvalidInput)
	}
	if input.RequestKey == "" || len(input.RequestKey) > 512 {
		return RoutePreviewInput{}, fmt.Errorf("%w: request key is invalid", ErrInvalidInput)
	}
	capabilities := make([]string, 0, len(input.RequiredCapabilities))
	seen := make(map[string]struct{}, len(input.RequiredCapabilities))
	for _, capability := range input.RequiredCapabilities {
		capability = strings.TrimSpace(capability)
		if capability == "" {
			continue
		}
		if len(capability) > 128 {
			return RoutePreviewInput{}, fmt.Errorf("%w: capability is too long", ErrInvalidInput)
		}
		if _, ok := seen[capability]; ok {
			continue
		}
		seen[capability] = struct{}{}
		capabilities = append(capabilities, capability)
	}
	if len(capabilities) > 32 {
		return RoutePreviewInput{}, fmt.Errorf("%w: too many required capabilities", ErrInvalidInput)
	}
	input.RequiredCapabilities = capabilities
	return input, nil
}

func capabilityEnabled(value any) bool {
	switch typed := value.(type) {
	case nil:
		return false
	case bool:
		return typed
	case string:
		value := strings.TrimSpace(strings.ToLower(typed))
		return value != "" && value != "false" && value != "0" && value != "no"
	case float64:
		return typed != 0
	default:
		return true
	}
}

func nodeHasCapabilities(node Node, required []string) bool {
	if len(required) == 0 {
		return true
	}
	var capabilities map[string]any
	if err := json.Unmarshal(node.Capabilities, &capabilities); err != nil {
		return false
	}
	for _, capability := range required {
		value, ok := capabilities[capability]
		if !ok || !capabilityEnabled(value) {
			return false
		}
	}
	return true
}

func orderRouteCandidates(nodes []Node, requestKey string) []Node {
	ordered := append([]Node(nil), nodes...)
	type candidateHash struct {
		digest [32]byte
		id     string
	}
	hashes := make(map[string]candidateHash, len(ordered))
	for _, node := range ordered {
		hashes[node.ID] = candidateHash{digest: sha256.Sum256([]byte(requestKey + "\x00" + node.ID)), id: node.ID}
	}
	sort.SliceStable(ordered, func(i, j int) bool {
		left, right := hashes[ordered[i].ID], hashes[ordered[j].ID]
		if string(left.digest[:]) == string(right.digest[:]) {
			return left.id < right.id
		}
		return string(left.digest[:]) < string(right.digest[:])
	})
	return ordered
}

func selectRouteCandidates(nodes []Node, requestKey string, required []string) []Node {
	filtered := make([]Node, 0, len(nodes))
	for _, node := range nodes {
		if node.Status != "online" || !node.Enabled || !nodeHasCapabilities(node, required) {
			continue
		}
		filtered = append(filtered, node)
	}
	return orderRouteCandidates(filtered, requestKey)
}

// PreviewRoute returns the online, capability-compatible nodes in stable
// request-key order. It is intentionally read-only and does not invoke an
// endpoint or enqueue a task.
func (s *Store) PreviewRoute(ctx context.Context, input RoutePreviewInput) (RoutePreview, error) {
	if s == nil || s.pool == nil {
		return RoutePreview{}, ErrStore
	}
	normalized, err := normalizeRoutePreviewInput(input)
	if err != nil {
		return RoutePreview{}, err
	}
	rows, err := s.pool.Query(ctx, nodeSelect+` WHERE n.pool_type=$1 AND n.pool_code=$2 AND n.enabled AND n.last_heartbeat_at IS NOT NULL AND n.last_heartbeat_at >= clock_timestamp()-interval '90 seconds' ORDER BY n.code`, normalized.PoolType, normalized.PoolCode)
	if err != nil {
		return RoutePreview{}, fmt.Errorf("%w: %v", ErrStore, err)
	}
	defer rows.Close()
	online := make([]Node, 0)
	for rows.Next() {
		item, scanErr := scanNode(rows)
		if scanErr != nil {
			return RoutePreview{}, fmt.Errorf("%w: %v", ErrStore, scanErr)
		}
		online = append(online, item)
	}
	if err := rows.Err(); err != nil {
		return RoutePreview{}, fmt.Errorf("%w: %v", ErrStore, err)
	}
	candidates := selectRouteCandidates(online, normalized.RequestKey, normalized.RequiredCapabilities)
	preview := RoutePreview{
		PoolType: normalized.PoolType, PoolCode: normalized.PoolCode,
		RequiredCapabilities: normalized.RequiredCapabilities, RequestKey: normalized.RequestKey,
		SelectionPolicy: "stable-hash(request_key,node_id)", Candidates: candidates,
	}
	switch {
	case len(online) == 0:
		preview.Reason = "no_online_node"
	case len(candidates) == 0:
		preview.Reason = "no_capable_node"
	default:
		preview.Reason = "selected_online_node"
		preview.Selected = &preview.Candidates[0]
	}
	return preview, nil
}

func (s *Store) Update(ctx context.Context, input UpdateInput) (Node, error) {
	if s == nil || s.pool == nil {
		return Node{}, ErrStore
	}
	input.ID = strings.TrimSpace(input.ID)
	input.Name = strings.TrimSpace(input.Name)
	input.Description = strings.TrimSpace(input.Description)
	input.Endpoint = strings.TrimSpace(input.Endpoint)
	input.NodeType = strings.TrimSpace(input.NodeType)
	input.PoolType = strings.TrimSpace(input.PoolType)
	input.PoolCode = strings.TrimSpace(input.PoolCode)
	if err := validateCommon("valid-node", input.Name, input.Description, input.Endpoint, input.NodeType, input.PoolType, input.PoolCode); err != nil {
		return Node{}, err
	}
	tag, err := s.pool.Exec(ctx, `UPDATE gateway.agent_nodes SET name=$2,description=$3,endpoint=$4,node_type=$5,pool_type=$6,pool_code=$7,enabled=$8,updated_at=clock_timestamp() WHERE id=$1`, input.ID, input.Name, input.Description, input.Endpoint, input.NodeType, input.PoolType, input.PoolCode, input.Enabled)
	if err != nil {
		return Node{}, fmt.Errorf("%w: %v", ErrStore, err)
	}
	if tag.RowsAffected() == 0 {
		return Node{}, ErrNotFound
	}
	return s.Get(ctx, input.ID)
}

func (s *Store) Delete(ctx context.Context, id string) error {
	if s == nil || s.pool == nil {
		return ErrStore
	}
	tag, err := s.pool.Exec(ctx, `DELETE FROM gateway.agent_nodes WHERE id=$1`, strings.TrimSpace(id))
	if err != nil {
		return fmt.Errorf("%w: %v", ErrStore, err)
	}
	if tag.RowsAffected() == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *Store) RotateToken(ctx context.Context, id string) (Node, string, error) {
	if s == nil || s.pool == nil {
		return Node{}, "", ErrStore
	}
	secret, prefix, digest, err := generateToken()
	if err != nil {
		return Node{}, "", err
	}
	tag, err := s.pool.Exec(ctx, `UPDATE gateway.agent_nodes SET token_prefix=$2,token_hash=$3,updated_at=clock_timestamp() WHERE id=$1`, strings.TrimSpace(id), prefix, digest)
	if err != nil {
		return Node{}, "", fmt.Errorf("%w: %v", ErrStore, err)
	}
	if tag.RowsAffected() == 0 {
		return Node{}, "", ErrNotFound
	}
	item, err := s.Get(ctx, id)
	return item, secret, err
}

func (s *Store) Heartbeat(ctx context.Context, code, token string, remoteIP net.IP, input HeartbeatInput) (Node, error) {
	if s == nil || s.pool == nil {
		return Node{}, ErrStore
	}
	code = strings.ToLower(strings.TrimSpace(code))
	token = strings.TrimSpace(token)
	if !nodeCodePattern.MatchString(code) || token == "" || len(token) > 512 || len(input.Version) > 128 || len(input.Error) > 4000 {
		return Node{}, ErrInvalidInput
	}
	capabilities, err := objectJSON(input.Capabilities)
	if err != nil {
		return Node{}, err
	}
	metadata, err := objectJSON(input.Metadata)
	if err != nil {
		return Node{}, err
	}
	digest := sha256.Sum256([]byte(token))
	ip := ""
	if remoteIP != nil {
		ip = remoteIP.String()
	}
	// 先按 code 取出令牌哈希,在 Go 侧做恒定时间比较:未知 code 与错误
	// 令牌返回同一个错误,避免通过 404/401 差异枚举有效节点;数据库端
	// bytea 比较可能提前短路,不做恒定时间保证。
	var storedHash []byte
	err = s.pool.QueryRow(ctx, `SELECT token_hash FROM gateway.agent_nodes WHERE code=$1 AND enabled`, code).Scan(&storedHash)
	if errors.Is(err, pgx.ErrNoRows) {
		return Node{}, ErrInvalidToken
	}
	if err != nil {
		return Node{}, fmt.Errorf("%w: %v", ErrStore, err)
	}
	if subtle.ConstantTimeCompare(digest[:], storedHash) != 1 {
		return Node{}, ErrInvalidToken
	}
	var id string
	err = s.pool.QueryRow(ctx, `UPDATE gateway.agent_nodes SET version=$3,capabilities=coalesce($4::jsonb,capabilities),metadata=coalesce($5::jsonb,metadata),last_error=$6,last_heartbeat_at=clock_timestamp(),last_heartbeat_ip=nullif($7,'')::inet,updated_at=clock_timestamp() WHERE code=$1 AND token_hash=$2 AND enabled RETURNING id::text`, code, digest[:], input.Version, capabilities, metadata, strings.TrimSpace(input.Error), ip).Scan(&id)
	if errors.Is(err, pgx.ErrNoRows) {
		return Node{}, ErrInvalidToken
	}
	if err != nil {
		return Node{}, fmt.Errorf("%w: %v", ErrStore, err)
	}
	return s.Get(ctx, id)
}