Support SDK-local custom tool preload

Author: dhawal1

python/composio/core/models/custom_tool.py

@@ -156,6 +156,7 @@ def _create_tool(
     execute: CustomToolExecuteFn,
     extends_toolkit: t.Optional[str] = None,
     output_params: t.Optional[t.Type[BaseModel]] = None,
+    preload: t.Optional[bool] = None,
 ) -> CustomTool:
     """Internal: create and validate a CustomTool."""
     context = "experimental.tool"
@@ -218,6 +219,7 @@ def _create_tool(
         output_schema=output_schema,
         input_params=input_params,
         execute=execute,
+        preload=preload,
     )
 
 
@@ -258,6 +260,7 @@ def _infer_tool_from_function(
     description: t.Optional[str] = None,
     extends_toolkit: t.Optional[str] = None,
     output_params: t.Optional[t.Type[BaseModel]] = None,
+    preload: t.Optional[bool] = None,
     annotation_locals: t.Optional[t.Mapping[str, t.Any]] = None,
 ) -> CustomTool:
     """Create a CustomTool by inferring metadata from a decorated function.
@@ -341,6 +344,7 @@ def execute(input: t.Any, ctx: t.Any) -> t.Dict[str, t.Any]:
         execute=execute,
         extends_toolkit=extends_toolkit,
         output_params=output_params,
+        preload=preload,
     )
 
 
@@ -372,7 +376,14 @@ def search_code(input: SearchInput, ctx):
             return {"results": []}
     """
 
-    def __init__(self, *, slug: str, name: str, description: str) -> None:
+    def __init__(
+        self,
+        *,
+        slug: str,
+        name: str,
+        description: str,
+        preload: t.Optional[bool] = None,
+    ) -> None:
         context = "experimental.Toolkit"
         _validate_slug(slug, context)
         if not name:
@@ -383,6 +394,7 @@ def __init__(self, *, slug: str, name: str, description: str) -> None:
         self.slug = slug
         self.name = name
         self.description = description
+        self.preload = preload
         self._tools: t.List[CustomTool] = []
 
     @property
@@ -400,6 +412,7 @@ def tool(
         name: t.Optional[str] = None,
         description: t.Optional[str] = None,
         output_params: t.Optional[t.Type[BaseModel]] = None,
+        preload: t.Optional[bool] = None,
     ) -> t.Callable[[t.Callable[..., t.Any]], CustomTool]: ...
 
     def tool(
@@ -410,6 +423,7 @@ def tool(
         name: t.Optional[str] = None,
         description: t.Optional[str] = None,
         output_params: t.Optional[t.Type[BaseModel]] = None,
+        preload: t.Optional[bool] = None,
     ) -> t.Union[CustomTool, t.Callable[[t.Callable[..., t.Any]], CustomTool]]:
         """Decorator to add a tool to this toolkit.
 
@@ -425,6 +439,7 @@ def decorator(f: t.Callable[..., t.Any]) -> CustomTool:
                 name=name,
                 description=description,
                 output_params=output_params,
+                preload=preload,
                 annotation_locals=annotation_locals,
                 # No extends_toolkit for toolkit tools
             )
@@ -441,6 +456,7 @@ def decorator(f: t.Callable[..., t.Any]) -> CustomTool:
                 name=name,
                 description=description,
                 output_params=output_params,
+                preload=preload,
                 annotation_locals=_get_caller_locals(),
             )
             _validate_slug_length(
@@ -479,6 +495,7 @@ def tool(
         description: t.Optional[str] = None,
         extends_toolkit: t.Optional[str] = None,
         output_params: t.Optional[t.Type[BaseModel]] = None,
+        preload: t.Optional[bool] = None,
     ) -> t.Callable[[t.Callable[..., t.Any]], CustomTool]: ...
 
     def tool(
@@ -490,6 +507,7 @@ def tool(
         description: t.Optional[str] = None,
         extends_toolkit: t.Optional[str] = None,
         output_params: t.Optional[t.Type[BaseModel]] = None,
+        preload: t.Optional[bool] = None,
     ) -> t.Union[CustomTool, t.Callable[[t.Callable[..., t.Any]], CustomTool]]:
         """Decorator to create a custom tool from a function.
 
@@ -526,6 +544,7 @@ def decorator(f: t.Callable[..., t.Any]) -> CustomTool:
                 description=description,
                 extends_toolkit=extends_toolkit,
                 output_params=output_params,
+                preload=preload,
                 annotation_locals=annotation_locals,
             )
 
@@ -537,6 +556,7 @@ def decorator(f: t.Callable[..., t.Any]) -> CustomTool:
                 description=description,
                 extends_toolkit=extends_toolkit,
                 output_params=output_params,
+                preload=preload,
                 annotation_locals=_get_caller_locals(),
             )
         return decorator
@@ -547,7 +567,28 @@ def decorator(f: t.Callable[..., t.Any]) -> CustomTool:
 # ────────────────────────────────────────────────────────────────
 
 
-def serialize_custom_tools(tools: t.List[CustomTool]) -> t.List[t.Dict[str, t.Any]]:
+def _serialized_preload_value(
+    preload: t.Optional[bool],
+    inherited_preload: t.Optional[bool],
+    default_preload: bool,
+) -> t.Optional[bool]:
+    resolved = (
+        preload
+        if preload is not None
+        else inherited_preload
+        if inherited_preload is not None
+        else default_preload
+    )
+    if preload is not None or inherited_preload is not None or default_preload:
+        return resolved
+    return None
+
+
+def serialize_custom_tools(
+    tools: t.List[CustomTool],
+    *,
+    default_preload: bool = False,
+) -> t.List[t.Dict[str, t.Any]]:
     """Serialize custom tools into the format expected by the backend."""
     result = []
     for tool in tools:
@@ -561,12 +602,19 @@ def serialize_custom_tools(tools: t.List[CustomTool]) -> t.List[t.Dict[str, t.An
             entry["output_schema"] = tool.output_schema
         if tool.extends_toolkit:
             entry["extends_toolkit"] = tool.extends_toolkit
+        preload = _serialized_preload_value(
+            tool.preload, inherited_preload=None, default_preload=default_preload
+        )
+        if preload is not None:
+            entry["preload"] = preload
         result.append(entry)
     return result
 
 
 def serialize_custom_toolkits(
     toolkits: t.Sequence[ExperimentalToolkit],
+    *,
+    default_preload: bool = False,
 ) -> t.List[t.Dict[str, t.Any]]:
     """Serialize custom toolkits into the format expected by the backend."""
     result = []
@@ -581,15 +629,26 @@ def serialize_custom_toolkits(
             }
             if tool.output_schema:
                 entry["output_schema"] = tool.output_schema
+            preload = _serialized_preload_value(
+                tool.preload,
+                inherited_preload=tk.preload,
+                default_preload=default_preload,
+            )
+            if preload is not None:
+                entry["preload"] = preload
             toolkit_tools.append(entry)
-        result.append(
-            {
-                "slug": tk.slug,
-                "name": tk.name,
-                "description": tk.description,
-                "tools": toolkit_tools,
-            }
+        toolkit_entry: t.Dict[str, t.Any] = {
+            "slug": tk.slug,
+            "name": tk.name,
+            "description": tk.description,
+            "tools": toolkit_tools,
+        }
+        preload = _serialized_preload_value(
+            tk.preload, inherited_preload=None, default_preload=default_preload
         )
+        if preload is not None:
+            toolkit_entry["preload"] = preload
+        result.append(toolkit_entry)
     return result
 
 
@@ -729,23 +788,63 @@ def find_custom_tool_map_entry_by_final_slug(
     return custom_tools_map.by_final_slug.get(slug.upper())
 
 
+def assert_no_custom_tool_slugs_in_preload(
+    preload_tools: t.Union[t.Sequence[str], t.Literal["all"], None],
+    custom_tools_map: t.Optional[CustomToolsMap],
+) -> None:
+    """Reject legacy top-level preload of custom tool slugs."""
+    if preload_tools is None or preload_tools == "all":
+        return
+
+    custom_preload_slugs = []
+    for slug in preload_tools:
+        normalized = slug.upper()
+        if normalized.startswith(LOCAL_TOOL_PREFIX) or (
+            custom_tools_map is not None
+            and (
+                normalized in custom_tools_map.by_original_slug
+                or normalized in custom_tools_map.by_final_slug
+            )
+        ):
+            custom_preload_slugs.append(slug)
+
+    if custom_preload_slugs:
+        raise ValidationError(
+            "Custom tool slugs are not supported in preload.tools: "
+            f"{', '.join(custom_preload_slugs)}. Set preload=True on the SDK "
+            "custom tool or custom toolkit definition instead."
+        )
+
+
 def get_preloaded_custom_tool_slugs(
-    tool_router_tools: t.Optional[t.Iterable[str]],
     custom_tools_map: t.Optional[CustomToolsMap],
+    *,
+    default_preload: bool = False,
 ) -> t.List[str]:
-    """Return final custom tool slugs selected by the backend for preload."""
-    if tool_router_tools is None or custom_tools_map is None:
+    """Return final custom tool slugs selected locally for preload."""
+    if custom_tools_map is None:
         return []
 
     seen: t.Set[str] = set()
     custom_tool_slugs: t.List[str] = []
 
-    for slug in tool_router_tools:
-        entry = find_custom_tool_map_entry_by_final_slug(
-            custom_tools_map,
-            str(slug),
+    for entry in custom_tools_map.by_final_slug.values():
+        toolkit = next(
+            (
+                tk
+                for tk in custom_tools_map.toolkits or []
+                if entry.toolkit and tk.slug.lower() == entry.toolkit.lower()
+            ),
+            None,
+        )
+        should_preload = (
+            entry.handle.preload
+            if entry.handle.preload is not None
+            else toolkit.preload
+            if toolkit is not None and toolkit.preload is not None
+            else default_preload
         )
-        if entry is None:
+        if not should_preload:
             continue
 
         final_slug_key = entry.final_slug.upper()

python/composio/core/models/custom_tool_types.py

@@ -106,6 +106,7 @@ class CustomTool:
     execute: CustomToolExecuteFn
     extends_toolkit: t.Optional[str] = None
     output_schema: t.Optional[t.Dict[str, t.Any]] = None
+    preload: t.Optional[bool] = None
 
 
 # ────────────────────────────────────────────────────────────────

python/composio/core/models/tool_router.py

@@ -19,6 +19,8 @@
 from composio.core.models.base import Resource
 from composio.core.models.custom_tool import (
     ExperimentalToolkit,
+    assert_no_custom_tool_slugs_in_preload,
+    build_custom_tools_map,
     build_custom_tools_map_from_response,
     get_preloaded_custom_tool_slugs,
     serialize_custom_tools,
@@ -627,6 +629,10 @@ def create(
                             - 'assistive_prompt' (dict): Configuration for assistive prompt generation.
                               - 'user_timezone' (str): IANA timezone identifier
                                 (e.g., "America/New_York", "Europe/London").
+                            - 'custom_tools' / 'custom_toolkits': SDK custom tools.
+                              Set preload=True on a custom tool or toolkit to expose
+                              it directly from session.tools(); otherwise custom tools
+                              remain search-only.
                             Example: {'assistive_prompt': {'user_timezone': 'America/New_York'}}
         :return: Tool router session object
 
@@ -716,6 +722,7 @@ def create(
             workbench=workbench,
             preload=preload,
         )
+        default_custom_preload = preload is not None and preload.get("tools") == "all"
 
         # Parse manage_connections config
         manage_connections = (
@@ -851,6 +858,8 @@ def create(
         # experimental.assistive_prompt_config.user_timezone
         custom_tools: t.Optional[t.List[CustomTool]] = None
         custom_toolkits: t.Optional[t.List[ExperimentalToolkit]] = None
+        local_custom_tools_map: t.Optional[CustomToolsMap] = None
+        inline_custom_tools_payload: t.Optional[t.Dict[str, t.Any]] = None
 
         if experimental is not None:
             experimental_payload: t.Dict[str, t.Any] = {}
@@ -866,16 +875,36 @@ def create(
             # Serialize custom tools and toolkits for the backend
             custom_tools = experimental.get("custom_tools")
             custom_toolkits = experimental.get("custom_toolkits")
+            if custom_tools or custom_toolkits:
+                local_custom_tools_map = build_custom_tools_map(
+                    custom_tools or [],
+                    custom_toolkits,
+                )
+                assert_no_custom_tool_slugs_in_preload(
+                    preload.get("tools") if preload is not None else None,
+                    local_custom_tools_map,
+                )
 
             if custom_tools:
                 experimental_payload["custom_tools"] = serialize_custom_tools(
-                    custom_tools
+                    custom_tools,
+                    default_preload=default_custom_preload,
                 )
             if custom_toolkits:
                 experimental_payload["custom_toolkits"] = serialize_custom_toolkits(
-                    custom_toolkits
+                    custom_toolkits,
+                    default_preload=default_custom_preload,
                 )
 
+            if (
+                "custom_tools" in experimental_payload
+                or "custom_toolkits" in experimental_payload
+            ):
+                inline_custom_tools_payload = {
+                    "custom_tools": experimental_payload.get("custom_tools"),
+                    "custom_toolkits": experimental_payload.get("custom_toolkits"),
+                }
+
             if experimental_payload:
                 create_params["experimental"] = experimental_payload
 
@@ -897,8 +926,8 @@ def create(
                 experimental=session.experimental,
             )
         preloaded_custom_tool_slugs = get_preloaded_custom_tool_slugs(
-            getattr(session, "tool_router_tools", None),
             custom_tools_map,
+            default_preload=default_custom_preload,
         )
 
         # Transform experimental response:
@@ -930,6 +959,7 @@ def create(
             user_id=user_id,
             preload=_session_preload_config(session),
             preloaded_custom_tool_slugs=preloaded_custom_tool_slugs,
+            inline_custom_tools_payload=inline_custom_tools_payload,
         )
 
     def use(self, session_id: str) -> ToolRouterSession[TTool, TToolCollection]: