diff --git a/cmd/createkey/main.go b/cmd/createkey/main.go new file mode 100644 index 0000000..3a18668 --- /dev/null +++ b/cmd/createkey/main.go @@ -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) +} diff --git a/configs/config.yaml b/configs/config.yaml index 36a5a05..3dedaba 100644 --- a/configs/config.yaml +++ b/configs/config.yaml @@ -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 diff --git a/gateway b/gateway new file mode 100755 index 0000000..071b5cd Binary files /dev/null and b/gateway differ diff --git a/go.mod b/go.mod index 59a96fe..8900873 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index bb2790f..56ace82 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/internal/adapter/adapter.go b/internal/adapter/adapter.go index d0017d4..cf228b3 100644 --- a/internal/adapter/adapter.go +++ b/internal/adapter/adapter.go @@ -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) diff --git a/internal/adapter/ollama.go b/internal/adapter/ollama.go index 6148bff..d60b5f5 100644 --- a/internal/adapter/ollama.go +++ b/internal/adapter/ollama.go @@ -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 { diff --git a/internal/adapter/vllm.go b/internal/adapter/vllm.go index cddceac..37f352d 100644 --- a/internal/adapter/vllm.go +++ b/internal/adapter/vllm.go @@ -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 { diff --git a/internal/auth/auth.go b/internal/auth/auth.go index ad13d75..421ee73 100644 --- a/internal/auth/auth.go +++ b/internal/auth/auth.go @@ -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 diff --git a/internal/config/config.go b/internal/config/config.go index 5800e48..d4c3904 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -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 } diff --git a/internal/context/adapter_summary.go b/internal/context/adapter_summary.go new file mode 100644 index 0000000..1355d2b --- /dev/null +++ b/internal/context/adapter_summary.go @@ -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()), + }, + } +} diff --git a/internal/context/assembler.go b/internal/context/assembler.go index 3676b74..60dd8dc 100644 --- a/internal/context/assembler.go +++ b/internal/context/assembler.go @@ -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...) diff --git a/internal/context/summarizer.go b/internal/context/summarizer.go new file mode 100644 index 0000000..65febc9 --- /dev/null +++ b/internal/context/summarizer.go @@ -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)) +} diff --git a/internal/handler/sse.go b/internal/handler/sse.go index c6fa6a2..930c3ad 100644 --- a/internal/handler/sse.go +++ b/internal/handler/sse.go @@ -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, diff --git a/internal/middleware/middleware.go b/internal/middleware/middleware.go index 5aa286d..015389b 100644 --- a/internal/middleware/middleware.go +++ b/internal/middleware/middleware.go @@ -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) + }) +} diff --git a/internal/observability/audit.go b/internal/observability/audit.go new file mode 100644 index 0000000..aef5107 --- /dev/null +++ b/internal/observability/audit.go @@ -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 +} diff --git a/internal/observability/metrics.go b/internal/observability/metrics.go index bb33e4e..46d543d 100644 --- a/internal/observability/metrics.go +++ b/internal/observability/metrics.go @@ -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) } diff --git a/internal/ratelimit/limiter.go b/internal/ratelimit/limiter.go new file mode 100644 index 0000000..5472707 --- /dev/null +++ b/internal/ratelimit/limiter.go @@ -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() + } +} diff --git a/internal/scheduler/circuit_breaker.go b/internal/scheduler/circuit_breaker.go index efeaaa6..7a9da54 100644 --- a/internal/scheduler/circuit_breaker.go +++ b/internal/scheduler/circuit_breaker.go @@ -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) diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index 1207129..5643193 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -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() diff --git a/internal/server/admin_handlers.go b/internal/server/admin_handlers.go new file mode 100644 index 0000000..c883a25 --- /dev/null +++ b/internal/server/admin_handlers.go @@ -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": "配置已热重载，模型映射已更新", + }) +} diff --git a/internal/server/handlers.go b/internal/server/handlers.go index 3565d53..d9f0f43 100644 --- a/internal/server/handlers.go +++ b/internal/server/handlers.go @@ -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) { diff --git a/internal/server/idempotency.go b/internal/server/idempotency.go new file mode 100644 index 0000000..b133898 --- /dev/null +++ b/internal/server/idempotency.go @@ -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() + } +} diff --git a/internal/server/server.go b/internal/server/server.go index 5c78fca..9e88e1e 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -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://") diff --git a/internal/session/session.go b/internal/session/session.go index ff69f30..b86f246 100644 --- a/internal/session/session.go +++ b/internal/session/session.go @@ -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 diff --git a/internal/usage/tracker.go b/internal/usage/tracker.go new file mode 100644 index 0000000..ca3597a --- /dev/null +++ b/internal/usage/tracker.go @@ -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 + } +} diff --git a/pkg/api/types.go b/pkg/api/types.go index 8f13837..c415b10 100644 --- a/pkg/api/types.go +++ b/pkg/api/types.go @@ -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"` }