Agent execution error in json format

[jngz-es]

Description

Return agent execution error in xcontent format. If xcontent building failed, use plain text.

Related Issues

Resolves #[Issue number to be closed when this PR is merged]

Check List

• New functionality includes testing.
• New functionality has been documented.
• API changes companion pull request created.
• Commits are signed per the DCO using --signoff.
• Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[dhrubo-os]

Difference between these two tests are only assertEquals("application/json; charset=UTF-8", response.contentType()); vs assertEquals("text/plain; charset=UTF-8", response.contentType());. Is that expected?

[jngz-es]

Yes, it is expected. Actually there is another difference when(channel.newBuilder()).thenReturn(XContentFactory.jsonBuilder());
Without this jsonBuilder, the xcontent building will fail, then downgrade to catch branch to send plain text response.

[dhrubo-os]

[nit] I see. In that case, what do you think about something like following:

public void testAgentExecutionResponse() throws Exception {
    // Test JSON response
    runAgentExecutionResponseTest(
        "application/json; charset=UTF-8",
        "{\"status\":400,\"error\":{\"type\":\"IllegalArgumentException\",\"reason\":\"Invalid Request\",\"details\":\"Illegal Argument Exception\"}}",
        true
    );

    // Test plain text response
    runAgentExecutionResponseTest(
        "text/plain; charset=UTF-8",
        "{\"error\":{\"reason\":\"Invalid Request\",\"details\":\"Illegal Argument Exception\",\"type\":\"IllegalArgumentException\"},\"status\":400}",
        false
    );
}

private void runAgentExecutionResponseTest(String expectedContentType, String expectedError, boolean isJsonResponse) throws Exception {
    RestRequest request = getExecuteAgentRestRequest();
    
    // Mock client behavior
    doAnswer(invocation -> {
        ActionListener<MLExecuteTaskResponse> actionListener = invocation.getArgument(2);
        actionListener.onFailure(new RemoteTransportException("Remote Transport Exception", new IllegalArgumentException("Illegal Argument Exception")));
        return null;
    }).when(client).execute(eq(MLExecuteTaskAction.INSTANCE), any(), any());
    
    doNothing().when(channel).sendResponse(any());

    // Conditionally mock JSON response building
    if (isJsonResponse) {
        when(channel.newBuilder()).thenReturn(XContentFactory.jsonBuilder());
    }

    restMLExecuteAction.handleRequest(request, channel, client);

    // Verify request processing
    ArgumentCaptor<MLExecuteTaskRequest> argumentCaptor = ArgumentCaptor.forClass(MLExecuteTaskRequest.class);
    verify(client, times(1)).execute(eq(MLExecuteTaskAction.INSTANCE), argumentCaptor.capture(), any());
    Input input = argumentCaptor.getValue().getInput();
    assertEquals(FunctionName.AGENT, input.getFunctionName());

    // Verify response
    ArgumentCaptor<RestResponse> restResponseArgumentCaptor = ArgumentCaptor.forClass(RestResponse.class);
    verify(channel, times(1)).sendResponse(restResponseArgumentCaptor.capture());
    BytesRestResponse response = (BytesRestResponse) restResponseArgumentCaptor.getValue();
    assertEquals(RestStatus.BAD_REQUEST, response.status());
    assertEquals(expectedContentType, response.contentType());
    assertEquals(expectedError, response.content().utf8ToString());
}

[jngz-es]

Yeah, this looks good. Actually I would prefer to refactor all tests in RestMLExecuteActionTests in next pr, as some existing test cases have the similar improvement to do.

[dhrubo-os]

Sure, approving the PR in expectation of a follow up PR to refactor the tests in RestMLExecuteActionTests