diff --git a/backend/internal/jobs/job.go b/backend/internal/jobs/job.go new file mode 100644 index 0000000..6dcb33b --- /dev/null +++ b/backend/internal/jobs/job.go @@ -0,0 +1,207 @@ +// Package jobs owns generation-job state, ownership, idempotency, and worker +// coordination independently of HTTP, provider, storage, and billing adapters. +package jobs + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "sort" + "strings" + "time" +) + +type Status string + +const ( + StatusQueued Status = "queued" + StatusRunning Status = "running" + StatusSucceeded Status = "succeeded" + StatusFailed Status = "failed" + StatusExpired Status = "expired" + StatusCancelled Status = "cancelled" +) + +func (status Status) Terminal() bool { + return status == StatusSucceeded || status == StatusFailed || status == StatusExpired || status == StatusCancelled +} + +func (status Status) Valid() bool { + return status == StatusQueued || status == StatusRunning || status.Terminal() +} + +type Job struct { + ID string `json:"id"` + OwnerID string `json:"ownerId"` + ExternalClientID string `json:"externalClientId,omitempty"` + Capability string `json:"capability"` + Provider string `json:"provider"` + ReqKey string `json:"reqKey"` + Status Status `json:"status"` + Prompt string `json:"prompt,omitempty"` + InputAssetIDs []string `json:"inputAssetIds"` + InputURLs []string `json:"inputUrls"` + OutputAssetIDs []string `json:"outputAssetIds"` + ProviderTaskID string `json:"providerTaskId,omitempty"` + RequestPayload json.RawMessage `json:"requestPayload"` + ResponsePayload json.RawMessage `json:"responsePayload,omitempty"` + Error *JobError `json:"error,omitempty"` + RetryOf string `json:"retryOf,omitempty"` + IdempotencyKey string `json:"idempotencyKey,omitempty"` + IdempotencyFingerprint string `json:"idempotencyFingerprint,omitempty"` + Priority int `json:"priority,omitempty"` + Attempts int `json:"attempts,omitempty"` + MaxAttempts int `json:"maxAttempts,omitempty"` + ScheduledAt time.Time `json:"scheduledAt,omitempty"` + LockedAt *time.Time `json:"lockedAt,omitempty"` + LockedBy string `json:"lockedBy,omitempty"` + StartedAt *time.Time `json:"startedAt,omitempty"` + CompletedAt *time.Time `json:"completedAt,omitempty"` + WebhookURL string `json:"webhookUrl,omitempty"` + WebhookAttempts int `json:"webhookAttempts,omitempty"` + WebhookLastStatus json.RawMessage `json:"webhookLastStatus,omitempty"` + UsageContext json.RawMessage `json:"usageContext,omitempty"` + Billing json.RawMessage `json:"billing,omitempty"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +type JobError struct { + Code any `json:"code,omitempty"` + Message string `json:"message"` + Retryable bool `json:"retryable,omitempty"` +} + +type Scope struct { + OwnerID string + ExternalClientID string +} + +func (scope Scope) Owns(job Job) bool { + if scope.OwnerID == "" || job.OwnerID != scope.OwnerID { + return false + } + return scope.ExternalClientID == "" || job.ExternalClientID == scope.ExternalClientID +} + +type ErrorKind string + +const ( + ErrorInvalid ErrorKind = "invalid" + ErrorNotFound ErrorKind = "not_found" + ErrorConflict ErrorKind = "conflict" +) + +type Error struct { + Kind ErrorKind + Status int + Message string +} + +func (err *Error) Error() string { return err.Message } + +var ErrUniqueIdempotency = errors.New("generation job idempotency key already exists") + +type ListFilter struct { + Scope Scope + Status Status + Capability string + Limit int + Before *time.Time +} + +type CreateCommand struct { + Job Job + IdempotencyBody map[string]any +} + +func NormalizePriority(value int) int { + if value < -100 { + return -100 + } + if value > 100 { + return 100 + } + return value +} + +func NormalizePublicLimit(value int) int { + if value == 0 { + return 50 + } + if value < 1 { + return 1 + } + if value > 200 { + return 200 + } + return value +} + +// Fingerprint excludes the compatibility body key idempotencyKey, recursively +// sorts object keys, preserves array order, and hashes the exact stable JSON. +func Fingerprint(body map[string]any) (string, error) { + source := make(map[string]any, len(body)) + for key, value := range body { + if key != "idempotencyKey" { + source[key] = value + } + } + stable, err := stableJSON(source) + if err != nil { + return "", err + } + sum := sha256.Sum256(stable) + return hex.EncodeToString(sum[:]), nil +} + +func stableJSON(value any) ([]byte, error) { + switch typed := value.(type) { + case map[string]any: + keys := make([]string, 0, len(typed)) + for key := range typed { + keys = append(keys, key) + } + sort.Strings(keys) + var output strings.Builder + output.WriteByte('{') + for index, key := range keys { + if index > 0 { + output.WriteByte(',') + } + encodedKey, _ := json.Marshal(key) + encodedValue, err := stableJSON(typed[key]) + if err != nil { + return nil, err + } + output.Write(encodedKey) + output.WriteByte(':') + output.Write(encodedValue) + } + output.WriteByte('}') + return []byte(output.String()), nil + case []any: + var output strings.Builder + output.WriteByte('[') + for index, item := range typed { + if index > 0 { + output.WriteByte(',') + } + encoded, err := stableJSON(item) + if err != nil { + return nil, err + } + output.Write(encoded) + } + output.WriteByte(']') + return []byte(output.String()), nil + default: + encoded, err := json.Marshal(value) + if err != nil { + return nil, fmt.Errorf("encode idempotency fingerprint: %w", err) + } + return encoded, nil + } +} diff --git a/backend/internal/jobs/jobs_test.go b/backend/internal/jobs/jobs_test.go new file mode 100644 index 0000000..f2ee4ad --- /dev/null +++ b/backend/internal/jobs/jobs_test.go @@ -0,0 +1,218 @@ +package jobs + +import ( + "context" + "encoding/json" + "errors" + "testing" + "time" +) + +func TestFingerprintMatchesSharedStableJSONVector(t *testing.T) { + body := map[string]any{ + "settings": map[string]any{"width": float64(1024), "height": float64(768)}, + "prompt": "hello", "idempotencyKey": "body-key", "capability": "image.generate", + } + got, err := Fingerprint(body) + if err != nil { + t.Fatal(err) + } + const want = "3ac3128766c0bdac36627b996cb5b3b1ed752794989d559c24f6621988e51a03" + if got != want { + t.Fatalf("Fingerprint = %q, want %q", got, want) + } +} + +func TestServiceReplaysMatchingIdempotencyAndRejectsDrift(t *testing.T) { + clock := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) + store := newMemoryJobStore() + service := NewService(store, func() time.Time { return clock }) + command := fixtureCreateCommand("hello") + + created, reused, err := service.Create(context.Background(), command) + if err != nil || reused { + t.Fatalf("first Create = (%#v,%v,%v)", created, reused, err) + } + replayed, reused, err := service.Create(context.Background(), command) + if err != nil || !reused || replayed.ID != created.ID { + t.Fatalf("replay Create = (%#v,%v,%v)", replayed, reused, err) + } + _, _, err = service.Create(context.Background(), fixtureCreateCommand("hello!")) + var conflict *Error + if !errors.As(err, &conflict) || conflict.Kind != ErrorConflict || conflict.Status != 409 { + t.Fatalf("drift error = %#v, want conflict/409", err) + } +} + +func TestServiceRecoversUniqueInsertRaceAsIdempotentReplay(t *testing.T) { + store := newMemoryJobStore() + command := fixtureCreateCommand("hello") + fingerprint, _ := Fingerprint(command.IdempotencyBody) + existing := command.Job + existing.IdempotencyFingerprint = fingerprint + store.race = &existing + service := NewService(store, func() time.Time { return time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) }) + + job, reused, err := service.Create(context.Background(), command) + if err != nil || !reused || job.ID != existing.ID { + t.Fatalf("Create race = (%#v,%v,%v)", job, reused, err) + } +} + +func TestCancelIsOwnerScopedIdempotentAndClearsLease(t *testing.T) { + store := newMemoryJobStore() + locked := time.Date(2026, 8, 13, 7, 59, 0, 0, time.UTC) + store.jobs["job-1"] = Job{ID: "job-1", OwnerID: "api:client", ExternalClientID: "client", Status: StatusRunning, LockedAt: &locked, LockedBy: "worker"} + service := NewService(store, func() time.Time { return time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) }) + refunds := &refundStub{} + + cancelled, err := service.Cancel(context.Background(), Scope{OwnerID: "api:client", ExternalClientID: "client"}, "job-1", refunds) + if err != nil || cancelled.Status != StatusCancelled || cancelled.LockedAt != nil || cancelled.LockedBy != "" || refunds.calls != 1 { + t.Fatalf("Cancel = %#v, %v; refunds=%d", cancelled, err, refunds.calls) + } + again, err := service.Cancel(context.Background(), Scope{OwnerID: "api:client", ExternalClientID: "client"}, "job-1", refunds) + if err != nil || again.Status != StatusCancelled || refunds.calls != 1 { + t.Fatalf("second Cancel = %#v, %v; refunds=%d", again, err, refunds.calls) + } + _, err = service.Cancel(context.Background(), Scope{OwnerID: "api:other", ExternalClientID: "other"}, "job-1", refunds) + var notFound *Error + if !errors.As(err, ¬Found) || notFound.Status != 404 { + t.Fatalf("cross-owner error = %#v", err) + } +} + +func TestWorkerSchedulesRetryReleasesRunningAndSettlesTerminal(t *testing.T) { + now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) + tests := []struct { + name string + job Job + want string + status Status + attempts int + scheduled time.Time + }{ + {name: "retryable failure", job: Job{ID: "retry", Status: StatusFailed, Error: &JobError{Message: "temporary", Retryable: true}, MaxAttempts: 3}, want: "retry_scheduled", status: StatusQueued, attempts: 1, scheduled: now.Add(10 * time.Second)}, + {name: "running release", job: Job{ID: "running", Status: StatusRunning}, want: "released", status: StatusRunning, scheduled: now.Add(5 * time.Second)}, + {name: "terminal failure", job: Job{ID: "terminal", Status: StatusFailed, Error: &JobError{Message: "fatal"}, MaxAttempts: 3}, want: "processed", status: StatusFailed, attempts: 1}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + store := newMemoryJobStore() + test.job.OwnerID = "owner" + test.job.Capability = "image.generate" + test.job.UpdatedAt = now + store.claimed = []Job{test.job} + store.jobs[test.job.ID] = test.job + processor := &processorStub{job: test.job} + refunds := &refundStub{} + worker := NewWorker(store, processor, refunds, nil, nil, WorkerConfig{}, func() time.Time { return now }) + result, err := worker.Tick(context.Background(), "worker-1") + if err != nil || len(result.Jobs) != 1 || result.Jobs[0].Action != test.want { + t.Fatalf("Tick = %#v, %v", result, err) + } + stored := store.jobs[test.job.ID] + if stored.Status != test.status || stored.Attempts != test.attempts { + t.Fatalf("stored = %#v, want status=%s attempts=%d", stored, test.status, test.attempts) + } + if !test.scheduled.IsZero() && !stored.ScheduledAt.Equal(test.scheduled) { + t.Fatalf("scheduled = %s, want %s", stored.ScheduledAt, test.scheduled) + } + if test.want == "processed" && refunds.calls != 1 { + t.Fatalf("refund calls = %d, want 1", refunds.calls) + } + }) + } +} + +func fixtureCreateCommand(prompt string) CreateCommand { + body := map[string]any{ + "capability": "image.generate", "prompt": prompt, + "settings": map[string]any{"height": float64(768), "width": float64(1024)}, + "idempotencyKey": "body-key", + } + return CreateCommand{Job: Job{ + ID: "job-1", OwnerID: "api:client", ExternalClientID: "client", Capability: "image.generate", + Provider: "mock", ReqKey: "fixture", Status: StatusQueued, IdempotencyKey: "idem-1", + RequestPayload: json.RawMessage(`{"prompt":"` + prompt + `"}`), + }, IdempotencyBody: body} +} + +type memoryJobStore struct { + jobs map[string]Job + claimed []Job + race *Job +} + +func newMemoryJobStore() *memoryJobStore { return &memoryJobStore{jobs: map[string]Job{}} } +func (store *memoryJobStore) ListJobs(context.Context, ListFilter) ([]Job, error) { return nil, nil } +func (store *memoryJobStore) FindJob(_ context.Context, scope Scope, id string) (Job, bool, error) { + job, ok := store.jobs[id] + return job, ok && scope.Owns(job), nil +} +func (store *memoryJobStore) FindIdempotentJob(_ context.Context, scope Scope, key string) (Job, bool, error) { + for _, job := range store.jobs { + if scope.Owns(job) && job.IdempotencyKey == key { + return job, true, nil + } + } + return Job{}, false, nil +} +func (store *memoryJobStore) CreateJob(_ context.Context, job Job) (Job, error) { + if store.race != nil { + store.jobs[store.race.ID] = *store.race + store.race = nil + return Job{}, ErrUniqueIdempotency + } + store.jobs[job.ID] = job + return job, nil +} +func (store *memoryJobStore) UpdateJob(_ context.Context, id string, patch Patch) (Job, error) { + job := store.jobs[id] + if patch.Status != nil { + job.Status = *patch.Status + } + if patch.Error != nil { + job.Error = patch.Error + } + if patch.Attempts != nil { + job.Attempts = *patch.Attempts + } + if patch.ScheduledAt != nil { + job.ScheduledAt = *patch.ScheduledAt + } + if patch.CompletedAt != nil { + job.CompletedAt = patch.CompletedAt + } + if patch.ClearLease { + job.LockedAt = nil + job.LockedBy = "" + } + if patch.ClearProviderTaskID { + job.ProviderTaskID = "" + } + if patch.WebhookAttempts != nil { + job.WebhookAttempts = *patch.WebhookAttempts + } + if patch.SetWebhookStatus { + job.WebhookLastStatus = patch.WebhookLastStatus + } + store.jobs[id] = job + return job, nil +} +func (store *memoryJobStore) ClaimJobs(context.Context, string, int, int) ([]Job, error) { + return store.claimed, nil +} + +type refundStub struct{ calls int } + +func (stub *refundStub) Refund(_ context.Context, job Job, _ string) (Job, error) { + stub.calls++ + return job, nil +} + +type processorStub struct { + job Job + err error +} + +func (stub *processorStub) Advance(context.Context, Job) (Job, error) { return stub.job, stub.err } diff --git a/backend/internal/jobs/service.go b/backend/internal/jobs/service.go new file mode 100644 index 0000000..a8fb6ba --- /dev/null +++ b/backend/internal/jobs/service.go @@ -0,0 +1,174 @@ +package jobs + +import ( + "context" + "fmt" + "time" +) + +type Store interface { + ListJobs(context.Context, ListFilter) ([]Job, error) + FindJob(context.Context, Scope, string) (Job, bool, error) + FindIdempotentJob(context.Context, Scope, string) (Job, bool, error) + CreateJob(context.Context, Job) (Job, error) + UpdateJob(context.Context, string, Patch) (Job, error) + ClaimJobs(context.Context, string, int, int) ([]Job, error) +} + +type Patch struct { + Status *Status + Error *JobError + Attempts *int + ScheduledAt *time.Time + CompletedAt *time.Time + ProviderTaskID *string + ClearProviderTaskID bool + ClearLease bool + WebhookAttempts *int + WebhookLastStatus []byte + SetWebhookStatus bool +} + +type Service struct { + store Store + now func() time.Time +} + +func NewService(store Store, now func() time.Time) *Service { + if now == nil { + now = time.Now + } + return &Service{store: store, now: now} +} + +func (service *Service) List(ctx context.Context, filter ListFilter) ([]Job, error) { + if filter.Scope.OwnerID == "" { + return nil, &Error{Kind: ErrorInvalid, Status: 400, Message: "owner is required"} + } + if filter.Status != "" && !filter.Status.Valid() { + return nil, &Error{Kind: ErrorInvalid, Status: 400, Message: fmt.Sprintf("Unsupported status filter: %s", filter.Status)} + } + if filter.Capability != "" && filter.Capability != "image.generate" && filter.Capability != "video.generate" { + return nil, &Error{Kind: ErrorInvalid, Status: 400, Message: fmt.Sprintf("Unsupported capability filter: %s", filter.Capability)} + } + filter.Limit = NormalizePublicLimit(filter.Limit) + return service.store.ListJobs(ctx, filter) +} + +func (service *Service) Get(ctx context.Context, scope Scope, id string) (Job, error) { + job, found, err := service.store.FindJob(ctx, scope, id) + if err != nil { + return Job{}, err + } + if !found { + return Job{}, &Error{Kind: ErrorNotFound, Status: 404, Message: "Job not found."} + } + return job, nil +} + +func (service *Service) Create(ctx context.Context, command CreateCommand) (Job, bool, error) { + job := command.Job + if job.OwnerID == "" || job.ID == "" || !job.Status.Valid() || job.Capability == "" || job.Provider == "" || job.ReqKey == "" { + return Job{}, false, &Error{Kind: ErrorInvalid, Status: 400, Message: "invalid generation job"} + } + job.Priority = NormalizePriority(job.Priority) + if job.MaxAttempts <= 0 { + job.MaxAttempts = 3 + } + now := service.now().UTC() + if job.CreatedAt.IsZero() { + job.CreatedAt = now + } + if job.UpdatedAt.IsZero() { + job.UpdatedAt = now + } + if job.ScheduledAt.IsZero() { + job.ScheduledAt = now + } + if job.InputAssetIDs == nil { + job.InputAssetIDs = []string{} + } + if job.InputURLs == nil { + job.InputURLs = []string{} + } + if job.OutputAssetIDs == nil { + job.OutputAssetIDs = []string{} + } + if len(job.RequestPayload) == 0 { + job.RequestPayload = []byte(`{}`) + } + + scope := Scope{OwnerID: job.OwnerID, ExternalClientID: job.ExternalClientID} + if job.IdempotencyKey != "" && job.ExternalClientID != "" { + fingerprint, err := Fingerprint(command.IdempotencyBody) + if err != nil { + return Job{}, false, err + } + job.IdempotencyFingerprint = fingerprint + if existing, found, err := service.store.FindIdempotentJob(ctx, scope, job.IdempotencyKey); err != nil { + return Job{}, false, err + } else if found { + return compareIdempotent(existing, fingerprint) + } + } + + created, err := service.store.CreateJob(ctx, job) + if err == ErrUniqueIdempotency && job.IdempotencyKey != "" { + existing, found, lookupErr := service.store.FindIdempotentJob(ctx, scope, job.IdempotencyKey) + if lookupErr != nil { + return Job{}, false, lookupErr + } + if found { + return compareIdempotent(existing, job.IdempotencyFingerprint) + } + } + return created, false, err +} + +func compareIdempotent(existing Job, fingerprint string) (Job, bool, error) { + if existing.IdempotencyFingerprint != fingerprint { + return Job{}, false, &Error{Kind: ErrorConflict, Status: 409, Message: "Idempotency key was already used with a different request body."} + } + return existing, true, nil +} + +type RefundPort interface { + Refund(context.Context, Job, string) (Job, error) +} + +func (service *Service) Cancel(ctx context.Context, scope Scope, id string, refunds RefundPort) (Job, error) { + job, err := service.Get(ctx, scope, id) + if err != nil || job.Status.Terminal() { + return job, err + } + now := service.now().UTC() + status := StatusCancelled + job, err = service.store.UpdateJob(ctx, job.ID, Patch{Status: &status, CompletedAt: &now}) + if err != nil { + return Job{}, err + } + if refunds != nil { + job, err = refunds.Refund(ctx, job, "任务已取消") + if err != nil { + return Job{}, err + } + } + return service.store.UpdateJob(ctx, job.ID, Patch{ClearLease: true}) +} + +func RetryDelay(attempts int, base, maximum time.Duration) time.Duration { + if base <= 0 { + base = 10 * time.Second + } + if maximum <= 0 { + maximum = 5 * time.Minute + } + delay := base + for index := 1; index < attempts && delay < maximum; index++ { + delay *= 2 + if delay > maximum { + return maximum + } + } + return delay +} diff --git a/backend/internal/jobs/worker.go b/backend/internal/jobs/worker.go new file mode 100644 index 0000000..131e106 --- /dev/null +++ b/backend/internal/jobs/worker.go @@ -0,0 +1,180 @@ +package jobs + +import ( + "context" + "encoding/json" + "fmt" + "time" +) + +type Processor interface { + Advance(context.Context, Job) (Job, error) +} + +type TerminalRefund interface { + Refund(context.Context, Job, string) (Job, error) +} + +type UsageRecorder interface { + Record(context.Context, Job) error +} + +type WebhookDelivery interface { + Deliver(context.Context, Job) (WebhookResult, error) +} + +type WebhookResult struct { + Attempts int + LastStatus any +} + +type WorkerConfig struct { + BatchSize int + LockTimeoutSeconds int + PollInterval time.Duration + RetryBase time.Duration + RetryMaximum time.Duration +} + +type Worker struct { + store Store + processor Processor + refunds TerminalRefund + usage UsageRecorder + webhooks WebhookDelivery + config WorkerConfig + now func() time.Time +} + +func NewWorker(store Store, processor Processor, refunds TerminalRefund, usage UsageRecorder, webhooks WebhookDelivery, config WorkerConfig, now func() time.Time) *Worker { + if now == nil { + now = time.Now + } + if config.BatchSize < 1 { + config.BatchSize = 3 + } + if config.BatchSize > 20 { + config.BatchSize = 20 + } + if config.LockTimeoutSeconds <= 0 { + config.LockTimeoutSeconds = 300 + } + if config.PollInterval <= 0 { + config.PollInterval = 5 * time.Second + } + return &Worker{store: store, processor: processor, refunds: refunds, usage: usage, webhooks: webhooks, config: config, now: now} +} + +type TickResult struct { + WorkerID string `json:"workerId"` + Claimed int `json:"claimed"` + Jobs []TickJob `json:"jobs"` +} + +type TickJob struct { + ID string `json:"id"` + Status Status `json:"status"` + Action string `json:"action"` + Error string `json:"error,omitempty"` +} + +func (worker *Worker) Tick(ctx context.Context, workerID string) (TickResult, error) { + claimed, err := worker.store.ClaimJobs(ctx, workerID, worker.config.BatchSize, worker.config.LockTimeoutSeconds) + if err != nil { + return TickResult{}, err + } + result := TickResult{WorkerID: workerID, Claimed: len(claimed), Jobs: make([]TickJob, 0, len(claimed))} + for _, job := range claimed { + advanced, advanceErr := worker.processor.Advance(ctx, job) + if advanceErr != nil { + status := StatusFailed + jobError := &JobError{Message: advanceErr.Error(), Retryable: true} + advanced, err = worker.store.UpdateJob(ctx, job.ID, Patch{Status: &status, Error: jobError}) + if err != nil { + return result, err + } + } + settled, action, err := worker.settle(ctx, advanced) + if err != nil { + return result, err + } + item := TickJob{ID: settled.ID, Status: settled.Status, Action: action} + if advanceErr != nil { + item.Action = "failed" + item.Error = advanceErr.Error() + } + result.Jobs = append(result.Jobs, item) + } + return result, nil +} + +func (worker *Worker) settle(ctx context.Context, job Job) (Job, string, error) { + now := worker.now().UTC() + if job.Status == StatusFailed && job.Error != nil && job.Error.Retryable && job.Attempts < maxAttempts(job) { + attempts := job.Attempts + 1 + scheduled := now.Add(RetryDelay(attempts, worker.config.RetryBase, worker.config.RetryMaximum)) + status := StatusQueued + returnPatch := Patch{Status: &status, Attempts: &attempts, ScheduledAt: &scheduled, ClearProviderTaskID: true, ClearLease: true} + retried, err := worker.store.UpdateJob(ctx, job.ID, returnPatch) + return retried, "retry_scheduled", err + } + + if !job.Status.Terminal() { + scheduled := now.Add(worker.config.PollInterval) + released, err := worker.store.UpdateJob(ctx, job.ID, Patch{ScheduledAt: &scheduled, ClearLease: true}) + return released, "released", err + } + + if job.Status != StatusSucceeded && worker.refunds != nil { + var err error + job, err = worker.refunds.Refund(ctx, job, terminalReason(job)) + if err != nil { + return Job{}, "", err + } + } + if job.Status == StatusSucceeded && worker.usage != nil { + if err := worker.usage.Record(ctx, job); err != nil { + return Job{}, "", err + } + } + attempts := job.Attempts + if job.Status == StatusFailed { + attempts++ + } + completed := now + job, err := worker.store.UpdateJob(ctx, job.ID, Patch{Attempts: &attempts, CompletedAt: &completed, ClearLease: true}) + if err != nil { + return Job{}, "", err + } + if worker.webhooks != nil { + delivery, err := worker.webhooks.Deliver(ctx, job) + if err != nil { + return Job{}, "", err + } + if delivery.LastStatus != nil { + lastStatus, err := json.Marshal(delivery.LastStatus) + if err != nil { + return Job{}, "", fmt.Errorf("encode webhook last status: %w", err) + } + job, err = worker.store.UpdateJob(ctx, job.ID, Patch{WebhookAttempts: &delivery.Attempts, WebhookLastStatus: lastStatus, SetWebhookStatus: true}) + if err != nil { + return Job{}, "", err + } + } + } + return job, "processed", nil +} + +func maxAttempts(job Job) int { + if job.MaxAttempts > 0 { + return job.MaxAttempts + } + return 3 +} + +func terminalReason(job Job) string { + if job.Error != nil && job.Error.Message != "" { + return job.Error.Message + } + return "任务" + string(job.Status) +} diff --git a/backend/internal/postgres/jobs.go b/backend/internal/postgres/jobs.go new file mode 100644 index 0000000..698b2f4 --- /dev/null +++ b/backend/internal/postgres/jobs.go @@ -0,0 +1,347 @@ +package postgres + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "strings" + "time" + + "git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs" +) + +const jobColumns = `id, owner_id, external_client_id, capability, provider, req_key, status, prompt, +input_asset_ids, input_urls, output_asset_ids, provider_task_id, request_payload, response_payload, +error, retry_of, idempotency_key, idempotency_fingerprint, priority, attempts, max_attempts, +scheduled_at, locked_at, locked_by, started_at, completed_at, webhook_url, webhook_attempts, +webhook_last_status, usage_context, billing, created_at, updated_at` + +const ClaimJobsSQL = `SELECT ` + jobColumns + ` FROM public.claim_generation_jobs($1::text, $2::integer, $3::integer)` + +func (db *Database) ListJobs(ctx context.Context, filter jobs.ListFilter) ([]jobs.Job, error) { + if err := db.requirePostgres("list generation jobs"); err != nil { + return nil, err + } + clauses := []string{"owner_id = $1::text"} + args := []any{filter.Scope.OwnerID} + if filter.Scope.ExternalClientID != "" { + args = append(args, filter.Scope.ExternalClientID) + clauses = append(clauses, fmt.Sprintf("external_client_id = $%d::text", len(args))) + } + if filter.Status != "" { + args = append(args, string(filter.Status)) + clauses = append(clauses, fmt.Sprintf("status = $%d::text", len(args))) + } + if filter.Capability != "" { + args = append(args, filter.Capability) + clauses = append(clauses, fmt.Sprintf("capability = $%d::text", len(args))) + } + if filter.Before != nil { + args = append(args, *filter.Before) + clauses = append(clauses, fmt.Sprintf("created_at < $%d::timestamptz", len(args))) + } + args = append(args, filter.Limit) + query := `SELECT ` + jobColumns + ` FROM public.generation_jobs WHERE ` + strings.Join(clauses, " AND ") + + fmt.Sprintf(" ORDER BY created_at DESC LIMIT $%d::integer", len(args)) + return db.queryJobs(ctx, query, args...) +} + +func (db *Database) FindJob(ctx context.Context, scope jobs.Scope, id string) (jobs.Job, bool, error) { + if err := db.requirePostgres("find generation job"); err != nil { + return jobs.Job{}, false, err + } + query := `SELECT ` + jobColumns + ` FROM public.generation_jobs WHERE id = $1::text AND owner_id = $2::text` + args := []any{id, scope.OwnerID} + if scope.ExternalClientID != "" { + query += ` AND external_client_id = $3::text` + args = append(args, scope.ExternalClientID) + } + query += ` LIMIT 1` + return db.queryOneJob(ctx, query, args...) +} + +func (db *Database) FindIdempotentJob(ctx context.Context, scope jobs.Scope, key string) (jobs.Job, bool, error) { + if err := db.requirePostgres("find idempotent generation job"); err != nil { + return jobs.Job{}, false, err + } + query := `SELECT ` + jobColumns + ` FROM public.generation_jobs +WHERE owner_id = $1::text AND external_client_id = $2::text AND idempotency_key = $3::text LIMIT 1` + return db.queryOneJob(ctx, query, scope.OwnerID, scope.ExternalClientID, key) +} + +func (db *Database) CreateJob(ctx context.Context, job jobs.Job) (jobs.Job, error) { + if err := db.requirePostgres("create generation job"); err != nil { + return jobs.Job{}, err + } + const columns = `id, owner_id, external_client_id, capability, provider, req_key, status, prompt, +input_asset_ids, input_urls, output_asset_ids, provider_task_id, request_payload, response_payload, +error, retry_of, idempotency_key, idempotency_fingerprint, priority, attempts, max_attempts, +scheduled_at, locked_at, locked_by, started_at, completed_at, webhook_url, webhook_attempts, +webhook_last_status, usage_context, billing, created_at, updated_at` + placeholders := make([]string, 33) + for index := range placeholders { + placeholders[index] = fmt.Sprintf("$%d", index+1) + } + query := `INSERT INTO public.generation_jobs (` + columns + `) VALUES (` + strings.Join(placeholders, ", ") + `) RETURNING ` + jobColumns + row, found, err := db.queryOneJob(ctx, query, jobArguments(job)...) + if err != nil { + if sqlState(err) == "23505" && job.ExternalClientID != "" && job.IdempotencyKey != "" { + return jobs.Job{}, jobs.ErrUniqueIdempotency + } + return jobs.Job{}, fmt.Errorf("create generation job: %w", err) + } + if !found { + return jobs.Job{}, fmt.Errorf("create generation job returned no row") + } + return row, nil +} + +func (db *Database) UpdateJob(ctx context.Context, id string, patch jobs.Patch) (jobs.Job, error) { + if err := db.requirePostgres("update generation job"); err != nil { + return jobs.Job{}, err + } + sets := make([]string, 0, 12) + args := []any{id} + add := func(column, cast string, value any) { + args = append(args, value) + sets = append(sets, fmt.Sprintf("%s = $%d%s", column, len(args), cast)) + } + if patch.Status != nil { + add("status", "::text", string(*patch.Status)) + } + if patch.Error != nil { + raw, err := json.Marshal(patch.Error) + if err != nil { + return jobs.Job{}, fmt.Errorf("encode generation job error: %w", err) + } + add("error", "::jsonb", raw) + } + if patch.Attempts != nil { + add("attempts", "::integer", *patch.Attempts) + } + if patch.ScheduledAt != nil { + add("scheduled_at", "::timestamptz", *patch.ScheduledAt) + } + if patch.CompletedAt != nil { + add("completed_at", "::timestamptz", *patch.CompletedAt) + } + if patch.ProviderTaskID != nil { + add("provider_task_id", "::text", *patch.ProviderTaskID) + } + if patch.ClearProviderTaskID { + sets = append(sets, "provider_task_id = NULL") + } + if patch.ClearLease { + sets = append(sets, "locked_at = NULL", "locked_by = NULL") + } + if patch.WebhookAttempts != nil { + add("webhook_attempts", "::integer", *patch.WebhookAttempts) + } + if patch.SetWebhookStatus { + add("webhook_last_status", "::jsonb", patch.WebhookLastStatus) + } + if len(sets) == 0 { + return db.mustFindJobByID(ctx, id) + } + sets = append(sets, "updated_at = now()") + query := `UPDATE public.generation_jobs SET ` + strings.Join(sets, ", ") + ` WHERE id = $1::text RETURNING ` + jobColumns + job, found, err := db.queryOneJob(ctx, query, args...) + if err != nil { + return jobs.Job{}, fmt.Errorf("update generation job: %w", err) + } + if !found { + return jobs.Job{}, fmt.Errorf("generation job not found: %s", id) + } + return job, nil +} + +func (db *Database) ClaimJobs(ctx context.Context, workerID string, limit, lockTimeoutSeconds int) ([]jobs.Job, error) { + if err := db.requirePostgres("claim generation jobs"); err != nil { + return nil, err + } + limit = max(1, min(limit, 20)) + if lockTimeoutSeconds <= 0 { + lockTimeoutSeconds = 300 + } + return db.queryJobs(ctx, ClaimJobsSQL, workerID, limit, lockTimeoutSeconds) +} + +func (db *Database) mustFindJobByID(ctx context.Context, id string) (jobs.Job, error) { + job, found, err := db.queryOneJob(ctx, `SELECT `+jobColumns+` FROM public.generation_jobs WHERE id = $1::text LIMIT 1`, id) + if err != nil { + return jobs.Job{}, err + } + if !found { + return jobs.Job{}, fmt.Errorf("generation job not found: %s", id) + } + return job, nil +} + +func (db *Database) queryJobs(ctx context.Context, query string, args ...any) ([]jobs.Job, error) { + rows, err := db.querier.Query(ctx, query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + items := make([]jobs.Job, 0) + for rows.Next() { + item, err := scanJob(rows) + if err != nil { + return nil, err + } + items = append(items, item) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +func (db *Database) queryOneJob(ctx context.Context, query string, args ...any) (jobs.Job, bool, error) { + rows, err := db.querier.Query(ctx, query, args...) + if err != nil { + return jobs.Job{}, false, err + } + defer rows.Close() + if !rows.Next() { + return jobs.Job{}, false, rows.Err() + } + job, err := scanJob(rows) + if err != nil { + return jobs.Job{}, false, err + } + if err := rows.Err(); err != nil { + return jobs.Job{}, false, err + } + return job, true, nil +} + +func scanJob(rows Rows) (jobs.Job, error) { + var job jobs.Job + var externalClientID, prompt, providerTaskID, retryOf, idempotencyKey, fingerprint sql.NullString + var lockedBy, webhookURL sql.NullString + var lockedAt, startedAt, completedAt sql.NullTime + var requestPayload, responsePayload, errorPayload, webhookStatus, usageContext, billing []byte + var status string + err := rows.Scan( + &job.ID, &job.OwnerID, &externalClientID, &job.Capability, &job.Provider, &job.ReqKey, &status, &prompt, + &job.InputAssetIDs, &job.InputURLs, &job.OutputAssetIDs, &providerTaskID, &requestPayload, &responsePayload, + &errorPayload, &retryOf, &idempotencyKey, &fingerprint, &job.Priority, &job.Attempts, &job.MaxAttempts, + &job.ScheduledAt, &lockedAt, &lockedBy, &startedAt, &completedAt, &webhookURL, &job.WebhookAttempts, + &webhookStatus, &usageContext, &billing, &job.CreatedAt, &job.UpdatedAt, + ) + if err != nil { + return jobs.Job{}, fmt.Errorf("scan generation job: %w", err) + } + job.Status = jobs.Status(status) + job.ExternalClientID = nullString(externalClientID) + job.Prompt = nullString(prompt) + job.ProviderTaskID = nullString(providerTaskID) + job.RetryOf = nullString(retryOf) + job.IdempotencyKey = nullString(idempotencyKey) + job.IdempotencyFingerprint = nullString(fingerprint) + job.LockedBy = nullString(lockedBy) + job.WebhookURL = nullString(webhookURL) + job.LockedAt = nullTimePointer(lockedAt) + job.StartedAt = nullTimePointer(startedAt) + job.CompletedAt = nullTimePointer(completedAt) + job.RequestPayload = normalizedJSON(requestPayload, `{}`) + job.ResponsePayload = normalizedOptionalJSON(responsePayload) + job.WebhookLastStatus = normalizedOptionalJSON(webhookStatus) + job.UsageContext = normalizedOptionalJSON(usageContext) + job.Billing = normalizedOptionalJSON(billing) + if len(errorPayload) > 0 && string(errorPayload) != "null" { + if err := json.Unmarshal(errorPayload, &job.Error); err != nil { + return jobs.Job{}, fmt.Errorf("decode generation job error: %w", err) + } + } + if job.InputAssetIDs == nil { + job.InputAssetIDs = []string{} + } + if job.InputURLs == nil { + job.InputURLs = []string{} + } + if job.OutputAssetIDs == nil { + job.OutputAssetIDs = []string{} + } + return job, nil +} + +func jobArguments(job jobs.Job) []any { + errorPayload, _ := json.Marshal(job.Error) + if job.Error == nil { + errorPayload = nil + } + return []any{ + job.ID, job.OwnerID, optionalDatabaseText(job.ExternalClientID), job.Capability, job.Provider, job.ReqKey, + string(job.Status), optionalDatabaseText(job.Prompt), job.InputAssetIDs, job.InputURLs, job.OutputAssetIDs, + optionalDatabaseText(job.ProviderTaskID), job.RequestPayload, optionalJSON(job.ResponsePayload), optionalJSON(errorPayload), + optionalDatabaseText(job.RetryOf), optionalDatabaseText(job.IdempotencyKey), optionalDatabaseText(job.IdempotencyFingerprint), + job.Priority, job.Attempts, job.MaxAttempts, job.ScheduledAt, job.LockedAt, optionalDatabaseText(job.LockedBy), + job.StartedAt, job.CompletedAt, optionalDatabaseText(job.WebhookURL), job.WebhookAttempts, + optionalJSON(job.WebhookLastStatus), optionalJSON(job.UsageContext), optionalJSON(job.Billing), job.CreatedAt, job.UpdatedAt, + } +} + +func (db *Database) requirePostgres(operation string) error { + if db == nil || db.config.Backend != BackendPostgres || db.querier == nil { + return fmt.Errorf("%s: PostgreSQL is unavailable", operation) + } + return nil +} + +func nullString(value sql.NullString) string { + if value.Valid { + return value.String + } + return "" +} + +func nullTimePointer(value sql.NullTime) *time.Time { + if !value.Valid { + return nil + } + copy := value.Time + return © +} + +func optionalJSON(value []byte) any { + if len(value) == 0 || string(value) == "null" { + return nil + } + return value +} + +func normalizedJSON(value []byte, fallback string) json.RawMessage { + if len(value) == 0 || string(value) == "null" { + return json.RawMessage(fallback) + } + return append(json.RawMessage(nil), value...) +} + +func normalizedOptionalJSON(value []byte) json.RawMessage { + if len(value) == 0 || string(value) == "null" { + return nil + } + return append(json.RawMessage(nil), value...) +} + +type sqlStateError interface{ SQLState() string } + +func sqlState(err error) string { + for err != nil { + if state, ok := err.(sqlStateError); ok { + return state.SQLState() + } + type unwrapper interface{ Unwrap() error } + wrapped, ok := err.(unwrapper) + if !ok { + break + } + err = wrapped.Unwrap() + } + return "" +} + +var _ jobs.Store = (*Database)(nil) diff --git a/backend/internal/postgres/jobs_test.go b/backend/internal/postgres/jobs_test.go new file mode 100644 index 0000000..44134c9 --- /dev/null +++ b/backend/internal/postgres/jobs_test.go @@ -0,0 +1,165 @@ +package postgres + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "reflect" + "strings" + "testing" + "time" + + "git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs" +) + +func TestJobsAdapterUsesFullClaimFunctionAndBoundsArguments(t *testing.T) { + now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) + querier := &jobQuerier{rows: &jobRows{rows: [][]any{jobRow(now)}}} + database := NewDatabase(Config{Backend: BackendPostgres}, querier) + + claimed, err := database.ClaimJobs(context.Background(), "worker-1", 99, 0) + if err != nil { + t.Fatal(err) + } + if querier.query != ClaimJobsSQL || !reflect.DeepEqual(querier.args, []any{"worker-1", 20, 300}) { + t.Fatalf("query=%q args=%#v", querier.query, querier.args) + } + if len(claimed) != 1 || claimed[0].ID != "job-1" || claimed[0].Status != jobs.StatusQueued || claimed[0].RequestPayload == nil { + t.Fatalf("claimed=%#v", claimed) + } + if strings.Contains(ClaimJobsSQL, "SELECT id FROM") || !strings.Contains(ClaimJobsSQL, "claim_generation_jobs") { + t.Fatalf("ClaimJobsSQL does not select the complete contract: %s", ClaimJobsSQL) + } +} + +func TestJobsAdapterScopesPublicFindAndListByOwnerAndClient(t *testing.T) { + now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) + querier := &jobQuerier{rows: &jobRows{rows: [][]any{jobRow(now)}}} + database := NewDatabase(Config{Backend: BackendPostgres}, querier) + scope := jobs.Scope{OwnerID: "api:client", ExternalClientID: "client"} + + _, found, err := database.FindJob(context.Background(), scope, "job-1") + if err != nil || !found { + t.Fatalf("FindJob found=%v err=%v", found, err) + } + if !strings.Contains(querier.query, "owner_id = $2::text") || !strings.Contains(querier.query, "external_client_id = $3::text") || !reflect.DeepEqual(querier.args, []any{"job-1", "api:client", "client"}) { + t.Fatalf("query=%q args=%#v", querier.query, querier.args) + } + + querier.rows = &jobRows{} + _, err = database.ListJobs(context.Background(), jobs.ListFilter{Scope: scope, Status: jobs.StatusRunning, Capability: "video.generate", Limit: 50}) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(querier.query, "owner_id = $1::text") || !strings.Contains(querier.query, "external_client_id = $2::text") || !strings.Contains(querier.query, "status = $3::text") || !strings.Contains(querier.query, "capability = $4::text") { + t.Fatalf("list query=%q", querier.query) + } +} + +func TestJobsAdapterUpdateClearsLeaseAndProviderTask(t *testing.T) { + now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) + row := jobRow(now) + row[11] = nil + row[22], row[23] = nil, nil + querier := &jobQuerier{rows: &jobRows{rows: [][]any{row}}} + database := NewDatabase(Config{Backend: BackendPostgres}, querier) + status := jobs.StatusQueued + attempts := 2 + + _, err := database.UpdateJob(context.Background(), "job-1", jobs.Patch{Status: &status, Attempts: &attempts, ClearLease: true, ClearProviderTaskID: true}) + if err != nil { + t.Fatal(err) + } + for _, fragment := range []string{"status = $2::text", "attempts = $3::integer", "provider_task_id = NULL", "locked_at = NULL", "locked_by = NULL"} { + if !strings.Contains(querier.query, fragment) { + t.Fatalf("update query missing %q: %s", fragment, querier.query) + } + } +} + +func TestJobsAdapterFailsClosedWithoutPostgres(t *testing.T) { + database := NewDatabase(Config{Backend: BackendLocal}, nil) + if _, err := database.ClaimJobs(context.Background(), "worker", 1, 300); err == nil { + t.Fatal("ClaimJobs error=nil") + } + if _, found, err := database.FindJob(context.Background(), jobs.Scope{OwnerID: "o"}, "j"); err == nil || found { + t.Fatalf("FindJob found=%v err=%v", found, err) + } +} + +func jobRow(now time.Time) []any { + return []any{ + "job-1", "api:client", "client", "image.generate", "mock", "fixture", "queued", "prompt", + []string{}, []string{}, []string{}, nil, []byte(`{"input":true}`), nil, nil, nil, + "idem-1", "fingerprint", 5, 0, 3, now, nil, nil, nil, nil, "https://hooks.example.test", 0, + nil, []byte(`{"source":"api","accountId":"client","displayName":"client"}`), nil, now, now, + } +} + +type jobQuerier struct { + rows *jobRows + err error + query string + args []any +} + +func (querier *jobQuerier) Query(_ context.Context, query string, args ...any) (Rows, error) { + querier.query, querier.args = query, args + return querier.rows, querier.err +} + +type jobRows struct { + rows [][]any + index int + err error +} + +func (rows *jobRows) Close() {} +func (rows *jobRows) Err() error { return rows.err } +func (rows *jobRows) Next() bool { return rows.index < len(rows.rows) } +func (rows *jobRows) Scan(dest ...any) error { + if rows.index >= len(rows.rows) { + return errors.New("scan past end") + } + row := rows.rows[rows.index] + rows.index++ + if len(row) != len(dest) { + return errors.New("scan arity mismatch") + } + for index, value := range row { + switch target := dest[index].(type) { + case *string: + if value != nil { + *target = value.(string) + } + case *int: + *target = value.(int) + case *time.Time: + *target = value.(time.Time) + case *sql.NullString: + if value != nil { + *target = sql.NullString{String: value.(string), Valid: true} + } + case *sql.NullTime: + if value != nil { + *target = sql.NullTime{Time: value.(time.Time), Valid: true} + } + case *[]string: + *target = append([]string(nil), value.([]string)...) + case *[]byte: + if value != nil { + *target = append([]byte(nil), value.([]byte)...) + } + default: + return errors.New("unsupported scan target") + } + } + return nil +} + +func TestJobFixtureJSONIsValid(t *testing.T) { + if !json.Valid([]byte(`{"status":"queued"}`)) { + t.Fatal("impossible") + } +} diff --git a/backend/internal/webhook/webhook.go b/backend/internal/webhook/webhook.go new file mode 100644 index 0000000..75953bd --- /dev/null +++ b/backend/internal/webhook/webhook.go @@ -0,0 +1,125 @@ +// Package webhook owns the generation-job callback byte and retry contract. +package webhook + +import ( + "context" + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "strings" + "time" + + "git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs" +) + +const ( + MaximumAttempts = 3 + UserAgent = "zhinian-aigc-webhook/1.0" +) + +type Payload struct { + JobID string `json:"jobId"` + Status jobs.Status `json:"status"` + Capability string `json:"capability"` + OutputAssetIDs []string `json:"outputAssetIds"` + Error *jobs.JobError `json:"error,omitempty"` + UpdatedAt string `json:"updatedAt"` +} + +func Body(job jobs.Job) ([]byte, error) { + outputIDs := job.OutputAssetIDs + if outputIDs == nil { + outputIDs = []string{} + } + payload := Payload{ + JobID: job.ID, Status: job.Status, Capability: job.Capability, + OutputAssetIDs: outputIDs, Error: job.Error, + // JavaScript Date#toISOString always emits exactly three fractional + // digits. Keep those bytes stable because the HMAC covers them. + UpdatedAt: job.UpdatedAt.UTC().Format("2006-01-02T15:04:05.000Z"), + } + return json.Marshal(payload) +} + +func Sign(body []byte, secret string) string { + secret = strings.TrimSpace(secret) + if secret == "" { + return "" + } + mac := hmac.New(sha256.New, []byte(secret)) + _, _ = mac.Write(body) + return "sha256=" + hex.EncodeToString(mac.Sum(nil)) +} + +type Request struct { + URL string + Body []byte + Headers map[string]string +} + +type Response struct { + Status int +} + +type Sender interface { + Send(context.Context, Request) (Response, error) +} + +type LastStatus struct { + OK bool `json:"ok"` + Status int `json:"status,omitempty"` + Error string `json:"error,omitempty"` + AttemptedAt string `json:"attemptedAt"` +} + +type Result struct { + Attempts int + LastStatus *LastStatus +} + +type Deliverer struct { + sender Sender + secret string + now func() time.Time +} + +func NewDeliverer(sender Sender, secret string, now func() time.Time) *Deliverer { + if now == nil { + now = time.Now + } + return &Deliverer{sender: sender, secret: secret, now: now} +} + +func (deliverer *Deliverer) Deliver(ctx context.Context, job jobs.Job) (Result, error) { + if job.WebhookURL == "" { + return Result{Attempts: job.WebhookAttempts}, nil + } + if deliverer == nil || deliverer.sender == nil { + return Result{}, fmt.Errorf("webhook sender is not configured") + } + body, err := Body(job) + if err != nil { + return Result{}, err + } + headers := map[string]string{"Content-Type": "application/json", "User-Agent": UserAgent} + if signature := Sign(body, deliverer.secret); signature != "" { + headers["X-Zhinian-Signature"] = signature + } + result := Result{Attempts: job.WebhookAttempts} + for result.Attempts < MaximumAttempts { + result.Attempts++ + attemptedAt := deliverer.now().UTC().Format(time.RFC3339Nano) + response, sendErr := deliverer.sender.Send(ctx, Request{URL: job.WebhookURL, Body: body, Headers: headers}) + if sendErr != nil { + result.LastStatus = &LastStatus{OK: false, Error: sendErr.Error(), AttemptedAt: attemptedAt} + continue + } + result.LastStatus = &LastStatus{OK: response.Status >= 200 && response.Status < 300, Status: response.Status, AttemptedAt: attemptedAt} + if result.LastStatus.OK { + break + } + } + return result, nil +} diff --git a/backend/internal/webhook/webhook_test.go b/backend/internal/webhook/webhook_test.go new file mode 100644 index 0000000..9d1b82c --- /dev/null +++ b/backend/internal/webhook/webhook_test.go @@ -0,0 +1,103 @@ +package webhook + +import ( + "context" + "encoding/json" + "errors" + "os" + "path/filepath" + "runtime" + "testing" + "time" + + "git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs" +) + +func TestBodyAndSignatureMatchSharedFixture(t *testing.T) { + fixture := loadWebhookFixture(t) + updatedAt, err := time.Parse(time.RFC3339Nano, fixture.Job.UpdatedAt) + if err != nil { + t.Fatal(err) + } + job := jobs.Job{ID: fixture.Job.ID, Status: jobs.Status(fixture.Job.Status), Capability: fixture.Job.Capability, OutputAssetIDs: fixture.Job.OutputAssetIDs, Error: fixture.Job.Error, UpdatedAt: updatedAt} + body, err := Body(job) + if err != nil { + t.Fatal(err) + } + if string(body) != fixture.Body { + t.Fatalf("body = %s, want %s", body, fixture.Body) + } + if got := Sign(body, fixture.Secret); got != fixture.Signature { + t.Fatalf("signature = %q, want %q", got, fixture.Signature) + } +} + +func TestDeliverRetriesUntilSuccessAndPreservesExactRequest(t *testing.T) { + clock := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) + sender := &senderStub{responses: []sendResult{{response: Response{Status: 503}}, {err: errors.New("network")}, {response: Response{Status: 204}}}} + deliverer := NewDeliverer(sender, " secret ", func() time.Time { return clock }) + job := jobs.Job{ID: "job", Status: jobs.StatusSucceeded, Capability: "image.generate", OutputAssetIDs: []string{}, UpdatedAt: clock, WebhookURL: "https://hooks.example.test/job"} + + result, err := deliverer.Deliver(context.Background(), job) + if err != nil || result.Attempts != 3 || result.LastStatus == nil || !result.LastStatus.OK || result.LastStatus.Status != 204 { + t.Fatalf("Deliver = %#v, %v", result, err) + } + if len(sender.requests) != 3 || sender.requests[0].Headers["Content-Type"] != "application/json" || sender.requests[0].Headers["User-Agent"] != UserAgent || sender.requests[0].Headers["X-Zhinian-Signature"] == "" { + t.Fatalf("requests = %#v", sender.requests) + } +} + +func TestDeliverRespectsPersistedAttemptCeiling(t *testing.T) { + sender := &senderStub{} + deliverer := NewDeliverer(sender, "", nil) + result, err := deliverer.Deliver(context.Background(), jobs.Job{WebhookURL: "https://hooks.example.test", WebhookAttempts: 3}) + if err != nil || result.Attempts != 3 || len(sender.requests) != 0 { + t.Fatalf("Deliver = %#v, %v requests=%d", result, err, len(sender.requests)) + } +} + +type sendResult struct { + response Response + err error +} +type senderStub struct { + responses []sendResult + requests []Request +} + +func (stub *senderStub) Send(_ context.Context, request Request) (Response, error) { + stub.requests = append(stub.requests, request) + result := stub.responses[len(stub.requests)-1] + return result.response, result.err +} + +type webhookFixture struct { + Secret string `json:"secret"` + Body string `json:"body"` + Signature string `json:"signature"` + Job struct { + ID string `json:"id"` + Status string `json:"status"` + Capability string `json:"capability"` + OutputAssetIDs []string `json:"outputAssetIds"` + Error *jobs.JobError `json:"error"` + UpdatedAt string `json:"updatedAt"` + } `json:"job"` +} + +func loadWebhookFixture(t *testing.T) webhookFixture { + t.Helper() + _, filename, _, ok := runtime.Caller(0) + if !ok { + t.Fatal("locate fixture") + } + raw, err := os.ReadFile(filepath.Join(filepath.Dir(filename), "..", "..", "..", "contracts", "webhook", "job-webhook-v1.json")) + if err != nil { + t.Fatal(err) + } + var fixture webhookFixture + if err := json.Unmarshal(raw, &fixture); err != nil { + t.Fatal(err) + } + return fixture +} diff --git a/contracts/jobs/jobs-v1.json b/contracts/jobs/jobs-v1.json new file mode 100644 index 0000000..451e2c9 --- /dev/null +++ b/contracts/jobs/jobs-v1.json @@ -0,0 +1,44 @@ +{ + "version": 1, + "statuses": ["queued", "running", "succeeded", "failed", "expired", "cancelled"], + "terminalStatuses": ["succeeded", "failed", "expired", "cancelled"], + "idempotency": { + "scope": ["ownerId", "externalClientId", "idempotencyKey"], + "body": { + "settings": { "width": 1024, "height": 768 }, + "prompt": "hello", + "idempotencyKey": "body-key", + "capability": "image.generate" + }, + "fingerprintSource": { + "capability": "image.generate", + "prompt": "hello", + "settings": { "height": 768, "width": 1024 } + }, + "fingerprint": "3ac3128766c0bdac36627b996cb5b3b1ed752794989d559c24f6621988e51a03", + "conflictStatus": 409, + "reusedStatus": 200, + "createdStatus": 202 + }, + "priority": { "minimum": -100, "maximum": 100, "default": 0 }, + "publicList": { "defaultLimit": 50, "minimumLimit": 1, "maximumLimit": 200 }, + "claim": { "minimumBatch": 1, "maximumBatch": 20, "defaultLockTimeoutSeconds": 300 }, + "retry": { + "defaultMaxAttempts": 3, + "baseDelayMilliseconds": 10000, + "maximumDelayMilliseconds": 300000, + "cases": [ + { "attemptsBefore": 0, "retryable": true, "delayMilliseconds": 10000, "action": "retry_scheduled" }, + { "attemptsBefore": 1, "retryable": true, "delayMilliseconds": 20000, "action": "retry_scheduled" }, + { "attemptsBefore": 2, "retryable": true, "delayMilliseconds": 40000, "action": "retry_scheduled" }, + { "attemptsBefore": 3, "retryable": true, "action": "processed" }, + { "attemptsBefore": 0, "retryable": false, "action": "processed" } + ] + }, + "cancel": { + "idempotentTerminal": true, + "activeResult": "cancelled", + "clearsLease": true, + "refundKey": "job-refund:{jobId}" + } +} diff --git a/contracts/webhook/job-webhook-v1.json b/contracts/webhook/job-webhook-v1.json new file mode 100644 index 0000000..bd41e79 --- /dev/null +++ b/contracts/webhook/job-webhook-v1.json @@ -0,0 +1,18 @@ +{ + "version": 1, + "maximumAttempts": 3, + "contentType": "application/json", + "userAgent": "zhinian-aigc-webhook/1.0", + "signatureHeader": "X-Zhinian-Signature", + "secret": "webhook-secret", + "job": { + "id": "job-1", + "status": "failed", + "capability": "image.generate", + "outputAssetIds": [], + "error": { "message": "provider unavailable", "retryable": true }, + "updatedAt": "2026-08-13T08:00:00.000Z" + }, + "body": "{\"jobId\":\"job-1\",\"status\":\"failed\",\"capability\":\"image.generate\",\"outputAssetIds\":[],\"error\":{\"message\":\"provider unavailable\",\"retryable\":true},\"updatedAt\":\"2026-08-13T08:00:00.000Z\"}", + "signature": "sha256=32b8fd385692f6f549c0b1e1ffc9dce404c7a3a9a0aa19ca45225ad257d94411" +} diff --git a/tests/jobs-go-contract.test.ts b/tests/jobs-go-contract.test.ts new file mode 100644 index 0000000..6104b4e --- /dev/null +++ b/tests/jobs-go-contract.test.ts @@ -0,0 +1,55 @@ +import { mkdtemp, readFile, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +import { afterEach, describe, expect, it } from "vitest"; + +import { createPublicGenerationJob, PublicApiConflictError } from "@/lib/server/public-api-jobs"; + +const fixtureUrl = new URL("../contracts/jobs/jobs-v1.json", import.meta.url); +let runtimeDir: string | undefined; + +afterEach(async () => { + if (runtimeDir) await rm(runtimeDir, { recursive: true, force: true }); + runtimeDir = undefined; +}); + +describe("jobs v1 cross-language contract", () => { + it("freezes statuses, bounds and retry behavior", async () => { + const fixture = JSON.parse(await readFile(fixtureUrl, "utf8")); + expect(fixture.version).toBe(1); + expect(fixture.statuses).toEqual(["queued", "running", "succeeded", "failed", "expired", "cancelled"]); + expect(fixture.terminalStatuses).toEqual(["succeeded", "failed", "expired", "cancelled"]); + expect(fixture.priority).toEqual({ minimum: -100, maximum: 100, default: 0 }); + expect(fixture.claim).toEqual({ minimumBatch: 1, maximumBatch: 20, defaultLockTimeoutSeconds: 300 }); + }); + + it("uses stable request fingerprints and rejects payload drift", async () => { + const fixture = JSON.parse(await readFile(fixtureUrl, "utf8")); + runtimeDir = await mkdtemp(join(tmpdir(), "zhinian-jobs-contract-")); + process.env.ZHINIAN_RUNTIME_DIR = runtimeDir; + process.env.ZHINIAN_DATA_BACKEND = "local"; + process.env.JIMENG_VISUAL_MOCK = "true"; + const request = new Request("http://local.test/api/v1/jobs", { + headers: { "Idempotency-Key": "header-key" }, + }); + const input = { + client: { id: "fixture-client", key: "secret" }, + request, + origin: "http://local.test", + body: fixture.idempotency.body, + }; + + const created = await createPublicGenerationJob(input); + expect(created.reused).toBe(false); + expect(created.job.idempotencyKey).toBe("header-key"); + expect(created.job.idempotencyFingerprint).toBe(fixture.idempotency.fingerprint); + const replay = await createPublicGenerationJob(input); + expect(replay.reused).toBe(true); + expect(replay.job.id).toBe(created.job.id); + await expect(createPublicGenerationJob({ + ...input, + body: { ...fixture.idempotency.body, prompt: "hello!" }, + })).rejects.toBeInstanceOf(PublicApiConflictError); + }); +}); diff --git a/tests/webhook-go-contract.test.ts b/tests/webhook-go-contract.test.ts new file mode 100644 index 0000000..6428836 --- /dev/null +++ b/tests/webhook-go-contract.test.ts @@ -0,0 +1,28 @@ +import { readFile } from "node:fs/promises"; + +import { describe, expect, it } from "vitest"; + +import type { GenerationJob } from "@/lib/types"; +import { buildJobWebhookPayload, signWebhookBody } from "@/lib/server/webhook"; + +const fixtureUrl = new URL("../contracts/webhook/job-webhook-v1.json", import.meta.url); + +describe("job webhook v1 cross-language contract", () => { + it("freezes exact body bytes and HMAC header", async () => { + const fixture = JSON.parse(await readFile(fixtureUrl, "utf8")); + const job = { + ...fixture.job, + ownerId: "owner-1", + provider: "mock", + reqKey: "fixture", + inputAssetIds: [], + inputUrls: [], + requestPayload: {}, + createdAt: fixture.job.updatedAt, + } as GenerationJob; + const body = JSON.stringify(buildJobWebhookPayload(job)); + expect(body).toBe(fixture.body); + expect(signWebhookBody(body, fixture.secret)).toBe(fixture.signature); + expect(fixture.maximumAttempts).toBe(3); + }); +});