From 48889607bc0a3162b3ac1ee29f6042f4943d01c0 Mon Sep 17 00:00:00 2001
From: Mike Pellegrini <mike.pellegrini@elastic.co>
Date: Mon, 7 Apr 2025 09:05:00 -0400
Subject: [PATCH] Update InferenceException to retain top-level message
 (#126345)

---
 .../xpack/inference/InferenceException.java   | 11 +++-
 .../inference/InferenceExceptionTests.java    | 51 +++++++++++++++++++
 2 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100644 x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/InferenceExceptionTests.java

diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceException.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceException.java
index 1ae9dac538d2b..3caa5bd06d058 100644
--- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceException.java
+++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceException.java
@@ -8,11 +8,18 @@
 package org.elasticsearch.xpack.inference;
 
 import org.elasticsearch.ElasticsearchException;
-import org.elasticsearch.ElasticsearchWrapperException;
+import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.rest.RestStatus;
 
-public class InferenceException extends ElasticsearchException implements ElasticsearchWrapperException {
+public class InferenceException extends ElasticsearchException {
     public InferenceException(String message, Throwable cause, Object... args) {
         super(message, cause, args);
     }
 
+    @Override
+    public RestStatus status() {
+        // Override status so that we get the status of the cause while retaining the message of the inference exception when emitting to
+        // XContent
+        return ExceptionsHelper.status(getCause());
+    }
 }
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/InferenceExceptionTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/InferenceExceptionTests.java
new file mode 100644
index 0000000000000..ba9daf95d17c2
--- /dev/null
+++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/InferenceExceptionTests.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference;
+
+import org.elasticsearch.ElasticsearchStatusException;
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xcontent.ToXContent;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentFactory;
+
+import static org.hamcrest.CoreMatchers.equalTo;
+
+public class InferenceExceptionTests extends ESTestCase {
+    public void testWrapException() throws Exception {
+        ElasticsearchStatusException cause = new ElasticsearchStatusException("test", RestStatus.BAD_REQUEST);
+        InferenceException testException = new InferenceException("test wrapper", cause);
+
+        XContentBuilder builder = XContentFactory.jsonBuilder();
+        builder.startObject();
+        testException.toXContent(builder, ToXContent.EMPTY_PARAMS);
+        builder.endObject();
+
+        assertThat(
+            Strings.toString(builder),
+            equalTo(
+                "{\"type\":\"inference_exception\",\"reason\":\"test wrapper\","
+                    + "\"caused_by\":{\"type\":\"status_exception\",\"reason\":\"test\"}}"
+            )
+        );
+        assertThat(testException.status(), equalTo(RestStatus.BAD_REQUEST));
+    }
+
+    public void testNullCause() throws Exception {
+        InferenceException testException = new InferenceException("test exception", null);
+
+        XContentBuilder builder = XContentFactory.jsonBuilder();
+        builder.startObject();
+        testException.toXContent(builder, ToXContent.EMPTY_PARAMS);
+        builder.endObject();
+
+        assertThat(Strings.toString(builder), equalTo("{\"type\":\"inference_exception\",\"reason\":\"test exception\"}"));
+        assertThat(testException.status(), equalTo(RestStatus.INTERNAL_SERVER_ERROR));
+    }
+}