Skip to content

Commit 883ff98

Browse files
doughaydenGWeale
authored andcommitted
fix(flows): terminate invocation at tool-level EUC
Merge #5638 ### Link to Issue or Description of Change **1. Link to an existing issue (if applicable):** - Closes: #5637 This change adds `invocation_context.end_invocation = True` after the auth event yield in `_postprocess_handle_function_calls_async`, mirroring the existing termination signal in `_resolve_toolset_auth`. Tool-level auth now terminates symmetrically with toolset-level auth at the EUC, instead of continuing for one more LLM call. ### Testing Plan **Unit Tests:** - [x] I have added or updated unit tests for my change. - [x] All unit tests pass locally. Three existing tests in `test_functions_request_euc.py` had assertions tied to the trailing post-EUC LLM call: - `test_function_request_euc`: adds `assert len(mock_model.requests) == 1` to anchor the new termination behavior. - `test_function_get_auth_response`: `events[-3]` → `events[-2]` for the auth event lookup, since the auth event is now second-to-last. - `test_function_get_auth_response_partial`: same `events[-3]` → `events[-2]` change, plus the two `len(mock_model.requests)` assertions drop by 1 (3 → 2 and 4 → 3). ``` $ pytest tests/unittests/flows/llm_flows/test_functions_request_euc.py ======================== 3 passed, 17 warnings in 1.31s ======================== $ pytest tests/unittests/ =============== 5695 passed, 2308 warnings in 122.89s (0:02:02) ================ ``` **Manual End-to-End (E2E) Tests:** A self-contained Runner-based reproduction is at https://github.com/doughayden/adk-issue-examples/tree/main/04-tool_level_auth_continuation. The agent definition (`agent.py`) wires up an `OpenAPIToolset` against a local OAuth2 test server. `main.py` constructs an `InMemoryRunner`, applies the workaround for #5327 (`get_auth_config = lambda: None`) at runtime to land on the tool-level auth path, and sends a tool-triggering prompt. The `--apply-fix` flag monkey-patches the proposed fix to demonstrate the resolution end-to-end. Without the fix: ``` 👤 User: What's the weather in San Francisco? 🌤️ Weather Assistant event stream: [function_call] get_weather by WeatherAssistant [auth_event] adk_request_credential by WeatherAssistant [function_response] get_weather by WeatherAssistant [post_euc_text] WeatherAssistant: "I'm sorry, I cannot retrieve the weather for San Francisco at the moment. It ..." Event counts: function_calls: 1 auth_events: 1 function_responses: 1 text_events: 1 post_euc_text_events: 1 ✅ Bug reproduced: 1 text event(s) after the EUC (agent loop continued past adk_request_credential). ``` With the fix: ``` 👤 User: What's the weather in San Francisco? 🌤️ Weather Assistant event stream: [function_call] get_weather by WeatherAssistant [auth_event] adk_request_credential by WeatherAssistant [function_response] get_weather by WeatherAssistant Event counts: function_calls: 1 auth_events: 1 function_responses: 1 text_events: 0 post_euc_text_events: 0 ✅ Fix verified: no LLM events after the EUC. ``` ### Checklist - [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document. - [x] I have performed a self-review of my own code. - [x] I have commented my code, particularly in hard-to-understand areas. - [x] I have added tests that prove my fix is effective or that my feature works. - [x] New and existing unit tests pass locally with my changes. - [x] I have manually tested my changes end-to-end. - [ ] Any dependent changes have been merged and published in downstream modules. ### Additional context **Alternative considered:** A reorder of the yields (yield `auth_event` last so `last_event.is_final_response()` returns True) would also fix the loop termination in a single iteration without needing the flag. I went with `end_invocation = True` to preserve the observable event order and to match the existing pattern in `_resolve_toolset_auth`. Happy to switch if maintainers prefer the reorder. **Related:** The same yield site at lines 1126-1130 also produces `tool_confirmation_event` for HITL with the same `long_running_tool_ids` shape and the same termination gap. This PR scopes to `auth_event` only. Happy to open a follow-up PR with the same fix for `tool_confirmation_event` if the team agrees with the approach here. Co-authored-by: George Weale <gweale@google.com> COPYBARA_INTEGRATE_REVIEW=#5638 from doughayden:fix/tool-level-auth-terminates-at-euc 0a04d30 PiperOrigin-RevId: 932731604
1 parent 2e28e5d commit 883ff98

4 files changed

Lines changed: 15 additions & 16 deletions

File tree

src/google/adk/flows/llm_flows/base_llm_flow.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1197,6 +1197,9 @@ async def _postprocess_handle_function_calls_async(
11971197
if auth_event:
11981198
yield auth_event
11991199

1200+
# Interrupt invocation (mirrors _resolve_toolset_auth behavior)
1201+
invocation_context.end_invocation = True
1202+
12001203
tool_confirmation_event = functions.generate_request_confirmation_event(
12011204
invocation_context, function_call_event, function_response_event
12021205
)

tests/integration/integrations/agent_identity/test_3lo_flow.py

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,6 @@ async def test_gcp_agent_identity_3lo_user_consent_flow() -> None:
168168
mock_model = testing_utils.MockModel.create(
169169
responses=[
170170
types.Part.from_function_call(name="dummy_tool", args={}),
171-
"I am waiting for your authorization.",
172171
"Tool executed successfully.",
173172
]
174173
)
@@ -276,11 +275,9 @@ def _find_auth_request_event(events):
276275

277276
# Validate requests received by the mock model
278277
requests = mock_model.requests
279-
# Events:
280-
# 1. User Input (Get me the token.)
281-
# 2. LLM (I am waiting for your authorization.)
282-
# 3. LLM (Tool executed successfully.)
283-
assert len(requests) == 3
278+
# Two LLM calls: the tool call in turn 1 (which ends at the EUC) and the
279+
# post-consent response in turn 2.
280+
assert len(requests) == 2
284281

285282
# Extract the function response from the prompt payload sent to the LLM
286283
last_request = requests[-1]

tests/integration/integrations/agent_identity/test_agent_identity_3lo_flow.py

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,6 @@ async def test_gcp_agent_identity_3lo_user_consent_flow() -> None:
159159
mock_model = testing_utils.MockModel.create(
160160
responses=[
161161
types.Part.from_function_call(name="dummy_tool", args={}),
162-
"I am waiting for your authorization.",
163162
"Tool executed successfully.",
164163
]
165164
)
@@ -266,11 +265,9 @@ def _find_auth_request_event(events):
266265

267266
# Validate requests received by the mock model
268267
requests = mock_model.requests
269-
# Events:
270-
# 1. User Input (Get me the token.)
271-
# 2. LLM (I am waiting for your authorization.)
272-
# 3. LLM (Tool executed successfully.)
273-
assert len(requests) == 3
268+
# Two LLM calls: the tool call in turn 1 (which ends at the EUC) and the
269+
# post-consent response in turn 2.
270+
assert len(requests) == 2
274271

275272
# Extract the function response from the prompt payload sent to the LLM
276273
last_request = requests[-1]

tests/unittests/flows/llm_flows/test_functions_request_euc.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,8 @@ def call_external_api2(tool_context: ToolContext) -> Optional[int]:
152152
== auth_configs[idx].raw_auth_credential
153153
)
154154

155+
assert len(mock_model.requests) == 1
156+
155157

156158
def test_function_get_auth_response():
157159
id_1 = 'id_1'
@@ -309,7 +311,7 @@ def call_external_api2(tool_context: ToolContext) -> int:
309311
)
310312
runner = testing_utils.InMemoryRunner(agent)
311313
runner.run('test')
312-
request_euc_function_call_event = runner.session.events[-3]
314+
request_euc_function_call_event = runner.session.events[-2]
313315
function_response1 = types.FunctionResponse(
314316
name=request_euc_function_call_event.content.parts[0].function_call.name,
315317
response=auth_response1.model_dump(),
@@ -505,7 +507,7 @@ def call_external_api2(tool_context: ToolContext) -> int:
505507
)
506508
runner = testing_utils.InMemoryRunner(agent)
507509
runner.run('test')
508-
request_euc_function_call_event = runner.session.events[-3]
510+
request_euc_function_call_event = runner.session.events[-2]
509511
function_response1 = types.FunctionResponse(
510512
name=request_euc_function_call_event.content.parts[0].function_call.name,
511513
response=auth_response1.model_dump(),
@@ -531,7 +533,7 @@ def call_external_api2(tool_context: ToolContext) -> int:
531533
)
532534

533535
assert function_invoked == 3
534-
assert len(mock_model.requests) == 3
536+
assert len(mock_model.requests) == 2
535537
request = mock_model.requests[-1]
536538
content = request.contents[-1]
537539
parts = content.parts
@@ -550,7 +552,7 @@ def call_external_api2(tool_context: ToolContext) -> int:
550552
),
551553
)
552554
assert function_invoked == 4
553-
assert len(mock_model.requests) == 4
555+
assert len(mock_model.requests) == 3
554556
request = mock_model.requests[-1]
555557
content = request.contents[-1]
556558
parts = content.parts

0 commit comments

Comments
 (0)