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:
2026-08-13 10:50:51 +08:00
parent b536672000
commit 9501751792
136 changed files with 8024 additions and 1476 deletions
+218
View File
@@ -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, "节点心跳处理失败")
}
}
+60
View File
@@ -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")
}
}
+502
View File
@@ -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)
}
}
+5 -4
View File
@@ -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})
}
+3 -1
View File
@@ -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.
+13 -2
View File
@@ -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)
}
+3 -2
View File
@@ -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)
}
+41 -8
View File
@@ -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)
}
+169 -4
View File
@@ -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)
}
+9 -5
View File
@@ -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
}
+3
View File
@@ -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)
+22 -3
View File
@@ -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
}
+6
View File
@@ -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")
+59 -11
View File
@@ -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
+3 -3
View File
@@ -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")
}
}
+5
View File
@@ -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
View File
@@ -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)
}
}
+12 -1
View File
@@ -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": {},
}
+11 -1
View File
@@ -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)
+18 -4
View File
@@ -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": "站内消息"}},
}},
}
}
+24 -7
View File
@@ -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
View File
@@ -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" {
+48 -22
View File
@@ -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
}
+28 -4
View File
@@ -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")
}
}
+41 -8
View File
@@ -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) {
+68 -8
View File
@@ -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
+10 -1
View File
@@ -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:]
+5
View File
@@ -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 {
+81 -14
View File
@@ -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
+1 -1
View File
@@ -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 {
+6 -3
View File
@@ -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
}
+7
View File
@@ -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})
+38 -7
View File
@@ -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())
}
+6 -29
View File
@@ -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)
+3
View File
@@ -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, "/")
+33
View File
@@ -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
+16 -2
View File
@@ -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
+82 -5
View File
@@ -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))
}
}
+131
View File
@@ -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 年内无匹配执行时间")
}
+35
View File
@@ -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)
}
}
}
+403
View File
@@ -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)
}
+173
View File
@@ -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)
}
}
+441
View File
@@ -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
}
+11
View File
@@ -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)
+166
View File
@@ -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"
}
+459
View File
@@ -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()
}
+87
View File
@@ -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)
}
}
+7 -2
View File
@@ -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
}
+6
View File
@@ -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)
+378
View File
@@ -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 != ""
}
+120
View File
@@ -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)
}
}
+89
View File
@@ -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})
}
+125
View File
@@ -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)
}
}
}
+33 -2
View File
@@ -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 {
+16 -4
View File
@@ -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()
+43 -3
View File
@@ -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
+4
View File
@@ -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
+41 -13
View File
@@ -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 }
+57 -10
View File
@@ -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")
}
}
+148
View File
@@ -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) {