From 75586c8233ff83ab0df403d13d5a674c216d0271 Mon Sep 17 00:00:00 2001
From: Steph Milovic <stephanie.milovic@elastic.co>
Date: Mon, 30 Oct 2023 17:31:47 -0500
Subject: [PATCH] [Security solution] AI connector changes  (#170150)

---
 .../server/connector_types/bedrock/bedrock.ts              | 2 ++
 .../server/connector_types/openai/openai.test.ts           | 7 +++++++
 .../server/connector_types/openai/openai.ts                | 2 ++
 3 files changed, 11 insertions(+)

diff --git a/x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts b/x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts
index 0e1235312a52c..ea0c72420b41e 100644
--- a/x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts
+++ b/x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts
@@ -93,6 +93,8 @@ export class BedrockConnector extends SubActionConnector<Config, Secrets> {
         },
         body,
         path,
+        // Despite AWS docs, this value does not always get inferred. We need to always send it
+        service: 'bedrock',
       },
       {
         secretAccessKey: this.secrets.secret,
diff --git a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts
index 0a4a6a2931d8d..e214047a5c6d5 100644
--- a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts
+++ b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts
@@ -85,6 +85,7 @@ describe('OpenAIConnector', () => {
         const response = await connector.runApi({ body: JSON.stringify(sampleOpenAiBody) });
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://api.openai.com/v1/chat/completions',
           method: 'post',
           responseSchema: RunActionResponseSchema,
@@ -102,6 +103,7 @@ describe('OpenAIConnector', () => {
         const response = await connector.runApi({ body: JSON.stringify(requestBody) });
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://api.openai.com/v1/chat/completions',
           method: 'post',
           responseSchema: RunActionResponseSchema,
@@ -118,6 +120,7 @@ describe('OpenAIConnector', () => {
         const response = await connector.runApi({ body: JSON.stringify(sampleOpenAiBody) });
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://api.openai.com/v1/chat/completions',
           method: 'post',
           responseSchema: RunActionResponseSchema,
@@ -148,6 +151,7 @@ describe('OpenAIConnector', () => {
         });
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://api.openai.com/v1/chat/completions',
           method: 'post',
           responseSchema: RunActionResponseSchema,
@@ -269,6 +273,7 @@ describe('OpenAIConnector', () => {
         const response = await connector.invokeAI(sampleOpenAiBody);
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://api.openai.com/v1/chat/completions',
           method: 'post',
           responseSchema: RunActionResponseSchema,
@@ -379,6 +384,7 @@ describe('OpenAIConnector', () => {
         const response = await connector.runApi({ body: JSON.stringify(sampleAzureAiBody) });
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://My-test-resource-123.openai.azure.com/openai/deployments/NEW-DEPLOYMENT-321/chat/completions?api-version=2023-05-15',
           method: 'post',
           responseSchema: RunActionResponseSchema,
@@ -405,6 +411,7 @@ describe('OpenAIConnector', () => {
         });
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://My-test-resource-123.openai.azure.com/openai/deployments/NEW-DEPLOYMENT-321/chat/completions?api-version=2023-05-15',
           method: 'post',
           responseSchema: RunActionResponseSchema,
diff --git a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts
index 7413ba56090a1..7680cae94db9a 100644
--- a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts
+++ b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts
@@ -114,6 +114,8 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
       method: 'post',
       responseSchema: RunActionResponseSchema,
       data: sanitizedBody,
+      // give up to 2 minutes for response
+      timeout: 120000,
       ...axiosOptions,
     });
     return response.data;