Add toolset as a provider module category (#70122)

airflow-core/src/airflow/provider.yaml.schema.json

@@ -373,6 +373,30 @@
                 ]
             }
         },
+        "toolsets": {
+            "type": "array",
+            "items": {
+                "type": "object",
+                "properties": {
+                    "integration-name": {
+                        "type": "string",
+                        "description": "Integration name. It must have a matching item in the 'integration' section of any provider."
+                    },
+                    "python-modules": {
+                        "description": "List of python modules containing the toolsets.",
+                        "type": "array",
+                        "items": {
+                            "type": "string"
+                        }
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "integration-name",
+                    "python-modules"
+                ]
+            }
+        },
         "connection-types": {
             "type": "array",
             "description": "Array of connection types mapped to hook class names",

dev/registry/registry_tools/types.py

@@ -102,6 +102,13 @@ MODULE_TYPES: dict[str, dict] = {
         "label": "Decorators",
         "icon": "@",
     },
+    "toolset": {
+        "yaml_key": "toolsets",
+        "level": "module",
+        "suffixes": ["Toolset"],
+        "label": "Toolsets",
+        "icon": "U",
+    },
 }
 
 # Runtime base class imports for issubclass checks (extract_parameters.py).
@@ -113,6 +120,7 @@ BASE_CLASS_IMPORTS: list[tuple[str, str]] = [
     ("hook", "airflow.sdk.bases.hook.BaseHook"),
     ("bundle", "airflow.dag_processing.bundles.base.BaseDagBundle"),
     ("operator", "airflow.sdk.bases.operator.BaseOperator"),
+    ("toolset", "pydantic_ai.toolsets.abstract.AbstractToolset"),
 ]
 
 # Derived lookups used by extraction scripts.

providers/common/ai/provider.yaml

@@ -451,6 +451,17 @@ operators:
       - airflow.providers.common.ai.operators.llamaindex_embedding
       - airflow.providers.common.ai.operators.llamaindex_retrieval
 
+toolsets:
+  - integration-name: Common AI
+    python-modules:
+      - airflow.providers.common.ai.toolsets.hook
+      - airflow.providers.common.ai.toolsets.sql
+      - airflow.providers.common.ai.toolsets.datafusion
+      - airflow.providers.common.ai.toolsets.logging
+      - airflow.providers.common.ai.toolsets.mcp
+      - airflow.providers.common.ai.toolsets.skills
+      - airflow.providers.common.ai.toolsets.langchain_bridge
+
 task-decorators:
   - class-name: airflow.providers.common.ai.decorators.agent.agent_task
     name: agent

registry/src/_data/types.json

@@ -53,5 +53,10 @@
     "id": "decorator",
     "label": "Decorators",
     "icon": "@"
+  },
+  {
+    "id": "toolset",
+    "label": "Toolsets",
+    "icon": "U"
   }
 ]

registry/src/css/main.css

@@ -3065,6 +3065,7 @@ main {
 .tab-icon.executor { background: rgb(from var(--color-executor) r g b / 0.2); color: var(--color-executor); }
 .tab-icon.bundle { background: rgb(from var(--color-bundle) r g b / 0.2); color: var(--color-bundle); }
 .tab-icon.decorator { background: rgb(from var(--color-decorator) r g b / 0.2); color: var(--color-decorator); }
+.tab-icon.toolset { background: rgb(from var(--color-toolset) r g b / 0.2); color: var(--color-toolset); }
 
 /* Modules Layout (sidebar + content) */
 .modules-layout {

registry/src/css/tokens.css

@@ -65,6 +65,7 @@
   --color-executor: #f97316;
   --color-bundle: #0ea5e9;
   --color-decorator: #d946ef;
+  --color-toolset: #84cc16;
 
   /* Additional Colors */
   --color-green-400: #4ade80;

scripts/in_container/run_provider_yaml_files_check.py

@@ -75,6 +75,17 @@ KNOWN_DEPRECATED_CLASSES = [
     "airflow.providers.amazon.aws.hooks.kinesis.FirehoseHook",
 ]
 
+# AbstractToolset subclasses that are internal implementation detail -- auto-applied by
+# a provider's own runtime rather than constructed by users -- and are intentionally not
+# part of the registry's public "toolsets" module category, so they are exempt from
+# check_all_provider_classes_are_registered's registration requirement. Contrast with
+# e.g. LoggingToolset, which is documented in the toolsets how-to guide and is registered.
+INTERNAL_UNREGISTERED_TOOLSET_CLASSES = {
+    # Wraps a toolset with per-step result caching for durable execution; applied
+    # automatically by AgentOperator, not part of the public toolsets how-to guide.
+    "airflow.providers.common.ai.durable.caching_toolset.CachingToolset",
+}
+
 if __name__ != "__main__":
     raise SystemExit(
         "This file is intended to be executed as an executable program. You cannot use it as a module."
@@ -368,14 +379,14 @@ def check_integration_duplicates(yaml_files: dict[str, dict]) -> tuple[int, int]
     return num_integrations, num_errors
 
 
-@run_check("Checking completeness of list of {sensors, hooks, operators, triggers, bundles}")
+@run_check("Checking completeness of list of {sensors, hooks, operators, triggers, bundles, toolsets}")
 def check_correctness_of_list_of_sensors_operators_hook_trigger_modules(
     yaml_files: dict[str, dict],
 ) -> tuple[int, int]:
     num_errors = 0
     num_modules = 0
     for (yaml_file_path, provider_data), resource_type in itertools.product(
-        yaml_files.items(), ["sensors", "operators", "hooks", "triggers", "bundles"]
+        yaml_files.items(), ["sensors", "operators", "hooks", "triggers", "bundles", "toolsets"]
     ):
         expected_modules, provider_package, resource_data = parse_module_data(
             provider_data, resource_type, yaml_file_path
@@ -407,14 +418,14 @@ def check_correctness_of_list_of_sensors_operators_hook_trigger_modules(
     return num_modules, num_errors
 
 
-@run_check("Checking for duplicates in list of {sensors, hooks, operators, triggers, bundles}")
+@run_check("Checking for duplicates in list of {sensors, hooks, operators, triggers, bundles, toolsets}")
 def check_duplicates_in_integrations_names_of_hooks_sensors_operators(
     yaml_files: dict[str, dict],
 ) -> tuple[int, int]:
     num_errors = 0
     num_integrations = 0
     for (yaml_file_path, provider_data), resource_type in itertools.product(
-        yaml_files.items(), ["sensors", "operators", "hooks", "triggers", "bundles"]
+        yaml_files.items(), ["sensors", "operators", "hooks", "triggers", "bundles", "toolsets"]
     ):
         resource_data = provider_data.get(resource_type, [])
         count_integrations = Counter(r.get("integration-name", "") for r in resource_data)
@@ -595,18 +606,20 @@ def check_hook_classes_with_conn_type_are_registered(yaml_files: dict[str, dict]
 
 
 @run_check(
-    "Checking that all provider Hook/Operator/Sensor/Trigger/Executor/Notifier"
+    "Checking that all provider Hook/Operator/Sensor/Trigger/Executor/Notifier/Toolset"
     " classes are registered in provider.yaml"
 )
 def check_all_provider_classes_are_registered(yaml_files: dict[str, dict]) -> tuple[int, int]:
     """
     Walk all provider source files, find Hook/Operator/Sensor/Trigger/Executor/Notifier/
-    SecretsBackend/AuthManager/LoggingHandler/DagBundle/DBManager subclasses, and verify
-    they are registered in the appropriate provider.yaml section.
+    SecretsBackend/AuthManager/LoggingHandler/DagBundle/DBManager/Toolset subclasses, and
+    verify they are registered in the appropriate provider.yaml section.
 
     This catches classes placed in non-standard directories or modules that were missed
     when updating provider.yaml.
     """
+    from pydantic_ai.toolsets.abstract import AbstractToolset
+
     from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager
     from airflow.dag_processing.bundles.base import BaseDagBundle
     from airflow.executors.base_executor import BaseExecutor
@@ -632,6 +645,7 @@ def check_all_provider_classes_are_registered(yaml_files: dict[str, dict]) -> tu
         (FileTaskHandler, "logging"),
         (BaseDagBundle, "bundles"),
         (BaseDBManager, "db-managers"),
+        (AbstractToolset, "toolsets"),
     ]
 
     # Resource types where registration is by class path (not module)
@@ -656,7 +670,7 @@ def check_all_provider_classes_are_registered(yaml_files: dict[str, dict]) -> tu
 
         # Collect all modules registered in provider.yaml across all resource types
         registered_modules: set[str] = set()
-        for resource_type in ("hooks", "operators", "sensors", "triggers", "bundles"):
+        for resource_type in ("hooks", "operators", "sensors", "triggers", "bundles", "toolsets"):
             for entry in provider_data.get(resource_type, []):
                 registered_modules.update(entry.get("python-modules", []))
         for entry in provider_data.get("transfers", []):
@@ -733,8 +747,10 @@ def check_all_provider_classes_are_registered(yaml_files: dict[str, dict]) -> tu
 
                 for base_class, resource_type in base_class_resource_map:
                     if issubclass(obj, base_class) and obj is not base_class:
-                        num_checks += 1
                         full_class_name = f"{module_name}.{attr_name}"
+                        if full_class_name in INTERNAL_UNREGISTERED_TOOLSET_CLASSES:
+                            break
+                        num_checks += 1
                         # Executors and notifications are registered by class path;
                         # other types are registered by module path.
                         if resource_type in class_level_resource_types:
@@ -889,7 +905,7 @@ def check_invalid_integration(yaml_files: dict[str, dict]) -> tuple[int, int]:
     num_errors = 0
     num_integrations = len(all_integration_names)
     for (yaml_file_path, provider_data), resource_type in itertools.product(
-        yaml_files.items(), ["sensors", "operators", "hooks", "triggers", "bundles"]
+        yaml_files.items(), ["sensors", "operators", "hooks", "triggers", "bundles", "toolsets"]
     ):
         resource_data = provider_data.get(resource_type, [])
         current_names = {r["integration-name"] for r in resource_data}