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

package agentnode

import (
	"encoding/json"
	"errors"
	"net"
	"net/http"
	"strings"

	"aigateway.local/core/internal/identity"
	"aigateway.local/core/internal/platform/apiresponse"
)

type HTTPHandler struct {
	store    *Store
	identity *identity.Service
	mux      *http.ServeMux
}

type nodeRequest struct {
	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"`
}

type routePreviewRequest struct {
	PoolType             string   `json:"pool_type"`
	PoolCode             string   `json:"pool_code"`
	RequiredCapabilities []string `json:"required_capabilities"`
	RequestKey           string   `json:"request_key"`
}

func NewHTTPHandler(store *Store, identityService *identity.Service) *HTTPHandler {
	h := &HTTPHandler{store: store, identity: identityService, mux: http.NewServeMux()}
	h.mux.HandleFunc("GET /api/v1/admin/agent-nodes", h.list)
	h.mux.HandleFunc("POST /api/v1/admin/agent-nodes", h.create)
	h.mux.HandleFunc("POST /api/v1/admin/agent-nodes/route-preview", h.routePreview)
	h.mux.HandleFunc("PUT /api/v1/admin/agent-nodes/{id}", h.update)
	h.mux.HandleFunc("DELETE /api/v1/admin/agent-nodes/{id}", h.delete)
	h.mux.HandleFunc("POST /api/v1/admin/agent-nodes/{id}/rotate-token", h.rotateToken)
	h.mux.HandleFunc("POST /api/v1/agent/nodes/{code}/heartbeat", h.heartbeat)
	return h
}

func (h *HTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.mux.ServeHTTP(w, r) }

func (h *HTTPHandler) require(w http.ResponseWriter, r *http.Request, permission string) (identity.Account, bool) {
	account, err := h.identity.Authenticate(r.Context(), identity.KindAdmin, r.Header.Get("Authorization"))
	if err != nil {
		apiresponse.Error(w, http.StatusUnauthorized, "登录状态无效或已过期")
		return identity.Account{}, false
	}
	if !identity.HasPermission(account, permission) {
		apiresponse.Error(w, http.StatusForbidden, "缺少智能体节点操作权限")
		return identity.Account{}, false
	}
	return account, true
}

func decodeJSON(w http.ResponseWriter, r *http.Request, target any) bool {
	decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(target); err != nil {
		apiresponse.Error(w, http.StatusBadRequest, "请求格式无效")
		return false
	}
	return true
}

func (h *HTTPHandler) list(w http.ResponseWriter, r *http.Request) {
	if _, ok := h.require(w, r, identity.PermissionAgentNodeRead); !ok {
		return
	}
	items, err := h.store.List(r.Context())
	if err != nil {
		writeError(w, err)
		return
	}
	apiresponse.OK(w, items)
}

func (h *HTTPHandler) routePreview(w http.ResponseWriter, r *http.Request) {
	if _, ok := h.require(w, r, identity.PermissionAgentNodeRead); !ok {
		return
	}
	var input routePreviewRequest
	if !decodeJSON(w, r, &input) {
		return
	}
	preview, err := h.store.PreviewRoute(r.Context(), RoutePreviewInput{
		PoolType: input.PoolType, PoolCode: input.PoolCode,
		RequiredCapabilities: input.RequiredCapabilities, RequestKey: input.RequestKey,
	})
	if err != nil {
		writeError(w, err)
		return
	}
	apiresponse.OK(w, preview)
}

func (h *HTTPHandler) create(w http.ResponseWriter, r *http.Request) {
	actor, ok := h.require(w, r, identity.PermissionAgentNodeManage)
	if !ok {
		return
	}
	var input nodeRequest
	if !decodeJSON(w, r, &input) {
		return
	}
	enabled := true
	if input.Enabled != nil {
		enabled = *input.Enabled
	}
	node, token, err := h.store.Create(r.Context(), CreateInput{Code: input.Code, Name: input.Name, Description: input.Description, Endpoint: input.Endpoint, NodeType: input.NodeType, PoolType: input.PoolType, PoolCode: input.PoolCode, Enabled: enabled}, actor.ID)
	if err != nil {
		writeError(w, err)
		return
	}
	apiresponse.OK(w, map[string]any{"node": node, "token": token, "warning": "令牌只显示一次，请立即安全保存并配置到节点"})
}

func (h *HTTPHandler) update(w http.ResponseWriter, r *http.Request) {
	if _, ok := h.require(w, r, identity.PermissionAgentNodeManage); !ok {
		return
	}
	var input nodeRequest
	if !decodeJSON(w, r, &input) {
		return
	}
	if input.Enabled == nil {
		apiresponse.Error(w, http.StatusBadRequest, "enabled 字段不能为空")
		return
	}
	node, err := h.store.Update(r.Context(), UpdateInput{ID: r.PathValue("id"), Name: input.Name, Description: input.Description, Endpoint: input.Endpoint, NodeType: input.NodeType, PoolType: input.PoolType, PoolCode: input.PoolCode, Enabled: *input.Enabled})
	if err != nil {
		writeError(w, err)
		return
	}
	apiresponse.OK(w, node)
}

func (h *HTTPHandler) delete(w http.ResponseWriter, r *http.Request) {
	if _, ok := h.require(w, r, identity.PermissionAgentNodeManage); !ok {
		return
	}
	if err := h.store.Delete(r.Context(), r.PathValue("id")); err != nil {
		writeError(w, err)
		return
	}
	apiresponse.OK(w, map[string]bool{"deleted": true})
}

func (h *HTTPHandler) rotateToken(w http.ResponseWriter, r *http.Request) {
	if _, ok := h.require(w, r, identity.PermissionAgentNodeManage); !ok {
		return
	}
	node, token, err := h.store.RotateToken(r.Context(), r.PathValue("id"))
	if err != nil {
		writeError(w, err)
		return
	}
	apiresponse.OK(w, map[string]any{"node": node, "token": token, "warning": "旧令牌已立即失效，新令牌只显示一次"})
}

func (h *HTTPHandler) heartbeat(w http.ResponseWriter, r *http.Request) {
	var input HeartbeatInput
	if !decodeJSON(w, r, &input) {
		return
	}
	remoteIP := parseRemoteIP(r.RemoteAddr)
	node, err := h.store.Heartbeat(r.Context(), r.PathValue("code"), r.Header.Get("X-Agent-Token"), remoteIP, input)
	if err != nil {
		writeHeartbeatError(w, err)
		return
	}
	apiresponse.OK(w, map[string]any{"accepted": true, "node": node})
}

func parseRemoteIP(remoteAddr string) net.IP {
	host, _, err := net.SplitHostPort(strings.TrimSpace(remoteAddr))
	if err != nil {
		host = strings.TrimSpace(remoteAddr)
	}
	return net.ParseIP(host)
}

func writeError(w http.ResponseWriter, err error) {
	switch {
	case errors.Is(err, ErrNotFound):
		apiresponse.Error(w, http.StatusNotFound, "智能体节点不存在")
	case errors.Is(err, ErrConflict):
		apiresponse.Error(w, http.StatusConflict, "节点编码已存在")
	case errors.Is(err, ErrInvalidInput):
		apiresponse.Error(w, http.StatusBadRequest, err.Error())
	case errors.Is(err, ErrStore):
		apiresponse.Error(w, http.StatusServiceUnavailable, "智能体节点服务暂不可用")
	default:
		apiresponse.Error(w, http.StatusInternalServerError, "智能体节点处理失败")
	}
}

func writeHeartbeatError(w http.ResponseWriter, err error) {
	switch {
	case errors.Is(err, ErrInvalidToken):
		apiresponse.Error(w, http.StatusUnauthorized, "节点令牌无效或节点已停用")
	case errors.Is(err, ErrInvalidInput):
		apiresponse.Error(w, http.StatusBadRequest, err.Error())
	case errors.Is(err, ErrStore):
		apiresponse.Error(w, http.StatusServiceUnavailable, "智能体节点服务暂不可用")
	default:
		apiresponse.Error(w, http.StatusInternalServerError, "节点心跳处理失败")
	}
}