From 6b6eb116438f055cb6344d510e37138d8b998ccb Mon Sep 17 00:00:00 2001
From: Bennet Bo Fenner <bennet@zed.dev>
Date: Tue, 19 Aug 2025 20:06:09 +0200
Subject: [PATCH] agent2: Fix tool schemas for Gemini (#36507)

Release Notes:

- N/A

---------

Co-authored-by: Agus Zubiaga <agus@zed.dev>
---
 crates/agent2/src/agent2.rs       |  1 +
 crates/agent2/src/thread.rs       |  6 ++---
 crates/agent2/src/tool_schema.rs  | 43 +++++++++++++++++++++++++++++++
 crates/google_ai/src/google_ai.rs |  2 +-
 4 files changed, 48 insertions(+), 4 deletions(-)
 create mode 100644 crates/agent2/src/tool_schema.rs

diff --git a/crates/agent2/src/agent2.rs b/crates/agent2/src/agent2.rs
index f13cd1bd67..8d18da7fe1 100644
--- a/crates/agent2/src/agent2.rs
+++ b/crates/agent2/src/agent2.rs
@@ -2,6 +2,7 @@ mod agent;
 mod native_agent_server;
 mod templates;
 mod thread;
+mod tool_schema;
 mod tools;
 
 #[cfg(test)]
diff --git a/crates/agent2/src/thread.rs b/crates/agent2/src/thread.rs
index aeb600e232..d90d0bd4f8 100644
--- a/crates/agent2/src/thread.rs
+++ b/crates/agent2/src/thread.rs
@@ -1732,8 +1732,8 @@ where
     fn initial_title(&self, input: Result<Self::Input, serde_json::Value>) -> SharedString;
 
     /// Returns the JSON schema that describes the tool's input.
-    fn input_schema(&self) -> Schema {
-        schemars::schema_for!(Self::Input)
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Schema {
+        crate::tool_schema::root_schema_for::<Self::Input>(format)
     }
 
     /// Some tools rely on a provider for the underlying billing or other reasons.
@@ -1819,7 +1819,7 @@ where
     }
 
     fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
-        let mut json = serde_json::to_value(self.0.input_schema())?;
+        let mut json = serde_json::to_value(self.0.input_schema(format))?;
         adapt_schema_to_format(&mut json, format)?;
         Ok(json)
     }
diff --git a/crates/agent2/src/tool_schema.rs b/crates/agent2/src/tool_schema.rs
new file mode 100644
index 0000000000..f608336b41
--- /dev/null
+++ b/crates/agent2/src/tool_schema.rs
@@ -0,0 +1,43 @@
+use language_model::LanguageModelToolSchemaFormat;
+use schemars::{
+    JsonSchema, Schema,
+    generate::SchemaSettings,
+    transform::{Transform, transform_subschemas},
+};
+
+pub(crate) fn root_schema_for<T: JsonSchema>(format: LanguageModelToolSchemaFormat) -> Schema {
+    let mut generator = match format {
+        LanguageModelToolSchemaFormat::JsonSchema => SchemaSettings::draft07().into_generator(),
+        LanguageModelToolSchemaFormat::JsonSchemaSubset => SchemaSettings::openapi3()
+            .with(|settings| {
+                settings.meta_schema = None;
+                settings.inline_subschemas = true;
+            })
+            .with_transform(ToJsonSchemaSubsetTransform)
+            .into_generator(),
+    };
+    generator.root_schema_for::<T>()
+}
+
+#[derive(Debug, Clone)]
+struct ToJsonSchemaSubsetTransform;
+
+impl Transform for ToJsonSchemaSubsetTransform {
+    fn transform(&mut self, schema: &mut Schema) {
+        // Ensure that the type field is not an array, this happens when we use
+        // Option<T>, the type will be [T, "null"].
+        if let Some(type_field) = schema.get_mut("type")
+            && let Some(types) = type_field.as_array()
+            && let Some(first_type) = types.first()
+        {
+            *type_field = first_type.clone();
+        }
+
+        // oneOf is not supported, use anyOf instead
+        if let Some(one_of) = schema.remove("oneOf") {
+            schema.insert("anyOf".to_string(), one_of);
+        }
+
+        transform_subschemas(self, schema);
+    }
+}
diff --git a/crates/google_ai/src/google_ai.rs b/crates/google_ai/src/google_ai.rs
index 95a6daa1d9..a1b5ca3a03 100644
--- a/crates/google_ai/src/google_ai.rs
+++ b/crates/google_ai/src/google_ai.rs
@@ -266,7 +266,7 @@ pub struct CitationMetadata {
 pub struct PromptFeedback {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub block_reason: Option<String>,
-    pub safety_ratings: Vec<SafetyRating>,
+    pub safety_ratings: Option<Vec<SafetyRating>>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub block_reason_message: Option<String>,
 }