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}) }