Resumable Workflow skips an interrupted ordinary tool call; direct Agent root replays it #7719

"""No provider, credentials, private adapters, or ADK patches required."""
import asyncio
import json
import platform
import uuid
from importlib.metadata import version

from google.adk.agents import Agent
from google.adk.apps import App, ResumabilityConfig
from google.adk.models.base_llm import BaseLlm
from google.adk.models.llm_response import LlmResponse
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.adk.tools import BaseTool
from google.adk.workflow import START, Workflow
from google.genai import types


class Model(BaseLlm):
    model: str = "scripted"
    generations: int = 0
    final_responses: int = 0

    async def generate_content_async(self, llm_request, stream=False):
        answered = any(p.function_response for c in llm_request.contents for p in c.parts or [])
        if answered:
            self.final_responses += 1
            parts = [types.Part(text="done")]
        else:
            self.generations += 1
            if self.generations > 1:
                raise RuntimeError("pending call regenerated")
            parts = [types.Part(function_call=types.FunctionCall(
                id="call_" + uuid.uuid4().hex, name="work", args={"value": 7}))]
        yield LlmResponse(content=types.Content(role="model", parts=parts))


class Tool(BaseTool):
    def __init__(self, block):
        super().__init__(name="work", description="Ordinary async tool")
        self.block = block
        self.entered = asyncio.Event()
        self.calls = []

    def _get_declaration(self):
        return types.FunctionDeclaration(name=self.name, description=self.description,
            parameters_json_schema={"type": "object", "properties": {"value": {"type": "integer"}},
                                    "required": ["value"]})

    async def run_async(self, *, args, tool_context):
        self.calls.append((tool_context.function_call_id, args))
        self.entered.set()
        if self.block:
            await asyncio.Event().wait()
        return {"value": args["value"]}


def summarize(event):
    return {"author": event.author, "path": event.node_info.path,
            "call_names": [c.name for c in event.get_function_calls()],
            "response_names": [r.name for r in event.get_function_responses()],
            "agent_state": event.actions.agent_state,
            "output": event.output,
            "text": [p.text for p in event.content.parts or [] if p.text] if event.content else []}


async def probe(wrapped):
    service = InMemorySessionService()
    await service.create_session(app_name="proof", user_id="u", session_id="s")
    model = Model()

    async def launch(tool):
        agent = Agent(name="worker", model=model, tools=[tool])
        root = Workflow(name="root", edges=[(START, agent)]) if wrapped else agent
        runner = Runner(app=App(name="proof", root_agent=root,
            resumability_config=ResumabilityConfig(is_resumable=True)), session_service=service)
        return [event async for event in runner.run_async(user_id="u", session_id="s",
            invocation_id="same-invocation", new_message=types.Content(role="user",
            parts=[types.Part(text="go")]))]

    old = Tool(block=True)
    task = asyncio.create_task(launch(old))
    try:
        await asyncio.wait_for(old.entered.wait(), 5)
    finally:
        task.cancel()
        await asyncio.gather(task, return_exceptions=True)
    saved = await service.get_session(app_name="proof", user_id="u", session_id="s")
    before = [summarize(e) for e in saved.events]
    replacement = Tool(block=False)
    resumed = await asyncio.wait_for(launch(replacement), 5)
    saved = await service.get_session(app_name="proof", user_id="u", session_id="s")
    result = {"workflow": wrapped, "same_call_and_args": old.calls == replacement.calls,
        "original_calls": len(old.calls), "replacement_calls": len(replacement.calls),
        "call_generations": model.generations, "final_responses": model.final_responses,
        "saved_before_resume": before, "resumed_events": [summarize(e) for e in resumed],
        "user_events": sum(e.author == "user" for e in saved.events),
        "function_calls": sum(len(e.get_function_calls()) for e in saved.events),
        "function_responses": sum(len(e.get_function_responses()) for e in saved.events)}
    print(json.dumps(result, default=str))
    return result


async def main():
    print("ADK", version("google-adk"), "Python", platform.python_version(), platform.system())
    control = await probe(False)
    workflow = await probe(True)
    assert control["same_call_and_args"] and control["function_responses"] == 1
    assert control["call_generations"] == control["final_responses"] == 1
    assert workflow["replacement_calls"] == workflow["function_responses"] == 0
    assert workflow["call_generations"] == 1 and workflow["final_responses"] == 0


asyncio.run(main())