Skip to content

Commit 98e7439

Browse files
committed
fix(gemini): resolve $ref, deep-merge allOf, align e2e fixtures
1 parent 69aa07d commit 98e7439

3 files changed

Lines changed: 150 additions & 10 deletions

File tree

apps/vscode-e2e/fixtures/gemini.json

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,29 +3,44 @@
33
{
44
"match": {
55
"model": "gemini-3.1-pro-preview",
6-
"userMessage": "gemini-e2e:reasoning-on: what is 2+2? Reply with only the number."
6+
"userMessage": "gemini-e2e:reasoning-high: what is 2+2? Reply with only the number."
77
},
88
"response": {
99
"toolCalls": [
1010
{
1111
"name": "attempt_completion",
1212
"arguments": "{\"result\":\"4\"}",
13-
"id": "call_gemini_reasoning_on_done"
13+
"id": "call_gemini_reasoning_high_done"
1414
}
1515
]
1616
}
1717
},
1818
{
1919
"match": {
2020
"model": "gemini-3.1-pro-preview",
21-
"userMessage": "gemini-e2e:reasoning-off: what is 2+2? Reply with only the number."
21+
"userMessage": "gemini-e2e:reasoning-low: what is 2+2? Reply with only the number."
2222
},
2323
"response": {
2424
"toolCalls": [
2525
{
2626
"name": "attempt_completion",
2727
"arguments": "{\"result\":\"4\"}",
28-
"id": "call_gemini_reasoning_off_done"
28+
"id": "call_gemini_reasoning_low_done"
29+
}
30+
]
31+
}
32+
},
33+
{
34+
"match": {
35+
"model": "gemini-3.1-pro-preview",
36+
"userMessage": "gemini-e2e:reasoning-disable: what is 2+2? Reply with only the number."
37+
},
38+
"response": {
39+
"toolCalls": [
40+
{
41+
"name": "attempt_completion",
42+
"arguments": "{\"result\":\"4\"}",
43+
"id": "call_gemini_reasoning_disable_done"
2944
}
3045
]
3146
}

src/api/providers/__tests__/gemini-handler.spec.ts

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -493,5 +493,99 @@ describe("GeminiHandler backend support", () => {
493493
},
494494
})
495495
})
496+
497+
it("should deep-merge allOf fragments instead of overwriting earlier properties", async () => {
498+
const options = { apiProvider: "gemini" } as ApiHandlerOptions
499+
const handler = new GeminiHandler(options)
500+
const stub = vi.fn().mockReturnValue((async function* () {})())
501+
// @ts-ignore access private client
502+
handler["client"].models.generateContentStream = stub
503+
504+
await handler
505+
.createMessage("test", [] as any, {
506+
taskId: "test-task",
507+
tools: [
508+
{
509+
type: "function",
510+
function: {
511+
name: "multi_allof_tool",
512+
description: "Tool with multi-fragment allOf",
513+
parameters: {
514+
allOf: [
515+
{
516+
type: "object",
517+
properties: { a: { type: "string" } },
518+
required: ["a"],
519+
},
520+
{
521+
type: "object",
522+
properties: { b: { type: "integer" } },
523+
required: ["b"],
524+
},
525+
],
526+
},
527+
},
528+
},
529+
],
530+
})
531+
.next()
532+
533+
const schema = stub.mock.calls[0][0].config.tools[0].functionDeclarations[0].parametersJsonSchema
534+
// Both property blocks must survive the merge — previously `b` overwrote `a`
535+
expect(schema.properties).toEqual({
536+
a: { type: "string" },
537+
b: { type: "integer" },
538+
})
539+
expect(schema.required).toEqual(expect.arrayContaining(["a", "b"]))
540+
})
541+
542+
it("should resolve $ref entries before dropping $defs", async () => {
543+
const options = { apiProvider: "gemini" } as ApiHandlerOptions
544+
const handler = new GeminiHandler(options)
545+
const stub = vi.fn().mockReturnValue((async function* () {})())
546+
// @ts-ignore access private client
547+
handler["client"].models.generateContentStream = stub
548+
549+
await handler
550+
.createMessage("test", [] as any, {
551+
taskId: "test-task",
552+
tools: [
553+
{
554+
type: "function",
555+
function: {
556+
name: "ref_tool",
557+
description: "Tool with $ref",
558+
parameters: {
559+
type: "object",
560+
$defs: {
561+
Config: {
562+
type: "object",
563+
properties: { timeout: { type: "integer" } },
564+
required: ["timeout"],
565+
},
566+
},
567+
properties: {
568+
cfg: { $ref: "#/$defs/Config" },
569+
name: { type: "string" },
570+
},
571+
required: ["cfg", "name"],
572+
},
573+
},
574+
},
575+
],
576+
})
577+
.next()
578+
579+
const schema = stub.mock.calls[0][0].config.tools[0].functionDeclarations[0].parametersJsonSchema
580+
// $defs must be gone, $ref must be inlined
581+
expect(JSON.stringify(schema)).not.toContain("$defs")
582+
expect(JSON.stringify(schema)).not.toContain("$ref")
583+
expect(schema.properties.cfg).toEqual({
584+
type: "object",
585+
properties: { timeout: { type: "integer" } },
586+
required: ["timeout"],
587+
})
588+
expect(schema.properties.name).toEqual({ type: "string" })
589+
})
496590
})
497591
})

src/api/providers/gemini.ts

Lines changed: 37 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,16 +46,33 @@ const GEMINI_SCHEMA_COMPATIBILITY_DROP_KEYS = new Set([
4646
"definitions",
4747
])
4848

49-
function sanitizeSchemaForGemini(schema: unknown): unknown {
49+
function sanitizeSchemaForGemini(schema: unknown, defs?: Record<string, unknown>): unknown {
5050
if (!schema || typeof schema !== "object") {
5151
return schema
5252
}
5353

5454
if (Array.isArray(schema)) {
55-
return schema.map((item) => sanitizeSchemaForGemini(item))
55+
return schema.map((item) => sanitizeSchemaForGemini(item, defs))
5656
}
5757

5858
const source = schema as Record<string, unknown>
59+
60+
// Extract $defs / definitions from the root schema on the first call so
61+
// they can be used to resolve $ref entries encountered deeper in the tree.
62+
const resolvedDefs = defs ?? ((source.$defs ?? source.definitions) as Record<string, unknown> | undefined)
63+
64+
// Resolve local JSON Pointer $ref before any other processing.
65+
// Without this, dropping $defs leaves dangling references that Gemini rejects.
66+
if (typeof source.$ref === "string" && resolvedDefs) {
67+
const match = source.$ref.match(/^#\/(?:\$defs|definitions)\/(.+)$/)
68+
if (match) {
69+
const resolved = resolvedDefs[match[1]]
70+
if (resolved !== undefined) {
71+
return sanitizeSchemaForGemini(resolved, resolvedDefs)
72+
}
73+
}
74+
}
75+
5976
const result: Record<string, unknown> = {}
6077
let nullable = source.nullable === true
6178

@@ -67,14 +84,28 @@ function sanitizeSchemaForGemini(schema: unknown): unknown {
6784
: true
6885
})
6986
nullable = nullable || variants.length < composition.length
70-
Object.assign(result, sanitizeSchemaForGemini(variants[0] ?? {}))
87+
Object.assign(result, sanitizeSchemaForGemini(variants[0] ?? {}, resolvedDefs))
7188
}
7289

7390
if (Array.isArray(source.allOf)) {
7491
for (const variant of source.allOf) {
75-
const sanitized = sanitizeSchemaForGemini(variant)
92+
const sanitized = sanitizeSchemaForGemini(variant, resolvedDefs)
7693
if (sanitized && typeof sanitized === "object" && !Array.isArray(sanitized)) {
77-
Object.assign(result, sanitized)
94+
const s = sanitized as Record<string, unknown>
95+
// Deep-merge properties so later allOf fragments don't overwrite
96+
// earlier ones (last-write-wins Object.assign drops prior keys).
97+
if (s.properties && typeof s.properties === "object") {
98+
result.properties = {
99+
...(result.properties as Record<string, unknown> | undefined),
100+
...(s.properties as Record<string, unknown>),
101+
}
102+
}
103+
if (Array.isArray(s.required)) {
104+
const existing = Array.isArray(result.required) ? (result.required as string[]) : []
105+
result.required = [...new Set([...existing, ...(s.required as string[])])]
106+
}
107+
const { properties: _p, required: _r, ...rest } = s
108+
Object.assign(result, rest)
78109
}
79110
}
80111
}
@@ -93,7 +124,7 @@ function sanitizeSchemaForGemini(schema: unknown): unknown {
93124
continue
94125
}
95126

96-
result[key] = sanitizeSchemaForGemini(value)
127+
result[key] = sanitizeSchemaForGemini(value, resolvedDefs)
97128
}
98129

99130
if (nullable) {

0 commit comments

Comments
 (0)