diff --git a/src/google/adk/flows/llm_flows/context/_compaction.py b/src/google/adk/flows/llm_flows/context/_compaction.py index ae294eb85d7..4807ac2cfb6 100644 --- a/src/google/adk/flows/llm_flows/context/_compaction.py +++ b/src/google/adk/flows/llm_flows/context/_compaction.py @@ -168,6 +168,10 @@ def _recover_compacted_function_calls( removed are re-injected too, so a sibling is not surfaced as a phantom pending call. + Call ids can be reused. Ownership is resolved in source order before any + events are restored: a response belongs to the latest preceding call with + its id, even when another call with that id survived compaction. + Args: events: The post-compaction events being assembled into request contents. source_events: The pre-compaction events to recover missing calls from. @@ -177,71 +181,68 @@ def _recover_compacted_function_calls( compacted sibling responses) re-injected; the original list is returned unchanged when nothing needs recovery. """ - call_ids_present: set[str] = set() - response_ids_present: set[str] = set() - for event in events: - for function_call in event.get_function_calls(): - if function_call.id: - call_ids_present.add(function_call.id) - for function_response in event.get_function_responses(): - if function_response.id: - response_ids_present.add(function_response.id) - - orphaned_ids = { - response_id - for response_id in response_ids_present - if response_id not in call_ids_present - } - if not orphaned_ids: + if not any(event.get_function_responses() for event in events): return events - call_event_by_id: dict[str, Event] = {} + present_events = {id(event) for event in events} + latest_call_by_id: dict[str, Event] = {} + call_by_response: dict[tuple[int, str], Event] = {} + response_by_call: dict[tuple[int, str], Event] = {} + surviving_responses: set[tuple[int, str]] = set() for event in source_events: for function_call in event.get_function_calls(): - if function_call.id in orphaned_ids: - call_event_by_id.setdefault(function_call.id, event) - - if not call_event_by_id: - return events - - # Keep the highest-timestamp response per id so a sibling that completed - # before being compacted contributes its real result, not its stale - # placeholder; ties fall back to source order. - response_event_by_id: dict[str, Event] = {} - for event in source_events: + if function_call.id: + latest_call_by_id[function_call.id] = event for function_response in event.get_function_responses(): - if not function_response.id: + response_id = function_response.id + if not response_id: continue - existing = response_event_by_id.get(function_response.id) + call_event = latest_call_by_id.get(response_id) + if call_event is None: + continue + call_by_response[id(event), response_id] = call_event + call_key = (id(call_event), response_id) + if id(event) in present_events: + surviving_responses.add(call_key) + # Keep the latest update for this occurrence, not a later call that + # happens to reuse the id. Equal timestamps retain source order. + existing = response_by_call.get(call_key) if existing is None or event.timestamp >= existing.timestamp: - response_event_by_id[function_response.id] = event + response_by_call[call_key] = event result: list[Event] = [] - reinjected_ids: set[str] = set() + reinjected_calls: set[int] = set() + reinjected_responses: set[int] = set() for event in events: for function_response in event.get_function_responses(): function_response_id = function_response.id if not function_response_id: continue - call_event = call_event_by_id.get(function_response_id) - if call_event is None or function_response_id in reinjected_ids: + call_event = call_by_response.get((id(event), function_response_id)) + if ( + call_event is None + or id(call_event) in present_events + or id(call_event) in reinjected_calls + ): continue result.append(call_event) - sibling_ids = [ - function_call.id - for function_call in call_event.get_function_calls() - if function_call.id - ] - reinjected_ids.update(sibling_ids) + reinjected_calls.add(id(call_event)) # Recover sibling responses that compaction removed so a parallel sibling # is not left looking like a pending call. - for sibling_id in sibling_ids: - if sibling_id not in response_ids_present: - sibling_response = response_event_by_id.get(sibling_id) - if sibling_response is not None: + for sibling_call in call_event.get_function_calls(): + if not sibling_call.id: + continue + sibling_key = (id(call_event), sibling_call.id) + if sibling_key not in surviving_responses: + sibling_response = response_by_call.get(sibling_key) + if ( + sibling_response is not None + and id(sibling_response) not in reinjected_responses + ): result.append(sibling_response) + reinjected_responses.add(id(sibling_response)) result.append(event) - return result + return result if reinjected_calls else events __all__ = [ diff --git a/tests/unittests/flows/llm_flows/context/test_content_compaction.py b/tests/unittests/flows/llm_flows/context/test_content_compaction.py index 05df0ce41e4..35931230440 100644 --- a/tests/unittests/flows/llm_flows/context/test_content_compaction.py +++ b/tests/unittests/flows/llm_flows/context/test_content_compaction.py @@ -14,12 +14,25 @@ """Unit tests for google.adk.flows.llm_flows._content_compaction.""" +from __future__ import annotations + +from google.adk.agents.llm_agent import LlmAgent +from google.adk.apps.app import App +from google.adk.apps.app import ResumabilityConfig from google.adk.events.event import Event from google.adk.events.event_actions import EventActions from google.adk.events.event_actions import EventCompaction +from google.adk.flows.llm_flows.context import _contents from google.adk.flows.llm_flows.context._compaction import _process_compaction_events from google.adk.flows.llm_flows.context._compaction import _recover_compacted_function_calls +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.sessions.sqlite_session_service import SqliteSessionService +from google.adk.tools.long_running_tool import LongRunningFunctionTool from google.genai import types +import pytest + +from .... import testing_utils def _long_running_call_event() -> Event: @@ -193,3 +206,266 @@ def _response_event( assert result == [summary_event, parallel_call, lr2_result, lr1_result] assert result[2].get_function_responses()[0].response == {"result": "done-2"} + + +def _call(call_id: str, item: str, timestamp: float, invocation: str) -> Event: + return Event( + invocation_id=invocation, + author="agent", + timestamp=timestamp, + long_running_tool_ids={call_id}, + content=types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + id=call_id, name="lookup", args={"item": item} + ), + thought_signature=b"call-signature", + ) + ], + ), + ) + + +def _response( + call_id: str, value: str, timestamp: float, invocation: str +) -> Event: + return Event( + invocation_id=invocation, + author="user", + timestamp=timestamp, + content=types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + id=call_id, name="lookup", response={"result": value} + ) + ) + ], + ), + ) + + +def _summary(start: float, end: float) -> Event: + return Event( + invocation_id="summary", + author="agent", + timestamp=end + 0.5, + actions=EventActions( + compaction=EventCompaction( + start_timestamp=start, + end_timestamp=end, + compacted_content=types.ModelContent("summary"), + ) + ), + ) + + +@pytest.mark.parametrize("reuse_id", [False, True]) +@pytest.mark.parametrize("compaction_start", [None, 1.0, 3.0]) +@pytest.mark.parametrize("same_invocation", [False, True]) +def test_resumed_response_keeps_its_call_after_compaction( + reuse_id, compaction_start, same_invocation +): + """A repeated call id must not give the resumed result an older call's args.""" + old_id = "shared" if reuse_id else "old" + old_invocation = "current-invocation" if same_invocation else "old-invocation" + events = [ + _call(old_id, "old-item", 1.0, old_invocation), + _response(old_id, "old-result", 2.0, old_invocation), + _call("shared", "current-item", 3.0, "current-invocation"), + _response("shared", "pending", 4.0, "current-invocation"), + ] + if compaction_start is not None: + events.append(_summary(compaction_start, 4.0)) + events.append( + _response("shared", "current-result", 6.0, "current-invocation") + ) + original = [event.model_dump() for event in events] + + result = _contents._get_contents( + None, events, agent_name="agent", preserve_function_call_ids=True + ) + + call_parts = [ + part + for content in result + for part in content.parts or [] + if part.function_call + ] + responses = [ + part.function_response + for content in result + for part in content.parts or [] + if part.function_response + ] + assert call_parts[-1].function_call.args == {"item": "current-item"} + assert call_parts[-1].thought_signature == b"call-signature" + assert responses[-1].response == {"result": "current-result"} + assert [event.model_dump() for event in events] == original + + +@pytest.mark.parametrize("backend", ["memory", "sqlite"]) +async def test_runner_resumes_compacted_call_with_reused_id(backend, tmp_path): + """Persisted resume history must reach the model with the current call's args.""" + service = ( + InMemorySessionService() + if backend == "memory" + else SqliteSessionService(str(tmp_path / "session.db")) + ) + executed = [] + + def lookup(item: str) -> None: + """Starts a lookup that will be answered later.""" + executed.append(item) + + model = testing_utils.MockModel.create( + responses=[ + types.Part( + function_call=types.FunctionCall( + id="shared", name="lookup", args={"item": "old-item"} + ) + ), + "old-ready", + types.Part( + function_call=types.FunctionCall( + id="shared", name="lookup", args={"item": "current-item"} + ) + ), + "resumed", + ] + ) + agent = LlmAgent( + name="agent", model=model, tools=[LongRunningFunctionTool(lookup)] + ) + runner = Runner( + app=App( + name="test_app", + root_agent=agent, + resumability_config=ResumabilityConfig(is_resumable=True), + ), + session_service=service, + ) + try: + session = await service.create_session(app_name="test_app", user_id="user") + # Generate the history through real tool execution and a prior resume. + for message in [ + types.UserContent("Start the old lookup"), + _response("shared", "old-result", 2.0, "old-invocation").content, + types.UserContent("Start the current lookup"), + ]: + async for _ in runner.run_async( + user_id="user", session_id=session.id, new_message=message + ): + pass + session = await service.get_session( + app_name="test_app", user_id="user", session_id=session.id + ) + assert session is not None + timestamps = [event.timestamp for event in session.events] + await service.append_event( + session=session, + event=Event( + author="agent", + invocation_id="summary", + actions=EventActions( + compaction=EventCompaction( + start_timestamp=min(timestamps), + end_timestamp=max(timestamps), + compacted_content=types.ModelContent("summary"), + ) + ), + ), + ) + before = [event.model_dump() for event in session.events] + + resumed = [ + event + async for event in runner.run_async( + user_id="user", + session_id=session.id, + new_message=_response( + "shared", "current-result", 6.0, "current-invocation" + ).content, + ) + ] + + assert executed == ["old-item", "current-item"] + assert len(model.requests) == 4 + parts = [ + p + for content in model.requests[-1].contents + for p in content.parts or [] + ] + calls = [p.function_call for p in parts if p.function_call] + responses = [p.function_response for p in parts if p.function_response] + assert calls[-1].args == {"item": "current-item"} + assert responses[-1].response == {"result": "current-result"} + assert any( + p.text == "resumed" + for e in resumed + for p in (e.content.parts if e.content else []) or [] + ) + saved = await service.get_session( + app_name="test_app", user_id="user", session_id=session.id + ) + assert [ + event.model_dump() for event in saved.events[: len(before)] + ] == before + finally: + await runner.close() + if isinstance(service, SqliteSessionService): + await service.close() + + +@pytest.mark.parametrize("reuse_sibling_id", [False, True]) +def test_recovered_parallel_call_keeps_its_own_sibling_result(reuse_sibling_id): + """A compacted sibling's result must not come from a later unrelated call.""" + parallel_call = _call("pending", "pending-item", 1.0, "invocation") + parallel_call.content.parts.append( + types.Part( + function_call=types.FunctionCall( + id="sibling", name="lookup", args={"item": "old-sibling-item"} + ) + ) + ) + newer_id = "sibling" if reuse_sibling_id else "other" + events = [ + parallel_call, + _response("pending", "pending", 2.0, "invocation"), + _response("sibling", "old-sibling-result", 3.0, "invocation"), + _call(newer_id, "new-item", 4.0, "invocation"), + _response(newer_id, "new-result", 5.0, "invocation"), + _summary(1.0, 5.0), + _response("pending", "done", 7.0, "invocation"), + ] + original = [event.model_dump() for event in events] + + result = _contents._get_contents( + None, events, agent_name="agent", preserve_function_call_ids=True + ) + + responses = { + part.function_response.id: part.function_response.response + for content in result + for part in content.parts or [] + if part.function_response + } + assert responses == { + "pending": {"result": "done"}, + "sibling": {"result": "old-sibling-result"}, + } + call_parts = [ + part + for content in result + for part in content.parts or [] + if part.function_call + ] + assert [part.function_call.args for part in call_parts] == [ + {"item": "pending-item"}, + {"item": "old-sibling-item"}, + ] + assert call_parts[0].thought_signature == b"call-signature" + assert [event.model_dump() for event in events] == original