Merge pull request #184 from quangdang46/fix/gemini-schema-and-skill-routing

fix(gemini): sanitize JSON-Schema metadata and inline $ref for Gemini (#126)

Author: quangdang46

src/provider/gemini.rs

@@ -921,22 +921,76 @@ pub(crate) fn build_tools(tools: &[ToolDefinition]) -> Option<Vec<GeminiTool>> {
 }
 
 fn gemini_compatible_schema(schema: &Value) -> Value {
+    // Gemini's generateContent uses an OpenAPI 3.0 schema subset and rejects
+    // standard JSON-Schema metadata ($defs, $ref, $schema, $id, $comment, etc.).
+    // MCP servers like notion and supabase emit schemas with $defs+$ref, which
+    // would otherwise be forwarded verbatim and cause 400-class failures.
+    //
+    // Extract $defs/definitions from the root and recursively inline $ref
+    // references while stripping metadata keywords Gemini doesn't accept.
+    // See issue #126 / upstream PR #162.
+    let defs = schema
+        .as_object()
+        .and_then(|m| m.get("$defs").or_else(|| m.get("definitions")))
+        .and_then(|v| v.as_object())
+        .cloned()
+        .unwrap_or_default();
+    sanitize_for_gemini(schema, &defs, 0)
+}
+
+const GEMINI_SCHEMA_MAX_DEPTH: usize = 24;
+
+fn sanitize_for_gemini(
+    schema: &Value,
+    defs: &serde_json::Map<String, Value>,
+    depth: usize,
+) -> Value {
+    if depth > GEMINI_SCHEMA_MAX_DEPTH {
+        // Avoid blow-up on circular schemas — return permissive empty object.
+        return Value::Object(serde_json::Map::new());
+    }
     match schema {
         Value::Object(map) => {
+            // Resolve $ref by replacing the entire object with the target definition.
+            if let Some(ref_str) = map.get("$ref").and_then(|v| v.as_str()) {
+                let name = ref_str
+                    .strip_prefix("#/$defs/")
+                    .or_else(|| ref_str.strip_prefix("#/definitions/"));
+                if let Some(target_name) = name
+                    && let Some(target) = defs.get(target_name)
+                {
+                    return sanitize_for_gemini(target, defs, depth + 1);
+                }
+                // Unresolvable $ref → permissive empty object.
+                return Value::Object(serde_json::Map::new());
+            }
+
             let mut out = serde_json::Map::new();
             for (key, value) in map {
+                // Strip JSON-Schema metadata not supported by Gemini.
+                if matches!(
+                    key.as_str(),
+                    "$defs" | "definitions" | "$schema" | "$id" | "$comment" | "$anchor" | "title"
+                ) {
+                    continue;
+                }
                 if key == "const" {
                     out.insert(
                         "enum".to_string(),
-                        Value::Array(vec![gemini_compatible_schema(value)]),
+                        Value::Array(vec![sanitize_for_gemini(value, defs, depth + 1)]),
                     );
                 } else {
-                    out.insert(key.clone(), gemini_compatible_schema(value));
+                    out.insert(key.clone(), sanitize_for_gemini(value, defs, depth + 1));
                 }
             }
             Value::Object(out)
         }
-        Value::Array(items) => Value::Array(items.iter().map(gemini_compatible_schema).collect()),
+        Value::Array(items) => Value::Array(
+            items
+                .iter()
+                .map(|i| sanitize_for_gemini(i, defs, depth + 1))
+                .collect(),
+        ),
         _ => schema.clone(),
     }
 }

src/provider/gemini_tests.rs

@@ -351,3 +351,128 @@ fn parses_candidate_finish_message() {
         Some("Response blocked by safety filters")
     );
 }
+
+// ---------------------------------------------------------------------------
+// Regression tests for issue #126 / upstream PR #162: Gemini's generateContent
+// rejects standard JSON-Schema metadata. MCP servers (notion, supabase, …)
+// emit schemas with $defs/$ref/$schema, which would otherwise crash the call.
+// `gemini_compatible_schema` must inline $ref targets and strip the metadata.
+// ---------------------------------------------------------------------------
+
+#[test]
+fn gemini_schema_inlines_refs_and_strips_metadata() {
+    let schema = serde_json::json!({
+        "$schema": "https://json-schema.org/draft-07",
+        "$id": "https://example.com/foo.json",
+        "title": "Foo",
+        "type": "object",
+        "$defs": {
+            "Inner": {
+                "type": "object",
+                "properties": {
+                    "kind": { "const": "leaf" }
+                }
+            }
+        },
+        "properties": {
+            "inner": { "$ref": "#/$defs/Inner" }
+        }
+    });
+
+    let out = gemini_compatible_schema(&schema);
+    let obj = out.as_object().expect("object root");
+
+    // Metadata keywords are stripped.
+    assert!(!obj.contains_key("$schema"), "$schema should be stripped");
+    assert!(!obj.contains_key("$id"), "$id should be stripped");
+    assert!(!obj.contains_key("title"), "title should be stripped");
+    assert!(
+        !obj.contains_key("$defs"),
+        "$defs should be stripped from root"
+    );
+
+    // $ref is inlined — `inner` should now look like the Inner def.
+    let inner = obj
+        .get("properties")
+        .and_then(|p| p.get("inner"))
+        .and_then(|v| v.as_object())
+        .expect("inner object");
+    assert_eq!(inner.get("type").and_then(|v| v.as_str()), Some("object"));
+    // const → enum conversion still happens after inlining.
+    let kind = inner
+        .get("properties")
+        .and_then(|p| p.get("kind"))
+        .and_then(|v| v.as_object())
+        .expect("kind object");
+    let enum_arr = kind.get("enum").and_then(|v| v.as_array()).expect("enum");
+    assert_eq!(enum_arr.len(), 1);
+    assert_eq!(enum_arr[0].as_str(), Some("leaf"));
+}
+
+#[test]
+fn gemini_schema_resolves_definitions_alias_too() {
+    // Older drafts use `definitions` instead of `$defs`. Both must work.
+    let schema = serde_json::json!({
+        "type": "object",
+        "definitions": {
+            "Bar": { "type": "string" }
+        },
+        "properties": {
+            "bar": { "$ref": "#/definitions/Bar" }
+        }
+    });
+
+    let out = gemini_compatible_schema(&schema);
+    let obj = out.as_object().expect("object root");
+    assert!(!obj.contains_key("definitions"));
+    let bar = obj
+        .get("properties")
+        .and_then(|p| p.get("bar"))
+        .and_then(|v| v.as_object())
+        .expect("bar object");
+    assert_eq!(bar.get("type").and_then(|v| v.as_str()), Some("string"));
+}
+
+#[test]
+fn gemini_schema_falls_back_to_empty_object_on_unresolved_ref() {
+    // Unresolvable $ref must produce a permissive empty object, not a crash
+    // and not a forwarded $ref Gemini would reject.
+    let schema = serde_json::json!({
+        "type": "object",
+        "properties": {
+            "ghost": { "$ref": "#/$defs/DoesNotExist" }
+        }
+    });
+    let out = gemini_compatible_schema(&schema);
+    let ghost = out
+        .get("properties")
+        .and_then(|p| p.get("ghost"))
+        .expect("ghost");
+    let ghost_obj = ghost.as_object().expect("ghost object");
+    assert!(
+        ghost_obj.is_empty(),
+        "unresolved $ref should become an empty object, got {ghost:?}"
+    );
+}
+
+#[test]
+fn gemini_schema_does_not_recurse_forever_on_cycles() {
+    // A self-referential schema must terminate (returns an empty object once
+    // depth limit is exceeded) instead of overflowing the stack.
+    let schema = serde_json::json!({
+        "type": "object",
+        "$defs": {
+            "Loop": {
+                "type": "object",
+                "properties": {
+                    "next": { "$ref": "#/$defs/Loop" }
+                }
+            }
+        },
+        "properties": {
+            "head": { "$ref": "#/$defs/Loop" }
+        }
+    });
+    // Should not panic / overflow.
+    let _ = gemini_compatible_schema(&schema);
+}