feat(ollama): 支持 Cloud 官方用量自动刷新

backend/internal/handler/admin/account_handler.go

@@ -63,6 +63,7 @@ type AccountHandler struct {
 	tokenCacheInvalidator   service.TokenCacheInvalidator
 	grokImportProber        grokImportProber
 	upstreamBillingProbe    *service.UpstreamBillingProbeService
+	ollamaCloudUsage        *service.OllamaCloudUsageService
 }
 
 // SetUpstreamBillingProbeService attaches the optional remote billing probe service.
@@ -70,6 +71,10 @@ func (h *AccountHandler) SetUpstreamBillingProbeService(probe *service.UpstreamB
 	h.upstreamBillingProbe = probe
 }
 
+func (h *AccountHandler) SetOllamaCloudUsageService(usage *service.OllamaCloudUsageService) {
+	h.ollamaCloudUsage = usage
+}
+
 // NewAccountHandler creates a new admin account handler
 func NewAccountHandler(
 	adminService service.AdminService,
@@ -208,9 +213,17 @@ type AccountSchedulerGroupScore struct {
 
 const accountListGroupUngroupedQueryValue = "ungrouped"
 
+func (h *AccountHandler) accountResponseFromService(account *service.Account) *dto.Account {
+	out := dto.AccountFromService(account)
+	if h != nil && h.ollamaCloudUsage != nil && out != nil {
+		h.ollamaCloudUsage.EnrichState(out.OllamaCloudUsage)
+	}
+	return out
+}
+
 func (h *AccountHandler) buildAccountResponseWithRuntime(ctx context.Context, account *service.Account) AccountWithConcurrency {
 	item := AccountWithConcurrency{
-		Account:            dto.AccountFromService(account),
+		Account:            h.accountResponseFromService(account),
 		CurrentConcurrency: 0,
 	}
 	if account == nil {
@@ -523,6 +536,16 @@ func (h *AccountHandler) List(c *gin.Context) {
 		response.ErrorFrom(c, err)
 		return
 	}
+	if h.ollamaCloudUsage != nil && len(accounts) > 0 {
+		accountPointers := make([]*service.Account, len(accounts))
+		for index := range accounts {
+			accountPointers[index] = &accounts[index]
+		}
+		if err := h.ollamaCloudUsage.ResolveAccounts(c.Request.Context(), accountPointers); err != nil {
+			response.ErrorFrom(c, err)
+			return
+		}
+	}
 
 	// Get current concurrency counts for all accounts
 	accountIDs := make([]int64, len(accounts))
@@ -626,7 +649,7 @@ func (h *AccountHandler) List(c *gin.Context) {
 	for i := range accounts {
 		acc := &accounts[i]
 		item := AccountWithConcurrency{
-			Account:            dto.AccountFromService(acc),
+			Account:            h.accountResponseFromService(acc),
 			CurrentConcurrency: concurrencyCounts[acc.ID],
 			SchedulerScore:     schedulerScores[acc.ID],
 			SchedulerScores:    schedulerGroupScores[acc.ID],
@@ -740,6 +763,12 @@ func (h *AccountHandler) GetByID(c *gin.Context) {
 		response.ErrorFrom(c, err)
 		return
 	}
+	if h.ollamaCloudUsage != nil {
+		if err := h.ollamaCloudUsage.ResolveAccounts(c.Request.Context(), []*service.Account{account}); err != nil {
+			response.ErrorFrom(c, err)
+			return
+		}
+	}
 
 	response.Success(c, h.buildAccountResponseWithRuntime(c.Request.Context(), account))
 }

backend/internal/handler/admin/account_ollama_cloud_usage.go

@@ -0,0 +1,162 @@
+package admin
+
+import (
+	"strconv"
+
+	"github.com/Wei-Shaw/sub2api/internal/pkg/response"
+	"github.com/Wei-Shaw/sub2api/internal/service"
+	"github.com/gin-gonic/gin"
+)
+
+type ollamaCloudUsageSessionRequest struct {
+	Session string `json:"session" binding:"required"`
+}
+
+type ollamaCloudUsageAutoRefreshRequest struct {
+	Enabled *bool `json:"enabled" binding:"required"`
+}
+
+func (h *AccountHandler) GetOllamaCloudUsageSettings(c *gin.Context) {
+	if h.ollamaCloudUsage == nil {
+		response.ErrorFrom(c, service.ErrOllamaCloudUsageUnavailable)
+		return
+	}
+	settings, err := h.ollamaCloudUsage.GetSettings(c.Request.Context())
+	if err != nil {
+		response.ErrorFrom(c, err)
+		return
+	}
+	response.Success(c, settings)
+}
+
+func (h *AccountHandler) UpdateOllamaCloudUsageSettings(c *gin.Context) {
+	if h.ollamaCloudUsage == nil {
+		response.ErrorFrom(c, service.ErrOllamaCloudUsageUnavailable)
+		return
+	}
+	var req service.OllamaCloudUsageSettings
+	if err := c.ShouldBindJSON(&req); err != nil {
+		response.BadRequest(c, "Invalid request: "+err.Error())
+		return
+	}
+	if err := h.ollamaCloudUsage.UpdateSettings(c.Request.Context(), &req); err != nil {
+		response.ErrorFrom(c, err)
+		return
+	}
+	settings, err := h.ollamaCloudUsage.GetSettings(c.Request.Context())
+	if err != nil {
+		response.ErrorFrom(c, err)
+		return
+	}
+	response.Success(c, settings)
+}
+
+func (h *AccountHandler) GetOllamaCloudUsage(c *gin.Context) {
+	if !h.requireOllamaCloudUsage(c) {
+		return
+	}
+	accountID, ok := ollamaCloudUsageAccountID(c)
+	if !ok {
+		return
+	}
+	state, err := h.ollamaCloudUsage.GetState(c.Request.Context(), accountID)
+	if err != nil {
+		response.ErrorFrom(c, err)
+		return
+	}
+	response.Success(c, state)
+}
+
+func (h *AccountHandler) SaveOllamaCloudUsageSession(c *gin.Context) {
+	if !h.requireOllamaCloudUsage(c) {
+		return
+	}
+	accountID, ok := ollamaCloudUsageAccountID(c)
+	if !ok {
+		return
+	}
+	var req ollamaCloudUsageSessionRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		response.BadRequest(c, "Invalid request: "+err.Error())
+		return
+	}
+	state, err := h.ollamaCloudUsage.SaveSession(c.Request.Context(), accountID, req.Session)
+	if err != nil {
+		response.ErrorFrom(c, err)
+		return
+	}
+	response.Success(c, state)
+}
+
+func (h *AccountHandler) DeleteOllamaCloudUsageSession(c *gin.Context) {
+	if !h.requireOllamaCloudUsage(c) {
+		return
+	}
+	accountID, ok := ollamaCloudUsageAccountID(c)
+	if !ok {
+		return
+	}
+	state, err := h.ollamaCloudUsage.DeleteSession(c.Request.Context(), accountID)
+	if err != nil {
+		response.ErrorFrom(c, err)
+		return
+	}
+	response.Success(c, state)
+}
+
+func (h *AccountHandler) SetOllamaCloudUsageAutoRefresh(c *gin.Context) {
+	if !h.requireOllamaCloudUsage(c) {
+		return
+	}
+	accountID, ok := ollamaCloudUsageAccountID(c)
+	if !ok {
+		return
+	}
+	var req ollamaCloudUsageAutoRefreshRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		response.BadRequest(c, "Invalid request: "+err.Error())
+		return
+	}
+	state, err := h.ollamaCloudUsage.SetAutoRefresh(c.Request.Context(), accountID, *req.Enabled)
+	if err != nil {
+		response.ErrorFrom(c, err)
+		return
+	}
+	response.Success(c, state)
+}
+
+func (h *AccountHandler) RefreshOllamaCloudUsage(c *gin.Context) {
+	if !h.requireOllamaCloudUsage(c) {
+		return
+	}
+	accountID, ok := ollamaCloudUsageAccountID(c)
+	if !ok {
+		return
+	}
+	state, err := h.ollamaCloudUsage.Refresh(c.Request.Context(), accountID)
+	if err != nil {
+		response.ErrorFrom(c, err)
+		return
+	}
+	response.Success(c, state)
+}
+
+func (h *AccountHandler) requireOllamaCloudUsage(c *gin.Context) bool {
+	if h != nil && h.ollamaCloudUsage != nil {
+		return true
+	}
+	response.ErrorFrom(c, service.ErrOllamaCloudUsageUnavailable)
+	return false
+}
+
+func ollamaCloudUsageAccountID(c *gin.Context) (int64, bool) {
+	if c == nil {
+		return 0, false
+	}
+	accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
+	if err != nil || accountID <= 0 {
+		response.BadRequest(c, "Invalid account ID")
+		return 0, false
+	}
+	return accountID, true
+}

backend/internal/handler/admin/account_ollama_cloud_usage_test.go

@@ -0,0 +1,276 @@
+package admin
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/service"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+)
+
+type ollamaCloudUsageHandlerTestRepo struct {
+	service.AccountRepository
+	account           *service.Account
+	accounts          []*service.Account
+	groupResolveCalls int
+}
+
+func (r *ollamaCloudUsageHandlerTestRepo) GetByID(_ context.Context, id int64) (*service.Account, error) {
+	if r.account != nil && r.account.ID == id {
+		return r.account, nil
+	}
+	for _, account := range r.accounts {
+		if account.ID == id {
+			return account, nil
+		}
+	}
+	return nil, service.ErrAccountNotFound
+}
+
+func (r *ollamaCloudUsageHandlerTestRepo) ListOllamaCloudUsageGroupAccounts(_ context.Context, _ []*service.Account) ([]service.Account, error) {
+	r.groupResolveCalls++
+	result := make([]service.Account, 0, len(r.accounts)+1)
+	if r.account != nil {
+		result = append(result, *r.account)
+	}
+	for _, account := range r.accounts {
+		result = append(result, *account)
+	}
+	return result, nil
+}
+
+func (r *ollamaCloudUsageHandlerTestRepo) SaveOllamaCloudUsageSession(context.Context, *service.Account, string, bool) error {
+	return nil
+}
+func (r *ollamaCloudUsageHandlerTestRepo) DeleteOllamaCloudUsageSession(context.Context, *service.Account) error {
+	return nil
+}
+func (r *ollamaCloudUsageHandlerTestRepo) SetOllamaCloudUsageAutoRefresh(context.Context, *service.Account, bool) error {
+	return nil
+}
+func (r *ollamaCloudUsageHandlerTestRepo) UpdateOllamaCloudUsageSnapshot(context.Context, *service.Account, *service.OllamaCloudUsageSnapshot) error {
+	return nil
+}
+func (r *ollamaCloudUsageHandlerTestRepo) DisableOllamaCloudUsageAutoRefresh(context.Context, *service.Account) error {
+	return nil
+}
+func (r *ollamaCloudUsageHandlerTestRepo) ListDueOllamaCloudUsageAccounts(context.Context, time.Time, int) ([]service.Account, error) {
+	return nil, nil
+}
+
+func newOllamaCloudUsageHandlerTestService(t *testing.T) *service.OllamaCloudUsageService {
+	t.Helper()
+	svc := service.NewOllamaCloudUsageService(nil, nil, nil, nil, false)
+	t.Cleanup(svc.Stop)
+	return svc
+}
+
+func newOllamaCloudUsageHandlerContext(method, target, body, id string) (*gin.Context, *httptest.ResponseRecorder) {
+	recorder := httptest.NewRecorder()
+	request := httptest.NewRequest(method, target, bytes.NewBufferString(body))
+	request.Header.Set("Content-Type", "application/json")
+	ctx, _ := gin.CreateTestContext(recorder)
+	ctx.Request = request
+	if id != "" {
+		ctx.Params = gin.Params{{Key: "id", Value: id}}
+	}
+	return ctx, recorder
+}
+
+func TestOllamaCloudUsageHandlersValidateRequestsAndDependencies(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	svc := newOllamaCloudUsageHandlerTestService(t)
+
+	t.Run("invalid account id", func(t *testing.T) {
+		ctx, recorder := newOllamaCloudUsageHandlerContext(http.MethodGet, "/admin/accounts/not-an-id/ollama-cloud-usage", "", "not-an-id")
+		(&AccountHandler{ollamaCloudUsage: svc}).GetOllamaCloudUsage(ctx)
+		require.Equal(t, http.StatusBadRequest, recorder.Code)
+	})
+
+	t.Run("empty session", func(t *testing.T) {
+		ctx, recorder := newOllamaCloudUsageHandlerContext(http.MethodPut, "/admin/accounts/7/ollama-cloud-usage/session", `{"session":""}`, "7")
+		(&AccountHandler{ollamaCloudUsage: svc}).SaveOllamaCloudUsageSession(ctx)
+		require.Equal(t, http.StatusBadRequest, recorder.Code)
+	})
+
+	t.Run("missing enabled", func(t *testing.T) {
+		ctx, recorder := newOllamaCloudUsageHandlerContext(http.MethodPut, "/admin/accounts/7/ollama-cloud-usage/auto-refresh", `{}`, "7")
+		(&AccountHandler{ollamaCloudUsage: svc}).SetOllamaCloudUsageAutoRefresh(ctx)
+		require.Equal(t, http.StatusBadRequest, recorder.Code)
+	})
+
+	t.Run("service unavailable", func(t *testing.T) {
+		ctx, recorder := newOllamaCloudUsageHandlerContext(http.MethodGet, "/admin/accounts/7/ollama-cloud-usage", "", "7")
+		(&AccountHandler{}).GetOllamaCloudUsage(ctx)
+		require.Equal(t, http.StatusServiceUnavailable, recorder.Code)
+		require.Contains(t, recorder.Body.String(), "OLLAMA_CLOUD_USAGE_UNAVAILABLE")
+	})
+}
+
+func TestOllamaCloudUsageEncryptionKeyStateConsistentAcrossAccountResponses(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	for _, configured := range []bool{false, true} {
+		t.Run("configured="+strconv.FormatBool(configured), func(t *testing.T) {
+			account := &service.Account{
+				ID:          7,
+				Name:        "ollama",
+				Platform:    service.PlatformOpenAI,
+				Type:        service.AccountTypeAPIKey,
+				Credentials: map[string]any{"base_url": "https://ollama.com", "api_key": "test-key"},
+				Extra:       map[string]any{},
+				Status:      service.StatusActive,
+			}
+			adminService := newStubAdminService()
+			adminService.accounts = []service.Account{*account}
+			adminService.getAccountResult = account
+			usageService := service.NewOllamaCloudUsageService(
+				&ollamaCloudUsageHandlerTestRepo{account: account}, nil, nil, nil, configured,
+			)
+			t.Cleanup(usageService.Stop)
+
+			handler := NewAccountHandler(adminService, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
+			handler.SetOllamaCloudUsageService(usageService)
+			router := gin.New()
+			router.GET("/accounts", handler.List)
+			router.GET("/accounts/:id", handler.GetByID)
+			router.GET("/accounts/:id/ollama-cloud-usage", handler.GetOllamaCloudUsage)
+
+			listRecorder := httptest.NewRecorder()
+			router.ServeHTTP(listRecorder, httptest.NewRequest(http.MethodGet, "/accounts?page=1&page_size=20", nil))
+			require.Equal(t, http.StatusOK, listRecorder.Code)
+			var listPayload struct {
+				Data struct {
+					Items []struct {
+						OllamaCloudUsage *service.OllamaCloudUsageState `json:"ollama_cloud_usage"`
+					} `json:"items"`
+				} `json:"data"`
+			}
+			require.NoError(t, json.Unmarshal(listRecorder.Body.Bytes(), &listPayload))
+			require.Len(t, listPayload.Data.Items, 1)
+			require.NotNil(t, listPayload.Data.Items[0].OllamaCloudUsage)
+
+			detailRecorder := httptest.NewRecorder()
+			router.ServeHTTP(detailRecorder, httptest.NewRequest(http.MethodGet, "/accounts/7", nil))
+			require.Equal(t, http.StatusOK, detailRecorder.Code)
+			var detailPayload struct {
+				Data struct {
+					OllamaCloudUsage *service.OllamaCloudUsageState `json:"ollama_cloud_usage"`
+				} `json:"data"`
+			}
+			require.NoError(t, json.Unmarshal(detailRecorder.Body.Bytes(), &detailPayload))
+			require.NotNil(t, detailPayload.Data.OllamaCloudUsage)
+
+			stateRecorder := httptest.NewRecorder()
+			router.ServeHTTP(stateRecorder, httptest.NewRequest(http.MethodGet, "/accounts/7/ollama-cloud-usage", nil))
+			require.Equal(t, http.StatusOK, stateRecorder.Code)
+			var statePayload struct {
+				Data service.OllamaCloudUsageState `json:"data"`
+			}
+			require.NoError(t, json.Unmarshal(stateRecorder.Body.Bytes(), &statePayload))
+
+			listConfigured := listPayload.Data.Items[0].OllamaCloudUsage.EncryptionKeyConfigured
+			detailConfigured := detailPayload.Data.OllamaCloudUsage.EncryptionKeyConfigured
+			require.Equal(t, configured, listConfigured)
+			require.Equal(t, statePayload.Data.EncryptionKeyConfigured, listConfigured)
+			require.Equal(t, statePayload.Data.EncryptionKeyConfigured, detailConfigured)
+		})
+	}
+}
+
+func TestOllamaCloudUsageSharedStateMatchesListDetailAndSpecialEndpointWithoutListNPlusOne(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	now := time.Now().UTC()
+	source := &service.Account{
+		ID: 7, Name: "source", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
+		Credentials: map[string]any{"base_url": "https://ollama.com", "api_key": "shared-secret-key"},
+		Extra: map[string]any{
+			service.OllamaCloudUsageSessionExtraKey:     "ciphertext-secret",
+			service.OllamaCloudUsageAutoRefreshExtraKey: true,
+			service.OllamaCloudUsageSnapshotExtraKey: &service.OllamaCloudUsageSnapshot{
+				Status: service.OllamaCloudUsageStatusOK, Data: &service.OllamaCloudUsageData{Plan: "pro"},
+				LastAttemptAt: now, NextRefreshAt: now.Add(time.Hour),
+			},
+		},
+		Status: service.StatusActive,
+	}
+	sibling := &service.Account{
+		ID: 8, Name: "sibling", Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey,
+		Credentials: map[string]any{"base_url": "HTTPS://WWW.OLLAMA.COM:443/v1", "api_key": "shared-secret-key"},
+		Extra:       map[string]any{}, Status: service.StatusActive,
+	}
+	repo := &ollamaCloudUsageHandlerTestRepo{accounts: []*service.Account{source, sibling}}
+	adminService := newStubAdminService()
+	adminService.accounts = []service.Account{*source, *sibling}
+	adminService.getAccountResult = sibling
+	usageService := service.NewOllamaCloudUsageService(repo, nil, nil, nil, true)
+	t.Cleanup(usageService.Stop)
+	handler := NewAccountHandler(adminService, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
+	handler.SetOllamaCloudUsageService(usageService)
+	router := gin.New()
+	router.GET("/accounts", handler.List)
+	router.GET("/accounts/:id", handler.GetByID)
+	router.GET("/accounts/:id/ollama-cloud-usage", handler.GetOllamaCloudUsage)
+
+	listRecorder := httptest.NewRecorder()
+	router.ServeHTTP(listRecorder, httptest.NewRequest(http.MethodGet, "/accounts?page=1&page_size=20", nil))
+	require.Equal(t, http.StatusOK, listRecorder.Code)
+	require.Equal(t, 1, repo.groupResolveCalls, "the full list page must use one group-resolution batch")
+	var listPayload struct {
+		Data struct {
+			Items []struct {
+				ID               int64                          `json:"id"`
+				OllamaCloudUsage *service.OllamaCloudUsageState `json:"ollama_cloud_usage"`
+			} `json:"items"`
+		} `json:"data"`
+	}
+	require.NoError(t, json.Unmarshal(listRecorder.Body.Bytes(), &listPayload))
+	require.Len(t, listPayload.Data.Items, 2)
+	for _, item := range listPayload.Data.Items {
+		require.True(t, item.OllamaCloudUsage.Configured)
+		require.Equal(t, "pro", item.OllamaCloudUsage.Snapshot.Data.Plan)
+	}
+
+	detailRecorder := httptest.NewRecorder()
+	router.ServeHTTP(detailRecorder, httptest.NewRequest(http.MethodGet, "/accounts/8", nil))
+	require.Equal(t, http.StatusOK, detailRecorder.Code)
+	var detailPayload struct {
+		Data struct {
+			OllamaCloudUsage *service.OllamaCloudUsageState `json:"ollama_cloud_usage"`
+		} `json:"data"`
+	}
+	require.NoError(t, json.Unmarshal(detailRecorder.Body.Bytes(), &detailPayload))
+
+	stateRecorder := httptest.NewRecorder()
+	router.ServeHTTP(stateRecorder, httptest.NewRequest(http.MethodGet, "/accounts/8/ollama-cloud-usage", nil))
+	require.Equal(t, http.StatusOK, stateRecorder.Code)
+	var statePayload struct {
+		Data service.OllamaCloudUsageState `json:"data"`
+	}
+	require.NoError(t, json.Unmarshal(stateRecorder.Body.Bytes(), &statePayload))
+	require.Equal(t, statePayload.Data.Configured, detailPayload.Data.OllamaCloudUsage.Configured)
+	require.Equal(t, statePayload.Data.Snapshot, detailPayload.Data.OllamaCloudUsage.Snapshot)
+	for _, body := range []string{listRecorder.Body.String(), detailRecorder.Body.String(), stateRecorder.Body.String()} {
+		require.NotContains(t, body, "shared-secret-key")
+		require.NotContains(t, body, "ciphertext-secret")
+	}
+}
+
+func TestGetOllamaCloudUsageSettingsHandlerSuccess(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	ctx, recorder := newOllamaCloudUsageHandlerContext(http.MethodGet, "/admin/accounts/ollama-cloud-usage/settings", "", "")
+	handler := &AccountHandler{ollamaCloudUsage: newOllamaCloudUsageHandlerTestService(t)}
+
+	handler.GetOllamaCloudUsageSettings(ctx)
+
+	require.Equal(t, http.StatusOK, recorder.Code)
+	require.Contains(t, recorder.Body.String(), `"enabled":false`)
+	require.Contains(t, recorder.Body.String(), `"interval_minutes":60`)
+}