diff --git a/src/google/adk/flows/llm_flows/context/_contents.py b/src/google/adk/flows/llm_flows/context/_contents.py index 488f7e75a54..edc200e3710 100644 --- a/src/google/adk/flows/llm_flows/context/_contents.py +++ b/src/google/adk/flows/llm_flows/context/_contents.py @@ -335,6 +335,21 @@ def _build_task_input_user_content( parts.append(types.Part(text=_SINGLE_TURN_NUDGE)) return types.Content(role='user', parts=parts) + if '@' in isolation_scope and '/' not in isolation_scope: + invocation_id = isolation_scope.rpartition('@')[2] + for event in all_events: + if ( + event.invocation_id == invocation_id + and event.author == 'user' + and event.content + and event.content.parts + and not any(p.function_response for p in event.content.parts) + ): + parts = list(event.content.parts) + if is_single_turn: + parts.append(types.Part(text=_SINGLE_TURN_NUDGE)) + return types.Content(role='user', parts=parts) + # Fallback: workflow-node task with no originating FC. Use the # node_input that the wrapper stamped onto ``ic.user_content``. if user_content and user_content.parts: diff --git a/src/google/adk/workflow/_dynamic_node_scheduler.py b/src/google/adk/workflow/_dynamic_node_scheduler.py index 18dcc4e135d..0f36215a231 100644 --- a/src/google/adk/workflow/_dynamic_node_scheduler.py +++ b/src/google/adk/workflow/_dynamic_node_scheduler.py @@ -348,7 +348,10 @@ async def _execute_step( override_isolation_scope is None and getattr(node, 'mode', None) == AgentMode.TASK ): - override_isolation_scope = node_path + if not curr_parent_path: + override_isolation_scope = f'{target_node_name}@{ctx.invocation_id}' + else: + override_isolation_scope = node_path # Rehydration chronological sequence barrier setup for the parent path if self._enable_replay and curr_parent_path: diff --git a/tests/unittests/a2a/integration/test_client_server_v1.py b/tests/unittests/a2a/integration/test_client_server_v1.py index 37d1c32726a..d0392ff9dc0 100644 --- a/tests/unittests/a2a/integration/test_client_server_v1.py +++ b/tests/unittests/a2a/integration/test_client_server_v1.py @@ -37,6 +37,7 @@ from .client import create_a2a_client from .client import create_client +from .server import agent_card from .server import create_server_app_v1 @@ -847,3 +848,91 @@ def _status_update(text: str, state: int) -> pb.StreamResponse: # the latest status is applied to the running task. final_task, _ = results[2] assert final_task.status.state == pb.TASK_STATE_COMPLETED + + +@pytest.mark.asyncio +async def test_a2a_task_clarification_and_resumption(): + """Tests that root task input preservation works across the A2A boundary. + + When a RemoteA2aAgent operates in mode="task", it acts as a sub-agent on the + client, but runs as a root agent on the server. The server-side Runner must + successfully associate the clarification turn with the paused task's scope + and reconstruct the strictly-alternating task history correctly. + """ + from google.adk.agents.llm_agent import LlmAgent + from tests.unittests import testing_utils + + model = testing_utils.MockModel.create( + responses=[ + "What is your specific question?", + "Task completed.", + ] + ) + server_agent = LlmAgent(name="my_task", model=model, mode="task") + + from google.adk.runners import Runner + from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor + from google.adk.a2a import _compat + from starlette.applications import Starlette + from a2a.server.tasks import InMemoryTaskStore as TaskStore + + server_runner = Runner( + app_name="ServerApp", + agent=server_agent, + session_service=InMemorySessionService(), + ) + executor = A2aAgentExecutor(runner=server_runner) + app = Starlette() + _compat.attach_a2a_routes_to_app( + app, + agent_card=agent_card, + agent_executor=executor, + task_store=TaskStore(), + ) + + async with app.router.lifespan_context(app): + client_agent = create_client(app, streaming=False, mode="task") + + session_service = InMemorySessionService() + await session_service.create_session( + app_name="ClientApp", user_id="test_user", session_id="test_session" + ) + client_runner = Runner( + app_name="ClientApp", + agent=client_agent, + session_service=session_service, + ) + + # Turn 1: Initial task request + new_message = types.Content( + parts=[types.Part(text="Do a task")], role="user" + ) + async for _ in client_runner.run_async( + user_id="test_user", session_id="test_session", new_message=new_message + ): + pass + + # Turn 2: Clarification response + new_message2 = types.Content( + parts=[types.Part(text="Here is the clarification")], role="user" + ) + async for _ in client_runner.run_async( + user_id="test_user", session_id="test_session", new_message=new_message2 + ): + pass + + # Verify server-side model requests to ensure correct history reconstruction + assert len(model.requests) == 2 + + req1 = model.requests[0] + parts1 = [p.text for c in req1.contents if c.role == "user" for p in c.parts] + assert "Do a task" in parts1[0] + # Check strict alternation on Turn 1 + assert [c.role for c in req1.contents] == ["user"] + + req2 = model.requests[1] + parts2 = [p.text for c in req2.contents if c.role == "user" for p in c.parts] + assert "Do a task" in parts2[0] + assert "Here is the clarification" in parts2[1] + # Check strict alternation on Turn 2 + assert [c.role for c in req2.contents] == ["user", "model", "user"] diff --git a/tests/unittests/flows/llm_flows/context/test_contents.py b/tests/unittests/flows/llm_flows/context/test_contents.py index 92ad80cbb25..63da323616e 100644 --- a/tests/unittests/flows/llm_flows/context/test_contents.py +++ b/tests/unittests/flows/llm_flows/context/test_contents.py @@ -2347,6 +2347,46 @@ def test_task_input_user_content_preserves_non_ascii(): assert "\\u" not in text +def test_task_input_user_content_nested_scope_does_not_match_root_reconstruction(): + """Nested workflow scopes (e.g. wf@1/node@2) must not enter root-task reconstruction.""" + events = [ + Event( + invocation_id="2", + author="user", + content=types.UserContent("unrelated user message from invocation 2"), + ), + ] + fallback_content = types.UserContent("node input content") + + # Nested scope contains '@' and '/' + content = contents._build_task_input_user_content( # pylint: disable=protected-access + events, + isolation_scope="wf@1/node@2", + user_content=fallback_content, + ) + + # Must fall back to user_content and not match invocation_id="2" + assert content is not None and content.parts + assert content.parts[0].text == "node input content" + + +def test_task_input_user_content_root_scope_matches_invocation(): + """Root task scopes (e.g. task@inv1) extract user content matching invocation_id.""" + events = [ + Event( + invocation_id="inv1", + author="user", + content=types.UserContent("root task prompt"), + ), + ] + content = contents._build_task_input_user_content( # pylint: disable=protected-access + events, + isolation_scope="task@inv1", + ) + assert content is not None and content.parts + assert content.parts[0].text == "root task prompt" + + @pytest.mark.asyncio async def test_include_contents_none_keeps_turn_across_long_running_resume(): """A tool result posted back by the caller does not restart the turn. diff --git a/tests/unittests/workflow/test_task_api_e2e.py b/tests/unittests/workflow/test_task_api_e2e.py index 5d67d0d46f5..71d30040793 100644 --- a/tests/unittests/workflow/test_task_api_e2e.py +++ b/tests/unittests/workflow/test_task_api_e2e.py @@ -29,7 +29,6 @@ from __future__ import annotations from typing import Any -from typing import AsyncGenerator from google.adk.agents.context import Context from google.adk.agents.llm_agent import LlmAgent @@ -939,3 +938,199 @@ async def driver(ctx): p.text or "" for c in beta_request.contents or [] for p in c.parts or [] ) assert "ALPHA_SECRET_CONVERSATION" not in rendered_beta_context + + +@pytest.mark.asyncio +async def test_root_task_survives_clarification( + request: pytest.FixtureRequest, +): + model = testing_utils.MockModel.create( + responses=["What is your specific question?", "Got it."] + ) + agent = LlmAgent(name="my_task", model=model, mode="task") + app = App(name=request.function.__name__, root_agent=agent) + runner = testing_utils.InMemoryRunner(app=app) + + await runner.run_async("Do a task") + await runner.run_async("Here is the clarification") + + assert len(model.requests) == 2 + req = model.requests[1] + parts = [p.text for c in req.contents if c.role == "user" for p in c.parts] + # Exactly two user-role text parts: original task + clarification answer. + # len(parts) == 2 guards against duplicated answers (regression for #7384). + assert len(parts) == 2 + assert "Do a task" in parts[0] + assert "Here is the clarification" in parts[1] + + +@pytest.mark.asyncio +async def test_two_sequential_root_tasks_clarification_no_cross_contamination( + request: pytest.FixtureRequest, +): + """Four-turn regression for #7384: second root task clarification isolation. + + Turn 1: Task A (root task) returns a clarification question. + Turn 2: User answers "Clarification for A" — Task A resumes and finishes. + Turn 3: Task B (a new root task in the same session) is invoked with + "Do task B" — Task B returns its own clarification question. + Turn 4: User answers "Clarification for B" — Task B resumes and finishes. + The test verifies both of Task B's LLM requests (calls 3 & 4) only + contain Task B's inputs and do not leak Task A's clarification answer. + """ + model_a = testing_utils.MockModel.create( + responses=[ + # Turn 1: Task A asks a clarification question. + "What details do you need for A?", + # Turn 2: Task A receives the answer and finishes. + _finish_part({"result": "Task A done"}), + ] + ) + model_b = testing_utils.MockModel.create( + responses=[ + # Turn 3 (call 3): Task B asks its own clarification question. + "What details do you need for B?", + # Turn 4 (call 4): Task B receives the answer and finishes. + _finish_part({"result": "Task B done"}), + ] + ) + agent_a = LlmAgent(name="task_a", model=model_a, mode="task") + agent_b = LlmAgent(name="task_b", model=model_b, mode="task") + + from google.adk.runners import Runner + from google.adk.sessions.in_memory_session_service import InMemorySessionService + + # Both runners share the same app_name so they can look up the shared session. + app_name = request.function.__name__ + session_service = InMemorySessionService() + session = await session_service.create_session( + app_name=app_name, + user_id="u", + ) + + app_a = App(name=app_name, root_agent=agent_a) + runner_a = Runner(app=app_a, session_service=session_service) + + app_b = App(name=app_name, root_agent=agent_b) + runner_b = Runner(app=app_b, session_service=session_service) + + # Turn 1: Task A is invoked and asks a clarification question. + async for _ in runner_a.run_async( + user_id="u", + session_id=session.id, + new_message=testing_utils.get_user_content("Do task A"), + ): + pass + + # Turn 2: User answers Task A's clarification question; Task A finishes. + async for _ in runner_a.run_async( + user_id="u", + session_id=session.id, + new_message=testing_utils.get_user_content("Clarification for A"), + ): + pass + + # Task A made 2 LLM calls. + assert len(model_a.requests) == 2 + + # Turn 3 (call 3): Task B is invoked fresh in the same session and asks + # its own clarification question. + async for _ in runner_b.run_async( + user_id="u", + session_id=session.id, + new_message=testing_utils.get_user_content("Do task B"), + ): + pass + + # Turn 4 (call 4): User answers Task B's clarification question; Task B finishes. + async for _ in runner_b.run_async( + user_id="u", + session_id=session.id, + new_message=testing_utils.get_user_content("Clarification for B"), + ): + pass + + # Task B made 2 LLM calls. + assert ( + len(model_b.requests) == 2 + ), f"Task B should have been called twice; got {len(model_b.requests)}" + + # Verify call 3 (Turn 3): Task B's first request should contain only "Do task B". + req3 = model_b.requests[0] + b_user_parts_turn3 = [ + p.text + for c in req3.contents + if c.role == "user" + for p in c.parts + if p.text + ] + assert len(b_user_parts_turn3) == 1, ( + "Expected exactly 1 user-role text part for Task B on Turn 3 (its own" + f" prompt), got {len(b_user_parts_turn3)}: {b_user_parts_turn3}" + ) + assert "Do task B" in b_user_parts_turn3[0], ( + "Task B's prompt not found in its first request; got:" + f" {b_user_parts_turn3}" + ) + assert ( + "Clarification for A" not in b_user_parts_turn3[0] + ), "Task A's clarification answer must not appear in Task B's context" + + # Verify call 4 (Turn 4): Task B's resumption request should contain its own + # original prompt and its clarification answer, but never Task A's answer. + req4 = model_b.requests[1] + b_user_parts_turn4 = [ + p.text + for c in req4.contents + if c.role == "user" + for p in c.parts + if p.text + ] + assert len(b_user_parts_turn4) == 2, ( + "Expected exactly 2 user-role text parts for Task B on Turn 4 " + "(original task + clarification answer), " + f"got {len(b_user_parts_turn4)}: {b_user_parts_turn4}" + ) + assert "Do task B" in b_user_parts_turn4[0], ( + "Task B's prompt not found in its resumption request; got:" + f" {b_user_parts_turn4}" + ) + assert "Clarification for B" in b_user_parts_turn4[1], ( + "Task B's clarification answer not found in its resumption request; got:" + f" {b_user_parts_turn4}" + ) + assert ( + "Clarification for A" not in b_user_parts_turn4[0] + ), "Task A's clarification answer must not appear in Task B's prompt" + assert "Clarification for A" not in b_user_parts_turn4[1], ( + "Task A's clarification answer must not appear in Task B's resumption" + " context" + ) + + +@pytest.mark.asyncio +async def test_separate_root_tasks_same_session_get_separate_scopes( + request: pytest.FixtureRequest, +): + model = testing_utils.MockModel.create( + responses=[ + _finish_part({"result": "Done A"}), + _finish_part({"result": "Done B"}), + ] + ) + agent = LlmAgent(name="my_task", model=model, mode="task") + app = App(name=request.function.__name__, root_agent=agent) + runner = testing_utils.InMemoryRunner(app=app) + + await runner.run_async("Run A") + await runner.run_async("Run B") + + events = runner.session.events + task_events = [e for e in events if e.author == "my_task" and e.output] + assert len(task_events) == 2 + # scopes should be different + scope_a = task_events[0].isolation_scope + scope_b = task_events[1].isolation_scope + assert scope_a != scope_b + assert scope_a is not None + assert scope_b is not None