This commit is contained in:
王性驊 2026-07-28 06:33:40 +00:00
parent ced11faf11
commit f2be5a436e
32 changed files with 1044 additions and 228 deletions

View File

@ -36,21 +36,128 @@ func main() {
defer client.Disconnect(context.Background()) defer client.Disconnect(context.Background())
db := client.Database(c.Mongo.Database) db := client.Database(c.Mongo.Database)
indexes := map[string][]mongo.IndexModel{ for collection, models := range indexModels() {
"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 {
if _, err := db.Collection(collection).Indexes().CreateMany(ctx, models); err != nil { if _, err := db.Collection(collection).Indexes().CreateMany(ctx, models); err != nil {
panic(fmt.Errorf("create %s indexes: %w", collection, err)) panic(fmt.Errorf("create %s indexes: %w", collection, err))
} }
} }
fmt.Println("database indexes initialized") 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")},
},
}
}

View File

@ -164,8 +164,15 @@ func main() {
ctx := context.Background() ctx := context.Background()
const outboxLockTTL = 3 * time.Minute const outboxLockTTL = 3 * time.Minute
outboxLockKey := strings.Trim(c.Redis.Namespace, ":") + ":worker:outbox" ns := strings.Trim(c.Redis.Namespace, ":")
outboxLock := redislock.NewWithClient(redisClient, outboxLockKey, workerID, outboxLockTTL) 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 { for {
select { select {
case <-sig: case <-sig:
@ -231,20 +238,50 @@ func main() {
} }
// 2) One worker owns an outbox tick and renews its lease while claims run. // 2) One worker owns an outbox tick and renews its lease while claims run.
processOutbox(ctx, studio, outboxLock, workerID, outboxLockTTL) 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 { if n, err := growthSvc.ObserveTick(ctx, 0); err != nil {
logx.Errorf("worker %s outcome observe: %v", workerID, err) logx.Errorf("worker %s outcome observe: %v", workerID, err)
} else if n > 0 { } else if n > 0 {
logx.Infof("worker %s outcome observe updated %d", workerID, n) logx.Infof("worker %s outcome observe updated %d", workerID, n)
} }
// 4) 終態任務超過 2 天自動清除
if purged, err := jobs.PurgeExpiredTerminal(ctx); err != nil { if purged, err := jobs.PurgeExpiredTerminal(ctx); err != nil {
logx.Errorf("worker %s purge jobs: %v", workerID, err) logx.Errorf("worker %s purge jobs: %v", workerID, err)
} else if purged > 0 { } else if purged > 0 {
logx.Infof("worker %s purged %d expired terminal job(s)", workerID, purged) 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) { func processOutbox(ctx context.Context, studio *studioUC.Service, lock *redislock.Lock, workerID string, ttl time.Duration) {

View File

@ -61,6 +61,11 @@ type (
Unlimited bool `json:"unlimited"` 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 { UsageSetPrefsReq {
Uid string `json:"uid"` Uid string `json:"uid"`
PlanId string `json:"plan_id,optional"` PlanId string `json:"plan_id,optional"`
@ -165,7 +170,7 @@ service gateway {
get /events (UsageEventsReq) returns (UsageEventsData) get /events (UsageEventsReq) returns (UsageEventsData)
@handler GetUsagePrefs @handler GetUsagePrefs
get /prefs returns (UsagePrefsData) get /prefs (UsagePrefsReq) returns (UsagePrefsData)
@handler PurchasePlan @handler PurchasePlan
post /purchase (UsagePurchaseReq) returns (UsagePurchasePublic) post /purchase (UsagePurchaseReq) returns (UsagePurchasePublic)

View File

@ -1265,35 +1265,6 @@ func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
rest.WithPrefix("/api/v1/utm"), 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( server.AddRoutes(
rest.WithMiddlewares( rest.WithMiddlewares(
[]rest.Middleware{serverCtx.AuthJWT}, []rest.Middleware{serverCtx.AuthJWT},
@ -1352,4 +1323,33 @@ func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
), ),
rest.WithPrefix("/api/v1/workspaces"), 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"),
)
} }

View File

@ -9,12 +9,20 @@ import (
"apps/backend/internal/logic/usage" "apps/backend/internal/logic/usage"
"apps/backend/internal/response" "apps/backend/internal/response"
"apps/backend/internal/svc" "apps/backend/internal/svc"
"apps/backend/internal/types"
"github.com/zeromicro/go-zero/rest/httpx"
) )
func GetUsagePrefsHandler(svcCtx *svc.ServiceContext) http.HandlerFunc { func GetUsagePrefsHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) { 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) l := usage.NewGetUsagePrefsLogic(r.Context(), svcCtx)
data, err := l.GetUsagePrefs() data, err := l.GetUsagePrefs(&req)
response.Write(r.Context(), w, data, err) response.Write(r.Context(), w, data, err)
} }
} }

View File

@ -22,7 +22,7 @@ func NewAICompleteLogic(ctx context.Context, svcCtx *svc.ServiceContext) *AIComp
return &AICompleteLogic{Logger: logx.WithContext(ctx), ctx: ctx, svcCtx: svcCtx} 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 { if l.svcCtx.Usage == nil || l.svcCtx.AI == nil || l.svcCtx.KeyResolver == nil {
return nil, response.Biz(503, 503001, "ai proxy not configured") 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 == "" { if meter == "" {
meter = usageDomain.MeterAICopy 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) mode, apiKey, err := l.svcCtx.KeyResolver.ResolveKey(l.ctx, uid, meter)
if err != nil { if err != nil {
return nil, err 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 model := req.Model
if model == "" { if model == "" {
model = "grok-3" model = "grok-3"
@ -51,6 +52,5 @@ func (l *AICompleteLogic) AIComplete(req *types.AICompleteReq) (*types.AIComplet
if err != nil { if err != nil {
return nil, response.Biz(502, 502001, err.Error()) 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 return &types.AICompleteData{Text: text, KeyMode: mode, Model: model}, nil
} }

View File

@ -22,27 +22,27 @@ func NewExaSearchLogic(ctx context.Context, svcCtx *svc.ServiceContext) *ExaSear
return &ExaSearchLogic{Logger: logx.WithContext(ctx), ctx: ctx, svcCtx: svcCtx} return &ExaSearchLogic{Logger: logx.WithContext(ctx), ctx: ctx, svcCtx: svcCtx}
} }
func (l *ExaSearchLogic) ExaSearch(req *types.SearchReq) (*types.SearchData, error) { func (l *ExaSearchLogic) ExaSearch(req *types.SearchReq) (_ *types.SearchData, err error) {
if l.svcCtx.Usage == nil || l.svcCtx.Search == nil { if l.svcCtx.Usage == nil || l.svcCtx.Search == nil || l.svcCtx.KeyResolver == nil {
return nil, response.Biz(503, 503001, "search proxy not configured") return nil, response.Biz(503, 503001, "search proxy not configured")
} }
uid, ok := middleware.UIDFrom(l.ctx) uid, ok := middleware.UIDFrom(l.ctx)
if !ok { if !ok {
return nil, response.Biz(401, 401001, "missing authorization") 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) mode, apiKey, err := l.svcCtx.KeyResolver.ResolveKey(l.ctx, uid, usageDomain.MeterWebSearch)
if err != nil { if err != nil {
return nil, err 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) hits, err := l.svcCtx.Search.Search(l.ctx, apiKey, req.Query, req.Limit)
if err != nil { if err != nil {
return nil, response.Biz(502, 502002, err.Error()) 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)) out := make([]types.SearchHit, 0, len(hits))
for _, h := range hits { for _, h := range hits {
out = append(out, types.SearchHit{Title: h.Title, Url: h.URL, Snippet: h.Snippet}) out = append(out, types.SearchHit{Title: h.Title, Url: h.URL, Snippet: h.Snippet})

View File

@ -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)
}
}

View File

@ -2,6 +2,7 @@ package threads
import ( import (
"context" "context"
"strings"
threadsDomain "apps/backend/internal/module/threads/domain" threadsDomain "apps/backend/internal/module/threads/domain"
"apps/backend/internal/svc" "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) { func (l *ThreadsOAuthCallbackLogic) ThreadsOAuthCallback(req *types.ThreadsOAuthCallbackReq) (*types.ThreadsOAuthCallbackData, error) {
if l.svcCtx.Threads == nil { 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 return &types.ThreadsOAuthCallbackData{Ok: false, Message: "not configured", RedirectUrl: redir}, nil
} }
_, err := l.svcCtx.Threads.OAuthCallback(l.ctx, req.Code, req.State, req.Error) _, err := l.svcCtx.Threads.OAuthCallback(l.ctx, req.Code, req.State, req.Error)

View File

@ -2,8 +2,11 @@ package usage
import ( import (
"context" "context"
"strconv"
"strings"
"apps/backend/internal/middleware" "apps/backend/internal/middleware"
"apps/backend/internal/module/permission"
"apps/backend/internal/response" "apps/backend/internal/response"
"apps/backend/internal/svc" "apps/backend/internal/svc"
"apps/backend/internal/types" "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} 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 { if l.svcCtx.Usage == nil {
return nil, response.Biz(503, 503001, "usage not configured") return nil, response.Biz(503, 503001, "usage not configured")
} }
@ -29,6 +32,20 @@ func (l *GetUsagePrefsLogic) GetUsagePrefs() (*types.UsagePrefsData, error) {
if !ok { if !ok {
return nil, response.Biz(401, 401001, "missing authorization") 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) p, err := l.svcCtx.Usage.GetPrefs(l.ctx, uid)
if err != nil { if err != nil {
return nil, err return nil, err

View File

@ -142,7 +142,11 @@ type Repository interface {
GetSubscription(context.Context, int64) (*SubscriptionState, error) GetSubscription(context.Context, int64) (*SubscriptionState, error)
FindUIDByCustomer(context.Context, string) (int64, error) FindUIDByCustomer(context.Context, string) (int64, error)
ApplyEntitlement(context.Context, EntitlementUpdate) (bool, error) ApplyEntitlement(context.Context, EntitlementUpdate) (bool, error)
WebhookProcessed(context.Context, string) (bool, error) // ClaimWebhook atomically records an event as being handled, reporting false when another
MarkWebhookProcessed(context.Context, string, string, int64) error // 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 UpsertPurchase(context.Context, *PurchaseAudit) error
} }

View File

@ -124,16 +124,28 @@ func (r *MongoRepository) ApplyEntitlement(ctx context.Context, u EntitlementUpd
return res.MatchedCount > 0 || res.UpsertedCount > 0, nil return res.MatchedCount > 0 || res.UpsertedCount > 0, nil
} }
func (r *MongoRepository) WebhookProcessed(ctx context.Context, id string) (bool, error) { func (r *MongoRepository) ClaimWebhook(ctx context.Context, id, eventType string, created int64) (bool, error) {
var out bson.M _, err := r.webhooks.InsertOne(ctx, bson.M{
err := r.webhooks.FindOne(ctx, &out, bson.M{"_id": id}) "_id": id,
if err == mon.ErrNotFound { "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 false, nil
} }
return err == nil, err if err != nil {
return false, err
}
return true, 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)) 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
}
return err return err
} }
func (r *MongoRepository) UpsertPurchase(ctx context.Context, p *PurchaseAudit) error { func (r *MongoRepository) UpsertPurchase(ctx context.Context, p *PurchaseAudit) error {

View File

@ -9,6 +9,7 @@ import (
"time" "time"
"github.com/google/uuid" "github.com/google/uuid"
"github.com/zeromicro/go-zero/core/logx"
) )
type Service struct { type Service struct {
@ -176,14 +177,22 @@ func (s *Service) HandleWebhook(ctx context.Context, payload []byte, signature s
if err != nil { if err != nil {
return err return err
} }
done, err := s.Repo.WebhookProcessed(ctx, e.ID) // Claim before processing, so two concurrent deliveries of the same event cannot both
if err != nil || done { // 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 return err
} }
if !claimed {
return nil
}
if err := s.processEvent(ctx, e); err != 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 err
} }
return s.Repo.MarkWebhookProcessed(ctx, e.ID, e.Type, e.Created) return nil
} }
func (s *Service) processEvent(ctx context.Context, e *Event) error { func (s *Service) processEvent(ctx context.Context, e *Event) error {

View File

@ -56,6 +56,13 @@ type fakeRepo struct {
lastEvent string lastEvent string
purchases map[string]*PurchaseAudit purchases map[string]*PurchaseAudit
entitlementErr error entitlementErr error
applyCalls int
}
func (r *fakeRepo) applyCallCount() int {
r.mu.Lock()
defer r.mu.Unlock()
return r.applyCalls
} }
func newFakeRepo() *fakeRepo { 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) { func (r *fakeRepo) ApplyEntitlement(_ context.Context, u EntitlementUpdate) (bool, error) {
r.mu.Lock() r.mu.Lock()
defer r.mu.Unlock() defer r.mu.Unlock()
r.applyCalls++
if r.entitlementErr != nil { if r.entitlementErr != nil {
return false, r.entitlementErr return false, r.entitlementErr
} }
@ -150,15 +158,19 @@ func (r *fakeRepo) ApplyEntitlement(_ context.Context, u EntitlementUpdate) (boo
} }
return true, nil 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
}
func (r *fakeRepo) MarkWebhookProcessed(_ context.Context, id, _ string, _ int64) error {
r.mu.Lock() r.mu.Lock()
defer r.mu.Unlock() defer r.mu.Unlock()
if r.processed[id] {
return false, nil
}
r.processed[id] = true r.processed[id] = true
return true, nil
}
func (r *fakeRepo) ReleaseWebhookClaim(_ context.Context, id string) error {
r.mu.Lock()
defer r.mu.Unlock()
delete(r.processed, id)
return nil return nil
} }
func (r *fakeRepo) UpsertPurchase(_ context.Context, p *PurchaseAudit) error { func (r *fakeRepo) UpsertPurchase(_ context.Context, p *PurchaseAudit) error {
@ -225,6 +237,58 @@ func TestSubscriptionRemainsReadableWhenStripeDisabled(t *testing.T) {
require.Equal(t, PlanFree, state.PlanID) 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) { func TestWebhookDuplicateOutOfOrderDowngradeAndAdminOverride(t *testing.T) {
r := newFakeRepo() r := newFakeRepo()
r.state = SubscriptionState{UID: 7, PlanID: PlanFree, CustomerID: "cus_1"} r.state = SubscriptionState{UID: 7, PlanID: PlanFree, CustomerID: "cus_1"}

View File

@ -6,11 +6,14 @@ import (
"encoding/hex" "encoding/hex"
"fmt" "fmt"
"math" "math"
"net/url"
"sort" "sort"
"strings" "strings"
"unicode" "unicode"
"apps/backend/internal/module/growth/domain" "apps/backend/internal/module/growth/domain"
"github.com/zeromicro/go-zero/core/logx"
) )
// --- Playbooks --- // --- Playbooks ---
@ -224,10 +227,28 @@ func tokenizeTC(s string) []string {
// --- UTM --- // --- 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) dest = strings.TrimSpace(dest)
if dest == "" || !strings.HasPrefix(dest, "http") { u, err := url.Parse(dest)
return nil, fmt.Errorf("%w: destination_url must be http(s)", domain.ErrValidation) 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] code := domain.NewID()[:10]
u := &domain.UtmLink{ u := &domain.UtmLink{
@ -259,10 +280,14 @@ func (s *Service) TrackUtmClick(ctx context.Context, code string) (dest string,
e.Status = domain.StatusConfirmed e.Status = domain.StatusConfirmed
} }
e.UpdatedAt = domain.NowNano() 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 --- // --- Benchmark ---

View File

@ -168,3 +168,55 @@ func TestWS_05_ExportReport(t *testing.T) {
require.Contains(t, fn, "2026-07") require.Contains(t, fn, "2026-07")
require.Contains(t, body, "月報") 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)
}

View File

@ -10,6 +10,9 @@ var (
ErrForbidden = errors.New("inspire forbidden") ErrForbidden = errors.New("inspire forbidden")
ErrValidation = errors.New("inspire validation") ErrValidation = errors.New("inspire validation")
ErrRemoved = errors.New("legacy inspire API removed") 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() } func NowNano() int64 { return time.Now().UTC().UnixNano() }

View File

@ -224,23 +224,27 @@ func TestIN_06_ResearchSearch(t *testing.T) {
require.NotEmpty(t, hits[0].URL) 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() svc := newInspire()
uid := int64(5_001_017) uid := int64(5_001_017)
setupInspireUID(svc, uid) setupInspireUID(svc, uid)
img, err := svc.GenerateImage(context.Background(), 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.NoError(t, err)
require.NotEmpty(t, img.URL) require.Zero(t, sum.Platform.CreditsUsed, "an unimplemented feature must not be billed")
require.NotEmpty(t, img.ID)
} }
func TestIN_08_UploadSeparate(t *testing.T) { func TestIN_08_GenerateImageStillValidatesPrompt(t *testing.T) {
svc := newInspire() svc := newInspire()
uid := int64(5_001_008) uid := int64(5_001_008)
setupInspireUID(svc, uid) setupInspireUID(svc, uid)
img, err := svc.GenerateImage(context.Background(), uid, "x") _, err := svc.GenerateImage(context.Background(), uid, " ")
require.NoError(t, err) require.ErrorIs(t, err, domain.ErrValidation)
require.True(t, len(img.URL) > 8)
} }
func TestIN_09_LegacyRemoved(t *testing.T) { func TestIN_09_LegacyRemoved(t *testing.T) {

View File

@ -16,6 +16,7 @@ import (
usageUC "apps/backend/internal/module/usage/usecase" usageUC "apps/backend/internal/module/usage/usecase"
"github.com/google/uuid" "github.com/google/uuid"
"github.com/zeromicro/go-zero/core/logx"
) )
type Service struct { type Service struct {
@ -54,14 +55,16 @@ func (s *Service) ListTrends(ctx context.Context, ownerUID int64) ([]*domain.Tre
// RefreshTrends — 手動「找靈感」:取得真實結果後才扣點並覆蓋快取。 // 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) list := s.tryLiveTrends(ctx, ownerUID)
if len(list) == 0 { if len(list) == 0 {
return nil, fmt.Errorf("找不到新的可用話題,已保留原本內容") 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 return nil, err
} }
defer charge.Settle(ctx, &err)
if err := s.Repo.SaveTrends(ctx, ownerUID, list); err != nil { if err := s.Repo.SaveTrends(ctx, ownerUID, list); err != nil {
return nil, err return nil, err
} }
@ -687,7 +690,7 @@ func generateUserLogMessage(notes, material string) string {
return "【產文】" + notes + "|素材:" + preview 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) message = strings.TrimSpace(message)
material = strings.TrimSpace(material) material = strings.TrimSpace(material)
if mode != "chat" && mode != "generate" { if mode != "chat" && mode != "generate" {
@ -714,9 +717,11 @@ func (s *Service) chatInternal(ctx context.Context, ownerUID int64, message stri
if mode == "generate" { if mode == "generate" {
label = "inspire rewrite" 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 return nil, err
} }
defer charge.Settle(ctx, &err)
// 回覆語言:人設指紋/範例優先,其次會員 UI 語系ctx 已由 Auth 注入) // 回覆語言:人設指紋/範例優先,其次會員 UI 語系ctx 已由 Auth 注入)
persona := s.resolvePersonaSnap(ctx, ownerUID, personaID) persona := s.resolvePersonaSnap(ctx, ownerUID, personaID)
if mode == "generate" && (persona == nil || persona.Status != "ready") { 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 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) query = strings.TrimSpace(query)
if query == "" { if query == "" {
return "", nil 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 return "", err
} }
defer charge.Settle(ctx, &err)
key := "fake" key := "fake"
if s.ResolveKey != nil { if s.ResolveKey != nil {
if _, resolved, err := s.ResolveKey(ctx, ownerUID, usageDomain.MeterWebSearch); err == nil && resolved != "" { 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) 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) query = strings.TrimSpace(query)
if query == "" { if query == "" {
return nil, fmt.Errorf("%w: empty query", domain.ErrValidation) 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 return nil, err
} }
defer charge.Settle(ctx, &err)
key := "fake" key := "fake"
if s.ResolveKey != nil { if s.ResolveKey != nil {
if _, k, err := s.ResolveKey(ctx, ownerUID, usageDomain.MeterWebSearch); err == nil && k != "" { 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 return out, nil
} }
func (s *Service) GenerateImage(ctx context.Context, ownerUID int64, prompt string) (*domain.GeneratedImage, error) { // GenerateImage has no image backend yet. It used to return a random avatar URL from a public
prompt = strings.TrimSpace(prompt) // placeholder service while charging a real ai_image credit, which billed members for a picture
if prompt == "" { // 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) 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, fmt.Errorf("%w: 圖片生成尚未接上生成服務", domain.ErrNotImplemented)
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
} }
// Legacy removed APIs // Legacy removed APIs
func (s *Service) LegacyRemoved() error { return domain.ErrRemoved } func (s *Service) LegacyRemoved() error { return domain.ErrRemoved }
func (s *Service) bill(ctx context.Context, uid int64, meter, label, source string) error { // inspireCharge is one reserved credit; see Settle/Release for the commit-on-success contract.
if s.Usage == nil { type inspireCharge struct {
return nil 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) mode, err := s.Usage.PrepareCall(ctx, uid, meter)
if err != nil { 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 { func emptySession(ownerUID int64) *domain.Session {

View File

@ -7,6 +7,7 @@ import (
libmongo "apps/backend/internal/lib/mongo" libmongo "apps/backend/internal/lib/mongo"
"apps/backend/internal/module/job/domain" "apps/backend/internal/module/job/domain"
"github.com/zeromicro/go-zero/core/logx"
"github.com/zeromicro/go-zero/core/stores/mon" "github.com/zeromicro/go-zero/core/stores/mon"
"go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo" "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 { if err := s.jobs.Find(ctx, &list, filter); err != nil {
return err 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 { for _, j := range list {
j.Status = domain.StatusCancelled j.Status = domain.StatusCancelled
j.ProgressSummary = "已由新的定期排程取代" j.ProgressSummary = "已由新的定期排程取代"
j.CompletedAt = now j.CompletedAt = now
j.UpdatedAt = 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 return nil
} }

View File

@ -128,7 +128,9 @@ func (s *Service) StartDemo(ctx context.Context, ownerUID int64) (*domain.Job, e
j.Status = domain.StatusQueued j.Status = domain.StatusQueued
j.ProgressSummary = "Demo 測試任務 · 等待 worker 領取" j.ProgressSummary = "Demo 測試任務 · 等待 worker 領取"
j.UpdatedAt = domain.NowNano() 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) s.notify(ctx, j)
return j, nil return j, nil
} }
@ -288,7 +290,10 @@ func (s *Service) ScheduleTokenRenew(ctx context.Context, ownerUID int64, accoun
if runAfter <= 0 { if runAfter <= 0 {
runAfter = now + domain.DefaultTokenRenewDelayNs 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{ j := &domain.Job{
ID: uuid.NewString(), OwnerUID: ownerUID, ID: uuid.NewString(), OwnerUID: ownerUID,
TemplateType: domain.TemplateThreadsTokenRenew, TemplateType: domain.TemplateThreadsTokenRenew,
@ -325,6 +330,15 @@ type PersonaAnalyzeTextPayload struct {
Lang string `json:"lang,omitempty"` 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同人設舊佇列取消 // SchedulePersonaAnalyzeAccount — 立刻可領ref=personaID同人設舊佇列取消
func (s *Service) SchedulePersonaAnalyzeAccount(ctx context.Context, ownerUID int64, personaID, username, lang string) (*domain.Job, error) { func (s *Service) SchedulePersonaAnalyzeAccount(ctx context.Context, ownerUID int64, personaID, username, lang string) (*domain.Job, error) {
if ownerUID <= 0 || personaID == "" { if ownerUID <= 0 || personaID == "" {
@ -334,8 +348,9 @@ func (s *Service) SchedulePersonaAnalyzeAccount(ctx context.Context, ownerUID in
if username == "" { if username == "" {
return nil, fmt.Errorf("empty username") return nil, fmt.Errorf("empty username")
} }
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeAccount, personaID) if err := s.cancelPersonaAnalyzeJobs(ctx, ownerUID, personaID); err != nil {
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeText, personaID) return nil, err
}
if s.MaxActiveAccountScrapes > 0 { if s.MaxActiveAccountScrapes > 0 {
active, err := s.Repo.CountActiveByOwnerAndTemplate(ctx, ownerUID, domain.TemplatePersonaAnalyzeAccount) active, err := s.Repo.CountActiveByOwnerAndTemplate(ctx, ownerUID, domain.TemplatePersonaAnalyzeAccount)
if err != nil { if err != nil {
@ -421,7 +436,9 @@ func (s *Service) SchedulePlayGenerateScript(ctx context.Context, ownerUID int64
return nil, domain.ErrForbidden return nil, domain.ErrForbidden
} }
playID = strings.TrimSpace(playID) 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}) body, _ := json.Marshal(PlayGenerateScriptPayload{PlayID: playID, OnlyEmpty: onlyEmpty})
now := domain.NowNano() now := domain.NowNano()
j := &domain.Job{ j := &domain.Job{
@ -476,8 +493,9 @@ func (s *Service) SchedulePersonaAnalyzeText(ctx context.Context, ownerUID int64
if rawText == "" { if rawText == "" {
return nil, fmt.Errorf("empty text") return nil, fmt.Errorf("empty text")
} }
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeAccount, personaID) if err := s.cancelPersonaAnalyzeJobs(ctx, ownerUID, personaID); err != nil {
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeText, personaID) return nil, err
}
body, _ := json.Marshal(PersonaAnalyzeTextPayload{ body, _ := json.Marshal(PersonaAnalyzeTextPayload{
RawText: rawText, SourceLabel: strings.TrimSpace(sourceLabel), Lang: lang, RawText: rawText, SourceLabel: strings.TrimSpace(sourceLabel), Lang: lang,
}) })

View File

@ -2,6 +2,8 @@ package usecase_test
import ( import (
"context" "context"
"errors"
"fmt"
"testing" "testing"
"time" "time"
@ -76,71 +78,29 @@ func newStudio() (*usecase.Service, *publish.FakeTransport, *memAccounts) {
return svc, tp, acc 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) { 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{ _ = 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, ok := svc.Usage.Resolver.(*usageUC.StaticResolver); ok {
if sr.Map == nil { if sr.Map == nil {
sr.Map = map[string]string{} sr.Map = map[string]string{}
} }
key := "" // StaticResolver keys on "uid:meter".
// StaticResolver uses fmt.Sprintf("%d:%s", uid, meter) sr.Map[fmt.Sprintf("%d:%s", uid, usageDomain.MeterAICopy)] = usageDomain.KeyModePlatform
key = sprintf("%d:%s", uid, usageDomain.MeterAICopy)
sr.Map[key] = usageDomain.KeyModePlatform
} }
} }
func sprintf(f string, a ...any) string { // platformCreditsUsed reports what the member has actually been charged this month.
return format(f, a...) func platformCreditsUsed(t *testing.T, svc *usecase.Service, uid int64) int {
} t.Helper()
sum, err := svc.Usage.GetSummary(context.Background(), uid, "")
// avoid importing fmt in helpers clutter — use strconv require.NoError(t, err)
func format(f string, a ...any) string { return sum.Platform.CreditsUsed
// 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:])
} }
func addAccount(acc *memAccounts, ownerUID int64, id, user string) { 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 ---- // ---- PE ----
func TestPE_01_CreateSave(t *testing.T) { func TestPE_01_CreateSave(t *testing.T) {
@ -915,6 +955,64 @@ func TestOP_11_GenerateFromFormulaRemoved(t *testing.T) {
require.ErrorIs(t, err, domain.ErrFormulaRemove) 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 { type fakeMentionsMedia struct {
hits []usecase.FetchedMention hits []usecase.FetchedMention
err error err error

View File

@ -16,7 +16,7 @@ import (
) )
// PersonaPreview依指紋一次 LLM 產主貼 + 回文(可抓新聞當話題靈感)。 // 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 { if ownerUID <= 0 {
return nil, domain.ErrForbidden 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) 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 return nil, err
} }
defer charge.Settle(ctx, &err)
topic = strings.TrimSpace(topic) topic = strings.TrimSpace(topic)
topicSource := "manual" topicSource := "manual"

View File

@ -279,7 +279,7 @@ func (s *Service) AnalyzeFromAccount(ctx context.Context, ownerUID int64, id, us
} }
// ExecuteAnalyzeFromText — worker測試扣費 + LLM 分析 + 存檔 ready。 // 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) { report := func(pct int, sum string) {
if onProgress != nil { if onProgress != nil {
onProgress(pct, sum) 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) return nil, fmt.Errorf("%w: 請至少貼 2 段參考文字(可用 --- 分隔),每段至少約 10 字", domain.ErrValidation)
} }
report(25, "人設分析 · 檢查文字樣本…") 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()) _ = s.markPersonaAnalyzeFailed(ctx, p, err.Error())
return nil, err return nil, err
} }
defer charge.Settle(ctx, &err)
p.Status = domain.PersonaAnalyzing p.Status = domain.PersonaAnalyzing
p.UpdatedAt = domain.NowNano() p.UpdatedAt = domain.NowNano()
_ = s.Repo.SavePersona(ctx, p) _ = s.Repo.SavePersona(ctx, p)
@ -324,7 +326,7 @@ func (s *Service) ExecuteAnalyzeFromText(ctx context.Context, ownerUID int64, id
} }
// ExecuteAnalyzeFromAccount — worker爬公開貼文 + LLM + 存檔 ready。 // 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) { report := func(pct int, sum string) {
if onProgress != nil { if onProgress != nil {
onProgress(pct, sum) 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) return nil, fmt.Errorf("%w: username 請只填帳號,例如 ultralab_tw", domain.ErrValidation)
} }
report(15, fmt.Sprintf("人設分析 · 準備爬取 @%s…", username)) 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()) _ = s.markPersonaAnalyzeFailed(ctx, p, err.Error())
return nil, err return nil, err
} }
defer charge.Settle(ctx, &err)
p.Status = domain.PersonaAnalyzing p.Status = domain.PersonaAnalyzing
p.Style.BenchmarkUsername = username p.Style.BenchmarkUsername = username
@ -807,7 +811,7 @@ func (s *Service) SubmitPlay(ctx context.Context, ownerUID int64, playID string)
// GeneratePlayScript — 一次 LLM 產完整劇本並寫回 play.stepsonlyEmpty=true 只填空白步)。 // GeneratePlayScript — 一次 LLM 產完整劇本並寫回 play.stepsonlyEmpty=true 只填空白步)。
// 回傳填入步數。供 worker job / 同步測試。 // 回傳填入步數。供 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) play, err := s.GetPlay(ctx, ownerUID, playID)
if err != nil { if err != nil {
return 0, err return 0, err
@ -841,9 +845,11 @@ func (s *Service) GeneratePlayScript(ctx context.Context, ownerUID int64, playID
if len(needs) == 0 { if len(needs) == 0 {
return 0, fmt.Errorf("%w: 沒有空白步驟可產(全部已有正文)", domain.ErrValidation) 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 return 0, err
} }
defer charge.Settle(ctx, &err)
// 主上下文:目標貼文/主題 // 主上下文:目標貼文/主題
targetCtx := strings.TrimSpace(play.Topic) targetCtx := strings.TrimSpace(play.Topic)
@ -1017,15 +1023,17 @@ func cleanPlayGeneratedText(kind, text string) string {
// GeneratePlayStep — 互回/串場劇本:依人設真 LLM 產一步正文。 // GeneratePlayStep — 互回/串場劇本:依人設真 LLM 產一步正文。
// 注意:與 compose mimic 相同OpenCode reasoning 模型可能 3090s 且偶發空白回覆。 // 注意:與 compose mimic 相同OpenCode reasoning 模型可能 3090s 且偶發空白回覆。
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) contextText = strings.TrimSpace(contextText)
topic = strings.TrimSpace(topic) topic = strings.TrimSpace(topic)
if contextText == "" && topic == "" { if contextText == "" && topic == "" {
return "", fmt.Errorf("%w: empty context", domain.ErrValidation) 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 return "", err
} }
defer charge.Settle(ctx, &err)
mode = strings.ToLower(strings.TrimSpace(mode)) mode = strings.ToLower(strings.TrimSpace(mode))
if mode == "" { if mode == "" {
mode = "reply" mode = "reply"
@ -1377,14 +1385,11 @@ func (s *Service) publishWithLease(ctx context.Context, bundleID, stepID, leaseO
// ---------- Compose (CP) ---------- // ---------- 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) sourceText = strings.TrimSpace(sourceText)
if sourceText == "" { if sourceText == "" {
return "", fmt.Errorf("%w: empty source", domain.ErrValidation) 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) personaID = strings.TrimSpace(personaID)
if personaID == "" { if personaID == "" {
return "", fmt.Errorf("%w: 請選擇已完成人設分析的人設", domain.ErrValidation) 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 { if p.Status != domain.PersonaReady {
return "", fmt.Errorf("%w: 選擇的人設尚未完成分析", domain.ErrValidation) 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) fp := personaExpressionFingerprintBlock(p)
// 指紋過長會拖慢 reasoning 模型 // 指紋過長會拖慢 reasoning 模型
if utf8.RuneCountInString(fp) > 800 { if utf8.RuneCountInString(fp) > 800 {
@ -1490,10 +1500,12 @@ func buildMimicPrompt(fp, sourceText, direction, notesBlock string) string {
`, direction, fp, sourceText, notesBlock)) `, direction, fp, sourceText, notesBlock))
} }
func (s *Service) AnalyzeViral(ctx context.Context, ownerUID int64, text string) (*domain.ViralAnalysis, error) { func (s *Service) AnalyzeViral(ctx context.Context, ownerUID int64, text string) (_ *domain.ViralAnalysis, err error) {
if err := s.billAI(ctx, ownerUID, "compose analyze viral", "compose.analyzeViral"); err != nil { charge, err := s.billAI(ctx, ownerUID, "compose analyze viral", "compose.analyzeViral")
if err != nil {
return nil, err return nil, err
} }
defer charge.Settle(ctx, &err)
return s.analyzeViralUnbilled(ctx, ownerUID, text) 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)) 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 var wg sync.WaitGroup
limit := make(chan struct{}, 4) limit := make(chan struct{}, 4)
for i, th := range threads { for i, th := range threads {
@ -1904,7 +1919,12 @@ func (s *Service) syncOwnPostsFromThreads(ctx context.Context, ownerUID int64, a
case <-ctx.Done(): case <-ctx.Done():
return 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) }(i, th.ID)
} }
wg.Wait() wg.Wait()
@ -1917,6 +1937,15 @@ func (s *Service) syncOwnPostsFromThreads(ctx context.Context, ownerUID int64, a
ins := insights[i] ins := insights[i]
prev := byMedia[th.ID] 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 id := "op_" + th.ID
if prev != nil && prev.ID != "" { if prev != nil && prev.ID != "" {
id = prev.ID id = prev.ID
@ -2104,7 +2133,7 @@ func mergeReplyStatus(prev, next []domain.OwnPostReply) []domain.OwnPostReply {
return next 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) post, err := s.getOwnPostOwned(ctx, ownerUID, postID)
if err != nil { if err != nil {
return "", err 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) 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 return "", err
} }
defer charge.Settle(ctx, &err)
persona := s.loadPersonaForGen(ctx, ownerUID, personaID) persona := s.loadPersonaForGen(ctx, ownerUID, personaID)
fp := personaExpressionFingerprintBlock(persona) fp := personaExpressionFingerprintBlock(persona)
@ -2230,7 +2261,7 @@ func (s *Service) SendReply(ctx context.Context, ownerUID int64, postID, replyID
return post, nil 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) post, err := s.getOwnPostOwned(ctx, ownerUID, postID)
if err != nil { if err != nil {
return nil, err return nil, err
@ -2238,9 +2269,11 @@ func (s *Service) AnalyzePost(ctx context.Context, ownerUID int64, postID string
if strings.TrimSpace(post.Text) == "" { if strings.TrimSpace(post.Text) == "" {
return nil, fmt.Errorf("%w: 貼文沒有文字可分析(圖片/純媒體貼可改貼上說明再分析)", domain.ErrValidation) 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 return nil, err
} }
defer charge.Settle(ctx, &err)
va, err := s.analyzeViralUnbilled(ctx, ownerUID, post.Text) va, err := s.analyzeViralUnbilled(ctx, ownerUID, post.Text)
if err != nil { if err != nil {
return nil, err return nil, err
@ -2371,14 +2404,16 @@ func mentionContextSnippet(hit FetchedMention) string {
return kind + " · " + text 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) m, err := s.getMentionOwned(ctx, ownerUID, id)
if err != nil { if err != nil {
return nil, err 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 return nil, err
} }
defer charge.Settle(ctx, &err)
persona := s.loadPersonaForGen(ctx, ownerUID, personaID) persona := s.loadPersonaForGen(ctx, ownerUID, personaID)
fp := personaExpressionFingerprintBlock(persona) fp := personaExpressionFingerprintBlock(persona)
prompt := buildMentionReplyPrompt(fp, m.FromUsername, m.Text, m.ContextSnippet) 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 ---------- // ---------- helpers ----------
func (s *Service) billAI(ctx context.Context, ownerUID int64, label, source string) error { // aiCharge is one reserved AI credit. Quota is reserved up front so an over-limit member is
if s.Usage == nil { // stopped before any work starts, but the audit event is only written once the work succeeded —
return nil // 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) mode, err := s.Usage.PrepareCall(ctx, ownerUID, usageDomain.MeterAICopy)
if err != nil { 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) { func (s *Service) getOwnPostOwned(ctx context.Context, ownerUID int64, id string) (*domain.OwnPost, error) {

View File

@ -280,11 +280,11 @@ func (s *Service) Disconnect(ctx context.Context, ownerUID int64, accountID stri
} }
// WebRedirect builds FE redirect after callback. // 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 { func (s *Service) WebRedirect(ok bool, msg string) string {
base := trimSlash(s.WebBase) base := trimSlash(s.WebBase)
if base == "" {
base = "http://127.0.0.1:5173"
}
if ok { if ok {
return base + "/app/crew?oauth=ok" return base + "/app/crew?oauth=ok"
} }

View File

@ -6,6 +6,10 @@ type Repository interface {
// ReservePlatform atomically guards both limits and increments one monthly counter document. // ReservePlatform atomically guards both limits and increments one monthly counter document.
// A negative limit means unlimited. BYOK must never call this operation. // 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 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) GetMonthlyCounter(ctx context.Context, uid int64, monthKey string) (*MonthlyCounter, error)
InsertEvent(ctx context.Context, e *Event) error InsertEvent(ctx context.Context, e *Event) error

View File

@ -58,6 +58,25 @@ func (s *MemoryStore) ReservePlatform(_ context.Context, uid int64, monthKey, me
return nil 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) { func (s *MemoryStore) GetMonthlyCounter(_ context.Context, uid int64, monthKey string) (*domain.MonthlyCounter, error) {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()

View File

@ -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") 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) { func (s *MonStore) GetMonthlyCounter(ctx context.Context, uid int64, monthKey string) (*domain.MonthlyCounter, error) {
var counter domain.MonthlyCounter var counter domain.MonthlyCounter
err := s.counters.FindOne(ctx, &counter, bson.M{"_id": monthlyCounterID(uid, monthKey)}) err := s.counters.FindOne(ctx, &counter, bson.M{"_id": monthlyCounterID(uid, monthKey)})

View File

@ -81,17 +81,39 @@ func (s *Service) PrepareCall(ctx context.Context, uid int64, meter string) (key
if !has { if !has {
return "", domain.ErrNoKey return "", domain.ErrNoKey
} }
if mode == domain.KeyModePlatform { if err := s.PrepareCallWithMode(ctx, uid, meter, mode); err != nil {
if s.Gate != nil && !s.Gate.AllowPlatform(meter) {
return "", domain.ErrPlatformCapacity
}
if err := s.reservePlatform(ctx, uid, meter); err != nil {
return "", err return "", err
} }
}
return mode, nil 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). // MarkPlatformLimited cools down platform keys (e.g. after upstream 429).
// BYOK users are unaffected. until zero clears the cool-down. // BYOK users are unaffected. until zero clears the cool-down.
func (s *Service) MarkPlatformLimited(until time.Time) { func (s *Service) MarkPlatformLimited(until time.Time) {

View File

@ -295,6 +295,66 @@ func TestUSG_14_MemberCannotSetOthersPrefs(t *testing.T) {
require.ErrorIs(t, err, domain.ErrForbidden) 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) { func TestCreditInviteBonus_AddsToPool(t *testing.T) {
svc := newUsage(42, domain.KeyModePlatform) svc := newUsage(42, domain.KeyModePlatform)
require.NoError(t, svc.CreditInviteBonus(context.Background(), 42, 100, "r1")) require.NoError(t, svc.CreditInviteBonus(context.Background(), 42, 100, "r1"))

View File

@ -192,6 +192,8 @@ func mapError(err error) (int, Envelope) {
return http.StatusBadRequest, Envelope{Code: 400060, Message: cleanBizMessage(err.Error())} return http.StatusBadRequest, Envelope{Code: 400060, Message: cleanBizMessage(err.Error())}
case errors.Is(err, inspireDomain.ErrRemoved), errors.Is(err, scoutDomain.ErrTopicRemoved): case errors.Is(err, inspireDomain.ErrRemoved), errors.Is(err, scoutDomain.ErrTopicRemoved):
return http.StatusGone, Envelope{Code: 410002, Message: err.Error()} 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): case errors.Is(err, scoutDomain.ErrNoCrawlerSession):
return http.StatusBadRequest, Envelope{Code: 400061, Message: "crawler session required when dev_mode enabled"} return http.StatusBadRequest, Envelope{Code: 400061, Message: "crawler session required when dev_mode enabled"}
case errors.Is(err, scoutDomain.ErrHasProducts): case errors.Is(err, scoutDomain.ErrHasProducts):

View File

@ -1680,6 +1680,10 @@ type UsagePrefsData struct {
Unlimited bool `json:"unlimited"` Unlimited bool `json:"unlimited"`
} }
type UsagePrefsReq struct {
Uid string `form:"uid,optional"`
}
type UsagePurchasePublic struct { type UsagePurchasePublic struct {
Id string `json:"id"` Id string `json:"id"`
PlanId string `json:"plan_id"` PlanId string `json:"plan_id"`