254 lines
9.7 KiB
Go
254 lines
9.7 KiB
Go
package orchestration
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"strings"
|
|
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/billing"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/identity"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/providers"
|
|
)
|
|
|
|
type CreationInput struct {
|
|
OwnerID, ExternalClientID string
|
|
Capability string
|
|
Body map[string]any
|
|
IdempotencyKey string
|
|
RetryOf string
|
|
}
|
|
|
|
type CreationBuilder interface {
|
|
Build(context.Context, string, string, string, string, map[string]any) (jobs.CreateCommand, error)
|
|
}
|
|
|
|
type JobCreator interface {
|
|
Create(context.Context, jobs.CreateCommand) (jobs.Job, bool, error)
|
|
}
|
|
|
|
type BillingQuoter interface {
|
|
Quote(context.Context, billing.QuoteCommand) (*billing.Quote, error)
|
|
}
|
|
|
|
type ChargeLedger interface {
|
|
Charge(context.Context, billing.ChargeRequest) (billing.WalletPosting, error)
|
|
}
|
|
|
|
type CreationStateWriter interface {
|
|
WriteBilling(context.Context, string, json.RawMessage) error
|
|
ActivateCreation(context.Context, string, json.RawMessage) error
|
|
FailCreation(context.Context, jobs.Job) error
|
|
}
|
|
|
|
type AtomicCreationCharger interface {
|
|
ChargeAndActivateCreation(context.Context, billing.ChargeRequest, json.RawMessage) (json.RawMessage, error)
|
|
}
|
|
|
|
// CreationCoordinator is the single creation transaction policy used by HTTP.
|
|
// Its persistence port intentionally exposes durable state transitions rather
|
|
// than wallet or SQL details.
|
|
type CreationCoordinator struct {
|
|
builder CreationBuilder
|
|
jobs JobCreator
|
|
quotes BillingQuoter
|
|
charges ChargeLedger
|
|
state CreationStateWriter
|
|
}
|
|
|
|
func NewCreationCoordinator(builder CreationBuilder, creator JobCreator, quotes BillingQuoter, charges ChargeLedger, state CreationStateWriter) *CreationCoordinator {
|
|
return &CreationCoordinator{builder: builder, jobs: creator, quotes: quotes, charges: charges, state: state}
|
|
}
|
|
|
|
func (c *CreationCoordinator) CreatePlatform(ctx context.Context, session identity.Session, input CreationInput) (jobs.Job, bool, error) {
|
|
if c == nil || c.builder == nil || c.jobs == nil || c.quotes == nil || c.charges == nil || c.state == nil {
|
|
return jobs.Job{}, false, errors.New("create generation job")
|
|
}
|
|
input.OwnerID = session.User.ID
|
|
input.ExternalClientID = ""
|
|
command, err := c.builder.Build(ctx, input.OwnerID, "", input.Capability, input.IdempotencyKey, input.Body)
|
|
if err != nil {
|
|
return jobs.Job{}, false, err
|
|
}
|
|
command.Job.RetryOf = input.RetryOf
|
|
use := usageContext{Source: "platform", AccountID: session.User.ID, Username: session.User.Username, DisplayName: session.User.DisplayName, Role: session.User.Role, TenantID: session.User.TenantID, OrganizationID: session.User.OrganizationID, OrganizationName: session.User.OrganizationName}
|
|
command.Job.UsageContext, _ = json.Marshal(use)
|
|
parameters := map[string]any{}
|
|
var prepared providers.Request
|
|
if json.Unmarshal(command.Job.RequestPayload, &prepared) == nil {
|
|
for key, value := range prepared.Settings {
|
|
parameters[key] = value
|
|
}
|
|
parameters["inputUrlCount"] = len(prepared.InputURLs)
|
|
}
|
|
quote, err := c.quotes.Quote(ctx, billing.QuoteCommand{AccountID: session.User.ID, OrganizationID: session.User.OrganizationID, OrganizationName: session.User.OrganizationName, Role: session.User.Role, Provider: command.Job.Provider, Capability: command.Job.Capability, ReqKey: command.Job.ReqKey, Parameters: parameters, Payload: input.Body})
|
|
if err != nil {
|
|
return jobs.Job{}, false, errors.New("quote generation charge")
|
|
}
|
|
if quote != nil {
|
|
state := quoteMap(*quote)
|
|
state["status"] = "pending"
|
|
command.Job.Billing, _ = json.Marshal(state)
|
|
command.HoldDispatch = true
|
|
}
|
|
created, reused, err := c.jobs.Create(ctx, command)
|
|
if err != nil {
|
|
return jobs.Job{}, false, err
|
|
}
|
|
return c.charge(ctx, created, reused)
|
|
}
|
|
|
|
func (c *CreationCoordinator) CreatePublic(ctx context.Context, input CreationInput) (jobs.Job, bool, error) {
|
|
if c == nil || c.builder == nil || c.jobs == nil {
|
|
return jobs.Job{}, false, errors.New("create generation job")
|
|
}
|
|
command, err := c.builder.Build(ctx, input.OwnerID, input.ExternalClientID, input.Capability, input.IdempotencyKey, input.Body)
|
|
if err != nil {
|
|
return jobs.Job{}, false, err
|
|
}
|
|
return c.jobs.Create(ctx, command)
|
|
}
|
|
|
|
// RetryPlatform mirrors the platform image retry contract: ownership and
|
|
// capability are checked by the HTTP adapter, while the new job is rebuilt
|
|
// through the normal builder/quote/charge pipeline with no inherited billing.
|
|
func (c *CreationCoordinator) RetryPlatform(ctx context.Context, session identity.Session, original jobs.Job) (jobs.Job, error) {
|
|
if original.OwnerID != session.User.ID || original.Capability != "image.generate" {
|
|
return jobs.Job{}, &jobs.Error{Kind: jobs.ErrorNotFound, Status: 404, Message: "任务不存在"}
|
|
}
|
|
var request providers.Request
|
|
if json.Unmarshal(original.RequestPayload, &request) != nil {
|
|
return jobs.Job{}, errors.New("retry generation job")
|
|
}
|
|
body := map[string]any{
|
|
"prompt": request.Prompt,
|
|
"inputUrls": request.InputURLs,
|
|
"inputAssetIds": append([]string(nil), original.InputAssetIDs...),
|
|
"materials": request.Materials,
|
|
"settings": request.Settings,
|
|
"priority": jobs.NormalizePriority(original.Priority),
|
|
}
|
|
if original.WebhookURL != "" {
|
|
body["webhookUrl"] = original.WebhookURL
|
|
}
|
|
if engine := retryEngine(original.Provider, original.Capability); engine != "" {
|
|
body["engine"] = engine
|
|
}
|
|
created, _, err := c.CreatePlatform(ctx, session, CreationInput{Capability: original.Capability, Body: body, RetryOf: original.ID})
|
|
if err != nil {
|
|
return jobs.Job{}, err
|
|
}
|
|
return created, nil
|
|
}
|
|
|
|
func (c *CreationCoordinator) charge(ctx context.Context, job jobs.Job, reused bool) (jobs.Job, bool, error) {
|
|
if len(job.Billing) == 0 {
|
|
return job, reused, nil
|
|
}
|
|
var charge billingState
|
|
if json.Unmarshal(job.Billing, &charge) != nil {
|
|
return jobs.Job{}, reused, errors.New("charge generation job")
|
|
}
|
|
if charge.Status == "charged" || charge.Status == "not_charged" || charge.Status == "refunded" {
|
|
return job, reused, nil
|
|
}
|
|
if charge.Status != "pending" {
|
|
return jobs.Job{}, reused, errors.New("charge generation job")
|
|
}
|
|
if charge.QuotaExempt {
|
|
charge.raw["status"] = "not_charged"
|
|
return c.writeCharge(ctx, job, reused, charge.raw)
|
|
}
|
|
var use usageContext
|
|
if json.Unmarshal(job.UsageContext, &use) != nil || use.OrganizationID == "" || charge.AmountFen <= 0 {
|
|
return c.failCharge(ctx, job, reused, errors.New("charge generation job"), charge.raw)
|
|
}
|
|
request := billing.ChargeRequest{OrganizationID: use.OrganizationID, AccountID: use.AccountID, JobID: job.ID, AmountFen: charge.AmountFen, Description: capabilityLabel(job.Capability) + " · " + job.ReqKey, Metadata: map[string]any{"quote": cloneMap(charge.raw), "accountName": use.DisplayName, "organizationName": use.OrganizationName}}
|
|
if atomic, ok := c.state.(AtomicCreationCharger); ok {
|
|
pending, encodeErr := json.Marshal(charge.raw)
|
|
if encodeErr != nil {
|
|
return jobs.Job{}, reused, errors.New("charge generation job")
|
|
}
|
|
encoded, atomicErr := atomic.ChargeAndActivateCreation(ctx, request, pending)
|
|
if atomicErr != nil {
|
|
if errors.Is(atomicErr, billing.ErrCommitOutcomeUnknown) {
|
|
return jobs.Job{}, reused, errors.New("charge generation job")
|
|
}
|
|
return c.failCharge(ctx, job, reused, safeBillingError(atomicErr), charge.raw)
|
|
}
|
|
job.Billing = encoded
|
|
return job, reused, nil
|
|
}
|
|
posting, err := c.charges.Charge(ctx, request)
|
|
if err != nil {
|
|
return c.failCharge(ctx, job, reused, safeBillingError(err), charge.raw)
|
|
}
|
|
charge.raw["status"] = "charged"
|
|
charge.raw["ledgerEntryId"] = posting.LedgerID
|
|
charge.raw["chargedAt"] = posting.CreatedAt.UTC().Format("2006-01-02T15:04:05.999999999Z07:00")
|
|
return c.writeCharge(ctx, job, reused, charge.raw)
|
|
}
|
|
|
|
func (c *CreationCoordinator) writeCharge(ctx context.Context, job jobs.Job, reused bool, state map[string]any) (jobs.Job, bool, error) {
|
|
encoded, err := json.Marshal(state)
|
|
if err != nil || c.state.ActivateCreation(ctx, job.ID, encoded) != nil {
|
|
return jobs.Job{}, reused, errors.New("persist generation charge")
|
|
}
|
|
job.Billing = encoded
|
|
return job, reused, nil
|
|
}
|
|
|
|
func (c *CreationCoordinator) failCharge(ctx context.Context, job jobs.Job, reused bool, cause error, state map[string]any) (jobs.Job, bool, error) {
|
|
state["status"] = "not_charged"
|
|
job.Billing, _ = json.Marshal(state)
|
|
job.Status = jobs.StatusFailed
|
|
job.Error = &jobs.JobError{Message: "generation charge failed", Retryable: false}
|
|
if err := c.state.FailCreation(ctx, job); err != nil {
|
|
return jobs.Job{}, reused, errors.New("persist failed generation charge")
|
|
}
|
|
return job, reused, cause
|
|
}
|
|
|
|
func quoteMap(quote billing.Quote) map[string]any {
|
|
raw, _ := json.Marshal(quote)
|
|
value := map[string]any{}
|
|
_ = json.Unmarshal(raw, &value)
|
|
return value
|
|
}
|
|
|
|
func safeBillingError(err error) error {
|
|
switch billing.HTTPStatus(err) {
|
|
case 402:
|
|
return &billing.StatusError{Status: 402, Err: billing.ErrInsufficientBalance}
|
|
case 409:
|
|
return &billing.StatusError{Status: 409, Err: billing.ErrIdempotencyConflict}
|
|
default:
|
|
message := err.Error()
|
|
if strings.Contains(message, "BILLING_INSUFFICIENT_BALANCE") {
|
|
return &billing.StatusError{Status: 402, Err: billing.ErrInsufficientBalance}
|
|
}
|
|
if strings.Contains(message, "BILLING_IDEMPOTENCY_PAYLOAD_MISMATCH") {
|
|
return &billing.StatusError{Status: 409, Err: billing.ErrIdempotencyConflict}
|
|
}
|
|
return errors.New("charge generation job")
|
|
}
|
|
}
|
|
|
|
func retryEngine(provider, capability string) string {
|
|
if capability == "video.generate" {
|
|
if provider == "seedance" || provider == "bailian" {
|
|
return provider
|
|
}
|
|
return ""
|
|
}
|
|
switch provider {
|
|
case "volcengine-visual":
|
|
return "jimeng"
|
|
case "evolink", "bailian":
|
|
return provider
|
|
default:
|
|
return ""
|
|
}
|
|
}
|