fix(agents): persist output_key when before_agent_callback short-circuits LlmAgent

Merge google#4838

- Closes: google#4837

**Problem:**
When `LlmAgent` uses both `output_key` and `before_agent_callback`, and the before-agent callback returns `types.Content`, execution is correctly short-circuited and the response is returned to the user, but `session.state[output_key]` is not updated because the normal `LlmAgent` output-saving path is bypassed.

**Solution:**
Override `_handle_before_agent_callback` in `LlmAgent` to call `__maybe_save_output_to_state()` on any returned event. This keeps the fix scoped to `LlmAgent`, which owns `output_key`, avoids introducing `BaseAgent` changes for subclass-specific behavior, and fixes both `run_async` and `run_live` through the shared callback path.

### Testing Plan

Added a regression test that runs `LlmAgent` with `output_key` and a short-circuiting `before_agent_callback`, then checks that `session.state["result"]` contains the returned text.

**Unit Tests:**

- [X] I have added or updated unit tests for my change.
- [X] All unit tests pass locally.

**Manual End-to-End (E2E) Tests:**
Run repro code in google#4837

In that example,  `session.state["result"]` should equal "cached answer"

### 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.
- [X] Any dependent changes have been merged and published in downstream modules.

Co-authored-by: Liang Wu <wuliang@google.com>
COPYBARA_INTEGRATE_REVIEW=google#4838 from gautamvarmadatla:fix/output-key-before-agent-callback 7f084e0
PiperOrigin-RevId: 934022327

Author: 2 people

src/google/adk/agents/llm_agent.py

@@ -491,6 +491,15 @@ class LlmAgent(BaseAgent, abc.ABC):
   """
   # Callbacks - End
 
+  @override
+  async def _handle_before_agent_callback(
+      self, ctx: InvocationContext
+  ) -> Optional[Event]:
+    event = await super()._handle_before_agent_callback(ctx)
+    if event is not None:
+      self.__maybe_save_output_to_state(event)
+    return event
+
   @override
   async def _run_async_impl(
       self, ctx: InvocationContext

tests/unittests/agents/test_llm_agent_output_save.py

@@ -17,13 +17,16 @@
 import logging
 from unittest.mock import patch
 
+from google.adk.agents.callback_context import CallbackContext
 from google.adk.agents.llm_agent import LlmAgent
 from google.adk.events.event import Event
 from google.adk.events.event_actions import EventActions
 from google.genai import types
 from pydantic import BaseModel
 import pytest
 
+from .. import testing_utils
+
 
 class MockOutputSchema(BaseModel):
   message: str
@@ -277,6 +280,27 @@ def test_maybe_save_output_to_state_handles_empty_final_chunk_with_schema(
     # should remain empty.
     assert len(event.actions.state_delta) == 0
 
+  @pytest.mark.asyncio
+  async def test_output_key_saved_when_before_agent_callback_short_circuits(
+      self,
+  ):
+    """Test that output_key is written to session state when
+    before_agent_callback short-circuits the agent."""
+
+    def cache_callback(callback_context: CallbackContext) -> types.Content:
+      return types.Content(parts=[types.Part.from_text(text="cached answer")])
+
+    agent = LlmAgent(
+        name="test_agent",
+        output_key="result",
+        before_agent_callback=cache_callback,
+    )
+
+    runner = testing_utils.InMemoryRunner(agent)
+    await runner.run_async("hello")
+
+    assert runner.session.state.get("result") == "cached answer"
+
   def test_maybe_save_output_to_state_skips_function_response_only_event(self):
     """Test that state_delta set by callback is not overwritten when event