Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 21 additions & 2 deletions admin/api_key_account_usage_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,13 @@ func TestAggregateAPIKeyAccountGroups(t *testing.T) {
AccountID: 2, Requests: 3, TotalTokens: 200, AccountBilled: 0.4, UserBilled: 0.6,
Groups: []database.APIKeyAccountGroup{{ID: 10, Name: "primary"}, {ID: 20, Name: "shared"}},
},
{
AccountID: 3, Requests: 1, TotalTokens: 50, AccountBilled: 0.1, UserBilled: 0.15,
},
}

groups, summary := aggregateAPIKeyAccountGroups(items)
if summary.Accounts != 2 || summary.Requests != 5 || summary.TotalTokens != 300 || math.Abs(summary.AccountBilled-0.6) > 1e-9 {
groups, summary, reconciliation := aggregateAPIKeyAccountGroups(items)
if summary.Accounts != 3 || summary.Requests != 6 || summary.TotalTokens != 350 || math.Abs(summary.AccountBilled-0.7) > 1e-9 || math.Abs(summary.UserBilled-1.05) > 1e-9 {
t.Fatalf("summary = %+v", summary)
}
if len(groups) != 2 {
Expand All @@ -32,4 +35,20 @@ func TestAggregateAPIKeyAccountGroups(t *testing.T) {
if groups[1].ID != 20 || groups[1].Accounts != 1 || groups[1].TotalTokens != 200 {
t.Fatalf("shared group = %+v", groups[1])
}
if reconciliation.UniqueGroupedAccounts != 2 || reconciliation.MultiGroupAccounts != 1 {
t.Fatalf("reconciliation account counts = %+v", reconciliation)
}
if reconciliation.GroupedTotal.Accounts != 3 || reconciliation.GroupedTotal.Requests != 8 || reconciliation.GroupedTotal.TotalTokens != 500 || math.Abs(reconciliation.GroupedTotal.UserBilled-1.5) > 1e-9 {
t.Fatalf("grouped total = %+v", reconciliation.GroupedTotal)
}
if reconciliation.Ungrouped.Accounts != 1 || reconciliation.Ungrouped.TotalTokens != 50 || math.Abs(reconciliation.Ungrouped.UserBilled-0.15) > 1e-9 {
t.Fatalf("ungrouped = %+v", reconciliation.Ungrouped)
}
if reconciliation.Duplicate.Accounts != 1 || reconciliation.Duplicate.Requests != 3 || reconciliation.Duplicate.TotalTokens != 200 || math.Abs(reconciliation.Duplicate.UserBilled-0.6) > 1e-9 {
t.Fatalf("duplicate = %+v", reconciliation.Duplicate)
}
reconciled := reconciliation.GroupedTotal.UserBilled + reconciliation.Ungrouped.UserBilled - reconciliation.Duplicate.UserBilled
if math.Abs(reconciled-summary.UserBilled) > 1e-9 {
t.Fatalf("reconciled billed = %f, summary = %f", reconciled, summary.UserBilled)
}
}
50 changes: 43 additions & 7 deletions admin/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -713,6 +713,8 @@ func (h *Handler) RegisterRoutes(r *gin.Engine) {
api.DELETE("/prompt-policy/risk-profiles/:subject_type/:subject_key/trust", h.RevokePromptRiskTrustPolicy)
api.POST("/prompt-policy/conversation-locks/:lock_key/unlock", h.UnlockPromptConversation)
api.POST("/prompt-filter/test", h.TestPromptFilter)
api.GET("/prompt-filter/review/keys", h.ListPromptReviewAPIKeys)
api.DELETE("/prompt-filter/review/keys/:key_id", h.DeletePromptReviewAPIKey)
api.POST("/prompt-filter/review/test", h.TestPromptReviewConnection)
api.POST("/prompt-filter/rules/test", h.TestPromptFilterRulePattern)
api.GET("/prompt-filter/rules", h.GetPromptFilterRules)
Expand Down Expand Up @@ -6047,10 +6049,11 @@ func (h *Handler) GetAPIKeyAccountStats(c *gin.Context) {

cacheKey := fmt.Sprintf("%d:%d:%d", id, rangeStart.Unix()/30, rangeEnd.Unix()/30)
type cachedResponse struct {
Items []database.APIKeyAccountStat `json:"items"`
Groups []apiKeyAccountGroupUsage `json:"groups"`
Summary apiKeyAccountUsageSummary `json:"summary"`
MembershipBasis string `json:"membership_basis"`
Items []database.APIKeyAccountStat `json:"items"`
Groups []apiKeyAccountGroupUsage `json:"groups"`
Summary apiKeyAccountUsageSummary `json:"summary"`
Reconciliation apiKeyAccountUsageReconciliation `json:"reconciliation"`
MembershipBasis string `json:"membership_basis"`
}
var response cachedResponse
if h.getRuntimeJSON(ctx, adminAPIKeyAccountsNamespace, cacheKey, &response) {
Expand All @@ -6067,7 +6070,7 @@ func (h *Handler) GetAPIKeyAccountStats(c *gin.Context) {
items = []database.APIKeyAccountStat{}
}
response.Items = items
response.Groups, response.Summary = aggregateAPIKeyAccountGroups(items)
response.Groups, response.Summary, response.Reconciliation = aggregateAPIKeyAccountGroups(items)
response.MembershipBasis = "current_and_deleted_last_membership"
h.setRuntimeJSON(ctx, adminAPIKeyAccountsNamespace, cacheKey, response, adminUsageRangeCacheTTL)
c.JSON(http.StatusOK, response)
Expand All @@ -6092,18 +6095,46 @@ type apiKeyAccountGroupUsage struct {
UserBilled float64 `json:"user_billed"`
}

type apiKeyAccountUsageReconciliation struct {
GroupedTotal apiKeyAccountUsageSummary `json:"grouped_total"`
Ungrouped apiKeyAccountUsageSummary `json:"ungrouped"`
Duplicate apiKeyAccountUsageSummary `json:"duplicate"`
UniqueGroupedAccounts int `json:"unique_grouped_accounts"`
MultiGroupAccounts int `json:"multi_group_accounts"`
}

// aggregateAPIKeyAccountGroups uses current memberships for active accounts and
// the retained last membership for recycle-bin accounts. If an account belongs
// to multiple groups, its usage is intentionally included in each group; the
// overall summary remains de-duplicated.
func aggregateAPIKeyAccountGroups(items []database.APIKeyAccountStat) ([]apiKeyAccountGroupUsage, apiKeyAccountUsageSummary) {
func aggregateAPIKeyAccountGroups(items []database.APIKeyAccountStat) ([]apiKeyAccountGroupUsage, apiKeyAccountUsageSummary, apiKeyAccountUsageReconciliation) {
groupMap := make(map[int64]*apiKeyAccountGroupUsage)
summary := apiKeyAccountUsageSummary{Accounts: len(items)}
reconciliation := apiKeyAccountUsageReconciliation{}
for _, item := range items {
summary.Requests += item.Requests
summary.TotalTokens += item.TotalTokens
summary.AccountBilled += item.AccountBilled
summary.UserBilled += item.UserBilled
groupCount := len(item.Groups)
if groupCount == 0 {
reconciliation.Ungrouped.Accounts++
reconciliation.Ungrouped.Requests += item.Requests
reconciliation.Ungrouped.TotalTokens += item.TotalTokens
reconciliation.Ungrouped.AccountBilled += item.AccountBilled
reconciliation.Ungrouped.UserBilled += item.UserBilled
} else {
reconciliation.UniqueGroupedAccounts++
}
if groupCount > 1 {
reconciliation.MultiGroupAccounts++
extraAssignments := int64(groupCount - 1)
reconciliation.Duplicate.Accounts += groupCount - 1
reconciliation.Duplicate.Requests += item.Requests * extraAssignments
reconciliation.Duplicate.TotalTokens += item.TotalTokens * extraAssignments
reconciliation.Duplicate.AccountBilled += item.AccountBilled * float64(extraAssignments)
reconciliation.Duplicate.UserBilled += item.UserBilled * float64(extraAssignments)
}
for _, group := range item.Groups {
total := groupMap[group.ID]
if total == nil {
Expand All @@ -6120,14 +6151,19 @@ func aggregateAPIKeyAccountGroups(items []database.APIKeyAccountStat) ([]apiKeyA
groups := make([]apiKeyAccountGroupUsage, 0, len(groupMap))
for _, group := range groupMap {
groups = append(groups, *group)
reconciliation.GroupedTotal.Accounts += group.Accounts
reconciliation.GroupedTotal.Requests += group.Requests
reconciliation.GroupedTotal.TotalTokens += group.TotalTokens
reconciliation.GroupedTotal.AccountBilled += group.AccountBilled
reconciliation.GroupedTotal.UserBilled += group.UserBilled
}
sort.Slice(groups, func(i, j int) bool {
if groups[i].UserBilled == groups[j].UserBilled {
return groups[i].TotalTokens > groups[j].TotalTokens
}
return groups[i].UserBilled > groups[j].UserBilled
})
return groups, summary
return groups, summary, reconciliation
}

// GetChartData 返回图表聚合数据(服务端分桶 + 内存缓存)
Expand Down
125 changes: 124 additions & 1 deletion admin/prompt_filter.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,9 @@ package admin

import (
"context"
"crypto/sha256"
"database/sql"
"encoding/hex"
"encoding/json"
"errors"
"net/http"
Expand Down Expand Up @@ -74,6 +76,8 @@ type promptReviewTestRequest struct {

type promptReviewKeyTestResult struct {
KeyIndex int `json:"key_index"`
KeyID string `json:"key_id,omitempty"`
KeyMasked string `json:"key_masked,omitempty"`
OK bool `json:"ok"`
Endpoint string `json:"endpoint,omitempty"`
Model string `json:"model,omitempty"`
Expand All @@ -90,6 +94,44 @@ type promptReviewKeyTestResult struct {
Error string `json:"error,omitempty"`
}

type promptReviewAPIKeyDescriptor struct {
ID string `json:"id"`
Index int `json:"index"`
Masked string `json:"masked"`
}

type promptReviewAPIKeysResponse struct {
Items []promptReviewAPIKeyDescriptor `json:"items"`
Count int `json:"count"`
}

func promptReviewAPIKeyID(key string) string {
sum := sha256.Sum256([]byte(strings.TrimSpace(key)))
return hex.EncodeToString(sum[:])
}

func maskPromptReviewAPIKey(key string) string {
key = strings.TrimSpace(key)
if len(key) <= 4 {
return "••••"
}
prefix := ""
if len(key) >= 3 {
prefix = key[:3]
}
return prefix + "••••" + key[len(key)-4:]
Comment on lines +113 to +122

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔒 Security & Privacy | 🟠 Major | ⚡ Quick win

Prevent short key disclosure.

For keys with 5 to 7 characters, the three-character prefix and four-character suffix overlap. The response then exposes every character of the key.

Return a fully masked value when the visible prefix and suffix would overlap. Add tests for keys with lengths 5 through 7.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@admin/prompt_filter.go` around lines 113 - 122, Update maskPromptReviewAPIKey
so keys of length 5 through 7 return the fully masked value instead of exposing
overlapping prefix and suffix characters. Preserve the existing masking behavior
for longer keys and add tests covering each length from 5 through 7.

}

func promptReviewAPIKeyDescriptors(keys []string) []promptReviewAPIKeyDescriptor {
items := make([]promptReviewAPIKeyDescriptor, 0, len(keys))
for index, key := range keys {
items = append(items, promptReviewAPIKeyDescriptor{
ID: promptReviewAPIKeyID(key), Index: index + 1, Masked: maskPromptReviewAPIKey(key),
})
}
return items
}

type promptReviewTestResponse struct {
OK bool `json:"ok"`
Endpoint string `json:"endpoint"`
Expand Down Expand Up @@ -535,6 +577,7 @@ func (h *Handler) TestPromptReviewConnection(c *gin.Context) {
}(index, key)
}
results := make([]promptReviewKeyTestResult, len(keys))
descriptors := promptReviewAPIKeyDescriptors(keys)
allOK := true
var first promptfilter.ReviewOutcome
for range keys {
Expand All @@ -543,7 +586,8 @@ func (h *Handler) TestPromptReviewConnection(c *gin.Context) {
first = item.outcome
}
result := promptReviewKeyTestResult{
KeyIndex: item.index + 1, OK: item.err == nil, Flagged: item.outcome.Flagged,
KeyIndex: item.index + 1, KeyID: descriptors[item.index].ID, KeyMasked: descriptors[item.index].Masked,
OK: item.err == nil, Flagged: item.outcome.Flagged,
Endpoint: item.outcome.Endpoint, Model: item.outcome.Model, Confidence: item.outcome.Confidence,
Reason: item.outcome.Reason, HighestCategory: item.outcome.HighestCategory,
DecisionCategory: item.outcome.DecisionCategory, DecisionScore: item.outcome.DecisionScore,
Expand Down Expand Up @@ -593,6 +637,85 @@ func (h *Handler) TestPromptReviewConnection(c *gin.Context) {
})
}

func (h *Handler) ListPromptReviewAPIKeys(c *gin.Context) {
if h == nil || h.store == nil {
writeError(c, http.StatusServiceUnavailable, "Prompt 审核配置不可用")
return
}
items := promptReviewAPIKeyDescriptors(h.store.GetPromptFilterConfig().Review.APIKeyList())
c.JSON(http.StatusOK, promptReviewAPIKeysResponse{Items: items, Count: len(items)})
}

func (h *Handler) DeletePromptReviewAPIKey(c *gin.Context) {
if h == nil || h.store == nil || h.db == nil {
writeError(c, http.StatusServiceUnavailable, "Prompt 审核配置不可用")
return
}
keyID := strings.ToLower(strings.TrimSpace(c.Param("key_id")))
if len(keyID) != sha256.Size*2 {
writeError(c, http.StatusBadRequest, "审查 Key 标识无效")
return
}
if _, err := hex.DecodeString(keyID); err != nil {
writeError(c, http.StatusBadRequest, "审查 Key 标识无效")
return
}

h.settingsUpdateMu.Lock()
defer h.settingsUpdateMu.Unlock()
settings, err := h.db.GetSystemSettings(c.Request.Context())
if err != nil {
writeInternalError(c, err)
return
}
if settings == nil {
writeError(c, http.StatusNotFound, "审查 Key 不存在")
return
}
currentRaw := strings.TrimSpace(settings.PromptFilterReviewAPIKey)
keys := (promptfilter.ReviewConfig{APIKey: currentRaw}).APIKeyList()
remaining := make([]string, 0, len(keys))
found := false
for _, key := range keys {
if promptReviewAPIKeyID(key) == keyID {
found = true
continue
}
remaining = append(remaining, key)
}
if !found {
writeError(c, http.StatusNotFound, "审查 Key 不存在或已被删除")
return
}
if settings.PromptFilterReviewEnabled && len(remaining) == 0 {
writeError(c, http.StatusConflict, "模型复核启用时不能删除最后一个审查 Key,请先关闭模型复核或添加替代 Key")
return
}
replacement := strings.Join(remaining, "\n")
runtimeCfg := h.store.GetPromptFilterConfig()
runtimeCfg.Review.APIKey = replacement
runtimeCfg = promptfilter.NormalizeConfig(runtimeCfg)
if err := promptfilter.ValidateReviewConfig(runtimeCfg.Review); err != nil {
writeError(c, http.StatusConflict, "删除后审查配置无效: "+err.Error())
return
}
swapped, err := h.db.CompareAndSwapPromptFilterReviewAPIKeys(c.Request.Context(), currentRaw, replacement)
if err != nil {
writeInternalError(c, err)
return
}
if !swapped {
writeError(c, http.StatusConflict, "审查 Key 列表已被其他操作修改,请刷新后重试")
return
}
if err := h.store.SetPromptFilterConfigWithAdvancedRaw(runtimeCfg, h.store.GetPromptFilterAdvancedConfig()); err != nil {
writeError(c, http.StatusInternalServerError, "审查 Key 已保存,但运行时配置更新失败")
return
Comment on lines +711 to +713

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔒 Security & Privacy | 🟠 Major | 🏗️ Heavy lift

Keep persisted and runtime key state consistent.

If SetPromptFilterConfigWithAdvancedRaw fails after the compare-and-swap succeeds, the database no longer contains the key but the running service continues to use it. This is critical when the deletion revokes a compromised credential.

Roll back the compare-and-swap before returning an error, or make the persistence and runtime update one atomic operation. Add a failure-path test.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@admin/prompt_filter.go` around lines 711 - 713, Update the handler around
SetPromptFilterConfigWithAdvancedRaw so a runtime configuration failure after
the compare-and-swap restores the previously persisted prompt-filter key state
before returning the error. Prefer rolling back the persisted change using the
prior configuration, or make both updates atomic, and add a test covering this
failure path.

}
items := promptReviewAPIKeyDescriptors(remaining)
c.JSON(http.StatusOK, promptReviewAPIKeysResponse{Items: items, Count: len(items)})
}

func (h *Handler) TestPromptFilterRulePattern(c *gin.Context) {
var req promptFilterRulePatternTestRequest
if err := c.ShouldBindJSON(&req); err != nil {
Expand Down
66 changes: 66 additions & 0 deletions admin/prompt_filter_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"encoding/json"
"net/http"
"net/http/httptest"
"path/filepath"
"strings"
"sync"
"sync/atomic"
Expand Down Expand Up @@ -131,6 +132,71 @@ func TestPromptReviewConnectionTestsAllKeysConcurrentlyWithoutReturningSecrets(t
t.Fatalf("response leaked key: %s", recorder.Body.String())
}
}
for index, result := range response.Results {
if result.KeyID == "" || result.KeyMasked == "" || result.KeyIndex != index+1 {
t.Fatalf("result %d has no stable redacted identity: %+v", index, result)
}
}
}

func TestDeletePromptReviewAPIKeyRemovesOnlySelectedKey(t *testing.T) {
gin.SetMode(gin.TestMode)
db, err := database.New("sqlite", filepath.Join(t.TempDir(), "review-keys.db"))
if err != nil {
t.Fatalf("New(sqlite): %v", err)
}
defer db.Close()
settings := &database.SystemSettings{
PromptFilterReviewEnabled: true, PromptFilterReviewAPIKey: "key-one\nkey-two",
PromptFilterReviewBaseURL: "https://review.example.com", PromptFilterReviewModel: "review-model",
}
if err := db.UpdateSystemSettings(context.Background(), settings); err != nil {
t.Fatalf("UpdateSystemSettings: %v", err)
}
store := auth.NewStore(nil, nil, settings)
t.Cleanup(store.Stop)
handler := &Handler{db: db, store: store}

deleteKey := func(key string) *httptest.ResponseRecorder {
t.Helper()
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Params = gin.Params{{Key: "key_id", Value: promptReviewAPIKeyID(key)}}
c.Request = httptest.NewRequest(http.MethodDelete, "/api/admin/prompt-filter/review/keys/selected", nil)
handler.DeletePromptReviewAPIKey(c)
return recorder
}

if recorder := deleteKey("key-one"); recorder.Code != http.StatusOK {
t.Fatalf("delete selected status=%d body=%s", recorder.Code, recorder.Body.String())
}
persisted, err := db.GetSystemSettings(context.Background())
if err != nil || persisted.PromptFilterReviewAPIKey != "key-two" {
t.Fatalf("persisted keys=%q err=%v, want key-two", persisted.PromptFilterReviewAPIKey, err)
}
if got := store.GetPromptFilterConfig().Review.APIKeyList(); len(got) != 1 || got[0] != "key-two" {
t.Fatalf("runtime keys=%v, want [key-two]", got)
}
if recorder := deleteKey("key-two"); recorder.Code != http.StatusConflict {
t.Fatalf("delete last enabled key status=%d body=%s", recorder.Code, recorder.Body.String())
}
}

func TestPromptReviewAPIKeyDescriptorsNeverExposeSecrets(t *testing.T) {
keys := []string{"sk-secret-alpha-1234", "opaque-secret-beta-9876"}
items := promptReviewAPIKeyDescriptors(keys)
if len(items) != len(keys) || items[0].ID == items[1].ID {
t.Fatalf("descriptors=%+v", items)
}
encoded, err := json.Marshal(items)
if err != nil {
t.Fatalf("Marshal: %v", err)
}
for _, key := range keys {
if strings.Contains(string(encoded), key) {
t.Fatalf("descriptor response leaked key %q: %s", key, encoded)
}
}
}

func TestPromptFilterTestEndpointUsesRealGuardPipelineMetadata(t *testing.T) {
Expand Down
Loading
Loading