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())
db := client.Database(c.Mongo.Database)
indexes := map[string][]mongo.IndexModel{
"jobs": {
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "run_after", Value: 1}, {Key: "created_at", Value: 1}}, Options: options.Index().SetName("claim_due_jobs")},
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "completed_at", Value: 1}}, Options: options.Index().SetName("purge_terminal_jobs")},
},
"studio_outbox": {
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "updated_at", Value: 1}}, Options: options.Index().SetName("worker_outbox_status_updated")},
{Keys: bson.D{{Key: "steps.status", Value: 1}, {Key: "steps.scheduled_at", Value: 1}, {Key: "steps.lease_expires_at", Value: 1}}, Options: options.Index().SetName("claim_due_outbox_steps")},
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "updated_at", Value: -1}}, Options: options.Index().SetName("owner_outbox_updated")},
},
}
for collection, models := range indexes {
for collection, models := range indexModels() {
if _, err := db.Collection(collection).Indexes().CreateMany(ctx, models); err != nil {
panic(fmt.Errorf("create %s indexes: %w", collection, err))
}
}
fmt.Println("database indexes initialized")
}
// ownerIndex covers the near-universal "list one member's rows, newest first" shape. The sort key
// belongs in the index too, otherwise Mongo fetches every matching row before sorting.
func ownerIndex(sortField, name string) mongo.IndexModel {
return mongo.IndexModel{
Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: sortField, Value: -1}},
Options: options.Index().SetName(name),
}
}
// indexModels lists every index the running code depends on. Owner-scoped collections were all
// unindexed, so each list endpoint was a full collection scan that grew with total rows across
// every member rather than with the member's own data.
//
// These are deliberately all non-unique. Adding a uniqueness constraint to existing data can make
// CreateMany fail and take a deploy down, so that belongs in its own migration with a duplicate
// pre-check rather than here.
func indexModels() map[string][]mongo.IndexModel {
return map[string][]mongo.IndexModel{
// members.uid is read on every authenticated request via the auth middleware, and it is a
// field lookup rather than _id, so without this it scanned the whole collection whenever
// the Redis cache missed.
"members": {
{Keys: bson.D{{Key: "uid", Value: 1}}, Options: options.Index().SetName("member_uid")},
{Keys: bson.D{{Key: "email", Value: 1}}, Options: options.Index().SetName("member_email")},
{Keys: bson.D{{Key: "phone", Value: 1}}, Options: options.Index().SetName("member_phone")},
},
"identities": {
{Keys: bson.D{{Key: "login_id", Value: 1}, {Key: "platform", Value: 1}}, Options: options.Index().SetName("identity_login_platform")},
},
"member_settings": {
{Keys: bson.D{{Key: "uid", Value: 1}}, Options: options.Index().SetName("settings_uid")},
},
"jobs": {
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "run_after", Value: 1}, {Key: "created_at", Value: 1}}, Options: options.Index().SetName("claim_due_jobs")},
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "completed_at", Value: 1}}, Options: options.Index().SetName("purge_terminal_jobs")},
// The jobs list is polled continuously by the UI and runs a count plus a find per poll.
ownerIndex("updated_at", "owner_jobs_updated"),
},
// One notification is written per job state change and nothing purges them, so this
// collection grows without bound and its scan cost grows with it.
"notifications": {
ownerIndex("created_at", "owner_notifications_created"),
},
"threads_accounts": {
ownerIndex("created_at", "owner_accounts_created"),
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "threads_user_id", Value: 1}}, Options: options.Index().SetName("owner_threads_user")},
},
"studio_outbox": {
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "updated_at", Value: 1}}, Options: options.Index().SetName("worker_outbox_status_updated")},
{Keys: bson.D{{Key: "steps.status", Value: 1}, {Key: "steps.scheduled_at", Value: 1}, {Key: "steps.lease_expires_at", Value: 1}}, Options: options.Index().SetName("claim_due_outbox_steps")},
ownerIndex("updated_at", "owner_outbox_updated"),
},
"studio_personas": {ownerIndex("updated_at", "owner_personas_updated")},
"studio_plays": {ownerIndex("updated_at", "owner_plays_updated")},
"studio_own_posts": {
ownerIndex("published_at", "owner_own_posts_published"),
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "account_id", Value: 1}, {Key: "published_at", Value: -1}}, Options: options.Index().SetName("owner_account_own_posts")},
},
"studio_mentions": {
ownerIndex("created_at", "owner_mentions_created"),
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "account_id", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("owner_account_mentions")},
},
"scout_brands": {{Keys: bson.D{{Key: "owner_uid", Value: 1}}, Options: options.Index().SetName("owner_brands")}},
"scout_products": {{Keys: bson.D{{Key: "owner_uid", Value: 1}}, Options: options.Index().SetName("owner_products")}},
"scout_posts": {
ownerIndex("created_at", "owner_posts_created"),
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "brand_id", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("owner_brand_posts")},
},
"scout_homework": {{Keys: bson.D{{Key: "owner_uid", Value: 1}}, Options: options.Index().SetName("owner_homework")}},
"inspire_elements": {ownerIndex("updated_at", "owner_elements_updated")},
"inspire_sessions": {ownerIndex("updated_at", "owner_sessions_updated")},
"usage_events": {
{Keys: bson.D{{Key: "uid", Value: 1}, {Key: "month_key", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("uid_month_events")},
{Keys: bson.D{{Key: "month_key", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("tenant_month_events")},
{Keys: bson.D{{Key: "created_at", Value: 1}}, Options: options.Index().SetName("events_created_range")},
},
"usage_prefs": {
{Keys: bson.D{{Key: "uid", Value: 1}}, Options: options.Index().SetName("prefs_uid")},
// Stripe webhooks arrive knowing only the customer, so this is the reverse lookup
// back to a member and it runs on every billing event.
{Keys: bson.D{{Key: "stripe_customer_id", Value: 1}}, Options: options.Index().SetName("prefs_stripe_customer")},
},
"usage_purchases": {
{Keys: bson.D{{Key: "uid", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("uid_purchases_created")},
{Keys: bson.D{{Key: "stripe_event_id", Value: 1}}, Options: options.Index().SetName("purchase_stripe_event")},
},
"growth_outcomes": {
// The observe worker scans by status across all tenants on every tick.
{Keys: bson.D{{Key: "status", Value: 1}}, Options: options.Index().SetName("observe_status")},
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "source_type", Value: 1}, {Key: "source_id", Value: 1}}, Options: options.Index().SetName("owner_outcome_source")},
ownerIndex("created_at", "owner_outcomes_created"),
},
"growth_checkups": {ownerIndex("created_at", "owner_checkups_created")},
"growth_account_health": {{Keys: bson.D{{Key: "owner_uid", Value: 1}}, Options: options.Index().SetName("owner_health")}},
"growth_workspaces": {{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "archived", Value: 1}}, Options: options.Index().SetName("owner_workspaces_archived")}},
"growth_draft_reviews": {
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "ref_type", Value: 1}, {Key: "ref_id", Value: 1}, {Key: "updated_at", Value: -1}}, Options: options.Index().SetName("owner_review_ref")},
},
"growth_invite_rewards": {
{Keys: bson.D{{Key: "inviter_uid", Value: 1}, {Key: "created_at", Value: -1}}, Options: options.Index().SetName("inviter_rewards_created")},
{Keys: bson.D{{Key: "invitee_uid", Value: 1}}, Options: options.Index().SetName("invitee_reward")},
},
"growth_utm_links": {
// Looked up by code on the unauthenticated public redirect.
{Keys: bson.D{{Key: "code", Value: 1}}, Options: options.Index().SetName("utm_code")},
ownerIndex("created_at", "owner_utm_created"),
},
"growth_ws_members": {
{Keys: bson.D{{Key: "workspace_id", Value: 1}, {Key: "uid", Value: 1}}, Options: options.Index().SetName("workspace_member")},
},
"billing_checkout_attempts": {
{Keys: bson.D{{Key: "uid", Value: 1}, {Key: "request_id", Value: 1}}, Options: options.Index().SetName("uid_request")},
{Keys: bson.D{{Key: "session_id", Value: 1}}, Options: options.Index().SetName("checkout_session")},
{Keys: bson.D{{Key: "customer_id", Value: 1}}, Options: options.Index().SetName("checkout_customer")},
},
}
}

View File

@ -164,8 +164,15 @@ func main() {
ctx := context.Background()
const outboxLockTTL = 3 * time.Minute
outboxLockKey := strings.Trim(c.Redis.Namespace, ":") + ":worker:outbox"
outboxLock := redislock.NewWithClient(redisClient, outboxLockKey, workerID, outboxLockTTL)
ns := strings.Trim(c.Redis.Namespace, ":")
outboxLock := redislock.NewWithClient(redisClient, ns+":worker:outbox", workerID, outboxLockTTL)
// 巡場維護(過期 outcome、清終態 job是全域掃描不像 job 領取有 guarded update 擋重複。
// 每個 worker 每 tick 各跑一次的話N 台就等於同一份掃描做 N 次。改成一次一台、且拉長間隔。
const maintenanceLockTTL = 2 * time.Minute
const maintenanceEvery = time.Minute
maintenanceLock := redislock.NewWithClient(redisClient, ns+":worker:maintenance", workerID, maintenanceLockTTL)
lastMaintenance := time.Time{}
for {
select {
case <-sig:
@ -231,22 +238,52 @@ func main() {
}
// 2) One worker owns an outbox tick and renews its lease while claims run.
processOutbox(ctx, studio, outboxLock, workerID, outboxLockTTL)
// 3) growth-loop: expire observing outcomes + light health refresh
if n, err := growthSvc.ObserveTick(ctx, 0); err != nil {
logx.Errorf("worker %s outcome observe: %v", workerID, err)
} else if n > 0 {
logx.Infof("worker %s outcome observe updated %d", workerID, n)
}
// 4) 終態任務超過 2 天自動清除
if purged, err := jobs.PurgeExpiredTerminal(ctx); err != nil {
logx.Errorf("worker %s purge jobs: %v", workerID, err)
} else if purged > 0 {
logx.Infof("worker %s purged %d expired terminal job(s)", workerID, purged)
// 3) 巡場維護:過期 outcome + 清終態 job單一 worker、低頻
if time.Since(lastMaintenance) >= maintenanceEvery {
if runMaintenance(ctx, growthSvc, jobs, maintenanceLock, workerID) {
lastMaintenance = time.Now()
}
}
}
}
}
// runMaintenance reports whether this worker actually did the sweep, so a worker that lost the
// lock retries on the next tick instead of waiting out the full interval.
func runMaintenance(
ctx context.Context,
growthSvc *growthUC.Service,
jobs *jobUC.Service,
lock *redislock.Lock,
workerID string,
) bool {
locked, err := lock.Acquire(ctx)
if err != nil {
logx.Errorf("worker %s maintenance lock acquire: %v", workerID, err)
return false
}
if !locked {
return false
}
defer func() {
if rerr := lock.Release(ctx); rerr != nil {
logx.Errorf("worker %s maintenance lock release: %v", workerID, rerr)
}
}()
if n, err := growthSvc.ObserveTick(ctx, 0); err != nil {
logx.Errorf("worker %s outcome observe: %v", workerID, err)
} else if n > 0 {
logx.Infof("worker %s outcome observe updated %d", workerID, n)
}
if purged, err := jobs.PurgeExpiredTerminal(ctx); err != nil {
logx.Errorf("worker %s purge jobs: %v", workerID, err)
} else if purged > 0 {
logx.Infof("worker %s purged %d expired terminal job(s)", workerID, purged)
}
return true
}
func processOutbox(ctx context.Context, studio *studioUC.Service, lock *redislock.Lock, workerID string, ttl time.Duration) {
locked, err := lock.Acquire(ctx)
if err != nil {

View File

@ -61,6 +61,11 @@ type (
Unlimited bool `json:"unlimited"`
}
UsagePrefsReq {
// Uid reads another member's plan; admin only. Empty means the caller's own prefs.
Uid string `form:"uid,optional"`
}
UsageSetPrefsReq {
Uid string `json:"uid"`
PlanId string `json:"plan_id,optional"`
@ -165,7 +170,7 @@ service gateway {
get /events (UsageEventsReq) returns (UsageEventsData)
@handler GetUsagePrefs
get /prefs returns (UsagePrefsData)
get /prefs (UsagePrefsReq) returns (UsagePrefsData)
@handler PurchasePlan
post /purchase (UsagePurchaseReq) returns (UsagePurchasePublic)

View File

@ -1265,35 +1265,6 @@ func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
rest.WithPrefix("/api/v1/utm"),
)
server.AddRoutes(
rest.WithMiddlewares(
[]rest.Middleware{serverCtx.AuthJWT},
[]rest.Route{
{
Method: http.MethodPut,
Path: "/:id/branding",
Handler: workspaces.UpdateWorkspaceBrandingHandler(serverCtx),
},
{
Method: http.MethodGet,
Path: "/:id/members",
Handler: workspaces.ListWorkspaceMembersHandler(serverCtx),
},
{
Method: http.MethodPost,
Path: "/:id/members",
Handler: workspaces.AddWorkspaceMemberHandler(serverCtx),
},
{
Method: http.MethodDelete,
Path: "/:id/members/:uid",
Handler: workspaces.RemoveWorkspaceMemberHandler(serverCtx),
},
}...,
),
rest.WithPrefix("/api/v1/workspaces"),
)
server.AddRoutes(
rest.WithMiddlewares(
[]rest.Middleware{serverCtx.AuthJWT},
@ -1352,4 +1323,33 @@ func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
),
rest.WithPrefix("/api/v1/workspaces"),
)
server.AddRoutes(
rest.WithMiddlewares(
[]rest.Middleware{serverCtx.AuthJWT},
[]rest.Route{
{
Method: http.MethodPut,
Path: "/:id/branding",
Handler: workspaces.UpdateWorkspaceBrandingHandler(serverCtx),
},
{
Method: http.MethodGet,
Path: "/:id/members",
Handler: workspaces.ListWorkspaceMembersHandler(serverCtx),
},
{
Method: http.MethodPost,
Path: "/:id/members",
Handler: workspaces.AddWorkspaceMemberHandler(serverCtx),
},
{
Method: http.MethodDelete,
Path: "/:id/members/:uid",
Handler: workspaces.RemoveWorkspaceMemberHandler(serverCtx),
},
}...,
),
rest.WithPrefix("/api/v1/workspaces"),
)
}

View File

@ -9,12 +9,20 @@ import (
"apps/backend/internal/logic/usage"
"apps/backend/internal/response"
"apps/backend/internal/svc"
"apps/backend/internal/types"
"github.com/zeromicro/go-zero/rest/httpx"
)
func GetUsagePrefsHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
var req types.UsagePrefsReq
if err := httpx.Parse(r, &req); err != nil {
response.Write(r.Context(), w, nil, response.WrapRequestError(err))
return
}
l := usage.NewGetUsagePrefsLogic(r.Context(), svcCtx)
data, err := l.GetUsagePrefs()
data, err := l.GetUsagePrefs(&req)
response.Write(r.Context(), w, data, err)
}
}

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}
}
func (l *AICompleteLogic) AIComplete(req *types.AICompleteReq) (*types.AICompleteData, error) {
func (l *AICompleteLogic) AIComplete(req *types.AICompleteReq) (_ *types.AICompleteData, err error) {
if l.svcCtx.Usage == nil || l.svcCtx.AI == nil || l.svcCtx.KeyResolver == nil {
return nil, response.Biz(503, 503001, "ai proxy not configured")
}
@ -34,14 +34,15 @@ func (l *AICompleteLogic) AIComplete(req *types.AICompleteReq) (*types.AIComplet
if meter == "" {
meter = usageDomain.MeterAICopy
}
mode, err := l.svcCtx.Usage.PrepareCall(l.ctx, uid, meter)
if err != nil {
return nil, err
}
mode, apiKey, err := l.svcCtx.KeyResolver.ResolveKey(l.ctx, uid, meter)
if err != nil {
return nil, err
}
charge, err := reserveQuota(l.ctx, l.svcCtx, uid, meter, mode, "AI complete", "proxy.ai")
if err != nil {
return nil, err
}
defer charge.Settle(l.ctx, &err)
model := req.Model
if model == "" {
model = "grok-3"
@ -51,6 +52,5 @@ func (l *AICompleteLogic) AIComplete(req *types.AICompleteReq) (*types.AIComplet
if err != nil {
return nil, response.Biz(502, 502001, err.Error())
}
_, _ = l.svcCtx.Usage.RecordCall(l.ctx, uid, meter, mode, "AI complete", "proxy.ai")
return &types.AICompleteData{Text: text, KeyMode: mode, Model: model}, nil
}

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

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 (
"context"
"strings"
threadsDomain "apps/backend/internal/module/threads/domain"
"apps/backend/internal/svc"
@ -22,7 +23,8 @@ func NewThreadsOAuthCallbackLogic(ctx context.Context, svcCtx *svc.ServiceContex
func (l *ThreadsOAuthCallbackLogic) ThreadsOAuthCallback(req *types.ThreadsOAuthCallbackReq) (*types.ThreadsOAuthCallbackData, error) {
if l.svcCtx.Threads == nil {
redir := "http://127.0.0.1:5173/app/crew?oauth=error&msg=not_configured"
// 相對路徑:讓瀏覽器留在目前站台,而不是被丟到某台機器的 loopback。
redir := strings.TrimRight(l.svcCtx.Config.PublicWebBase, "/") + "/app/crew?oauth=error&msg=not_configured"
return &types.ThreadsOAuthCallbackData{Ok: false, Message: "not configured", RedirectUrl: redir}, nil
}
_, err := l.svcCtx.Threads.OAuthCallback(l.ctx, req.Code, req.State, req.Error)

View File

@ -2,8 +2,11 @@ package usage
import (
"context"
"strconv"
"strings"
"apps/backend/internal/middleware"
"apps/backend/internal/module/permission"
"apps/backend/internal/response"
"apps/backend/internal/svc"
"apps/backend/internal/types"
@ -21,7 +24,7 @@ func NewGetUsagePrefsLogic(ctx context.Context, svcCtx *svc.ServiceContext) *Get
return &GetUsagePrefsLogic{Logger: logx.WithContext(ctx), ctx: ctx, svcCtx: svcCtx}
}
func (l *GetUsagePrefsLogic) GetUsagePrefs() (*types.UsagePrefsData, error) {
func (l *GetUsagePrefsLogic) GetUsagePrefs(req *types.UsagePrefsReq) (*types.UsagePrefsData, error) {
if l.svcCtx.Usage == nil {
return nil, response.Biz(503, 503001, "usage not configured")
}
@ -29,6 +32,20 @@ func (l *GetUsagePrefsLogic) GetUsagePrefs() (*types.UsagePrefsData, error) {
if !ok {
return nil, response.Biz(401, 401001, "missing authorization")
}
// Reading somebody else's plan is admin-only; the admin console needs it to show the plan of
// the member being edited rather than the admin's own.
if target := strings.TrimSpace(req.Uid); target != "" {
parsed, err := strconv.ParseInt(target, 10, 64)
if err != nil {
return nil, response.Biz(400, 400001, "invalid uid")
}
if parsed != uid {
if !permission.IsAdmin(middleware.RolesFrom(l.ctx)) {
return nil, response.Biz(403, 403002, "permission denied")
}
uid = parsed
}
}
p, err := l.svcCtx.Usage.GetPrefs(l.ctx, uid)
if err != nil {
return nil, err

View File

@ -142,7 +142,11 @@ type Repository interface {
GetSubscription(context.Context, int64) (*SubscriptionState, error)
FindUIDByCustomer(context.Context, string) (int64, error)
ApplyEntitlement(context.Context, EntitlementUpdate) (bool, error)
WebhookProcessed(context.Context, string) (bool, error)
MarkWebhookProcessed(context.Context, string, string, int64) error
// ClaimWebhook atomically records an event as being handled, reporting false when another
// delivery already claimed it. Stripe retries aggressively and can deliver the same event
// concurrently, so a read-then-write check would let two deliveries both proceed.
ClaimWebhook(ctx context.Context, id, eventType string, created int64) (claimed bool, err error)
// ReleaseWebhookClaim drops a claim whose processing failed, so the next delivery retries.
ReleaseWebhookClaim(ctx context.Context, id string) error
UpsertPurchase(context.Context, *PurchaseAudit) error
}

View File

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

View File

@ -9,6 +9,7 @@ import (
"time"
"github.com/google/uuid"
"github.com/zeromicro/go-zero/core/logx"
)
type Service struct {
@ -176,14 +177,22 @@ func (s *Service) HandleWebhook(ctx context.Context, payload []byte, signature s
if err != nil {
return err
}
done, err := s.Repo.WebhookProcessed(ctx, e.ID)
if err != nil || done {
// Claim before processing, so two concurrent deliveries of the same event cannot both
// fulfil it. Processing failures release the claim so Stripe's retry still gets through.
claimed, err := s.Repo.ClaimWebhook(ctx, e.ID, e.Type, e.Created)
if err != nil {
return err
}
if !claimed {
return nil
}
if err := s.processEvent(ctx, e); err != nil {
if rerr := s.Repo.ReleaseWebhookClaim(ctx, e.ID); rerr != nil {
logx.Errorf("billing release webhook claim %s: %v", e.ID, rerr)
}
return err
}
return s.Repo.MarkWebhookProcessed(ctx, e.ID, e.Type, e.Created)
return nil
}
func (s *Service) processEvent(ctx context.Context, e *Event) error {

View File

@ -56,6 +56,13 @@ type fakeRepo struct {
lastEvent string
purchases map[string]*PurchaseAudit
entitlementErr error
applyCalls int
}
func (r *fakeRepo) applyCallCount() int {
r.mu.Lock()
defer r.mu.Unlock()
return r.applyCalls
}
func newFakeRepo() *fakeRepo {
@ -135,6 +142,7 @@ func (r *fakeRepo) FindUIDByCustomer(_ context.Context, customer string) (int64,
func (r *fakeRepo) ApplyEntitlement(_ context.Context, u EntitlementUpdate) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.applyCalls++
if r.entitlementErr != nil {
return false, r.entitlementErr
}
@ -150,15 +158,19 @@ func (r *fakeRepo) ApplyEntitlement(_ context.Context, u EntitlementUpdate) (boo
}
return true, nil
}
func (r *fakeRepo) WebhookProcessed(_ context.Context, id string) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
return r.processed[id], nil
}
func (r *fakeRepo) MarkWebhookProcessed(_ context.Context, id, _ string, _ int64) error {
func (r *fakeRepo) ClaimWebhook(_ context.Context, id, _ string, _ int64) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
if r.processed[id] {
return false, nil
}
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
}
func (r *fakeRepo) UpsertPurchase(_ context.Context, p *PurchaseAudit) error {
@ -225,6 +237,58 @@ func TestSubscriptionRemainsReadableWhenStripeDisabled(t *testing.T) {
require.Equal(t, PlanFree, state.PlanID)
}
// Stripe can deliver the same event concurrently. The claim has to be atomic, otherwise both
// deliveries see "not processed yet" and fulfil the event twice.
func TestWebhookConcurrentDeliveriesProcessEventOnce(t *testing.T) {
r := newFakeRepo()
r.state = SubscriptionState{UID: 7, PlanID: PlanFree, CustomerID: "cus_1"}
p := &fakeProvider{event: &Event{
ID: "evt_concurrent", Type: "customer.subscription.updated", Created: 20,
ObjectID: "sub_1", CustomerID: "cus_1", PriceID: "price_pro", Status: "active",
}}
s := newService(r, p)
const deliveries = 8
var wg sync.WaitGroup
errs := make(chan error, deliveries)
for i := 0; i < deliveries; i++ {
wg.Add(1)
go func() {
defer wg.Done()
errs <- s.HandleWebhook(context.Background(), []byte("payload"), "sig")
}()
}
wg.Wait()
close(errs)
for err := range errs {
require.NoError(t, err)
}
require.Equal(t, 1, r.applyCallCount(), "concurrent duplicates must fulfil the event exactly once")
require.Equal(t, PlanPro, r.state.PlanID)
require.Len(t, r.purchases, 1)
}
// Claiming up front must not swallow the event when processing fails, or a transient outage
// would permanently lose a payment.
func TestWebhookFailedProcessingIsRetryable(t *testing.T) {
r := newFakeRepo()
r.state = SubscriptionState{UID: 7, PlanID: PlanFree, CustomerID: "cus_1"}
r.entitlementErr = errors.New("mongo unavailable")
p := &fakeProvider{event: &Event{
ID: "evt_retry", Type: "customer.subscription.updated", Created: 20,
ObjectID: "sub_1", CustomerID: "cus_1", PriceID: "price_pro", Status: "active",
}}
s := newService(r, p)
require.Error(t, s.HandleWebhook(context.Background(), []byte("payload"), "sig"))
require.Equal(t, PlanFree, r.state.PlanID)
// Stripe retries once the outage clears; the claim must have been released.
r.entitlementErr = nil
require.NoError(t, s.HandleWebhook(context.Background(), []byte("payload"), "sig"))
require.Equal(t, PlanPro, r.state.PlanID)
}
func TestWebhookDuplicateOutOfOrderDowngradeAndAdminOverride(t *testing.T) {
r := newFakeRepo()
r.state = SubscriptionState{UID: 7, PlanID: PlanFree, CustomerID: "cus_1"}

View File

@ -6,11 +6,14 @@ import (
"encoding/hex"
"fmt"
"math"
"net/url"
"sort"
"strings"
"unicode"
"apps/backend/internal/module/growth/domain"
"github.com/zeromicro/go-zero/core/logx"
)
// --- Playbooks ---
@ -224,10 +227,28 @@ func tokenizeTC(s string) []string {
// --- UTM ---
func (s *Service) CreateUtmLink(ctx context.Context, ownerUID int64, dest, outcomeID, label string) (*domain.UtmLink, error) {
// validateRedirectTarget keeps /u/{code} from becoming an open redirect that can point anywhere.
// A HasPrefix("http") check is not enough: it lets through "httpx://", scheme-relative "//evil",
// and anything unparseable, all of which the browser resolves differently than we assume.
func validateRedirectTarget(dest string) (string, error) {
dest = strings.TrimSpace(dest)
if dest == "" || !strings.HasPrefix(dest, "http") {
return nil, fmt.Errorf("%w: destination_url must be http(s)", domain.ErrValidation)
u, err := url.Parse(dest)
if err != nil {
return "", fmt.Errorf("%w: destination_url is not a valid URL", domain.ErrValidation)
}
if u.Scheme != "http" && u.Scheme != "https" {
return "", fmt.Errorf("%w: destination_url must be http or https", domain.ErrValidation)
}
if u.Host == "" {
return "", fmt.Errorf("%w: destination_url must include a host", domain.ErrValidation)
}
return u.String(), nil
}
func (s *Service) CreateUtmLink(ctx context.Context, ownerUID int64, dest, outcomeID, label string) (*domain.UtmLink, error) {
dest, err := validateRedirectTarget(dest)
if err != nil {
return nil, err
}
code := domain.NewID()[:10]
u := &domain.UtmLink{
@ -259,10 +280,14 @@ func (s *Service) TrackUtmClick(ctx context.Context, code string) (dest string,
e.Status = domain.StatusConfirmed
}
e.UpdatedAt = domain.NowNano()
_ = s.Repo.SaveOutcome(ctx, e)
if serr := s.Repo.SaveOutcome(ctx, e); serr != nil {
// 點擊已經發生redirect 照走;但歸因掉了要看得見,否則成效數字會默默偏低。
logx.WithContext(ctx).Errorf("utm click attribution not saved: outcome=%s err=%v", e.ID, serr)
}
}
}
return u.DestinationURL, nil
// 舊資料是在只檢查 "http" 前綴時寫入的,出門前再驗一次。
return validateRedirectTarget(u.DestinationURL)
}
// --- Benchmark ---

View File

@ -168,3 +168,55 @@ func TestWS_05_ExportReport(t *testing.T) {
require.Contains(t, fn, "2026-07")
require.Contains(t, body, "月報")
}
// /u/{code} 是免登入 302任何存進去的 destination 都會變成本站可用的跳板。
// 舊檢查只看 "http" 前綴,"httpx://" 與 "//evil" 都能過。
func TestUTM_RejectsDestinationsThatAreNotHttpURLs(t *testing.T) {
svc := New(repository.NewMemory())
ctx := context.Background()
for _, dest := range []string{
"javascript:alert(1)",
"httpx://evil.example",
"//evil.example/path",
"http://",
"data:text/html,<script>alert(1)</script>",
"",
" ",
} {
_, err := svc.CreateUtmLink(ctx, 1, dest, "", "label")
require.ErrorIsf(t, err, domain.ErrValidation, "destination %q should be rejected", dest)
}
}
func TestUTM_AcceptsHttpAndHttpsAndRedirectsThere(t *testing.T) {
svc := New(repository.NewMemory())
ctx := context.Background()
for _, dest := range []string{
"https://shop.example/product?utm_source=threads",
"http://shop.example",
} {
link, err := svc.CreateUtmLink(ctx, 1, dest, "", "label")
require.NoError(t, err)
got, err := svc.TrackUtmClick(ctx, link.Code)
require.NoError(t, err)
require.Equal(t, dest, got)
}
}
// 這個檢查是後來才加的,先前寫進 DB 的列可能還帶著壞值,所以出門前要再驗一次。
func TestUTM_RefusesToRedirectToAStoredBadDestination(t *testing.T) {
repo := repository.NewMemory()
svc := New(repo)
ctx := context.Background()
link, err := svc.CreateUtmLink(ctx, 1, "https://shop.example", "", "label")
require.NoError(t, err)
link.DestinationURL = "javascript:alert(1)"
require.NoError(t, repo.SaveUtmLink(ctx, link))
_, err = svc.TrackUtmClick(ctx, link.Code)
require.ErrorIs(t, err, domain.ErrValidation)
}

View File

@ -10,6 +10,9 @@ var (
ErrForbidden = errors.New("inspire forbidden")
ErrValidation = errors.New("inspire validation")
ErrRemoved = errors.New("legacy inspire API removed")
// ErrNotImplemented marks a capability that is wired but has no real backend yet. It must be
// returned instead of fabricated data, and must never be billed.
ErrNotImplemented = errors.New("not implemented")
)
func NowNano() int64 { return time.Now().UTC().UnixNano() }

View File

@ -224,23 +224,27 @@ func TestIN_06_ResearchSearch(t *testing.T) {
require.NotEmpty(t, hits[0].URL)
}
func TestIN_07_GenerateImage(t *testing.T) {
// There is no image backend yet, so the endpoint must say so rather than charge for a
// placeholder avatar from a public URL.
func TestIN_07_GenerateImageReportsMissingBackend(t *testing.T) {
svc := newInspire()
uid := int64(5_001_017)
setupInspireUID(svc, uid)
img, err := svc.GenerateImage(context.Background(), uid, "溫暖海邊插畫")
require.ErrorIs(t, err, domain.ErrNotImplemented)
require.Nil(t, img)
sum, err := svc.Usage.GetSummary(context.Background(), uid, "")
require.NoError(t, err)
require.NotEmpty(t, img.URL)
require.NotEmpty(t, img.ID)
require.Zero(t, sum.Platform.CreditsUsed, "an unimplemented feature must not be billed")
}
func TestIN_08_UploadSeparate(t *testing.T) {
func TestIN_08_GenerateImageStillValidatesPrompt(t *testing.T) {
svc := newInspire()
uid := int64(5_001_008)
setupInspireUID(svc, uid)
img, err := svc.GenerateImage(context.Background(), uid, "x")
require.NoError(t, err)
require.True(t, len(img.URL) > 8)
_, err := svc.GenerateImage(context.Background(), uid, " ")
require.ErrorIs(t, err, domain.ErrValidation)
}
func TestIN_09_LegacyRemoved(t *testing.T) {

View File

@ -16,6 +16,7 @@ import (
usageUC "apps/backend/internal/module/usage/usecase"
"github.com/google/uuid"
"github.com/zeromicro/go-zero/core/logx"
)
type Service struct {
@ -54,14 +55,16 @@ func (s *Service) ListTrends(ctx context.Context, ownerUID int64) ([]*domain.Tre
// RefreshTrends — 手動「找靈感」:取得真實結果後才扣點並覆蓋快取。
// 搜尋失敗時保留舊內容;不可用示意種子冒充本次刷新成功。
func (s *Service) RefreshTrends(ctx context.Context, ownerUID int64) ([]*domain.TrendItem, error) {
func (s *Service) RefreshTrends(ctx context.Context, ownerUID int64) (_ []*domain.TrendItem, err error) {
list := s.tryLiveTrends(ctx, ownerUID)
if len(list) == 0 {
return nil, fmt.Errorf("找不到新的可用話題,已保留原本內容")
}
if err := s.bill(ctx, ownerUID, usageDomain.MeterWebSearch, "找靈感話題", "inspire.refreshTrends"); err != nil {
charge, err := s.bill(ctx, ownerUID, usageDomain.MeterWebSearch, "找靈感話題", "inspire.refreshTrends")
if err != nil {
return nil, err
}
defer charge.Settle(ctx, &err)
if err := s.Repo.SaveTrends(ctx, ownerUID, list); err != nil {
return nil, err
}
@ -687,7 +690,7 @@ func generateUserLogMessage(notes, material string) string {
return "【產文】" + notes + "|素材:" + preview
}
func (s *Service) chatInternal(ctx context.Context, ownerUID int64, message string, pinnedIDs []string, mode, personaID, sessionID, material string, useWeb bool, onDelta func(chunk string) error) (*ChatOutcome, error) {
func (s *Service) chatInternal(ctx context.Context, ownerUID int64, message string, pinnedIDs []string, mode, personaID, sessionID, material string, useWeb bool, onDelta func(chunk string) error) (_ *ChatOutcome, err error) {
message = strings.TrimSpace(message)
material = strings.TrimSpace(material)
if mode != "chat" && mode != "generate" {
@ -714,9 +717,11 @@ func (s *Service) chatInternal(ctx context.Context, ownerUID int64, message stri
if mode == "generate" {
label = "inspire rewrite"
}
if err := s.bill(ctx, ownerUID, usageDomain.MeterAICopy, label, "inspire.chat"); err != nil {
charge, err := s.bill(ctx, ownerUID, usageDomain.MeterAICopy, label, "inspire.chat")
if err != nil {
return nil, err
}
defer charge.Settle(ctx, &err)
// 回覆語言:人設指紋/範例優先,其次會員 UI 語系ctx 已由 Auth 注入)
persona := s.resolvePersonaSnap(ctx, ownerUID, personaID)
if mode == "generate" && (persona == nil || persona.Status != "ready") {
@ -1153,14 +1158,16 @@ func (s *Service) buildInspirePrompt(ctx context.Context, ownerUID int64, pinned
return full, blocks, sections
}
func (s *Service) searchPromptContext(ctx context.Context, ownerUID int64, query string) (string, error) {
func (s *Service) searchPromptContext(ctx context.Context, ownerUID int64, query string) (_ string, err error) {
query = strings.TrimSpace(query)
if query == "" {
return "", nil
}
if err := s.bill(ctx, ownerUID, usageDomain.MeterWebSearch, "inspire web search", "inspire.chat"); err != nil {
charge, err := s.bill(ctx, ownerUID, usageDomain.MeterWebSearch, "inspire web search", "inspire.chat")
if err != nil {
return "", err
}
defer charge.Settle(ctx, &err)
key := "fake"
if s.ResolveKey != nil {
if _, resolved, err := s.ResolveKey(ctx, ownerUID, usageDomain.MeterWebSearch); err == nil && resolved != "" {
@ -1438,14 +1445,16 @@ func (s *Service) runInspireLLM(ctx context.Context, ownerUID int64, prompt stri
return "", fmt.Errorf("%w: 請到設定填寫 AI Key", domain.ErrValidation)
}
func (s *Service) ResearchSearch(ctx context.Context, ownerUID int64, query string) ([]*domain.ResearchHit, error) {
func (s *Service) ResearchSearch(ctx context.Context, ownerUID int64, query string) (_ []*domain.ResearchHit, err error) {
query = strings.TrimSpace(query)
if query == "" {
return nil, fmt.Errorf("%w: empty query", domain.ErrValidation)
}
if err := s.bill(ctx, ownerUID, usageDomain.MeterWebSearch, "research search", "research.search"); err != nil {
charge, err := s.bill(ctx, ownerUID, usageDomain.MeterWebSearch, "research search", "research.search")
if err != nil {
return nil, err
}
defer charge.Settle(ctx, &err)
key := "fake"
if s.ResolveKey != nil {
if _, k, err := s.ResolveKey(ctx, ownerUID, usageDomain.MeterWebSearch); err == nil && k != "" {
@ -1473,34 +1482,73 @@ func (s *Service) ResearchSearch(ctx context.Context, ownerUID int64, query stri
return out, nil
}
func (s *Service) GenerateImage(ctx context.Context, ownerUID int64, prompt string) (*domain.GeneratedImage, error) {
prompt = strings.TrimSpace(prompt)
if prompt == "" {
// GenerateImage has no image backend yet. It used to return a random avatar URL from a public
// placeholder service while charging a real ai_image credit, which billed members for a picture
// nobody asked for. Until a real generator is wired in, this reports the gap instead.
func (s *Service) GenerateImage(_ context.Context, _ int64, prompt string) (*domain.GeneratedImage, error) {
if strings.TrimSpace(prompt) == "" {
return nil, fmt.Errorf("%w: empty prompt", domain.ErrValidation)
}
if err := s.bill(ctx, ownerUID, usageDomain.MeterAIImage, "generate image", "media.generateImage"); err != nil {
return nil, err
}
// placeholder SVG data URL (live-complete with billable path; real SD later)
id := "img_" + uuid.NewString()[:10]
// minimal 1x1 png base64 is tiny; use dicebear-like public URL for preview
url := fmt.Sprintf("https://api.dicebear.com/9.x/shapes/svg?seed=%s", id)
return &domain.GeneratedImage{ID: id, URL: url, Prompt: prompt}, nil
return nil, fmt.Errorf("%w: 圖片生成尚未接上生成服務", domain.ErrNotImplemented)
}
// Legacy removed APIs
func (s *Service) LegacyRemoved() error { return domain.ErrRemoved }
func (s *Service) bill(ctx context.Context, uid int64, meter, label, source string) error {
if s.Usage == nil {
return nil
// inspireCharge is one reserved credit; see Settle/Release for the commit-on-success contract.
type inspireCharge struct {
svc *Service
uid int64
meter string
mode string
label string
source string
settled bool
}
// bill reserves credit for one metered call. Pair it with `defer charge.Settle(ctx, &err)` on a
// named error return so every failure path refunds rather than charging for nothing.
func (s *Service) bill(ctx context.Context, uid int64, meter, label, source string) (*inspireCharge, error) {
if s == nil || s.Usage == nil {
return &inspireCharge{settled: true}, nil
}
mode, err := s.Usage.PrepareCall(ctx, uid, meter)
if err != nil {
return err
return nil, err
}
return &inspireCharge{svc: s, uid: uid, meter: meter, mode: mode, label: label, source: source}, nil
}
func (c *inspireCharge) Settle(ctx context.Context, errp *error) {
if errp != nil && *errp != nil {
c.Release(ctx)
return
}
c.Commit(ctx)
}
// Commit writes the audit event. The member already got their result, so a failure here is
// logged rather than turned into an error.
func (c *inspireCharge) Commit(ctx context.Context) {
if c == nil || c.settled {
return
}
c.settled = true
if _, err := c.svc.Usage.RecordCall(ctx, c.uid, c.meter, c.mode, c.label, c.source); err != nil {
logx.Errorf("usage record uid=%d meter=%s source=%s: %v", c.uid, c.meter, c.source, err)
}
}
func (c *inspireCharge) Release(ctx context.Context) {
if c == nil || c.settled {
return
}
c.settled = true
rctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
defer cancel()
if err := c.svc.Usage.ReleaseCall(rctx, c.uid, c.meter, c.mode); err != nil {
logx.Errorf("usage release uid=%d meter=%s source=%s: %v", c.uid, c.meter, c.source, err)
}
_, err = s.Usage.RecordCall(ctx, uid, meter, mode, label, source)
return err
}
func emptySession(ownerUID int64) *domain.Session {

View File

@ -7,6 +7,7 @@ import (
libmongo "apps/backend/internal/lib/mongo"
"apps/backend/internal/module/job/domain"
"github.com/zeromicro/go-zero/core/logx"
"github.com/zeromicro/go-zero/core/stores/mon"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo"
@ -216,12 +217,22 @@ func (s *MonStore) CancelPendingByRef(ctx context.Context, ownerUID int64, templ
if err := s.jobs.Find(ctx, &list, filter); err != nil {
return err
}
// Update is optimistic-concurrency guarded, so a lost race means the job is still pending and
// will run alongside its replacement. Swallowing that error reported a clean cancel that never
// happened; a stale job racing the new one is worth failing the call over.
for _, j := range list {
j.Status = domain.StatusCancelled
j.ProgressSummary = "已由新的定期排程取代"
j.CompletedAt = now
j.UpdatedAt = now
_ = s.Update(ctx, j)
if err := s.Update(ctx, j); err != nil {
// 這一輪已經被別人搶去改(例如 worker 剛領走),不是我們該取消的目標。
if errors.Is(err, domain.ErrIllegalStatus) {
logx.WithContext(ctx).Infof("cancel pending job skipped, changed concurrently: job=%s", j.ID)
continue
}
return err
}
}
return nil
}

View File

@ -128,7 +128,9 @@ func (s *Service) StartDemo(ctx context.Context, ownerUID int64) (*domain.Job, e
j.Status = domain.StatusQueued
j.ProgressSummary = "Demo 測試任務 · 等待 worker 領取"
j.UpdatedAt = domain.NowNano()
_ = s.Repo.Update(ctx, j)
if err := s.Repo.Update(ctx, j); err != nil {
return nil, err
}
s.notify(ctx, j)
return j, nil
}
@ -288,7 +290,10 @@ func (s *Service) ScheduleTokenRenew(ctx context.Context, ownerUID int64, accoun
if runAfter <= 0 {
runAfter = now + domain.DefaultTokenRenewDelayNs
}
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplateThreadsTokenRenew, accountID)
// 取消沒成功就別排新的,否則同一個帳號會有兩個 renew job 一起跑。
if err := s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplateThreadsTokenRenew, accountID); err != nil {
return err
}
j := &domain.Job{
ID: uuid.NewString(), OwnerUID: ownerUID,
TemplateType: domain.TemplateThreadsTokenRenew,
@ -325,6 +330,15 @@ type PersonaAnalyzeTextPayload struct {
Lang string `json:"lang,omitempty"`
}
// cancelPersonaAnalyzeJobs 收掉同一個人設還在排隊的兩種分析,兩者只能有一個在跑,
// 否則兩份結果會互相覆寫成不確定的人設。
func (s *Service) cancelPersonaAnalyzeJobs(ctx context.Context, ownerUID int64, personaID string) error {
if err := s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeAccount, personaID); err != nil {
return err
}
return s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeText, personaID)
}
// SchedulePersonaAnalyzeAccount — 立刻可領ref=personaID同人設舊佇列取消
func (s *Service) SchedulePersonaAnalyzeAccount(ctx context.Context, ownerUID int64, personaID, username, lang string) (*domain.Job, error) {
if ownerUID <= 0 || personaID == "" {
@ -334,8 +348,9 @@ func (s *Service) SchedulePersonaAnalyzeAccount(ctx context.Context, ownerUID in
if username == "" {
return nil, fmt.Errorf("empty username")
}
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeAccount, personaID)
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeText, personaID)
if err := s.cancelPersonaAnalyzeJobs(ctx, ownerUID, personaID); err != nil {
return nil, err
}
if s.MaxActiveAccountScrapes > 0 {
active, err := s.Repo.CountActiveByOwnerAndTemplate(ctx, ownerUID, domain.TemplatePersonaAnalyzeAccount)
if err != nil {
@ -421,7 +436,9 @@ func (s *Service) SchedulePlayGenerateScript(ctx context.Context, ownerUID int64
return nil, domain.ErrForbidden
}
playID = strings.TrimSpace(playID)
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePlayGenerateScript, playID)
if err := s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePlayGenerateScript, playID); err != nil {
return nil, err
}
body, _ := json.Marshal(PlayGenerateScriptPayload{PlayID: playID, OnlyEmpty: onlyEmpty})
now := domain.NowNano()
j := &domain.Job{
@ -476,8 +493,9 @@ func (s *Service) SchedulePersonaAnalyzeText(ctx context.Context, ownerUID int64
if rawText == "" {
return nil, fmt.Errorf("empty text")
}
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeAccount, personaID)
_ = s.Repo.CancelPendingByRef(ctx, ownerUID, domain.TemplatePersonaAnalyzeText, personaID)
if err := s.cancelPersonaAnalyzeJobs(ctx, ownerUID, personaID); err != nil {
return nil, err
}
body, _ := json.Marshal(PersonaAnalyzeTextPayload{
RawText: rawText, SourceLabel: strings.TrimSpace(sourceLabel), Lang: lang,
})

View File

@ -2,6 +2,8 @@ package usecase_test
import (
"context"
"errors"
"fmt"
"testing"
"time"
@ -76,71 +78,29 @@ func newStudio() (*usecase.Service, *publish.FakeTransport, *memAccounts) {
return svc, tp, acc
}
func withUsage(svc *usecase.Service, uid int64) {
// unlimited + platform key
_ = svc.Usage.Repo.SavePrefs(context.Background(), &usageDomain.MemberPrefs{
UID: uid, PlanID: usageDomain.PlanPro, Unlimited: true, UpdatedAt: domain.NowNano(),
})
if sr, ok := svc.Usage.Resolver.(*usageUC.StaticResolver); ok {
if sr.Map == nil {
sr.Map = map[string]string{}
}
sr.Map[fmtUIDMeter(uid, usageDomain.MeterAICopy)] = usageDomain.KeyModePlatform
}
}
func fmtUIDMeter(uid int64, meter string) string {
return string(rune(0)) // placeholder — use fmt
}
func setupUID(svc *usecase.Service, uid int64) {
setupUIDWithPlan(svc, uid, usageDomain.PlanPro, true)
}
func setupUIDWithPlan(svc *usecase.Service, uid int64, planID string, unlimited bool) {
_ = svc.Usage.Repo.SavePrefs(context.Background(), &usageDomain.MemberPrefs{
UID: uid, PlanID: usageDomain.PlanPro, Unlimited: true, UpdatedAt: domain.NowNano(),
UID: uid, PlanID: planID, Unlimited: unlimited, UpdatedAt: domain.NowNano(),
})
if sr, ok := svc.Usage.Resolver.(*usageUC.StaticResolver); ok {
if sr.Map == nil {
sr.Map = map[string]string{}
}
key := ""
// StaticResolver uses fmt.Sprintf("%d:%s", uid, meter)
key = sprintf("%d:%s", uid, usageDomain.MeterAICopy)
sr.Map[key] = usageDomain.KeyModePlatform
// StaticResolver keys on "uid:meter".
sr.Map[fmt.Sprintf("%d:%s", uid, usageDomain.MeterAICopy)] = usageDomain.KeyModePlatform
}
}
func sprintf(f string, a ...any) string {
return format(f, a...)
}
// avoid importing fmt in helpers clutter — use strconv
func format(f string, a ...any) string {
// minimal for our key only
if f == "%d:%s" && len(a) == 2 {
return itoa(a[0].(int64)) + ":" + a[1].(string)
}
return f
}
func itoa(n int64) string {
if n == 0 {
return "0"
}
neg := n < 0
if neg {
n = -n
}
var b [32]byte
i := len(b)
for n > 0 {
i--
b[i] = byte('0' + n%10)
n /= 10
}
if neg {
i--
b[i] = '-'
}
return string(b[i:])
// platformCreditsUsed reports what the member has actually been charged this month.
func platformCreditsUsed(t *testing.T, svc *usecase.Service, uid int64) int {
t.Helper()
sum, err := svc.Usage.GetSummary(context.Background(), uid, "")
require.NoError(t, err)
return sum.Platform.CreditsUsed
}
func addAccount(acc *memAccounts, ownerUID int64, id, user string) {
@ -150,6 +110,86 @@ func addAccount(acc *memAccounts, ownerUID int64, id, user string) {
}
}
// ---- AI billing ----
// A request the service itself rejects must not consume credits. These used to charge the
// member up front, so every early return below was a silent debit for work never done.
func TestBilling_RejectedRequestIsNotCharged(t *testing.T) {
uid := int64(4_009_001)
cases := []struct {
name string
call func(svc *usecase.Service) error
}{
{
name: "mimic without persona",
call: func(svc *usecase.Service) error {
_, err := svc.Mimic(context.Background(), uid, "來源文字", "", "", "")
return err
},
},
{
name: "mimic with unknown persona",
call: func(svc *usecase.Service) error {
_, err := svc.Mimic(context.Background(), uid, "來源文字", "does-not-exist", "", "")
return err
},
},
{
name: "analyze post that has no text",
call: func(svc *usecase.Service) error {
post := &domain.OwnPost{
ID: "post-no-text", OwnerUID: uid, Text: "", PublishedAt: domain.NowNano(),
}
require.NoError(t, svc.Repo.SaveOwnPost(context.Background(), post))
_, err := svc.AnalyzePost(context.Background(), uid, post.ID)
return err
},
},
{
name: "generate play step without context",
call: func(svc *usecase.Service) error {
_, err := svc.GeneratePlayStep(context.Background(), uid, "", "", "", "", true, "root")
return err
},
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
svc, _, _ := newStudio()
// A metered plan, so any charge is visible.
setupUIDWithPlan(svc, uid, usageDomain.PlanPro, false)
require.Error(t, tc.call(svc))
require.Zero(t, platformCreditsUsed(t, svc, uid), "rejected request must not be charged")
})
}
}
// A persona analysis that cannot gather enough samples fails after billing has been reserved.
func TestBilling_FailedPersonaAnalysisIsRefunded(t *testing.T) {
svc, _, _ := newStudio()
uid := int64(4_009_002)
setupUIDWithPlan(svc, uid, usageDomain.PlanPro, false)
p, err := svc.SavePersona(context.Background(), uid, &domain.Persona{Name: "小海"})
require.NoError(t, err)
// One sample is below the two-sample minimum, so the analysis is rejected.
_, err = svc.ExecuteAnalyzeFromText(context.Background(), uid, p.ID, "只有一段文字", "手動", nil)
require.Error(t, err)
require.Zero(t, platformCreditsUsed(t, svc, uid))
}
// The success path must still charge exactly once.
func TestBilling_SuccessfulCallIsChargedOnce(t *testing.T) {
svc, _, _ := newStudio()
uid := int64(4_009_003)
setupUIDWithPlan(svc, uid, usageDomain.PlanPro, false)
_, err := svc.AnalyzeViral(context.Background(), uid, "這是一段要分析的貼文內容,講一個轉折。")
require.NoError(t, err)
require.Equal(t, usageDomain.MeterCost(usageDomain.MeterAICopy), platformCreditsUsed(t, svc, uid))
}
// ---- PE ----
func TestPE_01_CreateSave(t *testing.T) {
@ -915,6 +955,64 @@ func TestOP_11_GenerateFromFormulaRemoved(t *testing.T) {
require.ErrorIs(t, err, domain.ErrFormulaRemove)
}
// fakeInsightsMedia serves a fixed post list and lets a test make the insights call fail, which
// is what a Threads rate limit looks like mid-sync.
type fakeInsightsMedia struct {
threads []usecase.FetchedThread
insights usecase.FetchedInsights
insightsErr error
}
func (f *fakeInsightsMedia) ListThreads(context.Context, string, int) ([]usecase.FetchedThread, error) {
return f.threads, nil
}
func (f *fakeInsightsMedia) GetInsights(context.Context, string, string) (usecase.FetchedInsights, error) {
if f.insightsErr != nil {
return usecase.FetchedInsights{}, f.insightsErr
}
return f.insights, nil
}
func (f *fakeInsightsMedia) ListConversation(context.Context, string, string, int) ([]usecase.FetchedReply, error) {
return nil, nil
}
func (f *fakeInsightsMedia) ListMentions(context.Context, string, string, int) ([]usecase.FetchedMention, error) {
return nil, nil
}
func (f *fakeInsightsMedia) ListProfilePosts(context.Context, string, string, int) ([]usecase.FetchedThread, error) {
return nil, nil
}
// A single rate-limited insights call used to overwrite a post's real like/view counts with
// zeros, which then corrupted the insights summary, benchmark and account health downstream.
func TestSyncOwnPosts_FailedInsightsKeepsPreviousCounts(t *testing.T) {
svc, _, acc := newStudio()
uid := int64(4_004_020)
setupUID(svc, uid)
addAccount(acc, uid, "acc1", "me")
media := &fakeInsightsMedia{
threads: []usecase.FetchedThread{{ID: "m_1", Text: "貼文", PublishedAt: domain.NowNano()}},
insights: usecase.FetchedInsights{Views: 500, Likes: 42, Replies: 7, Status: "ok"},
}
svc.Media = media
first, err := svc.SyncOwnPosts(context.Background(), uid, "acc1")
require.NoError(t, err)
require.Len(t, first, 1)
require.Equal(t, 42, first[0].LikeCount)
require.Equal(t, 500, first[0].ViewCount)
// Threads starts rate limiting; the next sync must not zero the numbers out.
media.insightsErr = errors.New("rate limited")
second, err := svc.SyncOwnPosts(context.Background(), uid, "acc1")
require.NoError(t, err)
require.Len(t, second, 1)
require.Equal(t, 42, second[0].LikeCount, "a failed insights fetch must not erase real counts")
require.Equal(t, 500, second[0].ViewCount)
require.Equal(t, 7, second[0].ReplyCount)
require.Equal(t, "error", second[0].InsightsStatus, "but the failure must still be visible")
}
type fakeMentionsMedia struct {
hits []usecase.FetchedMention
err error

View File

@ -16,7 +16,7 @@ import (
)
// PersonaPreview依指紋一次 LLM 產主貼 + 回文(可抓新聞當話題靈感)。
func (s *Service) PersonaPreview(ctx context.Context, ownerUID int64, personaID, topic string, useNews bool) (*domain.PersonaPreview, error) {
func (s *Service) PersonaPreview(ctx context.Context, ownerUID int64, personaID, topic string, useNews bool) (_ *domain.PersonaPreview, err error) {
if ownerUID <= 0 {
return nil, domain.ErrForbidden
}
@ -36,9 +36,11 @@ func (s *Service) PersonaPreview(ctx context.Context, ownerUID int64, personaID,
return nil, fmt.Errorf("%w: 人設尚無指紋,請先完成分析或手動填寫指紋", domain.ErrValidation)
}
if err := s.billAI(ctx, ownerUID, "persona preview", "compose.personaPreview"); err != nil {
charge, err := s.billAI(ctx, ownerUID, "persona preview", "compose.personaPreview")
if err != nil {
return nil, err
}
defer charge.Settle(ctx, &err)
topic = strings.TrimSpace(topic)
topicSource := "manual"

View File

@ -279,7 +279,7 @@ func (s *Service) AnalyzeFromAccount(ctx context.Context, ownerUID int64, id, us
}
// ExecuteAnalyzeFromText — worker測試扣費 + LLM 分析 + 存檔 ready。
func (s *Service) ExecuteAnalyzeFromText(ctx context.Context, ownerUID int64, id, rawText, sourceLabel string, onProgress ProgressFn) (*domain.Persona, error) {
func (s *Service) ExecuteAnalyzeFromText(ctx context.Context, ownerUID int64, id, rawText, sourceLabel string, onProgress ProgressFn) (_ *domain.Persona, err error) {
report := func(pct int, sum string) {
if onProgress != nil {
onProgress(pct, sum)
@ -299,10 +299,12 @@ func (s *Service) ExecuteAnalyzeFromText(ctx context.Context, ownerUID int64, id
return nil, fmt.Errorf("%w: 請至少貼 2 段參考文字(可用 --- 分隔),每段至少約 10 字", domain.ErrValidation)
}
report(25, "人設分析 · 檢查文字樣本…")
if err := s.billAI(ctx, ownerUID, "persona analyze text", "personas.analyzeFromText"); err != nil {
charge, err := s.billAI(ctx, ownerUID, "persona analyze text", "personas.analyzeFromText")
if err != nil {
_ = s.markPersonaAnalyzeFailed(ctx, p, err.Error())
return nil, err
}
defer charge.Settle(ctx, &err)
p.Status = domain.PersonaAnalyzing
p.UpdatedAt = domain.NowNano()
_ = s.Repo.SavePersona(ctx, p)
@ -324,7 +326,7 @@ func (s *Service) ExecuteAnalyzeFromText(ctx context.Context, ownerUID int64, id
}
// ExecuteAnalyzeFromAccount — worker爬公開貼文 + LLM + 存檔 ready。
func (s *Service) ExecuteAnalyzeFromAccount(ctx context.Context, ownerUID int64, id, username string, onProgress ProgressFn) (*domain.Persona, error) {
func (s *Service) ExecuteAnalyzeFromAccount(ctx context.Context, ownerUID int64, id, username string, onProgress ProgressFn) (_ *domain.Persona, err error) {
report := func(pct int, sum string) {
if onProgress != nil {
onProgress(pct, sum)
@ -342,10 +344,12 @@ func (s *Service) ExecuteAnalyzeFromAccount(ctx context.Context, ownerUID int64,
return nil, fmt.Errorf("%w: username 請只填帳號,例如 ultralab_tw", domain.ErrValidation)
}
report(15, fmt.Sprintf("人設分析 · 準備爬取 @%s…", username))
if err := s.billAI(ctx, ownerUID, "persona analyze account", "personas.analyzeFromAccount"); err != nil {
charge, err := s.billAI(ctx, ownerUID, "persona analyze account", "personas.analyzeFromAccount")
if err != nil {
_ = s.markPersonaAnalyzeFailed(ctx, p, err.Error())
return nil, err
}
defer charge.Settle(ctx, &err)
p.Status = domain.PersonaAnalyzing
p.Style.BenchmarkUsername = username
@ -807,7 +811,7 @@ func (s *Service) SubmitPlay(ctx context.Context, ownerUID int64, playID string)
// GeneratePlayScript — 一次 LLM 產完整劇本並寫回 play.stepsonlyEmpty=true 只填空白步)。
// 回傳填入步數。供 worker job / 同步測試。
func (s *Service) GeneratePlayScript(ctx context.Context, ownerUID int64, playID string, onlyEmpty bool) (int, error) {
func (s *Service) GeneratePlayScript(ctx context.Context, ownerUID int64, playID string, onlyEmpty bool) (_ int, err error) {
play, err := s.GetPlay(ctx, ownerUID, playID)
if err != nil {
return 0, err
@ -841,9 +845,11 @@ func (s *Service) GeneratePlayScript(ctx context.Context, ownerUID int64, playID
if len(needs) == 0 {
return 0, fmt.Errorf("%w: 沒有空白步驟可產(全部已有正文)", domain.ErrValidation)
}
if err := s.billAI(ctx, ownerUID, "play generate script", "plays.generateScript"); err != nil {
charge, err := s.billAI(ctx, ownerUID, "play generate script", "plays.generateScript")
if err != nil {
return 0, err
}
defer charge.Settle(ctx, &err)
// 主上下文:目標貼文/主題
targetCtx := strings.TrimSpace(play.Topic)
@ -1017,15 +1023,17 @@ func cleanPlayGeneratedText(kind, text string) string {
// GeneratePlayStep — 互回/串場劇本:依人設真 LLM 產一步正文。
// 注意:與 compose mimic 相同OpenCode reasoning 模型可能 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)
topic = strings.TrimSpace(topic)
if contextText == "" && topic == "" {
return "", fmt.Errorf("%w: empty context", domain.ErrValidation)
}
if err := s.billAI(ctx, ownerUID, "play generate step", "plays.generateStep"); err != nil {
charge, err := s.billAI(ctx, ownerUID, "play generate step", "plays.generateStep")
if err != nil {
return "", err
}
defer charge.Settle(ctx, &err)
mode = strings.ToLower(strings.TrimSpace(mode))
if mode == "" {
mode = "reply"
@ -1377,14 +1385,11 @@ func (s *Service) publishWithLease(ctx context.Context, bundleID, stepID, leaseO
// ---------- Compose (CP) ----------
func (s *Service) Mimic(ctx context.Context, ownerUID int64, sourceText, personaID, direction, structureNotes string) (string, error) {
func (s *Service) Mimic(ctx context.Context, ownerUID int64, sourceText, personaID, direction, structureNotes string) (_ string, err error) {
sourceText = strings.TrimSpace(sourceText)
if sourceText == "" {
return "", fmt.Errorf("%w: empty source", domain.ErrValidation)
}
if err := s.billAI(ctx, ownerUID, "compose mimic", "compose.mimic"); err != nil {
return "", err
}
personaID = strings.TrimSpace(personaID)
if personaID == "" {
return "", fmt.Errorf("%w: 請選擇已完成人設分析的人設", domain.ErrValidation)
@ -1396,6 +1401,11 @@ func (s *Service) Mimic(ctx context.Context, ownerUID int64, sourceText, persona
if p.Status != domain.PersonaReady {
return "", fmt.Errorf("%w: 選擇的人設尚未完成分析", domain.ErrValidation)
}
charge, err := s.billAI(ctx, ownerUID, "compose mimic", "compose.mimic")
if err != nil {
return "", err
}
defer charge.Settle(ctx, &err)
fp := personaExpressionFingerprintBlock(p)
// 指紋過長會拖慢 reasoning 模型
if utf8.RuneCountInString(fp) > 800 {
@ -1490,10 +1500,12 @@ func buildMimicPrompt(fp, sourceText, direction, notesBlock string) string {
`, direction, fp, sourceText, notesBlock))
}
func (s *Service) AnalyzeViral(ctx context.Context, ownerUID int64, text string) (*domain.ViralAnalysis, error) {
if err := s.billAI(ctx, ownerUID, "compose analyze viral", "compose.analyzeViral"); err != nil {
func (s *Service) AnalyzeViral(ctx context.Context, ownerUID int64, text string) (_ *domain.ViralAnalysis, err error) {
charge, err := s.billAI(ctx, ownerUID, "compose analyze viral", "compose.analyzeViral")
if err != nil {
return nil, err
}
defer charge.Settle(ctx, &err)
return s.analyzeViralUnbilled(ctx, ownerUID, text)
}
@ -1889,6 +1901,9 @@ func (s *Service) syncOwnPostsFromThreads(ctx context.Context, ownerUID int64, a
}
insights := make([]FetchedInsights, len(threads))
// Each slot starts as "not fetched" so a rate-limited call, a timeout, or a cancelled sync is
// distinguishable from a post that genuinely has zero engagement.
insightFetched := make([]bool, len(threads))
var wg sync.WaitGroup
limit := make(chan struct{}, 4)
for i, th := range threads {
@ -1904,7 +1919,12 @@ func (s *Service) syncOwnPostsFromThreads(ctx context.Context, ownerUID int64, a
case <-ctx.Done():
return
}
insights[i], _ = s.Media.GetInsights(ctx, accessToken, mediaID)
got, err := s.Media.GetInsights(ctx, accessToken, mediaID)
if err != nil {
logx.Errorf("threads insights media=%s: %v", mediaID, err)
return
}
insights[i], insightFetched[i] = got, true
}(i, th.ID)
}
wg.Wait()
@ -1917,6 +1937,15 @@ func (s *Service) syncOwnPostsFromThreads(ctx context.Context, ownerUID int64, a
ins := insights[i]
prev := byMedia[th.ID]
// A failed insights call must keep whatever counts we already had. Writing zeros would
// destroy the real numbers, and they feed insights summary, benchmark and account health.
if !insightFetched[i] {
if prev != nil {
ins.Likes, ins.Replies, ins.Reposts = prev.LikeCount, prev.ReplyCount, prev.RepostCount
ins.Quotes, ins.Views, ins.Shares = prev.QuoteCount, prev.ViewCount, prev.ShareCount
}
ins.Status = "error"
}
id := "op_" + th.ID
if prev != nil && prev.ID != "" {
id = prev.ID
@ -2104,7 +2133,7 @@ func mergeReplyStatus(prev, next []domain.OwnPostReply) []domain.OwnPostReply {
return next
}
func (s *Service) GenerateReply(ctx context.Context, ownerUID int64, postID, replyID, personaID string) (string, error) {
func (s *Service) GenerateReply(ctx context.Context, ownerUID int64, postID, replyID, personaID string) (_ string, err error) {
post, err := s.getOwnPostOwned(ctx, ownerUID, postID)
if err != nil {
return "", err
@ -2127,9 +2156,11 @@ func (s *Service) GenerateReply(ctx context.Context, ownerUID int64, postID, rep
return "", fmt.Errorf("%w: reply not found", domain.ErrValidation)
}
}
if err := s.billAI(ctx, ownerUID, "own post reply draft", "ownPosts.generateReply"); err != nil {
charge, err := s.billAI(ctx, ownerUID, "own post reply draft", "ownPosts.generateReply")
if err != nil {
return "", err
}
defer charge.Settle(ctx, &err)
persona := s.loadPersonaForGen(ctx, ownerUID, personaID)
fp := personaExpressionFingerprintBlock(persona)
@ -2230,7 +2261,7 @@ func (s *Service) SendReply(ctx context.Context, ownerUID int64, postID, replyID
return post, nil
}
func (s *Service) AnalyzePost(ctx context.Context, ownerUID int64, postID string) (*domain.OwnPost, error) {
func (s *Service) AnalyzePost(ctx context.Context, ownerUID int64, postID string) (_ *domain.OwnPost, err error) {
post, err := s.getOwnPostOwned(ctx, ownerUID, postID)
if err != nil {
return nil, err
@ -2238,9 +2269,11 @@ func (s *Service) AnalyzePost(ctx context.Context, ownerUID int64, postID string
if strings.TrimSpace(post.Text) == "" {
return nil, fmt.Errorf("%w: 貼文沒有文字可分析(圖片/純媒體貼可改貼上說明再分析)", domain.ErrValidation)
}
if err := s.billAI(ctx, ownerUID, "own post structure analyze", "ownPosts.analyze"); err != nil {
charge, err := s.billAI(ctx, ownerUID, "own post structure analyze", "ownPosts.analyze")
if err != nil {
return nil, err
}
defer charge.Settle(ctx, &err)
va, err := s.analyzeViralUnbilled(ctx, ownerUID, post.Text)
if err != nil {
return nil, err
@ -2371,14 +2404,16 @@ func mentionContextSnippet(hit FetchedMention) string {
return kind + " · " + text
}
func (s *Service) GenerateMentionReply(ctx context.Context, ownerUID int64, id, personaID string) (*domain.Mention, error) {
func (s *Service) GenerateMentionReply(ctx context.Context, ownerUID int64, id, personaID string) (_ *domain.Mention, err error) {
m, err := s.getMentionOwned(ctx, ownerUID, id)
if err != nil {
return nil, err
}
if err := s.billAI(ctx, ownerUID, "mention reply draft", "mentions.generateReply"); err != nil {
charge, err := s.billAI(ctx, ownerUID, "mention reply draft", "mentions.generateReply")
if err != nil {
return nil, err
}
defer charge.Settle(ctx, &err)
persona := s.loadPersonaForGen(ctx, ownerUID, personaID)
fp := personaExpressionFingerprintBlock(persona)
prompt := buildMentionReplyPrompt(fp, m.FromUsername, m.Text, m.ContextSnippet)
@ -2652,16 +2687,67 @@ func (s *Service) SeedMention(ctx context.Context, m *domain.Mention) error {
// ---------- helpers ----------
func (s *Service) billAI(ctx context.Context, ownerUID int64, label, source string) error {
if s.Usage == nil {
return nil
// aiCharge is one reserved AI credit. Quota is reserved up front so an over-limit member is
// stopped before any work starts, but the audit event is only written once the work succeeded —
// the monthly counter is rebuilt from events if it is ever lost, so an event for a call that
// never happened would silently undo the refund.
type aiCharge struct {
svc *Service
ownerUID int64
label string
source string
mode string
settled bool
}
// billAI reserves credit for one AI call. Pair it with `defer charge.Settle(ctx, &err)` on a
// named error return so that every failure path — including ones added later — refunds.
func (s *Service) billAI(ctx context.Context, ownerUID int64, label, source string) (*aiCharge, error) {
if s == nil || s.Usage == nil {
return &aiCharge{settled: true}, nil
}
mode, err := s.Usage.PrepareCall(ctx, ownerUID, usageDomain.MeterAICopy)
if err != nil {
return err
return nil, err
}
return &aiCharge{svc: s, ownerUID: ownerUID, label: label, source: source, mode: mode}, nil
}
// Settle commits the charge when the operation succeeded, or refunds it when it failed.
func (c *aiCharge) Settle(ctx context.Context, errp *error) {
if errp != nil && *errp != nil {
c.Release(ctx)
return
}
c.Commit(ctx)
}
// Commit writes the audit event. A failure here is logged rather than surfaced: the member
// already received the AI result, so charging them is correct and losing the event must not
// turn a successful call into an error.
func (c *aiCharge) Commit(ctx context.Context) {
if c == nil || c.settled {
return
}
c.settled = true
if _, err := c.svc.Usage.RecordCall(ctx, c.ownerUID, usageDomain.MeterAICopy, c.mode, c.label, c.source); err != nil {
logx.Errorf("usage record uid=%d source=%s mode=%s: %v", c.ownerUID, c.source, c.mode, err)
}
}
// Release refunds a charge that was never committed. Safe to call more than once.
func (c *aiCharge) Release(ctx context.Context) {
if c == nil || c.settled {
return
}
c.settled = true
// The failure being compensated for is often a cancelled request, so the refund needs a
// context that outlives it.
rctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
defer cancel()
if err := c.svc.Usage.ReleaseCall(rctx, c.ownerUID, usageDomain.MeterAICopy, c.mode); err != nil {
logx.Errorf("usage release uid=%d source=%s mode=%s: %v", c.ownerUID, c.source, c.mode, err)
}
_, err = s.Usage.RecordCall(ctx, ownerUID, usageDomain.MeterAICopy, mode, label, source)
return err
}
func (s *Service) getOwnPostOwned(ctx context.Context, ownerUID int64, id string) (*domain.OwnPost, error) {

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 where the browser lands after the Threads OAuth round trip.
// An unset WebBase yields a relative path on purpose: falling back to a hardcoded loopback
// address would silently strand every production user on their own machine.
func (s *Service) WebRedirect(ok bool, msg string) string {
base := trimSlash(s.WebBase)
if base == "" {
base = "http://127.0.0.1:5173"
}
if ok {
return base + "/app/crew?oauth=ok"
}

View File

@ -6,6 +6,10 @@ type Repository interface {
// ReservePlatform atomically guards both limits and increments one monthly counter document.
// A negative limit means unlimited. BYOK must never call this operation.
ReservePlatform(ctx context.Context, uid int64, monthKey, meter string, cost, totalLimit, meterLimit int) error
// ReleasePlatform gives back a reservation whose call never happened. It is guarded so a
// counter can never drop below zero, and is a no-op when there is nothing left to refund
// (e.g. the month already rolled over), so callers may retry it safely.
ReleasePlatform(ctx context.Context, uid int64, monthKey, meter string, cost int) error
GetMonthlyCounter(ctx context.Context, uid int64, monthKey string) (*MonthlyCounter, error)
InsertEvent(ctx context.Context, e *Event) error

View File

@ -58,6 +58,25 @@ func (s *MemoryStore) ReservePlatform(_ context.Context, uid int64, monthKey, me
return nil
}
func (s *MemoryStore) ReleasePlatform(_ context.Context, uid int64, monthKey, meter string, cost int) error {
if cost <= 0 {
return nil
}
s.mu.Lock()
defer s.mu.Unlock()
counter := s.counters[memoryCounterID(uid, monthKey)]
if counter == nil {
return nil
}
if counter.TotalCredits < cost || counter.ByMeter[meter] < cost {
return nil
}
counter.TotalCredits -= cost
counter.ByMeter[meter] -= cost
counter.UpdatedAt = domain.NowNano()
return nil
}
func (s *MemoryStore) GetMonthlyCounter(_ context.Context, uid int64, monthKey string) (*domain.MonthlyCounter, error) {
s.mu.Lock()
defer s.mu.Unlock()

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")
}
func (s *MonStore) ReleasePlatform(ctx context.Context, uid int64, monthKey, meter string, cost int) error {
field, err := meterCounterField(meter)
if err != nil {
return err
}
if cost <= 0 {
return nil
}
// Both guards keep the refund from pushing either counter negative; a non-match means
// there is nothing to give back, which is not an error.
filter := bson.M{
"_id": monthlyCounterID(uid, monthKey),
"total_credits": bson.M{"$gte": cost},
field: bson.M{"$gte": cost},
}
_, err = s.counters.UpdateOne(ctx, filter, bson.M{
"$inc": bson.M{"total_credits": -cost, field: -cost},
"$set": bson.M{"updated_at": domain.NowNano()},
})
return err
}
func (s *MonStore) GetMonthlyCounter(ctx context.Context, uid int64, monthKey string) (*domain.MonthlyCounter, error) {
var counter domain.MonthlyCounter
err := s.counters.FindOne(ctx, &counter, bson.M{"_id": monthlyCounterID(uid, monthKey)})

View File

@ -81,17 +81,39 @@ func (s *Service) PrepareCall(ctx context.Context, uid int64, meter string) (key
if !has {
return "", domain.ErrNoKey
}
if mode == domain.KeyModePlatform {
if s.Gate != nil && !s.Gate.AllowPlatform(meter) {
return "", domain.ErrPlatformCapacity
}
if err := s.reservePlatform(ctx, uid, meter); err != nil {
return "", err
}
if err := s.PrepareCallWithMode(ctx, uid, meter, mode); err != nil {
return "", err
}
return mode, nil
}
// PrepareCallWithMode reserves quota for a call whose key mode the caller already resolved.
// Callers that also need the key material must use this rather than PrepareCall: resolving
// twice reads member settings twice, and if those reads disagree the credits get reserved
// against one mode while the call is billed as the other.
func (s *Service) PrepareCallWithMode(ctx context.Context, uid int64, meter, keyMode string) error {
if keyMode != domain.KeyModePlatform && keyMode != domain.KeyModeByok {
return fmt.Errorf("invalid key_mode")
}
if keyMode != domain.KeyModePlatform {
return nil
}
if s.Gate != nil && !s.Gate.AllowPlatform(meter) {
return domain.ErrPlatformCapacity
}
return s.reservePlatform(ctx, uid, meter)
}
// ReleaseCall gives back credits reserved by PrepareCall when the call it was paying for
// never completed. BYOK reserves nothing, so it is a no-op there. Callers should reach for
// this on every failure path between PrepareCall and RecordCall.
func (s *Service) ReleaseCall(ctx context.Context, uid int64, meter, keyMode string) error {
if keyMode != domain.KeyModePlatform {
return nil
}
return s.Repo.ReleasePlatform(ctx, uid, domain.CurrentMonthKey(), meter, domain.MeterCost(meter))
}
// MarkPlatformLimited cools down platform keys (e.g. after upstream 429).
// BYOK users are unaffected. until zero clears the cool-down.
func (s *Service) MarkPlatformLimited(until time.Time) {

View File

@ -295,6 +295,66 @@ func TestUSG_14_MemberCannotSetOthersPrefs(t *testing.T) {
require.ErrorIs(t, err, domain.ErrForbidden)
}
func TestReleaseCall_RefundsReservedPlatformCredit(t *testing.T) {
repo := repository.NewMemory()
uid := int64(1_000_010)
svc := usecase.New(repo, &usecase.StaticResolver{Map: map[string]string{
fmt.Sprintf("%d:%s", uid, domain.MeterAICopy): domain.KeyModePlatform,
}})
ctx := context.Background()
mode, err := svc.PrepareCall(ctx, uid, domain.MeterAICopy)
require.NoError(t, err)
counter, err := repo.GetMonthlyCounter(ctx, uid, domain.CurrentMonthKey())
require.NoError(t, err)
require.Equal(t, domain.MeterCost(domain.MeterAICopy), counter.TotalCredits)
require.NoError(t, svc.ReleaseCall(ctx, uid, domain.MeterAICopy, mode))
counter, err = repo.GetMonthlyCounter(ctx, uid, domain.CurrentMonthKey())
require.NoError(t, err)
require.Zero(t, counter.TotalCredits, "a call that never happened must not consume credits")
require.Zero(t, counter.ByMeter[domain.MeterAICopy])
// The member must still see the full quota available afterwards.
sum, err := svc.GetSummary(ctx, uid, "")
require.NoError(t, err)
require.Zero(t, sum.Platform.CreditsUsed)
}
func TestReleaseCall_NeverDrivesCounterNegative(t *testing.T) {
repo := repository.NewMemory()
uid := int64(1_000_011)
svc := usecase.New(repo, &usecase.StaticResolver{Map: map[string]string{
fmt.Sprintf("%d:%s", uid, domain.MeterAICopy): domain.KeyModePlatform,
}})
ctx := context.Background()
mode, err := svc.PrepareCall(ctx, uid, domain.MeterAICopy)
require.NoError(t, err)
// A retried or duplicated refund must stay idempotent.
for i := 0; i < 3; i++ {
require.NoError(t, svc.ReleaseCall(ctx, uid, domain.MeterAICopy, mode))
}
counter, err := repo.GetMonthlyCounter(ctx, uid, domain.CurrentMonthKey())
require.NoError(t, err)
require.Zero(t, counter.TotalCredits)
require.GreaterOrEqual(t, counter.ByMeter[domain.MeterAICopy], 0)
}
func TestReleaseCall_ByokIsNoOp(t *testing.T) {
repo := repository.NewMemory()
uid := int64(1_000_012)
svc := usecase.New(repo, &usecase.StaticResolver{Map: map[string]string{
fmt.Sprintf("%d:%s", uid, domain.MeterAICopy): domain.KeyModeByok,
}})
ctx := context.Background()
mode, err := svc.PrepareCall(ctx, uid, domain.MeterAICopy)
require.NoError(t, err)
require.NoError(t, svc.ReleaseCall(ctx, uid, domain.MeterAICopy, mode))
_, err = repo.GetMonthlyCounter(ctx, uid, domain.CurrentMonthKey())
require.ErrorIs(t, err, domain.ErrNotFound, "BYOK reserves nothing, so a refund must not create a counter")
}
func TestCreditInviteBonus_AddsToPool(t *testing.T) {
svc := newUsage(42, domain.KeyModePlatform)
require.NoError(t, svc.CreditInviteBonus(context.Background(), 42, 100, "r1"))

View File

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

View File

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