0.10.1: 安全与业务逻辑加固、新品牌与部署加固
三轮审查修复(60+ 项),相对远端 main(b536672)的关键变更:
- 安全: 数据面 SSRF 拨号防护(防 DNS rebinding)/上游凭据剥离/登录防枚举
与锁定态统一/可信代理(X-Forwarded-For)限流加固/会话版本失效机制/
撤销即时传播/弱密钥拒绝启动/脱敏字节级重写(保签名契约)
- 业务逻辑: 裸 body 上传 panic/bootstrap 审计管线卡死/定价通配符优先级/
全局工具可见性/调度器停机补跑/TOTP 挑战令牌消费顺序/熔断探针语义/
>4MB 响应 token 计量/管理员重置密码作废会话 等
- 前端: 新 logo(语枢 AI 网关主题)/Provider 凭据异常警示/删除入口/
后端错误消息透传/localStorage 敏感数据收敛
- 部署: CREDENTIAL_MASTER_KEY 持久化与弱值拒绝/Provider DELETE 接口/
nginx 安全头/worker 内存限制
- 新增迁移 000029(key_hash 索引)/000030(usage_daily 币种维度)
This commit is contained in:
@@ -0,0 +1,218 @@
|
||||
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, "节点心跳处理失败")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
package agentnode
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func routeTestNode(id, status string, enabled bool, capabilities map[string]any) Node {
|
||||
raw, _ := json.Marshal(capabilities)
|
||||
return Node{ID: id, Code: id, Status: status, Enabled: enabled, Capabilities: raw, LastHeartbeatAt: func() *time.Time { now := time.Now(); return &now }()}
|
||||
}
|
||||
|
||||
func TestSelectRouteCandidatesFiltersCapabilitiesAndStatus(t *testing.T) {
|
||||
nodes := []Node{
|
||||
routeTestNode("online-capable", "online", true, map[string]any{"tool_exec": true, "region": "cn"}),
|
||||
routeTestNode("online-disabled-capability", "online", true, map[string]any{"tool_exec": false}),
|
||||
routeTestNode("online-missing-capability", "online", true, map[string]any{"region": "cn"}),
|
||||
routeTestNode("offline-capable", "offline", true, map[string]any{"tool_exec": true}),
|
||||
routeTestNode("disabled-capable", "online", false, map[string]any{"tool_exec": true}),
|
||||
}
|
||||
selected := selectRouteCandidates(nodes, "request-1", []string{"tool_exec"})
|
||||
if len(selected) != 1 || selected[0].ID != "online-capable" {
|
||||
t.Fatalf("unexpected candidates: %#v", selected)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectRouteCandidatesIsDeterministic(t *testing.T) {
|
||||
nodes := []Node{
|
||||
routeTestNode("node-a", "online", true, nil),
|
||||
routeTestNode("node-b", "online", true, nil),
|
||||
routeTestNode("node-c", "online", true, nil),
|
||||
}
|
||||
first := selectRouteCandidates(nodes, "request-42", nil)
|
||||
second := selectRouteCandidates([]Node{nodes[2], nodes[0], nodes[1]}, "request-42", nil)
|
||||
if len(first) != len(second) || len(first) == 0 {
|
||||
t.Fatalf("candidate lengths differ: %d %d", len(first), len(second))
|
||||
}
|
||||
for index := range first {
|
||||
if first[index].ID != second[index].ID {
|
||||
t.Fatalf("selection order is not stable: first=%v second=%v", first, second)
|
||||
}
|
||||
}
|
||||
if first[0].ID == first[1].ID {
|
||||
t.Fatal("candidate order contains duplicates")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeRoutePreviewInput(t *testing.T) {
|
||||
input, err := normalizeRoutePreviewInput(RoutePreviewInput{
|
||||
PoolType: " PUBLIC ", PoolCode: "shared", RequestKey: " request-1 ",
|
||||
RequiredCapabilities: []string{"tool_exec", " tool_exec ", ""},
|
||||
})
|
||||
if err != nil || input.PoolType != "public" || input.RequestKey != "request-1" || len(input.RequiredCapabilities) != 1 {
|
||||
t.Fatalf("normalized input=%+v err=%v", input, err)
|
||||
}
|
||||
if _, err := normalizeRoutePreviewInput(RoutePreviewInput{PoolType: "public", PoolCode: "shared"}); err == nil {
|
||||
t.Fatal("empty request key must be rejected")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,502 @@
|
||||
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)
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package agentnode
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/platform/config"
|
||||
"aigateway.local/core/internal/platform/database"
|
||||
)
|
||||
|
||||
func TestAgentNodePostgreSQLLifecycle(t *testing.T) {
|
||||
databaseURL := os.Getenv("AGENT_NODE_TEST_DATABASE_URL")
|
||||
if databaseURL == "" {
|
||||
t.Skip("AGENT_NODE_TEST_DATABASE_URL is not set")
|
||||
}
|
||||
ctx := context.Background()
|
||||
pool, err := database.Open(ctx, config.Database{URL: databaseURL, MaxConns: 4, MinConns: 0})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer pool.Close()
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM gateway.agent_nodes WHERE code='node-integration'`)
|
||||
defer pool.Exec(ctx, `DELETE FROM gateway.agent_nodes WHERE code='node-integration'`)
|
||||
store := NewStore(pool)
|
||||
node, token, err := store.Create(ctx, CreateInput{Code: "node-integration", Name: "Integration Node", PoolType: "private", PoolCode: "test", Enabled: true}, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if token == "" || node.Status != "pending" || node.TokenPrefix == "" {
|
||||
t.Fatalf("created node=%+v token=%q", node, token)
|
||||
}
|
||||
updated, err := store.Update(ctx, UpdateInput{ID: node.ID, Name: "Integration Node v2", Description: "test", Endpoint: "https://node.invalid", NodeType: "executor", PoolType: "public", PoolCode: "shared", Enabled: true})
|
||||
if err != nil || updated.Name != "Integration Node v2" || updated.PoolType != "public" {
|
||||
t.Fatalf("updated node=%+v err=%v", updated, err)
|
||||
}
|
||||
online, err := store.Heartbeat(ctx, node.Code, token, net.ParseIP("192.0.2.10"), HeartbeatInput{Version: "0.10.0-node", Capabilities: map[string]any{"tool_exec": true}, Metadata: map[string]any{"region": "test"}})
|
||||
if err != nil || online.Status != "online" || online.Version != "0.10.0-node" || online.LastHeartbeatIP != "192.0.2.10" {
|
||||
t.Fatalf("heartbeat node=%+v err=%v", online, err)
|
||||
}
|
||||
preview, err := store.PreviewRoute(ctx, RoutePreviewInput{PoolType: "public", PoolCode: "shared", RequestKey: "integration-request", RequiredCapabilities: []string{"tool_exec"}})
|
||||
if err != nil || preview.Reason != "selected_online_node" || preview.Selected == nil || preview.Selected.ID != online.ID || len(preview.Candidates) != 1 {
|
||||
t.Fatalf("route preview=%+v err=%v", preview, err)
|
||||
}
|
||||
rotated, newToken, err := store.RotateToken(ctx, node.ID)
|
||||
if err != nil || newToken == token || rotated.TokenPrefix == node.TokenPrefix {
|
||||
t.Fatalf("rotated node=%+v token=%q err=%v", rotated, newToken, err)
|
||||
}
|
||||
if _, err = store.Heartbeat(ctx, node.Code, token, nil, HeartbeatInput{}); err != ErrInvalidToken {
|
||||
t.Fatalf("old token err=%v", err)
|
||||
}
|
||||
if _, err = store.Heartbeat(ctx, node.Code, newToken, nil, HeartbeatInput{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
items, err := store.List(ctx)
|
||||
if err != nil || len(items) == 0 || items[0].UpdatedAt.Before(time.Now().UTC().Add(-time.Minute)) {
|
||||
t.Fatalf("items=%+v err=%v", items, err)
|
||||
}
|
||||
}
|
||||
@@ -93,9 +93,10 @@ func (h *AdminHTTPHandler) updateLimits(writer http.ResponseWriter, request *htt
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
// 缓存失效失败不视为操作失败:数据库已生效(权威源),缓存最迟在 TTL 后
|
||||
// 自动过期;若在此报错,运维会误以为限流更新失败而重试。
|
||||
if err := h.authenticator.Invalidate(request.Context(), hash); err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
h.authenticator.logger.Warn("api key limits cache invalidation failed; key stays cached until TTL", "error", err)
|
||||
}
|
||||
response := publicRecord(record)
|
||||
if h.usage != nil {
|
||||
@@ -161,9 +162,9 @@ func (h *AdminHTTPHandler) revoke(writer http.ResponseWriter, request *http.Requ
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
// 同上:撤销已在数据库生效,缓存失效失败仅记录,不误报为撤销失败。
|
||||
if err := h.authenticator.Invalidate(request.Context(), hash); err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
h.authenticator.logger.Warn("api key revocation cache invalidation failed; key stays cached until TTL", "error", err)
|
||||
}
|
||||
apiresponse.OK(writer, map[string]bool{"revoked": true})
|
||||
}
|
||||
|
||||
@@ -24,7 +24,9 @@ type Authenticator struct {
|
||||
}
|
||||
|
||||
func NewAuthenticator(repository *Repository, client *redis.Client, bootstrap string) *Authenticator {
|
||||
return &Authenticator{repository: repository, redis: client, bootstrap: bootstrap, cacheTTL: 30 * time.Second}
|
||||
// cacheTTL 是撤销/限流变更的最坏传播窗口:数据库提交后 Invalidate 会
|
||||
// 立即删除缓存键,TTL 只是 Redis 故障时的兜底,因此保持短小。
|
||||
return &Authenticator{repository: repository, redis: client, bootstrap: bootstrap, cacheTTL: 10 * time.Second}
|
||||
}
|
||||
|
||||
// SetLogger wires an optional logger used for best-effort cache diagnostics.
|
||||
|
||||
@@ -21,17 +21,19 @@ type MaintenanceResult struct {
|
||||
DroppedPartitions []string `json:"dropped_partitions"`
|
||||
DeletedAuditRows int64 `json:"deleted_audit_rows"`
|
||||
DeletedUsageRows int64 `json:"deleted_usage_rows"`
|
||||
DeletedTraceRows int64 `json:"deleted_trace_rows"`
|
||||
}
|
||||
|
||||
type Maintenance struct {
|
||||
pool *pgxpool.Pool
|
||||
auditRetention time.Duration
|
||||
usageRetention time.Duration
|
||||
traceRetention time.Duration
|
||||
monthsAhead int
|
||||
}
|
||||
|
||||
func NewMaintenance(pool *pgxpool.Pool, auditRetention, usageRetention time.Duration, monthsAhead int) *Maintenance {
|
||||
return &Maintenance{pool: pool, auditRetention: auditRetention, usageRetention: usageRetention, monthsAhead: monthsAhead}
|
||||
func NewMaintenance(pool *pgxpool.Pool, auditRetention, usageRetention, traceRetention time.Duration, monthsAhead int) *Maintenance {
|
||||
return &Maintenance{pool: pool, auditRetention: auditRetention, usageRetention: usageRetention, traceRetention: traceRetention, monthsAhead: monthsAhead}
|
||||
}
|
||||
|
||||
func (m *Maintenance) Run(ctx context.Context, now time.Time) (MaintenanceResult, error) {
|
||||
@@ -105,6 +107,15 @@ func (m *Maintenance) Run(ctx context.Context, now time.Time) (MaintenanceResult
|
||||
return result, fmt.Errorf("apply usage retention: %w", err)
|
||||
}
|
||||
result.DeletedUsageRows = deleted.RowsAffected()
|
||||
// M9 Trace 保留:agent_trace_spans 通过 ON DELETE CASCADE 一并清理,
|
||||
// 防止 trace 表无界增长(每条应用/数字员工请求都会写 trace)。
|
||||
if m.traceRetention > 0 {
|
||||
deleted, err = tx.Exec(ctx, `DELETE FROM gateway.agent_traces WHERE started_at < $1`, now.Add(-m.traceRetention))
|
||||
if err != nil {
|
||||
return result, fmt.Errorf("apply trace retention: %w", err)
|
||||
}
|
||||
result.DeletedTraceRows = deleted.RowsAffected()
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return result, fmt.Errorf("commit audit maintenance: %w", err)
|
||||
}
|
||||
|
||||
@@ -29,14 +29,14 @@ func TestMaintenancePartitionsRetentionAndIdempotency(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
now := time.Date(2026, 8, 11, 0, 0, 0, 0, time.UTC)
|
||||
first, err := NewMaintenance(pool, 300*24*time.Hour, 730*24*time.Hour, 1).Run(ctx, now)
|
||||
first, err := NewMaintenance(pool, 300*24*time.Hour, 730*24*time.Hour, 30*24*time.Hour, 1).Run(ctx, now)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !contains(first.CreatedPartitions, "audit_events_202601") || !contains(first.CreatedPartitions, "audit_events_202608") {
|
||||
t.Fatalf("expected old and current partitions, got %#v", first.CreatedPartitions)
|
||||
}
|
||||
second, err := NewMaintenance(pool, 90*24*time.Hour, 730*24*time.Hour, 1).Run(ctx, now)
|
||||
second, err := NewMaintenance(pool, 90*24*time.Hour, 730*24*time.Hour, 30*24*time.Hour, 1).Run(ctx, now)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -54,7 +54,7 @@ func TestMaintenancePartitionsRetentionAndIdempotency(t *testing.T) {
|
||||
if err := pool.QueryRow(ctx, `SELECT count(*) FROM gateway.audit_events WHERE id=$1`, oldID).Scan(&oldCount); err != nil || oldCount != 0 {
|
||||
t.Fatalf("expired row still exists: count=%d err=%v", oldCount, err)
|
||||
}
|
||||
third, err := NewMaintenance(pool, 90*24*time.Hour, 730*24*time.Hour, 1).Run(ctx, now)
|
||||
third, err := NewMaintenance(pool, 90*24*time.Hour, 730*24*time.Hour, 30*24*time.Hour, 1).Run(ctx, now)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
@@ -97,6 +97,7 @@ type DailyUsageView struct {
|
||||
APIKeyName string `json:"api_key_name"`
|
||||
ProviderCode string `json:"provider_code"`
|
||||
Model string `json:"model"`
|
||||
Currency string `json:"currency"`
|
||||
Requests int64 `json:"requests"`
|
||||
FailedRequests int64 `json:"failed_requests"`
|
||||
PromptTokens int64 `json:"prompt_tokens"`
|
||||
@@ -130,7 +131,7 @@ func (s *QueryService) ListDailyUsage(ctx context.Context, filter UsageFilter) (
|
||||
if filter.Model != "" {
|
||||
add("u.model", filter.Model)
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `SELECT u.usage_date,u.api_key_id::text,k.name,u.provider_code,u.model,
|
||||
rows, err := s.pool.Query(ctx, `SELECT u.usage_date,u.api_key_id::text,k.name,u.provider_code,u.model,u.currency,
|
||||
u.requests,u.failed_requests,u.prompt_tokens,u.completion_tokens,u.cost_microunits
|
||||
FROM gateway.usage_daily u JOIN gateway.api_keys k ON k.id=u.api_key_id WHERE `+
|
||||
strings.Join(where, " AND ")+` ORDER BY u.usage_date DESC,k.name,u.provider_code,u.model`, args...)
|
||||
@@ -141,7 +142,7 @@ func (s *QueryService) ListDailyUsage(ctx context.Context, filter UsageFilter) (
|
||||
items := make([]DailyUsageView, 0)
|
||||
for rows.Next() {
|
||||
var item DailyUsageView
|
||||
if err := rows.Scan(&item.Date, &item.APIKeyID, &item.APIKeyName, &item.ProviderCode, &item.Model,
|
||||
if err := rows.Scan(&item.Date, &item.APIKeyID, &item.APIKeyName, &item.ProviderCode, &item.Model, &item.Currency,
|
||||
&item.Requests, &item.FailedRequests, &item.PromptTokens, &item.CompletionTokens, &item.CostMicrounits); err != nil {
|
||||
return nil, fmt.Errorf("scan daily usage: %w", err)
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"time"
|
||||
|
||||
platformid "aigateway.local/core/internal/platform/id"
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgtype"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
@@ -129,9 +130,25 @@ func (r *Recorder) Run(ctx context.Context) {
|
||||
case event := <-r.queue:
|
||||
pending = append(pending, event)
|
||||
default:
|
||||
flushCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
// 关闭前的最后一次落盘:数据库短暂不可用时重试有限次数,
|
||||
// 而不是只试一次就把整批审计事件静默丢弃。
|
||||
flushCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
if len(pending) > 0 {
|
||||
_ = r.flush(flushCtx, pending)
|
||||
lastErr := r.flush(flushCtx, pending)
|
||||
for attempt := 0; lastErr != nil && attempt < 3; attempt++ {
|
||||
select {
|
||||
case <-time.After(2 * time.Second):
|
||||
case <-flushCtx.Done():
|
||||
lastErr = flushCtx.Err()
|
||||
}
|
||||
if flushCtx.Err() != nil {
|
||||
break
|
||||
}
|
||||
lastErr = r.flush(flushCtx, pending)
|
||||
}
|
||||
if lastErr != nil && r.logger != nil {
|
||||
r.logger.Error("audit drain failed; events lost", "events", len(pending), "error", lastErr)
|
||||
}
|
||||
}
|
||||
cancel()
|
||||
return
|
||||
@@ -142,7 +159,7 @@ func (r *Recorder) Run(ctx context.Context) {
|
||||
}
|
||||
|
||||
type dailyKey struct {
|
||||
date, apiKeyID, provider, model string
|
||||
date, apiKeyID, provider, model, currency string
|
||||
}
|
||||
|
||||
type dailyValue struct {
|
||||
@@ -192,7 +209,18 @@ func (r *Recorder) flush(ctx context.Context, events []Event) error {
|
||||
nil, nil, labels, event.RecordedAt,
|
||||
})
|
||||
if event.APIKeyID != nil && *event.APIKeyID != "" {
|
||||
key := dailyKey{date: event.RecordedAt.UTC().Format("2006-01-02"), apiKeyID: *event.APIKeyID, provider: event.ProviderCode, model: event.Model}
|
||||
// 仅聚合合法 UUID 的 API Key:usage_daily.api_key_id 是
|
||||
// REFERENCES api_keys 的 uuid 列,bootstrap 等非 UUID 身份写入
|
||||
// 会让整批事务失败、审计管线永久卡死。审计事件本身仍落 audit_events。
|
||||
if _, uuidErr := uuid.Parse(strings.TrimSpace(*event.APIKeyID)); uuidErr != nil {
|
||||
continue
|
||||
}
|
||||
// 成本按币种独立聚合:不同货币的价格不能相加成单一数字。
|
||||
currency := strings.ToUpper(strings.TrimSpace(event.Currency))
|
||||
if currency == "" {
|
||||
currency = "USD"
|
||||
}
|
||||
key := dailyKey{date: event.RecordedAt.UTC().Format("2006-01-02"), apiKeyID: *event.APIKeyID, provider: event.ProviderCode, model: event.Model, currency: currency}
|
||||
value := daily[key]
|
||||
value.requests++
|
||||
if event.StatusCode >= 400 {
|
||||
@@ -214,15 +242,15 @@ func (r *Recorder) flush(ctx context.Context, events []Event) error {
|
||||
}
|
||||
batch := &pgx.Batch{}
|
||||
for key, value := range daily {
|
||||
batch.Queue(`INSERT INTO gateway.usage_daily(usage_date,api_key_id,provider_code,model,requests,failed_requests,prompt_tokens,completion_tokens,cost_microunits)
|
||||
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9)
|
||||
ON CONFLICT(usage_date,api_key_id,provider_code,model) DO UPDATE SET
|
||||
batch.Queue(`INSERT INTO gateway.usage_daily(usage_date,api_key_id,provider_code,model,currency,requests,failed_requests,prompt_tokens,completion_tokens,cost_microunits)
|
||||
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10)
|
||||
ON CONFLICT(usage_date,api_key_id,provider_code,model,currency) DO UPDATE SET
|
||||
requests=gateway.usage_daily.requests+EXCLUDED.requests,
|
||||
failed_requests=gateway.usage_daily.failed_requests+EXCLUDED.failed_requests,
|
||||
prompt_tokens=gateway.usage_daily.prompt_tokens+EXCLUDED.prompt_tokens,
|
||||
completion_tokens=gateway.usage_daily.completion_tokens+EXCLUDED.completion_tokens,
|
||||
cost_microunits=gateway.usage_daily.cost_microunits+EXCLUDED.cost_microunits,
|
||||
updated_at=clock_timestamp()`, key.date, key.apiKeyID, key.provider, key.model, value.requests, value.failed, value.prompt, value.completion, value.cost)
|
||||
updated_at=clock_timestamp()`, key.date, key.apiKeyID, key.provider, key.model, key.currency, value.requests, value.failed, value.prompt, value.completion, value.cost)
|
||||
}
|
||||
for _, alert := range alerts {
|
||||
batch.Queue(`INSERT INTO gateway.outbox_events(event_id,event_type,event_version,tenant_id,aggregate_type,aggregate_id,payload)
|
||||
@@ -257,6 +285,11 @@ func uuidPointer(value *string) any {
|
||||
if value == nil || strings.TrimSpace(*value) == "" {
|
||||
return nil
|
||||
}
|
||||
// 非 UUID 身份(bootstrap key 等)返回 nil:audit_events.api_key_id 允许
|
||||
// NULL,写零值/非法值只会掩盖问题。
|
||||
if _, err := uuid.Parse(strings.TrimSpace(*value)); err != nil {
|
||||
return nil
|
||||
}
|
||||
return uuidValue(*value)
|
||||
}
|
||||
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"net/http"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
@@ -176,6 +177,9 @@ func (e *Engine) Apply(request *http.Request, maxBody int64, apiKeyID string) (R
|
||||
}
|
||||
model := findModel(document)
|
||||
result := Result{}
|
||||
// 只收集"本请求命中"的 redact 策略的规则,供字节级重写使用;
|
||||
// 跨端点/模型/API Key 作用域的策略不得改写本请求。
|
||||
matchedRedactionRules := make([]compiledRule, 0)
|
||||
for _, policy := range e.current.Load().policies {
|
||||
if !policyApplies(policy, request.URL.Path, model, apiKeyID) {
|
||||
continue
|
||||
@@ -206,17 +210,178 @@ func (e *Engine) Apply(request *http.Request, maxBody int64, apiKeyID string) (R
|
||||
break
|
||||
}
|
||||
result.Redacted = result.Redacted || changed
|
||||
if policy.Action == "redact" && changed {
|
||||
matchedRedactionRules = append(matchedRedactionRules, policy.rules...)
|
||||
}
|
||||
}
|
||||
if result.Redacted && !result.Blocked {
|
||||
encoded, err := json.Marshal(document)
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
// 在原始请求字节上做字符串字面量级替换,而不是解码后重新
|
||||
// json.Marshal:重编码会改变键序、数字格式与 HTML 转义,破坏上游
|
||||
// 的请求签名/哈希与字节级契约,且对大 body 是双倍编解码开销。
|
||||
encoded, changed := redactBytes(body, matchedRedactionRules)
|
||||
if changed {
|
||||
restoreBody(request, encoded)
|
||||
} else {
|
||||
// 字节级替换未生效(如规则只匹配解码后文本但替换失败):
|
||||
// 不得谎报已脱敏,否则 X-Gateway-Content-Redacted 与审计
|
||||
// 都声称敏感信息已被移除,而实际请求原样发往上游。
|
||||
result.Redacted = false
|
||||
}
|
||||
restoreBody(request, encoded)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// redactBytes 逐字节扫描 JSON,仅对"位于 textual 字段(content/text/input/
|
||||
// prompt/instructions)下的值字符串"应用脱敏规则,其余字节(键名、空白、
|
||||
// 数字、布尔、结构符)原样保留。与解码后替换相比:
|
||||
// - 字节流与原始请求一致,除被替换的匹配段外零改动,不破坏上游签名/
|
||||
// 哈希与字节级契约,也不做大 body 的双倍编解码;
|
||||
// - 字符串内的 JSON 转义按原文匹配(如 \uXXXX),secret 模式通常不含
|
||||
// 需要转义的字符,实际影响可忽略。
|
||||
//
|
||||
// 返回替换后的字节与是否发生过替换。
|
||||
func redactBytes(raw []byte, rules []compiledRule) ([]byte, bool) {
|
||||
if len(rules) == 0 {
|
||||
return raw, false
|
||||
}
|
||||
type frame struct {
|
||||
inObject bool
|
||||
selected bool
|
||||
}
|
||||
changed := false
|
||||
out := make([]byte, 0, len(raw)+64)
|
||||
curSelected := false
|
||||
keySelected := false
|
||||
stack := make([]frame, 0, 8)
|
||||
for i := 0; i < len(raw); {
|
||||
ch := raw[i]
|
||||
switch ch {
|
||||
case '{', '[':
|
||||
stack = append(stack, frame{inObject: ch == '{', selected: curSelected})
|
||||
out = append(out, ch)
|
||||
i++
|
||||
continue
|
||||
case '}', ']':
|
||||
if len(stack) > 0 {
|
||||
curSelected = stack[len(stack)-1].selected
|
||||
stack = stack[:len(stack)-1]
|
||||
}
|
||||
out = append(out, ch)
|
||||
i++
|
||||
continue
|
||||
case ',':
|
||||
// 对象内逗号后是键:selected 由下一个键决定;数组内逗号后是
|
||||
// 元素,继承当前 selected。
|
||||
if len(stack) > 0 && stack[len(stack)-1].inObject {
|
||||
curSelected = false
|
||||
}
|
||||
out = append(out, ch)
|
||||
i++
|
||||
continue
|
||||
case ':':
|
||||
// 键后冒号:值字符串的 selected 由该键决定。
|
||||
curSelected = keySelected
|
||||
out = append(out, ch)
|
||||
i++
|
||||
continue
|
||||
case '"':
|
||||
// 定位字符串结束(处理反斜杠转义)。
|
||||
j := i + 1
|
||||
escaped := false
|
||||
for j < len(raw) {
|
||||
if escaped {
|
||||
escaped = false
|
||||
j++
|
||||
continue
|
||||
}
|
||||
if raw[j] == '\\' {
|
||||
escaped = true
|
||||
j++
|
||||
continue
|
||||
}
|
||||
if raw[j] == '"' {
|
||||
break
|
||||
}
|
||||
j++
|
||||
}
|
||||
if j >= len(raw) {
|
||||
// 截断/畸形 JSON:剩余字节原样保留。
|
||||
out = append(out, raw[i:]...)
|
||||
break
|
||||
}
|
||||
content := raw[i+1 : j]
|
||||
// 键还是值:字符串后第一个非空白字符是 ':' 即为对象键。
|
||||
k := j + 1
|
||||
for k < len(raw) && (raw[k] == ' ' || raw[k] == '\t' || raw[k] == '\n' || raw[k] == '\r') {
|
||||
k++
|
||||
}
|
||||
isKey := k < len(raw) && raw[k] == ':'
|
||||
if isKey {
|
||||
// textual 字段名传播到其值:父级 selected 或键名命中。
|
||||
parentSelected := false
|
||||
if len(stack) > 0 {
|
||||
parentSelected = stack[len(stack)-1].selected
|
||||
}
|
||||
keySelected = parentSelected || textualFields[strings.ToLower(string(content))]
|
||||
} else if curSelected {
|
||||
if replaced, hit := applyRules(content, rules); hit {
|
||||
changed = true
|
||||
out = append(out, '"')
|
||||
out = append(out, replaced...)
|
||||
out = append(out, '"')
|
||||
i = j + 1
|
||||
continue
|
||||
}
|
||||
}
|
||||
out = append(out, raw[i:j+1]...)
|
||||
i = j + 1
|
||||
continue
|
||||
default:
|
||||
out = append(out, ch)
|
||||
i++
|
||||
}
|
||||
}
|
||||
return out, changed
|
||||
}
|
||||
|
||||
func applyRules(value []byte, rules []compiledRule) ([]byte, bool) {
|
||||
changed := false
|
||||
text := string(value)
|
||||
for _, rule := range rules {
|
||||
if rule.expression.MatchString(text) {
|
||||
next := rule.expression.ReplaceAllString(text, jsonEscapeReplacement(rule.replacement))
|
||||
if next != text {
|
||||
text = next
|
||||
changed = true
|
||||
}
|
||||
continue
|
||||
}
|
||||
// 原文(含 JSON 转义)未命中,但解码后的文本可能命中
|
||||
// (Go 的 json.Marshal 会把 - & < > 等转义为 \u002d \u0026 ...):
|
||||
// 对解码文本应用替换后重新做 JSON 字符串转义,其余字节不变。
|
||||
var decoded string
|
||||
wrapped := append([]byte(`"`), value...)
|
||||
wrapped = append(wrapped, '"')
|
||||
if json.Unmarshal(wrapped, &decoded) == nil && decoded != text && rule.expression.MatchString(decoded) {
|
||||
next := rule.expression.ReplaceAllString(decoded, rule.replacement)
|
||||
quoted := strconv.Quote(next)
|
||||
text = quoted[1 : len(quoted)-1]
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
return []byte(text), changed
|
||||
}
|
||||
|
||||
// jsonEscapeReplacement 把替换文本转义成可安全嵌入 JSON 字符串字面量的
|
||||
// 形式:替换文本含引号/反斜杠/控制字符时直接插入会破坏 JSON 结构。
|
||||
func jsonEscapeReplacement(value string) string {
|
||||
if !strings.ContainsAny(value, "\"\\\n\r\t") && !strings.ContainsFunc(value, func(r rune) bool { return r < 0x20 }) {
|
||||
return value
|
||||
}
|
||||
quoted := strconv.Quote(value)
|
||||
return quoted[1 : len(quoted)-1]
|
||||
}
|
||||
|
||||
func policyApplies(policy compiledPolicy, path, model, apiKeyID string) bool {
|
||||
return matches(policy.Paths, path) && matches(policy.Models, model) && matches(policy.APIKeyIDs, apiKeyID)
|
||||
}
|
||||
|
||||
@@ -61,8 +61,9 @@ func NewEngine(pool *pgxpool.Pool, retriever EvidenceRetriever, logger *slog.Log
|
||||
// Check verifies one assistant answer against the configured knowledge bases
|
||||
// and persists a fact_check_events row. The returned Event has a zero ID when
|
||||
// fact-checking is not configured or no policy applies; callers should skip
|
||||
// quietly in that case.
|
||||
func (e *Engine) Check(ctx context.Context, requestID, question, answer string, verifier Verifier) (Event, error) {
|
||||
// quietly in that case. scope 形如 department:<uuid>(由调用方从资源部门推导);
|
||||
// 空 scope 时只应用 global 策略。
|
||||
func (e *Engine) Check(ctx context.Context, requestID, scope, question, answer string, verifier Verifier) (Event, error) {
|
||||
if e == nil || e.retriever == nil || verifier == nil || strings.TrimSpace(answer) == "" {
|
||||
return Event{}, nil
|
||||
}
|
||||
@@ -73,7 +74,7 @@ func (e *Engine) Check(ctx context.Context, requestID, question, answer string,
|
||||
if strings.TrimSpace(settings.Model) == "" {
|
||||
return Event{}, nil // not configured; skip without noise
|
||||
}
|
||||
policy, err := e.enabledPolicy(ctx)
|
||||
policy, err := e.enabledPolicy(ctx, scope)
|
||||
if err != nil {
|
||||
return Event{}, err
|
||||
}
|
||||
@@ -117,8 +118,11 @@ func (e *Engine) checkSettings(ctx context.Context) (Settings, error) {
|
||||
return x, err
|
||||
}
|
||||
|
||||
func (e *Engine) enabledPolicy(ctx context.Context) (Policy, error) {
|
||||
policy, err := scanPolicy(e.pool.QueryRow(ctx, policySelect+` WHERE enabled ORDER BY scope LIMIT 1`))
|
||||
// enabledPolicy 选择命中的策略:优先精确匹配调用方 scope(department:<uuid>
|
||||
// 等),否则回退 global。修复之前 ORDER BY scope LIMIT 1 只取字典序第一条
|
||||
// 的问题——多策略并存时其余部门的策略被静默忽略或张冠李戴。
|
||||
func (e *Engine) enabledPolicy(ctx context.Context, scope string) (Policy, error) {
|
||||
policy, err := scanPolicy(e.pool.QueryRow(ctx, policySelect+` WHERE enabled AND (scope='global' OR scope=$1) ORDER BY (scope='global'),scope LIMIT 1`, scope))
|
||||
if errors.Is(err, ErrNotFound) {
|
||||
return Policy{}, nil
|
||||
}
|
||||
|
||||
@@ -89,7 +89,10 @@ func (s *auditSpan) captureBody(body io.ReadCloser) io.ReadCloser {
|
||||
return body
|
||||
}
|
||||
capture := &captureReadCloser{ReadCloser: body, limit: auditRequestCaptureBytes}
|
||||
// 与 finish/setModel 的读取保持同一把锁,防止未来异步化审计时出现竞态。
|
||||
s.mu.Lock()
|
||||
s.capture = capture
|
||||
s.mu.Unlock()
|
||||
return capture
|
||||
}
|
||||
|
||||
|
||||
@@ -71,6 +71,7 @@ func TestContentPolicyAndPricingIntegration(t *testing.T) {
|
||||
go func() { recorder.Run(recordCtx); close(stopped) }()
|
||||
defer func() { cancel(); <-stopped }()
|
||||
proxy := NewProxy(adapter, "test-key", 1<<20, slog.Default())
|
||||
proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1
|
||||
proxy.SetAuditRecorder(recorder)
|
||||
proxy.SetContentPolicyEngine(engine)
|
||||
proxy.SetPricingService(prices)
|
||||
|
||||
@@ -7,7 +7,6 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httputil"
|
||||
"strconv"
|
||||
@@ -67,9 +66,11 @@ func NewProxyWithAuthenticator(adapter provider.Adapter, authenticator apikey.Ke
|
||||
}
|
||||
|
||||
func NewDynamicProxy(resolver AdapterResolver, authenticator apikey.KeyAuthenticator, maxBody int64, logger *slog.Logger) *Proxy {
|
||||
// 默认拒绝拨号到非公网地址:管理员未显式放行私网时,数据平面在拨号阶段
|
||||
// 复检目标地址,防止 DNS rebinding 把流量引到内网(169.254.169.254 等)。
|
||||
transport := &http.Transport{
|
||||
Proxy: http.ProxyFromEnvironment,
|
||||
DialContext: (&net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
|
||||
DialContext: provider.SafeDialContext(false, 5*time.Second, 30*time.Second),
|
||||
ForceAttemptHTTP2: true,
|
||||
MaxIdleConns: 512,
|
||||
MaxIdleConnsPerHost: 256,
|
||||
@@ -81,6 +82,15 @@ func NewDynamicProxy(resolver AdapterResolver, authenticator apikey.KeyAuthentic
|
||||
return &Proxy{resolver: resolver, auth: authenticator, maxBody: maxBody, logger: logger, transport: transport, resilience: DefaultResiliencePolicy()}
|
||||
}
|
||||
|
||||
// SetAllowPrivateProviderURLs 允许数据平面拨号到私网地址(与管理员配置的
|
||||
// ALLOW_PRIVATE_PROVIDER_URLS 保持一致);关闭时保持拨号阶段 SSRF 校验。
|
||||
func (p *Proxy) SetAllowPrivateProviderURLs(allow bool) {
|
||||
timeout := p.transport.ResponseHeaderTimeout
|
||||
p.transport = p.transport.Clone()
|
||||
p.transport.DialContext = provider.SafeDialContext(allow, 5*time.Second, 30*time.Second)
|
||||
p.transport.ResponseHeaderTimeout = timeout
|
||||
}
|
||||
|
||||
func (p *Proxy) SetAdmissionController(controller AdmissionController) {
|
||||
p.admission = controller
|
||||
}
|
||||
@@ -328,7 +338,16 @@ func (p *Proxy) proxyFor(resolved ResolvedAdapter) *httputil.ReverseProxy {
|
||||
}
|
||||
writeOpenAIError(writer, http.StatusBadGateway, "upstream_error", "upstream service is unavailable")
|
||||
}
|
||||
p.proxies.Store(resolved.Code, cachedProxy{key: key, proxy: reverseProxy})
|
||||
// 并发缓存 miss 时只保留一个胜出的代理,其余立即丢弃,避免重复构建。
|
||||
if actual, loaded := p.proxies.LoadOrStore(resolved.Code, cachedProxy{key: key, proxy: reverseProxy}); loaded {
|
||||
entry := actual.(cachedProxy)
|
||||
if entry.key == key {
|
||||
return entry.proxy
|
||||
}
|
||||
// 另一个 goroutine 写入了不同的 key(快照已前进):保留新条目。
|
||||
_ = reverseProxy
|
||||
return entry.proxy
|
||||
}
|
||||
return reverseProxy
|
||||
}
|
||||
|
||||
|
||||
@@ -72,6 +72,7 @@ func TestProxyRejectsInvalidKey(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
proxy := NewProxy(adapter, "gateway-secret", 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
request.Header.Set("Authorization", "Bearer wrong")
|
||||
@@ -89,6 +90,7 @@ func TestProxyRejectsKnownOversizedBody(t *testing.T) {
|
||||
defer upstream.Close()
|
||||
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
|
||||
proxy := NewProxy(adapter, "gateway-secret", 4, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader("12345"))
|
||||
request.Header.Set("Authorization", "Bearer gateway-secret")
|
||||
@@ -110,6 +112,7 @@ func TestProxyReplacesClientAuthorization(t *testing.T) {
|
||||
defer upstream.Close()
|
||||
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
|
||||
proxy := NewProxy(adapter, "gateway-secret", 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
request.Header.Set("Authorization", "Bearer gateway-secret")
|
||||
@@ -162,6 +165,7 @@ func TestProxyReconcilesReservedTokensWithUpstreamUsage(t *testing.T) {
|
||||
proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
|
||||
APIKeyID: "key-1", Scopes: []string{"gateway:invoke"}, MonthlyTokenQuota: 1000,
|
||||
}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1
|
||||
proxy.SetTokenQuotaController(quota)
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"test","max_tokens":20}`))
|
||||
@@ -218,6 +222,7 @@ func TestProxyRewritesModelAliasBeforeUpstream(t *testing.T) {
|
||||
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
|
||||
resolved := ResolvedAdapter{Code: "routed", Adapter: adapter, Capabilities: map[provider.Capability]bool{provider.CapabilityChat: true}}
|
||||
proxy := NewDynamicProxy(fixedRoutingResolver{adapter: resolved}, principalAuthenticator{principal: apikey.Principal{Scopes: []string{"gateway:invoke"}}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"public-chat","messages":[]}`))
|
||||
request.Header.Set("Authorization", "Bearer gateway-secret")
|
||||
response := httptest.NewRecorder()
|
||||
@@ -238,6 +243,7 @@ func TestProxyRecordsAuditWithoutBufferingWholeResponse(t *testing.T) {
|
||||
proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
|
||||
APIKeyID: "11111111-1111-4111-8111-111111111111", Scopes: []string{"gateway:invoke"},
|
||||
}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1
|
||||
proxy.SetAuditRecorder(recorder)
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"audit-model","messages":[]}`))
|
||||
request.Header.Set("Authorization", "Bearer gateway-secret")
|
||||
|
||||
@@ -43,25 +43,48 @@ func newCircuitBreaker(policy ...ResiliencePolicy) *circuitBreaker {
|
||||
return &circuitBreaker{threshold: settings.CircuitThreshold, openFor: settings.CircuitOpenDuration}
|
||||
}
|
||||
|
||||
func (c *circuitBreaker) allow(now time.Time) bool {
|
||||
// allow returns whether the request may proceed and whether it is the
|
||||
// half-open probe (the single request allowed through an open circuit to
|
||||
// test recovery).
|
||||
func (c *circuitBreaker) allow(now time.Time) (allowed bool, probe bool) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if c.openUntil.IsZero() {
|
||||
return true
|
||||
return true, false
|
||||
}
|
||||
if now.Before(c.openUntil) || c.halfOpenRun {
|
||||
return false
|
||||
return false, false
|
||||
}
|
||||
c.halfOpenRun = true
|
||||
return true
|
||||
return true, true
|
||||
}
|
||||
|
||||
func (c *circuitBreaker) success() {
|
||||
func (c *circuitBreaker) success(probe bool) {
|
||||
c.mu.Lock()
|
||||
c.failures = 0
|
||||
c.openUntil = time.Time{}
|
||||
defer c.mu.Unlock()
|
||||
if c.openUntil.IsZero() {
|
||||
// 关闭状态下普通成功:仅清零失败计数。
|
||||
c.failures = 0
|
||||
return
|
||||
}
|
||||
if probe && c.halfOpenRun {
|
||||
// 半开探针成功:关闭电路,恢复正常流量。
|
||||
c.failures = 0
|
||||
c.openUntil = time.Time{}
|
||||
c.halfOpenRun = false
|
||||
return
|
||||
}
|
||||
// 电路已打开而请求在打开前就通过 allow():陈旧成功不得关闭电路,
|
||||
// 否则刚触发熔断的上游被一个在途成功立即放行。
|
||||
}
|
||||
|
||||
// abortProbe 在探针请求被客户端取消(而非上游失败)时调用:
|
||||
// 既无成功也无失败的证据,释放探针名额但不改变电路状态,让下一次
|
||||
// allow() 重新发起探针。
|
||||
func (c *circuitBreaker) abortProbe() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.halfOpenRun = false
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
func (c *circuitBreaker) failure(now time.Time) {
|
||||
@@ -82,7 +105,8 @@ type resilientTransport struct {
|
||||
}
|
||||
|
||||
func (t *resilientTransport) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
if !t.circuit.allow(time.Now()) {
|
||||
allowed, probe := t.circuit.allow(time.Now())
|
||||
if !allowed {
|
||||
return nil, ErrCircuitOpen
|
||||
}
|
||||
replayable := request.Method == http.MethodGet || request.Method == http.MethodHead ||
|
||||
@@ -102,6 +126,11 @@ func (t *resilientTransport) RoundTrip(request *http.Request) (*http.Response, e
|
||||
}
|
||||
current = request.Clone(request.Context())
|
||||
if request.Body != nil && request.Body != http.NoBody {
|
||||
if request.GetBody == nil {
|
||||
// 请求体不可重放(如 GET/HEAD 携带未设置 GetBody 的 body):
|
||||
// 放弃重试,避免把已消费的空 body 重发或调用 nil 方法。
|
||||
break
|
||||
}
|
||||
body, bodyErr := request.GetBody()
|
||||
if bodyErr != nil {
|
||||
err = bodyErr
|
||||
@@ -119,14 +148,24 @@ func (t *resilientTransport) RoundTrip(request *http.Request) (*http.Response, e
|
||||
_ = response.Body.Close()
|
||||
}
|
||||
}
|
||||
if retryableResult(response, err) {
|
||||
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
// 客户端取消/超时不是上游故障证据:不计成败。若本请求是半开
|
||||
// 探针,释放探针名额让电路保持打开,由下一次请求重新探测。
|
||||
if probe {
|
||||
t.circuit.abortProbe()
|
||||
}
|
||||
return response, err
|
||||
}
|
||||
if failureResult(response, err) {
|
||||
t.circuit.failure(time.Now())
|
||||
} else {
|
||||
t.circuit.success()
|
||||
t.circuit.success(probe)
|
||||
}
|
||||
return response, err
|
||||
}
|
||||
|
||||
// retryableResult 决定是否值得重试:仅传输错误与 502/503/504 会重试,
|
||||
// 500 等其余 5xx 不做自动重试(响应可能已被上游处理,重试有副作用)。
|
||||
func retryableResult(response *http.Response, err error) bool {
|
||||
if err != nil {
|
||||
return !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded)
|
||||
@@ -134,6 +173,15 @@ func retryableResult(response *http.Response, err error) bool {
|
||||
return response != nil && (response.StatusCode == http.StatusBadGateway || response.StatusCode == http.StatusServiceUnavailable || response.StatusCode == http.StatusGatewayTimeout)
|
||||
}
|
||||
|
||||
// failureResult 决定是否计入熔断失败:所有 5xx 都视为上游故障。否则持续返回
|
||||
// 500 的上游永远不会触发熔断,而 success() 还会不断清零失败计数,熔断保护失效。
|
||||
func failureResult(response *http.Response, err error) bool {
|
||||
if err != nil {
|
||||
return !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded)
|
||||
}
|
||||
return response != nil && response.StatusCode >= http.StatusInternalServerError
|
||||
}
|
||||
|
||||
func waitBackoff(ctx context.Context, duration time.Duration) error {
|
||||
if duration <= 0 {
|
||||
return nil
|
||||
|
||||
@@ -49,13 +49,13 @@ func TestCircuitOpensAndAllowsSingleProbe(t *testing.T) {
|
||||
now := time.Now()
|
||||
circuit.failure(now)
|
||||
circuit.failure(now)
|
||||
if circuit.allow(now) {
|
||||
if allowed, _ := circuit.allow(now); allowed {
|
||||
t.Fatal("open circuit allowed request")
|
||||
}
|
||||
if !circuit.allow(now.Add(2 * time.Millisecond)) {
|
||||
if allowed, probe := circuit.allow(now.Add(2 * time.Millisecond)); !allowed || !probe {
|
||||
t.Fatal("circuit did not allow half-open probe")
|
||||
}
|
||||
if circuit.allow(now.Add(2 * time.Millisecond)) {
|
||||
if allowed, _ := circuit.allow(now.Add(2 * time.Millisecond)); allowed {
|
||||
t.Fatal("circuit allowed concurrent half-open probe")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -110,6 +110,11 @@ return {1, current}
|
||||
|
||||
const tokenCommitScript = `
|
||||
local delta = tonumber(ARGV[1])
|
||||
-- 预留与提交之间月份可能已翻转,计数器键已过期:此时直接放弃回写,
|
||||
-- 不能重建一个永不过期的残留键(旧月份数据已无意义)。
|
||||
if redis.call('EXISTS', KEYS[1]) == 0 then
|
||||
return 0
|
||||
end
|
||||
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
|
||||
local updated = current + delta
|
||||
if updated < 0 then updated = 0 end
|
||||
|
||||
+39
-11
@@ -107,17 +107,19 @@ type usageCollector struct {
|
||||
}
|
||||
|
||||
func (c *usageCollector) feed(chunk []byte) {
|
||||
if !c.sse {
|
||||
c.doc = append(c.doc, chunk...)
|
||||
// Keep only a bounded tail window. The "usage" member lives at the end
|
||||
// of a non-streaming response, so dropping the head (never the tail)
|
||||
// preserves accounting for arbitrarily large bodies at a fixed memory
|
||||
// cost instead of truncating usage away.
|
||||
if len(c.doc) > maxUsageDocumentBytes {
|
||||
c.doc = append([]byte(nil), c.doc[len(c.doc)-maxUsageDocumentBytes:]...)
|
||||
if !c.sse {
|
||||
c.doc = append(c.doc, chunk...)
|
||||
// Keep only a bounded tail window. The "usage" member lives at the end
|
||||
// of a non-streaming response, so dropping the head (never the tail)
|
||||
// preserves accounting for arbitrarily large bodies at a fixed memory
|
||||
// cost instead of truncating usage away. 仅在超过 2× 窗口时压缩一次,
|
||||
// 避免每个 32KiB 块都做 O(窗口) 的尾部拷贝(大响应下退化为 O(n²))。
|
||||
if len(c.doc) > 2*maxUsageDocumentBytes {
|
||||
copy(c.doc, c.doc[len(c.doc)-maxUsageDocumentBytes:])
|
||||
c.doc = c.doc[:maxUsageDocumentBytes]
|
||||
}
|
||||
return
|
||||
}
|
||||
return
|
||||
}
|
||||
c.pending = append(c.pending, chunk...)
|
||||
for {
|
||||
index := bytes.IndexByte(c.pending, '\n')
|
||||
@@ -164,7 +166,19 @@ func (c *usageCollector) usage() TokenUsage {
|
||||
// still counted when the buffered window starts mid-object.
|
||||
func (c *usageCollector) consumeUsageObject(doc []byte) {
|
||||
const key = `"usage"`
|
||||
// 只认 JSON 对象成员位置的 "usage"(前一个非空白字符是 '{' 或 ','),
|
||||
// 避免命中字符串值里的同名文本。
|
||||
index := bytes.LastIndex(doc, []byte(key))
|
||||
for index >= 0 {
|
||||
j := index - 1
|
||||
for j >= 0 && (doc[j] == ' ' || doc[j] == '\t' || doc[j] == '\n' || doc[j] == '\r') {
|
||||
j--
|
||||
}
|
||||
if j < 0 || doc[j] == '{' || doc[j] == ',' {
|
||||
break
|
||||
}
|
||||
index = bytes.LastIndex(doc[:index], []byte(key))
|
||||
}
|
||||
if index < 0 {
|
||||
return
|
||||
}
|
||||
@@ -209,7 +223,21 @@ func (c *usageCollector) consumeUsageObject(doc []byte) {
|
||||
}
|
||||
}
|
||||
if end > 0 {
|
||||
c.consumeJSON(rest[:end])
|
||||
// 提取出的是 usage 对象本身:必须按 inUsage=true 解析,否则其
|
||||
// 顶层 prompt_tokens/completion_tokens/total_tokens 不会被计数,
|
||||
// 大响应(>4MB 压缩后)的 token 计量静默丢失。
|
||||
c.consumeJSONAsUsage(rest[:end])
|
||||
}
|
||||
}
|
||||
|
||||
// consumeJSONAsUsage parses payload with the "inside usage" flag already set,
|
||||
// so top-level *_tokens keys are counted.
|
||||
func (c *usageCollector) consumeJSONAsUsage(payload []byte) {
|
||||
decoder := json.NewDecoder(bytes.NewReader(payload))
|
||||
decoder.UseNumber()
|
||||
var value any
|
||||
if decoder.Decode(&value) == nil {
|
||||
c.walk(value, true)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -75,6 +75,13 @@ const (
|
||||
PermissionMarketplaceManage = "marketplace:manage"
|
||||
PermissionFileRead = "file:read"
|
||||
PermissionFileManage = "file:manage"
|
||||
PermissionInboxRead = "inbox:read"
|
||||
PermissionInboxManage = "inbox:manage"
|
||||
PermissionScheduledTaskRead = "scheduled_task:read"
|
||||
PermissionScheduledTaskManage = "scheduled_task:manage"
|
||||
PermissionTraceRead = "trace:read"
|
||||
PermissionAgentNodeRead = "agent_node:read"
|
||||
PermissionAgentNodeManage = "agent_node:manage"
|
||||
)
|
||||
|
||||
var rolePermissions = map[string][]string{
|
||||
@@ -96,8 +103,12 @@ var rolePermissions = map[string][]string{
|
||||
PermissionDigitalEmployeeRead, PermissionDigitalEmployeeManage,
|
||||
PermissionMarketplaceRead, PermissionMarketplaceManage,
|
||||
PermissionFileRead, PermissionFileManage,
|
||||
PermissionInboxRead, PermissionInboxManage,
|
||||
PermissionScheduledTaskRead, PermissionScheduledTaskManage,
|
||||
PermissionTraceRead,
|
||||
PermissionAgentNodeRead, PermissionAgentNodeManage,
|
||||
},
|
||||
"auditor": {PermissionProviderRead, PermissionAPIKeyRead, PermissionAuditRead, PermissionUsageRead, PermissionOutboxRead, PermissionContentPolicyRead, PermissionPricingRead, PermissionPromptRead, PermissionKnowledgeRead, PermissionToolRead, PermissionApplicationRead, PermissionNotificationRead, PermissionMCPServerRead, PermissionSkillRead, PermissionDigitalEmployeeRead, PermissionMarketplaceRead, PermissionFileRead},
|
||||
"auditor": {PermissionProviderRead, PermissionAPIKeyRead, PermissionAuditRead, PermissionUsageRead, PermissionOutboxRead, PermissionContentPolicyRead, PermissionPricingRead, PermissionPromptRead, PermissionKnowledgeRead, PermissionToolRead, PermissionApplicationRead, PermissionNotificationRead, PermissionMCPServerRead, PermissionSkillRead, PermissionDigitalEmployeeRead, PermissionMarketplaceRead, PermissionFileRead, PermissionInboxRead, PermissionScheduledTaskRead, PermissionTraceRead, PermissionAgentNodeRead},
|
||||
"member": {},
|
||||
}
|
||||
|
||||
|
||||
@@ -79,11 +79,21 @@ func (h *ManagementHTTPHandler) updateDepartment(writer http.ResponseWriter, req
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
_, department, ok := decodeDepartment(writer, request)
|
||||
input, department, ok := decodeDepartment(writer, request)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
department.ID = request.PathValue("department_id")
|
||||
current, err := h.service.repository.GetDepartment(request.Context(), department.ID)
|
||||
if err != nil {
|
||||
h.writeDepartmentError(writer, err)
|
||||
return
|
||||
}
|
||||
// 部分更新语义:省略 active 时保留当前状态,避免"只改名称"的 PUT
|
||||
// 绕过停用保护把部门静默重新激活。
|
||||
if input.Active == nil {
|
||||
department.Active = current.Active
|
||||
}
|
||||
updated, err := h.service.repository.UpdateDepartment(request.Context(), department, actor.ID)
|
||||
if err != nil {
|
||||
h.writeDepartmentError(writer, err)
|
||||
|
||||
@@ -85,7 +85,7 @@ func (h *HTTPHandler) ServeHTTP(writer http.ResponseWriter, request *http.Reques
|
||||
func (h *HTTPHandler) login(kind Kind) http.HandlerFunc {
|
||||
return func(writer http.ResponseWriter, request *http.Request) {
|
||||
// 防爆破:按 IP 的滑动窗口限流,超限返回 429(与账号锁定叠加)。
|
||||
if !h.service.AllowLogin(request.Context(), ClientIP(request)) {
|
||||
if !h.service.AllowLogin(request.Context(), h.service.ClientIP(request)) {
|
||||
apiresponse.Error(writer, http.StatusTooManyRequests, "登录尝试过于频繁,请稍后再试")
|
||||
return
|
||||
}
|
||||
@@ -131,7 +131,7 @@ func (h *HTTPHandler) registerTOTP(kind Kind, prefix string) {
|
||||
func (h *HTTPHandler) completeTOTPLogin(kind Kind) http.HandlerFunc {
|
||||
return func(writer http.ResponseWriter, request *http.Request) {
|
||||
// 防爆破:TOTP 完成端点同样按 IP 限流。
|
||||
if !h.service.AllowLogin(request.Context(), ClientIP(request)) {
|
||||
if !h.service.AllowLogin(request.Context(), h.service.ClientIP(request)) {
|
||||
apiresponse.Error(writer, http.StatusTooManyRequests, "登录尝试过于频繁,请稍后再试")
|
||||
return
|
||||
}
|
||||
@@ -348,8 +348,8 @@ func adminMenus(account Account) []map[string]any {
|
||||
menus = append(menus, map[string]any{"name": "Gateway", "path": "/gateway", "component": "/index/index", "meta": map[string]any{"title": "网关接入", "icon": "ri:router-line"}, "children": gatewayChildren})
|
||||
}
|
||||
|
||||
// 安全与审计:审计用量、内容策略与模型治理。
|
||||
securityChildren := make([]map[string]any, 0, 3)
|
||||
// 安全与审计:审计用量、内容策略、模型治理、Trace、会话与节点。
|
||||
securityChildren := make([]map[string]any, 0, 6)
|
||||
if HasPermission(account, PermissionAuditRead) || HasPermission(account, PermissionUsageRead) {
|
||||
securityChildren = append(securityChildren, map[string]any{"name": "AuditUsage", "path": "audit-usage", "component": "/gateway/audit-usage", "meta": map[string]any{"title": "审计与用量"}})
|
||||
}
|
||||
@@ -359,6 +359,13 @@ func adminMenus(account Account) []map[string]any {
|
||||
if HasPermission(account, PermissionKnowledgeRead) || HasPermission(account, PermissionKnowledgeManage) {
|
||||
securityChildren = append(securityChildren, map[string]any{"name": "Governance", "path": "governance", "component": "/gateway/governance", "meta": map[string]any{"title": "模型治理"}})
|
||||
}
|
||||
if HasPermission(account, PermissionTraceRead) {
|
||||
securityChildren = append(securityChildren, map[string]any{"name": "Traces", "path": "traces", "component": "/gateway/traces", "meta": map[string]any{"title": "LLM Trace"}})
|
||||
securityChildren = append(securityChildren, map[string]any{"name": "AgentSessions", "path": "agent-sessions", "component": "/gateway/agent-sessions", "meta": map[string]any{"title": "智能体会话"}})
|
||||
}
|
||||
if HasPermission(account, PermissionAgentNodeRead) || HasPermission(account, PermissionAgentNodeManage) {
|
||||
securityChildren = append(securityChildren, map[string]any{"name": "AgentNodes", "path": "agent-nodes", "component": "/gateway/agent-nodes", "meta": map[string]any{"title": "智能体节点"}})
|
||||
}
|
||||
if len(securityChildren) > 0 {
|
||||
menus = append(menus, map[string]any{"name": "Security", "path": "/security", "component": "/index/index", "meta": map[string]any{"title": "安全与审计", "icon": "ri:shield-check-line"}, "children": securityChildren})
|
||||
}
|
||||
@@ -413,6 +420,12 @@ func adminMenus(account Account) []map[string]any {
|
||||
if HasPermission(account, PermissionNotificationRead) || HasPermission(account, PermissionNotificationManage) {
|
||||
systemChildren = append(systemChildren, map[string]any{"name": "Notifications", "path": "notifications", "component": "/gateway/notifications", "meta": map[string]any{"title": "通知中心"}})
|
||||
}
|
||||
if HasPermission(account, PermissionInboxRead) || HasPermission(account, PermissionInboxManage) {
|
||||
systemChildren = append(systemChildren, map[string]any{"name": "Inbox", "path": "inbox", "component": "/gateway/inbox", "meta": map[string]any{"title": "站内消息"}})
|
||||
}
|
||||
if HasPermission(account, PermissionScheduledTaskRead) || HasPermission(account, PermissionScheduledTaskManage) {
|
||||
systemChildren = append(systemChildren, map[string]any{"name": "ScheduledTasks", "path": "scheduled-tasks", "component": "/gateway/scheduled-tasks", "meta": map[string]any{"title": "定时任务"}})
|
||||
}
|
||||
if len(systemChildren) > 0 {
|
||||
menus = append(menus, map[string]any{"name": "System", "path": "/system", "component": "/index/index", "meta": map[string]any{"title": "系统管理", "icon": "ri:user-3-line"}, "children": systemChildren})
|
||||
}
|
||||
@@ -428,6 +441,7 @@ func portalMenus() []map[string]any {
|
||||
{"name": "PortalUsage", "path": "usage", "component": "/portal/usage", "meta": map[string]any{"title": "我的用量"}},
|
||||
{"name": "PortalAccess", "path": "access", "component": "/portal/access", "meta": map[string]any{"title": "模型权限"}},
|
||||
{"name": "PortalFiles", "path": "files", "component": "/portal/files", "meta": map[string]any{"title": "文件仓库"}},
|
||||
{"name": "PortalInbox", "path": "inbox", "component": "/portal/inbox", "meta": map[string]any{"title": "站内消息"}},
|
||||
}},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -105,7 +105,7 @@ func (h *ManagementHTTPHandler) update(kind Kind) http.HandlerFunc {
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
_, account, password, ok := h.decode(writer, request, kind, false)
|
||||
input, account, password, ok := h.decode(writer, request, kind, false)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
@@ -115,6 +115,14 @@ func (h *ManagementHTTPHandler) update(kind Kind) http.HandlerFunc {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
// 部分更新语义:未提供的字段保留当前值。否则"只改显示名"的 PUT 会
|
||||
// 把角色重置为默认 operator、把停用账号重新激活,造成意外的权限变更。
|
||||
if account.Role == "" {
|
||||
account.Role = current.Role
|
||||
}
|
||||
if input.Active == nil {
|
||||
account.Active = current.Active
|
||||
}
|
||||
if kind == KindAdmin && actor.ID == current.ID && (account.Role != current.Role || !account.Active) {
|
||||
apiresponse.Error(writer, http.StatusConflict, "不能停用自身账号或修改自身角色")
|
||||
return
|
||||
@@ -133,6 +141,11 @@ func (h *ManagementHTTPHandler) update(kind Kind) http.HandlerFunc {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
if password != "" {
|
||||
// 管理员重置密码属于凭据变更:立即作废该账号的全部既有会话,
|
||||
// 与自助改密/2FA 变更的语义保持一致。
|
||||
h.service.sessions.BumpAuthVersion(request.Context(), kind, account.ID)
|
||||
}
|
||||
apiresponse.OK(writer, managementView(updated))
|
||||
}
|
||||
}
|
||||
@@ -149,21 +162,25 @@ func (h *ManagementHTTPHandler) decode(writer http.ResponseWriter, request *http
|
||||
input.DisplayName = strings.TrimSpace(input.DisplayName)
|
||||
input.Role = strings.ToLower(strings.TrimSpace(input.Role))
|
||||
if input.Role == "" {
|
||||
if kind == KindPortal {
|
||||
input.Role = "member"
|
||||
} else {
|
||||
input.Role = "operator"
|
||||
// 创建时缺省角色;更新时留空表示"不修改该字段",
|
||||
// 由 update() 保留当前值,避免只改显示名就静默重置角色。
|
||||
if creating {
|
||||
if kind == KindPortal {
|
||||
input.Role = "member"
|
||||
} else {
|
||||
input.Role = "operator"
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(input.Login) < 2 || len(input.Login) > 128 || len(input.DisplayName) > 64 {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "账号或显示名称格式无效")
|
||||
return input, Account{}, "", false
|
||||
}
|
||||
if kind == KindAdmin && input.Role != "superadmin" && input.Role != "operator" && input.Role != "auditor" {
|
||||
if kind == KindAdmin && input.Role != "" && input.Role != "superadmin" && input.Role != "operator" && input.Role != "auditor" {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "管理员角色无效")
|
||||
return input, Account{}, "", false
|
||||
}
|
||||
if kind == KindPortal && input.Role != "member" {
|
||||
if kind == KindPortal && input.Role != "" && input.Role != "member" {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "门户角色无效")
|
||||
return input, Account{}, "", false
|
||||
}
|
||||
|
||||
+48
-13
@@ -471,6 +471,51 @@ func newOIDCHTTPClient(allowPrivate bool) *http.Client {
|
||||
}
|
||||
}
|
||||
|
||||
// jwksCacheTTL 是 JWKS 缓存有效期,以 IdP 最短键轮换周期为界。
|
||||
const jwksCacheTTL = 5 * time.Minute
|
||||
|
||||
type jwksKey struct{ Kid, Kty, N, E string }
|
||||
|
||||
type jwksDocument struct {
|
||||
Keys []jwksKey
|
||||
}
|
||||
|
||||
type jwksCacheEntry struct {
|
||||
doc jwksDocument
|
||||
fetched time.Time
|
||||
}
|
||||
|
||||
// jwksFor 返回 IdP 的 JWKS,带数分钟缓存(provider 数量有限,map 无需淘汰)。
|
||||
func (s *Service) jwksFor(ctx context.Context, jwksURI string) (jwksDocument, error) {
|
||||
s.jwksMu.Lock()
|
||||
if s.jwks == nil {
|
||||
s.jwks = make(map[string]jwksCacheEntry)
|
||||
}
|
||||
if entry, ok := s.jwks[jwksURI]; ok && time.Since(entry.fetched) < jwksCacheTTL {
|
||||
s.jwksMu.Unlock()
|
||||
return entry.doc, nil
|
||||
}
|
||||
s.jwksMu.Unlock()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, jwksURI, nil)
|
||||
if err != nil {
|
||||
return jwksDocument{}, err
|
||||
}
|
||||
response, err := s.oidcHTTPClient().Do(req)
|
||||
if err != nil {
|
||||
return jwksDocument{}, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
payload, _ := io.ReadAll(io.LimitReader(response.Body, oidcMaxResponse+1))
|
||||
var doc jwksDocument
|
||||
if response.StatusCode/100 != 2 || len(payload) > oidcMaxResponse || json.Unmarshal(payload, &doc) != nil {
|
||||
return jwksDocument{}, errors.New("jwks")
|
||||
}
|
||||
s.jwksMu.Lock()
|
||||
s.jwks[jwksURI] = jwksCacheEntry{doc: doc, fetched: time.Now()}
|
||||
s.jwksMu.Unlock()
|
||||
return doc, nil
|
||||
}
|
||||
|
||||
func (s *Service) verifyIDToken(ctx context.Context, d oidcDiscovery, clientID, nonce, token string) (oidcClaims, error) {
|
||||
parts := strings.Split(token, ".")
|
||||
if len(parts) != 3 {
|
||||
@@ -484,22 +529,12 @@ func (s *Service) verifyIDToken(ctx context.Context, d oidcDiscovery, clientID,
|
||||
if json.Unmarshal(headerBytes, &header) != nil || header.Alg != "RS256" || header.Kid == "" {
|
||||
return oidcClaims{}, errors.New("jwt header")
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, d.JWKSURI, nil)
|
||||
// JWKS 按 provider 缓存数分钟:每次登录都拉 discovery + JWKS 是 2-3 个
|
||||
// 同步往返;缓存以 IdP 最短键轮换周期为界(默认 5 分钟)。
|
||||
keys, err := s.jwksFor(ctx, d.JWKSURI)
|
||||
if err != nil {
|
||||
return oidcClaims{}, err
|
||||
}
|
||||
response, err := s.oidcHTTPClient().Do(req)
|
||||
if err != nil {
|
||||
return oidcClaims{}, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
payload, _ := io.ReadAll(io.LimitReader(response.Body, oidcMaxResponse+1))
|
||||
var keys struct {
|
||||
Keys []struct{ Kid, Kty, N, E string }
|
||||
}
|
||||
if response.StatusCode/100 != 2 || len(payload) > oidcMaxResponse || json.Unmarshal(payload, &keys) != nil {
|
||||
return oidcClaims{}, errors.New("jwks")
|
||||
}
|
||||
var key *rsa.PublicKey
|
||||
for _, j := range keys.Keys {
|
||||
if j.Kid == header.Kid && j.Kty == "RSA" {
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -24,13 +25,55 @@ import (
|
||||
// If Redis is unavailable or not configured the limiter fails OPEN (login
|
||||
// proceeds, account lockout still applies) rather than locking every user out.
|
||||
type LoginLimiter struct {
|
||||
client *redis.Client
|
||||
max int
|
||||
window time.Duration
|
||||
client *redis.Client
|
||||
max int
|
||||
window time.Duration
|
||||
trusted []netip.Prefix
|
||||
}
|
||||
|
||||
func NewLoginLimiter(client *redis.Client, max int, window time.Duration) *LoginLimiter {
|
||||
return &LoginLimiter{client: client, max: max, window: window}
|
||||
func NewLoginLimiter(client *redis.Client, max int, window time.Duration, trustedProxies []netip.Prefix) *LoginLimiter {
|
||||
return &LoginLimiter{client: client, max: max, window: window, trusted: trustedProxies}
|
||||
}
|
||||
|
||||
// ClientIP 提取用于登录限流的客户端 IP。仅当直连对端(RemoteAddr)属于可信
|
||||
// 代理网段时才采信 X-Forwarded-For;否则任何公网客户端都可以伪造该头,把
|
||||
// 每 IP 滑动窗口的键旋转掉,彻底绕过登录限流。
|
||||
func (l *LoginLimiter) ClientIP(r *http.Request) string {
|
||||
if l != nil && len(l.trusted) > 0 {
|
||||
peer, err := netip.ParseAddr(peerHost(r.RemoteAddr))
|
||||
if err == nil {
|
||||
peer = peer.Unmap()
|
||||
trustedPeer := false
|
||||
for _, prefix := range l.trusted {
|
||||
if prefix.Contains(peer) {
|
||||
trustedPeer = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if trustedPeer {
|
||||
if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
|
||||
if first := strings.TrimSpace(strings.Split(fwd, ",")[0]); first != "" {
|
||||
return first
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return peerHost(r.RemoteAddr)
|
||||
}
|
||||
|
||||
func peerHost(remoteAddr string) string {
|
||||
host, _, err := net.SplitHostPort(remoteAddr)
|
||||
if err != nil {
|
||||
host = remoteAddr
|
||||
}
|
||||
return host
|
||||
}
|
||||
|
||||
// ClientIP 兼容旧签名:未配置可信代理时退化为"仅信任内网对端"。
|
||||
// 保留给直接调用方;HTTP 处理路径统一走 LoginLimiter.ClientIP。
|
||||
func ClientIP(r *http.Request) string {
|
||||
return NewLoginLimiter(nil, 0, 0, nil).ClientIP(r)
|
||||
}
|
||||
|
||||
// Allow reports whether a login attempt from ip may proceed.
|
||||
@@ -85,20 +128,3 @@ redis.call('ZADD', key, now, ARGV[4])
|
||||
redis.call('EXPIRE', key, ARGV[5])
|
||||
return {0, count + 1}
|
||||
`)
|
||||
|
||||
// ClientIP extracts the caller's IP for login rate limiting. X-Forwarded-For
|
||||
// is trusted here because nginx is the only ingress and overwrites the header
|
||||
// on every proxy hop; the first value is the client address. Falls back to
|
||||
// RemoteAddr for direct connections.
|
||||
func ClientIP(r *http.Request) string {
|
||||
if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
|
||||
if first := strings.TrimSpace(strings.Split(fwd, ",")[0]); first != "" {
|
||||
return first
|
||||
}
|
||||
}
|
||||
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
||||
if err != nil {
|
||||
host = r.RemoteAddr
|
||||
}
|
||||
return host
|
||||
}
|
||||
|
||||
@@ -3,21 +3,34 @@ package identity
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/netip"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestClientIP(t *testing.T) {
|
||||
// 可信代理:环回 + RFC1918(docker-compose nginx 同网段场景)。
|
||||
trusted := []netip.Prefix{
|
||||
netip.MustParsePrefix("127.0.0.0/8"),
|
||||
netip.MustParsePrefix("10.0.0.0/8"),
|
||||
netip.MustParsePrefix("172.16.0.0/12"),
|
||||
netip.MustParsePrefix("192.168.0.0/16"),
|
||||
}
|
||||
limiter := NewLoginLimiter(nil, 0, 0, trusted)
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
remoteAddr string
|
||||
xfwd string
|
||||
want string
|
||||
}{
|
||||
{"xfwd first value", "10.0.0.1:52341", "203.0.113.9, 10.0.0.2", "203.0.113.9"},
|
||||
{"xfwd single", "10.0.0.1:52341", "198.51.100.7", "198.51.100.7"},
|
||||
{"xfwd with spaces", "10.0.0.1:52341", " 192.0.2.5 ", "192.0.2.5"},
|
||||
{"xfwd first value from trusted proxy", "10.0.0.1:52341", "203.0.113.9, 10.0.0.2", "203.0.113.9"},
|
||||
{"xfwd single from trusted proxy", "10.0.0.1:52341", "198.51.100.7", "198.51.100.7"},
|
||||
{"xfwd with spaces from trusted proxy", "10.0.0.1:52341", " 192.0.2.5 ", "192.0.2.5"},
|
||||
{"xfwd ignored from untrusted public peer", "203.0.113.9:8080", "198.51.100.7", "203.0.113.9"},
|
||||
{"xfwd ignored from CGNAT peer outside trust", "100.64.0.5:8080", "198.51.100.7", "100.64.0.5"},
|
||||
{"no xfwd falls back to remote", "203.0.113.9:8080", "", "203.0.113.9"},
|
||||
{"no xfwd and no port", "[2001:db8::1]:443", "", "2001:db8::1"},
|
||||
{"trusted peer without xfwd uses peer", "172.20.0.2:8080", "", "172.20.0.2"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
@@ -26,9 +39,20 @@ func TestClientIP(t *testing.T) {
|
||||
if tc.xfwd != "" {
|
||||
req.Header.Set("X-Forwarded-For", tc.xfwd)
|
||||
}
|
||||
if got := ClientIP(req); got != tc.want {
|
||||
if got := limiter.ClientIP(req); got != tc.want {
|
||||
t.Fatalf("ClientIP() = %q, want %q", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientIPNoTrustedProxies(t *testing.T) {
|
||||
// 未配置可信代理时(如网关端口直接暴露):任何对端的 XFF 都被忽略。
|
||||
limiter := NewLoginLimiter(nil, 0, 0, nil)
|
||||
req := httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
req.RemoteAddr = "203.0.113.9:8080"
|
||||
req.Header.Set("X-Forwarded-For", "198.51.100.7")
|
||||
if got := limiter.ClientIP(req); got != "203.0.113.9" {
|
||||
t.Fatalf("ClientIP() = %q, want %q", got, "203.0.113.9")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -56,6 +56,8 @@ type Service struct {
|
||||
oidcClient *http.Client
|
||||
samlMetadataMu sync.RWMutex
|
||||
samlMetadata map[string]samlMetadataCacheEntry
|
||||
jwksMu sync.Mutex
|
||||
jwks map[string]jwksCacheEntry
|
||||
now func() time.Time
|
||||
}
|
||||
|
||||
@@ -66,7 +68,7 @@ func (s *Service) SetIdentityProviderCipher(cipher cryptox.Cipher, allowPrivate
|
||||
}
|
||||
|
||||
func NewService(repository *Repository, sessions *SessionStore, limiter *LoginLimiter, cfg config.Auth, totpCipher cryptox.Cipher) *Service {
|
||||
return &Service{repository: repository, sessions: sessions, limiter: limiter, hasher: PasswordHasher{}, config: cfg, totpCipher: totpCipher, oidcClient: newOIDCHTTPClient(false), samlMetadata: make(map[string]samlMetadataCacheEntry), now: time.Now}
|
||||
return &Service{repository: repository, sessions: sessions, limiter: limiter, hasher: PasswordHasher{}, config: cfg, totpCipher: totpCipher, oidcClient: newOIDCHTTPClient(false), samlMetadata: make(map[string]samlMetadataCacheEntry), jwks: make(map[string]jwksCacheEntry), now: time.Now}
|
||||
}
|
||||
|
||||
// AllowLogin reports whether a login attempt from ip may proceed. When the
|
||||
@@ -75,6 +77,15 @@ func (s *Service) AllowLogin(ctx context.Context, ip string) bool {
|
||||
return s.limiter == nil || s.limiter.Allow(ctx, ip)
|
||||
}
|
||||
|
||||
// ClientIP 提取登录限流使用的客户端 IP:仅在直连对端是可信代理时采信
|
||||
// X-Forwarded-For,否则直接用对端地址,防止伪造头绕过限流。
|
||||
func (s *Service) ClientIP(r *http.Request) string {
|
||||
if s.limiter == nil {
|
||||
return peerHost(r.RemoteAddr)
|
||||
}
|
||||
return s.limiter.ClientIP(r)
|
||||
}
|
||||
|
||||
func (s *Service) Login(ctx context.Context, kind Kind, login, password string) (LoginResult, error) {
|
||||
account, err := s.findByLogin(ctx, kind, login)
|
||||
if errors.Is(err, ErrNotFound) {
|
||||
@@ -84,11 +95,11 @@ func (s *Service) Login(ctx context.Context, kind Kind, login, password string)
|
||||
if err != nil {
|
||||
return LoginResult{}, err
|
||||
}
|
||||
if !account.Active {
|
||||
return LoginResult{}, ErrAccountDisabled
|
||||
}
|
||||
if account.Locked(s.now()) {
|
||||
return LoginResult{}, LockedError{Until: *account.LockedUntil}
|
||||
if !account.Active || account.Locked(s.now()) {
|
||||
// 防枚举:停用/锁定账号与"账号不存在/口令错误"返回完全相同的
|
||||
// 错误与耗时(dummy 哈希),避免通过响应差异或计时差异探测账号状态。
|
||||
_ = s.hasher.Verify(password, dummyPasswordHash)
|
||||
return LoginResult{}, ErrInvalidCredentials
|
||||
}
|
||||
if account.PasswordHash == "" || !s.hasher.Verify(password, account.PasswordHash) {
|
||||
lockedUntil, recordErr := s.repository.RecordFailure(ctx, account, s.config.MaxFailures, s.config.LockDuration)
|
||||
@@ -130,6 +141,8 @@ func (s *Service) Login(ctx context.Context, kind Kind, login, password string)
|
||||
}
|
||||
|
||||
func (s *Service) CompleteTOTPLogin(ctx context.Context, kind Kind, tempToken, code, backupCode string) (LoginResult, error) {
|
||||
// 先只读取(不消费)挑战令牌:验证码输错时令牌保留,用户可用同一
|
||||
// 令牌重试,而不是每个笔误都强制重新走完整登录。
|
||||
principal, err := s.sessions.AuthenticatePending(ctx, tempToken, kind)
|
||||
if err != nil {
|
||||
return LoginResult{}, err
|
||||
@@ -161,6 +174,11 @@ func (s *Service) CompleteTOTPLogin(ctx context.Context, kind Kind, tempToken, c
|
||||
}
|
||||
return LoginResult{}, ErrInvalidTOTP
|
||||
}
|
||||
// 验证通过后才原子消费令牌(GetDel):并发请求用同一令牌时只有一个
|
||||
// 能铸出会话,同时避免令牌在验证失败时被白白烧掉。
|
||||
if _, err := s.sessions.ConsumePending(ctx, tempToken, kind); err != nil {
|
||||
return LoginResult{}, err
|
||||
}
|
||||
token, err := s.sessions.Create(ctx, principalFor(account))
|
||||
if err != nil {
|
||||
return LoginResult{}, err
|
||||
@@ -217,6 +235,8 @@ func (s *Service) ConfirmTOTP(ctx context.Context, account Account, code string)
|
||||
if err := s.repository.EnableTOTP(ctx, account, step, records); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// 启用 2FA 属于凭据变更:作废启用前签发的所有会话。
|
||||
s.sessions.BumpAuthVersion(ctx, account.Kind, account.ID)
|
||||
return codes, nil
|
||||
}
|
||||
|
||||
@@ -238,7 +258,12 @@ func (s *Service) DisableTOTP(ctx context.Context, account Account, password, co
|
||||
if !valid {
|
||||
return ErrInvalidTOTP
|
||||
}
|
||||
return s.repository.DisableTOTP(ctx, account)
|
||||
if err := s.repository.DisableTOTP(ctx, account); err != nil {
|
||||
return err
|
||||
}
|
||||
// 停用 2FA 属于凭据变更:作废既有会话,强制重新走完整登录。
|
||||
s.sessions.BumpAuthVersion(ctx, account.Kind, account.ID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) RegenerateBackupCodes(ctx context.Context, account Account, password, code, backupCode string) ([]string, error) {
|
||||
@@ -266,6 +291,9 @@ func (s *Service) RegenerateBackupCodes(ctx context.Context, account Account, pa
|
||||
if err := s.repository.ReplaceBackupCodes(ctx, account, records); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// 备用码重生成:旧备用码全部作废,同步作废既有会话(已失效的备用码
|
||||
// 不应继续与旧会话组合使用)。
|
||||
s.sessions.BumpAuthVersion(ctx, account.Kind, account.ID)
|
||||
return codes, nil
|
||||
}
|
||||
|
||||
@@ -340,7 +368,12 @@ func (s *Service) ChangePassword(ctx context.Context, account Account, oldPasswo
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.repository.SetPassword(ctx, account, hash)
|
||||
if err := s.repository.SetPassword(ctx, account, hash); err != nil {
|
||||
return err
|
||||
}
|
||||
// 改密后立即作废既有会话,被盗会话无法在凭据轮换后继续存活。
|
||||
s.sessions.BumpAuthVersion(ctx, account.Kind, account.ID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) findByLogin(ctx context.Context, kind Kind, login string) (Account, error) {
|
||||
|
||||
@@ -25,6 +25,41 @@ type Principal struct {
|
||||
Role string `json:"role,omitempty"`
|
||||
Purpose string `json:"purpose"`
|
||||
IssuedAt int64 `json:"issued_at"`
|
||||
// AuthVersion 是签发会话时账号的凭据版本;凭据变更(改密/2FA 变更)会
|
||||
// 递增该版本,旧版本会话在 Authenticate 时被拒绝,被盗会话无法在
|
||||
// 凭据轮换后继续存活。
|
||||
AuthVersion int64 `json:"auth_version,omitempty"`
|
||||
}
|
||||
|
||||
// authVersionTTL 必须严格大于会话 TTL(上限 7 天,由 config 校验保证):
|
||||
// 版本键过期时所有旧会话已自然过期,凭据变更后旧会话不会复活。
|
||||
const authVersionTTL = 14 * 24 * time.Hour
|
||||
|
||||
func authVersionKey(kind Kind, subjectID string) string {
|
||||
return "gateway:auth-version:" + string(kind) + ":" + subjectID
|
||||
}
|
||||
|
||||
// AuthVersion 返回账号当前凭据版本;从未变更过则为 0。
|
||||
func (s *SessionStore) AuthVersion(ctx context.Context, kind Kind, subjectID string) int64 {
|
||||
if s.client == nil {
|
||||
return 0
|
||||
}
|
||||
value, err := s.client.Get(ctx, authVersionKey(kind, subjectID)).Int64()
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
// BumpAuthVersion 使账号的全部既有会话失效(改密、2FA 启用/停用等凭据变更后调用)。
|
||||
// Redis 不可用时静默失败:会话仍按 TTL 自然过期,凭据变更的即时失效降级为延迟生效。
|
||||
func (s *SessionStore) BumpAuthVersion(ctx context.Context, kind Kind, subjectID string) {
|
||||
if s.client == nil {
|
||||
return
|
||||
}
|
||||
key := authVersionKey(kind, subjectID)
|
||||
_ = s.client.Incr(ctx, key).Err()
|
||||
_ = s.client.Expire(ctx, key, authVersionTTL).Err()
|
||||
}
|
||||
|
||||
type SessionStore struct {
|
||||
@@ -56,6 +91,7 @@ func (s *SessionStore) create(ctx context.Context, principal Principal, ttl time
|
||||
}
|
||||
token := base64.RawURLEncoding.EncodeToString(random)
|
||||
principal.IssuedAt = time.Now().Unix()
|
||||
principal.AuthVersion = s.AuthVersion(ctx, principal.Kind, principal.SubjectID)
|
||||
payload, err := json.Marshal(principal)
|
||||
if err != nil {
|
||||
return "", err
|
||||
@@ -85,18 +121,16 @@ func (s *SessionStore) Authenticate(ctx context.Context, authorization string, e
|
||||
if err := json.Unmarshal(payload, &principal); err != nil || principal.Kind != expected || principal.SubjectID == "" || principal.Purpose != "session" {
|
||||
return Principal{}, ErrInvalidSession
|
||||
}
|
||||
return principal, nil
|
||||
}
|
||||
|
||||
func (s *SessionStore) AuthenticatePending(ctx context.Context, token string, expected Kind) (Principal, error) {
|
||||
principal, err := s.authenticateToken(ctx, token)
|
||||
if err != nil || principal.Kind != expected || principal.SubjectID == "" || principal.Purpose != "totp_pending" {
|
||||
// 凭据版本不匹配:改密/2FA 变更后旧会话一律失效。
|
||||
if principal.AuthVersion != s.AuthVersion(ctx, principal.Kind, principal.SubjectID) {
|
||||
return Principal{}, ErrInvalidSession
|
||||
}
|
||||
return principal, nil
|
||||
}
|
||||
|
||||
func (s *SessionStore) authenticateToken(ctx context.Context, token string) (Principal, error) {
|
||||
// AuthenticatePending 只读取(不消费)挑战令牌,供 CompleteTOTPLogin 在
|
||||
// 验证前解析 principal;验证码输错时令牌保留可重试。
|
||||
func (s *SessionStore) AuthenticatePending(ctx context.Context, token string, expected Kind) (Principal, error) {
|
||||
if s.client == nil {
|
||||
return Principal{}, ErrUnavailable
|
||||
}
|
||||
@@ -108,7 +142,33 @@ func (s *SessionStore) authenticateToken(ctx context.Context, token string) (Pri
|
||||
return Principal{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
|
||||
}
|
||||
var principal Principal
|
||||
if err := json.Unmarshal(payload, &principal); err != nil {
|
||||
if err := json.Unmarshal(payload, &principal); err != nil || principal.Kind != expected || principal.SubjectID == "" || principal.Purpose != "totp_pending" {
|
||||
return Principal{}, ErrInvalidSession
|
||||
}
|
||||
if principal.AuthVersion != s.AuthVersion(ctx, principal.Kind, principal.SubjectID) {
|
||||
return Principal{}, ErrInvalidSession
|
||||
}
|
||||
return principal, nil
|
||||
}
|
||||
|
||||
// ConsumePending 原子消费挑战令牌(GetDel):验证通过后调用,并发请求用同一
|
||||
// 令牌时只有一个能成功,防止一次 2FA 挑战铸出两个会话。
|
||||
func (s *SessionStore) ConsumePending(ctx context.Context, token string, expected Kind) (Principal, error) {
|
||||
if s.client == nil {
|
||||
return Principal{}, ErrUnavailable
|
||||
}
|
||||
payload, err := s.client.GetDel(ctx, sessionKey(strings.TrimSpace(token))).Bytes()
|
||||
if errors.Is(err, redis.Nil) {
|
||||
return Principal{}, ErrInvalidSession
|
||||
}
|
||||
if err != nil {
|
||||
return Principal{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
|
||||
}
|
||||
var principal Principal
|
||||
if err := json.Unmarshal(payload, &principal); err != nil || principal.Kind != expected || principal.SubjectID == "" || principal.Purpose != "totp_pending" {
|
||||
return Principal{}, ErrInvalidSession
|
||||
}
|
||||
if principal.AuthVersion != s.AuthVersion(ctx, principal.Kind, principal.SubjectID) {
|
||||
return Principal{}, ErrInvalidSession
|
||||
}
|
||||
return principal, nil
|
||||
|
||||
@@ -101,7 +101,16 @@ func GenerateBackupCodes() ([]string, []BackupCodeRecord, error) {
|
||||
return nil, nil, err
|
||||
}
|
||||
for index := range random {
|
||||
random[index] = backupAlphabet[int(random[index])%len(backupAlphabet)]
|
||||
// 拒绝采样消除取模偏差:256 % 31 = 8,直接取模会让前 8 个字符
|
||||
// 的概率略高于其余字符。丢弃 248..255 的取值后分布均匀。
|
||||
value := random[index]
|
||||
for value >= 248 {
|
||||
if _, err := rand.Read(random[index : index+1]); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
value = random[index]
|
||||
}
|
||||
random[index] = backupAlphabet[int(value)%len(backupAlphabet)]
|
||||
}
|
||||
raw := string(random)
|
||||
code := raw[:4] + "-" + raw[4:]
|
||||
|
||||
@@ -82,6 +82,11 @@ func (w *Worker) runBatch(ctx context.Context) (int, error) {
|
||||
}
|
||||
|
||||
func retryDelay(attempt int, maximum time.Duration) time.Duration {
|
||||
// attempt ≥ 35 时 2^(attempt-1) 秒会溢出 int64 纳秒,得到负 duration,
|
||||
// 使 MarkFailed 把 available_at 设到过去,事件立即被重新认领形成热循环。
|
||||
if attempt > 30 {
|
||||
return maximum
|
||||
}
|
||||
seconds := math.Pow(2, float64(max(attempt-1, 0)))
|
||||
delay := time.Duration(seconds * float64(time.Second))
|
||||
if delay > maximum {
|
||||
|
||||
@@ -3,6 +3,7 @@ package config
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"os"
|
||||
"strconv"
|
||||
@@ -11,20 +12,22 @@ import (
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
Environment string
|
||||
Server Server
|
||||
Database Database
|
||||
Redis Redis
|
||||
Security Security
|
||||
Auth Auth
|
||||
Credentials Credentials
|
||||
Upstream Upstream
|
||||
Audit Audit
|
||||
Outbox Outbox
|
||||
RuntimeData RuntimeData
|
||||
Environment string
|
||||
Server Server
|
||||
Database Database
|
||||
Redis Redis
|
||||
Security Security
|
||||
Auth Auth
|
||||
Credentials Credentials
|
||||
Upstream Upstream
|
||||
Audit Audit
|
||||
Outbox Outbox
|
||||
RuntimeData RuntimeData
|
||||
Shadow Shadow
|
||||
ObjectStorage ObjectStorage
|
||||
Embeddings Embeddings
|
||||
Inbox Inbox
|
||||
Scheduler Scheduler
|
||||
}
|
||||
|
||||
type Server struct {
|
||||
@@ -58,6 +61,7 @@ type Auth struct {
|
||||
LockDuration time.Duration
|
||||
LoginRateLimitMax int // 单 IP 滑动窗口内的最大登录尝试次数
|
||||
LoginRateLimitWindow time.Duration // 登录限流滑动窗口
|
||||
TrustedProxies []netip.Prefix // 可信反向代理网段;仅来自这些对端的 X-Forwarded-For 被采信
|
||||
}
|
||||
|
||||
type Credentials struct {
|
||||
@@ -87,6 +91,7 @@ type Audit struct {
|
||||
FlushInterval time.Duration
|
||||
Retention time.Duration
|
||||
UsageRetention time.Duration
|
||||
TraceRetention time.Duration
|
||||
PartitionMonthsAhead int
|
||||
MaintenanceInterval time.Duration
|
||||
}
|
||||
@@ -137,6 +142,20 @@ type Embeddings struct {
|
||||
Timeout time.Duration
|
||||
}
|
||||
|
||||
// Inbox 配置站内消息(M8 P4)。Channel 是通知 worker 落库后 PUBLISH 的 Redis 频道,
|
||||
// 供未来实时推送订阅;未读数以 PostgreSQL 为权威源,不依赖 Redis。
|
||||
type Inbox struct {
|
||||
Channel string
|
||||
}
|
||||
|
||||
type Scheduler struct {
|
||||
GatewayBaseURL string
|
||||
PollInterval time.Duration
|
||||
ExecutionTimeout time.Duration
|
||||
BatchSize int
|
||||
MaxAttempts int
|
||||
}
|
||||
|
||||
func Load() (Config, error) {
|
||||
cfg := Config{
|
||||
Environment: env("APP_ENV", "local"),
|
||||
@@ -167,6 +186,7 @@ func Load() (Config, error) {
|
||||
LockDuration: duration("LOGIN_LOCK_DURATION", 15*time.Minute),
|
||||
LoginRateLimitMax: intValue("LOGIN_RATE_LIMIT_MAX", 30),
|
||||
LoginRateLimitWindow: duration("LOGIN_RATE_LIMIT_WINDOW", 5*time.Minute),
|
||||
TrustedProxies: parsePrefixList(env("TRUSTED_PROXIES", "127.0.0.0/8,::1,10.0.0.0/8,172.16.0.0/12,192.168.0.0/16,fc00::/7")),
|
||||
},
|
||||
Credentials: Credentials{
|
||||
MasterKey: strings.TrimSpace(os.Getenv("CREDENTIAL_MASTER_KEY")),
|
||||
@@ -186,7 +206,8 @@ func Load() (Config, error) {
|
||||
Audit: Audit{
|
||||
QueueSize: intValue("AUDIT_QUEUE_SIZE", 4096), BatchSize: intValue("AUDIT_BATCH_SIZE", 200),
|
||||
FlushInterval: duration("AUDIT_FLUSH_INTERVAL", time.Second), Retention: duration("AUDIT_RETENTION", 90*24*time.Hour),
|
||||
UsageRetention: duration("USAGE_RETENTION", 730*24*time.Hour), PartitionMonthsAhead: intValue("AUDIT_PARTITION_MONTHS_AHEAD", 3),
|
||||
UsageRetention: duration("USAGE_RETENTION", 730*24*time.Hour), TraceRetention: duration("TRACE_RETENTION", 90*24*time.Hour),
|
||||
PartitionMonthsAhead: intValue("AUDIT_PARTITION_MONTHS_AHEAD", 3),
|
||||
MaintenanceInterval: duration("AUDIT_MAINTENANCE_INTERVAL", 6*time.Hour),
|
||||
},
|
||||
Outbox: Outbox{
|
||||
@@ -222,6 +243,16 @@ func Load() (Config, error) {
|
||||
BatchSize: intValue("EMBEDDING_BATCH_SIZE", 64),
|
||||
Timeout: duration("EMBEDDING_TIMEOUT", 120*time.Second),
|
||||
},
|
||||
Inbox: Inbox{
|
||||
Channel: env("INBOX_CHANNEL", "gateway:inbox:events"),
|
||||
},
|
||||
Scheduler: Scheduler{
|
||||
GatewayBaseURL: strings.TrimRight(env("SCHEDULER_GATEWAY_BASE_URL", "http://gateway-api:8080"), "/"),
|
||||
PollInterval: duration("SCHEDULER_POLL_INTERVAL", 5*time.Second),
|
||||
ExecutionTimeout: duration("SCHEDULER_EXECUTION_TIMEOUT", 5*time.Minute),
|
||||
BatchSize: intValue("SCHEDULER_BATCH_SIZE", 10),
|
||||
MaxAttempts: intValue("SCHEDULER_MAX_ATTEMPTS", 3),
|
||||
},
|
||||
}
|
||||
|
||||
return cfg, cfg.Validate()
|
||||
@@ -235,7 +266,7 @@ func (c Config) Validate() error {
|
||||
if c.Database.MinConns < 0 || c.Database.MaxConns < 1 || c.Database.MinConns > c.Database.MaxConns {
|
||||
errs = append(errs, errors.New("database pool sizes are invalid"))
|
||||
}
|
||||
if c.Auth.SessionTTL < 5*time.Minute || c.Auth.TOTPChallengeTTL < time.Minute || c.Auth.TOTPChallengeTTL > 15*time.Minute || c.Auth.MaxFailures < 1 || c.Auth.LockDuration < time.Minute {
|
||||
if c.Auth.SessionTTL < 5*time.Minute || c.Auth.SessionTTL > 7*24*time.Hour || c.Auth.TOTPChallengeTTL < time.Minute || c.Auth.TOTPChallengeTTL > 15*time.Minute || c.Auth.MaxFailures < 1 || c.Auth.LockDuration < time.Minute {
|
||||
errs = append(errs, errors.New("authentication limits are invalid"))
|
||||
}
|
||||
if c.Auth.LoginRateLimitMax < 1 || c.Auth.LoginRateLimitWindow < time.Second {
|
||||
@@ -256,7 +287,7 @@ func (c Config) Validate() error {
|
||||
if c.Audit.QueueSize < 100 || c.Audit.QueueSize > 1_000_000 || c.Audit.BatchSize < 1 || c.Audit.BatchSize > c.Audit.QueueSize || c.Audit.FlushInterval < 100*time.Millisecond || c.Audit.FlushInterval > time.Minute {
|
||||
errs = append(errs, errors.New("audit buffering settings are invalid"))
|
||||
}
|
||||
if c.Audit.Retention < 24*time.Hour || c.Audit.Retention > 10*365*24*time.Hour || c.Audit.UsageRetention < c.Audit.Retention || c.Audit.UsageRetention > 10*365*24*time.Hour || c.Audit.PartitionMonthsAhead < 1 || c.Audit.PartitionMonthsAhead > 24 || c.Audit.MaintenanceInterval < time.Hour || c.Audit.MaintenanceInterval > 7*24*time.Hour {
|
||||
if c.Audit.Retention < 24*time.Hour || c.Audit.Retention > 10*365*24*time.Hour || c.Audit.UsageRetention < c.Audit.Retention || c.Audit.UsageRetention > 10*365*24*time.Hour || c.Audit.TraceRetention < 24*time.Hour || c.Audit.TraceRetention > 10*365*24*time.Hour || c.Audit.PartitionMonthsAhead < 1 || c.Audit.PartitionMonthsAhead > 24 || c.Audit.MaintenanceInterval < time.Hour || c.Audit.MaintenanceInterval > 7*24*time.Hour {
|
||||
errs = append(errs, errors.New("audit retention settings are invalid"))
|
||||
}
|
||||
if !strings.Contains(c.Outbox.Stream, "{outbox}") || c.Outbox.BatchSize < 1 || c.Outbox.BatchSize > 1000 || c.Outbox.PollInterval < 50*time.Millisecond || c.Outbox.PollInterval > time.Minute || c.Outbox.Lease < 5*time.Second || c.Outbox.Lease > 10*time.Minute || c.Outbox.MaxAttempts < 1 || c.Outbox.MaxAttempts > 100 || c.Outbox.MaxBackoff < time.Second || c.Outbox.MaxBackoff > time.Hour || c.Outbox.StreamMaxLength < 1000 || c.Outbox.StreamMaxLength > 100_000_000 || c.Outbox.MarkerTTL < 24*time.Hour || c.Outbox.MarkerTTL > 365*24*time.Hour {
|
||||
@@ -312,6 +343,12 @@ func (c Config) Validate() error {
|
||||
errs = append(errs, errors.New("EMBEDDING_TIMEOUT must be between 1s and 30m"))
|
||||
}
|
||||
}
|
||||
if err := validateHTTPURL(c.Scheduler.GatewayBaseURL); err != nil {
|
||||
errs = append(errs, fmt.Errorf("SCHEDULER_GATEWAY_BASE_URL: %w", err))
|
||||
}
|
||||
if c.Scheduler.PollInterval < time.Second || c.Scheduler.PollInterval > time.Minute || c.Scheduler.ExecutionTimeout < time.Minute || c.Scheduler.ExecutionTimeout > time.Hour || c.Scheduler.BatchSize < 1 || c.Scheduler.BatchSize > 100 || c.Scheduler.MaxAttempts < 1 || c.Scheduler.MaxAttempts > 10 {
|
||||
errs = append(errs, errors.New("scheduler settings are invalid"))
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
@@ -342,6 +379,15 @@ func (c Config) ValidateRuntime() error {
|
||||
errs = append(errs, errors.New("CREDENTIAL_MASTER_KEY is required in production"))
|
||||
}
|
||||
}
|
||||
// 拒绝已知弱默认密钥(所有环境,含本地 compose):deploy/docker-compose.yml
|
||||
// 曾把全零密钥作为默认值,凡是用该值加密的 Provider 凭据/TOTP 密钥/
|
||||
// Webhook 签名密钥,任何拿到仓库的人都能解密。
|
||||
if c.Credentials.MasterKey != "" {
|
||||
switch strings.ToLower(strings.TrimSpace(c.Credentials.MasterKey)) {
|
||||
case "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa=", "change-me", "changeme", "password", "secret":
|
||||
errs = append(errs, errors.New("CREDENTIAL_MASTER_KEY is set to a known weak default; generate a strong random key with: openssl rand -base64 32"))
|
||||
}
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
@@ -356,6 +402,27 @@ func validateHTTPURL(raw string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// parsePrefixList 解析逗号分隔的 IP/CIDR 列表;非法项跳过并返回 nil 表示不信任任何代理。
|
||||
func parsePrefixList(raw string) []netip.Prefix {
|
||||
var prefixes []netip.Prefix
|
||||
for _, part := range strings.Split(raw, ",") {
|
||||
part = strings.TrimSpace(part)
|
||||
if part == "" {
|
||||
continue
|
||||
}
|
||||
prefix, err := netip.ParsePrefix(part)
|
||||
if err != nil {
|
||||
if addr, addrErr := netip.ParseAddr(part); addrErr == nil {
|
||||
prefix = netip.PrefixFrom(addr, addr.BitLen())
|
||||
} else {
|
||||
continue
|
||||
}
|
||||
}
|
||||
prefixes = append(prefixes, prefix.Masked())
|
||||
}
|
||||
return prefixes
|
||||
}
|
||||
|
||||
func env(key, fallback string) string {
|
||||
if value := strings.TrimSpace(os.Getenv(key)); value != "" {
|
||||
return value
|
||||
|
||||
@@ -33,7 +33,7 @@ func TestProductionCanDisableBootstrapCompatibility(t *testing.T) {
|
||||
t.Setenv("GATEWAY_BOOTSTRAP_API_KEY", "")
|
||||
t.Setenv("GATEWAY_BOOTSTRAP_API_KEY_ENABLED", "false")
|
||||
t.Setenv("UPSTREAM_FALLBACK_ENABLED", "false")
|
||||
t.Setenv("CREDENTIAL_MASTER_KEY", "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=")
|
||||
t.Setenv("CREDENTIAL_MASTER_KEY", "yP7sK9xR2mV4nQ8wT1uB3cE5fG6hJ0kL=")
|
||||
t.Setenv("UPSTREAM_BASE_URL", "https://example.com")
|
||||
cfg, err := Load()
|
||||
if err != nil {
|
||||
|
||||
@@ -136,7 +136,7 @@ func (s *Service) appendMessage(ctx context.Context, conversationID, role, conte
|
||||
return ConversationMessage{Sequence: sequence, Role: role, Content: content, CreatedAt: created}, nil
|
||||
}
|
||||
|
||||
func (s *Service) callApplication(ctx context.Context, appCode, secret string, messages []ConversationMessage, variables map[string]any) (map[string]any, string, error) {
|
||||
func (s *Service) callApplication(ctx context.Context, appCode, secret string, messages []ConversationMessage, variables map[string]any, conversationID string) (map[string]any, string, error) {
|
||||
payloadMessages := make([]map[string]any, 0, len(messages))
|
||||
for _, m := range messages {
|
||||
payloadMessages = append(payloadMessages, map[string]any{"role": m.Role, "content": m.Content})
|
||||
@@ -145,6 +145,9 @@ func (s *Service) callApplication(ctx context.Context, appCode, secret string, m
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/applications/"+appCode+"/chat/completions", bytes.NewReader(payload)).WithContext(gateway.WithRequestID(ctx, "portal-"+time.Now().UTC().Format("20060102150405.000000000")))
|
||||
request.Header.Set("Authorization", "Bearer "+secret)
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
if strings.TrimSpace(conversationID) != "" {
|
||||
request.Header.Set("X-Gateway-Conversation-ID", conversationID)
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
s.runtime.ServeHTTP(recorder, request)
|
||||
var response map[string]any
|
||||
@@ -180,7 +183,7 @@ func (s *Service) Chat(ctx context.Context, account identity.Account, code, mess
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
response, _, err := s.callApplication(ctx, app.Code, secret, []ConversationMessage{{Role: "user", Content: message}}, variables)
|
||||
response, _, err := s.callApplication(ctx, app.Code, secret, []ConversationMessage{{Role: "user", Content: message}}, variables, "")
|
||||
return response, err
|
||||
}
|
||||
|
||||
@@ -223,7 +226,7 @@ func (s *Service) AppendConversationMessage(ctx context.Context, account identit
|
||||
history := make([]ConversationMessage, len(conversation.Messages)+1)
|
||||
copy(history, conversation.Messages)
|
||||
history[len(conversation.Messages)] = ConversationMessage{Role: "user", Content: message}
|
||||
response, answer, err := s.callApplication(ctx, code, secret, history, variables)
|
||||
response, answer, err := s.callApplication(ctx, code, secret, history, variables, id)
|
||||
if err != nil {
|
||||
return response, err
|
||||
}
|
||||
|
||||
@@ -88,6 +88,13 @@ func (s *Service) Reload(ctx context.Context) error {
|
||||
if exactI != exactJ {
|
||||
return exactI
|
||||
}
|
||||
// 通配符之间按前缀长度降序:更具体的模式(gpt-4o*)必须先于宽泛模式
|
||||
// (gpt-4*)命中,否则 gpt-4o-mini 会按 gpt-4* 的价格错误计费。
|
||||
// 同长度再按生效时间(新价格优先)。
|
||||
lengthI, lengthJ := len(active[i].ModelPattern), len(active[j].ModelPattern)
|
||||
if !exactI && lengthI != lengthJ {
|
||||
return lengthI > lengthJ
|
||||
}
|
||||
return active[i].EffectiveFrom.After(active[j].EffectiveFrom)
|
||||
})
|
||||
s.current.Store(&priceSnapshot{prices: active})
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
@@ -23,6 +24,7 @@ type AdminHTTPHandler struct {
|
||||
changeHook func(context.Context) error
|
||||
operations AdminOperations
|
||||
mux *http.ServeMux
|
||||
logger *slog.Logger
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) SetChangeHook(hook func(context.Context) error) {
|
||||
@@ -63,11 +65,12 @@ type modelRouteConditions struct {
|
||||
func NewAdminHTTPHandler(repository *Repository, cipher *CredentialCipher, identityService *identity.Service, allowPrivate bool) *AdminHTTPHandler {
|
||||
handler := &AdminHTTPHandler{
|
||||
repository: repository, cipher: cipher, identity: identityService,
|
||||
allowPrivate: allowPrivate, mux: http.NewServeMux(),
|
||||
allowPrivate: allowPrivate, mux: http.NewServeMux(), logger: slog.Default(),
|
||||
}
|
||||
handler.mux.HandleFunc("GET /api/v1/admin/providers", handler.list)
|
||||
handler.mux.HandleFunc("POST /api/v1/admin/providers", handler.create)
|
||||
handler.mux.HandleFunc("PUT /api/v1/admin/providers/{provider_id}", handler.update)
|
||||
handler.mux.HandleFunc("DELETE /api/v1/admin/providers/{provider_id}", handler.delete)
|
||||
handler.mux.HandleFunc("POST /api/v1/admin/providers/{provider_id}/test", handler.testConnection)
|
||||
handler.mux.HandleFunc("GET /api/v1/admin/providers/{provider_id}/models", handler.listModels)
|
||||
handler.mux.HandleFunc("POST /api/v1/admin/providers/{provider_id}/models/sync", handler.syncModels)
|
||||
@@ -138,6 +141,19 @@ func (h *AdminHTTPHandler) create(writer http.ResponseWriter, request *http.Requ
|
||||
apiresponse.OK(writer, view)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) delete(writer http.ResponseWriter, request *http.Request) {
|
||||
actor, ok := h.requirePermission(writer, request, identity.PermissionProviderManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := h.repository.Delete(request.Context(), request.PathValue("provider_id"), actor.ID); err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
h.propagateChange(request.Context())
|
||||
apiresponse.OK(writer, map[string]bool{"deleted": true})
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) update(writer http.ResponseWriter, request *http.Request) {
|
||||
actor, ok := h.requirePermission(writer, request, identity.PermissionProviderManage)
|
||||
if !ok {
|
||||
@@ -244,20 +260,30 @@ func (h *AdminHTTPHandler) setCredentials(record *Record, apiKey string) error {
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) view(record Record) (map[string]any, error) {
|
||||
// 凭据解密失败(KEK 轮换后旧记录、数据损坏)不得让整个列表接口报错:
|
||||
// 降级为"未确认"状态并附警示,管理端仍可编辑/删除该记录恢复。
|
||||
keyConfigured := false
|
||||
masked := ""
|
||||
credentialError := ""
|
||||
plaintext, err := h.cipher.Decrypt(record.EncryptedCredentials, record.CredentialKEKVersion)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var credentials Credentials
|
||||
if err := json.Unmarshal(plaintext, &credentials); err != nil {
|
||||
return nil, err
|
||||
credentialError = "凭据无法解密(加密密钥不匹配或数据损坏),请重新保存凭据"
|
||||
} else {
|
||||
var credentials Credentials
|
||||
if json.Unmarshal(plaintext, &credentials) != nil {
|
||||
credentialError = "凭据数据格式无效,请重新保存凭据"
|
||||
} else {
|
||||
keyConfigured = credentials.APIKey != ""
|
||||
masked = maskSecret(credentials.APIKey)
|
||||
}
|
||||
}
|
||||
return map[string]any{
|
||||
"id": record.ID, "code": record.Code, "adapter": record.Adapter,
|
||||
"base_url": record.BaseURL, "capabilities": record.Capabilities,
|
||||
"config": record.Config, "enabled": record.Enabled, "revision": record.Revision,
|
||||
"credential_kek_version": record.CredentialKEKVersion,
|
||||
"key_configured": credentials.APIKey != "", "api_key_masked": maskSecret(credentials.APIKey),
|
||||
"key_configured": keyConfigured, "api_key_masked": masked,
|
||||
"credential_error": credentialError,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -482,6 +508,11 @@ func (h *AdminHTTPHandler) writeError(writer http.ResponseWriter, err error) {
|
||||
apiresponse.Error(writer, http.StatusServiceUnavailable, "供应商配置服务暂不可用")
|
||||
case errors.Is(err, ErrProviderUpstream):
|
||||
apiresponse.Error(writer, http.StatusBadGateway, "无法从上游供应商获取模型信息")
|
||||
case errors.Is(err, ErrBlockedAddress):
|
||||
// 原始错误含解析出的地址(如 "blocked address 10.0.0.1"),泄露内网
|
||||
// 拓扑;细节只进服务端日志,客户端返回通用提示。
|
||||
h.logger.Warn("provider URL blocked", "error", err)
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "供应商地址不允许访问内网或保留网段")
|
||||
default:
|
||||
apiresponse.Error(writer, http.StatusBadRequest, err.Error())
|
||||
}
|
||||
|
||||
@@ -128,7 +128,10 @@ func (s *Service) RotateCredentials(ctx context.Context, actorID string) (provid
|
||||
}
|
||||
plaintext, err := s.cipher.Decrypt(record.EncryptedCredentials, record.CredentialKEKVersion)
|
||||
if err != nil {
|
||||
return provider.CredentialRotationResult{}, fmt.Errorf("provider %s credentials cannot be decrypted: %w", record.Code, err)
|
||||
// 单条损坏(如 KEK 版本被删除)不阻塞其余 Provider 的轮换:
|
||||
// 跳过并计数,管理端从 Skipped 明细中定位问题记录。
|
||||
result.Skipped++
|
||||
continue
|
||||
}
|
||||
ciphertext, version, err := s.cipher.Encrypt(plaintext)
|
||||
if err != nil {
|
||||
@@ -229,34 +232,8 @@ func hasCapability(capabilities []string, expected string) bool {
|
||||
}
|
||||
|
||||
func safeDialContext(allowPrivate bool) func(context.Context, string, string) (net.Conn, error) {
|
||||
dialer := &net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}
|
||||
if allowPrivate {
|
||||
return dialer.DialContext
|
||||
}
|
||||
return func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
host, port, err := net.SplitHostPort(address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(addresses) == 0 {
|
||||
return nil, errors.New("upstream host did not resolve")
|
||||
}
|
||||
for _, address := range addresses {
|
||||
if !isPublicAddress(address.IP) {
|
||||
return nil, fmt.Errorf("upstream resolved to blocked address %s", address.IP)
|
||||
}
|
||||
}
|
||||
return dialer.DialContext(ctx, network, net.JoinHostPort(addresses[0].IP.String(), port))
|
||||
}
|
||||
}
|
||||
|
||||
func isPublicAddress(ip net.IP) bool {
|
||||
return ip != nil && !ip.IsPrivate() && !ip.IsLoopback() && !ip.IsLinkLocalUnicast() &&
|
||||
!ip.IsLinkLocalMulticast() && !ip.IsMulticast() && !ip.IsUnspecified()
|
||||
// 统一复用 provider.IsPublicAddress 的完整网段判定(含 CGNAT/6to4/NAT64 等)。
|
||||
return provider.SafeDialContext(allowPrivate, 5*time.Second, 30*time.Second)
|
||||
}
|
||||
|
||||
var _ provider.AdminOperations = (*Service)(nil)
|
||||
|
||||
@@ -39,6 +39,9 @@ func (a *Adapter) Capabilities() []provider.Capability {
|
||||
}
|
||||
|
||||
func (a *Adapter) Prepare(request *http.Request) {
|
||||
// 无条件剥离客户端凭据:客户端携带的网关 API Key 绝不能转发给上游。
|
||||
// 只有配置了 Provider 自身凭据时才注入 Authorization。
|
||||
request.Header.Del("Authorization")
|
||||
request.Header.Del("X-Gateway-API-Key")
|
||||
request.Header.Del("X-Gateway-Provider")
|
||||
basePath := strings.TrimRight(a.target.Path, "/")
|
||||
|
||||
@@ -250,6 +250,39 @@ func (r *Repository) Get(ctx context.Context, id string) (Record, error) {
|
||||
return record, mapProviderError(err)
|
||||
}
|
||||
|
||||
// Delete removes a provider and its cascaded model routes / synced models
|
||||
// (FKs are ON DELETE CASCADE), emitting a provider.deleted outbox event in the
|
||||
// same transaction.
|
||||
func (r *Repository) Delete(ctx context.Context, id, actorID string) error {
|
||||
if r.pool == nil {
|
||||
return ErrProviderStore
|
||||
}
|
||||
eventID, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tx, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
result, err := tx.Exec(ctx, `DELETE FROM gateway.providers WHERE id=$1 AND tenant_id IS NULL`, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
if result.RowsAffected() == 0 {
|
||||
return ErrProviderNotFound
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{"provider_id": id, "actor_id": actorID})
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway.outbox_events
|
||||
(event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
|
||||
VALUES ($1, 'provider.deleted', 1, 'provider', $2, $3)`, eventID, id, payload); err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
func (r *Repository) Create(ctx context.Context, record Record, actorID string) (Record, error) {
|
||||
if r.pool == nil {
|
||||
return Record{}, ErrProviderStore
|
||||
|
||||
@@ -107,13 +107,27 @@ func (r *Resolver) ResolveModelRoute(query gateway.ModelRouteQuery) (gateway.Mod
|
||||
}
|
||||
total := 0
|
||||
for _, route := range matched {
|
||||
// 防御:weight<=0 的行(绕过管理端校验直接入库的脏数据)不得参与
|
||||
// 权重池,否则 total=0 时取模除零 panic,同一请求 ID 将永远 500。
|
||||
if route.weight <= 0 {
|
||||
continue
|
||||
}
|
||||
total += route.weight
|
||||
}
|
||||
weighted := make([]compiledRoute, 0, len(matched))
|
||||
for _, route := range matched {
|
||||
if route.weight > 0 {
|
||||
weighted = append(weighted, route)
|
||||
}
|
||||
}
|
||||
if len(weighted) == 0 {
|
||||
return gateway.ModelRouteResult{Known: known}, nil
|
||||
}
|
||||
hasher := fnv.New64a()
|
||||
_, _ = hasher.Write([]byte(query.Seed + "\x00" + query.Model))
|
||||
selected := int(hasher.Sum64() % uint64(total))
|
||||
chosen := matched[len(matched)-1]
|
||||
for _, route := range matched {
|
||||
chosen := weighted[len(weighted)-1]
|
||||
for _, route := range weighted {
|
||||
if selected < route.weight {
|
||||
chosen = route
|
||||
break
|
||||
|
||||
@@ -5,10 +5,16 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ErrBlockedAddress 标记 base_url 解析到被禁止的网段(私网/特殊用途网段)。
|
||||
// 该错误携带解析出的地址细节,只应记录在服务端日志,不得原样返回给客户端。
|
||||
var ErrBlockedAddress = errors.New("base_url resolves to a blocked address")
|
||||
|
||||
func ValidateBaseURL(ctx context.Context, raw string, allowPrivate bool) (string, error) {
|
||||
parsed, err := url.Parse(strings.TrimSpace(raw))
|
||||
if err != nil {
|
||||
@@ -29,8 +35,8 @@ func ValidateBaseURL(ctx context.Context, raw string, allowPrivate bool) (string
|
||||
return "", errors.New("base_url host did not resolve")
|
||||
}
|
||||
for _, address := range addresses {
|
||||
if !isPublicAddress(address.IP) {
|
||||
return "", fmt.Errorf("base_url resolves to blocked address %s", address.IP)
|
||||
if !IsPublicAddress(address.IP) {
|
||||
return "", fmt.Errorf("%w %s", ErrBlockedAddress, address.IP)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -38,7 +44,78 @@ func ValidateBaseURL(ctx context.Context, raw string, allowPrivate bool) (string
|
||||
return parsed.String(), nil
|
||||
}
|
||||
|
||||
func isPublicAddress(ip net.IP) bool {
|
||||
return ip != nil && !ip.IsPrivate() && !ip.IsLoopback() && !ip.IsLinkLocalUnicast() &&
|
||||
!ip.IsLinkLocalMulticast() && !ip.IsMulticast() && !ip.IsUnspecified()
|
||||
// specialPurposePrefixes 是 Go netip 内建分类(loopback/private/link-local/
|
||||
// multicast/unspecified)之外、但绝不应作为出站上游的 IANA 特殊用途网段。
|
||||
// 内网服务常部署在 CGNAT(100.64/10)与 benchmark(198.18/15)段,而 6to4/NAT64
|
||||
// 前缀可以把 IPv6 地址桥接回内网 IPv4,因此必须一并拦截。
|
||||
var specialPurposePrefixes = []netip.Prefix{
|
||||
// IPv4 特殊用途网段(RFC 6890 及其更新)。
|
||||
netip.MustParsePrefix("100.64.0.0/10"), // CGNAT 共享地址空间 RFC 6598
|
||||
netip.MustParsePrefix("192.0.0.0/24"), // IETF 协议保留
|
||||
netip.MustParsePrefix("192.0.2.0/24"), // TEST-NET-1 文档
|
||||
netip.MustParsePrefix("192.88.99.0/24"), // 6to4 中继任播(已弃用)
|
||||
netip.MustParsePrefix("198.18.0.0/15"), // 基准测试 RFC 2544
|
||||
netip.MustParsePrefix("198.51.100.0/24"), // TEST-NET-2 文档
|
||||
netip.MustParsePrefix("203.0.113.0/24"), // TEST-NET-3 文档
|
||||
netip.MustParsePrefix("240.0.0.0/4"), // 保留(含广播地址)
|
||||
// IPv6 特殊用途网段。
|
||||
netip.MustParsePrefix("2001:db8::/32"), // 文档地址
|
||||
netip.MustParsePrefix("2001:10::/28"), // ORCHID
|
||||
netip.MustParsePrefix("2002::/16"), // 6to4:内嵌 IPv4,可桥接回内网
|
||||
netip.MustParsePrefix("64:ff9b::/96"), // NAT64 知名前缀
|
||||
netip.MustParsePrefix("64:ff9b:1::/48"), // NAT64 本地使用前缀
|
||||
}
|
||||
|
||||
// IsPublicAddress 报告 ip 是否为可安全访问的公网单播地址。IPv4-mapped
|
||||
// IPv6(::ffff:a.b.c.d)先解映射为 IPv4 再判断,防止绕过。SSRF 防护统一使用
|
||||
// 本函数,写入校验与拨号时校验共用同一份判定。
|
||||
func IsPublicAddress(ip net.IP) bool {
|
||||
if ip == nil {
|
||||
return false
|
||||
}
|
||||
addr, ok := netip.AddrFromSlice(ip)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
addr = addr.Unmap()
|
||||
if !addr.IsValid() || addr.IsUnspecified() || addr.IsLoopback() || addr.IsMulticast() ||
|
||||
addr.IsLinkLocalUnicast() || addr.IsLinkLocalMulticast() || addr.IsInterfaceLocalMulticast() ||
|
||||
addr.IsPrivate() {
|
||||
return false
|
||||
}
|
||||
for _, prefix := range specialPurposePrefixes {
|
||||
if prefix.Contains(addr) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// SafeDialContext 构造拨号函数:allowPrivate 为 false 时,在拨号前对解析出的
|
||||
// 全部地址做 IsPublicAddress 校验,并按校验通过的地址直连(不再二次解析,
|
||||
// 缓解 DNS rebinding)。gateway 数据平面与控制面客户端共用此实现。
|
||||
func SafeDialContext(allowPrivate bool, timeout, keepAlive time.Duration) func(context.Context, string, string) (net.Conn, error) {
|
||||
dialer := &net.Dialer{Timeout: timeout, KeepAlive: keepAlive}
|
||||
if allowPrivate {
|
||||
return dialer.DialContext
|
||||
}
|
||||
return func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
host, port, err := net.SplitHostPort(address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(addresses) == 0 {
|
||||
return nil, errors.New("upstream host did not resolve")
|
||||
}
|
||||
for _, address := range addresses {
|
||||
if !IsPublicAddress(address.IP) {
|
||||
return nil, fmt.Errorf("upstream resolved to blocked address %s", address.IP)
|
||||
}
|
||||
}
|
||||
return dialer.DialContext(ctx, network, net.JoinHostPort(addresses[0].IP.String(), port))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Cron implements the standard five-field cron format: minute, hour,
|
||||
// day-of-month, month, day-of-week. Lists, ranges and steps are supported.
|
||||
type Cron struct {
|
||||
minute, hour, day, month, weekday field
|
||||
dayWildcard, weekdayWildcard bool
|
||||
}
|
||||
|
||||
type field struct {
|
||||
min, max int
|
||||
values map[int]bool
|
||||
}
|
||||
|
||||
func ParseCron(expression string) (Cron, error) {
|
||||
parts := strings.Fields(expression)
|
||||
if len(parts) != 5 {
|
||||
return Cron{}, errors.New("cron 表达式必须包含 5 段: 分 时 日 月 周")
|
||||
}
|
||||
definitions := [5][2]int{{0, 59}, {0, 23}, {1, 31}, {1, 12}, {0, 7}}
|
||||
fields := make([]field, 5)
|
||||
for i, part := range parts {
|
||||
parsed, err := parseField(part, definitions[i][0], definitions[i][1], i == 4)
|
||||
if err != nil {
|
||||
return Cron{}, fmt.Errorf("cron 第 %d 段无效: %w", i+1, err)
|
||||
}
|
||||
fields[i] = parsed
|
||||
}
|
||||
return Cron{minute: fields[0], hour: fields[1], day: fields[2], month: fields[3], weekday: fields[4], dayWildcard: parts[2] == "*", weekdayWildcard: parts[4] == "*"}, nil
|
||||
}
|
||||
|
||||
func parseField(raw string, minimum, maximum int, weekday bool) (field, error) {
|
||||
result := field{min: minimum, max: maximum, values: map[int]bool{}}
|
||||
for _, item := range strings.Split(raw, ",") {
|
||||
item = strings.TrimSpace(item)
|
||||
if item == "" {
|
||||
return field{}, errors.New("存在空列表项")
|
||||
}
|
||||
base, stepText, hasStep := strings.Cut(item, "/")
|
||||
step := 1
|
||||
if hasStep {
|
||||
var err error
|
||||
step, err = strconv.Atoi(stepText)
|
||||
if err != nil || step < 1 || step > maximum-minimum+1 {
|
||||
return field{}, errors.New("步长无效")
|
||||
}
|
||||
}
|
||||
start, end := minimum, maximum
|
||||
switch {
|
||||
case base == "*":
|
||||
case strings.Contains(base, "-"):
|
||||
left, right, _ := strings.Cut(base, "-")
|
||||
var err error
|
||||
start, err = cronNumber(left, minimum, maximum, weekday)
|
||||
if err != nil {
|
||||
return field{}, err
|
||||
}
|
||||
end, err = cronNumber(right, minimum, maximum, weekday)
|
||||
if err != nil || start > end {
|
||||
return field{}, errors.New("范围无效")
|
||||
}
|
||||
default:
|
||||
var err error
|
||||
start, err = cronNumber(base, minimum, maximum, weekday)
|
||||
if err != nil {
|
||||
return field{}, err
|
||||
}
|
||||
end = start
|
||||
if hasStep {
|
||||
end = maximum
|
||||
}
|
||||
}
|
||||
for value := start; value <= end; value += step {
|
||||
if weekday && value == 7 {
|
||||
value = 0
|
||||
result.values[value] = true
|
||||
break
|
||||
}
|
||||
result.values[value] = true
|
||||
}
|
||||
}
|
||||
if len(result.values) == 0 {
|
||||
return field{}, errors.New("没有可用取值")
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func cronNumber(raw string, minimum, maximum int, weekday bool) (int, error) {
|
||||
value, err := strconv.Atoi(raw)
|
||||
if err != nil || value < minimum || value > maximum {
|
||||
return 0, fmt.Errorf("%q 超出 %d-%d", raw, minimum, maximum)
|
||||
}
|
||||
if weekday && value == 7 {
|
||||
return 7, nil
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func (c Cron) Matches(value time.Time) bool {
|
||||
dayMatch := c.day.values[value.Day()]
|
||||
weekdayMatch := c.weekday.values[int(value.Weekday())]
|
||||
calendarMatch := dayMatch && weekdayMatch
|
||||
// Vixie cron semantics: when both day fields are restricted, either may match.
|
||||
if !c.dayWildcard && !c.weekdayWildcard {
|
||||
calendarMatch = dayMatch || weekdayMatch
|
||||
}
|
||||
return c.minute.values[value.Minute()] && c.hour.values[value.Hour()] && c.month.values[int(value.Month())] && calendarMatch
|
||||
}
|
||||
|
||||
func (c Cron) Next(after time.Time, location *time.Location) (time.Time, error) {
|
||||
if location == nil {
|
||||
location = time.UTC
|
||||
}
|
||||
candidate := after.UTC().Truncate(time.Minute).Add(time.Minute)
|
||||
deadline := candidate.AddDate(5, 0, 0)
|
||||
for candidate.Before(deadline) {
|
||||
if c.Matches(candidate.In(location)) {
|
||||
return candidate, nil
|
||||
}
|
||||
candidate = candidate.Add(time.Minute)
|
||||
}
|
||||
return time.Time{}, errors.New("未来 5 年内无匹配执行时间")
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCronNext(t *testing.T) {
|
||||
cases := []struct {
|
||||
expression, after, want string
|
||||
}{
|
||||
{"*/15 * * * *", "2026-08-12T10:07:00Z", "2026-08-12T10:15:00Z"},
|
||||
{"0 9 * * 1-5", "2026-08-14T09:01:00Z", "2026-08-17T09:00:00Z"},
|
||||
{"30 8 1 * *", "2026-08-12T00:00:00Z", "2026-09-01T08:30:00Z"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
schedule, err := ParseCron(tc.expression)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
after, _ := time.Parse(time.RFC3339, tc.after)
|
||||
got, err := schedule.Next(after, time.UTC)
|
||||
if err != nil || got.Format(time.RFC3339) != tc.want {
|
||||
t.Errorf("%s next=%s err=%v want=%s", tc.expression, got.Format(time.RFC3339), err, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCronRejectsInvalid(t *testing.T) {
|
||||
for _, expression := range []string{"* * *", "60 * * * *", "*/0 * * * *", "* 24 * * *"} {
|
||||
if _, err := ParseCron(expression); err == nil {
|
||||
t.Errorf("expected %q to be rejected", expression)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,403 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
platformid "aigateway.local/core/internal/platform/id"
|
||||
)
|
||||
|
||||
const maxExecutionResponseBytes = 2 << 20
|
||||
|
||||
type Engine struct {
|
||||
service *Service
|
||||
baseURL string
|
||||
client *http.Client
|
||||
workerID string
|
||||
batchSize int
|
||||
maxAttempts int
|
||||
executionTTL time.Duration
|
||||
logger *slog.Logger
|
||||
}
|
||||
|
||||
func NewEngine(service *Service, baseURL, workerID string, batchSize, maxAttempts int, executionTTL time.Duration, logger *slog.Logger) *Engine {
|
||||
if batchSize < 1 {
|
||||
batchSize = 10
|
||||
}
|
||||
if maxAttempts < 1 {
|
||||
maxAttempts = 3
|
||||
}
|
||||
if executionTTL < time.Minute {
|
||||
executionTTL = 5 * time.Minute
|
||||
}
|
||||
if logger == nil {
|
||||
logger = slog.Default()
|
||||
}
|
||||
return &Engine{service: service, baseURL: strings.TrimRight(baseURL, "/"), client: &http.Client{Timeout: executionTTL}, workerID: workerID, batchSize: batchSize, maxAttempts: maxAttempts, executionTTL: executionTTL, logger: logger}
|
||||
}
|
||||
|
||||
// Tick coalesces overdue schedules into durable pending runs, reclaims stale
|
||||
// executions, and processes a bounded batch. SKIP LOCKED makes it safe for
|
||||
// multiple scheduler replicas to call Tick concurrently.
|
||||
func (e *Engine) Tick(ctx context.Context, now time.Time) (int, error) {
|
||||
if err := e.scheduleDue(ctx, now.UTC()); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if err := e.reclaim(ctx, now.UTC()); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
runs, err := e.claim(ctx)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
for _, run := range runs {
|
||||
task, getErr := e.service.Get(ctx, run.TaskID)
|
||||
if getErr != nil {
|
||||
_ = e.complete(ctx, run, Task{ID: run.TaskID, Code: run.TaskCode}, nil, getErr)
|
||||
continue
|
||||
}
|
||||
response, executeErr := e.execute(ctx, task, run)
|
||||
if executeErr != nil && run.Attempts < e.maxAttempts {
|
||||
if retryErr := e.retry(ctx, run, executeErr); retryErr != nil {
|
||||
return len(runs), retryErr
|
||||
}
|
||||
continue
|
||||
}
|
||||
if completeErr := e.complete(ctx, run, task, response, executeErr); completeErr != nil {
|
||||
return len(runs), completeErr
|
||||
}
|
||||
}
|
||||
return len(runs), nil
|
||||
}
|
||||
|
||||
func (e *Engine) scheduleDue(ctx context.Context, now time.Time) error {
|
||||
tx, err := e.service.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
rows, err := tx.Query(ctx, `SELECT id::text,cron_expression,timezone,next_run_at FROM gateway.scheduled_tasks WHERE enabled AND next_run_at <= $1 ORDER BY next_run_at,id FOR UPDATE SKIP LOCKED LIMIT $2`, now, e.batchSize)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
type due struct {
|
||||
id, expression, timezone string
|
||||
scheduledFor time.Time
|
||||
}
|
||||
items := []due{}
|
||||
for rows.Next() {
|
||||
var item due
|
||||
if err = rows.Scan(&item.id, &item.expression, &item.timezone, &item.scheduledFor); err != nil {
|
||||
rows.Close()
|
||||
return err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
rows.Close()
|
||||
if err = rows.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, item := range items {
|
||||
schedule, parseErr := ParseCron(item.expression)
|
||||
location, locationErr := time.LoadLocation(item.timezone)
|
||||
if parseErr != nil || locationErr != nil {
|
||||
_, err = tx.Exec(ctx, `UPDATE gateway.scheduled_tasks SET enabled=false,next_run_at=NULL,last_status='failed',last_error='cron 或时区配置无效',updated_at=clock_timestamp() WHERE id=$1`, item.id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
continue
|
||||
}
|
||||
// 补跑停机期间漏掉的执行:从旧的 next_run_at 起逐个 occurrence
|
||||
// 落一条 pending run,直到越过 now,而不是只补最新一次。否则调度器
|
||||
// 宕机超过一个周期后,中间所有计划执行被静默丢弃。
|
||||
runID, idErr := platformid.NewUUID()
|
||||
if idErr != nil {
|
||||
return idErr
|
||||
}
|
||||
if _, err = tx.Exec(ctx, `INSERT INTO gateway.scheduled_task_runs(id,task_id,trigger_type,scheduled_for) VALUES($1,$2,'scheduled',$3) ON CONFLICT(task_id,trigger_type,scheduled_for) DO NOTHING`, runID, item.id, item.scheduledFor); err != nil {
|
||||
return err
|
||||
}
|
||||
nextRun := item.scheduledFor
|
||||
var nextErr error
|
||||
inserted := 1
|
||||
for nextRun.Before(now) || nextRun.Equal(now) {
|
||||
if inserted >= catchUpLimit {
|
||||
break
|
||||
}
|
||||
nextRun, nextErr = schedule.Next(nextRun, location)
|
||||
if nextErr != nil {
|
||||
return nextErr
|
||||
}
|
||||
if nextRun.After(now) {
|
||||
break
|
||||
}
|
||||
runID, idErr = platformid.NewUUID()
|
||||
if idErr != nil {
|
||||
return idErr
|
||||
}
|
||||
if _, err = tx.Exec(ctx, `INSERT INTO gateway.scheduled_task_runs(id,task_id,trigger_type,scheduled_for) VALUES($1,$2,'scheduled',$3) ON CONFLICT(task_id,trigger_type,scheduled_for) DO NOTHING`, runID, item.id, nextRun); err != nil {
|
||||
return err
|
||||
}
|
||||
inserted++
|
||||
}
|
||||
next, nextErr := schedule.Next(now, location)
|
||||
if nextErr != nil {
|
||||
return nextErr
|
||||
}
|
||||
if _, err = tx.Exec(ctx, `UPDATE gateway.scheduled_tasks SET next_run_at=$2,updated_at=clock_timestamp() WHERE id=$1`, item.id, next); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
// catchUpLimit 是单任务单次补跑的最大执行数;超过部分丢弃,防止调度器长期
|
||||
// 停机后瞬间插入海量补跑记录。
|
||||
const catchUpLimit = 100
|
||||
|
||||
func (e *Engine) reclaim(ctx context.Context, now time.Time) error {
|
||||
cutoff := now.Add(-e.executionTTL)
|
||||
tx, err := e.service.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
if _, err = tx.Exec(ctx, `UPDATE gateway.scheduled_task_runs SET status='pending',worker_id='',started_at=NULL,error='上次执行超时,已回收重试' WHERE status='running' AND started_at < $1 AND attempts < $2`, cutoff, e.maxAttempts); err != nil {
|
||||
return err
|
||||
}
|
||||
rows, err := tx.Query(ctx, `WITH picked AS (SELECT id FROM gateway.scheduled_task_runs WHERE status='running' AND started_at < $1 AND attempts >= $2 ORDER BY started_at,id FOR UPDATE SKIP LOCKED LIMIT $3) UPDATE gateway.scheduled_task_runs r SET worker_id=$4,started_at=$5 FROM picked WHERE r.id=picked.id RETURNING r.id::text`, cutoff, e.maxAttempts, e.batchSize, e.workerID, now)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ids := []string{}
|
||||
for rows.Next() {
|
||||
var id string
|
||||
if err = rows.Scan(&id); err != nil {
|
||||
rows.Close()
|
||||
return err
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
rows.Close()
|
||||
if err = rows.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, id := range ids {
|
||||
run, getErr := e.service.getRun(ctx, id)
|
||||
if getErr != nil {
|
||||
return getErr
|
||||
}
|
||||
task, taskErr := e.service.Get(ctx, run.TaskID)
|
||||
if taskErr != nil {
|
||||
task = Task{ID: run.TaskID, Code: run.TaskCode}
|
||||
}
|
||||
if err = e.complete(ctx, run, task, nil, errors.New("执行超时且达到最大重试次数")); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (e *Engine) claim(ctx context.Context) ([]Run, error) {
|
||||
tx, err := e.service.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
rows, err := tx.Query(ctx, `WITH picked AS (SELECT id FROM gateway.scheduled_task_runs WHERE status='pending' ORDER BY created_at,id FOR UPDATE SKIP LOCKED LIMIT $1) UPDATE gateway.scheduled_task_runs r SET status='running',attempts=attempts+1,worker_id=$2,started_at=clock_timestamp(),error='' FROM picked WHERE r.id=picked.id RETURNING r.id::text`, e.batchSize, e.workerID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ids := []string{}
|
||||
for rows.Next() {
|
||||
var id string
|
||||
if err = rows.Scan(&id); err != nil {
|
||||
rows.Close()
|
||||
return nil, err
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
rows.Close()
|
||||
if err = rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
runs := make([]Run, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
run, getErr := e.service.getRun(ctx, id)
|
||||
if getErr != nil {
|
||||
return nil, getErr
|
||||
}
|
||||
runs = append(runs, run)
|
||||
}
|
||||
return runs, nil
|
||||
}
|
||||
|
||||
func (e *Engine) conversationMessages(ctx context.Context, task Task, currentRunID string) ([]map[string]any, error) {
|
||||
messages := []map[string]any{}
|
||||
if task.ConversationID != "" {
|
||||
rows, err := e.service.pool.Query(ctx, `SELECT response FROM gateway.scheduled_task_runs WHERE task_id=$1 AND id<>$2 AND status='success' AND response IS NOT NULL ORDER BY created_at DESC LIMIT 5`, task.ID, currentRunID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
answers := []string{}
|
||||
for rows.Next() {
|
||||
var response json.RawMessage
|
||||
if err = rows.Scan(&response); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if answer := responseAnswer(response); answer != "" {
|
||||
answers = append(answers, answer)
|
||||
}
|
||||
}
|
||||
for i := len(answers) - 1; i >= 0; i-- {
|
||||
messages = append(messages, map[string]any{"role": "user", "content": task.Prompt}, map[string]any{"role": "assistant", "content": answers[i]})
|
||||
}
|
||||
}
|
||||
messages = append(messages, map[string]any{"role": "user", "content": task.Prompt})
|
||||
return messages, nil
|
||||
}
|
||||
|
||||
func responseAnswer(raw json.RawMessage) string {
|
||||
var response map[string]any
|
||||
if json.Unmarshal(raw, &response) != nil {
|
||||
return ""
|
||||
}
|
||||
choices, _ := response["choices"].([]any)
|
||||
if len(choices) == 0 {
|
||||
return ""
|
||||
}
|
||||
choice, _ := choices[0].(map[string]any)
|
||||
message, _ := choice["message"].(map[string]any)
|
||||
answer, _ := message["content"].(string)
|
||||
return answer
|
||||
}
|
||||
|
||||
func (e *Engine) execute(ctx context.Context, task Task, run Run) (json.RawMessage, error) {
|
||||
secret, err := e.service.decryptAPIKey(task)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("解密执行 API Key: %w", err)
|
||||
}
|
||||
messages, err := e.conversationMessages(ctx, task, run.ID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var variables map[string]any
|
||||
if json.Unmarshal(task.Variables, &variables) != nil {
|
||||
variables = map[string]any{}
|
||||
}
|
||||
payload := map[string]any{"messages": messages, "variables": variables}
|
||||
path := "/v1/applications/" + task.TargetCode + "/chat/completions"
|
||||
if task.TargetType == "digital_employee" {
|
||||
path = "/v1/digital-employees/" + task.TargetCode + "/chat/completions"
|
||||
payload["skill_ids"] = task.SkillIDs
|
||||
payload["mcp_server_ids"] = task.MCPServerIDs
|
||||
}
|
||||
body, _ := json.Marshal(payload)
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodPost, e.baseURL+path, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
request.Header.Set("X-Gateway-API-Key", secret)
|
||||
request.Header.Set("X-Request-ID", "scheduled-"+run.ID)
|
||||
response, err := e.client.Do(request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
data, err := io.ReadAll(io.LimitReader(response.Body, maxExecutionResponseBytes+1))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(data) > maxExecutionResponseBytes {
|
||||
return nil, errors.New("模型响应超过 2MB 上限")
|
||||
}
|
||||
if !json.Valid(data) {
|
||||
return nil, errors.New("模型响应不是合法 JSON")
|
||||
}
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("模型调用失败(HTTP %d): %s", response.StatusCode, truncate(string(data), 1000))
|
||||
}
|
||||
return json.RawMessage(data), nil
|
||||
}
|
||||
|
||||
func truncate(value string, maximum int) string {
|
||||
value = strings.TrimSpace(value)
|
||||
if len(value) <= maximum {
|
||||
return value
|
||||
}
|
||||
cut := value[:maximum]
|
||||
// 按字节截断可能把多字节 rune 切半,产生无效 UTF-8;PostgreSQL text 列
|
||||
// 会拒绝写入,导致任务永远停在重试循环。回退到最近的 rune 边界。
|
||||
for len(cut) > 0 && !utf8.RuneStart(cut[len(cut)-1]) {
|
||||
cut = cut[:len(cut)-1]
|
||||
}
|
||||
return cut
|
||||
}
|
||||
|
||||
func (e *Engine) retry(ctx context.Context, run Run, executeErr error) error {
|
||||
errorText := truncate(executeErr.Error(), 4000)
|
||||
tag, err := e.service.pool.Exec(ctx, `UPDATE gateway.scheduled_task_runs SET status='pending',worker_id='',started_at=NULL,response=NULL,error=$2 WHERE id=$1 AND status='running' AND worker_id=$3`, run.ID, errorText, e.workerID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return errors.New("定时任务执行租约已失效")
|
||||
}
|
||||
e.logger.Warn("scheduled task execution will retry", "task", run.TaskCode, "run_id", run.ID, "attempt", run.Attempts, "error", executeErr)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (e *Engine) complete(ctx context.Context, run Run, task Task, response json.RawMessage, executeErr error) error {
|
||||
status := "success"
|
||||
errorText := ""
|
||||
eventType := "scheduled_task.completed"
|
||||
if executeErr != nil {
|
||||
status = "failed"
|
||||
eventType = "scheduled_task.failed"
|
||||
errorText = truncate(executeErr.Error(), 4000)
|
||||
e.logger.Warn("scheduled task execution failed", "task", task.Code, "run_id", run.ID, "error", executeErr)
|
||||
}
|
||||
tx, err := e.service.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
tag, err := tx.Exec(ctx, `UPDATE gateway.scheduled_task_runs SET status=$2,response=$3,error=$4,finished_at=clock_timestamp() WHERE id=$1 AND status='running' AND worker_id=$5`, run.ID, status, response, errorText, e.workerID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return errors.New("定时任务执行租约已失效")
|
||||
}
|
||||
_, err = tx.Exec(ctx, `UPDATE gateway.scheduled_tasks SET last_run_at=clock_timestamp(),last_status=$2,last_error=$3,updated_at=clock_timestamp() WHERE id=$1`, task.ID, status, errorText)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
eventID, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{"scheduled_task_id": task.ID, "task_code": task.Code, "run_id": run.ID, "status": status, "error": errorText, "actor_id": task.CreatedBy, "notification_channel_id": task.NotificationChannelID})
|
||||
_, err = tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,$2,1,'scheduled_task',$3,$4)`, eventID, eventType, run.ID, payload)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
@@ -0,0 +1,173 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"aigateway.local/core/internal/identity"
|
||||
"aigateway.local/core/internal/platform/apiresponse"
|
||||
)
|
||||
|
||||
type AdminHTTPHandler struct {
|
||||
service *Service
|
||||
identity *identity.Service
|
||||
mux *http.ServeMux
|
||||
}
|
||||
|
||||
func NewAdminHTTPHandler(service *Service, identityService *identity.Service) *AdminHTTPHandler {
|
||||
h := &AdminHTTPHandler{service: service, identity: identityService, mux: http.NewServeMux()}
|
||||
h.mux.HandleFunc("GET /api/v1/admin/scheduled-tasks", h.list)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/scheduled-tasks", h.create)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/scheduled-tasks/{id}", h.get)
|
||||
h.mux.HandleFunc("PUT /api/v1/admin/scheduled-tasks/{id}", h.update)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/scheduled-tasks/{id}", h.delete)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/scheduled-tasks/{id}/start", h.start)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/scheduled-tasks/{id}/pause", h.pause)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/scheduled-tasks/{id}/run", h.runNow)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/scheduled-tasks/{id}/runs", h.runs)
|
||||
return h
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.mux.ServeHTTP(w, r) }
|
||||
|
||||
func (h *AdminHTTPHandler) 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 decodeTask(w http.ResponseWriter, r *http.Request, target any) bool {
|
||||
decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, 2<<20))
|
||||
decoder.DisallowUnknownFields()
|
||||
if err := decoder.Decode(target); err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "请求格式无效")
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func taskError(w http.ResponseWriter, err error) {
|
||||
if errors.Is(err, ErrNotFound) {
|
||||
apiresponse.Error(w, http.StatusNotFound, "定时任务不存在")
|
||||
return
|
||||
}
|
||||
apiresponse.Error(w, http.StatusBadRequest, err.Error())
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) list(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionScheduledTaskRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.service.List(r.Context())
|
||||
if err != nil {
|
||||
taskError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) get(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionScheduledTaskRead); !ok {
|
||||
return
|
||||
}
|
||||
item, err := h.service.Get(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
taskError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) create(w http.ResponseWriter, r *http.Request) {
|
||||
account, ok := h.require(w, r, identity.PermissionScheduledTaskManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input TaskInput
|
||||
if !decodeTask(w, r, &input) {
|
||||
return
|
||||
}
|
||||
item, err := h.service.Save(r.Context(), "", input, account.ID)
|
||||
if err != nil {
|
||||
taskError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) update(w http.ResponseWriter, r *http.Request) {
|
||||
account, ok := h.require(w, r, identity.PermissionScheduledTaskManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input TaskInput
|
||||
if !decodeTask(w, r, &input) {
|
||||
return
|
||||
}
|
||||
item, err := h.service.Save(r.Context(), r.PathValue("id"), input, account.ID)
|
||||
if err != nil {
|
||||
taskError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) delete(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionScheduledTaskManage); !ok {
|
||||
return
|
||||
}
|
||||
if err := h.service.Delete(r.Context(), r.PathValue("id")); err != nil {
|
||||
taskError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"deleted": true})
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) setEnabled(w http.ResponseWriter, r *http.Request, enabled bool) {
|
||||
if _, ok := h.require(w, r, identity.PermissionScheduledTaskManage); !ok {
|
||||
return
|
||||
}
|
||||
item, err := h.service.SetEnabled(r.Context(), r.PathValue("id"), enabled)
|
||||
if err != nil {
|
||||
taskError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) start(w http.ResponseWriter, r *http.Request) { h.setEnabled(w, r, true) }
|
||||
func (h *AdminHTTPHandler) pause(w http.ResponseWriter, r *http.Request) { h.setEnabled(w, r, false) }
|
||||
|
||||
func (h *AdminHTTPHandler) runNow(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionScheduledTaskManage); !ok {
|
||||
return
|
||||
}
|
||||
run, err := h.service.QueueManual(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
taskError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, run)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) runs(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionScheduledTaskRead); !ok {
|
||||
return
|
||||
}
|
||||
limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
|
||||
items, err := h.service.Runs(r.Context(), r.PathValue("id"), limit)
|
||||
if err != nil {
|
||||
taskError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/platform/config"
|
||||
"aigateway.local/core/internal/platform/cryptox"
|
||||
"aigateway.local/core/internal/platform/database"
|
||||
)
|
||||
|
||||
func TestSchedulerPostgreSQLLifecycle(t *testing.T) {
|
||||
databaseURL := os.Getenv("SCHEDULER_TEST_DATABASE_URL")
|
||||
if databaseURL == "" {
|
||||
t.Skip("SCHEDULER_TEST_DATABASE_URL is not set")
|
||||
}
|
||||
ctx := context.Background()
|
||||
pool, err := database.Open(ctx, config.Database{URL: databaseURL, MaxConns: 8})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer pool.Close()
|
||||
adminID := "64444444-4444-4444-8444-444444444444"
|
||||
appID := "65555555-5555-4555-8555-555555555555"
|
||||
versionID := "66666666-6666-4666-8666-666666666666"
|
||||
cleanup := func() {
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM gateway.scheduled_tasks WHERE code='scheduler_test_task'`)
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM gateway.applications WHERE id=$1`, appID)
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM gateway.admin_accounts WHERE id=$1`, adminID)
|
||||
}
|
||||
cleanup()
|
||||
defer cleanup()
|
||||
if _, err = pool.Exec(ctx, `INSERT INTO gateway.admin_accounts(id,username,password_hash,role,active) VALUES($1,'scheduler-test-admin','test','superadmin',true)`, adminID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tx, err := pool.Begin(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
configJSON := `{"model":"test-model","knowledge_base_ids":[],"tool_ids":[],"retrieval_top_k":4,"temperature":0.2,"max_tool_rounds":1}`
|
||||
if _, err = tx.Exec(ctx, `INSERT INTO gateway.applications(id,code,name,status,draft_config,created_by) VALUES($1,'scheduler_test_app','Scheduler Test App','active',$2,$3)`, appID, configJSON, adminID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = tx.Exec(ctx, `INSERT INTO gateway.application_versions(id,application_id,version,config,published_by) VALUES($1,$2,1,$3,$4)`, versionID, appID, configJSON, adminID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = tx.Exec(ctx, `UPDATE gateway.applications SET published_version=1 WHERE id=$1`, appID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
requestCount := 0
|
||||
failRequests := false
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requestCount++
|
||||
if r.URL.Path != "/v1/applications/scheduler_test_app/chat/completions" || r.Header.Get("X-Gateway-API-Key") != "gw_scheduler_test" {
|
||||
http.Error(w, "unexpected request", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
if failRequests {
|
||||
http.Error(w, `{"error":"temporary failure"}`, http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"scheduled answer"}}]}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
cipher, err := cryptox.NewKeyring("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=", 1, "", "scheduled-task-api-key")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
service := NewService(pool, cipher)
|
||||
task, err := service.Save(ctx, "", TaskInput{Code: "scheduler_test_task", Name: "Scheduler Test", CronExpression: "*/5 * * * *", Timezone: "UTC", TargetType: "application", TargetCode: "scheduler_test_app", Prompt: "create report", Variables: json.RawMessage(`{"scope":"daily"}`), APIKey: "gw_scheduler_test"}, adminID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !task.HasAPIKey || task.Enabled {
|
||||
t.Fatalf("unexpected task: %+v", task)
|
||||
}
|
||||
if _, err = service.QueueManual(ctx, task.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
engine := NewEngine(service, server.URL, "integration-worker", 10, 3, time.Minute, nil)
|
||||
processed, err := engine.Tick(ctx, time.Now())
|
||||
if err != nil || processed != 1 || requestCount != 1 {
|
||||
t.Fatalf("tick processed=%d requests=%d err=%v", processed, requestCount, err)
|
||||
}
|
||||
runs, err := service.Runs(ctx, task.ID, 10)
|
||||
if err != nil || len(runs) != 1 || runs[0].Status != "success" || responseAnswer(runs[0].Response) != "scheduled answer" {
|
||||
t.Fatalf("runs=%+v err=%v", runs, err)
|
||||
}
|
||||
var completed bool
|
||||
if err = pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.outbox_events WHERE aggregate_type='scheduled_task' AND aggregate_id=$1 AND event_type='scheduled_task.completed')`, runs[0].ID).Scan(&completed); err != nil || !completed {
|
||||
t.Fatalf("completion event=%v err=%v", completed, err)
|
||||
}
|
||||
|
||||
// Force a due schedule in the past. Tick must enqueue it once, execute it,
|
||||
// and advance next_run_at beyond now rather than replay every missed slot.
|
||||
if _, err = service.SetEnabled(ctx, task.ID, true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
forcedNow := time.Now().UTC()
|
||||
if _, err = pool.Exec(ctx, `UPDATE gateway.scheduled_tasks SET next_run_at=$2 WHERE id=$1`, task.ID, forcedNow.Add(-time.Minute)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
processed, err = engine.Tick(ctx, forcedNow)
|
||||
if err != nil || processed != 1 || requestCount != 2 {
|
||||
t.Fatalf("due tick processed=%d requests=%d err=%v", processed, requestCount, err)
|
||||
}
|
||||
task, err = service.Get(ctx, task.ID)
|
||||
if err != nil || task.NextRunAt == nil || !task.NextRunAt.After(forcedNow) {
|
||||
t.Fatalf("next run was not advanced: %+v err=%v", task.NextRunAt, err)
|
||||
}
|
||||
|
||||
// Ordinary gateway failures remain pending until the configured attempt
|
||||
// limit, then become failed and emit exactly one terminal event.
|
||||
if _, err = service.SetEnabled(ctx, task.ID, false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
failRequests = true
|
||||
failedRun, err := service.QueueManual(ctx, task.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for attempt := 1; attempt <= 3; attempt++ {
|
||||
processed, err = engine.Tick(ctx, time.Now())
|
||||
if err != nil || processed != 1 {
|
||||
t.Fatalf("retry tick attempt=%d processed=%d err=%v", attempt, processed, err)
|
||||
}
|
||||
current, getErr := service.getRun(ctx, failedRun.ID)
|
||||
wantStatus := "pending"
|
||||
if attempt == 3 {
|
||||
wantStatus = "failed"
|
||||
}
|
||||
if getErr != nil || current.Status != wantStatus || current.Attempts != attempt {
|
||||
t.Fatalf("retry attempt=%d run=%+v err=%v", attempt, current, getErr)
|
||||
}
|
||||
}
|
||||
var failedEvents int
|
||||
if err = pool.QueryRow(ctx, `SELECT count(*) FROM gateway.outbox_events WHERE aggregate_type='scheduled_task' AND aggregate_id=$1 AND event_type='scheduled_task.failed'`, failedRun.ID).Scan(&failedEvents); err != nil || failedEvents != 1 {
|
||||
t.Fatalf("failure events=%d err=%v", failedEvents, err)
|
||||
}
|
||||
|
||||
// A worker lease that stays running past its timeout is finalized through
|
||||
// the same failed-run and outbox path once it reaches the attempt limit.
|
||||
staleRun, err := service.QueueManual(ctx, task.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = pool.Exec(ctx, `UPDATE gateway.scheduled_task_runs SET status='running',attempts=3,worker_id='dead-worker',started_at=$2 WHERE id=$1`, staleRun.ID, time.Now().Add(-2*time.Minute)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
processed, err = engine.Tick(ctx, time.Now())
|
||||
if err != nil || processed != 0 {
|
||||
t.Fatalf("stale tick processed=%d err=%v", processed, err)
|
||||
}
|
||||
staleRun, err = service.getRun(ctx, staleRun.ID)
|
||||
if err != nil || staleRun.Status != "failed" || staleRun.Error != "执行超时且达到最大重试次数" {
|
||||
t.Fatalf("stale run=%+v err=%v", staleRun, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,441 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/apikey"
|
||||
"aigateway.local/core/internal/platform/cryptox"
|
||||
platformid "aigateway.local/core/internal/platform/id"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrNotFound = errors.New("scheduled task not found")
|
||||
codePattern = regexp.MustCompile(`^[a-z][a-z0-9_-]{1,63}$`)
|
||||
uuidPattern = regexp.MustCompile(`^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-5][0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}$`)
|
||||
)
|
||||
|
||||
type Task struct {
|
||||
ID string `json:"id"`
|
||||
Code string `json:"code"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
CronExpression string `json:"cron_expression"`
|
||||
Timezone string `json:"timezone"`
|
||||
TargetType string `json:"target_type"`
|
||||
TargetCode string `json:"target_code"`
|
||||
Prompt string `json:"prompt"`
|
||||
Variables json.RawMessage `json:"variables"`
|
||||
SkillIDs []string `json:"skill_ids"`
|
||||
MCPServerIDs []string `json:"mcp_server_ids"`
|
||||
ConversationID string `json:"conversation_id"`
|
||||
NotificationChannelID *string `json:"notification_channel_id,omitempty"`
|
||||
HasAPIKey bool `json:"has_api_key"`
|
||||
Enabled bool `json:"enabled"`
|
||||
NextRunAt *time.Time `json:"next_run_at,omitempty"`
|
||||
LastRunAt *time.Time `json:"last_run_at,omitempty"`
|
||||
LastStatus string `json:"last_status"`
|
||||
LastError string `json:"last_error"`
|
||||
CreatedBy string `json:"created_by"`
|
||||
Revision int64 `json:"revision"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
encryptedAPIKey []byte
|
||||
apiKeyKEKVersion int
|
||||
}
|
||||
|
||||
type TaskInput struct {
|
||||
Code string `json:"code"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
CronExpression string `json:"cron_expression"`
|
||||
Timezone string `json:"timezone"`
|
||||
TargetType string `json:"target_type"`
|
||||
TargetCode string `json:"target_code"`
|
||||
Prompt string `json:"prompt"`
|
||||
Variables json.RawMessage `json:"variables"`
|
||||
SkillIDs []string `json:"skill_ids"`
|
||||
MCPServerIDs []string `json:"mcp_server_ids"`
|
||||
ConversationID string `json:"conversation_id"`
|
||||
NotificationChannelID *string `json:"notification_channel_id"`
|
||||
APIKey string `json:"api_key"`
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
|
||||
type Run struct {
|
||||
ID string `json:"id"`
|
||||
TaskID string `json:"task_id"`
|
||||
TaskCode string `json:"task_code"`
|
||||
TriggerType string `json:"trigger_type"`
|
||||
ScheduledFor time.Time `json:"scheduled_for"`
|
||||
Status string `json:"status"`
|
||||
Attempts int `json:"attempts"`
|
||||
WorkerID string `json:"worker_id"`
|
||||
StartedAt *time.Time `json:"started_at,omitempty"`
|
||||
FinishedAt *time.Time `json:"finished_at,omitempty"`
|
||||
Response json.RawMessage `json:"response,omitempty"`
|
||||
Error string `json:"error"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
type Service struct {
|
||||
pool *pgxpool.Pool
|
||||
cipher cryptox.Cipher
|
||||
now func() time.Time
|
||||
}
|
||||
|
||||
func NewService(pool *pgxpool.Pool, cipher cryptox.Cipher) *Service {
|
||||
return &Service{pool: pool, cipher: cipher, now: time.Now}
|
||||
}
|
||||
|
||||
const taskSelect = `SELECT id::text,code,name,description,cron_expression,timezone,target_type,target_code,prompt,variables,skill_ids::text[],mcp_server_ids::text[],conversation_id,notification_channel_id::text,encrypted_api_key,api_key_kek_version,enabled,next_run_at,last_run_at,last_status,last_error,created_by::text,revision,created_at,updated_at FROM gateway.scheduled_tasks`
|
||||
|
||||
func scanTask(row pgx.Row) (Task, error) {
|
||||
var task Task
|
||||
err := row.Scan(&task.ID, &task.Code, &task.Name, &task.Description, &task.CronExpression, &task.Timezone, &task.TargetType, &task.TargetCode, &task.Prompt, &task.Variables, &task.SkillIDs, &task.MCPServerIDs, &task.ConversationID, &task.NotificationChannelID, &task.encryptedAPIKey, &task.apiKeyKEKVersion, &task.Enabled, &task.NextRunAt, &task.LastRunAt, &task.LastStatus, &task.LastError, &task.CreatedBy, &task.Revision, &task.CreatedAt, &task.UpdatedAt)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return Task{}, ErrNotFound
|
||||
}
|
||||
task.HasAPIKey = len(task.encryptedAPIKey) > 0
|
||||
if task.SkillIDs == nil {
|
||||
task.SkillIDs = []string{}
|
||||
}
|
||||
if task.MCPServerIDs == nil {
|
||||
task.MCPServerIDs = []string{}
|
||||
}
|
||||
return task, err
|
||||
}
|
||||
|
||||
func (s *Service) List(ctx context.Context) ([]Task, error) {
|
||||
rows, err := s.pool.Query(ctx, taskSelect+` ORDER BY updated_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []Task{}
|
||||
for rows.Next() {
|
||||
item, scanErr := scanTask(rows)
|
||||
if scanErr != nil {
|
||||
return nil, scanErr
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Service) Get(ctx context.Context, id string) (Task, error) {
|
||||
return scanTask(s.pool.QueryRow(ctx, taskSelect+` WHERE id=$1`, id))
|
||||
}
|
||||
|
||||
func normalizeIDs(values []string, maximum int) ([]string, error) {
|
||||
seen := map[string]bool{}
|
||||
result := []string{}
|
||||
for _, value := range values {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" || seen[value] {
|
||||
continue
|
||||
}
|
||||
if !uuidPattern.MatchString(value) {
|
||||
return nil, errors.New("资源 ID 格式无效")
|
||||
}
|
||||
seen[value] = true
|
||||
result = append(result, value)
|
||||
}
|
||||
if len(result) > maximum {
|
||||
return nil, fmt.Errorf("资源绑定最多允许 %d 项", maximum)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *Service) validate(ctx context.Context, input *TaskInput, current *Task) (time.Time, []byte, int, error) {
|
||||
input.Code = strings.ToLower(strings.TrimSpace(input.Code))
|
||||
input.Name = strings.TrimSpace(input.Name)
|
||||
input.Description = strings.TrimSpace(input.Description)
|
||||
input.CronExpression = strings.TrimSpace(input.CronExpression)
|
||||
input.Timezone = strings.TrimSpace(input.Timezone)
|
||||
input.TargetType = strings.TrimSpace(input.TargetType)
|
||||
input.TargetCode = strings.ToLower(strings.TrimSpace(input.TargetCode))
|
||||
input.Prompt = strings.TrimSpace(input.Prompt)
|
||||
input.ConversationID = strings.TrimSpace(input.ConversationID)
|
||||
if !codePattern.MatchString(input.Code) || input.Name == "" || len(input.Name) > 128 || len(input.Description) > 4000 {
|
||||
return time.Time{}, nil, 0, errors.New("任务编码、名称或描述格式无效")
|
||||
}
|
||||
if len(input.Prompt) < 1 || len(input.Prompt) > 100000 || len(input.ConversationID) > 128 {
|
||||
return time.Time{}, nil, 0, errors.New("提示词或会话 ID 格式无效")
|
||||
}
|
||||
if input.Timezone == "" {
|
||||
input.Timezone = "UTC"
|
||||
}
|
||||
location, err := time.LoadLocation(input.Timezone)
|
||||
if err != nil {
|
||||
return time.Time{}, nil, 0, errors.New("时区名称无效")
|
||||
}
|
||||
schedule, err := ParseCron(input.CronExpression)
|
||||
if err != nil {
|
||||
return time.Time{}, nil, 0, err
|
||||
}
|
||||
next, err := schedule.Next(s.now(), location)
|
||||
if err != nil {
|
||||
return time.Time{}, nil, 0, err
|
||||
}
|
||||
if len(input.Variables) == 0 {
|
||||
input.Variables = json.RawMessage(`{}`)
|
||||
}
|
||||
var variables map[string]any
|
||||
if json.Unmarshal(input.Variables, &variables) != nil {
|
||||
return time.Time{}, nil, 0, errors.New("变量必须是 JSON 对象")
|
||||
}
|
||||
input.Variables, _ = json.Marshal(variables)
|
||||
if input.SkillIDs, err = normalizeIDs(input.SkillIDs, 100); err != nil {
|
||||
return time.Time{}, nil, 0, err
|
||||
}
|
||||
if input.MCPServerIDs, err = normalizeIDs(input.MCPServerIDs, 100); err != nil {
|
||||
return time.Time{}, nil, 0, err
|
||||
}
|
||||
if err = s.validateTarget(ctx, input); err != nil {
|
||||
return time.Time{}, nil, 0, err
|
||||
}
|
||||
if input.NotificationChannelID != nil {
|
||||
trimmed := strings.TrimSpace(*input.NotificationChannelID)
|
||||
if trimmed == "" {
|
||||
input.NotificationChannelID = nil
|
||||
} else {
|
||||
var exists bool
|
||||
if err = s.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.notification_channels WHERE id=$1 AND enabled)`, trimmed).Scan(&exists); err != nil || !exists {
|
||||
return time.Time{}, nil, 0, errors.New("通知渠道不存在或未启用")
|
||||
}
|
||||
input.NotificationChannelID = &trimmed
|
||||
}
|
||||
}
|
||||
secret := strings.TrimSpace(input.APIKey)
|
||||
if secret == "" && current == nil {
|
||||
return time.Time{}, nil, 0, errors.New("首次创建必须填写执行 API Key")
|
||||
}
|
||||
if secret == "" {
|
||||
return next, current.encryptedAPIKey, current.apiKeyKEKVersion, nil
|
||||
}
|
||||
if len(secret) > 512 {
|
||||
return time.Time{}, nil, 0, errors.New("执行 API Key 过长")
|
||||
}
|
||||
encrypted, version, err := s.cipher.Encrypt([]byte(secret))
|
||||
if err != nil {
|
||||
return time.Time{}, nil, 0, fmt.Errorf("加密执行 API Key: %w", err)
|
||||
}
|
||||
return next, encrypted, version, nil
|
||||
}
|
||||
|
||||
func subset(selected, allowed []string) bool {
|
||||
set := map[string]bool{}
|
||||
for _, id := range allowed {
|
||||
set[id] = true
|
||||
}
|
||||
for _, id := range selected {
|
||||
if !set[id] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *Service) validateTarget(ctx context.Context, input *TaskInput) error {
|
||||
switch input.TargetType {
|
||||
case "application":
|
||||
if len(input.SkillIDs) > 0 || len(input.MCPServerIDs) > 0 {
|
||||
return errors.New("应用任务不支持额外绑定 Skill 或 MCP")
|
||||
}
|
||||
var exists bool
|
||||
err := s.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.applications WHERE code=$1 AND status='active' AND published_version IS NOT NULL)`, input.TargetCode).Scan(&exists)
|
||||
if err != nil || !exists {
|
||||
return errors.New("目标应用不存在或未发布")
|
||||
}
|
||||
var departmentIDs []string
|
||||
_ = s.pool.QueryRow(ctx, `SELECT COALESCE(department_ids,'{}'::text[]) FROM gateway.applications WHERE code=$1`, input.TargetCode).Scan(&departmentIDs)
|
||||
if len(departmentIDs) > 0 {
|
||||
if err := s.requireKeyTenant(ctx, input.APIKey, departmentIDs); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
case "digital_employee":
|
||||
var skillIDs, mcpIDs, departmentIDs []string
|
||||
var enabled bool
|
||||
var status string
|
||||
err := s.pool.QueryRow(ctx, `SELECT skill_ids::text[],mcp_server_ids::text[],enabled,status FROM gateway.digital_employees WHERE code=$1`, input.TargetCode).Scan(&skillIDs, &mcpIDs, &enabled, &status)
|
||||
if err != nil || !enabled || status != "published" {
|
||||
return errors.New("目标数字员工不存在或未发布")
|
||||
}
|
||||
if !subset(input.SkillIDs, skillIDs) || !subset(input.MCPServerIDs, mcpIDs) {
|
||||
return errors.New("任务选择的 Skill/MCP 必须已绑定到目标数字员工")
|
||||
}
|
||||
_ = s.pool.QueryRow(ctx, `SELECT COALESCE(department_ids,'{}'::text[]) FROM gateway.digital_employees WHERE code=$1`, input.TargetCode).Scan(&departmentIDs)
|
||||
if len(departmentIDs) > 0 {
|
||||
if err := s.requireKeyTenant(ctx, input.APIKey, departmentIDs); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
default:
|
||||
return errors.New("目标类型必须是 application 或 digital_employee")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// requireKeyTenant 校验任务 API Key 的部门归属能访问部门限定目标,在保存
|
||||
// 阶段就失败,而不是让任务创建成功后永远执行失败(执行 key 无 tenant 时
|
||||
// 运行时对部门限定资源一律不可见)。secret 为空(更新时沿用旧 key)跳过。
|
||||
func (s *Service) requireKeyTenant(ctx context.Context, secret string, departmentIDs []string) error {
|
||||
secret = strings.TrimSpace(secret)
|
||||
if secret == "" {
|
||||
return nil
|
||||
}
|
||||
hash, _ := apikey.Digest(secret)
|
||||
var tenant *string
|
||||
err := s.pool.QueryRow(ctx, `SELECT tenant_id::text FROM gateway.api_keys WHERE key_hash=$1 AND enabled`, hash).Scan(&tenant)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return errors.New("执行 API Key 不存在或已停用")
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tenant == nil {
|
||||
return errors.New("目标资源按部门限定,但执行 API Key 未绑定部门;请使用该部门下的 API Key")
|
||||
}
|
||||
for _, id := range departmentIDs {
|
||||
if id == *tenant {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return errors.New("执行 API Key 所属部门与目标资源部门不匹配")
|
||||
}
|
||||
|
||||
func (s *Service) Save(ctx context.Context, id string, input TaskInput, actorID string) (Task, error) {
|
||||
var current *Task
|
||||
if id != "" {
|
||||
item, err := s.Get(ctx, id)
|
||||
if err != nil {
|
||||
return Task{}, err
|
||||
}
|
||||
current = &item
|
||||
}
|
||||
next, encrypted, version, err := s.validate(ctx, &input, current)
|
||||
if err != nil {
|
||||
return Task{}, err
|
||||
}
|
||||
if id == "" {
|
||||
id, err = platformid.NewUUID()
|
||||
if err != nil {
|
||||
return Task{}, err
|
||||
}
|
||||
_, err = s.pool.Exec(ctx, `INSERT INTO gateway.scheduled_tasks(id,code,name,description,cron_expression,timezone,target_type,target_code,prompt,variables,skill_ids,mcp_server_ids,conversation_id,notification_channel_id,encrypted_api_key,api_key_kek_version,enabled,next_run_at,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19)`, id, input.Code, input.Name, input.Description, input.CronExpression, input.Timezone, input.TargetType, input.TargetCode, input.Prompt, input.Variables, input.SkillIDs, input.MCPServerIDs, input.ConversationID, input.NotificationChannelID, encrypted, version, input.Enabled, nullableNext(input.Enabled, next), actorID)
|
||||
} else {
|
||||
_, err = s.pool.Exec(ctx, `UPDATE gateway.scheduled_tasks SET code=$2,name=$3,description=$4,cron_expression=$5,timezone=$6,target_type=$7,target_code=$8,prompt=$9,variables=$10,skill_ids=$11,mcp_server_ids=$12,conversation_id=$13,notification_channel_id=$14,encrypted_api_key=$15,api_key_kek_version=$16,enabled=$17,next_run_at=$18,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, input.Code, input.Name, input.Description, input.CronExpression, input.Timezone, input.TargetType, input.TargetCode, input.Prompt, input.Variables, input.SkillIDs, input.MCPServerIDs, input.ConversationID, input.NotificationChannelID, encrypted, version, input.Enabled, nullableNext(input.Enabled, next))
|
||||
}
|
||||
if err != nil {
|
||||
return Task{}, err
|
||||
}
|
||||
return s.Get(ctx, id)
|
||||
}
|
||||
|
||||
func nullableNext(enabled bool, next time.Time) any {
|
||||
if !enabled {
|
||||
return nil
|
||||
}
|
||||
return next
|
||||
}
|
||||
|
||||
func (s *Service) SetEnabled(ctx context.Context, id string, enabled bool) (Task, error) {
|
||||
task, err := s.Get(ctx, id)
|
||||
if err != nil {
|
||||
return Task{}, err
|
||||
}
|
||||
var next any
|
||||
if enabled {
|
||||
location, _ := time.LoadLocation(task.Timezone)
|
||||
schedule, _ := ParseCron(task.CronExpression)
|
||||
value, nextErr := schedule.Next(s.now(), location)
|
||||
if nextErr != nil {
|
||||
return Task{}, nextErr
|
||||
}
|
||||
next = value
|
||||
}
|
||||
tag, err := s.pool.Exec(ctx, `UPDATE gateway.scheduled_tasks SET enabled=$2,next_run_at=$3,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, enabled, next)
|
||||
if err != nil || tag.RowsAffected() == 0 {
|
||||
return Task{}, ErrNotFound
|
||||
}
|
||||
return s.Get(ctx, id)
|
||||
}
|
||||
|
||||
func (s *Service) Delete(ctx context.Context, id string) error {
|
||||
tag, err := s.pool.Exec(ctx, `DELETE FROM gateway.scheduled_tasks WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) QueueManual(ctx context.Context, id string) (Run, error) {
|
||||
if _, err := s.Get(ctx, id); err != nil {
|
||||
return Run{}, err
|
||||
}
|
||||
runID, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return Run{}, err
|
||||
}
|
||||
now := s.now().UTC()
|
||||
_, err = s.pool.Exec(ctx, `INSERT INTO gateway.scheduled_task_runs(id,task_id,trigger_type,scheduled_for) VALUES($1,$2,'manual',$3)`, runID, id, now)
|
||||
if err != nil {
|
||||
return Run{}, err
|
||||
}
|
||||
return s.getRun(ctx, runID)
|
||||
}
|
||||
|
||||
const runSelect = `SELECT r.id::text,r.task_id::text,t.code,r.trigger_type,r.scheduled_for,r.status,r.attempts,r.worker_id,r.started_at,r.finished_at,r.response,r.error,r.created_at FROM gateway.scheduled_task_runs r JOIN gateway.scheduled_tasks t ON t.id=r.task_id`
|
||||
|
||||
func scanRun(row pgx.Row) (Run, error) {
|
||||
var run Run
|
||||
err := row.Scan(&run.ID, &run.TaskID, &run.TaskCode, &run.TriggerType, &run.ScheduledFor, &run.Status, &run.Attempts, &run.WorkerID, &run.StartedAt, &run.FinishedAt, &run.Response, &run.Error, &run.CreatedAt)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return Run{}, ErrNotFound
|
||||
}
|
||||
return run, err
|
||||
}
|
||||
|
||||
func (s *Service) getRun(ctx context.Context, id string) (Run, error) {
|
||||
return scanRun(s.pool.QueryRow(ctx, runSelect+` WHERE r.id=$1`, id))
|
||||
}
|
||||
|
||||
func (s *Service) Runs(ctx context.Context, taskID string, limit int) ([]Run, error) {
|
||||
if limit < 1 || limit > 500 {
|
||||
limit = 100
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, runSelect+` WHERE r.task_id=$1 ORDER BY r.created_at DESC LIMIT $2`, taskID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []Run{}
|
||||
for rows.Next() {
|
||||
item, scanErr := scanRun(rows)
|
||||
if scanErr != nil {
|
||||
return nil, scanErr
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Service) decryptAPIKey(task Task) (string, error) {
|
||||
plain, err := s.cipher.Decrypt(task.encryptedAPIKey, task.apiKeyKEKVersion)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(plain), nil
|
||||
}
|
||||
@@ -52,8 +52,19 @@ func (m *Middleware) Wrap(next http.Handler) http.Handler {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
completed := false
|
||||
defer func() {
|
||||
if !completed {
|
||||
// 主 handler 抛 panic(如代理中途客户端断开引发的
|
||||
// http.ErrAbortHandler)时,shadow 比较 goroutine 从未启动,
|
||||
// 必须在此归还并发槽位,否则连续 N 次断开后 shadow 流量
|
||||
// 被永久静默关闭。
|
||||
<-m.limit
|
||||
}
|
||||
}()
|
||||
capture := &captureWriter{ResponseWriter: w, limit: m.config.MaxBodyBytes}
|
||||
next.ServeHTTP(capture, r)
|
||||
completed = true
|
||||
primaryStatus := capture.Status()
|
||||
primaryBody := append([]byte(nil), capture.body...)
|
||||
headers := shadowHeaders(r.Header)
|
||||
|
||||
@@ -0,0 +1,166 @@
|
||||
package trace
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/identity"
|
||||
"aigateway.local/core/internal/platform/apiresponse"
|
||||
)
|
||||
|
||||
type AdminHTTPHandler struct {
|
||||
store *Store
|
||||
identity *identity.Service
|
||||
mux *http.ServeMux
|
||||
}
|
||||
|
||||
func NewAdminHTTPHandler(store *Store, identityService *identity.Service) *AdminHTTPHandler {
|
||||
h := &AdminHTTPHandler{store: store, identity: identityService, mux: http.NewServeMux()}
|
||||
h.mux.HandleFunc("GET /api/v1/admin/traces", h.list)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/traces/{id}", h.get)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/agent-sessions", h.listSessions)
|
||||
return h
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.mux.ServeHTTP(w, r) }
|
||||
|
||||
func (h *AdminHTTPHandler) require(w http.ResponseWriter, r *http.Request) bool {
|
||||
account, err := h.identity.Authenticate(r.Context(), identity.KindAdmin, r.Header.Get("Authorization"))
|
||||
if err != nil {
|
||||
status := http.StatusUnauthorized
|
||||
if !errors.Is(err, identity.ErrInvalidSession) && !errors.Is(err, identity.ErrNotFound) {
|
||||
status = http.StatusServiceUnavailable
|
||||
}
|
||||
apiresponse.Error(w, status, "登录状态无效或身份服务暂不可用")
|
||||
return false
|
||||
}
|
||||
if !identity.HasPermission(account, identity.PermissionTraceRead) {
|
||||
apiresponse.Error(w, http.StatusForbidden, "缺少 LLM Trace 查看权限")
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) list(w http.ResponseWriter, r *http.Request) {
|
||||
if !h.require(w, r) {
|
||||
return
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
from, err := queryTime(r, "from", now.Add(-24*time.Hour))
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "from 时间无效")
|
||||
return
|
||||
}
|
||||
to, err := queryTime(r, "to", now.Add(time.Second))
|
||||
if err != nil || !to.After(from) || to.Sub(from) > 366*24*time.Hour {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "Trace 查询时间范围无效或超过 366 天")
|
||||
return
|
||||
}
|
||||
limit := 50
|
||||
if value := strings.TrimSpace(r.URL.Query().Get("limit")); value != "" {
|
||||
parsed, scanErr := strconv.Atoi(value)
|
||||
if scanErr != nil || parsed < 1 || parsed > 200 {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "limit 必须在 1 到 200 之间")
|
||||
return
|
||||
}
|
||||
limit = parsed
|
||||
}
|
||||
status := strings.TrimSpace(r.URL.Query().Get("status"))
|
||||
if status != "" && !validStatus(status) {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "Trace 状态无效")
|
||||
return
|
||||
}
|
||||
traceType := strings.TrimSpace(r.URL.Query().Get("trace_type"))
|
||||
if traceType != "" && !validTraceType(traceType) {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "Trace 类型无效")
|
||||
return
|
||||
}
|
||||
items, err := h.store.List(r.Context(), Filter{From: from, To: to, TraceType: traceType, TargetCode: strings.TrimSpace(r.URL.Query().Get("target_code")), RequestID: strings.TrimSpace(r.URL.Query().Get("request_id")), Status: status, Limit: limit})
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusServiceUnavailable, "Trace 查询服务暂不可用")
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) listSessions(w http.ResponseWriter, r *http.Request) {
|
||||
if !h.require(w, r) {
|
||||
return
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
from, err := queryTime(r, "from", now.Add(-30*24*time.Hour))
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "from 时间无效")
|
||||
return
|
||||
}
|
||||
to, err := queryTime(r, "to", now.Add(time.Second))
|
||||
if err != nil || !to.After(from) || to.Sub(from) > 366*24*time.Hour {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "会话查询时间范围无效或超过 366 天")
|
||||
return
|
||||
}
|
||||
limit := 50
|
||||
if value := strings.TrimSpace(r.URL.Query().Get("limit")); value != "" {
|
||||
parsed, scanErr := strconv.Atoi(value)
|
||||
if scanErr != nil || parsed < 1 || parsed > 200 {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "limit 必须在 1 到 200 之间")
|
||||
return
|
||||
}
|
||||
limit = parsed
|
||||
}
|
||||
traceType := strings.TrimSpace(r.URL.Query().Get("trace_type"))
|
||||
if traceType != "" && !validTraceType(traceType) {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "会话类型无效")
|
||||
return
|
||||
}
|
||||
targetCode := strings.TrimSpace(r.URL.Query().Get("target_code"))
|
||||
if len(targetCode) > 128 {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "目标编码过长")
|
||||
return
|
||||
}
|
||||
sessionID := strings.TrimSpace(r.URL.Query().Get("session_id"))
|
||||
if len(sessionID) > 512 {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "会话 ID 过长")
|
||||
return
|
||||
}
|
||||
items, err := h.store.ListSessions(r.Context(), SessionFilter{From: from, To: to, TraceType: traceType, TargetCode: targetCode, SessionID: sessionID, Limit: limit})
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusServiceUnavailable, "智能体会话查询服务暂不可用")
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) get(w http.ResponseWriter, r *http.Request) {
|
||||
if !h.require(w, r) {
|
||||
return
|
||||
}
|
||||
item, err := h.store.Get(r.Context(), r.PathValue("id"))
|
||||
if errors.Is(err, ErrNotFound) {
|
||||
apiresponse.Error(w, http.StatusNotFound, "Trace 不存在")
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusServiceUnavailable, "Trace 查询服务暂不可用")
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func queryTime(r *http.Request, name string, fallback time.Time) (time.Time, error) {
|
||||
value := strings.TrimSpace(r.URL.Query().Get(name))
|
||||
if value == "" {
|
||||
return fallback, nil
|
||||
}
|
||||
return time.Parse(time.RFC3339, value)
|
||||
}
|
||||
|
||||
func validTraceType(value string) bool {
|
||||
return value == "application" || value == "digital_employee"
|
||||
}
|
||||
|
||||
func validStatus(value string) bool {
|
||||
return value == "running" || value == "success" || value == "error"
|
||||
}
|
||||
@@ -0,0 +1,459 @@
|
||||
package trace
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
platformid "aigateway.local/core/internal/platform/id"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
var ErrNotFound = errors.New("trace not found")
|
||||
|
||||
type Store struct{ pool *pgxpool.Pool }
|
||||
|
||||
func NewStore(pool *pgxpool.Pool) *Store { return &Store{pool: pool} }
|
||||
|
||||
type StartInput struct {
|
||||
RequestID string
|
||||
APIKeyID string
|
||||
TenantID *string
|
||||
TraceType string
|
||||
TargetID string
|
||||
TargetCode string
|
||||
ConversationID string
|
||||
Metadata map[string]any
|
||||
}
|
||||
|
||||
type Trace struct {
|
||||
ID string `json:"id"`
|
||||
RequestID string `json:"request_id"`
|
||||
APIKeyID *string `json:"api_key_id,omitempty"`
|
||||
TenantID *string `json:"tenant_id,omitempty"`
|
||||
TraceType string `json:"trace_type"`
|
||||
TargetID *string `json:"target_id,omitempty"`
|
||||
TargetCode string `json:"target_code"`
|
||||
ConversationID string `json:"conversation_id"`
|
||||
Status string `json:"status"`
|
||||
StartedAt time.Time `json:"started_at"`
|
||||
FinishedAt *time.Time `json:"finished_at,omitempty"`
|
||||
LatencyMS *int `json:"latency_ms,omitempty"`
|
||||
RetrievalCount int `json:"retrieval_count"`
|
||||
ModelCallCount int `json:"model_call_count"`
|
||||
ToolCallCount int `json:"tool_call_count"`
|
||||
Error string `json:"error"`
|
||||
Metadata json.RawMessage `json:"metadata"`
|
||||
Spans []Span `json:"spans,omitempty"`
|
||||
}
|
||||
|
||||
type SpanInput struct {
|
||||
TraceID string
|
||||
ParentID string
|
||||
SpanType string
|
||||
Name string
|
||||
Round int
|
||||
ProviderCode string
|
||||
Model string
|
||||
Metadata map[string]any
|
||||
}
|
||||
|
||||
type Span struct {
|
||||
ID string `json:"id"`
|
||||
TraceID string `json:"trace_id"`
|
||||
ParentID *string `json:"parent_id,omitempty"`
|
||||
SpanType string `json:"span_type"`
|
||||
Name string `json:"name"`
|
||||
Status string `json:"status"`
|
||||
StartedAt time.Time `json:"started_at"`
|
||||
FinishedAt *time.Time `json:"finished_at,omitempty"`
|
||||
LatencyMS *int `json:"latency_ms,omitempty"`
|
||||
ProviderCode string `json:"provider_code,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
Round int `json:"round"`
|
||||
Error string `json:"error"`
|
||||
Metadata json.RawMessage `json:"metadata"`
|
||||
}
|
||||
|
||||
type FinishInput struct {
|
||||
Status string
|
||||
Error string
|
||||
RetrievalCount int
|
||||
ModelCallCount int
|
||||
ToolCallCount int
|
||||
Metadata map[string]any
|
||||
}
|
||||
|
||||
type SpanFinishInput struct {
|
||||
Status string
|
||||
Error string
|
||||
InputTokens int64
|
||||
OutputTokens int64
|
||||
ProviderCode string
|
||||
Model string
|
||||
Metadata map[string]any
|
||||
}
|
||||
|
||||
type Filter struct {
|
||||
From time.Time
|
||||
To time.Time
|
||||
TraceType string
|
||||
TargetCode string
|
||||
RequestID string
|
||||
Status string
|
||||
Limit int
|
||||
}
|
||||
|
||||
// Session is a metadata-only aggregation of traces that share a conversation
|
||||
// ID. Stateless requests use a request-derived key so they remain visible in
|
||||
// the session center without pretending to be part of a persistent chat.
|
||||
type Session struct {
|
||||
ID string `json:"id"`
|
||||
TraceType string `json:"trace_type"`
|
||||
TargetCode string `json:"target_code"`
|
||||
TraceCount int `json:"trace_count"`
|
||||
LatestTraceID string `json:"latest_trace_id"`
|
||||
LatestStatus string `json:"latest_status"`
|
||||
StartedAt time.Time `json:"started_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
RetrievalCount int `json:"retrieval_count"`
|
||||
ModelCallCount int `json:"model_call_count"`
|
||||
ToolCallCount int `json:"tool_call_count"`
|
||||
}
|
||||
|
||||
type SessionFilter struct {
|
||||
From time.Time
|
||||
To time.Time
|
||||
TraceType string
|
||||
TargetCode string
|
||||
SessionID string
|
||||
Limit int
|
||||
}
|
||||
|
||||
func metadataJSON(value map[string]any) []byte {
|
||||
if value == nil {
|
||||
return []byte(`{}`)
|
||||
}
|
||||
raw, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
return []byte(`{}`)
|
||||
}
|
||||
return raw
|
||||
}
|
||||
|
||||
func normalizeMetadata(raw []byte) json.RawMessage {
|
||||
if len(raw) == 0 || !json.Valid(raw) {
|
||||
return json.RawMessage(`{}`)
|
||||
}
|
||||
return raw
|
||||
}
|
||||
|
||||
func validateStart(input StartInput) error {
|
||||
if strings.TrimSpace(input.RequestID) == "" || strings.TrimSpace(input.TargetCode) == "" {
|
||||
return errors.New("trace request_id 和 target_code 不能为空")
|
||||
}
|
||||
if input.TraceType != "application" && input.TraceType != "digital_employee" {
|
||||
return errors.New("trace 类型无效")
|
||||
}
|
||||
if len(input.TargetCode) > 128 || len(input.ConversationID) > 128 {
|
||||
return errors.New("trace 目标或会话 ID 过长")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateSpan(input SpanInput) error {
|
||||
if strings.TrimSpace(input.TraceID) == "" || strings.TrimSpace(input.Name) == "" {
|
||||
return errors.New("trace span 标识不能为空")
|
||||
}
|
||||
if input.SpanType != "model" && input.SpanType != "tool" && input.SpanType != "retrieval" {
|
||||
return errors.New("trace span 类型无效")
|
||||
}
|
||||
if input.Round < 0 || len(input.Name) > 256 {
|
||||
return errors.New("trace span 参数无效")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) Start(ctx context.Context, input StartInput) (Trace, error) {
|
||||
if s == nil || s.pool == nil {
|
||||
return Trace{}, errors.New("trace store unavailable")
|
||||
}
|
||||
input.RequestID = strings.TrimSpace(input.RequestID)
|
||||
input.TargetCode = strings.TrimSpace(input.TargetCode)
|
||||
input.ConversationID = strings.TrimSpace(input.ConversationID)
|
||||
if err := validateStart(input); err != nil {
|
||||
return Trace{}, err
|
||||
}
|
||||
id, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return Trace{}, err
|
||||
}
|
||||
_, err = s.pool.Exec(ctx, `INSERT INTO gateway.agent_traces(id,request_id,api_key_id,tenant_id,trace_type,target_id,target_code,conversation_id,metadata) VALUES($1,$2,nullif($3,'')::uuid,nullif($4,'')::uuid,$5,nullif($6,'')::uuid,$7,$8,$9)`, id, input.RequestID, input.APIKeyID, valueOrEmpty(input.TenantID), input.TraceType, input.TargetID, input.TargetCode, input.ConversationID, metadataJSON(input.Metadata))
|
||||
if err != nil {
|
||||
return Trace{}, fmt.Errorf("start trace: %w", err)
|
||||
}
|
||||
return s.Get(ctx, id)
|
||||
}
|
||||
|
||||
func valueOrEmpty(value *string) string {
|
||||
if value == nil {
|
||||
return ""
|
||||
}
|
||||
return *value
|
||||
}
|
||||
|
||||
func (s *Store) StartSpan(ctx context.Context, input SpanInput) (Span, error) {
|
||||
if s == nil || s.pool == nil {
|
||||
return Span{}, errors.New("trace store unavailable")
|
||||
}
|
||||
input.Name = strings.TrimSpace(input.Name)
|
||||
if err := validateSpan(input); err != nil {
|
||||
return Span{}, err
|
||||
}
|
||||
id, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return Span{}, err
|
||||
}
|
||||
_, err = s.pool.Exec(ctx, `INSERT INTO gateway.agent_trace_spans(id,trace_id,parent_id,span_type,name,round,provider_code,model,metadata) VALUES($1,$2,nullif($3,'')::uuid,$4,$5,$6,nullif($7,''),nullif($8,''),$9)`, id, input.TraceID, input.ParentID, input.SpanType, input.Name, input.Round, input.ProviderCode, input.Model, metadataJSON(input.Metadata))
|
||||
if err != nil {
|
||||
return Span{}, fmt.Errorf("start trace span: %w", err)
|
||||
}
|
||||
return s.GetSpan(ctx, id)
|
||||
}
|
||||
|
||||
func (s *Store) Finish(ctx context.Context, id string, input FinishInput) error {
|
||||
if s == nil || s.pool == nil {
|
||||
return errors.New("trace store unavailable")
|
||||
}
|
||||
status := normalizeStatus(input.Status)
|
||||
errorText := truncate(input.Error, 4000)
|
||||
metadata := metadataJSON(input.Metadata)
|
||||
tag, err := s.pool.Exec(ctx, `UPDATE gateway.agent_traces SET status=$2,error=$3,retrieval_count=$4,model_call_count=$5,tool_call_count=$6,metadata=$7,finished_at=clock_timestamp(),latency_ms=(extract(epoch FROM (clock_timestamp()-started_at))*1000)::integer WHERE id=$1 AND status='running'`, id, status, errorText, max(input.RetrievalCount, 0), max(input.ModelCallCount, 0), max(input.ToolCallCount, 0), metadata)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) FinishSpan(ctx context.Context, id string, input SpanFinishInput) error {
|
||||
if s == nil || s.pool == nil {
|
||||
return errors.New("trace store unavailable")
|
||||
}
|
||||
status := normalizeStatus(input.Status)
|
||||
tag, err := s.pool.Exec(ctx, `UPDATE gateway.agent_trace_spans SET status=$2,error=$3,input_tokens=$4,output_tokens=$5,provider_code=coalesce(nullif($6,''),provider_code),model=coalesce(nullif($7,''),model),metadata=$8,finished_at=clock_timestamp(),latency_ms=(extract(epoch FROM (clock_timestamp()-started_at))*1000)::integer WHERE id=$1 AND status='running'`, id, status, truncate(input.Error, 4000), max(input.InputTokens, 0), max(input.OutputTokens, 0), input.ProviderCode, input.Model, metadataJSON(input.Metadata))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func normalizeStatus(value string) string {
|
||||
if value == "success" {
|
||||
return "success"
|
||||
}
|
||||
return "error"
|
||||
}
|
||||
|
||||
func truncate(value string, limit int) string {
|
||||
value = strings.TrimSpace(value)
|
||||
if len(value) <= limit {
|
||||
return value
|
||||
}
|
||||
cut := value[:limit]
|
||||
// 按字节截断可能切半多字节 rune,产生无效 UTF-8 使 trace 落库失败;
|
||||
// 回退到最近一个完整 rune 的边界。
|
||||
for len(cut) > 0 && !utf8.RuneStart(cut[len(cut)-1]) {
|
||||
cut = cut[:len(cut)-1]
|
||||
}
|
||||
return cut
|
||||
}
|
||||
|
||||
const traceSelect = `SELECT id::text,request_id,api_key_id::text,tenant_id::text,trace_type,target_id::text,target_code,conversation_id,status,started_at,finished_at,latency_ms,retrieval_count,model_call_count,tool_call_count,error,metadata FROM gateway.agent_traces`
|
||||
|
||||
func scanTrace(row pgx.Row) (Trace, error) {
|
||||
var item Trace
|
||||
err := row.Scan(&item.ID, &item.RequestID, &item.APIKeyID, &item.TenantID, &item.TraceType, &item.TargetID, &item.TargetCode, &item.ConversationID, &item.Status, &item.StartedAt, &item.FinishedAt, &item.LatencyMS, &item.RetrievalCount, &item.ModelCallCount, &item.ToolCallCount, &item.Error, &item.Metadata)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return Trace{}, ErrNotFound
|
||||
}
|
||||
item.Metadata = normalizeMetadata(item.Metadata)
|
||||
return item, err
|
||||
}
|
||||
|
||||
func scanSpan(row pgx.Row) (Span, error) {
|
||||
var item Span
|
||||
err := row.Scan(&item.ID, &item.TraceID, &item.ParentID, &item.SpanType, &item.Name, &item.Status, &item.StartedAt, &item.FinishedAt, &item.LatencyMS, &item.ProviderCode, &item.Model, &item.InputTokens, &item.OutputTokens, &item.Round, &item.Error, &item.Metadata)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return Span{}, ErrNotFound
|
||||
}
|
||||
item.Metadata = normalizeMetadata(item.Metadata)
|
||||
return item, err
|
||||
}
|
||||
|
||||
func (s *Store) Get(ctx context.Context, id string) (Trace, error) {
|
||||
if s == nil || s.pool == nil {
|
||||
return Trace{}, errors.New("trace store unavailable")
|
||||
}
|
||||
item, err := scanTrace(s.pool.QueryRow(ctx, traceSelect+` WHERE id=$1`, id))
|
||||
if err != nil {
|
||||
return Trace{}, err
|
||||
}
|
||||
spans, err := s.listSpans(ctx, id)
|
||||
if err != nil {
|
||||
return Trace{}, err
|
||||
}
|
||||
item.Spans = spans
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func (s *Store) GetSpan(ctx context.Context, id string) (Span, error) {
|
||||
if s == nil || s.pool == nil {
|
||||
return Span{}, errors.New("trace store unavailable")
|
||||
}
|
||||
return scanSpan(s.pool.QueryRow(ctx, `SELECT id::text,trace_id::text,parent_id::text,span_type,name,status,started_at,finished_at,latency_ms,coalesce(provider_code,''),coalesce(model,''),input_tokens,output_tokens,round,error,metadata FROM gateway.agent_trace_spans WHERE id=$1`, id))
|
||||
}
|
||||
|
||||
func (s *Store) listSpans(ctx context.Context, traceID string) ([]Span, error) {
|
||||
if s == nil || s.pool == nil {
|
||||
return nil, errors.New("trace store unavailable")
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `SELECT id::text,trace_id::text,parent_id::text,span_type,name,status,started_at,finished_at,latency_ms,coalesce(provider_code,''),coalesce(model,''),input_tokens,output_tokens,round,error,metadata FROM gateway.agent_trace_spans WHERE trace_id=$1 ORDER BY started_at,id`, traceID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []Span{}
|
||||
for rows.Next() {
|
||||
item, scanErr := scanSpan(rows)
|
||||
if scanErr != nil {
|
||||
return nil, scanErr
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) List(ctx context.Context, filter Filter) ([]Trace, error) {
|
||||
if s == nil || s.pool == nil {
|
||||
return nil, errors.New("trace store unavailable")
|
||||
}
|
||||
if filter.Limit < 1 || filter.Limit > 200 {
|
||||
filter.Limit = 50
|
||||
}
|
||||
where := []string{"started_at >= $1", "started_at < $2"}
|
||||
args := []any{filter.From, filter.To}
|
||||
add := func(condition string, value any) {
|
||||
args = append(args, value)
|
||||
where = append(where, fmt.Sprintf(condition, len(args)))
|
||||
}
|
||||
if filter.TraceType != "" {
|
||||
add("trace_type = $%d", filter.TraceType)
|
||||
}
|
||||
if filter.TargetCode != "" {
|
||||
add("target_code = $%d", filter.TargetCode)
|
||||
}
|
||||
if filter.RequestID != "" {
|
||||
add("request_id = $%d", filter.RequestID)
|
||||
}
|
||||
if filter.Status != "" {
|
||||
add("status = $%d", filter.Status)
|
||||
}
|
||||
args = append(args, filter.Limit)
|
||||
query := traceSelect + ` WHERE ` + strings.Join(where, " AND ") + ` ORDER BY started_at DESC,id DESC LIMIT $` + strconv.Itoa(len(args))
|
||||
rows, err := s.pool.Query(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []Trace{}
|
||||
for rows.Next() {
|
||||
item, scanErr := scanTrace(rows)
|
||||
if scanErr != nil {
|
||||
return nil, scanErr
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
const sessionGroupSelect = `
|
||||
WITH grouped AS (
|
||||
SELECT trace_type,
|
||||
target_code,
|
||||
coalesce(nullif(conversation_id, ''), 'request:' || request_id) AS session_key,
|
||||
count(*)::int AS trace_count,
|
||||
(array_agg(id::text ORDER BY started_at DESC, id DESC))[1] AS latest_trace_id,
|
||||
(array_agg(status ORDER BY started_at DESC, id DESC))[1] AS latest_status,
|
||||
min(started_at) AS started_at,
|
||||
max(coalesce(finished_at, started_at)) AS updated_at,
|
||||
sum(retrieval_count)::int AS retrieval_count,
|
||||
sum(model_call_count)::int AS model_call_count,
|
||||
sum(tool_call_count)::int AS tool_call_count
|
||||
FROM gateway.agent_traces`
|
||||
|
||||
func scanSession(row pgx.Row) (Session, error) {
|
||||
var item Session
|
||||
err := row.Scan(&item.ID, &item.TraceType, &item.TargetCode, &item.TraceCount, &item.LatestTraceID, &item.LatestStatus, &item.StartedAt, &item.UpdatedAt, &item.RetrievalCount, &item.ModelCallCount, &item.ToolCallCount)
|
||||
return item, err
|
||||
}
|
||||
|
||||
func (s *Store) ListSessions(ctx context.Context, filter SessionFilter) ([]Session, error) {
|
||||
if s == nil || s.pool == nil {
|
||||
return nil, errors.New("trace store unavailable")
|
||||
}
|
||||
if filter.Limit < 1 || filter.Limit > 200 {
|
||||
filter.Limit = 50
|
||||
}
|
||||
innerWhere := []string{"started_at >= $1", "started_at < $2"}
|
||||
args := []any{filter.From, filter.To}
|
||||
addInner := func(condition string, value any) {
|
||||
args = append(args, value)
|
||||
innerWhere = append(innerWhere, fmt.Sprintf(condition, len(args)))
|
||||
}
|
||||
if filter.TraceType != "" {
|
||||
addInner("trace_type = $%d", filter.TraceType)
|
||||
}
|
||||
if filter.TargetCode != "" {
|
||||
addInner("target_code = $%d", filter.TargetCode)
|
||||
}
|
||||
outerWhere := []string{}
|
||||
if filter.SessionID != "" {
|
||||
args = append(args, filter.SessionID)
|
||||
outerWhere = append(outerWhere, fmt.Sprintf("trace_type || ':' || target_code || ':' || session_key = $%d", len(args)))
|
||||
}
|
||||
args = append(args, filter.Limit)
|
||||
limitArg := strconv.Itoa(len(args))
|
||||
query := sessionGroupSelect + ` WHERE ` + strings.Join(innerWhere, " AND ") + ` GROUP BY trace_type,target_code,session_key) SELECT trace_type || ':' || target_code || ':' || session_key AS id,trace_type,target_code,trace_count,latest_trace_id,latest_status,started_at,updated_at,retrieval_count,model_call_count,tool_call_count FROM grouped`
|
||||
if len(outerWhere) > 0 {
|
||||
query += ` WHERE ` + strings.Join(outerWhere, " AND ")
|
||||
}
|
||||
query += ` ORDER BY updated_at DESC,id DESC LIMIT $` + limitArg
|
||||
rows, err := s.pool.Query(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []Session{}
|
||||
for rows.Next() {
|
||||
item, scanErr := scanSession(rows)
|
||||
if scanErr != nil {
|
||||
return nil, scanErr
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
package trace
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/platform/config"
|
||||
"aigateway.local/core/internal/platform/database"
|
||||
)
|
||||
|
||||
func TestTracePostgreSQLLifecycle(t *testing.T) {
|
||||
databaseURL := os.Getenv("TRACE_TEST_DATABASE_URL")
|
||||
if databaseURL == "" {
|
||||
t.Skip("TRACE_TEST_DATABASE_URL is not set")
|
||||
}
|
||||
ctx := context.Background()
|
||||
pool, err := database.Open(ctx, config.Database{URL: databaseURL, MaxConns: 4, MinConns: 0})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer pool.Close()
|
||||
store := NewStore(pool)
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM gateway.agent_traces WHERE request_id LIKE 'trace-integration%'`)
|
||||
item, err := store.Start(ctx, StartInput{RequestID: "trace-integration", TraceType: "application", TargetID: "77777777-7777-4777-8777-777777777777", TargetCode: "trace_app", ConversationID: "conversation-1", Metadata: map[string]any{"version": 1}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer pool.Exec(ctx, `DELETE FROM gateway.agent_traces WHERE request_id LIKE 'trace-integration%'`)
|
||||
model, err := store.StartSpan(ctx, SpanInput{TraceID: item.ID, SpanType: "model", Name: "chat.completions", Round: 0, ProviderCode: "test", Model: "test-model"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = store.FinishSpan(ctx, model.ID, SpanFinishInput{Status: "success", InputTokens: 12, OutputTokens: 8, Metadata: map[string]any{"http_status": 200}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tool, err := store.StartSpan(ctx, SpanInput{TraceID: item.ID, SpanType: "tool", Name: "lookup", Round: 0})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = store.FinishSpan(ctx, tool.ID, SpanFinishInput{Status: "error", Error: "upstream timeout"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = store.Finish(ctx, item.ID, FinishInput{Status: "success", RetrievalCount: 2, ModelCallCount: 1, ToolCallCount: 1}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
loaded, err := store.Get(ctx, item.ID)
|
||||
if err != nil || loaded.Status != "success" || loaded.ModelCallCount != 1 || len(loaded.Spans) != 2 {
|
||||
t.Fatalf("loaded=%+v err=%v", loaded, err)
|
||||
}
|
||||
if loaded.Spans[0].ProviderCode != "test" || loaded.Spans[0].Model != "test-model" {
|
||||
t.Fatalf("model span route metadata was not preserved: %+v", loaded.Spans[0])
|
||||
}
|
||||
if loaded.LatencyMS == nil || loaded.StartedAt.After(time.Now().UTC().Add(time.Second)) {
|
||||
t.Fatalf("invalid timing: %+v", loaded)
|
||||
}
|
||||
items, err := store.List(ctx, Filter{From: time.Now().Add(-time.Minute), To: time.Now().Add(time.Minute), TargetCode: "trace_app", Limit: 10})
|
||||
if err != nil || len(items) != 1 || items[0].ID != item.ID {
|
||||
t.Fatalf("list=%+v err=%v", items, err)
|
||||
}
|
||||
sessions, err := store.ListSessions(ctx, SessionFilter{From: time.Now().Add(-time.Minute), To: time.Now().Add(time.Minute), SessionID: "application:trace_app:conversation-1", Limit: 10})
|
||||
if err != nil || len(sessions) != 1 || sessions[0].TraceCount != 1 || sessions[0].LatestTraceID != item.ID || sessions[0].ModelCallCount != 1 {
|
||||
t.Fatalf("sessions=%+v err=%v", sessions, err)
|
||||
}
|
||||
digital, err := store.Start(ctx, StartInput{RequestID: "trace-integration-digital", TraceType: "digital_employee", TargetID: "88888888-8888-4888-8888-888888888888", TargetCode: "trace_employee", ConversationID: "conversation-1"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = store.Finish(ctx, digital.ID, FinishInput{Status: "success", ModelCallCount: 1}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
allSessions, err := store.ListSessions(ctx, SessionFilter{From: time.Now().Add(-time.Minute), To: time.Now().Add(time.Minute), Limit: 10})
|
||||
if err != nil || len(allSessions) < 2 {
|
||||
t.Fatalf("all sessions=%+v err=%v", allSessions, err)
|
||||
}
|
||||
seenTypes := map[string]bool{}
|
||||
for _, session := range allSessions {
|
||||
if strings.HasPrefix(session.ID, "application:trace_app:") || strings.HasPrefix(session.ID, "digital_employee:trace_employee:") {
|
||||
seenTypes[session.TraceType] = true
|
||||
}
|
||||
}
|
||||
if !seenTypes["application"] || !seenTypes["digital_employee"] {
|
||||
t.Fatalf("session types=%v", seenTypes)
|
||||
}
|
||||
}
|
||||
@@ -56,8 +56,13 @@ func normalizeApplicationConfig(config *ApplicationConfig, requireModel bool) er
|
||||
if config.Temperature < 0 || config.Temperature > 2 {
|
||||
return errors.New("temperature 应在 0-2 之间")
|
||||
}
|
||||
if config.MaxToolRounds < 0 || config.MaxToolRounds > 8 {
|
||||
return errors.New("max_tool_rounds 应在 0-8 之间")
|
||||
// 缺省 max_tool_rounds(0)时按 5 轮处理,与数字员工一致;否则运行时
|
||||
// round >= 0 在第一次工具调用前就判定"已达上限",应用永远无法完成工具调用。
|
||||
if config.MaxToolRounds == 0 {
|
||||
config.MaxToolRounds = 5
|
||||
}
|
||||
if config.MaxToolRounds < 1 || config.MaxToolRounds > 8 {
|
||||
return errors.New("max_tool_rounds 应在 1-8 之间")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -93,6 +93,10 @@ func (h *FilesAdminHTTPHandler) upload(w http.ResponseWriter, r *http.Request) {
|
||||
contentType = ct
|
||||
}
|
||||
}
|
||||
// 非 multipart 请求(body 为 nil 时)按原始请求体上传(?filename= 指定文件名)。
|
||||
if body == nil {
|
||||
body = r.Body
|
||||
}
|
||||
if originalName == "" {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "缺少文件名")
|
||||
return
|
||||
@@ -150,6 +154,8 @@ func serveFileContent(w http.ResponseWriter, r *http.Request, files *FileService
|
||||
}
|
||||
defer reader.Close()
|
||||
w.Header().Set("Content-Disposition", "attachment; filename*=UTF-8''"+url.PathEscape(obj.OriginalName))
|
||||
// Content-Type 来自用户上传,回显前必须禁 MIME 嗅探,防止存储型 XSS。
|
||||
w.Header().Set("X-Content-Type-Options", "nosniff")
|
||||
w.Header().Set("Content-Type", obj.ContentType)
|
||||
w.Header().Set("Content-Length", strconv.FormatInt(size, 10))
|
||||
_, _ = io.Copy(w, reader)
|
||||
|
||||
@@ -0,0 +1,378 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// InboxMessage 是站内消息的一行:一条消息对应一个收件人(admin 或 portal)。
|
||||
// 管理员广播时按目标用户逐行落库,因此每行独立维护 read_at 已读回执。
|
||||
type InboxMessage struct {
|
||||
ID string `json:"id"`
|
||||
RecipientKind string `json:"recipient_kind"`
|
||||
RecipientUserID string `json:"recipient_user_id"`
|
||||
SenderType string `json:"sender_type"`
|
||||
Category string `json:"category"`
|
||||
Title string `json:"title"`
|
||||
Body string `json:"body"`
|
||||
Link string `json:"link"`
|
||||
Payload json.RawMessage `json:"payload,omitempty"`
|
||||
ReadAt *time.Time `json:"read_at,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
type InboxInput struct {
|
||||
RecipientKind string `json:"recipient_kind"`
|
||||
Category string `json:"category"`
|
||||
Title string `json:"title"`
|
||||
Body string `json:"body"`
|
||||
Link string `json:"link"`
|
||||
Payload json.RawMessage `json:"payload,omitempty"`
|
||||
}
|
||||
|
||||
// InboxService 负责将 outbox 事件物化为站内消息,并对外提供收件箱读写。
|
||||
// 未读数以 PostgreSQL 为权威源(部分索引 COUNT 快速),Redis 仅作实时 PUBLISH
|
||||
// 提示(为未来 SSE 预留);因此未读计数不会因 Redis 抖动或消息漂移而失真。
|
||||
type InboxService struct {
|
||||
assets *Service
|
||||
redis *redis.Client
|
||||
channel string
|
||||
}
|
||||
|
||||
func NewInboxService(assets *Service, client *redis.Client, channel string) *InboxService {
|
||||
return &InboxService{assets: assets, redis: client, channel: channel}
|
||||
}
|
||||
|
||||
const inboxSelect = `SELECT id::text,recipient_kind,recipient_user_id::text,sender_type,category,title,body,link,payload,read_at,created_at FROM gateway.inbox_messages`
|
||||
|
||||
func scanInboxMessage(row pgx.Row) (InboxMessage, error) {
|
||||
var m InboxMessage
|
||||
err := row.Scan(&m.ID, &m.RecipientKind, &m.RecipientUserID, &m.SenderType, &m.Category, &m.Title, &m.Body, &m.Link, &m.Payload, &m.ReadAt, &m.CreatedAt)
|
||||
return m, err
|
||||
}
|
||||
|
||||
// inboxDraft 描述一个 outbox 事件要落成的一条站内消息(收件人解析方式不同)。
|
||||
type inboxDraft struct {
|
||||
RecipientKind string // admin | portal
|
||||
Category string
|
||||
Title string
|
||||
Body string
|
||||
Link string
|
||||
UserID string // 直接收件人(从 payload 取),空串表示需额外解析
|
||||
AllAdmins bool // 收件人 = 全部启用管理员
|
||||
RequestUser bool // 收件人 = model_access_requests.portal_user_id(payload.request_id)
|
||||
}
|
||||
|
||||
func payloadValue(payload json.RawMessage, key string) string {
|
||||
var values map[string]any
|
||||
if err := json.Unmarshal(payload, &values); err != nil {
|
||||
return ""
|
||||
}
|
||||
value, ok := values[key]
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
switch v := value.(type) {
|
||||
case string:
|
||||
return v
|
||||
case float64:
|
||||
return strconv.FormatFloat(v, 'f', -1, 64)
|
||||
}
|
||||
encoded, _ := json.Marshal(value)
|
||||
return string(encoded)
|
||||
}
|
||||
|
||||
// inboxPlan 把事件类型映射成站内消息草稿(纯函数,便于单测)。未知事件返回 nil,
|
||||
// 不构成错误——并非所有 outbox 事件都需要站内信。
|
||||
func inboxPlan(eventType string, payload json.RawMessage) []inboxDraft {
|
||||
switch eventType {
|
||||
case "model_access.requested":
|
||||
return []inboxDraft{{RecipientKind: "admin", Category: "approval", Title: "新的模型访问申请", Body: "用户申请访问模型 " + payloadValue(payload, "model"), Link: "/security/governance", AllAdmins: true}}
|
||||
case "model_access.decided":
|
||||
text := "已批准"
|
||||
if payloadValue(payload, "status") == "rejected" {
|
||||
text = "已驳回"
|
||||
}
|
||||
return []inboxDraft{{RecipientKind: "portal", Category: "approval", Title: "模型申请已处理", Body: "您的模型访问申请已被" + text, Link: "/portal/access", RequestUser: true}}
|
||||
case "marketplace.installed":
|
||||
return []inboxDraft{{RecipientKind: "portal", Category: "resource", Title: "资源已安装", Body: "资源 " + payloadValue(payload, "code") + " 已安装到您的工作区", Link: "/portal/marketplace", UserID: payloadValue(payload, "portal_user_id")}}
|
||||
case "knowledge_document.ready":
|
||||
return []inboxDraft{{RecipientKind: "admin", Category: "system", Title: "知识文档已入库", Body: "文档已分块入库(" + payloadValue(payload, "chunk_count") + " 切片)", Link: "/assets/knowledge", UserID: payloadValue(payload, "actor_id")}}
|
||||
case "knowledge_document.reprocessed":
|
||||
return []inboxDraft{{RecipientKind: "admin", Category: "system", Title: "知识文档已重新处理", Body: "文档已重新分块入库", Link: "/assets/knowledge", UserID: payloadValue(payload, "actor_id")}}
|
||||
case "knowledge_document.embedding_failed":
|
||||
return []inboxDraft{{RecipientKind: "admin", Category: "system", Title: "知识文档向量化失败", Body: "文档已入库但向量化失败,检索将回退全文检索;请检查 Ollama 后重新处理", Link: "/assets/knowledge", UserID: payloadValue(payload, "actor_id")}}
|
||||
case "scheduled_task.completed":
|
||||
return []inboxDraft{{RecipientKind: "admin", Category: "task_result", Title: "定时任务已执行", Body: "定时任务 " + payloadValue(payload, "task_code") + " 已完成", Link: "/system/scheduled-tasks", UserID: payloadValue(payload, "actor_id")}}
|
||||
case "scheduled_task.failed":
|
||||
return []inboxDraft{{RecipientKind: "admin", Category: "task_result", Title: "定时任务执行失败", Body: "定时任务 " + payloadValue(payload, "task_code") + " 执行失败: " + payloadValue(payload, "error"), Link: "/system/scheduled-tasks", UserID: payloadValue(payload, "actor_id")}}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *InboxService) resolveRecipients(ctx context.Context, draft inboxDraft, payload json.RawMessage) ([]string, error) {
|
||||
switch {
|
||||
case draft.UserID != "":
|
||||
return []string{draft.UserID}, nil
|
||||
case draft.RequestUser:
|
||||
var userID string
|
||||
if err := s.assets.pool.QueryRow(ctx, `SELECT portal_user_id::text FROM gateway.model_access_requests WHERE id=$1`, payloadValue(payload, "request_id")).Scan(&userID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []string{userID}, nil
|
||||
case draft.AllAdmins:
|
||||
rows, err := s.assets.pool.Query(ctx, `SELECT id::text FROM gateway.admin_accounts WHERE active`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
ids := []string{}
|
||||
for rows.Next() {
|
||||
var id string
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids, rows.Err()
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Materialize 把一条 outbox 事件落成站内消息。以 (source_event_id, 收件人) 幂等:
|
||||
// 事件重放时 ON CONFLICT DO NOTHING,不产生重复消息,也不报错。
|
||||
func (s *InboxService) Materialize(ctx context.Context, eventID, eventType string, payload json.RawMessage) error {
|
||||
for _, draft := range inboxPlan(eventType, payload) {
|
||||
recipients, err := s.resolveRecipients(ctx, draft, payload)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, userID := range recipients {
|
||||
if err := s.notify(ctx, eventID, draft, userID, payload); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *InboxService) notify(ctx context.Context, sourceEventID string, draft inboxDraft, userID string, payload json.RawMessage) error {
|
||||
id, err := newUUID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// 数据库 CHECK 按字符数(length)校验,Go 的 len() 是字节数;多字节文本下
|
||||
// 字节校验通过但字符数超限,INSERT 会失败并把事件永远卡在 pending。
|
||||
// 落库前按 rune 数截断,保证任何语言正文都能入库。
|
||||
title := runeTruncate(draft.Title, 256)
|
||||
body := runeTruncate(draft.Body, 4000)
|
||||
link := runeTruncate(draft.Link, 512)
|
||||
tag, err := s.assets.pool.Exec(ctx, `INSERT INTO gateway.inbox_messages(id,source_event_id,recipient_kind,recipient_user_id,sender_type,category,title,body,link,payload)
|
||||
VALUES($1,$2,$3,$4,'system',$5,$6,$7,$8,$9)
|
||||
ON CONFLICT (source_event_id, recipient_kind, recipient_user_id) WHERE source_event_id IS NOT NULL DO NOTHING`,
|
||||
id, sourceEventID, draft.RecipientKind, userID, draft.Category, title, body, link, payload)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return nil // 同一来源事件对同一收件人的重放,跳过
|
||||
}
|
||||
s.publish(draft.RecipientKind, userID)
|
||||
return nil
|
||||
}
|
||||
|
||||
// runeTruncate 按 rune(字符)数截断;超出 limit 个字符时截到第 limit 个
|
||||
// 完整 rune,绝不产生无效 UTF-8。
|
||||
func runeTruncate(value string, limit int) string {
|
||||
if utf8.RuneCountInString(value) <= limit {
|
||||
return value
|
||||
}
|
||||
runes := []rune(value)
|
||||
return string(runes[:limit])
|
||||
}
|
||||
|
||||
// publish 仅作实时提示(未来 SSE 可订阅);收件箱未读数以 DB 为准。
|
||||
func (s *InboxService) publish(kind, userID string) {
|
||||
if s.redis == nil {
|
||||
return
|
||||
}
|
||||
_ = s.redis.Publish(context.WithoutCancel(context.Background()), s.channel, kind+":"+userID).Err()
|
||||
}
|
||||
|
||||
// List 返回某个收件人的收件箱(倒序)。
|
||||
func (s *InboxService) List(ctx context.Context, kind, userID string, limit int) ([]InboxMessage, error) {
|
||||
if limit < 1 {
|
||||
limit = 50
|
||||
}
|
||||
if limit > 200 {
|
||||
limit = 200
|
||||
}
|
||||
rows, err := s.assets.pool.Query(ctx, inboxSelect+` WHERE recipient_kind=$1 AND recipient_user_id=$2 ORDER BY created_at DESC LIMIT $3`, kind, userID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []InboxMessage{}
|
||||
for rows.Next() {
|
||||
m, err := scanInboxMessage(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, m)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
// UnreadCount 以 PostgreSQL 为权威源统计未读消息(部分索引快速扫描)。
|
||||
func (s *InboxService) UnreadCount(ctx context.Context, kind, userID string) (int, error) {
|
||||
var count int
|
||||
err := s.assets.pool.QueryRow(ctx, `SELECT count(*) FROM gateway.inbox_messages WHERE recipient_kind=$1 AND recipient_user_id=$2 AND read_at IS NULL`, kind, userID).Scan(&count)
|
||||
return count, err
|
||||
}
|
||||
|
||||
// MarkRead 把某条消息标记为已读;只允许收件人本人操作。
|
||||
func (s *InboxService) MarkRead(ctx context.Context, id, kind, userID string) (bool, error) {
|
||||
tag, err := s.assets.pool.Exec(ctx, `UPDATE gateway.inbox_messages SET read_at=coalesce(read_at,clock_timestamp()) WHERE id=$1 AND recipient_kind=$2 AND recipient_user_id=$3`, id, kind, userID)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return tag.RowsAffected() > 0, nil
|
||||
}
|
||||
|
||||
// MarkAllRead 把某收件人的全部未读标记为已读,返回标记条数。
|
||||
func (s *InboxService) MarkAllRead(ctx context.Context, kind, userID string) (int, error) {
|
||||
tag, err := s.assets.pool.Exec(ctx, `UPDATE gateway.inbox_messages SET read_at=clock_timestamp() WHERE recipient_kind=$1 AND recipient_user_id=$2 AND read_at IS NULL`, kind, userID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return int(tag.RowsAffected()), nil
|
||||
}
|
||||
|
||||
// AdminList 返回管理端消息中心列表。scope=mine 看发给自己(admin)的消息;
|
||||
// scope=broadcasts 看管理员发起的广播(portal 收件)。其余 scope 全部返回。
|
||||
func (s *InboxService) AdminList(ctx context.Context, adminID, scope string, limit int) ([]InboxMessage, error) {
|
||||
if limit < 1 {
|
||||
limit = 50
|
||||
}
|
||||
if limit > 500 {
|
||||
limit = 500
|
||||
}
|
||||
query, args := inboxSelect+` WHERE`, []any{}
|
||||
switch scope {
|
||||
case "mine":
|
||||
query += ` recipient_kind='admin' AND recipient_user_id=$1`
|
||||
args = append(args, adminID)
|
||||
case "broadcasts":
|
||||
query += ` sender_type='admin' AND recipient_kind='portal'`
|
||||
default:
|
||||
query += ` recipient_kind='admin' AND recipient_user_id=$1`
|
||||
args = append(args, adminID)
|
||||
}
|
||||
args = append(args, limit)
|
||||
query += ` ORDER BY created_at DESC LIMIT $` + strconv.Itoa(len(args))
|
||||
rows, err := s.assets.pool.Query(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []InboxMessage{}
|
||||
for rows.Next() {
|
||||
m, err := scanInboxMessage(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, m)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
// Broadcast 由管理员向 portal(可按部门过滤)或全部 admin 发广播;逐收件人落库,
|
||||
// 返回实际送达用户数。sender_type='admin',无 source_event_id(不与事件幂等键冲突)。
|
||||
func (s *InboxService) Broadcast(ctx context.Context, input InboxInput, departmentIDs []string, actorID string) (int, error) {
|
||||
input.RecipientKind = strings.TrimSpace(input.RecipientKind)
|
||||
input.Category = strings.TrimSpace(input.Category)
|
||||
input.Title = strings.TrimSpace(input.Title)
|
||||
input.Body = strings.TrimSpace(input.Body)
|
||||
input.Link = strings.TrimSpace(input.Link)
|
||||
if input.RecipientKind != "portal" && input.RecipientKind != "admin" {
|
||||
return 0, errors.New("广播对象必须是 admin 或 portal")
|
||||
}
|
||||
if input.Title == "" || utf8.RuneCountInString(input.Title) > 256 || utf8.RuneCountInString(input.Body) > 4000 || utf8.RuneCountInString(input.Link) > 512 {
|
||||
return 0, errors.New("广播标题或正文格式无效")
|
||||
}
|
||||
if !validInboxLink(input.Link) {
|
||||
return 0, errors.New("跳转链接仅支持站内绝对路径或 http(s) 地址")
|
||||
}
|
||||
if input.Category != "system" && input.Category != "approval" && input.Category != "task_result" && input.Category != "resource" {
|
||||
input.Category = "system"
|
||||
}
|
||||
var query string
|
||||
var args []any
|
||||
if input.RecipientKind == "admin" {
|
||||
query = `SELECT id::text FROM gateway.admin_accounts WHERE active`
|
||||
} else {
|
||||
query = `SELECT id::text FROM gateway.portal_users WHERE active AND (array_length($1::uuid[],1) IS NULL OR department_id = ANY($1::uuid[]))`
|
||||
args = append(args, departmentIDs)
|
||||
}
|
||||
rows, err := s.assets.pool.Query(ctx, query, args...)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer rows.Close()
|
||||
recipients := []string{}
|
||||
for rows.Next() {
|
||||
var id string
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
recipients = append(recipients, id)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
// 广播整体包在单个事务里:中途失败不留半套消息,收件人数目与落库一致。
|
||||
tx, err := s.assets.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
count := 0
|
||||
for _, userID := range recipients {
|
||||
id, idErr := newUUID()
|
||||
if idErr != nil {
|
||||
return 0, idErr
|
||||
}
|
||||
tag, insertErr := tx.Exec(ctx, `INSERT INTO gateway.inbox_messages(id,recipient_kind,recipient_user_id,sender_type,category,title,body,link,payload) VALUES($1,$2,$3,'admin',$4,$5,$6,$7,$8)`, id, input.RecipientKind, userID, input.Category, input.Title, input.Body, input.Link, input.Payload)
|
||||
if insertErr != nil {
|
||||
return 0, insertErr
|
||||
}
|
||||
count += int(tag.RowsAffected())
|
||||
s.publish(input.RecipientKind, userID)
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
// validInboxLink blocks executable and browser-special schemes. Empty links are allowed;
|
||||
// internal links must be root-relative, while external links are limited to HTTP(S).
|
||||
func validInboxLink(link string) bool {
|
||||
if link == "" {
|
||||
return true
|
||||
}
|
||||
if strings.HasPrefix(link, "/") && !strings.HasPrefix(link, "//") {
|
||||
return true
|
||||
}
|
||||
parsed, err := url.ParseRequestURI(link)
|
||||
return err == nil && (parsed.Scheme == "http" || parsed.Scheme == "https") && parsed.Host != ""
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
|
||||
"aigateway.local/core/internal/identity"
|
||||
"aigateway.local/core/internal/platform/apiresponse"
|
||||
)
|
||||
|
||||
// InboxAdminHTTPHandler 向管理端暴露站内消息:消息中心、未读数、广播与读回执。
|
||||
type InboxAdminHTTPHandler struct {
|
||||
inbox *InboxService
|
||||
identity *identity.Service
|
||||
mux *http.ServeMux
|
||||
}
|
||||
|
||||
func NewInboxAdminHTTPHandler(inbox *InboxService, identityService *identity.Service) *InboxAdminHTTPHandler {
|
||||
h := &InboxAdminHTTPHandler{inbox: inbox, identity: identityService, mux: http.NewServeMux()}
|
||||
h.mux.HandleFunc("GET /api/v1/admin/inbox", h.list)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/inbox/unread", h.unread)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/inbox/broadcast", h.broadcast)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/inbox/read-all", h.readAll)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/inbox/{id}/read", h.read)
|
||||
return h
|
||||
}
|
||||
|
||||
func (h *InboxAdminHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.mux.ServeHTTP(w, r) }
|
||||
|
||||
func (h *InboxAdminHTTPHandler) 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 (h *InboxAdminHTTPHandler) list(w http.ResponseWriter, r *http.Request) {
|
||||
account, ok := h.require(w, r, identity.PermissionInboxRead)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.inbox.AdminList(r.Context(), account.ID, r.URL.Query().Get("scope"), 100)
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
|
||||
func (h *InboxAdminHTTPHandler) unread(w http.ResponseWriter, r *http.Request) {
|
||||
account, ok := h.require(w, r, identity.PermissionInboxRead)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
count, err := h.inbox.UnreadCount(r.Context(), "admin", account.ID)
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]int{"unread": count})
|
||||
}
|
||||
|
||||
func (h *InboxAdminHTTPHandler) broadcast(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionInboxManage); !ok {
|
||||
return
|
||||
}
|
||||
var input struct {
|
||||
RecipientKind string `json:"recipient_kind"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
Category string `json:"category"`
|
||||
Title string `json:"title"`
|
||||
Body string `json:"body"`
|
||||
Link string `json:"link"`
|
||||
Payload json.RawMessage `json:"payload,omitempty"`
|
||||
}
|
||||
if !decodeAsset(w, r, &input) {
|
||||
return
|
||||
}
|
||||
count, err := h.inbox.Broadcast(r.Context(), InboxInput{
|
||||
RecipientKind: input.RecipientKind, Category: input.Category,
|
||||
Title: input.Title, Body: input.Body, Link: input.Link, Payload: input.Payload,
|
||||
}, input.DepartmentIDs, "")
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]any{"sent": count, "ok": true})
|
||||
}
|
||||
|
||||
func (h *InboxAdminHTTPHandler) read(w http.ResponseWriter, r *http.Request) {
|
||||
account, ok := h.require(w, r, identity.PermissionInboxRead)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
changed, err := h.inbox.MarkRead(r.Context(), r.PathValue("id"), "admin", account.ID)
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"read": changed})
|
||||
}
|
||||
|
||||
func (h *InboxAdminHTTPHandler) readAll(w http.ResponseWriter, r *http.Request) {
|
||||
account, ok := h.require(w, r, identity.PermissionInboxRead)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
count, err := h.inbox.MarkAllRead(r.Context(), "admin", account.ID)
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]any{"read_all": count})
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"aigateway.local/core/internal/platform/config"
|
||||
"aigateway.local/core/internal/platform/database"
|
||||
)
|
||||
|
||||
// TestInboxMaterializeAndBroadcast 验证站内消息核心链路:事件物化 → 未读计数 →
|
||||
// 重放幂等 → 已读回执 → 管理员广播。Redis 传 nil,走 DB 权威未读路径。
|
||||
func TestInboxMaterializeAndBroadcast(t *testing.T) {
|
||||
databaseURL := os.Getenv("WORKBENCH_TEST_DATABASE_URL")
|
||||
if databaseURL == "" {
|
||||
t.Skip("WORKBENCH_TEST_DATABASE_URL is not set")
|
||||
}
|
||||
ctx := context.Background()
|
||||
pool, err := database.Open(ctx, config.Database{URL: databaseURL, MaxConns: 8, MinConns: 0})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer pool.Close()
|
||||
|
||||
adminID := "44444444-4444-4444-4444-444444444444"
|
||||
portalID := "55555555-5555-5555-5555-555555555555"
|
||||
cleanup := func() {
|
||||
// 消息按收件人删:materialize 落给 admin 收件人,广播落给 portal 收件人(payload 为空,
|
||||
// 不能只按 payload 匹配,否则广播消息残留导致重跑未读数累加)。
|
||||
// 注意:同一 $1 同时比较 uuid 列与 jsonb text 提取,须显式 ::uuid / ::text,
|
||||
// 否则 PG 无法推断参数类型报 "text = uuid"(SQLSTATE 42883)。
|
||||
if _, cErr := pool.Exec(ctx, `DELETE FROM gateway.inbox_messages WHERE recipient_user_id=$1::uuid OR recipient_user_id=$2::uuid OR payload->>'actor_id'=$1::text OR payload->>'portal_user_id'=$1::text`, adminID, portalID); cErr != nil {
|
||||
t.Logf("cleanup inbox DELETE failed: %v", cErr)
|
||||
}
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM gateway.admin_accounts WHERE id=$1::uuid`, adminID)
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM gateway.portal_users WHERE id=$1::uuid OR lower(account)='m8-inbox-portal'`, portalID)
|
||||
}
|
||||
cleanup()
|
||||
defer cleanup()
|
||||
_, err = pool.Exec(ctx, `INSERT INTO gateway.admin_accounts(id,username,password_hash,role,active) VALUES($1,'m8-inbox-admin','test','superadmin',true) ON CONFLICT(id) DO NOTHING`, adminID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = pool.Exec(ctx, `INSERT INTO gateway.portal_users(id,account,password_hash,active) VALUES($1,'m8-inbox-portal','test',true) ON CONFLICT(id) DO NOTHING`, portalID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
svc := NewInboxService(NewService(pool), nil, "")
|
||||
|
||||
// 1) 事件物化:knowledge_document.ready → admin 收件人
|
||||
eventID := "99999999-9999-4999-8999-999999999999"
|
||||
payload, _ := json.Marshal(map[string]any{"chunk_count": "8", "actor_id": adminID})
|
||||
if err := svc.Materialize(ctx, eventID, "knowledge_document.ready", payload); err != nil {
|
||||
t.Fatalf("materialize failed: %v", err)
|
||||
}
|
||||
if count, err := svc.UnreadCount(ctx, "admin", adminID); err != nil || count != 1 {
|
||||
t.Fatalf("unread after materialize = %d, err=%v; want 1", count, err)
|
||||
}
|
||||
|
||||
// 2) 同事件重放幂等:不新增行、不报错
|
||||
if err := svc.Materialize(ctx, eventID, "knowledge_document.ready", payload); err != nil {
|
||||
t.Fatalf("replay failed: %v", err)
|
||||
}
|
||||
if count, _ := svc.UnreadCount(ctx, "admin", adminID); count != 1 {
|
||||
t.Fatalf("unread after replay = %d, want 1 (idempotent)", count)
|
||||
}
|
||||
|
||||
// 3) 收件箱列出 + 已读回执
|
||||
items, err := svc.List(ctx, "admin", adminID, 10)
|
||||
if err != nil || len(items) != 1 {
|
||||
t.Fatalf("list = %d items, err=%v; want 1", len(items), err)
|
||||
}
|
||||
if items[0].SenderType != "system" || items[0].Category != "system" || items[0].Title != "知识文档已入库" {
|
||||
t.Fatalf("unexpected message shape: %+v", items[0])
|
||||
}
|
||||
changed, err := svc.MarkRead(ctx, items[0].ID, "admin", adminID)
|
||||
if err != nil || !changed {
|
||||
t.Fatalf("mark read changed=%v err=%v; want true", changed, err)
|
||||
}
|
||||
if count, _ := svc.UnreadCount(ctx, "admin", adminID); count != 0 {
|
||||
t.Fatalf("unread after mark-read = %d, want 0", count)
|
||||
}
|
||||
|
||||
// 4) 管理员广播到全部 portal 用户(至少命中测试门户用户)
|
||||
sent, err := svc.Broadcast(ctx, InboxInput{RecipientKind: "portal", Category: "system", Title: "m8 升级公告", Body: "新增站内消息功能", Link: "/portal/inbox"}, nil, adminID)
|
||||
if err != nil {
|
||||
t.Fatalf("broadcast failed: %v", err)
|
||||
}
|
||||
if sent < 1 {
|
||||
t.Fatalf("broadcast sent = %d, want >=1", sent)
|
||||
}
|
||||
if count, _ := svc.UnreadCount(ctx, "portal", portalID); count != 1 {
|
||||
t.Fatalf("portal unread after broadcast = %d, want 1", count)
|
||||
}
|
||||
|
||||
// 5) AdminList scope=broadcasts 能看到这条管理员广播。broadcasts 是全局视角
|
||||
// (所有 admin→portal 广播),不能假设列表恰好 1 条,改为在其中找到本测试广播。
|
||||
broadcasts, err := svc.AdminList(ctx, adminID, "broadcasts", 100)
|
||||
if err != nil {
|
||||
t.Fatalf("broadcasts list failed: %v", err)
|
||||
}
|
||||
found := false
|
||||
for _, b := range broadcasts {
|
||||
if b.SenderType == "admin" && b.Body == "新增站内消息功能" {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("test broadcast not found in broadcasts list (%d items)", len(broadcasts))
|
||||
}
|
||||
|
||||
// 6) MarkAllRead 清空门户未读
|
||||
marked, err := svc.MarkAllRead(ctx, "portal", portalID)
|
||||
if err != nil || marked != 1 {
|
||||
t.Fatalf("mark-all-read = %d, err=%v; want 1", marked, err)
|
||||
}
|
||||
if count, _ := svc.UnreadCount(ctx, "portal", portalID); count != 0 {
|
||||
t.Fatalf("portal unread after mark-all = %d, want 0", count)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"aigateway.local/core/internal/identity"
|
||||
"aigateway.local/core/internal/platform/apiresponse"
|
||||
)
|
||||
|
||||
// InboxPortalHTTPHandler 向门户端暴露个人收件箱与未读徽标。
|
||||
type InboxPortalHTTPHandler struct {
|
||||
inbox *InboxService
|
||||
identity *identity.Service
|
||||
mux *http.ServeMux
|
||||
}
|
||||
|
||||
func NewInboxPortalHTTPHandler(inbox *InboxService, identityService *identity.Service) *InboxPortalHTTPHandler {
|
||||
h := &InboxPortalHTTPHandler{inbox: inbox, identity: identityService, mux: http.NewServeMux()}
|
||||
h.mux.HandleFunc("GET /api/v1/portal/inbox", h.list)
|
||||
h.mux.HandleFunc("GET /api/v1/portal/inbox/unread", h.unread)
|
||||
h.mux.HandleFunc("POST /api/v1/portal/inbox/read-all", h.readAll)
|
||||
h.mux.HandleFunc("POST /api/v1/portal/inbox/{id}/read", h.read)
|
||||
return h
|
||||
}
|
||||
|
||||
func (h *InboxPortalHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.mux.ServeHTTP(w, r) }
|
||||
|
||||
func (h *InboxPortalHTTPHandler) 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 *InboxPortalHTTPHandler) list(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.account(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
|
||||
items, err := h.inbox.List(r.Context(), "portal", a.ID, limit)
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
|
||||
func (h *InboxPortalHTTPHandler) unread(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.account(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
count, err := h.inbox.UnreadCount(r.Context(), "portal", a.ID)
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]int{"unread": count})
|
||||
}
|
||||
|
||||
func (h *InboxPortalHTTPHandler) read(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.account(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
changed, err := h.inbox.MarkRead(r.Context(), r.PathValue("id"), "portal", a.ID)
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"read": changed})
|
||||
}
|
||||
|
||||
func (h *InboxPortalHTTPHandler) readAll(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.account(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
count, err := h.inbox.MarkAllRead(r.Context(), "portal", a.ID)
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]any{"read_all": count})
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestInboxPlanMapsEvents 覆盖 inboxPlan 纯函数:每个支持的事件类型都要产出
|
||||
// 预期类别 / 收件人类别 / 文案关键词,未知事件返回 nil。
|
||||
func TestInboxPlanMapsEvents(t *testing.T) {
|
||||
payload := func(values map[string]any) json.RawMessage {
|
||||
encoded, err := json.Marshal(values)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return encoded
|
||||
}
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
eventType string
|
||||
values map[string]any
|
||||
wantKind string // recipient_kind
|
||||
wantCategory string
|
||||
wantTitle string
|
||||
wantUserID string
|
||||
wantAll bool
|
||||
wantRequest bool
|
||||
}{
|
||||
{name: "model_access.requested 通知全部管理员审批", eventType: "model_access.requested", values: map[string]any{"model": "gpt-5"}, wantKind: "admin", wantCategory: "approval", wantTitle: "新的模型访问申请", wantAll: true},
|
||||
{name: "model_access.decided 已批准回执给申请用户", eventType: "model_access.decided", values: map[string]any{"status": "approved"}, wantKind: "portal", wantCategory: "approval", wantTitle: "模型申请已处理", wantRequest: true},
|
||||
{name: "model_access.decided 已驳回文案", eventType: "model_access.decided", values: map[string]any{"status": "rejected"}, wantKind: "portal", wantCategory: "approval", wantTitle: "模型申请已处理", wantRequest: true},
|
||||
{name: "marketplace.installed 发给安装用户", eventType: "marketplace.installed", values: map[string]any{"code": "report-bot", "portal_user_id": "11111111-1111-1111-1111-111111111111"}, wantKind: "portal", wantCategory: "resource", wantTitle: "资源已安装", wantUserID: "11111111-1111-1111-1111-111111111111"},
|
||||
{name: "knowledge_document.ready 发给执行管理员", eventType: "knowledge_document.ready", values: map[string]any{"chunk_count": "12", "actor_id": "22222222-2222-2222-2222-222222222222"}, wantKind: "admin", wantCategory: "system", wantTitle: "知识文档已入库", wantUserID: "22222222-2222-2222-2222-222222222222"},
|
||||
{name: "knowledge_document.reprocessed 发给执行管理员", eventType: "knowledge_document.reprocessed", values: map[string]any{"actor_id": "22222222-2222-2222-2222-222222222222"}, wantKind: "admin", wantCategory: "system", wantTitle: "知识文档已重新处理", wantUserID: "22222222-2222-2222-2222-222222222222"},
|
||||
{name: "knowledge_document.embedding_failed 降级提示", eventType: "knowledge_document.embedding_failed", values: map[string]any{"actor_id": "22222222-2222-2222-2222-222222222222"}, wantKind: "admin", wantCategory: "system", wantTitle: "知识文档向量化失败", wantUserID: "22222222-2222-2222-2222-222222222222"},
|
||||
{name: "scheduled_task.completed 发给创建者", eventType: "scheduled_task.completed", values: map[string]any{"task_code": "daily-report", "actor_id": "33333333-3333-3333-3333-333333333333"}, wantKind: "admin", wantCategory: "task_result", wantTitle: "定时任务已执行", wantUserID: "33333333-3333-3333-3333-333333333333"},
|
||||
{name: "scheduled_task.failed 发给创建者", eventType: "scheduled_task.failed", values: map[string]any{"task_code": "daily-report", "error": "timeout", "actor_id": "33333333-3333-3333-3333-333333333333"}, wantKind: "admin", wantCategory: "task_result", wantTitle: "定时任务执行失败", wantUserID: "33333333-3333-3333-3333-333333333333"},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
drafts := inboxPlan(tc.eventType, payload(tc.values))
|
||||
if len(drafts) != 1 {
|
||||
t.Fatalf("expected exactly one draft, got %d", len(drafts))
|
||||
}
|
||||
draft := drafts[0]
|
||||
if draft.RecipientKind != tc.wantKind {
|
||||
t.Errorf("recipient_kind = %q, want %q", draft.RecipientKind, tc.wantKind)
|
||||
}
|
||||
if draft.Category != tc.wantCategory {
|
||||
t.Errorf("category = %q, want %q", draft.Category, tc.wantCategory)
|
||||
}
|
||||
if draft.Title != tc.wantTitle {
|
||||
t.Errorf("title = %q, want %q", draft.Title, tc.wantTitle)
|
||||
}
|
||||
if draft.UserID != tc.wantUserID {
|
||||
t.Errorf("user_id = %q, want %q", draft.UserID, tc.wantUserID)
|
||||
}
|
||||
if draft.AllAdmins != tc.wantAll {
|
||||
t.Errorf("all_admins = %v, want %v", draft.AllAdmins, tc.wantAll)
|
||||
}
|
||||
if draft.RequestUser != tc.wantRequest {
|
||||
t.Errorf("request_user = %v, want %v", draft.RequestUser, tc.wantRequest)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
if drafts := inboxPlan("some.unknown.event", payload(map[string]any{})); drafts != nil {
|
||||
t.Fatalf("unknown event should map to no drafts, got %+v", drafts)
|
||||
}
|
||||
}
|
||||
|
||||
// TestInboxPlanModelAccessRejectedBody 校验驳回与批准的不同正文文案。
|
||||
func TestInboxPlanModelAccessRejectedBody(t *testing.T) {
|
||||
values := func(status string) json.RawMessage {
|
||||
encoded, _ := json.Marshal(map[string]any{"status": status})
|
||||
return encoded
|
||||
}
|
||||
approved := inboxPlan("model_access.decided", values("approved"))
|
||||
rejected := inboxPlan("model_access.decided", values("rejected"))
|
||||
if !contains(approved[0].Body, "已批准") {
|
||||
t.Errorf("approved body should mention 已批准, got %q", approved[0].Body)
|
||||
}
|
||||
if !contains(rejected[0].Body, "已驳回") {
|
||||
t.Errorf("rejected body should mention 已驳回, got %q", rejected[0].Body)
|
||||
}
|
||||
}
|
||||
|
||||
func contains(haystack, needle string) bool {
|
||||
for i := 0; i+len(needle) <= len(haystack); i++ {
|
||||
if haystack[i:i+len(needle)] == needle {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// TestPayloadValue 校验 payloadValue 对字符串与数字两类取值的兼容(事件载荷
|
||||
// 中数字可能被 JSON 解码为 float64)。
|
||||
func TestPayloadValue(t *testing.T) {
|
||||
payload := json.RawMessage(`{"model":"gpt-5","chunk_count":12,"active":true}`)
|
||||
if got := payloadValue(payload, "model"); got != "gpt-5" {
|
||||
t.Errorf("string key = %q, want gpt-5", got)
|
||||
}
|
||||
if got := payloadValue(payload, "chunk_count"); got != "12" {
|
||||
t.Errorf("numeric key = %q, want 12", got)
|
||||
}
|
||||
if got := payloadValue(payload, "missing"); got != "" {
|
||||
t.Errorf("missing key = %q, want empty", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidInboxLink(t *testing.T) {
|
||||
cases := map[string]bool{
|
||||
"": true, "/portal/inbox": true, "https://example.com/notice": true,
|
||||
"http://example.com": true, "javascript:alert(1)": false,
|
||||
"data:text/html,x": false, "//example.com/path": false, "portal/inbox": false,
|
||||
}
|
||||
for link, want := range cases {
|
||||
if got := validInboxLink(link); got != want {
|
||||
t.Errorf("validInboxLink(%q) = %v, want %v", link, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -258,7 +258,19 @@ func (s *MarketplaceService) Detail(ctx context.Context, resourceType, code stri
|
||||
return item, nil, getErr
|
||||
}
|
||||
item = MarketItem{Type: "mcp_server", Code: server.Code, Name: server.Name, Description: server.Description, CategoryID: server.CategoryID, CategoryName: server.CategoryName, Tags: server.Tags, DepartmentIDs: server.DepartmentIDs, UpdatedAt: server.UpdatedAt}
|
||||
detail = server
|
||||
// 门户详情不得暴露 endpoint_url 与 has_secret_headers:内网服务拓扑和
|
||||
// 密钥状态仅管理端可见(运行时列表同样省略该字段)。
|
||||
raw, err := json.Marshal(server)
|
||||
if err != nil {
|
||||
return item, nil, err
|
||||
}
|
||||
var sanitized map[string]any
|
||||
if err := json.Unmarshal(raw, &sanitized); err != nil {
|
||||
return item, nil, err
|
||||
}
|
||||
delete(sanitized, "endpoint_url")
|
||||
delete(sanitized, "has_secret_headers")
|
||||
detail = sanitized
|
||||
case "skill":
|
||||
skill, getErr := s.skills.GetPublishedByCode(ctx, code)
|
||||
if getErr != nil {
|
||||
@@ -320,7 +332,9 @@ func (s *MarketplaceService) Install(ctx context.Context, resourceType, code, po
|
||||
}
|
||||
|
||||
func (s *MarketplaceService) Uninstall(ctx context.Context, resourceType, code, portalUserID string) error {
|
||||
resourceID, err := s.publishedResourceID(ctx, resourceType, code)
|
||||
// 卸载不受 enabled/status 限制:管理员停用或归档资源后,用户仍能
|
||||
// 移除自己的安装,否则安装行永久卡死、列表永远显示。
|
||||
resourceID, err := s.resourceIDByCode(ctx, resourceType, code)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -416,6 +430,23 @@ func (s *MarketplaceService) publishedResourceID(ctx context.Context, resourceTy
|
||||
return id, mapNotFound(err)
|
||||
}
|
||||
|
||||
// resourceIDByCode 按 code 解析资源 ID,不限制启用/发布状态(卸载场景使用)。
|
||||
func (s *MarketplaceService) resourceIDByCode(ctx context.Context, resourceType, code string) (string, error) {
|
||||
var id string
|
||||
var err error
|
||||
switch resourceType {
|
||||
case "mcp_server":
|
||||
err = s.assets.pool.QueryRow(ctx, `SELECT id::text FROM gateway.mcp_servers WHERE code=$1`, code).Scan(&id)
|
||||
case "skill":
|
||||
err = s.assets.pool.QueryRow(ctx, `SELECT id::text FROM gateway.skills WHERE code=$1`, code).Scan(&id)
|
||||
case "digital_employee":
|
||||
err = s.assets.pool.QueryRow(ctx, `SELECT id::text FROM gateway.digital_employees WHERE code=$1`, code).Scan(&id)
|
||||
default:
|
||||
return "", errors.New("未知的资源类型")
|
||||
}
|
||||
return id, mapNotFound(err)
|
||||
}
|
||||
|
||||
func (s *MarketplaceService) resourceByID(ctx context.Context, resourceType, id string) (MarketItem, bool, error) {
|
||||
var item MarketItem
|
||||
switch resourceType {
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -84,9 +85,14 @@ func (c *MCPClient) DiscoverTools(ctx context.Context, server MCPServer, headers
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// 读缓存与写入共用同一把锁,避免 -race 下 tools/toolsAt 的无锁读。
|
||||
c.mu.Lock()
|
||||
if state.tools != nil && time.Since(state.toolsAt) < c.cacheTTL {
|
||||
return state.tools, nil
|
||||
snapshot := state.tools
|
||||
c.mu.Unlock()
|
||||
return snapshot, nil
|
||||
}
|
||||
c.mu.Unlock()
|
||||
result, err := c.call(ctx, server, headers, "tools/list", map[string]any{})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -157,13 +163,19 @@ func (c *MCPClient) CallTool(ctx context.Context, server MCPServer, headers map[
|
||||
return MCPToolResult{Content: text.String()}, nil
|
||||
}
|
||||
|
||||
// cacheKey 以服务器 revision 参与缓存键:管理端编辑 endpoint/请求头后
|
||||
// revision 递增,旧会话与旧工具列表立即失效,不会把过期会话发往新端点。
|
||||
func cacheKey(server MCPServer) string {
|
||||
return server.ID + ":" + strconv.FormatInt(server.Revision, 10)
|
||||
}
|
||||
|
||||
// ensureInitialized performs the MCP initialize handshake for a server if its
|
||||
// session has lapsed (or no cached tools exist yet), then acknowledges with
|
||||
// notifications/initialized. The handshake is guarded by the per-server cache
|
||||
// so a burst of calls does not re-initialize every request.
|
||||
func (c *MCPClient) ensureInitialized(ctx context.Context, server MCPServer, headers map[string]string) (*mcpServerState, error) {
|
||||
c.mu.Lock()
|
||||
state, ok := c.states[server.ID]
|
||||
state, ok := c.states[cacheKey(server)]
|
||||
if ok && time.Since(state.initAt) < c.cacheTTL {
|
||||
c.mu.Unlock()
|
||||
return state, nil
|
||||
@@ -182,7 +194,7 @@ func (c *MCPClient) ensureInitialized(ctx context.Context, server MCPServer, hea
|
||||
|
||||
c.mu.Lock()
|
||||
state = &mcpServerState{initAt: time.Now(), sessionID: sessionID}
|
||||
c.states[server.ID] = state
|
||||
c.states[cacheKey(server)] = state
|
||||
c.mu.Unlock()
|
||||
|
||||
// Best-effort acknowledgment; servers that require it will reject later
|
||||
@@ -233,7 +245,7 @@ func (c *MCPClient) sendNotification(ctx context.Context, server MCPServer, head
|
||||
func (c *MCPClient) call(ctx context.Context, server MCPServer, headers map[string]string, method string, params any) (json.RawMessage, error) {
|
||||
c.mu.Lock()
|
||||
sessionID := ""
|
||||
if state, ok := c.states[server.ID]; ok {
|
||||
if state, ok := c.states[cacheKey(server)]; ok {
|
||||
sessionID = state.sessionID
|
||||
}
|
||||
c.mu.Unlock()
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"aigateway.local/core/internal/platform/cryptox"
|
||||
"aigateway.local/core/internal/provider"
|
||||
@@ -253,7 +254,13 @@ func (s *NotificationService) deliver(ctx context.Context, channel NotificationC
|
||||
message = fmt.Sprintf("Webhook 返回 HTTP %d", status)
|
||||
}
|
||||
if len(message) > 1000 {
|
||||
message = message[:1000]
|
||||
// 按字节截断可能切半多字节 rune;无效 UTF-8 会被 PostgreSQL 拒绝,
|
||||
// 使投递记录无法更新,事件永远重试。
|
||||
cut := message[:1000]
|
||||
for len(cut) > 0 && !utf8.RuneStart(cut[len(cut)-1]) {
|
||||
cut = cut[:len(cut)-1]
|
||||
}
|
||||
message = cut
|
||||
}
|
||||
_, dbErr := s.assets.pool.Exec(context.WithoutCancel(ctx), `UPDATE gateway.notification_deliveries SET status=$2,attempts=attempts+1,response_status=nullif($3,0),last_error=$4,delivered_at=CASE WHEN $2='delivered' THEN clock_timestamp() ELSE delivered_at END,updated_at=clock_timestamp() WHERE id=$1`, delivery.ID, map[bool]string{true: "delivered", false: "failed"}[success], status, message)
|
||||
if dbErr != nil {
|
||||
@@ -280,6 +287,7 @@ func (s *NotificationService) RetryDelivery(ctx context.Context, id string) erro
|
||||
|
||||
type NotificationDispatcher struct {
|
||||
service *NotificationService
|
||||
inbox *InboxService
|
||||
redis *redis.Client
|
||||
stream, group, consumer string
|
||||
logger *slog.Logger
|
||||
@@ -288,6 +296,10 @@ type NotificationDispatcher struct {
|
||||
func NewNotificationDispatcher(service *NotificationService, client *redis.Client, stream, consumer string, logger *slog.Logger) *NotificationDispatcher {
|
||||
return &NotificationDispatcher{service: service, redis: client, stream: stream, group: "gateway-notifications-v1", consumer: consumer, logger: logger}
|
||||
}
|
||||
|
||||
// SetInbox wires the in-app inbox materializer (M8 P4). 为 nil 时站内消息不落库,
|
||||
// Webhook 投递不受影响。handle 内幂等:inbox 以 (source_event_id, 收件人) 去重。
|
||||
func (d *NotificationDispatcher) SetInbox(inbox *InboxService) { d.inbox = inbox }
|
||||
func (d *NotificationDispatcher) Run(ctx context.Context) error {
|
||||
if err := d.redis.XGroupCreateMkStream(ctx, d.stream, d.group, "$").Err(); err != nil && !strings.Contains(err.Error(), "BUSYGROUP") {
|
||||
return err
|
||||
@@ -414,11 +426,25 @@ func (d *NotificationDispatcher) handle(ctx context.Context, message redis.XMess
|
||||
eventID := fmt.Sprint(message.Values["event_id"])
|
||||
eventType := fmt.Sprint(message.Values["event_type"])
|
||||
payload := json.RawMessage(fmt.Sprint(message.Values["payload"]))
|
||||
// M8 P4:物化站内消息。失败与 webhook 同语义——事件留在 pending,由 reclaim 重试;
|
||||
// inbox 幂等(ON CONFLICT)保证重放不产生重复消息。
|
||||
if d.inbox != nil {
|
||||
if err := d.inbox.Materialize(ctx, eventID, eventType, payload); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
channels, err := d.service.ListChannels(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
selectedChannelID := payloadValue(payload, "notification_channel_id")
|
||||
if selectedChannelID == "null" {
|
||||
selectedChannelID = ""
|
||||
}
|
||||
for _, channel := range channels {
|
||||
if selectedChannelID != "" && channel.ID != selectedChannelID {
|
||||
continue
|
||||
}
|
||||
if !channel.Enabled || !matchesEvent(channel.EventPatterns, eventType) {
|
||||
continue
|
||||
}
|
||||
@@ -429,9 +455,23 @@ func (d *NotificationDispatcher) handle(ctx context.Context, message redis.XMess
|
||||
if delivery.Status == "delivered" {
|
||||
continue
|
||||
}
|
||||
if deliveryErr = d.service.deliver(ctx, channel, delivery); deliveryErr != nil && d.logger != nil {
|
||||
d.logger.Warn("webhook delivery failed", "channel", channel.Name, "event_id", eventID, "error", deliveryErr)
|
||||
if deliveryErr = d.service.deliver(ctx, channel, delivery); deliveryErr != nil {
|
||||
// 投递失败必须让事件留在 pending 列表由 reclaim 重试;但重试预算
|
||||
// 耗尽后放弃自动重试(投递记录保留 failed 状态,管理端可人工重试),
|
||||
// 否则永久失败的 Webhook 会让事件无限期卡在 pending,阻塞该事件
|
||||
// 的其它通道投递与站内消息。
|
||||
if delivery.Attempts >= webhookMaxAttempts {
|
||||
if d.logger != nil {
|
||||
d.logger.Error("webhook delivery exhausted retries; manual retry available in admin",
|
||||
"channel", channel.Name, "event_id", eventID, "attempts", delivery.Attempts, "error", deliveryErr)
|
||||
}
|
||||
continue
|
||||
}
|
||||
return deliveryErr
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// webhookMaxAttempts 是单条投递记录的自动重试上限;deliver 每次失败 attempts+1。
|
||||
const webhookMaxAttempts = 10
|
||||
|
||||
@@ -63,6 +63,10 @@ func (h *FilesPortalHTTPHandler) upload(w http.ResponseWriter, r *http.Request)
|
||||
contentType = ct
|
||||
}
|
||||
}
|
||||
// 非 multipart 请求(body 为 nil 时)按原始请求体上传(?filename= 指定文件名)。
|
||||
if body == nil {
|
||||
body = r.Body
|
||||
}
|
||||
if originalName == "" {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "缺少文件名")
|
||||
return
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
"aigateway.local/core/internal/factcheck"
|
||||
"aigateway.local/core/internal/gateway"
|
||||
"aigateway.local/core/internal/platform/apiresponse"
|
||||
tracepkg "aigateway.local/core/internal/trace"
|
||||
)
|
||||
|
||||
type RuntimeHTTPHandler struct {
|
||||
@@ -25,6 +26,7 @@ type RuntimeHTTPHandler struct {
|
||||
auth apikey.PrincipalAuthenticator
|
||||
gateway http.Handler
|
||||
factCheck *factcheck.Engine
|
||||
traces *tracepkg.Store
|
||||
logger *slog.Logger
|
||||
mux *http.ServeMux
|
||||
market MarketplaceDeps
|
||||
@@ -70,6 +72,11 @@ func (h *RuntimeHTTPHandler) SetLogger(logger *slog.Logger) {
|
||||
// conversations. When nil (the default) fact-checking is skipped entirely.
|
||||
func (h *RuntimeHTTPHandler) SetFactCheckEngine(engine *factcheck.Engine) { h.factCheck = engine }
|
||||
|
||||
// SetTraceStore enables metadata-only LLM Trace recording for application and
|
||||
// digital-employee runs. Trace persistence is best effort and never changes
|
||||
// the runtime response when the database is unavailable.
|
||||
func (h *RuntimeHTTPHandler) SetTraceStore(store *tracepkg.Store) { h.traces = store }
|
||||
|
||||
// factCheckRetriever adapts the workbench Retriever to the fact-check engine's
|
||||
// EvidenceRetriever interface, reusing the same knowledge-base search path that
|
||||
// application prompts already use.
|
||||
@@ -155,11 +162,17 @@ func (h *RuntimeHTTPHandler) principal(w http.ResponseWriter, r *http.Request) (
|
||||
return principal, true
|
||||
}
|
||||
func visible(departments []string, principal apikey.Principal, secure bool) bool {
|
||||
// fail-closed:无 APIKeyID 的匿名主体不视为"可见一切"。
|
||||
// 今天认证器总是返回 bootstrap 或真实 key ID,但任何未来认证路径的
|
||||
// 变化都不应静默放开所有部门作用域资产。
|
||||
if principal.APIKeyID == "" {
|
||||
return true
|
||||
return false
|
||||
}
|
||||
// 无部门限定的资源是全局资源:所有已认证主体可见。secure 只标记
|
||||
// "执行敏感能力"类资源,不改变可见性规则——否则全局工具/MCP 对
|
||||
// 所有人不可见,绑定它们的应用会在运行时失败。
|
||||
if len(departments) == 0 {
|
||||
return !secure
|
||||
return true
|
||||
}
|
||||
if principal.TenantID == nil {
|
||||
return false
|
||||
@@ -321,13 +334,18 @@ func (h *RuntimeHTTPHandler) runApplication(w http.ResponseWriter, r *http.Reque
|
||||
runError := ""
|
||||
retrievalCount := 0
|
||||
toolCount := 0
|
||||
modelCallCount := 0
|
||||
conversationID := strings.TrimSpace(r.Header.Get("X-Gateway-Conversation-ID"))
|
||||
traceID := h.beginTrace(r.Context(), principal, "application", app.ID, app.Code, conversationID)
|
||||
defer func() {
|
||||
traceCtx := context.WithoutCancel(r.Context())
|
||||
h.finishTrace(traceCtx, traceID, status, runError, retrievalCount, modelCallCount, toolCount)
|
||||
runID, idErr := newUUID()
|
||||
if idErr == nil {
|
||||
_, _ = h.service.pool.Exec(context.WithoutCancel(r.Context()), `INSERT INTO gateway.application_runs(id,application_id,version,api_key_id,request_id,status,latency_ms,retrieval_count,tool_count,error) VALUES($1,$2,nullif($3,0),nullif($4,'')::uuid,$5,$6,$7,$8,$9,$10)`, runID, app.ID, valueOrZero(app.PublishedVersion), principal.APIKeyID, gateway.RequestID(r.Context()), status, time.Since(started).Milliseconds(), retrievalCount, toolCount, runError)
|
||||
_, _ = h.service.pool.Exec(traceCtx, `INSERT INTO gateway.application_runs(id,application_id,version,api_key_id,trace_id,request_id,status,latency_ms,retrieval_count,tool_count,error) VALUES($1,$2,nullif($3,0),nullif($4,'')::uuid,nullif($5,'')::uuid,$6,$7,$8,$9,$10,$11)`, runID, app.ID, valueOrZero(app.PublishedVersion), principal.APIKeyID, traceID, gateway.RequestID(r.Context()), status, time.Since(started).Milliseconds(), retrievalCount, toolCount, runError)
|
||||
}
|
||||
}()
|
||||
payload, toolsByCode, prepareErr := h.prepareApplication(r.Context(), app, input, principal, &retrievalCount)
|
||||
payload, toolsByCode, prepareErr := h.prepareApplication(r.Context(), app, input, principal, &retrievalCount, traceID)
|
||||
if prepareErr != nil {
|
||||
runError = prepareErr.Error()
|
||||
runtimeError(w, 400, runError)
|
||||
@@ -338,7 +356,8 @@ func (h *RuntimeHTTPHandler) runApplication(w http.ResponseWriter, r *http.Reque
|
||||
var responseHeaders http.Header
|
||||
var statusCode int
|
||||
for round := 0; ; round++ {
|
||||
statusCode, responseHeaders, response, err = h.callGateway(r, payload)
|
||||
modelCallCount++
|
||||
statusCode, responseHeaders, response, err = h.callGatewayWithTrace(r, payload, traceID, round)
|
||||
if err != nil {
|
||||
runError = err.Error()
|
||||
copyHeaders(w.Header(), responseHeaders)
|
||||
@@ -367,7 +386,9 @@ func (h *RuntimeHTTPHandler) runApplication(w http.ResponseWriter, r *http.Reque
|
||||
if json.Unmarshal([]byte(call.Arguments), &args) != nil {
|
||||
args = map[string]any{}
|
||||
}
|
||||
result, executeErr := h.tools.Execute(r.Context(), tool, args, principal.APIKeyID, gateway.RequestID(r.Context()))
|
||||
result, executeErr := h.executeToolWithTrace(r.Context(), traceID, call.Name, call.ID, round, func() (map[string]any, error) {
|
||||
return h.tools.Execute(r.Context(), tool, args, principal.APIKeyID, gateway.RequestID(r.Context()))
|
||||
})
|
||||
if executeErr != nil {
|
||||
runError = executeErr.Error()
|
||||
runtimeError(w, 502, runError)
|
||||
@@ -379,7 +400,7 @@ func (h *RuntimeHTTPHandler) runApplication(w http.ResponseWriter, r *http.Reque
|
||||
}
|
||||
}
|
||||
if h.factCheck != nil {
|
||||
h.applyFactCheck(r, input, response)
|
||||
h.applyFactCheck(r, input, response, app.DepartmentIDs)
|
||||
}
|
||||
response["application"] = map[string]any{"code": app.Code, "name": app.Name, "version": valueOrZero(app.PublishedVersion), "retrieval_count": retrievalCount, "tool_calls": toolCount}
|
||||
status = "success"
|
||||
@@ -389,17 +410,22 @@ func (h *RuntimeHTTPHandler) runApplication(w http.ResponseWriter, r *http.Reque
|
||||
|
||||
// applyFactCheck verifies the assistant answer against configured knowledge
|
||||
// bases and applies the policy action. It must never fail the chat: any error
|
||||
// is logged and the answer is returned unchanged.
|
||||
func (h *RuntimeHTTPHandler) applyFactCheck(r *http.Request, input applicationRequest, response map[string]any) {
|
||||
// is logged and the answer is returned unchanged. departments 用于选择
|
||||
// department:<uuid> 作用域的策略,空列表只应用 global 策略。
|
||||
func (h *RuntimeHTTPHandler) applyFactCheck(r *http.Request, input applicationRequest, response map[string]any, departments []string) {
|
||||
answer, _ := assistantAnswer(response)
|
||||
lastQuestion := lastUserMessage(input.Messages)
|
||||
if strings.TrimSpace(answer) == "" || strings.TrimSpace(lastQuestion) == "" {
|
||||
return
|
||||
}
|
||||
scope := ""
|
||||
if len(departments) > 0 {
|
||||
scope = "department:" + departments[0]
|
||||
}
|
||||
verifier := func(ctx context.Context, model, system, user string, timeout time.Duration) (string, error) {
|
||||
return h.VerifyFactCheck(ctx, r, model, system, user, timeout)
|
||||
}
|
||||
event, err := h.factCheck.Check(r.Context(), gateway.RequestID(r.Context()), lastQuestion, answer, factcheck.VerifierFunc(verifier))
|
||||
event, err := h.factCheck.Check(r.Context(), gateway.RequestID(r.Context()), scope, lastQuestion, answer, factcheck.VerifierFunc(verifier))
|
||||
if err != nil {
|
||||
h.logger.Warn("fact-check skipped", "request_id", gateway.RequestID(r.Context()), "error", err)
|
||||
return
|
||||
@@ -446,7 +472,7 @@ func overrideAnswer(response map[string]any, content string) {
|
||||
}
|
||||
}
|
||||
|
||||
func (h *RuntimeHTTPHandler) prepareApplication(ctx context.Context, app Application, input applicationRequest, principal apikey.Principal, retrievalCount *int) (map[string]any, map[string]Tool, error) {
|
||||
func (h *RuntimeHTTPHandler) prepareApplication(ctx context.Context, app Application, input applicationRequest, principal apikey.Principal, retrievalCount *int, traceID string) (map[string]any, map[string]Tool, error) {
|
||||
config := *app.PublishedConfig
|
||||
messages := make([]map[string]any, 0, len(input.Messages)+2)
|
||||
total := 0
|
||||
@@ -490,7 +516,7 @@ func (h *RuntimeHTTPHandler) prepareApplication(ctx context.Context, app Applica
|
||||
if err != nil || !kb.Enabled || !visible(kb.DepartmentIDs, principal, false) {
|
||||
return nil, nil, fmt.Errorf("应用绑定的知识库 %s 当前不可用", kbID)
|
||||
}
|
||||
hits, searchErr := h.retriever.Search(ctx, kbID, lastQuestion, config.RetrievalTopK)
|
||||
hits, searchErr := h.searchWithTrace(ctx, traceID, kbID, lastQuestion, config.RetrievalTopK)
|
||||
if searchErr != nil {
|
||||
continue
|
||||
}
|
||||
@@ -588,7 +614,9 @@ type boundedRecorder struct {
|
||||
}
|
||||
|
||||
func newBoundedRecorder() *boundedRecorder {
|
||||
return &boundedRecorder{code: http.StatusOK, header: make(http.Header)}
|
||||
// code 初始为 0:WriteHeader 只在首次调用时生效,若网关从未调用
|
||||
// WriteHeader,则 Write 时默认回退 200。
|
||||
return &boundedRecorder{header: make(http.Header)}
|
||||
}
|
||||
|
||||
func (r *boundedRecorder) Header() http.Header { return r.header }
|
||||
|
||||
@@ -131,8 +131,10 @@ func (h *RuntimeHTTPHandler) invokeMCPTool(w http.ResponseWriter, r *http.Reques
|
||||
}
|
||||
|
||||
type digitalEmployeeRequest struct {
|
||||
Messages []map[string]any `json:"messages"`
|
||||
Variables map[string]any `json:"variables"`
|
||||
Messages []map[string]any `json:"messages"`
|
||||
Variables map[string]any `json:"variables"`
|
||||
SkillIDs []string `json:"skill_ids"`
|
||||
MCPServerIDs []string `json:"mcp_server_ids"`
|
||||
}
|
||||
|
||||
func (h *RuntimeHTTPHandler) runDigitalEmployee(w http.ResponseWriter, r *http.Request) {
|
||||
@@ -168,13 +170,18 @@ func (h *RuntimeHTTPHandler) runDigitalEmployee(w http.ResponseWriter, r *http.R
|
||||
runError := ""
|
||||
retrievalCount := 0
|
||||
toolCount := 0
|
||||
modelCallCount := 0
|
||||
conversationID := strings.TrimSpace(r.Header.Get("X-Gateway-Conversation-ID"))
|
||||
traceID := h.beginTrace(r.Context(), principal, "digital_employee", employee.ID, employee.Code, conversationID)
|
||||
defer func() {
|
||||
traceCtx := context.WithoutCancel(r.Context())
|
||||
h.finishTrace(traceCtx, traceID, status, runError, retrievalCount, modelCallCount, toolCount)
|
||||
runID, idErr := newUUID()
|
||||
if idErr == nil {
|
||||
_, _ = h.service.pool.Exec(context.WithoutCancel(r.Context()), `INSERT INTO gateway.digital_employee_runs(id,digital_employee_id,api_key_id,request_id,status,latency_ms,retrieval_count,tool_count,error) VALUES($1,$2,nullif($3,'')::uuid,$4,$5,$6,$7,$8,$9)`, runID, employee.ID, principal.APIKeyID, gateway.RequestID(r.Context()), status, time.Since(started).Milliseconds(), retrievalCount, toolCount, runError)
|
||||
_, _ = h.service.pool.Exec(traceCtx, `INSERT INTO gateway.digital_employee_runs(id,digital_employee_id,api_key_id,trace_id,request_id,status,latency_ms,retrieval_count,tool_count,error) VALUES($1,$2,nullif($3,'')::uuid,nullif($4,'')::uuid,$5,$6,$7,$8,$9,$10)`, runID, employee.ID, principal.APIKeyID, traceID, gateway.RequestID(r.Context()), status, time.Since(started).Milliseconds(), retrievalCount, toolCount, runError)
|
||||
}
|
||||
}()
|
||||
executors, payload, prepareErr := h.prepareDigitalEmployee(r.Context(), employee, input, principal, &retrievalCount, portalUserID)
|
||||
executors, payload, prepareErr := h.prepareDigitalEmployee(r.Context(), employee, input, principal, &retrievalCount, portalUserID, traceID)
|
||||
if prepareErr != nil {
|
||||
runError = prepareErr.Error()
|
||||
runtimeError(w, 400, runError)
|
||||
@@ -184,7 +191,8 @@ func (h *RuntimeHTTPHandler) runDigitalEmployee(w http.ResponseWriter, r *http.R
|
||||
var responseHeaders http.Header
|
||||
var statusCode int
|
||||
for round := 0; ; round++ {
|
||||
statusCode, responseHeaders, response, err = h.callGateway(r, payload)
|
||||
modelCallCount++
|
||||
statusCode, responseHeaders, response, err = h.callGatewayWithTrace(r, payload, traceID, round)
|
||||
if err != nil {
|
||||
runError = err.Error()
|
||||
copyHeaders(w.Header(), responseHeaders)
|
||||
@@ -213,7 +221,9 @@ func (h *RuntimeHTTPHandler) runDigitalEmployee(w http.ResponseWriter, r *http.R
|
||||
if json.Unmarshal([]byte(call.Arguments), &args) != nil {
|
||||
args = map[string]any{}
|
||||
}
|
||||
result, executeErr := exec(r.Context(), args)
|
||||
result, executeErr := h.executeToolWithTrace(r.Context(), traceID, call.Name, call.ID, round, func() (map[string]any, error) {
|
||||
return exec(r.Context(), args)
|
||||
})
|
||||
if executeErr != nil {
|
||||
runError = executeErr.Error()
|
||||
runtimeError(w, 502, runError)
|
||||
@@ -224,6 +234,10 @@ func (h *RuntimeHTTPHandler) runDigitalEmployee(w http.ResponseWriter, r *http.R
|
||||
toolCount++
|
||||
}
|
||||
}
|
||||
// 事实核验与普通应用一致:block 策略不能因走数字员工入口而被绕过。
|
||||
if h.factCheck != nil {
|
||||
h.applyFactCheck(r, applicationRequest{Messages: input.Messages, Variables: input.Variables}, response, employee.DepartmentIDs)
|
||||
}
|
||||
response["digital_employee"] = map[string]any{"code": employee.Code, "name": employee.Name, "persona": employee.Persona, "retrieval_count": retrievalCount, "tool_calls": toolCount}
|
||||
status = "success"
|
||||
copyHeaders(w.Header(), responseHeaders)
|
||||
@@ -234,7 +248,7 @@ func (h *RuntimeHTTPHandler) runDigitalEmployee(w http.ResponseWriter, r *http.R
|
||||
// persona + rendered skills as system context, knowledge RAG evidence, and the
|
||||
// union of bound tools (regular + MCP) exposed to the model. It returns the
|
||||
// tool executors keyed by the exact schema name the model may call.
|
||||
func (h *RuntimeHTTPHandler) prepareDigitalEmployee(ctx context.Context, employee DigitalEmployee, input digitalEmployeeRequest, principal apikey.Principal, retrievalCount *int, portalUserID string) (map[string]toolExecutor, map[string]any, error) {
|
||||
func (h *RuntimeHTTPHandler) prepareDigitalEmployee(ctx context.Context, employee DigitalEmployee, input digitalEmployeeRequest, principal apikey.Principal, retrievalCount *int, portalUserID, traceID string) (map[string]toolExecutor, map[string]any, error) {
|
||||
messages := make([]map[string]any, 0, len(input.Messages)+3)
|
||||
total := 0
|
||||
lastQuestion := ""
|
||||
@@ -260,8 +274,16 @@ func (h *RuntimeHTTPHandler) prepareDigitalEmployee(ctx context.Context, employe
|
||||
if strings.TrimSpace(employee.Persona) != "" {
|
||||
system = append(system, employee.Persona)
|
||||
}
|
||||
selectedSkills, err := selectedBindings(input.SkillIDs, employee.SkillIDs, "Skill")
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
selectedMCPServers, err := selectedBindings(input.MCPServerIDs, employee.MCPServerIDs, "MCP")
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
skills := map[string]Skill{}
|
||||
for _, skillID := range employee.SkillIDs {
|
||||
for _, skillID := range selectedSkills {
|
||||
skill, err := h.market.Skills.Get(ctx, skillID)
|
||||
if err != nil || !skill.Enabled || !visible(skill.DepartmentIDs, principal, false) {
|
||||
return nil, nil, fmt.Errorf("绑定的 Skill %s 当前不可用", skillID)
|
||||
@@ -286,7 +308,7 @@ func (h *RuntimeHTTPHandler) prepareDigitalEmployee(ctx context.Context, employe
|
||||
if err != nil || !kb.Enabled || !visible(kb.DepartmentIDs, principal, false) {
|
||||
return fmt.Errorf("绑定的知识库 %s 当前不可用", kbID)
|
||||
}
|
||||
hits, searchErr := h.retriever.Search(ctx, kbID, lastQuestion, employee.RetrievalTopK)
|
||||
hits, searchErr := h.searchWithTrace(ctx, traceID, kbID, lastQuestion, employee.RetrievalTopK)
|
||||
if searchErr != nil {
|
||||
return nil
|
||||
}
|
||||
@@ -382,7 +404,7 @@ func (h *RuntimeHTTPHandler) prepareDigitalEmployee(ctx context.Context, employe
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, serverID := range employee.MCPServerIDs {
|
||||
for _, serverID := range selectedMCPServers {
|
||||
if err := addMCP(serverID); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
@@ -405,6 +427,31 @@ func (h *RuntimeHTTPHandler) prepareDigitalEmployee(ctx context.Context, employe
|
||||
return executors, payload, nil
|
||||
}
|
||||
|
||||
// selectedBindings lets scheduled tasks restrict a digital employee to a
|
||||
// subset of its published Skill/MCP bindings. Empty means the employee's full
|
||||
// published binding set; callers can never add resources it does not own.
|
||||
func selectedBindings(selected, allowed []string, label string) ([]string, error) {
|
||||
if len(selected) == 0 {
|
||||
return allowed, nil
|
||||
}
|
||||
allowedSet := make(map[string]bool, len(allowed))
|
||||
for _, id := range allowed {
|
||||
allowedSet[id] = true
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
result := make([]string, 0, len(selected))
|
||||
for _, id := range selected {
|
||||
if !allowedSet[id] {
|
||||
return nil, fmt.Errorf("请求的 %s %s 未绑定到数字员工", label, id)
|
||||
}
|
||||
if !seen[id] {
|
||||
seen[id] = true
|
||||
result = append(result, id)
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// portalUserID resolves the portal user behind an API key, if any. Resource
|
||||
// marketplace installs are scoped to portal users.
|
||||
func (h *RuntimeHTTPHandler) portalUserID(ctx context.Context, principal apikey.Principal) (string, error) {
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
package workbench
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestSelectedBindings(t *testing.T) {
|
||||
allowed := []string{"a", "b", "c"}
|
||||
all, err := selectedBindings(nil, allowed, "Skill")
|
||||
if err != nil || len(all) != 3 {
|
||||
t.Fatalf("empty selection should use all bindings: %#v err=%v", all, err)
|
||||
}
|
||||
selected, err := selectedBindings([]string{"b", "b"}, allowed, "Skill")
|
||||
if err != nil || len(selected) != 1 || selected[0] != "b" {
|
||||
t.Fatalf("selection should be deduplicated: %#v err=%v", selected, err)
|
||||
}
|
||||
if _, err := selectedBindings([]string{"outside"}, allowed, "Skill"); err == nil {
|
||||
t.Fatal("unbound selection should be rejected")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"aigateway.local/core/internal/apikey"
|
||||
"aigateway.local/core/internal/gateway"
|
||||
tracepkg "aigateway.local/core/internal/trace"
|
||||
)
|
||||
|
||||
func (h *RuntimeHTTPHandler) beginTrace(ctx context.Context, principal apikey.Principal, traceType, targetID, targetCode, conversationID string) string {
|
||||
if h.traces == nil {
|
||||
return ""
|
||||
}
|
||||
conversationID = strings.TrimSpace(conversationID)
|
||||
input := tracepkg.StartInput{RequestID: gateway.RequestID(ctx), APIKeyID: principal.APIKeyID, TenantID: principal.TenantID, TraceType: traceType, TargetID: targetID, TargetCode: targetCode, ConversationID: conversationID}
|
||||
item, err := h.traces.Start(context.WithoutCancel(ctx), input)
|
||||
if err != nil {
|
||||
h.logger.Warn("llm trace start failed", "request_id", input.RequestID, "target", targetCode, "error", err)
|
||||
return ""
|
||||
}
|
||||
return item.ID
|
||||
}
|
||||
|
||||
func (h *RuntimeHTTPHandler) finishTrace(ctx context.Context, traceID, status, errorText string, retrievalCount, modelCallCount, toolCallCount int) {
|
||||
if h.traces == nil || traceID == "" {
|
||||
return
|
||||
}
|
||||
if err := h.traces.Finish(context.WithoutCancel(ctx), traceID, tracepkg.FinishInput{Status: status, Error: errorText, RetrievalCount: retrievalCount, ModelCallCount: modelCallCount, ToolCallCount: toolCallCount}); err != nil {
|
||||
h.logger.Warn("llm trace finish failed", "trace_id", traceID, "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (h *RuntimeHTTPHandler) callGatewayWithTrace(original *http.Request, payload map[string]any, traceID string, round int) (int, http.Header, map[string]any, error) {
|
||||
spanID := ""
|
||||
model, _ := payload["model"].(string)
|
||||
if h.traces != nil && traceID != "" {
|
||||
span, err := h.traces.StartSpan(context.WithoutCancel(original.Context()), tracepkg.SpanInput{TraceID: traceID, SpanType: "model", Name: "chat.completions", Round: round, Model: model, Metadata: map[string]any{"endpoint": "/v1/chat/completions"}})
|
||||
if err != nil {
|
||||
h.logger.Warn("llm model span start failed", "trace_id", traceID, "error", err)
|
||||
} else {
|
||||
spanID = span.ID
|
||||
}
|
||||
}
|
||||
statusCode, headers, response, callErr := h.callGateway(original, payload)
|
||||
if spanID != "" {
|
||||
inputTokens, outputTokens := responseUsage(response)
|
||||
spanStatus := "success"
|
||||
if callErr != nil || statusCode < 200 || statusCode >= 300 {
|
||||
spanStatus = "error"
|
||||
}
|
||||
metadata := map[string]any{"http_status": statusCode, "round": round}
|
||||
providerCode := ""
|
||||
spanModel := model
|
||||
if provider := headers.Get("X-Gateway-Provider"); provider != "" {
|
||||
providerCode = provider
|
||||
metadata["provider"] = provider
|
||||
}
|
||||
if resolvedModel := strings.TrimSpace(headers.Get("X-Gateway-Model")); resolvedModel != "" {
|
||||
spanModel = resolvedModel
|
||||
}
|
||||
if err := h.traces.FinishSpan(context.WithoutCancel(original.Context()), spanID, tracepkg.SpanFinishInput{Status: spanStatus, Error: errorString(callErr), InputTokens: inputTokens, OutputTokens: outputTokens, ProviderCode: providerCode, Model: spanModel, Metadata: metadata}); err != nil {
|
||||
h.logger.Warn("llm model span finish failed", "span_id", spanID, "error", err)
|
||||
}
|
||||
}
|
||||
return statusCode, headers, response, callErr
|
||||
}
|
||||
|
||||
func (h *RuntimeHTTPHandler) executeToolWithTrace(ctx context.Context, traceID, name, callID string, round int, execute func() (map[string]any, error)) (map[string]any, error) {
|
||||
spanID := ""
|
||||
if h.traces != nil && traceID != "" {
|
||||
span, err := h.traces.StartSpan(context.WithoutCancel(ctx), tracepkg.SpanInput{TraceID: traceID, SpanType: "tool", Name: name, Round: round, Metadata: map[string]any{"tool_call_id": callID}})
|
||||
if err != nil {
|
||||
h.logger.Warn("llm tool span start failed", "trace_id", traceID, "tool", name, "error", err)
|
||||
} else {
|
||||
spanID = span.ID
|
||||
}
|
||||
}
|
||||
result, executeErr := execute()
|
||||
if spanID != "" {
|
||||
status := "success"
|
||||
if executeErr != nil {
|
||||
status = "error"
|
||||
}
|
||||
if err := h.traces.FinishSpan(context.WithoutCancel(ctx), spanID, tracepkg.SpanFinishInput{Status: status, Error: errorString(executeErr), Metadata: map[string]any{"tool_call_id": callID}}); err != nil {
|
||||
h.logger.Warn("llm tool span finish failed", "span_id", spanID, "error", err)
|
||||
}
|
||||
}
|
||||
return result, executeErr
|
||||
}
|
||||
|
||||
func (h *RuntimeHTTPHandler) searchWithTrace(ctx context.Context, traceID, knowledgeBaseID, query string, topK int) ([]SearchHit, error) {
|
||||
spanID := ""
|
||||
if h.traces != nil && traceID != "" {
|
||||
span, err := h.traces.StartSpan(context.WithoutCancel(ctx), tracepkg.SpanInput{TraceID: traceID, SpanType: "retrieval", Name: "knowledge.search", Metadata: map[string]any{"knowledge_base_id": knowledgeBaseID, "top_k": topK}})
|
||||
if err != nil {
|
||||
h.logger.Warn("llm retrieval span start failed", "trace_id", traceID, "error", err)
|
||||
} else {
|
||||
spanID = span.ID
|
||||
}
|
||||
}
|
||||
hits, searchErr := h.retriever.Search(ctx, knowledgeBaseID, query, topK)
|
||||
if spanID != "" {
|
||||
status := "success"
|
||||
if searchErr != nil {
|
||||
status = "error"
|
||||
}
|
||||
metadata := map[string]any{"knowledge_base_id": knowledgeBaseID, "hit_count": len(hits)}
|
||||
if err := h.traces.FinishSpan(context.WithoutCancel(ctx), spanID, tracepkg.SpanFinishInput{Status: status, Error: errorString(searchErr), Metadata: metadata}); err != nil {
|
||||
h.logger.Warn("llm retrieval span finish failed", "span_id", spanID, "error", err)
|
||||
}
|
||||
}
|
||||
return hits, searchErr
|
||||
}
|
||||
|
||||
func responseUsage(response map[string]any) (int64, int64) {
|
||||
if response == nil {
|
||||
return 0, 0
|
||||
}
|
||||
usage, _ := response["usage"].(map[string]any)
|
||||
return numberValue(usage["prompt_tokens"], usage["input_tokens"]), numberValue(usage["completion_tokens"], usage["output_tokens"])
|
||||
}
|
||||
|
||||
func numberValue(values ...any) int64 {
|
||||
for _, value := range values {
|
||||
switch number := value.(type) {
|
||||
case float64:
|
||||
return int64(number)
|
||||
case float32:
|
||||
return int64(number)
|
||||
case int:
|
||||
return int64(number)
|
||||
case int64:
|
||||
return number
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func errorString(err error) string {
|
||||
if err == nil {
|
||||
return ""
|
||||
}
|
||||
return fmt.Sprint(err)
|
||||
}
|
||||
@@ -10,12 +10,14 @@ import (
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/apikey"
|
||||
"aigateway.local/core/internal/gateway"
|
||||
"aigateway.local/core/internal/platform/config"
|
||||
"aigateway.local/core/internal/platform/cryptox"
|
||||
"aigateway.local/core/internal/platform/database"
|
||||
tracepkg "aigateway.local/core/internal/trace"
|
||||
)
|
||||
|
||||
func TestWorkbenchPostgreSQLLifecycle(t *testing.T) {
|
||||
@@ -35,7 +37,7 @@ func TestWorkbenchPostgreSQLLifecycle(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cleanup := func() {
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM gateway.notification_channels WHERE name='m4-webhook'; DELETE FROM gateway.applications WHERE code='m4_app'; DELETE FROM gateway.tool_definitions WHERE code='m4_lookup'; DELETE FROM gateway.knowledge_bases WHERE name='m4-integration-kb'; DELETE FROM gateway.prompt_templates WHERE name='m4-integration-prompt'`)
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM gateway.agent_traces WHERE request_id='m4-runtime'; DELETE FROM gateway.notification_channels WHERE name='m4-webhook'; DELETE FROM gateway.applications WHERE code='m4_app'; DELETE FROM gateway.tool_definitions WHERE code='m4_lookup'; DELETE FROM gateway.knowledge_bases WHERE name='m4-integration-kb'; DELETE FROM gateway.prompt_templates WHERE name='m4-integration-prompt'`)
|
||||
}
|
||||
cleanup()
|
||||
defer cleanup()
|
||||
@@ -117,6 +119,8 @@ func TestWorkbenchPostgreSQLLifecycle(t *testing.T) {
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{"message": map[string]any{"role": "assistant", "content": "完成"}}}})
|
||||
})
|
||||
runtime := NewRuntimeHTTPHandler(assets, tools, NewRetriever(assets, nil), staticPrincipalAuthenticator{}, fakeGateway, MarketplaceDeps{})
|
||||
traceStore := tracepkg.NewStore(pool)
|
||||
runtime.SetTraceStore(traceStore)
|
||||
runtimeRequest := httptest.NewRequest(http.MethodPost, "/v1/applications/m4_app/chat/completions", bytes.NewBufferString(`{"messages":[{"role":"user","content":"不可变运行时快照是什么?"}],"variables":{"question":"架构"}}`))
|
||||
runtimeRequest.Header.Set("Authorization", "Bearer test")
|
||||
runtimeRequest = runtimeRequest.WithContext(gateway.WithRequestID(runtimeRequest.Context(), "m4-runtime"))
|
||||
@@ -125,6 +129,14 @@ func TestWorkbenchPostgreSQLLifecycle(t *testing.T) {
|
||||
if runtimeResponse.Code != http.StatusOK || !governed || !strings.Contains(runtimeResponse.Body.String(), `"application"`) {
|
||||
t.Fatalf("runtime status=%d governed=%v body=%s", runtimeResponse.Code, governed, runtimeResponse.Body.String())
|
||||
}
|
||||
traces, err := traceStore.List(ctx, tracepkg.Filter{From: time.Now().Add(-time.Minute), To: time.Now().Add(time.Minute), RequestID: "m4-runtime", Limit: 10})
|
||||
if err != nil || len(traces) != 1 || traces[0].TraceType != "application" || traces[0].ModelCallCount != 1 {
|
||||
t.Fatalf("runtime trace=%+v err=%v", traces, err)
|
||||
}
|
||||
detail, err := traceStore.Get(ctx, traces[0].ID)
|
||||
if err != nil || len(detail.Spans) < 2 {
|
||||
t.Fatalf("runtime trace detail=%+v err=%v", detail, err)
|
||||
}
|
||||
|
||||
signed := ""
|
||||
webhook := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
Reference in New Issue
Block a user