From a5f0d333d7f26f2966ed511d5d9def7a1933f0c2 Mon Sep 17 00:00:00 2001
From: Xuan Yang <xygoogle@google.com>
Date: Tue, 13 Jan 2026 09:50:26 -0800
Subject: [PATCH] feat: Use json schema for RestApiTool declaration when
 feature enabled

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 855767527
---
 .../openapi_spec_parser/rest_api_tool.py      | 17 ++++++--
 .../openapi_spec_parser/test_rest_api_tool.py | 41 +++++++++++++++++++
 2 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py b/src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py
index 27c6acda..a2340b95 100644
--- a/src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py
+++ b/src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py
@@ -33,6 +33,8 @@ from typing_extensions import override
 from ....agents.readonly_context import ReadonlyContext
 from ....auth.auth_credential import AuthCredential
 from ....auth.auth_schemes import AuthScheme
+from ....features import FeatureName
+from ....features import is_feature_enabled
 from ..._gemini_schema_util import _to_gemini_schema
 from ..._gemini_schema_util import _to_snake_case
 from ...base_tool import BaseTool
@@ -221,10 +223,17 @@ class RestApiTool(BaseTool):
   def _get_declaration(self) -> FunctionDeclaration:
     """Returns the function declaration in the Gemini Schema format."""
     schema_dict = self._operation_parser.get_json_schema()
-    parameters = _to_gemini_schema(schema_dict)
-    function_decl = FunctionDeclaration(
-        name=self.name, description=self.description, parameters=parameters
-    )
+    if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
+      function_decl = FunctionDeclaration(
+          name=self.name,
+          description=self.description,
+          parameters_json_schema=schema_dict,
+      )
+    else:
+      parameters = _to_gemini_schema(schema_dict)
+      function_decl = FunctionDeclaration(
+          name=self.name, description=self.description, parameters=parameters
+      )
     return function_decl
 
   def configure_auth_scheme(
diff --git a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_rest_api_tool.py b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_rest_api_tool.py
index ddf09aeb..309d7c37 100644
--- a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_rest_api_tool.py
+++ b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_rest_api_tool.py
@@ -29,6 +29,8 @@ from google.adk.auth.auth_credential import AuthCredential
 from google.adk.auth.auth_credential import AuthCredentialTypes
 from google.adk.auth.auth_credential import HttpAuth
 from google.adk.auth.auth_credential import HttpCredentials
+from google.adk.features import FeatureName
+from google.adk.features._feature_registry import temporary_feature_override
 from google.adk.sessions.state import State
 from google.adk.tools.openapi_tool.auth.auth_helpers import token_to_scheme_credential
 from google.adk.tools.openapi_tool.common.common import ApiParameter
@@ -204,6 +206,45 @@ class TestRestApiTool:
     assert declaration.description == "Test description"
     assert isinstance(declaration.parameters, Schema)
 
+  def test_get_declaration_with_json_schema_feature_enabled(
+      self, sample_endpoint, sample_operation
+  ):
+    """Test that _get_declaration uses parameters_json_schema when feature is enabled."""
+    mock_parser = MagicMock(spec=OperationParser)
+    mock_parser.get_json_schema.return_value = {
+        "type": "object",
+        "properties": {
+            "test_param": {"type": "string"},
+        },
+        "required": ["test_param"],
+    }
+
+    tool = RestApiTool(
+        name="test_tool",
+        description="Test description",
+        endpoint=sample_endpoint,
+        operation=sample_operation,
+        should_parse_operation=False,
+    )
+    tool._operation_parser = mock_parser
+
+    with temporary_feature_override(
+        FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True
+    ):
+      declaration = tool._get_declaration()
+
+    assert isinstance(declaration, FunctionDeclaration)
+    assert declaration.name == "test_tool"
+    assert declaration.description == "Test description"
+    assert declaration.parameters is None
+    assert declaration.parameters_json_schema == {
+        "type": "object",
+        "properties": {
+            "test_param": {"type": "string"},
+        },
+        "required": ["test_param"],
+    }
+
   @patch(
       "google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool.requests.request"
   )