airbyte-agent-amazon-ads 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

airbyte_agent_amazon_ads/_vendored/connector_sdk/connector_model_loader.py

@@ -62,6 +62,53 @@ class TokenExtractValidationError(ConnectorModelLoaderError):
     pass
 
 
+# Expected auth_mapping keys for each auth type.
+# These are the auth parameters that each security scheme expects, NOT the user's credential field names.
+EXPECTED_AUTH_MAPPING_KEYS: dict[AuthType, set[str]] = {
+    AuthType.BEARER: {"token"},
+    AuthType.BASIC: {"username", "password"},
+    AuthType.API_KEY: {"api_key"},
+    AuthType.OAUTH2: {"access_token", "refresh_token", "client_id", "client_secret"},
+}
+
+
+def _validate_auth_mapping_keys(
+    auth_type: AuthType,
+    auth_config: AirbyteAuthConfig | None,
+    scheme_name: str = "default",
+) -> None:
+    """Validate that auth_mapping keys match expected parameters for the auth type.
+
+    The auth_mapping keys must be the parameters expected by the security scheme
+    (e.g., "token" for bearer), not the user's credential field names.
+
+    Args:
+        auth_type: The authentication type
+        auth_config: The x-airbyte-auth-config containing auth_mapping
+        scheme_name: Name of the security scheme for error messages
+
+    Raises:
+        InvalidOpenAPIError: If auth_mapping keys don't match expected parameters
+    """
+    if auth_config is None or auth_config.auth_mapping is None:
+        return  # No explicit auth_mapping, will use defaults
+
+    expected_keys = EXPECTED_AUTH_MAPPING_KEYS.get(auth_type)
+    if expected_keys is None:
+        return  # Unknown auth type, skip validation
+
+    actual_keys = set(auth_config.auth_mapping.keys())
+    invalid_keys = actual_keys - expected_keys
+
+    if invalid_keys:
+        raise InvalidOpenAPIError(
+            f"Invalid auth_mapping keys for {auth_type.value} auth in scheme '{scheme_name}': {invalid_keys}. "
+            f"Expected keys for {auth_type.value}: {sorted(expected_keys)}. "
+            f"Note: auth_mapping keys must be the auth parameters (e.g., 'token' for bearer), "
+            f'not your credential field names. Use template syntax to map: token: "${{your_field}}"'
+        )
+
+
 def extract_path_params(path: str) -> list[str]:
     """Extract parameter names from path template.
 
@@ -145,6 +192,87 @@ def _deproxy_schema(obj: Any) -> Any:
         return obj
 
 
+def _type_includes(type_value: Any, target: str) -> bool:
+    if isinstance(type_value, list):
+        return target in type_value
+    return type_value == target
+
+
+def _flatten_cache_properties(properties: dict[str, Any], prefix: str) -> list[str]:
+    entries: list[str] = []
+    for prop_name, prop in properties.items():
+        path = f"{prefix}{prop_name}" if prefix else prop_name
+        entries.append(path)
+
+        prop_type = getattr(prop, "type", None) if not isinstance(prop, dict) else prop.get("type")
+        prop_properties = getattr(prop, "properties", None) if not isinstance(prop, dict) else prop.get("properties")
+
+        if _type_includes(prop_type, "array"):
+            array_path = f"{path}[]"
+            entries.append(array_path)
+            if isinstance(prop_properties, dict):
+                entries.extend(_flatten_cache_properties(prop_properties, prefix=f"{array_path}."))
+        elif isinstance(prop_properties, dict):
+            entries.extend(_flatten_cache_properties(prop_properties, prefix=f"{path}."))
+
+    return entries
+
+
+def _flatten_cache_field_paths(field: Any) -> list[str]:
+    field_name = getattr(field, "name", None) if not isinstance(field, dict) else field.get("name")
+    if not isinstance(field_name, str) or not field_name:
+        return []
+
+    field_type = getattr(field, "type", None) if not isinstance(field, dict) else field.get("type")
+    field_properties = getattr(field, "properties", None) if not isinstance(field, dict) else field.get("properties")
+
+    entries = [field_name]
+    if _type_includes(field_type, "array"):
+        array_path = f"{field_name}[]"
+        entries.append(array_path)
+        if isinstance(field_properties, dict):
+            entries.extend(_flatten_cache_properties(field_properties, prefix=f"{array_path}."))
+    elif isinstance(field_properties, dict):
+        entries.extend(_flatten_cache_properties(field_properties, prefix=f"{field_name}."))
+
+    return entries
+
+
+def _dedupe_strings(values: list[str]) -> list[str]:
+    seen: set[str] = set()
+    ordered: list[str] = []
+    for value in values:
+        if value not in seen:
+            seen.add(value)
+            ordered.append(value)
+    return ordered
+
+
+def _extract_search_field_paths(spec: OpenAPIConnector) -> dict[str, list[str]]:
+    cache_config = getattr(spec.info, "x_airbyte_cache", None)
+    entities = getattr(cache_config, "entities", None)
+    if not isinstance(entities, list):
+        return {}
+
+    search_fields: dict[str, list[str]] = {}
+    for entity in entities:
+        entity_name = getattr(entity, "entity", None) if not isinstance(entity, dict) else entity.get("entity")
+        if not isinstance(entity_name, str) or not entity_name:
+            continue
+
+        fields = getattr(entity, "fields", None) if not isinstance(entity, dict) else entity.get("fields")
+        if not isinstance(fields, list):
+            continue
+
+        field_paths: list[str] = []
+        for field in fields:
+            field_paths.extend(_flatten_cache_field_paths(field))
+
+        search_fields[entity_name] = _dedupe_strings(field_paths)
+
+    return search_fields
+
+
 def parse_openapi_spec(raw_config: dict) -> OpenAPIConnector:
     """Parse OpenAPI specification from YAML.
 
@@ -434,6 +562,8 @@ def convert_openapi_to_connector_model(spec: OpenAPIConnector) -> ConnectorModel
     if not connector_id:
         raise InvalidOpenAPIError("Missing required x-airbyte-connector-id field")
 
+    search_field_paths = _extract_search_field_paths(spec)
+
     # Create ConnectorModel
     model = ConnectorModel(
         id=connector_id,
@@ -444,6 +574,7 @@ def convert_openapi_to_connector_model(spec: OpenAPIConnector) -> ConnectorModel
         entities=entities,
         openapi_spec=spec,
         retry_config=retry_config,
+        search_field_paths=search_field_paths,
     )
 
     return model
@@ -840,6 +971,9 @@ def _parse_single_security_scheme(scheme: Any) -> AuthConfig:
         oauth2_config = _parse_oauth2_config(scheme)
         # Use explicit x-airbyte-auth-config if present, otherwise generate default
         auth_config_obj = scheme.x_airbyte_auth_config or _generate_default_auth_config(AuthType.OAUTH2)
+        # Validate auth_mapping keys if explicitly provided
+        if scheme.x_airbyte_auth_config:
+            _validate_auth_mapping_keys(AuthType.OAUTH2, scheme.x_airbyte_auth_config)
         return AuthConfig(
             type=AuthType.OAUTH2,
             config=oauth2_config,
@@ -850,6 +984,10 @@ def _parse_single_security_scheme(scheme: Any) -> AuthConfig:
     # Use explicit x-airbyte-auth-config if present, otherwise generate default
     auth_config_obj = scheme.x_airbyte_auth_config or _generate_default_auth_config(auth_type)
 
+    # Validate auth_mapping keys if explicitly provided
+    if scheme.x_airbyte_auth_config:
+        _validate_auth_mapping_keys(auth_type, scheme.x_airbyte_auth_config)
+
     return AuthConfig(
         type=auth_type,
         config=auth_config,

airbyte_agent_amazon_ads/_vendored/connector_sdk/introspection.py

@@ -18,6 +18,185 @@ from typing import Any, Protocol
 MAX_EXAMPLE_QUESTIONS = 5  # Maximum number of example questions to include in description
 
 
+def _type_includes(type_value: Any, target: str) -> bool:
+    if isinstance(type_value, list):
+        return target in type_value
+    return type_value == target
+
+
+def _is_object_schema(schema: dict[str, Any]) -> bool:
+    if "properties" in schema:
+        return True
+    return _type_includes(schema.get("type"), "object")
+
+
+def _is_array_schema(schema: dict[str, Any]) -> bool:
+    if "items" in schema:
+        return True
+    return _type_includes(schema.get("type"), "array")
+
+
+def _dedupe_param_entries(entries: list[tuple[str, bool]]) -> list[tuple[str, bool]]:
+    seen: dict[str, bool] = {}
+    ordered: list[str] = []
+    for name, required in entries:
+        if name not in seen:
+            seen[name] = required
+            ordered.append(name)
+        else:
+            seen[name] = seen[name] or required
+    return [(name, seen[name]) for name in ordered]
+
+
+def _flatten_schema_params(
+    schema: dict[str, Any],
+    prefix: str = "",
+    parent_required: bool = True,
+    seen_stack: set[int] | None = None,
+) -> list[tuple[str, bool]]:
+    if not isinstance(schema, dict):
+        return []
+
+    if seen_stack is None:
+        seen_stack = set()
+
+    schema_id = id(schema)
+    if schema_id in seen_stack:
+        return []
+
+    seen_stack.add(schema_id)
+    try:
+        entries: list[tuple[str, bool]] = []
+
+        for subschema in schema.get("allOf", []) or []:
+            if isinstance(subschema, dict):
+                entries.extend(_flatten_schema_params(subschema, prefix, parent_required, seen_stack))
+
+        for keyword in ("anyOf", "oneOf"):
+            for subschema in schema.get(keyword, []) or []:
+                if isinstance(subschema, dict):
+                    entries.extend(_flatten_schema_params(subschema, prefix, False, seen_stack))
+
+        properties = schema.get("properties")
+        if isinstance(properties, dict):
+            required_fields = set(schema.get("required", [])) if isinstance(schema.get("required"), list) else set()
+            for prop_name, prop_schema in properties.items():
+                path = f"{prefix}{prop_name}" if prefix else prop_name
+                is_required = parent_required and prop_name in required_fields
+                entries.append((path, is_required))
+
+                if isinstance(prop_schema, dict):
+                    if _is_array_schema(prop_schema):
+                        array_path = f"{path}[]"
+                        entries.append((array_path, is_required))
+                        items = prop_schema.get("items")
+                        if isinstance(items, dict):
+                            entries.extend(_flatten_schema_params(items, prefix=f"{array_path}.", parent_required=is_required, seen_stack=seen_stack))
+                    if _is_object_schema(prop_schema):
+                        entries.extend(_flatten_schema_params(prop_schema, prefix=f"{path}.", parent_required=is_required, seen_stack=seen_stack))
+
+        return _dedupe_param_entries(entries)
+    finally:
+        seen_stack.remove(schema_id)
+
+
+def _cache_field_value(field: Any, key: str) -> Any:
+    if isinstance(field, dict):
+        return field.get(key)
+    return getattr(field, key, None)
+
+
+def _flatten_cache_properties(properties: dict[str, Any], prefix: str) -> list[str]:
+    entries: list[str] = []
+    for prop_name, prop in properties.items():
+        path = f"{prefix}{prop_name}" if prefix else prop_name
+        entries.append(path)
+
+        prop_type = _cache_field_value(prop, "type")
+        prop_properties = _cache_field_value(prop, "properties")
+
+        if _type_includes(prop_type, "array"):
+            array_path = f"{path}[]"
+            entries.append(array_path)
+            if isinstance(prop_properties, dict):
+                entries.extend(_flatten_cache_properties(prop_properties, prefix=f"{array_path}."))
+        elif isinstance(prop_properties, dict):
+            entries.extend(_flatten_cache_properties(prop_properties, prefix=f"{path}."))
+
+    return entries
+
+
+def _flatten_cache_field_paths(field: Any) -> list[str]:
+    field_name = _cache_field_value(field, "name")
+    if not isinstance(field_name, str) or not field_name:
+        return []
+
+    field_type = _cache_field_value(field, "type")
+    field_properties = _cache_field_value(field, "properties")
+
+    entries = [field_name]
+    if _type_includes(field_type, "array"):
+        array_path = f"{field_name}[]"
+        entries.append(array_path)
+        if isinstance(field_properties, dict):
+            entries.extend(_flatten_cache_properties(field_properties, prefix=f"{array_path}."))
+    elif isinstance(field_properties, dict):
+        entries.extend(_flatten_cache_properties(field_properties, prefix=f"{field_name}."))
+
+    return entries
+
+
+def _dedupe_strings(values: list[str]) -> list[str]:
+    seen: set[str] = set()
+    ordered: list[str] = []
+    for value in values:
+        if value not in seen:
+            seen.add(value)
+            ordered.append(value)
+    return ordered
+
+
+def _collect_search_field_paths(model: ConnectorModelProtocol) -> dict[str, list[str]]:
+    search_field_paths = getattr(model, "search_field_paths", None)
+    if isinstance(search_field_paths, dict) and search_field_paths:
+        normalized: dict[str, list[str]] = {}
+        for entity, fields in search_field_paths.items():
+            if not isinstance(entity, str) or not entity:
+                continue
+            if isinstance(fields, list):
+                normalized[entity] = _dedupe_strings([field for field in fields if isinstance(field, str) and field])
+        return normalized
+
+    openapi_spec = getattr(model, "openapi_spec", None)
+    info = getattr(openapi_spec, "info", None)
+    cache_config = getattr(info, "x_airbyte_cache", None)
+    entities = getattr(cache_config, "entities", None)
+    if not isinstance(entities, list):
+        return {}
+
+    search_fields: dict[str, list[str]] = {}
+    for entity in entities:
+        entity_name = _cache_field_value(entity, "entity")
+        if not isinstance(entity_name, str) or not entity_name:
+            continue
+
+        fields = _cache_field_value(entity, "fields") or []
+        if not isinstance(fields, list):
+            continue
+        field_paths: list[str] = []
+        for field in fields:
+            field_paths.extend(_flatten_cache_field_paths(field))
+
+        search_fields[entity_name] = _dedupe_strings(field_paths)
+
+    return search_fields
+
+
+def _format_search_param_signature() -> str:
+    params = ["query*", "limit?", "cursor?", "fields?"]
+    return f"({', '.join(params)})"
+
+
 class EndpointProtocol(Protocol):
     """Protocol defining the expected interface for endpoint parameters.
 
@@ -54,6 +233,9 @@ class ConnectorModelProtocol(Protocol):
     @property
     def openapi_spec(self) -> Any: ...
 
+    @property
+    def search_field_paths(self) -> dict[str, list[str]] | None: ...
+
 
 def format_param_signature(endpoint: EndpointProtocol) -> str:
     """Format parameter signature for an endpoint action.
@@ -86,9 +268,12 @@ def format_param_signature(endpoint: EndpointProtocol) -> str:
         required = schema.get("required", False)
         params.append(f"{name}{'*' if required else '?'}")
 
-    # Body fields
-    if request_schema:
-        required_fields = set(request_schema.get("required", []))
+    # Body fields (include nested params from schema when available)
+    if isinstance(request_schema, dict):
+        for name, required in _flatten_schema_params(request_schema):
+            params.append(f"{name}{'*' if required else '?'}")
+    elif request_schema:
+        required_fields = set(request_schema.get("required", [])) if isinstance(request_schema, dict) else set()
         for name in body_fields:
             params.append(f"{name}{'*' if name in required_fields else '?'}")
 
@@ -99,7 +284,7 @@ def describe_entities(model: ConnectorModelProtocol) -> list[dict[str, Any]]:
     """Generate entity descriptions from ConnectorModel.
 
     Returns a list of entity descriptions with detailed parameter information
-    for each action. This is used by generated connectors' describe() method.
+    for each action. This is used by generated connectors' list_entities() method.
 
     Args:
         model: Object conforming to ConnectorModelProtocol (e.g., ConnectorModel)
@@ -203,8 +388,8 @@ def generate_tool_description(model: ConnectorModelProtocol) -> str:
     - Response structure documentation with pagination hints
     - Example questions if available in the OpenAPI spec
 
-    This is used by the Connector.describe class method decorator to populate
-    function docstrings for AI framework integration.
+    This is used by the Connector.tool_utils decorator to populate function
+    docstrings for AI framework integration.
 
     Args:
         model: Object conforming to ConnectorModelProtocol (e.g., ConnectorModel)
@@ -213,8 +398,11 @@ def generate_tool_description(model: ConnectorModelProtocol) -> str:
         Formatted description string suitable for AI tool documentation
     """
     lines = []
+    # NOTE: Do not insert blank lines in the docstring; pydantic-ai parsing truncates
+    # at the first empty line and only keeps the initial section.
 
     # Entity/action parameter details (including pagination params like limit, starting_after)
+    search_field_paths = _collect_search_field_paths(model)
     lines.append("ENTITIES AND PARAMETERS:")
     for entity in model.entities:
         lines.append(f"  {entity.name}:")
@@ -228,14 +416,41 @@ def generate_tool_description(model: ConnectorModelProtocol) -> str:
                 lines.append(f"    - {action_str}{param_sig}")
             else:
                 lines.append(f"    - {action_str}()")
+        if entity.name in search_field_paths:
+            search_sig = _format_search_param_signature()
+            lines.append(f"    - search{search_sig}")
 
     # Response structure (brief, includes pagination hint)
-    lines.append("")
     lines.append("RESPONSE STRUCTURE:")
     lines.append("  - list/api_search: {data: [...], meta: {has_more: bool}}")
     lines.append("  - get: Returns entity directly (no envelope)")
     lines.append("  To paginate: pass starting_after=<last_id> while has_more is true")
 
+    lines.append("GUIDELINES:")
+    lines.append('  - Prefer cached search over direct API calls when using execute(): action="search" whenever possible.')
+    lines.append("  - Direct API actions (list/get/download) are slower and should be used only if search cannot answer the query.")
+    lines.append("  - Keep results small: use params.fields, params.query.filter, small params.limit, and cursor pagination.")
+    lines.append("  - If output is too large, refine the query with tighter filters/fields/limit.")
+
+    if search_field_paths:
+        lines.append("SEARCH (PREFERRED):")
+        lines.append('  execute(entity, action="search", params={')
+        lines.append('    "query": {"filter": <condition>, "sort": [{"field": "asc|desc"}, ...]},')
+        lines.append('    "limit": <int>, "cursor": <str>, "fields": ["field", "nested.field", ...]')
+        lines.append("  })")
+        lines.append('  Example: {"query": {"filter": {"eq": {"title": "Intro to Airbyte | Miinto"}}}, "limit": 1,')
+        lines.append('            "fields": ["id", "title", "started", "primaryUserId"]}')
+        lines.append("  Conditions are composable:")
+        lines.append("    - eq, neq, gt, gte, lt, lte, in, like, fuzzy, keyword, contains, any")
+        lines.append('    - and/or/not to combine conditions (e.g., {"and": [cond1, cond2]})')
+
+        lines.append("SEARCHABLE FIELDS:")
+        for entity_name, field_paths in search_field_paths.items():
+            if field_paths:
+                lines.append(f"  {entity_name}: {', '.join(field_paths)}")
+            else:
+                lines.append(f"  {entity_name}: (no fields listed)")
+
     # Add example questions if available in openapi_spec
     openapi_spec = getattr(model, "openapi_spec", None)
     if openapi_spec:
@@ -245,18 +460,15 @@ def generate_tool_description(model: ConnectorModelProtocol) -> str:
             if example_questions:
                 supported = getattr(example_questions, "supported", None)
                 if supported:
-                    lines.append("")
                     lines.append("EXAMPLE QUESTIONS:")
                     for q in supported[:MAX_EXAMPLE_QUESTIONS]:
                         lines.append(f"  - {q}")
 
     # Generic parameter description for function signature
-    lines.append("")
     lines.append("FUNCTION PARAMETERS:")
     lines.append("  - entity: Entity name (string)")
     lines.append("  - action: Operation to perform (string)")
     lines.append("  - params: Operation parameters (dict) - see entity details above")
-    lines.append("")
     lines.append("Parameter markers: * = required, ? = optional")
 
     return "\n".join(lines)

airbyte_agent_amazon_ads/_vendored/connector_sdk/types.py

@@ -252,3 +252,4 @@ class ConnectorModel(BaseModel):
     entities: list[EntityDefinition]
     openapi_spec: Any | None = None  # Optional reference to OpenAPIConnector
     retry_config: RetryConfig | None = None  # Optional retry configuration
+    search_field_paths: dict[str, list[str]] | None = None

airbyte_agent_amazon_ads/connector.py

@@ -4,8 +4,11 @@ Amazon-Ads connector.
 
 from __future__ import annotations
 
+import inspect
+import json
 import logging
-from typing import TYPE_CHECKING, Any, Callable, TypeVar, overload
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Callable, Mapping, TypeVar, overload
 try:
     from typing import Literal
 except ImportError:
@@ -46,6 +49,38 @@ from .models import (
 # TypeVar for decorator type preservation
 _F = TypeVar("_F", bound=Callable[..., Any])
 
+DEFAULT_MAX_OUTPUT_CHARS = 50_000  # ~50KB default, configurable per-tool
+
+
+def _raise_output_too_large(message: str) -> None:
+    try:
+        from pydantic_ai import ModelRetry  # type: ignore[import-not-found]
+    except Exception as exc:
+        raise RuntimeError(message) from exc
+    raise ModelRetry(message)
+
+
+def _check_output_size(result: Any, max_chars: int | None, tool_name: str) -> Any:
+    if max_chars is None or max_chars <= 0:
+        return result
+
+    try:
+        serialized = json.dumps(result, default=str)
+    except (TypeError, ValueError):
+        return result
+
+    if len(serialized) > max_chars:
+        truncated_preview = serialized[:500] + "..." if len(serialized) > 500 else serialized
+        _raise_output_too_large(
+            f"Tool '{tool_name}' output too large ({len(serialized):,} chars, limit {max_chars:,}). "
+            "Please narrow your query by: using the 'fields' parameter to select only needed fields, "
+            "adding filters, or reducing the 'limit'. "
+            f"Preview: {truncated_preview}"
+        )
+
+    return result
+
+
 
 
 class AmazonAdsConnector:
@@ -225,15 +260,15 @@ class AmazonAdsConnector:
     async def execute(
         self,
         entity: str,
-        action: str,
-        params: dict[str, Any]
+        action: Literal["list", "get", "search"],
+        params: Mapping[str, Any]
     ) -> AmazonAdsExecuteResult[Any] | AmazonAdsExecuteResultWithMeta[Any, Any] | Any: ...
 
     async def execute(
         self,
         entity: str,
-        action: str,
-        params: dict[str, Any] | None = None
+        action: Literal["list", "get", "search"],
+        params: Mapping[str, Any] | None = None
     ) -> Any:
         """
         Execute an entity operation with full type safety.
@@ -261,16 +296,17 @@ class AmazonAdsConnector:
         from ._vendored.connector_sdk.executor import ExecutionConfig
 
         # Remap parameter names from snake_case (TypedDict keys) to API parameter names
-        if params:
+        resolved_params = dict(params) if params is not None else None
+        if resolved_params:
             param_map = self._PARAM_MAP.get((entity, action), {})
             if param_map:
-                params = {param_map.get(k, k): v for k, v in params.items()}
+                resolved_params = {param_map.get(k, k): v for k, v in resolved_params.items()}
 
         # Use ExecutionConfig for both local and hosted executors
         config = ExecutionConfig(
             entity=entity,
             action=action,
-            params=params
+            params=resolved_params
         )
 
         result = await self._executor.execute(config)
@@ -297,41 +333,67 @@ class AmazonAdsConnector:
     # ===== INTROSPECTION METHODS =====
 
     @classmethod
-    def describe(cls, func: _F) -> _F:
+    def tool_utils(
+        cls,
+        func: _F | None = None,
+        *,
+        update_docstring: bool = True,
+        max_output_chars: int | None = DEFAULT_MAX_OUTPUT_CHARS,
+    ) -> _F | Callable[[_F], _F]:
         """
-        Decorator that populates a function's docstring with connector capabilities.
-
-        This class method can be used as a decorator to automatically generate
-        comprehensive documentation for AI tool functions.
+        Decorator that adds tool utilities like docstring augmentation and output limits.
 
         Usage:
             @mcp.tool()
-            @AmazonAdsConnector.describe
+            @AmazonAdsConnector.tool_utils
             async def execute(entity: str, action: str, params: dict):
-                '''Execute operations.'''
                 ...
 
-        The decorated function's __doc__ will be updated with:
-        - Available entities and their actions
-        - Parameter signatures with required (*) and optional (?) markers
-        - Response structure documentation
-        - Example questions (if available in OpenAPI spec)
+            @mcp.tool()
+            @AmazonAdsConnector.tool_utils(update_docstring=False, max_output_chars=None)
+            async def execute(entity: str, action: str, params: dict):
+                ...
 
         Args:
-            func: The function to decorate
-
-        Returns:
-            The same function with updated __doc__
+            update_docstring: When True, append connector capabilities to __doc__.
+            max_output_chars: Max serialized output size before raising. Use None to disable.
         """
-        description = generate_tool_description(AmazonAdsConnectorModel)
 
-        original_doc = func.__doc__ or ""
-        if original_doc.strip():
-            func.__doc__ = f"{original_doc.strip()}\n{description}"
-        else:
-            func.__doc__ = description
+        def decorate(inner: _F) -> _F:
+            if update_docstring:
+                description = generate_tool_description(AmazonAdsConnectorModel)
+                original_doc = inner.__doc__ or ""
+                if original_doc.strip():
+                    full_doc = f"{original_doc.strip()}\n{description}"
+                else:
+                    full_doc = description
+            else:
+                full_doc = ""
+
+            if inspect.iscoroutinefunction(inner):
+
+                @wraps(inner)
+                async def aw(*args: Any, **kwargs: Any) -> Any:
+                    result = await inner(*args, **kwargs)
+                    return _check_output_size(result, max_output_chars, inner.__name__)
+
+                wrapped = aw
+            else:
+
+                @wraps(inner)
+                def sw(*args: Any, **kwargs: Any) -> Any:
+                    result = inner(*args, **kwargs)
+                    return _check_output_size(result, max_output_chars, inner.__name__)
+
+                wrapped = sw
+
+            if update_docstring:
+                wrapped.__doc__ = full_doc
+            return wrapped  # type: ignore[return-value]
 
-        return func
+        if func is not None:
+            return decorate(func)
+        return decorate
 
     def list_entities(self) -> list[dict[str, Any]]:
         """