364 lines
9.1 KiB
Go
364 lines
9.1 KiB
Go
|
package application
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"encoding/json"
|
|||
|
|
"fmt"
|
|||
|
|
"net/url"
|
|||
|
|
"regexp"
|
|||
|
|
"strings"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"img_infinite_canvas/internal/domain/design"
|
|||
|
|
|
|||
|
|
"github.com/zeromicro/go-zero/core/logx"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
assetCleanupWorkers = 5
|
|||
|
|
assetCleanupQueueSize = 4096
|
|||
|
|
assetCleanupMaxRetries = 3
|
|||
|
|
assetCleanupTimeout = 60 * time.Second
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
var (
|
|||
|
|
promptImageTokenPattern = regexp.MustCompile(`\[@image:#\d+:[^:\]]+:(?:\[[^\]]+\]\(([^)]+)\)|([^\]]+))\]`)
|
|||
|
|
promptImageDirectivePattern = regexp.MustCompile(`(?i)(?:参考图|Reference image)\s*\d?\s*(?:[（(][^）)\n]+[）)])?\s*[:：]\s*(https?://[^\s\]]+)`)
|
|||
|
|
messageImageURLPattern = regexp.MustCompile(`(?i)(https?://[^\s"'<>)\]]*\.(?:png|jpe?g|webp|gif|svg)(?:\?[^\s"'<>)\]]*)?|/[^\s"'<>)\]]*\.(?:png|jpe?g|webp|gif|svg)(?:\?[^\s"'<>)\]]*)?)`)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
type assetCleanupQueue struct {
|
|||
|
|
storage design.AssetStorage
|
|||
|
|
jobs chan string
|
|||
|
|
maxRetries int
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func newAssetCleanupQueue(storage design.AssetStorage) *assetCleanupQueue {
|
|||
|
|
if storage == nil {
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
queue := &assetCleanupQueue{
|
|||
|
|
storage: storage,
|
|||
|
|
jobs: make(chan string, assetCleanupQueueSize),
|
|||
|
|
maxRetries: assetCleanupMaxRetries,
|
|||
|
|
}
|
|||
|
|
for i := 0; i < assetCleanupWorkers; i++ {
|
|||
|
|
go queue.run()
|
|||
|
|
}
|
|||
|
|
return queue
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (s *DesignService) enqueueAssetDeletion(publicURLs []string) {
|
|||
|
|
publicURLs = normalizeAssetURLs(publicURLs)
|
|||
|
|
if len(publicURLs) == 0 {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
if s.jobs != nil {
|
|||
|
|
if err := s.jobs.Enqueue(context.Background(), design.Job{
|
|||
|
|
Kind: design.JobAssetCleanup,
|
|||
|
|
PublicURLs: publicURLs,
|
|||
|
|
}); err == nil {
|
|||
|
|
return
|
|||
|
|
} else {
|
|||
|
|
logx.Errorf("enqueue asset cleanup failed: %v", err)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
if s.assetCleanup == nil {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
s.assetCleanup.enqueue(publicURLs)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (q *assetCleanupQueue) enqueue(publicURLs []string) {
|
|||
|
|
for _, publicURL := range normalizeAssetURLs(publicURLs) {
|
|||
|
|
select {
|
|||
|
|
case q.jobs <- publicURL:
|
|||
|
|
default:
|
|||
|
|
go q.enqueueSlow(publicURL)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (q *assetCleanupQueue) enqueueSlow(publicURL string) {
|
|||
|
|
select {
|
|||
|
|
case q.jobs <- publicURL:
|
|||
|
|
case <-time.After(2 * time.Second):
|
|||
|
|
logx.Errorf("asset cleanup queue is full, dropping delete job for %s", publicURL)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (q *assetCleanupQueue) run() {
|
|||
|
|
for publicURL := range q.jobs {
|
|||
|
|
if err := q.deleteWithRetry(publicURL); err != nil {
|
|||
|
|
logx.Errorf("asset cleanup failed for %s: %v", publicURL, err)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (q *assetCleanupQueue) deleteWithRetry(publicURL string) error {
|
|||
|
|
var err error
|
|||
|
|
for attempt := 1; attempt <= q.maxRetries; attempt++ {
|
|||
|
|
ctx, cancel := context.WithTimeout(context.Background(), assetCleanupTimeout)
|
|||
|
|
err = q.storage.DeleteObject(ctx, publicURL)
|
|||
|
|
cancel()
|
|||
|
|
if err == nil {
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
time.Sleep(time.Duration(attempt*attempt) * 200 * time.Millisecond)
|
|||
|
|
}
|
|||
|
|
return err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (s *DesignService) deleteAssetURLsWithRetry(publicURLs []string) error {
|
|||
|
|
if s.assets == nil {
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
var failures []string
|
|||
|
|
for _, publicURL := range normalizeAssetURLs(publicURLs) {
|
|||
|
|
queue := &assetCleanupQueue{
|
|||
|
|
storage: s.assets,
|
|||
|
|
maxRetries: assetCleanupMaxRetries,
|
|||
|
|
}
|
|||
|
|
if err := queue.deleteWithRetry(publicURL); err != nil {
|
|||
|
|
failures = append(failures, fmt.Sprintf("%s: %v", publicURL, err))
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
if len(failures) > 0 {
|
|||
|
|
return fmt.Errorf("asset cleanup failed: %s", strings.Join(failures, "; "))
|
|||
|
|
}
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (s *DesignService) assetReferencedByProjects(ctx context.Context, publicURL string) (bool, error) {
|
|||
|
|
if strings.TrimSpace(publicURL) == "" {
|
|||
|
|
return false, nil
|
|||
|
|
}
|
|||
|
|
summaries, err := s.repo.List(ctx)
|
|||
|
|
if err != nil {
|
|||
|
|
return false, err
|
|||
|
|
}
|
|||
|
|
for _, summary := range summaries {
|
|||
|
|
project, err := s.repo.Get(ctx, summary.ID)
|
|||
|
|
if err != nil {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
if projectReferencesAssetURL(project, publicURL) {
|
|||
|
|
return true, nil
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return false, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func normalizeAssetURLs(publicURLs []string) []string {
|
|||
|
|
seen := make(map[string]struct{})
|
|||
|
|
urls := make([]string, 0, len(publicURLs))
|
|||
|
|
for _, publicURL := range publicURLs {
|
|||
|
|
publicURL = strings.TrimSpace(publicURL)
|
|||
|
|
if publicURL == "" {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
if _, ok := seen[publicURL]; ok {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
seen[publicURL] = struct{}{}
|
|||
|
|
urls = append(urls, publicURL)
|
|||
|
|
}
|
|||
|
|
return urls
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func projectReferencesAssetURL(project design.Project, publicURL string) bool {
|
|||
|
|
protectedURLs := canvasAssetURLSet(project.Nodes)
|
|||
|
|
for _, url := range canvasSnapshotAssetURLs(project.Canvas) {
|
|||
|
|
addAssetURLKeys(protectedURLs, url)
|
|||
|
|
}
|
|||
|
|
for _, message := range project.Messages {
|
|||
|
|
for _, url := range messageAssetURLs(message) {
|
|||
|
|
addAssetURLKeys(protectedURLs, url)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return assetURLSetContains(protectedURLs, publicURL)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func deletedCanvasAssetURLs(previous []design.Node, next []design.Node, messages []design.Message) []string {
|
|||
|
|
protectedURLs := canvasAssetURLSet(next)
|
|||
|
|
for _, message := range messages {
|
|||
|
|
for _, url := range messageAssetURLs(message) {
|
|||
|
|
addAssetURLKeys(protectedURLs, url)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
seen := make(map[string]struct{})
|
|||
|
|
var deleted []string
|
|||
|
|
for _, node := range previous {
|
|||
|
|
url := nodeAssetURL(node)
|
|||
|
|
if url == "" {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
if assetURLSetContains(protectedURLs, url) {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
if _, ok := seen[url]; ok {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
seen[url] = struct{}{}
|
|||
|
|
deleted = append(deleted, url)
|
|||
|
|
}
|
|||
|
|
return deleted
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func projectAssetURLs(project design.Project) []string {
|
|||
|
|
seen := make(map[string]struct{})
|
|||
|
|
var urls []string
|
|||
|
|
add := func(url string) {
|
|||
|
|
url = strings.TrimSpace(url)
|
|||
|
|
if url == "" {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
if _, ok := seen[url]; ok {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
seen[url] = struct{}{}
|
|||
|
|
urls = append(urls, url)
|
|||
|
|
}
|
|||
|
|
for _, node := range project.Nodes {
|
|||
|
|
add(nodeAssetURL(node))
|
|||
|
|
}
|
|||
|
|
for _, url := range canvasSnapshotAssetURLs(project.Canvas) {
|
|||
|
|
add(url)
|
|||
|
|
}
|
|||
|
|
for _, message := range project.Messages {
|
|||
|
|
for _, url := range messageAssetURLs(message) {
|
|||
|
|
add(url)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return urls
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func messageAssetURLs(message design.Message) []string {
|
|||
|
|
var urls []string
|
|||
|
|
content := strings.TrimSpace(message.Content)
|
|||
|
|
if content == "" {
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
var events []struct {
|
|||
|
|
EventData struct {
|
|||
|
|
Artifact []struct {
|
|||
|
|
ImageURL string `json:"image_url"`
|
|||
|
|
} `json:"artifact"`
|
|||
|
|
} `json:"event_data"`
|
|||
|
|
}
|
|||
|
|
if err := json.Unmarshal([]byte(content), &events); err == nil {
|
|||
|
|
for _, event := range events {
|
|||
|
|
for _, artifact := range event.EventData.Artifact {
|
|||
|
|
if strings.TrimSpace(artifact.ImageURL) != "" {
|
|||
|
|
urls = append(urls, artifact.ImageURL)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
urls = append(urls, promptMessageAssetURLs(content)...)
|
|||
|
|
urls = append(urls, looseMessageImageURLs(content)...)
|
|||
|
|
return urls
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func canvasSnapshotAssetURLs(canvas string) []string {
|
|||
|
|
snapshot, err := design.DecodeCanvasSnapshot(canvas)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
urls := make([]string, 0, len(snapshot.Nodes))
|
|||
|
|
for _, node := range snapshot.Nodes {
|
|||
|
|
if url := nodeAssetURL(node); url != "" {
|
|||
|
|
urls = append(urls, url)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return urls
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func canvasAssetURLSet(nodes []design.Node) map[string]struct{} {
|
|||
|
|
items := make(map[string]struct{})
|
|||
|
|
for _, node := range nodes {
|
|||
|
|
if url := nodeAssetURL(node); url != "" {
|
|||
|
|
addAssetURLKeys(items, url)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return items
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func nodeAssetURL(node design.Node) string {
|
|||
|
|
if node.Type != design.NodeTypeImage && node.Type != design.NodeTypeFrame {
|
|||
|
|
return ""
|
|||
|
|
}
|
|||
|
|
content := strings.TrimSpace(node.Content)
|
|||
|
|
if content == "" || strings.HasPrefix(content, "data:image/") || !isGeneratedImageContent(content) {
|
|||
|
|
return ""
|
|||
|
|
}
|
|||
|
|
return content
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func promptMessageAssetURLs(content string) []string {
|
|||
|
|
var urls []string
|
|||
|
|
for _, match := range promptImageTokenPattern.FindAllStringSubmatch(content, -1) {
|
|||
|
|
if len(match) > 1 && strings.TrimSpace(match[1]) != "" {
|
|||
|
|
urls = append(urls, strings.TrimSpace(match[1]))
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
if len(match) > 2 && strings.TrimSpace(match[2]) != "" {
|
|||
|
|
urls = append(urls, strings.TrimSpace(match[2]))
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
for _, match := range promptImageDirectivePattern.FindAllStringSubmatch(content, -1) {
|
|||
|
|
if len(match) > 1 && strings.TrimSpace(match[1]) != "" {
|
|||
|
|
urls = append(urls, strings.TrimSpace(match[1]))
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return urls
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func looseMessageImageURLs(content string) []string {
|
|||
|
|
matches := messageImageURLPattern.FindAllString(content, -1)
|
|||
|
|
urls := make([]string, 0, len(matches))
|
|||
|
|
for _, match := range matches {
|
|||
|
|
url := strings.TrimRight(strings.TrimSpace(match), ".,，。;；")
|
|||
|
|
if url != "" {
|
|||
|
|
urls = append(urls, url)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return urls
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func assetURLSetContains(items map[string]struct{}, publicURL string) bool {
|
|||
|
|
for key := range assetURLKeys(publicURL) {
|
|||
|
|
if _, ok := items[key]; ok {
|
|||
|
|
return true
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return false
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func addAssetURLKeys(items map[string]struct{}, publicURL string) {
|
|||
|
|
for key := range assetURLKeys(publicURL) {
|
|||
|
|
items[key] = struct{}{}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func assetURLKeys(publicURL string) map[string]struct{} {
|
|||
|
|
publicURL = strings.TrimSpace(publicURL)
|
|||
|
|
keys := make(map[string]struct{})
|
|||
|
|
if publicURL == "" || strings.HasPrefix(publicURL, "data:image/") || strings.HasPrefix(publicURL, "blob:") {
|
|||
|
|
return keys
|
|||
|
|
}
|
|||
|
|
keys[publicURL] = struct{}{}
|
|||
|
|
if queryIndex := strings.Index(publicURL, "?"); queryIndex > 0 {
|
|||
|
|
keys[publicURL[:queryIndex]] = struct{}{}
|
|||
|
|
}
|
|||
|
|
parsed, err := url.Parse(publicURL)
|
|||
|
|
if err == nil && parsed.Path != "" {
|
|||
|
|
keys[parsed.Path] = struct{}{}
|
|||
|
|
pathWithoutQuery := parsed.Path
|
|||
|
|
if parsed.RawQuery != "" {
|
|||
|
|
keys[pathWithoutQuery+"?"+parsed.RawQuery] = struct{}{}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return keys
|
|||
|
|
}
|