Skip to content

Commit f706a1e

Browse files
gezwcopybara-github
authored andcommitted
fix: use branch-scoped events for auth responses
Merge #5982 This ensures auth handling follows the same branch visibility semantics used by the rest of the invocation context. PiperOrigin-RevId: 941326155
1 parent 3911010 commit f706a1e

2 files changed

Lines changed: 38 additions & 3 deletions

File tree

src/google/adk/auth/auth_preprocessor.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ async def run_async(
134134
agent = invocation_context.agent
135135
if not hasattr(agent, 'canonical_tools'):
136136
return
137-
events = invocation_context.session.events
137+
events = invocation_context._get_events(current_branch=True)
138138
if not events:
139139
return
140140

tests/unittests/auth/test_auth_preprocessor.py

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ def mock_invocation_context(self, mock_llm_agent, mock_session):
6969
context = Mock(spec=InvocationContext)
7070
context.agent = mock_llm_agent
7171
context.session = mock_session
72+
context._get_events.side_effect = lambda **_: context.session.events
7273
return context
7374

7475
@pytest.fixture
@@ -165,7 +166,8 @@ async def test_non_llm_agent_returns_early(
165166
):
166167
"""Test that non-LLM agents return early."""
167168
mock_context = Mock(spec=InvocationContext)
168-
mock_context.agent = Mock()
169+
# Using spec=[] ensures hasattr(agent, 'canonical_tools') returns False.
170+
mock_context.agent = Mock(spec=[])
169171
mock_context.agent.__class__.__name__ = 'BaseAgent'
170172
mock_context.session = mock_session
171173

@@ -273,6 +275,38 @@ async def test_last_event_no_auth_responses_returns_early(
273275

274276
assert result == []
275277

278+
@pytest.mark.asyncio
279+
@patch('google.adk.auth.auth_preprocessor.AuthHandler')
280+
@patch('google.adk.auth.auth_tool.AuthConfig.model_validate')
281+
async def test_ignores_auth_responses_outside_current_branch(
282+
self,
283+
mock_auth_config_validate,
284+
mock_auth_handler_class,
285+
processor,
286+
mock_invocation_context,
287+
mock_llm_request,
288+
mock_user_event_with_auth_response,
289+
):
290+
"""Test auth responses hidden by branch filtering are ignored."""
291+
mock_invocation_context.session.events = [
292+
mock_user_event_with_auth_response
293+
]
294+
mock_invocation_context._get_events.side_effect = None
295+
mock_invocation_context._get_events.return_value = []
296+
297+
result = []
298+
async for event in processor.run_async(
299+
mock_invocation_context, mock_llm_request
300+
):
301+
result.append(event)
302+
303+
mock_invocation_context._get_events.assert_called_once_with(
304+
current_branch=True
305+
)
306+
mock_auth_config_validate.assert_not_called()
307+
mock_auth_handler_class.assert_not_called()
308+
assert result == []
309+
276310
@pytest.mark.asyncio
277311
@patch('google.adk.auth.auth_preprocessor.AuthHandler')
278312
@patch('google.adk.auth.auth_tool.AuthConfig.model_validate')
@@ -536,7 +570,8 @@ async def test_isinstance_check_for_llm_agent(
536570

537571
# Create a mock that fails isinstance check
538572
mock_context = Mock(spec=InvocationContext)
539-
mock_context.agent = Mock() # This will fail isinstance(agent, LlmAgent)
573+
# This will fail isinstance(agent, LlmAgent)
574+
mock_context.agent = Mock(spec=[])
540575
mock_context.session = mock_session
541576

542577
result = []

0 commit comments

Comments
 (0)