feat: 十轮网关优化 - 安全加固/可观测性/性能/可靠性
- 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:
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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()),
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -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...)
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 设置熔断器状态 gauge(0=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)
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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/keys(GET 列出所有 Key,POST 创建新 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/usage(GET 返回所有 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-breaker(GET 返回熔断器状态)。
|
||||
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/scheduler(GET 返回调度器状态)。
|
||||
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/reload(POST 触发配置热重载)。
|
||||
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
@@ -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) {
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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"`
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user