After a tool confirmation resumes a Workflow, an LlmAgent node's model receives the workflow's user message instead of the node's input #7650

"""An LlmAgent node of a Workflow whose tool requires confirmation: after the confirmation resumes
the run, the model's conversation starts with the workflow's user message instead of the node input."""
import asyncio
from importlib.metadata import version

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 Runner
from google.adk.sessions import InMemorySessionService
from google.adk.tools import FunctionTool
from google.adk.workflow import START, FunctionNode, Workflow
from google.genai import types

seen = []


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

    async def generate_content_async(self, llm_request, stream=False):
        first = [part.text for part in llm_request.contents[0].parts if part.text]
        seen.append(first)
        answered = any(part.function_response for content in llm_request.contents for part in content.parts or [])
        part = (types.Part(text="done") if answered else
                types.Part(function_call=types.FunctionCall(id="call-1", name="add", args={"a": 2, "b": 3})))
        yield LlmResponse(content=types.Content(role="model", parts=[part]))


def add(a: int, b: int) -> int:
    return a + b


def prepare(node_input):
    return "node input: add 2 and 3"


agent = LlmAgent(name="agent", model=Model(), tools=[FunctionTool(add, require_confirmation=True)])
prepared = FunctionNode(func=prepare, name="prepare")
workflow = Workflow(name="root", edges=[(START, prepared), (prepared, agent)])


async def main():
    sessions = InMemorySessionService()
    runner = Runner(app_name="app", node=workflow, session_service=sessions)
    session = await sessions.create_session(app_name="app", user_id="u")
    message = types.Content(role="user", parts=[types.Part(text="workflow message")])
    while True:
        pending = []
        async for event in runner.run_async(user_id="u", session_id=session.id, new_message=message):
            for part in event.content.parts or [] if event.content else []:
                if part.function_call and part.function_call.id in (event.long_running_tool_ids or ()):
                    pending.append(part.function_call)
        if not pending:
            break
        message = types.Content(role="user", parts=[types.Part(function_response=types.FunctionResponse(
            id=call.id, name=call.name, response={"confirmed": True})) for call in pending])
    print("google-adk", version("google-adk"))
    for index, first in enumerate(seen):
        print(f"model call {index}: conversation starts with {first}")


asyncio.run(main())