openai: express nullable arrays and objects as anyOf in strict tool schemas

pkg/model/provider/openai/schema.go

@@ -2,6 +2,7 @@ package openai
 
 import (
 	"iter"
+	"reflect"
 	"slices"
 	"strconv"
 	"strings"
@@ -26,7 +27,25 @@ func ConvertParametersToSchema(params any) (shared.FunctionParameters, bool, err
 	}
 
 	strict := isStrictCompatible(p)
-	return fixSchemaArrayItems(removeFormatFields(ensureTypeFields(makeAllRequired(p)))), strict, nil
+	normalized := fixSchemaArrayItems(removeFormatFields(ensureTypeFields(makeAllRequired(p))))
+	return normalized, strict && !hasUnresolvedLocalRef(normalized), nil
+}
+
+// hasUnresolvedLocalRef is the output-side counterpart of the $ref check in
+// isStrictCompatible: normalization relocates schemas, and a strict schema
+// must never be sent with a local $ref that no longer resolves.
+func hasUnresolvedLocalRef(root map[string]any) bool {
+	unresolved := false
+	walkSchema(root, func(node map[string]any) {
+		ref, ok := node["$ref"].(string)
+		if !ok || (ref != "#" && !strings.HasPrefix(ref, "#/")) {
+			return
+		}
+		if !jsonPointerResolves(root, ref) {
+			unresolved = true
+		}
+	})
+	return unresolved
 }
 
 // childSchemas yields every direct sub-schema of node. Both the strict
@@ -133,25 +152,25 @@ func isStrictCompatibleRef(root, node map[string]any, ref any) bool {
 	if refStr != "#" && !strings.HasPrefix(refStr, "#/") {
 		return false
 	}
-	if _, ok := resolveJSONPointer(root, refStr); !ok {
+	if !jsonPointerResolves(root, refStr) {
 		return false
 	}
 	// No sibling keywords: the node must contain nothing but "$ref".
 	return len(node) == 1
 }
 
-// resolveJSONPointer resolves a local JSON pointer (e.g. "#/$defs/thing" or
-// "#") against root, per RFC 6901. It only descends through map[string]any
+// jsonPointerResolves reports whether a local JSON pointer (e.g. "#/$defs/thing" or
+// "#") resolves to a schema object within root, per RFC 6901. It only descends through map[string]any
 // and []any nodes; anything else, or a missing key/out-of-range index, fails
 // resolution rather than panicking.
-func resolveJSONPointer(root map[string]any, pointer string) (map[string]any, bool) {
+func jsonPointerResolves(root map[string]any, pointer string) bool {
 	pointer = strings.TrimPrefix(pointer, "#")
 	if pointer == "" {
-		return root, true
+		return true
 	}
 	rest, found := strings.CutPrefix(pointer, "/")
 	if !found {
-		return nil, false
+		return false
 	}
 
 	var current any = root
@@ -163,22 +182,22 @@ func resolveJSONPointer(root map[string]any, pointer string) (map[string]any, bo
 		case map[string]any:
 			v, ok := node[tok]
 			if !ok {
-				return nil, false
+				return false
 			}
 			current = v
 		case []any:
 			idx, err := strconv.Atoi(tok)
 			if err != nil || idx < 0 || idx >= len(node) {
-				return nil, false
+				return false
 			}
 			current = node[idx]
 		default:
-			return nil, false
+			return false
 		}
 	}
 
-	m, ok := current.(map[string]any)
-	return m, ok
+	_, ok := current.(map[string]any)
+	return ok
 }
 
 // walkSchema calls fn on the given schema node, then recursively walks into
@@ -202,12 +221,18 @@ func walkSchema(schema map[string]any, fn func(map[string]any)) {
 // $ref nodes are left untouched (no additionalProperties/type injection): a
 // newly-required $ref property is instead wrapped as
 // {"anyOf": [<$ref>, {"type": "null"}]}, OpenAI's documented pattern for an
-// optional reference, so the model isn't forced to always emit it.
+// optional reference, so the model isn't forced to always emit it. Nullable
+// array and object properties use the same anyOf form, whether they were
+// newly made nullable or already declared as a ["array", "null"] type pair:
+// strict mode rejects the type-pair form when the object lives behind a
+// $ref/$defs indirection. Scalar properties keep the ["T", "null"] form.
 func makeAllRequired(schema shared.FunctionParameters) shared.FunctionParameters {
 	if schema == nil {
 		schema = map[string]any{"type": "object", "properties": map[string]any{}}
 	}
 
+	wrappers := nullableWrappers{}
+
 	walkSchema(schema, func(node map[string]any) {
 		if _, ok := node["$ref"]; ok {
 			return
@@ -255,25 +280,154 @@ func makeAllRequired(schema shared.FunctionParameters) shared.FunctionParameters
 		propertyOrder := tools.SchemaPropertyNames(node)
 		for _, propName := range propertyOrder {
 			newRequired = append(newRequired, propName)
-			if !originallyRequired[propName] {
-				if propMap, ok := properties[propName].(map[string]any); ok {
-					if _, isRef := propMap["$ref"]; isRef {
-						properties[propName] = map[string]any{
-							"anyOf": []any{propMap, map[string]any{"type": "null"}},
-						}
-					} else if t, ok := propMap["type"].(string); ok {
-						propMap["type"] = []string{t, "null"}
-					}
-				}
+
+			propMap, ok := properties[propName].(map[string]any)
+			if !ok {
+				continue
+			}
+
+			if containerType, isNullableContainer := nullableContainerType(propMap); isNullableContainer {
+				propMap["type"] = containerType
+				properties[propName] = wrappers.wrap(propMap, false)
+				continue
+			}
+
+			if originallyRequired[propName] {
+				continue
+			}
+
+			_, isRef := propMap["$ref"]
+			t, isScalarType := propMap["type"].(string)
+			switch {
+			case isRef || t == "array" || t == "object":
+				properties[propName] = wrappers.wrap(propMap, true)
+			case isScalarType:
+				propMap["type"] = []string{t, "null"}
 			}
 		}
 
 		node["required"] = newRequired
 	})
 
+	wrappers.retargetLocalRefs(schema)
+
 	return schema
 }
 
+// nullableContainerType reports the container type ("array" or "object") of a
+// schema whose "type" is exactly that container type plus "null", in either
+// order and in any of the slice representations produced by schema generators.
+func nullableContainerType(schema map[string]any) (string, bool) {
+	var types []string
+	switch t := schema["type"].(type) {
+	case []string:
+		types = t
+	case []any:
+		for _, v := range t {
+			s, ok := v.(string)
+			if !ok {
+				return "", false
+			}
+			types = append(types, s)
+		}
+	default:
+		return "", false
+	}
+
+	if len(types) != 2 {
+		return "", false
+	}
+	for i, t := range types {
+		other := types[1-i]
+		if t == "null" && (other == "array" || other == "object") {
+			return other, true
+		}
+	}
+	return "", false
+}
+
+// nullableWrappers tracks the {"anyOf": [schema, {"type": "null"}]} wrappers
+// created while making properties nullable. Wrapping relocates the original
+// schema to <path>/anyOf/0, so local $refs pointing at or into it have to be
+// rewritten. The value says whether a $ref targeting the wrapped node itself
+// should keep seeing the non-null variant: true for properties that were not
+// nullable before, false for ones already declared as a nullable type pair.
+type nullableWrappers map[uintptr]bool
+
+func (w nullableWrappers) wrap(schema map[string]any, referrersSeeNonNull bool) map[string]any {
+	wrapper := map[string]any{
+		"anyOf": []any{schema, map[string]any{"type": "null"}},
+	}
+	w[reflect.ValueOf(wrapper).Pointer()] = referrersSeeNonNull
+	return wrapper
+}
+
+func (w nullableWrappers) variantOf(node any) (wrapped map[string]any, referrersSeeNonNull, isWrapper bool) {
+	wrapper, ok := node.(map[string]any)
+	if !ok {
+		return nil, false, false
+	}
+	referrersSeeNonNull, isWrapper = w[reflect.ValueOf(wrapper).Pointer()]
+	if !isWrapper {
+		return nil, false, false
+	}
+	wrapped = wrapper["anyOf"].([]any)[0].(map[string]any)
+	return wrapped, referrersSeeNonNull, true
+}
+
+func (w nullableWrappers) retargetLocalRefs(root map[string]any) {
+	if len(w) == 0 {
+		return
+	}
+	walkSchema(root, func(node map[string]any) {
+		if ref, ok := node["$ref"].(string); ok {
+			node["$ref"] = w.retarget(root, ref)
+		}
+	})
+}
+
+// retarget returns ref rewritten so it addresses the same schema it addressed
+// before wrapping. Refs that cannot be resolved are returned unchanged.
+func (w nullableWrappers) retarget(root map[string]any, ref string) string {
+	rest, isLocalPointer := strings.CutPrefix(ref, "#/")
+	if !isLocalPointer {
+		return ref
+	}
+
+	var current any = root
+	rewritten := []string{}
+	for token := range strings.SplitSeq(rest, "/") {
+		if wrapped, _, isWrapper := w.variantOf(current); isWrapper {
+			rewritten = append(rewritten, "anyOf", "0")
+			current = wrapped
+		}
+
+		key := strings.ReplaceAll(strings.ReplaceAll(token, "~1", "/"), "~0", "~")
+		switch node := current.(type) {
+		case map[string]any:
+			child, ok := node[key]
+			if !ok {
+				return ref
+			}
+			current = child
+		case []any:
+			index, err := strconv.Atoi(key)
+			if err != nil || index < 0 || index >= len(node) {
+				return ref
+			}
+			current = node[index]
+		default:
+			return ref
+		}
+		rewritten = append(rewritten, token)
+	}
+
+	if _, referrersSeeNonNull, isWrapper := w.variantOf(current); isWrapper && referrersSeeNonNull {
+		rewritten = append(rewritten, "anyOf", "0")
+	}
+	return "#/" + strings.Join(rewritten, "/")
+}
+
 // isCompositionNode reports whether node describes its shape via $ref or a
 // composition keyword (anyOf/oneOf/allOf) rather than its own "type". A
 // "type" sibling combines with these via AND, so injecting one is only safe
@@ -346,21 +500,34 @@ func fixSchemaArrayItems(schema shared.FunctionParameters) shared.FunctionParame
 			continue
 		}
 
-		checkForMissingItems := false
-		switch t := prop["type"].(type) {
-		case string:
-			checkForMissingItems = t == "array"
-		case []string:
-			checkForMissingItems = slices.Contains(t, "array")
-		}
-		if !checkForMissingItems {
-			continue
-		}
-
-		if _, ok := prop["items"]; !ok {
-			prop["items"] = map[string]any{"type": "object"}
+		addMissingArrayItems(prop)
+		if variants, ok := prop["anyOf"].([]any); ok {
+			for _, variant := range variants {
+				if variantMap, ok := variant.(map[string]any); ok {
+					addMissingArrayItems(variantMap)
+				}
+			}
 		}
 	}
 
 	return schema
 }
+
+// addMissingArrayItems gives an array schema without "items" a default
+// {"type": "object"} items schema, which OpenAI requires.
+func addMissingArrayItems(schema map[string]any) {
+	isArray := false
+	switch t := schema["type"].(type) {
+	case string:
+		isArray = t == "array"
+	case []string:
+		isArray = slices.Contains(t, "array")
+	}
+	if !isArray {
+		return
+	}
+
+	if _, ok := schema["items"]; !ok {
+		schema["items"] = map[string]any{"type": "object"}
+	}
+}