From 1779dcb4d1f8999d53e0346943a12624ed1881d7 Mon Sep 17 00:00:00 2001
From: Hugo Aguirre Parra <aguirreparra@google.com>
Date: Fri, 3 Oct 2025 21:58:51 +0000
Subject: [PATCH] fix(go/plugins/googlegenai): parse anyOf in schemas

---
 go/plugins/googlegenai/gemini.go      | 41 +++++++++++++++++++++++----
 go/plugins/googlegenai/gemini_test.go | 12 ++++++++
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/go/plugins/googlegenai/gemini.go b/go/plugins/googlegenai/gemini.go
index 9f8bb5ca59..ebf7ca5de0 100644
--- a/go/plugins/googlegenai/gemini.go
+++ b/go/plugins/googlegenai/gemini.go
@@ -456,17 +456,48 @@ func toGeminiSchema(originalSchema map[string]any, genkitSchema map[string]any)
 		return nil, nil
 	}
 	if v, ok := genkitSchema["$ref"]; ok {
-		ref := v.(string)
+		ref, ok := v.(string)
+		if !ok {
+			return nil, fmt.Errorf("invalid $ref value: not a string")
+		}
 		return toGeminiSchema(originalSchema, resolveRef(originalSchema, ref))
 	}
+
+	// Handle "anyOf" subschemas by finding the first valid schema definition
+	if v, ok := genkitSchema["anyOf"]; ok {
+		if anyOfList, isList := v.([]map[string]any); isList {
+			for _, subSchema := range anyOfList {
+				if subSchemaType, hasType := subSchema["type"]; hasType {
+					if typeStr, isString := subSchemaType.(string); isString && typeStr != "null" {
+						if title, ok := genkitSchema["title"]; ok {
+							subSchema["title"] = title
+						}
+						if description, ok := genkitSchema["description"]; ok {
+							subSchema["description"] = description
+						}
+						// Found a schema like: {"type": "string"}
+						return toGeminiSchema(originalSchema, subSchema)
+					}
+				}
+			}
+		}
+	}
+
 	schema := &genai.Schema{}
+	typeVal, ok := genkitSchema["type"]
+	if !ok {
+		return nil, fmt.Errorf("schema is missing the 'type' field: %#v", genkitSchema)
+	}
+
+	typeStr, ok := typeVal.(string)
+	if !ok {
+		return nil, fmt.Errorf("schema 'type' field is not a string, but %T", typeVal)
+	}
 
-	switch genkitSchema["type"].(string) {
+	switch typeStr {
 	case "string":
 		schema.Type = genai.TypeString
-	case "float64":
-		schema.Type = genai.TypeNumber
-	case "number":
+	case "float64", "number":
 		schema.Type = genai.TypeNumber
 	case "integer":
 		schema.Type = genai.TypeInteger
diff --git a/go/plugins/googlegenai/gemini_test.go b/go/plugins/googlegenai/gemini_test.go
index c58f56fcfc..1686306b42 100644
--- a/go/plugins/googlegenai/gemini_test.go
+++ b/go/plugins/googlegenai/gemini_test.go
@@ -75,6 +75,18 @@ func TestConvertRequest(t *testing.T) {
 					"object": map[string]any{
 						"type": string("object"),
 					},
+					"domain": map[string]any{
+						"anyOf": []map[string]any{
+							{
+								"type": string("string"),
+							},
+							{
+								"type": string("null"),
+							},
+						},
+						"default": "null",
+						"title":   string("Domain"),
+					},
 				},
 			},
 		},