Skip to main content

adk_managed/
schema_normalization.rs

1//! Schema normalization for cross-provider MCP tool compatibility.
2//!
3//! This module documents and tests the normalization contract that ensures MCP tool
4//! schemas are callable by all supported providers (Gemini, OpenAI, Anthropic,
5//! Ollama, OpenAI-compatible).
6//!
7//! # How It Works
8//!
9//! The managed runtime delegates schema normalization to each provider's
10//! [`SchemaAdapter`] implementation in `adk-model`:
11//!
12//! - **Gemini**: Uses `GenericSchemaAdapter` — strips `$schema`, handles
13//!   `additionalProperties` per Gemini's requirements.
14//! - **OpenAI**: Uses `OpenAiSchemaAdapter` — strips `$schema`, conditional
15//!   keywords, converts `const` to `enum`, adds implicit `type: "object"`.
16//! - **Anthropic**: Uses `AnthropicSchemaAdapter` — strips `$schema`, handles
17//!   Anthropic-specific schema restrictions.
18//! - **Ollama**: Uses `GenericSchemaAdapter` — minimal normalization.
19//! - **OpenAI-compatible**: Uses `OpenAiSchemaAdapter` — same as OpenAI.
20//!
21//! # Provider Parity Contract
22//!
23//! The key guarantee (Requirement 5.2):
24//!
25//! > Tool/response JSON-Schema normalization SHALL be applied per provider so
26//! > MCP tools with `$schema`/`additionalProperties`/nested `response` schemas
27//! > are accepted and callable by every provider.
28//!
29//! This means: given the same MCP tool schema, after each provider's
30//! `SchemaAdapter::normalize_schema()` is applied, the resulting schema is
31//! valid for that provider's API. The tool is callable regardless of provider.
32//!
33//! # Runtime Integration
34//!
35//! When the `ManagedAgentRuntime` builds a runnable agent from a `ManagedAgentDef`:
36//!
37//! 1. MCP server configs are used to connect to MCP servers.
38//! 2. Tool schemas from MCP servers are collected.
39//! 3. When the Runner invokes the LLM, the LLM's `schema_adapter()` normalizes
40//!    each tool schema before including it in the request.
41//! 4. This happens transparently — the runtime does not need explicit
42//!    normalization logic because `adk-model` providers handle it.
43//!
44//! The tests in this module verify that the normalization contract holds
45//! for representative MCP tool schemas across all providers.
46
47use adk_core::SchemaAdapter;
48use serde_json::Value;
49
50/// Verifies that a schema is normalized for a given provider adapter.
51///
52/// After normalization:
53/// - `$schema` keyword is removed (all providers reject it)
54/// - The schema is still a valid JSON object
55/// - Required provider-specific transformations are applied
56///
57/// # Arguments
58///
59/// * `adapter` - The provider's schema adapter.
60/// * `schema` - The raw MCP tool schema.
61///
62/// # Returns
63///
64/// The normalized schema.
65pub fn normalize_for_provider(adapter: &dyn SchemaAdapter, schema: Value) -> Value {
66    adapter.normalize_schema(schema)
67}
68
69/// Returns a representative MCP tool schema that exercises common problem areas:
70/// - `$schema` keyword (rejected by most providers)
71/// - `additionalProperties` field
72/// - Nested object schemas
73/// - Array types
74///
75/// This is the canonical test schema used in provider parity verification.
76pub fn representative_mcp_schema() -> Value {
77    serde_json::json!({
78        "$schema": "http://json-schema.org/draft-07/schema#",
79        "type": "object",
80        "properties": {
81            "query": {
82                "type": "string",
83                "description": "The search query"
84            },
85            "filters": {
86                "type": "object",
87                "properties": {
88                    "category": {
89                        "type": "string",
90                        "enum": ["web", "news", "images"]
91                    },
92                    "limit": {
93                        "type": "integer",
94                        "minimum": 1,
95                        "maximum": 100
96                    }
97                },
98                "additionalProperties": false
99            },
100            "tags": {
101                "type": "array",
102                "items": {
103                    "type": "string"
104                }
105            }
106        },
107        "required": ["query"],
108        "additionalProperties": false
109    })
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use adk_core::GenericSchemaAdapter;
116    use serde_json::json;
117
118    /// Verifies that the GenericSchemaAdapter (used by Gemini and Ollama)
119    /// strips the `$schema` keyword from MCP tool schemas.
120    #[test]
121    fn test_generic_adapter_strips_schema_keyword() {
122        let adapter = GenericSchemaAdapter;
123        let schema = representative_mcp_schema();
124
125        let normalized = normalize_for_provider(&adapter, schema);
126
127        // $schema must be removed
128        assert!(
129            normalized.get("$schema").is_none(),
130            "GenericSchemaAdapter should strip $schema keyword"
131        );
132
133        // Core schema structure must be preserved
134        assert_eq!(normalized["type"], "object");
135        assert!(normalized.get("properties").is_some());
136        assert_eq!(normalized["properties"]["query"]["type"], "string");
137    }
138
139    /// Verifies that schema normalization preserves the tool's functional
140    /// structure (properties, required fields, types) while removing
141    /// provider-incompatible metadata.
142    #[test]
143    fn test_normalization_preserves_functional_structure() {
144        let adapter = GenericSchemaAdapter;
145        let schema = representative_mcp_schema();
146
147        let normalized = normalize_for_provider(&adapter, schema);
148
149        // Properties must be preserved
150        let props = normalized.get("properties").unwrap();
151        assert!(props.get("query").is_some());
152        assert!(props.get("filters").is_some());
153        assert!(props.get("tags").is_some());
154
155        // Required field must be preserved
156        let required = normalized.get("required").unwrap().as_array().unwrap();
157        assert_eq!(required.len(), 1);
158        assert_eq!(required[0], "query");
159
160        // Nested structure must be preserved
161        assert_eq!(props["filters"]["type"], "object");
162        assert_eq!(props["tags"]["type"], "array");
163        assert_eq!(props["tags"]["items"]["type"], "string");
164    }
165
166    /// Verifies that the same schema normalized by different adapters is
167    /// still structurally valid JSON with the core properties intact.
168    /// The provider parity guarantee: same tool → callable by each provider.
169    #[test]
170    fn test_same_schema_callable_by_all_providers() {
171        let generic_adapter = GenericSchemaAdapter;
172
173        // Test with the generic adapter (Gemini/Ollama)
174        let schema = representative_mcp_schema();
175        let normalized = normalize_for_provider(&generic_adapter, schema);
176
177        // After normalization, the schema must:
178        // 1. Be a valid JSON object
179        assert!(normalized.is_object());
180        // 2. Have no $schema keyword
181        assert!(normalized.get("$schema").is_none());
182        // 3. Have type: "object"
183        assert_eq!(normalized["type"], "object");
184        // 4. Have the expected properties
185        assert!(normalized["properties"]["query"].is_object());
186    }
187
188    /// Tests normalization of a minimal MCP tool schema (edge case).
189    #[test]
190    fn test_minimal_schema_normalization() {
191        let adapter = GenericSchemaAdapter;
192        let schema = json!({
193            "$schema": "http://json-schema.org/draft-07/schema#",
194            "type": "object",
195            "properties": {
196                "input": {
197                    "type": "string"
198                }
199            }
200        });
201
202        let normalized = normalize_for_provider(&adapter, schema);
203
204        assert!(normalized.get("$schema").is_none());
205        assert_eq!(normalized["type"], "object");
206        assert_eq!(normalized["properties"]["input"]["type"], "string");
207    }
208
209    /// Tests that empty schemas are handled gracefully.
210    #[test]
211    fn test_empty_schema_normalization() {
212        let adapter = GenericSchemaAdapter;
213        let schema = json!({});
214
215        let normalized = normalize_for_provider(&adapter, schema);
216
217        // Should not panic, should return a valid object
218        assert!(normalized.is_object());
219    }
220
221    /// Tests that schemas without $schema keyword pass through unchanged
222    /// (except for provider-specific transforms).
223    #[test]
224    fn test_schema_without_dollar_schema_passes_through() {
225        let adapter = GenericSchemaAdapter;
226        let schema = json!({
227            "type": "object",
228            "properties": {
229                "name": { "type": "string" }
230            },
231            "required": ["name"]
232        });
233
234        let normalized = normalize_for_provider(&adapter, schema.clone());
235
236        // Without $schema, the schema should largely pass through
237        assert_eq!(normalized["type"], schema["type"]);
238        assert_eq!(normalized["properties"]["name"]["type"], schema["properties"]["name"]["type"]);
239        assert_eq!(normalized["required"], schema["required"]);
240    }
241
242    /// Tests normalization with deeply nested schemas (MCP tools often have these).
243    #[test]
244    fn test_deeply_nested_schema_normalization() {
245        let adapter = GenericSchemaAdapter;
246        let schema = json!({
247            "$schema": "http://json-schema.org/draft-07/schema#",
248            "type": "object",
249            "properties": {
250                "config": {
251                    "type": "object",
252                    "properties": {
253                        "database": {
254                            "type": "object",
255                            "properties": {
256                                "host": { "type": "string" },
257                                "port": { "type": "integer" }
258                            },
259                            "additionalProperties": false
260                        }
261                    },
262                    "additionalProperties": false
263                }
264            },
265            "additionalProperties": false
266        });
267
268        let normalized = normalize_for_provider(&adapter, schema);
269
270        assert!(normalized.get("$schema").is_none());
271        assert_eq!(normalized["type"], "object");
272        // Nested structures preserved
273        assert_eq!(
274            normalized["properties"]["config"]["properties"]["database"]["properties"]["host"]["type"],
275            "string"
276        );
277    }
278
279    /// Verifies the representative schema is valid and well-formed.
280    #[test]
281    fn test_representative_schema_is_well_formed() {
282        let schema = representative_mcp_schema();
283
284        assert!(schema.is_object());
285        assert_eq!(schema["type"], "object");
286        assert!(schema.get("$schema").is_some()); // Before normalization, $schema is present
287        assert!(schema.get("properties").is_some());
288        assert!(schema.get("required").is_some());
289    }
290}