From 094cf55d1f35985ab4e4132656b6f49d160b0d94 Mon Sep 17 00:00:00 2001 From: naaa760 Date: Wed, 30 Jul 2025 21:52:24 +0530 Subject: [PATCH 1/4] added the new test file to the test suite by including --- .../apis/ai_assistant/index.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/index.ts b/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/index.ts index ee3473aed0954..24d35a7ec9908 100644 --- a/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/index.ts +++ b/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/index.ts @@ -38,6 +38,7 @@ export default function aiAssistantApiIntegrationTests({ // Misc. loadTestFile(require.resolve('./chat/chat.spec.ts')); + loadTestFile(require.resolve('./elasticsearch_tool/elasticsearch_tool.spec.ts')); loadTestFile(require.resolve('./complete/complete.spec.ts')); loadTestFile(require.resolve('./index_assets/index_assets.spec.ts')); loadTestFile(require.resolve('./connectors/connectors.spec.ts')); From 071a7903b2f98969665ed75061f37b99a78a66f4 Mon Sep 17 00:00:00 2001 From: naaa760 Date: Wed, 30 Jul 2025 21:58:06 +0530 Subject: [PATCH 2/4] Added the isOperationAllowed helper function --- .../server/functions/elasticsearch.ts | 26 ++++++++++++++----- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/elasticsearch.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/elasticsearch.ts index d53003431c17d..2d89237a78603 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/elasticsearch.ts +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/elasticsearch.ts @@ -9,6 +9,24 @@ import type { FunctionRegistrationParameters } from '.'; export const ELASTICSEARCH_FUNCTION_NAME = 'elasticsearch'; +/** + * Determines if the given HTTP method and path are allowed for the Elasticsearch tool. + * Only GET requests and POST requests to the "_search" endpoint are permitted. + * + * @param method - The HTTP method (GET, POST, PUT, DELETE, PATCH) + * @param path - The Elasticsearch API path + * @returns true if the operation is allowed, false otherwise + */ +function isOperationAllowed(method: string, path: string): boolean { + // Allowlist: (1) all GET requests, (2) POST requests whose *final* path segment is exactly "_search". + const [pathWithoutQuery] = path.split('?'); + const pathSegments = pathWithoutQuery.replace(/^\//, '').split('/'); + const lastPathSegment = pathSegments[pathSegments.length - 1]; + const isSearchEndpoint = lastPathSegment === '_search'; + + return method === 'GET' || (method === 'POST' && isSearchEndpoint); +} + export function registerElasticsearchFunction({ functions, resources, @@ -52,13 +70,7 @@ export function registerElasticsearchFunction({ }, }, async ({ arguments: { method, path, body } }) => { - // Allowlist: (1) all GET requests, (2) POST requests whose *final* path segment is exactly "_search". - const [pathWithoutQuery] = path.split('?'); - const pathSegments = pathWithoutQuery.replace(/^\//, '').split('/'); - const lastPathSegment = pathSegments[pathSegments.length - 1]; - const isSearchEndpoint = lastPathSegment === '_search'; - - if (method !== 'GET' && !(method === 'POST' && isSearchEndpoint)) { + if (!isOperationAllowed(method, path)) { throw new Error( 'Only GET requests or POST requests to the "_search" endpoint are permitted through this assistant function.' ); From 0c54a2c0e5963c126670e311492d048e3254b386 Mon Sep 17 00:00:00 2001 From: naaa760 Date: Wed, 30 Jul 2025 21:58:49 +0530 Subject: [PATCH 3/4] This is the new API test file that covers various scenarios for the Elasticsearch tool restrictions --- .../elasticsearch_tool.spec.ts | 426 ++++++++++++++++++ 1 file changed, 426 insertions(+) create mode 100644 x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/elasticsearch_tool/elasticsearch_tool.spec.ts diff --git a/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/elasticsearch_tool/elasticsearch_tool.spec.ts b/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/elasticsearch_tool/elasticsearch_tool.spec.ts new file mode 100644 index 0000000000000..591a7b5242389 --- /dev/null +++ b/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/elasticsearch_tool/elasticsearch_tool.spec.ts @@ -0,0 +1,426 @@ +import expect from '@kbn/expect'; +import { MessageRole, type Message } from '@kbn/observability-ai-assistant-plugin/common'; +import { ELASTICSEARCH_FUNCTION_NAME } from '@kbn/observability-ai-assistant-plugin/server/functions/elasticsearch'; +import { createLlmProxy, LlmProxy } from '../utils/create_llm_proxy'; +import type { DeploymentAgnosticFtrProviderContext } from '../../../ftr_provider_context'; + +const SYSTEM_MESSAGE = `You are an AI assistant that can help with Elasticsearch operations. You have access to the ${ELASTICSEARCH_FUNCTION_NAME} tool.`; + +export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) { + const log = getService('log'); + const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi'); + + const messages: Message[] = [ + { + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.User, + content: 'Test message', + }, + }, + ]; + + describe('/internal/observability_ai_assistant/chat - Elasticsearch Tool Restrictions', function () { + // Fails on MKI: https://github.com/elastic/kibana/issues/205581 + this.tags(['skipCloud']); + let proxy: LlmProxy; + let connectorId: string; + + before(async () => { + proxy = await createLlmProxy(log); + connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({ + port: proxy.getPort(), + }); + }); + + after(async () => { + proxy.close(); + await observabilityAIAssistantAPIClient.deleteActionConnector({ + actionId: connectorId, + }); + }); + + describe('Allowed Operations', () => { + it('should allow GET requests to any endpoint', async () => { + const simulatorPromise = proxy.interceptWithFunctionRequest({ + name: ELASTICSEARCH_FUNCTION_NAME, + arguments: () => JSON.stringify({ + method: 'GET', + path: '/_cluster/health', + }), + }); + + await observabilityAIAssistantAPIClient.editor({ + endpoint: 'POST /internal/observability_ai_assistant/chat', + params: { + body: { + name: 'test_get_request', + systemMessage: SYSTEM_MESSAGE, + messages: [ + { + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.User, + content: 'Get cluster health', + }, + }, + ], + connectorId, + functions: [ELASTICSEARCH_FUNCTION_NAME], + scopes: ['all'], + }, + }, + }); + + await proxy.waitForAllInterceptorsToHaveBeenCalled(); + const simulator = await simulatorPromise; + const requestBody = simulator.requestBody; + expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + }); + + it('should allow POST requests to _search endpoint', async () => { + const simulatorPromise = proxy.interceptWithFunctionRequest({ + name: ELASTICSEARCH_FUNCTION_NAME, + arguments: () => JSON.stringify({ + method: 'POST', + path: '/_search', + body: { query: { match_all: {} } }, + }), + }); + + await observabilityAIAssistantAPIClient.editor({ + endpoint: 'POST /internal/observability_ai_assistant/chat', + params: { + body: { + name: 'test_post_search_request', + systemMessage: SYSTEM_MESSAGE, + messages: [ + { + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.User, + content: 'Search for all documents', + }, + }, + ], + connectorId, + functions: [ELASTICSEARCH_FUNCTION_NAME], + scopes: ['all'], + }, + }, + }); + + await proxy.waitForAllInterceptorsToHaveBeenCalled(); + const simulator = await simulatorPromise; + const requestBody = simulator.requestBody; + expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + }); + + it('should allow POST requests to nested _search endpoints', async () => { + const simulatorPromise = proxy.interceptWithFunctionRequest({ + name: ELASTICSEARCH_FUNCTION_NAME, + arguments: () => JSON.stringify({ + method: 'POST', + path: '/my-index/_search', + body: { query: { match_all: {} } }, + }), + }); + + await observabilityAIAssistantAPIClient.editor({ + endpoint: 'POST /internal/observability_ai_assistant/chat', + params: { + body: { + name: 'test_post_nested_search_request', + systemMessage: SYSTEM_MESSAGE, + messages: [ + { + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.User, + content: 'Search in my-index', + }, + }, + ], + connectorId, + functions: [ELASTICSEARCH_FUNCTION_NAME], + scopes: ['all'], + }, + }, + }); + + await proxy.waitForAllInterceptorsToHaveBeenCalled(); + const simulator = await simulatorPromise; + const requestBody = simulator.requestBody; + expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + }); + }); + + describe('Disallowed Operations', () => { + it('should reject PUT requests', async () => { + const simulatorPromise = proxy.interceptWithFunctionRequest({ + name: ELASTICSEARCH_FUNCTION_NAME, + arguments: () => JSON.stringify({ + method: 'PUT', + path: '/my-index', + body: { mappings: { properties: {} } }, + }), + }); + + await observabilityAIAssistantAPIClient.editor({ + endpoint: 'POST /internal/observability_ai_assistant/chat', + params: { + body: { + name: 'test_put_request', + systemMessage: SYSTEM_MESSAGE, + messages: [ + { + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.User, + content: 'Create an index', + }, + }, + ], + connectorId, + functions: [ELASTICSEARCH_FUNCTION_NAME], + scopes: ['all'], + }, + }, + }); + + await proxy.waitForAllInterceptorsToHaveBeenCalled(); + const simulator = await simulatorPromise; + const requestBody = simulator.requestBody; + expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + }); + + it('should reject DELETE requests', async () => { + const simulatorPromise = proxy.interceptWithFunctionRequest({ + name: ELASTICSEARCH_FUNCTION_NAME, + arguments: () => JSON.stringify({ + method: 'DELETE', + path: '/my-index', + }), + }); + + await observabilityAIAssistantAPIClient.editor({ + endpoint: 'POST /internal/observability_ai_assistant/chat', + params: { + body: { + name: 'test_delete_request', + systemMessage: SYSTEM_MESSAGE, + messages: [ + { + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.User, + content: 'Delete an index', + }, + }, + ], + connectorId, + functions: [ELASTICSEARCH_FUNCTION_NAME], + scopes: ['all'], + }, + }, + }); + + await proxy.waitForAllInterceptorsToHaveBeenCalled(); + const simulator = await simulatorPromise; + const requestBody = simulator.requestBody; + expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + }); + + it('should reject PATCH requests', async () => { + const simulatorPromise = proxy.interceptWithFunctionRequest({ + name: ELASTICSEARCH_FUNCTION_NAME, + arguments: () => JSON.stringify({ + method: 'PATCH', + path: '/my-index/_doc/1', + body: { doc: { field: 'value' } }, + }), + }); + + await observabilityAIAssistantAPIClient.editor({ + endpoint: 'POST /internal/observability_ai_assistant/chat', + params: { + body: { + name: 'test_patch_request', + systemMessage: SYSTEM_MESSAGE, + messages: [ + { + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.User, + content: 'Update a document', + }, + }, + ], + connectorId, + functions: [ELASTICSEARCH_FUNCTION_NAME], + scopes: ['all'], + }, + }, + }); + + await proxy.waitForAllInterceptorsToHaveBeenCalled(); + const simulator = await simulatorPromise; + const requestBody = simulator.requestBody; + expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + }); + + it('should reject POST requests to non-search endpoints', async () => { + const simulatorPromise = proxy.interceptWithFunctionRequest({ + name: ELASTICSEARCH_FUNCTION_NAME, + arguments: () => JSON.stringify({ + method: 'POST', + path: '/my-index/_doc', + body: { field: 'value' }, + }), + }); + + await observabilityAIAssistantAPIClient.editor({ + endpoint: 'POST /internal/observability_ai_assistant/chat', + params: { + body: { + name: 'test_post_non_search_request', + systemMessage: SYSTEM_MESSAGE, + messages: [ + { + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.User, + content: 'Index a document', + }, + }, + ], + connectorId, + functions: [ELASTICSEARCH_FUNCTION_NAME], + scopes: ['all'], + }, + }, + }); + + await proxy.waitForAllInterceptorsToHaveBeenCalled(); + const simulator = await simulatorPromise; + const requestBody = simulator.requestBody; + expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + }); + + it('should reject POST requests to endpoints ending with _search but not being _search', async () => { + const simulatorPromise = proxy.interceptWithFunctionRequest({ + name: ELASTICSEARCH_FUNCTION_NAME, + arguments: () => JSON.stringify({ + method: 'POST', + path: '/my-index/custom_search', + body: { query: { match_all: {} } }, + }), + }); + + await observabilityAIAssistantAPIClient.editor({ + endpoint: 'POST /internal/observability_ai_assistant/chat', + params: { + body: { + name: 'test_post_custom_search_request', + systemMessage: SYSTEM_MESSAGE, + messages: [ + { + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.User, + content: 'Search with custom endpoint', + }, + }, + ], + connectorId, + functions: [ELASTICSEARCH_FUNCTION_NAME], + scopes: ['all'], + }, + }, + }); + + await proxy.waitForAllInterceptorsToHaveBeenCalled(); + const simulator = await simulatorPromise; + const requestBody = simulator.requestBody; + expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + }); + }); + + describe('Edge Cases', () => { + it('should handle paths with query parameters correctly', async () => { + const simulatorPromise = proxy.interceptWithFunctionRequest({ + name: ELASTICSEARCH_FUNCTION_NAME, + arguments: () => JSON.stringify({ + method: 'GET', + path: '/_cluster/health?pretty=true', + }), + }); + + await observabilityAIAssistantAPIClient.editor({ + endpoint: 'POST /internal/observability_ai_assistant/chat', + params: { + body: { + name: 'test_get_request_with_query_params', + systemMessage: SYSTEM_MESSAGE, + messages: [ + { + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.User, + content: 'Get cluster health with pretty format', + }, + }, + ], + connectorId, + functions: [ELASTICSEARCH_FUNCTION_NAME], + scopes: ['all'], + }, + }, + }); + + await proxy.waitForAllInterceptorsToHaveBeenCalled(); + const simulator = await simulatorPromise; + const requestBody = simulator.requestBody; + expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + }); + + it('should handle POST _search with query parameters', async () => { + const simulatorPromise = proxy.interceptWithFunctionRequest({ + name: ELASTICSEARCH_FUNCTION_NAME, + arguments: () => JSON.stringify({ + method: 'POST', + path: '/_search?scroll=1m', + body: { query: { match_all: {} } }, + }), + }); + + await observabilityAIAssistantAPIClient.editor({ + endpoint: 'POST /internal/observability_ai_assistant/chat', + params: { + body: { + name: 'test_post_search_with_query_params', + systemMessage: SYSTEM_MESSAGE, + messages: [ + { + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.User, + content: 'Search with scroll parameter', + }, + }, + ], + connectorId, + functions: [ELASTICSEARCH_FUNCTION_NAME], + scopes: ['all'], + }, + }, + }); + + await proxy.waitForAllInterceptorsToHaveBeenCalled(); + const simulator = await simulatorPromise; + const requestBody = simulator.requestBody; + expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + }); + }); + }); +} From 279413875d7928e11f47306b1e5fe42b643071fe Mon Sep 17 00:00:00 2001 From: naaa760 Date: Thu, 7 Aug 2025 12:50:30 +0530 Subject: [PATCH 4/4] fix: update elasticsearch tool tests to properly validate function restrictions - Replace system message validation with actual function execution testing - Use invokeChatCompleteWithFunctionRequest helper as suggested by reviewer - Test that allowed operations (GET, POST to _search) return success responses - Test that disallowed operations (PUT, DELETE, PATCH, POST to non-search) return proper error messages - Verify actual security restrictions rather than just system configuration - Addresses reviewer feedback about validating function call responses vs system messages --- .../elasticsearch_tool.spec.ts | 542 +++++++----------- 1 file changed, 222 insertions(+), 320 deletions(-) diff --git a/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/elasticsearch_tool/elasticsearch_tool.spec.ts b/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/elasticsearch_tool/elasticsearch_tool.spec.ts index 591a7b5242389..54a037890658e 100644 --- a/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/elasticsearch_tool/elasticsearch_tool.spec.ts +++ b/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/elasticsearch_tool/elasticsearch_tool.spec.ts @@ -2,25 +2,17 @@ import expect from '@kbn/expect'; import { MessageRole, type Message } from '@kbn/observability-ai-assistant-plugin/common'; import { ELASTICSEARCH_FUNCTION_NAME } from '@kbn/observability-ai-assistant-plugin/server/functions/elasticsearch'; import { createLlmProxy, LlmProxy } from '../utils/create_llm_proxy'; +import { + getMessageAddedEvents, + invokeChatCompleteWithFunctionRequest, +} from '../utils/conversation'; import type { DeploymentAgnosticFtrProviderContext } from '../../../ftr_provider_context'; -const SYSTEM_MESSAGE = `You are an AI assistant that can help with Elasticsearch operations. You have access to the ${ELASTICSEARCH_FUNCTION_NAME} tool.`; - export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) { const log = getService('log'); const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi'); - const messages: Message[] = [ - { - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.User, - content: 'Test message', - }, - }, - ]; - - describe('/internal/observability_ai_assistant/chat - Elasticsearch Tool Restrictions', function () { + describe('tool: elasticsearch tool restrictions', function () { // Fails on MKI: https://github.com/elastic/kibana/issues/205581 this.tags(['skipCloud']); let proxy: LlmProxy; @@ -42,384 +34,294 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon describe('Allowed Operations', () => { it('should allow GET requests to any endpoint', async () => { - const simulatorPromise = proxy.interceptWithFunctionRequest({ - name: ELASTICSEARCH_FUNCTION_NAME, - arguments: () => JSON.stringify({ - method: 'GET', - path: '/_cluster/health', - }), - }); - - await observabilityAIAssistantAPIClient.editor({ - endpoint: 'POST /internal/observability_ai_assistant/chat', - params: { - body: { - name: 'test_get_request', - systemMessage: SYSTEM_MESSAGE, - messages: [ - { - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.User, - content: 'Get cluster health', - }, - }, - ], - connectorId, - functions: [ELASTICSEARCH_FUNCTION_NAME], - scopes: ['all'], - }, + void proxy.interceptWithResponse('Hello from LLM Proxy'); + + const responseBody = await invokeChatCompleteWithFunctionRequest({ + connectorId, + observabilityAIAssistantAPIClient, + functionCall: { + name: ELASTICSEARCH_FUNCTION_NAME, + trigger: MessageRole.User, + arguments: JSON.stringify({ + method: 'GET', + path: '/_cluster/health', + }), }, }); await proxy.waitForAllInterceptorsToHaveBeenCalled(); - const simulator = await simulatorPromise; - const requestBody = simulator.requestBody; - expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + + const events = getMessageAddedEvents(responseBody); + const esFunctionResponse = events[0]; + + expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME); + // Function should succeed - response should contain actual data, not an error + const content = JSON.parse(esFunctionResponse.message.message.content!); + expect(content).to.have.property('response'); + expect(content.response).to.not.have.property('error'); }); it('should allow POST requests to _search endpoint', async () => { - const simulatorPromise = proxy.interceptWithFunctionRequest({ - name: ELASTICSEARCH_FUNCTION_NAME, - arguments: () => JSON.stringify({ - method: 'POST', - path: '/_search', - body: { query: { match_all: {} } }, - }), - }); - - await observabilityAIAssistantAPIClient.editor({ - endpoint: 'POST /internal/observability_ai_assistant/chat', - params: { - body: { - name: 'test_post_search_request', - systemMessage: SYSTEM_MESSAGE, - messages: [ - { - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.User, - content: 'Search for all documents', - }, - }, - ], - connectorId, - functions: [ELASTICSEARCH_FUNCTION_NAME], - scopes: ['all'], - }, + void proxy.interceptWithResponse('Hello from LLM Proxy'); + + const responseBody = await invokeChatCompleteWithFunctionRequest({ + connectorId, + observabilityAIAssistantAPIClient, + functionCall: { + name: ELASTICSEARCH_FUNCTION_NAME, + trigger: MessageRole.User, + arguments: JSON.stringify({ + method: 'POST', + path: '/_search', + body: { query: { match_all: {} } }, + }), }, }); await proxy.waitForAllInterceptorsToHaveBeenCalled(); - const simulator = await simulatorPromise; - const requestBody = simulator.requestBody; - expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + + const events = getMessageAddedEvents(responseBody); + const esFunctionResponse = events[0]; + + expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME); + // Function should succeed - response should contain actual data, not an error + const content = JSON.parse(esFunctionResponse.message.message.content!); + expect(content).to.have.property('response'); + expect(content.response).to.not.have.property('error'); }); it('should allow POST requests to nested _search endpoints', async () => { - const simulatorPromise = proxy.interceptWithFunctionRequest({ - name: ELASTICSEARCH_FUNCTION_NAME, - arguments: () => JSON.stringify({ - method: 'POST', - path: '/my-index/_search', - body: { query: { match_all: {} } }, - }), - }); - - await observabilityAIAssistantAPIClient.editor({ - endpoint: 'POST /internal/observability_ai_assistant/chat', - params: { - body: { - name: 'test_post_nested_search_request', - systemMessage: SYSTEM_MESSAGE, - messages: [ - { - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.User, - content: 'Search in my-index', - }, - }, - ], - connectorId, - functions: [ELASTICSEARCH_FUNCTION_NAME], - scopes: ['all'], - }, + void proxy.interceptWithResponse('Hello from LLM Proxy'); + + const responseBody = await invokeChatCompleteWithFunctionRequest({ + connectorId, + observabilityAIAssistantAPIClient, + functionCall: { + name: ELASTICSEARCH_FUNCTION_NAME, + trigger: MessageRole.User, + arguments: JSON.stringify({ + method: 'POST', + path: '/my-index/_search', + body: { query: { match_all: {} } }, + }), }, }); await proxy.waitForAllInterceptorsToHaveBeenCalled(); - const simulator = await simulatorPromise; - const requestBody = simulator.requestBody; - expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + + const events = getMessageAddedEvents(responseBody); + const esFunctionResponse = events[0]; + + expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME); + // Function should succeed - response should contain actual data, not an error + const content = JSON.parse(esFunctionResponse.message.message.content!); + expect(content).to.have.property('response'); + expect(content.response).to.not.have.property('error'); }); }); describe('Disallowed Operations', () => { it('should reject PUT requests', async () => { - const simulatorPromise = proxy.interceptWithFunctionRequest({ - name: ELASTICSEARCH_FUNCTION_NAME, - arguments: () => JSON.stringify({ - method: 'PUT', - path: '/my-index', - body: { mappings: { properties: {} } }, - }), - }); - - await observabilityAIAssistantAPIClient.editor({ - endpoint: 'POST /internal/observability_ai_assistant/chat', - params: { - body: { - name: 'test_put_request', - systemMessage: SYSTEM_MESSAGE, - messages: [ - { - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.User, - content: 'Create an index', - }, - }, - ], - connectorId, - functions: [ELASTICSEARCH_FUNCTION_NAME], - scopes: ['all'], - }, + void proxy.interceptWithResponse('Hello from LLM Proxy'); + + const responseBody = await invokeChatCompleteWithFunctionRequest({ + connectorId, + observabilityAIAssistantAPIClient, + functionCall: { + name: ELASTICSEARCH_FUNCTION_NAME, + trigger: MessageRole.User, + arguments: JSON.stringify({ + method: 'PUT', + path: '/my-index', + body: { mappings: { properties: {} } }, + }), }, }); await proxy.waitForAllInterceptorsToHaveBeenCalled(); - const simulator = await simulatorPromise; - const requestBody = simulator.requestBody; - expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + + const events = getMessageAddedEvents(responseBody); + const esFunctionResponse = events[0]; + + expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME); + // Function should fail - response should contain an error + const content = JSON.parse(esFunctionResponse.message.message.content!); + expect(content).to.have.property('error'); + expect(content.error).to.contain('Only GET requests or POST requests to the "_search" endpoint are permitted'); }); it('should reject DELETE requests', async () => { - const simulatorPromise = proxy.interceptWithFunctionRequest({ - name: ELASTICSEARCH_FUNCTION_NAME, - arguments: () => JSON.stringify({ - method: 'DELETE', - path: '/my-index', - }), - }); - - await observabilityAIAssistantAPIClient.editor({ - endpoint: 'POST /internal/observability_ai_assistant/chat', - params: { - body: { - name: 'test_delete_request', - systemMessage: SYSTEM_MESSAGE, - messages: [ - { - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.User, - content: 'Delete an index', - }, - }, - ], - connectorId, - functions: [ELASTICSEARCH_FUNCTION_NAME], - scopes: ['all'], - }, + void proxy.interceptWithResponse('Hello from LLM Proxy'); + + const responseBody = await invokeChatCompleteWithFunctionRequest({ + connectorId, + observabilityAIAssistantAPIClient, + functionCall: { + name: ELASTICSEARCH_FUNCTION_NAME, + trigger: MessageRole.User, + arguments: JSON.stringify({ + method: 'DELETE', + path: '/my-index', + }), }, }); await proxy.waitForAllInterceptorsToHaveBeenCalled(); - const simulator = await simulatorPromise; - const requestBody = simulator.requestBody; - expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + + const events = getMessageAddedEvents(responseBody); + const esFunctionResponse = events[0]; + + expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME); + // Function should fail - response should contain an error + const content = JSON.parse(esFunctionResponse.message.message.content!); + expect(content).to.have.property('error'); + expect(content.error).to.contain('Only GET requests or POST requests to the "_search" endpoint are permitted'); }); it('should reject PATCH requests', async () => { - const simulatorPromise = proxy.interceptWithFunctionRequest({ - name: ELASTICSEARCH_FUNCTION_NAME, - arguments: () => JSON.stringify({ - method: 'PATCH', - path: '/my-index/_doc/1', - body: { doc: { field: 'value' } }, - }), - }); - - await observabilityAIAssistantAPIClient.editor({ - endpoint: 'POST /internal/observability_ai_assistant/chat', - params: { - body: { - name: 'test_patch_request', - systemMessage: SYSTEM_MESSAGE, - messages: [ - { - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.User, - content: 'Update a document', - }, - }, - ], - connectorId, - functions: [ELASTICSEARCH_FUNCTION_NAME], - scopes: ['all'], - }, + void proxy.interceptWithResponse('Hello from LLM Proxy'); + + const responseBody = await invokeChatCompleteWithFunctionRequest({ + connectorId, + observabilityAIAssistantAPIClient, + functionCall: { + name: ELASTICSEARCH_FUNCTION_NAME, + trigger: MessageRole.User, + arguments: JSON.stringify({ + method: 'PATCH', + path: '/my-index/_doc/1', + body: { doc: { field: 'value' } }, + }), }, }); await proxy.waitForAllInterceptorsToHaveBeenCalled(); - const simulator = await simulatorPromise; - const requestBody = simulator.requestBody; - expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + + const events = getMessageAddedEvents(responseBody); + const esFunctionResponse = events[0]; + + expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME); + // Function should fail - response should contain an error + const content = JSON.parse(esFunctionResponse.message.message.content!); + expect(content).to.have.property('error'); + expect(content.error).to.contain('Only GET requests or POST requests to the "_search" endpoint are permitted'); }); it('should reject POST requests to non-search endpoints', async () => { - const simulatorPromise = proxy.interceptWithFunctionRequest({ - name: ELASTICSEARCH_FUNCTION_NAME, - arguments: () => JSON.stringify({ - method: 'POST', - path: '/my-index/_doc', - body: { field: 'value' }, - }), - }); - - await observabilityAIAssistantAPIClient.editor({ - endpoint: 'POST /internal/observability_ai_assistant/chat', - params: { - body: { - name: 'test_post_non_search_request', - systemMessage: SYSTEM_MESSAGE, - messages: [ - { - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.User, - content: 'Index a document', - }, - }, - ], - connectorId, - functions: [ELASTICSEARCH_FUNCTION_NAME], - scopes: ['all'], - }, + void proxy.interceptWithResponse('Hello from LLM Proxy'); + + const responseBody = await invokeChatCompleteWithFunctionRequest({ + connectorId, + observabilityAIAssistantAPIClient, + functionCall: { + name: ELASTICSEARCH_FUNCTION_NAME, + trigger: MessageRole.User, + arguments: JSON.stringify({ + method: 'POST', + path: '/my-index/_doc', + body: { field: 'value' }, + }), }, }); await proxy.waitForAllInterceptorsToHaveBeenCalled(); - const simulator = await simulatorPromise; - const requestBody = simulator.requestBody; - expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + + const events = getMessageAddedEvents(responseBody); + const esFunctionResponse = events[0]; + + expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME); + // Function should fail - response should contain an error + const content = JSON.parse(esFunctionResponse.message.message.content!); + expect(content).to.have.property('error'); + expect(content.error).to.contain('Only GET requests or POST requests to the "_search" endpoint are permitted'); }); it('should reject POST requests to endpoints ending with _search but not being _search', async () => { - const simulatorPromise = proxy.interceptWithFunctionRequest({ - name: ELASTICSEARCH_FUNCTION_NAME, - arguments: () => JSON.stringify({ - method: 'POST', - path: '/my-index/custom_search', - body: { query: { match_all: {} } }, - }), - }); - - await observabilityAIAssistantAPIClient.editor({ - endpoint: 'POST /internal/observability_ai_assistant/chat', - params: { - body: { - name: 'test_post_custom_search_request', - systemMessage: SYSTEM_MESSAGE, - messages: [ - { - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.User, - content: 'Search with custom endpoint', - }, - }, - ], - connectorId, - functions: [ELASTICSEARCH_FUNCTION_NAME], - scopes: ['all'], - }, + void proxy.interceptWithResponse('Hello from LLM Proxy'); + + const responseBody = await invokeChatCompleteWithFunctionRequest({ + connectorId, + observabilityAIAssistantAPIClient, + functionCall: { + name: ELASTICSEARCH_FUNCTION_NAME, + trigger: MessageRole.User, + arguments: JSON.stringify({ + method: 'POST', + path: '/my-index/custom_search', + body: { query: { match_all: {} } }, + }), }, }); await proxy.waitForAllInterceptorsToHaveBeenCalled(); - const simulator = await simulatorPromise; - const requestBody = simulator.requestBody; - expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + + const events = getMessageAddedEvents(responseBody); + const esFunctionResponse = events[0]; + + expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME); + // Function should fail - response should contain an error + const content = JSON.parse(esFunctionResponse.message.message.content!); + expect(content).to.have.property('error'); + expect(content.error).to.contain('Only GET requests or POST requests to the "_search" endpoint are permitted'); }); }); describe('Edge Cases', () => { it('should handle paths with query parameters correctly', async () => { - const simulatorPromise = proxy.interceptWithFunctionRequest({ - name: ELASTICSEARCH_FUNCTION_NAME, - arguments: () => JSON.stringify({ - method: 'GET', - path: '/_cluster/health?pretty=true', - }), - }); - - await observabilityAIAssistantAPIClient.editor({ - endpoint: 'POST /internal/observability_ai_assistant/chat', - params: { - body: { - name: 'test_get_request_with_query_params', - systemMessage: SYSTEM_MESSAGE, - messages: [ - { - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.User, - content: 'Get cluster health with pretty format', - }, - }, - ], - connectorId, - functions: [ELASTICSEARCH_FUNCTION_NAME], - scopes: ['all'], - }, + void proxy.interceptWithResponse('Hello from LLM Proxy'); + + const responseBody = await invokeChatCompleteWithFunctionRequest({ + connectorId, + observabilityAIAssistantAPIClient, + functionCall: { + name: ELASTICSEARCH_FUNCTION_NAME, + trigger: MessageRole.User, + arguments: JSON.stringify({ + method: 'GET', + path: '/_cluster/health?pretty=true', + }), }, }); await proxy.waitForAllInterceptorsToHaveBeenCalled(); - const simulator = await simulatorPromise; - const requestBody = simulator.requestBody; - expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + + const events = getMessageAddedEvents(responseBody); + const esFunctionResponse = events[0]; + + expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME); + // Function should succeed - response should contain actual data, not an error + const content = JSON.parse(esFunctionResponse.message.message.content!); + expect(content).to.have.property('response'); + expect(content.response).to.not.have.property('error'); }); it('should handle POST _search with query parameters', async () => { - const simulatorPromise = proxy.interceptWithFunctionRequest({ - name: ELASTICSEARCH_FUNCTION_NAME, - arguments: () => JSON.stringify({ - method: 'POST', - path: '/_search?scroll=1m', - body: { query: { match_all: {} } }, - }), - }); - - await observabilityAIAssistantAPIClient.editor({ - endpoint: 'POST /internal/observability_ai_assistant/chat', - params: { - body: { - name: 'test_post_search_with_query_params', - systemMessage: SYSTEM_MESSAGE, - messages: [ - { - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.User, - content: 'Search with scroll parameter', - }, - }, - ], - connectorId, - functions: [ELASTICSEARCH_FUNCTION_NAME], - scopes: ['all'], - }, + void proxy.interceptWithResponse('Hello from LLM Proxy'); + + const responseBody = await invokeChatCompleteWithFunctionRequest({ + connectorId, + observabilityAIAssistantAPIClient, + functionCall: { + name: ELASTICSEARCH_FUNCTION_NAME, + trigger: MessageRole.User, + arguments: JSON.stringify({ + method: 'POST', + path: '/_search?scroll=1m', + body: { query: { match_all: {} } }, + }), }, }); await proxy.waitForAllInterceptorsToHaveBeenCalled(); - const simulator = await simulatorPromise; - const requestBody = simulator.requestBody; - expect(requestBody.messages[0].content).to.eql(SYSTEM_MESSAGE); + + const events = getMessageAddedEvents(responseBody); + const esFunctionResponse = events[0]; + + expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME); + // Function should succeed - response should contain actual data, not an error + const content = JSON.parse(esFunctionResponse.message.message.content!); + expect(content).to.have.property('response'); + expect(content.response).to.not.have.property('error'); }); }); });