geo/server/internal/tenant/app/generation_task_runtime.go

package app

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/geo-platform/tenant-api/internal/shared/llm"
	"github.com/geo-platform/tenant-api/internal/tenant/repository"
)

type ClaimedGenerationTask struct {
	ID                 int64
	TenantID           int64
	OperatorID         int64
	ArticleID          int64
	QuotaReservationID int64
	TaskType           string
	InputParams        map[string]interface{}
}

const generationCleanupTimeout = 30 * time.Second

type promptRuleTaskContext struct {
	PromptRuleID int64
	BrandID      int64
}

func newGenerationCleanupContext() (context.Context, context.CancelFunc) {
	return context.WithTimeout(context.Background(), generationCleanupTimeout)
}

func rollbackGenerationTx(tx pgx.Tx) {
	if tx == nil {
		return
	}

	ctx, cancel := newGenerationCleanupContext()
	defer cancel()

	_ = tx.Rollback(ctx)
}

func HandleClaimedGenerationTaskFailure(
	ctx context.Context,
	pool *pgxpool.Pool,
	templateService *TemplateService,
	promptRuleService *PromptRuleGenerationService,
	task ClaimedGenerationTask,
	taskErr error,
) error {
	cleanupCtx, cancel := newGenerationCleanupContext()
	defer cancel()

	errMsg := formatGenerationFailure("task_load", taskErr)
	if taskErr != nil && strings.Contains(taskErr.Error(), "[worker_recovery]") {
		errMsg = strings.TrimSpace(taskErr.Error())
	}
	now := time.Now()

	auditRepo := repository.NewAuditRepository(pool)
	articleRepo := repository.NewArticleRepository(pool)
	quotaRepo := repository.NewQuotaRepository(pool)

	var failureErr error
	if err := auditRepo.UpdateGenerationTaskStatus(cleanupCtx, repository.GenerationTaskStatusInput{
		ID:           task.ID,
		TenantID:     task.TenantID,
		Status:       "failed",
		ErrorMessage: &errMsg,
		CompletedAt:  &now,
	}); err != nil {
		failureErr = errors.Join(failureErr, fmt.Errorf("update generation task failed status: %w", err))
	}
	if task.ArticleID > 0 {
		if err := articleRepo.UpdateArticleGenerateStatus(cleanupCtx, task.ArticleID, task.TenantID, "failed"); err != nil {
			failureErr = errors.Join(failureErr, fmt.Errorf("update article failed status: %w", err))
		}
	}

	if task.QuotaReservationID > 0 {
		var balance int
		balanceErr := error(nil)
		if task.OperatorID > 0 {
			balance, balanceErr = quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
		}
		refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, task.QuotaReservationID, task.TenantID)
		if err != nil {
			failureErr = errors.Join(failureErr, fmt.Errorf("refund generation reservation: %w", err))
		} else if refunded {
			if balanceErr != nil {
				failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
			} else if task.OperatorID > 0 {
				reason := "generation_refund"
				referenceType := "generation_task"
				taskReferenceID := task.ID
				if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
					TenantID:      task.TenantID,
					OperatorID:    task.OperatorID,
					QuotaType:     "article_generation",
					Delta:         1,
					BalanceAfter:  balance + 1,
					Reason:        &reason,
					ReferenceType: &referenceType,
					ReferenceID:   &taskReferenceID,
				}); err != nil {
					failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err))
				}
			}
		}
	}

	if task.ArticleID > 0 {
		switch {
		case templateService != nil && templateService.cache != nil:
			invalidateArticleCaches(cleanupCtx, templateService.cache, task.TenantID, &task.ArticleID)
		case promptRuleService != nil && promptRuleService.cache != nil:
			invalidateArticleCaches(cleanupCtx, promptRuleService.cache, task.TenantID, &task.ArticleID)
		}
	}

	title := resolveArticleTitle(task.InputParams, "")
	if strings.TrimSpace(extractString(task.InputParams, "task_name")) != "" {
		title = strings.TrimSpace(extractString(task.InputParams, "task_name"))
	}
	if title == "" {
		title = promptRuleGeneratingTitlePlaceholder
	}

	if task.TaskType == "template" && templateService != nil &&
		templateService.runtimeConfig().StreamEnabled && templateService.streams != nil {
		templateService.streams.Fail(task.ArticleID, task.ID, title, errMsg)
	}
	if task.TaskType == "custom_generation" && promptRuleService != nil &&
		promptRuleService.runtimeConfig().StreamEnabled && promptRuleService.streams != nil {
		promptRuleService.streams.Fail(task.ArticleID, task.ID, title, errMsg)
	}

	return failureErr
}

func (s *TemplateService) ProcessQueuedTask(ctx context.Context, task ClaimedGenerationTask) {
	job := generationJob{
		OperatorID:    task.OperatorID,
		TenantID:      task.TenantID,
		ArticleID:     task.ArticleID,
		TaskID:        task.ID,
		ReservationID: task.QuotaReservationID,
		Params:        cloneInputParams(task.InputParams),
		InitialTitle:  resolveArticleTitle(task.InputParams, ""),
		WebSearch: llm.WebSearchOptions{
			Enabled: extractBool(task.InputParams, "enable_web_search"),
		},
	}
	if limit := extractInt(task.InputParams, "web_search_limit"); limit > 0 {
		job.WebSearch.Limit = int32(limit)
	}

	templateKey, templateName, promptTemplate := extractTemplateSnapshot(task.InputParams)
	if templateKey == "" || templateName == "" {
		record, err := s.loadTemplateRecordForTask(ctx, task)
		if err != nil {
			s.failGeneration(context.Background(), job, job.InitialTitle, "task_load", err)
			return
		}
		templateKey = record.TemplateKey
		templateName = record.TemplateName
		promptTemplate = record.PromptTemplate
	}
	job.TemplateKey = templateKey
	job.TemplateName = templateName
	job.PromptTemplate = promptTemplate

	s.executeGeneration(ctx, job)
}

func (s *TemplateService) loadTemplateRecordForTask(ctx context.Context, task ClaimedGenerationTask) (*repository.TemplateRecord, error) {
	var templateID int64
	if err := s.pool.QueryRow(ctx, `
		SELECT template_id
		FROM articles
		WHERE id = $1
		  AND tenant_id = $2
		  AND source_type = 'template'
		  AND deleted_at IS NULL
	`, task.ArticleID, task.TenantID).Scan(&templateID); err != nil {
		return nil, fmt.Errorf("load task template id: %w", err)
	}

	record, err := s.templates.GetTemplateByID(ctx, templateID, task.TenantID)
	if err != nil {
		return nil, fmt.Errorf("load task template: %w", err)
	}
	return record, nil
}

func (s *PromptRuleGenerationService) ProcessQueuedTask(ctx context.Context, task ClaimedGenerationTask) {
	inputParams := cloneInputParams(task.InputParams)
	initialTitle := strings.TrimSpace(extractString(inputParams, "task_name"))
	if initialTitle == "" {
		initialTitle = promptRuleGeneratingTitlePlaceholder
	}

	job := promptRuleGenerationJob{
		OperatorID:    task.OperatorID,
		TenantID:      task.TenantID,
		ArticleID:     task.ArticleID,
		TaskID:        task.ID,
		ReservationID: task.QuotaReservationID,
		InputParams:   inputParams,
		InitialTitle:  initialTitle,
		WebSearch: llm.WebSearchOptions{
			Enabled: extractBool(inputParams, "enable_web_search"),
		},
	}

	rule := extractPromptRuleSnapshot(inputParams)
	if rule == nil || int64(extractInt(inputParams, "prompt_rule_id")) <= 0 || int64(extractInt(inputParams, "brand_id")) <= 0 {
		taskContext, err := s.resolvePromptRuleTaskContext(ctx, task)
		if err != nil {
			s.failGeneration(context.Background(), job, initialTitle, "task_load", err)
			return
		}
		inputParams["prompt_rule_id"] = taskContext.PromptRuleID
		inputParams["brand_id"] = taskContext.BrandID
	}
	if rule == nil {
		promptRuleID := int64(extractInt(inputParams, "prompt_rule_id"))
		brandID := int64(extractInt(inputParams, "brand_id"))
		var err error
		rule, err = s.loadPromptRule(ctx, task.TenantID, brandID, promptRuleID)
		if err != nil {
			s.failGeneration(context.Background(), job, initialTitle, "task_load", err)
			return
		}
	}

	knowledgePrompt := strings.TrimSpace(extractString(inputParams, generationPayloadKnowledgePrompt))
	if knowledgePrompt == "" {
		if knowledgeGroupIDs := extractKnowledgeGroupIDs(inputParams["knowledge_group_ids"]); len(knowledgeGroupIDs) > 0 {
			if s.knowledge == nil {
				s.failGeneration(context.Background(), job, initialTitle, "knowledge_resolve", fmt.Errorf("knowledge service is not configured"))
				return
			}

			knowledgeContext, err := s.knowledge.ResolveContext(
				ctx,
				task.TenantID,
				knowledgeGroupIDs,
				buildPromptRuleKnowledgeQuery(rule, inputParams),
			)
			if err != nil {
				s.failGeneration(context.Background(), job, initialTitle, "knowledge_resolve", err)
				return
			}
			knowledgePrompt = s.knowledge.RenderPromptSection(knowledgeContext)
		}
	}

	job.Prompt = buildPromptRuleGenerationPrompt(rule, inputParams, knowledgePrompt)
	s.executeGeneration(ctx, job)
}

func (s *PromptRuleGenerationService) resolvePromptRuleTaskContext(ctx context.Context, task ClaimedGenerationTask) (promptRuleTaskContext, error) {
	resolved := promptRuleTaskContext{
		PromptRuleID: int64(extractInt(task.InputParams, "prompt_rule_id")),
		BrandID:      int64(extractInt(task.InputParams, "brand_id")),
	}
	if resolved.PromptRuleID > 0 && resolved.BrandID > 0 {
		return resolved, nil
	}

	if task.ArticleID <= 0 {
		return resolved, fmt.Errorf("article context missing for prompt rule task")
	}

	var articlePromptRuleID int64
	var articleBrandID int64
	if err := s.pool.QueryRow(ctx, `
		SELECT COALESCE(prompt_rule_id, 0), brand_id
		FROM articles
		WHERE id = $1
		  AND tenant_id = $2
		  AND source_type = 'custom_generation'
		  AND deleted_at IS NULL
	`, task.ArticleID, task.TenantID).Scan(&articlePromptRuleID, &articleBrandID); err != nil {
		return resolved, fmt.Errorf("load task prompt rule context: %w", err)
	}

	if resolved.PromptRuleID <= 0 {
		resolved.PromptRuleID = articlePromptRuleID
	}
	if resolved.BrandID <= 0 {
		resolved.BrandID = articleBrandID
	}
	if resolved.PromptRuleID <= 0 {
		return resolved, fmt.Errorf("prompt rule context missing for prompt rule task")
	}
	if resolved.BrandID <= 0 {
		return resolved, fmt.Errorf("brand context missing for prompt rule task")
	}
	return resolved, nil
}

func (s *PromptRuleGenerationService) resolvePromptRuleIDForTask(ctx context.Context, task ClaimedGenerationTask) (int64, error) {
	if promptRuleID := int64(extractInt(task.InputParams, "prompt_rule_id")); promptRuleID > 0 {
		return promptRuleID, nil
	}

	var promptRuleID int64
	if err := s.pool.QueryRow(ctx, `
		SELECT prompt_rule_id
		FROM articles
		WHERE id = $1
		  AND tenant_id = $2
		  AND source_type = 'custom_generation'
		  AND deleted_at IS NULL
	`, task.ArticleID, task.TenantID).Scan(&promptRuleID); err != nil {
		return 0, fmt.Errorf("load task prompt rule id: %w", err)
	}
	return promptRuleID, nil
}