From 699db81856eaf9a3666f4f46fa58ee4619b3b17c Mon Sep 17 00:00:00 2001
From: Andrew Macri <andrew.macri@elastic.co>
Date: Tue, 17 Sep 2024 15:46:48 -0400
Subject: [PATCH] [Security Solution] [Attack discovery] Updates Gemini
 connector to ignore unknown properties in responses (#192915)

## [Security Solution] [Attack discovery] Updates Gemini connector to ignore unknown properties in responses

### Summary

This PR updates the Gemini connector's `RunApiResponseSchema` to ignore unknown properties in responses.

It fixes an issue where the recent addition of a `modelVersion` in Gemini responses causes Attack discovery to fail with the following error:

```
Attack discovery generation failed

ActionsClientLlm: action result status is error: an error occurred while running the action - Response validation failed (Error: [modelVersion]: definition for this key is missing)
```

shown in the screenshot below:

![error](https://github.com/user-attachments/assets/8b706290-27ca-42d9-b972-3d8613af690c)

### Desk testing

To reproduce:

1) Navigate to Security > Attack discovery

2) Select a Gemini connector

3) Click Generate

**Expected result**

- Attack discoveries are generated

**Actual result**

- The following error is displayed:

```
Attack discovery generation failed

ActionsClientLlm: action result status is error: an error occurred while running the action - Response validation failed (Error: [modelVersion]: definition for this key is missing)
```

![error](https://github.com/user-attachments/assets/8b706290-27ca-42d9-b972-3d8613af690c)
---
 .../stack_connectors/common/gemini/schema.ts  | 19 +++++++-----
 .../connector_types/gemini/gemini.test.ts     | 29 +++++++++++++++++++
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/x-pack/plugins/stack_connectors/common/gemini/schema.ts b/x-pack/plugins/stack_connectors/common/gemini/schema.ts
index 43bf7337626e3e..1171cdc2a037af 100644
--- a/x-pack/plugins/stack_connectors/common/gemini/schema.ts
+++ b/x-pack/plugins/stack_connectors/common/gemini/schema.ts
@@ -29,14 +29,17 @@ export const RunActionParamsSchema = schema.object({
   raw: schema.maybe(schema.boolean()),
 });
 
-export const RunApiResponseSchema = schema.object({
-  candidates: schema.any(),
-  usageMetadata: schema.object({
-    promptTokenCount: schema.number(),
-    candidatesTokenCount: schema.number(),
-    totalTokenCount: schema.number(),
-  }),
-});
+export const RunApiResponseSchema = schema.object(
+  {
+    candidates: schema.any(),
+    usageMetadata: schema.object({
+      promptTokenCount: schema.number(),
+      candidatesTokenCount: schema.number(),
+      totalTokenCount: schema.number(),
+    }),
+  },
+  { unknowns: 'ignore' } // unknown keys will NOT fail validation, but will be removed
+);
 
 export const RunActionResponseSchema = schema.object(
   {
diff --git a/x-pack/plugins/stack_connectors/server/connector_types/gemini/gemini.test.ts b/x-pack/plugins/stack_connectors/server/connector_types/gemini/gemini.test.ts
index 94dd7aa0d153cd..5b0bff4098a208 100644
--- a/x-pack/plugins/stack_connectors/server/connector_types/gemini/gemini.test.ts
+++ b/x-pack/plugins/stack_connectors/server/connector_types/gemini/gemini.test.ts
@@ -141,6 +141,35 @@ describe('GeminiConnector', () => {
 
         expect(response).toEqual(connectorResponse);
       });
+
+      describe('RunApiResponseSchema', () => {
+        it('successfully validates a response that only has known properties', () => {
+          const onlyKnownProperties = {
+            ...defaultResponse.data,
+          };
+
+          expect(RunApiResponseSchema.validate(onlyKnownProperties)).toEqual(onlyKnownProperties);
+        });
+
+        it('fails validation when the response does NOT conform to the schema', () => {
+          const missingRequiredFields = {
+            // missing candidates and usageMetadata
+          };
+
+          expect(() => RunApiResponseSchema.validate(missingRequiredFields)).toThrowError();
+        });
+
+        it('removes unknown properties, but does NOT fail validation when they are present', () => {
+          const hasUnknownProperties = {
+            ...defaultResponse.data,
+            modelVersion: '1.0.0', // <-- an unknown property
+          };
+
+          expect(RunApiResponseSchema.validate(hasUnknownProperties)).toEqual({
+            ...defaultResponse.data,
+          });
+        });
+      });
     });
 
     describe('invokeAI', () => {