ai-gateway-go/internal/workbench/agentpolicy.go

package workbench

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

	"aigateway.local/core/internal/identity"
	"aigateway.local/core/internal/platform/apiresponse"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// AgentPolicy 是个人智能体安全策略。
type AgentPolicy struct {
	AutoApproveTools    bool `json:"auto_approve_tools"`
	RateLimitMultiplier int  `json:"rate_limit_multiplier"`
}

// AgentPolicyService 管理个人智能体安全策略。
type AgentPolicyService struct {
	pool *pgxpool.Pool
}

func NewAgentPolicyService(pool *pgxpool.Pool) *AgentPolicyService {
	return &AgentPolicyService{pool: pool}
}

// Get 返回用户的策略(未配置时返回默认值)。
func (s *AgentPolicyService) Get(ctx context.Context, portalUserID string) (AgentPolicy, error) {
	var policy AgentPolicy
	if s == nil || s.pool == nil {
		return policy, errors.New("安全策略服务不可用")
	}
	err := s.pool.QueryRow(ctx, `SELECT auto_approve_tools,rate_limit_multiplier FROM gateway.portal_agent_policies WHERE portal_user_id=$1`, portalUserID).Scan(&policy.AutoApproveTools, &policy.RateLimitMultiplier)
	if errors.Is(err, pgx.ErrNoRows) {
		policy.RateLimitMultiplier = 1
		return policy, nil
	}
	if err != nil {
		return policy, err
	}
	return policy, nil
}

// Set 更新用户的策略。
func (s *AgentPolicyService) Set(ctx context.Context, portalUserID string, policy AgentPolicy) error {
	if s == nil || s.pool == nil {
		return errors.New("安全策略服务不可用")
	}
	if policy.RateLimitMultiplier < 1 || policy.RateLimitMultiplier > 100 {
		return errors.New("限流倍数必须在 1-100 之间")
	}
	_, err := s.pool.Exec(ctx, `INSERT INTO gateway.portal_agent_policies(portal_user_id,auto_approve_tools,rate_limit_multiplier) VALUES($1,$2,$3)
		ON CONFLICT(portal_user_id) DO UPDATE SET auto_approve_tools=$2,rate_limit_multiplier=$3,updated_at=clock_timestamp()`,
		portalUserID, policy.AutoApproveTools, policy.RateLimitMultiplier)
	return err
}

// AgentPolicyHTTPHandler 门户个人智能体安全策略。
type AgentPolicyHTTPHandler struct {
	service  *AgentPolicyService
	identity *identity.Service
	mux      *http.ServeMux
}

func NewAgentPolicyHTTPHandler(service *AgentPolicyService, identityService *identity.Service) *AgentPolicyHTTPHandler {
	h := &AgentPolicyHTTPHandler{service: service, identity: identityService, mux: http.NewServeMux()}
	h.mux.HandleFunc("GET /api/v1/portal/agent-policy", h.get)
	h.mux.HandleFunc("PUT /api/v1/portal/agent-policy", h.put)
	return h
}

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

func (h *AgentPolicyHTTPHandler) account(w http.ResponseWriter, r *http.Request) (identity.Account, bool) {
	account, err := h.identity.Authenticate(r.Context(), identity.KindPortal, r.Header.Get("Authorization"))
	if err != nil {
		apiresponse.Error(w, http.StatusUnauthorized, "登录状态无效或已过期")
		return identity.Account{}, false
	}
	return account, true
}

func (h *AgentPolicyHTTPHandler) get(w http.ResponseWriter, r *http.Request) {
	a, ok := h.account(w, r)
	if !ok {
		return
	}
	policy, err := h.service.Get(r.Context(), a.ID)
	if err != nil {
		apiresponse.Error(w, http.StatusServiceUnavailable, "安全策略查询失败")
		return
	}
	apiresponse.OK(w, policy)
}

func (h *AgentPolicyHTTPHandler) put(w http.ResponseWriter, r *http.Request) {
	a, ok := h.account(w, r)
	if !ok {
		return
	}
	var input struct {
		AutoApproveTools    bool `json:"auto_approve_tools"`
		RateLimitMultiplier int  `json:"rate_limit_multiplier"`
	}
	decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20))
	decoder.DisallowUnknownFields()
	if decoder.Decode(&input) != nil {
		apiresponse.Error(w, http.StatusBadRequest, "请求格式无效")
		return
	}
	if err := h.service.Set(r.Context(), a.ID, AgentPolicy{AutoApproveTools: input.AutoApproveTools, RateLimitMultiplier: input.RateLimitMultiplier}); err != nil {
		apiresponse.Error(w, http.StatusBadRequest, strings.TrimSpace(err.Error()))
		return
	}
	apiresponse.OK(w, map[string]bool{"saved": true})
}