fix: Set invocation_id and branch for event generated when both output_schema and tools are used

fixes #2631

PiperOrigin-RevId: 797843460

Author: seanzhougoogle

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

@@ -78,7 +78,11 @@ def create_final_model_response_event(
   from google.genai import types
 
   # Create a proper model response event
-  final_event = Event(author=invocation_context.agent.name)
+  final_event = Event(
+      author=invocation_context.agent.name,
+      invocation_id=invocation_context.invocation_id,
+      branch=invocation_context.branch,
+  )
   final_event.content = types.Content(
       role='model', parts=[types.Part(text=json_response)]
   )

tests/unittests/flows/llm_flows/test_output_schema_processor.py

@@ -230,6 +230,8 @@ async def test_output_schema_helper_functions():
   # Test create_final_model_response_event function
   final_event = create_final_model_response_event(invocation_context, test_json)
   assert final_event.author == 'test_agent'
+  assert final_event.invocation_id == invocation_context.invocation_id
+  assert final_event.branch == invocation_context.branch
   assert final_event.content.role == 'model'
   assert final_event.content.parts[0].text == test_json
 
@@ -348,6 +350,8 @@ async def test_flow_yields_both_events_for_set_model_response():
   # Second event should be the final model response with JSON
   second_event = events[1]
   assert second_event.author == 'test_agent'
+  assert second_event.invocation_id == invocation_context.invocation_id
+  assert second_event.branch == invocation_context.branch
   assert second_event.content.role == 'model'
   assert (
       second_event.content.parts[0].text