From 8d5e7db18b2e3bfae0c413d40aa42a6794ef189b Mon Sep 17 00:00:00 2001
From: robert-cronin <robert.owen.cronin@gmail.com>
Date: Fri, 17 Oct 2025 01:29:22 +0000
Subject: [PATCH] fix: handle multi-type arrays in JSON schema to prevent panic

Signed-off-by: robert-cronin <robert.owen.cronin@gmail.com>
---
 pkg/functions/grammars/json_schema.go      | 50 +++++++++++++---
 pkg/functions/grammars/json_schema_test.go | 69 ++++++++++++++++++++++
 2 files changed, 111 insertions(+), 8 deletions(-)

diff --git a/pkg/functions/grammars/json_schema.go b/pkg/functions/grammars/json_schema.go
index df05a914de15..a2fe3768ce3d 100644
--- a/pkg/functions/grammars/json_schema.go
+++ b/pkg/functions/grammars/json_schema.go
@@ -61,8 +61,26 @@ func (sc *JSONSchemaConverter) addRule(name, rule string) string {
 func (sc *JSONSchemaConverter) visit(schema map[string]interface{}, name string, rootSchema map[string]interface{}) (string, error) {
 	st, existType := schema["type"]
 	var schemaType string
+	var schemaTypes []string
 	if existType {
-		schemaType = st.(string)
+		// Handle both single type strings and arrays of types (e.g., ["string", "null"])
+		switch v := st.(type) {
+		case string:
+			// Single type: "type": "string"
+			schemaType = v
+			schemaTypes = []string{v}
+		case []interface{}:
+			// Multiple types: "type": ["string", "null"]
+			for _, item := range v {
+				if typeStr, ok := item.(string); ok {
+					schemaTypes = append(schemaTypes, typeStr)
+				}
+			}
+			// Use the first type as the primary schema type for compatibility
+			if len(schemaTypes) > 0 {
+				schemaType = schemaTypes[0]
+			}
+		}
 	}
 	ruleName := name
 	if name == "" {
@@ -176,14 +194,30 @@ func (sc *JSONSchemaConverter) visit(schema map[string]interface{}, name string,
 		rule := `"{" space "}" space`
 		return sc.addRule(ruleName, rule), nil
 	} else {
-		primitiveRule, exists := PRIMITIVE_RULES[schemaType]
-		if !exists {
-			return "", fmt.Errorf("unrecognized schema: %v (type: %s)", schema, schemaType)
-		}
-		if ruleName == "root" {
-			schemaType = "root"
+		// Handle primitive types, including multi-type arrays like ["string", "null"]
+		if len(schemaTypes) > 1 {
+			// Generate a union of multiple primitive types
+			var typeRules []string
+			for _, t := range schemaTypes {
+				primitiveRule, exists := PRIMITIVE_RULES[t]
+				if !exists {
+					return "", fmt.Errorf("unrecognized type in multi-type schema: %s (schema: %v)", t, schema)
+				}
+				typeRules = append(typeRules, primitiveRule)
+			}
+			rule := "(" + strings.Join(typeRules, " | ") + ")"
+			return sc.addRule(ruleName, rule), nil
+		} else {
+			// Single type
+			primitiveRule, exists := PRIMITIVE_RULES[schemaType]
+			if !exists {
+				return "", fmt.Errorf("unrecognized schema: %v (type: %s)", schema, schemaType)
+			}
+			if ruleName == "root" {
+				schemaType = "root"
+			}
+			return sc.addRule(schemaType, primitiveRule), nil
 		}
-		return sc.addRule(schemaType, primitiveRule), nil
 	}
 }
 func (sc *JSONSchemaConverter) resolveReference(ref string, rootSchema map[string]interface{}) (map[string]interface{}, error) {
diff --git a/pkg/functions/grammars/json_schema_test.go b/pkg/functions/grammars/json_schema_test.go
index 2b2c5d22bba6..ac8c1d4bc4a9 100644
--- a/pkg/functions/grammars/json_schema_test.go
+++ b/pkg/functions/grammars/json_schema_test.go
@@ -476,5 +476,74 @@ realvalue
 			Expect(err).To(BeNil())
 			Expect(grammar).To(ContainSubstring(`root ::= "{" space "}" space`))
 		})
+
+		It("handles multi-type array definitions like [string, null]", func() {
+			// Type defined as an array should not panic
+			multiTypeSchema := `{
+				"type": "object",
+				"properties": {
+					"street": {
+						"description": "The given street name where the company resides.",
+						"type": ["string", "null"]
+					},
+					"city": {
+						"description": "The given city where the company resides.",
+						"type": ["string", "null"]
+					}
+				}
+			}`
+
+			grammar, err := NewJSONSchemaConverter("").GrammarFromBytes([]byte(multiTypeSchema))
+			Expect(err).To(BeNil())
+			// The grammar should contain rules for both string and null types
+			Expect(grammar).To(ContainSubstring("string"))
+			Expect(grammar).To(ContainSubstring("null"))
+			// Should not panic and should generate valid grammar
+			Expect(grammar).ToNot(BeEmpty())
+		})
+
+		It("handles complex nested schema with multi-type arrays (issue #5572)", func() {
+			complexSchema := `{
+				"type": "object",
+				"properties": {
+					"companylist": {
+						"type": "array",
+						"items": {
+							"type": "object",
+							"properties": {
+								"companyname": {
+									"description": "The given name of the company.",
+									"type": "string"
+								},
+								"street": {
+									"description": "The given street name where the company resides.",
+									"type": ["string", "null"]
+								},
+								"city": {
+									"description": "The given city where the company resides.",
+									"type": ["string", "null"]
+								}
+							},
+							"additionalProperties": false,
+							"required": ["companyname", "street", "city"]
+						}
+					},
+					"filter": {
+						"description": "The type we should filter the list of companies by.",
+						"type": "string"
+					}
+				},
+				"required": ["companylist", "filter"],
+				"additionalProperties": false
+			}`
+
+			grammar, err := NewJSONSchemaConverter("").GrammarFromBytes([]byte(complexSchema))
+			Expect(err).To(BeNil())
+			// The grammar should be generated without panic
+			Expect(grammar).ToNot(BeEmpty())
+			// Should contain object and array structures
+			Expect(grammar).To(ContainSubstring("{"))
+			Expect(grammar).To(ContainSubstring("["))
+		})
 	})
 })