diff --git a/cmd/dobotai/main.go b/cmd/dobotai/main.go index 8587a16..6ec818e 100644 --- a/cmd/dobotai/main.go +++ b/cmd/dobotai/main.go @@ -322,6 +322,8 @@ func main() { SettingsRouter: &dingTalkSettingsRouter{controller: settingsRuntime.controller}, ReasoningRouter: &dingTalkReasoningRouter{service: reasoningService}, LanguageResolver: settingsRuntime.repository, + + OutboundRewriteTimeout: outboundImages.rewriteTimeout, }, ) if err != nil { diff --git a/cmd/dobotai/outbound_images.go b/cmd/dobotai/outbound_images.go index 8f1c570..6512619 100644 --- a/cmd/dobotai/outbound_images.go +++ b/cmd/dobotai/outbound_images.go @@ -19,6 +19,10 @@ const ( outboundImageMaxBytes = 10 << 20 outboundImagePrefix = "dobotai/im-images" outboundImageRoute = "/dingtalk/i/" + // outboundImageRewriteTimeoutDefault is the budget for one answer's image + // rewrite. Manual screenshots are multi-megabyte uploads, so the historical + // 15s delivery timeout truncated long answers. + outboundImageRewriteTimeoutDefault = 60 * time.Second ) // outboundImageMaxCount resolves how many images one answer may rewrite. @@ -36,6 +40,22 @@ func outboundImageMaxCount() int { return n } +// outboundImageRewriteTimeout resolves the budget for the image rewrite stage. +// DOBOTAI_OUTBOUND_IMAGE_REWRITE_TIMEOUT accepts a Go duration ("90s", "2m"); +// invalid or non-positive values fall back to the default. +func outboundImageRewriteTimeout() time.Duration { + raw := strings.TrimSpace(os.Getenv("DOBOTAI_OUTBOUND_IMAGE_REWRITE_TIMEOUT")) + if raw == "" { + return outboundImageRewriteTimeoutDefault + } + duration, err := time.ParseDuration(raw) + if err != nil || duration <= 0 { + log.Printf("outbound image: invalid DOBOTAI_OUTBOUND_IMAGE_REWRITE_TIMEOUT=%q; using default %s", raw, outboundImageRewriteTimeoutDefault) + return outboundImageRewriteTimeoutDefault + } + return duration +} + type outboundImageConfig struct { WeKnoraBaseURL string WeKnoraAPIKey string @@ -51,6 +71,8 @@ type outboundImageRuntime struct { rewriter dingtalk.OutboundContentRewriter publisher *outboundimage.StableHashPublisher handler http.Handler + // rewriteTimeout is the budget the channel gives the image rewrite stage. + rewriteTimeout time.Duration // publicURLPrefix is the prefix of every published image URL, used by the // DingTalk card image uploader to recognise (and refetch) our own images. publicURLPrefix string @@ -106,6 +128,7 @@ func buildOutboundImageRuntime(cfg outboundImageConfig) (outboundImageRuntime, e rewriter: outboundimage.NewRewriter(fetcher, publisher, outboundImageMaxCount()), publisher: publisher, handler: handler, + rewriteTimeout: outboundImageRewriteTimeout(), publicURLPrefix: strings.TrimRight(cfg.PublicBaseURL, "/") + "/", }, nil } diff --git a/internal/channels/dingtalk/handler.go b/internal/channels/dingtalk/handler.go index d8fc2e5..b942abb 100644 --- a/internal/channels/dingtalk/handler.go +++ b/internal/channels/dingtalk/handler.go @@ -163,6 +163,11 @@ type HandlerOptions struct { SettingsRouter SettingsRouter ReasoningRouter ReasoningRouter LanguageResolver LanguagePreferenceResolver + // OutboundRewriteTimeout bounds the outbound image rewrite stage (fetch + + // publish). Answers carrying many large screenshots need more than the + // delivery timeout, otherwise the remaining images are dropped. Zero keeps + // finalDeliveryTimeout. + OutboundRewriteTimeout time.Duration } type Handler struct { @@ -181,6 +186,9 @@ type Handler struct { now func() time.Time cardUpdateInterval time.Duration finalDeliveryTimeout time.Duration + // outboundRewriteTimeout applies to the image rewrite stage only; zero + // falls back to finalDeliveryTimeout. + outboundRewriteTimeout time.Duration } // NewHandler accepts at most one image downloader. A nil/omitted downloader @@ -228,6 +236,8 @@ func NewHandlerWithOptions(cfg TenantConfig, qa QAAsker, reply MarkdownSender, o now: time.Now, cardUpdateInterval: 500 * time.Millisecond, finalDeliveryTimeout: defaultHTTPTimeout, + + outboundRewriteTimeout: opts.OutboundRewriteTimeout, }, nil } @@ -653,7 +663,10 @@ func (h *Handler) rewriteOutboundContent(ctx context.Context, content string) st if h.outboundRewriter == nil { return normalizeWideMarkdownTables(sanitizeDisplayAnswer(content)) } - timeout := h.finalDeliveryTimeout + timeout := h.outboundRewriteTimeout + if timeout <= 0 { + timeout = h.finalDeliveryTimeout + } if timeout <= 0 { timeout = defaultHTTPTimeout } diff --git a/internal/outboundimage/rewriter.go b/internal/outboundimage/rewriter.go index 0942c95..533a805 100644 --- a/internal/outboundimage/rewriter.go +++ b/internal/outboundimage/rewriter.go @@ -2,11 +2,15 @@ package outboundimage import ( "context" + "log" "net" "net/url" "path" "regexp" + "sort" "strings" + "sync" + "time" ) var ( @@ -44,35 +48,34 @@ func NewRewriter(fetcher Fetcher, publisher Publisher, maxImages int) *Rewriter return &Rewriter{fetcher: fetcher, publisher: publisher, maxImages: maxImages} } +// rewriteConcurrency bounds how many images are fetched and published in +// parallel. Manual screenshots are multi-megabyte objects whose upload time +// dominates the delivery budget, so a serial pass cannot publish a long answer +// before the delivery deadline and silently loses the remaining images. +const rewriteConcurrency = 4 + func (r *Rewriter) Rewrite(ctx context.Context, markdown string) string { if r == nil || r.fetcher == nil || r.publisher == nil || r.maxImages <= 0 { markdown = normalizePublicHTMLImages(stripResourceImages(markdown)) markdown = resourceReferenceRE.ReplaceAllString(resourceHTMLUnsafeImageRE.ReplaceAllString(markdown, ""), "") return sanitizePublicMarkdownImages(markdown) } - urls := map[string]string{} - attempted := 0 + started := time.Now() + urls := r.resolveReferences(ctx, markdown) resolve := func(reference string) string { - if publishedURL, ok := urls[reference]; ok { - return publishedURL + return urls[reference] + } + if len(urls) > 0 { + published, dropped := 0, 0 + for _, publishedURL := range urls { + if publishedURL == "" { + dropped++ + continue + } + published++ } - if attempted >= r.maxImages { - urls[reference] = "" - return "" - } - attempted++ - data, contentType, err := r.fetcher.Fetch(ctx, reference) - if err != nil { - urls[reference] = "" - return "" - } - publishedURL, err := r.publisher.Publish(ctx, data, contentType) - if err != nil || !strings.HasPrefix(publishedURL, "https://") { - urls[reference] = "" - return "" - } - urls[reference] = publishedURL - return publishedURL + log.Printf("outbound image: rewrite finished references=%d published=%d dropped=%d elapsed=%s", + len(urls), published, dropped, time.Since(started).Round(time.Millisecond)) } markdown = resourceMarkdownImageRE.ReplaceAllStringFunc(markdown, func(image string) string { @@ -94,6 +97,97 @@ func (r *Rewriter) Rewrite(ctx context.Context, markdown string) string { return sanitizePublicMarkdownImages(markdown) } +// resolveReferences returns reference -> public URL for every image reference +// in content, in document order and deduplicated, capped at maxImages. Each +// image is fetched and published with at most rewriteConcurrency workers; a +// reference that cannot be published maps to "". +func (r *Rewriter) resolveReferences(ctx context.Context, content string) map[string]string { + references := r.collectReferences(content) + if len(references) == 0 { + return map[string]string{} + } + results := make(map[string]string, len(references)) + var mutex sync.Mutex + var budgetOnce sync.Once + slots := make(chan struct{}, rewriteConcurrency) + var wait sync.WaitGroup + for _, reference := range references { + wait.Add(1) + go func(reference string) { + defer wait.Done() + slots <- struct{}{} + defer func() { <-slots }() + publishedURL := "" + if err := ctx.Err(); err != nil { + budgetOnce.Do(func() { + log.Printf("outbound image: delivery budget exhausted, dropping remaining images: %v", err) + }) + } else { + publishedURL = r.publishReference(ctx, reference) + } + mutex.Lock() + results[reference] = publishedURL + mutex.Unlock() + }(reference) + } + wait.Wait() + return results +} + +// collectReferences gathers the resource references that carry a real image, +// in document order, dropping duplicates and anything beyond maxImages. Bare or +// unquoted references are never fetched. +func (r *Rewriter) collectReferences(content string) []string { + type hit struct { + at int + reference string + } + hits := make([]hit, 0, 16) + for _, pattern := range []*regexp.Regexp{resourceMarkdownImageRE, resourceHTMLDoubleImageRE, resourceHTMLSingleImageRE} { + for _, match := range pattern.FindAllStringSubmatchIndex(content, -1) { + if len(match) < 4 || match[2] < 0 || match[3] > len(content) { + continue + } + hits = append(hits, hit{at: match[0], reference: content[match[2]:match[3]]}) + } + } + sort.SliceStable(hits, func(i, j int) bool { return hits[i].at < hits[j].at }) + references := make([]string, 0, len(hits)) + seen := make(map[string]bool, len(hits)) + for _, candidate := range hits { + if seen[candidate.reference] { + continue + } + seen[candidate.reference] = true + if len(references) >= r.maxImages { + break + } + references = append(references, candidate.reference) + } + return references +} + +// publishReference fetches one image from WeKnora and publishes it to our own +// object storage. Failures are logged with their reason instead of being +// silently dropped from the answer. +func (r *Rewriter) publishReference(ctx context.Context, reference string) string { + started := time.Now() + data, contentType, err := r.fetcher.Fetch(ctx, reference) + if err != nil { + log.Printf("outbound image: fetch failed, image dropped reference=%s error=%v", reference, err) + return "" + } + publishedURL, err := r.publisher.Publish(ctx, data, contentType) + if err != nil || !strings.HasPrefix(publishedURL, "https://") { + log.Printf("outbound image: publish failed, image dropped reference=%s bytes=%d contentType=%s error=%v", + reference, len(data), contentType, err) + return "" + } + log.Printf("outbound image: published reference=%s bytes=%d contentType=%s elapsed=%s", + reference, len(data), contentType, time.Since(started).Round(time.Millisecond)) + return publishedURL +} + func sanitizePublicMarkdownImages(markdown string) string { markdown = markdownImageRE.ReplaceAllStringFunc(markdown, func(image string) string { match := markdownImageRE.FindStringSubmatch(image) diff --git a/internal/outboundimage/rewriter_test.go b/internal/outboundimage/rewriter_test.go index d33c632..e4c6aa5 100644 --- a/internal/outboundimage/rewriter_test.go +++ b/internal/outboundimage/rewriter_test.go @@ -4,15 +4,19 @@ import ( "context" "errors" "strings" + "sync" "testing" ) type fakeFetcher struct { + mutex sync.Mutex calls map[string]int fail map[string]bool } func (f *fakeFetcher) Fetch(_ context.Context, ref string) ([]byte, string, error) { + f.mutex.Lock() + defer f.mutex.Unlock() if f.calls == nil { f.calls = map[string]int{} } @@ -24,10 +28,13 @@ func (f *fakeFetcher) Fetch(_ context.Context, ref string) ([]byte, string, erro } type fakePublisher struct { + mutex sync.Mutex calls int } func (p *fakePublisher) Publish(_ context.Context, data []byte, contentType string) (string, error) { + p.mutex.Lock() + defer p.mutex.Unlock() p.calls++ if contentType != "image/png" || len(data) == 0 { return "", errors.New("bad image") @@ -375,3 +382,43 @@ func TestDocumentImageDowngradePreservesImageAndLinkContracts(t *testing.T) { t.Fatalf("unsafe document image became a link: %s", got) } } + +func TestRewriteDropsImagesWhenDeliveryBudgetIsAlreadyExhausted(t *testing.T) { + fetcher := &fakeFetcher{} + publisher := &fakePublisher{} + r := NewRewriter(fetcher, publisher, 10) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + input := "正文\n![一](resource://FirstRefAbCdEfGhIjKlMnO)\n![二](resource://SecondRefAbCdEfGhIjKlMn)" + + got := r.Rewrite(ctx, input) + + if strings.Contains(got, "resource://") { + t.Fatalf("exhausted budget leaked a resource reference: %s", got) + } + if !strings.Contains(got, "正文") { + t.Fatalf("text was lost: %s", got) + } + if len(fetcher.calls) != 0 || publisher.calls != 0 { + t.Fatalf("no image may be fetched or published without budget: fetch=%v publish=%d", fetcher.calls, publisher.calls) + } +} + +func TestRewritePublishesEveryImageWithinTheLimit(t *testing.T) { + fetcher := &fakeFetcher{} + publisher := &fakePublisher{} + r := NewRewriter(fetcher, publisher, 12) + input := "正文" + for i := 0; i < 12; i++ { + input += "\n![图](resource://Ref" + string(rune('A'+i)) + "bCdEfGhIjKlMnO)" + } + + got := r.Rewrite(context.Background(), input) + + if strings.Contains(got, "resource://") { + t.Fatalf("resource reference remained: %s", got) + } + if publisher.calls != 12 || len(fetcher.calls) != 12 { + t.Fatalf("calls fetch=%d publish=%d, want 12", len(fetcher.calls), publisher.calls) + } +}