|
28 | 28 | from unittest.mock import patch |
29 | 29 |
|
30 | 30 | from google.adk import runners |
| 31 | +from google.adk.agents._caller_principal import CallerPrincipal |
31 | 32 | from google.adk.agents.base_agent import BaseAgent |
32 | 33 | from google.adk.agents.context import Context |
33 | 34 | from google.adk.agents.context_cache_config import ContextCacheConfig |
@@ -4971,3 +4972,42 @@ async def on_event_callback(self, *, invocation_context, event): |
4971 | 4972 |
|
4972 | 4973 | if __name__ == "__main__": |
4973 | 4974 | pytest.main([__file__]) |
| 4975 | + |
| 4976 | + |
| 4977 | +@pytest.mark.asyncio |
| 4978 | +async def test_runner_root_llm_agent_carries_caller_principal_on_node_path(): |
| 4979 | + """run_async(caller_principal=...) reaches a root LlmAgent's invocation context. |
| 4980 | +
|
| 4981 | + A root LlmAgent runs through the node runtime, which builds its own |
| 4982 | + invocation context. The principal a serving layer set has to survive that |
| 4983 | + hop, or the confirmation gate downstream never sees it. |
| 4984 | + """ |
| 4985 | + seen = [] |
| 4986 | + |
| 4987 | + def capture(callback_context): |
| 4988 | + seen.append(callback_context._invocation_context.caller_principal) |
| 4989 | + return None |
| 4990 | + |
| 4991 | + session_service = InMemorySessionService() |
| 4992 | + agent = LlmAgent( |
| 4993 | + name="chat_agent", |
| 4994 | + model=testing_utils.MockModel.create(responses=["hi"]), |
| 4995 | + before_agent_callback=capture, |
| 4996 | + ) |
| 4997 | + runner = Runner( |
| 4998 | + app_name=TEST_APP_ID, agent=agent, session_service=session_service |
| 4999 | + ) |
| 5000 | + await session_service.create_session( |
| 5001 | + app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID |
| 5002 | + ) |
| 5003 | + principal = CallerPrincipal(authenticated=False, source="test") |
| 5004 | + |
| 5005 | + async for _ in runner.run_async( |
| 5006 | + user_id=TEST_USER_ID, |
| 5007 | + session_id=TEST_SESSION_ID, |
| 5008 | + new_message=types.Content(role="user", parts=[types.Part(text="hello")]), |
| 5009 | + caller_principal=principal, |
| 5010 | + ): |
| 5011 | + pass |
| 5012 | + |
| 5013 | + assert seen == [principal] |
0 commit comments