feat: 十轮网关优化 - 安全加固/可观测性/性能/可靠性
CI / lint (push) Has been cancelled
CI / test (push) Has been cancelled
CI / build (push) Has been cancelled
CI / security-scan (push) Has been cancelled

- SSE Keepalive Ping (15s心跳防止代理断连)
- Timing HTTP 头 (X-Timing-Queue/Inference/Total-Ms)
- Adapter Request-ID 传播到后端
- Session 清理日志回调
- Server 安全加固 (ReadHeaderTimeout/MaxHeaderBytes 防 slowloris)
- Usage Tracker 数据保留清理 (retentionDays + 定期清理)
- Config Reload 后 Adapter Registry 更新 (RegisterIfAbsent + RWMutex)
- Rate Limiter 空闲 Bucket 清理 (30分钟过期)
- Shutdown Drain 超时可配置 (ShutdownDrainSeconds)
- Config 模型字段校验增强 (provider/endpoint/actual_model)
- Auth 过期 Key 自动清理 (5分钟扫描)
- Admin API Rate Limiting
- Adapter Health Check 独立超时 (每个 adapter 3s)
- TCP 连接阶段超时 (DialContext 5s + KeepAlive 30s)
- 幂等键缓存、审计日志、Gzip 中间件、CORS Expose Headers
- Backpressure 响应头、熔断器 Prometheus 指标
- 连接池优化、Trace-ID 全链路传播
This commit is contained in:
selfrelease
2026-08-03 15:43:11 +08:00
parent e7e98271d4
commit da9c8334d8
27 changed files with 3002 additions and 309 deletions
+66
View File
@@ -0,0 +1,66 @@
package main
import (
"flag"
"fmt"
"os"
"strings"
"github.com/edgeai/gateway/internal/auth"
"github.com/edgeai/gateway/internal/observability"
)
func main() {
dbPath := flag.String("db", "/tmp/edgeai-data/auth.db", "auth database path")
appID := flag.String("app", "test", "application ID")
tenantID := flag.String("tenant", "default", "tenant ID")
name := flag.String("name", "test", "key name")
models := flag.String("models", "", "allowed models (comma-separated, empty=all)")
priorities := flag.String("priorities", "0,1,2,3", "allowed priorities (comma-separated)")
isAdmin := flag.Bool("admin", false, "is admin key")
keyValue := flag.String("key", "", "specific API key value (auto-generated if empty)")
flag.Parse()
observability.SetLogLevel("info")
logger := observability.GetLogger()
a, err := auth.NewAuthenticator(*dbPath, logger)
if err != nil {
fmt.Fprintf(os.Stderr, "error: %v\n", err)
os.Exit(1)
}
defer a.Close()
var allowedModels []string
if *models != "" {
allowedModels = strings.Split(*models, ",")
}
var allowedPriorities []int
for _, p := range strings.Split(*priorities, ",") {
p = strings.TrimSpace(p)
if p != "" {
var v int
fmt.Sscanf(p, "%d", &v)
allowedPriorities = append(allowedPriorities, v)
}
}
apiKey := *keyValue
if apiKey == "" {
apiKey = auth.GenerateAPIKey()
}
identity := &auth.AppIdentity{
AppID: *appID,
TenantID: *tenantID,
Name: *name,
AllowedModels: allowedModels,
AllowedPriorities: allowedPriorities,
IsAdmin: *isAdmin,
}
if err := a.AddKey(apiKey, identity); err != nil {
fmt.Fprintf(os.Stderr, "AddKey error: %v\n", err)
os.Exit(1)
}
fmt.Printf("API key created successfully:\n key: %s\n app: %s\n name: %s\n admin: %v\n", apiKey, *appID, *name, *isAdmin)
}
+12 -15
View File
@@ -1,14 +1,18 @@
server:
host: "0.0.0.0"
port: 8080
port: 39000
admin_port: 8081
max_request_body_mb: 20
cors_allowed_origins: ["*"]
auth:
enabled: true
methods: [api_key]
jwt_issuer: edge-ai-gateway
jwt_secret_env: EDGEAI_JWT_SECRET
rate_limit_per_minute: 60
rate_limit_burst: 10
usage_window_minutes: 60
scheduler:
max_running_tasks: 8
@@ -21,9 +25,9 @@ timeouts:
default_connect_ms: 5000
default_queue_ms: 5000
default_first_token_ms: 10000
default_inference_ms: 60000
default_inference_ms: 120000
default_idle_ms: 15000
default_total_ms: 90000
default_total_ms: 180000
cancel_grace_period_ms: 3000
context:
@@ -32,11 +36,14 @@ context:
max_session_messages: 200
session_idle_ttl_minutes: 60
enable_prompt_persistence: false
enable_llm_summary: true
summary_max_tokens: 256
summary_timeout_seconds: 15
models:
general-chat:
provider: ollama
actual_model: deepseek-r1:1.5b
actual_model: qwen2:latest
endpoint: http://127.0.0.1:11434
context_window: 32768
max_output_tokens: 4096
@@ -46,7 +53,7 @@ models:
fast-chat:
provider: ollama
actual_model: deepseek-r1:1.5b
actual_model: qwen2:latest
endpoint: http://127.0.0.1:11434
context_window: 16384
max_output_tokens: 2048
@@ -54,16 +61,6 @@ models:
residency: on_demand
idle_unload_seconds: 600
vllm-chat:
provider: vllm
actual_model: deepseek-r1:1.5b
endpoint: http://127.0.0.1:8000
context_window: 32768
max_output_tokens: 4096
max_concurrency: 4
residency: always
cancel_supported: true
routing:
sensitive_data_local_only: true
allow_cloud_fallback_by_default: false
Executable
BIN
View File
Binary file not shown.
+1 -1
View File
@@ -7,4 +7,4 @@ require (
gopkg.in/yaml.v3 v3.0.1
)
require github.com/mattn/go-sqlite3 v1.14.49
require github.com/mattn/go-sqlite3 v1.14.22
+2 -2
View File
@@ -1,7 +1,7 @@
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/mattn/go-sqlite3 v1.14.49 h1:B8jBHC3xhxZgxztrgruTuLucebnULQnx4W7cF7SAE9w=
github.com/mattn/go-sqlite3 v1.14.49/go.mod h1:6JTjA44L93a0QCyJef5YvlPoKXntQPjzWv5gtm9sB6w=
github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU=
github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+42 -19
View File
@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"io"
"sync"
"github.com/edgeai/gateway/pkg/api"
)
@@ -31,33 +32,33 @@ type ModelAdapter interface {
// ChatRequest is the internal request sent to an adapter.
type ChatRequest struct {
RequestID string
Model string // actual model name
Messages []api.Message
MaxTokens int
Temperature *float64
TopP *float64
Stream bool
CancelCh <-chan struct{}
RequestID string
Model string // actual model name
Messages []api.Message
MaxTokens int
Temperature *float64
TopP *float64
Stream bool
CancelCh <-chan struct{}
}
// ChatResponse is the internal response from an adapter.
type ChatResponse struct {
Content string
FinishReason string
InputTokens int
OutputTokens int
ActualModel string
Content string
FinishReason string
InputTokens int
OutputTokens int
ActualModel string
}
// StreamChunk represents a single chunk in a streaming response.
type StreamChunk struct {
Delta string
FinishReason string
InputTokens int
OutputTokens int
Error error
Done bool
Delta string
FinishReason string
InputTokens int
OutputTokens int
Error error
Done bool
}
// ModelInfo describes a model available in the engine.
@@ -68,6 +69,7 @@ type ModelInfo struct {
// Registry manages model adapters by provider name.
type Registry struct {
mu sync.RWMutex
adapters map[string]ModelAdapter
}
@@ -75,11 +77,29 @@ func NewRegistry() *Registry {
return &Registry{adapters: make(map[string]ModelAdapter)}
}
// Register 注册一个 adapter,如果同名 adapter 已存在则覆盖。
func (r *Registry) Register(name string, adapter ModelAdapter) {
r.mu.Lock()
defer r.mu.Unlock()
r.adapters[name] = adapter
}
// RegisterIfAbsent 注册一个 adapter,仅当该 provider 尚未注册时才创建。
// 返回 true 表示新注册,false 表示已存在。
func (r *Registry) RegisterIfAbsent(name string, factory func() ModelAdapter) bool {
r.mu.Lock()
defer r.mu.Unlock()
if _, exists := r.adapters[name]; exists {
return false
}
r.adapters[name] = factory()
return true
}
// Get 返回指定名称的 adapter。
func (r *Registry) Get(name string) (ModelAdapter, error) {
r.mu.RLock()
defer r.mu.RUnlock()
a, ok := r.adapters[name]
if !ok {
return nil, fmt.Errorf("adapter not found: %s", name)
@@ -87,7 +107,10 @@ func (r *Registry) Get(name string) (ModelAdapter, error) {
return a, nil
}
// Names 返回所有已注册 adapter 的名称列表。
func (r *Registry) Names() []string {
r.mu.RLock()
defer r.mu.RUnlock()
names := make([]string, 0, len(r.adapters))
for n := range r.adapters {
names = append(names, n)
+37 -14
View File
@@ -6,11 +6,30 @@ import (
"encoding/json"
"fmt"
"io"
"net"
"net/http"
"strings"
"time"
)
// newHTTPClient 创建带连接池优化的 HTTP 客户端,复用 TCP 连接以减少延迟。
// timeout 为整体请求超时,connectTimeout 为 TCP 连接阶段超时。
func newHTTPClient(timeout time.Duration) *http.Client {
return &http.Client{
Timeout: timeout,
Transport: &http.Transport{
MaxIdleConns: 100,
MaxIdleConnsPerHost: 20,
IdleConnTimeout: 90 * time.Second,
MaxConnsPerHost: 0, // 不限制
DialContext: (&net.Dialer{
Timeout: 5 * time.Second,
KeepAlive: 30 * time.Second,
}).DialContext,
},
}
}
// OllamaAdapter implements ModelAdapter for Ollama inference engine.
type OllamaAdapter struct {
endpoint string
@@ -20,10 +39,8 @@ type OllamaAdapter struct {
// NewOllamaAdapter creates a new Ollama adapter.
func NewOllamaAdapter(endpoint string) *OllamaAdapter {
return &OllamaAdapter{
endpoint: strings.TrimRight(endpoint, "/"),
httpClient: &http.Client{
Timeout: 120 * time.Second,
},
endpoint: strings.TrimRight(endpoint, "/"),
httpClient: newHTTPClient(120 * time.Second),
}
}
@@ -52,20 +69,20 @@ type ollamaOptions struct {
// ollamaChatResponse is the Ollama /api/chat non-streaming response.
type ollamaChatResponse struct {
Model string `json:"model"`
Message ollamaMsg `json:"message"`
Done bool `json:"done"`
PromptEvalCount int `json:"prompt_eval_count"`
EvalCount int `json:"eval_count"`
Model string `json:"model"`
Message ollamaMsg `json:"message"`
Done bool `json:"done"`
PromptEvalCount int `json:"prompt_eval_count"`
EvalCount int `json:"eval_count"`
}
// ollamaChatStreamResponse is a single chunk in Ollama streaming response.
type ollamaChatStreamResponse struct {
Model string `json:"model"`
Message ollamaMsg `json:"message"`
Done bool `json:"done"`
PromptEvalCount int `json:"prompt_eval_count,omitempty"`
EvalCount int `json:"eval_count,omitempty"`
Model string `json:"model"`
Message ollamaMsg `json:"message"`
Done bool `json:"done"`
PromptEvalCount int `json:"prompt_eval_count,omitempty"`
EvalCount int `json:"eval_count,omitempty"`
}
func (a *OllamaAdapter) ChatCompletion(ctx context.Context, req *ChatRequest) (*ChatResponse, error) {
@@ -81,6 +98,9 @@ func (a *OllamaAdapter) ChatCompletion(ctx context.Context, req *ChatRequest) (*
return nil, fmt.Errorf("create ollama request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
if req.RequestID != "" {
httpReq.Header.Set("X-Request-ID", req.RequestID)
}
resp, err := a.httpClient.Do(httpReq)
if err != nil {
@@ -120,6 +140,9 @@ func (a *OllamaAdapter) ChatCompletionStream(ctx context.Context, req *ChatReque
return nil, fmt.Errorf("create ollama stream request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
if req.RequestID != "" {
httpReq.Header.Set("X-Request-ID", req.RequestID)
}
resp, err := a.httpClient.Do(httpReq)
if err != nil {
+17 -13
View File
@@ -24,10 +24,8 @@ type VLLMAdapter struct {
// NewVLLMAdapter creates a new vLLM adapter.
func NewVLLMAdapter(endpoint string) *VLLMAdapter {
return &VLLMAdapter{
endpoint: strings.TrimRight(endpoint, "/"),
httpClient: &http.Client{
Timeout: 120 * time.Second,
},
endpoint: strings.TrimRight(endpoint, "/"),
httpClient: newHTTPClient(120 * time.Second),
}
}
@@ -37,12 +35,12 @@ func (a *VLLMAdapter) Name() string {
// vllmChatRequest is the OpenAI-compatible chat request for vLLM.
type vllmChatRequest struct {
Model string `json:"model"`
Messages []vllmMsg `json:"messages"`
Stream bool `json:"stream"`
MaxTokens int `json:"max_tokens,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
Model string `json:"model"`
Messages []vllmMsg `json:"messages"`
Stream bool `json:"stream"`
MaxTokens int `json:"max_tokens,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
}
type vllmMsg struct {
@@ -55,9 +53,9 @@ type vllmChatResponse struct {
ID string `json:"id"`
Model string `json:"model"`
Choices []struct {
Index int `json:"index"`
Index int `json:"index"`
Message vllmMsg `json:"message"`
FinishReason string `json:"finish_reason"`
FinishReason string `json:"finish_reason"`
} `json:"choices"`
Usage struct {
PromptTokens int `json:"prompt_tokens"`
@@ -71,7 +69,7 @@ type vllmStreamChunk struct {
ID string `json:"id"`
Model string `json:"model"`
Choices []struct {
Index int `json:"index"`
Index int `json:"index"`
Delta vllmMsg `json:"delta"`
FinishReason *string `json:"finish_reason"`
} `json:"choices"`
@@ -95,6 +93,9 @@ func (a *VLLMAdapter) ChatCompletion(ctx context.Context, req *ChatRequest) (*Ch
return nil, fmt.Errorf("create vllm request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
if req.RequestID != "" {
httpReq.Header.Set("X-Request-ID", req.RequestID)
}
resp, err := a.httpClient.Do(httpReq)
if err != nil {
@@ -145,6 +146,9 @@ func (a *VLLMAdapter) ChatCompletionStream(ctx context.Context, req *ChatRequest
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Accept", "text/event-stream")
if req.RequestID != "" {
httpReq.Header.Set("X-Request-ID", req.RequestID)
}
resp, err := a.httpClient.Do(httpReq)
if err != nil {
+220 -7
View File
@@ -2,6 +2,7 @@ package auth
import (
"context"
"crypto/rand"
"crypto/sha256"
"database/sql"
"encoding/hex"
@@ -10,6 +11,7 @@ import (
"net/http"
"strings"
"sync"
"time"
"github.com/edgeai/gateway/internal/handler"
"github.com/edgeai/gateway/internal/middleware"
@@ -25,6 +27,7 @@ type AppIdentity struct {
AllowedModels []string
AllowedPriorities []int
IsAdmin bool
ExpiresAt *time.Time // 过期时间,nil 表示永不过期
}
type contextKey string
@@ -39,6 +42,7 @@ type Authenticator struct {
keys map[string]*AppIdentity // hashed_key -> identity
db *sql.DB
logger *observability.Logger
stopCh chan struct{} // 停止清理 goroutine
}
// NewAuthenticator creates a new Authenticator with SQLite storage.
@@ -56,12 +60,15 @@ func NewAuthenticator(dbPath string, logger *observability.Logger) (*Authenticat
keys: make(map[string]*AppIdentity),
db: db,
logger: logger,
stopCh: make(chan struct{}),
}
if err := a.loadKeys(); err != nil {
return nil, fmt.Errorf("load api keys: %w", err)
}
go a.cleanupExpiredKeys()
return a, nil
}
@@ -85,7 +92,7 @@ func initAuthDB(db *sql.DB) error {
}
func (a *Authenticator) loadKeys() error {
rows, err := a.db.Query(`SELECT key_hash, app_id, tenant_id, name, allowed_models, allowed_priorities, is_admin FROM api_keys WHERE enabled = 1`)
rows, err := a.db.Query(`SELECT key_hash, app_id, tenant_id, name, allowed_models, allowed_priorities, is_admin, expires_at FROM api_keys WHERE enabled = 1`)
if err != nil {
return err
}
@@ -94,7 +101,8 @@ func (a *Authenticator) loadKeys() error {
for rows.Next() {
var hash, appID, tenantID, name, allowedModelsJSON, allowedPrioritiesJSON string
var isAdmin int
if err := rows.Scan(&hash, &appID, &tenantID, &name, &allowedModelsJSON, &allowedPrioritiesJSON, &isAdmin); err != nil {
var expiresAtStr sql.NullString
if err := rows.Scan(&hash, &appID, &tenantID, &name, &allowedModelsJSON, &allowedPrioritiesJSON, &isAdmin, &expiresAtStr); err != nil {
return err
}
@@ -104,6 +112,11 @@ func (a *Authenticator) loadKeys() error {
Name: name,
IsAdmin: isAdmin == 1,
}
if expiresAtStr.Valid && expiresAtStr.String != "" {
if t, err := time.Parse(time.RFC3339, expiresAtStr.String); err == nil {
identity.ExpiresAt = &t
}
}
if allowedModelsJSON != "" && allowedModelsJSON != "null" {
json.Unmarshal([]byte(allowedModelsJSON), &identity.AllowedModels)
}
@@ -116,6 +129,13 @@ func (a *Authenticator) loadKeys() error {
return rows.Err()
}
// GenerateAPIKey 生成一个随机的 API Key(前缀 edgeai- + 32 字节随机十六进制)。
func GenerateAPIKey() string {
b := make([]byte, 32)
rand.Read(b)
return "edgeai-" + hex.EncodeToString(b)
}
// hashKey hashes an API key with SHA-256.
func hashKey(key string) string {
h := sha256.Sum256([]byte(key))
@@ -123,6 +143,7 @@ func hashKey(key string) string {
}
// Authenticate validates an API key and returns the AppIdentity.
// 检查 Key 是否存在且未过期。
func (a *Authenticator) Authenticate(apiKey string) (*AppIdentity, bool) {
hash := hashKey(apiKey)
a.mu.RLock()
@@ -131,19 +152,29 @@ func (a *Authenticator) Authenticate(apiKey string) (*AppIdentity, bool) {
if !ok {
return nil, false
}
// 检查过期
if identity.ExpiresAt != nil && time.Now().After(*identity.ExpiresAt) {
return nil, false
}
return identity, true
}
// AddKey adds a new API key (for management API).
// 如果 identity.ExpiresAt 不为 nil,则设置过期时间。
func (a *Authenticator) AddKey(apiKey string, identity *AppIdentity) error {
hash := hashKey(apiKey)
allowedModelsJSON, _ := json.Marshal(identity.AllowedModels)
allowedPrioritiesJSON, _ := json.Marshal(identity.AllowedPriorities)
var expiresAt interface{}
if identity.ExpiresAt != nil {
expiresAt = identity.ExpiresAt.Format(time.RFC3339)
}
_, err := a.db.Exec(
`INSERT INTO api_keys (app_id, tenant_id, name, key_hash, allowed_models, allowed_priorities, is_admin, enabled)
VALUES (?, ?, ?, ?, ?, ?, ?, 1)`,
identity.AppID, identity.TenantID, identity.Name, hash, string(allowedModelsJSON), string(allowedPrioritiesJSON), isAdminInt(identity.IsAdmin),
`INSERT INTO api_keys (app_id, tenant_id, name, key_hash, allowed_models, allowed_priorities, is_admin, enabled, expires_at)
VALUES (?, ?, ?, ?, ?, ?, ?, 1, ?)`,
identity.AppID, identity.TenantID, identity.Name, hash, string(allowedModelsJSON), string(allowedPrioritiesJSON), isAdminInt(identity.IsAdmin), expiresAt,
)
if err != nil {
return err
@@ -155,6 +186,164 @@ func (a *Authenticator) AddKey(apiKey string, identity *AppIdentity) error {
return nil
}
// KeyInfo 是 API Key 的元信息(不含哈希值)。
// 用于管理 API 列出所有 Key。
type KeyInfo struct {
ID int64 `json:"id"`
AppID string `json:"app_id"`
TenantID string `json:"tenant_id"`
Name string `json:"name"`
AllowedModels []string `json:"allowed_models"`
AllowedPriorities []int `json:"allowed_priorities"`
IsAdmin bool `json:"is_admin"`
Enabled bool `json:"enabled"`
CreatedAt string `json:"created_at"`
ExpiresAt string `json:"expires_at,omitempty"`
}
// ListKeys 列出所有 API Key 的元信息(不含哈希值)。
func (a *Authenticator) ListKeys() ([]KeyInfo, error) {
rows, err := a.db.Query(
`SELECT id, app_id, tenant_id, name, allowed_models, allowed_priorities, is_admin, enabled, created_at, expires_at FROM api_keys ORDER BY id DESC`,
)
if err != nil {
return nil, fmt.Errorf("query api keys: %w", err)
}
defer rows.Close()
var keys []KeyInfo
for rows.Next() {
var ki KeyInfo
var allowedModelsJSON, allowedPrioritiesJSON string
var isAdmin, enabled int
var expiresAt sql.NullString
if err := rows.Scan(&ki.ID, &ki.AppID, &ki.TenantID, &ki.Name, &allowedModelsJSON, &allowedPrioritiesJSON, &isAdmin, &enabled, &ki.CreatedAt, &expiresAt); err != nil {
return nil, err
}
ki.IsAdmin = isAdmin == 1
ki.Enabled = enabled == 1
if expiresAt.Valid && expiresAt.String != "" {
ki.ExpiresAt = expiresAt.String
}
if allowedModelsJSON != "" && allowedModelsJSON != "null" {
json.Unmarshal([]byte(allowedModelsJSON), &ki.AllowedModels)
}
if allowedPrioritiesJSON != "" && allowedPrioritiesJSON != "null" {
json.Unmarshal([]byte(allowedPrioritiesJSON), &ki.AllowedPriorities)
}
keys = append(keys, ki)
}
return keys, rows.Err()
}
// DisableKey 禁用一个 API Key(通过数据库 ID)。
func (a *Authenticator) DisableKey(id int64) error {
// 先查出 hash 以便从内存中移除
var hash string
err := a.db.QueryRow(`SELECT key_hash FROM api_keys WHERE id = ?`, id).Scan(&hash)
if err == sql.ErrNoRows {
return fmt.Errorf("api key not found: %d", id)
}
if err != nil {
return fmt.Errorf("query api key: %w", err)
}
_, err = a.db.Exec(`UPDATE api_keys SET enabled = 0 WHERE id = ?`, id)
if err != nil {
return fmt.Errorf("disable api key: %w", err)
}
a.mu.Lock()
delete(a.keys, hash)
a.mu.Unlock()
return nil
}
// DeleteKey 彻底删除一个 API Key(通过数据库 ID)。
func (a *Authenticator) DeleteKey(id int64) error {
var hash string
err := a.db.QueryRow(`SELECT key_hash FROM api_keys WHERE id = ?`, id).Scan(&hash)
if err == sql.ErrNoRows {
return fmt.Errorf("api key not found: %d", id)
}
if err != nil {
return fmt.Errorf("query api key: %w", err)
}
_, err = a.db.Exec(`DELETE FROM api_keys WHERE id = ?`, id)
if err != nil {
return fmt.Errorf("delete api key: %w", err)
}
a.mu.Lock()
delete(a.keys, hash)
a.mu.Unlock()
return nil
}
// RotateKey 轮换 API Key:生成新 Key,禁用旧 Key,返回新 Key 值。
func (a *Authenticator) RotateKey(id int64) (string, error) {
var oldHash string
var appID, tenantID, name, allowedModelsJSON, allowedPrioritiesJSON string
var isAdmin int
var expiresAt sql.NullString
err := a.db.QueryRow(
`SELECT key_hash, app_id, tenant_id, name, allowed_models, allowed_priorities, is_admin, expires_at FROM api_keys WHERE id = ?`,
id,
).Scan(&oldHash, &appID, &tenantID, &name, &allowedModelsJSON, &allowedPrioritiesJSON, &isAdmin, &expiresAt)
if err == sql.ErrNoRows {
return "", fmt.Errorf("api key not found: %d", id)
}
if err != nil {
return "", fmt.Errorf("query api key: %w", err)
}
newKey := GenerateAPIKey()
newHash := hashKey(newKey)
// 插入新 Key,继承旧 Key 的所有属性
_, err = a.db.Exec(
`INSERT INTO api_keys (app_id, tenant_id, name, key_hash, allowed_models, allowed_priorities, is_admin, enabled, expires_at)
VALUES (?, ?, ?, ?, ?, ?, ?, 1, ?)`,
appID, tenantID, name+" (rotated)", newHash, allowedModelsJSON, allowedPrioritiesJSON, isAdmin, expiresAt,
)
if err != nil {
return "", fmt.Errorf("insert rotated key: %w", err)
}
// 禁用旧 Key
_, err = a.db.Exec(`UPDATE api_keys SET enabled = 0 WHERE id = ?`, id)
if err != nil {
return "", fmt.Errorf("disable old key: %w", err)
}
// 更新内存:移除旧 Key,添加新 Key
identity := &AppIdentity{
AppID: appID,
TenantID: tenantID,
Name: name + " (rotated)",
IsAdmin: isAdmin == 1,
}
if allowedModelsJSON != "" && allowedModelsJSON != "null" {
json.Unmarshal([]byte(allowedModelsJSON), &identity.AllowedModels)
}
if allowedPrioritiesJSON != "" && allowedPrioritiesJSON != "null" {
json.Unmarshal([]byte(allowedPrioritiesJSON), &identity.AllowedPriorities)
}
if expiresAt.Valid && expiresAt.String != "" {
if t, err := time.Parse(time.RFC3339, expiresAt.String); err == nil {
identity.ExpiresAt = &t
}
}
a.mu.Lock()
delete(a.keys, oldHash)
a.keys[newHash] = identity
a.mu.Unlock()
return newKey, nil
}
// Middleware returns an HTTP middleware that enforces API Key authentication.
func (a *Authenticator) Middleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -219,12 +408,13 @@ func RequireAdmin(next http.Handler) http.Handler {
}
// CheckModelPermission verifies the app can access the given model.
// 支持通配符 "*" 匹配所有模型。
func CheckModelPermission(identity *AppIdentity, model string) bool {
if len(identity.AllowedModels) == 0 {
return true // empty = all models allowed
}
for _, m := range identity.AllowedModels {
if m == model {
if m == model || m == "*" {
return true
}
}
@@ -251,10 +441,33 @@ func isAdminInt(b bool) int {
return 0
}
// Close closes the database connection.
// Close closes the database connection and stops background cleanup.
func (a *Authenticator) Close() error {
close(a.stopCh)
return a.db.Close()
}
// cleanupExpiredKeys 定期清理内存中已过期的 API Key,防止内存泄漏。
func (a *Authenticator) cleanupExpiredKeys() {
ticker := time.NewTicker(5 * time.Minute)
defer ticker.Stop()
for {
select {
case <-ticker.C:
now := time.Now()
a.mu.Lock()
for hash, identity := range a.keys {
if identity.ExpiresAt != nil && now.After(*identity.ExpiresAt) {
delete(a.keys, hash)
}
}
a.mu.Unlock()
case <-a.stopCh:
return
}
}
}
// Ensure middleware import is used.
var _ = middleware.GetRequestID
+106 -39
View File
@@ -12,35 +12,39 @@ import (
// Config is the root configuration structure.
type Config struct {
Server ServerConfig `yaml:"server"`
Auth AuthConfig `yaml:"auth"`
Scheduler SchedulerConfig `yaml:"scheduler"`
Timeouts TimeoutConfig `yaml:"timeouts"`
Context ContextConfig `yaml:"context"`
Server ServerConfig `yaml:"server"`
Auth AuthConfig `yaml:"auth"`
Scheduler SchedulerConfig `yaml:"scheduler"`
Timeouts TimeoutConfig `yaml:"timeouts"`
Context ContextConfig `yaml:"context"`
Models map[string]ModelConfig `yaml:"models"`
Routing RoutingConfig `yaml:"routing"`
CircuitBreaker CircuitBreakerConfig `yaml:"circuit_breaker"`
Backpressure BackpressureConfig `yaml:"backpressure"`
Observability ObservabilityConfig `yaml:"observability"`
Storage StorageConfig `yaml:"storage"`
Routing RoutingConfig `yaml:"routing"`
CircuitBreaker CircuitBreakerConfig `yaml:"circuit_breaker"`
Backpressure BackpressureConfig `yaml:"backpressure"`
Observability ObservabilityConfig `yaml:"observability"`
Storage StorageConfig `yaml:"storage"`
}
type ServerConfig struct {
Host string `yaml:"host"`
Port int `yaml:"port"`
AdminPort int `yaml:"admin_port"`
MaxRequestBodyMB int `yaml:"max_request_body_mb"`
Host string `yaml:"host"`
Port int `yaml:"port"`
AdminPort int `yaml:"admin_port"`
MaxRequestBodyMB int `yaml:"max_request_body_mb"`
CORSAllowedOrigins []string `yaml:"cors_allowed_origins"`
}
type AuthConfig struct {
Enabled bool `yaml:"enabled"`
Methods []string `yaml:"methods"`
JWTIssuer string `yaml:"jwt_issuer"`
JWTSecretEnv string `yaml:"jwt_secret_env"`
Enabled bool `yaml:"enabled"`
Methods []string `yaml:"methods"`
JWTIssuer string `yaml:"jwt_issuer"`
JWTSecretEnv string `yaml:"jwt_secret_env"`
RateLimitPerMinute int `yaml:"rate_limit_per_minute"`
RateLimitBurst int `yaml:"rate_limit_burst"`
UsageWindowMinutes int `yaml:"usage_window_minutes"`
}
type SchedulerConfig struct {
MaxRunningTasks int `yaml:"max_running_tasks"`
MaxRunningTasks int `yaml:"max_running_tasks"`
MaxQueuedTasks int `yaml:"max_queued_tasks"`
Fairness string `yaml:"fairness"`
PriorityAgingSeconds int `yaml:"priority_aging_seconds"`
@@ -48,21 +52,25 @@ type SchedulerConfig struct {
}
type TimeoutConfig struct {
DefaultConnectMs int `yaml:"default_connect_ms"`
DefaultQueueMs int `yaml:"default_queue_ms"`
DefaultFirstTokenMs int `yaml:"default_first_token_ms"`
DefaultInferenceMs int `yaml:"default_inference_ms"`
DefaultIdleMs int `yaml:"default_idle_ms"`
DefaultTotalMs int `yaml:"default_total_ms"`
CancelGracePeriodMs int `yaml:"cancel_grace_period_ms"`
DefaultConnectMs int `yaml:"default_connect_ms"`
DefaultQueueMs int `yaml:"default_queue_ms"`
DefaultFirstTokenMs int `yaml:"default_first_token_ms"`
DefaultInferenceMs int `yaml:"default_inference_ms"`
DefaultIdleMs int `yaml:"default_idle_ms"`
DefaultTotalMs int `yaml:"default_total_ms"`
CancelGracePeriodMs int `yaml:"cancel_grace_period_ms"`
ShutdownDrainSeconds int `yaml:"shutdown_drain_seconds"` // 优雅关机时等待运行中任务的超时时间
}
type ContextConfig struct {
SafetyMarginRatio float64 `yaml:"safety_margin_ratio"`
DefaultPolicy string `yaml:"default_policy"`
MaxSessionMessages int `yaml:"max_session_messages"`
SessionIdleTTLMinutes int `yaml:"session_idle_ttl_minutes"`
EnablePromptPersistence bool `yaml:"enable_prompt_persistence"`
SafetyMarginRatio float64 `yaml:"safety_margin_ratio"`
DefaultPolicy string `yaml:"default_policy"`
MaxSessionMessages int `yaml:"max_session_messages"`
SessionIdleTTLMinutes int `yaml:"session_idle_ttl_minutes"`
EnablePromptPersistence bool `yaml:"enable_prompt_persistence"`
EnableLLMSummary bool `yaml:"enable_llm_summary"`
SummaryMaxTokens int `yaml:"summary_max_tokens"`
SummaryTimeoutSeconds int `yaml:"summary_timeout_seconds"`
}
type ModelConfig struct {
@@ -80,7 +88,7 @@ type ModelConfig struct {
type RoutingConfig struct {
SensitiveDataLocalOnly bool `yaml:"sensitive_data_local_only"`
AllowCloudFallbackByDefault bool `yaml:"allow_cloud_fallback_by_default"`
OverloadStrategy []string `yaml:"overload_strategy"`
OverloadStrategy []string `yaml:"overload_strategy"`
}
type CircuitBreakerConfig struct {
@@ -98,12 +106,12 @@ type BackpressureConfig struct {
}
type ObservabilityConfig struct {
MetricsEnabled bool `yaml:"metrics_enabled"`
MetricsPath string `yaml:"metrics_path"`
TracingEnabled bool `yaml:"tracing_enabled"`
PromptLogging string `yaml:"prompt_logging"`
AuditRetentionDays int `yaml:"audit_retention_days"`
LogLevel string `yaml:"log_level"`
MetricsEnabled bool `yaml:"metrics_enabled"`
MetricsPath string `yaml:"metrics_path"`
TracingEnabled bool `yaml:"tracing_enabled"`
PromptLogging string `yaml:"prompt_logging"`
AuditRetentionDays int `yaml:"audit_retention_days"`
LogLevel string `yaml:"log_level"`
}
type StorageConfig struct {
@@ -153,7 +161,7 @@ func applyDefaults(cfg *Config) {
cfg.Server.Host = "0.0.0.0"
}
if cfg.Server.Port == 0 {
cfg.Server.Port = 8080
cfg.Server.Port = 39000
}
if cfg.Server.AdminPort == 0 {
cfg.Server.AdminPort = 8081
@@ -194,6 +202,9 @@ func applyDefaults(cfg *Config) {
if cfg.Timeouts.CancelGracePeriodMs == 0 {
cfg.Timeouts.CancelGracePeriodMs = 3000
}
if cfg.Timeouts.ShutdownDrainSeconds == 0 {
cfg.Timeouts.ShutdownDrainSeconds = 10
}
if cfg.Context.SafetyMarginRatio == 0 {
cfg.Context.SafetyMarginRatio = 0.08
}
@@ -206,6 +217,12 @@ func applyDefaults(cfg *Config) {
if cfg.Context.SessionIdleTTLMinutes == 0 {
cfg.Context.SessionIdleTTLMinutes = 60
}
if cfg.Context.SummaryMaxTokens == 0 {
cfg.Context.SummaryMaxTokens = 256
}
if cfg.Context.SummaryTimeoutSeconds == 0 {
cfg.Context.SummaryTimeoutSeconds = 15
}
if cfg.CircuitBreaker.ErrorRateThreshold == 0 {
cfg.CircuitBreaker.ErrorRateThreshold = 0.1
}
@@ -242,6 +259,15 @@ func applyDefaults(cfg *Config) {
if cfg.Observability.AuditRetentionDays == 0 {
cfg.Observability.AuditRetentionDays = 180
}
if cfg.Auth.RateLimitPerMinute <= 0 {
cfg.Auth.RateLimitPerMinute = 60
}
if cfg.Auth.RateLimitBurst <= 0 {
cfg.Auth.RateLimitBurst = 10
}
if cfg.Auth.UsageWindowMinutes <= 0 {
cfg.Auth.UsageWindowMinutes = 60
}
if cfg.Storage.SessionDB == "" {
cfg.Storage.SessionDB = "sqlite:///var/lib/edgeai/sessions.db"
}
@@ -251,6 +277,18 @@ func applyDefaults(cfg *Config) {
}
func validate(cfg *Config) error {
if cfg.Server.Port <= 0 || cfg.Server.Port > 65535 {
return fmt.Errorf("server.port must be in [1, 65535]")
}
if cfg.Server.AdminPort < 0 || cfg.Server.AdminPort > 65535 {
return fmt.Errorf("server.admin_port must be in [0, 65535]")
}
if cfg.Server.AdminPort > 0 && cfg.Server.AdminPort == cfg.Server.Port {
return fmt.Errorf("server.admin_port must differ from server.port")
}
if cfg.Server.MaxRequestBodyMB <= 0 {
return fmt.Errorf("server.max_request_body_mb must be positive")
}
if cfg.Scheduler.MaxRunningTasks <= 0 {
return fmt.Errorf("scheduler.max_running_tasks must be positive")
}
@@ -260,12 +298,41 @@ func validate(cfg *Config) error {
if cfg.Context.SafetyMarginRatio < 0 || cfg.Context.SafetyMarginRatio >= 1 {
return fmt.Errorf("context.safety_margin_ratio must be in [0, 1)")
}
if cfg.Context.MaxSessionMessages <= 0 {
return fmt.Errorf("context.max_session_messages must be positive")
}
if cfg.Backpressure.Level1Threshold >= cfg.Backpressure.Level2Threshold {
return fmt.Errorf("backpressure level1 threshold must be less than level2")
}
if cfg.Backpressure.Level2Threshold >= cfg.Backpressure.Level3Threshold {
return fmt.Errorf("backpressure level2 threshold must be less than level3")
}
if cfg.Auth.RateLimitPerMinute <= 0 {
return fmt.Errorf("auth.rate_limit_per_minute must be positive")
}
if cfg.Auth.RateLimitBurst <= 0 {
return fmt.Errorf("auth.rate_limit_burst must be positive")
}
if len(cfg.Models) == 0 {
return fmt.Errorf("models must not be empty")
}
for name, mc := range cfg.Models {
if mc.Provider == "" {
return fmt.Errorf("models.%s.provider must not be empty", name)
}
if mc.Endpoint == "" {
return fmt.Errorf("models.%s.endpoint must not be empty", name)
}
if mc.ActualModel == "" {
return fmt.Errorf("models.%s.actual_model must not be empty", name)
}
if mc.ContextWindow < 0 {
return fmt.Errorf("models.%s.context_window must not be negative", name)
}
if mc.MaxOutputTokens < 0 {
return fmt.Errorf("models.%s.max_output_tokens must not be negative", name)
}
}
return nil
}
+79
View File
@@ -0,0 +1,79 @@
package context
import (
"context"
"fmt"
"strings"
"github.com/edgeai/gateway/internal/adapter"
"github.com/edgeai/gateway/pkg/api"
)
// AdapterSummaryGenerator 使用 adapter 调用 LLM 生成对话摘要。
type AdapterSummaryGenerator struct {
adapter adapter.ModelAdapter
model string
maxTokens int
temperature float64
}
// NewAdapterSummaryGenerator 创建基于 adapter 的摘要生成器。
// model 为实际模型名(如 qwen2:latest),maxTokens 限制摘要长度。
func NewAdapterSummaryGenerator(adp adapter.ModelAdapter, model string, maxTokens int) *AdapterSummaryGenerator {
temp := 0.3
return &AdapterSummaryGenerator{
adapter: adp,
model: model,
maxTokens: maxTokens,
temperature: temp,
}
}
// GenerateSummary 调用 LLM 对历史消息生成摘要。
func (g *AdapterSummaryGenerator) GenerateSummary(ctx context.Context, messages []api.Message) (string, error) {
if len(messages) == 0 {
return "", nil
}
// 构建摘要请求的 prompt
prompt := buildSummaryPrompt(messages)
req := &adapter.ChatRequest{
RequestID: "summary-" + fmt.Sprintf("%d", len(messages)),
Model: g.model,
Messages: prompt,
MaxTokens: g.maxTokens,
Temperature: &g.temperature,
Stream: false,
}
resp, err := g.adapter.ChatCompletion(ctx, req)
if err != nil {
return "", fmt.Errorf("summary generation failed: %w", err)
}
return resp.Content, nil
}
// buildSummaryPrompt 构建用于生成摘要的消息列表。
func buildSummaryPrompt(messages []api.Message) []api.Message {
// 将历史消息拼接为文本
var sb strings.Builder
for _, m := range messages {
role := m.Role
content, _ := m.Content.(string)
sb.WriteString(fmt.Sprintf("%s: %s\n", role, content))
}
return []api.Message{
{
Role: "system",
Content: "你是一个对话摘要助手。请将以下对话历史压缩为简洁的摘要," +
"保留关键信息、用户意图和重要结论。摘要不超过200字。",
},
{
Role: "user",
Content: fmt.Sprintf("请总结以下对话历史:\n\n%s", sb.String()),
},
}
}
+18 -4
View File
@@ -7,8 +7,9 @@ import (
// Assembler assembles context messages for a chat request.
type Assembler struct {
estimator *TokenEstimator
cfg *config.ContextConfig
estimator *TokenEstimator
cfg *config.ContextConfig
summarizer *Summarizer
}
// NewAssembler creates a new context assembler.
@@ -19,6 +20,11 @@ func NewAssembler(cfg *config.ContextConfig) *Assembler {
}
}
// SetSummarizer 注入摘要器,启用后 trimSummaryAndRecent 策略将生成真实 LLM 摘要。
func (a *Assembler) SetSummarizer(s *Summarizer) {
a.summarizer = s
}
// AssembleResult contains the assembled messages and metadata.
type AssembleResult struct {
Messages []api.Message
@@ -143,13 +149,21 @@ func (a *Assembler) trimSummaryAndRecent(messages []api.Message, budget int) tri
recentTokens += msgTokens
}
// Add summary placeholder if we trimmed anything
// Add summary if we trimmed anything
result := make([]api.Message, 0, len(systemMsgs)+1+len(recentMsgs))
result = append(result, systemMsgs...)
if len(recentMsgs) < len(rest) {
// 被裁剪的老消息
trimmedMsgs := rest[:len(rest)-len(recentMsgs)]
var summaryText string
if a.summarizer != nil {
summaryText = a.summarizer.Summarize(trimmedMsgs)
} else {
summaryText = "[Earlier conversation history has been summarized and omitted.]"
}
result = append(result, api.Message{
Role: "system",
Content: "[Earlier conversation history has been summarized and omitted.]",
Content: summaryText,
})
}
result = append(result, recentMsgs...)
+48
View File
@@ -0,0 +1,48 @@
package context
import (
"context"
"fmt"
"strings"
"time"
"github.com/edgeai/gateway/pkg/api"
)
// SummaryGenerator 定义摘要生成接口,允许注入不同的实现。
type SummaryGenerator interface {
// GenerateSummary 对给定消息生成对话摘要,返回摘要文本。
GenerateSummary(ctx context.Context, messages []api.Message) (string, error)
}
// Summarizer 使用 LLM 对历史消息生成对话摘要。
type Summarizer struct {
generator SummaryGenerator
timeout time.Duration
}
// NewSummarizer 创建一个新的摘要器。
func NewSummarizer(generator SummaryGenerator, timeout time.Duration) *Summarizer {
return &Summarizer{
generator: generator,
timeout: timeout,
}
}
// Summarize 对被裁剪的老消息生成摘要。
// 如果摘要生成失败,回退到占位符文本。
func (s *Summarizer) Summarize(messages []api.Message) string {
if s == nil || s.generator == nil || len(messages) == 0 {
return "[Earlier conversation history has been summarized and omitted.]"
}
ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
defer cancel()
summary, err := s.generator.GenerateSummary(ctx, messages)
if err != nil || strings.TrimSpace(summary) == "" {
return "[Earlier conversation history has been summarized and omitted.]"
}
return fmt.Sprintf("[Earlier conversation summary: %s]", strings.TrimSpace(summary))
}
+85 -50
View File
@@ -4,6 +4,8 @@ import (
"encoding/json"
"fmt"
"net/http"
"strings"
"time"
"github.com/edgeai/gateway/internal/adapter"
"github.com/edgeai/gateway/pkg/api"
@@ -47,77 +49,110 @@ func (s *SSEWriter) WriteDone() {
s.flusher.Flush()
}
// WritePing 发送 SSE 注释行作为心跳,防止代理/负载均衡器因空闲超时断开连接。
// 注释行以冒号开头,客户端会忽略,不影响事件流。
func (s *SSEWriter) WritePing() {
fmt.Fprintf(s.w, ": keepalive\n\n")
s.flusher.Flush()
}
// StreamChatCompletion streams chunks from an adapter to the client in OpenAI SSE format.
func StreamChatCompletion(sse *SSEWriter, ch <-chan adapter.StreamChunk, requestID, taskID, model string) (int, int, error) {
// Returns input tokens, output tokens, full content string, and error.
// onFirstToken is called when the first content chunk is received (may be nil).
// 每 15 秒发送一次 keepalive ping,防止连接被代理断开。
func StreamChatCompletion(sse *SSEWriter, ch <-chan adapter.StreamChunk, requestID, taskID, model string, onFirstToken func()) (int, int, string, error) {
inputTokens := 0
outputTokens := 0
var contentBuilder strings.Builder
firstTokenSent := false
for chunk := range ch {
if chunk.Error != nil {
return inputTokens, outputTokens, chunk.Error
}
// 启动 keepalive 心跳定时器
keepalive := time.NewTicker(15 * time.Second)
defer keepalive.Stop()
if chunk.Done {
if chunk.InputTokens > 0 {
inputTokens = chunk.InputTokens
for {
select {
case chunk, ok := <-ch:
if !ok {
return inputTokens, outputTokens, contentBuilder.String(), nil
}
if chunk.OutputTokens > 0 {
outputTokens = chunk.OutputTokens
if chunk.Error != nil {
return inputTokens, outputTokens, contentBuilder.String(), chunk.Error
}
// Write final chunk with finish_reason
sseChunk := map[string]any{
"id": requestID,
"object": "chat.completion.chunk",
"model": model,
"choices": []map[string]any{
{
"index": 0,
"delta": map[string]any{},
"finish_reason": chunk.FinishReason,
if chunk.Done {
if chunk.InputTokens > 0 {
inputTokens = chunk.InputTokens
}
if chunk.OutputTokens > 0 {
outputTokens = chunk.OutputTokens
}
// Write final chunk with finish_reason
sseChunk := map[string]any{
"id": requestID,
"object": "chat.completion.chunk",
"model": model,
"choices": []map[string]any{
{
"index": 0,
"delta": map[string]any{},
"finish_reason": chunk.FinishReason,
},
},
},
}
if inputTokens > 0 || outputTokens > 0 {
sseChunk["usage"] = map[string]int{
"input_tokens": inputTokens,
"output_tokens": outputTokens,
"total_tokens": inputTokens + outputTokens,
}
}
sse.WriteChunk(sseChunk)
sse.WriteDone()
return inputTokens, outputTokens, contentBuilder.String(), nil
}
if inputTokens > 0 || outputTokens > 0 {
sseChunk["usage"] = map[string]int{
"input_tokens": inputTokens,
"output_tokens": outputTokens,
"total_tokens": inputTokens + outputTokens,
// 首 token 延迟回调
if !firstTokenSent && chunk.Delta != "" {
firstTokenSent = true
if onFirstToken != nil {
onFirstToken()
}
}
sse.WriteChunk(sseChunk)
sse.WriteDone()
return inputTokens, outputTokens, nil
}
// Write content delta
sseChunk := map[string]any{
"id": requestID,
"object": "chat.completion.chunk",
"model": model,
"choices": []map[string]any{
{
"index": 0,
"delta": map[string]any{
"content": chunk.Delta,
// Write content delta
sseChunk := map[string]any{
"id": requestID,
"object": "chat.completion.chunk",
"model": model,
"choices": []map[string]any{
{
"index": 0,
"delta": map[string]any{
"content": chunk.Delta,
},
"finish_reason": nil,
},
"finish_reason": nil,
},
},
}
sse.WriteChunk(sseChunk)
}
}
sse.WriteChunk(sseChunk)
contentBuilder.WriteString(chunk.Delta)
return inputTokens, outputTokens, nil
case <-keepalive.C:
// 发送心跳,保持连接活跃
sse.WritePing()
}
}
}
// BuildChatResponse creates a non-streaming ChatResponse from adapter result.
func BuildChatResponse(requestID, taskID, logicalModel string, resp *adapter.ChatResponse) api.ChatResponse {
return api.ChatResponse{
RequestID: requestID,
TaskID: taskID,
Status: "completed",
Model: logicalModel,
RequestID: requestID,
TaskID: taskID,
Status: "completed",
Model: logicalModel,
Choices: []api.Choice{
{
Index: 0,
+159 -12
View File
@@ -1,10 +1,12 @@
package middleware
import (
"compress/gzip"
"context"
"fmt"
"net/http"
"runtime/debug"
"strings"
"time"
"github.com/edgeai/gateway/internal/observability"
@@ -17,7 +19,39 @@ const (
RequestIDKey contextKey = "request_id"
)
// CORS middleware 添加跨域响应头,支持浏览器客户端直接调用 API。
// allowedOrigins 为允许的源列表,"*" 表示允许所有源。
func CORS(allowedOrigins []string) func(http.Handler) http.Handler {
allowed := map[string]bool{}
for _, o := range allowedOrigins {
allowed[o] = true
}
allowAll := allowed["*"]
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
origin := r.Header.Get("Origin")
if origin != "" {
if allowAll || allowed[origin] {
w.Header().Set("Access-Control-Allow-Origin", origin)
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE, OPTIONS")
w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type, X-Request-ID")
w.Header().Set("Access-Control-Expose-Headers", "X-Request-ID, X-Trace-Id, X-RateLimit-Limit, X-RateLimit-Remaining, Retry-After, X-Backpressure-Level, X-Timing-Queue-Ms, X-Timing-Inference-Ms, X-Timing-Total-Ms")
w.Header().Set("Access-Control-Max-Age", "3600")
}
}
// 处理预检请求
if r.Method == http.MethodOptions {
w.WriteHeader(http.StatusNoContent)
return
}
next.ServeHTTP(w, r)
})
}
}
// RequestID middleware generates a unique request ID and sets it in context and response header.
// 同时设置 X-Trace-Id 响应头,便于客户端和分布式追踪系统关联请求。
func RequestID(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requestID := r.Header.Get("X-Request-ID")
@@ -25,6 +59,7 @@ func RequestID(next http.Handler) http.Handler {
requestID = uuid.New().String()
}
w.Header().Set("X-Request-ID", requestID)
w.Header().Set("X-Trace-Id", requestID)
ctx := context.WithValue(r.Context(), RequestIDKey, requestID)
next.ServeHTTP(w, r.WithContext(ctx))
})
@@ -38,7 +73,7 @@ func BodyLimit(maxMB int) func(http.Handler) http.Handler {
if r.ContentLength > maxBytes {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusBadRequest)
fmt.Fprintf(w, `{"error":{"code":"INVALID_REQUEST","message":"request body exceeds %dMB limit"}}`, maxMB)
fmt.Fprintf(w, `{"error":{"code":"INVALID_REQUEST","message":"request body exceeds %dMB limit","request_id":"%s"}}`, maxMB, r.Header.Get("X-Request-ID"))
return
}
r.Body = http.MaxBytesReader(w, r.Body, maxBytes)
@@ -57,8 +92,9 @@ func Recovery(logger *observability.Logger) func(http.Handler) http.Handler {
observability.F().Event("panic").
RequestID(r.Header.Get("X-Request-ID")).
Reason(fmt.Sprintf("%v\n%s", rec, debug.Stack())))
http.Error(w, `{"error":{"code":"INTERNAL_ERROR","message":"internal server error"}}`,
http.StatusInternalServerError)
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusInternalServerError)
fmt.Fprintf(w, `{"error":{"code":"INTERNAL_ERROR","message":"internal server error","request_id":"%s"}}`, r.Header.Get("X-Request-ID"))
}
}()
next.ServeHTTP(w, r)
@@ -73,21 +109,26 @@ func Logging(logger *observability.Logger) func(http.Handler) http.Handler {
start := time.Now()
rw := &responseWriter{ResponseWriter: w, status: 200}
next.ServeHTTP(rw, r)
logger.Info("http request",
observability.F().
Event("http_request").
RequestID(r.Header.Get("X-Request-ID")).
Set("method", r.Method).
Set("path", r.URL.Path).
Set("status", rw.status).
Set("duration_ms", time.Since(start).Milliseconds()))
fields := observability.F().
Event("http_request").
RequestID(r.Header.Get("X-Request-ID")).
Set("method", r.Method).
Set("path", r.URL.Path).
Set("status", rw.status).
Set("duration_ms", time.Since(start).Milliseconds()).
Set("response_bytes", rw.bytesWritten)
if origin := r.Header.Get("Origin"); origin != "" {
fields.Set("origin", origin)
}
logger.Info("http request", fields)
})
}
}
type responseWriter struct {
http.ResponseWriter
status int
status int
bytesWritten int
}
func (rw *responseWriter) WriteHeader(code int) {
@@ -95,6 +136,12 @@ func (rw *responseWriter) WriteHeader(code int) {
rw.ResponseWriter.WriteHeader(code)
}
func (rw *responseWriter) Write(b []byte) (int, error) {
n, err := rw.ResponseWriter.Write(b)
rw.bytesWritten += n
return n, err
}
func (rw *responseWriter) Flush() {
if f, ok := rw.ResponseWriter.(http.Flusher); ok {
f.Flush()
@@ -108,3 +155,103 @@ func GetRequestID(ctx context.Context) string {
}
return ""
}
// RateLimitFn 是限流检查函数的类型。
// 返回 allowed, limit, remaining。
type RateLimitFn func(appID string) (allowed bool, limit int, remaining int)
// RateLimit middleware 对所有认证请求执行 per-app 限流,并设置 X-RateLimit-* 响应头。
// skipPaths 中的路径不限流(如 /health、/ready)。
func RateLimit(checkFn RateLimitFn, getIdentity func(*http.Request) string, skipPaths map[string]bool) func(http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if skipPaths[r.URL.Path] {
next.ServeHTTP(w, r)
return
}
appID := getIdentity(r)
if appID == "" {
next.ServeHTTP(w, r)
return
}
allowed, limit, remaining := checkFn(appID)
w.Header().Set("X-RateLimit-Limit", fmt.Sprintf("%d", limit))
w.Header().Set("X-RateLimit-Remaining", fmt.Sprintf("%d", remaining))
if !allowed {
w.Header().Set("Retry-After", "1")
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusTooManyRequests)
fmt.Fprintf(w, `{"error":{"code":"RATE_LIMITED","message":"rate limit exceeded for this application","request_id":"%s"}}`, r.Header.Get("X-Request-ID"))
return
}
next.ServeHTTP(w, r)
})
}
}
// gzipResponseWriter 包装 ResponseWriter,对响应体进行 gzip 压缩。
// 仅当客户端发送 Accept-Encoding: gzip 且响应 Content-Type 为 JSON 时启用。
type gzipResponseWriter struct {
http.ResponseWriter
gz *gzip.Writer
contentType string
gzStarted bool
}
func (g *gzipResponseWriter) WriteHeader(code int) {
ct := g.ResponseWriter.Header().Get("Content-Type")
g.contentType = ct
// SSE 流式响应不压缩
if strings.Contains(ct, "text/event-stream") {
g.ResponseWriter.WriteHeader(code)
return
}
// 对 JSON 等可压缩内容启用 gzip
if shouldCompress(ct) {
g.ResponseWriter.Header().Set("Content-Encoding", "gzip")
g.ResponseWriter.Header().Del("Content-Length")
g.gzStarted = true
}
g.ResponseWriter.WriteHeader(code)
}
func (g *gzipResponseWriter) Write(b []byte) (int, error) {
if g.gzStarted {
return g.gz.Write(b)
}
return g.ResponseWriter.Write(b)
}
func (g *gzipResponseWriter) Flush() {
if g.gzStarted {
g.gz.Flush()
}
if f, ok := g.ResponseWriter.(http.Flusher); ok {
f.Flush()
}
}
// shouldCompress 判断 Content-Type 是否值得压缩。
func shouldCompress(ct string) bool {
return strings.Contains(ct, "json") ||
strings.Contains(ct, "text") ||
strings.Contains(ct, "javascript") ||
strings.Contains(ct, "xml")
}
// Gzip 中间件对响应体进行 gzip 压缩,减少网络传输量。
// 仅当客户端支持 gzip 且响应内容类型可压缩时启用。
func Gzip(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
next.ServeHTTP(w, r)
return
}
gz := gzip.NewWriter(w)
defer gz.Close()
gw := &gzipResponseWriter{ResponseWriter: w, gz: gz}
next.ServeHTTP(gw, r)
})
}
+100
View File
@@ -0,0 +1,100 @@
package observability
import (
"encoding/json"
"fmt"
"net/http"
"os"
"path/filepath"
"sync"
"time"
)
// AuditEntry 审计日志条目,记录关键管理操作的详细信息。
type AuditEntry struct {
Timestamp string `json:"timestamp"`
Actor string `json:"actor"` // 操作者 app_id
Action string `json:"action"` // 操作类型
Resource string `json:"resource"` // 操作资源
ResourceID string `json:"resource_id"` // 资源 ID
Method string `json:"method"` // HTTP 方法
Path string `json:"path"` // 请求路径
IP string `json:"ip"` // 客户端 IP
Status int `json:"status"` // HTTP 响应码
Details map[string]interface{} `json:"details,omitempty"`
}
// AuditLogger 审计日志记录器,将关键操作写入独立日志文件。
type AuditLogger struct {
mu sync.Mutex
file *os.File
logger *Logger
}
// NewAuditLogger 创建审计日志记录器,日志写入指定目录下的 audit.log 文件。
func NewAuditLogger(dir string, logger *Logger) *AuditLogger {
auditPath := filepath.Join(dir, "audit.log")
file, err := os.OpenFile(auditPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
if err != nil {
logger.Warn("failed to open audit log file, audit logs will go to stderr",
F().Event("audit_log_init_failed").Reason(err.Error()))
return &AuditLogger{file: nil, logger: logger}
}
return &AuditLogger{file: file, logger: logger}
}
// Record 记录一条审计日志。
func (a *AuditLogger) Record(entry AuditEntry) {
if entry.Timestamp == "" {
entry.Timestamp = time.Now().UTC().Format(time.RFC3339Nano)
}
data, err := json.Marshal(entry)
if err != nil {
a.logger.Error("audit log marshal error", F().Event("audit_marshal_error").Reason(err.Error()))
return
}
a.mu.Lock()
defer a.mu.Unlock()
if a.file != nil {
fmt.Fprintln(a.file, string(data))
} else {
fmt.Fprintln(os.Stderr, string(data))
}
}
// RecordFromRequest 从 HTTP 请求中提取信息并记录审计日志。
func (a *AuditLogger) RecordFromRequest(r *http.Request, actor, action, resource, resourceID string, status int, details map[string]interface{}) {
entry := AuditEntry{
Actor: actor,
Action: action,
Resource: resource,
ResourceID: resourceID,
Method: r.Method,
Path: r.URL.Path,
IP: extractIP(r),
Status: status,
Details: details,
}
a.Record(entry)
}
// Close 关闭审计日志文件。
func (a *AuditLogger) Close() {
if a.file != nil {
a.file.Close()
}
}
// extractIP 从请求中提取客户端 IP 地址。
func extractIP(r *http.Request) string {
if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
return xff
}
if xri := r.Header.Get("X-Real-IP"); xri != "" {
return xri
}
return r.RemoteAddr
}
+49
View File
@@ -27,6 +27,7 @@ type Metrics struct {
runningTasks int64
activeSessions int64
backpressureLevel int64
breakerState int64 // 0=closed, 1=open, 2=half_open
// Histograms (simplified as buckets)
gatewayLatencyBuckets map[string]int64
@@ -110,6 +111,11 @@ func (m *Metrics) SetBackpressureLevel(level int) {
atomic.StoreInt64(&m.backpressureLevel, int64(level))
}
// SetBreakerState 设置熔断器状态 gauge0=closed, 1=open, 2=half_open)。
func (m *Metrics) SetBreakerState(state int) {
atomic.StoreInt64(&m.breakerState, int64(state))
}
// ObserveGatewayLatency records gateway latency in a histogram bucket.
func (m *Metrics) ObserveGatewayLatency(ms int64) {
bucket := latencyBucket(ms)
@@ -142,34 +148,77 @@ func (m *Metrics) Handler() http.HandlerFunc {
w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
// Counters
fmt.Fprintf(w, "# HELP edgeai_requests_total Total requests by status\n")
fmt.Fprintf(w, "# TYPE edgeai_requests_total counter\n")
m.mu.RLock()
for key, val := range m.requestsTotal {
fmt.Fprintf(w, "edgeai_requests_total{%s} %d\n", key, val)
}
fmt.Fprintf(w, "# HELP edgeai_tasks_total Total tasks by final state\n")
fmt.Fprintf(w, "# TYPE edgeai_tasks_total counter\n")
for key, val := range m.tasksTotal {
fmt.Fprintf(w, "edgeai_tasks_total{%s} %d\n", key, val)
}
m.mu.RUnlock()
fmt.Fprintf(w, "# HELP edgeai_tokens_input_total Total input tokens processed\n")
fmt.Fprintf(w, "# TYPE edgeai_tokens_input_total counter\n")
fmt.Fprintf(w, "edgeai_tokens_input_total %d\n", atomic.LoadInt64(&m.tokensInputTotal))
fmt.Fprintf(w, "# HELP edgeai_tokens_output_total Total output tokens generated\n")
fmt.Fprintf(w, "# TYPE edgeai_tokens_output_total counter\n")
fmt.Fprintf(w, "edgeai_tokens_output_total %d\n", atomic.LoadInt64(&m.tokensOutputTotal))
fmt.Fprintf(w, "# HELP edgeai_cancellations_total Total cancelled requests\n")
fmt.Fprintf(w, "# TYPE edgeai_cancellations_total counter\n")
fmt.Fprintf(w, "edgeai_cancellations_total %d\n", atomic.LoadInt64(&m.cancellationsTotal))
fmt.Fprintf(w, "# HELP edgeai_queue_timeouts_total Total queue timeouts\n")
fmt.Fprintf(w, "# TYPE edgeai_queue_timeouts_total counter\n")
fmt.Fprintf(w, "edgeai_queue_timeouts_total %d\n", atomic.LoadInt64(&m.queueTimeoutsTotal))
fmt.Fprintf(w, "# HELP edgeai_first_token_timeouts_total Total first token timeouts\n")
fmt.Fprintf(w, "# TYPE edgeai_first_token_timeouts_total counter\n")
fmt.Fprintf(w, "edgeai_first_token_timeouts_total %d\n", atomic.LoadInt64(&m.firstTokenTimeoutsTotal))
fmt.Fprintf(w, "# HELP edgeai_inference_timeouts_total Total inference timeouts\n")
fmt.Fprintf(w, "# TYPE edgeai_inference_timeouts_total counter\n")
fmt.Fprintf(w, "edgeai_inference_timeouts_total %d\n", atomic.LoadInt64(&m.inferenceTimeoutsTotal))
fmt.Fprintf(w, "# HELP edgeai_degraded_requests_total Total requests served with degradation\n")
fmt.Fprintf(w, "# TYPE edgeai_degraded_requests_total counter\n")
fmt.Fprintf(w, "edgeai_degraded_requests_total %d\n", atomic.LoadInt64(&m.degradedRequestsTotal))
// Gauges
fmt.Fprintf(w, "# HELP edgeai_queue_length Current queue length\n")
fmt.Fprintf(w, "# TYPE edgeai_queue_length gauge\n")
fmt.Fprintf(w, "edgeai_queue_length %d\n", atomic.LoadInt64(&m.queueLength))
fmt.Fprintf(w, "# HELP edgeai_running_tasks Current running tasks\n")
fmt.Fprintf(w, "# TYPE edgeai_running_tasks gauge\n")
fmt.Fprintf(w, "edgeai_running_tasks %d\n", atomic.LoadInt64(&m.runningTasks))
fmt.Fprintf(w, "# HELP edgeai_active_sessions Current active sessions\n")
fmt.Fprintf(w, "# TYPE edgeai_active_sessions gauge\n")
fmt.Fprintf(w, "edgeai_active_sessions %d\n", atomic.LoadInt64(&m.activeSessions))
fmt.Fprintf(w, "# HELP edgeai_backpressure_level Current backpressure level (0-3)\n")
fmt.Fprintf(w, "# TYPE edgeai_backpressure_level gauge\n")
fmt.Fprintf(w, "edgeai_backpressure_level %d\n", atomic.LoadInt64(&m.backpressureLevel))
fmt.Fprintf(w, "# HELP edgeai_circuit_breaker_state Circuit breaker state (0=closed, 1=open, 2=half_open)\n")
fmt.Fprintf(w, "# TYPE edgeai_circuit_breaker_state gauge\n")
fmt.Fprintf(w, "edgeai_circuit_breaker_state %d\n", atomic.LoadInt64(&m.breakerState))
// Histograms
fmt.Fprintf(w, "# HELP edgeai_gateway_latency_bucket Gateway latency distribution in milliseconds\n")
fmt.Fprintf(w, "# TYPE edgeai_gateway_latency_bucket histogram\n")
m.mu.RLock()
for bucket, count := range m.gatewayLatencyBuckets {
fmt.Fprintf(w, "edgeai_gateway_latency_bucket{%s} %d\n", bucket, count)
}
fmt.Fprintf(w, "# HELP edgeai_first_token_latency_bucket First token latency distribution in milliseconds\n")
fmt.Fprintf(w, "# TYPE edgeai_first_token_latency_bucket histogram\n")
for bucket, count := range m.firstTokenLatencyBuckets {
fmt.Fprintf(w, "edgeai_first_token_latency_bucket{%s} %d\n", bucket, count)
}
+185
View File
@@ -0,0 +1,185 @@
package ratelimit
import (
"sync"
"time"
)
// TokenBucket 是一个简单的令牌桶限流器。
type TokenBucket struct {
capacity float64 // 桶容量
refillRate float64 // 每秒补充令牌数
tokens float64 // 当前令牌数
lastRefill time.Time // 上次补充时间
lastAccess time.Time // 最后访问时间,用于空闲清理
mu sync.Mutex
}
// NewTokenBucket 创建一个令牌桶。
// capacity 为桶容量(突发上限),refillPerSecond 为每秒补充速率。
func NewTokenBucket(capacity float64, refillPerSecond float64) *TokenBucket {
now := time.Now()
return &TokenBucket{
capacity: capacity,
refillRate: refillPerSecond,
tokens: capacity, // 初始满桶
lastRefill: now,
lastAccess: now,
}
}
// Allow 尝试消耗 1 个令牌,返回是否允许。
func (tb *TokenBucket) Allow() bool {
return AllowN(tb, 1)
}
// AllowN 尝试消耗 n 个令牌,返回是否允许。
func AllowN(tb *TokenBucket, n float64) bool {
tb.mu.Lock()
defer tb.mu.Unlock()
now := time.Now()
elapsed := now.Sub(tb.lastRefill).Seconds()
tb.tokens += elapsed * tb.refillRate
if tb.tokens > tb.capacity {
tb.tokens = tb.capacity
}
tb.lastRefill = now
tb.lastAccess = now
if tb.tokens >= n {
tb.tokens -= n
return true
}
return false
}
// Tokens 返回当前可用令牌数(近似值)。
func (tb *TokenBucket) Tokens() float64 {
tb.mu.Lock()
defer tb.mu.Unlock()
return tb.tokens
}
// Limiter 管理 per-app 令牌桶限流器。
type Limiter struct {
mu sync.RWMutex
buckets map[string]*TokenBucket // app_id -> bucket
capacity float64
refill float64
}
// NewLimiter 创建一个 per-app 限流管理器。
// capacity 为每个 app 的桶容量,refillPerSecond 为每秒补充速率。
// idleTimeout 为空闲 app bucket 的过期时间,超过此时间未被访问的 bucket 将被清理。
func NewLimiter(capacity float64, refillPerSecond float64) *Limiter {
l := &Limiter{
buckets: make(map[string]*TokenBucket),
capacity: capacity,
refill: refillPerSecond,
}
go l.cleanupIdle()
return l
}
// Allow 检查指定 app 是否被限流。
func (l *Limiter) Allow(appID string) bool {
l.mu.RLock()
bucket, ok := l.buckets[appID]
l.mu.RUnlock()
if !ok {
bucket = NewTokenBucket(l.capacity, l.refill)
l.mu.Lock()
// 双检查,防止竞态
if existing, ok := l.buckets[appID]; ok {
bucket = existing
} else {
l.buckets[appID] = bucket
}
l.mu.Unlock()
}
return bucket.Allow()
}
// RateLimitInfo 包含限流信息,用于响应头。
type RateLimitInfo struct {
Allowed bool
Limit int
Remaining int
}
// AllowWithInfo 检查限流并返回详情(用于 X-RateLimit-* 响应头)。
func (l *Limiter) AllowWithInfo(appID string) RateLimitInfo {
l.mu.RLock()
bucket, ok := l.buckets[appID]
l.mu.RUnlock()
if !ok {
bucket = NewTokenBucket(l.capacity, l.refill)
l.mu.Lock()
if existing, ok := l.buckets[appID]; ok {
bucket = existing
} else {
l.buckets[appID] = bucket
}
l.mu.Unlock()
}
allowed := bucket.Allow()
tokens := int(bucket.Tokens())
return RateLimitInfo{
Allowed: allowed,
Limit: int(l.capacity),
Remaining: tokens,
}
}
// SetLimit 为指定 app 设置自定义限流参数。
func (l *Limiter) SetLimit(appID string, capacity, refillPerSecond float64) {
l.mu.Lock()
defer l.mu.Unlock()
l.buckets[appID] = NewTokenBucket(capacity, refillPerSecond)
}
// GetLimit 返回指定 app 的限流配置(burst 和 rate_per_minute)。
func (l *Limiter) GetLimit(appID string) (burst int, ratePerMinute int) {
l.mu.RLock()
bucket, ok := l.buckets[appID]
l.mu.RUnlock()
if !ok {
return int(l.capacity), int(l.refill * 60)
}
return int(bucket.capacity), int(bucket.refillRate * 60)
}
// Remove 移除指定 app 的限流器。
func (l *Limiter) Remove(appID string) {
l.mu.Lock()
defer l.mu.Unlock()
delete(l.buckets, appID)
}
// cleanupIdle 定期清理空闲超过 30 分钟的 app bucket,防止内存泄漏。
func (l *Limiter) cleanupIdle() {
ticker := time.NewTicker(5 * time.Minute)
defer ticker.Stop()
idleThreshold := 30 * time.Minute
for range ticker.C {
cutoff := time.Now().Add(-idleThreshold)
l.mu.Lock()
for appID, bucket := range l.buckets {
bucket.mu.Lock()
idle := bucket.lastAccess.Before(cutoff)
bucket.mu.Unlock()
if idle {
delete(l.buckets, appID)
}
}
l.mu.Unlock()
}
}
+56 -9
View File
@@ -7,27 +7,27 @@ import (
// CircuitBreaker implements a sliding-window circuit breaker for adapter health.
type CircuitBreaker struct {
mu sync.Mutex
mu sync.Mutex
errorRateThreshold float64
minRequests int
windowSeconds int
openDuration time.Duration
halfOpenMax int
minRequests int
windowSeconds int
openDuration time.Duration
halfOpenMax int
// sliding window state
requests []time.Time
requests []time.Time
errors []time.Time
// breaker state
state breakerState
openedAt time.Time
state breakerState
openedAt time.Time
halfOpenCount int
}
type breakerState int
const (
breakerClosed breakerState = iota
breakerClosed breakerState = iota
breakerOpen
breakerHalfOpen
)
@@ -126,6 +126,53 @@ func (cb *CircuitBreaker) State() string {
return "unknown"
}
// BreakerStats 熔断器统计信息。
type BreakerStats struct {
State string `json:"state"`
TotalRequests int `json:"total_requests"`
TotalErrors int `json:"total_errors"`
ErrorRate float64 `json:"error_rate"`
WindowSeconds int `json:"window_seconds"`
OpenDuration string `json:"open_duration"`
Threshold float64 `json:"error_rate_threshold"`
MinRequests int `json:"min_requests"`
}
// Stats 返回熔断器的详细统计信息。
func (cb *CircuitBreaker) Stats() BreakerStats {
cb.mu.Lock()
defer cb.mu.Unlock()
now := time.Now()
cb.prune(now)
total := len(cb.requests)
errors := len(cb.errors)
var errorRate float64
if total > 0 {
errorRate = float64(errors) / float64(total)
}
stateName := "closed"
switch cb.state {
case breakerOpen:
stateName = "open"
case breakerHalfOpen:
stateName = "half_open"
}
return BreakerStats{
State: stateName,
TotalRequests: total,
TotalErrors: errors,
ErrorRate: errorRate,
WindowSeconds: cb.windowSeconds,
OpenDuration: cb.openDuration.String(),
Threshold: cb.errorRateThreshold,
MinRequests: cb.minRequests,
}
}
// prune removes entries outside the sliding window.
func (cb *CircuitBreaker) prune(now time.Time) {
cutoff := now.Add(-time.Duration(cb.windowSeconds) * time.Second)
+83 -17
View File
@@ -14,29 +14,33 @@ import (
// Scheduler manages task queuing and execution with priority-based scheduling.
type Scheduler struct {
mu sync.Mutex
queue *priorityQueue
running map[string]*task.Task
maxRunning int
maxQueued int
notifyCh chan struct{}
logger *observability.Logger
ctx context.Context
cancel context.CancelFunc
mu sync.Mutex
queue *priorityQueue
running map[string]*task.Task
maxRunning int
maxQueued int
notifyCh chan struct{}
logger *observability.Logger
ctx context.Context
cancel context.CancelFunc
agingSeconds int // 优先级老化阈值(秒),0 表示禁用
agingMinPrio int // 老化生效的最低优先级(仅 P2-P4 老化,P0/P1 不老化)
}
// NewScheduler creates a new scheduler.
func NewScheduler(cfg *config.SchedulerConfig, logger *observability.Logger) *Scheduler {
ctx, cancel := context.WithCancel(context.Background())
s := &Scheduler{
queue: &priorityQueue{},
running: make(map[string]*task.Task),
maxRunning: cfg.MaxRunningTasks,
maxQueued: cfg.MaxQueuedTasks,
notifyCh: make(chan struct{}, 1),
logger: logger,
ctx: ctx,
cancel: cancel,
queue: &priorityQueue{},
running: make(map[string]*task.Task),
maxRunning: cfg.MaxRunningTasks,
maxQueued: cfg.MaxQueuedTasks,
notifyCh: make(chan struct{}, 1),
logger: logger,
ctx: ctx,
cancel: cancel,
agingSeconds: cfg.PriorityAgingSeconds,
agingMinPrio: 2, // P2 及以上优先级才会老化
}
heap.Init(s.queue)
return s
@@ -71,6 +75,10 @@ func (s *Scheduler) Submit(t *task.Task) error {
func (s *Scheduler) GetNext(ctx context.Context) (*task.Task, error) {
for {
s.mu.Lock()
// 优先级老化:提升等待过久的低优先级任务
if s.agingSeconds > 0 {
s.applyPriorityAging()
}
if s.queue.Len() > 0 && len(s.running) < s.maxRunning {
t := heap.Pop(s.queue).(*task.Task)
s.running[t.ID] = t
@@ -88,6 +96,35 @@ func (s *Scheduler) GetNext(ctx context.Context) (*task.Task, error) {
}
}
// applyPriorityAging 对队列中等待超过 agingSeconds 的低优先级任务提升一级优先级。
// 必须在持有 s.mu 锁的情况下调用。
func (s *Scheduler) applyPriorityAging() {
if s.agingSeconds <= 0 || s.queue.Len() == 0 {
return
}
now := time.Now()
aged := 0
for i := 0; i < s.queue.Len(); i++ {
t := (*s.queue)[i]
if int(t.Priority) < s.agingMinPrio {
continue // P0/P1 不老化
}
waitSec := int(now.Sub(t.CreatedAt).Seconds())
if waitSec >= s.agingSeconds {
t.Priority-- // 提升一级(数值越小优先级越高)
if t.Priority < 0 {
t.Priority = 0
}
aged++
}
}
if aged > 0 {
heap.Init(s.queue) // 重新堆化
s.logger.Info("priority aging applied",
observability.F().Event("priority_aging").Set("aged_count", aged))
}
}
// Complete marks a task as completed and removes it from running.
func (s *Scheduler) Complete(taskID string) {
s.mu.Lock()
@@ -130,6 +167,35 @@ func (s *Scheduler) GetTask(taskID string) (*task.Task, bool) {
return nil, false
}
// ListTasks 返回所有运行中和排队中的任务(支持分页)。
// status 过滤:running、queued、空字符串表示全部。
func (s *Scheduler) ListTasks(status string, limit, offset int) ([]*task.Task, int) {
s.mu.Lock()
defer s.mu.Unlock()
var all []*task.Task
if status == "" || status == "running" {
for _, t := range s.running {
all = append(all, t)
}
}
if status == "" || status == "queued" {
for i := 0; i < s.queue.Len(); i++ {
all = append(all, (*s.queue)[i])
}
}
total := len(all)
if offset >= total {
return []*task.Task{}, total
}
end := offset + limit
if end > total {
end = total
}
return all[offset:end], total
}
// Stop shuts down the scheduler.
func (s *Scheduler) Stop() {
s.cancel()
+333
View File
@@ -0,0 +1,333 @@
package server
import (
"encoding/json"
"fmt"
"net/http"
"strconv"
"strings"
"time"
"github.com/edgeai/gateway/internal/adapter"
"github.com/edgeai/gateway/internal/auth"
"github.com/edgeai/gateway/internal/config"
"github.com/edgeai/gateway/internal/handler"
"github.com/edgeai/gateway/internal/middleware"
"github.com/edgeai/gateway/internal/observability"
)
// handleAdminKeys 处理 /v1/admin/keysGET 列出所有 KeyPOST 创建新 Key)。
func (s *Server) handleAdminKeys(w http.ResponseWriter, r *http.Request) {
requestID := middleware.GetRequestID(r.Context())
switch r.Method {
case http.MethodGet:
keys, err := s.auth.ListKeys()
if err != nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInternalError, err.Error(), requestID))
return
}
handler.WriteJSON(w, http.StatusOK, map[string]any{"keys": keys})
case http.MethodPost:
var req struct {
AppID string `json:"app_id"`
TenantID string `json:"tenant_id"`
Name string `json:"name"`
AllowedModels []string `json:"allowed_models"`
AllowedPriorities []int `json:"allowed_priorities"`
IsAdmin bool `json:"is_admin"`
ExpiresAt string `json:"expires_at"` // RFC3339 格式,空表示永不过期
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "invalid JSON body", requestID))
return
}
if req.AppID == "" {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "app_id is required", requestID))
return
}
if req.Name == "" {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "name is required", requestID))
return
}
apiKey := auth.GenerateAPIKey()
identity := &auth.AppIdentity{
AppID: req.AppID,
TenantID: req.TenantID,
Name: req.Name,
AllowedModels: req.AllowedModels,
AllowedPriorities: req.AllowedPriorities,
IsAdmin: req.IsAdmin,
}
if req.ExpiresAt != "" {
if t, err := time.Parse(time.RFC3339, req.ExpiresAt); err == nil {
identity.ExpiresAt = &t
}
}
if err := s.auth.AddKey(apiKey, identity); err != nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInternalError, err.Error(), requestID))
return
}
// 审计日志
actor := ""
if id := auth.GetAppIdentityFromRequest(r); id != nil {
actor = id.AppID
}
s.audit.RecordFromRequest(r, actor, "create_key", "api_key", req.AppID, http.StatusCreated, map[string]interface{}{
"app_id": req.AppID,
"name": req.Name,
"is_admin": req.IsAdmin,
"allowed_models": req.AllowedModels,
})
handler.WriteJSON(w, http.StatusCreated, map[string]any{
"api_key": apiKey,
"app_id": req.AppID,
"name": req.Name,
"is_admin": req.IsAdmin,
"allowed_models": req.AllowedModels,
"expires_at": req.ExpiresAt,
"message": "请妥善保存此 API Key,之后将无法再次查看",
})
default:
handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
}
}
// handleAdminKeyByID 处理 /v1/admin/keys/{id}DELETE 删除,PATCH 禁用,PUT 轮换)。
func (s *Server) handleAdminKeyByID(w http.ResponseWriter, r *http.Request) {
requestID := middleware.GetRequestID(r.Context())
idStr := strings.TrimPrefix(r.URL.Path, "/v1/admin/keys/")
id, err := strconv.ParseInt(idStr, 10, 64)
if err != nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "invalid key id", requestID))
return
}
switch r.Method {
case http.MethodDelete:
if err := s.auth.DeleteKey(id); err != nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, err.Error(), requestID))
return
}
// 审计日志
actor := ""
if ident := auth.GetAppIdentityFromRequest(r); ident != nil {
actor = ident.AppID
}
s.audit.RecordFromRequest(r, actor, "delete_key", "api_key", strconv.FormatInt(id, 10), http.StatusOK, nil)
handler.WriteJSON(w, http.StatusOK, map[string]string{"status": "deleted"})
case http.MethodPatch:
if err := s.auth.DisableKey(id); err != nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, err.Error(), requestID))
return
}
handler.WriteJSON(w, http.StatusOK, map[string]string{"status": "disabled"})
case http.MethodPut:
// Key 轮换:生成新 Key,禁用旧 Key
newKey, err := s.auth.RotateKey(id)
if err != nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, err.Error(), requestID))
return
}
// 审计日志
actor := ""
if ident := auth.GetAppIdentityFromRequest(r); ident != nil {
actor = ident.AppID
}
s.audit.RecordFromRequest(r, actor, "rotate_key", "api_key", strconv.FormatInt(id, 10), http.StatusOK, nil)
handler.WriteJSON(w, http.StatusOK, map[string]any{
"status": "rotated",
"api_key": newKey,
"message": "旧 Key 已禁用,请妥善保存新 Key",
})
default:
handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
}
}
// handleAdminUsage 处理 /v1/admin/usageGET 返回所有 app 的用量统计)。
func (s *Server) handleAdminUsage(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
return
}
usage := s.usageTracker.GetAll()
handler.WriteJSON(w, http.StatusOK, map[string]any{"usage": usage})
}
// handleAdminUsageByApp 处理 /v1/admin/usage/{app_id}GET 返回指定 app 的用量统计,支持 ?history=hours 查询历史)。
func (s *Server) handleAdminUsageByApp(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
return
}
appID := strings.TrimPrefix(r.URL.Path, "/v1/admin/usage/")
if appID == "" {
handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "app_id is required"))
return
}
// 支持 ?history=24 查询历史用量
if historyHours := r.URL.Query().Get("history"); historyHours != "" {
hours := 24
fmt.Sscanf(historyHours, "%d", &hours)
history, err := s.usageTracker.GetHistory(appID, hours)
if err != nil {
handler.WriteError(w, handler.NewGatewayError(handler.ErrInternalError, err.Error()))
return
}
handler.WriteJSON(w, http.StatusOK, map[string]any{
"app_id": appID,
"hours": hours,
"history": history,
})
return
}
usage := s.usageTracker.Get(appID)
handler.WriteJSON(w, http.StatusOK, usage)
}
// handleAdminRateLimit 处理 /v1/admin/ratelimit/{app_id}GET 查看限流,PUT 设置自定义限流)。
func (s *Server) handleAdminRateLimit(w http.ResponseWriter, r *http.Request) {
requestID := middleware.GetRequestID(r.Context())
appID := strings.TrimPrefix(r.URL.Path, "/v1/admin/ratelimit/")
if appID == "" {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "app_id is required", requestID))
return
}
switch r.Method {
case http.MethodGet:
burst, ratePerMin := s.rateLimiter.GetLimit(appID)
handler.WriteJSON(w, http.StatusOK, map[string]any{
"app_id": appID,
"burst": burst,
"rate_per_minute": ratePerMin,
})
case http.MethodPut:
var req struct {
Burst int `json:"burst"`
RatePerMinute int `json:"rate_per_minute"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "invalid JSON body", requestID))
return
}
if req.Burst <= 0 || req.RatePerMinute <= 0 {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "burst and rate_per_minute must be positive", requestID))
return
}
refillPerSec := float64(req.RatePerMinute) / 60.0
s.rateLimiter.SetLimit(appID, float64(req.Burst), refillPerSec)
handler.WriteJSON(w, http.StatusOK, map[string]any{
"app_id": appID,
"burst": req.Burst,
"rate_per_minute": req.RatePerMinute,
"status": "updated",
})
default:
handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
}
}
// handleAdminCircuitBreaker 处理 /v1/admin/circuit-breakerGET 返回熔断器状态)。
func (s *Server) handleAdminCircuitBreaker(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
return
}
stats := s.breaker.Stats()
handler.WriteJSON(w, http.StatusOK, stats)
}
// handleAdminScheduler 处理 /v1/admin/schedulerGET 返回调度器状态)。
func (s *Server) handleAdminScheduler(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
return
}
stats := map[string]any{
"running_tasks": s.scheduler.RunningCount(),
"queued_tasks": s.scheduler.QueueLength(),
"max_running": s.cfg.Scheduler.MaxRunningTasks,
"max_queued": s.cfg.Scheduler.MaxQueuedTasks,
"priority_aging_seconds": s.cfg.Scheduler.PriorityAgingSeconds,
"fairness": s.cfg.Scheduler.Fairness,
}
handler.WriteJSON(w, http.StatusOK, stats)
}
// handleAdminConfigReload 处理 /v1/admin/config/reloadPOST 触发配置热重载)。
func (s *Server) handleAdminConfigReload(w http.ResponseWriter, r *http.Request) {
requestID := middleware.GetRequestID(r.Context())
if r.Method != http.MethodPost {
handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
return
}
// 重新加载配置文件
cfgPath := config.ConfigPath()
newCfg, err := config.Load(cfgPath)
if err != nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInternalError, fmt.Sprintf("reload config failed: %v", err), requestID))
return
}
// 热更新模型映射
s.modelMap.Update(newCfg)
// 注册新增 adapter(已有 adapter 不重复注册)
for _, mc := range newCfg.Models {
provider := mc.Provider
endpoint := mc.Endpoint
switch provider {
case "ollama":
s.registry.RegisterIfAbsent(provider, func() adapter.ModelAdapter {
return adapter.NewOllamaAdapter(endpoint)
})
case "vllm":
s.registry.RegisterIfAbsent(provider, func() adapter.ModelAdapter {
return adapter.NewVLLMAdapter(endpoint)
})
}
}
// 更新 Server 持有的配置引用
s.cfg = newCfg
s.logger.Info("config hot-reloaded",
observability.F().
Event("config_reload").
Set("config_path", cfgPath))
// 审计日志
actor := ""
if ident := auth.GetAppIdentityFromRequest(r); ident != nil {
actor = ident.AppID
}
s.audit.RecordFromRequest(r, actor, "config_reload", "config", cfgPath, http.StatusOK, map[string]interface{}{
"models": len(newCfg.Models),
})
handler.WriteJSON(w, http.StatusOK, map[string]any{
"status": "reloaded",
"models": len(newCfg.Models),
"message": "配置已热重载,模型映射已更新",
})
}
+376 -25
View File
@@ -5,12 +5,14 @@ import (
"encoding/json"
"fmt"
"net/http"
"strconv"
"strings"
"time"
"github.com/edgeai/gateway/internal/adapter"
"github.com/edgeai/gateway/internal/auth"
"github.com/edgeai/gateway/internal/config"
"github.com/edgeai/gateway/internal/connector"
ctxasm "github.com/edgeai/gateway/internal/context"
"github.com/edgeai/gateway/internal/handler"
"github.com/edgeai/gateway/internal/middleware"
@@ -47,12 +49,39 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
return
}
// 参数范围校验
if req.Temperature != nil && (*req.Temperature < 0 || *req.Temperature > 2) {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "temperature must be between 0 and 2", requestID))
return
}
if req.TopP != nil && (*req.TopP < 0 || *req.TopP > 1) {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "top_p must be between 0 and 1", requestID))
return
}
if req.MaxOutputTokens < 0 {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "max_output_tokens must be non-negative", requestID))
return
}
if len(req.Messages) > s.cfg.Context.MaxSessionMessages {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrContextTooLarge, fmt.Sprintf("messages count exceeds max %d", s.cfg.Context.MaxSessionMessages), requestID))
return
}
// 幂等性检查(仅非流式请求支持幂等)
if !req.Stream && req.IdempotencyKey != "" && identity != nil {
if s.checkIdempotency(w, req.IdempotencyKey, identity.AppID) {
return
}
}
// Check model permission
if identity != nil && !auth.CheckModelPermission(identity, req.Model) {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrPermissionDenied, "model not allowed for this application", requestID))
return
}
// Per-app 限流由全局 RateLimit 中间件处理,此处不再重复检查
// Resolve logical model
target, err := s.modelMap.Resolve(req.Model)
if err != nil {
@@ -60,6 +89,13 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
return
}
// 校验 max_output_tokens 不超过模型配置上限
if req.MaxOutputTokens > 0 && target.MaxOutputTokens > 0 && req.MaxOutputTokens > target.MaxOutputTokens {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest,
fmt.Sprintf("max_output_tokens %d exceeds model limit %d", req.MaxOutputTokens, target.MaxOutputTokens), requestID))
return
}
// Get adapter
adapterInst, err := s.registry.Get(target.Provider)
if err != nil {
@@ -75,6 +111,7 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
// Backpressure check: reject low-priority requests under high load
s.backpressure.Update(s.scheduler.RunningCount(), s.scheduler.QueueLength())
w.Header().Set("X-Backpressure-Level", fmt.Sprintf("%d", s.backpressure.Level()))
if !s.backpressure.ShouldAccept(priority) {
reason := s.backpressure.RejectReason(priority)
s.metrics.IncRequest("backpressure_rejected")
@@ -85,9 +122,19 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
// Circuit breaker check
if !s.breaker.AllowRequest() {
s.metrics.IncRequest("circuit_open")
s.metrics.SetBreakerState(1) // open
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrModelUnavailable, "circuit breaker open, please retry later", requestID))
return
}
// 同步熔断器状态到 Prometheus 指标
switch s.breaker.State() {
case "open":
s.metrics.SetBreakerState(1)
case "half_open":
s.metrics.SetBreakerState(2)
default:
s.metrics.SetBreakerState(0)
}
// Create task
taskID := uuid.New().String()
@@ -103,9 +150,21 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
s.metrics.SetQueueLength(s.scheduler.QueueLength())
// 解析 per-request 超时覆盖
var overrides map[string]int
if req.Timeouts != nil {
overrides = map[string]int{
"queue_ms": req.Timeouts.QueueMs,
"first_token_ms": req.Timeouts.FirstTokenMs,
"inference_ms": req.Timeouts.InferenceMs,
"total_ms": req.Timeouts.TotalMs,
}
}
resolvedTimeouts := s.timeoutMgr.ResolveTimeouts(&s.cfg.Timeouts, overrides)
// Wait for task to be dequeued
queueStart := time.Now()
ctx, cancel := context.WithTimeout(r.Context(), time.Duration(s.cfg.Timeouts.DefaultQueueMs)*time.Millisecond)
ctx, cancel := context.WithTimeout(r.Context(), time.Duration(resolvedTimeouts.QueueMs)*time.Millisecond)
defer cancel()
dequeued, err := s.scheduler.GetNext(ctx)
@@ -127,12 +186,25 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
// Build adapter request — assemble context if session_id provided
messages := req.Messages
if req.SessionID != "" {
sess, err := s.sessions.Get(req.SessionID)
if err == nil && sess != nil {
// 会话访问隔离:校验 session 属于当前 app
var sess *session.Session
var sessErr error
if identity != nil {
sess, sessErr = s.sessions.GetForApp(req.SessionID, identity.AppID)
} else {
sess, sessErr = s.sessions.Get(req.SessionID)
}
if sessErr == nil && sess != nil {
// Load session history and assemble context
history := s.loadSessionHistory(sess)
if len(history) > 0 {
assembler := ctxasm.NewAssembler(&s.cfg.Context)
// 注入摘要器,启用真实 LLM 摘要(受配置控制)
if s.cfg.Context.EnableLLMSummary {
summaryGen := ctxasm.NewAdapterSummaryGenerator(adapterInst, target.ActualModel, s.cfg.Context.SummaryMaxTokens)
summarizer := ctxasm.NewSummarizer(summaryGen, time.Duration(s.cfg.Context.SummaryTimeoutSeconds)*time.Second)
assembler.SetSummarizer(summarizer)
}
result := assembler.Assemble(history, req.Messages, target.ContextWindow, target.MaxOutputTokens, req.ContextPolicy)
messages = result.Messages
}
@@ -155,13 +227,20 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
}
if req.Stream {
s.handleStreaming(w, r, adapterInst, adapterReq, dequeued, requestID, target, req.Model)
s.handleStreaming(w, r, adapterInst, adapterReq, dequeued, requestID, target, req.Model, req.SessionID, req.Messages, identity.AppID, resolvedTimeouts)
} else {
s.handleNonStreaming(w, r, adapterInst, adapterReq, dequeued, requestID, target, req.Model)
// 非流式请求:如果带有幂等键,捕获响应用于缓存
if req.IdempotencyKey != "" && identity != nil {
crw := &captureResponseWriter{ResponseWriter: w}
s.handleNonStreaming(crw, r, adapterInst, adapterReq, dequeued, requestID, target, req.Model, req.SessionID, req.Messages, identity.AppID, resolvedTimeouts)
s.storeIdempotencyResult(req.IdempotencyKey, identity.AppID, crw.statusCode, crw.body)
} else {
s.handleNonStreaming(w, r, adapterInst, adapterReq, dequeued, requestID, target, req.Model, req.SessionID, req.Messages, identity.AppID, resolvedTimeouts)
}
}
}
func (s *Server) handleStreaming(w http.ResponseWriter, r *http.Request, adapterInst adapter.ModelAdapter, req *adapter.ChatRequest, tk *task.Task, requestID string, _ *router.ModelTarget, logicalModel string) {
func (s *Server) handleStreaming(w http.ResponseWriter, r *http.Request, adapterInst adapter.ModelAdapter, req *adapter.ChatRequest, tk *task.Task, requestID string, target *router.ModelTarget, logicalModel string, sessionID string, originalMessages []api.Message, appID string, tc *connector.TimeoutConfig) {
sse := handler.NewSSEWriter(w)
if sse == nil {
s.scheduler.Complete(tk.ID)
@@ -172,21 +251,120 @@ func (s *Server) handleStreaming(w http.ResponseWriter, r *http.Request, adapter
tk.Transition(task.StateStreaming)
inferenceStart := time.Now()
ch, err := adapterInst.ChatCompletionStream(r.Context(), req)
// 流式超时控制:使用 per-request 解析后的 inference timeout
streamCtx, streamCancel := context.WithTimeout(r.Context(), time.Duration(tc.InferenceMs)*time.Millisecond)
defer streamCancel()
ch, err := adapterInst.ChatCompletionStream(streamCtx, req)
if err != nil {
s.scheduler.Complete(tk.ID)
tk.Transition(task.StateFailed)
s.metrics.IncTask("failed")
// 发送 SSE 错误事件
sse.WriteChunk(map[string]any{
"id": requestID,
"object": "chat.completion.chunk",
"model": logicalModel,
"choices": []map[string]any{
{
"index": 0,
"delta": map[string]any{},
"finish_reason": "error",
},
},
"error": map[string]string{
"code": handler.ErrModelUnavailable,
"message": err.Error(),
},
})
sse.WriteDone()
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrModelUnavailable, err.Error(), requestID))
return
}
inputTokens, outputTokens, err := handler.StreamChatCompletion(sse, ch, requestID, tk.ID, logicalModel)
inputTokens, outputTokens, fullContent, err := handler.StreamChatCompletion(sse, ch, requestID, tk.ID, logicalModel, func() {
tk.FirstTokenMs = int(time.Since(inferenceStart) / time.Millisecond)
s.metrics.ObserveFirstTokenLatency(int64(tk.FirstTokenMs))
})
if err != nil {
s.logger.Error("streaming error", observability.F().Event("stream_error").TaskID(tk.ID).Reason(err.Error()))
tk.Transition(task.StateFailed)
s.metrics.IncTask("failed")
s.breaker.RecordError()
// 检查是否客户端断开
if r.Context().Err() != nil {
s.logger.Info("client disconnected during streaming",
observability.F().Event("client_disconnect").TaskID(tk.ID))
tk.Cancel("client disconnected")
s.metrics.IncCancellation()
} else {
s.logger.Error("streaming error", observability.F().Event("stream_error").TaskID(tk.ID).Reason(err.Error()))
// 尝试非流式 fallback 降级
fallbackResp, fbErr := s.tryOverloadFallback(r.Context(), req, target, logicalModel)
if fbErr == nil && fallbackResp != nil {
tk.Degraded = true
tk.InferenceMs = int(time.Since(inferenceStart) / time.Millisecond)
tk.TotalMs = int(time.Since(tk.CreatedAt) / time.Millisecond)
s.metrics.ObserveGatewayLatency(int64(tk.TotalMs))
tk.Transition(task.StateCompleted)
s.metrics.IncTask("completed")
s.metrics.IncRequest("stream_fallback")
s.metrics.IncDegraded()
s.metrics.AddTokens(fallbackResp.InputTokens, fallbackResp.OutputTokens)
s.breaker.RecordSuccess()
s.scheduler.Complete(tk.ID)
s.metrics.SetRunningTasks(s.scheduler.RunningCount())
s.metrics.SetQueueLength(s.scheduler.QueueLength())
s.usageTracker.Record(appID, fallbackResp.InputTokens, fallbackResp.OutputTokens, false)
// 将降级结果作为单个 SSE chunk 发送
sse.WriteChunk(map[string]any{
"id": requestID,
"object": "chat.completion.chunk",
"model": logicalModel,
"choices": []map[string]any{
{
"index": 0,
"delta": map[string]any{
"content": fallbackResp.Content,
},
"finish_reason": fallbackResp.FinishReason,
},
},
"usage": map[string]int{
"input_tokens": fallbackResp.InputTokens,
"output_tokens": fallbackResp.OutputTokens,
"total_tokens": fallbackResp.InputTokens + fallbackResp.OutputTokens,
},
"degraded": true,
})
sse.WriteDone()
if sessionID != "" && fallbackResp.Content != "" {
s.saveSessionMessages(sessionID, originalMessages, fallbackResp.Content)
}
return
}
// fallback 也失败,发送 SSE 错误事件
sse.WriteChunk(map[string]any{
"id": requestID,
"object": "chat.completion.chunk",
"model": logicalModel,
"choices": []map[string]any{
{
"index": 0,
"delta": map[string]any{},
"finish_reason": "error",
},
},
"error": map[string]string{
"code": handler.ErrInferenceTimeout,
"message": err.Error(),
},
})
sse.WriteDone()
tk.Transition(task.StateFailed)
s.metrics.IncTask("failed")
s.breaker.RecordError()
}
} else {
tk.Transition(task.StateCompleted)
s.metrics.IncTask("completed")
@@ -196,15 +374,30 @@ func (s *Server) handleStreaming(w http.ResponseWriter, r *http.Request, adapter
tk.InferenceMs = int(time.Since(inferenceStart) / time.Millisecond)
tk.TotalMs = int(time.Since(tk.CreatedAt) / time.Millisecond)
// 记录延迟指标
s.metrics.ObserveGatewayLatency(int64(tk.TotalMs))
s.metrics.AddTokens(inputTokens, outputTokens)
s.scheduler.Complete(tk.ID)
s.metrics.SetRunningTasks(s.scheduler.RunningCount())
s.metrics.SetQueueLength(s.scheduler.QueueLength())
s.metrics.IncRequest("stream_ok")
if err != nil {
s.metrics.IncRequest("stream_error")
} else {
s.metrics.IncRequest("stream_ok")
}
// 记录 per-app 用量
s.usageTracker.Record(appID, inputTokens, outputTokens, err != nil)
// 保存对话到 session
if sessionID != "" && fullContent != "" {
s.saveSessionMessages(sessionID, originalMessages, fullContent)
}
}
func (s *Server) handleNonStreaming(w http.ResponseWriter, r *http.Request, adapterInst adapter.ModelAdapter, req *adapter.ChatRequest, tk *task.Task, requestID string, target *router.ModelTarget, logicalModel string) {
ctx, cancel := context.WithTimeout(r.Context(), time.Duration(s.cfg.Timeouts.DefaultInferenceMs)*time.Millisecond)
func (s *Server) handleNonStreaming(w http.ResponseWriter, r *http.Request, adapterInst adapter.ModelAdapter, req *adapter.ChatRequest, tk *task.Task, requestID string, target *router.ModelTarget, logicalModel string, sessionID string, originalMessages []api.Message, appID string, tc *connector.TimeoutConfig) {
ctx, cancel := context.WithTimeout(r.Context(), time.Duration(tc.InferenceMs)*time.Millisecond)
defer cancel()
inferenceStart := time.Now()
@@ -228,9 +421,16 @@ func (s *Server) handleNonStreaming(w http.ResponseWriter, r *http.Request, adap
s.metrics.SetRunningTasks(s.scheduler.RunningCount())
s.metrics.SetQueueLength(s.scheduler.QueueLength())
// 记录 per-app 用量(降级请求)
s.usageTracker.Record(appID, fallbackResp.InputTokens, fallbackResp.OutputTokens, false)
chatResp := handler.BuildChatResponse(requestID, tk.ID, logicalModel, fallbackResp)
chatResp.Degraded = true
chatResp.Timing = s.buildTiming(tk)
// 通过 HTTP 头暴露关键 timing 指标
w.Header().Set("X-Timing-Queue-Ms", strconv.Itoa(tk.QueueMs))
w.Header().Set("X-Timing-Inference-Ms", strconv.Itoa(tk.InferenceMs))
w.Header().Set("X-Timing-Total-Ms", strconv.Itoa(tk.TotalMs))
handler.WriteJSON(w, http.StatusOK, chatResp)
return
}
@@ -253,6 +453,9 @@ func (s *Server) handleNonStreaming(w http.ResponseWriter, r *http.Request, adap
tk.InferenceMs = int(time.Since(inferenceStart) / time.Millisecond)
tk.TotalMs = int(time.Since(tk.CreatedAt) / time.Millisecond)
// 记录延迟指标
s.metrics.ObserveGatewayLatency(int64(tk.TotalMs))
tk.Transition(task.StateCompleted)
s.metrics.IncTask("completed")
s.metrics.IncRequest("ok")
@@ -262,9 +465,21 @@ func (s *Server) handleNonStreaming(w http.ResponseWriter, r *http.Request, adap
s.metrics.SetRunningTasks(s.scheduler.RunningCount())
s.metrics.SetQueueLength(s.scheduler.QueueLength())
// 记录 per-app 用量
s.usageTracker.Record(appID, resp.InputTokens, resp.OutputTokens, false)
chatResp := handler.BuildChatResponse(requestID, tk.ID, logicalModel, resp)
chatResp.Timing = s.buildTiming(tk)
// 通过 HTTP 头暴露关键 timing 指标,便于客户端和监控系统采集
w.Header().Set("X-Timing-Queue-Ms", strconv.Itoa(tk.QueueMs))
w.Header().Set("X-Timing-Inference-Ms", strconv.Itoa(tk.InferenceMs))
w.Header().Set("X-Timing-Total-Ms", strconv.Itoa(tk.TotalMs))
handler.WriteJSON(w, http.StatusOK, chatResp)
// 保存对话到 session
if sessionID != "" {
s.saveSessionMessages(sessionID, originalMessages, resp.Content)
}
}
// tryOverloadFallback attempts fallback strategies when the primary model fails.
@@ -314,10 +529,18 @@ func (s *Server) handleModels(w http.ResponseWriter, r *http.Request) {
models := s.modelMap.List()
data := make([]api.ModelInfo, len(models))
for i, m := range models {
target, err := s.modelMap.Resolve(m)
if err != nil {
data[i] = api.ModelInfo{ID: m, Object: "model", OwnedBy: "edgeai-gateway"}
continue
}
data[i] = api.ModelInfo{
ID: m,
Object: "model",
OwnedBy: "edgeai-gateway",
ID: m,
Object: "model",
OwnedBy: "edgeai-gateway",
Provider: target.Provider,
ContextWindow: target.ContextWindow,
MaxOutputTokens: target.MaxOutputTokens,
}
}
@@ -366,8 +589,46 @@ func (s *Server) handleSessions(w http.ResponseWriter, r *http.Request) {
handler.WriteJSON(w, http.StatusCreated, resp)
case http.MethodGet:
// List sessions (simplified: return empty for now)
handler.WriteJSON(w, http.StatusOK, map[string]any{"sessions": []any{}})
// 列出当前 app 的会话(支持分页)
appID := ""
if identity != nil {
appID = identity.AppID
}
page := 1
pageSize := 20
if v := r.URL.Query().Get("page"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 {
page = n
}
}
if v := r.URL.Query().Get("page_size"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 && n <= 100 {
pageSize = n
}
}
offset := (page - 1) * pageSize
sessions, err := s.sessions.ListByApp(appID, pageSize, offset)
if err != nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInternalError, err.Error(), requestID))
return
}
// 脱敏:不返回 messages 全文,只返回元信息
items := make([]map[string]any, 0, len(sessions))
for _, sess := range sessions {
items = append(items, map[string]any{
"session_id": sess.ID,
"application_id": sess.ApplicationID,
"user_id": sess.UserID,
"message_count": len(sess.Messages),
"created_at": sess.CreatedAt.Format(time.RFC3339),
"last_active": sess.LastActive.Format(time.RFC3339),
})
}
handler.WriteJSON(w, http.StatusOK, map[string]any{
"sessions": items,
"page": page,
"page_size": pageSize,
})
default:
handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
@@ -376,11 +637,19 @@ func (s *Server) handleSessions(w http.ResponseWriter, r *http.Request) {
func (s *Server) handleSessionByID(w http.ResponseWriter, r *http.Request) {
requestID := middleware.GetRequestID(r.Context())
identity := auth.GetAppIdentityFromRequest(r)
sessionID := strings.TrimPrefix(r.URL.Path, "/v1/sessions/")
switch r.Method {
case http.MethodGet:
sess, err := s.sessions.Get(sessionID)
// 会话访问隔离:校验 session 属于当前 app
var sess *session.Session
var err error
if identity != nil {
sess, err = s.sessions.GetForApp(sessionID, identity.AppID)
} else {
sess, err = s.sessions.Get(sessionID)
}
if err != nil || sess == nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, "session not found", requestID))
return
@@ -388,9 +657,17 @@ func (s *Server) handleSessionByID(w http.ResponseWriter, r *http.Request) {
handler.WriteJSON(w, http.StatusOK, sess)
case http.MethodDelete:
if err := s.sessions.Delete(sessionID); err != nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInternalError, err.Error(), requestID))
return
// 会话访问隔离:校验 session 属于当前 app
if identity != nil {
if err := s.sessions.DeleteForApp(sessionID, identity.AppID); err != nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInvalidRequest, err.Error(), requestID))
return
}
} else {
if err := s.sessions.Delete(sessionID); err != nil {
handler.WriteError(w, handler.NewGatewayErrorWithID(handler.ErrInternalError, err.Error(), requestID))
return
}
}
handler.WriteJSON(w, http.StatusOK, map[string]string{"status": "deleted"})
@@ -407,6 +684,28 @@ func (s *Server) loadSessionHistory(sess *session.Session) []api.Message {
return sess.Messages
}
// saveSessionMessages 将用户消息和助手回复保存到会话历史。
// 当配置了 MaxSessionMessages 时,自动裁剪最旧消息以防止无限增长。
func (s *Server) saveSessionMessages(sessionID string, userMessages []api.Message, assistantContent string) {
if sessionID == "" {
return
}
maxMsgs := s.cfg.Context.MaxSessionMessages
for _, msg := range userMessages {
if err := s.sessions.AddMessageWithLimit(sessionID, msg, maxMsgs); err != nil {
s.logger.Error("failed to save user message to session",
observability.F().Event("session_save_error").Reason(err.Error()))
}
}
if assistantContent != "" {
assistantMsg := api.Message{Role: "assistant", Content: assistantContent}
if err := s.sessions.AddMessageWithLimit(sessionID, assistantMsg, maxMsgs); err != nil {
s.logger.Error("failed to save assistant message to session",
observability.F().Event("session_save_error").Reason(err.Error()))
}
}
}
// buildTiming constructs Timing metadata from a task.
func (s *Server) buildTiming(tk *task.Task) *api.Timing {
return &api.Timing{
@@ -422,7 +721,59 @@ func (s *Server) handleTasks(w http.ResponseWriter, r *http.Request) {
handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "method not allowed"))
return
}
handler.WriteJSON(w, http.StatusOK, map[string]any{"tasks": []any{}})
// 分页参数
page := 1
pageSize := 20
if v := r.URL.Query().Get("page"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 {
page = n
}
}
if v := r.URL.Query().Get("page_size"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 && n <= 100 {
pageSize = n
}
}
offset := (page - 1) * pageSize
// 状态过滤
status := r.URL.Query().Get("status")
tasks, total := s.scheduler.ListTasks(status, pageSize, offset)
items := make([]map[string]any, 0, len(tasks))
for _, tk := range tasks {
item := map[string]any{
"task_id": tk.ID,
"request_id": tk.RequestID,
"session_id": tk.SessionID,
"state": string(tk.GetState()),
"priority": int(tk.Priority),
"logical_model": tk.LogicalModel,
"stream": tk.Stream,
"created_at": tk.CreatedAt.Format(time.RFC3339),
"degraded": tk.Degraded,
}
if tk.StartedAt != nil {
item["started_at"] = tk.StartedAt.Format(time.RFC3339)
}
if tk.CompletedAt != nil {
item["completed_at"] = tk.CompletedAt.Format(time.RFC3339)
}
if tk.ErrorMessage != "" {
item["error_message"] = tk.ErrorMessage
}
item["timing"] = s.buildTiming(tk)
items = append(items, item)
}
handler.WriteJSON(w, http.StatusOK, map[string]any{
"tasks": items,
"total": total,
"page": page,
"page_size": pageSize,
})
}
func (s *Server) handleTaskByID(w http.ResponseWriter, r *http.Request) {
+147
View File
@@ -0,0 +1,147 @@
package server
import (
"net/http"
"time"
"github.com/edgeai/gateway/internal/handler"
)
// idempotencyEntry 幂等键缓存条目。
// 记录请求处理状态和响应,用于在重复请求时返回缓存结果。
type idempotencyEntry struct {
status int
response []byte
createdAt time.Time
inFlight bool // 正在处理中
}
const (
// idempotencyTTL 幂等键缓存存活时间。
idempotencyTTL = 10 * time.Minute
// idempotencyCleanupInterval 清理间隔。
idempotencyCleanupInterval = 5 * time.Minute
)
// checkIdempotency 检查幂等键,如果重复请求则返回缓存的响应。
// 如果是首次请求,返回 nil 表示可以继续处理。
// 如果是正在处理中的重复请求,返回 409 Conflict。
func (s *Server) checkIdempotency(w http.ResponseWriter, key, appID string) bool {
if key == "" {
return false
}
cacheKey := appID + ":" + key
s.idempotencyMu.Lock()
entry, exists := s.idempotencyCache[cacheKey]
if exists {
if entry.inFlight {
// 正在处理中,返回 409
s.idempotencyMu.Unlock()
handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "duplicate idempotency key, request already in progress"))
return true
}
// 检查是否过期
if time.Since(entry.createdAt) > idempotencyTTL {
delete(s.idempotencyCache, cacheKey)
} else {
// 返回缓存的响应
s.idempotencyMu.Unlock()
w.Header().Set("Content-Type", "application/json")
w.Header().Set("X-Idempotent-Replay", "true")
w.WriteHeader(entry.status)
w.Write(entry.response)
return true
}
}
// 标记为正在处理中
s.idempotencyCache[cacheKey] = &idempotencyEntry{
inFlight: true,
createdAt: time.Now(),
}
s.idempotencyMu.Unlock()
return false
}
// storeIdempotencyResult 存储幂等请求的响应结果。
func (s *Server) storeIdempotencyResult(key, appID string, status int, response []byte) {
if key == "" {
return
}
cacheKey := appID + ":" + key
s.idempotencyMu.Lock()
s.idempotencyCache[cacheKey] = &idempotencyEntry{
status: status,
response: response,
createdAt: time.Now(),
inFlight: false,
}
s.idempotencyMu.Unlock()
}
// cleanupIdempotencyCache 清理过期的幂等键缓存。
func (s *Server) cleanupIdempotencyCache() {
s.idempotencyMu.Lock()
defer s.idempotencyMu.Unlock()
now := time.Now()
for k, entry := range s.idempotencyCache {
if now.Sub(entry.createdAt) > idempotencyTTL {
delete(s.idempotencyCache, k)
}
}
}
// startIdempotencyCleanup 启动后台清理任务,定期清理过期的幂等键。
// 返回停止函数。
func (s *Server) startIdempotencyCleanup() func() {
ticker := time.NewTicker(idempotencyCleanupInterval)
stopCh := make(chan struct{})
go func() {
for {
select {
case <-ticker.C:
s.cleanupIdempotencyCache()
case <-stopCh:
ticker.Stop()
return
}
}
}()
return func() {
close(stopCh)
}
}
// captureResponseWriter 捕获响应状态和响应体,用于幂等性缓存。
type captureResponseWriter struct {
http.ResponseWriter
statusCode int
body []byte
}
func (crw *captureResponseWriter) WriteHeader(code int) {
crw.statusCode = code
crw.ResponseWriter.WriteHeader(code)
}
func (crw *captureResponseWriter) Write(b []byte) (int, error) {
if crw.statusCode == 0 {
crw.statusCode = 200
}
crw.body = append(crw.body, b...)
return crw.ResponseWriter.Write(b)
}
func (crw *captureResponseWriter) Flush() {
if f, ok := crw.ResponseWriter.(http.Flusher); ok {
f.Flush()
}
}
+296 -36
View File
@@ -7,33 +7,50 @@ import (
"os"
"path/filepath"
"strings"
"sync"
"time"
"github.com/edgeai/gateway/internal/adapter"
"github.com/edgeai/gateway/internal/auth"
"github.com/edgeai/gateway/internal/config"
"github.com/edgeai/gateway/internal/connector"
"github.com/edgeai/gateway/internal/handler"
"github.com/edgeai/gateway/internal/middleware"
"github.com/edgeai/gateway/internal/observability"
"github.com/edgeai/gateway/internal/ratelimit"
"github.com/edgeai/gateway/internal/router"
"github.com/edgeai/gateway/internal/scheduler"
"github.com/edgeai/gateway/internal/session"
"github.com/edgeai/gateway/internal/usage"
)
// Server is the main HTTP server for the AI gateway.
type Server struct {
cfg *config.Config
logger *observability.Logger
metrics *observability.Metrics
HTTPSrv *http.Server
auth *auth.Authenticator
registry *adapter.Registry
modelMap *router.LogicalModelMapping
scheduler *scheduler.Scheduler
sessions *session.Store
breaker *scheduler.CircuitBreaker
backpressure *scheduler.BackpressureManager
overload *router.OverloadResolver
cfg *config.Config
logger *observability.Logger
metrics *observability.Metrics
HTTPSrv *http.Server
auth *auth.Authenticator
registry *adapter.Registry
modelMap *router.LogicalModelMapping
scheduler *scheduler.Scheduler
sessions *session.Store
breaker *scheduler.CircuitBreaker
backpressure *scheduler.BackpressureManager
overload *router.OverloadResolver
rateLimiter *ratelimit.Limiter
usageTracker *usage.Tracker
audit *observability.AuditLogger
timeoutMgr *connector.TimeoutManager
sessionCleanup func() // 停止会话清理任务的函数
adminSrv *http.Server // Admin API 独立端口服务器
healthCache map[string]bool // adapter 健康缓存
healthCacheTime time.Time // 缓存时间
healthCacheMu sync.RWMutex // 缓存锁
idempotencyCache map[string]*idempotencyEntry // 幂等键缓存
idempotencyMu sync.Mutex // 幂等缓存锁
}
// New creates a new Server instance with all components wired.
@@ -105,36 +122,115 @@ func New(cfg *config.Config, logger *observability.Logger) (*Server, error) {
// Initialize metrics
metrics := observability.NewMetrics()
s := &Server{
cfg: cfg,
logger: logger,
metrics: metrics,
auth: authenticator,
registry: registry,
modelMap: modelMap,
scheduler: sched,
sessions: sessionStore,
breaker: breaker,
backpressure: bp,
overload: overloadResolver,
// Initialize per-app rate limiter (令牌桶)
refillPerSec := float64(cfg.Auth.RateLimitPerMinute) / 60.0
rateLimiter := ratelimit.NewLimiter(float64(cfg.Auth.RateLimitBurst), refillPerSec)
// Initialize per-app usage tracker (SQLite 持久化)
usageDBPath := filepath.Join(filepath.Dir(dbPath), "usage.db")
usageTracker, err := usage.NewTracker(time.Duration(cfg.Auth.UsageWindowMinutes)*time.Minute, usageDBPath, cfg.Observability.AuditRetentionDays)
if err != nil {
return nil, fmt.Errorf("init usage tracker: %w", err)
}
timeoutMgr := connector.NewTimeoutManager(&cfg.Timeouts)
auditLogger := observability.NewAuditLogger(filepath.Dir(dbPath), logger)
s := &Server{
cfg: cfg,
logger: logger,
metrics: metrics,
auth: authenticator,
registry: registry,
modelMap: modelMap,
scheduler: sched,
sessions: sessionStore,
breaker: breaker,
backpressure: bp,
overload: overloadResolver,
rateLimiter: rateLimiter,
usageTracker: usageTracker,
audit: auditLogger,
timeoutMgr: timeoutMgr,
healthCache: make(map[string]bool),
idempotencyCache: make(map[string]*idempotencyEntry),
}
// 启动会话 TTL 自动清理任务(带日志回调)
s.sessionCleanup = sessionStore.StartCleanupTask(cfg.Context.SessionIdleTTLMinutes, 10, func(deleted int64) {
logger.Info("session cleanup completed",
observability.F().
Event("session_cleanup").
Set("deleted_sessions", deleted))
})
// 启动幂等键缓存清理任务
idempotencyCleanup := s.startIdempotencyCleanup()
_ = idempotencyCleanup // 进程退出时自动回收
mux := http.NewServeMux()
s.registerRoutes(mux)
// Apply middleware chain (order: Recovery → Logging → RequestID → BodyLimit → Auth → handler)
// Apply middleware chain (order: Recovery → Logging → CORS → Auth → RateLimit → RequestID → BodyLimit → handler)
h := middleware.RequestID(mux)
h = middleware.BodyLimit(cfg.Server.MaxRequestBodyMB)(h)
h = middleware.RateLimit(func(appID string) (bool, int, int) {
info := s.rateLimiter.AllowWithInfo(appID)
return info.Allowed, info.Limit, info.Remaining
}, func(r *http.Request) string {
if id := auth.GetAppIdentityFromRequest(r); id != nil {
return id.AppID
}
return ""
}, map[string]bool{"/health": true, "/ready": true, "/metrics": true})(h)
h = s.auth.Middleware(h)
h = middleware.CORS(cfg.Server.CORSAllowedOrigins)(h)
h = middleware.Gzip(h)
h = middleware.Logging(logger)(h)
h = middleware.Recovery(logger)(h)
s.HTTPSrv = &http.Server{
Addr: fmt.Sprintf("%s:%d", cfg.Server.Host, cfg.Server.Port),
Handler: h,
ReadTimeout: 30 * time.Second,
WriteTimeout: 0, // no write timeout for SSE
IdleTimeout: 120 * time.Second,
Addr: fmt.Sprintf("%s:%d", cfg.Server.Host, cfg.Server.Port),
Handler: h,
ReadTimeout: 30 * time.Second,
ReadHeaderTimeout: 10 * time.Second, // 防止 slowloris 攻击
WriteTimeout: 0, // no write timeout for SSE
IdleTimeout: 120 * time.Second,
MaxHeaderBytes: 1 << 20, // 1MB 限制请求头大小
}
// Admin API 独立端口(如果配置了 admin_port)
if cfg.Server.AdminPort > 0 && cfg.Server.AdminPort != cfg.Server.Port {
adminMux := http.NewServeMux()
s.registerAdminRoutes(adminMux)
// Admin 中间件链(与主服务器相同,但额外加 CORS)
adminH := middleware.RequestID(adminMux)
adminH = middleware.BodyLimit(cfg.Server.MaxRequestBodyMB)(adminH)
adminH = middleware.RateLimit(func(appID string) (bool, int, int) {
info := s.rateLimiter.AllowWithInfo(appID)
return info.Allowed, info.Limit, info.Remaining
}, func(r *http.Request) string {
if id := auth.GetAppIdentityFromRequest(r); id != nil {
return id.AppID
}
return ""
}, nil)(adminH)
adminH = s.auth.Middleware(adminH)
adminH = middleware.CORS(cfg.Server.CORSAllowedOrigins)(adminH)
adminH = middleware.Logging(logger)(adminH)
adminH = middleware.Recovery(logger)(adminH)
s.adminSrv = &http.Server{
Addr: fmt.Sprintf("%s:%d", cfg.Server.Host, cfg.Server.AdminPort),
Handler: adminH,
ReadTimeout: 30 * time.Second,
ReadHeaderTimeout: 10 * time.Second, // 防止 slowloris 攻击
WriteTimeout: 30 * time.Second,
IdleTimeout: 120 * time.Second,
MaxHeaderBytes: 1 << 20, // 1MB 限制请求头大小
}
}
return s, nil
@@ -161,53 +257,182 @@ func (s *Server) registerRoutes(mux *http.ServeMux) {
mux.HandleFunc("/v1/tasks/", s.handleTaskByID)
}
// registerAdminRoutes 注册管理 API 路由到独立的 mux(运行在 admin_port)。
func (s *Server) registerAdminRoutes(mux *http.ServeMux) {
// Admin API: API Key 管理(需要 admin 权限)
mux.Handle("/v1/admin/keys", auth.RequireAdmin(http.HandlerFunc(s.handleAdminKeys)))
mux.Handle("/v1/admin/keys/", auth.RequireAdmin(http.HandlerFunc(s.handleAdminKeyByID)))
// Admin API: 用量统计(需要 admin 权限)
mux.Handle("/v1/admin/usage", auth.RequireAdmin(http.HandlerFunc(s.handleAdminUsage)))
mux.Handle("/v1/admin/usage/", auth.RequireAdmin(http.HandlerFunc(s.handleAdminUsageByApp)))
// Admin API: per-app 限流管理(需要 admin 权限)
mux.Handle("/v1/admin/ratelimit/", auth.RequireAdmin(http.HandlerFunc(s.handleAdminRateLimit)))
// Admin API: 熔断器状态(需要 admin 权限)
mux.Handle("/v1/admin/circuit-breaker", auth.RequireAdmin(http.HandlerFunc(s.handleAdminCircuitBreaker)))
// Admin API: 调度器状态(需要 admin 权限)
mux.Handle("/v1/admin/scheduler", auth.RequireAdmin(http.HandlerFunc(s.handleAdminScheduler)))
// Admin API: 配置热重载(需要 admin 权限)
mux.Handle("/v1/admin/config/reload", auth.RequireAdmin(http.HandlerFunc(s.handleAdminConfigReload)))
// Admin API: 健康检查(复用主服务器逻辑)
mux.HandleFunc("/health", s.handleHealth)
mux.HandleFunc("/ready", s.handleReady)
}
// Authenticator returns the authenticator instance (for testing/management).
func (s *Server) Authenticator() *auth.Authenticator {
return s.auth
}
// Start begins listening for HTTP requests.
// 如果配置了 admin_port,同时启动 Admin API 服务器。
// 启动时对所有已注册 adapter 执行健康预检,不可达的 adapter 记录警告但不阻止启动。
func (s *Server) Start() error {
s.logger.Info("http server starting", observability.F().
Event("server_start").
Set("addr", s.HTTPSrv.Addr))
// Adapter 启动健康预检
s.checkAdaptersHealth()
// 启动 Admin API 服务器(独立 goroutine
if s.adminSrv != nil {
s.logger.Info("admin api server starting", observability.F().
Event("admin_server_start").
Set("addr", s.adminSrv.Addr))
go func() {
if err := s.adminSrv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
s.logger.Error("admin api server error", observability.F().
Event("admin_server_error").
Reason(err.Error()))
}
}()
}
return s.HTTPSrv.ListenAndServe()
}
// checkAdaptersHealth 对所有已注册 adapter 执行健康预检。
// 不可达的 adapter 记录警告但不阻止启动,允许部分降级运行。
// 每个 adapter 使用独立超时,避免单个 adapter 阻塞拖累其他。
func (s *Server) checkAdaptersHealth() {
for _, name := range s.registry.Names() {
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
a, _ := s.registry.Get(name)
if err := a.HealthCheck(ctx); err != nil {
s.logger.Warn("adapter health check failed on startup",
observability.F().
Event("adapter_health_check_failed").
Set("adapter", name).
Reason(err.Error()))
} else {
s.logger.Info("adapter health check passed",
observability.F().
Event("adapter_health_check_ok").
Set("adapter", name))
}
cancel()
}
}
// Shutdown gracefully shuts down the server.
// 等待最多 30 秒让进行中的请求完成,记录未完成任务状态。
func (s *Server) Shutdown() error {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
// 记录关机时的任务状态
runningCount := s.scheduler.RunningCount()
queueLength := s.scheduler.QueueLength()
s.logger.Info("server shutting down",
observability.F().
Event("server_shutdown").
Set("running_tasks", runningCount).
Set("queued_tasks", queueLength))
// 等待运行中任务完成(超时时间从配置读取)
drainSeconds := s.cfg.Timeouts.ShutdownDrainSeconds
if drainSeconds <= 0 {
drainSeconds = 10
}
if runningCount > 0 {
s.logger.Info("waiting for running tasks to complete",
observability.F().
Event("shutdown_drain").
Set("running_tasks", runningCount).
Set("drain_timeout_seconds", drainSeconds))
for i := 0; i < drainSeconds*10 && s.scheduler.RunningCount() > 0; i++ {
time.Sleep(100 * time.Millisecond)
}
remaining := s.scheduler.RunningCount()
if remaining > 0 {
s.logger.Warn("shutdown: tasks still running after drain timeout",
observability.F().
Event("shutdown_drain_timeout").
Set("remaining_tasks", remaining))
} else {
s.logger.Info("all running tasks completed",
observability.F().Event("shutdown_drain_complete"))
}
}
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
s.scheduler.Stop()
if s.sessionCleanup != nil {
s.sessionCleanup()
}
if s.sessions != nil {
s.sessions.Close()
}
if s.auth != nil {
s.auth.Close()
}
if s.usageTracker != nil {
s.usageTracker.Close()
}
if s.audit != nil {
s.audit.Close()
}
if s.adminSrv != nil {
s.adminSrv.Shutdown(ctx)
}
return s.HTTPSrv.Shutdown(ctx)
}
func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
handler.WriteJSON(w, http.StatusOK, map[string]string{"status": "ok"})
// Liveness 探针:进程存活即返回 200
handler.WriteJSON(w, http.StatusOK, map[string]string{"status": "alive"})
}
func (s *Server) handleReady(w http.ResponseWriter, r *http.Request) {
// Readiness 探针:检查 adapter 健康状态 + DB 连通性
// 使用 5 秒 TTL 缓存,避免高频探针请求打满 adapter
w.Header().Set("Content-Type", "application/json")
ready := true
reasons := []string{}
for _, name := range s.registry.Names() {
a, _ := s.registry.Get(name)
if err := a.HealthCheck(r.Context()); err != nil {
// 检查所有 adapter 健康状态(带缓存)
adapterHealth := s.getCachedAdapterHealth(r.Context())
for name, ok := range adapterHealth {
if !ok {
ready = false
reasons = append(reasons, fmt.Sprintf("%s: %v", name, err))
reasons = append(reasons, fmt.Sprintf("adapter %s: unhealthy", name))
}
}
// 检查 scheduler 是否正常
if s.scheduler.RunningCount() >= s.cfg.Scheduler.MaxRunningTasks {
ready = false
reasons = append(reasons, "scheduler: at max capacity")
}
if ready {
w.WriteHeader(http.StatusOK)
w.Write([]byte(`{"status":"ready"}`))
@@ -217,6 +442,41 @@ func (s *Server) handleReady(w http.ResponseWriter, r *http.Request) {
}
}
// getCachedAdapterHealth 返回 adapter 健康状态,使用 5 秒 TTL 缓存。
// 缓存过期后异步刷新,首次请求或缓存过期时同步调用 adapter HealthCheck。
func (s *Server) getCachedAdapterHealth(ctx context.Context) map[string]bool {
const healthCacheTTL = 5 * time.Second
s.healthCacheMu.RLock()
if time.Since(s.healthCacheTime) < healthCacheTTL && len(s.healthCache) > 0 {
result := make(map[string]bool, len(s.healthCache))
for k, v := range s.healthCache {
result[k] = v
}
s.healthCacheMu.RUnlock()
return result
}
s.healthCacheMu.RUnlock()
// 缓存过期,同步执行健康检查
result := make(map[string]bool)
for _, name := range s.registry.Names() {
a, _ := s.registry.Get(name)
if err := a.HealthCheck(ctx); err != nil {
result[name] = false
} else {
result[name] = true
}
}
s.healthCacheMu.Lock()
s.healthCache = result
s.healthCacheTime = time.Now()
s.healthCacheMu.Unlock()
return result
}
func extractDBPath(connStr string) string {
if strings.HasPrefix(connStr, "sqlite://") {
return strings.TrimPrefix(connStr, "sqlite://")
+153 -10
View File
@@ -128,26 +128,84 @@ func (s *Store) Get(id string) (*Session, error) {
return session, nil
}
// AddMessage appends a message to the session and updates last_active.
func (s *Store) AddMessage(id string, msg api.Message) error {
// GetForApp 检索会话并校验 application_id 是否匹配。
// 如果会话不属于该 app,返回 nil(权限隔离)。
func (s *Store) GetForApp(id, appID string) (*Session, error) {
sess, err := s.Get(id)
if err != nil {
return nil, err
}
if sess == nil {
return nil, nil
}
if sess.ApplicationID != appID {
return nil, nil
}
return sess, nil
}
// DeleteForApp 删除会话并校验 application_id 是否匹配。
// 如果会话不属于该 app,返回错误。
func (s *Store) DeleteForApp(id, appID string) error {
s.mu.Lock()
defer s.mu.Unlock()
session, err := s.Get(id)
if err != nil {
return err
}
if session == nil {
var existingAppID string
err := s.db.QueryRow(`SELECT application_id FROM sessions WHERE id = ?`, id).Scan(&existingAppID)
if err == sql.ErrNoRows {
return fmt.Errorf("session not found: %s", id)
}
if err != nil {
return fmt.Errorf("query session: %w", err)
}
if existingAppID != appID {
return fmt.Errorf("session does not belong to application: %s", appID)
}
session.Messages = append(session.Messages, msg)
msgsJSON, _ := json.Marshal(session.Messages)
_, err = s.db.Exec(`DELETE FROM sessions WHERE id = ?`, id)
return err
}
// AddMessage 追加消息到会话并更新 last_active。
// 当消息数超过 maxMessages 时,自动裁剪最旧的消息。
// 注意:不能调用 s.Get(),因为 Get 会获取读锁,而此处已持有写锁,会导致死锁。
func (s *Store) AddMessage(id string, msg api.Message) error {
return s.AddMessageWithLimit(id, msg, 0)
}
// AddMessageWithLimit 追加消息到会话,当 maxMessages > 0 时限制消息总数。
// 超出上限时自动裁剪最旧的消息,保留最新的 maxMessages 条。
func (s *Store) AddMessageWithLimit(id string, msg api.Message, maxMessages int) error {
s.mu.Lock()
defer s.mu.Unlock()
var msgsJSON string
err := s.db.QueryRow(`SELECT messages FROM sessions WHERE id = ?`, id).Scan(&msgsJSON)
if err == sql.ErrNoRows {
return fmt.Errorf("session not found: %s", id)
}
if err != nil {
return fmt.Errorf("query session messages: %w", err)
}
var msgs []api.Message
if err := json.Unmarshal([]byte(msgsJSON), &msgs); err != nil {
return fmt.Errorf("unmarshal session messages: %w", err)
}
msgs = append(msgs, msg)
// 消息数上限保护:裁剪最旧消息
if maxMessages > 0 && len(msgs) > maxMessages {
msgs = msgs[len(msgs)-maxMessages:]
}
newMsgsJSON, _ := json.Marshal(msgs)
now := time.Now().Format(time.RFC3339)
_, err = s.db.Exec(
`UPDATE sessions SET messages = ?, last_active = ? WHERE id = ?`,
string(msgsJSON), now, id,
string(newMsgsJSON), now, id,
)
return err
}
@@ -161,11 +219,96 @@ func (s *Store) Delete(id string) error {
return err
}
// ListByApp 列出指定 app 的会话(分页)。
func (s *Store) ListByApp(appID string, limit, offset int) ([]*Session, error) {
s.mu.RLock()
defer s.mu.RUnlock()
if limit <= 0 || limit > 100 {
limit = 20
}
rows, err := s.db.Query(
`SELECT id, application_id, tenant_id, user_id, messages, config, created_at, last_active
FROM sessions WHERE application_id = ? ORDER BY last_active DESC LIMIT ? OFFSET ?`,
appID, limit, offset,
)
if err != nil {
return nil, fmt.Errorf("list sessions: %w", err)
}
defer rows.Close()
var sessions []*Session
for rows.Next() {
var (
sid, sAppID, sTenantID, sUserID, msgsJSON, configJSON, createdAt, lastActive string
)
if err := rows.Scan(&sid, &sAppID, &sTenantID, &sUserID, &msgsJSON, &configJSON, &createdAt, &lastActive); err != nil {
return nil, err
}
sess := &Session{
ID: sid,
ApplicationID: sAppID,
TenantID: sTenantID,
UserID: sUserID,
CreatedAt: parseTime(createdAt),
LastActive: parseTime(lastActive),
}
json.Unmarshal([]byte(msgsJSON), &sess.Messages)
json.Unmarshal([]byte(configJSON), &sess.Config)
sessions = append(sessions, sess)
}
return sessions, rows.Err()
}
// Close closes the database connection.
func (s *Store) Close() error {
return s.db.Close()
}
// CleanupExpired 清理超过 TTL 未活跃的会话。
// ttlMinutes 为会话空闲超时(分钟),超过此时间未活跃的会话将被删除。
// 返回删除的会话数量。
func (s *Store) CleanupExpired(ttlMinutes int) (int64, error) {
s.mu.Lock()
defer s.mu.Unlock()
cutoff := time.Now().Add(-time.Duration(ttlMinutes) * time.Minute).Format(time.RFC3339)
result, err := s.db.Exec(`DELETE FROM sessions WHERE last_active < ?`, cutoff)
if err != nil {
return 0, fmt.Errorf("cleanup expired sessions: %w", err)
}
n, _ := result.RowsAffected()
return n, nil
}
// StartCleanupTask 启动后台定时清理任务。
// ttlMinutes 为会话空闲超时(分钟),intervalMinutes 为清理间隔(分钟)。
// onCleanup 为清理完成后的回调(可传 nil),参数为清理的会话数量。
// 返回停止函数,调用以停止清理任务。
func (s *Store) StartCleanupTask(ttlMinutes, intervalMinutes int, onCleanup func(deleted int64)) func() {
ticker := time.NewTicker(time.Duration(intervalMinutes) * time.Minute)
stop := make(chan struct{})
go func() {
for {
select {
case <-ticker.C:
if n, err := s.CleanupExpired(ttlMinutes); err == nil && n > 0 {
if onCleanup != nil {
onCleanup(n)
}
}
case <-stop:
ticker.Stop()
return
}
}
}()
return func() { close(stop) }
}
func parseTime(s string) time.Time {
t, _ := time.Parse(time.RFC3339, s)
return t
+293
View File
@@ -0,0 +1,293 @@
package usage
import (
"database/sql"
"fmt"
"sync"
"time"
_ "github.com/mattn/go-sqlite3"
)
// AppUsage 记录单个应用的用量统计。
type AppUsage struct {
AppID string `json:"app_id"`
RequestCount int64 `json:"request_count"`
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
TotalTokens int64 `json:"total_tokens"`
ErrorCount int64 `json:"error_count"`
WindowStart string `json:"window_start"`
WindowEnd string `json:"window_end"`
}
// Tracker 管理 per-app 用量统计(内存滑动窗口 + SQLite 持久化)。
type Tracker struct {
mu sync.RWMutex
apps map[string]*appCounters
window time.Duration
db *sql.DB
logCh chan usageLogEntry // 异步写入 channel
stopCh chan struct{}
retentionDays int // 历史数据保留天数
}
type appCounters struct {
requestCount int64
inputTokens int64
outputTokens int64
errorCount int64
windowStart time.Time
}
type usageLogEntry struct {
appID string
inputTokens int
outputTokens int
isError bool
recordedAt time.Time
}
// NewTracker 创建一个用量统计器。
// window 为统计窗口时长(如 1 小时),窗口到期后自动重置。
// dbPath 为 SQLite 持久化路径,空字符串则仅用内存。
// retentionDays 为历史数据保留天数,超过此天数的数据自动清理(0 表示不清理)。
func NewTracker(window time.Duration, dbPath string, retentionDays int) (*Tracker, error) {
t := &Tracker{
apps: make(map[string]*appCounters),
window: window,
logCh: make(chan usageLogEntry, 1024),
stopCh: make(chan struct{}),
retentionDays: retentionDays,
}
if dbPath != "" {
db, err := sql.Open("sqlite3", dbPath)
if err != nil {
return nil, fmt.Errorf("open usage db: %w", err)
}
if err := initUsageDB(db); err != nil {
return nil, fmt.Errorf("init usage db: %w", err)
}
t.db = db
go t.logWriter() // 启动异步写入 goroutine
}
return t, nil
}
func initUsageDB(db *sql.DB) error {
schema := `
CREATE TABLE IF NOT EXISTS usage_logs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
app_id TEXT NOT NULL,
input_tokens INTEGER NOT NULL DEFAULT 0,
output_tokens INTEGER NOT NULL DEFAULT 0,
is_error INTEGER NOT NULL DEFAULT 0,
recorded_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_usage_app ON usage_logs(app_id);
CREATE INDEX IF NOT EXISTS idx_usage_time ON usage_logs(recorded_at);`
_, err := db.Exec(schema)
return err
}
// Record 记录一次请求的用量。
func (t *Tracker) Record(appID string, inputTokens, outputTokens int, isError bool) {
now := time.Now()
// 更新内存计数器
t.mu.Lock()
c, ok := t.apps[appID]
if !ok {
c = &appCounters{windowStart: now}
t.apps[appID] = c
}
// 窗口过期则重置
if now.Sub(c.windowStart) > t.window {
c.requestCount = 0
c.inputTokens = 0
c.outputTokens = 0
c.errorCount = 0
c.windowStart = now
}
c.requestCount++
c.inputTokens += int64(inputTokens)
c.outputTokens += int64(outputTokens)
if isError {
c.errorCount++
}
t.mu.Unlock()
// 异步写入 SQLite(非阻塞,channel 满则丢弃)
if t.db != nil {
select {
case t.logCh <- usageLogEntry{appID, inputTokens, outputTokens, isError, now}:
default: // channel 满则丢弃,避免阻塞主流程
}
}
}
// logWriter 后台 goroutine,从 channel 读取用量日志并批量写入 SQLite。
func (t *Tracker) logWriter() {
batch := make([]usageLogEntry, 0, 64)
ticker := time.NewTicker(5 * time.Second)
defer ticker.Stop()
flush := func() {
if len(batch) == 0 {
return
}
tx, err := t.db.Begin()
if err != nil {
batch = batch[:0]
return
}
stmt, _ := tx.Prepare(`INSERT INTO usage_logs (app_id, input_tokens, output_tokens, is_error, recorded_at) VALUES (?, ?, ?, ?, ?)`)
if stmt != nil {
for _, e := range batch {
errInt := 0
if e.isError {
errInt = 1
}
stmt.Exec(e.appID, e.inputTokens, e.outputTokens, errInt, e.recordedAt.Format(time.RFC3339))
}
stmt.Close()
}
tx.Commit()
batch = batch[:0]
}
for {
select {
case e := <-t.logCh:
batch = append(batch, e)
if len(batch) >= 64 {
flush()
}
case <-ticker.C:
flush()
// 定期清理过期数据
if t.retentionDays > 0 {
t.cleanupOldLogs()
}
case <-t.stopCh:
// 排空 channel 后 flush
for len(t.logCh) > 0 {
batch = append(batch, <-t.logCh)
}
flush()
return
}
}
}
// Get 返回指定应用的当前用量快照。
func (t *Tracker) Get(appID string) AppUsage {
t.mu.RLock()
defer t.mu.RUnlock()
c, ok := t.apps[appID]
if !ok {
return AppUsage{AppID: appID}
}
return AppUsage{
AppID: appID,
RequestCount: c.requestCount,
InputTokens: c.inputTokens,
OutputTokens: c.outputTokens,
TotalTokens: c.inputTokens + c.outputTokens,
ErrorCount: c.errorCount,
WindowStart: c.windowStart.Format(time.RFC3339),
WindowEnd: c.windowStart.Add(t.window).Format(time.RFC3339),
}
}
// GetAll 返回所有应用的用量快照。
func (t *Tracker) GetAll() []AppUsage {
t.mu.RLock()
defer t.mu.RUnlock()
result := make([]AppUsage, 0, len(t.apps))
for appID, c := range t.apps {
result = append(result, AppUsage{
AppID: appID,
RequestCount: c.requestCount,
InputTokens: c.inputTokens,
OutputTokens: c.outputTokens,
TotalTokens: c.inputTokens + c.outputTokens,
ErrorCount: c.errorCount,
WindowStart: c.windowStart.Format(time.RFC3339),
WindowEnd: c.windowStart.Add(t.window).Format(time.RFC3339),
})
}
return result
}
// GetHistory 返回指定应用的历史用量(从 SQLite 查询,按小时聚合)。
// hours 为查询的时间范围(最近 N 小时)。
func (t *Tracker) GetHistory(appID string, hours int) ([]map[string]any, error) {
if t.db == nil {
return nil, fmt.Errorf("persistence not enabled")
}
if hours <= 0 {
hours = 24
}
since := time.Now().Add(-time.Duration(hours) * time.Hour).Format(time.RFC3339)
rows, err := t.db.Query(
`SELECT
substr(recorded_at, 1, 13) as hour_bucket,
COUNT(*) as request_count,
SUM(input_tokens) as input_tokens,
SUM(output_tokens) as output_tokens,
SUM(is_error) as error_count
FROM usage_logs
WHERE app_id = ? AND recorded_at >= ?
GROUP BY hour_bucket
ORDER BY hour_bucket DESC`,
appID, since,
)
if err != nil {
return nil, fmt.Errorf("query usage history: %w", err)
}
defer rows.Close()
var result []map[string]any
for rows.Next() {
var hourBucket string
var reqCount, inputTokens, outputTokens, errCount int64
if err := rows.Scan(&hourBucket, &reqCount, &inputTokens, &outputTokens, &errCount); err != nil {
return nil, err
}
result = append(result, map[string]any{
"hour": hourBucket,
"request_count": reqCount,
"input_tokens": inputTokens,
"output_tokens": outputTokens,
"total_tokens": inputTokens + outputTokens,
"error_count": errCount,
})
}
return result, rows.Err()
}
// Close 关闭数据库连接。
func (t *Tracker) Close() error {
close(t.stopCh)
if t.db != nil {
return t.db.Close()
}
return nil
}
// cleanupOldLogs 清理超过保留期限的用量日志。
func (t *Tracker) cleanupOldLogs() {
cutoff := time.Now().Add(-time.Duration(t.retentionDays) * 24 * time.Hour).Format(time.RFC3339)
if _, err := t.db.Exec(`DELETE FROM usage_logs WHERE recorded_at < ?`, cutoff); err != nil {
// 清理失败不影响主流程,下次再试
return
}
}
+39 -36
View File
@@ -2,19 +2,19 @@ package api
// ChatRequest is the request body for POST /v1/chat/completions.
type ChatRequest struct {
Model string `json:"model"`
Messages []Message `json:"messages"`
Stream bool `json:"stream,omitempty"`
SessionID string `json:"session_id,omitempty"`
IdempotencyKey string `json:"idempotency_key,omitempty"`
Priority string `json:"priority,omitempty"`
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
ContextPolicy string `json:"context_policy,omitempty"`
Timeouts *RequestTimeouts `json:"timeouts,omitempty"`
Routing *RoutingOptions `json:"routing,omitempty"`
Metadata map[string]string `json:"metadata,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
Model string `json:"model"`
Messages []Message `json:"messages"`
Stream bool `json:"stream,omitempty"`
SessionID string `json:"session_id,omitempty"`
IdempotencyKey string `json:"idempotency_key,omitempty"`
Priority string `json:"priority,omitempty"`
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
ContextPolicy string `json:"context_policy,omitempty"`
Timeouts *RequestTimeouts `json:"timeouts,omitempty"`
Routing *RoutingOptions `json:"routing,omitempty"`
Metadata map[string]string `json:"metadata,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
}
type Message struct {
@@ -36,18 +36,18 @@ type RoutingOptions struct {
// ChatResponse is the non-streaming response.
type ChatResponse struct {
RequestID string `json:"request_id"`
TaskID string `json:"task_id"`
SessionID string `json:"session_id,omitempty"`
Status string `json:"status"`
Model string `json:"model"`
Choices []Choice `json:"choices"`
LogicalModel string `json:"logical_model"`
ActualModel string `json:"actual_model"`
NodeID string `json:"node_id,omitempty"`
Usage *Usage `json:"usage,omitempty"`
Timing *Timing `json:"timing,omitempty"`
Degraded bool `json:"degraded,omitempty"`
RequestID string `json:"request_id"`
TaskID string `json:"task_id"`
SessionID string `json:"session_id,omitempty"`
Status string `json:"status"`
Model string `json:"model"`
Choices []Choice `json:"choices"`
LogicalModel string `json:"logical_model"`
ActualModel string `json:"actual_model"`
NodeID string `json:"node_id,omitempty"`
Usage *Usage `json:"usage,omitempty"`
Timing *Timing `json:"timing,omitempty"`
Degraded bool `json:"degraded,omitempty"`
}
type Choice struct {
@@ -88,23 +88,26 @@ type ModelListResponse struct {
}
type ModelInfo struct {
ID string `json:"id"`
Object string `json:"object"`
OwnedBy string `json:"owned_by"`
ID string `json:"id"`
Object string `json:"object"`
OwnedBy string `json:"owned_by"`
Provider string `json:"provider,omitempty"`
ContextWindow int `json:"context_window,omitempty"`
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
}
// SessionRequest is the request body for POST /v1/sessions.
type SessionRequest struct {
ApplicationID string `json:"application_id"`
UserID string `json:"user_id,omitempty"`
Config map[string]any `json:"config,omitempty"`
ApplicationID string `json:"application_id"`
UserID string `json:"user_id,omitempty"`
Config map[string]any `json:"config,omitempty"`
}
// SessionResponse is the response for session operations.
type SessionResponse struct {
SessionID string `json:"session_id"`
ApplicationID string `json:"application_id"`
UserID string `json:"user_id,omitempty"`
CreatedAt string `json:"created_at"`
LastActive string `json:"last_active"`
SessionID string `json:"session_id"`
ApplicationID string `json:"application_id"`
UserID string `json:"user_id,omitempty"`
CreatedAt string `json:"created_at"`
LastActive string `json:"last_active"`
}