LlmAgent `output_key` is not written to session state when `before_agent_callback` short-circuits execution

[gautamvarmadatla]

Describe the Bug:
When LlmAgent uses both output_key and before_agent_callback, a types.Content returned from the callback correctly short-circuits execution and is shown to the user, but session.state[output_key] is never updated. This would silently break downstream agents, tools, or prompt templates that rely on output_key , since the UI shows the correct response while the backend state remains stale or empty.

Steps to Reproduce:
We can reproduce using this repro code

 from google.adk.agents import LlmAgent
 from google.adk.agents.callback_context import CallbackContext
 from google.adk.runners import InMemoryRunner
 from google.genai import types

 def cache_callback(callback_context: CallbackContext):
 return types.Content(parts=[types.Part(text="cached answer")])

 agent = LlmAgent(
 name="my_agent",
 model="gemini-2.0-flash",
 instruction="Answer questions.",
 output_key="result",
 before_agent_callback=cache_callback,
 )
 runner = InMemoryRunner(agent=agent, app_name="test")
 session = await runner.session_service.create_session(app_name="test", user_id="user1")
 async for _ in runner.run_async(
 user_id="user1", session_id=session.id,
 new_message=types.Content(parts=[types.Part(text="hello")]),
 ):
 pass
 session = await runner.session_service.get_session(app_name="test", user_id="user1", session_id=session.id)
 print(session.state.get("result")) 

Expected Behavior:
session.state["result"] should equal "cached answer", the same text the callback returned to the user.

Observed Behavior:
session.state.get("result") returns None. No error or warning is raised.

Environment Details:

• ADK Library Version (pip show google-adk): 1.27.1
• Desktop OS: Windows
• Python Version (python -V): 3.12.12

How often has this issue occurred?:

• Always (100%)

[llalitkumarrr]

Hello @gautamvarmadatla ,

This is expected behavior. When before_agent_callback is used and content is present the agent execution is skipped and the provided content is returned. Because output_key extract the agent reply and no agent reply was generated in this scenario hence it is None.

You might want to look into using before_model_callback instead. Could you check if that better suits your requirements?

[gautamvarmadatla]

hi @llalitkumarrr,

tysm for the response. I did look into before_model_callback, but it runs at the LLM boundary after the agent’s preprocessing/setup has already happened, so it does not fit this use case where we need to skip the agent entirely.

And before_agent_callback can return the final response event. In this case, it returns content that ADK emits as an agent-authored response. The issue is just that this short-circuit path was missing the output_key save step.

[llalitkumarrr]

Hello @gautamvarmadatla ,

The reason for issue with output_key is mentioned in my previous comment. You might to figure out a work around or adjust your workflow.

[gautamvarmadatla]

Hi @llalitkumarrr,

I understand that before_agent_callback skips the normal LlmAgent flow. My point is that it can still return the actual response shown to the user. For example, a common pattern is to check a cache in before_agent_callback and return "cached answer" immediately if there is a hit. In that case, the user still receives "cached answer" as the agent response. If output_key="result", I would expect session.state["result"] to also contain "cached answer". That is why I see this as a missing output_key save on the short-circuit path, rather than something we need to work around.

[llalitkumarrr]

Hello @wyf7107 ,

Could you please take a look at this and share your inputs?