fix bugs
This commit is contained in:
parent
ced11faf11
commit
f2be5a436e
|
|
@ -36,21 +36,128 @@ func main() {
|
|||
defer client.Disconnect(context.Background())
|
||||
|
||||
db := client.Database(c.Mongo.Database)
|
||||
indexes := map[string][]mongo.IndexModel{
|
||||
"jobs": {
|
||||
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "run_after", Value: 1}, {Key: "created_at", Value: 1}}, Options: options.Index().SetName("claim_due_jobs")},
|
||||
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "completed_at", Value: 1}}, Options: options.Index().SetName("purge_terminal_jobs")},
|
||||
},
|
||||
"studio_outbox": {
|
||||
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "updated_at", Value: 1}}, Options: options.Index().SetName("worker_outbox_status_updated")},
|
||||
{Keys: bson.D{{Key: "steps.status", Value: 1}, {Key: "steps.scheduled_at", Value: 1}, {Key: "steps.lease_expires_at", Value: 1}}, Options: options.Index().SetName("claim_due_outbox_steps")},
|
||||
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "updated_at", Value: -1}}, Options: options.Index().SetName("owner_outbox_updated")},
|
||||
},
|
||||
}
|
||||
for collection, models := range indexes {
|
||||
for collection, models := range indexModels() {
|
||||
if _, err := db.Collection(collection).Indexes().CreateMany(ctx, models); err != nil {
|
||||
panic(fmt.Errorf("create %s indexes: %w", collection, err))
|
||||
}
|
||||
}
|
||||
fmt.Println("database indexes initialized")
|
||||
}
|
||||
|
||||
// ownerIndex covers the near-universal "list one member's rows, newest first" shape. The sort key
|
||||
// belongs in the index too, otherwise Mongo fetches every matching row before sorting.
|
||||
func ownerIndex(sortField, name string) mongo.IndexModel {
|
||||
return mongo.IndexModel{
|
||||
Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: sortField, Value: -1}},
|
||||
Options: options.Index().SetName(name),
|
||||
}
|
||||
}
|
||||
|
||||
// indexModels lists every index the running code depends on. Owner-scoped collections were all
|
||||
// unindexed, so each list endpoint was a full collection scan that grew with total rows across
|
||||
// every member rather than with the member's own data.
|
||||
//
|
||||
// These are deliberately all non-unique. Adding a uniqueness constraint to existing data can make
|
||||
// CreateMany fail and take a deploy down, so that belongs in its own migration with a duplicate
|
||||
// pre-check rather than here.
|
||||
func indexModels() map[string][]mongo.IndexModel {
|
||||
return map[string][]mongo.IndexModel{
|
||||
// members.uid is read on every authenticated request via the auth middleware, and it is a
|
||||
// field lookup rather than _id, so without this it scanned the whole collection whenever
|
||||
// the Redis cache missed.
|
||||
"members": {
|
||||
{Keys: bson.D{{Key: "uid", Value: 1}}, Options: options.Index().SetName("member_uid")},
|
||||
{Keys: bson.D{{Key: "email", Value: 1}}, Options: options.Index().SetName("member_email")},
|
||||
{Keys: bson.D{{Key: "phone", Value: 1}}, Options: options.Index().SetName("member_phone")},
|
||||
},
|
||||
"identities": {
|
||||
{Keys: bson.D{{Key: "login_id", Value: 1}, {Key: "platform", Value: 1}}, Options: options.Index().SetName("identity_login_platform")},
|
||||
},
|
||||
"member_settings": {
|
||||
{Keys: bson.D{{Key: "uid", Value: 1}}, Options: options.Index().SetName("settings_uid")},
|
||||
},
|
||||
"jobs": {
|
||||
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "run_after", Value: 1}, {Key: "created_at", Value: 1}}, Options: options.Index().SetName("claim_due_jobs")},
|
||||
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "completed_at", Value: 1}}, Options: options.Index().SetName("purge_terminal_jobs")},
|
||||
// The jobs list is polled continuously by the UI and runs a count plus a find per poll.
|
||||
ownerIndex("updated_at", "owner_jobs_updated"),
|
||||
},
|
||||
// One notification is written per job state change and nothing purges them, so this
|
||||
// collection grows without bound and its scan cost grows with it.
|
||||
"notifications": {
|
||||
ownerIndex("created_at", "owner_notifications_created"),
|
||||
},
|
||||
"threads_accounts": {
|
||||
ownerIndex("created_at", "owner_accounts_created"),
|
||||
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "threads_user_id", Value: 1}}, Options: options.Index().SetName("owner_threads_user")},
|
||||
},
|
||||
"studio_outbox": {
|
||||
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "updated_at", Value: 1}}, Options: options.Index().SetName("worker_outbox_status_updated")},
|
||||
{Keys: bson.D{{Key: "steps.status", Value: 1}, {Key: "steps.scheduled_at", Value: 1}, {Key: "steps.lease_expires_at", Value: 1}}, Options: options.Index().SetName("claim_due_outbox_steps")},
|
||||
ownerIndex("updated_at", "owner_outbox_updated"),
|
||||
},
|
||||
"studio_personas": {ownerIndex("updated_at", "owner_personas_updated")},
|
||||
"studio_plays": {ownerIndex("updated_at", "owner_plays_updated")},
|
||||
"studio_own_posts": {
|
||||
ownerIndex("published_at", "owner_own_posts_published"),
|
||||
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "account_id", Value: 1}, {Key: "published_at", Value: -1}}, Options: options.Index().SetName("owner_account_own_posts")},
|
||||
},
|
||||
"studio_mentions": {
|
||||
ownerIndex("created_at", "owner_mentions_created"),
|
||||
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "account_id", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("owner_account_mentions")},
|
||||
},
|
||||
"scout_brands": {{Keys: bson.D{{Key: "owner_uid", Value: 1}}, Options: options.Index().SetName("owner_brands")}},
|
||||
"scout_products": {{Keys: bson.D{{Key: "owner_uid", Value: 1}}, Options: options.Index().SetName("owner_products")}},
|
||||
"scout_posts": {
|
||||
ownerIndex("created_at", "owner_posts_created"),
|
||||
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "brand_id", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("owner_brand_posts")},
|
||||
},
|
||||
"scout_homework": {{Keys: bson.D{{Key: "owner_uid", Value: 1}}, Options: options.Index().SetName("owner_homework")}},
|
||||
"inspire_elements": {ownerIndex("updated_at", "owner_elements_updated")},
|
||||
"inspire_sessions": {ownerIndex("updated_at", "owner_sessions_updated")},
|
||||
"usage_events": {
|
||||
{Keys: bson.D{{Key: "uid", Value: 1}, {Key: "month_key", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("uid_month_events")},
|
||||
{Keys: bson.D{{Key: "month_key", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("tenant_month_events")},
|
||||
{Keys: bson.D{{Key: "created_at", Value: 1}}, Options: options.Index().SetName("events_created_range")},
|
||||
},
|
||||
"usage_prefs": {
|
||||
{Keys: bson.D{{Key: "uid", Value: 1}}, Options: options.Index().SetName("prefs_uid")},
|
||||
// Stripe webhooks arrive knowing only the customer, so this is the reverse lookup
|
||||
// back to a member and it runs on every billing event.
|
||||
{Keys: bson.D{{Key: "stripe_customer_id", Value: 1}}, Options: options.Index().SetName("prefs_stripe_customer")},
|
||||
},
|
||||
"usage_purchases": {
|
||||
{Keys: bson.D{{Key: "uid", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("uid_purchases_created")},
|
||||
{Keys: bson.D{{Key: "stripe_event_id", Value: 1}}, Options: options.Index().SetName("purchase_stripe_event")},
|
||||
},
|
||||
"growth_outcomes": {
|
||||
// The observe worker scans by status across all tenants on every tick.
|
||||
{Keys: bson.D{{Key: "status", Value: 1}}, Options: options.Index().SetName("observe_status")},
|
||||
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "source_type", Value: 1}, {Key: "source_id", Value: 1}}, Options: options.Index().SetName("owner_outcome_source")},
|
||||
ownerIndex("created_at", "owner_outcomes_created"),
|
||||
},
|
||||
"growth_checkups": {ownerIndex("created_at", "owner_checkups_created")},
|
||||
"growth_account_health": {{Keys: bson.D{{Key: "owner_uid", Value: 1}}, Options: options.Index().SetName("owner_health")}},
|
||||
"growth_workspaces": {{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "archived", Value: 1}}, Options: options.Index().SetName("owner_workspaces_archived")}},
|
||||
"growth_draft_reviews": {
|
||||
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "ref_type", Value: 1}, {Key: "ref_id", Value: 1}, {Key: "updated_at", Value: -1}}, Options: options.Index().SetName("owner_review_ref")},
|
||||
},
|
||||
"growth_invite_rewards": {
|
||||
{Keys: bson.D{{Key: "inviter_uid", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("inviter_rewards_created")},
|
||||
{Keys: bson.D{{Key: "invitee_uid", Value: 1}}, Options: options.Index().SetName("invitee_reward")},
|
||||
},
|
||||
"growth_utm_links": {
|
||||
// Looked up by code on the unauthenticated public redirect.
|
||||
{Keys: bson.D{{Key: "code", Value: 1}}, Options: options.Index().SetName("utm_code")},
|
||||
ownerIndex("created_at", "owner_utm_created"),
|
||||
},
|
||||
"growth_ws_members": {
|
||||
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "uid", Value: 1}}, Options: options.Index().SetName("workspace_member")},
|
||||
},
|
||||
"billing_checkout_attempts": {
|
||||
{Keys: bson.D{{Key: "uid", Value: 1}, {Key: "request_id", Value: 1}}, Options: options.Index().SetName("uid_request")},
|
||||
{Keys: bson.D{{Key: "session_id", Value: 1}}, Options: options.Index().SetName("checkout_session")},
|
||||
{Keys: bson.D{{Key: "customer_id", Value: 1}}, Options: options.Index().SetName("checkout_customer")},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -164,8 +164,15 @@ func main() {
|
|||
|
||||
ctx := context.Background()
|
||||
const outboxLockTTL = 3 * time.Minute
|
||||
outboxLockKey := strings.Trim(c.Redis.Namespace, ":") + ":worker:outbox"
|
||||
outboxLock := redislock.NewWithClient(redisClient, outboxLockKey, workerID, outboxLockTTL)
|
||||
ns := strings.Trim(c.Redis.Namespace, ":")
|
||||
outboxLock := redislock.NewWithClient(redisClient, ns+":worker:outbox", workerID, outboxLockTTL)
|
||||
|
||||
// 巡場維護(過期 outcome、清終態 job)是全域掃描,不像 job 領取有 guarded update 擋重複。
|
||||
// 每個 worker 每 tick 各跑一次的話,N 台就等於同一份掃描做 N 次。改成一次一台、且拉長間隔。
|
||||
const maintenanceLockTTL = 2 * time.Minute
|
||||
const maintenanceEvery = time.Minute
|
||||
maintenanceLock := redislock.NewWithClient(redisClient, ns+":worker:maintenance", workerID, maintenanceLockTTL)
|
||||
lastMaintenance := time.Time{}
|
||||
for {
|
||||
select {
|
||||
case <-sig:
|
||||
|
|
@ -231,20 +238,50 @@ func main() {
|
|||
}
|
||||
// 2) One worker owns an outbox tick and renews its lease while claims run.
|
||||
processOutbox(ctx, studio, outboxLock, workerID, outboxLockTTL)
|
||||
// 3) growth-loop: expire observing outcomes + light health refresh
|
||||
// 3) 巡場維護:過期 outcome + 清終態 job(單一 worker、低頻)
|
||||
if time.Since(lastMaintenance) >= maintenanceEvery {
|
||||
if runMaintenance(ctx, growthSvc, jobs, maintenanceLock, workerID) {
|
||||
lastMaintenance = time.Now()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// runMaintenance reports whether this worker actually did the sweep, so a worker that lost the
|
||||
// lock retries on the next tick instead of waiting out the full interval.
|
||||
func runMaintenance(
|
||||
ctx context.Context,
|
||||
growthSvc *growthUC.Service,
|
||||
jobs *jobUC.Service,
|
||||
lock *redislock.Lock,
|
||||
workerID string,
|
||||
) bool {
|
||||
locked, err := lock.Acquire(ctx)
|
||||
if err != nil {
|
||||
logx.Errorf("worker %s maintenance lock acquire: %v", workerID, err)
|
||||
return false
|
||||
}
|
||||
if !locked {
|
||||
return false
|
||||
}
|
||||
defer func() {
|
||||
if rerr := lock.Release(ctx); rerr != nil {
|
||||
logx.Errorf("worker %s maintenance lock release: %v", workerID, rerr)
|
||||
}
|
||||
}()
|
||||
|
||||
if n, err := growthSvc.ObserveTick(ctx, 0); err != nil {
|
||||
logx.Errorf("worker %s outcome observe: %v", workerID, err)
|
||||
} else if n > 0 {
|
||||
logx.Infof("worker %s outcome observe updated %d", workerID, n)
|
||||
}
|
||||
// 4) 終態任務超過 2 天自動清除
|
||||
if purged, err := jobs.PurgeExpiredTerminal(ctx); err != nil {
|
||||
logx.Errorf("worker %s purge jobs: %v", workerID, err)
|
||||
} else if purged > 0 {
|
||||
logx.Infof("worker %s purged %d expired terminal job(s)", workerID, purged)
|
||||
}
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func processOutbox(ctx context.Context, studio *studioUC.Service, lock *redislock.Lock, workerID string, ttl time.Duration) {
|
||||
|
|
|
|||
|
|
@ -61,6 +61,11 @@ type (
|
|||
Unlimited bool `json:"unlimited"`
|
||||
}
|
||||
|
||||
UsagePrefsReq {
|
||||
// Uid reads another member's plan; admin only. Empty means the caller's own prefs.
|
||||
Uid string `form:"uid,optional"`
|
||||
}
|
||||
|
||||
UsageSetPrefsReq {
|
||||
Uid string `json:"uid"`
|
||||
PlanId string `json:"plan_id,optional"`
|
||||
|
|
@ -165,7 +170,7 @@ service gateway {
|
|||
get /events (UsageEventsReq) returns (UsageEventsData)
|
||||
|
||||
@handler GetUsagePrefs
|
||||
get /prefs returns (UsagePrefsData)
|
||||
get /prefs (UsagePrefsReq) returns (UsagePrefsData)
|
||||
|
||||
@handler PurchasePlan
|
||||
post /purchase (UsagePurchaseReq) returns (UsagePurchasePublic)
|
||||
|
|
|
|||
|
|
@ -1265,35 +1265,6 @@ func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
|
|||
rest.WithPrefix("/api/v1/utm"),
|
||||
)
|
||||
|
||||
server.AddRoutes(
|
||||
rest.WithMiddlewares(
|
||||
[]rest.Middleware{serverCtx.AuthJWT},
|
||||
[]rest.Route{
|
||||
{
|
||||
Method: http.MethodPut,
|
||||
Path: "/:id/branding",
|
||||
Handler: workspaces.UpdateWorkspaceBrandingHandler(serverCtx),
|
||||
},
|
||||
{
|
||||
Method: http.MethodGet,
|
||||
Path: "/:id/members",
|
||||
Handler: workspaces.ListWorkspaceMembersHandler(serverCtx),
|
||||
},
|
||||
{
|
||||
Method: http.MethodPost,
|
||||
Path: "/:id/members",
|
||||
Handler: workspaces.AddWorkspaceMemberHandler(serverCtx),
|
||||
},
|
||||
{
|
||||
Method: http.MethodDelete,
|
||||
Path: "/:id/members/:uid",
|
||||
Handler: workspaces.RemoveWorkspaceMemberHandler(serverCtx),
|
||||
},
|
||||
}...,
|
||||
),
|
||||
rest.WithPrefix("/api/v1/workspaces"),
|
||||
)
|
||||
|
||||
server.AddRoutes(
|
||||
rest.WithMiddlewares(
|
||||
[]rest.Middleware{serverCtx.AuthJWT},
|
||||
|
|
@ -1352,4 +1323,33 @@ func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
|
|||
),
|
||||
rest.WithPrefix("/api/v1/workspaces"),
|
||||
)
|
||||
|
||||
server.AddRoutes(
|
||||
rest.WithMiddlewares(
|
||||
[]rest.Middleware{serverCtx.AuthJWT},
|
||||
[]rest.Route{
|
||||
{
|
||||
Method: http.MethodPut,
|
||||
Path: "/:id/branding",
|
||||
Handler: workspaces.UpdateWorkspaceBrandingHandler(serverCtx),
|
||||
},
|
||||
{
|
||||
Method: http.MethodGet,
|
||||
Path: "/:id/members",
|
||||
Handler: workspaces.ListWorkspaceMembersHandler(serverCtx),
|
||||
},
|
||||
{
|
||||
Method: http.MethodPost,
|
||||
Path: "/:id/members",
|
||||
Handler: workspaces.AddWorkspaceMemberHandler(serverCtx),
|
||||
},
|
||||
{
|
||||
Method: http.MethodDelete,
|
||||
Path: "/:id/members/:uid",
|
||||
Handler: workspaces.RemoveWorkspaceMemberHandler(serverCtx),
|
||||
},
|
||||
}...,
|
||||
),
|
||||
rest.WithPrefix("/api/v1/workspaces"),
|
||||
)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -9,12 +9,20 @@ import (
|
|||
"apps/backend/internal/logic/usage"
|
||||
"apps/backend/internal/response"
|
||||
"apps/backend/internal/svc"
|
||||
"apps/backend/internal/types"
|
||||
"github.com/zeromicro/go-zero/rest/httpx"
|
||||
)
|
||||
|
||||
func GetUsagePrefsHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
var req types.UsagePrefsReq
|
||||
if err := httpx.Parse(r, &req); err != nil {
|
||||
response.Write(r.Context(), w, nil, response.WrapRequestError(err))
|
||||
return
|
||||
}
|
||||
|
||||
l := usage.NewGetUsagePrefsLogic(r.Context(), svcCtx)
|
||||
data, err := l.GetUsagePrefs()
|
||||
data, err := l.GetUsagePrefs(&req)
|
||||
response.Write(r.Context(), w, data, err)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ func NewAICompleteLogic(ctx context.Context, svcCtx *svc.ServiceContext) *AIComp
|
|||
return &AICompleteLogic{Logger: logx.WithContext(ctx), ctx: ctx, svcCtx: svcCtx}
|
||||
}
|
||||
|
||||
func (l *AICompleteLogic) AIComplete(req *types.AICompleteReq) (*types.AICompleteData, error) {
|
||||
func (l *AICompleteLogic) AIComplete(req *types.AICompleteReq) (_ *types.AICompleteData, err error) {
|
||||
if l.svcCtx.Usage == nil || l.svcCtx.AI == nil || l.svcCtx.KeyResolver == nil {
|
||||
return nil, response.Biz(503, 503001, "ai proxy not configured")
|
||||
}
|
||||
|
|
@ -34,14 +34,15 @@ func (l *AICompleteLogic) AIComplete(req *types.AICompleteReq) (*types.AIComplet
|
|||
if meter == "" {
|
||||
meter = usageDomain.MeterAICopy
|
||||
}
|
||||
mode, err := l.svcCtx.Usage.PrepareCall(l.ctx, uid, meter)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
mode, apiKey, err := l.svcCtx.KeyResolver.ResolveKey(l.ctx, uid, meter)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
charge, err := reserveQuota(l.ctx, l.svcCtx, uid, meter, mode, "AI complete", "proxy.ai")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer charge.Settle(l.ctx, &err)
|
||||
model := req.Model
|
||||
if model == "" {
|
||||
model = "grok-3"
|
||||
|
|
@ -51,6 +52,5 @@ func (l *AICompleteLogic) AIComplete(req *types.AICompleteReq) (*types.AIComplet
|
|||
if err != nil {
|
||||
return nil, response.Biz(502, 502001, err.Error())
|
||||
}
|
||||
_, _ = l.svcCtx.Usage.RecordCall(l.ctx, uid, meter, mode, "AI complete", "proxy.ai")
|
||||
return &types.AICompleteData{Text: text, KeyMode: mode, Model: model}, nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -22,27 +22,27 @@ func NewExaSearchLogic(ctx context.Context, svcCtx *svc.ServiceContext) *ExaSear
|
|||
return &ExaSearchLogic{Logger: logx.WithContext(ctx), ctx: ctx, svcCtx: svcCtx}
|
||||
}
|
||||
|
||||
func (l *ExaSearchLogic) ExaSearch(req *types.SearchReq) (*types.SearchData, error) {
|
||||
if l.svcCtx.Usage == nil || l.svcCtx.Search == nil {
|
||||
func (l *ExaSearchLogic) ExaSearch(req *types.SearchReq) (_ *types.SearchData, err error) {
|
||||
if l.svcCtx.Usage == nil || l.svcCtx.Search == nil || l.svcCtx.KeyResolver == nil {
|
||||
return nil, response.Biz(503, 503001, "search proxy not configured")
|
||||
}
|
||||
uid, ok := middleware.UIDFrom(l.ctx)
|
||||
if !ok {
|
||||
return nil, response.Biz(401, 401001, "missing authorization")
|
||||
}
|
||||
mode, err := l.svcCtx.Usage.PrepareCall(l.ctx, uid, usageDomain.MeterWebSearch)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
mode, apiKey, err := l.svcCtx.KeyResolver.ResolveKey(l.ctx, uid, usageDomain.MeterWebSearch)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
charge, err := reserveQuota(l.ctx, l.svcCtx, uid, usageDomain.MeterWebSearch, mode, "search", "proxy.search")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer charge.Settle(l.ctx, &err)
|
||||
hits, err := l.svcCtx.Search.Search(l.ctx, apiKey, req.Query, req.Limit)
|
||||
if err != nil {
|
||||
return nil, response.Biz(502, 502002, err.Error())
|
||||
}
|
||||
_, _ = l.svcCtx.Usage.RecordCall(l.ctx, uid, usageDomain.MeterWebSearch, mode, "search", "proxy.search")
|
||||
out := make([]types.SearchHit, 0, len(hits))
|
||||
for _, h := range hits {
|
||||
out = append(out, types.SearchHit{Title: h.Title, Url: h.URL, Snippet: h.Snippet})
|
||||
|
|
|
|||
|
|
@ -0,0 +1,71 @@
|
|||
package proxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"apps/backend/internal/svc"
|
||||
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
)
|
||||
|
||||
// proxyCharge tracks one reserved quota unit for a proxied upstream call.
|
||||
//
|
||||
// The key mode is resolved once by the caller and reused for both the reservation and the audit
|
||||
// event, because resolving separately reads member settings twice and a member editing their
|
||||
// keys in between would reserve against one mode and bill the other.
|
||||
type proxyCharge struct {
|
||||
svcCtx *svc.ServiceContext
|
||||
uid int64
|
||||
meter string
|
||||
mode string
|
||||
label string
|
||||
source string
|
||||
settled bool
|
||||
}
|
||||
|
||||
// reserveQuota holds quota for an upcoming call. Pair it with `defer charge.Settle(ctx, &err)`
|
||||
// on a named error return so that failures refund instead of charging for nothing.
|
||||
func reserveQuota(ctx context.Context, svcCtx *svc.ServiceContext, uid int64, meter, mode, label, source string) (*proxyCharge, error) {
|
||||
if err := svcCtx.Usage.PrepareCallWithMode(ctx, uid, meter, mode); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &proxyCharge{
|
||||
svcCtx: svcCtx, uid: uid, meter: meter, mode: mode, label: label, source: source,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Settle commits the charge on success and refunds it on failure.
|
||||
func (c *proxyCharge) Settle(ctx context.Context, errp *error) {
|
||||
if errp != nil && *errp != nil {
|
||||
c.release(ctx)
|
||||
return
|
||||
}
|
||||
c.commit(ctx)
|
||||
}
|
||||
|
||||
// commit writes the audit event. The upstream call already succeeded, so a failure here is
|
||||
// logged rather than returned: the member has their result and charging them is correct, but
|
||||
// losing the event must not be silent either.
|
||||
func (c *proxyCharge) commit(ctx context.Context) {
|
||||
if c == nil || c.settled {
|
||||
return
|
||||
}
|
||||
c.settled = true
|
||||
if _, err := c.svcCtx.Usage.RecordCall(ctx, c.uid, c.meter, c.mode, c.label, c.source); err != nil {
|
||||
logx.Errorf("usage record uid=%d meter=%s source=%s: %v", c.uid, c.meter, c.source, err)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *proxyCharge) release(ctx context.Context) {
|
||||
if c == nil || c.settled {
|
||||
return
|
||||
}
|
||||
c.settled = true
|
||||
// The request context is often already cancelled on the path being compensated for.
|
||||
rctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
|
||||
defer cancel()
|
||||
if err := c.svcCtx.Usage.ReleaseCall(rctx, c.uid, c.meter, c.mode); err != nil {
|
||||
logx.Errorf("usage release uid=%d meter=%s source=%s: %v", c.uid, c.meter, c.source, err)
|
||||
}
|
||||
}
|
||||
|
|
@ -2,6 +2,7 @@ package threads
|
|||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
threadsDomain "apps/backend/internal/module/threads/domain"
|
||||
"apps/backend/internal/svc"
|
||||
|
|
@ -22,7 +23,8 @@ func NewThreadsOAuthCallbackLogic(ctx context.Context, svcCtx *svc.ServiceContex
|
|||
|
||||
func (l *ThreadsOAuthCallbackLogic) ThreadsOAuthCallback(req *types.ThreadsOAuthCallbackReq) (*types.ThreadsOAuthCallbackData, error) {
|
||||
if l.svcCtx.Threads == nil {
|
||||
redir := "http://127.0.0.1:5173/app/crew?oauth=error&msg=not_configured"
|
||||
// 相對路徑:讓瀏覽器留在目前站台,而不是被丟到某台機器的 loopback。
|
||||
redir := strings.TrimRight(l.svcCtx.Config.PublicWebBase, "/") + "/app/crew?oauth=error&msg=not_configured"
|
||||
return &types.ThreadsOAuthCallbackData{Ok: false, Message: "not configured", RedirectUrl: redir}, nil
|
||||
}
|
||||
_, err := l.svcCtx.Threads.OAuthCallback(l.ctx, req.Code, req.State, req.Error)
|
||||
|
|
|
|||
|
|
@ -2,8 +2,11 @@ package usage
|
|||
|
||||
import (
|
||||
"context"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"apps/backend/internal/middleware"
|
||||
"apps/backend/internal/module/permission"
|
||||
"apps/backend/internal/response"
|
||||
"apps/backend/internal/svc"
|
||||
"apps/backend/internal/types"
|
||||
|
|
@ -21,7 +24,7 @@ func NewGetUsagePrefsLogic(ctx context.Context, svcCtx *svc.ServiceContext) *Get
|
|||
return &GetUsagePrefsLogic{Logger: logx.WithContext(ctx), ctx: ctx, svcCtx: svcCtx}
|
||||
}
|
||||
|
||||
func (l *GetUsagePrefsLogic) GetUsagePrefs() (*types.UsagePrefsData, error) {
|
||||
func (l *GetUsagePrefsLogic) GetUsagePrefs(req *types.UsagePrefsReq) (*types.UsagePrefsData, error) {
|
||||
if l.svcCtx.Usage == nil {
|
||||
return nil, response.Biz(503, 503001, "usage not configured")
|
||||
}
|
||||
|
|
@ -29,6 +32,20 @@ func (l *GetUsagePrefsLogic) GetUsagePrefs() (*types.UsagePrefsData, error) {
|
|||
if !ok {
|
||||
return nil, response.Biz(401, 401001, "missing authorization")
|
||||
}
|
||||
// Reading somebody else's plan is admin-only; the admin console needs it to show the plan of
|
||||
// the member being edited rather than the admin's own.
|
||||
if target := strings.TrimSpace(req.Uid); target != "" {
|
||||
parsed, err := strconv.ParseInt(target, 10, 64)
|
||||
if err != nil {
|
||||
return nil, response.Biz(400, 400001, "invalid uid")
|
||||
}
|
||||
if parsed != uid {
|
||||
if !permission.IsAdmin(middleware.RolesFrom(l.ctx)) {
|
||||
return nil, response.Biz(403, 403002, "permission denied")
|
||||
}
|
||||
uid = parsed
|
||||
}
|
||||
}
|
||||
p, err := l.svcCtx.Usage.GetPrefs(l.ctx, uid)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
|
|
|||
|
|
@ -142,7 +142,11 @@ type Repository interface {
|
|||
GetSubscription(context.Context, int64) (*SubscriptionState, error)
|
||||
FindUIDByCustomer(context.Context, string) (int64, error)
|
||||
ApplyEntitlement(context.Context, EntitlementUpdate) (bool, error)
|
||||
WebhookProcessed(context.Context, string) (bool, error)
|
||||
MarkWebhookProcessed(context.Context, string, string, int64) error
|
||||
// ClaimWebhook atomically records an event as being handled, reporting false when another
|
||||
// delivery already claimed it. Stripe retries aggressively and can deliver the same event
|
||||
// concurrently, so a read-then-write check would let two deliveries both proceed.
|
||||
ClaimWebhook(ctx context.Context, id, eventType string, created int64) (claimed bool, err error)
|
||||
// ReleaseWebhookClaim drops a claim whose processing failed, so the next delivery retries.
|
||||
ReleaseWebhookClaim(ctx context.Context, id string) error
|
||||
UpsertPurchase(context.Context, *PurchaseAudit) error
|
||||
}
|
||||
|
|
|
|||
|
|
@ -124,16 +124,28 @@ func (r *MongoRepository) ApplyEntitlement(ctx context.Context, u EntitlementUpd
|
|||
return res.MatchedCount > 0 || res.UpsertedCount > 0, nil
|
||||
}
|
||||
|
||||
func (r *MongoRepository) WebhookProcessed(ctx context.Context, id string) (bool, error) {
|
||||
var out bson.M
|
||||
err := r.webhooks.FindOne(ctx, &out, bson.M{"_id": id})
|
||||
if err == mon.ErrNotFound {
|
||||
func (r *MongoRepository) ClaimWebhook(ctx context.Context, id, eventType string, created int64) (bool, error) {
|
||||
_, err := r.webhooks.InsertOne(ctx, bson.M{
|
||||
"_id": id,
|
||||
"stripe_event_id": id,
|
||||
"type": eventType,
|
||||
"stripe_created_at": created * int64(time.Second),
|
||||
"processed_at": time.Now().UTC().UnixNano(),
|
||||
})
|
||||
if mongo.IsDuplicateKeyError(err) {
|
||||
return false, nil
|
||||
}
|
||||
return err == nil, err
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *MongoRepository) ReleaseWebhookClaim(ctx context.Context, id string) error {
|
||||
_, err := r.webhooks.DeleteOne(ctx, bson.M{"_id": id})
|
||||
if err == mon.ErrNotFound {
|
||||
return nil
|
||||
}
|
||||
func (r *MongoRepository) MarkWebhookProcessed(ctx context.Context, id, eventType string, created int64) error {
|
||||
_, err := r.webhooks.ReplaceOne(ctx, bson.M{"_id": id}, bson.M{"_id": id, "stripe_event_id": id, "type": eventType, "stripe_created_at": created * int64(time.Second), "processed_at": time.Now().UTC().UnixNano()}, options.Replace().SetUpsert(true))
|
||||
return err
|
||||
}
|
||||
func (r *MongoRepository) UpsertPurchase(ctx context.Context, p *PurchaseAudit) error {
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import (
|
|||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
|
|
@ -176,14 +177,22 @@ func (s *Service) HandleWebhook(ctx context.Context, payload []byte, signature s
|
|||
if err != nil {
|
||||
return err
|
||||
}
|
||||
done, err := s.Repo.WebhookProcessed(ctx, e.ID)
|
||||
if err != nil || done {
|
||||
// Claim before processing, so two concurrent deliveries of the same event cannot both
|
||||
// fulfil it. Processing failures release the claim so Stripe's retry still gets through.
|
||||
claimed, err := s.Repo.ClaimWebhook(ctx, e.ID, e.Type, e.Created)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !claimed {
|
||||
return nil
|
||||
}
|
||||
if err := s.processEvent(ctx, e); err != nil {
|
||||
if rerr := s.Repo.ReleaseWebhookClaim(ctx, e.ID); rerr != nil {
|
||||
logx.Errorf("billing release webhook claim %s: %v", e.ID, rerr)
|
||||
}
|
||||
return err
|
||||
}
|
||||
return s.Repo.MarkWebhookProcessed(ctx, e.ID, e.Type, e.Created)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) processEvent(ctx context.Context, e *Event) error {
|
||||
|
|
|
|||
|
|
@ -56,6 +56,13 @@ type fakeRepo struct {
|
|||
lastEvent string
|
||||
purchases map[string]*PurchaseAudit
|
||||
entitlementErr error
|
||||
applyCalls int
|
||||
}
|
||||
|
||||
func (r *fakeRepo) applyCallCount() int {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.applyCalls
|
||||
}
|
||||
|
||||
func newFakeRepo() *fakeRepo {
|
||||
|
|
@ -135,6 +142,7 @@ func (r *fakeRepo) FindUIDByCustomer(_ context.Context, customer string) (int64,
|
|||
func (r *fakeRepo) ApplyEntitlement(_ context.Context, u EntitlementUpdate) (bool, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.applyCalls++
|
||||
if r.entitlementErr != nil {
|
||||
return false, r.entitlementErr
|
||||
}
|
||||
|
|
@ -150,15 +158,19 @@ func (r *fakeRepo) ApplyEntitlement(_ context.Context, u EntitlementUpdate) (boo
|
|||
}
|
||||
return true, nil
|
||||
}
|
||||
func (r *fakeRepo) WebhookProcessed(_ context.Context, id string) (bool, error) {
|
||||
func (r *fakeRepo) ClaimWebhook(_ context.Context, id, _ string, _ int64) (bool, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.processed[id], nil
|
||||
if r.processed[id] {
|
||||
return false, nil
|
||||
}
|
||||
func (r *fakeRepo) MarkWebhookProcessed(_ context.Context, id, _ string, _ int64) error {
|
||||
r.processed[id] = true
|
||||
return true, nil
|
||||
}
|
||||
func (r *fakeRepo) ReleaseWebhookClaim(_ context.Context, id string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.processed[id] = true
|
||||
delete(r.processed, id)
|
||||
return nil
|
||||
}
|
||||
func (r *fakeRepo) UpsertPurchase(_ context.Context, p *PurchaseAudit) error {
|
||||
|
|
@ -225,6 +237,58 @@ func TestSubscriptionRemainsReadableWhenStripeDisabled(t *testing.T) {
|
|||
require.Equal(t, PlanFree, state.PlanID)
|
||||
}
|
||||
|
||||
// Stripe can deliver the same event concurrently. The claim has to be atomic, otherwise both
|
||||
// deliveries see "not processed yet" and fulfil the event twice.
|
||||
func TestWebhookConcurrentDeliveriesProcessEventOnce(t *testing.T) {
|
||||
r := newFakeRepo()
|
||||
r.state = SubscriptionState{UID: 7, PlanID: PlanFree, CustomerID: "cus_1"}
|
||||
p := &fakeProvider{event: &Event{
|
||||
ID: "evt_concurrent", Type: "customer.subscription.updated", Created: 20,
|
||||
ObjectID: "sub_1", CustomerID: "cus_1", PriceID: "price_pro", Status: "active",
|
||||
}}
|
||||
s := newService(r, p)
|
||||
|
||||
const deliveries = 8
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, deliveries)
|
||||
for i := 0; i < deliveries; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
errs <- s.HandleWebhook(context.Background(), []byte("payload"), "sig")
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.Equal(t, 1, r.applyCallCount(), "concurrent duplicates must fulfil the event exactly once")
|
||||
require.Equal(t, PlanPro, r.state.PlanID)
|
||||
require.Len(t, r.purchases, 1)
|
||||
}
|
||||
|
||||
// Claiming up front must not swallow the event when processing fails, or a transient outage
|
||||
// would permanently lose a payment.
|
||||
func TestWebhookFailedProcessingIsRetryable(t *testing.T) {
|
||||
r := newFakeRepo()
|
||||
r.state = SubscriptionState{UID: 7, PlanID: PlanFree, CustomerID: "cus_1"}
|
||||
r.entitlementErr = errors.New("mongo unavailable")
|
||||
p := &fakeProvider{event: &Event{
|
||||
ID: "evt_retry", Type: "customer.subscription.updated", Created: 20,
|
||||
ObjectID: "sub_1", CustomerID: "cus_1", PriceID: "price_pro", Status: "active",
|
||||
}}
|
||||
s := newService(r, p)
|
||||
|
||||
require.Error(t, s.HandleWebhook(context.Background(), []byte("payload"), "sig"))
|
||||
require.Equal(t, PlanFree, r.state.PlanID)
|
||||
|
||||
// Stripe retries once the outage clears; the claim must have been released.
|
||||
r.entitlementErr = nil
|
||||
require.NoError(t, s.HandleWebhook(context.Background(), []byte("payload"), "sig"))
|
||||
require.Equal(t, PlanPro, r.state.PlanID)
|
||||
}
|
||||
|
||||
func TestWebhookDuplicateOutOfOrderDowngradeAndAdminOverride(t *testing.T) {
|
||||
r := newFakeRepo()
|
||||
r.state = SubscriptionState{UID: 7, PlanID: PlanFree, CustomerID: "cus_1"}
|
||||
|
|
|
|||
|
|
@ -6,11 +6,14 @@ import (
|
|||
"encoding/hex"
|
||||
"fmt"
|
||||
"math"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strings"
|
||||
"unicode"
|
||||
|
||||
"apps/backend/internal/module/growth/domain"
|
||||
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
)
|
||||
|
||||
// --- Playbooks ---
|
||||
|
|
@ -224,10 +227,28 @@ func tokenizeTC(s string) []string {
|
|||
|
||||
// --- UTM ---
|
||||
|
||||
func (s *Service) CreateUtmLink(ctx context.Context, ownerUID int64, dest, outcomeID, label string) (*domain.UtmLink, error) {
|
||||
// validateRedirectTarget keeps /u/{code} from becoming an open redirect that can point anywhere.
|
||||
// A HasPrefix("http") check is not enough: it lets through "httpx://", scheme-relative "//evil",
|
||||
// and anything unparseable, all of which the browser resolves differently than we assume.
|
||||
func validateRedirectTarget(dest string) (string, error) {
|
||||
dest = strings.TrimSpace(dest)
|
||||
if dest == "" || !strings.HasPrefix(dest, "http") {
|
||||
return nil, fmt.Errorf("%w: destination_url must be http(s)", domain.ErrValidation)
|
||||
u, err := url.Parse(dest)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: destination_url is not a valid URL", domain.ErrValidation)
|
||||
}
|
||||
if u.Scheme != "http" && u.Scheme != "https" {
|
||||
return "", fmt.Errorf("%w: destination_url must be http or https", domain.ErrValidation)
|
||||
}
|
||||
if u.Host == "" {
|
||||
return "", fmt.Errorf("%w: destination_url must include a host", domain.ErrValidation)
|
||||
}
|
||||
return u.String(), nil
|
||||
}
|
||||
|
||||
func (s *Service) CreateUtmLink(ctx context.Context, ownerUID int64, dest, outcomeID, label string) (*domain.UtmLink, error) {
|
||||
dest, err := validateRedirectTarget(dest)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
code := domain.NewID()[:10]
|
||||
u := &domain.UtmLink{
|
||||
|
|
@ -259,10 +280,14 @@ func (s *Service) TrackUtmClick(ctx context.Context, code string) (dest string,
|
|||
e.Status = domain.StatusConfirmed
|
||||
}
|
||||
e.UpdatedAt = domain.NowNano()
|
||||
_ = s.Repo.SaveOutcome(ctx, e)
|
||||
if serr := s.Repo.SaveOutcome(ctx, e); serr != nil {
|
||||
// 點擊已經發生,redirect 照走;但歸因掉了要看得見,否則成效數字會默默偏低。
|
||||
logx.WithContext(ctx).Errorf("utm click attribution not saved: outcome=%s err=%v", e.ID, serr)
|
||||
}
|
||||
}
|
||||
return u.DestinationURL, nil
|
||||
}
|
||||
// 舊資料是在只檢查 "http" 前綴時寫入的,出門前再驗一次。
|
||||
return validateRedirectTarget(u.DestinationURL)
|
||||
}
|
||||
|
||||
// --- Benchmark ---
|
||||
|
|
|
|||
|
|
@ -168,3 +168,55 @@ func TestWS_05_ExportReport(t *testing.T) {
|
|||
require.Contains(t, fn, "2026-07")
|
||||
require.Contains(t, body, "月報")
|
||||
}
|
||||
|
||||
// /u/{code} 是免登入 302,任何存進去的 destination 都會變成本站可用的跳板。
|
||||
// 舊檢查只看 "http" 前綴,"httpx://" 與 "//evil" 都能過。
|
||||
func TestUTM_RejectsDestinationsThatAreNotHttpURLs(t *testing.T) {
|
||||
svc := New(repository.NewMemory())
|
||||
ctx := context.Background()
|
||||
|
||||
for _, dest := range []string{
|
||||
"javascript:alert(1)",
|
||||
"httpx://evil.example",
|
||||
"//evil.example/path",
|
||||
"http://",
|
||||
"data:text/html,<script>alert(1)</script>",
|
||||
"",
|
||||
" ",
|
||||
} {
|
||||
_, err := svc.CreateUtmLink(ctx, 1, dest, "", "label")
|
||||
require.ErrorIsf(t, err, domain.ErrValidation, "destination %q should be rejected", dest)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUTM_AcceptsHttpAndHttpsAndRedirectsThere(t *testing.T) {
|
||||
svc := New(repository.NewMemory())
|
||||
ctx := context.Background()
|
||||
|
||||
for _, dest := range []string{
|
||||
"https://shop.example/product?utm_source=threads",
|
||||
"http://shop.example",
|
||||
} {
|
||||
link, err := svc.CreateUtmLink(ctx, 1, dest, "", "label")
|
||||
require.NoError(t, err)
|
||||
|
||||
got, err := svc.TrackUtmClick(ctx, link.Code)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, dest, got)
|
||||
}
|
||||
}
|
||||
|
||||
// 這個檢查是後來才加的,先前寫進 DB 的列可能還帶著壞值,所以出門前要再驗一次。
|
||||
func TestUTM_RefusesToRedirectToAStoredBadDestination(t *testing.T) {
|
||||
repo := repository.NewMemory()
|
||||
svc := New(repo)
|
||||
ctx := context.Background()
|
||||
|
||||
link, err := svc.CreateUtmLink(ctx, 1, "https://shop.example", "", "label")
|
||||
require.NoError(t, err)
|
||||
link.DestinationURL = "javascript:alert(1)"
|
||||
require.NoError(t, repo.SaveUtmLink(ctx, link))
|
||||
|
||||
_, err = svc.TrackUtmClick(ctx, link.Code)
|
||||
require.ErrorIs(t, err, domain.ErrValidation)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,6 +10,9 @@ var (
|
|||
ErrForbidden = errors.New("inspire forbidden")
|
||||
ErrValidation = errors.New("inspire validation")
|
||||
ErrRemoved = errors.New("legacy inspire API removed")
|
||||
// ErrNotImplemented marks a capability that is wired but has no real backend yet. It must be
|
||||
// returned instead of fabricated data, and must never be billed.
|
||||
ErrNotImplemented = errors.New("not implemented")
|
||||
)
|
||||
|
||||
func NowNano() int64 { return time.Now().UTC().UnixNano() }
|
||||
|
|
|
|||
|
|
@ -224,23 +224,27 @@ func TestIN_06_ResearchSearch(t *testing.T) {
|
|||
require.NotEmpty(t, hits[0].URL)
|
||||
}
|
||||
|
||||
func TestIN_07_GenerateImage(t *testing.T) {
|
||||
// There is no image backend yet, so the endpoint must say so rather than charge for a
|
||||
// placeholder avatar from a public URL.
|
||||
func TestIN_07_GenerateImageReportsMissingBackend(t *testing.T) {
|
||||
svc := newInspire()
|
||||
uid := int64(5_001_017)
|
||||
setupInspireUID(svc, uid)
|
||||
img, err := svc.GenerateImage(context.Background(), uid, "溫暖海邊插畫")
|
||||
require.ErrorIs(t, err, domain.ErrNotImplemented)
|
||||
require.Nil(t, img)
|
||||
|
||||
sum, err := svc.Usage.GetSummary(context.Background(), uid, "")
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, img.URL)
|
||||
require.NotEmpty(t, img.ID)
|
||||
require.Zero(t, sum.Platform.CreditsUsed, "an unimplemented feature must not be billed")
|
||||
}
|
||||
|
||||
func TestIN_08_UploadSeparate(t *testing.T) {
|
||||
func TestIN_08_GenerateImageStillValidatesPrompt(t *testing.T) {
|
||||
svc := newInspire()
|
||||
uid := int64(5_001_008)
|
||||
setupInspireUID(svc, uid)
|
||||
img, err := svc.GenerateImage(context.Background(), uid, "x")
|
||||
require.NoError(t, err)
|
||||
require.True(t, len(img.URL) > 8)
|
||||
_, err := svc.GenerateImage(context.Background(), uid, " ")
|
||||
require.ErrorIs(t, err, domain.ErrValidation)
|
||||
}
|
||||
|
||||
func TestIN_09_LegacyRemoved(t *testing.T) {
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ import (
|
|||
usageUC "apps/backend/internal/module/usage/usecase"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
|
|
@ -54,14 +55,16 @@ func (s *Service) ListTrends(ctx context.Context, ownerUID int64) ([]*domain.Tre
|
|||
|
||||
// RefreshTrends — 手動「找靈感」:取得真實結果後才扣點並覆蓋快取。
|
||||
// 搜尋失敗時保留舊內容;不可用示意種子冒充本次刷新成功。
|
||||
func (s *Service) RefreshTrends(ctx context.Context, ownerUID int64) ([]*domain.TrendItem, error) {
|
||||
func (s *Service) RefreshTrends(ctx context.Context, ownerUID int64) (_ []*domain.TrendItem, err error) {
|
||||
list := s.tryLiveTrends(ctx, ownerUID)
|
||||
if len(list) == 0 {
|
||||
return nil, fmt.Errorf("找不到新的可用話題,已保留原本內容")
|
||||
}
|
||||
if err := s.bill(ctx, ownerUID, usageDomain.MeterWebSearch, "找靈感話題", "inspire.refreshTrends"); err != nil {
|
||||
charge, err := s.bill(ctx, ownerUID, usageDomain.MeterWebSearch, "找靈感話題", "inspire.refreshTrends")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
if err := s.Repo.SaveTrends(ctx, ownerUID, list); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
|
@ -687,7 +690,7 @@ func generateUserLogMessage(notes, material string) string {
|
|||
return "【產文】" + notes + "|素材:" + preview
|
||||
}
|
||||
|
||||
func (s *Service) chatInternal(ctx context.Context, ownerUID int64, message string, pinnedIDs []string, mode, personaID, sessionID, material string, useWeb bool, onDelta func(chunk string) error) (*ChatOutcome, error) {
|
||||
func (s *Service) chatInternal(ctx context.Context, ownerUID int64, message string, pinnedIDs []string, mode, personaID, sessionID, material string, useWeb bool, onDelta func(chunk string) error) (_ *ChatOutcome, err error) {
|
||||
message = strings.TrimSpace(message)
|
||||
material = strings.TrimSpace(material)
|
||||
if mode != "chat" && mode != "generate" {
|
||||
|
|
@ -714,9 +717,11 @@ func (s *Service) chatInternal(ctx context.Context, ownerUID int64, message stri
|
|||
if mode == "generate" {
|
||||
label = "inspire rewrite"
|
||||
}
|
||||
if err := s.bill(ctx, ownerUID, usageDomain.MeterAICopy, label, "inspire.chat"); err != nil {
|
||||
charge, err := s.bill(ctx, ownerUID, usageDomain.MeterAICopy, label, "inspire.chat")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
// 回覆語言:人設指紋/範例優先,其次會員 UI 語系(ctx 已由 Auth 注入)
|
||||
persona := s.resolvePersonaSnap(ctx, ownerUID, personaID)
|
||||
if mode == "generate" && (persona == nil || persona.Status != "ready") {
|
||||
|
|
@ -1153,14 +1158,16 @@ func (s *Service) buildInspirePrompt(ctx context.Context, ownerUID int64, pinned
|
|||
return full, blocks, sections
|
||||
}
|
||||
|
||||
func (s *Service) searchPromptContext(ctx context.Context, ownerUID int64, query string) (string, error) {
|
||||
func (s *Service) searchPromptContext(ctx context.Context, ownerUID int64, query string) (_ string, err error) {
|
||||
query = strings.TrimSpace(query)
|
||||
if query == "" {
|
||||
return "", nil
|
||||
}
|
||||
if err := s.bill(ctx, ownerUID, usageDomain.MeterWebSearch, "inspire web search", "inspire.chat"); err != nil {
|
||||
charge, err := s.bill(ctx, ownerUID, usageDomain.MeterWebSearch, "inspire web search", "inspire.chat")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
key := "fake"
|
||||
if s.ResolveKey != nil {
|
||||
if _, resolved, err := s.ResolveKey(ctx, ownerUID, usageDomain.MeterWebSearch); err == nil && resolved != "" {
|
||||
|
|
@ -1438,14 +1445,16 @@ func (s *Service) runInspireLLM(ctx context.Context, ownerUID int64, prompt stri
|
|||
return "", fmt.Errorf("%w: 請到設定填寫 AI Key", domain.ErrValidation)
|
||||
}
|
||||
|
||||
func (s *Service) ResearchSearch(ctx context.Context, ownerUID int64, query string) ([]*domain.ResearchHit, error) {
|
||||
func (s *Service) ResearchSearch(ctx context.Context, ownerUID int64, query string) (_ []*domain.ResearchHit, err error) {
|
||||
query = strings.TrimSpace(query)
|
||||
if query == "" {
|
||||
return nil, fmt.Errorf("%w: empty query", domain.ErrValidation)
|
||||
}
|
||||
if err := s.bill(ctx, ownerUID, usageDomain.MeterWebSearch, "research search", "research.search"); err != nil {
|
||||
charge, err := s.bill(ctx, ownerUID, usageDomain.MeterWebSearch, "research search", "research.search")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
key := "fake"
|
||||
if s.ResolveKey != nil {
|
||||
if _, k, err := s.ResolveKey(ctx, ownerUID, usageDomain.MeterWebSearch); err == nil && k != "" {
|
||||
|
|
@ -1473,34 +1482,73 @@ func (s *Service) ResearchSearch(ctx context.Context, ownerUID int64, query stri
|
|||
return out, nil
|
||||
}
|
||||
|
||||
func (s *Service) GenerateImage(ctx context.Context, ownerUID int64, prompt string) (*domain.GeneratedImage, error) {
|
||||
prompt = strings.TrimSpace(prompt)
|
||||
if prompt == "" {
|
||||
// GenerateImage has no image backend yet. It used to return a random avatar URL from a public
|
||||
// placeholder service while charging a real ai_image credit, which billed members for a picture
|
||||
// nobody asked for. Until a real generator is wired in, this reports the gap instead.
|
||||
func (s *Service) GenerateImage(_ context.Context, _ int64, prompt string) (*domain.GeneratedImage, error) {
|
||||
if strings.TrimSpace(prompt) == "" {
|
||||
return nil, fmt.Errorf("%w: empty prompt", domain.ErrValidation)
|
||||
}
|
||||
if err := s.bill(ctx, ownerUID, usageDomain.MeterAIImage, "generate image", "media.generateImage"); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// placeholder SVG data URL (live-complete with billable path; real SD later)
|
||||
id := "img_" + uuid.NewString()[:10]
|
||||
// minimal 1x1 png base64 is tiny; use dicebear-like public URL for preview
|
||||
url := fmt.Sprintf("https://api.dicebear.com/9.x/shapes/svg?seed=%s", id)
|
||||
return &domain.GeneratedImage{ID: id, URL: url, Prompt: prompt}, nil
|
||||
return nil, fmt.Errorf("%w: 圖片生成尚未接上生成服務", domain.ErrNotImplemented)
|
||||
}
|
||||
|
||||
// Legacy removed APIs
|
||||
func (s *Service) LegacyRemoved() error { return domain.ErrRemoved }
|
||||
|
||||
func (s *Service) bill(ctx context.Context, uid int64, meter, label, source string) error {
|
||||
if s.Usage == nil {
|
||||
return nil
|
||||
// inspireCharge is one reserved credit; see Settle/Release for the commit-on-success contract.
|
||||
type inspireCharge struct {
|
||||
svc *Service
|
||||
uid int64
|
||||
meter string
|
||||
mode string
|
||||
label string
|
||||
source string
|
||||
settled bool
|
||||
}
|
||||
|
||||
// bill reserves credit for one metered call. Pair it with `defer charge.Settle(ctx, &err)` on a
|
||||
// named error return so every failure path refunds rather than charging for nothing.
|
||||
func (s *Service) bill(ctx context.Context, uid int64, meter, label, source string) (*inspireCharge, error) {
|
||||
if s == nil || s.Usage == nil {
|
||||
return &inspireCharge{settled: true}, nil
|
||||
}
|
||||
mode, err := s.Usage.PrepareCall(ctx, uid, meter)
|
||||
if err != nil {
|
||||
return err
|
||||
return nil, err
|
||||
}
|
||||
return &inspireCharge{svc: s, uid: uid, meter: meter, mode: mode, label: label, source: source}, nil
|
||||
}
|
||||
|
||||
func (c *inspireCharge) Settle(ctx context.Context, errp *error) {
|
||||
if errp != nil && *errp != nil {
|
||||
c.Release(ctx)
|
||||
return
|
||||
}
|
||||
c.Commit(ctx)
|
||||
}
|
||||
|
||||
// Commit writes the audit event. The member already got their result, so a failure here is
|
||||
// logged rather than turned into an error.
|
||||
func (c *inspireCharge) Commit(ctx context.Context) {
|
||||
if c == nil || c.settled {
|
||||
return
|
||||
}
|
||||
c.settled = true
|
||||
if _, err := c.svc.Usage.RecordCall(ctx, c.uid, c.meter, c.mode, c.label, c.source); err != nil {
|
||||
logx.Errorf("usage record uid=%d meter=%s source=%s: %v", c.uid, c.meter, c.source, err)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *inspireCharge) Release(ctx context.Context) {
|
||||
if c == nil || c.settled {
|
||||
return
|
||||
}
|
||||
c.settled = true
|
||||
rctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
|
||||
defer cancel()
|
||||
if err := c.svc.Usage.ReleaseCall(rctx, c.uid, c.meter, c.mode); err != nil {
|
||||
logx.Errorf("usage release uid=%d meter=%s source=%s: %v", c.uid, c.meter, c.source, err)
|
||||
}
|
||||
_, err = s.Usage.RecordCall(ctx, uid, meter, mode, label, source)
|
||||
return err
|
||||
}
|
||||
|
||||
func emptySession(ownerUID int64) *domain.Session {
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import (
|
|||
libmongo "apps/backend/internal/lib/mongo"
|
||||
"apps/backend/internal/module/job/domain"
|
||||
|
||||
"github.com/zeromicro/go-zero/core/logx"
|
||||
"github.com/zeromicro/go-zero/core/stores/mon"
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
|
|
@ -216,12 +217,22 @@ func (s *MonStore) CancelPendingByRef(ctx context.Context, ownerUID int64, templ
|
|||
if err := s.jobs.Find(ctx, &list, filter); err != nil {
|
||||
return err
|
||||
}
|
||||
// Update is optimistic-concurrency guarded, so a lost race means the job is still pending and
|
||||
// will run alongside its replacement. Swallowing that error reported a clean cancel that never
|
||||
// happened; a stale job racing the new one is worth failing the call over.
|
||||
for _, j := range list {
|
||||
j.Status = domain.StatusCancelled
|
||||
j.ProgressSummary = "已由新的定期排程取代"
|
||||
j.CompletedAt = now
|
||||
j.UpdatedAt = now
|
||||
_ = s.Update(ctx, j)
|
||||
if err := s.Update(ctx, j); err != nil {
|
||||
// 這一輪已經被別人搶去改(例如 worker 剛領走),不是我們該取消的目標。
|
||||
if errors.Is(err, domain.ErrIllegalStatus) {
|
||||
logx.WithContext(ctx).Infof("cancel pending job skipped, changed concurrently: job=%s", j.ID)
|
||||
continue
|
||||
}
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -128,7 +128,9 @@ func (s *Service) StartDemo(ctx context.Context, ownerUID int64) (*domain.Job, e
|
|||
j.Status = domain.StatusQueued
|
||||
j.ProgressSummary = "Demo 測試任務 · 等待 worker 領取"
|
||||
j.UpdatedAt = domain.NowNano()
|
||||
_ = s.Repo.Update(ctx, j)
|
||||
if err := s.Repo.Update(ctx, j); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
s.notify(ctx, j)
|
||||
return j, nil
|
||||
}
|
||||
|
|
@ -288,7 +290,10 @@ func (s *Service) ScheduleTokenRenew(ctx context.Context, ownerUID int64, accoun
|
|||
if runAfter <= 0 {
|
||||
runAfter = now + domain.DefaultTokenRenewDelayNs
|
||||
}
|
||||
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplateThreadsTokenRenew, accountID)
|
||||
// 取消沒成功就別排新的,否則同一個帳號會有兩個 renew job 一起跑。
|
||||
if err := s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplateThreadsTokenRenew, accountID); err != nil {
|
||||
return err
|
||||
}
|
||||
j := &domain.Job{
|
||||
ID: uuid.NewString(), OwnerUID: ownerUID,
|
||||
TemplateType: domain.TemplateThreadsTokenRenew,
|
||||
|
|
@ -325,6 +330,15 @@ type PersonaAnalyzeTextPayload struct {
|
|||
Lang string `json:"lang,omitempty"`
|
||||
}
|
||||
|
||||
// cancelPersonaAnalyzeJobs 收掉同一個人設還在排隊的兩種分析,兩者只能有一個在跑,
|
||||
// 否則兩份結果會互相覆寫成不確定的人設。
|
||||
func (s *Service) cancelPersonaAnalyzeJobs(ctx context.Context, ownerUID int64, personaID string) error {
|
||||
if err := s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeAccount, personaID); err != nil {
|
||||
return err
|
||||
}
|
||||
return s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeText, personaID)
|
||||
}
|
||||
|
||||
// SchedulePersonaAnalyzeAccount — 立刻可領;ref=personaID;同人設舊佇列取消
|
||||
func (s *Service) SchedulePersonaAnalyzeAccount(ctx context.Context, ownerUID int64, personaID, username, lang string) (*domain.Job, error) {
|
||||
if ownerUID <= 0 || personaID == "" {
|
||||
|
|
@ -334,8 +348,9 @@ func (s *Service) SchedulePersonaAnalyzeAccount(ctx context.Context, ownerUID in
|
|||
if username == "" {
|
||||
return nil, fmt.Errorf("empty username")
|
||||
}
|
||||
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeAccount, personaID)
|
||||
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeText, personaID)
|
||||
if err := s.cancelPersonaAnalyzeJobs(ctx, ownerUID, personaID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if s.MaxActiveAccountScrapes > 0 {
|
||||
active, err := s.Repo.CountActiveByOwnerAndTemplate(ctx, ownerUID, domain.TemplatePersonaAnalyzeAccount)
|
||||
if err != nil {
|
||||
|
|
@ -421,7 +436,9 @@ func (s *Service) SchedulePlayGenerateScript(ctx context.Context, ownerUID int64
|
|||
return nil, domain.ErrForbidden
|
||||
}
|
||||
playID = strings.TrimSpace(playID)
|
||||
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePlayGenerateScript, playID)
|
||||
if err := s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePlayGenerateScript, playID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body, _ := json.Marshal(PlayGenerateScriptPayload{PlayID: playID, OnlyEmpty: onlyEmpty})
|
||||
now := domain.NowNano()
|
||||
j := &domain.Job{
|
||||
|
|
@ -476,8 +493,9 @@ func (s *Service) SchedulePersonaAnalyzeText(ctx context.Context, ownerUID int64
|
|||
if rawText == "" {
|
||||
return nil, fmt.Errorf("empty text")
|
||||
}
|
||||
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeAccount, personaID)
|
||||
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeText, personaID)
|
||||
if err := s.cancelPersonaAnalyzeJobs(ctx, ownerUID, personaID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body, _ := json.Marshal(PersonaAnalyzeTextPayload{
|
||||
RawText: rawText, SourceLabel: strings.TrimSpace(sourceLabel), Lang: lang,
|
||||
})
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@ package usecase_test
|
|||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
|
|
@ -76,71 +78,29 @@ func newStudio() (*usecase.Service, *publish.FakeTransport, *memAccounts) {
|
|||
return svc, tp, acc
|
||||
}
|
||||
|
||||
func withUsage(svc *usecase.Service, uid int64) {
|
||||
// unlimited + platform key
|
||||
_ = svc.Usage.Repo.SavePrefs(context.Background(), &usageDomain.MemberPrefs{
|
||||
UID: uid, PlanID: usageDomain.PlanPro, Unlimited: true, UpdatedAt: domain.NowNano(),
|
||||
})
|
||||
if sr, ok := svc.Usage.Resolver.(*usageUC.StaticResolver); ok {
|
||||
if sr.Map == nil {
|
||||
sr.Map = map[string]string{}
|
||||
}
|
||||
sr.Map[fmtUIDMeter(uid, usageDomain.MeterAICopy)] = usageDomain.KeyModePlatform
|
||||
}
|
||||
}
|
||||
|
||||
func fmtUIDMeter(uid int64, meter string) string {
|
||||
return string(rune(0)) // placeholder — use fmt
|
||||
}
|
||||
|
||||
func setupUID(svc *usecase.Service, uid int64) {
|
||||
setupUIDWithPlan(svc, uid, usageDomain.PlanPro, true)
|
||||
}
|
||||
|
||||
func setupUIDWithPlan(svc *usecase.Service, uid int64, planID string, unlimited bool) {
|
||||
_ = svc.Usage.Repo.SavePrefs(context.Background(), &usageDomain.MemberPrefs{
|
||||
UID: uid, PlanID: usageDomain.PlanPro, Unlimited: true, UpdatedAt: domain.NowNano(),
|
||||
UID: uid, PlanID: planID, Unlimited: unlimited, UpdatedAt: domain.NowNano(),
|
||||
})
|
||||
if sr, ok := svc.Usage.Resolver.(*usageUC.StaticResolver); ok {
|
||||
if sr.Map == nil {
|
||||
sr.Map = map[string]string{}
|
||||
}
|
||||
key := ""
|
||||
// StaticResolver uses fmt.Sprintf("%d:%s", uid, meter)
|
||||
key = sprintf("%d:%s", uid, usageDomain.MeterAICopy)
|
||||
sr.Map[key] = usageDomain.KeyModePlatform
|
||||
// StaticResolver keys on "uid:meter".
|
||||
sr.Map[fmt.Sprintf("%d:%s", uid, usageDomain.MeterAICopy)] = usageDomain.KeyModePlatform
|
||||
}
|
||||
}
|
||||
|
||||
func sprintf(f string, a ...any) string {
|
||||
return format(f, a...)
|
||||
}
|
||||
|
||||
// avoid importing fmt in helpers clutter — use strconv
|
||||
func format(f string, a ...any) string {
|
||||
// minimal for our key only
|
||||
if f == "%d:%s" && len(a) == 2 {
|
||||
return itoa(a[0].(int64)) + ":" + a[1].(string)
|
||||
}
|
||||
return f
|
||||
}
|
||||
|
||||
func itoa(n int64) string {
|
||||
if n == 0 {
|
||||
return "0"
|
||||
}
|
||||
neg := n < 0
|
||||
if neg {
|
||||
n = -n
|
||||
}
|
||||
var b [32]byte
|
||||
i := len(b)
|
||||
for n > 0 {
|
||||
i--
|
||||
b[i] = byte('0' + n%10)
|
||||
n /= 10
|
||||
}
|
||||
if neg {
|
||||
i--
|
||||
b[i] = '-'
|
||||
}
|
||||
return string(b[i:])
|
||||
// platformCreditsUsed reports what the member has actually been charged this month.
|
||||
func platformCreditsUsed(t *testing.T, svc *usecase.Service, uid int64) int {
|
||||
t.Helper()
|
||||
sum, err := svc.Usage.GetSummary(context.Background(), uid, "")
|
||||
require.NoError(t, err)
|
||||
return sum.Platform.CreditsUsed
|
||||
}
|
||||
|
||||
func addAccount(acc *memAccounts, ownerUID int64, id, user string) {
|
||||
|
|
@ -150,6 +110,86 @@ func addAccount(acc *memAccounts, ownerUID int64, id, user string) {
|
|||
}
|
||||
}
|
||||
|
||||
// ---- AI billing ----
|
||||
|
||||
// A request the service itself rejects must not consume credits. These used to charge the
|
||||
// member up front, so every early return below was a silent debit for work never done.
|
||||
func TestBilling_RejectedRequestIsNotCharged(t *testing.T) {
|
||||
uid := int64(4_009_001)
|
||||
cases := []struct {
|
||||
name string
|
||||
call func(svc *usecase.Service) error
|
||||
}{
|
||||
{
|
||||
name: "mimic without persona",
|
||||
call: func(svc *usecase.Service) error {
|
||||
_, err := svc.Mimic(context.Background(), uid, "來源文字", "", "", "")
|
||||
return err
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "mimic with unknown persona",
|
||||
call: func(svc *usecase.Service) error {
|
||||
_, err := svc.Mimic(context.Background(), uid, "來源文字", "does-not-exist", "", "")
|
||||
return err
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "analyze post that has no text",
|
||||
call: func(svc *usecase.Service) error {
|
||||
post := &domain.OwnPost{
|
||||
ID: "post-no-text", OwnerUID: uid, Text: "", PublishedAt: domain.NowNano(),
|
||||
}
|
||||
require.NoError(t, svc.Repo.SaveOwnPost(context.Background(), post))
|
||||
_, err := svc.AnalyzePost(context.Background(), uid, post.ID)
|
||||
return err
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "generate play step without context",
|
||||
call: func(svc *usecase.Service) error {
|
||||
_, err := svc.GeneratePlayStep(context.Background(), uid, "", "", "", "", true, "root")
|
||||
return err
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
svc, _, _ := newStudio()
|
||||
// A metered plan, so any charge is visible.
|
||||
setupUIDWithPlan(svc, uid, usageDomain.PlanPro, false)
|
||||
require.Error(t, tc.call(svc))
|
||||
require.Zero(t, platformCreditsUsed(t, svc, uid), "rejected request must not be charged")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// A persona analysis that cannot gather enough samples fails after billing has been reserved.
|
||||
func TestBilling_FailedPersonaAnalysisIsRefunded(t *testing.T) {
|
||||
svc, _, _ := newStudio()
|
||||
uid := int64(4_009_002)
|
||||
setupUIDWithPlan(svc, uid, usageDomain.PlanPro, false)
|
||||
|
||||
p, err := svc.SavePersona(context.Background(), uid, &domain.Persona{Name: "小海"})
|
||||
require.NoError(t, err)
|
||||
|
||||
// One sample is below the two-sample minimum, so the analysis is rejected.
|
||||
_, err = svc.ExecuteAnalyzeFromText(context.Background(), uid, p.ID, "只有一段文字", "手動", nil)
|
||||
require.Error(t, err)
|
||||
require.Zero(t, platformCreditsUsed(t, svc, uid))
|
||||
}
|
||||
|
||||
// The success path must still charge exactly once.
|
||||
func TestBilling_SuccessfulCallIsChargedOnce(t *testing.T) {
|
||||
svc, _, _ := newStudio()
|
||||
uid := int64(4_009_003)
|
||||
setupUIDWithPlan(svc, uid, usageDomain.PlanPro, false)
|
||||
|
||||
_, err := svc.AnalyzeViral(context.Background(), uid, "這是一段要分析的貼文內容,講一個轉折。")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, usageDomain.MeterCost(usageDomain.MeterAICopy), platformCreditsUsed(t, svc, uid))
|
||||
}
|
||||
|
||||
// ---- PE ----
|
||||
|
||||
func TestPE_01_CreateSave(t *testing.T) {
|
||||
|
|
@ -915,6 +955,64 @@ func TestOP_11_GenerateFromFormulaRemoved(t *testing.T) {
|
|||
require.ErrorIs(t, err, domain.ErrFormulaRemove)
|
||||
}
|
||||
|
||||
// fakeInsightsMedia serves a fixed post list and lets a test make the insights call fail, which
|
||||
// is what a Threads rate limit looks like mid-sync.
|
||||
type fakeInsightsMedia struct {
|
||||
threads []usecase.FetchedThread
|
||||
insights usecase.FetchedInsights
|
||||
insightsErr error
|
||||
}
|
||||
|
||||
func (f *fakeInsightsMedia) ListThreads(context.Context, string, int) ([]usecase.FetchedThread, error) {
|
||||
return f.threads, nil
|
||||
}
|
||||
func (f *fakeInsightsMedia) GetInsights(context.Context, string, string) (usecase.FetchedInsights, error) {
|
||||
if f.insightsErr != nil {
|
||||
return usecase.FetchedInsights{}, f.insightsErr
|
||||
}
|
||||
return f.insights, nil
|
||||
}
|
||||
func (f *fakeInsightsMedia) ListConversation(context.Context, string, string, int) ([]usecase.FetchedReply, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeInsightsMedia) ListMentions(context.Context, string, string, int) ([]usecase.FetchedMention, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeInsightsMedia) ListProfilePosts(context.Context, string, string, int) ([]usecase.FetchedThread, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// A single rate-limited insights call used to overwrite a post's real like/view counts with
|
||||
// zeros, which then corrupted the insights summary, benchmark and account health downstream.
|
||||
func TestSyncOwnPosts_FailedInsightsKeepsPreviousCounts(t *testing.T) {
|
||||
svc, _, acc := newStudio()
|
||||
uid := int64(4_004_020)
|
||||
setupUID(svc, uid)
|
||||
addAccount(acc, uid, "acc1", "me")
|
||||
|
||||
media := &fakeInsightsMedia{
|
||||
threads: []usecase.FetchedThread{{ID: "m_1", Text: "貼文", PublishedAt: domain.NowNano()}},
|
||||
insights: usecase.FetchedInsights{Views: 500, Likes: 42, Replies: 7, Status: "ok"},
|
||||
}
|
||||
svc.Media = media
|
||||
|
||||
first, err := svc.SyncOwnPosts(context.Background(), uid, "acc1")
|
||||
require.NoError(t, err)
|
||||
require.Len(t, first, 1)
|
||||
require.Equal(t, 42, first[0].LikeCount)
|
||||
require.Equal(t, 500, first[0].ViewCount)
|
||||
|
||||
// Threads starts rate limiting; the next sync must not zero the numbers out.
|
||||
media.insightsErr = errors.New("rate limited")
|
||||
second, err := svc.SyncOwnPosts(context.Background(), uid, "acc1")
|
||||
require.NoError(t, err)
|
||||
require.Len(t, second, 1)
|
||||
require.Equal(t, 42, second[0].LikeCount, "a failed insights fetch must not erase real counts")
|
||||
require.Equal(t, 500, second[0].ViewCount)
|
||||
require.Equal(t, 7, second[0].ReplyCount)
|
||||
require.Equal(t, "error", second[0].InsightsStatus, "but the failure must still be visible")
|
||||
}
|
||||
|
||||
type fakeMentionsMedia struct {
|
||||
hits []usecase.FetchedMention
|
||||
err error
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ import (
|
|||
)
|
||||
|
||||
// PersonaPreview:依指紋一次 LLM 產主貼 + 回文(可抓新聞當話題靈感)。
|
||||
func (s *Service) PersonaPreview(ctx context.Context, ownerUID int64, personaID, topic string, useNews bool) (*domain.PersonaPreview, error) {
|
||||
func (s *Service) PersonaPreview(ctx context.Context, ownerUID int64, personaID, topic string, useNews bool) (_ *domain.PersonaPreview, err error) {
|
||||
if ownerUID <= 0 {
|
||||
return nil, domain.ErrForbidden
|
||||
}
|
||||
|
|
@ -36,9 +36,11 @@ func (s *Service) PersonaPreview(ctx context.Context, ownerUID int64, personaID,
|
|||
return nil, fmt.Errorf("%w: 人設尚無指紋,請先完成分析或手動填寫指紋", domain.ErrValidation)
|
||||
}
|
||||
|
||||
if err := s.billAI(ctx, ownerUID, "persona preview", "compose.personaPreview"); err != nil {
|
||||
charge, err := s.billAI(ctx, ownerUID, "persona preview", "compose.personaPreview")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
|
||||
topic = strings.TrimSpace(topic)
|
||||
topicSource := "manual"
|
||||
|
|
|
|||
|
|
@ -279,7 +279,7 @@ func (s *Service) AnalyzeFromAccount(ctx context.Context, ownerUID int64, id, us
|
|||
}
|
||||
|
||||
// ExecuteAnalyzeFromText — worker/測試:扣費 + LLM 分析 + 存檔 ready。
|
||||
func (s *Service) ExecuteAnalyzeFromText(ctx context.Context, ownerUID int64, id, rawText, sourceLabel string, onProgress ProgressFn) (*domain.Persona, error) {
|
||||
func (s *Service) ExecuteAnalyzeFromText(ctx context.Context, ownerUID int64, id, rawText, sourceLabel string, onProgress ProgressFn) (_ *domain.Persona, err error) {
|
||||
report := func(pct int, sum string) {
|
||||
if onProgress != nil {
|
||||
onProgress(pct, sum)
|
||||
|
|
@ -299,10 +299,12 @@ func (s *Service) ExecuteAnalyzeFromText(ctx context.Context, ownerUID int64, id
|
|||
return nil, fmt.Errorf("%w: 請至少貼 2 段參考文字(可用 --- 分隔),每段至少約 10 字", domain.ErrValidation)
|
||||
}
|
||||
report(25, "人設分析 · 檢查文字樣本…")
|
||||
if err := s.billAI(ctx, ownerUID, "persona analyze text", "personas.analyzeFromText"); err != nil {
|
||||
charge, err := s.billAI(ctx, ownerUID, "persona analyze text", "personas.analyzeFromText")
|
||||
if err != nil {
|
||||
_ = s.markPersonaAnalyzeFailed(ctx, p, err.Error())
|
||||
return nil, err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
p.Status = domain.PersonaAnalyzing
|
||||
p.UpdatedAt = domain.NowNano()
|
||||
_ = s.Repo.SavePersona(ctx, p)
|
||||
|
|
@ -324,7 +326,7 @@ func (s *Service) ExecuteAnalyzeFromText(ctx context.Context, ownerUID int64, id
|
|||
}
|
||||
|
||||
// ExecuteAnalyzeFromAccount — worker:爬公開貼文 + LLM + 存檔 ready。
|
||||
func (s *Service) ExecuteAnalyzeFromAccount(ctx context.Context, ownerUID int64, id, username string, onProgress ProgressFn) (*domain.Persona, error) {
|
||||
func (s *Service) ExecuteAnalyzeFromAccount(ctx context.Context, ownerUID int64, id, username string, onProgress ProgressFn) (_ *domain.Persona, err error) {
|
||||
report := func(pct int, sum string) {
|
||||
if onProgress != nil {
|
||||
onProgress(pct, sum)
|
||||
|
|
@ -342,10 +344,12 @@ func (s *Service) ExecuteAnalyzeFromAccount(ctx context.Context, ownerUID int64,
|
|||
return nil, fmt.Errorf("%w: username 請只填帳號,例如 ultralab_tw", domain.ErrValidation)
|
||||
}
|
||||
report(15, fmt.Sprintf("人設分析 · 準備爬取 @%s…", username))
|
||||
if err := s.billAI(ctx, ownerUID, "persona analyze account", "personas.analyzeFromAccount"); err != nil {
|
||||
charge, err := s.billAI(ctx, ownerUID, "persona analyze account", "personas.analyzeFromAccount")
|
||||
if err != nil {
|
||||
_ = s.markPersonaAnalyzeFailed(ctx, p, err.Error())
|
||||
return nil, err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
|
||||
p.Status = domain.PersonaAnalyzing
|
||||
p.Style.BenchmarkUsername = username
|
||||
|
|
@ -807,7 +811,7 @@ func (s *Service) SubmitPlay(ctx context.Context, ownerUID int64, playID string)
|
|||
|
||||
// GeneratePlayScript — 一次 LLM 產完整劇本並寫回 play.steps(onlyEmpty=true 只填空白步)。
|
||||
// 回傳填入步數。供 worker job / 同步測試。
|
||||
func (s *Service) GeneratePlayScript(ctx context.Context, ownerUID int64, playID string, onlyEmpty bool) (int, error) {
|
||||
func (s *Service) GeneratePlayScript(ctx context.Context, ownerUID int64, playID string, onlyEmpty bool) (_ int, err error) {
|
||||
play, err := s.GetPlay(ctx, ownerUID, playID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
|
|
@ -841,9 +845,11 @@ func (s *Service) GeneratePlayScript(ctx context.Context, ownerUID int64, playID
|
|||
if len(needs) == 0 {
|
||||
return 0, fmt.Errorf("%w: 沒有空白步驟可產(全部已有正文)", domain.ErrValidation)
|
||||
}
|
||||
if err := s.billAI(ctx, ownerUID, "play generate script", "plays.generateScript"); err != nil {
|
||||
charge, err := s.billAI(ctx, ownerUID, "play generate script", "plays.generateScript")
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
|
||||
// 主上下文:目標貼文/主題
|
||||
targetCtx := strings.TrimSpace(play.Topic)
|
||||
|
|
@ -1017,15 +1023,17 @@ func cleanPlayGeneratedText(kind, text string) string {
|
|||
|
||||
// GeneratePlayStep — 互回/串場劇本:依人設真 LLM 產一步正文。
|
||||
// 注意:與 compose mimic 相同,OpenCode reasoning 模型可能 30~90s 且偶發空白回覆。
|
||||
func (s *Service) GeneratePlayStep(ctx context.Context, ownerUID int64, personaID, contextText, topic, speakerLabel string, isLead bool, mode string) (string, error) {
|
||||
func (s *Service) GeneratePlayStep(ctx context.Context, ownerUID int64, personaID, contextText, topic, speakerLabel string, isLead bool, mode string) (_ string, err error) {
|
||||
contextText = strings.TrimSpace(contextText)
|
||||
topic = strings.TrimSpace(topic)
|
||||
if contextText == "" && topic == "" {
|
||||
return "", fmt.Errorf("%w: empty context", domain.ErrValidation)
|
||||
}
|
||||
if err := s.billAI(ctx, ownerUID, "play generate step", "plays.generateStep"); err != nil {
|
||||
charge, err := s.billAI(ctx, ownerUID, "play generate step", "plays.generateStep")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
mode = strings.ToLower(strings.TrimSpace(mode))
|
||||
if mode == "" {
|
||||
mode = "reply"
|
||||
|
|
@ -1377,14 +1385,11 @@ func (s *Service) publishWithLease(ctx context.Context, bundleID, stepID, leaseO
|
|||
|
||||
// ---------- Compose (CP) ----------
|
||||
|
||||
func (s *Service) Mimic(ctx context.Context, ownerUID int64, sourceText, personaID, direction, structureNotes string) (string, error) {
|
||||
func (s *Service) Mimic(ctx context.Context, ownerUID int64, sourceText, personaID, direction, structureNotes string) (_ string, err error) {
|
||||
sourceText = strings.TrimSpace(sourceText)
|
||||
if sourceText == "" {
|
||||
return "", fmt.Errorf("%w: empty source", domain.ErrValidation)
|
||||
}
|
||||
if err := s.billAI(ctx, ownerUID, "compose mimic", "compose.mimic"); err != nil {
|
||||
return "", err
|
||||
}
|
||||
personaID = strings.TrimSpace(personaID)
|
||||
if personaID == "" {
|
||||
return "", fmt.Errorf("%w: 請選擇已完成人設分析的人設", domain.ErrValidation)
|
||||
|
|
@ -1396,6 +1401,11 @@ func (s *Service) Mimic(ctx context.Context, ownerUID int64, sourceText, persona
|
|||
if p.Status != domain.PersonaReady {
|
||||
return "", fmt.Errorf("%w: 選擇的人設尚未完成分析", domain.ErrValidation)
|
||||
}
|
||||
charge, err := s.billAI(ctx, ownerUID, "compose mimic", "compose.mimic")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
fp := personaExpressionFingerprintBlock(p)
|
||||
// 指紋過長會拖慢 reasoning 模型
|
||||
if utf8.RuneCountInString(fp) > 800 {
|
||||
|
|
@ -1490,10 +1500,12 @@ func buildMimicPrompt(fp, sourceText, direction, notesBlock string) string {
|
|||
`, direction, fp, sourceText, notesBlock))
|
||||
}
|
||||
|
||||
func (s *Service) AnalyzeViral(ctx context.Context, ownerUID int64, text string) (*domain.ViralAnalysis, error) {
|
||||
if err := s.billAI(ctx, ownerUID, "compose analyze viral", "compose.analyzeViral"); err != nil {
|
||||
func (s *Service) AnalyzeViral(ctx context.Context, ownerUID int64, text string) (_ *domain.ViralAnalysis, err error) {
|
||||
charge, err := s.billAI(ctx, ownerUID, "compose analyze viral", "compose.analyzeViral")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
return s.analyzeViralUnbilled(ctx, ownerUID, text)
|
||||
}
|
||||
|
||||
|
|
@ -1889,6 +1901,9 @@ func (s *Service) syncOwnPostsFromThreads(ctx context.Context, ownerUID int64, a
|
|||
}
|
||||
|
||||
insights := make([]FetchedInsights, len(threads))
|
||||
// Each slot starts as "not fetched" so a rate-limited call, a timeout, or a cancelled sync is
|
||||
// distinguishable from a post that genuinely has zero engagement.
|
||||
insightFetched := make([]bool, len(threads))
|
||||
var wg sync.WaitGroup
|
||||
limit := make(chan struct{}, 4)
|
||||
for i, th := range threads {
|
||||
|
|
@ -1904,7 +1919,12 @@ func (s *Service) syncOwnPostsFromThreads(ctx context.Context, ownerUID int64, a
|
|||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
insights[i], _ = s.Media.GetInsights(ctx, accessToken, mediaID)
|
||||
got, err := s.Media.GetInsights(ctx, accessToken, mediaID)
|
||||
if err != nil {
|
||||
logx.Errorf("threads insights media=%s: %v", mediaID, err)
|
||||
return
|
||||
}
|
||||
insights[i], insightFetched[i] = got, true
|
||||
}(i, th.ID)
|
||||
}
|
||||
wg.Wait()
|
||||
|
|
@ -1917,6 +1937,15 @@ func (s *Service) syncOwnPostsFromThreads(ctx context.Context, ownerUID int64, a
|
|||
ins := insights[i]
|
||||
|
||||
prev := byMedia[th.ID]
|
||||
// A failed insights call must keep whatever counts we already had. Writing zeros would
|
||||
// destroy the real numbers, and they feed insights summary, benchmark and account health.
|
||||
if !insightFetched[i] {
|
||||
if prev != nil {
|
||||
ins.Likes, ins.Replies, ins.Reposts = prev.LikeCount, prev.ReplyCount, prev.RepostCount
|
||||
ins.Quotes, ins.Views, ins.Shares = prev.QuoteCount, prev.ViewCount, prev.ShareCount
|
||||
}
|
||||
ins.Status = "error"
|
||||
}
|
||||
id := "op_" + th.ID
|
||||
if prev != nil && prev.ID != "" {
|
||||
id = prev.ID
|
||||
|
|
@ -2104,7 +2133,7 @@ func mergeReplyStatus(prev, next []domain.OwnPostReply) []domain.OwnPostReply {
|
|||
return next
|
||||
}
|
||||
|
||||
func (s *Service) GenerateReply(ctx context.Context, ownerUID int64, postID, replyID, personaID string) (string, error) {
|
||||
func (s *Service) GenerateReply(ctx context.Context, ownerUID int64, postID, replyID, personaID string) (_ string, err error) {
|
||||
post, err := s.getOwnPostOwned(ctx, ownerUID, postID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
|
|
@ -2127,9 +2156,11 @@ func (s *Service) GenerateReply(ctx context.Context, ownerUID int64, postID, rep
|
|||
return "", fmt.Errorf("%w: reply not found", domain.ErrValidation)
|
||||
}
|
||||
}
|
||||
if err := s.billAI(ctx, ownerUID, "own post reply draft", "ownPosts.generateReply"); err != nil {
|
||||
charge, err := s.billAI(ctx, ownerUID, "own post reply draft", "ownPosts.generateReply")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
persona := s.loadPersonaForGen(ctx, ownerUID, personaID)
|
||||
fp := personaExpressionFingerprintBlock(persona)
|
||||
|
||||
|
|
@ -2230,7 +2261,7 @@ func (s *Service) SendReply(ctx context.Context, ownerUID int64, postID, replyID
|
|||
return post, nil
|
||||
}
|
||||
|
||||
func (s *Service) AnalyzePost(ctx context.Context, ownerUID int64, postID string) (*domain.OwnPost, error) {
|
||||
func (s *Service) AnalyzePost(ctx context.Context, ownerUID int64, postID string) (_ *domain.OwnPost, err error) {
|
||||
post, err := s.getOwnPostOwned(ctx, ownerUID, postID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
|
@ -2238,9 +2269,11 @@ func (s *Service) AnalyzePost(ctx context.Context, ownerUID int64, postID string
|
|||
if strings.TrimSpace(post.Text) == "" {
|
||||
return nil, fmt.Errorf("%w: 貼文沒有文字可分析(圖片/純媒體貼可改貼上說明再分析)", domain.ErrValidation)
|
||||
}
|
||||
if err := s.billAI(ctx, ownerUID, "own post structure analyze", "ownPosts.analyze"); err != nil {
|
||||
charge, err := s.billAI(ctx, ownerUID, "own post structure analyze", "ownPosts.analyze")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
va, err := s.analyzeViralUnbilled(ctx, ownerUID, post.Text)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
|
@ -2371,14 +2404,16 @@ func mentionContextSnippet(hit FetchedMention) string {
|
|||
return kind + " · " + text
|
||||
}
|
||||
|
||||
func (s *Service) GenerateMentionReply(ctx context.Context, ownerUID int64, id, personaID string) (*domain.Mention, error) {
|
||||
func (s *Service) GenerateMentionReply(ctx context.Context, ownerUID int64, id, personaID string) (_ *domain.Mention, err error) {
|
||||
m, err := s.getMentionOwned(ctx, ownerUID, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.billAI(ctx, ownerUID, "mention reply draft", "mentions.generateReply"); err != nil {
|
||||
charge, err := s.billAI(ctx, ownerUID, "mention reply draft", "mentions.generateReply")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer charge.Settle(ctx, &err)
|
||||
persona := s.loadPersonaForGen(ctx, ownerUID, personaID)
|
||||
fp := personaExpressionFingerprintBlock(persona)
|
||||
prompt := buildMentionReplyPrompt(fp, m.FromUsername, m.Text, m.ContextSnippet)
|
||||
|
|
@ -2652,16 +2687,67 @@ func (s *Service) SeedMention(ctx context.Context, m *domain.Mention) error {
|
|||
|
||||
// ---------- helpers ----------
|
||||
|
||||
func (s *Service) billAI(ctx context.Context, ownerUID int64, label, source string) error {
|
||||
if s.Usage == nil {
|
||||
return nil
|
||||
// aiCharge is one reserved AI credit. Quota is reserved up front so an over-limit member is
|
||||
// stopped before any work starts, but the audit event is only written once the work succeeded —
|
||||
// the monthly counter is rebuilt from events if it is ever lost, so an event for a call that
|
||||
// never happened would silently undo the refund.
|
||||
type aiCharge struct {
|
||||
svc *Service
|
||||
ownerUID int64
|
||||
label string
|
||||
source string
|
||||
mode string
|
||||
settled bool
|
||||
}
|
||||
|
||||
// billAI reserves credit for one AI call. Pair it with `defer charge.Settle(ctx, &err)` on a
|
||||
// named error return so that every failure path — including ones added later — refunds.
|
||||
func (s *Service) billAI(ctx context.Context, ownerUID int64, label, source string) (*aiCharge, error) {
|
||||
if s == nil || s.Usage == nil {
|
||||
return &aiCharge{settled: true}, nil
|
||||
}
|
||||
mode, err := s.Usage.PrepareCall(ctx, ownerUID, usageDomain.MeterAICopy)
|
||||
if err != nil {
|
||||
return err
|
||||
return nil, err
|
||||
}
|
||||
return &aiCharge{svc: s, ownerUID: ownerUID, label: label, source: source, mode: mode}, nil
|
||||
}
|
||||
|
||||
// Settle commits the charge when the operation succeeded, or refunds it when it failed.
|
||||
func (c *aiCharge) Settle(ctx context.Context, errp *error) {
|
||||
if errp != nil && *errp != nil {
|
||||
c.Release(ctx)
|
||||
return
|
||||
}
|
||||
c.Commit(ctx)
|
||||
}
|
||||
|
||||
// Commit writes the audit event. A failure here is logged rather than surfaced: the member
|
||||
// already received the AI result, so charging them is correct and losing the event must not
|
||||
// turn a successful call into an error.
|
||||
func (c *aiCharge) Commit(ctx context.Context) {
|
||||
if c == nil || c.settled {
|
||||
return
|
||||
}
|
||||
c.settled = true
|
||||
if _, err := c.svc.Usage.RecordCall(ctx, c.ownerUID, usageDomain.MeterAICopy, c.mode, c.label, c.source); err != nil {
|
||||
logx.Errorf("usage record uid=%d source=%s mode=%s: %v", c.ownerUID, c.source, c.mode, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Release refunds a charge that was never committed. Safe to call more than once.
|
||||
func (c *aiCharge) Release(ctx context.Context) {
|
||||
if c == nil || c.settled {
|
||||
return
|
||||
}
|
||||
c.settled = true
|
||||
// The failure being compensated for is often a cancelled request, so the refund needs a
|
||||
// context that outlives it.
|
||||
rctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
|
||||
defer cancel()
|
||||
if err := c.svc.Usage.ReleaseCall(rctx, c.ownerUID, usageDomain.MeterAICopy, c.mode); err != nil {
|
||||
logx.Errorf("usage release uid=%d source=%s mode=%s: %v", c.ownerUID, c.source, c.mode, err)
|
||||
}
|
||||
_, err = s.Usage.RecordCall(ctx, ownerUID, usageDomain.MeterAICopy, mode, label, source)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Service) getOwnPostOwned(ctx context.Context, ownerUID int64, id string) (*domain.OwnPost, error) {
|
||||
|
|
|
|||
|
|
@ -280,11 +280,11 @@ func (s *Service) Disconnect(ctx context.Context, ownerUID int64, accountID stri
|
|||
}
|
||||
|
||||
// WebRedirect builds FE redirect after callback.
|
||||
// WebRedirect builds where the browser lands after the Threads OAuth round trip.
|
||||
// An unset WebBase yields a relative path on purpose: falling back to a hardcoded loopback
|
||||
// address would silently strand every production user on their own machine.
|
||||
func (s *Service) WebRedirect(ok bool, msg string) string {
|
||||
base := trimSlash(s.WebBase)
|
||||
if base == "" {
|
||||
base = "http://127.0.0.1:5173"
|
||||
}
|
||||
if ok {
|
||||
return base + "/app/crew?oauth=ok"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,6 +6,10 @@ type Repository interface {
|
|||
// ReservePlatform atomically guards both limits and increments one monthly counter document.
|
||||
// A negative limit means unlimited. BYOK must never call this operation.
|
||||
ReservePlatform(ctx context.Context, uid int64, monthKey, meter string, cost, totalLimit, meterLimit int) error
|
||||
// ReleasePlatform gives back a reservation whose call never happened. It is guarded so a
|
||||
// counter can never drop below zero, and is a no-op when there is nothing left to refund
|
||||
// (e.g. the month already rolled over), so callers may retry it safely.
|
||||
ReleasePlatform(ctx context.Context, uid int64, monthKey, meter string, cost int) error
|
||||
GetMonthlyCounter(ctx context.Context, uid int64, monthKey string) (*MonthlyCounter, error)
|
||||
|
||||
InsertEvent(ctx context.Context, e *Event) error
|
||||
|
|
|
|||
|
|
@ -58,6 +58,25 @@ func (s *MemoryStore) ReservePlatform(_ context.Context, uid int64, monthKey, me
|
|||
return nil
|
||||
}
|
||||
|
||||
func (s *MemoryStore) ReleasePlatform(_ context.Context, uid int64, monthKey, meter string, cost int) error {
|
||||
if cost <= 0 {
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
counter := s.counters[memoryCounterID(uid, monthKey)]
|
||||
if counter == nil {
|
||||
return nil
|
||||
}
|
||||
if counter.TotalCredits < cost || counter.ByMeter[meter] < cost {
|
||||
return nil
|
||||
}
|
||||
counter.TotalCredits -= cost
|
||||
counter.ByMeter[meter] -= cost
|
||||
counter.UpdatedAt = domain.NowNano()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *MemoryStore) GetMonthlyCounter(_ context.Context, uid int64, monthKey string) (*domain.MonthlyCounter, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
|
|
|||
|
|
@ -119,6 +119,28 @@ func (s *MonStore) ReservePlatform(ctx context.Context, uid int64, monthKey, met
|
|||
return fmt.Errorf("platform quota reservation did not match counter")
|
||||
}
|
||||
|
||||
func (s *MonStore) ReleasePlatform(ctx context.Context, uid int64, monthKey, meter string, cost int) error {
|
||||
field, err := meterCounterField(meter)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if cost <= 0 {
|
||||
return nil
|
||||
}
|
||||
// Both guards keep the refund from pushing either counter negative; a non-match means
|
||||
// there is nothing to give back, which is not an error.
|
||||
filter := bson.M{
|
||||
"_id": monthlyCounterID(uid, monthKey),
|
||||
"total_credits": bson.M{"$gte": cost},
|
||||
field: bson.M{"$gte": cost},
|
||||
}
|
||||
_, err = s.counters.UpdateOne(ctx, filter, bson.M{
|
||||
"$inc": bson.M{"total_credits": -cost, field: -cost},
|
||||
"$set": bson.M{"updated_at": domain.NowNano()},
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *MonStore) GetMonthlyCounter(ctx context.Context, uid int64, monthKey string) (*domain.MonthlyCounter, error) {
|
||||
var counter domain.MonthlyCounter
|
||||
err := s.counters.FindOne(ctx, &counter, bson.M{"_id": monthlyCounterID(uid, monthKey)})
|
||||
|
|
|
|||
|
|
@ -81,17 +81,39 @@ func (s *Service) PrepareCall(ctx context.Context, uid int64, meter string) (key
|
|||
if !has {
|
||||
return "", domain.ErrNoKey
|
||||
}
|
||||
if mode == domain.KeyModePlatform {
|
||||
if s.Gate != nil && !s.Gate.AllowPlatform(meter) {
|
||||
return "", domain.ErrPlatformCapacity
|
||||
}
|
||||
if err := s.reservePlatform(ctx, uid, meter); err != nil {
|
||||
if err := s.PrepareCallWithMode(ctx, uid, meter, mode); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
return mode, nil
|
||||
}
|
||||
|
||||
// PrepareCallWithMode reserves quota for a call whose key mode the caller already resolved.
|
||||
// Callers that also need the key material must use this rather than PrepareCall: resolving
|
||||
// twice reads member settings twice, and if those reads disagree the credits get reserved
|
||||
// against one mode while the call is billed as the other.
|
||||
func (s *Service) PrepareCallWithMode(ctx context.Context, uid int64, meter, keyMode string) error {
|
||||
if keyMode != domain.KeyModePlatform && keyMode != domain.KeyModeByok {
|
||||
return fmt.Errorf("invalid key_mode")
|
||||
}
|
||||
if keyMode != domain.KeyModePlatform {
|
||||
return nil
|
||||
}
|
||||
if s.Gate != nil && !s.Gate.AllowPlatform(meter) {
|
||||
return domain.ErrPlatformCapacity
|
||||
}
|
||||
return s.reservePlatform(ctx, uid, meter)
|
||||
}
|
||||
|
||||
// ReleaseCall gives back credits reserved by PrepareCall when the call it was paying for
|
||||
// never completed. BYOK reserves nothing, so it is a no-op there. Callers should reach for
|
||||
// this on every failure path between PrepareCall and RecordCall.
|
||||
func (s *Service) ReleaseCall(ctx context.Context, uid int64, meter, keyMode string) error {
|
||||
if keyMode != domain.KeyModePlatform {
|
||||
return nil
|
||||
}
|
||||
return s.Repo.ReleasePlatform(ctx, uid, domain.CurrentMonthKey(), meter, domain.MeterCost(meter))
|
||||
}
|
||||
|
||||
// MarkPlatformLimited cools down platform keys (e.g. after upstream 429).
|
||||
// BYOK users are unaffected. until zero clears the cool-down.
|
||||
func (s *Service) MarkPlatformLimited(until time.Time) {
|
||||
|
|
|
|||
|
|
@ -295,6 +295,66 @@ func TestUSG_14_MemberCannotSetOthersPrefs(t *testing.T) {
|
|||
require.ErrorIs(t, err, domain.ErrForbidden)
|
||||
}
|
||||
|
||||
func TestReleaseCall_RefundsReservedPlatformCredit(t *testing.T) {
|
||||
repo := repository.NewMemory()
|
||||
uid := int64(1_000_010)
|
||||
svc := usecase.New(repo, &usecase.StaticResolver{Map: map[string]string{
|
||||
fmt.Sprintf("%d:%s", uid, domain.MeterAICopy): domain.KeyModePlatform,
|
||||
}})
|
||||
ctx := context.Background()
|
||||
|
||||
mode, err := svc.PrepareCall(ctx, uid, domain.MeterAICopy)
|
||||
require.NoError(t, err)
|
||||
counter, err := repo.GetMonthlyCounter(ctx, uid, domain.CurrentMonthKey())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, domain.MeterCost(domain.MeterAICopy), counter.TotalCredits)
|
||||
|
||||
require.NoError(t, svc.ReleaseCall(ctx, uid, domain.MeterAICopy, mode))
|
||||
counter, err = repo.GetMonthlyCounter(ctx, uid, domain.CurrentMonthKey())
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, counter.TotalCredits, "a call that never happened must not consume credits")
|
||||
require.Zero(t, counter.ByMeter[domain.MeterAICopy])
|
||||
|
||||
// The member must still see the full quota available afterwards.
|
||||
sum, err := svc.GetSummary(ctx, uid, "")
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, sum.Platform.CreditsUsed)
|
||||
}
|
||||
|
||||
func TestReleaseCall_NeverDrivesCounterNegative(t *testing.T) {
|
||||
repo := repository.NewMemory()
|
||||
uid := int64(1_000_011)
|
||||
svc := usecase.New(repo, &usecase.StaticResolver{Map: map[string]string{
|
||||
fmt.Sprintf("%d:%s", uid, domain.MeterAICopy): domain.KeyModePlatform,
|
||||
}})
|
||||
ctx := context.Background()
|
||||
mode, err := svc.PrepareCall(ctx, uid, domain.MeterAICopy)
|
||||
require.NoError(t, err)
|
||||
|
||||
// A retried or duplicated refund must stay idempotent.
|
||||
for i := 0; i < 3; i++ {
|
||||
require.NoError(t, svc.ReleaseCall(ctx, uid, domain.MeterAICopy, mode))
|
||||
}
|
||||
counter, err := repo.GetMonthlyCounter(ctx, uid, domain.CurrentMonthKey())
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, counter.TotalCredits)
|
||||
require.GreaterOrEqual(t, counter.ByMeter[domain.MeterAICopy], 0)
|
||||
}
|
||||
|
||||
func TestReleaseCall_ByokIsNoOp(t *testing.T) {
|
||||
repo := repository.NewMemory()
|
||||
uid := int64(1_000_012)
|
||||
svc := usecase.New(repo, &usecase.StaticResolver{Map: map[string]string{
|
||||
fmt.Sprintf("%d:%s", uid, domain.MeterAICopy): domain.KeyModeByok,
|
||||
}})
|
||||
ctx := context.Background()
|
||||
mode, err := svc.PrepareCall(ctx, uid, domain.MeterAICopy)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, svc.ReleaseCall(ctx, uid, domain.MeterAICopy, mode))
|
||||
_, err = repo.GetMonthlyCounter(ctx, uid, domain.CurrentMonthKey())
|
||||
require.ErrorIs(t, err, domain.ErrNotFound, "BYOK reserves nothing, so a refund must not create a counter")
|
||||
}
|
||||
|
||||
func TestCreditInviteBonus_AddsToPool(t *testing.T) {
|
||||
svc := newUsage(42, domain.KeyModePlatform)
|
||||
require.NoError(t, svc.CreditInviteBonus(context.Background(), 42, 100, "r1"))
|
||||
|
|
|
|||
|
|
@ -192,6 +192,8 @@ func mapError(err error) (int, Envelope) {
|
|||
return http.StatusBadRequest, Envelope{Code: 400060, Message: cleanBizMessage(err.Error())}
|
||||
case errors.Is(err, inspireDomain.ErrRemoved), errors.Is(err, scoutDomain.ErrTopicRemoved):
|
||||
return http.StatusGone, Envelope{Code: 410002, Message: err.Error()}
|
||||
case errors.Is(err, inspireDomain.ErrNotImplemented):
|
||||
return http.StatusNotImplemented, Envelope{Code: 501000, Message: cleanBizMessage(err.Error())}
|
||||
case errors.Is(err, scoutDomain.ErrNoCrawlerSession):
|
||||
return http.StatusBadRequest, Envelope{Code: 400061, Message: "crawler session required when dev_mode enabled"}
|
||||
case errors.Is(err, scoutDomain.ErrHasProducts):
|
||||
|
|
|
|||
|
|
@ -1680,6 +1680,10 @@ type UsagePrefsData struct {
|
|||
Unlimited bool `json:"unlimited"`
|
||||
}
|
||||
|
||||
type UsagePrefsReq struct {
|
||||
Uid string `form:"uid,optional"`
|
||||
}
|
||||
|
||||
type UsagePurchasePublic struct {
|
||||
Id string `json:"id"`
|
||||
PlanId string `json:"plan_id"`
|
||||
|
|
|
|||
Loading…
Reference in New Issue