diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 40c47354f5..1ae1f42d2e 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -893,6 +893,8 @@ async def _send_to_model( if live_request.content: content = live_request.content + if content.parts and any(p.function_call for p in content.parts): + raise ValueError('User message cannot contain function calls.') # Persist user text content to session (similar to non-live mode) # Skip function responses - they are already handled separately is_function_response = content.parts and any( @@ -1362,6 +1364,16 @@ def _get_agent_to_run( agent_to_run = root_agent.find_agent(agent_name) if not agent_to_run: raise ValueError(f'Agent {agent_name} not found in the agent tree.') + + from ...agents.llm_agent import LlmAgent + + if ( + isinstance(invocation_context.agent, LlmAgent) + and invocation_context.agent.disallow_transfer_to_peers + and agent_to_run.parent_agent == invocation_context.agent.parent_agent + and agent_to_run != invocation_context.agent + ): + raise ValueError(f'Transfer to sibling agent {agent_name} is disallowed.') return agent_to_run async def _call_llm_async( diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 397bb3aca4..2030977dc3 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -915,6 +915,9 @@ async def _append_new_message_to_session( if not new_message.parts: raise ValueError('No parts in the new_message.') + if any(p.function_call for p in new_message.parts): + raise ValueError('User message cannot contain function calls.') + if self.artifact_service and save_input_blobs_as_artifacts: # Issue deprecation warning warnings.warn( diff --git a/tests/unittests/flows/llm_flows/test_base_llm_flow.py b/tests/unittests/flows/llm_flows/test_base_llm_flow.py index cb4de478d0..2ad178edc1 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow.py +++ b/tests/unittests/flows/llm_flows/test_base_llm_flow.py @@ -20,6 +20,7 @@ from google.adk.agents.live_request_queue import LiveRequestQueue from google.adk.agents.llm_agent import Agent +from google.adk.agents.loop_agent import LoopAgent from google.adk.agents.run_config import RunConfig from google.adk.events.event import Event from google.adk.flows.llm_flows.base_llm_flow import _handle_after_model_callback @@ -2010,3 +2011,143 @@ async def test_postprocess_live_skips_none_function_response_event(): ] assert all(event is not None for event in events) + + +@pytest.mark.asyncio +async def test_send_to_model_rejects_function_call(): + """Test that _send_to_model raises ValueError if user message contains function calls.""" + agent = Agent(name='test_agent') + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + invocation_context.live_request_queue = LiveRequestQueue() + + # Put a malicious content request in the queue + from google.adk.agents.live_request_queue import LiveRequest + + malicious_request = LiveRequest( + content=types.Content( + role='user', + parts=[ + types.Part( + function_call=types.FunctionCall( + name='some_tool', + args={'key': 'value'}, + ) + ) + ], + ) + ) + invocation_context.live_request_queue.send(malicious_request) + # Close the queue so that _send_to_model returns instead of blocking on the + # next request if the malicious one is ever accepted. + invocation_context.live_request_queue.close() + + flow = BaseLlmFlowForTesting() + mock_connection = mock.AsyncMock() + + with pytest.raises( + ValueError, match='User message cannot contain function calls' + ): + await flow._send_to_model(mock_connection, invocation_context) + + +def _make_agent_tree(): + root = Agent(name='root') + child1 = Agent(name='child1') + child2 = Agent(name='child2') + + child1.parent_agent = root + child2.parent_agent = root + root.sub_agents = [child1, child2] + return root, child1, child2 + + +@pytest.mark.asyncio +async def test_transfer_to_sibling_disallowed_raises_value_error(): + """Transfer to sibling raises ValueError when disallow_transfer_to_peers is True.""" + # Arrange + root, child1, child2 = _make_agent_tree() + caller = child1 + caller.disallow_transfer_to_peers = True + ctx = await testing_utils.create_invocation_context(caller) + flow = BaseLlmFlowForTesting() + + # Act & Assert + with pytest.raises( + ValueError, match='Transfer to sibling agent child2 is disallowed' + ): + flow._get_agent_to_run(ctx, 'child2') + + +@pytest.mark.asyncio +async def test_transfer_to_sibling_allowed_returns_agent(): + """Transfer to sibling returns the agent when disallow_transfer_to_peers is False.""" + # Arrange + root, child1, child2 = _make_agent_tree() + caller = child1 + caller.disallow_transfer_to_peers = False + ctx = await testing_utils.create_invocation_context(caller) + flow = BaseLlmFlowForTesting() + + # Act + agent = flow._get_agent_to_run(ctx, 'child2') + + # Assert + assert agent is not None + assert agent.name == 'child2' + + +@pytest.mark.asyncio +async def test_transfer_to_unknown_agent_raises_value_error(): + """Transfer to unknown agent name raises ValueError.""" + # Arrange + root, child1, child2 = _make_agent_tree() + caller = child1 + ctx = await testing_utils.create_invocation_context(caller) + flow = BaseLlmFlowForTesting() + + # Act & Assert + with pytest.raises(ValueError, match='not found in the agent tree'): + flow._get_agent_to_run(ctx, 'not_in_tree') + + +@pytest.mark.asyncio +async def test_transfer_to_self_allowed_when_peers_disallowed(): + """Transfer to self is allowed even when disallow_transfer_to_peers is True.""" + # Arrange + root, child1, child2 = _make_agent_tree() + caller = child1 + caller.disallow_transfer_to_peers = True + ctx = await testing_utils.create_invocation_context(caller) + flow = BaseLlmFlowForTesting() + + # Act + agent = flow._get_agent_to_run(ctx, 'child1') + + # Assert + assert agent is not None + assert agent.name == 'child1' + + +@pytest.mark.asyncio +async def test_transfer_to_sibling_from_non_llm_agent_allowed(): + """Transfer to sibling is allowed when the caller is not an LlmAgent.""" + # Arrange + root = Agent(name='root') + child1 = LoopAgent(name='child1') + child2 = Agent(name='child2') + + child1.parent_agent = root + child2.parent_agent = root + root.sub_agents = [child1, child2] + + ctx = await testing_utils.create_invocation_context(child1) + flow = BaseLlmFlowForTesting() + + # Act + agent = flow._get_agent_to_run(ctx, 'child2') + + # Assert + assert agent is not None + assert agent.name == 'child2' diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index aa3fc030f3..4292f97cf5 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import aclosing import importlib from pathlib import Path import sys @@ -1540,5 +1541,41 @@ async def test_get_session_config_limits_events(): assert len(limited_session.events) == 3 +@pytest.mark.asyncio +async def test_run_async_rejects_user_function_call(): + """Verify that runner rejects user-authored messages with function calls.""" + session_service = InMemorySessionService() + runner = Runner( + app_name=TEST_APP_ID, + agent=MockAgent("test_agent"), + session_service=session_service, + artifact_service=InMemoryArtifactService(), + auto_create_session=True, + ) + + malicious_message = types.Content( + role="user", + parts=[ + types.Part( + function_call=types.FunctionCall( + name="some_tool", + args={"key": "value"}, + ) + ) + ], + ) + + agen = runner.run_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + new_message=malicious_message, + ) + + with pytest.raises(ValueError, match="cannot contain function calls"): + async with aclosing(agen) as a: + async for _ in a: + pass + + if __name__ == "__main__": pytest.main([__file__])