thread-master/apps/backend/cmd/worker/main.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"apps/backend/internal/config"
	redislock "apps/backend/internal/lib/redislock"
	"apps/backend/internal/lib/securelog"
	"apps/backend/internal/module/ai"
	appnotifRepo "apps/backend/internal/module/appnotif/repository"
	appnotifUC "apps/backend/internal/module/appnotif/usecase"
	fsDomain "apps/backend/internal/module/filestorage/domain"
	"apps/backend/internal/module/filestorage/noop"
	"apps/backend/internal/module/filestorage/s3store"
	jobDomain "apps/backend/internal/module/job/domain"
	jobRepo "apps/backend/internal/module/job/repository"
	jobUC "apps/backend/internal/module/job/usecase"
	memberDomain "apps/backend/internal/module/member/domain"
	memberRepo "apps/backend/internal/module/member/repository"
	scoutDomain "apps/backend/internal/module/scout/domain"
	scoutRepo "apps/backend/internal/module/scout/repository"
	scoutUC "apps/backend/internal/module/scout/usecase"
	studioPublish "apps/backend/internal/module/studio/publish"
	studioRepo "apps/backend/internal/module/studio/repository"
	studioUC "apps/backend/internal/module/studio/usecase"
	threadsDomain "apps/backend/internal/module/threads/domain"
	threadsProv "apps/backend/internal/module/threads/provider"
	threadsRepo "apps/backend/internal/module/threads/repository"
	threadsUC "apps/backend/internal/module/threads/usecase"
	usageDomain "apps/backend/internal/module/usage/domain"
	usageRepo "apps/backend/internal/module/usage/repository"
	usageUC "apps/backend/internal/module/usage/usecase"

	"github.com/zeromicro/go-zero/core/conf"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/core/stores/redis"
)

// Worker：demo + threads_token_renew + persona_analyze_* + outbox publish
//
//	cd apps/backend && go build -o bin/worker ./cmd/worker
//	./bin/worker -f etc/gateway.yaml
func main() {
	securelog.SilenceMongo()

	configFile := flag.String("f", "etc/gateway.yaml", "config file")
	flag.Parse()

	var c config.Config
	conf.MustLoad(*configFile, &c)
	c.ApplyEnv()
	if err := c.ValidateProductionDependencies(); err != nil {
		logx.Must(err)
	}
	redisClient := redis.MustNewRedis(c.CacheRedis[0].RedisConf)
	pingCtx, cancelPing := context.WithTimeout(context.Background(), 3*time.Second)
	redisReady := redisClient.PingCtx(pingCtx)
	cancelPing()
	if !redisReady {
		logx.Must(fmt.Errorf("redis startup ping failed"))
	}

	workerID := c.Worker.ID
	if workerID == "" || workerID == "worker-1" || workerID == "worker-local-1" || workerID == "demo-worker-1" {
		host, _ := os.Hostname()
		if host == "" {
			host = "worker"
		}
		workerID = fmt.Sprintf("%s-%d", host, os.Getpid())
	}
	interval := time.Duration(c.Worker.PollIntervalMs) * time.Millisecond
	if interval <= 0 {
		interval = 2 * time.Second
	}

	appN := appnotifUC.New(appnotifRepo.NewMonStore(c.Mongo.URI, c.Mongo.Database))
	jobs := jobUC.New(jobRepo.NewMonStore(c.Mongo.URI, c.Mongo.Database))
	jobs.Notifier = appN

	var provider threadsDomain.Provider
	if strings.TrimSpace(c.Platform.ThreadsAppId) != "" && strings.TrimSpace(c.Platform.ThreadsAppSecret) != "" {
		provider = threadsProv.NewMeta(c.Platform.ThreadsAppId, c.Platform.ThreadsAppSecret)
		logx.Info("worker threads: meta provider")
	} else {
		provider = threadsProv.NewUnavailable("set THREADS_APP_ID and THREADS_APP_SECRET")
		logx.Info("worker threads: Threads operations disabled until credentials are configured")
	}
	threadsSvc := threadsUC.New(threadsRepo.NewMonStore(c.Mongo.URI, c.Mongo.Database), provider, c.Auth.AccessSecret)
	threadsSvc.Renew = jobs

	// usage + AI keys（與 gateway 對齊，persona LLM 需真 key）
	settingsCodec, err := c.NewMemberSettingsCodec()
	if err != nil {
		logx.Must(err)
	}
	members := memberRepo.NewMoncStore(c.Mongo.URI, c.Mongo.Database, c.CacheRedis, c.CacheRedis[0].RedisConf, c.Redis.Namespace, settingsCodec)
	keyRes := &usageUC.SettingsResolver{
		Members:          members,
		PlatformAI:       c.Platform.AIKey,
		PlatformXAI:      c.Platform.XAIKey,
		PlatformOpenCode: c.Platform.OpenCodeKey,
		PlatformExa:      c.Platform.ExaKey,
	}
	usageSvc := usageUC.New(usageRepo.NewMonStore(c.Mongo.URI, c.Mongo.Database), keyRes)
	usageSvc.Gate = usageUC.NewRedisPlatformGate(c.CacheRedis[0].RedisConf, 120, 60, c.Redis.Namespace)
	modelsCache := ai.NewRedisModelsCache(c.CacheRedis[0].RedisConf, c.Redis.Namespace)
	aiRegistry := ai.NewRegistry(modelsCache, c.Redis.AIModelCacheFingerprintSecret)

	var pub studioPublish.Transport
	if strings.TrimSpace(c.Platform.ThreadsAppId) != "" && strings.TrimSpace(c.Platform.ThreadsAppSecret) != "" {
		pub = studioPublish.NewMeta("https://graph.threads.net")
		logx.Info("worker studio: meta publish transport (real Threads)")
	} else {
		pub = studioPublish.NewUnavailable("set THREADS_APP_ID and THREADS_APP_SECRET")
	}
	studio := studioUC.New(studioRepo.NewMonStore(c.Mongo.URI, c.Mongo.Database), pub)
	studio.OutboxWorkerID = workerID
	studio.Accounts = threadsSvc
	studio.Usage = usageSvc
	studio.AI = &ai.FakeClient{}
	studio.AIRegistry = aiRegistry
	studio.Keys = &workerAIKeys{Members: members, Resolver: keyRes}
	// 發文暫存圖：Meta 發成功後刪 temp/*
	store := workerStorage(c)
	studio.Storage = store
	studio.StoragePublicBase = strings.TrimSpace(c.ObjectStorage.PublicBaseURL)

	scoutSvc := scoutUC.New(scoutRepo.NewMonStore(c.Mongo.URI, c.Mongo.Database))
	scoutSvc.Settings = &workerDevMode{Members: members}
	scoutSvc.Provider = scoutUC.NewExaThreadsProvider(c.Platform.ExaKey)
	scoutSvc.SessionSecret = c.Scout.SessionSecret
	scoutSvc.Crawler = scoutUC.NewHTTPCrawlerProvider(c.Scout.CrawlerEndpoint, c.Scout.CrawlerToken)
	studio.Crawler = scoutSvc
	// worker 執行 job 本體，不經 API 再入列
	studio.Jobs = nil

	logx.Infof("haixun worker started id=%s interval=%s (demo + token_renew + persona_analyze + compose_mimic + outbox)", workerID, interval)
	fmt.Printf("worker running id=%s interval=%s (Ctrl+C to stop)\n", workerID, interval)

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	tick := time.NewTicker(interval)
	defer tick.Stop()

	ctx := context.Background()
	const outboxLockTTL = 3 * time.Minute
	outboxLockKey := strings.Trim(c.Redis.Namespace, ":") + ":worker:outbox"
	outboxLock := redislock.NewWithClient(redisClient, outboxLockKey, workerID, outboxLockTTL)
	for {
		select {
		case <-sig:
			logx.Infof("worker %s shutting down", workerID)
			return
		case <-tick.C:
			// 1) claim one due job
			j, err := jobs.ClaimNext(ctx, workerID)
			if err == nil {
				switch j.TemplateType {
				case "", jobDomain.TemplateDemo:
					logx.Infof("worker %s demo job %s", workerID, j.ID)
					if _, err := jobs.RunDemoToSuccess(ctx, j.ID); err != nil {
						logx.Errorf("worker %s job %s failed: %v", workerID, j.ID, err)
						_, _ = jobs.FailJob(ctx, j.ID, err.Error())
					} else {
						logx.Infof("worker %s job %s succeeded", workerID, j.ID)
					}
				case jobDomain.TemplateThreadsTokenRenew:
					logx.Infof("worker %s token renew job %s account=%s", workerID, j.ID, j.RefID)
					if err := runTokenRenew(ctx, jobs, threadsSvc, j); err != nil {
						logx.Errorf("worker %s renew %s failed: %v", workerID, j.ID, err)
						_, _ = jobs.FailJob(ctx, j.ID, err.Error())
					} else {
						logx.Infof("worker %s renew %s ok", workerID, j.ID)
					}
				case jobDomain.TemplatePersonaAnalyzeAccount, jobDomain.TemplatePersonaAnalyzeText:
					logx.Infof("worker %s persona analyze job %s type=%s persona=%s", workerID, j.ID, j.TemplateType, j.RefID)
					if err := runPersonaAnalyze(ctx, jobs, studio, j); err != nil {
						logx.Errorf("worker %s persona %s failed: %v", workerID, j.ID, err)
						_, _ = jobs.FailJob(ctx, j.ID, err.Error())
					} else {
						logx.Infof("worker %s persona %s ok", workerID, j.ID)
					}
				case jobDomain.TemplateComposeMimic:
					logx.Infof("worker %s compose_mimic job %s", workerID, j.ID)
					if err := runComposeMimic(ctx, jobs, studio, j); err != nil {
						logx.Errorf("worker %s compose_mimic %s failed: %v", workerID, j.ID, err)
						_, _ = jobs.FailJob(ctx, j.ID, err.Error())
					} else {
						logx.Infof("worker %s compose_mimic %s ok", workerID, j.ID)
					}
				case jobDomain.TemplatePlayGenerateScript:
					logx.Infof("worker %s play_generate_script job %s play=%s", workerID, j.ID, j.RefID)
					if err := runPlayGenerateScript(ctx, jobs, studio, j); err != nil {
						logx.Errorf("worker %s play_generate_script %s failed: %v", workerID, j.ID, err)
						_, _ = jobs.FailJob(ctx, j.ID, err.Error())
					} else {
						logx.Infof("worker %s play_generate_script %s ok", workerID, j.ID)
					}
				case jobDomain.TemplateScoutScan:
					if err := runScoutScan(ctx, jobs, scoutSvc, j); err != nil {
						logx.Errorf("worker %s scout scan %s failed: %v", workerID, j.ID, err)
						_, _ = jobs.FailJob(ctx, j.ID, err.Error())
					}
				default:
					logx.Infof("worker %s unknown template %s job %s — fail", workerID, j.TemplateType, j.ID)
					_, _ = jobs.FailJob(ctx, j.ID, "unknown template: "+j.TemplateType)
				}
			}
			// 2) One worker owns an outbox tick and renews its lease while claims run.
			processOutbox(ctx, studio, outboxLock, workerID, outboxLockTTL)
			// 3) 終態任務超過 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)
			}
		}
	}
}

func processOutbox(ctx context.Context, studio *studioUC.Service, lock *redislock.Lock, workerID string, ttl time.Duration) {
	locked, err := lock.Acquire(ctx)
	if err != nil {
		logx.Errorf("worker %s outbox lock acquire: %v", workerID, err)
		return
	}
	if !locked {
		return
	}

	runCtx, cancel := context.WithCancel(ctx)
	renewFailure := make(chan error, 1)
	renewStopped := make(chan struct{})
	go func() {
		defer close(renewStopped)
		ticker := time.NewTicker(ttl / 3)
		defer ticker.Stop()
		for {
			select {
			case <-runCtx.Done():
				return
			case <-ticker.C:
				renewed, renewErr := lock.Renew(runCtx)
				if renewErr != nil {
					renewFailure <- fmt.Errorf("renew: %w", renewErr)
					cancel()
					return
				}
				if !renewed {
					renewFailure <- errors.New("lease ownership lost")
					cancel()
					return
				}
			}
		}
	}()
	defer func() {
		cancel()
		<-renewStopped
		releaseCtx, releaseCancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer releaseCancel()
		if err := lock.Release(releaseCtx); err != nil {
			logx.Errorf("worker %s outbox lock release: %v", workerID, err)
		}
	}()

	n, processErr := studio.ProcessDueSteps(runCtx, 0)
	select {
	case renewErr := <-renewFailure:
		logx.Errorf("worker %s outbox lock renewal failed; claims canceled: %v", workerID, renewErr)
		return
	default:
	}
	if processErr != nil {
		logx.Errorf("worker %s outbox tick: %v", workerID, processErr)
	} else if n > 0 {
		logx.Infof("worker %s outbox published %d step(s)", workerID, n)
	}
}

func runScoutScan(ctx context.Context, jobs *jobUC.Service, scout *scoutUC.Service, j *jobDomain.Job) error {
	var brief scoutDomain.RunBrief
	if err := json.Unmarshal([]byte(j.Payload), &brief); err != nil {
		return fmt.Errorf("invalid scout scan payload: %w", err)
	}
	if _, err := jobs.MarkRunningProgress(ctx, j.ID, 15, "海巡 · 準備搜尋來源"); err != nil {
		return err
	}
	posts, err := scout.RunScanFromBrief(ctx, j.OwnerUID, &brief)
	if err != nil {
		return err
	}
	if _, err := jobs.MarkRunningProgress(ctx, j.ID, 90, fmt.Sprintf("海巡 · 已寫入 %d 筆候選", len(posts))); err != nil {
		return err
	}
	_, err = jobs.SucceedJob(ctx, j.ID, fmt.Sprintf("海巡完成 · 命中 %d 筆", len(posts)))
	return err
}

func runTokenRenew(ctx context.Context, jobs *jobUC.Service, threads *threadsUC.Service, j *jobDomain.Job) error {
	if j.RefID == "" {
		return fmt.Errorf("缺少帳號 ref_id")
	}
	_, _ = jobs.MarkRunningProgress(ctx, j.ID, 15, "定期延長 Threads token · 準備中")
	_, _ = jobs.MarkRunningProgress(ctx, j.ID, 40, "定期延長 Threads token · 向平台刷新中")
	acc, err := threads.Refresh(ctx, j.OwnerUID, j.RefID)
	if err != nil {
		return err
	}
	_, _ = jobs.MarkRunningProgress(ctx, j.ID, 75, "定期延長 Threads token · 排程下一次")
	// 下次再約 30 天
	next := jobDomain.NowNano() + jobDomain.DefaultTokenRenewDelayNs
	if acc != nil && acc.SessionExpiresAt > 0 {
		early := acc.SessionExpiresAt - int64(30*24*time.Hour)
		if early > jobDomain.NowNano() {
			next = early
		}
	}
	_ = jobs.ScheduleTokenRenew(ctx, j.OwnerUID, j.RefID, next)
	_, _ = jobs.MarkRunningProgress(ctx, j.ID, 90, "定期延長 Threads token · 寫入完成")
	_, err = jobs.SucceedJob(ctx, j.ID, "定期延長完成 · 已排程約 30 天後再執行")
	return err
}

func runPersonaAnalyze(ctx context.Context, jobs *jobUC.Service, studio *studioUC.Service, j *jobDomain.Job) error {
	if j.RefID == "" {
		return fmt.Errorf("缺少人設 ref_id")
	}
	startSum := "人設分析 · 文字分析中…"
	if j.TemplateType == jobDomain.TemplatePersonaAnalyzeAccount {
		startSum = "人設分析 · 爬取公開貼文 + AI 分析中…（可離開頁面）"
	}
	_, _ = jobs.MarkRunningProgress(ctx, j.ID, 10, startSum)

	err := studio.ExecutePersonaAnalyzeJob(ctx, j.TemplateType, j.OwnerUID, j.RefID, j.Payload, func(pct int, sum string) {
		if pct < 10 {
			pct = 10
		}
		if pct > 95 {
			pct = 95
		}
		_, _ = jobs.MarkRunningProgress(ctx, j.ID, pct, sum)
	})
	if err != nil {
		return err
	}
	_, err = jobs.SucceedJob(ctx, j.ID, "人設分析完成 · 已寫入指紋／範本")
	return err
}

func runPlayGenerateScript(ctx context.Context, jobs *jobUC.Service, studio *studioUC.Service, j *jobDomain.Job) error {
	var pl jobUC.PlayGenerateScriptPayload
	if err := json.Unmarshal([]byte(j.Payload), &pl); err != nil {
		return fmt.Errorf("invalid play_generate_script payload: %w", err)
	}
	playID := strings.TrimSpace(pl.PlayID)
	if playID == "" {
		playID = strings.TrimSpace(j.RefID)
	}
	if playID == "" {
		return fmt.Errorf("empty play_id")
	}
	_, _ = jobs.MarkRunningProgress(ctx, j.ID, 15, "劇本產文 · 準備 AI（一次產全部）…")
	llmCtx, cancel := context.WithTimeout(ctx, 4*time.Minute)
	defer cancel()
	_, _ = jobs.MarkRunningProgress(ctx, j.ID, 40, "劇本產文 · 呼叫模型中（可離開頁面）…")
	n, err := studio.GeneratePlayScript(llmCtx, j.OwnerUID, playID, pl.OnlyEmpty)
	if err != nil {
		return err
	}
	pl.Filled = n
	body, _ := json.Marshal(pl)
	_, _ = jobs.MarkRunningProgress(ctx, j.ID, 90, fmt.Sprintf("劇本產文 · 已寫入 %d 步…", n))
	_, err = jobs.SucceedJobWithPayload(ctx, j.ID, fmt.Sprintf("劇本產文完成 · 已填 %d 步，請回互回方案查看", n), string(body))
	return err
}

func runComposeMimic(ctx context.Context, jobs *jobUC.Service, studio *studioUC.Service, j *jobDomain.Job) error {
	var pl jobUC.ComposeMimicPayload
	if err := json.Unmarshal([]byte(j.Payload), &pl); err != nil {
		return fmt.Errorf("invalid compose_mimic payload: %w", err)
	}
	if strings.TrimSpace(pl.SourceText) == "" {
		return fmt.Errorf("empty source_text")
	}
	_, _ = jobs.MarkRunningProgress(ctx, j.ID, 15, "仿寫貼文 · 準備 AI…")
	// worker 不綁 gateway 120s；給模型較長時間（reasoning 模型需要）
	llmCtx, cancel := context.WithTimeout(ctx, 4*time.Minute)
	defer cancel()
	_, _ = jobs.MarkRunningProgress(ctx, j.ID, 35, "仿寫貼文 · 呼叫模型中（可離開頁面）…")
	text, err := studio.Mimic(llmCtx, j.OwnerUID, pl.SourceText, pl.PersonaID, pl.StructureNotes)
	if err != nil {
		return err
	}
	_, _ = jobs.MarkRunningProgress(ctx, j.ID, 90, "仿寫貼文 · 寫入結果…")
	pl.ResultText = text
	body, _ := json.Marshal(pl)
	_, err = jobs.SucceedJobWithPayload(ctx, j.ID, "仿寫完成 · 已可套用到正文", string(body))
	return err
}

// workerAIKeys — 與 gateway studioAIKeys 對齊
type workerAIKeys struct {
	Members  memberDomain.Repository
	Resolver *usageUC.SettingsResolver
}

func (k *workerAIKeys) ResolveAI(ctx context.Context, ownerUID int64) (provider, model, apiKey string, err error) {
	st, err := k.Members.GetSettings(ctx, ownerUID)
	if err != nil {
		return "", "", "", err
	}
	if st == nil {
		st = memberDomain.DefaultSettings(ownerUID)
	}
	// 完全依會員 AI 設定（provider / model / key），不寫死模型
	provider = ai.NormalizeProvider(st.Provider)
	model = strings.TrimSpace(st.Model)
	if model == "" {
		return provider, "", "", fmt.Errorf("請到設定選擇 AI 模型")
	}
	if k.Resolver != nil {
		_, key, rerr := k.Resolver.ResolveKey(ctx, ownerUID, usageDomain.MeterAICopy)
		if rerr != nil && !errors.Is(rerr, usageDomain.ErrNoKey) {
			return provider, model, "", rerr
		}
		if rerr == nil && strings.TrimSpace(key) != "" {
			return provider, model, key, nil
		}
	}
	if key := st.KeyForProvider(provider); key != "" {
		return provider, model, key, nil
	}
	return provider, model, "", fmt.Errorf("no AI key（請到設定填寫 Key，或確認平台已配置）")
}

type workerDevMode struct {
	Members memberDomain.Repository
}

func (d *workerDevMode) DevModeEnabled(ctx context.Context, uid int64) (bool, error) {
	if d == nil || d.Members == nil {
		return false, nil
	}
	st, err := d.Members.GetSettings(ctx, uid)
	if err != nil {
		return false, err
	}
	if st == nil {
		return false, nil
	}
	return st.DevModeEnabled, nil
}

func workerStorage(c config.Config) fsDomain.Storage {
	os := c.ObjectStorage
	if strings.TrimSpace(os.Endpoint) == "" {
		logx.Info("worker storage: disabled (no endpoint)")
		return noop.New()
	}
	s, err := s3store.New(s3store.Config{
		Endpoint: os.Endpoint, Region: os.Region, Bucket: os.Bucket,
		AccessKey: os.AccessKey, SecretKey: os.SecretKey,
		UsePathStyle: os.UsePathStyle, PublicBaseURL: os.PublicBaseURL,
	})
	if err != nil {
		logx.Errorf("worker storage init failed: %v", err)
		return noop.New()
	}
	logx.Infof("worker storage: minio/s3 bucket=%s (ephemeral publish images cleanup enabled)", os.Bucket)
	return s
}