Python: Add tool concurrency groups and sequential execution order for same-message calls

python/packages/core/agent_framework/_tools.py

@@ -320,6 +320,7 @@ def __init__(
         func: Callable[..., Any] | None = None,
         input_model: type[BaseModel] | Mapping[str, Any] | None = None,
         result_parser: Callable[[Any], str | list[Content]] | _SkipParsingSentinel | None = None,
+        concurrency_group: str | None = None,
         **kwargs: Any,
     ) -> None:
         """Initialize the FunctionTool.
@@ -334,10 +335,11 @@ def __init__(
             max_invocations: The maximum number of times this function can be invoked
                 across the **lifetime of this tool instance**. If None (default),
                 there is no limit. Should be at least 1. If the tool is called multiple
-                times in one iteration, those will execute, after that it will stop working. For example,
-                if max_invocations is 3 and the tool is called 5 times in a single iteration,
-                these will complete, but any subsequent calls to the tool (in the same or future iterations)
-                will raise a ToolException.
+                times in one iteration, those will execute, after that it will stop
+                working. For example, if max_invocations is 3 and the tool is called 5
+                times in a single iteration, these will complete, but any subsequent
+                calls to the tool (in the same or future iterations) will raise a
+                ToolException.
 
                 .. note::
                     This counter lives on the tool instance and is never automatically
@@ -348,30 +350,37 @@ def __init__(
                     ``FunctionInvocationConfiguration["max_function_calls"]``
                     for per-request limits instead.
 
-            max_invocation_exceptions: The maximum number of exceptions allowed during invocations.
-                If None, there is no limit. Should be at least 1.
+            max_invocation_exceptions: The maximum number of exceptions allowed
+                during invocations. If None, there is no limit. Should be at least 1.
             additional_properties: Additional properties to set on the function.
-            func: The function to wrap. When ``None``, creates a declaration-only tool
-                that has no implementation. Declaration-only tools are useful when you want
-                the agent to reason about tool usage without executing them, or when the
-                actual implementation exists elsewhere (e.g., client-side rendering).
-            input_model: The Pydantic model that defines the input parameters for the function.
-                This can also be a JSON schema dictionary.
-                If not provided and ``func`` is not ``None``, it will be inferred from
-                the function signature. When ``func`` is ``None`` and ``input_model`` is
-                not provided, the tool will use an empty input model (no parameters) in
-                its JSON schema. For declaration-only tools that should declare
-                parameters, explicitly provide ``input_model`` (either a Pydantic
-                ``BaseModel`` or a JSON schema dictionary) so the model can reason about
-                the expected arguments.
-            result_parser: An optional callable with signature ``Callable[[Any], str]`` that
-                overrides the default result parsing behavior. When provided, this callable
-                is used to convert the raw function return value to a string instead of the
-                built-in :meth:`parse_result` logic. Pass the :data:`SKIP_PARSING` sentinel
-                instead of a callable to opt out of parsing entirely; in that case
-                :meth:`invoke` returns the wrapped function's raw return value. Depending
-                on your function, it may be easiest to just do the serialization directly
-                in the function body rather than providing a custom ``result_parser``.
+            func: The function to wrap. When ``None``, creates a declaration-only
+                tool that has no implementation. Declaration-only tools are useful
+                when you want the agent to reason about tool usage without executing
+                them, or when the actual implementation exists elsewhere (e.g.,
+                client-side rendering).
+            input_model: The Pydantic model that defines the input parameters for the
+                function. This can also be a JSON schema dictionary.
+                If not provided and ``func`` is not ``None``, it will be inferred
+                from the function signature. When ``func`` is ``None`` and
+                ``input_model`` is not provided, the tool will use an empty input
+                model (no parameters) in its JSON schema. For declaration-only tools
+                that should declare parameters, explicitly provide ``input_model``
+                (either a Pydantic ``BaseModel`` or a JSON schema dictionary) so the
+                model can reason about the expected arguments.
+            result_parser: An optional callable with signature ``Callable[[Any], str]``
+                that overrides the default result parsing behavior. When provided,
+                this callable is used to convert the raw function return value to a
+                string instead of the built-in :meth:`parse_result` logic. Pass the
+                :data:`SKIP_PARSING` sentinel instead of a callable to opt out of
+                parsing entirely; in that case :meth:`invoke` returns the wrapped
+                function's raw return value. Depending on your function, it may be
+                easiest to just do the serialization directly in the function body
+                rather than providing a custom ``result_parser``.
+            concurrency_group: If provided, tool calls with the same
+                concurrency_group will execute sequentially in the order they were
+                invoked by the model. Tools without a group, or with different
+                groups, will execute concurrently. Useful for stateful tools with
+                write->read dependencies to prevent race conditions.
             **kwargs: Additional keyword arguments.
         """
         # Core attributes (formerly from BaseTool)
@@ -416,6 +425,7 @@ def __init__(
         self._invocation_duration_histogram = _default_histogram()
         self.type: Literal["function_tool"] = "function_tool"
         self.result_parser = result_parser
+        self.concurrency_group = concurrency_group
 
     def _discover_injected_parameters(self) -> None:
         """Inspect the wrapped function for runtime injection parameters."""
@@ -899,9 +909,13 @@ def to_json_schema_spec(self) -> dict[str, Any]:
     @override
     def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
         as_dict = super().to_dict(exclude=exclude, exclude_none=exclude_none)
+        if (not exclude or "concurrency_group" not in exclude) and (
+            not exclude_none or self.concurrency_group is not None
+        ):
+            as_dict["concurrency_group"] = self.concurrency_group
         if (exclude and "input_model" in exclude) or not self.input_model:
             return as_dict
-        as_dict["input_model"] = self.parameters()  # Use cached parameters()
+        as_dict["input_model"] = self.parameters()
         return as_dict
 
 
@@ -1138,6 +1152,7 @@ def tool(
     max_invocations: int | None = None,
     max_invocation_exceptions: int | None = None,
     additional_properties: dict[str, Any] | None = None,
+    concurrency_group: str | None = None,
     result_parser: Callable[[Any], str | list[Content]] | _SkipParsingSentinel | None = None,
 ) -> FunctionTool: ...
 
@@ -1154,6 +1169,7 @@ def tool(
     max_invocations: int | None = None,
     max_invocation_exceptions: int | None = None,
     additional_properties: dict[str, Any] | None = None,
+    concurrency_group: str | None = None,
     result_parser: Callable[[Any], str | list[Content]] | _SkipParsingSentinel | None = None,
 ) -> Callable[[Callable[..., Any]], FunctionTool]: ...
 
@@ -1169,6 +1185,7 @@ def tool(
     max_invocations: int | None = None,
     max_invocation_exceptions: int | None = None,
     additional_properties: dict[str, Any] | None = None,
+    concurrency_group: str | None = None,
     result_parser: Callable[[Any], str | list[Content]] | _SkipParsingSentinel | None = None,
 ) -> FunctionTool | Callable[[Callable[..., Any]], FunctionTool]:
     """Decorate a function to turn it into a FunctionTool that can be passed to models and executed automatically.
@@ -1213,6 +1230,11 @@ def tool(
         max_invocation_exceptions: The maximum number of exceptions allowed during invocations.
             If None, there is no limit, should be at least 1.
         additional_properties: Additional properties to set on the function.
+        concurrency_group: If provided, tool calls with the same
+            concurrency_group will execute sequentially in the order they were
+            invoked by the model. Tools without a group, or with different
+            groups, will execute concurrently. Useful for stateful tools with
+            write->read dependencies to prevent race conditions.
         result_parser: An optional callable with signature ``Callable[[Any], str]`` that
             overrides the default result parsing. When provided, this callable converts the
             raw function return value to a string instead of using the built-in
@@ -1313,6 +1335,7 @@ def wrapper(f: Callable[..., Any]) -> FunctionTool:
                 func=f,
                 input_model=schema,
                 result_parser=result_parser,
+                concurrency_group=concurrency_group,
             )
 
         return wrapper(func)
@@ -1378,6 +1401,7 @@ class FunctionInvocationConfiguration(TypedDict, total=False):
     terminate_on_unknown_calls: bool
     additional_tools: Sequence[FunctionTool]
     include_detailed_errors: bool
+    tool_execution_order: Literal["parallel", "sequential"]
 
 
 def normalize_function_invocation_configuration(
@@ -1391,6 +1415,7 @@ def normalize_function_invocation_configuration(
         "terminate_on_unknown_calls": False,
         "additional_tools": [],
         "include_detailed_errors": False,
+        "tool_execution_order": "parallel",
     }
     if config:
         normalized.update(config)
@@ -1762,7 +1787,8 @@ async def _try_execute_function_call_groups(
     has_declaration_only_call = False
     # A user-input pause takes precedence over unknown-call termination in mixed batches.
     for function_call in actionable_calls:
-        function_name = function_call.name
+        function_name = _underlying_function_call(function_call).name
+
         logger.debug(
             "Checking function call: type=%s, name=%s, in approval_tools=%s",
             function_call.type,
@@ -1829,25 +1855,56 @@ async def _try_execute_function_call_groups(
     # Only a fully executable batch reaches this point; run calls concurrently but retain per-call result groups.
     # Create each task inside a copied context so the active agent span is
     # preserved for every parallel tool invocation.
-    execution_tasks = [
-        contextvars.copy_context().run(
-            asyncio.create_task,
-            _execute_single_function_call(
-                function_call,
-                custom_args=custom_args,
-                config=config,
-                tool_map=tool_map,
-                invocation_session=invocation_session,
-                middleware_pipeline=middleware_pipeline,
-                live_tools=live_tools,
-            ),
-        )
-        for function_call in function_calls
-    ]
-    execution_results = await asyncio.gather(*execution_tasks)
+    execution_order = config.get("tool_execution_order", "parallel")
+
+    groups: dict[str, list[int]] = {}
+    for idx, function_call in enumerate(function_calls):
+        group_key: str | None = None
+
+        if execution_order == "parallel":
+            tool_name = _underlying_function_call(function_call).name
+            if tool_name is not None:
+                tool = tool_map.get(tool_name)
+                if tool is not None:
+                    group_key = getattr(tool, "concurrency_group", None)
+
+        if group_key is None:
+            group_key = "__sequential_all__" if execution_order == "sequential" else f"__ungrouped_{idx}"
+
+        if group_key not in groups:
+            groups[group_key] = []
+        groups[group_key].append(idx)
+
+    ordered_results: list[tuple[list[Content], bool] | None] = [None] * len(function_calls)
+
+    async def _execute_group(indices: list[int]) -> None:
+        for idx in indices:
+            call = function_calls[idx]
+            ctx = contextvars.copy_context()
+            task = ctx.run(
+                asyncio.create_task,
+                _execute_single_function_call(
+                    call,
+                    custom_args=custom_args,
+                    config=config,
+                    tool_map=tool_map,
+                    invocation_session=invocation_session,
+                    middleware_pipeline=middleware_pipeline,
+                    live_tools=live_tools,
+                ),
+            )
+            res = await task
+            ordered_results[idx] = res
 
-    should_terminate = any(terminate for _, terminate in execution_results)
-    return [result_contents for result_contents, _ in execution_results], should_terminate
+    execution_tasks = [asyncio.create_task(_execute_group(indices)) for indices in groups.values()]
+
+    await asyncio.gather(*execution_tasks)
+
+    if any(result is None for result in ordered_results):
+        raise RuntimeError("Internal error: missing tool execution result(s).")
+    completed_results = cast(list[tuple[list[Content], bool]], ordered_results)
+    should_terminate = any(terminate for _, terminate in completed_results)
+    return [result_contents for result_contents, _ in completed_results], should_terminate
 
 
 @dataclass
@@ -1882,6 +1939,11 @@ async def _execute_function_calls(
     invocation_session: AgentSession | None = None,
     middleware_pipeline: FunctionMiddlewarePipeline | None = None,
 ) -> _FunctionExecutionBatch:
+
+    run_config = cast("FunctionInvocationConfiguration", dict(config) if config else {})
+    if custom_args and "tool_execution_order" in custom_args:
+        run_config["tool_execution_order"] = custom_args["tool_execution_order"]
+
     tools = _extract_tools(options)
     if not tools:
         return _FunctionExecutionBatch(result_groups=[])
@@ -1891,7 +1953,7 @@ async def _execute_function_calls(
         tools=tools,
         invocation_session=invocation_session,
         middleware_pipeline=middleware_pipeline,
-        config=config,
+        config=run_config,
     )
     return _FunctionExecutionBatch(
         result_groups=result_groups,
@@ -3274,17 +3336,23 @@ def get_response(
         invocation_session = raw_session if isinstance(raw_session, _AgentSession) else None
 
         # Bind one executor with the run's custom arguments, middleware, configuration, and session.
+        options = dict(options) if options else {}
+        run_config = cast ("FunctionInvocationConfiguration", dict(self.function_invocation_configuration) if self.function_invocation_configuration else {})
+
+
+        if tool_exec_order := options.pop("tool_execution_order", None):
+            run_config["tool_execution_order"] = tool_exec_order
+
+        mutable_options: dict[str, Any] = dict(options)
+
         execute_function_calls = partial(
             _execute_function_calls,
             custom_args=additional_function_arguments,
-            config=self.function_invocation_configuration,
+            config=run_config,
             invocation_session=invocation_session,
             middleware_pipeline=function_middleware_pipeline,
         )
 
-        # Give the loop private mutable options and one shared run-local tool list for progressive tool changes.
-        # Make options mutable so we can update conversation_id during function invocation loop
-        mutable_options: dict[str, Any] = dict(options) if options else {}
         # Remove additional_function_arguments from options passed to underlying chat client
         # It's for tool invocation only and not recognized by chat service APIs
         mutable_options.pop("additional_function_arguments", None)

python/packages/core/agent_framework/_types.py

@@ -3705,6 +3705,10 @@ class _ChatOptionsBase(TypedDict, total=False):
     tool_choice: ToolMode | Literal["auto", "required", "none"]
     allow_multiple_tool_calls: bool
 
+    # Dictates whether multiple tool calls in a single message batch
+    # are executed concurrently (parallel) or one-by-one (sequential).
+    tool_execution_order: Literal["parallel", "sequential"]
+
     # Response configuration
     response_format: type[BaseModel] | Mapping[str, Any] | None