return 404 instead of 500 for missing agent config name

Signed-off-by: Yulong Ruan <[email protected]>

server/routes/agent_routes.test.ts

@@ -125,7 +125,7 @@ describe('test execute agent route', () => {
         "headers": Object {},
         "payload": Object {
           "error": "Internal Server Error",
-          "message": "Execute agent failed!",
+          "message": "An internal server error occurred.",
           "statusCode": 500,
         },
         "statusCode": 500,

server/routes/agent_routes.ts

@@ -7,6 +7,7 @@ import { schema } from '@osd/config-schema';
 import { IRouter } from '../../../../src/core/server';
 import { AGENT_API } from '../../common/constants/llm';
 import { AssistantServiceSetup } from '../services/assistant_service';
+import { handleError } from './error_handler';
 
 export function registerAgentRoutes(router: IRouter, assistantService: AssistantServiceSetup) {
   router.post(
@@ -39,18 +40,7 @@ export function registerAgentRoutes(router: IRouter, assistantService: Assistant
         );
         return res.ok({ body: response });
       } catch (e) {
-        context.assistant_plugin.logger.error('Execute agent failed!', e);
-        if (e.statusCode >= 400 && e.statusCode <= 499) {
-          return res.customError({
-            body: { message: typeof e.body === 'string' ? e.body : JSON.stringify(e.body) },
-            statusCode: e.statusCode,
-          });
-        } else {
-          return res.customError({
-            body: 'Execute agent failed!',
-            statusCode: 500,
-          });
-        }
+        return handleError(e, res, context.assistant_plugin.logger);
       }
     })
   );

server/routes/error_handler.test.ts

@@ -0,0 +1,83 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { handleError } from './error_handler';
+import { loggerMock } from '../../../../src/core/server/logging/logger.mock';
+import { AgentNotFoundError } from './errors';
+import { opensearchDashboardsResponseFactory } from '../../../../src/core/server';
+
+describe('Error handler', () => {
+  it('should return 404 not found response if error is AgentNotFoundError', () => {
+    const mockedLogger = loggerMock.create();
+    const error = handleError(
+      new AgentNotFoundError('test error'),
+      opensearchDashboardsResponseFactory,
+      mockedLogger
+    );
+    expect(error.status).toBe(404);
+    expect(error.options.body).toMatchInlineSnapshot('"Agent not found"');
+  });
+
+  it('should return 4xx with original error body', () => {
+    const mockedLogger = loggerMock.create();
+    const error = handleError(
+      {
+        statusCode: 429,
+        body: {
+          status: 429,
+          error: {
+            type: 'OpenSearchStatusException',
+            reason: 'System Error',
+            details: 'Request is throttled at model level.',
+          },
+        },
+      },
+      opensearchDashboardsResponseFactory,
+      mockedLogger
+    );
+    expect(error.status).toBe(429);
+    expect(error.options.body).toMatchInlineSnapshot(`
+    Object {
+      "message": "{\\"status\\":429,\\"error\\":{\\"type\\":\\"OpenSearchStatusException\\",\\"reason\\":\\"System Error\\",\\"details\\":\\"Request is throttled at model level.\\"}}",
+    }
+    `);
+  });
+
+  it('shuld return generic 5xx error', () => {
+    const mockedLogger = loggerMock.create();
+    const error = handleError(
+      {
+        statusCode: 502,
+        body: {
+          status: 502,
+          error: {
+            type: 'OpenSearchStatusException',
+            reason: 'System Error',
+            details: 'Some bad thing happened',
+          },
+        },
+      },
+      opensearchDashboardsResponseFactory,
+      mockedLogger
+    );
+    expect(error.status).toBe(502);
+
+    // No extra info should returned
+    expect(error.payload).toBe(undefined);
+    expect(error.options.body).toBe(undefined);
+  });
+
+  it('should return generic internalError for unhandled server-side issues', () => {
+    const mockedLogger = loggerMock.create();
+    const error = handleError(
+      new Error('Arbitrary Error'),
+      opensearchDashboardsResponseFactory,
+      mockedLogger
+    );
+    expect(error.status).toBe(500);
+    expect(error.payload).toEqual('Internal Error');
+    expect(error.options).toMatchInlineSnapshot('Object {}');
+  });
+});

server/routes/error_handler.ts

@@ -0,0 +1,33 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { Logger, OpenSearchDashboardsResponseFactory } from '../../../../src/core/server';
+import { AgentNotFoundError } from './errors';
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export const handleError = (e: any, res: OpenSearchDashboardsResponseFactory, logger: Logger) => {
+  logger.error('Error occurred', e);
+  // Handle specific type of Errors
+  if (e instanceof AgentNotFoundError) {
+    return res.notFound({ body: 'Agent not found' });
+  }
+
+  // handle http response error of calling backend API
+  if (e.statusCode) {
+    if (e.statusCode >= 400 && e.statusCode <= 499) {
+      return res.customError({
+        body: { message: typeof e.body === 'string' ? e.body : JSON.stringify(e.body) },
+        statusCode: e.statusCode,
+      });
+    } else {
+      return res.customError({
+        statusCode: e.statusCode,
+      });
+    }
+  }
+
+  // Return an general internalError for unhandled server-side issues
+  return res.internalError();
+};

server/routes/errors.ts

@@ -0,0 +1,11 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export class AgentNotFoundError extends Error {
+  constructor(message: string) {
+    super(message);
+    this.message = message;
+  }
+}

server/routes/get_agent.ts

@@ -5,64 +5,55 @@
 
 import { OpenSearchClient } from '../../../../src/core/server';
 import { ML_COMMONS_BASE_API } from '../utils/constants';
+import { AgentNotFoundError } from './errors';
 
-/**
- *
- */
 export const getAgentIdByConfigName = async (
   configName: string,
   client: OpenSearchClient['transport']
 ): Promise<string> => {
-  try {
-    const path = `${ML_COMMONS_BASE_API}/config/${configName}`;
-    const response = await client.request({
-      method: 'GET',
-      path,
-    });
+  const path = `${ML_COMMONS_BASE_API}/config/${configName}`;
+  const response = await client.request({
+    method: 'GET',
+    path,
+  });
 
-    if (
-      !response ||
-      !(response.body.ml_configuration?.agent_id || response.body.configuration?.agent_id)
-    ) {
-      throw new Error(`cannot get agent ${configName} by calling the api: ${path}`);
-    }
-    return response.body.ml_configuration?.agent_id || response.body.configuration.agent_id;
-  } catch (error) {
-    const errorMessage = JSON.stringify(error.meta?.body) || error;
-    throw new Error(`get agent ${configName} failed, reason: ${errorMessage}`);
+  if (
+    !response ||
+    !(response.body.ml_configuration?.agent_id || response.body.configuration?.agent_id)
+  ) {
+    throw new AgentNotFoundError(
+      `cannot get agent by config name ${configName} by calling the api: ${path}`
+    );
   }
+  return response.body.ml_configuration?.agent_id || response.body.configuration.agent_id;
 };
 
 export const searchAgent = async (
   { name }: { name: string },
   client: OpenSearchClient['transport']
 ) => {
-  try {
-    const requestParams = {
-      query: {
-        term: {
-          'name.keyword': name,
-        },
-      },
-      _source: ['_id'],
-      sort: {
-        created_time: 'desc',
+  const requestParams = {
+    query: {
+      term: {
+        'name.keyword': name,
       },
-      size: 1,
-    };
+    },
+    _source: ['_id'],
+    sort: {
+      created_time: 'desc',
+    },
+    size: 1,
+  };
 
-    const response = await client.request({
-      method: 'GET',
-      path: `${ML_COMMONS_BASE_API}/agents/_search`,
-      body: requestParams,
-    });
+  const path = `${ML_COMMONS_BASE_API}/agents/_search`;
+  const response = await client.request({
+    method: 'GET',
+    path,
+    body: requestParams,
+  });
 
-    if (!response || response.body.hits.total.value === 0) {
-      return undefined;
-    }
-    return response.body.hits.hits[0]._id;
-  } catch (error) {
-    const errorMessage = JSON.stringify(error.meta?.body) || error;
-    throw new Error(`search ${name} agent failed, reason: ` + errorMessage);
+  if (!response || response.body.hits.total.value === 0) {
+    throw new AgentNotFoundError(`cannot find agent by name ${name} by calling the api: ${path}`);
   }
+  return response.body.hits.hits[0]._id as string;
 };

server/routes/summary_routes.test.ts

@@ -13,6 +13,7 @@ import { SUMMARY_ASSISTANT_API } from '../../common/constants/llm';
 import { registerData2SummaryRoutes, registerSummaryAssistantRoutes } from './summary_routes';
 import { AssistantClient } from '../services/assistant_client';
 import { RequestHandlerContext } from '../../../../src/core/server';
+import * as AgentHelpers from './get_agent';
 const mockedLogger = loggerMock.create();
 
 export const createMockedAssistantClient = (
@@ -147,7 +148,7 @@ describe('test summary route', () => {
         "headers": Object {},
         "payload": Object {
           "error": "Internal Server Error",
-          "message": "Execute agent failed!",
+          "message": "An internal server error occurred.",
           "statusCode": 500,
         },
         "statusCode": 500,
@@ -167,9 +168,10 @@ describe('test summary route', () => {
         },
       },
     });
+    const spy = jest.spyOn(AgentHelpers, 'getAgentIdByConfigName').mockResolvedValue('test_agent');
     const result = (await insightRequest({
       summaryType: 'alerts',
-      insightType: 'test',
+      insightType: 'os_insight',
       summary: 'summary',
       question: 'Please summarize this alert, do not use any tool.',
       context: 'context',
@@ -185,16 +187,18 @@ describe('test summary route', () => {
         "statusCode": 429,
       }
     `);
+    spy.mockRestore();
   });
 
   it('return 4xx when executeAgent throws 4xx error in string format for insight API', async () => {
     mockedAssistantClient.executeAgent = jest.fn().mockRejectedValue({
       statusCode: 429,
       body: 'Request is throttled at model level',
     });
+    const spy = jest.spyOn(AgentHelpers, 'getAgentIdByConfigName').mockResolvedValue('test_agent');
     const result = (await insightRequest({
       summaryType: 'alerts',
-      insightType: 'test',
+      insightType: 'os_insight',
       summary: 'summary',
       question: 'Please summarize this alert, do not use any tool.',
       context: 'context',
@@ -210,16 +214,18 @@ describe('test summary route', () => {
         "statusCode": 429,
       }
     `);
+    spy.mockRestore();
   });
 
   it('return 5xx when executeAgent throws 5xx for insight API', async () => {
     mockedAssistantClient.executeAgent = jest.fn().mockRejectedValue({
       statusCode: 500,
       body: 'Server error',
     });
+    const spy = jest.spyOn(AgentHelpers, 'getAgentIdByConfigName').mockResolvedValue('test_agent');
     const result = (await insightRequest({
       summaryType: 'alerts',
-      insightType: 'test',
+      insightType: 'os_insight',
       summary: 'summary',
       question: 'Please summarize this alert, do not use any tool.',
       context: 'context',
@@ -229,12 +235,13 @@ describe('test summary route', () => {
         "headers": Object {},
         "payload": Object {
           "error": "Internal Server Error",
-          "message": "Execute agent failed!",
+          "message": "An internal server error occurred.",
           "statusCode": 500,
         },
         "statusCode": 500,
       }
     `);
+    spy.mockRestore();
   });
 
   it('return 4xx when execute agent throws 4xx error for data2Summary API', async () => {
@@ -312,7 +319,7 @@ describe('test summary route', () => {
         "headers": Object {},
         "payload": Object {
           "error": "Internal Server Error",
-          "message": "Execute agent failed!",
+          "message": "An internal server error occurred.",
           "statusCode": 500,
         },
         "statusCode": 500,