refactor(routes): replace oneof cartesian product with allOf+oneOf and inline component install

Replace O(∏fields) Cartesian product oneof generation with O(∑fields)
allOf+oneOf-per-group approach, and inline _install_components into
add_a2a_routes_to_fastapi to remove a single-use wrapper function.

Author: martimfasantos

src/a2a/server/routes/helpers/_proto_schema.py

@@ -5,10 +5,7 @@
 from google.protobuf.descriptor import Descriptor, FieldDescriptor
 from google.protobuf.message import Message
 
-from a2a.types.a2a_pb2 import (
-    SendMessageRequest,
-    TaskPushNotificationConfig,
-)
+from a2a.types.a2a_pb2 import SendMessageRequest, TaskPushNotificationConfig
 
 
 REST_BODY_TYPES: dict[tuple[str, str], type[Message]] = {
@@ -100,21 +97,22 @@ def message_schema(
         components[name] = {'type': 'object', 'properties': base_properties}
         return ref
 
-    variants: list[dict[str, Any]] = [{}]
-    for oneof in real_oneofs:
-        variants = [
-            {**variant, f.name: field_schema(f, components)}
-            for variant in variants
-            for f in oneof.fields
-        ]
-    components[name] = {
-        'oneOf': [
-            {
-                'type': 'object',
-                'properties': {**base_properties, **variant_props},
-                'required': sorted(set(variant_props) & oneof_field_names),
-            }
-            for variant_props in variants
-        ],
-    }
+    oneof_constraints = [
+        {
+            'oneOf': [
+                {
+                    'type': 'object',
+                    'properties': {f.name: field_schema(f, components)},
+                    'required': [f.name],
+                }
+                for f in oneof.fields
+            ]
+        }
+        for oneof in real_oneofs
+    ]
+    parts: list[dict[str, Any]] = []
+    if base_properties:
+        parts.append({'type': 'object', 'properties': base_properties})
+    parts.extend(oneof_constraints)
+    components[name] = parts[0] if len(parts) == 1 else {'allOf': parts}
     return ref

src/a2a/server/routes/helpers/fastapi.py

@@ -101,23 +101,6 @@ def _attach_route(
     )
 
 
-def _install_components(app: 'FastAPI', schemas: dict[str, Any]) -> None:
-    original_openapi = app.openapi
-
-    def _openapi() -> dict[str, Any]:
-        if app.openapi_schema:
-            return app.openapi_schema
-        schema = original_openapi()
-        component_schemas = schema.setdefault('components', {}).setdefault(
-            'schemas', {}
-        )
-        for name, sub_schema in schemas.items():
-            component_schemas.setdefault(name, sub_schema)
-        return schema
-
-    app.openapi = _openapi  # type: ignore[method-assign]
-
-
 def add_a2a_routes_to_fastapi(
     app: 'FastAPI',
     *,
@@ -197,4 +180,17 @@ def add_a2a_routes_to_fastapi(
             require_version_header=True,
         )
 
-    _install_components(app, components)
+    original_openapi = app.openapi
+
+    def _openapi() -> dict[str, Any]:
+        if app.openapi_schema:
+            return app.openapi_schema
+        schema = original_openapi()
+        component_schemas = schema.setdefault('components', {}).setdefault(
+            'schemas', {}
+        )
+        for name, sub_schema in components.items():
+            component_schemas.setdefault(name, sub_schema)
+        return schema
+
+    app.openapi = _openapi  # type: ignore[method-assign]

tests/server/routes/helpers/test_proto_schema.py

@@ -39,22 +39,68 @@ def test_message_schema_well_known_type_inline():
     assert 'Struct' not in components
 
 
-def test_message_schema_oneof_becomes_one_of():
+def test_message_schema_oneof_becomes_allof_with_one_of_constraint():
     components = {}
     message_schema(Part.DESCRIPTOR, components)
     schema = components['Part']
-    assert 'oneOf' in schema
-    oneof_keys = {list(v['properties'])[-1] for v in schema['oneOf']}
+    assert 'allOf' in schema
+    one_of_constraint = next(p for p in schema['allOf'] if 'oneOf' in p)
+    oneof_keys = {list(v['properties'])[0] for v in one_of_constraint['oneOf']}
     assert {'text', 'raw', 'url', 'data'} <= oneof_keys
 
 
 def test_message_schema_oneof_variants_have_required():
     components = {}
     message_schema(Part.DESCRIPTOR, components)
-    for variant in components['Part']['oneOf']:
+    one_of_constraint = next(
+        p for p in components['Part']['allOf'] if 'oneOf' in p
+    )
+    for variant in one_of_constraint['oneOf']:
         assert len(variant['required']) == 1
 
 
+def test_message_schema_multiple_oneofs_use_allof_not_cartesian_product():
+    # Simulate a descriptor with two oneofs: verify allOf has one constraint
+    # per oneof rather than a flat list of cross-product variants.
+    from unittest.mock import MagicMock
+
+    def _make_field(name):
+        f = MagicMock()
+        f.name = name
+        f.message_type = None
+        f.type = 9  # TYPE_STRING
+        f.is_repeated = False
+        return f
+
+    def _make_oneof(fields):
+        o = MagicMock()
+        o.fields = fields
+        return o
+
+    f_a, f_b = _make_field('a'), _make_field('b')
+    f_x, f_y = _make_field('x'), _make_field('y')
+    oneof1 = _make_oneof([f_a, f_b])
+    oneof2 = _make_oneof([f_x, f_y])
+
+    descriptor = MagicMock()
+    descriptor.full_name = 'test.MultiOneof'
+    descriptor.name = 'MultiOneof'
+    descriptor.oneofs = [oneof1, oneof2]
+    descriptor.fields = [f_a, f_b, f_x, f_y]
+
+    components = {}
+    message_schema(descriptor, components)
+    schema = components['MultiOneof']
+
+    # Should be allOf with two oneOf constraints (one per oneof group),
+    # NOT a flat oneOf with 2*2=4 Cartesian-product variants.
+    assert 'allOf' in schema
+    one_of_constraints = [p for p in schema['allOf'] if 'oneOf' in p]
+    assert len(one_of_constraints) == 2
+    assert len(one_of_constraints[0]['oneOf']) == 2
+    assert len(one_of_constraints[1]['oneOf']) == 2
+
+
 def test_field_schema_repeated_wraps_in_array():
     components = {}
     msg_descriptor = SendMessageRequest.DESCRIPTOR.fields_by_name[