RemoteA2aAgent(mode="task"): a remote that completes with only the finish_task FunctionResponse, as the docstring says, closes the delegation with an empty output #7746

"""RemoteA2aAgent(mode="task") + a custom A2A server that signals completion
the way the `mode` docstring says: a FunctionResponse named finish_task."""

import asyncio
import sys
import warnings

import httpx
from a2a.helpers.proto_helpers import new_task, new_text_part
from a2a.server.agent_execution import AgentExecutor
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.routes import (
    add_a2a_routes_to_fastapi,
    create_agent_card_routes,
    create_jsonrpc_routes,
)
from a2a.server.tasks import InMemoryTaskStore, TaskUpdater
from a2a.types import AgentCapabilities, AgentCard, AgentInterface, TaskState
from fastapi import FastAPI
from google.adk.a2a.converters.part_converter import convert_genai_part_to_a2a_part
from google.adk.agents import LlmAgent
from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
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.genai import types

warnings.filterwarnings("ignore")
peer_inputs, coordinator_calls = [], []


CONTROL = "--control" in sys.argv  # also send the finish_task call


class PlainExecutor(AgentExecutor):  # answers, then signals completion as documented
    async def execute(self, context, event_queue):
        text = context.get_user_input()
        peer_inputs.append(
            "<delegation from coordinator>" if "called tool `math`" in text else text
        )
        await event_queue.enqueue_event(
            new_task(
                context.task_id, context.context_id, TaskState.TASK_STATE_SUBMITTED
            )
        )
        updater = TaskUpdater(event_queue, context.task_id, context.context_id)
        parts = [new_text_part("17 * 23 = 391")]
        if CONTROL:
            call = types.FunctionCall(
                id="ft1", name="finish_task", args={"result": "17 * 23 = 391"}
            )
            parts.append(convert_genai_part_to_a2a_part(types.Part(function_call=call)))
        done = types.FunctionResponse(
            id="ft1", name="finish_task", response={"result": "Task completed."}
        )
        parts.append(convert_genai_part_to_a2a_part(types.Part(function_response=done)))
        msg = updater.new_agent_message(parts)
        await updater.complete(message=msg)

    async def cancel(self, context, event_queue): ...


class CoordinatorModel(BaseLlm):  # delegates to `math`, then reports what came back
    async def generate_content_async(self, llm_request, stream=False):
        last = llm_request.contents[-1].parts[-1]
        coordinator_calls.append(
            last.function_response.name if last.function_response else last.text
        )
        if last.function_response:
            part = types.Part(text=f"math said: {last.function_response.response}")
        else:
            part = types.Part(
                function_call=types.FunctionCall(name="math", args={"request": "17*23"})
            )
        yield LlmResponse(content=types.Content(role="model", parts=[part]))


async def main():
    card = AgentCard(
        name="math",
        description="math",
        version="1",
        capabilities=AgentCapabilities(streaming=True),
        default_input_modes=["text/plain"],
        default_output_modes=["text/plain"],
        supported_interfaces=[
            AgentInterface(url="http://peer.test/a2a", protocol_binding="JSONRPC")
        ],
    )
    peer = FastAPI()
    add_a2a_routes_to_fastapi(
        peer,
        agent_card_routes=create_agent_card_routes(card),
        jsonrpc_routes=create_jsonrpc_routes(
            DefaultRequestHandler(PlainExecutor(), InMemoryTaskStore(), card),
            rpc_url="/a2a",
        ),
    )
    client = httpx.AsyncClient(transport=httpx.ASGITransport(app=peer))

    math = RemoteA2aAgent(
        name="math",
        description="math",
        agent_card=card,
        mode="task",
        httpx_client=client,
    )
    coordinator = LlmAgent(
        name="coordinator", model=CoordinatorModel(model="fake"), sub_agents=[math]
    )
    runner = Runner(
        app_name="app",
        agent=coordinator,
        session_service=InMemorySessionService(),
        auto_create_session=True,
    )

    for text in ["what is 17*23?", "thanks! now say hi"]:
        msg = types.Content(role="user", parts=[types.Part(text=text)])
        events = await asyncio.wait_for(
            _collect(runner.run_async(user_id="u", session_id="s", new_message=msg)),
            timeout=20,
        )
        print(f"user: {text!r}")
        for e in events:
            for p in e.content.parts if e.content and e.content.parts else []:
                if p.text:
                    print(f"  {e.author}: text {p.text!r}")
                if p.function_call:
                    print(f"  {e.author}: call {p.function_call.name}")
                if p.function_response:
                    fr = p.function_response
                    print(f"  {e.author}: response {fr.name}={fr.response}")
    print("coordinator model called with:", coordinator_calls)
    print("plain peer received:", peer_inputs)


async def _collect(agen):
    return [e async for e in agen]


asyncio.run(main())