Aborting during parallel tool calls records tools that already finished as "aborted" and drops their state changes #7958

import asyncio

from google.adk.agents import LlmAgent
from google.adk.models.base_llm import BaseLlm
from google.adk.models.llm_response import LlmResponse
from google.adk.runners import InMemoryRunner
from google.adk.tools.tool_context import ToolContext
from google.genai import types

sent = []
slow_started = asyncio.Event()


def send_email(to: str, tool_context: ToolContext) -> dict:
  """Sends an email."""
  sent.append(to)
  tool_context.state["emails_sent"] = len(sent)
  return {"status": "sent", "to": to}


async def lookup_order(order_id: str) -> dict:
  """Looks up an order (slow)."""
  slow_started.set()
  await asyncio.sleep(30)
  return {"order_id": order_id}


class ScriptedModel(BaseLlm):
  model: str = "scripted"

  async def generate_content_async(self, llm_request, stream=False):
    yield LlmResponse(
        content=types.ModelContent(
            parts=[
                types.Part.from_function_call(
                    name="send_email", args={"to": "a@example.com"}
                ),
                types.Part.from_function_call(
                    name="lookup_order", args={"order_id": "42"}
                ),
            ]
        )
    )


async def main():
  agent = LlmAgent(
      name="agent", model=ScriptedModel(), tools=[send_email, lookup_order]
  )
  runner = InMemoryRunner(agent=agent)
  session = await runner.session_service.create_session(
      app_name=runner.app_name, user_id="u"
  )
  abort_signal = asyncio.Event()

  async def run():
    async for _ in runner.run_async(
        user_id="u",
        session_id=session.id,
        new_message=types.UserContent("email a@example.com and check order 42"),
        abort_signal=abort_signal,
    ):
      pass

  task = asyncio.create_task(run())
  await slow_started.wait()
  abort_signal.set()  # e.g. the user presses stop
  await task

  session = await runner.session_service.get_session(
      app_name=runner.app_name, user_id="u", session_id=session.id
  )
  print("send_email actually ran:", sent)
  print("state emails_sent:", session.state.get("emails_sent"))
  for event in session.events:
    for response in event.get_function_responses():
      print("recorded response:", response.name, response.response)


asyncio.run(main())