diff --git a/internal/billing/collector.go b/internal/billing/collector.go index 0882fac..589600b 100644 --- a/internal/billing/collector.go +++ b/internal/billing/collector.go @@ -9,6 +9,7 @@ import ( "io" "math" "net/http" + "strconv" "strings" "time" "unicode/utf8" @@ -166,6 +167,8 @@ func fetchSub2API(ctx context.Context, item *upstream.Upstream) (Result, error) type newAPIUsageResponse struct { Data struct { Object string `json:"object"` + Group string `json:"group"` + UserGroup string `json:"user_group"` TotalUsed float64 `json:"total_used"` TotalAvailable float64 `json:"total_available"` UnlimitedQuota bool `json:"unlimited_quota"` @@ -199,15 +202,63 @@ type newAPIGroupResponse struct { } `json:"data"` } +// newAPIUserResponse 是 OneAPI/New API 的当前用户信息。分组日志可能因 +// 限流、保留期或权限暂时不可用，但 self 接口仍能提供当前分组名称。 +type newAPIUserResponse struct { + Data struct { + Group string `json:"group"` + User struct { + Group string `json:"group"` + } `json:"user"` + } `json:"data"` +} + func decodeNewAPILogBilling(raw json.RawMessage, target *newAPILogBilling) bool { if len(raw) == 0 { return false } - if json.Unmarshal(raw, target) == nil { - return true - } var encoded string - return json.Unmarshal(raw, &encoded) == nil && json.Unmarshal([]byte(encoded), target) == nil + if json.Unmarshal(raw, &encoded) == nil { + raw = json.RawMessage(encoded) + } + var fields map[string]json.RawMessage + if json.Unmarshal(raw, &fields) != nil { + return false + } + target.GroupRatio = jsonFloat(fields["group_ratio"]) + target.UserGroupRatio = jsonFloat(fields["user_group_ratio"]) + return target.GroupRatio != nil || target.UserGroupRatio != nil +} + +func jsonFloat(raw json.RawMessage) *float64 { + if len(raw) == 0 { + return nil + } + var value float64 + if json.Unmarshal(raw, &value) == nil && math.IsNaN(value) == false && math.IsInf(value, 0) == false { + return &value + } + var text string + if json.Unmarshal(raw, &text) != nil { + return nil + } + value, err := strconv.ParseFloat(strings.TrimSpace(text), 64) + if err != nil || math.IsNaN(value) || math.IsInf(value, 0) { + return nil + } + return &value +} + +func fetchNewAPIUserGroup(ctx context.Context, item *upstream.Upstream) (string, error) { + var profile newAPIUserResponse + if err := getJSON(ctx, item, "/api/user/self", &profile); err != nil { + return "", err + } + group := strings.TrimSpace(profile.Data.Group) + if group == "" { + group = strings.TrimSpace(profile.Data.User.Group) + } + return group, nil } func fetchNewAPI(ctx context.Context, item *upstream.Upstream) (Result, error) { @@ -233,11 +284,18 @@ func fetchNewAPI(ctx context.Context, item *upstream.Upstream) (Result, error) { } actual := usage.Data.TotalUsed / status.Data.QuotaPerUnit result.ReportedActualCost = &actual + // Some New API forks include the token's selected group in the usage + // response. It is available with the relay key even when token logs are + // rate-limited, so prefer it as the first fallback for group detection. + result.BillingGroup = strings.TrimSpace(usage.Data.Group) + if result.BillingGroup == "" { + result.BillingGroup = strings.TrimSpace(usage.Data.UserGroup) + } var logs newAPILogResponse - if err := getJSON(ctx, item, "/api/log/token", &logs); err != nil { - result.Warning = err.Error() - return result, nil + logErr := getJSON(ctx, item, "/api/log/token", &logs) + if logErr != nil { + result.Warning = logErr.Error() } // 分组名取最新一条日志（错误日志也带 group，且反映当前分组归属）； // user_group_ratio 是**个人议价倍率**(>=0 时才是真实扣费)，取最新扣费日志里的。 @@ -253,7 +311,7 @@ func fetchNewAPI(ctx context.Context, item *upstream.Upstream) (Result, error) { result.BillingGroup = group } var detail newAPILogBilling - if !decodeNewAPILogBilling(entry.Other, &detail) || detail.GroupRatio == nil { + if !decodeNewAPILogBilling(entry.Other, &detail) || (detail.GroupRatio == nil && detail.UserGroupRatio == nil) { continue } // 分组变更后，旧分组的扣费不能套到当前分组上。 @@ -267,7 +325,19 @@ func fetchNewAPI(ctx context.Context, item *upstream.Upstream) (Result, error) { break } if result.BillingGroup == "" { - result.Warning = "New API has no recent token log for billing group detection" + // 日志接口不可用或尚无消费记录时，当前用户资料仍能给出分组。 + // 这允许继续读取当前公示倍率，同时保留日志不可用的 partial 状态。 + group, err := fetchNewAPIUserGroup(ctx, item) + if err == nil && group != "" { + result.BillingGroup = group + } else if result.Warning == "" { + result.Warning = "New API has no recent token log for billing group detection" + } + } + if result.BillingGroup == "" { + if result.Warning == "" { + result.Warning = "New API billing group is unavailable" + } return result, nil } @@ -279,10 +349,9 @@ func fetchNewAPI(ctx context.Context, item *upstream.Upstream) (Result, error) { return result, nil } if group, ok := groups.Data[result.BillingGroup]; ok { - var ratio float64 - if json.Unmarshal(group.Ratio, &ratio) == nil { - result.GroupMultiplier = &ratio - result.EffectiveMultiplier = &ratio + if ratio := jsonFloat(group.Ratio); ratio != nil && *ratio > 0 { + result.GroupMultiplier = ratio + result.EffectiveMultiplier = ratio } } // 有个人议价则覆盖 effective，group 仍是公示价(供审计对比) diff --git a/internal/billing/collector_test.go b/internal/billing/collector_test.go index 62abc85..f467a78 100644 --- a/internal/billing/collector_test.go +++ b/internal/billing/collector_test.go @@ -267,3 +267,68 @@ func TestFetchNewAPIPartialWithoutLogs(t *testing.T) { t.Fatalf("partial New API data should retain balance: %+v", result) } } + +func TestFetchNewAPIFallsBackToCurrentUserGroupWithoutLogs(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/usage/token/": + w.Write([]byte(`{"data":{"object":"token_usage","total_used":100,"unlimited_quota":true}}`)) + case "/api/status": + w.Write([]byte(`{"data":{"quota_per_unit":500000}}`)) + case "/api/log/token": + http.Error(w, "rate limited", http.StatusTooManyRequests) + case "/api/user/self": + w.Write([]byte(`{"success":true,"data":{"group":"kiro-high"}}`)) + case "/api/user/groups": + w.Write([]byte(`{"success":true,"data":{"kiro-high":{"ratio":"0.125"}}}`)) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + + result, err := Fetch(context.Background(), &upstream.Upstream{ + BaseURL: server.URL, APIKey: "sk-test", BillingType: upstream.BillingNewAPI, + }) + if err != nil { + t.Fatal(err) + } + if result.BillingGroup != "kiro-high" { + t.Fatalf("expected current user group fallback, got %+v", result) + } + if result.EffectiveMultiplier == nil || *result.EffectiveMultiplier != 0.125 { + t.Fatalf("expected numeric-string public multiplier, got %v", result.EffectiveMultiplier) + } + if result.Warning == "" { + t.Fatal("log rate-limit should remain visible as a partial warning") + } +} + +func TestFetchNewAPIUsesUsageGroupWhenLogsAreUnavailable(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/usage/token/": + w.Write([]byte(`{"data":{"object":"token_usage","group":"kiro-high","total_used":100,"unlimited_quota":true}}`)) + case "/api/status": + w.Write([]byte(`{"data":{"quota_per_unit":500000}}`)) + case "/api/log/token": + http.Error(w, "rate limited", http.StatusTooManyRequests) + case "/api/user/groups": + w.Write([]byte(`{"success":true,"data":{"kiro-high":{"ratio":"0.125"}}}`)) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + + result, err := Fetch(context.Background(), &upstream.Upstream{ + BaseURL: server.URL, APIKey: "sk-test", BillingType: upstream.BillingNewAPI, + }) + if err != nil { + t.Fatal(err) + } + if result.BillingGroup != "kiro-high" || result.EffectiveMultiplier == nil || + *result.EffectiveMultiplier != 0.125 { + t.Fatalf("expected usage group and current multiplier, got %+v", result) + } +} diff --git a/internal/store/billing.go b/internal/store/billing.go index b373d03..97bbfbc 100644 --- a/internal/store/billing.go +++ b/internal/store/billing.go @@ -124,7 +124,7 @@ func (s *Store) SaveBillingSuccess(state BillingStatus) error { VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?) ON CONFLICT(upstream_id) DO UPDATE SET currency=excluded.currency,remaining=excluded.remaining,unlimited=excluded.unlimited, - billing_group=COALESCE(excluded.billing_group,upstream_billing_status.billing_group), + billing_group=CASE WHEN excluded.billing_group <> '' THEN excluded.billing_group ELSE upstream_billing_status.billing_group END, group_multiplier=COALESCE(excluded.group_multiplier,upstream_billing_status.group_multiplier), effective_multiplier=COALESCE(excluded.effective_multiplier,upstream_billing_status.effective_multiplier), reported_list_cost=excluded.reported_list_cost, diff --git a/internal/store/billing_test.go b/internal/store/billing_test.go index 56a3ec6..4f33384 100644 --- a/internal/store/billing_test.go +++ b/internal/store/billing_test.go @@ -83,6 +83,25 @@ func TestBillingStatusAndSnapshots(t *testing.T) { t.Fatalf("unexpected billing state: %+v", got) } + // A partial provider response may omit the group while still carrying a + // balance. It must not erase the last known group used by routing/UI. + if err := st.SaveBillingSuccess(BillingStatus{ + UpstreamID: u.ID, Currency: "USD", Remaining: billingFloat(24.5), + ObservedAt: 1_700_000_050, RefreshedAt: 1_700_000_051, Status: "partial", + Error: "token log rate limited", + }); err != nil { + t.Fatal(err) + } + partial, err := st.GetBillingStatus(u.ID) + if err != nil { + t.Fatal(err) + } + if partial.BillingGroup != "pro" || partial.GroupMultiplier == nil || + *partial.GroupMultiplier != 0.155 || partial.EffectiveMultiplier == nil || + *partial.EffectiveMultiplier != 0.155 { + t.Fatalf("partial refresh must preserve known billing fields: %+v", partial) + } + if err := st.SaveBillingFailure(u.ID, "timeout", 1_700_000_100); err != nil { t.Fatal(err) } @@ -91,13 +110,13 @@ func TestBillingStatusAndSnapshots(t *testing.T) { t.Fatal(err) } if failed.Status != "error" || failed.Error != "timeout" || failed.Remaining == nil || - *failed.Remaining != 24.93 || failed.LastSuccessAt != state.RefreshedAt { + *failed.Remaining != 24.5 || failed.LastSuccessAt != partial.RefreshedAt { t.Fatalf("failure should preserve the last successful values: %+v", failed) } snapshots, err := st.ListBillingSnapshots(u.ID, 10) - if err != nil || len(snapshots) != 1 || snapshots[0].ReportedActualCost == nil || - *snapshots[0].ReportedActualCost != 15.03 { + if err != nil || len(snapshots) != 2 || snapshots[1].ReportedActualCost == nil || + *snapshots[1].ReportedActualCost != 15.03 { t.Fatalf("unexpected billing snapshots: %+v, err=%v", snapshots, err) } statuses, err := st.ListBillingStatuses() diff --git a/internal/upstream/upstream.go b/internal/upstream/upstream.go index d1f7cea..b04683c 100644 --- a/internal/upstream/upstream.go +++ b/internal/upstream/upstream.go @@ -4,6 +4,7 @@ package upstream import ( "context" "encoding/json" + "fmt" "io" "net/http" "net/url" @@ -257,32 +258,71 @@ func (u *Upstream) FetchModels(ctx context.Context, timeout time.Duration) ([]st if resp.StatusCode >= 400 { return nil, resp.StatusCode, &HTTPError{Status: resp.StatusCode, Body: strings.TrimSpace(string(body))} } - var parsed struct { - Data []struct { - ID string `json:"id"` - } `json:"data"` - Models []struct { - Name string `json:"name"` - } `json:"models"` - } + var parsed any if err := json.Unmarshal(body, &parsed); err != nil { return nil, resp.StatusCode, err } if IsErrorPayload(body) { return nil, resp.StatusCode, &HTTPError{Status: resp.StatusCode, Body: strings.TrimSpace(string(body))} } - models := make([]string, 0, len(parsed.Data)) - for _, m := range parsed.Data { - if m.ID != "" { - models = append(models, m.ID) + models := modelIDs(parsed) + if len(models) == 0 { + return nil, resp.StatusCode, fmt.Errorf("upstream returned no models from %s", modelsPath) + } + return models, resp.StatusCode, nil +} + +// modelIDs accepts the response shapes used by OpenAI-compatible, Claude +// relays and Gemini-compatible gateways: {data: [...]}, {models: [...]}, +// {items: [...]}, nested envelopes, and a bare array. A successful HTTP +// response with no recognizable model entry is handled by FetchModels as an +// error instead of being presented as a healthy empty list. +func modelIDs(value any) []string { + seen := make(map[string]struct{}) + var ids []string + var visit func(any) + visit = func(current any) { + switch item := current.(type) { + case string: + id := strings.TrimPrefix(strings.TrimSpace(item), "models/") + if id != "" { + if _, ok := seen[id]; !ok { + seen[id] = struct{}{} + ids = append(ids, id) + } + } + case []any: + for _, entry := range item { + visit(entry) + } + case map[string]any: + if id := modelID(item); id != "" { + if _, ok := seen[id]; !ok { + seen[id] = struct{}{} + ids = append(ids, id) + } + } + for _, key := range []string{"data", "models", "items", "results"} { + if nested, ok := item[key]; ok { + visit(nested) + } + } } } - for _, m := range parsed.Models { - if id := strings.TrimPrefix(strings.TrimSpace(m.Name), "models/"); id != "" { - models = append(models, id) + visit(value) + return ids +} + +func modelID(value map[string]any) string { + for _, key := range []string{"id", "name"} { + if raw, ok := value[key].(string); ok { + id := strings.TrimPrefix(strings.TrimSpace(raw), "models/") + if id != "" { + return id + } } } - return models, resp.StatusCode, nil + return "" } // HTTPError 上游返回非 2xx 时携带状态码与响应体片段。 diff --git a/internal/upstream/upstream_test.go b/internal/upstream/upstream_test.go index 2cee6ed..ad766c6 100644 --- a/internal/upstream/upstream_test.go +++ b/internal/upstream/upstream_test.go @@ -53,6 +53,47 @@ func TestFetchModelsRejectsSuccessfulErrorEnvelope(t *testing.T) { } } +func TestFetchModelsAcceptsCommonEnvelopeShapes(t *testing.T) { + cases := []struct { + name string + body string + want []string + }{ + {name: "items", body: `{"items":[{"id":"gpt-5.6"}]}`, want: []string{"gpt-5.6"}}, + {name: "string models", body: `{"models":["models/gpt-5.6"]}`, want: []string{"gpt-5.6"}}, + {name: "bare array", body: `[{"name":"models/claude-sonnet"}]`, want: []string{"claude-sonnet"}}, + {name: "nested results", body: `{"data":{"results":[{"id":"gpt-5.5"}]}}`, want: []string{"gpt-5.5"}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(tc.body)) + })) + defer server.Close() + u := &Upstream{BaseURL: server.URL, APIKey: "key", Protocol: "claude"} + models, status, err := u.FetchModels(context.Background(), time.Second) + if err != nil || status != http.StatusOK || len(models) != 1 || models[0] != tc.want[0] { + t.Fatalf("models=%v status=%d err=%v", models, status, err) + } + }) + } +} + +func TestFetchModelsRejectsSuccessfulEmptyEnvelope(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"data":[]}`)) + })) + defer server.Close() + + u := &Upstream{BaseURL: server.URL, APIKey: "key", Protocol: "claude"} + models, status, err := u.FetchModels(context.Background(), time.Second) + if err == nil || status != http.StatusOK || models != nil { + t.Fatalf("empty successful envelope must fail discovery: models=%v status=%d err=%v", models, status, err) + } +} + func TestIsFailureStatus(t *testing.T) { // 应触发故障切换/熔断的状态码 fail := []int{