Parallel tool call containing both ToolReturn and deferred tool does not work

[yf-yang]

Initial Checks

• I confirm that I'm using the latest version of Pydantic AI
• I confirm that I searched for my issue in https://github.com/pydantic/pydantic-ai/issues before opening this issue

Description

Related: #2360

@DouweM I am afraid #2365 is not enough to fix the issue.

The cause:
If parallel tool calls containing both tool A (whose return value is a ToolReturn) and a deferred tool B, then
When the agent exits, the last parts of message history is:
-2: ToolReturnPart (A.return_value)
-1: UserPromptPart (A.content)

The order breaks most LLM's order, as the tool return of the deferred tool call is not yet added.

Example Code

# Modified from https://ai.pydantic.dev/toolsets/#deferred-toolset
# `get_user_name` return value is modified
# prompt is modified to force a parallel tool call

import pprint

from pydantic import BaseModel
from pydantic_ai import Agent
from pydantic_ai.messages import (
  ModelMessage,
  ModelRequest,
  RetryPromptPart,
  ToolReturn,
  ToolReturnPart,
  UserPromptPart,
)
from pydantic_ai.output import DeferredToolCalls
from pydantic_ai.tools import ToolDefinition
from pydantic_ai.toolsets import DeferredToolset
from pydantic_ai.toolsets.function import FunctionToolset

toolset = FunctionToolset()


@toolset.tool
def get_default_language():
  return "en-US"


@toolset.tool
def get_user_name():
  return ToolReturn(
    return_value="David",
    content="The user's name is David",
    metadata="David",
  )


class PersonalizedGreeting(BaseModel):
  greeting: str
  language_code: str


agent = Agent("anthropic:claude-sonnet-4-0", toolsets=[toolset], output_type=PersonalizedGreeting)


def run_agent(
  messages: list[ModelMessage] = [], frontend_tools: list[ToolDefinition] = {}
) -> tuple[PersonalizedGreeting | DeferredToolCalls, list[ModelMessage]]:
  deferred_toolset = DeferredToolset(frontend_tools)
  result = agent.run_sync(
    toolsets=[deferred_toolset],
    output_type=[agent.output_type, DeferredToolCalls],
    message_history=messages,
  )
  return result.output, result.new_messages()


frontend_tool_definitions = [
  ToolDefinition(
    name="get_preferred_language",
    parameters_json_schema={"type": "object", "properties": {"default_language": {"type": "string"}}},
    description="Get the user's preferred language from their browser",
  )
]


def get_preferred_language(default_language: str) -> str:
  return "es-MX"


frontend_tool_functions = {"get_preferred_language": get_preferred_language}

messages: list[ModelMessage] = [
  ModelRequest(
    parts=[UserPromptPart(content="Call get_preferred_language and get_user_name in parallel")]
  )
]

final_output = None
while True:
  output, new_messages = run_agent(messages, frontend_tool_definitions)
  messages += new_messages

  if not isinstance(output, DeferredToolCalls):
    final_output = output
    break

  print(output.tool_calls)
  for tool_call in output.tool_calls:
    if function := frontend_tool_functions.get(tool_call.tool_name):
      part = ToolReturnPart(
        tool_name=tool_call.tool_name,
        content=function(**tool_call.args_as_dict()),
        tool_call_id=tool_call.tool_call_id,
      )
    else:
      part = RetryPromptPart(
        tool_name=tool_call.tool_name,
        content=f"Unknown tool {tool_call.tool_name!r}",
        tool_call_id=tool_call.tool_call_id,
      )
    messages.append(ModelRequest(parts=[part]))
  print("Current messages:")
  pprint.pprint(messages)

print(repr(final_output))

Python, Pydantic AI & LLM client version

latest

[DouweM]

@yf-yang That's a good point, the deferred tool return part would need to be inserted at the appropriate place in the message history, ahead of the extra user prompt part, which is currently up to the user code.

In #1995 (comment), I mentioned we could add a deferred_tool_results argument to run/iter, which could automatically add it where it belongs. Is that the sort of solution you had in mind?

[yf-yang]

@DouweM Thanks, I don't quite understand what deferred_tool_results does. Is it an explicit param just like user_prompt, so that the agent can check if it needs to be moved forward before it begins to execute?

I am not so sure. Sounds OK but maybe I'd prefer merging user_prompt and the proposed param into one param, as people will only use at most one of them.

Alternatively, is it a good idea to design a fix_message_history function that reorders it? Maybe not.

[DouweM]

@yf-yang Yeah I'm imagining an explicit param to make it clear that when the first run ends on a DeferredToolCalls object, there's a native way to then start the next run with a list of ToolReturnPart | RetryPromptParts that need to have matching tool_call_ids and can then be inserted at the appropriate place in the message history.

I agree this couldn't be used along with user_prompt, but I'm not sure if merging the params into one with a ton of supported types would be as clear as an explicit deferred_tool_calls, as otherwise the generically named param taking ToolReturnPart | RetryPromptPart would suggest that those can always be passed, while it's really only for deferred tool call scenarios.

Alternatively, is it a good idea to design a fix_message_history function that reorders it? Maybe not.

I feel like that'd be working around our technically incorrect documentation of "just add the tool return parts at the end of the message history", so I'd rather fix our documentation (and API) to not require the user to first do something that's technically incorrect.

[yf-yang]

I agree this couldn't be used along with user_prompt, but I'm not sure if merging the params into one with a ton of supported types would be as clear as an explicit deferred_tool_calls, as otherwise the generically named param taking ToolReturnPart | RetryPromptPart would suggest that those can always be passed, while it's really only for deferred tool call scenarios.

@DouweM From my very personal perspective, first I am OK with something like langgraph's graph.stream(message), where the message could be any type that could be later executed. Second, since both client side tool call and user message can initiate a new agent run, I am creating the same function to reuse across different input types. If they are split into multiple params, then I'll need to add another if else to determine how to convert the input type to correct argument name