fix(ingestion): scope extractor timeout to single LLM call (#21003)

internal/ingestion/component/extractor.go

@@ -25,9 +25,13 @@
 //   - CONCURRENCY: auto extraction tasks (keywords, questions, summary, metadata)
 //     run concurrently across chunks via a bounded worker pool (extractorPool).
 //
-//   - TIMEOUT / ELAPSED: the call is wrapped in
-//     runtime.WithTimeout(600s) and runtime.TrackElapsed so the
-//     upstream pipeline gets _created_time / _elapsed_time stamps.
+//   - TIMEOUT / ELAPSED: each LLM call carries its own budget
+//     (extractorChatTimeout) instead of one node-level deadline, so the
+//     total work scales with the chunk count — a 3000-chunk document
+//     legitimately needs longer than a 3-chunk one. A call that exceeds
+//     its budget skips its chunk instead of failing the document; task
+//     cancellation still aborts via ctx. Outputs are stamped with
+//     _created_time / _elapsed_time by runtime.TrackElapsed upstream.
 package component
 
 import (
@@ -62,11 +66,25 @@ import (
 
 const componentNameExtractor = "Extractor"
 
-// extractorTimeout bounds one LLM chat call. Matches the python
-// `@timeout(60)` default at rag/flow/base.py:60. The pipeline
-// orchestrator (Phase 3) overrides this if a stage-level ceiling
-// is configured.
-const extractorTimeout = 600 * time.Second
+// defaultExtractorChatTimeout bounds one LLM chat call. It is a hung-call
+// ceiling, not a pacing target: a fast model answers in seconds and never
+// approaches it, so per-model speed variance needs no tuning — only a hung
+// provider burns the full budget. Larger than the tag path's 60s because a
+// thinking model legitimately needs longer on open-ended prompts, and the
+// driver-level 20m stays as the absolute backstop.
+const defaultExtractorChatTimeout = 180 * time.Second
+
+// extractorChatTimeout resolves the per-call LLM budget from
+// EXTRACTOR_CHAT_TIMEOUT (seconds), falling back to
+// defaultExtractorChatTimeout on unset/invalid values.
+func extractorChatTimeout() time.Duration {
+	if v := os.Getenv("EXTRACTOR_CHAT_TIMEOUT"); v != "" {
+		if n, err := strconv.Atoi(strings.TrimSpace(v)); err == nil && n > 0 {
+			return time.Duration(n) * time.Second
+		}
+	}
+	return defaultExtractorChatTimeout
+}
 
 // extractorRetryMax and extractorRetryDelay are package-level vars
 // so tests can override them (extractorRetryDelay → time.Millisecond)
@@ -680,24 +698,30 @@ func (c *ExtractorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map
 	remainingOn := c.Param.Questions.Enabled || c.Param.Summary.Enabled || c.Param.Metadata.Enabled
 	phaseWindows := progress.NewSplit(keywordsOn, tagsOn, remainingOn)
 
-	if err := runtime.WithTimeout(ctx, extractorTimeout, func(timeoutCtx context.Context) error {
-		// Phase 1: Keywords extraction (if enabled), running across chunks via extractorPool.
-		if err := c.runAutoKeywordsPool(timeoutCtx, db, in, phaseWindows.Take(keywordsOn)); err != nil {
-			return err
-		}
+	// No node-level wall-clock budget (see the package doc): each LLM call
+	// carries its own budget, and an expired call skips its chunk instead
+	// of failing the document. Task cancellation still aborts via ctx.
+	if err := c.runAutoKeywordsPool(ctx, db, in, phaseWindows.Take(keywordsOn)); err != nil {
+		return nil, fmt.Errorf("extractor: %w", err)
+	}
 
-		// Phase 2: Tag phase (if enabled), benefiting from title and freshly extracted keywords.
-		if tagsOn {
-			tagged, tagErr := c.runAutoTags(timeoutCtx, db, in, phaseWindows.Take(tagsOn))
-			if tagErr != nil {
-				return tagErr
-			}
-			in.chunks = tagged
+	// Phase 2: Tag phase (if enabled), benefiting from title and freshly extracted keywords.
+	if tagsOn {
+		tagged, tagErr := c.runAutoTags(ctx, db, in, phaseWindows.Take(tagsOn))
+		if tagErr != nil {
+			return nil, fmt.Errorf("extractor: %w", tagErr)
 		}
+		in.chunks = tagged
+	}
 
-		// Phase 3: Remaining extractions (Questions, Summary, Metadata) via extractorPool.
-		return c.runRemainingExtractions(timeoutCtx, db, in, phaseWindows.Take(remainingOn))
-	}); err != nil {
+	// Phase 3: Remaining extractions (Questions, Summary, Metadata) via extractorPool.
+	if err := c.runRemainingExtractions(ctx, db, in, phaseWindows.Take(remainingOn)); err != nil {
+		return nil, fmt.Errorf("extractor: %w", err)
+	}
+	// The tag phase swallows per-chunk failures, so a task stopped mid-run
+	// could otherwise report success. This mirrors the old node-level
+	// WithTimeout post-check: never report success on a dead context.
+	if err := ctx.Err(); err != nil {
 		return nil, fmt.Errorf("extractor: %w", err)
 	}
 	common.Debug("extractor stage",
@@ -900,12 +924,60 @@ func (c *ExtractorComponent) runAutoSummary(ctx context.Context, db *gorm.DB, in
 	return nil
 }
 
+// resolveChunkCallError maps one chunk job's error to the pool outcome:
+//   - parent ctx done (task stop): propagate so cancellation stays fast.
+//   - call budget expiry with a live parent: skip the chunk, keep the
+//     document. errors.Is covers both our per-call budget and a
+//     driver-level deadline. completedAny reports whether an earlier op
+//     already succeeded for this chunk: a later timeout then keeps the
+//     partial result instead of counting toward the all-timed-out failure.
+//   - anything else: hard failure (existing fail-fast semantics).
+func resolveChunkCallError(ctx context.Context, idx int, phase string, err error, completedAny bool, skipped *atomic.Int64) error {
+	if err == nil {
+		return nil
+	}
+	if ctx.Err() != nil {
+		return err
+	}
+	if errors.Is(err, context.DeadlineExceeded) {
+		if completedAny {
+			common.Warn("extractor: chunk extraction partially timed out; keeping completed fields",
+				zap.String("component", componentNameExtractor),
+				zap.String("phase", phase),
+				zap.Int("chunk_index", idx),
+			)
+			return nil
+		}
+		skipped.Add(1)
+		common.Warn("extractor: chunk extraction timed out; skipping chunk",
+			zap.String("component", componentNameExtractor),
+			zap.String("phase", phase),
+			zap.Int("chunk_index", idx),
+		)
+		return nil
+	}
+	return err
+}
+
+// checkAllChunksSkipped fails the phase when every submitted chunk timed
+// out with zero completed extractions: all-hung means the model is down,
+// and silently producing a document with no extractions would hide the
+// outage. Chunks that timed out after a partial extraction succeed with
+// the per-chunk warnings above and do not count here.
+func checkAllChunksSkipped(phase string, total int, skipped int64) error {
+	if total > 0 && skipped == int64(total) {
+		return fmt.Errorf("extractor: %s: all %d chunks exceeded the per-call LLM timeout; the model is likely hung or unreachable", phase, total)
+	}
+	return nil
+}
+
 // runAutoKeywordsPool dispatches keyword extraction across all chunks concurrently
 // using extractorPool before the tagging stage.
 func (c *ExtractorComponent) runAutoKeywordsPool(ctx context.Context, db *gorm.DB, in extractorInputs, w progress.Window) error {
 	if !c.Param.Keywords.Enabled || len(in.chunks) == 0 {
 		return nil
 	}
+	var skipped atomic.Int64
 	futs := make([]utility.WorkerPoolFuture[extractorJob, struct{}], 0, len(in.chunks))
 	for i, ck := range in.chunks {
 		i, ck := i, ck
@@ -914,10 +986,12 @@ func (c *ExtractorComponent) runAutoKeywordsPool(ctx context.Context, db *gorm.D
 			continue
 		}
 		fn := func() error {
-			if err := c.runAutoKeywords(ctx, db, in, ck, text); err != nil {
-				return fmt.Errorf("chunk %d keywords: %w", i, err)
+			err := c.runAutoKeywords(ctx, db, in, ck, text)
+			completed := err == nil
+			if err != nil {
+				err = fmt.Errorf("chunk %d keywords: %w", i, err)
 			}
-			return nil
+			return resolveChunkCallError(ctx, i, "keywords", err, completed, &skipped)
 		}
 		f, err := extractorPool.Submit(ctx, fn)
 		if err != nil {
@@ -925,7 +999,10 @@ func (c *ExtractorComponent) runAutoKeywordsPool(ctx context.Context, db *gorm.D
 		}
 		futs = append(futs, f)
 	}
-	return awaitFutures(ctx, futs, w)
+	if err := awaitFutures(ctx, futs, w); err != nil {
+		return err
+	}
+	return checkAllChunksSkipped("keywords", len(futs), skipped.Load())
 }
 
 // runRemainingExtractions dispatches auto questions / summary / metadata
@@ -934,6 +1011,7 @@ func (c *ExtractorComponent) runRemainingExtractions(ctx context.Context, db *go
 	if !c.Param.Questions.Enabled && !c.Param.Summary.Enabled && !c.Param.Metadata.Enabled {
 		return nil
 	}
+	var skipped atomic.Int64
 	futs := make([]utility.WorkerPoolFuture[extractorJob, struct{}], 0, len(in.chunks))
 	for i, ck := range in.chunks {
 		i, ck := i, ck
@@ -941,34 +1019,45 @@ func (c *ExtractorComponent) runRemainingExtractions(ctx context.Context, db *go
 		if strings.TrimSpace(text) == "" {
 			continue
 		}
-		fn := c.remainingExtractionJob(ctx, db, in, i, ck, text)
+		fn := c.remainingExtractionJob(ctx, db, in, i, ck, text, &skipped)
 		f, err := extractorPool.Submit(ctx, fn)
 		if err != nil {
 			return err
 		}
 		futs = append(futs, f)
 	}
-	return awaitFutures(ctx, futs, w)
+	if err := awaitFutures(ctx, futs, w); err != nil {
+		return err
+	}
+	return checkAllChunksSkipped("remaining", len(futs), skipped.Load())
 }
 
-func (c *ExtractorComponent) remainingExtractionJob(ctx context.Context, db *gorm.DB, in extractorInputs, idx int, ck map[string]any, chunkText string) extractorJob {
+func (c *ExtractorComponent) remainingExtractionJob(ctx context.Context, db *gorm.DB, in extractorInputs, idx int, ck map[string]any, chunkText string, skipped *atomic.Int64) extractorJob {
 	return func() error {
+		var err error
+		var completed bool
 		if c.Param.Questions.Enabled {
-			if err := c.runAutoQuestions(ctx, db, in, ck, chunkText); err != nil {
-				return fmt.Errorf("chunk %d questions: %w", idx, err)
+			if qErr := c.runAutoQuestions(ctx, db, in, ck, chunkText); qErr != nil {
+				err = fmt.Errorf("chunk %d questions: %w", idx, qErr)
+			} else {
+				completed = true
 			}
 		}
-		if c.Param.Summary.Enabled {
-			if err := c.runAutoSummary(ctx, db, in, ck, chunkText); err != nil {
-				return fmt.Errorf("chunk %d summary: %w", idx, err)
+		if err == nil && c.Param.Summary.Enabled {
+			if sErr := c.runAutoSummary(ctx, db, in, ck, chunkText); sErr != nil {
+				err = fmt.Errorf("chunk %d summary: %w", idx, sErr)
+			} else {
+				completed = true
 			}
 		}
-		if c.Param.Metadata.Enabled {
-			if err := c.runEnableMetadata(ctx, db, in, ck, chunkText); err != nil {
-				return fmt.Errorf("chunk %d metadata: %w", idx, err)
+		if err == nil && c.Param.Metadata.Enabled {
+			if mErr := c.runEnableMetadata(ctx, db, in, ck, chunkText); mErr != nil {
+				err = fmt.Errorf("chunk %d metadata: %w", idx, mErr)
+			} else {
+				completed = true
 			}
 		}
-		return nil
+		return resolveChunkCallError(ctx, idx, "remaining", err, completed, skipped)
 	}
 }
 
@@ -1223,8 +1312,15 @@ func (c *ExtractorComponent) callRaw(ctx context.Context, db *gorm.DB, in extrac
 		req.Temperature = &temp
 	}
 	var resp *extractorChatResponse
-	if err := common.RetryWithBackoff(ctx, extractorRetryMax, extractorRetryDelay, func() error {
-		r, e := inv.Chat(ctx, req)
+	// One hung-call ceiling per LLM request (see extractorChatTimeout).
+	// RetryWithBackoff treats a budget expiry as terminal (no retry), so a
+	// hung call burns at most one attempt. The pool jobs map the expiry to
+	// a per-chunk skip; anything else keeps fail-fast semantics.
+	budget := extractorChatTimeout()
+	callCtx, cancel := context.WithTimeout(ctx, budget)
+	defer cancel()
+	if err := common.RetryWithBackoff(callCtx, extractorRetryMax, extractorRetryDelay, func() error {
+		r, e := inv.Chat(callCtx, req)
 		resp = r
 		return e
 	}, isRetryableLLMError); err != nil {

internal/ingestion/component/extractor_call_timeout_test.go

@@ -0,0 +1,252 @@
+//
+//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+package component
+
+import (
+	"context"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"ragflow/internal/ingestion/component/schema"
+)
+
+// routeStubChatInvoker answers immediately unless a message contains
+// "slow-chunk", in which case it hangs until ctx cancellation. It lets tests
+// mix healthy and hung chunks deterministically despite pool concurrency.
+type routeStubChatInvoker struct {
+	calls    atomic.Int32
+	mu       sync.Mutex
+	requests []extractorChatRequest
+}
+
+func (s *routeStubChatInvoker) Chat(ctx context.Context, req extractorChatRequest) (*extractorChatResponse, error) {
+	s.calls.Add(1)
+	s.mu.Lock()
+	s.requests = append(s.requests, req)
+	s.mu.Unlock()
+	for _, m := range req.Messages {
+		if strings.Contains(m.Content, "slow-chunk") {
+			<-ctx.Done()
+			return nil, ctx.Err()
+		}
+	}
+	return &extractorChatResponse{Content: "kw1, kw2"}, nil
+}
+
+func withRouteStubChatInvoker(t *testing.T) *routeStubChatInvoker {
+	t.Helper()
+	prev := defaultExtractorChatInvoker
+	stub := &routeStubChatInvoker{}
+	SetExtractorChatInvoker(stub)
+	t.Cleanup(func() {
+		SetExtractorChatInvoker(prev)
+	})
+	return stub
+}
+
+func keywordsOnlyExtractor() *ExtractorComponent {
+	return &ExtractorComponent{Param: schema.ExtractorParam{
+		LLMID:    "test-model",
+		Keywords: schema.KeywordExtractConfig{Enabled: true, TopN: 3},
+	}}
+}
+
+// TestExtractorChatTimeoutSkipsSlowChunk: one hung chunk must not fail the
+// document. The healthy chunk keeps its extraction, the hung one is skipped.
+func TestExtractorChatTimeoutSkipsSlowChunk(t *testing.T) {
+	t.Setenv("EXTRACTOR_CHAT_TIMEOUT", "1")
+	stub := withRouteStubChatInvoker(t)
+
+	c := keywordsOnlyExtractor()
+	out, err := c.Invoke(t.Context(), nil, map[string]any{
+		"chunks": []map[string]any{
+			{"text": "fast chunk one"},
+			{"text": "slow-chunk two hangs"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("Invoke: %v", err)
+	}
+	chunks, _ := out["chunks"].([]map[string]any)
+	if len(chunks) != 2 {
+		t.Fatalf("chunks = %d, want 2", len(chunks))
+	}
+	if _, ok := chunks[0]["important_kwd"]; !ok {
+		t.Error("fast chunk lost its keywords")
+	}
+	if _, ok := chunks[1]["important_kwd"]; ok {
+		t.Error("slow chunk should have been skipped without keywords")
+	}
+	if got := stub.calls.Load(); got != 2 {
+		t.Errorf("LLM calls = %d, want 2 (hung call burns one budget, no retry)", got)
+	}
+}
+
+// TestExtractorChatTimeoutAllHangFails: every chunk hung means the model is
+// down — the run must fail loudly, not silently produce empty extractions.
+func TestExtractorChatTimeoutAllHangFails(t *testing.T) {
+	t.Setenv("EXTRACTOR_CHAT_TIMEOUT", "1")
+	withRouteStubChatInvoker(t)
+
+	c := keywordsOnlyExtractor()
+	_, err := c.Invoke(t.Context(), nil, map[string]any{
+		"chunks": []map[string]any{
+			{"text": "slow-chunk one hangs"},
+			{"text": "slow-chunk two hangs"},
+		},
+	})
+	if err == nil {
+		t.Fatal("Invoke succeeded, want all-chunks-timed-out error")
+	}
+	if !strings.Contains(err.Error(), "all 2 chunks") {
+		t.Errorf("error = %q, want all-chunks-timed-out", err.Error())
+	}
+}
+
+// TestExtractorParentCancelPropagates: task stop must still fail fast, never
+// degrade into per-chunk skips.
+func TestExtractorParentCancelPropagates(t *testing.T) {
+	t.Setenv("EXTRACTOR_CHAT_TIMEOUT", "1")
+	withRouteStubChatInvoker(t)
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	c := keywordsOnlyExtractor()
+	_, err := c.Invoke(ctx, nil, map[string]any{
+		"chunks": []map[string]any{
+			{"text": "fast chunk one"},
+			{"text": "fast chunk two"},
+		},
+	})
+	if err == nil {
+		t.Fatal("Invoke succeeded on canceled ctx, want error")
+	}
+}
+
+// TestExtractorChatTimeoutEnv: budget resolution and fallback.
+func TestExtractorChatTimeoutEnv(t *testing.T) {
+	for _, tt := range []struct {
+		name  string
+		value string
+		want  time.Duration
+	}{
+		{name: "unset falls back to default", value: "", want: defaultExtractorChatTimeout},
+		{name: "override wins", value: "1", want: time.Second},
+		{name: "zero falls back", value: "0", want: defaultExtractorChatTimeout},
+		{name: "negative falls back", value: "-5", want: defaultExtractorChatTimeout},
+		{name: "non-numeric falls back", value: "abc", want: defaultExtractorChatTimeout},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Setenv("EXTRACTOR_CHAT_TIMEOUT", tt.value)
+			if got := extractorChatTimeout(); got != tt.want {
+				t.Errorf("extractorChatTimeout() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+	if defaultExtractorChatTimeout != 180*time.Second {
+		t.Errorf("defaultExtractorChatTimeout = %v, want 180s", defaultExtractorChatTimeout)
+	}
+}
+
+// summaryHangStub answers every call except summary ones, which hang until
+// ctx cancellation. It models a model that serves some extraction types
+// while another hangs.
+type summaryHangStub struct {
+	calls atomic.Int32
+}
+
+func (s *summaryHangStub) Chat(ctx context.Context, req extractorChatRequest) (*extractorChatResponse, error) {
+	s.calls.Add(1)
+	for _, m := range req.Messages {
+		if strings.Contains(m.Content, "faithful text summarizer") {
+			<-ctx.Done()
+			return nil, ctx.Err()
+		}
+	}
+	return &extractorChatResponse{Content: "q1?\nq2?"}, nil
+}
+
+func withSummaryHangStub(t *testing.T) {
+	t.Helper()
+	prev := defaultExtractorChatInvoker
+	SetExtractorChatInvoker(&summaryHangStub{})
+	t.Cleanup(func() {
+		SetExtractorChatInvoker(prev)
+	})
+}
+
+// TestExtractorPartialTimeoutKeepsCompleted: questions succeed while
+// summaries hang for every chunk. The phase must succeed and keep the
+// questions instead of failing as all-timed-out.
+func TestExtractorPartialTimeoutKeepsCompleted(t *testing.T) {
+	t.Setenv("EXTRACTOR_CHAT_TIMEOUT", "1")
+	withSummaryHangStub(t)
+
+	c := &ExtractorComponent{Param: schema.ExtractorParam{
+		LLMID:     "test-model",
+		Questions: schema.QuestionExtractConfig{Enabled: true, TopN: 2},
+		Summary:   schema.SummaryExtractConfig{Enabled: true},
+	}}
+	out, err := c.Invoke(t.Context(), nil, map[string]any{
+		"chunks": []map[string]any{
+			{"text": "first piece"},
+			{"text": "second piece"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("Invoke: %v", err)
+	}
+	chunks, _ := out["chunks"].([]map[string]any)
+	if len(chunks) != 2 {
+		t.Fatalf("chunks = %d, want 2", len(chunks))
+	}
+	for i, ck := range chunks {
+		if _, ok := ck["question_kwd"]; !ok {
+			t.Errorf("chunk %d lost its completed questions", i)
+		}
+		if _, ok := ck["summary"]; ok {
+			t.Errorf("chunk %d should have no summary (hung)", i)
+		}
+	}
+}
+
+// TestExtractorTagsOnlyCancelFails: the tag phase swallows per-chunk
+// failures, so a stop during a tags-only run must still surface as an
+// error rather than a false success.
+func TestExtractorTagsOnlyCancelFails(t *testing.T) {
+	withRouteStubChatInvoker(t)
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	c := &ExtractorComponent{Param: schema.ExtractorParam{
+		LLMID: "test-model",
+		Tags:  schema.TagExtractConfig{Enabled: true, TopN: 3},
+	}}
+	_, err := c.Invoke(ctx, nil, map[string]any{
+		"chunks": []map[string]any{
+			{"text": "some chunk"},
+		},
+	})
+	if err == nil {
+		t.Fatal("Invoke succeeded on canceled ctx, want error")
+	}
+}