From bd1e6c329e56d9a8b845ab7fe54dce0b7081c2ee Mon Sep 17 00:00:00 2001
From: Jim Ferenczi <jim.ferenczi@elastic.co>
Date: Thu, 25 Sep 2025 12:25:05 +0100
Subject: [PATCH] Avoid UOE when fetching the _inference metadata field

The fields list in the context can be frozen so this change ensures that we recreate the fetch fields context when adding the _inference metadata field.
---
 .../search/fetch/FetchPhase.java              | 11 +++---
 .../inference/30_semantic_text_inference.yml  | 39 +++++++++++++++++++
 2 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java b/server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java
index f15ce02e23ba4..14c392b675a65 100644
--- a/server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java
+++ b/server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java
@@ -137,12 +137,13 @@ private SearchHits buildSearchHits(SearchContext context, int[] docIdsToLoad, Pr
         if (lookup.inferenceFields().isEmpty() == false
             && shouldExcludeInferenceFieldsFromSource(context.indexShard().indexSettings(), context.fetchSourceContext()) == false) {
             // Rehydrate the inference fields into the {@code _source} because they were explicitly requested.
-            var fetchFieldsContext = context.fetchFieldsContext();
-            if (fetchFieldsContext == null) {
-                fetchFieldsContext = new FetchFieldsContext(new ArrayList<>());
+            var oldFetchFieldsContext = context.fetchFieldsContext();
+            var newFetchFieldsContext = new FetchFieldsContext(new ArrayList<>());
+            if (oldFetchFieldsContext != null) {
+                newFetchFieldsContext.fields().addAll(oldFetchFieldsContext.fields());
             }
-            fetchFieldsContext.fields().add(new FieldAndFormat(InferenceMetadataFieldsMapper.NAME, null));
-            context.fetchFieldsContext(fetchFieldsContext);
+            newFetchFieldsContext.fields().add(new FieldAndFormat(InferenceMetadataFieldsMapper.NAME, null));
+            context.fetchFieldsContext(newFetchFieldsContext);
         }
 
         SourceLoader sourceLoader = context.newSourceLoader(res.v2());
diff --git a/x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/30_semantic_text_inference.yml b/x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/30_semantic_text_inference.yml
index fc09ef35a3304..bb4c00acf39b7 100644
--- a/x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/30_semantic_text_inference.yml
+++ b/x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/30_semantic_text_inference.yml
@@ -1347,3 +1347,42 @@ setup:
 
   - match: { hits.total.value: 1 }
 
+---
+"Search can rehydrate the inference metadata fields in _source":
+  - do:
+      index:
+        index: test-index
+        id: doc_1
+        refresh: true
+        body:
+          sparse_field: "inference test"
+          dense_field: "another inference test"
+          non_inference_field: "non inference test"
+
+  - do:
+      search:
+        index: test-index
+        body:
+          _source:
+            exclude_vectors: false
+          fields: ["sparse_field", "dense_field", "non_inference_field"]
+          query:
+            term:
+              _id: doc_1
+
+  - match: { hits.total.value: 1 }
+  - match: { hits.hits.0.fields.sparse_field: ["inference test"]}
+  - match: { hits.hits.0.fields.dense_field: ["another inference test"]}
+  - match: { hits.hits.0.fields.non_inference_field: ["non inference test"]}
+  - not_exists: hits.hits.0.fields._inference_fields
+  - length: { hits.hits.0._source._inference_fields.sparse_field.inference.chunks: 1 }
+  - length: { hits.hits.0._source._inference_fields.sparse_field.inference.chunks.sparse_field: 1 }
+  - exists: hits.hits.0._source._inference_fields.sparse_field.inference.chunks.sparse_field.0.embeddings
+  - match: { hits.hits.0._source._inference_fields.sparse_field.inference.chunks.sparse_field.0.start_offset: 0 }
+  - match: { hits.hits.0._source._inference_fields.sparse_field.inference.chunks.sparse_field.0.end_offset: 14 }
+  - length: { hits.hits.0._source._inference_fields.dense_field.inference.chunks.dense_field: 1 }
+  - exists: hits.hits.0._source._inference_fields.dense_field.inference.chunks.dense_field.0.embeddings
+  - match: { hits.hits.0._source._inference_fields.dense_field.inference.chunks.dense_field.0.start_offset: 0 }
+  - match: { hits.hits.0._source._inference_fields.dense_field.inference.chunks.dense_field.0.end_offset: 22 }
+
+