Skip to content

Commit 1e8945a

Browse files
MichaelDoyleMichaelDoyle
committed
fix(go): implement recursive schema reference resolution
This change fixes "Reference must be canonical" validation errors when running actions through the JSON path (e.g., via the Dev UI). The Go framework now: 1. Recursively resolves 'genkit:' schema references within nested objects and arrays in ResolveSchema. 2. Ensures schema resolution is performed in runJSONWithTelemetry before input validation. This ensures that strict JSON schema validators see fully-resolved definitions rather than internal Genkit reference strings.
1 parent 266c123 commit 1e8945a

3 files changed

Lines changed: 237 additions & 20 deletions

File tree

‎go/core/action.go‎

Lines changed: 26 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -237,6 +237,10 @@ func (a *Action[In, Out, Stream]) Run(ctx context.Context, input In, cb StreamCa
237237
// inject a per-call one-shot adapter; spanInit, when non-nil, is recorded as
238238
// the span's genkit:init attribute.
239239
func (a *Action[In, Out, Stream]) runWithTelemetry(ctx context.Context, input In, cb StreamCallback[Stream], fn StreamingFunc[In, Out, Stream], spanInit any) (output api.ActionRunResult[Out], err error) {
240+
return a.runWithTelemetryInternal(ctx, input, cb, fn, spanInit, nil, nil)
241+
}
242+
243+
func (a *Action[In, Out, Stream]) runWithTelemetryInternal(ctx context.Context, input In, cb StreamCallback[Stream], fn StreamingFunc[In, Out, Stream], spanInit any, inputSchema, outputSchema map[string]any) (output api.ActionRunResult[Out], err error) {
240244
logger.FromContext(ctx).Debug("Action.Run", "name", a.Name())
241245
defer func() {
242246
logger.FromContext(ctx).Debug("Action.Run",
@@ -255,16 +259,18 @@ func (a *Action[In, Out, Stream]) runWithTelemetry(ctx context.Context, input In
255259
start := time.Now()
256260
defer func() { recordActionMetrics(ctx, a.desc.Name, start, err) }()
257261

258-
var inputSchema map[string]any
259-
inputSchema, err = ResolveSchema(a.registry, a.desc.InputSchema)
260-
if err != nil {
261-
return base.Zero[Out](), NewError(INVALID_ARGUMENT, "invalid input schema for action %q: %v", a.desc.Key, err)
262+
if inputSchema == nil {
263+
inputSchema, err = ResolveSchema(a.registry, a.desc.InputSchema)
264+
if err != nil {
265+
return base.Zero[Out](), NewError(INVALID_ARGUMENT, "invalid input schema for action %q: %v", a.desc.Key, err)
266+
}
262267
}
263268

264-
var outputSchema map[string]any
265-
outputSchema, err = a.resolveOutputSchema()
266-
if err != nil {
267-
return base.Zero[Out](), err
269+
if outputSchema == nil {
270+
outputSchema, err = a.resolveOutputSchema()
271+
if err != nil {
272+
return base.Zero[Out](), err
273+
}
268274
}
269275

270276
if err = base.ValidateValue(input, inputSchema); err != nil {
@@ -350,7 +356,17 @@ func (a *Action[In, Out, Stream]) RunJSONWithTelemetry(ctx context.Context, inpu
350356
// runJSONWithTelemetry is the shared JSON execution path. fn and spanInit
351357
// follow the same contract as runWithTelemetry.
352358
func (a *Action[In, Out, Stream]) runJSONWithTelemetry(ctx context.Context, input json.RawMessage, cb StreamCallback[json.RawMessage], fn StreamingFunc[In, Out, Stream], spanInit any) (*api.ActionRunResult[json.RawMessage], error) {
353-
i, err := base.UnmarshalAndNormalize[In](input, a.desc.InputSchema)
359+
inputSchema, err := ResolveSchema(a.registry, a.desc.InputSchema)
360+
if err != nil {
361+
return nil, NewError(INVALID_ARGUMENT, "invalid input schema for action %q: %v", a.desc.Key, err)
362+
}
363+
364+
outputSchema, err := a.resolveOutputSchema()
365+
if err != nil {
366+
return nil, err
367+
}
368+
369+
i, err := base.UnmarshalAndNormalize[In](input, inputSchema)
354370
if err != nil {
355371
return nil, NewSchemaValidationError(a.desc.Key, err)
356372
}
@@ -366,7 +382,7 @@ func (a *Action[In, Out, Stream]) runJSONWithTelemetry(ctx context.Context, inpu
366382
}
367383
}
368384

369-
r, err := a.runWithTelemetry(ctx, i, scb, fn, spanInit)
385+
r, err := a.runWithTelemetryInternal(ctx, i, scb, fn, spanInit, inputSchema, outputSchema)
370386
if err != nil {
371387
return &api.ActionRunResult[json.RawMessage]{
372388
TraceId: r.TraceId,

‎go/core/core.go‎

Lines changed: 116 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -66,22 +66,130 @@ func SchemaRef(name string) map[string]any {
6666
// Returns the original schema if no $ref is present, or the resolved schema if found.
6767
// Returns an error if the schema reference cannot be resolved.
6868
func ResolveSchema(r api.Registry, schema map[string]any) (map[string]any, error) {
69+
return resolveSchema(r, schema, nil, 0)
70+
}
71+
72+
const maxSchemaDepth = 50
73+
74+
func resolveSchema(r api.Registry, schema map[string]any, seen map[uintptr]bool, depth int) (map[string]any, error) {
6975
if schema == nil {
7076
return nil, nil
7177
}
72-
ref, ok := schema["$ref"].(string)
73-
if !ok {
78+
if depth > maxSchemaDepth {
79+
return nil, fmt.Errorf("schema reference too deep (possible cycle)")
80+
}
81+
82+
if ref, ok := schema["$ref"].(string); ok {
83+
if schemaName, found := strings.CutPrefix(ref, "genkit:"); found {
84+
resolved := r.LookupSchema(schemaName)
85+
if resolved == nil {
86+
return nil, fmt.Errorf("schema %q not found", schemaName)
87+
}
88+
89+
if seen[reflect.ValueOf(resolved).Pointer()] {
90+
return schema, nil
91+
}
92+
93+
newSeen := make(map[uintptr]bool, len(seen)+1)
94+
for k, v := range seen {
95+
newSeen[k] = v
96+
}
97+
newSeen[reflect.ValueOf(schema).Pointer()] = true
98+
99+
// Recursive call to resolve any refs within the looked-up schema.
100+
return resolveSchema(r, resolved, newSeen, depth+1)
101+
}
102+
// If it's a non-genkit $ref, we return the schema as-is to allow
103+
// the underlying validator to handle it.
74104
return schema, nil
75105
}
76-
schemaName, found := strings.CutPrefix(ref, "genkit:")
77-
if !found {
106+
107+
ptr := reflect.ValueOf(schema).Pointer()
108+
if seen[ptr] {
78109
return schema, nil
79110
}
80-
resolved := r.LookupSchema(schemaName)
81-
if resolved == nil {
82-
return nil, fmt.Errorf("schema %q not found", schemaName)
111+
112+
newSeen := make(map[uintptr]bool, len(seen)+1)
113+
for k, v := range seen {
114+
newSeen[k] = v
115+
}
116+
newSeen[ptr] = true
117+
seen = newSeen
118+
119+
// Iterate and recursively resolve any nested maps or arrays.
120+
// We only clone the schema if a change is actually made.
121+
var newSchema map[string]any
122+
for k, v := range schema {
123+
resolved, err := resolveValue(r, v, seen, depth+1)
124+
if err != nil {
125+
return nil, err
126+
}
127+
128+
if newSchema == nil && !isSame(v, resolved) {
129+
newSchema = make(map[string]any, len(schema))
130+
for k2, v2 := range schema {
131+
newSchema[k2] = v2
132+
}
133+
}
134+
if newSchema != nil {
135+
newSchema[k] = resolved
136+
}
137+
}
138+
139+
if newSchema == nil {
140+
return schema, nil
141+
}
142+
return newSchema, nil
143+
}
144+
145+
func resolveValue(r api.Registry, v any, seen map[uintptr]bool, depth int) (any, error) {
146+
switch val := v.(type) {
147+
case map[string]any:
148+
return resolveSchema(r, val, seen, depth)
149+
case []any:
150+
return resolveArray(r, val, seen, depth)
151+
default:
152+
return v, nil
153+
}
154+
}
155+
156+
func resolveArray(r api.Registry, arr []any, seen map[uintptr]bool, depth int) ([]any, error) {
157+
if depth > maxSchemaDepth {
158+
return nil, fmt.Errorf("schema reference too deep (possible cycle)")
159+
}
160+
var newArray []any
161+
for i, v := range arr {
162+
resolved, err := resolveValue(r, v, seen, depth+1)
163+
if err != nil {
164+
return nil, err
165+
}
166+
167+
if newArray == nil && !isSame(v, resolved) {
168+
newArray = make([]any, len(arr))
169+
copy(newArray, arr[:i])
170+
}
171+
if newArray != nil {
172+
newArray[i] = resolved
173+
}
174+
}
175+
if newArray == nil {
176+
return arr, nil
177+
}
178+
return newArray, nil
179+
}
180+
181+
func isSame(a, b any) bool {
182+
if reflect.TypeOf(a) != reflect.TypeOf(b) {
183+
return false
184+
}
185+
va := reflect.ValueOf(a)
186+
vb := reflect.ValueOf(b)
187+
switch va.Kind() {
188+
case reflect.Map, reflect.Slice:
189+
return va.Pointer() == vb.Pointer()
190+
default:
191+
return a == b
83192
}
84-
return resolved, nil
85193
}
86194

87195
// InferSchemaMap infers a JSON schema from a Go value and converts it to a map.

‎go/core/core_test.go‎

Lines changed: 95 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package core
1818

1919
import (
20+
"reflect"
2021
"testing"
2122

2223
"github.com/firebase/genkit/go/internal/registry"
@@ -194,18 +195,110 @@ func TestResolveSchema(t *testing.T) {
194195
}
195196
})
196197

198+
t.Run("resolves multi-level genkit ref", func(t *testing.T) {
199+
r := registry.New()
200+
innerInnerSchema := map[string]any{
201+
"type": "string",
202+
}
203+
r.RegisterSchema("InnerInner", innerInnerSchema)
204+
205+
innerSchema := map[string]any{
206+
"$ref": "genkit:InnerInner",
207+
}
208+
r.RegisterSchema("Inner", innerSchema)
209+
210+
outerSchema := map[string]any{
211+
"type": "object",
212+
"properties": map[string]any{
213+
"field": map[string]any{
214+
"$ref": "genkit:Inner",
215+
},
216+
},
217+
}
218+
219+
resolved, err := ResolveSchema(r, outerSchema)
220+
221+
if err != nil {
222+
t.Fatalf("unexpected error: %v", err)
223+
}
224+
225+
props := resolved["properties"].(map[string]any)
226+
field := props["field"].(map[string]any)
227+
if diff := cmp.Diff(innerInnerSchema, field); diff != "" {
228+
t.Errorf("multi-level schema mismatch (-want +got):\n%s", diff)
229+
}
230+
})
231+
197232
t.Run("returns error for missing schema", func(t *testing.T) {
198233
r := registry.New()
199234
refSchema := map[string]any{
200235
"$ref": "genkit:NonExistent",
201236
}
202-
203237
_, err := ResolveSchema(r, refSchema)
204-
205238
if err == nil {
206239
t.Error("expected error for missing schema, got nil")
207240
}
208241
})
242+
243+
t.Run("handles circular genkit ref without stack overflow", func(t *testing.T) {
244+
r := registry.New()
245+
nodeSchema := map[string]any{
246+
"type": "object",
247+
"properties": map[string]any{
248+
"child": map[string]any{
249+
"$ref": "genkit:Node",
250+
},
251+
},
252+
}
253+
r.RegisterSchema("Node", nodeSchema)
254+
255+
resolved, err := ResolveSchema(r, nodeSchema)
256+
if err != nil {
257+
t.Fatalf("unexpected error: %v", err)
258+
}
259+
props := resolved["properties"].(map[string]any)
260+
child := props["child"].(map[string]any)
261+
if ref, ok := child["$ref"].(string); !ok || ref != "genkit:Node" {
262+
t.Errorf("expected child to remain a ref, got %v", child)
263+
}
264+
})
265+
266+
t.Run("resolves ref in nested array", func(t *testing.T) {
267+
r := registry.New()
268+
r.RegisterSchema("Item", map[string]any{"type": "string"})
269+
schema := map[string]any{
270+
"type": "array",
271+
"items": []any{
272+
map[string]any{"$ref": "genkit:Item"},
273+
},
274+
}
275+
resolved, err := ResolveSchema(r, schema)
276+
if err != nil {
277+
t.Fatalf("unexpected error: %v", err)
278+
}
279+
items := resolved["items"].([]any)
280+
item := items[0].(map[string]any)
281+
if item["type"] != "string" {
282+
t.Errorf("expected type string, got %v", item["type"])
283+
}
284+
})
285+
286+
t.Run("returns same instance when no resolution needed", func(t *testing.T) {
287+
r := registry.New()
288+
schema := map[string]any{
289+
"type": "object",
290+
"properties": map[string]any{
291+
"foo": map[string]any{"type": "string"},
292+
},
293+
}
294+
resolved, err := ResolveSchema(r, schema)
295+
if err != nil {
296+
t.Fatalf("unexpected error: %v", err)
297+
}
298+
if reflect.ValueOf(schema).Pointer() != reflect.ValueOf(resolved).Pointer() {
299+
t.Error("expected same schema instance, but it was cloned")
300+
}
301+
})
209302
}
210303

211304
func TestInferSchemaMap(t *testing.T) {

0 commit comments

Comments
 (0)